根据 pyspark 中的条件聚合值

Posted 2023-04-15

【中文标题】根据 pyspark 中的条件聚合值

【英文标题】：Aggregate values based upon conditions in pyspark

【发布时间】：2020-06-29 13:31:12

【问题描述】：

我是 Spark 的新手，我需要一些关于值聚合的帮助。

 +--------------------+--------------------+-----+
|              amount|    transaction_code|Total|
+--------------------+--------------------+-----+
|[10, 20, 30, 40, ...|[buy, buy, sell, ...|210.0|
+--------------------+--------------------+-----+

如果我在 transaction_code 中看到“购买”，我需要在此数据框中添加一个新列 例如，我添加了 10 和 20，因为它们的 transaction_code 是“购买”。

我知道如何完全聚合它们，下面是我编写的代码。

df2extract = df2extract.select(
    'amount',
    'transaction_code',
   F.expr('AGGREGATE(amount, cast(0 as float), (acc, x) -> acc + x)').alias('Total')
 ).show()

我发现我们可以使用 if 函数，但我无法确定如何初始化它们以及如何跟踪数量。请在这件事上帮助我。非常感谢！

【参考方案1】：

您可以使用array_zip 和filter。

    from pyspark.sql import SparkSession
    from pyspark.sql import functions as F

    spark = SparkSession.builder \
        .appName('SO')\
        .getOrCreate()

    sc= spark.sparkContext

    df = sc.parallelize([
        ([10, 20, 30, 40], ["buy", "buy", "sell"])]).toDF(["amount", "transaction_code"])

    df.show()

    # +----------------+----------------+
    # |          amount|transaction_code|
    # +----------------+----------------+
    # |[10, 20, 30, 40]|[buy, buy, sell]|
    # +----------------+----------------+

    df1 = df.withColumn("zip", F.arrays_zip(F.col('amount'),F.col('transaction_code')))

    df2 = df1.withColumn("buy_filter", F.expr('''filter(zip, x-> x.transaction_code == 'buy')'''))

    df3 = df2.select("amount", "transaction_code", F.col("buy_filter.amount").alias("buy_values"))

    df3.select("amount", "transaction_code", F.expr('AGGREGATE(buy_values, cast(0 as float), (acc, x) -> acc + x)').alias('total')).show()

    # +----------------+----------------+-----+
    # |          amount|transaction_code|total|
    # +----------------+----------------+-----+
    # |[10, 20, 30, 40]|[buy, buy, sell]| 30.0|
    # +----------------+----------------+-----+