`open_groups` for zarr backends by eni-awowale · Pull Request #9469 · pydata/xarray

eni-awowale · 2024-09-09T21:45:19Z

Part of #8572

Closesopen_groups function to open any zarr file containing nested groups #9430
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

eni-awowale · 2024-09-09T23:18:23Z

cc @TomNicholas, @flamingbear, @owenlittlejohns, @keewis and @shoyer !

eni-awowale · 2024-09-10T16:15:27Z

+ roundtrip_dict = open_groups(filepath, engine="zarr")
+ roundtrip_dt = DataTree.from_dict(roundtrip_dict)
+
+ assert open_datatree(filepath, engine="zarr").identical(roundtrip_dt)


.identical is going to be different to check if the coordinates are defined on disk.

Specifically see #9473.

Okay so from looking at this and testing this locally I don't think this will affect these tests.

TomNicholas · 2024-09-10T18:17:22Z

+ return groups_dict


 def _iter_zarr_groups(root, parent="/"):


Let's try to add type hints (though I don't know what the type of root is)
Suggested change
def_iter_zarr_groups(root, parent="/"):
def_iter_zarr_groups(root: ZarrGroup?, parent: str="/")->Iterable[str]:

I am not sure what the appropriate TypeHint would be here that wouldn't break mypy.

I had a play with this. I traced root back and found it is indeed a zarr.Group, so I think the type hint above works (along with adding from zarr import Group as ZarrGroup in the if TYPE_CHECKING block at the top of the file).
The bit that was causing an issue the assignment of a NodePath object to parent, which is a str from the function signature. (Apologies if I'm just playing catch up here from what you guys have been saying all along)
I could make mypy happy if I did the following in the function:
def _iter_zarr_groups(root: ZarrGroup, parent: str = "/") -> Iterable[str]: from xarray.core.treenode import NodePath parent_nodepath = NodePath(parent) yield str(parent_nodepath) for path, group in root.groups(): gpath = parent_nodepath / path yield str(gpath) yield from _iter_zarr_groups(group, parent=gpath)
Maybe that's an option? (Unless I'm missing a trick with this breaking the recursion?)

Sounds good!
(No need for from xarray.core.treenode import NodePath to live inside the function though)

Thanks Owen! That did the trick!

TomNicholas · 2024-09-10T18:22:53Z

 with pytest.raises(ValueError):
- open_datatree(filepath)
+ open_datatree(unaligned_datatree_nc)


We should definitely use
pytest.raises(ValueError, match=...):
to match a much more specific error. We are expecting this test to fail because of an alignment error, not any other type of ValueError.
(Maybe we should create a new exception AlignmentError(ValueError), given that we do already have MergeError(ValueError)? @shoyer)

