Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from collections import defaultdict
4from typing import (
5 TYPE_CHECKING,
6 Any,
7)
8import warnings
10import numpy as np
12from pandas._libs import (
13 lib,
14 parsers,
15)
16from pandas.compat._optional import import_optional_dependency
17from pandas.errors import DtypeWarning
18from pandas.util._exceptions import find_stack_level
20from pandas.core.dtypes.common import pandas_dtype
21from pandas.core.dtypes.concat import (
22 concat_compat,
23 union_categoricals,
24)
25from pandas.core.dtypes.dtypes import CategoricalDtype
27from pandas.core.arrays.arrow.array import ArrowExtensionArray
28from pandas.core.arrays.string_ import StringDtype
29from pandas.core.arrays.string_arrow import ArrowStringArray
30from pandas.core.indexes.api import ensure_index_from_sequences
32from pandas.io.common import (
33 dedup_names,
34 is_potential_multi_index,
35)
36from pandas.io.parsers.base_parser import (
37 ParserBase,
38 ParserError,
39 date_converter,
40 evaluate_callable_usecols,
41 is_index_col,
42 validate_parse_dates_presence,
43)
45if TYPE_CHECKING:
46 from collections.abc import (
47 Hashable,
48 Mapping,
49 Sequence,
50 )
52 from pandas._typing import (
53 AnyArrayLike,
54 ArrayLike,
55 DtypeArg,
56 DtypeObj,
57 ReadCsvBuffer,
58 SequenceT,
59 )
61 from pandas import (
62 Index,
63 MultiIndex,
64 )
67class CParserWrapper(ParserBase):
68 low_memory: bool
69 _reader: parsers.TextReader
70 _exhausted: bool
71 # When False, read() leaves deferred string columns as raw pending
72 # handles for the caller to materialize (one chunked ExtensionArray per
73 # column, e.g. post-gather in a parallel read).
74 wrap_deferred: bool = True
76 def __init__(self, src: ReadCsvBuffer[str], **kwds) -> None:
77 super().__init__(kwds)
78 self.kwds = kwds
79 kwds = kwds.copy()
81 self.low_memory = kwds.pop("low_memory", False)
82 self._exhausted = False
84 # #2442
85 kwds["allow_leading_cols"] = self.index_col is not False
87 # GH20529, validate usecol arg before TextReader
88 kwds["usecols"] = self.usecols
90 # Have to pass int, would break tests using TextReader directly otherwise :(
91 kwds["on_bad_lines"] = self.on_bad_lines.value
93 for key in (
94 "storage_options",
95 "encoding",
96 "memory_map",
97 "compression",
98 ):
99 kwds.pop(key, None)
101 kwds["dtype"] = ensure_dtype_objs(kwds.get("dtype", None))
102 if "dtype_backend" not in kwds or kwds["dtype_backend"] is lib.no_default:
103 kwds["dtype_backend"] = "numpy"
104 if kwds["dtype_backend"] == "pyarrow":
105 # Fail here loudly instead of in cython after reading
106 import_optional_dependency("pyarrow")
107 self._reader = parsers.TextReader(src, **kwds)
108 # Let the pyarrow string fast path return raw pending-column handles;
109 # read() wraps them into one ExtensionArray per column at the end.
110 self._reader.defer_pa_wrap = True
112 self.unnamed_cols = self._reader.unnamed_cols
114 passed_names = self.names is None
116 if self._reader.header is None:
117 self.names = None
118 else:
119 (
120 self.names,
121 self.index_names,
122 self.col_names,
123 passed_names,
124 ) = self._extract_multi_indexer_columns(
125 self._reader.header,
126 self.index_names,
127 passed_names,
128 )
130 if self.names is None:
131 self.names = list(range(self._reader.table_width))
133 # gh-9755
134 #
135 # need to set orig_names here first
136 # so that proper indexing can be done
137 # with _set_noconvert_columns
138 #
139 # once names has been filtered, we will
140 # then set orig_names again to names
141 self.orig_names = self.names[:]
143 if self.usecols:
144 usecols = evaluate_callable_usecols(self.usecols, self.orig_names)
146 # GH 14671
147 # assert for mypy, orig_names is List or None, None would error in issubset
148 assert self.orig_names is not None
149 if self.usecols_dtype == "string" and not set(usecols).issubset(
150 self.orig_names
151 ):
152 self._validate_usecols_names(usecols, self.orig_names)
154 if len(self.names) > len(usecols):
155 self.names = [
156 n
157 for i, n in enumerate(self.names)
158 if (i in usecols or n in usecols)
159 ]
161 if len(self.names) < len(usecols):
162 self._validate_usecols_names(
163 usecols,
164 self.names,
165 )
167 validate_parse_dates_presence(self.parse_dates, self.names)
168 self._set_noconvert_columns()
170 self.orig_names = self.names
172 if self._reader.leading_cols == 0 and is_index_col(self.index_col):
173 (
174 index_names,
175 self.names,
176 self.index_col,
177 ) = self._clean_index_names(
178 self.names,
179 self.index_col,
180 )
182 if self.index_names is None:
183 self.index_names = index_names
185 if self._reader.header is None and not passed_names:
186 assert self.index_names is not None
187 self.index_names = [None] * len(self.index_names)
189 self._implicit_index = self._reader.leading_cols > 0
191 def close(self) -> None:
192 # close handles opened by C parser
193 try:
194 self._reader.close()
195 except ValueError:
196 pass
198 def _set_noconvert_columns(self) -> None:
199 """
200 Set the columns that should not undergo dtype conversions.
202 Currently, any column that is involved with date parsing will not
203 undergo such conversions.
204 """
205 assert self.orig_names is not None
206 # error: Cannot determine type of 'names'
208 # much faster than using orig_names.index(x) xref GH#44106
209 names_dict = {x: i for i, x in enumerate(self.orig_names)} # pyright: ignore[reportOptionalIterable]
210 col_indices = [names_dict[x] for x in self.names] # pyright: ignore[reportOptionalIterable]
211 noconvert_columns = self._set_noconvert_dtype_columns(
212 col_indices,
213 self.names,
214 )
215 for col in noconvert_columns:
216 self._reader.set_noconvert(col)
218 # Mark parse_dates columns so the C parser can try to emit datetime64
219 # directly, skipping the Python-str ndarray round-trip. In low_memory
220 # mode the C parser keeps raw-byte receipts per chunk so a fallback in
221 # any chunk restores the whole column to the object-string path.
222 # With a list parse_dates, noconvert_columns is exactly the resolved
223 # parse_dates targets (integer colspecs usecols-relative, matching
224 # _do_date_conversions).
225 # Not with implicit-index files (leading_cols > 0): noconvert indices
226 # are not shifted by leading_cols, so the fastpath gate in
227 # _convert_tokens cannot line up with them.
228 if isinstance(self.parse_dates, list) and self._reader.leading_cols == 0:
229 for col_idx in noconvert_columns:
230 name = self.orig_names[col_idx]
231 require_consistent = self._parse_dates_fastpath_strictness(name)
232 if require_consistent is not None:
233 self._reader.set_datetime_convert(col_idx, require_consistent)
235 def _parse_dates_fastpath_strictness(self, name: Hashable) -> bool | None:
236 """
237 Decide whether a parse_dates column is a candidate for the direct
238 char-buffer -> datetime64 fastpath, returning None if not, otherwise
239 whether the fastpath must require a single consistent format.
241 With date_format=None, to_datetime infers one format from the first
242 value and rejects rows that deviate, so the fastpath must do the same;
243 with date_format="ISO8601" mixed ISO8601 layouts are allowed.
244 """
245 if self.dayfirst:
246 return None
247 date_format = self.date_format
248 if isinstance(date_format, dict):
249 date_format = date_format.get(name)
250 if date_format is None:
251 return True
252 if date_format == "ISO8601":
253 return False
254 # If a specific format is supplied, let the existing path handle it.
255 # A future extension could detect ISO-shaped formats here.
256 return None
258 def _low_memory_column_labels(
259 self, chunks: list[dict[int, ArrayLike]]
260 ) -> dict[int, Hashable]:
261 """
262 Map the field positions ``read_low_memory`` keys its chunks by to labels.
264 The positions are those of the source row, so they skip whatever
265 ``usecols`` dropped and start past a leading implicit index column.
266 ``orig_names`` holds exactly the named columns that remain, in the same
267 order, which is how ``read`` renames these keys further down. A leading
268 implicit index column has no name of its own and maps to ``None``.
269 """
270 assert self.orig_names is not None
271 positions = sorted(chunks[0])
272 named = dict(
273 zip(
274 positions[self._reader.leading_cols :],
275 self.orig_names,
276 strict=False,
277 )
278 )
279 return {position: named.get(position) for position in positions}
281 def read(
282 self,
283 nrows: int | None = None,
284 ) -> tuple[
285 Index | MultiIndex | None,
286 Sequence[Hashable] | MultiIndex,
287 Mapping[Hashable, AnyArrayLike],
288 ]:
289 index: Index | MultiIndex | None
290 column_names: Sequence[Hashable] | MultiIndex
291 if self._exhausted:
292 # Exhausting the reader closed it, so calling into the C reader
293 # again would raise instead of signalling the end of the data.
294 raise StopIteration
295 try:
296 if self.low_memory:
297 chunks = self._reader.read_low_memory(nrows)
298 # destructive to chunks
299 data = _concatenate_chunks(
300 chunks, self._low_memory_column_labels(chunks)
301 )
302 # GH#56044 category-dtype inference is deferred until after
303 # concatenation so all chunks agree
304 self._reader._maybe_infer_categoricals(data)
305 else:
306 data = self._reader.read(nrows)
307 if self.wrap_deferred:
308 data = {
309 key: _wrap_deferred_pa(values) for key, values in data.items()
310 }
311 except StopIteration:
312 if self._first_chunk:
313 self._first_chunk = False
314 # assert for mypy, orig_names is List or None, None would error in
315 # list(...) in dedup_names
316 assert self.orig_names is not None
317 names = dedup_names(
318 self.orig_names,
319 is_potential_multi_index(self.orig_names, self.index_col),
320 )
321 index, columns, col_dict = self._get_empty_meta(
322 names,
323 dtype=self.dtype,
324 )
325 # error: Incompatible types in assignment (expression has type
326 # "list[Hashable] | MultiIndex", variable has type "list[Hashable]")
327 columns = self._maybe_make_multi_index_columns( # type: ignore[assignment]
328 columns, self.col_names
329 )
331 columns = _filter_usecols(self.usecols, columns)
332 columns_set = set(columns)
334 col_dict = {k: v for k, v in col_dict.items() if k in columns_set}
336 return index, columns, col_dict
338 else:
339 self._exhausted = True
340 self.close()
341 raise
343 # Done with first read, next time raise StopIteration
344 self._first_chunk = False
346 names = self.names
348 if self._reader.leading_cols:
349 # implicit index, no index names
350 arrays = []
352 if self.index_col and self._reader.leading_cols != len(self.index_col):
353 raise ParserError(
354 "Could not construct index. Requested to use "
355 f"{len(self.index_col)} number of columns, but "
356 f"{self._reader.leading_cols} left to parse."
357 )
359 for i in range(self._reader.leading_cols):
360 if self.index_col is None:
361 values = data.pop(i)
362 else:
363 values = data.pop(self.index_col[i])
365 if self._should_parse_dates(i):
366 values = date_converter(
367 values,
368 col=(
369 self.index_names[i]
370 if self.index_names is not None
371 else None
372 ),
373 dayfirst=self.dayfirst,
374 cache_dates=self.cache_dates,
375 date_format=self.date_format,
376 )
377 arrays.append(values)
379 index = ensure_index_from_sequences(arrays)
381 names = _filter_usecols(self.usecols, names)
383 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
385 # rename dict keys
386 data_tups = sorted(data.items())
387 data = {k: v for k, (i, v) in zip(names, data_tups, strict=True)}
389 date_data = self._do_date_conversions(names, data)
391 # maybe create a mi on the columns
392 column_names = self._maybe_make_multi_index_columns(names, self.col_names)
394 else:
395 # rename dict keys
396 data_tups = sorted(data.items())
398 # ugh, mutation
400 # assert for mypy, orig_names is List or None, None would error in list(...)
401 assert self.orig_names is not None
402 names = list(self.orig_names)
403 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
405 names = _filter_usecols(self.usecols, names)
407 # columns as list
408 alldata = [x[1] for x in data_tups]
409 if self.usecols is None:
410 self._check_data_length(names, alldata)
412 data = {k: v for k, (i, v) in zip(names, data_tups, strict=False)}
414 date_data = self._do_date_conversions(names, data)
415 index, column_names = self._make_index(alldata, names)
417 return index, column_names, date_data
420def _filter_usecols(usecols, names: SequenceT) -> SequenceT | list[Hashable]:
421 # hackish
422 usecols = evaluate_callable_usecols(usecols, names)
423 if usecols is not None and len(names) != len(usecols):
424 return [name for i, name in enumerate(names) if i in usecols or name in usecols]
425 return names
428def _pa_arrays_to_ea(arrs) -> ArrayLike:
429 """
430 Build one ExtensionArray from the pending string columns returned by
431 the deferred TextReader string fast path (``defer_pa_wrap``).
433 The fast path emits ``large_string`` for the default string dtype and
434 ``string`` for ``dtype_backend="pyarrow"``, so the arrow type determines
435 the target ExtensionArray.
436 """
437 import pyarrow as pa
439 chunked = pa.chunked_array([pending.materialize() for pending in arrs])
440 if pa.types.is_large_string(chunked.type):
441 return ArrowStringArray(chunked, dtype=StringDtype(na_value=np.nan))
442 return ArrowExtensionArray(chunked)
445def _wrap_deferred_pa(values):
446 """Wrap a pending string column from the deferred fast path; pass through
447 everything else."""
448 if isinstance(values, parsers._PendingStringColumn):
449 return _pa_arrays_to_ea([values])
450 return values
453def _harmonize_empty_categories(arrs: list[ArrayLike]) -> list[ArrayLike]:
454 """
455 Give chunks that inferred no categories the others' categories dtype.
457 GH#56044 a chunk that is entirely NA yields empty object-dtype categories
458 where a populated chunk yields string ones, which union_categoricals
459 rejects. When every chunk is empty there is nothing to match, so an
460 all-NA column keeps its object categories.
461 """
462 populated = {arr.categories.dtype for arr in arrs if len(arr.categories)} # type: ignore[union-attr]
463 if len(populated) != 1:
464 return arrs
465 cat_dtype = populated.pop()
466 return [
467 arr
468 if len(arr.categories) # type: ignore[union-attr]
469 else arr.set_categories(arr.categories.astype(cat_dtype)) # type: ignore[union-attr]
470 for arr in arrs
471 ]
474def _concatenate_chunks(
475 chunks: list[dict[int, ArrayLike]],
476 column_names: Sequence[Hashable] | Mapping[int, Hashable],
477 warn_mixed: bool = True,
478) -> dict[Any, ArrayLike]:
479 """
480 Concatenate chunks of data read with low_memory=True.
482 The tricky part is handling Categoricals, where different chunks
483 may have different inferred categories.
485 ``column_names`` names the keys of ``chunks`` for the mixed-dtype warning,
486 so it has to be indexable by whatever those keys are: the low_memory reader
487 keys its chunks by field position and passes a mapping keyed the same way,
488 while the parallel reader keys its own by name.
490 ``warn_mixed=False`` suppresses the mixed-dtype warning; the parallel
491 reader gathers byte-range chunks (not low_memory row chunks), so a
492 column's cross-chunk dtype has already been reconciled by the caller.
493 """
494 names = list(chunks[0].keys())
495 warning_columns = []
497 result: dict[Any, ArrayLike] = {}
498 for name in names:
499 arrs = [chunk.pop(name) for chunk in chunks]
501 # The homogeneous case (every chunk a pending string column from the
502 # deferred fast path) combines zero-copy into a single chunked
503 # ExtensionArray; mixed cases (e.g. earlier chunks inferred numeric)
504 # wrap each pending column and fall through to the regular concat
505 # below.
506 if isinstance(arrs[0], parsers._PendingStringColumn):
507 if all(isinstance(arr, parsers._PendingStringColumn) for arr in arrs):
508 result[name] = _pa_arrays_to_ea(arrs)
509 continue
510 arrs = [_wrap_deferred_pa(arr) for arr in arrs]
511 elif any(isinstance(arr, parsers._PendingStringColumn) for arr in arrs):
512 arrs = [_wrap_deferred_pa(arr) for arr in arrs]
514 # Check each arr for consistent types.
515 dtypes = {a.dtype for a in arrs}
516 non_cat_dtypes = {x for x in dtypes if not isinstance(x, CategoricalDtype)}
518 dtype = dtypes.pop()
519 if isinstance(dtype, CategoricalDtype):
520 arrs = _harmonize_empty_categories(arrs)
521 result[name] = union_categoricals(arrs, sort_categories=False) # type: ignore[arg-type]
522 else:
523 result[name] = concat_compat(arrs)
524 if len(non_cat_dtypes) > 1 and result[name].dtype == np.dtype(object):
525 warning_columns.append((name, column_names[name]))
527 if warning_columns and warn_mixed:
528 warning_names = ", ".join(
529 [
530 f"{position}: {label}" if label is not None else f"{position}"
531 for position, label in warning_columns
532 ]
533 )
534 warning_message = " ".join(
535 [
536 f"Columns ({warning_names}) have mixed types. "
537 f"Specify dtype option on import or set low_memory=False."
538 ]
539 )
540 warnings.warn(warning_message, DtypeWarning, stacklevel=find_stack_level())
541 return result
544def ensure_dtype_objs(
545 dtype: DtypeArg | dict[Hashable, DtypeArg] | None,
546) -> DtypeObj | dict[Hashable, DtypeObj] | None:
547 """
548 Ensure we have either None, a dtype object, or a dictionary mapping to
549 dtype objects.
550 """
551 if isinstance(dtype, defaultdict):
552 # "None" not callable [misc]
553 default_dtype = pandas_dtype(dtype.default_factory()) # type: ignore[misc]
554 dtype_converted: defaultdict[Hashable, DtypeObj] = defaultdict(
555 lambda: default_dtype
556 )
557 for key in dtype.keys():
558 dtype_converted[key] = pandas_dtype(dtype[key])
559 return dtype_converted
560 elif isinstance(dtype, dict):
561 return {k: pandas_dtype(dtype[k]) for k in dtype}
562 elif dtype is not None:
563 return pandas_dtype(dtype)
564 return dtype