如何在spark数据帧的不同列中应用许多操作并使用新别名保存它们

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'])

我想对url进行groupby并在df上执行以下操作，并将结果分配给新列：

• pdate min （min_pdate作为别名）
• pdate max （max_pdate作为别名）
• imp sum （sum_imp作为别名）
• imp weighted_mean （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|
+---+----------+----------+-------+-----------------+