数组列的火花聚集

Question

我有一个带有数组列的数据框。

val json = """[
{"id": 1, "value": [11, 12, 18]},
{"id": 2, "value": [23, 21, 29]}
]"""

val df = spark.read.json(Seq(json).toDS)

scala> df.show
+---+------------+
| id|       value|
+---+------------+
|  1|[11, 12, 18]|
|  2|[23, 21, 29]|
+---+------------+

现在，我需要将不同的聚合函数应用于value列。 我可以打电话给explode和groupBy

df.select($"id", explode($"value").as("value")).groupBy($"id").agg(max("value"), avg("value")).show

+---+----------+------------------+
| id|max(value)|        avg(value)|
+---+----------+------------------+
|  1|        18|13.666666666666666|
|  2|        29|24.333333333333332|
+---+----------+------------------+

在这里令我困扰的是，我将DataFrame分解成一个更大的数据，然后将其简化为原始调用groupBy。

是否有更好（即更有效）的方法来调用数组列上的聚合函数？也许我可以实现UDF，但是我不想自己实现所有聚合UDF。

编辑。有人引用了this SO question，但在我的情况下不起作用。 size运行正常

scala> df.select($"id", size($"value")).show
+---+-----------+
| id|size(value)|
+---+-----------+
|  1|          3|
|  2|          3|
+---+-----------+

但是avg或max不起作用。

Answer 1

简短的回答是“否” ，您必须实现自己的UDF才能在数组列上进行汇总。至少在最新版本的Spark（撰写本文时为2.3.1）中。正如您正确断言的那样，它效率不高，因为它会迫使您爆炸行或支付在Dataset API中工作的序列化和反序列化成本。

对于可能会发现此问题的其他人，要使用数据集以类型安全的方式编写聚合，则可以使用Aggregator API，该API公认的文档不多，并且作为类型签名使用时非常混乱变得很冗长。

更长的答案是，此功能在Apache Spark 2.4中很快就会出现。

父刊SPARK-23899添加：

• array_max
• array_min
• 汇总
• 地图
• array_distinct
• array_remove
• array_join

还有许多其他

本次演讲“ Extending Spark SQL API with Easier to Use Array Types Operations”在2018年6月的Spark + AI峰会上进行了介绍，涵盖了新功能。

如果已发布该版本，使您可以像示例中那样使用max函数，那么average则比较棘手。 奇怪的是，不存在array_sum，但是它可以通过aggregate函数来构建。可能看起来像这样：

def sum_array(array_col: Column) = aggregate($"my_array_col", 0, (s, x) => s + x, s => s) df.select(sum_array($"my_array_col") 零值是聚合缓冲区的初始状态。

正如您所指出的，size已经可以获取数组的长度，这意味着可以计算平均值。