Add xbeam.Dataset.head(), for quick testing purposes

PiperOrigin-RevId: 812923318

Author: shoyer

xarray_beam/_src/dataset.py

@@ -332,6 +332,18 @@ def consolidate_variables(self) -> Dataset:
     ptransform = self.ptransform | label >> rechunk.ConsolidateVariables()
     return type(self)(self.template, self.chunks, split_vars, ptransform)
 
+  _head = _whole_dataset_method('head')
+
+  def head(self, **indexers_kwargs: int) -> Dataset:
+    """Return a Dataset with the first N elements of each dimension."""
+    if not isinstance(self.ptransform, core.DatasetToChunks):
+      raise ValueError(
+          'head() is only supported on untransformed datasets, with '
+          'ptransform=DatasetToChunks. This dataset has '
+          f'ptransform={self.ptransform}'
+      )
+    return self._head(**indexers_kwargs)
+
   # TODO(shoyer): implement merge, rename, mean, etc
 
   # thin wrappers around xarray methods

xarray_beam/_src/dataset_test.py

@@ -179,6 +179,25 @@ def test_lazy_methods(self, call):
       actual = result.collect_with_direct_runner()
       xarray.testing.assert_identical(expected, actual)
 
+  def test_head(self):
+    ds = xarray.Dataset({'foo': ('x', np.arange(10))})
+    beam_ds = xbeam.Dataset.from_xarray(ds, {'x': 5})
+
+    head_ds = beam_ds.head(x=2)
+    self.assertRegex(head_ds.ptransform.label, r'^from_xarray_\d+|head_\d+$')
+    expected = ds.head(x=2)
+    actual = head_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(expected, actual)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        re.escape(
+            'head() is only supported on untransformed datasets, with '
+            'ptransform=DatasetToChunks. This dataset has ptransform='
+        ),
+    ):
+      beam_ds.map_blocks(lambda x: x).head(x=2)
+
   @parameterized.named_parameters(
       dict(
           testcase_name='no_chunking',