聚合 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|
+---+----+