Add template xarray object kwarg to map_blocks #3816

dcherian · 2020-03-02T07:20:19Z

Closes map_blocks output inference problems #3575
Tests added
Passes isort -rc . && black . && mypy . && flake8
Fully documented, including whats-new.rst for all changes and api.rst for new API

This PR adds a template kwarg to map_blocks so that we can do more complicated things where automated inference of the template fails. template is expected to be an xarray object that looks like the result of the map_blocks computation.

@jhamman To me, this seems a lot easier than defining a dict based schema. With dask variables, the memory cost shouldn't be high. It's easy to use standard xarray operations to make something that looks like the result dataset. Here's a notebook prototyping a to_schema/from_schema approach: https://gist.github.com/dcherian/130ba22d0fbadb616837deb914eaa67e#file-map_blocks_for_metsim_test-ipynb

ds = xr.tutorial.load_dataset("air_temperature").chunk({"lat": 5})
template = ds.isel(lat=slice(1, None, 5))
result = xr.map_blocks(lambda x: x.isel(lat=[1]), ds) # fails
result = xr.map_blocks(lambda x: x.isel(lat=[1]), ds, template=template) # yay
xr.testing.assert_equal(result, template)  # passes

Todo:

add tests
what if the function returns a variable not in the template?
should we do anything with template.attrs?

This accounts for dimension sizes being changed by the applied function.

…ed object

xarray/core/parallel.py

This reverts commit 045ae2b.

pep8speaks · 2020-03-06T11:59:59Z

Hello @dcherian! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2020-05-05 18:13:17 UTC

xarray/tests/test_dask.py

jhamman

Thanks @dcherian! This is looking great. I ran this on my test problem and it worked right out of the box. The main TODOs I see are:

port to DataArray and Dataset methods
add a few more tests (mostly around checking expected outputs with and without template)
add an example to the docs
get some additional eyes on this (probably @shoyer and/or @crusaderky)

xarray/core/parallel.py

xarray/tests/test_dask.py

doc/dask.rst

