Add xarray_beam.Dataset.tail

PiperOrigin-RevId: 822727787

Author: shoyer

xarray_beam/_src/dataset.py

@@ -974,6 +974,18 @@ def head(self, **indexers_kwargs: int) -> Dataset:
       )
     return self._head(**indexers_kwargs)
 
+  _tail = _whole_dataset_method('tail')
+
+  def tail(self, **indexers_kwargs: int) -> Dataset:
+    """Return a Dataset with the last N elements of each dimension."""
+    if not isinstance(self.ptransform, core.DatasetToChunks):
+      raise ValueError(
+          'tail() is only supported on untransformed datasets, with '
+          'ptransform=DatasetToChunks. This dataset has '
+          f'ptransform={self.ptransform}'
+      )
+    return self._tail(**indexers_kwargs)
+
   # thin wrappers around xarray methods
   __getitem__ = _whole_dataset_method('__getitem__')
   transpose = _whole_dataset_method('transpose')

xarray_beam/_src/dataset_test.py

@@ -849,6 +849,17 @@ def test_head(self):
     ):
       beam_ds.map_blocks(lambda x: x).head(x=2)
 
+  def test_tail(self):
+    ds = xarray.Dataset({'foo': ('x', np.arange(10))})
+    beam_ds = xbeam.Dataset.from_xarray(ds, {'x': 5})
+
+    tail_ds = beam_ds.tail(x=2)
+    self.assertEqual(tail_ds.chunks, {'x': 2})
+    self.assertRegex(tail_ds.ptransform.label, r'^from_xarray_\d+|tail_\d+$')
+    expected = ds.tail(x=2)
+    actual = tail_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(expected, actual)
+
   @parameterized.named_parameters(
       dict(
           testcase_name='no_chunking',