使用 withColumn 向现有 DataFrame 添加两列答案

【问题标题】：Adding two columns to existing DataFrame using withColumn使用 withColumn 向现有 DataFrame 添加两列
【发布时间】：2017-04-18 23:37:28
【问题描述】：

我有一个 DataFrame 有几列。现在我想在现有的 DataFrame 中再添加两列。

目前我正在使用 DataFrame 中的withColumn 方法进行此操作。

例如：

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

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

无论如何，我可以有效地做到这一点吗？在这里使用explode 是一个不错的选择吗？

即使我必须使用explode，我也必须使用一次withColumn，然后将列值返回为Array[String]，然后使用explode，再创建两列。

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

【问题讨论】：

标签： scala dataframe apache-spark-sql

【解决方案1】：

AFAIk 你需要调用withColumn 两次（每个新列一次）。但是，如果您的 udf 计算量很大，您可以避免调用它两次，将“复杂”结果存储在临时列中，然后“解包”结果，例如使用列的apply 方法（它可以访问数组元素）。请注意，有时需要缓存中间结果（以防止在解包期间每行调用两次 UDF），有时则不需要。这似乎取决于 spark 优化计划的方式：

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")

【讨论】：