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 cast,
7)
8import warnings
10from pandas._config import using_string_dtype
12from pandas._libs import lib
13from pandas.compat._optional import import_optional_dependency
14from pandas.compat.pyarrow import pa_version_under25p0
15from pandas.errors import (
16 EmptyDataError,
17 Pandas4Warning,
18 ParserError,
19 ParserWarning,
20)
21from pandas.util._exceptions import (
22 find_stack_level,
23)
25from pandas.core.dtypes.common import (
26 pandas_dtype,
27)
28from pandas.core.dtypes.dtypes import CategoricalDtype
29from pandas.core.dtypes.inference import is_integer
31from pandas.core.arrays.arrow.array import to_pyarrow_type
33from pandas.io._util import arrow_table_to_pandas
34from pandas.io.common import mangle_dupe_names
35from pandas.io.parsers.base_parser import (
36 ParserBase,
37 date_converter,
38)
40if TYPE_CHECKING:
41 from collections.abc import Hashable
43 import pyarrow as pa
45 from pandas._typing import ReadBuffer
47 from pandas import DataFrame
50def _pyarrow_parse_type(dtype) -> pa.DataType | None:
51 """
52 The pyarrow type a column should be parsed as in order to faithfully
53 apply the requested pandas dtype afterwards, or None to leave the
54 column to pyarrow's type inference.
56 GH#57666 dtypes whose conversion must start from the original text --
57 string dtypes and object -- are parsed as strings: converting an
58 inferred column instead is lossy (e.g. leading zeros are lost once a
59 column has been inferred as integer). All other dtypes are left to
60 inference and cast afterwards in _finalize_dtype, which is lossless for
61 them and keeps behavior consistent with the other engines (e.g. a column
62 of "True"/"False" values with dtype=float is inferred as boolean and
63 then cast, which pyarrow's parse-time conversion would reject).
64 """
65 pa = import_optional_dependency("pyarrow")
67 dtype = pandas_dtype(dtype)
68 if isinstance(dtype, CategoricalDtype):
69 if dtype.categories is None:
70 # GH#56044 the categories are inferred like any other column,
71 # so leave the parsing to inference
72 return None
73 # explicitly typed categories behave like their own dtype
74 dtype = dtype.categories.dtype
75 if dtype == object:
76 # match the other engines: object dtype holds the raw strings
77 return pa.large_string()
78 # pass the scalar type: to_pyarrow_type accepts numpy dtypes and
79 # scalar types but not extension dtype instances like StringDtype
80 # (pa.from_numpy_dtype raises TypeError for those)
81 target_dtype = to_pyarrow_type(dtype.type)
82 if target_dtype is not None and (
83 pa.types.is_string(target_dtype) or pa.types.is_large_string(target_dtype)
84 ):
85 # the default string dtype is backed by large_string, so parsing
86 # directly as large_string avoids a cast in the conversion to pandas
87 return pa.large_string()
88 return None
91class ArrowParserWrapper(ParserBase):
92 """
93 Wrapper for the pyarrow engine for read_csv()
94 """
96 def __init__(self, src: ReadBuffer[bytes], **kwds) -> None:
97 super().__init__(kwds)
98 self.kwds = kwds
99 self.src = src
101 self._parse_kwds()
103 def _parse_kwds(self) -> None:
104 """
105 Validates keywords before passing to pyarrow.
106 """
107 encoding: str | None = self.kwds.get("encoding")
108 self.encoding = "utf-8" if encoding is None else encoding
110 # GH#66317 validate the single-character options up front so the user
111 # gets the same informative errors as with the other engines instead
112 # of cryptic errors from pyarrow option conversion.
113 delimiter = self.kwds.get("delimiter")
114 if delimiter is not None and len(delimiter) != 1:
115 # longer separators already raise in _clean_options; this catches
116 # e.g. sep=""
117 raise ValueError("Only length-1 delimiters supported")
119 quotechar = self.kwds.get("quotechar", '"')
120 if quotechar is not None and not isinstance(quotechar, (str, bytes)):
121 raise TypeError(
122 f'"quotechar" must be string, not {type(quotechar).__name__}'
123 )
124 if not quotechar:
125 # None or zero-length; pyarrow only supports quoting=QUOTE_MINIMAL,
126 # so quoting is always enabled
127 raise TypeError("quotechar must be set if quoting enabled")
128 if len(quotechar) > 1:
129 raise TypeError('"quotechar" must be a 1-character string')
131 escapechar = self.kwds.get("escapechar")
132 if escapechar is not None and (
133 not isinstance(escapechar, (str, bytes)) or len(escapechar) != 1
134 ):
135 raise ValueError("Only length-1 escapes supported")
137 decimal = self.kwds.get("decimal", ".")
138 if not isinstance(decimal, (str, bytes)) or len(decimal) != 1:
139 raise ValueError("Only length-1 decimal markers supported")
141 na_values = self.kwds["na_values"]
142 if isinstance(na_values, dict):
143 raise ValueError(
144 "The pyarrow engine doesn't support passing a dict for na_values"
145 )
146 self.na_values = list(self.kwds["na_values"])
148 def _get_pyarrow_options(self) -> None:
149 """
150 Rename some arguments to pass to pyarrow
151 """
152 mapping = {
153 "usecols": "include_columns",
154 "na_values": "null_values",
155 "escapechar": "escape_char",
156 "skip_blank_lines": "ignore_empty_lines",
157 "decimal": "decimal_point",
158 "quotechar": "quote_char",
159 }
160 for pandas_name, pyarrow_name in mapping.items():
161 if pandas_name in self.kwds and self.kwds.get(pandas_name) is not None:
162 self.kwds[pyarrow_name] = self.kwds.pop(pandas_name)
164 # Date format handling
165 # If we get a string, we need to convert it into a list for pyarrow
166 # If we get a dict, we want to parse those separately
167 date_format = self.date_format
168 if isinstance(date_format, str):
169 date_format = [date_format]
170 else:
171 # In case of dict, we don't want to propagate through, so
172 # just set to pyarrow default of None
174 # Ideally, in future we disable pyarrow dtype inference (read in as string)
175 # to prevent misreads.
176 date_format = None
177 self.kwds["timestamp_parsers"] = date_format
179 self.parse_options = {
180 option_name: option_value
181 for option_name, option_value in self.kwds.items()
182 if option_value is not None
183 and option_name
184 in ("delimiter", "quote_char", "escape_char", "ignore_empty_lines")
185 }
187 on_bad_lines = self.kwds.get("on_bad_lines")
188 if on_bad_lines is not None:
189 if callable(on_bad_lines):
190 self.parse_options["invalid_row_handler"] = on_bad_lines
191 elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_ERROR:
192 self.parse_options["invalid_row_handler"] = (
193 None # PyArrow raises an exception by default
194 )
195 elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_WARN:
197 def handle_warning(invalid_row) -> str:
198 warnings.warn(
199 f"Expected {invalid_row.expected_columns} columns, but found "
200 f"{invalid_row.actual_columns}: {invalid_row.text}",
201 ParserWarning,
202 stacklevel=find_stack_level(),
203 )
204 return "skip"
206 self.parse_options["invalid_row_handler"] = handle_warning
207 elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_SKIP:
208 self.parse_options["invalid_row_handler"] = lambda _: "skip"
210 self.convert_options = {
211 option_name: option_value
212 for option_name, option_value in self.kwds.items()
213 if option_value is not None
214 and option_name
215 in (
216 "include_columns",
217 "null_values",
218 "true_values",
219 "false_values",
220 "decimal_point",
221 "timestamp_parsers",
222 )
223 }
224 self.convert_options["strings_can_be_null"] = "" in self.kwds["null_values"]
225 # autogenerated column names are prefixed with 'f' in pyarrow.csv
226 if self.header is None and "include_columns" in self.convert_options:
227 self.convert_options["include_columns"] = [
228 f"f{n}" for n in self.convert_options["include_columns"]
229 ]
231 if self.dtype is not None:
232 # GH#57666 dtypes whose conversion must start from the original
233 # text are applied while parsing; see _pyarrow_parse_type
234 if isinstance(self.dtype, dict):
235 column_types = {}
236 for col, col_dtype in self.dtype.items():
237 if not isinstance(col, str):
238 # non-string labels are legal (e.g. positional keys,
239 # or names=[0, 1]), but pyarrow's column_types
240 # requires string column names; these are still
241 # applied afterwards in _finalize_dtype
242 continue
243 target_dtype = _pyarrow_parse_type(col_dtype)
244 if target_dtype is not None:
245 column_types[col] = target_dtype
246 if column_types:
247 self.convert_options["column_types"] = column_types
248 if (
249 isinstance(self.dtype, defaultdict)
250 and self.dtype.default_factory is not None
251 and not pa_version_under25p0
252 ):
253 # the default of a defaultdict applies to every column
254 # not explicitly listed, like a scalar dtype; explicit
255 # column_types entries take precedence over
256 # default_column_type
257 target_dtype = _pyarrow_parse_type(self.dtype.default_factory())
258 if target_dtype is not None:
259 self.convert_options["default_column_type"] = target_dtype
260 elif not pa_version_under25p0:
261 # a scalar dtype applies to every column, which pyarrow
262 # supports via ConvertOptions(default_column_type=...) from
263 # pyarrow 25.0 (apache/arrow#47663). With older versions the
264 # dtype is still applied afterwards in _finalize_dtype,
265 # which does not preserve the original text.
266 target_dtype = _pyarrow_parse_type(self.dtype)
267 if target_dtype is not None:
268 self.convert_options["default_column_type"] = target_dtype
270 header = self.header
271 if isinstance(header, list):
272 # GH#66059 pyarrow.csv cannot produce multi-row/MultiIndex headers
273 if len(header) == 1:
274 header = header[0]
275 else:
276 raise ValueError(
277 "The 'pyarrow' engine does not support a list of integers "
278 "for the 'header' argument (MultiIndex columns are not "
279 "supported)."
280 )
282 # GH#48507 a negative skiprows skips nothing in the other engines;
283 # pyarrow rejects it.
284 skiprows = max(self.kwds["skiprows"] or 0, 0)
286 if header is not None and self.names is not None:
287 # GH#65862 pyarrow reads with autogenerated column names in this
288 # case, so the user's usecols cannot be mapped to them.
289 # Non-string usecols instead raise in _get_convert_options.
290 usecols = self.convert_options.get("include_columns")
291 if usecols and all(isinstance(col, str) for col in usecols):
292 raise ValueError(
293 "The 'pyarrow' engine does not support 'usecols' together "
294 "with 'names' when 'header' is an integer"
295 )
296 # GH#65862 names replace the header row, which is discarded as
297 # the other engines do. Use skip_rows_after_names rather than
298 # skip_rows since the latter cannot skip past a final line
299 # lacking a trailing newline.
300 skip_rows_after_names = 1
301 else:
302 skip_rows_after_names = 0
304 # GH#48507 pyarrow skips rows before reading the header, so the header
305 # is at row skiprows + header.
306 skip_rows = skiprows + (header or 0)
307 self.read_options = {
308 "autogenerate_column_names": header is None or self.names is not None,
309 "skip_rows": skip_rows,
310 "skip_rows_after_names": skip_rows_after_names,
311 "encoding": self.encoding,
312 }
314 def _index_col_should_parse_dates(self, name, position: int) -> bool:
315 """
316 Whether the index column with the given name/position should be parsed
317 as dates. Unlike the other engines, the pyarrow path sets the index after
318 reading, so parse_dates handling for index columns happens here.
319 """
320 if isinstance(self.parse_dates, bool):
321 return self.parse_dates
322 return name in self.parse_dates or position in self.parse_dates
324 def _get_convert_options(self):
325 pyarrow_csv = import_optional_dependency("pyarrow.csv")
327 try:
328 convert_options = pyarrow_csv.ConvertOptions(**self.convert_options)
329 except TypeError as err:
330 include = self.convert_options.get("include_columns", None)
331 if include is not None:
332 self._validate_usecols(include)
334 nulls = self.convert_options.get("null_values", set())
335 if not lib.is_list_like(nulls) or not all(
336 isinstance(x, str) for x in nulls
337 ):
338 raise TypeError(
339 "The 'pyarrow' engine requires all na_values to be strings"
340 ) from err
342 raise
344 return convert_options
346 def _dedup_column_names(
347 self, raw_names: list[str]
348 ) -> tuple[list[Hashable], set[Hashable]]:
349 """
350 Match other engines' header handling: empty names become
351 "Unnamed: {i}" and duplicated names are de-duplicated, mirroring the
352 algorithm in pandas._libs.parsers.TextReader.
354 Returns the new names along with the set of placeholder names.
355 """
356 names: list[Hashable] = []
357 unnamed_col_indices = []
358 for i, name in enumerate(raw_names):
359 if name == "":
360 name = f"Unnamed: {i}"
361 unnamed_col_indices.append(i)
362 names.append(name)
364 names = mangle_dupe_names(names, unnamed_col_indices, self.dtype)
365 return names, {names[i] for i in unnamed_col_indices}
367 def _adjust_column_names(self, table: pa.Table) -> bool:
368 num_cols = len(table.columns)
369 multi_index_named = True
370 self._implicit_index_count = 0
371 if self.header is None and self.names is None:
372 self.names = range(num_cols)
373 if self.names is not None and len(self.names) != num_cols:
374 # usecols is passed through to pyarrow, so num_cols here already
375 # reflects it
376 usecols = self.kwds.get("include_columns")
377 if (self.header is not None and len(self.names) > num_cols) or (
378 usecols is not None and len(self.names) != len(usecols)
379 ):
380 # GH#5156, GH#65862
381 raise ValueError(
382 "Number of passed names did not match number of "
383 "header fields in the file"
384 )
385 if len(self.names) < num_cols:
386 # GH#65862 Leading columns not covered by names form the index,
387 # as with other engines. Pad the names to the expected length;
388 # the padded columns are moved to the index in _finalize_index.
389 self._implicit_index_count = num_cols - len(self.names)
390 columns_prefix = [str(x) for x in range(self._implicit_index_count)]
391 self.names = columns_prefix + list(self.names)
392 multi_index_named = False
393 return multi_index_named
395 def _finalize_index(self, frame: DataFrame, multi_index_named: bool) -> DataFrame:
396 if self.index_col is None and self._implicit_index_count:
397 self.index_col = list(range(self._implicit_index_count))
398 if self.index_col is not None:
399 index_to_set = self.index_col.copy()
400 for i, item in enumerate(self.index_col):
401 if is_integer(item):
402 col_name = frame.columns[item]
403 index_to_set[i] = col_name
404 position = int(item)
405 # String case
406 elif item not in frame.columns:
407 raise ValueError(f"Index {item} invalid")
408 else:
409 col_name = item
410 position = cast("int", frame.columns.get_loc(item))
412 # parse_dates for index columns: the other engines parse these
413 # while building the index, but the pyarrow path sets the index
414 # after reading, so do the conversion here.
415 if self._index_col_should_parse_dates(col_name, position):
416 frame[col_name] = date_converter(
417 frame[col_name],
418 col=col_name,
419 dayfirst=self.dayfirst,
420 cache_dates=self.cache_dates,
421 date_format=self.date_format,
422 )
424 # Process dtype for index_col and drop from dtypes
425 if isinstance(self.dtype, dict):
426 key, new_dtype = (
427 (item, self.dtype.get(item))
428 if self.dtype.get(item) is not None
429 else (frame.columns[item], self.dtype.get(frame.columns[item]))
430 )
431 if new_dtype is not None:
432 frame[key] = frame[key].astype(new_dtype)
433 del self.dtype[key]
435 if self.dtype is not None and not isinstance(self.dtype, dict):
436 # GH#45801 match the c engine: an index column left as object
437 # by the scalar dtype (i.e. the dtype resolved to object) is
438 # re-inferred, as if no dtype had been passed
439 for key in index_to_set:
440 if frame[key].dtype == object:
441 new_values, _ = self._infer_types(
442 frame[key].to_numpy(copy=True),
443 set(),
444 no_dtype_specified=True,
445 )
446 frame[key] = new_values
447 frame[key] = frame[key].infer_objects()
449 frame = frame.set_index(index_to_set, drop=True)
450 # Clear names if no name given for the padded leading columns
451 if not multi_index_named:
452 frame.index.names = [None] * len(frame.index.names)
453 elif self.unnamed_cols:
454 # GH#13017 match other engines: empty header fields used as
455 # an index produce an unnamed index level
456 frame.index.names = [
457 None if name in self.unnamed_cols else name
458 for name in frame.index.names
459 ]
461 return frame
463 def _finalize_dtype(self, frame: DataFrame) -> DataFrame:
464 if self.dtype is not None:
465 # Ignore non-existent columns from dtype mapping
466 # like other parsers do
467 if isinstance(self.dtype, dict):
468 self.dtype = {
469 k: pandas_dtype(v)
470 for k, v in self.dtype.items()
471 if k in frame.columns
472 }
473 else:
474 # GH#34066 a scalar dtype must not clobber the dtype produced by
475 # date parsing, so exclude parse_dates columns from the cast.
476 scalar_dtype = pandas_dtype(self.dtype)
477 parse_dates_cols = (
478 set(self.parse_dates)
479 if isinstance(self.parse_dates, list)
480 else set()
481 )
482 self.dtype = {
483 col: scalar_dtype
484 for col in frame.columns
485 if col not in parse_dates_cols
486 }
487 try:
488 frame = frame.astype(self.dtype)
489 except TypeError as err:
490 # GH#44901 reraise to keep api consistent
491 raise ValueError(str(err)) from err
492 return frame
494 def _finalize_pandas_output(
495 self, frame: DataFrame, multi_index_named: bool
496 ) -> DataFrame:
497 """
498 Processes data read in based on kwargs.
500 Parameters
501 ----------
502 frame : DataFrame
503 The DataFrame to process.
504 multi_index_named : bool
506 Returns
507 -------
508 DataFrame
509 The processed DataFrame.
510 """
511 frame = self._do_date_conversions(frame.columns, frame)
512 frame = self._maybe_restore_string_dtype(frame)
513 frame = self._finalize_index(frame, multi_index_named)
514 frame = self._finalize_dtype(frame)
515 # GH#65862 tuples passed via names imply MultiIndex columns,
516 # as with other engines
517 frame.columns = self._maybe_make_multi_index_columns(
518 list(frame.columns), self.col_names
519 )
520 return frame
522 def _maybe_restore_string_dtype(self, frame: DataFrame) -> DataFrame:
523 # When parse_dates fails to parse a column it falls back to the original
524 # strings; keep the default string dtype rather than object, matching
525 # the other engines.
526 if not (using_string_dtype() and isinstance(self.parse_dates, list)):
527 return frame
528 for col in self.parse_dates:
529 if col in frame.columns and frame[col].dtype == object:
530 frame[col] = frame[col].astype("str")
531 return frame
533 def _validate_usecols(self, usecols) -> None:
534 if lib.is_list_like(usecols) and not all(isinstance(x, str) for x in usecols):
535 raise ValueError(
536 "The pyarrow engine does not allow 'usecols' to be integer "
537 "column positions. Pass a list of string column names instead."
538 )
539 elif callable(usecols):
540 raise ValueError(
541 "The pyarrow engine does not allow 'usecols' to be a callable."
542 )
544 def read(self) -> DataFrame:
545 """
546 Reads the contents of a CSV file into a DataFrame and
547 processes it according to the kwargs passed in the
548 constructor.
550 Returns
551 -------
552 DataFrame
553 The DataFrame created from the CSV file.
554 """
555 pa = import_optional_dependency("pyarrow")
556 pyarrow_csv = import_optional_dependency("pyarrow.csv")
557 self._get_pyarrow_options()
558 convert_options = self._get_convert_options()
560 try:
561 table = pyarrow_csv.read_csv(
562 self.src,
563 read_options=pyarrow_csv.ReadOptions(**self.read_options),
564 parse_options=pyarrow_csv.ParseOptions(**self.parse_options),
565 convert_options=convert_options,
566 )
567 except pa.ArrowInvalid as err:
568 # pyarrow reports "Empty CSV file or block" when it cannot extract
569 # any columns from the source; re-raise as EmptyDataError so the
570 # pyarrow engine matches the "c" and "python" engines.
571 if "Empty CSV file" in str(err):
572 raise EmptyDataError("No columns to parse from file") from err
573 raise ParserError(err) from err
574 except TypeError as err:
575 # pyarrow can only read from a binary buffer; a text-mode file
576 # handle raises a bare "a bytes-like object is required" TypeError.
577 if "a bytes-like object is required" in str(err):
578 raise TypeError(
579 "The 'pyarrow' engine can only read from a binary file "
580 "object; the given file was opened in text mode"
581 ) from err
582 raise
584 if self.usecols_dtype == "empty":
585 # GH#66056 pyarrow ignores an empty ``include_columns`` and returns
586 # every column; match the other engines by returning an empty frame.
587 table = table.select([]).slice(0, 0)
589 dtype_backend = self.kwds["dtype_backend"]
591 # Convert all pa.null() cols -> float64 (non nullable)
592 # else Int64 (nullable case, see below)
593 if dtype_backend is lib.no_default:
594 new_schema = table.schema
595 new_type = pa.float64()
596 for i, arrow_type in enumerate(table.schema.types):
597 if pa.types.is_null(arrow_type):
598 new_schema = new_schema.set(
599 i, new_schema.field(i).with_type(new_type)
600 )
602 table = table.cast(new_schema)
604 if self.header is not None and self.names is None:
605 new_names, self.unnamed_cols = self._dedup_column_names(table.column_names)
606 if new_names != table.column_names:
607 table = table.rename_columns(new_names)
609 multi_index_named = self._adjust_column_names(table)
611 if isinstance(self.dtype, defaultdict):
612 # GH#41574 materialize the factory default over the actual columns
613 # so missing keys get the default, matching the other engines
614 if self.names is not None:
615 # always set by _adjust_column_names when header is None
616 columns = list(self.names)
617 else:
618 columns = table.column_names
619 self.dtype = {col: self.dtype[col] for col in columns}
621 with warnings.catch_warnings():
622 warnings.filterwarnings(
623 "ignore",
624 "make_block is deprecated",
625 Pandas4Warning,
626 )
627 frame = arrow_table_to_pandas(
628 table,
629 dtype_backend=dtype_backend,
630 null_to_int64=True,
631 dtype=self.dtype,
632 names=self.names,
633 )
635 if self.names is not None:
636 frame.columns = self.names
638 return self._finalize_pandas_output(frame, multi_index_named)