Spark SQL 使用 foldLeft 和 withColumn 替代 groupby/pivot/agg/collect_list 以提高性能

Question

我有一个由三列组成的 Spark DataFrame：

 id | col1 | col2 
-----------------
 x  |  p1  |  a1  
-----------------
 x  |  p2  |  b1
-----------------
 y  |  p2  |  b2
-----------------
 y  |  p2  |  b3
-----------------
 y  |  p3  |  c1

应用df.groupBy("id").pivot("col1").agg(collect_list("col2"))我得到以下数据帧（aggDF）：

+---+----+--------+----+
| id|  p1|      p2|  p3|
+---+----+--------+----+
|  x|[a1]|    [b1]|  []|
|  y|  []|[b2, b3]|[c1]|
+---+----+--------+----+

然后我找到除id列之外的列的名称。

val cols = aggDF.columns.filter(x => x != "id")

之后我使用cols.foldLeft(aggDF)((df, x) => df.withColumn(x, when(size(col(x)) > 0, col(x)).otherwise(lit(null))))用null替换空数组。 当列数增加时，此代码的性能会变差。 另外，我有字符串列的名称val stringColumns = Array("p1","p3") 。 我想获得以下最终数据框：

+---+----+--------+----+
| id|  p1|      p2|  p3|
+---+----+--------+----+
|  x| a1 |    [b1]|null|
|  y|null|[b2, b3]| c1 |
+---+----+--------+----+

为了实现最终的数据框，有没有更好的解决方案来解决这个问题？

Answer 1

您当前的代码按结构支付了 2 项性能成本：

• 正如 Alexandros 所提到的，您为每个 DataFrame 转换支付 1 次催化剂分析，因此如果您循环其他数百或数千列，您会注意到在实际提交作业之前花费了一些时间在驱动程序上。 如果这对您来说是一个关键问题，您可以在 withColumns 上使用单个 select 语句而不是 foldLeft 但这不会因为下一点而真正改变执行时间

• 当您在可以优化为单个 select 语句的列上使用诸如 when().otherwise() 之类的表达式时，代码生成器将生成一个处理所有列的大型方法。 如果您有几百个以上的列，很可能默认情况下 JVM 不会对生成的方法进行 JIT 编译，从而导致执行性能非常慢（Hotspot 中的最大 JIT-able 方法是 8k 字节码）。

您可以通过检查执行程序日志来检测是否遇到第二个问题，并检查是否在无法 JIT 的太大方法上看到警告。

如何尝试解决这个问题？

1 - 改变逻辑

您可以使用窗口变换过滤枢轴之前的空单元格

import org.apache.spark.sql.expressions.Window

val finalDf = df
  .withColumn("count", count('col2) over Window.partitionBy('id,'col1)) 
  .filter('count > 0)
  .groupBy("id").pivot("col1").agg(collect_list("col2"))

这可能会更快，也可能不会更快，具体取决于实际数据集，因为数据透视本身也会生成一个大型 select 语句表达式，因此如果遇到超过大约 500 个 col1 值，它可能会达到大型方法阈值。 您可能还想将此与选项 2 结合使用。

2 - 尝试并完善 JVM

你可以在你的执行器上添加一个 extraJavaOption 来要求 JVM 尝试和 JIT 大于 8k 的热方法。

例如，在 spark-submit 上添加选项--conf "spark.executor.extraJavaOptions=-XX:-DontCompileHugeMethods"并查看它如何影响数据透视执行时间。

如果没有关于真实数据集的更多详细信息，很难保证大幅提高速度，但绝对值得一试。

Answer 2

如果您查看https://medium.com/@manuzhang/the-hidden-cost-of-spark-withcolumn-8ffea517c015，那么您会看到带有 foldLeft 的 withColumn 存在已知的性能问题。 Select 是一种替代方法，如下所示 - 使用可变参数。

不相信 collect_list 是一个问题。 我也保留了第一组逻辑。 pivot 启动一个 Job 以获得不同的值进行旋转。 这是一种公认​​的方法imo。 尝试自己动手对我来说似乎毫无意义，但其他答案可能证明我错了，或者 Spark 2.4 已得到改进。

import spark.implicits._ 
import org.apache.spark.sql.functions._

// Your code & assumig id is only col of interest as in THIS question. More elegant than 1st posting.
val df = Seq( ("x","p1","a1"), ("x","p2","b1"), ("y","p2","b2"), ("y","p2","b3"), ("y","p3","c1")).toDF("id", "col1", "col2")
val aggDF = df.groupBy("id").pivot("col1").agg(collect_list("col2")) 
//aggDF.show(false)

val colsToSelect = aggDF.columns  // All in this case, 1st col id handled by head & tail

val aggDF2 = aggDF.select((col(colsToSelect.head) +: colsToSelect.tail.map
    (col => when(size(aggDF(col)) === 0,lit(null)).otherwise(aggDF(col)).as(s"$col"))):_*)
aggDF2.show(false)

返回：

+---+----+--------+----+
|id |p1  |p2      |p3  |
+---+----+--------+----+
|x  |[a1]|[b1]    |null|
|y  |null|[b2, b3]|[c1]|
+---+----+--------+----+

BTW 也是一本不错的读物： https : //lansalo.com/2018/05/13/spark-how-to-add-multiple-columns-in-dataframes-and-how-not-to/ 。 随着列数的增加，效果变得更加明显。 最后，读者提出了一个相关的观点。

我认为当列数较多时，选择方法的性能更好。

UPD：在假期期间，我在 Spark 2.4.x 上试用了这两种方法，在多达 1000 列的情况下几乎没有观察到差异。 这让我很困惑。