Bucketize the pyspark dataframe rows into window duration of variable length based on values in a column

user_3541078 0
Viewed 48

I have a dataframe in pyspark (sample_data), and I want to bucketize the rows based on a column window_length in the dataframe. The original table has about 500 million rows, and distinct values in org column are of the order of several thousand (~10K).

spark = SparkSession.builder.appName("test_app").getOrCreate()

sample_data = (\
    ("2021-09-30 14:00:00", "Shiping", 3000, "60 seconds"), \
    ("2021-09-30 14:01:00", "Shiping", 4600, "60 seconds"),  \
    ("2021-09-30 14:02:00", "Shiping", 4100, "60 seconds"),   \
    ("2021-09-30 14:03:00", "Shiping", 3000, "60 seconds"),    \
    ("2021-09-30 14:00:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:01:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:02:00", "Finance", 2700, "105 seconds"),\
    ("2021-09-30 14:03:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:04:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:05:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:06:00", "Finance", 3000, "105 seconds"),  \
    ("2021-09-30 14:00:00", "Marketing", 2000, "120 seconds"), \
    ("2021-09-30 14:01:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:02:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:03:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:04:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:05:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:06:00", "Marketing", 1200, "120 seconds"), \
    ("2021-09-30 14:07:00", "Marketing", 1200, "120 seconds"), \
             )
columns= ["sale_time", "org", "value", "window_length"]
df = spark.createDataFrame(data = sample_data, schema = columns)
df = df.withColumn("sale_time",F.to_utc_timestamp("sale_time", "UTC"))

I can aggregate this dataframe over a constant time period window_duration (60 seconds), as shown below:

#aggregation query for 60 seconds window duration
df.groupBy('org' , F.window("sale_time", "60 second", startTime="0 seconds").alias("ts"))\
  .agg(F.avg("value").alias("avg_val")).orderBy(["org","ts"]).show(truncate=False)

which gives the following result:

enter image description here

Now, instead of 60 seconds, I want to write a query which can run the above aggregation query for the window_duration (60 sec for Shipping, 105 sec for Finance, and 120 sec for Marketing ) from window_length column of the sample_data dataframe. Please note, window_length is constant for a value in org column but window_length can have different values if the value in org field is different.

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

The only way I was able to think was to loop on all org separately and combine , Something like below

dfc=df.filter("1!=1").groupBy('org' , F.window("sale_time","1 seconds" , startTime="0 seconds").alias("ts"))\
  .agg(F.avg("value").alias("avg_val")) #Creating empty dataframe to hold combined data
fe=df.select(F.collect_set(F.concat('org',F.lit(","),'window_length')).alias('org')).first()['org'] # getting different org and time
for og in fe:
  df1=df.filter(f"org='{og.split(',')[0]}'").groupBy('org' , F.window("sale_time",og.split(',')[1] , startTime="0 seconds").alias("ts"))\
  .agg(F.avg("value").alias("avg_val")).orderBy(["org","ts"])
  dfc=dfc.unionAll(df1) # Combining result dataframe
  
dfc.show()

#output
+---------+------------------------------------------+-------+
|org      |ts                                        |avg_val|
+---------+------------------------------------------+-------+
|Marketing|{2021-09-30 14:00:00, 2021-09-30 14:02:00}|1600.0 |
|Marketing|{2021-09-30 14:02:00, 2021-09-30 14:04:00}|1200.0 |
|Marketing|{2021-09-30 14:04:00, 2021-09-30 14:06:00}|1200.0 |
|Marketing|{2021-09-30 14:06:00, 2021-09-30 14:08:00}|1200.0 |
|Finance  |{2021-09-30 14:00:00, 2021-09-30 14:01:45}|3000.0 |
|Finance  |{2021-09-30 14:01:45, 2021-09-30 14:03:30}|2850.0 |
|Finance  |{2021-09-30 14:03:30, 2021-09-30 14:05:15}|3000.0 |
|Finance  |{2021-09-30 14:05:15, 2021-09-30 14:07:00}|3000.0 |
|Shiping  |{2021-09-30 14:00:00, 2021-09-30 14:01:00}|3000.0 |
|Shiping  |{2021-09-30 14:01:00, 2021-09-30 14:02:00}|4600.0 |
|Shiping  |{2021-09-30 14:02:00, 2021-09-30 14:03:00}|4100.0 |
|Shiping  |{2021-09-30 14:03:00, 2021-09-30 14:04:00}|3000.0 |
+---------+------------------------------------------+-------+
user_190757710
Related