如何创建 Pyspark UDF 以向数据框添加新列

Posted 2023-04-14

【中文标题】如何创建 Pyspark UDF 以向数据框添加新列

【英文标题】：How to create a Pyspark UDF for adding new columns to a dataframe

【发布时间】：2017-08-17 10:29:41

【问题描述】：

我有大约 5 列要添加到数据框中。 (A - E) 这些列的值存储在 (a - e) 变量中。

而不是使用

 df.withColumn("A", a).withColumn("B", b).withColumn..... etc 

我们可以用 udf 做到这一点吗？

目前我已经命名了函数：

     def add_col(df_name,newCol,value):
         df = df_name
         df = df.withColumn(newCol, value)
         return df

但我无法理解如何将其转换为 UDF 并使用它。请帮忙。

【问题讨论】：

一个 UDF 不能产生超过 1 列...但您可以返回一个复杂列（数组或结构类型）。但显然你正在使用常量（文字）值，所以我看不出使用 UDF 的意义

【参考方案1】：

如果你想添加多个列，你可以使用select 和*：

df.select("*", some_column, another_column, ...)

【讨论】：

如何在同一步骤中也插入值？

【参考方案2】：

您不应该使用 UDF，它们不能创建多个结果。

但是您可以在其他答案中编写类似于此的选择语句：

df.select(col("*"), lit(a).as("a"), lit(b).as("b"), ...)

您也可以自动添加：

val fieldsMap = Map("a" -> a, "b" -> b)
 df.select(Array(col("*")) ++ fieldsMap.map(e => lit(e._2).as(e._1)) : _*)