Retain DataFrame's index when using groupby apply to generate Series

Question

When calling a function using groupby + apply, I want to go from a DataFrame to a Series groupby object, apply a function to each group that takes a Series as input and returns a Series as output, and then assign the output from the groupby + apply call as a field in the DataFrame. The default behavior is to have the output from groupby + apply indexed by the grouping fields, which prevents me from assigning it back to the DataFrame cleanly. I'd prefer to have the function I call with apply take a Series as input and return a Series as output; I think it's a bit cleaner than DataFrame to DataFrame. (This isn't the best way of getting to the result for this example; the real application is pretty different.)

import pandas as pd
df = pd.DataFrame({
 'A': [999, 999, 111, 111],
 'B': [1, 2, 3, 4],
 'C': [1, 3, 1, 3]
})
def less_than_two(series):
  # Intended for series of length 1 in this case
  # But not intended for many-to-one generally
  return series.iloc[0] < 2
output = df.groupby(['A', 'B'])['C'].apply(less_than_two)

I want the index on output to be the same as df, otherwise I cant assign to df (cleanly):

df['Less_Than_Two'] = output

Something like output.index = df.index seems too ugly, and using the group_keys argument doesn't seem to work:

output = df.groupby(['A', 'B'], group_keys = False)['C'].apply(less_than_two)
df['Less_Than_Two'] = output

Answer 1

transform returns the results with the original index, just as you've asked for. It will broadcast the same result across all elements of a group. Caveat, beware that the dtype may be inferred to be something else. You may have to cast it yourself.

In this case, in order to add another column, I'd use assign

df.assign(
    Less_Than_Two=df.groupby(['A', 'B'])['C'].transform(less_than_two).astype(bool))

     A  B  C Less_Than_Two
0  999  1  1          True
1  999  2  3         False
2  111  3  1          True
3  111  4  3         False