Агрегирование искрового фрейма данных и подсчет на основе наличия значения в столбце типа массива

Question

У меня есть искровой фрейм данных со столбцом типа массива:

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