Repository: dask/dask
Pull request: Fix project CSV columns when selecting
- [x] Closes #9518
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


Problem statement:
KeyError when supplying multiple paths to dask.dataframe.read_csv() AND including a path column AND using Dask compute
When using Dask compute in a Dask dataframe with `include_path_column=True` a `KeyError` is returned.

I would have expected that columns in lazy Dask collection will be the same in their memory equivalent. But unsure if that is the problem.

Example 1:

```py
import os
import pandas as pd
import dask.dataframe as dd

d = {'col1': [i for i in range(0, 100)], 'col2': [i for i in range(100, 200)]}
df = pd.DataFrame(data=d)

temp_path = "temp/"
os.makedirs(temp_path, exist_ok = False)
for i in range(6):
    df.to_csv(f"{temp_path}file_{i}.csv", index=False)
    
ddf = dd.read_csv(temp_path + "*.csv", include_path_column=True)

df.groupby("path").col1.mean().compute() # <-- this produces the error
```

Traceback 1:

<details>

```py
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In [4], line 1
----> 1 df.groupby("path").col1.mean().compute()

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/frame.py:8392, in DataFrame.groupby(self, by, axis, level, as_index, sort, group_keys, squeeze, observed, dropna)
   8389     raise TypeError("You have to supply one of 'by' and 'level'")
   8390 axis = self._get_axis_number(axis)
-> 8392 return DataFrameGroupBy(
   8393     obj=self,
   8394     keys=by,
   8395     axis=axis,
   8396     level=level,
   8397     as_index=as_index,
   8398     sort=sort,
   8399     group_keys=group_keys,
   8400     squeeze=squeeze,
   8401     observed=observed,
   8402     dropna=dropna,
   8403 )

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/groupby/groupby.py:959, in GroupBy.__init__(self, obj, keys, axis, level, grouper, exclusions, selection, as_index, sort, group_keys, squeeze, observed, mutated, dropna)
    956 if grouper is None:
    957     from pandas.core.groupby.grouper import get_grouper
--> 959     grouper, exclusions, obj = get_grouper(
    960         obj,
    961         keys,
    962         axis=axis,
    963         level=level,
    964         sort=sort,
    965         observed=observed,
    966         mutated=self.mutated,
    967         dropna=self.dropna,
    968     )
    970 self.obj = obj
    971 self.axis = obj._get_axis_number(axis)

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/groupby/grouper.py:889, in get_grouper(obj, key, axis, level, sort, observed, mutated, validate, dropna)
    887         in_axis, level, gpr = False, gpr, None
    888     else:
--> 889         raise KeyError(gpr)
    890 elif isinstance(gpr, Grouper) and gpr.key is not None:
    891     # Add key to exclusions
    892     exclusions.add(gpr.key)

KeyError: 'path'
```

</details>

Example 2:

```py
ddf.col1.value_counts().compute() # <-- this also produces the error
```

Traceback 2:

<details>

