I have a large dataframe (40 billion rows+) which can be grouped by key, I want to apply a custom calculation on few fields of each group and derive a single value for that group. eg, below dataframe has group_key and I want to derive a single value from col1 and col2 for the group_key. This calculation involves complex looping for col1 and col2. I have written a pyspark UDF which works fine and returns me the desired output for each group. However, the row-by-row UDF runs for few thousand rows which is slow, but it definitely does not run for 40 billion rows.
group_key col1 col2
123 a 5
123 a 6
123 b 6
123 cd 3
123 d 2
123 ab 9
456 d 4
456 ad 6
456 ce 7
456 a 4
456 s 3
desired output
group_key output
123 9.2
456 7.3
I know pyspark UDF is the worst choice since its row-by-row operation. For vectorized or batch operation I tried using pandas_udf (PandasUDFType: GROUPED_MAP) on each group. If I try to use for loop inside output_calc_udf it throws error. My logic needs to traverse each row of the dataframe passed to the pandas_udf and create a list out of it which I need to traverse few times until I keep pop-ing list elements and the list is empty. Its quite complex. But as I am understanding I cannot really do looping in vectorized operation ? How else I can execute the UDF for 40 billion rows ?
schema = StructType([StructField("output", DecimalType(), True)])
@F.pandas_udf(schema, F.PandasUDFType.GROUPED_MAP)
def output_calc_udf(pdf):
pdf=pdf[['output']]
return pdf
df_prob_new2.groupBy("group_key").apply(output_calc_udf).show(truncate=False)