在pandas数据框中应用返回多个值的函数

99 投票
7 回答
146037 浏览
提问于 2025-04-18 06:43

我有一个数据框,里面有时间索引和三列,分别表示一个三维向量的坐标:

                         x             y             z
ts
2014-05-15 10:38         0.120117      0.987305      0.116211
2014-05-15 10:39         0.117188      0.984375      0.122070
2014-05-15 10:40         0.119141      0.987305      0.119141
2014-05-15 10:41         0.116211      0.984375      0.120117
2014-05-15 10:42         0.119141      0.983398      0.118164

我想对每一行进行一个变换,这个变换也会返回一个向量

def myfunc(a, b, c):
    do something
    return e, f, g

但是如果我这样做:

df.apply(myfunc, axis=1)

最后我得到的是一个Pandas系列,其中的元素是元组。这是因为apply函数会直接返回myfunc的结果,而不会拆开它。请问我该如何修改myfunc,以便得到一个新的数据框,里面有三列?

编辑:

下面的所有解决方案都有效。使用Series的方案可以保留列名,而使用列表的方案似乎执行得更快。

def myfunc1(args):
    e=args[0] + 2*args[1]
    f=args[1]*args[2] +1
    g=args[2] + args[0] * args[1]
    return pd.Series([e,f,g], index=['a', 'b', 'c'])

def myfunc2(args):
    e=args[0] + 2*args[1]
    f=args[1]*args[2] +1
    g=args[2] + args[0] * args[1]
    return [e,f,g]

%timeit df.apply(myfunc1 ,axis=1)

100 loops, best of 3: 4.51 ms per loop

%timeit df.apply(myfunc2 ,axis=1)

100 loops, best of 3: 2.75 ms per loop

7 个回答

20

我尝试返回一个元组(我使用了像 scipy.stats.pearsonr 这样的函数,它们返回这种结构),但结果却返回了一个一维的序列,而不是我预期的数据框。如果我手动创建一个序列,性能还更差,所以我按照官方文档的说明,使用了 result_type 来解决这个问题:

在函数内部返回一个序列就像是传递 result_type='expand'。结果的列名将会是这个序列的索引。

所以你可以这样修改你的代码:

def myfunc(a, b, c):
    # do something
    return (e, f, g)

df.apply(myfunc, axis=1, result_type='expand')
29

有些人的回答有错误,所以我把它们总结了一下。下面是最好的答案。

首先准备好数据集。这里使用的pandas版本是 1.1.5

import numpy as np
import pandas as pd
import timeit

# check pandas version
print(pd.__version__)
# 1.1.5

# prepare DataFrame
df = pd.DataFrame({
    'x': [0.120117, 0.117188, 0.119141, 0.116211, 0.119141],
    'y': [0.987305, 0.984375, 0.987305, 0.984375, 0.983398],
    'z': [0.116211, 0.122070, 0.119141, 0.120117, 0.118164]},
    index=[
        '2014-05-15 10:38',
        '2014-05-15 10:39',
        '2014-05-15 10:40',
        '2014-05-15 10:41',
        '2014-05-15 10:42'],
    columns=['x', 'y', 'z'])
df.index.name = 'ts'
#                          x         y         z
# ts                                            
# 2014-05-15 10:38  0.120117  0.987305  0.116211
# 2014-05-15 10:39  0.117188  0.984375  0.122070
# 2014-05-15 10:40  0.119141  0.987305  0.119141
# 2014-05-15 10:41  0.116211  0.984375  0.120117
# 2014-05-15 10:42  0.119141  0.983398  0.118164

解决方案 01.

在应用函数中返回 pd.Series

def myfunc1(args):
    e = args[0] + 2*args[1]
    f = args[1]*args[2] + 1
    g = args[2] + args[0] * args[1]
    return pd.Series([e, f, g])

df[['e', 'f', 'g']] = df.apply(myfunc1, axis=1)
#                          x         y         z         e         f         g
# ts
# 2014-05-15 10:38  0.120117  0.987305  0.116211  2.094727  1.114736  0.234803
# 2014-05-15 10:39  0.117188  0.984375  0.122070  2.085938  1.120163  0.237427
# 2014-05-15 10:40  0.119141  0.987305  0.119141  2.093751  1.117629  0.236770
# 2014-05-15 10:41  0.116211  0.984375  0.120117  2.084961  1.118240  0.234512
# 2014-05-15 10:42  0.119141  0.983398  0.118164  2.085937  1.116202  0.235327

t1 = timeit.timeit(
    'df.apply(myfunc1, axis=1)',
    globals=dict(df=df, myfunc1=myfunc1), number=10000)
print(round(t1, 3), 'seconds')
# 14.571 seconds

解决方案 02.

在应用时使用 result_type ='expand'

