使用withColumn将两列添加到现有DataFrame

Question

我有DataFrame一些列.现在我想在现有的DataFrame中再添加两列.

目前我正在使用withColumnDataFrame中的方法.

例如:

df.withColumn("newColumn1", udf(col("somecolumn")))
  .withColumn("newColumn2", udf(col("somecolumn")))

实际上我可以使用Array [String]在单个UDF方法中返回两个newcoOlumn值.但目前这就是我的做法.

无论如何,我能有效地做到这一点吗？使用explode是不错的选择？

即使我必须使用explode,我必须使用withColumn一次,然后返回列值Array[String],然后使用explode,再创建两列.

哪一个有效？还是有其他选择吗？

Answer 1

AFAIk你需要调用withColumn两次(每个新列一次).但是如果你的udf在计算上很昂贵,你可以避免在将"复杂"结果存储到临时列中然后"解压缩"结果,例如使用apply列的方法(它可以访问数组元素)时调用它两次.请注意,有时需要缓存中间结果(以防止在解包期间每行调用两次UDF),有时则不需要.这似乎取决于优化计划的火花:

val myUDf = udf((s:String) => Array(s.toUpperCase(),s.toLowerCase()))

val df = sc.parallelize(Seq("Peter","John")).toDF("name")

val newDf = df
  .withColumn("udfResult",myUDf(col("name"))).cache 
  .withColumn("uppercaseColumn", col("udfResult")(0))
  .withColumn("lowercaseColumn", col("udfResult")(1))
  .drop("udfResult")

newDf.show()

给

+-----+---------------+---------------+
| name|uppercaseColumn|lowercaseColumn|
+-----+---------------+---------------+
|Peter|          PETER|          peter|
| John|           JOHN|           john|
+-----+---------------+---------------+

使用UDF返回元组,解压缩将如下所示:

val newDf = df
    .withColumn("udfResult",myUDf(col("name"))).cache
    .withColumn("lowercaseColumn", col("udfResult._1"))
    .withColumn("uppercaseColumn", col("udfResult._2"))
    .drop("udfResult")

Answer 2

2023 年 5 月：现在可以使用新的withColumns（注意最后的“s”）方法向现有 Spark 数据帧添加多个列，而无需多次调用withColumn。你只需要一张地图Map[String, Column]。给定本示例的两个 UDF udf1，udf2您可以像这样使用这个新方法：

val dfNew=df.withColumns(Map("newCol1"->udf1(col("oldCol1")),"newCol2"->udf2(col("oldCol2"))))

有关这方面的更多信息现在可以在官方文档中找到。