使用应用功能在 Dask 中添加两列

Question

我有一个 Dask 函数，可以将一列添加到现有的 Dask 数据框，这很好用：

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98],
    'weight': [150, 126, 133, 164, 203]
})

df_dask = dd.from_pandas(df, npartitions=2) 


s = """
obj.weight + 100
"""

df_dask['new_weight'] = df_dask.apply(lambda obj: eval(s), meta=dict, axis=1)

现在，我想添加两列而不是一列：

s = """
obj.weight + 100, obj.weight + 200
"""

df_dask['new_weight','new_weight2'] = df_dask.apply(lambda obj: eval(s), meta=dict, axis=1)

但是我明白了

NotImplementedError: Item assignment with <class 'tuple'> not supported

这是否意味着不支持或我做错了什么？如果不支持，是否有解决方法？我需要的是 return 一个浮动列表。

Answer 1

首先，确保基于 pandas 的解决方案提供预期结果将很有用：

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98],
    'weight': [150, 126, 133, 164, 203]
})

s = """
obj.weight + 100, obj.weight + 200
"""

df['new_weight'], df['new_weight2'] = zip(*df.apply(lambda obj: eval(s), axis=1))

print(df)
#    height  weight  new_weight  new_weight2
# 0    6.21     150       250.0        350.0
# 1    5.12     126       226.0        326.0
# 2    5.85     133       233.0        333.0
# 3    5.78     164       264.0        364.0
# 4    5.98     203       303.0        403.0

现在，dask 解决方案。由于每个分区都是一个 pandas 数据帧，最简单的解决方案（对于基于行的转换）是将 pandas 代码包装到一个函数中并将其插入 map_partitions:

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98],
    'weight': [150, 126, 133, 164, 203]
})

s = """
obj.weight + 100, obj.weight + 200
"""

df_dask = dd.from_pandas(df, npartitions=2) 

def new_cols(df):
    df = df.copy()
    df['new_weight'], df['new_weight2'] = zip(*df.apply(lambda obj: eval(s), axis=1))
    return df

df_dask = df_dask.map_partitions(new_cols)

print(df_dask.compute())
#    height  weight  new_weight  new_weight2
# 0    6.21     150       250.0        350.0
# 1    5.12     126       226.0        326.0
# 2    5.85     133       233.0        333.0
# 3    5.78     164       264.0        364.0
# 4    5.98     203       303.0        403.0

根据您的用例的具体情况，上述模式可能具有提高效率的空间。

使用应用功能在 Dask 中添加两列

Adding two columns in Dask with apply function

python

pandas

dask