Add check to validate_chunks() that input datasets are not Dask chunked.

This is generally a bad idea for chunks used in xarray-beam pipelines, because dask arrays may have arbitrarily large associated task graph, which may be very expensive to serialize.

PiperOrigin-RevId: 753317208

Author: shoyer

xarray_beam/_src/core.py

@@ -32,7 +32,6 @@
 import immutabledict
 import numpy as np
 import xarray
-
 from xarray_beam._src import threadmap
 
 _DEFAULT = object()
@@ -473,10 +472,25 @@ def expand(self, pcoll):
 
 
 def validate_chunk(key: Key, datasets: DatasetOrDatasets) -> None:
-  """Verify that keys correspond to Dataset properties."""
+  """Verify that a key and dataset(s) are valid for xarray-beam transforms."""
   if isinstance(datasets, xarray.Dataset):
     datasets: list[xarray.Dataset] = [datasets]
+
   for dataset in datasets:
+    # Verify that no variables are chunked with Dask
+    for var_name, variable in dataset.variables.items():
+      if variable.chunks is not None:
+        raise ValueError(
+            f"Dataset variable {var_name!r} corresponding to key {key} is"
+            " chunked with Dask. Datasets passed to validate_chunk must be"
+            f" fully computed (not chunked): {dataset}\nThis typically arises"
+            " with datasets originating with `xarray.open_zarr()`, which by"
+            " default use Dask. If this is the case, you can fix it by passing"
+            " `chunks=None` or xarray_beam.open_zarr(). Alternatively, you"
+            " can load datasets explicitly into memory with `.compute()`."
+        )
+
+    # Validate key offsets
     missing_keys = [
         repr(k) for k in key.offsets.keys() if k not in dataset.dims
     ]
@@ -486,22 +500,20 @@ def validate_chunk(key: Key, datasets: DatasetOrDatasets) -> None:
           f" Dataset dimensions: {dataset!r}"
       )
 
-    if key.vars is None:
-      continue
-
-    missing_vars = [repr(v) for v in key.vars if v not in dataset.data_vars]
-    if missing_vars:
-      raise ValueError(
-          f"Key var(s) {', '.join(missing_vars)} in {key} not found in Dataset"
-          f" data variables: {dataset!r}"
-      )
+    # Validate key vars
+    if key.vars is not None:
+      missing_vars = [repr(v) for v in key.vars if v not in dataset.data_vars]
+      if missing_vars:
+        raise ValueError(
+            f"Key var(s) {', '.join(missing_vars)} in {key} not found in"
+            f" Dataset data variables: {dataset!r}"
+        )
 
 
 class ValidateEachChunk(beam.PTransform):
-  """Check that keys match the dataset for each key, dataset tuple."""
+  """Check that keys and dataset(s) are valid for xarray-beam transforms."""
 
   def _validate(self, key, dataset):
-    # Other checks may come later...
     validate_chunk(key, dataset)
     return key, dataset
 

xarray_beam/_src/core_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for xarray_beam._src.core."""
 
+import re
 from absl.testing import absltest
 from absl.testing import parameterized
 import apache_beam as beam
@@ -509,48 +510,59 @@ def test_validate(self):
 
 class ValidateEachChunkTest(test_util.TestCase):
 
+  def test_validate_chunk_raises_on_dask_chunked(self):
+    dataset = xarray.Dataset({'foo': ('x', np.arange(6))}).chunk()
+    key = xbeam.Key({'x': 0})
+
+    with self.assertRaisesRegex(
+        ValueError,
+        re.escape(
+            "Dataset variable 'foo' corresponding to key Key(offsets={'x': 0},"
+            ' vars=None) is chunked with Dask. Datasets passed to'
+            ' validate_chunk must be fully computed (not chunked):'
+        ),
+    ):
+      core.validate_chunk(key, dataset)
+
   def test_unmatched_dimension_raises_error(self):
     dataset = xarray.Dataset({'foo': ('x', np.arange(6))})
-    with self.assertRaises(ValueError) as e:
-      ([(xbeam.Key({'x': 0, 'y': 0}), dataset)] | xbeam.ValidateEachChunk())
-    self.assertIn(
-        (
+    key = xbeam.Key({'x': 0, 'y': 0})
+    with self.assertRaisesRegex(
+        ValueError,
+        re.escape(
             "Key offset(s) 'y' in Key(offsets={'x': 0, 'y': 0}, vars=None) not "
             'found in Dataset dimensions'
         ),
-        e.exception.args[0],
-    )
+    ):
+      core.validate_chunk(key, dataset)
 
-  def test_unmatched_variables_raises_error(self):
+  def test_unmatched_variables_raises_error_core(self):
     dataset = xarray.Dataset({'foo': ('x', np.arange(6))})
-    with self.assertRaises(ValueError) as e:
-      ([(xbeam.Key({'x': 0}, {'bar'}), dataset)] | xbeam.ValidateEachChunk())
-    self.assertIn(
-        (
+    key = xbeam.Key({'x': 0}, {'bar'})
+    with self.assertRaisesRegex(
+        ValueError,
+        re.escape(
             "Key var(s) 'bar' in Key(offsets={'x': 0}, vars={'bar'}) not found"
             ' in Dataset data variables'
         ),
-        e.exception.args[0],
-    )
+    ):
+      core.validate_chunk(key, dataset)
 
-  def test_unmatched_variables_multiple_datasets_raises_error(self):
+  def test_unmatched_variables_multiple_datasets_raises_error_core(self):
     datasets = [
         xarray.Dataset({'foo': ('x', i + np.arange(6))}) for i in range(11)
     ]
     datasets[5] = datasets[5].rename({'foo': 'bar'})
+    key = xbeam.Key({'x': 0}, vars={'foo'})
 
-    with self.assertRaisesWithLiteralMatch(
+    with self.assertRaisesRegex(
         ValueError,
-        (
+        re.escape(
             "Key var(s) 'foo' in Key(offsets={'x': 0}, vars={'foo'}) "
-            f'not found in Dataset data variables: {datasets[5]} '
-            "[while running 'ValidateEachChunk/MapTuple(_validate)']"
+            f'not found in Dataset data variables: {datasets[5]}'
         ),
-    ) as e:
-      (
-          [(xbeam.Key({'x': 0}, vars={'foo'}), datasets)]
-          | xbeam.ValidateEachChunk()
-      )
+    ):
+      core.validate_chunk(key, datasets)
 
   def test_validate_chunks_compose_in_pipeline(self):
     dataset = xarray.Dataset({'foo': ('x', np.arange(6))})