Filtering within groupBy(x).agg to create averages on different subsets of the Pyspark Dataframe

user_7450545 0
Viewed 256

Normally when I have to make aggregations, I use something like the following code in PySpark:

import pyspark.sql.functions as f

df_aggregate = df.groupBy('id')\
                 .agg(f.mean('value_col').alias('value_col_mean'))

Now I actually want to compute the average or mean on multiple subsets of the dataframe df (i.e. on different time windows, for example a mean for the last year, a mean for the last 2 years, etc.). I understand I could do df.filter(f.col(filter_col) >= condition).groupBy.... for every subset, but instead I would prefer to do this in one 'go'.

Is it possible to apply the filtering within the .agg(..) part of PySpark?

Edit

Example data for one id looks like (the real data contains many values for id):

enter image description here

pyspark aggregate-functions dataframe
1 Answers

You can put the conditions inside a when statement, and put them all inside .agg:

import pyspark.sql.functions as f

df_aggregate = df.withColumn('value_col', f.regexp_replace('value_col', ',', '.'))\
                 .groupBy('id')\
                 .agg(f.mean(f.when(last_year_condition, f.col('value_col'))).alias('value_col_mean_last_year'),
                      f.mean(f.when(last_two_years_condition, f.col('value_col'))).alias('value_col_mean_last_two_years')
                     )
user_141657300
Related