从pandas apply() 返回多个列

Question

我有一个叫做 df_test 的 pandas 数据框，它里面有一列叫 'size'，表示大小，单位是字节。我用下面的代码计算了 KB、MB 和 GB：

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

[2 rows x 5 columns]

我在超过 120,000 行的数据上运行这个，结果每一列大约花费 2.97 秒，乘以 3 列就是大约 9 秒，根据 %timeit 的结果来看。

有没有办法可以让这个过程更快一点？比如说，能不能一次性返回三列，而不是每次只返回一列，这样就可以减少运行的次数，把结果直接插回原来的数据框里？

我找到的其他问题都是想从多个值中得到一个值，而我想要的是从一个值中得到多个列。

Answer 1

这是一种更易读的方式。这个代码会添加三个新列，并为这些列赋值，返回的结果是一个系列，而在使用apply函数时不需要参数。

def sizes(s):

    val_kb = locale.format("%.1f", s['size'] / 1024.0, grouping=True) + ' KB'
    val_mb = locale.format("%.1f", s['size'] / 1024.0 ** 2, grouping=True) + ' MB'
    val_gb = locale.format("%.1f", s['size'] / 1024.0 ** 3, grouping=True) + ' GB'
    return pd.Series([val_kb,val_mb,val_gb],index=['size_kb','size_mb','size_gb'])

df[['size_kb','size_mb','size_gb']] = df.apply(lambda x: sizes(x) , axis=1)

这是一个来自于的通用示例：https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.apply.html

df.apply(lambda x: pd.Series([1, 2], index=['foo', 'bar']), axis=1)

#foo  bar
#0    1    2
#1    1    2
#2    1    2

Answer 2

真的很棒的回答！谢谢Jesse和jaumebonet！我想分享一些关于以下内容的观察：

• zip(* ...
• ... result_type="expand")

虽然使用expand看起来更优雅（有点像“pandas风格”），但是**zip至少快了2倍。在下面这个简单的例子中，我发现它快了4倍。

import pandas as pd

dat = [ [i, 10*i] for i in range(1000)]

df = pd.DataFrame(dat, columns = ["a","b"])

def add_and_sub(row):
    add = row["a"] + row["b"]
    sub = row["a"] - row["b"]
    return add, sub

df[["add", "sub"]] = df.apply(add_and_sub, axis=1, result_type="expand")
# versus
df["add"], df["sub"] = zip(*df.apply(add_and_sub, axis=1))

Answer 3

现在有些回复的方法很好用，但我想提供一个可能更“适合pandas”的选项。这在我使用的pandas 0.23版本中有效（不确定在之前的版本是否也能用）：

import pandas as pd

df_test = pd.DataFrame([
  {'dir': '/Users/uname1', 'size': 994933},
  {'dir': '/Users/uname2', 'size': 109338711},
])

def sizes(s):
  a = locale.format_string("%.1f", s['size'] / 1024.0, grouping=True) + ' KB'
  b = locale.format_string("%.1f", s['size'] / 1024.0 ** 2, grouping=True) + ' MB'
  c = locale.format_string("%.1f", s['size'] / 1024.0 ** 3, grouping=True) + ' GB'
  return a, b, c

df_test[['size_kb', 'size_mb', 'size_gb']] = df_test.apply(sizes, axis=1, result_type="expand")

注意，这个技巧在于apply函数的result_type参数，它会把结果扩展成一个可以直接赋值给新列或旧列的DataFrame。

Answer 4

使用apply和zip的方法比用Series的方式快3倍。

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))

测试结果是：

Separate df.apply(): 

    100 loops, best of 3: 1.43 ms per loop

Return Series: 

    100 loops, best of 3: 2.61 ms per loop

Return tuple:

    1000 loops, best of 3: 819 µs per loop

Answer 5

你可以从应用的函数中返回一个包含新数据的系列，这样就不需要重复处理三次了。把 axis=1 传给 apply 函数，可以把 sizes 这个函数应用到数据表的每一行，然后返回一个系列，这个系列可以用来添加到一个新的数据表中。这个系列 s 包含了新值，还有原始数据。

def sizes(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.append(rows_list)
df_test = df_test.apply(sizes, axis=1)