PySpark и метод цепочки

Question

Я новичок в PySpark, и мне было интересно, как бы вы использовали там цепочку методов. В pandas я бы использовал assign с lambda, например,

df = pd.DataFrame({'number':[1,2,3],'date':['31-dec-19','02-jan-18','14-mar-20']})

df = (df.assign(number_plus_one = lambda x: x.number + 1)
        .assign(date = lambda x: pd.to_datetime(x.date))
        .loc[lambda x: x.number_plus_one.isin([2,3])]
        .drop(columns=['number','number_plus_one'])
      )

Как бы вы написали тот же код в PySpark без преобразования его в pandas фрейм данных? Я думаю, вы могли бы использовать filter, withColumn и drop, но как именно вы бы это сделали с цепочкой методов?

Answer 1

В Spark вы можете сделать то же самое, связав вызовы аналогичным образом. Вот пример:

sc.parallelize([Row(number=1, date='31-dec-19'),
                     Row(number=1, date='31-dec-19'),
                     Row(number=1, date='31-dec-19')])\
.toDF()\
.withColumn('number_plus_one', f.col('number') + 1)\
.filter(f.col('number_plus_one').cast(IntegerType()).isin(f.lit(2), f.lit(3)) )\
.drop('number','number_plus_one')\
.show()

Результат

+---------+
|     date|
+---------+
|31-dec-19|
|02-jan-18|
+---------+