Pandas 一次 Apply 返回多列结果的方法

pandas.DataFrame 中的 apply 方法可以用于生成新的数据。

基础方法(1对1,N对1)

1列输入-1列输出 最简单:

1
2
3
4
5
def foo(x):
y = x + 1
return y

df['new_column'] = df['column_1'].apply(foo)

N列输入-1列输出:

1
2
3
4
5
6
7
8
def foo_multiple_input(x):
if x['column_1'] > 0:
y = x['column_2'] + x['column_3']
else:
y = x['column_4'] * x['column_5']
return y

df['new_column'] = df.apply(foo_multiple_input,axis=1)

复杂场景(1对N)

1列输入-N列输出,N对N的处理方法类似:

这里给出 StackOverflow上相关问答

问题

有人提出这么一个问题:

想把size一列转换为KB/MB/GB 3列。
他自己的实现方式为,执行3次1列输入-1列输出的操作。
想知道有没有更好的实现方式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
df_test = pd.DataFrame([
{'dir': '/Users/uname1', 'size': 994933},
{'dir': '/Users/uname2', 'size': 109338711},
])

df_test['size_kb'] = df_test['size'].astype(int).apply(lambda x: locale.format("%.1f", x / 1024.0, grouping=True) + ' KB')
df_test['size_mb'] = df_test['size'].astype(int).apply(lambda x: locale.format("%.1f", x / 1024.0 ** 2, grouping=True) + ' MB')
df_test['size_gb'] = df_test['size'].astype(int).apply(lambda x: locale.format("%.1f", x / 1024.0 ** 3, grouping=True) + ' GB')

df_test


dir size size_kb size_mb size_gb
0 /Users/uname1 994933 971.6 KB 0.9 MB 0.0 GB
1 /Users/uname2 109338711 106,776.1 KB 104.3 MB 0.1 GB

该原始方案,我们暂且记做one-by-one

解决方案

有两个大神给出了两种方案:

方案1:返回Series

1
2
3
4
5
6
7
def sizes_series(s):
s['size_kb'] = locale.format("%.1f", s['size'] / 1024.0, grouping=True) + ' KB'
s['size_mb'] = locale.format("%.1f", s['size'] / 1024.0 ** 2, grouping=True) + ' MB'
s['size_gb'] = locale.format("%.1f", s['size'] / 1024.0 ** 3, grouping=True) + ' GB'
return s

df_test = df_test.apply(sizes_serias, axis=1)

该方案的原理为:

  • 对每一行记录返回1个 Series,这个Series包含原始的所有列
  • 汇总的结果还是整个Dataframe且包含新增列

方案2:使用tuple

1
2
3
4
5
6
def sizes(s):    
return locale.format("%.1f", s / 1024.0, grouping=True) + ' KB', \
locale.format("%.1f", s / 1024.0 ** 2, grouping=True) + ' MB', \
locale.format("%.1f", s / 1024.0 ** 3, grouping=True) + ' GB'

df_test['size_kb'], df_test['size_mb'], df_test['size_gb'] = zip(*df_test['size'].apply(sizes))

这里解释一下tuple方案的原理:

  • 先使用apply对每一返回一个tuple,apply的结果为 m*3 的 array (m为行数)
  • 再使用zip(*array)将数据转为 3 个 list 组成的 tuple
  • 用3个DataFrame[column]分别接收各个结果

补充一下zip的用法说明

1
2
3
4
5
6
7
8
9
a = [1,2,3]
b = ['a','b','c']

ziped_data = list(zip(a,b))
print(ziped_data) # [(1, 'a'), (2, 'b'), (3, 'c')]

# unzip
unziped_data = list(zip(*ziped_data))
print(unziped_data) # [(1, 2, 3), ('a', 'b', 'c')]

效率对比

从运行时间看:

  • tuple 方案运行效率最高,但可读性相对较差
  • 原始的one-by-one方案,效率可接受。特别是在数据量大的情况下,与tuple差异不明显。
  • series方案,运行效率明显较低。不推荐。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 2 行数据
## one-by-one apply
2.32 ms ± 121 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
## series
16.9 ms ± 75.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
## zip
1.29 ms ± 21 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# 10000 行数据
## one-by-one apply
580 ms ± 13.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
## series
4.32 s ± 29.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
## zip
565 ms ± 7.29 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

结论

  • 如果新生成列数不是很多,且逐列生成过程中没有重复计算的情况下,建议使用可读性更高的one-by-one方案。
  • 其他场景,推荐使用 tuple + (un)zip 的方案,效率更高。