Сгруппировать и найти количество, прежде чем делать поворот искры

Question

У меня есть кадр данных, как показано ниже

A   B   C       D
foo one small   1
foo one large   2
foo one large   2
foo two small   3

Мне нужно groupBy на основе A и B pivot для столбца C и sum столбец D

Яв состоянии сделать это, используя

df.groupBy("A", "B").pivot("C").sum("D") 

Однако мне нужно также найти count после groupBy, если я пытаюсь что-то вроде

df.groupBy("A", "B").pivot("C").agg(sum("D"), count)

, я получаю вывод вроде

A   B   large   small large_count small_count

Есть ли способ получить только один count после groupBy перед выполнением pivot

Answer 1

Это то, что вы ожидаете?

val df = Seq(("foo", "one", "small",   1),
("foo", "one", "large",   2),
("foo", "one", "large",   2),
("foo", "two", "small",   3)).toDF("A","B","C","D")

scala> df.show
+---+---+-----+---+
|  A|  B|    C|  D|
+---+---+-----+---+
|foo|one|small|  1|
|foo|one|large|  2|
|foo|one|large|  2|
|foo|two|small|  3|
+---+---+-----+---+

scala> val df2 = df.groupBy('A,'B).pivot("C").sum("D")
df2: org.apache.spark.sql.DataFrame = [A: string, B: string ... 2 more fields]

scala> val df3 = df.groupBy('A as "A1",'B as "B1").agg(sum('D) as "sumd")
df3: org.apache.spark.sql.DataFrame = [A1: string, B1: string ... 1 more field]

scala> df3.join(df2,'A==='A1 and 'B==='B1,"inner").select("A","B","sumd","large","small").show
+---+---+----+-----+-----+
|  A|  B|sumd|large|small|
+---+---+----+-----+-----+
|foo|one|   5|    4|    1|
|foo|two|   3| null|    3|
+---+---+----+-----+-----+


scala>

Answer 2

На выходе попробуйте

output.withColumn ("count", $ "large_count" + $ "small_count"). Show

Вы можете удалить два столбца count, если хотите

Чтобы сделать это до разворота, попробуйте df.groupBy ("A", "B"). Agg (count ("C"))