Aggregate values based upon conditions in pyspark

user_13834671 0
Viewed 454

I am new to Spark and I need some help regarding aggregation of values.

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

I need to add a new column to this dataframe where I add the values present in amount if I see 'buy' in transaction_code For example, I add 10 and 20 as the transaction_code for them is 'buy'.

I know how to aggregate them totally and below is the code I have written.

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

I found we can use if function but I am unable to determine how to initialize them and how to keep track of the amount . Please help me in this matter.Thanks a lot !

pyspark apache-spark-sql apache-spark hadoop
1 Answers

you can use array_zip and 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|
    # +----------------+----------------+-----+
user_86658430
Related