Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix writing of DataTree subgroups to zarr or netCDF #9677

Merged
merged 9 commits into from
Nov 4, 2024
Merged

Fix writing of DataTree subgroups to zarr or netCDF #9677

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
1709721
Fix writing of DataTree subgroups to zarr or netCDF
shoyer Oct 24, 2024
85711b1
Merge branch 'main' into datatree-write-subgroup
shoyer Oct 24, 2024
7e5ccd2
Add write_inherited_coords for additional control in DataTree.to_zarr
shoyer Oct 24, 2024
c51811e
Merge branch 'main' into datatree-write-subgroup
shoyer Oct 28, 2024
5646c61
Switch write_inherited_coords default to false
shoyer Oct 28, 2024
9cf885f
Merge branch 'main' into datatree-write-subgroup
shoyer Oct 28, 2024
0311614
add whats new
shoyer Oct 28, 2024
b7995f1
Merge branch 'main' into datatree-write-subgroup
shoyer Nov 4, 2024
6df6817
remove unused import
shoyer Nov 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions xarray/core/datatree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1573,6 +1573,7 @@ def to_netcdf(
format: T_DataTreeNetcdfTypes | None = None,
engine: T_DataTreeNetcdfEngine | None = None,
group: str | None = None,
write_inherited_coords: bool = True,
compute: bool = True,
**kwargs,
):
Expand Down Expand Up @@ -1609,6 +1610,11 @@ def to_netcdf(
group : str, optional
Path to the netCDF4 group in the given file to open as the root group
of the ``DataTree``. Currently, specifying a group is not supported.
write_inherited_coords : bool, default: True
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what happens at read-time? Do we normalize and remove duplicated coords by default? If so, I think we should change so that does not happen by default.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what happens at read-time? Do we normalize and remove duplicated coords by default?

Yes, this is what happens.

If so, I think we should change so that does not happen by default.

Can you clarify what you mean here?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIUC the default is to denormalize at write-time and normalize at read-time, where "normalizing" means that coordinates are de-duplicated and inherited from parent groups where possible.

To me, that's confusing in that you really have to opt-in to round-trip to disk exactly what you have in memory. Also, anyone reading a CF-compliant store with inherited coordinates and writing it out will be surprised when opening it up with another library. FWIW it seems to me that we usually regret this kind of convenience/consistency tradeoff in the long-run.

I don't feel too strongly though. At the very least we should add some docs on how to roundtrip exactly, and how to open and write CF-compliant datasets with coordinate inheritance.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The advantage of writing things out this way is that you get the same DataTree from reading out a particular group, i.e.,

tree.to_zarr(path)
child_from_disk = open_datatree(path, group=child)
xarray.testing.assert_equal(tree[child], child_from_disk)

I think this will be a little less surprising to users, but overall I agree that it does not matter too much.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd expect when handling CF compliant data to roundtrip accordingly.

This might lead to issues on the user side.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've reverted changing this behavior -- the default is now write_inherited_coords=False.

If true, replicate inherited coordinates on all descendant nodes.
Otherwise, only write coordinates at the level at which they are
originally defined. This saves disk space, but requires opening the
full tree to load inherited coordinates.
compute : bool, default: True
If true compute immediately, otherwise return a
``dask.delayed.Delayed`` object that can be computed later.
Expand All @@ -1632,6 +1638,7 @@ def to_netcdf(
format=format,
engine=engine,
group=group,
write_inherited_coords=write_inherited_coords,
compute=compute,
**kwargs,
)
Expand All @@ -1643,6 +1650,7 @@ def to_zarr(
encoding=None,
consolidated: bool = True,
group: str | None = None,
write_inherited_coords: bool = True,
compute: Literal[True] = True,
**kwargs,
):
Expand All @@ -1668,6 +1676,11 @@ def to_zarr(
after writing metadata for all groups.
group : str, optional
Group path. (a.k.a. `path` in zarr terminology.)
write_inherited_coords : bool, default: True
If true, replicate inherited coordinates on all descendant nodes.
Otherwise, only write coordinates at the level at which they are
originally defined. This saves disk space, but requires opening the
full tree to load inherited coordinates.
compute : bool, default: True
If true compute immediately, otherwise return a
``dask.delayed.Delayed`` object that can be computed later. Metadata
Expand All @@ -1690,6 +1703,7 @@ def to_zarr(
encoding=encoding,
consolidated=consolidated,
group=group,
write_inherited_coords=write_inherited_coords,
compute=compute,
**kwargs,
)
Expand Down
6 changes: 4 additions & 2 deletions xarray/core/datatree_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ def _datatree_to_netcdf(
format: T_DataTreeNetcdfTypes | None = None,
engine: T_DataTreeNetcdfEngine | None = None,
group: str | None = None,
write_inherited_coords: bool = True,
compute: bool = True,
**kwargs,
):
Expand Down Expand Up @@ -59,7 +60,7 @@ def _datatree_to_netcdf(

for node in dt.subtree:
at_root = node is dt
ds = node.to_dataset(inherit=at_root)
ds = node.to_dataset(inherit=write_inherited_coords or at_root)
group_path = None if at_root else "/" + node.relative_to(dt)
ds.to_netcdf(
filepath,
Expand All @@ -82,6 +83,7 @@ def _datatree_to_zarr(
encoding: Mapping[str, Any] | None = None,
consolidated: bool = True,
group: str | None = None,
write_inherited_coords: bool = True,
compute: Literal[True] = True,
**kwargs,
):
Expand Down Expand Up @@ -114,7 +116,7 @@ def _datatree_to_zarr(

for node in dt.subtree:
at_root = node is dt
ds = node.to_dataset(inherit=at_root)
ds = node.to_dataset(inherit=write_inherited_coords or at_root)
group_path = None if at_root else "/" + node.relative_to(dt)
ds.to_zarr(
store,
Expand Down
19 changes: 19 additions & 0 deletions xarray/tests/test_backends_datatree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -592,3 +592,22 @@ def test_write_subgroup(self, tmpdir):
with open_datatree(filepath, engine="zarr") as roundtrip_dt:
assert_equal(original_dt, roundtrip_dt)
assert_identical(expected_dt, roundtrip_dt)

def test_write_inherited_coords(self, tmpdir):
original_dt = DataTree.from_dict(
{
"/": xr.Dataset(coords={"x": [1, 2, 3]}),
"/child": xr.Dataset({"foo": ("x", [4, 5, 6])}),
}
)

filepath = tmpdir / "test.zarr"
original_dt.to_zarr(filepath, write_inherited_coords=False)

with open_datatree(filepath, engine="zarr") as roundtrip_dt:
assert_identical(original_dt, roundtrip_dt)

expected_child = original_dt.children["child"].copy(inherit=False)
expected_child.name = None
with open_datatree(filepath, group="child", engine="zarr") as roundtrip_child:
assert_identical(expected_child, roundtrip_child)
Loading