如何将 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|
+---------+----------+----------+-----+