Flexible Indexes: Avoid len(index) in map_blocks

pydata · Aug 3, 2021 · e65820c · e65820c
1 parent 8f5b4a1
commit e65820c
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py
@@ -295,9 +295,10 @@ def _wrapper(
         # check that index lengths and values are as expected
         for name, index in result.xindexes.items():
             if name in expected["shapes"]:
-                if len(index) != expected["shapes"][name]:
+                if result.sizes[name] != expected["shapes"][name]:
                     raise ValueError(
-                        f"Received dimension {name!r} of length {len(index)}. Expected length {expected['shapes'][name]}."
+                        f"Received dimension {name!r} of length {result.sizes[name]}. "
+                        f"Expected length {expected['shapes'][name]}."
                     )
             if name in expected["indexes"]:
                 expected_index = expected["indexes"][name]
@@ -568,8 +569,8 @@ def subset_dataset_to_block(
         for dim in dims:
             if dim in output_chunks:
                 var_chunks.append(output_chunks[dim])
-            elif dim in indexes:
-                var_chunks.append((len(indexes[dim]),))
+            elif dim in result.xindexes:
+                var_chunks.append((result.sizes[dim],))
             elif dim in template.dims:
                 # new unindexed dimension
                 var_chunks.append((template.sizes[dim],))