How about something like this but in the treenode.py module?
class AlignmentError(ValueError): """Raised when a child node has coordinates that are not aligned with its parents."""
If we do this I think we might have to change datatree.py to raise AlignmentError in
`_check_alignment

Alignment is a core xarray concept separate of datatree. My suggestion was to change the ValueErrors raised in alignment.py to become AlignmentErrors, then datatree.py would automatically throw AlignmentError when _check_alignment is called (because it calls xr.align internally).
This suggestion is outside the scope of this PR though - you should go ahead with just using match to match a specific ValueError and then we can maybe do AlignmentError in a separate PR.

Updated in the last commit 🚀

for more information, see https://pre-commit.ci

TomNicholas

Thanks @eni-awowale !

owenlittlejohns

Looks good to me. Nice work!

I looked at the type hint stuff, and might have an answer 🤞.

owenlittlejohns · 2024-09-12T00:50:33Z

+ return groups_dict


 def _iter_zarr_groups(root, parent="/"):


I had a play with this. I traced root back and found it is indeed a zarr.Group, so I think the type hint above works (along with adding from zarr import Group as ZarrGroup in the if TYPE_CHECKING block at the top of the file).
The bit that was causing an issue the assignment of a NodePath object to parent, which is a str from the function signature. (Apologies if I'm just playing catch up here from what you guys have been saying all along)
I could make mypy happy if I did the following in the function:
def _iter_zarr_groups(root: ZarrGroup, parent: str = "/") -> Iterable[str]: from xarray.core.treenode import NodePath parent_nodepath = NodePath(parent) yield str(parent_nodepath) for path, group in root.groups(): gpath = parent_nodepath / path yield str(gpath) yield from _iter_zarr_groups(group, parent=gpath)
Maybe that's an option? (Unless I'm missing a trick with this breaking the recursion?)

owenlittlejohns · 2024-09-12T00:52:34Z

- subgroup1_var[:] = np.array([[0.1, 0.2]])
+ with pytest.raises(
+ ValueError,
+ match=(


Nice! (I know there's already been a conversation on using the match kwarg, but wanted to 👍. This is definitely a good way to go to make sure the right ValueError is being raised.)

owenlittlejohns · 2024-09-12T00:58:21Z

+ + ".*"
+ ),
+ ):
+ open_datatree(unaligned_datatree_nc)


I'm not sure the order of merging to expect, but if @flamingbear merges #9033 first, then we can clean up the imports and use xr.open_datatree. If it's the other way around, then it'll probably be best for one of us to take a quick skim over the tests and tidy them up in a separate PR.
My guess is that this PR will merge before #9033, so we'll likely end up with a separate tiny PR for this.

Tidying up these imports isn't a big deal, and can be left for another PR

* main: (26 commits) Forbid modifying names of DataTree objects with parents (pydata#9494) DAS-2155 - Merge datatree documentation into main docs. (pydata#9033) Make illegal path-like variable names when constructing a DataTree from a Dataset (pydata#9378) Ensure TreeNode doesn't copy in-place (pydata#9482) `open_groups` for zarr backends (pydata#9469) Update pyproject.toml (pydata#9484) New whatsnew section (pydata#9483) Release notes for v2024.09.0 (pydata#9480) Fix `DataTree.coords.__setitem__` by adding `DataTreeCoordinates` class (pydata#9451) Rename DataTree's "ds" and "data" to "dataset" (pydata#9476) Update DataTree repr to indicate inheritance (pydata#9470) Bump pypa/gh-action-pypi-publish in the actions group (pydata#9460) Repo checker (pydata#9450) Add days_in_year and decimal_year to dt accessor (pydata#9105) remove parent argument from DataTree.__init__ (pydata#9465) Fix inheritance in DataTree.copy() (pydata#9457) Implement `DataTree.__delitem__` (pydata#9453) Add ASV for datatree.from_dict (pydata#9459) Make the first argument in DataTree.from_dict positional only (pydata#9446) Fix typos across the code, doc and comments (pydata#9443) ...

* main: Opt out of floor division for float dtype time encoding (pydata#9497) fixed formatting for whats-new (pydata#9493) Forbid modifying names of DataTree objects with parents (pydata#9494) DAS-2155 - Merge datatree documentation into main docs. (pydata#9033) Make illegal path-like variable names when constructing a DataTree from a Dataset (pydata#9378) Ensure TreeNode doesn't copy in-place (pydata#9482) `open_groups` for zarr backends (pydata#9469) Update pyproject.toml (pydata#9484) New whatsnew section (pydata#9483)

keewis · 2024-09-20T15:43:47Z

+ re.escape(
+ "group '/Group1/subgroup1' is not aligned with its parents:\nGroup:\n"
+ )


I wonder if this could be just an r string? Something like
Suggested change
re.escape(
"group '/Group1/subgroup1' is not aligned with its parents:\nGroup:\n"
)
r"group '/Group1/subgroup1' is not aligned with its parents:[\n]Group:[\n].*"

* open groups zarr initial commit * added tests * Added requested changes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * TypeHint for zarr groups * update for parent --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

eni-awowaleand others added 3 commits September 7, 2024 15:17

open groups zarr initial commit
7e9d5b0

added tests
746977a

Merge branch 'main' into open-groups-zarr
7a6fd8c

eni-awowale commented Sep 9, 2024
View reviewed changes

Comment threadxarray/backends/zarr.py
Comment threadxarray/backends/zarr.py Outdated
Comment threadxarray/tests/test_backends_datatree.py

eni-awowale marked this pull request as ready for review September 9, 2024 23:13

TomNicholas added topic-backends topic-zarr Related to zarr storage library topic-DataTree Related to the implementation of a DataTree class labels Sep 10, 2024

eni-awowale commented Sep 10, 2024
View reviewed changes

TomNicholas mentioned this pull request Sep 10, 2024
adding 'open_groups_as_dict' fucntion to open any zarr file containing nested groups #9431
Closed
2 tasks

TomNicholas reviewed Sep 10, 2024
View reviewed changes

eni-awowaleand others added 3 commits September 10, 2024 16:54

Added requested changes
dac8c32

Merge branch 'main' into open-groups-zarr
633a1f5

[pre-commit.ci] auto fixes from pre-commit.com hooks
25f171e
for more information, see https://pre-commit.ci

TomNicholas added the plan to merge Final call for comments label Sep 11, 2024

TomNicholas approved these changes Sep 11, 2024
View reviewed changes

owenlittlejohns approved these changes Sep 12, 2024
View reviewed changes

eni-awowale added 2 commits September 12, 2024 07:56

TypeHint for zarr groups
f2c169b

update for parent
1f17d95

TomNicholas merged commit aeaa082 into pydata:mainSep 12, 2024

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

keewis reviewed Sep 20, 2024
View reviewed changes

	def_iter_zarr_groups(root, parent="/"):
	def_iter_zarr_groups(root: ZarrGroup?, parent: str="/")->Iterable[str]:

Uh oh!

Uh oh!

Conversation

eni-awowale commented Sep 9, 2024• edited by TomNicholas Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

eni-awowale commented Sep 9, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomNicholas left a comment

Choose a reason for hiding this comment

Uh oh!

owenlittlejohns left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

eni-awowale commented Sep 9, 2024•
edited by TomNicholas
Loading