```py
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/indexes/base.py:3800, in Index.get_loc(self, key, method, tolerance)
   3799 try:
-> 3800     return self._engine.get_loc(casted_key)
   3801 except KeyError as err:

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/_libs/index.pyx:138, in pandas._libs.index.IndexEngine.get_loc()

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/_libs/index.pyx:165, in pandas._libs.index.IndexEngine.get_loc()

File pandas/_libs/hashtable_class_helper.pxi:5745, in pandas._libs.hashtable.PyObjectHashTable.get_item()

File pandas/_libs/hashtable_class_helper.pxi:5753, in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'path'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In [3], line 1
----> 1 ddf.col1.value_counts().compute()

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/base.py:315, in DaskMethodsMixin.compute(self, **kwargs)
    291 def compute(self, **kwargs):
    292     """Compute this dask collection
    293 
    294     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    313     dask.base.compute
    314     """
--> 315     (result,) = compute(self, traverse=False, **kwargs)
    316     return result

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/base.py:600, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    597     keys.append(x.__dask_keys__())
    598     postcomputes.append(x.__dask_postcompute__())
--> 600 results = schedule(dsk, keys, **kwargs)
    601 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/threaded.py:89, in get(dsk, keys, cache, num_workers, pool, **kwargs)
     86     elif isinstance(pool, multiprocessing.pool.Pool):
     87         pool = MultiprocessingPoolExecutor(pool)
---> 89 results = get_async(
     90     pool.submit,
     91     pool._max_workers,
     92     dsk,
     93     keys,
     94     cache=cache,
     95     get_id=_thread_get_id,
     96     pack_exception=pack_exception,
     97     **kwargs,
     98 )
    100 # Cleanup pools associated to dead threads
    101 with pools_lock:

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/local.py:511, in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    509         _execute_task(task, data)  # Re-execute locally
    510     else:
--> 511         raise_exception(exc, tb)
    512 res, worker_id = loads(res_info)
    513 state["cache"][key] = res

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/local.py:319, in reraise(exc, tb)
    317 if exc.__traceback__ is not tb:
    318     raise exc.with_traceback(tb)
--> 319 raise exc

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/local.py:224, in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    222 try:
    223     task, data = loads(task_info)
--> 224     result = _execute_task(task, data)
    225     id = get_id()
    226     result = dumps((result, id))

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in _execute_task(arg, cache, dsk)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/optimization.py:990, in SubgraphCallable.__call__(self, *args)
    988 if not len(args) == len(self.inkeys):
    989     raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
--> 990 return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:149, in get(dsk, out, cache)
    147 for key in toposort(dsk):
    148     task = dsk[key]
--> 149     result = _execute_task(task, cache)
    150     cache[key] = result
    151 result = _execute_task(out, cache)

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in _execute_task(arg, cache, dsk)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in <genexpr>(.0)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:113, in _execute_task(arg, cache, dsk)
     83 """Do the actual work of collecting data and executing a function
     84 
     85 Examples
   (...)
    110 'foo'
    111 """
    112 if isinstance(arg, list):
--> 113     return [_execute_task(a, cache) for a in arg]
    114 elif istask(arg):
    115     func, args = arg[0], arg[1:]

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:113, in <listcomp>(.0)
     83 """Do the actual work of collecting data and executing a function
     84 
     85 Examples
   (...)
    110 'foo'
    111 """
    112 if isinstance(arg, list):
--> 113     return [_execute_task(a, cache) for a in arg]
    114 elif istask(arg):
    115     func, args = arg[0], arg[1:]

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in _execute_task(arg, cache, dsk)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in <genexpr>(.0)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/core.py:119, in _execute_task(arg, cache, dsk)
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/dask/dataframe/io/csv.py:98, in CSVFunctionWrapper.__call__(self, part)
     93 # Construct `path_info`
     94 if path is not None:
     95     path_info = (
     96         self.colname,
     97         path,
---> 98         sorted(list(self.head[self.colname].cat.categories)),
     99     )
    100 else:
    101     path_info = None

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/frame.py:3805, in DataFrame.__getitem__(self, key)
   3803 if self.columns.nlevels > 1:
   3804     return self._getitem_multilevel(key)
-> 3805 indexer = self.columns.get_loc(key)
   3806 if is_integer(indexer):
   3807     indexer = [indexer]

File /data/fishid-pfs-tools/env/lib/python3.8/site-packages/pandas/core/indexes/base.py:3802, in Index.get_loc(self, key, method, tolerance)
   3800     return self._engine.get_loc(casted_key)
   3801 except KeyError as err:
-> 3802     raise KeyError(key) from err
   3803 except TypeError:
   3804     # If we have a listlike key, _check_indexing_error will raise
   3805     #  InvalidIndexError. Otherwise we fall through and re-raise
   3806     #  the TypeError.
   3807     self._check_indexing_error(key)

KeyError: 'path'
```

</details>

**Environment**:

- Dask version: 2022.9.1
- Python version: 3.8.0
- Operating System: Ubuntu
- Install method (conda, pip, source): pip

Hints from issue discussion:

Commit message: add comment

