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

Question

我需要创建一个基于dask数据帧的某些条件的列.在熊猫中它是相当简单的:

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 dataframe的Immutable对象吗？

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)