聚合spark数据帧并根据数组类型列中是否存在值进行计数

Question

我有一个带有数组类型列的spark数据框：

scala> mydf.printSchema
root
 |-- arraycol: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- id: integer (nullable = false)

我现在需要通过“ id ”聚合该数据框，并根据数组中是否存在特定值进行计数。我正在尝试这样做：

val aggdata = mydf.groupBy("id").
  agg(
    count(when($"arraycol" contains "someval", $"arraycol")).as("aggval"))

这似乎不起作用。任何输入我该怎么做？

Answer 1

有一种array_contains方法可以测试条件：

val df = Seq((1, Seq("a", "b")), (1, Seq("b")), (2, Seq("b"))).toDF("id", "arrayCol")
// df: org.apache.spark.sql.DataFrame = [id: int, arrayCol: array<string>]

df.show
+---+--------+
| id|arrayCol|
+---+--------+
|  1|  [a, b]|
|  1|     [b]|
|  2|     [b]|
+---+--------+


df.groupBy("id").agg(
  count(when(array_contains($"arrayCol", "a"), $"arrayCol")).as("hasA")
).show
+---+----+
| id|hasA|
+---+----+
|  1|   1|
|  2|   0|
+---+----+

或使用sum：

df.groupBy("id").agg(
  sum(when(array_contains($"arrayCol", "a"), 1).otherwise(0)).as("hasA")
).show
+---+----+
| id|hasA|
+---+----+
|  1|   1|
|  2|   0|
+---+----+