The `map_partitions` kwargs and args support delayed objects, but `reductions` doesn't
- 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
Assessment
This issue has not been assessed yet.