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

问题：

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

Answer 1

您可以使用：

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 2

答案：

你现在做的差不多就ok了。在您准备好接受最终答复之前，您无需致电 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))

从版本 0.10.2 开始，您现在可以将列直接插入 dask.dataframes
```
ddf['col'] = ddf.shop_week.apply(f)
```

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

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

Create an if-else condition column in dask dataframe

python

pandas

dask