DataFrame使用Groupby, Apply后返回不同数据类型的问题

Groupby是使用Pandas进行数据处理时常用的函数，功能强大。特别是在group后还能跟上apply函数对数据自定义的进行转换，实在是非常方便。但有时候我们会发现，apply可能返回的数据类型有两种，为什么会这样呢？

注意：我这里用的是apply的返回值，而不是groupby的返回值。因为groupby的返回值类型其实是一个确定的：DataFrameGroupBy 对象。在 DataFrameGroupBy 对象上继续调用apply时才可能会返回不同的数据类型。

认识GroupBy及其返回值

我们先来看代码，假设有一个 DataFrame, 包含 key, v1, v2 三列：

data = {
    'key': ['GD', 'GD', 'GX', 'GD', 'GX'],
    'v1': [1, 2, 3, 4, 5],
    'v2': [6, 7, 8, 9, 10]
    }

f = DataFrame(data)

如果从控制台输出这个 DataFrame，我们能看到如下的二维表形式：

  key  v1  v2
0  GD   1   6
1  GD   2   7
2  GX   3   8
3  GD   4   9
4  GX   5  10

现在，我们以”key”为分组条件进行分组，可以得到一个 DataFrameGroupBy 对象, 如果我们再接着取该分组对象的属性 groups ,可以得到包含的分组详情

g = f.groupby(['key'])

print (type(g))

print (g.groups)

输出如下:

1 2	<class 'pandas.core.groupby.generic.DataFrameGroupBy'> {'GD': Int64Index([0, 1, 3], dtype='int64'), 'GX': Int64Index([2, 4], dtype='int64')}

可以看到，DataFrameGroupBy 这个对象其实时包含了分类的条件和每个条件对应数据项在 DataFrame 中的索引。

我们来看一下运用不同的 lambda 表达式时，不同的返回值（类型）

代码1:

1
2
3

r = g.apply(lambda x: x.v1 + 1)
print (type(r))
print (r)

输出1:

<class 'pandas.core.series.Series'>
key
GD   0    2
     1    3
     3    5
GX   2    4
     4    6
Name: v1, dtype: int64

可以看到，现在返回的是 Series

代码2:

1
2
3

r1 = g.apply(lambda x: x.v1 + x.v2)
print (type(r1))
print (r1)

输出2:

<class 'pandas.core.series.Series'>
key
GD   0     7
     1     9
     3    13
GX   2    11
     4    15
dtype: int64

可以看到，返回仍然还是 Series

代码3:

1
2
3

r2 = g.apply(lambda x: x[1:3])
print (type(r2))
print (r2)

输出3:

<class 'pandas.core.frame.DataFrame'>
      key  v1  v2
key
GD  1  GD   2   7
    3  GD   4   9
GX  4  GX   5  10

可以看到，这次返回时 DataFrame 了。

如果我们比较一下上面的三段代码就会发现，apply函数的返回值类型其实时依赖与 lambda 表达式返回值的。当 lambda 返回的是一列值是，会包装为 Series, 当返回值是多列时，会包装为 DataFrame。

直接用 Series 提供的方法 to_frame() 就可以对结果进行转换，代码如下:

1
2
3

r4 = g.apply(lambda x: x.v1 + 1).to_frame()
print (type(r4))
print (r4)

参考输出：

<class 'pandas.core.frame.DataFrame'>
       v1
key
GD  0   2
    1   3
    3   5
GX  2   4
    4   6

比较上面的”输出1”的结果，可以看到数据是一致的，只是类型由 Series 转换为 DataFrame。