python - PySpark 和方法链

Question

我是 PySpark 的新手，想知道您将如何在那里使用方法链接。在 pandas 中，我会使用 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|
+---------+