scala - 数组列的 Spark 聚合

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]|
+---+------------+

现在我需要对值列应用不同的聚合函数。
我可以打电话 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，诚然没有很好的文档记录并且由于类型签名变得非常冗长而使用起来非常困惑。

更长的答案是此功能即将在 Apache Spark 2.4 中推出(？)。

父问题 SPARK-23899补充:

array_max

array_min

聚合

map

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可以已经获得数组的长度，这意味着可以计算平均值。