implement `dask` methods on `DataTree` by keewis · Pull Request #9670 · pydata/xarray

keewis · 2024-10-24T14:15:33Z

For now this only contains compute and load, but in implementing those I've found a bug in copy: we don't actually (shallow) copy the variables, which means that my implementation of compute would modify the original tree.

follow-up to support chunks in open_groups and open_datatree #9660
Tests added

TomNicholas · 2024-10-24T16:11:26Z

+ data = self._to_dataset_view(rebuild_dims=False, inherit=inherit)._copy(
+deep=deep, memo=memo
+ )


Good catch!

keewis · 2024-10-24T17:19:06Z

we're still missing an implementation for persist, which I skipped mostly because I wasn't sure how to test that. @TomNicholas, would you like to help me with that? If so, feel free to directly push to this PR.

Co-authored-by: Tom Nicholas <tom@cworthy.org>

keewis · 2024-10-24T17:38:28Z

+ {
+ dim: size
+ for dim, size in combined_chunks.items()
+ if dim in node.dataset.dims


I had to use node.dataset.dims to avoid including inherited dims (which we can't chunk anyways because we only inherit indexed dims, and there is no chunked index so far)

otherwise I just saw _node_dims in _get_all_dims, so that might have less overhead. I'll go ahead and use that instead.

That will also be more explicit as it will avoid using "rebuilt" dims (which doesn't really matter anyway because we can't chunk indexes).

TomNicholas · 2024-10-24T17:50:14Z

we're still missing an implementation for persist, which I skipped mostly because I wasn't sure how to test that. @TomNicholas, would you like to help me with that? If so, feel free to directly push to this PR.

I'm happy to help, but I think that (a) persist is much more niche than load and compute, so isn't as important to be in there before release, and (b) it's a separate method rather than a bugfix to this PR so it would make more sense for me to add that in a follow-up PR and we just merge this one without persist.

keewis · 2024-10-24T17:52:18Z

that would be fine with me.

That leaves us with the question whether we should guard against chunk with non-existing dims, and after resolving that we should be ready for a final review.

TomNicholas · 2024-10-24T17:53:56Z

whether we should guard against chunk with non-existing dims

We should, see #9670 (comment)

keewis · 2024-10-24T18:02:21Z

We should, see #9670 (comment)

yeah, I saw that one after posting. Should be done now, though.

TomNicholas · 2024-10-24T18:14:37Z

+ Mapping from group paths to a mapping of dimension names to block lengths for this datatree's data, or None if
+ the underlying data is not a dask array.


Apparently this docstring (and the one for .chunksizes on Dataset etc.) are wrong because this method will never return a None.

keewis · 2024-10-24T18:17:17Z

actually, I just tried to compute a DataTree object without chunked arrays (while trying to investigated your comment on chunksizes), and I got an ugly error. I'm investigating, but should not take too long to fix.

keewis · 2024-10-24T18:39:30Z

Tell me what you think about the reworded docstrings of chunksizes / Dataset.chunks, but otherwise this should be good to go, too.

TomNicholas · 2024-10-24T18:40:06Z

Looks good! Let's merge.

keewis · 2024-10-24T18:41:08Z

(I just noticed that DataArray.chunksizes has the same issue, so I'll change that, too)

TomNicholas · 2024-10-24T18:46:34Z

NamedArray.chunksizes too hahaha

keewis · 2024-10-24T18:49:53Z

should be good now, hopefully

* main: Add `DataTree.persist` (pydata#9682) Typing annotations for arithmetic overrides (e.g., DataArray + Dataset) (pydata#9688) Raise `ValueError` for unmatching chunks length in `DataArray.chunk()` (pydata#9689) Fix inadvertent deep-copying of child data in DataTree (pydata#9684) new blank whatsnew (pydata#9679) v2024.10.0 release summary (pydata#9678) drop the length from `numpy`'s fixed-width string dtypes (pydata#9586) fixing behaviour for group parameter in `open_datatree` (pydata#9666) Use zarr v3 dimension_names (pydata#9669) fix(zarr): use inplace array.resize for zarr 2 and 3 (pydata#9673) implement `dask` methods on `DataTree` (pydata#9670) support `chunks` in `open_groups` and `open_datatree` (pydata#9660) Compatibility for zarr-python 3.x (pydata#9552) Update to_dataframe doc to match current behavior (pydata#9662) Reduce graph size through writing indexes directly into graph for ``map_blocks`` (pydata#9658)

* main: (85 commits) Refactor out utility functions from to_zarr (pydata#9695) Use the same function to floatize coords in polyfit and polyval (pydata#9691) Add `DataTree.persist` (pydata#9682) Typing annotations for arithmetic overrides (e.g., DataArray + Dataset) (pydata#9688) Raise `ValueError` for unmatching chunks length in `DataArray.chunk()` (pydata#9689) Fix inadvertent deep-copying of child data in DataTree (pydata#9684) new blank whatsnew (pydata#9679) v2024.10.0 release summary (pydata#9678) drop the length from `numpy`'s fixed-width string dtypes (pydata#9586) fixing behaviour for group parameter in `open_datatree` (pydata#9666) Use zarr v3 dimension_names (pydata#9669) fix(zarr): use inplace array.resize for zarr 2 and 3 (pydata#9673) implement `dask` methods on `DataTree` (pydata#9670) support `chunks` in `open_groups` and `open_datatree` (pydata#9660) Compatibility for zarr-python 3.x (pydata#9552) Update to_dataframe doc to match current behavior (pydata#9662) Reduce graph size through writing indexes directly into graph for ``map_blocks`` (pydata#9658) Add close() method to DataTree and use it to clean-up open files in tests (pydata#9651) Change URL for pydap test (pydata#9655) Fix multiple grouping with missing groups (pydata#9650) ...

keewis added 2 commits October 24, 2024 16:04

implement compute and load
f625f5d

also shallow-copy variables
507fb7d

keewis added topic-DataTree Related to the implementation of a DataTree class topic-chunked-arrays Managing different chunked backends, e.g. dask labels Oct 24, 2024

keewis commented Oct 24, 2024
View reviewed changes

Comment threadxarray/core/datatree.py

keewis added 6 commits October 24, 2024 16:35

implement chunksizes
ce1683a

add tests for load
f2a4683

add tests for chunksizes
f0ff30f

improve the load tests using DataTree.chunksizes
d12203c

add a test for compute
dda02ed

un-xfail a xpassing test
329c689

keewis commented Oct 24, 2024
View reviewed changes

Comment threadxarray/core/datatree.py Outdated

implement and test DataTree.chunk
c9fb461

keewis commented Oct 24, 2024
View reviewed changes

Comment threadxarray/core/datatree.py Outdated

TomNicholas reviewed Oct 24, 2024
View reviewed changes

Comment threadxarray/core/datatree.py

keewisand others added 3 commits October 24, 2024 19:20

link to Dataset.load
0305fc5
Co-authored-by: Tom Nicholas <tom@cworthy.org>

use tree.subtree to get absolute paths
e2a3a14

filter out missing dims before delegating to Dataset.chunk
7f57ffa

keewis commented Oct 24, 2024
View reviewed changes

keewis added 4 commits October 24, 2024 19:41

fix the type hints for DataTree.chunksizes
900701b

try using self.from_dict instead
d45dbd0

type-hint intermediate test variables
5f88937

Merge branch 'main' into datatree-dask-methods
b515972

keewis added 2 commits October 24, 2024 19:58

use _node_dims instead
39d95f6

raise on unknown chunk dim
73b4466

check that errors in chunk are raised properly
8b35676

adapt the docstrings of the new methods
d70f4f0

TomNicholas reviewed Oct 24, 2024
View reviewed changes

TomNicholas approved these changes Oct 24, 2024
View reviewed changes

TomNicholas mentioned this pull request Oct 24, 2024
Track merging datatree into xarray #8572
Closed
27 tasks

keewis added 3 commits October 24, 2024 20:21

allow computing / loading unchunked trees
3e9745f

reword the chunksizes properties
b6c5f9a

also freeze the top-level chunk sizes
da8df36

also reword DataArray.chunksizes
b11a1ef

keewis enabled auto-merge (squash) October 24, 2024 18:41

TomNicholas disabled auto-merge October 24, 2024 18:46

keewis added 2 commits October 24, 2024 20:48

fix a copy-paste error
53c0897

same for NamedArray.chunksizes
f7e31b4

TomNicholas enabled auto-merge (squash) October 24, 2024 18:52

TomNicholas merged commit 1f5889a into pydata:mainOct 24, 2024

keewis deleted the datatree-dask-methods branch October 24, 2024 19:16

TomNicholas mentioned this pull request Oct 24, 2024
Implement DataTree.persist#9675
Closed

		Mapping from group paths to a mapping of dimension names to block lengths for this datatree's data, or None if
		the underlying data is not a dask array.

Uh oh!

Uh oh!

Conversation

keewis commented Oct 24, 2024• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomNicholasOct 24, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

keewis commented Oct 24, 2024

Uh oh!

keewisOct 24, 2024

Choose a reason for hiding this comment

Uh oh!

keewisOct 24, 2024

Choose a reason for hiding this comment

Uh oh!

TomNicholasOct 24, 2024

Choose a reason for hiding this comment

Uh oh!

TomNicholas commented Oct 24, 2024

Uh oh!

keewis commented Oct 24, 2024• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

TomNicholas commented Oct 24, 2024

Uh oh!

keewis commented Oct 24, 2024

Uh oh!

TomNicholasOct 24, 2024

Choose a reason for hiding this comment

Uh oh!

keewis commented Oct 24, 2024

Uh oh!

keewis commented Oct 24, 2024

Uh oh!

TomNicholas commented Oct 24, 2024

Uh oh!

keewis commented Oct 24, 2024• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

TomNicholas commented Oct 24, 2024

Uh oh!

keewis commented Oct 24, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

keewis commented Oct 24, 2024•
edited
Loading

keewis commented Oct 24, 2024•
edited
Loading

keewis commented Oct 24, 2024•
edited
Loading