bug fixes for xbeam.Dataset

PiperOrigin-RevId: 813392384

Author: shoyer

xarray_beam/_src/dataset.py

@@ -139,7 +139,7 @@ def method(self: Dataset, *args, **kwargs) -> Dataset:
       )
     else:
       ptransform = self.ptransform | label >> beam.MapTuple(
-          functools.partial(_apply_to_each_chunk, func)
+          functools.partial(_apply_to_each_chunk, func, self.chunks, chunks)
       )
     return Dataset(template, chunks, self.split_vars, ptransform)
 
@@ -417,4 +417,4 @@ def head(self, **indexers_kwargs: int) -> Dataset:
 
   def pipe(self, func, *args, **kwargs):
     """Apply a function to this dataset, like xarray.Dataset.pipe()."""
-    return func(*args, **kwargs)
+    return func(self, *args, **kwargs)

xarray_beam/_src/dataset_test.py

@@ -212,7 +212,7 @@ def test_lazy_methods(self, call):
       xarray.testing.assert_identical(expected, actual)
 
     with self.subTest('already_transformed'):
-      result = beam_ds.map_blocks(lambda x: x).pipe(call, beam_ds)
+      result = beam_ds.map_blocks(lambda x: x).pipe(call)
       actual = result.collect_with_direct_runner()
       xarray.testing.assert_identical(expected, actual)
 
@@ -284,6 +284,14 @@ def test_infer_new_chunks_uneven_new_size_error(self):
           old_sizes={'x': 10}, old_chunks={'x': 5}, new_sizes={'x': 3}
       )
 
+  def test_pipe(self):
+    source = xarray.Dataset({'foo': ('x', np.arange(10))})
+    source_ds = xbeam.Dataset.from_xarray(source, {'x': 5})
+    mapped_ds = source_ds.pipe(xbeam.Dataset.map_blocks, lambda ds: 2 * ds)
+    expected = 2 * source
+    actual = mapped_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(actual, expected)
+
 
 class MapBlocksTest(test_util.TestCase):