python - 在 dask 数据框中创建 if-else 条件列

Question

我需要创建一个基于 dask 数据框的某些条件的列。在 pandas 中，它相当简单：

ddf['TEST_VAR'] = ['THIS' if x == 200607 else  
              'NOT THIS' if x == 200608 else 
              'THAT' if x == 200609 else 'NONE'  
              for x in ddf['shop_week'] ]

在 dask 中，我必须做同样的事情，如下所示：

def f(x):
    if x == 200607:
         y= 'THIS'
    elif x == 200608 :
         y= 'THAT'
    else :
         y= 1 
    return y

ddf1 = ddf.assign(col1 = list(ddf.shop_week.apply(f).compute()))
ddf1.compute()

问题：

1. 有没有更好/更直接的方法来实现它？
2. 我无法修改第一个数据帧 ddf，我需要创建 ddf1 以确保更改是 dask 数据帧不可变对象？

Answer 1

答案：

1. 你现在做的几乎没问题。compute在您准备好获得最终答复之前，您无需致电。

   # ddf1 = ddf.assign(col1 = list(ddf.shop_week.apply(f).compute()))
   ddf1 = ddf.assign(col1 = ddf.shop_week.apply(f))
   

   在某些情况下dd.Series.where可能很合适

   ddf1 = ddf.assign(col1 = ddf.shop_week.where(cond=ddf.balance > 0, other=0))
   

2. 从 0.10.2 版开始，您现在可以将列直接插入 dask.dataframes

   ddf['col'] = ddf.shop_week.apply(f)
   

Answer 2

你可以只使用：

f = lambda x: 'THIS' if x == 200607 else 'NOT THIS' if x == 200608 else 'THAT' if x == 200609 else 'NONE'

接着：

ddf1 = ddf.assign(col1 = list(ddf.shop_week.apply(f).compute()))

不幸的是，我没有第二个问题的答案，或者我不明白......

Answer 3

更好的方法可能是将列作为 dask 数组拉出，然后where在将其添加回数据框之前执行一些嵌套操作：

import dask.array as da

x = ddf['shop_week'].to_dask_array()

df['TEST_VAR'] = \
    da.where(x == 200607, 'THIS',
    da.where(x == 200608, 'NOT THIS',
    da.where(x == 200609, 'THAT', 'NONE')))

df['TEST_VAR'].compute()