Repository: dask/dask

Issue:
dask dataframe read_csv + filtering fails when using comment kwarg
**What happened**:  A key error on a queried column that exists in the dataframe.  Specific to reading a csv using dask read_csv and specifying a comment kwarg to account for a header in the csv.

**What you expected to happen**:  The query should return a dataframe with the filter on column, without the keyerror occuring.

**Minimal Complete Verifiable Example**:

```
In [1]: import numpy as np

In [2]: import pandas as pd


In [3]: import dask.dataframe as dd

In [4]:

In [4]: testdata = pd.DataFrame(np.random.random((5000000, 5)), columns=list('abcde'))

In [5]: testdata
Out[5]:
                a         b         c         d         e
0        0.714239  0.197572  0.317503  0.676065  0.226038
1        0.164100  0.027354  0.574734  0.887295  0.171373
2        0.437257  0.005123  0.039310  0.634022  0.276112
3        0.808366  0.259773  0.358868  0.972922  0.400185
4        0.622027  0.764787  0.677192  0.257554  0.074227
...           ...       ...       ...       ...       ...
4999995  0.153502  0.417704  0.010464  0.715816  0.395864
4999996  0.412132  0.758808  0.191512  0.015014  0.712972
4999997  0.944817  0.666780  0.473993  0.955058  0.926175
4999998  0.048076  0.676805  0.246336  0.159064  0.653335
4999999  0.431097  0.003650  0.201192  0.330820  0.614423

[5000000 rows x 5 columns]

In [6]:

In [6]: with open('testfile.csv', 'w') as f:
   ...:     f.write('#Header\n')
   ...:     testdata.to_csv(f, index=False, sep='\t')
   ...:

In [7]: pd.read_csv('testfile.csv', sep='\t', comment='#', nrows=5)
Out[7]:
          a         b         c         d         e
0  0.714239  0.197572  0.317503  0.676065  0.226038
1  0.164100  0.027354  0.574734  0.887295  0.171373
2  0.437257  0.005123  0.039310  0.634022  0.276112
3  0.808366  0.259773  0.358868  0.972922  0.400185
4  0.622027  0.764787  0.677192  0.257554  0.074227

In [8]: ddf = dd.read_csv('testfile.csv', sep='\t', comment='#')

In [9]: ddf[ddf['a'] > 0.5].compute()

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3360             try:
-> 3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'a'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
<ipython-input-9-f8c106bd3221> in <module>
----> 1 ddf[ddf['a'] > 0.5].compute()

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/base.py in compute(self, **kwargs)
    286         dask.base.compute
    287         """
--> 288         (result,) = compute(self, traverse=False, **kwargs)
    289         return result
    290

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/base.py in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    569         postcomputes.append(x.__dask_postcompute__())
    570
--> 571     results = schedule(dsk, keys, **kwargs)
    572     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    573

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     85         get_id=_thread_get_id,
     86         pack_exception=pack_exception,
---> 87         **kwargs,
     88     )
     89

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    505                             _execute_task(task, data)  # Re-execute locally
    506                         else:
--> 507                             raise_exception(exc, tb)
    508                     res, worker_id = loads(res_info)
    509                     state["cache"][key] = res

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/local.py in reraise(exc, tb)
    313     if exc.__traceback__ is not tb:
    314         raise exc.with_traceback(tb)
--> 315     raise exc
    316
    317

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    218     try:
    219         task, data = loads(task_info)
--> 220         result = _execute_task(task, data)
    221         id = get_id()
    222         result = dumps((result, id))

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/optimization.py in __call__(self, *args)
    967         if not len(args) == len(self.inkeys):
    968             raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
--> 969         return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
    970
    971     def __reduce__(self):

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py 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)

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in <genexpr>(.0)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in <genexpr>(.0)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    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

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/pandas/core/frame.py in __getitem__(self, key)
   3456             if self.columns.nlevels > 1:
   3457                 return self._getitem_multilevel(key)
-> 3458             indexer = self.columns.get_loc(key)
   3459             if is_integer(indexer):
   3460                 indexer = [indexer]

/juno/work/shah/users/mcphera1/repos/signaturesanalysis/notebooks/venv/lib/python3.7/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:
-> 3363                 raise KeyError(key) from err
   3364
   3365         if is_scalar(key) and isna(key) and not self.hasnans:

KeyError: 'a'
```

**Environment**:

- Dask version: '2021.11.2'
- Python version: 3.7.3 (default, Mar 27 2019, 22:11:17)
- Pandas version: '1.3.4'
- Operating System: Linux
- Install method (conda, pip, source): pip



Discussion:
Thanks for writing this up and sorry for the slow response. I am looking at the codebase and it seems like the comment kwarg should work ok for all cases except when the line is the first in the file. In that case the header information gets kind of messed up. I'll have a PR open shortly with a fix.

Resolving patch:
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
--- a/dask/dataframe/io/csv.py
+++ b/dask/dataframe/io/csv.py
@@ -572,7 +572,23 @@ def read_pandas(
     names = kwargs.get("names", None)
     header = kwargs.get("header", "infer" if names is None else None)
     need = 1 if header is None else 2
-    parts = b_sample.split(b_lineterminator, lastskiprow + need)
+
+    if kwargs.get("comment"):
+        # if comment is provided, step through lines of b_sample and strip out comments
+        parts = []
+        for part in b_sample.split(b_lineterminator):
+            split_comment = part.decode().split(kwargs.get("comment"))
+            if len(split_comment) > 1:
+                # if line starts with comment, don't include that line in parts.
+                if len(split_comment[0]) > 0:
+                    parts.append(split_comment[0].strip().encode())
+            else:
+                parts.append(part)
+            if len(parts) > need:
+                break
+    else:
+        parts = b_sample.split(b_lineterminator, lastskiprow + need)
+
     # If the last partition is empty, don't count it
     nparts = 0 if not parts else len(parts) - int(not parts[-1])
 


Test patch:
diff --git a/dask/dataframe/io/tests/test_csv.py b/dask/dataframe/io/tests/test_csv.py
--- a/dask/dataframe/io/tests/test_csv.py
+++ b/dask/dataframe/io/tests/test_csv.py
@@ -249,6 +249,23 @@ def test_skiprows(dd_read, pd_read, files):
         assert_eq(df, expected_df, check_dtype=False)
 
 
+@pytest.mark.parametrize(
+    "dd_read,pd_read,files",
+    [(dd.read_csv, pd.read_csv, csv_files), (dd.read_table, pd.read_table, tsv_files)],
+)
+def test_comment(dd_read, pd_read, files):
+    files = {
+        name: comment_header
+        + b"\n"
+        + content.replace(b"\n", b"  # just some comment\n", 1)
+        for name, content in files.items()
+    }
+    with filetexts(files, mode="b"):
+        df = dd_read("2014-01-*.csv", comment="#")
+        expected_df = pd.concat([pd_read(n, comment="#") for n in sorted(files)])
+        assert_eq(df, expected_df, check_dtype=False)
+
+
 @pytest.mark.parametrize(
     "dd_read,pd_read,files",
     [(dd.read_csv, pd.read_csv, csv_files), (dd.read_table, pd.read_table, tsv_files)],