Create an if-else condition column in dask dataframe

Question

I need to create a column which is based on some condition on dask dataframe. In pandas it is fairly straightforward:

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'] ]

While in dask I have to do same thing like below:

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

Questions:

Is there a better/more straightforward way to achieve it?
I can't modify the first dataframe ddf, i need to create ddf1 to se the changes is dask dataframe Immutable object?

MRocklin · Accepted Answer · 2016-07-27T12:54:05.133

7

Answers:

What you're doing now is almost ok. You don't need to call compute until you're ready for your final answer.

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

For some cases dd.Series.where might be a good fit

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

As of version 0.10.2 you can now insert columns directly into dask.dataframes
```
ddf['col'] = ddf.shop_week.apply(f)
```

edited Jul 27 '16 at 12:54

answered Jul 27 '16 at 12:48

MRocklin

55,641
23
163
235

The dask dataframes are mutable by default or they are mutable post release of version 0.10.2? – Puneet Tripathi Jul 28 '16 at 04:53
1

In the most recent version of Dask, version 0.10.2, dask.dataframes support column assignment – MRocklin Jul 28 '16 at 11:57
Thanks for the answer @MRocklin – Puneet Tripathi Jul 28 '16 at 19:20

Ohumeronen · Answer 2 · 2016-07-27T09:36:30.053

1

You could just use:

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

And then:

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

Unfortunately I don't have an answer to the second question or I don't understand it...

edited Jul 27 '16 at 09:36

answered Jul 27 '16 at 09:29

Ohumeronen

1,769
2
14
28

score 0 · Answer 3 · answered Dec 19 '20 at 16:58

A better approach might be pull out the column as a dask array and then perform some nested where operations before adding it back to the dataframe:

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

Create an if-else condition column in dask dataframe

3 Answers3

Linked