fix dask-sql bug caused by blockwise fusion (#8989)

rjzamora · web-flow · commit e429efae3d65 · 2022-04-28T09:23:30.000-05:00
diff --git a/dask/blockwise.py b/dask/blockwise.py
@@ -185,7 +185,19 @@ def produces_keys(self) -> bool:
         return self._produces_keys
 
     def __getitem__(self, idx: tuple[int, ...]) -> Any:
-        return self.mapping[idx]
+        try:
+            return self.mapping[idx]
+        except KeyError as err:
+            # If a DataFrame collection was converted
+            # to an Array collection, the dimesion of
+            # `idx` may not agree with the keys in
+            # `self.mapping`. In this case, we can
+            # use `self.numblocks` to check for a key
+            # match in the leading elements of `idx`
+            flat_idx = idx[: len(self.numblocks)]
+            if flat_idx in self.mapping:
+                return self.mapping[flat_idx]
+            raise err
 
     def __dask_distributed_pack__(
         self, required_indices: tuple | list[tuple[int, ...]] | None = None
@@ -204,7 +216,7 @@ def __dask_distributed_pack__(
             },
             "numblocks": self.numblocks,
             "produces_tasks": self.produces_tasks,
-            "produces_keys": self._produces_keys,
+            "produces_keys": self.produces_keys,
         }
 
     @classmethod
diff --git a/dask/dataframe/io/tests/test_io.py b/dask/dataframe/io/tests/test_io.py
@@ -700,6 +700,30 @@ def test_from_delayed_optimize_fusion():
     assert len(optimize(ddf.dask, ddf.__dask_keys__()).layers) == 1
 
 
+def test_from_delayed_to_dask_array():
+    # Check that `from_delayed`` can be followed
+    # by `to_dask_array` without breaking
+    # optimization behavior
+    # See: https://github.com/dask-contrib/dask-sql/issues/497
+    from dask.blockwise import optimize_blockwise
+
+    dfs = [delayed(pd.DataFrame)(np.ones((3, 2))) for i in range(3)]
+    ddf = dd.from_delayed(dfs)
+    arr = ddf.to_dask_array()
+
+    # If we optimize this graph without calling
+    # `fuse_roots`, the underlying `BlockwiseDep`
+    # `mapping` keys will be 1-D (e.g. `(4,)`),
+    # while the collection keys will be 2-D
+    # (e.g. `(4, 0)`)
+    keys = [k[0] for k in arr.__dask_keys__()]
+    dsk = optimize_blockwise(arr.dask, keys=keys)
+    dsk.cull(keys)
+
+    result = arr.compute()
+    assert result.shape == (9, 2)
+
+
 def test_from_delayed_preserves_hlgs():
     df = pd.DataFrame(data=np.random.normal(size=(10, 4)), columns=list("abcd"))
     parts = [df.iloc[:1], df.iloc[1:3], df.iloc[3:6], df.iloc[6:10]]
