Transform pandas DataFrames groups - completely, not just a Series

user_4533188 0
Viewed 108

I want to transform each group in a pandas' DataFrame. By group I mean not a single column of the DataFrame, but the entire group. Here is an example of what I mean:

df = pd.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar', 'foo', 'bar'],
                   'B' : ['one', 'one', 'two', 'two', 'one', 'two'],
                   'C' : [    1,     5,     5,     2,     6,     5],
                   'D' : [  2.0,    5.,    8.,    1.,     2.,   9.]})

def transformation_function(group: pd.DataFrame) -> pd.DataFrame:
    group = group.copy()
    if all(group.B == 'one'):
        group.D[group.C>2] = group.D[group.C>2] + 1
    else:
        group.A = 'new'
    return group

df.groupby('B').transform(transformation_function)

where I would expect

     pd.DataFrame({'A' : ['foo', 'bar', 'new', 'new', 'foo', 'new'],
                   'B' : ['one', 'one', 'two', 'two', 'one', 'two'],
                   'C' : [    1,     5,     5,     2,     5,     5],
                   'D' : [  2.0,    6.,    8.,    1.,    3.,    9.]})

as a result. Now, I get the

AttributeError: 'Series' object has no attribute 'B'

which does not make sense to me, because the documentation explicitly states

Call function producing a like-indexed DataFrame on each group and return a DataFrame having the same indexes as the original object filled with the transformed values

I am aware the all examples work Series-based, like df.groupby('B')['a_column_name'].transform(change_fct), but something like that is not possible if one needs all columns for the transformation function.

So, how do I get what I expect, using pandas' method chaining?

pandas dataframe method-chaining python
1 Answers

In your solution transform function working with each column separately, so not possible select columns by names. Need GroupBy.apply:

df = df.groupby('B').apply(transformation_function)
    
print (df)
         A    B  C    D
B                      
one 0  foo  one  1  2.0
    1  bar  one  5  6.0
    4  foo  one  6  3.0
two 2  new  two  5  8.0
    3  new  two  2  1.0
    5  new  two  5  9.0

How are processing groups is possible see e.g. by print:

df.groupby('B').transform(lambda x: print (x))
0    foo
1    bar
4    foo
Name: A, dtype: object
0    1
1    5
4    6
Name: C, dtype: int64
0    2.0
1    5.0
4    2.0
Name: D, dtype: float64
     A  C    D
0  foo  1  2.0
1  bar  5  5.0
4  foo  6  2.0
2    foo
3    bar
5    bar
Name: A, dtype: object
2    5
3    2
5    5
Name: C, dtype: int64
2    8.0
3    1.0
5    9.0
Name: D, dtype: float64
     A  C    D
2  foo  5  8.0
3  bar  2  1.0
5  bar  5  9.0

df.groupby('B').apply(lambda x: print (x))

     A    B  C    D
0  foo  one  1  2.0
1  bar  one  5  5.0
4  foo  one  6  2.0
     A    B  C    D
2  foo  two  5  8.0
3  bar  two  2  1.0
5  bar  two  5  9.0
user_29010020
Related