Add rechunking methods to xarray_beam.Dataset

PiperOrigin-RevId: 812868604

Author: shoyer

xarray_beam/_src/dataset.py

@@ -246,6 +246,58 @@ def map_blocks(
     )
     return type(self)(template, chunks, self.split_vars, ptransform)
 
+  # rechunking methods
+
+  def rechunk(
+      self,
+      chunks: dict[str, int],
+      min_mem: int | None = None,
+      max_mem: int = 2**30,
+  ) -> Dataset:
+    """Rechunk this Dataset.
+
+    Args:
+      chunks: new chunk sizes, as a dict mapping from dimension name to chunk
+        size. -1 is interpreted as a "full chunk".
+      min_mem: optional minimum memory usage for rechunking.
+      max_mem: optional maximum memory usage for rechunking.
+
+    Returns:
+      New Dataset with updated chunks.
+    """
+    # TODO(shoyer): support human readable strings for chunksizes like dask,
+    # e.g., chunks={"time": "10 MB"}.
+    chunks = rechunk.normalize_chunks(chunks, self.sizes)  # pytype: disable=wrong-arg-types
+    if self.split_vars:
+      itemsize = max(v.dtype.itemsize for v in self.template.values())
+    else:
+      itemsize = sum(v.dtype.itemsize for v in self.template.values())
+    rechunk_transform = rechunk.Rechunk(
+        self.sizes,
+        self.chunks,
+        chunks,
+        itemsize=itemsize,
+        min_mem=min_mem,
+        max_mem=max_mem,
+    )
+    label = _get_label('rechunk')
+    ptransform = self.ptransform | label >> rechunk_transform
+    return type(self)(self.template, chunks, self.split_vars, ptransform)
+
+  def split_variables(self) -> Dataset:
+    """Split variables in this Dataset into separate chunks."""
+    split_vars = True
+    label = _get_label('split_vars')
+    ptransform = self.ptransform | label >> rechunk.SplitVariables()
+    return type(self)(self.template, self.chunks, split_vars, ptransform)
+
+  def consolidate_variables(self) -> Dataset:
+    """Consolidate variables in this Dataset into a single chunk."""
+    split_vars = False
+    label = _get_label('consolidate_vars')
+    ptransform = self.ptransform | label >> rechunk.ConsolidateVariables()
+    return type(self)(self.template, self.chunks, split_vars, ptransform)
+
   # TODO(shoyer): implement merge, rename, mean, etc
 
   # thin wrappers around xarray methods

xarray_beam/_src/dataset_test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import re
-import textwrap
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -185,10 +184,10 @@ def test_map_blocks_new_vars_and_dims(self):
     source = xarray.Dataset({'foo': ('x', np.arange(10))})
     source_ds = xbeam.Dataset.from_xarray(source, {'x': 5})
     mapped_ds = source_ds.map_blocks(
-        lambda ds: ds.assign(bar=2*ds.foo.expand_dims('y'))
+        lambda ds: ds.assign(bar=2 * ds.foo.expand_dims('y'))
     )
     self.assertEqual(mapped_ds.chunks, {'x': 5, 'y': 1})
-    expected = source.assign(bar=2*source.foo.expand_dims('y'))
+    expected = source.assign(bar=2 * source.foo.expand_dims('y'))
     actual = mapped_ds.collect_with_direct_runner()
     xarray.testing.assert_identical(actual, expected)
 
@@ -240,5 +239,61 @@ def test_map_blocks_explicit_template(self):
     xarray.testing.assert_identical(actual, source)
 
 
+class RechunkingTest(test_util.TestCase):
+
+  def test_split_variables(self):
+    source = xarray.Dataset(
+        {'foo': ('x', np.arange(10)), 'bar': ('x', np.arange(10))}
+    )
+    beam_ds = xbeam.Dataset.from_xarray(source, {'x': 5}, split_vars=False)
+    self.assertFalse(beam_ds.split_vars)
+    split_ds = beam_ds.split_variables()
+    self.assertTrue(split_ds.split_vars)
+    self.assertRegex(
+        split_ds.ptransform.label, r'^from_xarray_\d+\|split_vars_\d+$'
+    )
+    actual = split_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(actual, source)
+
+  def test_consolidate_variables(self):
+    source = xarray.Dataset(
+        {'foo': ('x', np.arange(10)), 'bar': ('x', np.arange(10))}
+    )
+    beam_ds = xbeam.Dataset.from_xarray(source, {'x': 5}, split_vars=True)
+    self.assertTrue(beam_ds.split_vars)
+    consolidated_ds = beam_ds.consolidate_variables()
+    self.assertFalse(consolidated_ds.split_vars)
+    self.assertRegex(
+        consolidated_ds.ptransform.label,
+        r'^from_xarray_\d+\|consolidate_vars_\d+$',
+    )
+    actual = consolidated_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(actual, source)
+
+  def test_rechunk(self):
+    source_chunks = {'x': 5, 'y': 1}
+    target_chunks = {'x': 2, 'y': -1}
+    source = xarray.Dataset({'foo': (('x', 'y'), np.arange(40).reshape(10, 4))})
+    beam_ds = xbeam.Dataset.from_xarray(source, source_chunks)
+    rechunked_ds = beam_ds.rechunk(target_chunks)
+
+    self.assertEqual(rechunked_ds.chunks, {'x': 2, 'y': 4})
+    actual = rechunked_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(actual, source)
+
+  def test_rechunk_split_vars(self):
+    source = xarray.Dataset({
+        'foo': (('x', 'y'), np.arange(20).reshape(10, 2)),
+        'bar': ('x', np.arange(10)),
+    })
+    beam_ds = xbeam.Dataset.from_xarray(
+        source, {'x': 5, 'y': 2}, split_vars=True
+    )
+    rechunked_ds = beam_ds.rechunk({'x': 2, 'y': 1})
+    self.assertEqual(rechunked_ds.chunks, {'x': 2, 'y': 1})
+    actual = rechunked_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(actual, source)
+
+
 if __name__ == '__main__':
   absltest.main()