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 TYPE_CHECKING
5import warnings
7import numpy as np
9from pandas._libs import (
10 lib,
11 parsers,
12)
13from pandas.compat._optional import import_optional_dependency
14from pandas.errors import DtypeWarning
15from pandas.util._exceptions import find_stack_level
17from pandas.core.dtypes.common import pandas_dtype
18from pandas.core.dtypes.concat import (
19 concat_compat,
20 union_categoricals,
21)
22from pandas.core.dtypes.dtypes import CategoricalDtype
24from pandas.core.indexes.api import ensure_index_from_sequences
26from pandas.io.common import (
27 dedup_names,
28 is_potential_multi_index,
29)
30from pandas.io.parsers.base_parser import (
31 ParserBase,
32 ParserError,
33 date_converter,
34 evaluate_callable_usecols,
35 is_index_col,
36 validate_parse_dates_presence,
37)
39if TYPE_CHECKING:
40 from collections.abc import (
41 Hashable,
42 Mapping,
43 Sequence,
44 )
46 from pandas._typing import (
47 AnyArrayLike,
48 ArrayLike,
49 DtypeArg,
50 DtypeObj,
51 ReadCsvBuffer,
52 SequenceT,
53 )
55 from pandas import (
56 Index,
57 MultiIndex,
58 )
61class CParserWrapper(ParserBase):
62 low_memory: bool
63 _reader: parsers.TextReader
65 def __init__(self, src: ReadCsvBuffer[str], **kwds) -> None:
66 super().__init__(kwds)
67 self.kwds = kwds
68 kwds = kwds.copy()
70 self.low_memory = kwds.pop("low_memory", False)
72 # #2442
73 kwds["allow_leading_cols"] = self.index_col is not False
75 # GH20529, validate usecol arg before TextReader
76 kwds["usecols"] = self.usecols
78 # Have to pass int, would break tests using TextReader directly otherwise :(
79 kwds["on_bad_lines"] = self.on_bad_lines.value
81 for key in (
82 "storage_options",
83 "encoding",
84 "memory_map",
85 "compression",
86 ):
87 kwds.pop(key, None)
89 kwds["dtype"] = ensure_dtype_objs(kwds.get("dtype", None))
90 if "dtype_backend" not in kwds or kwds["dtype_backend"] is lib.no_default:
91 kwds["dtype_backend"] = "numpy"
92 if kwds["dtype_backend"] == "pyarrow":
93 # Fail here loudly instead of in cython after reading
94 import_optional_dependency("pyarrow")
95 self._reader = parsers.TextReader(src, **kwds)
97 self.unnamed_cols = self._reader.unnamed_cols
99 passed_names = self.names is None
101 if self._reader.header is None:
102 self.names = None
103 else:
104 (
105 self.names,
106 self.index_names,
107 self.col_names,
108 passed_names,
109 ) = self._extract_multi_indexer_columns(
110 self._reader.header,
111 self.index_names,
112 passed_names,
113 )
115 if self.names is None:
116 self.names = list(range(self._reader.table_width))
118 # gh-9755
119 #
120 # need to set orig_names here first
121 # so that proper indexing can be done
122 # with _set_noconvert_columns
123 #
124 # once names has been filtered, we will
125 # then set orig_names again to names
126 self.orig_names = self.names[:]
128 if self.usecols:
129 usecols = evaluate_callable_usecols(self.usecols, self.orig_names)
131 # GH 14671
132 # assert for mypy, orig_names is List or None, None would error in issubset
133 assert self.orig_names is not None
134 if self.usecols_dtype == "string" and not set(usecols).issubset(
135 self.orig_names
136 ):
137 self._validate_usecols_names(usecols, self.orig_names)
139 if len(self.names) > len(usecols):
140 self.names = [
141 n
142 for i, n in enumerate(self.names)
143 if (i in usecols or n in usecols)
144 ]
146 if len(self.names) < len(usecols):
147 self._validate_usecols_names(
148 usecols,
149 self.names,
150 )
152 validate_parse_dates_presence(self.parse_dates, self.names)
153 self._set_noconvert_columns()
155 self.orig_names = self.names
157 if self._reader.leading_cols == 0 and is_index_col(self.index_col):
158 (
159 index_names,
160 self.names,
161 self.index_col,
162 ) = self._clean_index_names(
163 self.names,
164 self.index_col,
165 )
167 if self.index_names is None:
168 self.index_names = index_names
170 if self._reader.header is None and not passed_names:
171 assert self.index_names is not None
172 self.index_names = [None] * len(self.index_names)
174 self._implicit_index = self._reader.leading_cols > 0
176 def close(self) -> None:
177 # close handles opened by C parser
178 try:
179 self._reader.close()
180 except ValueError:
181 pass
183 def _set_noconvert_columns(self) -> None:
184 """
185 Set the columns that should not undergo dtype conversions.
187 Currently, any column that is involved with date parsing will not
188 undergo such conversions.
189 """
190 assert self.orig_names is not None
191 # error: Cannot determine type of 'names'
193 # much faster than using orig_names.index(x) xref GH#44106
194 names_dict = {x: i for i, x in enumerate(self.orig_names)}
195 col_indices = [names_dict[x] for x in self.names]
196 noconvert_columns = self._set_noconvert_dtype_columns(
197 col_indices,
198 self.names,
199 )
200 for col in noconvert_columns:
201 self._reader.set_noconvert(col)
203 def _low_memory_column_labels(
204 self, chunks: list[dict[int, ArrayLike]]
205 ) -> dict[int, Hashable]:
206 """
207 Map the field positions ``read_low_memory`` keys its chunks by to labels.
209 The positions are those of the source row, so they skip whatever
210 ``usecols`` dropped and start past a leading implicit index column.
211 ``orig_names`` holds exactly the named columns that remain, in the same
212 order, which is how ``read`` renames these keys further down. A leading
213 implicit index column has no name of its own and maps to ``None``.
214 """
215 assert self.orig_names is not None
216 positions = sorted(chunks[0])
217 named = dict(
218 zip(
219 positions[self._reader.leading_cols :],
220 self.orig_names,
221 strict=False,
222 )
223 )
224 return {position: named.get(position) for position in positions}
226 def read(
227 self,
228 nrows: int | None = None,
229 ) -> tuple[
230 Index | MultiIndex | None,
231 Sequence[Hashable] | MultiIndex,
232 Mapping[Hashable, AnyArrayLike],
233 ]:
234 index: Index | MultiIndex | None
235 column_names: Sequence[Hashable] | MultiIndex
236 try:
237 if self.low_memory:
238 chunks = self._reader.read_low_memory(nrows)
239 # destructive to chunks
240 data = _concatenate_chunks(
241 chunks, self._low_memory_column_labels(chunks)
242 )
243 else:
244 data = self._reader.read(nrows)
245 except StopIteration:
246 if self._first_chunk:
247 self._first_chunk = False
248 # assert for mypy, orig_names is List or None, None would error in
249 # list(...) in dedup_names
250 assert self.orig_names is not None
251 names = dedup_names(
252 self.orig_names,
253 is_potential_multi_index(self.orig_names, self.index_col),
254 )
255 index, columns, col_dict = self._get_empty_meta(
256 names,
257 dtype=self.dtype,
258 )
259 # error: Incompatible types in assignment (expression has type
260 # "list[Hashable] | MultiIndex", variable has type "list[Hashable]")
261 columns = self._maybe_make_multi_index_columns( # type: ignore[assignment]
262 columns, self.col_names
263 )
265 columns = _filter_usecols(self.usecols, columns)
266 columns_set = set(columns)
268 col_dict = {k: v for k, v in col_dict.items() if k in columns_set}
270 return index, columns, col_dict
272 else:
273 self.close()
274 raise
276 # Done with first read, next time raise StopIteration
277 self._first_chunk = False
279 names = self.names
281 if self._reader.leading_cols:
282 # implicit index, no index names
283 arrays = []
285 if self.index_col and self._reader.leading_cols != len(self.index_col):
286 raise ParserError(
287 "Could not construct index. Requested to use "
288 f"{len(self.index_col)} number of columns, but "
289 f"{self._reader.leading_cols} left to parse."
290 )
292 for i in range(self._reader.leading_cols):
293 if self.index_col is None:
294 values = data.pop(i)
295 else:
296 values = data.pop(self.index_col[i])
298 if self._should_parse_dates(i):
299 values = date_converter(
300 values,
301 col=(
302 self.index_names[i]
303 if self.index_names is not None
304 else None
305 ),
306 dayfirst=self.dayfirst,
307 cache_dates=self.cache_dates,
308 date_format=self.date_format,
309 )
310 arrays.append(values)
312 index = ensure_index_from_sequences(arrays)
314 names = _filter_usecols(self.usecols, names)
316 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
318 # rename dict keys
319 data_tups = sorted(data.items())
320 data = {k: v for k, (i, v) in zip(names, data_tups, strict=True)}
322 date_data = self._do_date_conversions(names, data)
324 # maybe create a mi on the columns
325 column_names = self._maybe_make_multi_index_columns(names, self.col_names)
327 else:
328 # rename dict keys
329 data_tups = sorted(data.items())
331 # ugh, mutation
333 # assert for mypy, orig_names is List or None, None would error in list(...)
334 assert self.orig_names is not None
335 names = list(self.orig_names)
336 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
338 names = _filter_usecols(self.usecols, names)
340 # columns as list
341 alldata = [x[1] for x in data_tups]
342 if self.usecols is None:
343 self._check_data_length(names, alldata)
345 data = {k: v for k, (i, v) in zip(names, data_tups, strict=False)}
347 date_data = self._do_date_conversions(names, data)
348 index, column_names = self._make_index(alldata, names)
350 return index, column_names, date_data
353def _filter_usecols(usecols, names: SequenceT) -> SequenceT | list[Hashable]:
354 # hackish
355 usecols = evaluate_callable_usecols(usecols, names)
356 if usecols is not None and len(names) != len(usecols):
357 return [name for i, name in enumerate(names) if i in usecols or name in usecols]
358 return names
361def _concatenate_chunks(
362 chunks: list[dict[int, ArrayLike]],
363 column_names: Sequence[Hashable] | Mapping[int, Hashable],
364) -> dict:
365 """
366 Concatenate chunks of data read with low_memory=True.
368 The tricky part is handling Categoricals, where different chunks
369 may have different inferred categories.
371 ``column_names`` names the keys of ``chunks`` for the mixed-dtype warning,
372 so it has to be indexable by whatever those keys are: the low_memory reader
373 keys its chunks by field position and passes a mapping keyed the same way,
374 while the parallel reader keys its own by name.
375 """
376 names = list(chunks[0].keys())
377 warning_columns = []
379 result: dict = {}
380 for name in names:
381 arrs = [chunk.pop(name) for chunk in chunks]
382 # Check each arr for consistent types.
383 dtypes = {a.dtype for a in arrs}
384 non_cat_dtypes = {x for x in dtypes if not isinstance(x, CategoricalDtype)}
386 dtype = dtypes.pop()
387 if isinstance(dtype, CategoricalDtype):
388 result[name] = union_categoricals(arrs, sort_categories=False)
389 else:
390 result[name] = concat_compat(arrs)
391 if len(non_cat_dtypes) > 1 and result[name].dtype == np.dtype(object):
392 warning_columns.append((name, column_names[name]))
394 if warning_columns:
395 warning_names = ", ".join(
396 [
397 f"{position}: {label}" if label is not None else f"{position}"
398 for position, label in warning_columns
399 ]
400 )
401 warning_message = " ".join(
402 [
403 f"Columns ({warning_names}) have mixed types. "
404 f"Specify dtype option on import or set low_memory=False."
405 ]
406 )
407 warnings.warn(warning_message, DtypeWarning, stacklevel=find_stack_level())
408 return result
411def ensure_dtype_objs(
412 dtype: DtypeArg | dict[Hashable, DtypeArg] | None,
413) -> DtypeObj | dict[Hashable, DtypeObj] | None:
414 """
415 Ensure we have either None, a dtype object, or a dictionary mapping to
416 dtype objects.
417 """
418 if isinstance(dtype, defaultdict):
419 # "None" not callable [misc]
420 default_dtype = pandas_dtype(dtype.default_factory()) # type: ignore[misc]
421 dtype_converted: defaultdict = defaultdict(lambda: default_dtype)
422 for key in dtype.keys():
423 dtype_converted[key] = pandas_dtype(dtype[key])
424 return dtype_converted
425 elif isinstance(dtype, dict):
426 return {k: pandas_dtype(dtype[k]) for k in dtype}
427 elif dtype is not None:
428 return pandas_dtype(dtype)
429 return dtype