如何将 PySpark 函数的返回类型指定为数据帧？

Question

我最近正在处理一些编码挑战，涉及将 Spark 数据帧传递到 Python 函数并返回一个新的数据帧。 我记得的语法是这样的：

def sampleFunction(df: Dataframe) -> Dataframe:
    * do stuff *
    return newDF

我现在正在尝试创建自己的示例，但无法将数据帧指定为输入/输出类型。 我假设我需要导入一些东西才能使数据帧成为可接受的类型，但是过去一个小时我一直在谷歌上搜索这个，而且我找不到一个关于如何在 PySpark 中进行这项工作的示例。

Answer 1

您需要在代码中导入DataFrame类型，并且还使用数据类型名称作为DataFrame而不是您在代码中提到的Dataframe

>>> from pyspark.sql import DataFrame
>>> def sampleFunction(df: DataFrame) -> DataFrame:
...     df1 = df.withColumn("necol", lit(1))
...     return df1
... 
>>> from pyspark.sql.functions  import *
>>> df.show()
+---------+----------+----------+
|DEVICE_ID|  MIN_DATE|  MAX_DATE|
+---------+----------+----------+
|        1|2019-08-29|2019-08-31|
|        2|2019-08-27|2019-09-02|
+---------+----------+----------+

>>> newdf = sampleFunction(df)
>>> newdf.show()
+---------+----------+----------+-----+
|DEVICE_ID|  MIN_DATE|  MAX_DATE|necol|
+---------+----------+----------+-----+
|        1|2019-08-29|2019-08-31|    1|
|        2|2019-08-27|2019-09-02|    1|
+---------+----------+----------+-----+