Разделение окна по количеству агрегации

Question

Я хочу сделать подсчет за окном. Результат подсчета агрегации должен быть сохранен в новом столбце:

Входные данные:

    val df = Seq(("N1", "M1","1"),("N1", "M1","2"),("N1", "M2","3")).toDF("NetworkID", "Station","value")

    +---------+-------+-----+
    |NetworkID|Station|value|
    +---------+-------+-----+
    |       N1|     M1|    1|
    |       N1|     M1|    2|
    |       N1|     M2|    3|
    +---------+-------+-----+

    val w = Window.partitionBy(df("NetworkID"))

Результат, который у меня пока есть:

        df.withColumn("count", count("Station").over(w)).show()
        +---------+-------+-----+-----+
        |NetworkID|Station|value|count|
        +---------+-------+-----+-----+
        |       N1|     M2|    3|    3|
        |       N1|     M1|    1|    3|
        |       N1|     M1|    2|    3|
        +---------+-------+-----+-----+

Результат, который я хотел бы получить:

+---------+-------+-----+-----+

|NetworkID|Station|value|count|

+---------+-------+-----+-----+

|       N1|     M2|    3|    2|

|       N1|     M1|    1|    2|

|       N1|     M1|    2|    2|

+---------+-------+-----+-----+

Поскольку количество станций для NetworkID N1 равно 2 (M1 и M2).

Я знаю, что смогу сделать это, создав новый фрейм данных, выбрав 2 столбца NetworkID и Station, и сделаю groupBy и присоединюсь к первому.

Но у меня есть много совокупного подсчета для разных столбцов в моем фрейме данных, и я должен избегать объединений.

Заранее спасибо

Answer 1

Ваш раздел также нужен для столбца «Станция», поскольку вы подсчитываете Станции для каждого NetworkID.

scala> val df = Seq(("N1", "M1","1"),("N1", "M1","2"),("N1", "M2","3"),("N2", "M1", "4"), ("N2", "M2", "2")).toDF("NetworkID", "Station", "value")
df: org.apache.spark.sql.DataFrame = [NetworkID: string, Station: string ... 1 more field]

scala> val w = Window.partitionBy("NetworkID", "Station")
w: org.apache.spark.sql.expressions.WindowSpec = org.apache.spark.sql.expressions.WindowSpec@5b481d77

scala> df.withColumn("count", count("Station").over(w)).show()
+---------+-------+-----+-----+
|NetworkID|Station|value|count|
+---------+-------+-----+-----+
|       N2|     M2|    2|    1|
|       N1|     M2|    3|    1|
|       N2|     M1|    4|    1|
|       N1|     M1|    1|    2|
|       N1|     M1|    2|    2|
+---------+-------+-----+-----+