Fix split_variables with different dimensions per variable

PiperOrigin-RevId: 536400546

Author: shoyer

setup.py

@@ -42,7 +42,7 @@
 
 setuptools.setup(
     name='xarray-beam',
-    version='0.6.1',
+    version='0.6.2',
     license='Apache 2.0',
     author='Google LLC',
     author_email='[email protected]',

xarray_beam/__init__.py

@@ -45,4 +45,4 @@
     DatasetToZarr,
 )
 
-__version__ = '0.6.1'
+__version__ = '0.6.2'

xarray_beam/_src/rechunk.py

@@ -460,7 +460,10 @@ def split_variables(
   # TODO(shoyer): add support for partial splitting, into explicitly provided
   # sets of variables
   for var_name in dataset:
-    yield key.replace(vars={var_name}), dataset[[var_name]]
+    new_dataset = dataset[[var_name]]
+    offsets = {k: v for k, v in key.offsets.items() if k in new_dataset.dims}
+    new_key = core.Key(offsets, vars={var_name})
+    yield new_key, new_dataset
 
 
 @dataclasses.dataclass

xarray_beam/_src/rechunk_test.py

@@ -257,6 +257,33 @@ def test_consolidate_and_split_variables_only(self):
       actual = consolidated | xbeam.SplitVariables()
       self.assertIdenticalChunks(actual, split)
 
+  def test_split_variables_with_different_dims(self):
+    inputs = [
+        (
+            xbeam.Key({'x': 0, 'y': 0}, vars=None),
+            xarray.Dataset({
+                'foo': ('x', np.array([1, 2])),
+                'bar': (('x', 'y'), np.array([[1, 2, 3], [4, 5, 6]])),
+            }),
+        ),
+    ]
+    expected = [
+        (
+            xbeam.Key({'x': 0}, vars={'foo'}),
+            xarray.Dataset({
+                'foo': ('x', np.array([1, 2])),
+            }),
+        ),
+        (
+            xbeam.Key({'x': 0, 'y': 0}, vars={'bar'}),
+            xarray.Dataset({
+                'bar': (('x', 'y'), np.array([[1, 2, 3], [4, 5, 6]])),
+            }),
+        ),
+    ]
+    actual = inputs | xbeam.SplitVariables()
+    self.assertIdenticalChunks(actual, expected)
+
   def test_consolidate_chunks_not_fully_shared_dims(self):
     inputs = [
         (