from __future__ import annotations from collections import defaultdict from typing import ( TYPE_CHECKING, cast, ) import warnings from pandas._config import using_string_dtype from pandas._libs import lib from pandas.compat._optional import import_optional_dependency from pandas.compat.pyarrow import pa_version_under25p0 from pandas.errors import ( EmptyDataError, Pandas4Warning, ParserError, ParserWarning, ) from pandas.util._exceptions import ( find_stack_level, ) from pandas.core.dtypes.common import ( pandas_dtype, ) from pandas.core.dtypes.dtypes import CategoricalDtype from pandas.core.dtypes.inference import is_integer from pandas.core.arrays.arrow.array import to_pyarrow_type from pandas.io._util import arrow_table_to_pandas from pandas.io.common import mangle_dupe_names from pandas.io.parsers.base_parser import ( ParserBase, date_converter, ) if TYPE_CHECKING: from collections.abc import Hashable import pyarrow as pa from pandas._typing import ReadBuffer from pandas import DataFrame def _pyarrow_parse_type(dtype) -> pa.DataType | None: """ The pyarrow type a column should be parsed as in order to faithfully apply the requested pandas dtype afterwards, or None to leave the column to pyarrow's type inference. GH#57666 dtypes whose conversion must start from the original text -- string dtypes and object -- are parsed as strings: converting an inferred column instead is lossy (e.g. leading zeros are lost once a column has been inferred as integer). All other dtypes are left to inference and cast afterwards in _finalize_dtype, which is lossless for them and keeps behavior consistent with the other engines (e.g. a column of "True"/"False" values with dtype=float is inferred as boolean and then cast, which pyarrow's parse-time conversion would reject). """ pa = import_optional_dependency("pyarrow") dtype = pandas_dtype(dtype) if isinstance(dtype, CategoricalDtype): if dtype.categories is None: # GH#56044 the categories are inferred like any other column, # so leave the parsing to inference return None # explicitly typed categories behave like their own dtype dtype = dtype.categories.dtype if dtype == object: # match the other engines: object dtype holds the raw strings return pa.large_string() # pass the scalar type: to_pyarrow_type accepts numpy dtypes and # scalar types but not extension dtype instances like StringDtype # (pa.from_numpy_dtype raises TypeError for those) target_dtype = to_pyarrow_type(dtype.type) if target_dtype is not None and ( pa.types.is_string(target_dtype) or pa.types.is_large_string(target_dtype) ): # the default string dtype is backed by large_string, so parsing # directly as large_string avoids a cast in the conversion to pandas return pa.large_string() return None class ArrowParserWrapper(ParserBase): """ Wrapper for the pyarrow engine for read_csv() """ def __init__(self, src: ReadBuffer[bytes], **kwds) -> None: super().__init__(kwds) self.kwds = kwds self.src = src self._parse_kwds() def _parse_kwds(self) -> None: """ Validates keywords before passing to pyarrow. """ encoding: str | None = self.kwds.get("encoding") self.encoding = "utf-8" if encoding is None else encoding # GH#66317 validate the single-character options up front so the user # gets the same informative errors as with the other engines instead # of cryptic errors from pyarrow option conversion. delimiter = self.kwds.get("delimiter") if delimiter is not None and len(delimiter) != 1: # longer separators already raise in _clean_options; this catches # e.g. sep="" raise ValueError("Only length-1 delimiters supported") quotechar = self.kwds.get("quotechar", '"') if quotechar is not None and not isinstance(quotechar, (str, bytes)): raise TypeError( f'"quotechar" must be string, not {type(quotechar).__name__}' ) if not quotechar: # None or zero-length; pyarrow only supports quoting=QUOTE_MINIMAL, # so quoting is always enabled raise TypeError("quotechar must be set if quoting enabled") if len(quotechar) > 1: raise TypeError('"quotechar" must be a 1-character string') escapechar = self.kwds.get("escapechar") if escapechar is not None and ( not isinstance(escapechar, (str, bytes)) or len(escapechar) != 1 ): raise ValueError("Only length-1 escapes supported") decimal = self.kwds.get("decimal", ".") if not isinstance(decimal, (str, bytes)) or len(decimal) != 1: raise ValueError("Only length-1 decimal markers supported") na_values = self.kwds["na_values"] if isinstance(na_values, dict): raise ValueError( "The pyarrow engine doesn't support passing a dict for na_values" ) self.na_values = list(self.kwds["na_values"]) def _get_pyarrow_options(self) -> None: """ Rename some arguments to pass to pyarrow """ mapping = { "usecols": "include_columns", "na_values": "null_values", "escapechar": "escape_char", "skip_blank_lines": "ignore_empty_lines", "decimal": "decimal_point", "quotechar": "quote_char", } for pandas_name, pyarrow_name in mapping.items(): if pandas_name in self.kwds and self.kwds.get(pandas_name) is not None: self.kwds[pyarrow_name] = self.kwds.pop(pandas_name) # Date format handling # If we get a string, we need to convert it into a list for pyarrow # If we get a dict, we want to parse those separately date_format = self.date_format if isinstance(date_format, str): date_format = [date_format] else: # In case of dict, we don't want to propagate through, so # just set to pyarrow default of None # Ideally, in future we disable pyarrow dtype inference (read in as string) # to prevent misreads. date_format = None self.kwds["timestamp_parsers"] = date_format self.parse_options = { option_name: option_value for option_name, option_value in self.kwds.items() if option_value is not None and option_name in ("delimiter", "quote_char", "escape_char", "ignore_empty_lines") } on_bad_lines = self.kwds.get("on_bad_lines") if on_bad_lines is not None: if callable(on_bad_lines): self.parse_options["invalid_row_handler"] = on_bad_lines elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_ERROR: self.parse_options["invalid_row_handler"] = ( None # PyArrow raises an exception by default ) elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_WARN: def handle_warning(invalid_row) -> str: warnings.warn( f"Expected {invalid_row.expected_columns} columns, but found " f"{invalid_row.actual_columns}: {invalid_row.text}", ParserWarning, stacklevel=find_stack_level(), ) return "skip" self.parse_options["invalid_row_handler"] = handle_warning elif on_bad_lines == ParserBase.BadLineHandleMethod.BLHM_SKIP: self.parse_options["invalid_row_handler"] = lambda _: "skip" self.convert_options = { option_name: option_value for option_name, option_value in self.kwds.items() if option_value is not None and option_name in ( "include_columns", "null_values", "true_values", "false_values", "decimal_point", "timestamp_parsers", ) } self.convert_options["strings_can_be_null"] = "" in self.kwds["null_values"] # autogenerated column names are prefixed with 'f' in pyarrow.csv if self.header is None and "include_columns" in self.convert_options: self.convert_options["include_columns"] = [ f"f{n}" for n in self.convert_options["include_columns"] ] if self.dtype is not None: # GH#57666 dtypes whose conversion must start from the original # text are applied while parsing; see _pyarrow_parse_type if isinstance(self.dtype, dict): column_types = {} for col, col_dtype in self.dtype.items(): if not isinstance(col, str): # non-string labels are legal (e.g. positional keys, # or names=[0, 1]), but pyarrow's column_types # requires string column names; these are still # applied afterwards in _finalize_dtype continue target_dtype = _pyarrow_parse_type(col_dtype) if target_dtype is not None: column_types[col] = target_dtype if column_types: self.convert_options["column_types"] = column_types if ( isinstance(self.dtype, defaultdict) and self.dtype.default_factory is not None and not pa_version_under25p0 ): # the default of a defaultdict applies to every column # not explicitly listed, like a scalar dtype; explicit # column_types entries take precedence over # default_column_type target_dtype = _pyarrow_parse_type(self.dtype.default_factory()) if target_dtype is not None: self.convert_options["default_column_type"] = target_dtype elif not pa_version_under25p0: # a scalar dtype applies to every column, which pyarrow # supports via ConvertOptions(default_column_type=...) from # pyarrow 25.0 (apache/arrow#47663). With older versions the # dtype is still applied afterwards in _finalize_dtype, # which does not preserve the original text. target_dtype = _pyarrow_parse_type(self.dtype) if target_dtype is not None: self.convert_options["default_column_type"] = target_dtype header = self.header if isinstance(header, list): # GH#66059 pyarrow.csv cannot produce multi-row/MultiIndex headers if len(header) == 1: header = header[0] else: raise ValueError( "The 'pyarrow' engine does not support a list of integers " "for the 'header' argument (MultiIndex columns are not " "supported)." ) # GH#48507 a negative skiprows skips nothing in the other engines; # pyarrow rejects it. skiprows = max(self.kwds["skiprows"] or 0, 0) if header is not None and self.names is not None: # GH#65862 pyarrow reads with autogenerated column names in this # case, so the user's usecols cannot be mapped to them. # Non-string usecols instead raise in _get_convert_options. usecols = self.convert_options.get("include_columns") if usecols and all(isinstance(col, str) for col in usecols): raise ValueError( "The 'pyarrow' engine does not support 'usecols' together " "with 'names' when 'header' is an integer" ) # GH#65862 names replace the header row, which is discarded as # the other engines do. Use skip_rows_after_names rather than # skip_rows since the latter cannot skip past a final line # lacking a trailing newline. skip_rows_after_names = 1 else: skip_rows_after_names = 0 # GH#48507 pyarrow skips rows before reading the header, so the header # is at row skiprows + header. skip_rows = skiprows + (header or 0) self.read_options = { "autogenerate_column_names": header is None or self.names is not None, "skip_rows": skip_rows, "skip_rows_after_names": skip_rows_after_names, "encoding": self.encoding, } def _index_col_should_parse_dates(self, name, position: int) -> bool: """ Whether the index column with the given name/position should be parsed as dates. Unlike the other engines, the pyarrow path sets the index after reading, so parse_dates handling for index columns happens here. """ if isinstance(self.parse_dates, bool): return self.parse_dates return name in self.parse_dates or position in self.parse_dates def _get_convert_options(self): pyarrow_csv = import_optional_dependency("pyarrow.csv") try: convert_options = pyarrow_csv.ConvertOptions(**self.convert_options) except TypeError as err: include = self.convert_options.get("include_columns", None) if include is not None: self._validate_usecols(include) nulls = self.convert_options.get("null_values", set()) if not lib.is_list_like(nulls) or not all( isinstance(x, str) for x in nulls ): raise TypeError( "The 'pyarrow' engine requires all na_values to be strings" ) from err raise return convert_options def _dedup_column_names( self, raw_names: list[str] ) -> tuple[list[Hashable], set[Hashable]]: """ Match other engines' header handling: empty names become "Unnamed: {i}" and duplicated names are de-duplicated, mirroring the algorithm in pandas._libs.parsers.TextReader. Returns the new names along with the set of placeholder names. """ names: list[Hashable] = [] unnamed_col_indices = [] for i, name in enumerate(raw_names): if name == "": name = f"Unnamed: {i}" unnamed_col_indices.append(i) names.append(name) names = mangle_dupe_names(names, unnamed_col_indices, self.dtype) return names, {names[i] for i in unnamed_col_indices} def _adjust_column_names(self, table: pa.Table) -> bool: num_cols = len(table.columns) multi_index_named = True self._implicit_index_count = 0 if self.header is None and self.names is None: self.names = range(num_cols) if self.names is not None and len(self.names) != num_cols: # usecols is passed through to pyarrow, so num_cols here already # reflects it usecols = self.kwds.get("include_columns") if (self.header is not None and len(self.names) > num_cols) or ( usecols is not None and len(self.names) != len(usecols) ): # GH#5156, GH#65862 raise ValueError( "Number of passed names did not match number of " "header fields in the file" ) if len(self.names) < num_cols: # GH#65862 Leading columns not covered by names form the index, # as with other engines. Pad the names to the expected length; # the padded columns are moved to the index in _finalize_index. self._implicit_index_count = num_cols - len(self.names) columns_prefix = [str(x) for x in range(self._implicit_index_count)] self.names = columns_prefix + list(self.names) multi_index_named = False return multi_index_named def _finalize_index(self, frame: DataFrame, multi_index_named: bool) -> DataFrame: if self.index_col is None and self._implicit_index_count: self.index_col = list(range(self._implicit_index_count)) if self.index_col is not None: index_to_set = self.index_col.copy() for i, item in enumerate(self.index_col): if is_integer(item): col_name = frame.columns[item] index_to_set[i] = col_name position = int(item) # String case elif item not in frame.columns: raise ValueError(f"Index {item} invalid") else: col_name = item position = cast("int", frame.columns.get_loc(item)) # parse_dates for index columns: the other engines parse these # while building the index, but the pyarrow path sets the index # after reading, so do the conversion here. if self._index_col_should_parse_dates(col_name, position): frame[col_name] = date_converter( frame[col_name], col=col_name, dayfirst=self.dayfirst, cache_dates=self.cache_dates, date_format=self.date_format, ) # Process dtype for index_col and drop from dtypes if isinstance(self.dtype, dict): key, new_dtype = ( (item, self.dtype.get(item)) if self.dtype.get(item) is not None else (frame.columns[item], self.dtype.get(frame.columns[item])) ) if new_dtype is not None: frame[key] = frame[key].astype(new_dtype) del self.dtype[key] if self.dtype is not None and not isinstance(self.dtype, dict): # GH#45801 match the c engine: an index column left as object # by the scalar dtype (i.e. the dtype resolved to object) is # re-inferred, as if no dtype had been passed for key in index_to_set: if frame[key].dtype == object: new_values, _ = self._infer_types( frame[key].to_numpy(copy=True), set(), no_dtype_specified=True, ) frame[key] = new_values frame[key] = frame[key].infer_objects() frame = frame.set_index(index_to_set, drop=True) # Clear names if no name given for the padded leading columns if not multi_index_named: frame.index.names = [None] * len(frame.index.names) elif self.unnamed_cols: # GH#13017 match other engines: empty header fields used as # an index produce an unnamed index level frame.index.names = [ None if name in self.unnamed_cols else name for name in frame.index.names ] return frame def _finalize_dtype(self, frame: DataFrame) -> DataFrame: if self.dtype is not None: # Ignore non-existent columns from dtype mapping # like other parsers do if isinstance(self.dtype, dict): self.dtype = { k: pandas_dtype(v) for k, v in self.dtype.items() if k in frame.columns } else: # GH#34066 a scalar dtype must not clobber the dtype produced by # date parsing, so exclude parse_dates columns from the cast. scalar_dtype = pandas_dtype(self.dtype) parse_dates_cols = ( set(self.parse_dates) if isinstance(self.parse_dates, list) else set() ) self.dtype = { col: scalar_dtype for col in frame.columns if col not in parse_dates_cols } try: frame = frame.astype(self.dtype) except TypeError as err: # GH#44901 reraise to keep api consistent raise ValueError(str(err)) from err return frame def _finalize_pandas_output( self, frame: DataFrame, multi_index_named: bool ) -> DataFrame: """ Processes data read in based on kwargs. Parameters ---------- frame : DataFrame The DataFrame to process. multi_index_named : bool Returns ------- DataFrame The processed DataFrame. """ frame = self._do_date_conversions(frame.columns, frame) frame = self._maybe_restore_string_dtype(frame) frame = self._finalize_index(frame, multi_index_named) frame = self._finalize_dtype(frame) # GH#65862 tuples passed via names imply MultiIndex columns, # as with other engines frame.columns = self._maybe_make_multi_index_columns( list(frame.columns), self.col_names ) return frame def _maybe_restore_string_dtype(self, frame: DataFrame) -> DataFrame: # When parse_dates fails to parse a column it falls back to the original # strings; keep the default string dtype rather than object, matching # the other engines. if not (using_string_dtype() and isinstance(self.parse_dates, list)): return frame for col in self.parse_dates: if col in frame.columns and frame[col].dtype == object: frame[col] = frame[col].astype("str") return frame def _validate_usecols(self, usecols) -> None: if lib.is_list_like(usecols) and not all(isinstance(x, str) for x in usecols): raise ValueError( "The pyarrow engine does not allow 'usecols' to be integer " "column positions. Pass a list of string column names instead." ) elif callable(usecols): raise ValueError( "The pyarrow engine does not allow 'usecols' to be a callable." ) def read(self) -> DataFrame: """ Reads the contents of a CSV file into a DataFrame and processes it according to the kwargs passed in the constructor. Returns ------- DataFrame The DataFrame created from the CSV file. """ pa = import_optional_dependency("pyarrow") pyarrow_csv = import_optional_dependency("pyarrow.csv") self._get_pyarrow_options() convert_options = self._get_convert_options() try: table = pyarrow_csv.read_csv( self.src, read_options=pyarrow_csv.ReadOptions(**self.read_options), parse_options=pyarrow_csv.ParseOptions(**self.parse_options), convert_options=convert_options, ) except pa.ArrowInvalid as err: # pyarrow reports "Empty CSV file or block" when it cannot extract # any columns from the source; re-raise as EmptyDataError so the # pyarrow engine matches the "c" and "python" engines. if "Empty CSV file" in str(err): raise EmptyDataError("No columns to parse from file") from err raise ParserError(err) from err except TypeError as err: # pyarrow can only read from a binary buffer; a text-mode file # handle raises a bare "a bytes-like object is required" TypeError. if "a bytes-like object is required" in str(err): raise TypeError( "The 'pyarrow' engine can only read from a binary file " "object; the given file was opened in text mode" ) from err raise if self.usecols_dtype == "empty": # GH#66056 pyarrow ignores an empty ``include_columns`` and returns # every column; match the other engines by returning an empty frame. table = table.select([]).slice(0, 0) dtype_backend = self.kwds["dtype_backend"] # Convert all pa.null() cols -> float64 (non nullable) # else Int64 (nullable case, see below) if dtype_backend is lib.no_default: new_schema = table.schema new_type = pa.float64() for i, arrow_type in enumerate(table.schema.types): if pa.types.is_null(arrow_type): new_schema = new_schema.set( i, new_schema.field(i).with_type(new_type) ) table = table.cast(new_schema) if self.header is not None and self.names is None: new_names, self.unnamed_cols = self._dedup_column_names(table.column_names) if new_names != table.column_names: table = table.rename_columns(new_names) multi_index_named = self._adjust_column_names(table) if isinstance(self.dtype, defaultdict): # GH#41574 materialize the factory default over the actual columns # so missing keys get the default, matching the other engines if self.names is not None: # always set by _adjust_column_names when header is None columns = list(self.names) else: columns = table.column_names self.dtype = {col: self.dtype[col] for col in columns} with warnings.catch_warnings(): warnings.filterwarnings( "ignore", "make_block is deprecated", Pandas4Warning, ) frame = arrow_table_to_pandas( table, dtype_backend=dtype_backend, null_to_int64=True, dtype=self.dtype, names=self.names, ) if self.names is not None: frame.columns = self.names return self._finalize_pandas_output(frame, multi_index_named)