dask / dask/distributed

The `map_partitions` kwargs and args support delayed objects, but `reductions` doesn't

Open
#3,376 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

When using dataframe `map_partitions`, passing in a delayed object in the kwargs results in the mapping `func` getting the computed form of the argument.

However when using `reductions`, passing in a delayed object in the kwargs results in the chunk `func` not getting the computed form of the argument. This API is inconsistent, and for alot of interesting usecases, it needs to get the computed form of the argument. I imagine this is similar for the aggregate kwargs and also merge kwargs.

Here's the proof:

```py
import dask
import pandas as pd
import dask.dataframe as dd
import fs

df = pd.DataFrame({
'a': [1,2,3,4],
'b': [5,6,7,8]
})

df = dd.from_pandas(df, npartitions=1)

def doSomething(partition, images_fs=None, masks_fs=None):
print('IN MAP PARTITIONS')
print(images_fs)
print(type(images_fs))
print(masks_fs)
print(type(masks_fs))
return partition

def openstuff(p1, p2):
print('OPENING THE FS')
return (fs.open_fs(p1), fs.open_fs(p2))

delayed_open_fs = dask.delayed(openstuff, nout=2)

(delayed_tmp1, delayed_tmp2) = delayed_open_fs('osfs://tmp', 'osfs://tmp2')

print(delayed_tmp1)
print(delayed_tmp2)

result_df_meta = dd.utils.make_meta(df)

result_df = df.map_partitions(
doSomething,
meta=result_df_meta,
images_fs=delayed_tmp1,
masks_fs=delayed_tmp2
)

def doSomething2(partition, images_fs, masks_fs):
print('IN REDUCTION')
print(images_fs)
print(type(images_fs))
print(masks_fs)
print(type(masks_fs))
return partition

result_df_2 = result_df.reduction(
doSomething2,
meta=result_df_meta,
images_fs=delayed_tmp1,
masks_fs=delayed_tmp2
)
```

Running it with `python -i`.

```
>>> result_df_2.compute()
OPENING THE FS
IN MAP PARTITIONS

IN REDUCTION
Delayed('getitem-4bbb47a02073b955af52bd2adf118d05')

Delayed('getitem-3a8d55bf52ce0c22f82d9c93b60869cc')

IN REDUCTION
Delayed('getitem-4bbb47a02073b955af52bd2adf118d05')

Delayed('getitem-3a8d55bf52ce0c22f82d9c93b60869cc')

a b
0 1 5
1 2 6
2 3 7
3 4 8
```

Notice how the `images_fs` and `masks_fs` in the mapping `func` is the computed `` whereas in the chunk function it's still a delayed object.

For now I have to run `compute` on those parameters, but this makes the whole code quite ugly as this is a hidden assumption on the parameters.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.