python - 如何在 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'])

我想groupby url并对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|
+---+----------+----------+-------+-----------------+