Как группировать по столбцам и агрегировать остальные столбцы

Question

Я новичок в pyspark, поэтому буду благодарен, если кто-нибудь сможет помочь решить проблему.

Предположим, у меня есть датафрейм в pyspark:

+----+----+----+----+----+
|col1|col2|col3|col4|col5|
+----+----+----+----+----+
|   A|2001|   2|   5|   6|
|   A|2001|   3|   6|  10|
|   A|2001|   3|   6|  10|
|   A|2002|   4|   5|   2|
|   B|2001|   2|   9|   4|
|   B|2001|   2|   4|   3|
|   B|2001|   2|   3|   4|
|   B|2001|   3|  95|   7|
+----+----+----+----+----+

Я хочу получить среднее значение col4, если соответствующие значения в col1, col2 и col3 совпадают, а затем избавиться от строк с повторяющимися значениями в первых 3 столбцах ,

Например, значения col1, col2, col3 в двух первых столбцах одинаковы, поэтому мы хотим исключить один из них и обновить значение col4 как среднее значение col4 и col5. Результат должен быть:

+----+----+----+----+----+
|col1|col2|col3|col4|col5|
+----+----+----+----+----+
|   A|2001|   2| 4.5|   7|
|   A|2001|   3|   6|  10|
|   A|2002|   4|   5|   2|
|   B|2001|   2|5.33|3.67|
|   B|2001|   3|  95|   7|
+----+----+----+----+----+

Аналогичный вопрос задавался, но в панде. Этот вопрос задается в фрейме pyspark

Answer 1

Шаг 1: Создание указанного DataFrame -

values = [('A',2001,2,5,6),('A',2001,2,4,8),('A',2001,3,6,10),('A',2002,4,5,2),
          ('B',2001,2,9,4),('B',2001,2,4,3),('B',2001,2,3,4),('B',2001,3,95,7)]
df = sqlContext.createDataFrame(values,['col1','col2','col3','col4','col5'])
df.show()
+----+----+----+----+----+
|col1|col2|col3|col4|col5|
+----+----+----+----+----+
|   A|2001|   2|   5|   6|
|   A|2001|   2|   4|   8|
|   A|2001|   3|   6|  10|
|   A|2002|   4|   5|   2|
|   B|2001|   2|   9|   4|
|   B|2001|   2|   4|   3|
|   B|2001|   2|   3|   4|
|   B|2001|   3|  95|   7|
+----+----+----+----+----+

Шаг 2: Агрегирование столбца col4 & col5.

df = df.groupby(['col1','col2','col3']).agg(avg('col4').alias('col4'),avg('col5').alias('col5'))
df.show()
+----+----+----+-----------------+------------------+
|col1|col2|col3|             col4|              col5|
+----+----+----+-----------------+------------------+
|   A|2001|   3|              6.0|              10.0|
|   A|2002|   4|              5.0|               2.0|
|   B|2001|   2|5.333333333333333|3.6666666666666665|
|   A|2001|   2|              4.5|               7.0|
|   B|2001|   3|             95.0|               7.0|
+----+----+----+-----------------+------------------+