Как указать тип возвращаемого значения функции PySpark в качестве фрейма данных?

Question

Недавно я работал над некоторыми проблемами кодирования, которые включали передачу фрейма данных Spark в функцию Python и возвращение нового фрейма данных. Синтаксис, который я помню, был примерно таким:

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

Сейчас я пытаюсь создать свои собственные примеры, но не могу указать фрейм данных в качестве типа ввода / вывода. Я предполагаю, что есть что-то, что мне нужно импортировать, чтобы сделать фрейм данных приемлемым типом, но я безостановочно гуглил в течение последнего часа и не могу найти ни одного примера, как заставить это работать в PySpark.

Answer 1

Вам необходимо импортировать 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|
+---------+----------+----------+-----+