* upstream/master: (54 commits) Limit repr of arrays containing long strings (pydata#3900) expose a few zarr backend functions as semi-public api (pydata#3897) Use drawstyle instead of linestyle in plot.step. (pydata#3274) Implementation of polyfit and polyval (pydata#3733) misplaced quote in whatsnew (pydata#3889) Rename ordered_dict_intersection -> compat_dict_intersection (pydata#3887) Control attrs of result in `merge()`, `concat()`, `combine_by_coords()` and `combine_nested()` (pydata#3877) xfail test_uamiv_format_write (pydata#3885) Use `fixes` in PR template (pydata#3886) Tweaks to "how_to_release" (pydata#3882) whatsnew section for 0.16.0 Release v0.15.1 whatsnew for 0.15.1 (pydata#3879) update panel documentation (pydata#3880) reword the whats-new entry for unit support (pydata#3878) Raise error when assigning to IndexVariable.values & IndexVariable.data (pydata#3862) Re-enable tests xfailed in pydata#3808 and fix new CFTimeIndex failures due to upstream changes (pydata#3874) add spacing in the versions section of the issue report (pydata#3876) map_blocks: allow user function to add new unindexed dimension. (pydata#3817) Delete associated indexes when deleting coordinate variables. (pydata#3840) ...

…ap-blocks-schema * 'map-blocks-schema' of github.com:dcherian/xarray: Update doc/dask.rst

doc/dask.rst

xarray/core/parallel.py

xarray/tests/test_dask.py

xarray/core/parallel.py

shoyer · 2020-04-30T04:10:16Z

xarray/core/dataarray.py

+            the function will be first run on mocked-up data, that looks like 'obj' but
+            has sizes 0, to determine properties of the returned object such as dtype,
+            variable names, new dimensions and new indexes (if any).
+            'template' must be provided if the function changes the size of existing dimensions.


I assume attrs are also copied from the template and ignored from the computed chunks?

Yes. But I don't know that this is right.

With auto-inferred templates, attrs are set by the user function. I think it would be nice to preserve that behaviour. What do you think?

Thinking about this some more... there's no way to update attrs after computation is there?

Correct, attrs needs to be set when the resulting Dataset is created, before we do the the computation.

shoyer · 2020-04-30T04:11:22Z

Looks really nice! I like this design, just a few minor concerns.

* upstream/master: (39 commits) Pint support for DataArray (pydata#3643) Apply blackdoc to the documentation (pydata#4012) ensure Variable._repr_html_ works (pydata#3973) Fix handling of abbreviated units like msec (pydata#3998) full_like: error on non-scalar fill_value (pydata#3979) Fix some code quality and bug-risk issues (pydata#3999) DOC: add pandas.DataFrame.to_xarray (pydata#3994) Better chunking error messages for zarr backend (pydata#3983) Silence sphinx warnings (pydata#3990) Fix distributed tests on upstream-dev (pydata#3989) Add multi-dimensional extrapolation example and mention different behavior of kwargs in interp (pydata#3956) keep attrs in interpolate_na (pydata#3970) actually use preformatted text in the details summary (pydata#3978) facetgrid: Ensure that colormap params are only determined once. (pydata#3915) RasterioDeprecationWarning (pydata#3964) Empty line missing for DataArray.assign_coords doc (pydata#3963) New coords to existing dim (doc) (pydata#3958) implement a more threadsafe call to colorbar (pydata#3944) Fix wrong order of coordinate converted from pd.series with MultiIndex (pydata#3953) Updated list of core developers (pydata#3943) ...

TomAugspurger

This looks quite nice.

Question on the name template. I think in dask.dataframe and dask.array we might call this meta. Is that keyword already used elsewhere in xarray? template is also a fine name though.

doc/dask.rst

dcherian · 2020-04-30T20:35:13Z

Thanks for the review @TomAugspurger

Question on the name template. I think in dask.dataframe and dask.array we might call this meta. Is that keyword already used elsewhere in xarray? template is also a fine name though.

I added the meta kwarg to apply_ufunc so that users could pass that down to dask i.e. that meta = dask's meta = np.ndarray or something like that. So I'd like to avoid reusing meta here where it would exclusively be an xarray object ≠ dask's meta

BUT it seems to me like there's a better name than template. Any ideas?

TomAugspurger · 2020-04-30T21:38:21Z

Makes sense. template seems fine.

…

On Thu, Apr 30, 2020 at 3:35 PM Deepak Cherian ***@***.***> wrote: Thanks for the review @TomAugspurger <https://github.com/TomAugspurger> Question on the name template. I think in dask.dataframe and dask.array we might call this meta. Is that keyword already used elsewhere in xarray? template is also a fine name though. I added the meta kwarg to apply_ufunc so that users could pass that down to dask i.e. that meta = dask's meta = np.ndarray or something like that. So I'd like to avoid reusing meta here where it would exclusively be an xarray object ≠ dask's meta BUT it seems to me like there's a better name than template. Any ideas? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3816 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKAOIQWE7DGYAOSJWLG5F3RPHOI7ANCNFSM4K7ODDRA> .

shoyer · 2020-04-30T22:05:50Z

I think it's also a good idea to use a different name from dask's meta because template is used differently. It's an object with the full shape of the desired output, not size 0 dimensions.

…f chunks

dcherian

Made all requested changes. ~~The only thing remaining is attrs.~~

EDIT: docs now mention that attrs is copied over from template.

xarray/core/parallel.py

xarray/tests/test_dask.py

shoyer

Looks good to me!

(Though we should probably add that warning about attrs first)

shoyer · 2020-05-05T01:21:55Z

xarray/core/dataarray.py

+            the function will be first run on mocked-up data, that looks like 'obj' but
+            has sizes 0, to determine properties of the returned object such as dtype,
+            variable names, new dimensions and new indexes (if any).
+            'template' must be provided if the function changes the size of existing dimensions.


Correct, attrs needs to be set when the resulting Dataset is created, before we do the the computation.

bradyrx · 2020-05-05T16:28:26Z

I missed this originally @dcherian, but thanks for the great work here. The docs changes are a great help.

dcherian · 2020-05-05T18:14:06Z

Added docs on attrs issue. This should be good to go.

Thanks for the kind words, @bradyrx

dcherian · 2020-05-06T16:41:48Z

Thanks for the reviews everyone.

…k-issues * upstream/master: (22 commits) support darkmode (pydata#4036) Use literal syntax instead of function calls to create the data structure (pydata#4038) Add template xarray object kwarg to map_blocks (pydata#3816) Transpose coords by default (pydata#3824) Remove broken test for Panel with to_pandas() (pydata#4028) Allow warning with cartopy in docs plotting build (pydata#4032) Support overriding existing variables in to_zarr() without appending (pydata#4029) chore: Remove unnecessary comprehension (pydata#4026) fix to_netcdf docstring typo (pydata#4021) Pint support for DataArray (pydata#3643) Apply blackdoc to the documentation (pydata#4012) ensure Variable._repr_html_ works (pydata#3973) Fix handling of abbreviated units like msec (pydata#3998) full_like: error on non-scalar fill_value (pydata#3979) Fix some code quality and bug-risk issues (pydata#3999) DOC: add pandas.DataFrame.to_xarray (pydata#3994) Better chunking error messages for zarr backend (pydata#3983) Silence sphinx warnings (pydata#3990) Fix distributed tests on upstream-dev (pydata#3989) Add multi-dimensional extrapolation example and mention different behavior of kwargs in interp (pydata#3956) ...

dcherian added 7 commits March 2, 2020 11:26

Allow providing template dataset to map_blocks.

064b583

Update dimension shape check.

1f14b11

This accounts for dimension sizes being changed by the applied function.

Allow user function to add new unindexed dimension.

045ae2b

Add docstring for template.

0e37b63

renaming

5704e84

Raise nice error if adding a new chunked dimension,

2c458e1

Raise nice error message when expected dimension is missing on return…

dced076

…ed object

dcherian commented Mar 2, 2020

View reviewed changes

xarray/core/parallel.py Show resolved Hide resolved

Revert "Allow user function to add new unindexed dimension."

717d900

This reverts commit 045ae2b.

dcherian mentioned this pull request Mar 2, 2020

map_blocks: Allow passing dask-backed objects in args #3818

Merged

3 tasks

dcherian added 3 commits March 3, 2020 22:27

Add test + fix output_chunks for dataarray template

42a9070

typing

64ba31f

fix test

a68cb41

dcherian changed the title ~~[WIP] Provide template xarray object to map_blocks~~ [WIP] Add template xarray object kwarg to map_blocks Mar 3, 2020

dcherian mentioned this pull request Mar 4, 2020

map_blocks output inference problems #3575

Closed

Add nice error messages when result doesn't match template.

0bc3754

blacken

ee66c88

dcherian requested a review from crusaderky March 9, 2020 09:40

jhamman reviewed Mar 19, 2020

View reviewed changes

xarray/tests/test_dask.py Show resolved Hide resolved

jhamman reviewed Mar 19, 2020

View reviewed changes

xarray/core/parallel.py Show resolved Hide resolved

Add template kwarg to DataArray.map_blocks & Dataset.map_blocks

d52dfd6

dcherian commented Mar 19, 2020

View reviewed changes

xarray/tests/test_dask.py Outdated Show resolved Hide resolved

dcherian added 4 commits March 19, 2020 05:35

minor error message fixes.

8ef47f6

docstring updates.

376f242

bugfix for expected shapes when template is not specified

d9029eb

Add map_blocks docs.

6f69955

dcherian commented Mar 19, 2020

View reviewed changes

doc/dask.rst Show resolved Hide resolved

aaronspring reviewed Mar 19, 2020

View reviewed changes

doc/dask.rst Show resolved Hide resolved

aaronspring reviewed Mar 19, 2020

View reviewed changes

doc/dask.rst Outdated Show resolved Hide resolved

dcherian added 2 commits March 28, 2020 08:10

Merge branch 'map-blocks-schema' of github.com:dcherian/xarray into m…

66fe4c4

…ap-blocks-schema * 'map-blocks-schema' of github.com:dcherian/xarray: Update doc/dask.rst

jhamman mentioned this pull request Apr 6, 2020

apply_ufunc(dask='parallelized') output_dtypes for datasets #1699

Open

shoyer reviewed Apr 30, 2020

View reviewed changes

TomAugspurger reviewed Apr 30, 2020

View reviewed changes

doc/dask.rst Outdated Show resolved Hide resolved

dcherian added 5 commits May 1, 2020 10:44

refactor out slicer for chunks

0f741e2

Check expected index values.

869e62d

Raise nice error when template object does not have required number o…

334e8d3

…f chunks

doc updates.

4dd699e

more review comments.

27eb873

dcherian commented May 1, 2020

View reviewed changes

xarray/core/parallel.py Show resolved Hide resolved

xarray/core/parallel.py Show resolved Hide resolved

xarray/core/parallel.py Show resolved Hide resolved

xarray/tests/test_dask.py Show resolved Hide resolved

Mention that attrs are taken from template.

04a2c3c

shoyer approved these changes May 5, 2020

View reviewed changes

dcherian mentioned this pull request May 5, 2020

0.16.0 release #4031

Closed

23 tasks

Add test and explicit point out that attrs is copied from template

4b92168

shoyer approved these changes May 5, 2020

View reviewed changes

jhamman approved these changes May 6, 2020

View reviewed changes

dcherian merged commit 0b6e22f into pydata:master May 6, 2020

dcherian deleted the map-blocks-schema branch May 6, 2020 16:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add template xarray object kwarg to map_blocks #3816

Add template xarray object kwarg to map_blocks #3816

dcherian commented Mar 2, 2020 •

edited

Loading

pep8speaks commented Mar 6, 2020 •

edited

Loading

jhamman left a comment •

edited by dcherian

Loading

shoyer Apr 30, 2020

dcherian May 1, 2020 •

edited

Loading

dcherian May 1, 2020

shoyer May 5, 2020

shoyer commented Apr 30, 2020

TomAugspurger left a comment

dcherian commented Apr 30, 2020

TomAugspurger commented Apr 30, 2020 via email

shoyer commented Apr 30, 2020

dcherian left a comment •

edited

Loading

shoyer left a comment

shoyer May 5, 2020

bradyrx commented May 5, 2020

dcherian commented May 5, 2020

dcherian commented May 6, 2020

Add template xarray object kwarg to map_blocks #3816

Add template xarray object kwarg to map_blocks #3816

Conversation

dcherian commented Mar 2, 2020 • edited Loading

pep8speaks commented Mar 6, 2020 • edited Loading

Comment last updated at 2020-05-05 18:13:17 UTC

jhamman left a comment • edited by dcherian Loading

Choose a reason for hiding this comment

shoyer Apr 30, 2020

Choose a reason for hiding this comment

dcherian May 1, 2020 • edited Loading

Choose a reason for hiding this comment

dcherian May 1, 2020

Choose a reason for hiding this comment

shoyer May 5, 2020

Choose a reason for hiding this comment

shoyer commented Apr 30, 2020

TomAugspurger left a comment

Choose a reason for hiding this comment

dcherian commented Apr 30, 2020

TomAugspurger commented Apr 30, 2020 via email

shoyer commented Apr 30, 2020

dcherian left a comment • edited Loading

Choose a reason for hiding this comment

shoyer left a comment

Choose a reason for hiding this comment

shoyer May 5, 2020

Choose a reason for hiding this comment

bradyrx commented May 5, 2020

dcherian commented May 5, 2020

dcherian commented May 6, 2020

dcherian commented Mar 2, 2020 •

edited

Loading

pep8speaks commented Mar 6, 2020 •

edited

Loading

jhamman left a comment •

edited by dcherian

Loading

dcherian May 1, 2020 •

edited

Loading

dcherian left a comment •

edited

Loading