Pyspark aggregation - pyspark

Question

I am trying to aggregate pyspark dataframe. Sample looks like this:
+---+-------------------+
| id| struct|
+---+-------------------+
|id1| [foo, true, true]|
|id1| [foo, true, false]|
|id1|[foo, false, false]|
|id1| [bar, true, true]|
|id1| [bar, true, false]|
|id1|[bar, false, false]|
|id2| [foo, true, true]|
|id2|[foo, false, false]|
|id2| [bar, true, true]|
|id2|[bar, false, false]|
+---+-------------------+
ID column will have max 1500 unique IDs, struct.name will have 5 unique values.
Here is my code, which compute what I want:
from pyspark.sql.types import *
from shared.spark import start_spark
import pyspark.sql.functions as F
spark = start_spark('app')
schema = StructType([StructField('id', StringType()),
StructField('struct', StructType(
[StructField('name', StringType()),
StructField('param1', BooleanType()),
StructField('param2', BooleanType()),
]
))])
data = [['id1', ['foo', True, True]],
['id1', ['foo', True, False]],
['id1', ['foo', False, False]],
['id1', ['bar', True, True]],
['id1', ['bar', True, False]],
['id1', ['bar', False, False]],
['id2', ['foo', True, True]],
['id2', ['foo', False, False]],
['id2', ['bar', True, True]],
['id2', ['bar', False, False]]
]
df = spark.createDataFrame(data, schema)
df.groupby('id')\
.agg(F.count(F.when((df['struct.name']=='foo') &
(df['struct.param1']) &
(df['struct.param2']), 1)).alias('foo_cond1'),
F.count(F.when((df['struct.name'] == 'foo') &
(df['struct.param1']) &
(df['struct.param2']==False), 1)).alias('foo_cond2'),
F.count(F.when((df['struct.name'] == 'foo') &
(df['struct.param1']==False) &
(df['struct.param2']==False), 1)).alias('foo_cond3'),
F.count(F.when((df['struct.name']=='bar') &
(df['struct.param1']) &
(df['struct.param2']), 1)).alias('bar_cond1'),
F.count(F.when((df['struct.name'] == 'bar') &
(df['struct.param1']) &
(df['struct.param2']==False), 1)).alias('bar_cond2'),
F.count(F.when((df['struct.name'] == 'bar') &
(df['struct.param1']==False) &
(df['struct.param2']==False), 1)).alias('bar_cond3'),
) \
.withColumn('foo', F.struct(F.col('foo_cond1').alias('cond1'),
F.col('foo_cond2').alias('cond2'),
F.col('foo_cond3').alias('cond3')
)
) \
.withColumn('bar', F.struct(F.col('bar_cond1').alias('cond1'),
F.col('bar_cond2').alias('cond2'),
F.col('bar_cond3').alias('cond3')
)
) \
.select('id', 'foo', 'bar') \
.show()
And the result is as follows:
+---+---------+---------+
| id| foo| bar|
+---+---------+---------+
|id1|[1, 1, 1]|[1, 1, 1]|
|id2|[1, 0, 1]|[1, 0, 1]|
+---+---------+---------+
Is there a better way to do such aggregation, which will perform better with less code? Maybe using pandas UDAF? Appreciate every comment. Thanks

Answer 1

