Repository: dask/dask

Issue:
KeyError tracing to optimize_dataframe_getitem
The following code from a [StackOverflow question](https://stackoverflow.com/questions/67645088) is raising an error with dask `2021.5.0` but not with `2021.4.1`:

```python
import dask
import dask.dataframe as dd
df = dask.datasets.timeseries()
df2 = df[df.y > 0]
df2.compute()
```

Traceback:
<details>

```pytb
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-8-a8eb713bc07c> in <module>
----> 1 df2.compute()

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/base.py in compute(self, **kwargs)
    283         dask.base.compute
    284         """
--> 285         (result,) = compute(self, traverse=False, **kwargs)
    286         return result
    287

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/base.py in compute(*args, **kwargs)
    559     )
    560
--> 561     dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
    562     keys, postcomputes = [], []
    563     for x in collections:

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/base.py in collections_to_dsk(collections, optimize_graph, optimizations, **kwargs)
    335         for opt, val in groups.items():
    336             dsk, keys = _extract_graph_and_keys(val)
--> 337             dsk = opt(dsk, keys, **kwargs)
    338
    339             for opt in optimizations:

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/dataframe/optimize.py in optimize(dsk, keys, **kwargs)
     20     else:
     21         # Perform Blockwise optimizations for HLG input
---> 22         dsk = optimize_dataframe_getitem(dsk, keys=keys)
     23         dsk = optimize_blockwise(dsk, keys=keys)
     24         dsk = fuse_roots(dsk, keys=keys)

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/dataframe/optimize.py in optimize_dataframe_getitem(dsk, keys)
    103         # Project columns and update blocks
    104         old = layers[k]
--> 105         new = old.project_columns(columns)[0]
    106         if new.name != old.name:
    107             columns = list(columns)

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/layers.py in project_columns(self, columns)
    941             # Apply column projection in IO function
    942             try:
--> 943                 io_func = self.io_func.project_columns(list(columns))
    944             except AttributeError:
    945                 io_func = self.io_func

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/dataframe/io/demo.py in project_columns(self, columns)
     87         func = copy.deepcopy(self)
     88         func.columns = columns
---> 89         func.dtypes = {c: self.dtypes[c] for c in columns}
     90         return func
     91

~/miniconda/envs/dask2/lib/python3.9/site-packages/dask/dataframe/io/demo.py in <dictcomp>(.0)
     87         func = copy.deepcopy(self)
     88         func.columns = columns
---> 89         func.dtypes = {c: self.dtypes[c] for c in columns}
     90         return func
     91

KeyError: 'gt-846fc0648f9e5e57119c74d81b65378c'
```
</details>

**Environment**:

- Dask version: 2021.5.0  (this error is not raised with dask 2021.4.1)
- Python version: 3.9.4
- Operating System: MacOS BigSur
- Install method (conda, pip, source): conda



Discussion:

Resolving patch:
diff --git a/dask/dataframe/optimize.py b/dask/dataframe/optimize.py
--- a/dask/dataframe/optimize.py
+++ b/dask/dataframe/optimize.py
@@ -63,11 +63,19 @@ def optimize_dataframe_getitem(dsk, keys):
 
     layers = dsk.layers.copy()
     dependencies = dsk.dependencies.copy()
-
     for k in dataframe_blockwise:
         columns = set()
         update_blocks = {}
 
+        if any(layers[k].name == x[0] for x in keys if isinstance(x, tuple)):
+            # ... but bail on the optimization if the dataframe_blockwise layer is in
+            # the requested keys, because we cannot change the name anymore.
+            # These keys are structured like [('getitem-<token>', 0), ...]
+            # so we check for the first item of the tuple.
+            # See https://github.com/dask/dask/issues/5893
+            return dsk
+
+        column_projection = True
         for dep in dsk.dependents[k]:
             block = dsk.layers[dep]
 
@@ -84,49 +92,47 @@ def optimize_dataframe_getitem(dsk, keys):
                 # ... where this value is __getitem__...
                 return dsk
 
-            if any(layers[k].name == x[0] for x in keys if isinstance(x, tuple)):
-                # ... but bail on the optimization if the dataframe_blockwise layer is in
-                # the requested keys, because we cannot change the name anymore.
-                # These keys are structured like [('getitem-<token>', 0), ...]
-                # so we check for the first item of the tuple.
-                # See https://github.com/dask/dask/issues/5893
-                return dsk
-
             block_columns = block.indices[1][0]
-            if isinstance(block_columns, str) or np.issubdtype(
-                type(block_columns), np.integer
-            ):
+            if isinstance(block_columns, str):
+                if block_columns in layers.keys():
+                    # Not a column selection if the getitem
+                    # key is a collection key
+                    column_projection = False
+                    break
+                block_columns = [block_columns]
+            elif np.issubdtype(type(block_columns), np.integer):
                 block_columns = [block_columns]
             columns |= set(block_columns)
             update_blocks[dep] = block
 
         # Project columns and update blocks
-        old = layers[k]
-        new = old.project_columns(columns)[0]
-        if new.name != old.name:
-            columns = list(columns)
-            assert len(update_blocks)
-            for block_key, block in update_blocks.items():
-                # (('read-parquet-old', (.,)), ( ... )) ->
-                # (('read-parquet-new', (.,)), ( ... ))
-                new_indices = ((new.name, block.indices[0][1]), block.indices[1])
-                numblocks = {new.name: block.numblocks[old.name]}
-                new_block = Blockwise(
-                    block.output,
-                    block.output_indices,
-                    block.dsk,
-                    new_indices,
-                    numblocks,
-                    block.concatenate,
-                    block.new_axes,
-                )
-                layers[block_key] = new_block
-                dependencies[block_key] = {new.name}
-            dependencies[new.name] = dependencies.pop(k)
-
-        layers[new.name] = new
-        if new.name != old.name:
-            del layers[old.name]
+        if column_projection:
+            old = layers[k]
+            new = old.project_columns(columns)[0]
+            if new.name != old.name:
+                columns = list(columns)
+                assert len(update_blocks)
+                for block_key, block in update_blocks.items():
+                    # (('read-parquet-old', (.,)), ( ... )) ->
+                    # (('read-parquet-new', (.,)), ( ... ))
+                    new_indices = ((new.name, block.indices[0][1]), block.indices[1])
+                    numblocks = {new.name: block.numblocks[old.name]}
+                    new_block = Blockwise(
+                        block.output,
+                        block.output_indices,
+                        block.dsk,
+                        new_indices,
+                        numblocks,
+                        block.concatenate,
+                        block.new_axes,
+                    )
+                    layers[block_key] = new_block
+                    dependencies[block_key] = {new.name}
+                dependencies[new.name] = dependencies.pop(k)
+
+            layers[new.name] = new
+            if new.name != old.name:
+                del layers[old.name]
 
     new_hlg = HighLevelGraph(layers, dependencies)
     return new_hlg


Test patch:
diff --git a/dask/dataframe/io/tests/test_demo.py b/dask/dataframe/io/tests/test_demo.py
--- a/dask/dataframe/io/tests/test_demo.py
+++ b/dask/dataframe/io/tests/test_demo.py
@@ -167,3 +167,13 @@ def test_make_timeseries_fancy_keywords():
 
     assert 100 < aa <= 10000000
     assert 1 < bb <= 100
+
+
+def test_make_timeseries_getitem_compute():
+    # See https://github.com/dask/dask/issues/7692
+
+    df = dd.demo.make_timeseries()
+    df2 = df[df.y > 0]
+    df3 = df2.compute()
+    assert df3["y"].min() > 0
+    assert list(df.columns) == list(df3.columns)