как применить много операций в разных столбцах в искровом фрейме и сохранить их с новым псевдонимом

Question

У меня есть следующий spark фрейм данных

df = spark.createDataFrame([['2017-04-01', 'A',1 , 1],
                             ['2017-04-01', 'B',2,3],
                             ['2017-04-01', 'B',3,4],
                             ['2017-04-01', 'A',5,5]], schema=['pdate', 'url', 'weight', 'imp'])

Я бы хотел groupby url и выполнить следующие операции над df и присвоить результат новым столбцам:

• min из pdate (min_pdate как псевдоним)
• max из pdate (max_pdate в качестве псевдонима)
• sum из imp (sum_imp в качестве псевдонима)
• weighted_mean из imp (wmean_imp в качестве псевдонима)

Есть ли удобный способ сделать это с помощью pyspark?

Answer 1

Просто используйте функцию agg, чтобы применить множество функций к groupBy

import pyspark.sql.functions as f

from pyspark.shell import spark

df = spark.createDataFrame([['2017-03-01', 'A', 1, 1],
                            ['2017-04-01', 'B', 2, 3],
                            ['2017-05-01', 'B', 3, 4],
                            ['2017-06-01', 'A', 5, 5]], schema=['pdate', 'url', 'weight', 'imp'])

df = df \
    .groupBy(f.col('url')) \
    .agg(f.min('pdate').alias('min_pdate'),
         f.max('pdate').alias('max_pdate'),
         f.sum('imp').alias('sum_imp'),
         (f.sum(f.col('imp') * f.col('weight')) / f.sum(f.col('weight'))).alias('wmean_imp'))
df.show()

Выход:

+---+----------+----------+-------+-----------------+
|url| min_pdate| max_pdate|sum_imp|        wmean_imp|
+---+----------+----------+-------+-----------------+
|  B|2017-04-01|2017-05-01|      7|              3.6|
|  A|2017-03-01|2017-06-01|      6|4.333333333333333|
+---+----------+----------+-------+-----------------+