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 read(
204 self,
205 nrows: int | None = None,
206 ) -> tuple[
207 Index | MultiIndex | None,
208 Sequence[Hashable] | MultiIndex,
209 Mapping[Hashable, AnyArrayLike],
210 ]:
211 index: Index | MultiIndex | None
212 column_names: Sequence[Hashable] | MultiIndex
213 try:
214 if self.low_memory:
215 chunks = self._reader.read_low_memory(nrows)
216 # destructive to chunks
217 data = _concatenate_chunks(chunks, self.names)
218 else:
219 data = self._reader.read(nrows)
220 except StopIteration:
221 if self._first_chunk:
222 self._first_chunk = False
223 # assert for mypy, orig_names is List or None, None would error in
224 # list(...) in dedup_names
225 assert self.orig_names is not None
226 names = dedup_names(
227 self.orig_names,
228 is_potential_multi_index(self.orig_names, self.index_col),
229 )
230 index, columns, col_dict = self._get_empty_meta(
231 names,
232 dtype=self.dtype,
233 )
234 # error: Incompatible types in assignment (expression has type
235 # "list[Hashable] | MultiIndex", variable has type "list[Hashable]")
236 columns = self._maybe_make_multi_index_columns( # type: ignore[assignment]
237 columns, self.col_names
238 )
240 columns = _filter_usecols(self.usecols, columns)
241 columns_set = set(columns)
243 col_dict = {k: v for k, v in col_dict.items() if k in columns_set}
245 return index, columns, col_dict
247 else:
248 self.close()
249 raise
251 # Done with first read, next time raise StopIteration
252 self._first_chunk = False
254 names = self.names
256 if self._reader.leading_cols:
257 # implicit index, no index names
258 arrays = []
260 if self.index_col and self._reader.leading_cols != len(self.index_col):
261 raise ParserError(
262 "Could not construct index. Requested to use "
263 f"{len(self.index_col)} number of columns, but "
264 f"{self._reader.leading_cols} left to parse."
265 )
267 for i in range(self._reader.leading_cols):
268 if self.index_col is None:
269 values = data.pop(i)
270 else:
271 values = data.pop(self.index_col[i])
273 if self._should_parse_dates(i):
274 values = date_converter(
275 values,
276 col=(
277 self.index_names[i]
278 if self.index_names is not None
279 else None
280 ),
281 dayfirst=self.dayfirst,
282 cache_dates=self.cache_dates,
283 date_format=self.date_format,
284 )
285 arrays.append(values)
287 index = ensure_index_from_sequences(arrays)
289 names = _filter_usecols(self.usecols, names)
291 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
293 # rename dict keys
294 data_tups = sorted(data.items())
295 data = {k: v for k, (i, v) in zip(names, data_tups, strict=True)}
297 date_data = self._do_date_conversions(names, data)
299 # maybe create a mi on the columns
300 column_names = self._maybe_make_multi_index_columns(names, self.col_names)
302 else:
303 # rename dict keys
304 data_tups = sorted(data.items())
306 # ugh, mutation
308 # assert for mypy, orig_names is List or None, None would error in list(...)
309 assert self.orig_names is not None
310 names = list(self.orig_names)
311 names = dedup_names(names, is_potential_multi_index(names, self.index_col))
313 names = _filter_usecols(self.usecols, names)
315 # columns as list
316 alldata = [x[1] for x in data_tups]
317 if self.usecols is None:
318 self._check_data_length(names, alldata)
320 data = {k: v for k, (i, v) in zip(names, data_tups, strict=False)}
322 date_data = self._do_date_conversions(names, data)
323 index, column_names = self._make_index(alldata, names)
325 return index, column_names, date_data
328def _filter_usecols(usecols, names: SequenceT) -> SequenceT | list[Hashable]:
329 # hackish
330 usecols = evaluate_callable_usecols(usecols, names)
331 if usecols is not None and len(names) != len(usecols):
332 return [name for i, name in enumerate(names) if i in usecols or name in usecols]
333 return names
336def _concatenate_chunks(
337 chunks: list[dict[int, ArrayLike]], column_names: list[str]
338) -> dict:
339 """
340 Concatenate chunks of data read with low_memory=True.
342 The tricky part is handling Categoricals, where different chunks
343 may have different inferred categories.
344 """
345 names = list(chunks[0].keys())
346 warning_columns = []
348 result: dict = {}
349 for name in names:
350 arrs = [chunk.pop(name) for chunk in chunks]
351 # Check each arr for consistent types.
352 dtypes = {a.dtype for a in arrs}
353 non_cat_dtypes = {x for x in dtypes if not isinstance(x, CategoricalDtype)}
355 dtype = dtypes.pop()
356 if isinstance(dtype, CategoricalDtype):
357 result[name] = union_categoricals(arrs, sort_categories=False)
358 else:
359 result[name] = concat_compat(arrs)
360 if len(non_cat_dtypes) > 1 and result[name].dtype == np.dtype(object):
361 warning_columns.append(column_names[name])
363 if warning_columns:
364 warning_names = ", ".join(
365 [f"{index}: {name}" for index, name in enumerate(warning_columns)]
366 )
367 warning_message = " ".join(
368 [
369 f"Columns ({warning_names}) have mixed types. "
370 f"Specify dtype option on import or set low_memory=False."
371 ]
372 )
373 warnings.warn(warning_message, DtypeWarning, stacklevel=find_stack_level())
374 return result
377def ensure_dtype_objs(
378 dtype: DtypeArg | dict[Hashable, DtypeArg] | None,
379) -> DtypeObj | dict[Hashable, DtypeObj] | None:
380 """
381 Ensure we have either None, a dtype object, or a dictionary mapping to
382 dtype objects.
383 """
384 if isinstance(dtype, defaultdict):
385 # "None" not callable [misc]
386 default_dtype = pandas_dtype(dtype.default_factory()) # type: ignore[misc]
387 dtype_converted: defaultdict = defaultdict(lambda: default_dtype)
388 for key in dtype.keys():
389 dtype_converted[key] = pandas_dtype(dtype[key])
390 return dtype_converted
391 elif isinstance(dtype, dict):
392 return {k: pandas_dtype(dtype[k]) for k in dtype}
393 elif dtype is not None:
394 return pandas_dtype(dtype)
395 return dtype