def myfunc2(args):
    e = args[0] + 2*args[1]
    f = args[1]*args[2] + 1
    g = args[2] + args[0] * args[1]
    return [e, f, g]

df[['e', 'f', 'g']] = df.apply(myfunc2, axis=1, result_type='expand')
#                          x         y         z         e         f         g
# ts                                                                          
# 2014-05-15 10:38  0.120117  0.987305  0.116211  2.094727  1.114736  0.234803
# 2014-05-15 10:39  0.117188  0.984375  0.122070  2.085938  1.120163  0.237427
# 2014-05-15 10:40  0.119141  0.987305  0.119141  2.093751  1.117629  0.236770
# 2014-05-15 10:41  0.116211  0.984375  0.120117  2.084961  1.118240  0.234512
# 2014-05-15 10:42  0.119141  0.983398  0.118164  2.085937  1.116202  0.235327

t2 = timeit.timeit(
    "df.apply(myfunc2, axis=1, result_type='expand')",
    globals=dict(df=df, myfunc2=myfunc2), number=10000)
print(round(t2, 3), 'seconds')
# 9.907 seconds

解决方案 03.

如果你想让它更快,可以使用 np.vectorize。注意,当使用 np.vectorize 时,参数不能只有一个。

def myfunc3(args0, args1, args2):
    e = args0 + 2*args1
    f = args1*args2 + 1
    g = args2 + args0 * args1
    return [e, f, g]

df[['e', 'f', 'g']] = pd.DataFrame(np.row_stack(np.vectorize(myfunc3, otypes=['O'])(df['x'], df['y'], df['z'])), index=df.index)
#                          x         y         z         e         f         g
# ts                                                                          
# 2014-05-15 10:38  0.120117  0.987305  0.116211  2.094727  1.114736  0.234803
# 2014-05-15 10:39  0.117188  0.984375  0.122070  2.085938  1.120163  0.237427
# 2014-05-15 10:40  0.119141  0.987305  0.119141  2.093751  1.117629  0.236770
# 2014-05-15 10:41  0.116211  0.984375  0.120117  2.084961  1.118240  0.234512
# 2014-05-15 10:42  0.119141  0.983398  0.118164  2.085937  1.116202  0.235327

t3 = timeit.timeit(
    "pd.DataFrame(np.row_stack(np.vectorize(myfunc3, otypes=['O'])(df['x'], df['y'], df['z'])), index=df.index)",
    globals=dict(pd=pd, np=np, df=df, myfunc3=myfunc3), number=10000)
print(round(t3, 3), 'seconds')
# 1.598 seconds
44

根据@U2EF1的精彩回答,我创建了一个实用的函数,这个函数可以把一个返回元组的指定函数应用到数据框的某个字段上,然后再把结果扩展回数据框里。

def apply_and_concat(dataframe, field, func, column_names):
    return pd.concat((
        dataframe,
        dataframe[field].apply(
            lambda cell: pd.Series(func(cell), index=column_names))), axis=1)

用法:

df = pd.DataFrame([1, 2, 3], index=['a', 'b', 'c'], columns=['A'])
print df
   A
a  1
b  2
c  3

def func(x):
    return x*x, x*x*x

print apply_and_concat(df, 'A', func, ['x^2', 'x^3'])

   A  x^2  x^3
a  1    1    1
b  2    4    8
c  3    9   27

希望这对某些人有帮助。

93

返回一个Series,这样它会把这些数据放到一个数据表(DataFrame)里。

def myfunc(a, b, c):
    do something
    return pd.Series([e, f, g])

这样做的好处是,你可以给每一列起个名字。如果你返回的是一个数据表(DataFrame),它只会为这个组插入多行数据。

13

直接返回一个列表,而不是元组。

In [81]: df
Out[81]: 
                            x         y         z
ts                                               
2014-05-15 10:38:00  0.120117  0.987305  0.116211
2014-05-15 10:39:00  0.117188  0.984375  0.122070
2014-05-15 10:40:00  0.119141  0.987305  0.119141
2014-05-15 10:41:00  0.116211  0.984375  0.120117
2014-05-15 10:42:00  0.119141  0.983398  0.118164

[5 rows x 3 columns]

In [82]: def myfunc(args):
   ....:        e=args[0] + 2*args[1]
   ....:        f=args[1]*args[2] +1
   ....:        g=args[2] + args[0] * args[1]
   ....:        return [e,f,g]
   ....: 

In [83]: df.apply(myfunc ,axis=1)
Out[83]: 
                            x         y         z
ts                                               
2014-05-15 10:38:00  2.094727  1.114736  0.234803
2014-05-15 10:39:00  2.085938  1.120163  0.237427
2014-05-15 10:40:00  2.093751  1.117629  0.236770
2014-05-15 10:41:00  2.084961  1.118240  0.234512
2014-05-15 10:42:00  2.085937  1.116202  0.235327

撰写回答