I was able to use pandasUDFType, but it seems that running time had increased on more that 30%. But I only use mentioned sample data.
from pyspark.sql.types import *
from shared.spark import start_spark
import pyspark.sql.functions as F
from pyspark.sql.functions import pandas_udf, PandasUDFType
spark = start_spark('app')
schema = StructType([StructField('id', StringType()),
StructField('struct', StructType(
[StructField('name', StringType()),
StructField('param1', BooleanType()),
StructField('param2', BooleanType()),
]
))])
schema_udf = StructType(
[StructField('id', StringType()),
StructField('foo1', DoubleType()),
StructField('foo2', DoubleType()),
StructField('foo3', DoubleType()),
StructField('bar1', DoubleType()),
StructField('bar2', DoubleType()),
StructField('bar3', DoubleType()),
])
data = [['id1', ['foo', True, True]],
['id1', ['foo', True, False]],
['id1', ['foo', False, False]],
['id1', ['bar', True, True]],
['id1', ['bar', True, False]],
['id1', ['bar', False, False]],
['id2', ['foo', True, True]],
['id2', ['foo', False, False]],
['id2', ['bar', True, True]],
['id2', ['bar', False, False]]
]
df = spark.createDataFrame(data, schema)
#pandas_udf(schema_udf, PandasUDFType.GROUPED_MAP)
def myGroupby(df_group):
def countComb(df):
def countCombinations(param1, param2):
cond1, cond2, cond3 = 0, 0, 0
if param1:
if param2:
cond1 += 1
else:
cond2 += 1
else:
cond3 += 1
return cond1, cond2, cond3
if df['name']=='foo':
df['foo1'], df['foo2'], df['foo3'] = countCombinations(df.param1, df.param2)
if df['name']=='bar':
df['bar1'], df['bar2'], df['bar3'] = countCombinations(df.param1, df.param2)
return df
df_result = df_group.apply(countComb, axis=1)
return df_result[['id', 'foo1', 'foo2', 'foo3', 'bar1', 'bar2', 'bar3']].groupby('id').sum().reset_index()
df \
.select('id', 'struct.name', 'struct.param1', 'struct.param2') \
.groupby("id") \
.apply(myGroupby) \
.withColumn('foo', F.struct(F.col('foo1').alias('cond1'),
F.col('foo2').alias('cond2'),
F.col('foo3').alias('cond3')
)
) \
.withColumn('bar', F.struct(F.col('bar1').alias('cond1'),
F.col('bar2').alias('cond2'),
F.col('bar3').alias('cond3')
)
) \
.select('id', 'foo', 'bar') \
.show()
So what is the best practise for writing pyspark aggregations? What if I want to aggregate on lot of columns, it is better to write a lot of conditions (like in original question) or divide aggregation to smaller parts and then join dataframes?

Answer 2

How about this one?
Split the Struc into separate columns and use window to compute the count.
from pyspark.sql.types import *
import pyspark.sql.functions as F
schema = StructType([StructField('id', StringType()),
StructField('struct', StructType(
[StructField('name', StringType()),
StructField('param1', BooleanType()),
StructField('param2', BooleanType()),
]
))])
data = [['id1', ['foo', True, True]],
['id1', ['foo', True, False]],
['id1', ['foo', False, False]],
['id1', ['bar', True, True]],
['id1', ['bar', True, False]],
['id1', ['bar', False, False]],
['id2', ['foo', True, True]],
['id2', ['foo', False, False]],
['id2', ['bar', True, True]],
['id2', ['bar', False, False]]
]
df = spark.createDataFrame(data, schema)
df = df.withColumn('name', F.col('struct').getField('name'))
df = df.withColumn('param1', F.col('struct').getField('param1'))
df = df.withColumn('param2', F.col('struct').getField('param2'))
w = Window.partitionBy(['id', 'name'])
df = df.withColumn('c1', F.count(F.when((df['param1']==True)&(df['param2']==True), 1)).over(w))
df = df.withColumn('c2', F.count(F.when((df['param1']==True)&(df['param2']==False), 1)).over(w))
df = df.withColumn('c3', F.count(F.when((df['param1']==False)&(df['param2']==False), 1)).over(w))
df = df.withColumn('result', F.array(['c1', 'c2', 'c3']))
df.show()
+---+-------------------+----+------+------+---+---+---+---------+
| id| struct|name|param1|param2| c1| c2| c3| result|
+---+-------------------+----+------+------+---+---+---+---------+
|id2| [bar, true, true]| bar| true| true| 1| 0| 1|[1, 0, 1]|
|id2|[bar, false, false]| bar| false| false| 1| 0| 1|[1, 0, 1]|
|id1| [foo, true, true]| foo| true| true| 1| 1| 1|[1, 1, 1]|
|id1|[foo, false, false]| foo| false| false| 1| 1| 1|[1, 1, 1]|
|id1| [foo, true, false]| foo| true| false| 1| 1| 1|[1, 1, 1]|
|id1| [bar, true, true]| bar| true| true| 1| 1| 1|[1, 1, 1]|
|id1| [bar, true, false]| bar| true| false| 1| 1| 1|[1, 1, 1]|
|id1|[bar, false, false]| bar| false| false| 1| 1| 1|[1, 1, 1]|
|id2|[foo, false, false]| foo| false| false| 1| 0| 1|[1, 0, 1]|
|id2| [foo, true, true]| foo| true| true| 1| 0| 1|[1, 0, 1]|
+---+-------------------+----+------+------+---+---+---+---------+
We then use pivot:
df = df.groupby('id').pivot('name').agg(F.first('result'))
df.show()
+---+---------+---------+
| id| bar| foo|
+---+---------+---------+
|id1|[1, 1, 1]|[1, 1, 1]|
|id2|[1, 0, 1]|[1, 0, 1]|
+---+---------+---------+