How can I get same result in both Pandas and Pyspark during groupby().agg() operation for categorical columns?

user_11018563 0
Viewed 200

I am using groupby() and agg() function together to find the group results. I used both pandas and pyspark. Pandas result and pyspark result is different. Seems pyspark result is not correct. I also noticed that if I don't change columns to category type pandas and pyspark results are similar. How can i fix the pyspark code so that i will get the same result in both cases?

Here is my code which I tried:

import pandas as pd
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()


pivot_agg_func= {'Max_Speed': 'min', 'Delay': 'sum'}
cat_col_all = ['Date', 'TrainID', 'Traintype']

df = pd.read_csv('groupby_example.csv', index_col=0)

print('Original DF')
print(df)

for col in cat_col_all:
    df[col] = df[col].astype('category')

df_groupby = df.groupby(cat_col_all, as_index=True).agg(pivot_agg_func).reset_index()
print('Pandas group by result')
print(df_groupby)

# reading spark dataframe
df_sp = spark.read.option("header", True).csv('groupby_example.csv').drop('_c0')
df = df_sp.groupby(*cat_col_all).agg(pivot_agg_func)
print('Pyspark group by result')
df.show()

Here is my output from console:

Original DF

      TrainID Traintype  Max_Speed  Delay  Date
2017   1000.0        IC      280.0   11.0  2017
2017   1000.0       ICE      280.0    2.0  2017
2018   1002.0        IC        NaN    NaN  2018
2019   1002.0        IC      260.0    3.0  2019
2019   1002.0        IC      220.0    9.0  2019
2019   1000.0       ICE      270.0    NaN  2019
2020   1000.0        IC      280.0    3.0  2020
2020   1000.0        IC      260.0   55.0  2020
2020   1002.0       ICE      280.0    NaN  2020
2020   1000.0        IC      220.0    4.0  2020

Pandas group by and agg result:

    Date TrainID Traintype  Max_Speed  Delay
0   2017  1000.0        IC      280.0   11.0
1   2017  1000.0       ICE      280.0    2.0
2   2017  1002.0        IC        NaN    NaN
3   2017  1002.0       ICE        NaN    NaN
4   2018  1000.0        IC        NaN    NaN
5   2018  1000.0       ICE        NaN    NaN
6   2018  1002.0        IC        NaN    0.0
7   2018  1002.0       ICE        NaN    NaN
8   2019  1000.0        IC        NaN    NaN
9   2019  1000.0       ICE      270.0    0.0
10  2019  1002.0        IC      220.0   12.0
11  2019  1002.0       ICE        NaN    NaN
12  2020  1000.0        IC      220.0   62.0
13  2020  1000.0       ICE        NaN    NaN
14  2020  1002.0        IC        NaN    NaN
15  2020  1002.0       ICE      280.0    0.0

Pyspark group by and agg result:

+----+-------+---------+--------------+----------+
|Date|TrainID|Traintype|min(Max_Speed)|sum(Delay)|
+----+-------+---------+--------------+----------+
|2019| 1000.0|      ICE|         270.0|      null|
|2017| 1000.0|      ICE|         280.0|       2.0|
|2020| 1000.0|       IC|         220.0|      62.0|
|2018| 1002.0|       IC|          null|      null|
|2017| 1000.0|       IC|         280.0|      11.0|
|2020| 1002.0|      ICE|         280.0|      null|
|2019| 1002.0|       IC|         220.0|      12.0|
+----+-------+---------+--------------+----------+
pyspark pandas dataframe aggregate group-by
1 Answers

I have a dirty solution, but it may not work depending on the quantity of data you have.

from functools import reduce

from pyspark.sql import functions as F


df_sp.join(
    reduce(
        lambda a, b: a.crossJoin(b),
        [F.broadcast(df_sp.select(col).distinct()) for col in cat_col_all],
    ),
    how="full",
    on=cat_col_all,
).groupby(cat_col_all).agg(pivot_agg_func).show()

+----+-------+---------+--------------+----------+
|Date|TrainID|Traintype|min(Max_Speed)|sum(Delay)|
+----+-------+---------+--------------+----------+
|2019| 1000.0|       IC|          null|      null|
|2018| 1000.0|      ICE|          null|      null|
|2018| 1002.0|       IC|          null|      null|
|2017| 1002.0|      ICE|          null|      null|
|2020| 1002.0|       IC|          null|      null|
|2017| 1000.0|      ICE|         280.0|       2.0|
|2017| 1000.0|       IC|         280.0|      11.0|
|2017| 1002.0|       IC|          null|      null|
|2019| 1000.0|      ICE|         270.0|      null|
|2020| 1002.0|      ICE|         280.0|      null|
|2018| 1002.0|      ICE|          null|      null|
|2020| 1000.0|      ICE|          null|      null|
|2019| 1002.0|      ICE|          null|      null|
|2020| 1000.0|       IC|         220.0|      62.0|
|2019| 1002.0|       IC|         220.0|      12.0|
|2018| 1000.0|       IC|          null|      null|
+----+-------+---------+--------------+----------+
user_50137520
Related