Pandas GroupBy.apply produces MultiIndex and RangeIndex, depending on the usage of head

Viewed 79

I encountered some behavior that I don't understand. When you take a look at the code below: using apply_func_rangeindex() in apply produces a result with a RangeIndex, using apply_func_multiindex() in apply results in a multiindex. But I don't understand why. Calling head on a dataframe should return a dataframe with the same specifications, so there shouldn't be a difference. And I expect in each case a multiindex. Does anyone know why they produce different results? (pandas version 1.3.3/Windows)

import pandas as pd
print(f"Using pandas version: {pd.__version__}")

df = pd.DataFrame(
    {
        "key1":["k1", "k2", "k3"],
        "key2":["a", "b", "c"],
        "value1":["v1", "v2", "v3"]
    }
)
print(df)

def apply_func_rangeindex(rows):
    """Simply returns the df results in a rangeindex. Even if len(rows) == 1"""
    return rows[["value1"]]

def apply_func_multiindex(rows):
    """Returns the first row results in a multiindex"""
    return rows.head(1)[["value1"]]

# Keys to group by
keys = ["key1", "key2"]
grouped = df.groupby(by=keys, as_index=True, dropna=False)


# The results
print("\n\n======= RangeIndex =========")
result_rangeindex = grouped.apply(apply_func_rangeindex)
print(f"Shape  : {result_rangeindex.shape}")
print(f"Index  : {result_rangeindex.index}")
print(f"Columns: {result_rangeindex.columns}")


print("\n\n======= MultiIndex =========")
result_multiindex = grouped.apply(apply_func_multiindex)
print(f"Shape  : {result_multiindex.shape}")
print(f"Index  : {result_multiindex.index}")
print(f"Columns: {result_multiindex.columns}")

This is the output:

Using pandas version: 1.3.3
  key1 key2 value1
0   k1    a     v1
1   k2    b     v2
2   k3    c     v3


======= RangeIndex =========
Shape  : (3, 1)
Index  : RangeIndex(start=0, stop=3, step=1)
Columns: Index(['value1'], dtype='object')


======= MultiIndex =========
Shape  : (3, 1)
Index  : MultiIndex([('k1', 'a', 0),
            ('k2', 'b', 1),
            ('k3', 'c', 2)],
           names=['key1', 'key2', None])
Columns: Index(['value1'], dtype='object')
0 Answers
Related