Raise ValueError in map_blocks if split_vars changes the set of variables.

When `split_vars=True`, `map_blocks` expects the set of variables with chunks to be invariant under the transformation. This change adds a check to ensure that the variables present in the template before and after the `map_blocks` function are the same, preventing unexpected behavior.

PiperOrigin-RevId: 813398449

Author: shoyer

xarray_beam/_src/dataset.py

@@ -332,6 +332,21 @@ def map_blocks(
         ) from e
     template = zarr.make_template(template)  # ensure template is lazy
 
+    if self.split_vars:
+      old_vars = {
+          k for k, v in self.template.variables.items() if v.chunks is not None
+      }
+      new_vars = {
+          k for k, v in template.variables.items() if v.chunks is not None
+      }
+      if old_vars != new_vars:
+        raise ValueError(
+            'cannot use map_blocks on a dataset with split_vars=True if '
+            'the transformation returns a different set of variables.\n'
+            f'Old split variables: {old_vars}\n'
+            f'New split variables: {new_vars}'
+        )
+
     if chunks is None:
       chunks = _infer_new_chunks(
           old_sizes=self.sizes,

xarray_beam/_src/dataset_test.py

@@ -364,6 +364,19 @@ def test_map_blocks_explicit_template(self):
     actual = mapped_ds.collect_with_direct_runner()
     xarray.testing.assert_identical(actual, source)
 
+  def test_map_blocks_new_split_vars_fails(self):
+    source = xarray.Dataset({'foo': ('x', np.arange(10))})
+    source_ds = xbeam.Dataset.from_xarray(source, {'x': 5}, split_vars=True)
+    func = lambda ds: ds.rename({'foo': 'bar'})
+    with self.assertRaisesWithLiteralMatch(
+        ValueError,
+        'cannot use map_blocks on a dataset with split_vars=True if the '
+        'transformation returns a different set of variables.\n'
+        "Old split variables: {'foo'}\n"
+        "New split variables: {'bar'}",
+    ):
+      source_ds.map_blocks(func)
+
 
 class RechunkingTest(test_util.TestCase):