Как я могу сосчитать два разных вхождения в кадре данных за один проход?

Question

Скажем, у меня есть кадр данных, который выглядит следующим образом:

+---------+------+
|  Col1   | Col2 |
+---------+------+
| Value 1 | A    |
| Value 2 | D    |
| Value 3 | B    |
| Value 4 | C    |
| Value 5 | A    |
| Value 6 | B    |
+---------+------+

Мне нужно посчитать, сколько раз появляется A, а также сколько раз появляется B.В настоящее время я просто делаю:

val aCount = dataframe.where("Col2 = A").count()
val bCount = dataframe.where("Col2 = B").count()

Но это приводит к тому, что весь DAG выполняется дважды, один раз для каждого счета.Данные, с которыми я работаю, огромны, поэтому я не могу просто кэшировать dataframe, и выполнение этого дважды занимает слишком много времени.
Есть ли способ, которым я могу выполнить оба подсчета за один проход?

Answer 1

Вы можете сделать это следующим образом за 1 проход:

val countMap = dataframe
  .where($"Col2".isin("A","B"))
  .groupBy($"Col2")
  .count()
  .as[(String,Long)].collect()
  .toMap

val aCount = countMap("A")
val bCount = countMap("B")

или только с одним утверждением:

val Array(aCount,bCount) = dataframe
  .where($"Col2".isin("A","B"))
  .groupBy($"Col2")
  .count()
  .orderBy($"Col2")
  .select($"count").as[Long]
  .collect()

Answer 2

Учитывая

val data = spark.createDataFrame(Seq(
  ("Value 1", "A"), 
  ("Value 2", "D"),
  ("Value 3", "B"),
  ("Value 4", "C"), 
  ("Value 5", "A"), 
  ("Value 6", "B"))
).toDF("Col1", "Col2")

Это то, что вы ищете?

data.filter($"Col2".isin("A", "B")).groupBy("Col2").count.show
+----+-----+
|Col2|count|
+----+-----+
|   B|    2|
|   A|    2|
+----+-----+

Затем вы могли бы потенциально .collect и извлечь пары ключ / значение