Add end-to-end tests for the xbeam.Dataset API.

This change removes the "not fully implemented yet!" note from the `xbeam.Dataset` docstring and adds a test case that runs the example pipeline from the docstring, verifying the output and chunking.

PiperOrigin-RevId: 812912985

Author: shoyer

xarray_beam/_src/dataset_test.py

@@ -357,5 +357,76 @@ def test_rechunk_split_vars(self):
     xarray.testing.assert_identical(actual, source)
 
 
+class EndToEndTest(test_util.TestCase):
+
+  def test_docstring_example(self):
+    input_path = self.create_tempdir('source').full_path
+    output_path = self.create_tempdir('output').full_path
+
+    source_ds = test_util.dummy_era5_surface_dataset(times=365, freq='24H')
+    source_ds.chunk({'time': 90}).to_zarr(input_path)
+
+    transform = (
+        xbeam.Dataset.from_zarr(input_path)
+        .rechunk({'time': -1, 'latitude': 10, 'longitude': 10})
+        .map_blocks(lambda x: x.median('time'))
+        .to_zarr(output_path)
+    )
+    test_util.EagerPipeline() | transform
+
+    expected = source_ds.median('time')
+    actual, chunks = xbeam.open_zarr(output_path)
+    xarray.testing.assert_identical(expected, actual)
+    self.assertEqual(chunks, {'latitude': 10, 'longitude': 10})
+
+  def test_climatology(self):
+    input_path = self.create_tempdir('source').full_path
+    output_path = self.create_tempdir('output').full_path
+
+    source_ds = test_util.dummy_era5_surface_dataset(times=365, freq='24H')
+    source_ds.chunk({'time': 90}).to_zarr(input_path)
+
+    transform = (
+        xbeam.Dataset.from_zarr(input_path)
+        .rechunk({'time': -1, 'latitude': 10, 'longitude': 10})
+        .map_blocks(lambda x: x.groupby('time.month').mean())
+        .to_zarr(output_path)
+    )
+    test_util.EagerPipeline() | transform
+
+    expected = source_ds.groupby('time.month').mean()
+    actual, chunks = xbeam.open_zarr(output_path)
+    xarray.testing.assert_identical(expected, actual)
+    self.assertEqual(chunks, {'month': 12, 'latitude': 10, 'longitude': 10})
+
+  def test_resample(self):
+    input_path = self.create_tempdir('source').full_path
+    output_path = self.create_tempdir('output').full_path
+
+    source_ds = test_util.dummy_era5_surface_dataset(
+        latitudes=73, longitudes=144, times=365, freq='24H'
+    )
+    source_ds.chunk({'time': 90}).to_zarr(input_path)
+
+    transform = (
+        xbeam.Dataset.from_zarr(input_path)
+        .rechunk({'time': -1, 'latitude': 10, 'longitude': 10})
+        .map_blocks(lambda x: x.resample(time='10D').mean())
+        .rechunk({'time': 20, 'latitude': -1, 'longitude': -1})
+        .to_zarr(
+            output_path,
+            zarr_chunks={'time': 10, 'latitude': -1, 'longitude': -1},
+            zarr_shards={'time': 20, 'latitude': -1, 'longitude': -1},
+            zarr_format=3,
+        )
+    )
+    test_util.EagerPipeline() | transform
+
+    expected = source_ds.resample(time='10D').mean()
+    actual, chunks = xbeam.open_zarr(output_path)
+    xarray.testing.assert_identical(expected, actual)
+    self.assertEqual(chunks, {'time': 10, 'latitude': 73, 'longitude': 144})
+
+
 if __name__ == '__main__':
   absltest.main()

xarray_beam/_src/zarr_test.py

@@ -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, 25, 53)
+    data = np.random.RandomState(0).randn(2920//100, 25, 53)
     ds = xarray.Dataset({'temperature': (('time', 'lat', 'lon'), data)})
-    chunks = {'time': 1000, 'lat': 25, 'lon': 53}
+    chunks = {'time': 1000//100, 'lat': 25, 'lon': 53}
     temp_dir = self.create_tempdir().full_path
     (
         test_util.EagerPipeline()