add group property in pandas within a chain (analogous to dplyr group_by - mutate in R)

Question

I would like to add some group properties as a new column to a pandas dataframe but without beaking the chain. I know this is possible in R using dplyr but I cannot get it to work in pandas.

The dplyr code would be (for adding max of column B per group in column A):

df %>%
   group_by(A) %>%
   mutate(max = max(B)) %>%
   ungroup() %>%
   ... more operations

The only way I can get it to work in pandas is:

df['max'] = df.groupby('A')['B'].transform('max')

but this requires a seperate line to assign the new column while I would like to do it inside a chain. Any help would be appreciated.

Answer 1

df.assign(max=df.groupby('A')['B'].transform('max'))....more operations

Answer 2

Now you are able to do it smoothly with datar

from datar import f
from datar.base import max
from datar.dplyr import group_by, mutate, ungroup

df >> \
   group_by(f.A) >> \
   mutate(max = max(f.B)) >> \
   ungroup() # >> 
   # ... more operations

I am the author of the package. Feel free to submit issues if you have any questions.