Pandas `.assign`用于单个函数中的多列

Question

我有一个Pandas数据帧，它包含一列val，还有一个函数func，它接受一个值并吐出一个固定长度的列表(假设为4(。我还有一个由4个字符串组成的列表cols。我想将func应用于每个单元格，并添加4个新列，根据我的列表进行标记。

似乎起作用的是这样的东西：

import pandas as pd
df = pd.DataFrame({'val': [1, 2, 4, 18, 9, 1]})
cols = ["X", "Y", "Z", "hello"]
func = lambda x: [2**x, str(x), x+1, "world"]
df[cols] = df['val'].apply(lambda val: pd.Series(func(val)))

由于我看到每个人都建议不要使用apply，所以我想尝试使用assign。我尝试将func的输出分配给一个临时列tmp，然后逐个提取各个值，如下所示：

import pandas as pd
df = pd.DataFrame({'val': [1, 2, 4, 18, 9, 1]})
cols = ["X", "Y", "Z", "hello"]
func = lambda x: [2**x, str(x), x+1, "world"]
kwargs = {name: (lambda x: x.tmp[idx]) for idx, name in enumerate(cols)}
df[cols] = df.assign(tmp=lambda x: pd.Series(func(x.val)), **kwargs)

但这带来了一些错误，我不知道如何解释ValueError: Columns must be same length as key。请注意，.assign[1]的文档中说允许这种自引用，请参阅最后一个示例。

[1]https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.assign.html

编辑：为了澄清起见，在我的实际应用程序中，对func的调用非常昂贵，我不想每行调用四次。它也不容易像我的例子那样分成四个子组件。

Answer 1

我认为关键是您需要使用lambda函数来应用于输入列，指定result_type='expand'选项和轴。然后，您可以定义任意数量的输出列，将结果分配给这些列。

以下是我创建的一个简单示例，其中包含一个输入列和两个输出列，大意如下：https://gist.github.com/84adam/29ff5fd1286a30d904540bf78e37f883

语法示例：

df[['output1','output2']] = df.apply(lambda x: func(x['input1']), axis=1, result_type='expand')

我认为这应该奏效。

Answer 2

您必须进行一些测试，看看原始函数是否比下面的assign方法更具性能。

df = pd.DataFrame({'val': [1, 2, 4, 18, 9, 1]})
df = df.assign(X=2**df['val'],
Z=df['val']+1,
Y=df['val'].astype('str'),
world='hello')
val       X   Z   Y  world
0    1       2   2   1  hello
1    2       4   3   2  hello
2    4      16   5   4  hello
3   18  262144  19  18  hello
4    9     512  10   9  hello
5    1       2   2   1  hello

Answer 3

import pandas as pd
df = pd.DataFrame({'val': [1, 2, 4, 18, 9, 1]})
cols = ["X", "Y", "Z", "hello"]
func = lambda x: [2**x, str(x), x+1, "world"]
df[cols] = df['val'].apply(lambda val: pd.Series(func(val)), result_type='expand')

我认为expand的加法将给出正确的结果。

更新：

使用assign:

import pandas as pd
df = pd.DataFrame({'val': [1, 2, 4, 18, 9, 1]})
cols = ["X", "Y", "Z", "hello"]
func1 = lambda x: 2**x
func2 = lambda x: str(x)
func3 = lambda x: x+1
func4 = lambda x: "world"
df.assign(X=lambda x: pd.Series(func1(x['val'])),
Y=df['val'].astype(str),
Z=lambda x: pd.Series(func3(x['val'])),
hello='world',
)

输出：

val X   Y   Z   hello
0   1   2   1   2   world
1   2   4   2   3   world
2   4   16  4   5   world
3   18  262144  18  19  world
4   9   512 9   10  world
5   1   2   1   2   world