Показать все столбцы pyspark после группы и агг

Question

I sh, чтобы сгруппировать столбец, а затем найти максимум другого столбца. Наконец, показать все столбцы на основе этого условия. Однако, когда я использовал свои коды, он показывает только 2 столбца, а не все.

# Normal way of creating dataframe in pyspark
sdataframe_temp = spark.createDataFrame([
    (2,2,'0-2'),
    (2,23,'22-24')],
    ['a', 'b', 'c']
)

sdataframe_temp2 = spark.createDataFrame([
    (4,6,'4-6'),
    (5,7,'6-8')],
    ['a', 'b', 'c']
)
# Concat two different pyspark dataframe
sdataframe_union_1_2 = sdataframe_temp.union(sdataframe_temp2)

sdataframe_union_1_2_g = sdataframe_union_1_2.groupby('a').agg({'b':'max'})

sdataframe_union_1_2_g.show()

output:

+---+------+
|  a|max(b)|
+---+------+
|  5|     7|
|  2|    23|
|  4|     6|
+---+------+

Ожидаемый результат:

+---+------+-----+
|  a|max(b)| c   |
+---+------+-----+
|  5|     7|6-8  |
|  2|    23|22-24|
|  4|     6|4-6  |
+---+------+---+

Answer 1

Вы можете использовать Window function, чтобы заставить его работать:

Метод 1: Использование оконной функции

import pyspark.sql.functions as F
from pyspark.sql.window import Window

w = Window().partitionBy("a").orderBy(F.desc("b"))

(sdataframe_union_1_2
.withColumn('max_val', F.row_number().over(w) == 1)
.where("max_val == True")
.drop("max_val")
.show())

+---+---+-----+
|  a|  b|    c|
+---+---+-----+
|  5|  7|  6-8|
|  2| 23|22-24|
|  4|  6|  4-6|
+---+---+-----+

Пояснение

1. Window функции полезны, когда мы хотим присоединить новый столбец к существующему набору столбцов.
2. В этом случае , Я говорю Window функция сгруппировать partitionBy('a') столбец и сортировать столбец b в порядке убывания F.desc(b). Это делает первое значение в b в каждой группе своим максимальным значением.
3. Затем мы используем F.row_number() для фильтрации максимальных значений, где номер строки равен 1.
4. Наконец, мы отбрасываем новый столбец, поскольку он не используется после фильтрации фрейма данных.

Метод 2: Использование groupby + внутреннее объединение

f = sdataframe_union_1_2.groupby('a').agg(F.max('b').alias('b'))

sdataframe_union_1_2.join(f, on=['a','b'], how='inner').show()

+---+---+-----+
|  a|  b|    c|
+---+---+-----+
|  2| 23|22-24|
|  5|  7|  6-8|
|  4|  6|  4-6|
+---+---+-----+