Patch to review:
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
index 9276a06fc40..4d2f458a17d 100644
--- a/dask/dataframe/io/csv.py
+++ b/dask/dataframe/io/csv.py
@@ -63,7 +63,11 @@ def __init__(
 
     @property
     def columns(self):
-        return self.full_columns if self._columns is None else self._columns
+        if self._columns is None:
+            return self.full_columns
+        if self.colname:
+            return self._columns + [self.colname]
+        return self._columns
 
     def project_columns(self, columns):
         """Return a new CSVFunctionWrapper object with
@@ -73,11 +77,17 @@ def project_columns(self, columns):
         columns = [c for c in self.head.columns if c in columns]
         if columns == self.columns:
             return self
+        if self.colname and self.colname not in columns:
+            # when path-as-column is on, we must keep it at IO
+            # whatever the selection
+            head = self.head[columns + [self.colname]]
+        else:
+            head = self.head[columns]
         return CSVFunctionWrapper(
             self.full_columns,
             columns,
             self.colname,
-            self.head[columns],
+            head,
             self.header,
             self.reader,
             {c: self.dtypes[c] for c in columns},
diff --git a/dask/dataframe/io/tests/test_csv.py b/dask/dataframe/io/tests/test_csv.py
index 71cb458f53b..adf72b458d1 100644
--- a/dask/dataframe/io/tests/test_csv.py
+++ b/dask/dataframe/io/tests/test_csv.py
@@ -1770,3 +1770,19 @@ def test_csv_name_should_be_different_even_if_head_is_same(tmpdir):
     )
 
     assert new_df.dask.keys() != old_df.dask.keys()
+
+
+def test_select_with_path_optimization(tmpdir):
+    # https://github.com/dask/dask/issues/9518
+
+    d = {"col1": [i for i in range(0, 100)], "col2": [i for i in range(100, 200)]}
+    df = pd.DataFrame(data=d)
+
+    temp_path = str(tmpdir) + "/"
+    for i in range(6):
+        df.to_csv(f"{temp_path}file_{i}.csv", index=False)
+
+    ddf = dd.read_csv(temp_path + "*.csv", include_path_column=True)
+
+    assert_eq(ddf.col1, pd.concat([df.col1] * 6))
+    assert ddf.groupby("path").col1.mean().compute().iloc[0] == 49.5


Code review comments:
diff_hunk: @@ -1770,3 +1770,19 @@ def test_csv_name_should_be_different_even_if_head_is_same(tmpdir):
     )
 
     assert new_df.dask.keys() != old_df.dask.keys()
+
+
+def test_select_with_path_optimization(tmpdir):
+    # https://github.com/dask/dask/issues/9518
+
+    d = {"col1": [i for i in range(0, 100)], "col2": [i for i in range(100, 200)]}
+    df = pd.DataFrame(data=d)
+
+    temp_path = str(tmpdir) + "/"
+    for i in range(6):
+        df.to_csv(f"{temp_path}file_{i}.csv", index=False)
+
+    ddf = dd.read_csv(temp_path + "*.csv", include_path_column=True)
+
+    assert_eq(ddf.col1, pd.concat([df.col1] * 6))
+    assert ddf.groupby("path").col1.mean().compute().iloc[0] == 49.5
path: dask/dataframe/io/tests/test_csv.py
text: @user1:
Using `49.5` here seems correct, but fairly specific. Could this be updated to use `assert_eq` and compare `ddf.groupby("path").col1.mean()` to the expected pandas object? 

@author:
Not so easily, because we'd have to build an index based on the pathname.

@user1:
Yeah, fair point. We could still use `df.col1.mean()` though. On second thought, do we need this second assert? When I run this test on `main` I trigger the originally reported error on the `assert_eq(ddf.col1, pd.concat([df.col1] * 6))` line above this one. 

@author:
Happy to remove

@author:
```suggestion
```

diff_hunk: @@ -1770,3 +1770,19 @@ def test_csv_name_should_be_different_even_if_head_is_same(tmpdir):
     )
 
     assert new_df.dask.keys() != old_df.dask.keys()
+
+
+def test_select_with_path_optimization(tmpdir):
path: dask/dataframe/io/tests/test_csv.py
text: @user1:
Including `include_path_column` in the test name seems appropriate. Also I expect it might come in handy in the future when grepping the codebase 

```suggestion
def test_select_with_include_path_column(tmpdir):
```