Add xbeam.Dataset.mean

Also fix a bug in writing scalar values to Zarr that this turned up.

PiperOrigin-RevId: 813438224

Author: shoyer

docs/api.md

@@ -100,6 +100,7 @@ guarantees.
     Dataset.rechunk
     Dataset.split_variables
     Dataset.consolidate_variables
+    Dataset.mean
     Dataset.head
     Dataset.pipe
 ```

xarray_beam/_src/dataset.py

@@ -41,7 +41,9 @@
 from typing import Any, Callable, Literal
 
 import apache_beam as beam
+import numpy.typing as npt
 import xarray
+from xarray_beam._src import combiners
 from xarray_beam._src import core
 from xarray_beam._src import rechunk
 from xarray_beam._src import zarr
@@ -407,6 +409,42 @@ def consolidate_variables(self) -> Dataset:
     ptransform = self.ptransform | label >> rechunk.ConsolidateVariables()
     return type(self)(self.template, self.chunks, split_vars, ptransform)
 
+  def mean(
+      self,
+      dim: str | list[str] | tuple[str, ...] | None = None,
+      *,
+      skipna: bool | None = None,
+      dtype: npt.DTypeLike | None = None,
+      fanout: int | None = None,
+  ) -> Dataset:
+    """Compute the mean of this Dataset using Beam combiners.
+
+    Args:
+      dim: dimension(s) to compute the mean over.
+      skipna: whether to skip missing data when computing the mean.
+      dtype: the desired dtype of the resulting Dataset.
+      fanout: size of an intermediate fanout stage for Beam combiners.
+
+    Returns:
+      New Dataset with the mean computed.
+    """
+    # TODO(shoyer): use heuristics to pick a default fanout size.
+    if dim is None:
+      dims = list(self.template.dims)
+    elif isinstance(dim, str):
+      dims = [dim]
+    else:
+      dims = dim
+    template = zarr.make_template(
+        self.template.mean(dim=dims, skipna=skipna, dtype=dtype)
+    )
+    chunks = {k: v for k, v in self.chunks.items() if k not in dims}
+    label = _get_label(f"mean_{'_'.join(dims)}")
+    ptransform = self.ptransform | label >> combiners.Mean(
+        dim=dims, skipna=skipna, dtype=dtype, fanout=fanout
+    )
+    return type(self)(template, chunks, self.split_vars, ptransform)
+
   _head = _whole_dataset_method('head')
 
   def head(self, **indexers_kwargs: int) -> Dataset:
@@ -419,8 +457,6 @@ def head(self, **indexers_kwargs: int) -> Dataset:
       )
     return self._head(**indexers_kwargs)
 
-  # TODO(shoyer): implement merge, rename, mean, etc
-
   # thin wrappers around xarray methods
   __getitem__ = _whole_dataset_method('__getitem__')
   transpose = _whole_dataset_method('transpose')

xarray_beam/_src/dataset_test.py

@@ -530,5 +530,26 @@ def test_resample(self):
     self.assertEqual(chunks, {'time': 10, 'latitude': 73, 'longitude': 144})
 
 
+class MeanTest(test_util.TestCase):
+
+  @parameterized.named_parameters(
+      dict(testcase_name='x', dim='x', skipna=True, fanout=None),
+      dict(testcase_name='y', dim='y', skipna=True, fanout=None),
+      dict(testcase_name='two_dims', dim=['x', 'y'], skipna=True, fanout=None),
+      dict(testcase_name='all_dims', dim=None, skipna=True, fanout=None),
+      dict(testcase_name='skipna_false', dim='y', skipna=False, fanout=None),
+      dict(testcase_name='with_fanout', dim='y', skipna=True, fanout=2),
+  )
+  def test_mean(self, dim, skipna, fanout):
+    source_ds = xarray.Dataset(
+        {'foo': (('x', 'y'), np.array([[1, 2, np.nan], [4, np.nan, 6]]))}
+    )
+    beam_ds = xbeam.Dataset.from_xarray(source_ds, chunks={'x': 1})
+    actual = beam_ds.mean(dim=dim, skipna=skipna, fanout=fanout)
+    expected = source_ds.mean(dim=dim, skipna=skipna)
+    actual_collected = actual.collect_with_direct_runner()
+    xarray.testing.assert_identical(expected, actual_collected)
+
+
 if __name__ == '__main__':
   absltest.main()

xarray_beam/_src/zarr.py

@@ -268,7 +268,7 @@ def expand(self, pcoll):
 
 def _verify_template_is_lazy(template: xarray.Dataset):
   """Verify that a Dataset is suitable for use as a Zarr template."""
-  if not template.chunks:
+  if all(var.chunks is None for var in template.variables.values()):
     # We require at least one chunked variable with Dask. Otherwise, there would
     # be no data to write as part of the Beam pipeline.
     raise ValueError(
@@ -555,7 +555,9 @@ def write_chunk_to_zarr(
     # setup_zarr.
     future = writable_chunk.to_zarr(
         store,
-        region=region,
+        # Xarray has a bug where it does not support region={}. This will be
+        # fixed upstream in https://github.com/pydata/xarray/pull/10796
+        region=region if region else None,
         compute=False,
         consolidated=True,
         mode='r+',

xarray_beam/_src/zarr_test.py

@@ -416,7 +416,7 @@ def test_chunks_to_zarr_with_invalid_shards(self):
     temp_dir = self.create_tempdir().full_path
     with self.assertRaisesWithLiteralMatch(
         ValueError,
-        "shard sizes are not all evenly divisible by chunk sizes: "
+        'shard sizes are not all evenly divisible by chunk sizes: '
         "shards={'x': 1, 'y': 1}, chunks={'x': 2, 'y': 3}",
     ):
       xbeam.ChunksToZarr(
@@ -729,9 +729,9 @@ def test_zarr_from_dask_chunks(self):
 
   def test_chunks_to_zarr_docs_demo(self):
     # verify that the ChunksToChunk demo from our docs works
-    data = np.random.RandomState(0).randn(2920//100, 25, 53)
+    data = np.random.RandomState(0).randn(2920 // 100, 25, 53)
     ds = xarray.Dataset({'temperature': (('time', 'lat', 'lon'), data)})
-    chunks = {'time': 1000//100, 'lat': 25, 'lon': 53}
+    chunks = {'time': 1000 // 100, 'lat': 25, 'lon': 53}
     temp_dir = self.create_tempdir().full_path
     (
         test_util.EagerPipeline()
@@ -743,6 +743,15 @@ def test_chunks_to_zarr_docs_demo(self):
     result = xarray.open_zarr(temp_dir)
     xarray.testing.assert_identical(result, ds)
 
+  def test_chunks_to_zarr_scalar_variable(self):
+    dataset = xarray.Dataset({'foo': da.zeros(())})
+    temp_dir = self.create_tempdir().full_path
+    [(xbeam.Key({}), dataset.compute())] | xbeam.ChunksToZarr(
+        temp_dir, template=dataset
+    )
+    actual = xarray.open_zarr(temp_dir, consolidated=True)
+    xarray.testing.assert_identical(actual, dataset)
+
 
 if __name__ == '__main__':
   absltest.main()