python - PySpark 相当于 Pandas UDF 中的 lambda 函数

Question

我在PySpark中用Pandas UDF编写了数据预处理代码。我正在使用 lambda 函数从列的所有记录中提取文本的一部分。

这是我的代码的样子:

@pandas_udf("string", PandasUDFType.SCALAR)
def get_X(col):
      return col.apply(lambda x: x.split(',')[-1] if len(x.split(',')) > 0 else x)

df = df.withColumn('X', get_first_name(df.Y))

这工作正常并给出了预期的结果。但我需要在 Spark 等效代码中编写相同的逻辑。有办法做到吗？谢谢。

Answer 1

我认为有一个功能substring_index足以完成此特定任务:

from pyspark.sql.functions import substring_index

df = spark.createDataFrame([(x,) for x in ['f,l', 'g', 'a,b,cd']], ['c1'])

df2.withColumn('c2', substring_index('c1', ',', -1)).show()                                                                 
+------+---+
|    c1| c2|
+------+---+
|   f,l|  l|
|     g|  g|
|a,b,cd| cd|
+------+---+