DataFrame.GroupBy.apply()

I’m continuously missing this behavior when converting my pandas scripts to koalas:

df = pd.DataFrame({"timestamp":[0.0, 0.5, 1.0, 0.0, 0.5],
              "car_id": ['A','A','A','B','B'], 
              "battery_charge": [100, 90, 80,100,90]
              })

print(df)

def calc_battery_usage(x):
    start_bat = x.sort_values('timestamp').iloc[0]['battery_charge']
    stop_bat = x.sort_values('timestamp').iloc[-1]['battery_charge']
    return start_bat - stop_bat

print(df.groupby('car_id').apply(calc_battery_usage))

That prints out:

   timestamp car_id  battery_charge
0        0.0      A             100
1        0.5      A              90
2        1.0      A              80
3        0.0      B             100
4        0.5      B              90

car_id
A    20
B    10

On koalas I just get:

PandasNotImplementedError: The method `pd.groupby.GroupBy.apply()` is not implemented yet.

I can take over and start implementing it myself, but would probably need some general guidelines on how to do it (probably using pysparks pandas_udf).

[icexelloss]

Hi @patryk-oleniuk! I think it’d great to have groupby apply in koalas.

Spark’s groupby apply and Pandas’s groupby apply are pretty similar. However, two main difference that I can think of are:

• Spark’s groupby apply need full return type specification on the udf definition
• Spark’s groupby apply doesn’t keep index, while pandas groupby apply use the grouping key as index.

As a start, I’d probably try this something like:

@pandas_wraps(full_return_schema):
def calc_battery_usage(x):
    start_bat = x.sort_values('timestamp').iloc[0]['battery_charge']
    stop_bat = x.sort_values('timestamp').iloc[-1]['battery_charge']
    return start_bat - stop_bat

kdf.groupby('car_id').apply(calc_battery_usage)

And just wrap it with Spark’s groupby apply + pandas_udf

Does that make sense?

[devarshml]

Thank you guys for the replies. I will surely check out all of the suggestions by tomorrow evening…!!