Получение количества сгруппированных столбцов

Question

У меня есть фрейм данных, который содержит 3 сети, и в каждой сети много станций. То, что я хотел бы сделать, это получить общее количество станций для каждой сети. Кадр данных должен по-прежнему включать имя сети и станции, поэтому он должен выглядеть следующим образом:

Network Station Total
XMN     DIS     3     
XMN     CNN     3
XMN     JFK     3
ALK     DIS     2
ALK     CNN     2

Как бы я поступил так? Я предполагаю, что мне нужно сгруппировать по столбцам, а затем использовать функцию окна для разделения по сети и станции, чтобы получить общее количество? Я не уверен, но как мне это сделать?

Answer 1

Window.partitionBy делает именно это:

df = spark_session.createDataFrame([
    Row(Network='XMN', Station='DIS'),
    Row(Network='XMN', Station='CNN'),
    Row(Network='XMN', Station='JFK'),
    Row(Network='ALK', Station='DIS'),
    Row(Network='ALK', Station='CNN')
])

df.select("Network", "Station", count("*").over(Window.partitionBy("Network")).alias("Total")).show()

Вывод:

.+-------+-------+-----+
|Network|Station|Total|
+-------+-------+-----+
|    XMN|    DIS|    3|
|    XMN|    CNN|    3|
|    XMN|    JFK|    3|
|    ALK|    DIS|    2|
|    ALK|    CNN|    2|
+-------+-------+-----+

Answer 2

Вам нужно сгруппироватьBy, получить счет и присоединиться к вашему исходному фрейму данных

scala> val df = Seq(("XMN", "DIS"), ("XMN", "CNN"), ("XMN", "JFK"), ("ALK", "DIS"), ("ALK", "CNN")).toDF("Network", "Station")
df: org.apache.spark.sql.DataFrame = [Network: string, Station: string]

scala> df.show
+-------+-------+
|Network|Station|
+-------+-------+
|    XMN|    DIS|
|    XMN|    CNN|
|    XMN|    JFK|
|    ALK|    DIS|
|    ALK|    CNN|
+-------+-------+


scala> val grpCountDF = df.groupBy("Network").count
grpCountDF: org.apache.spark.sql.DataFrame = [Network: string, count: bigint]

scala> grpCountDF.show
+-------+-----+
|Network|count|
+-------+-----+
|    XMN|    3|
|    ALK|    2|
+-------+-----+


scala> val outputDF = df.join(grpCountDF, "Network")
outputDF: org.apache.spark.sql.DataFrame = [Network: string, Station: string ... 1 more field]

scala> outputDF.show
+-------+-------+-----+
|Network|Station|count|
+-------+-------+-----+
|    XMN|    DIS|    3|
|    XMN|    CNN|    3|
|    XMN|    JFK|    3|
|    ALK|    DIS|    2|
|    ALK|    CNN|    2|
+-------+-------+-----+