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 typing import TYPE_CHECKING
4import warnings
6from pandas._libs import lib
7from pandas.compat._optional import import_optional_dependency
8from pandas.errors import (
9 Pandas4Warning,
10 ParserError,
11 ParserWarning,
12)
13from pandas.util._exceptions import (
14 find_stack_level,
15)
17from pandas.core.dtypes.common import (
18 pandas_dtype,
19)
20from pandas.core.dtypes.inference import is_integer
22from pandas.io._util import arrow_table_to_pandas
23from pandas.io.parsers.base_parser import ParserBase
25if TYPE_CHECKING:
26 import pyarrow as pa
28 from pandas._typing import ReadBuffer
30 from pandas import DataFrame
33class ArrowParserWrapper(ParserBase):
34 """
35 Wrapper for the pyarrow engine for read_csv()
36 """
38 def __init__(self, src: ReadBuffer[bytes], **kwds) -> None:
39 super().__init__(kwds)
40 self.kwds = kwds
41 self.src = src
43 self._parse_kwds()
45 def _parse_kwds(self) -> None:
46 """
47 Validates keywords before passing to pyarrow.
48 """
49 encoding: str | None = self.kwds.get("encoding")
50 self.encoding = "utf-8" if encoding is None else encoding
52 na_values = self.kwds["na_values"]
53 if isinstance(na_values, dict):
54 raise ValueError(
55 "The pyarrow engine doesn't support passing a dict for na_values"
56 )
57 self.na_values = list(self.kwds["na_values"])
59 def _get_pyarrow_options(self) -> None:
60 """
61 Rename some arguments to pass to pyarrow
62 """
63 mapping = {
64 "usecols": "include_columns",
65 "na_values": "null_values",
66 "escapechar": "escape_char",
67 "skip_blank_lines": "ignore_empty_lines",
68 "decimal": "decimal_point",
69 "quotechar": "quote_char",
70 }
71 for pandas_name, pyarrow_name in mapping.items():
72 if pandas_name in self.kwds and self.kwds.get(pandas_name) is not None:
73 self.kwds[pyarrow_name] = self.kwds.pop(pandas_name)
75 # Date format handling
76 # If we get a string, we need to convert it into a list for pyarrow
77 # If we get a dict, we want to parse those separately
78 date_format = self.date_format
79 if isinstance(date_format, str):
80 date_format = [date_format]
81 else:
82 # In case of dict, we don't want to propagate through, so
83 # just set to pyarrow default of None
85 # Ideally, in future we disable pyarrow dtype inference (read in as string)
86 # to prevent misreads.
87 date_format = None
88 self.kwds["timestamp_parsers"] = date_format
90 self.parse_options = {
91 option_name: option_value
92 for option_name, option_value in self.kwds.items()
93 if option_value is not None
94 and option_name
95 in ("delimiter", "quote_char", "escape_char", "ignore_empty_lines")
96 }
98 on_bad_lines = self.kwds.get("on_bad_lines")
99 if on_bad_lines is not None:
100 if callable(on_bad_lines):
101 self.parse_options["invalid_row_handler"] = on_bad_lines
102 elif on_bad_lines == ParserBase.BadLineHandleMethod.ERROR:
103 self.parse_options["invalid_row_handler"] = (
104 None # PyArrow raises an exception by default
105 )
106 elif on_bad_lines == ParserBase.BadLineHandleMethod.WARN:
108 def handle_warning(invalid_row) -> str:
109 warnings.warn(
110 f"Expected {invalid_row.expected_columns} columns, but found "
111 f"{invalid_row.actual_columns}: {invalid_row.text}",
112 ParserWarning,
113 stacklevel=find_stack_level(),
114 )
115 return "skip"
117 self.parse_options["invalid_row_handler"] = handle_warning
118 elif on_bad_lines == ParserBase.BadLineHandleMethod.SKIP:
119 self.parse_options["invalid_row_handler"] = lambda _: "skip"
121 self.convert_options = {
122 option_name: option_value
123 for option_name, option_value in self.kwds.items()
124 if option_value is not None
125 and option_name
126 in (
127 "include_columns",
128 "null_values",
129 "true_values",
130 "false_values",
131 "decimal_point",
132 "timestamp_parsers",
133 )
134 }
135 self.convert_options["strings_can_be_null"] = "" in self.kwds["null_values"]
136 # autogenerated column names are prefixed with 'f' in pyarrow.csv
137 if self.header is None and "include_columns" in self.convert_options:
138 self.convert_options["include_columns"] = [
139 f"f{n}" for n in self.convert_options["include_columns"]
140 ]
142 self.read_options = {
143 "autogenerate_column_names": self.header is None,
144 "skip_rows": self.header
145 if self.header is not None
146 else self.kwds["skiprows"],
147 "encoding": self.encoding,
148 }
150 def _get_convert_options(self):
151 pyarrow_csv = import_optional_dependency("pyarrow.csv")
153 try:
154 convert_options = pyarrow_csv.ConvertOptions(**self.convert_options)
155 except TypeError as err:
156 include = self.convert_options.get("include_columns", None)
157 if include is not None:
158 self._validate_usecols(include)
160 nulls = self.convert_options.get("null_values", set())
161 if not lib.is_list_like(nulls) or not all(
162 isinstance(x, str) for x in nulls
163 ):
164 raise TypeError(
165 "The 'pyarrow' engine requires all na_values to be strings"
166 ) from err
168 raise
170 return convert_options
172 def _adjust_column_names(self, table: pa.Table) -> bool:
173 num_cols = len(table.columns)
174 multi_index_named = True
175 if self.header is None:
176 if self.names is None:
177 self.names = range(num_cols)
178 if len(self.names) != num_cols:
179 # usecols is passed through to pyarrow, we only handle index col here
180 # The only way self.names is not the same length as number of cols is
181 # if we have int index_col. We should just pad the names(they will get
182 # removed anyways) to expected length then.
183 columns_prefix = [str(x) for x in range(num_cols - len(self.names))]
184 self.names = columns_prefix + self.names
185 multi_index_named = False
186 return multi_index_named
188 def _finalize_index(self, frame: DataFrame, multi_index_named: bool) -> DataFrame:
189 if self.index_col is not None:
190 index_to_set = self.index_col.copy()
191 for i, item in enumerate(self.index_col):
192 if is_integer(item):
193 index_to_set[i] = frame.columns[item]
194 # String case
195 elif item not in frame.columns:
196 raise ValueError(f"Index {item} invalid")
198 # Process dtype for index_col and drop from dtypes
199 if self.dtype is not None:
200 key, new_dtype = (
201 (item, self.dtype.get(item))
202 if self.dtype.get(item) is not None
203 else (frame.columns[item], self.dtype.get(frame.columns[item]))
204 )
205 if new_dtype is not None:
206 frame[key] = frame[key].astype(new_dtype)
207 del self.dtype[key]
209 frame.set_index(index_to_set, drop=True, inplace=True)
210 # Clear names if headerless and no name given
211 if self.header is None and not multi_index_named:
212 frame.index.names = [None] * len(frame.index.names)
214 return frame
216 def _finalize_dtype(self, frame: DataFrame) -> DataFrame:
217 if self.dtype is not None:
218 # Ignore non-existent columns from dtype mapping
219 # like other parsers do
220 if isinstance(self.dtype, dict):
221 self.dtype = {
222 k: pandas_dtype(v)
223 for k, v in self.dtype.items()
224 if k in frame.columns
225 }
226 else:
227 self.dtype = pandas_dtype(self.dtype)
228 try:
229 frame = frame.astype(self.dtype)
230 except TypeError as err:
231 # GH#44901 reraise to keep api consistent
232 raise ValueError(str(err)) from err
233 return frame
235 def _finalize_pandas_output(
236 self, frame: DataFrame, multi_index_named: bool
237 ) -> DataFrame:
238 """
239 Processes data read in based on kwargs.
241 Parameters
242 ----------
243 frame : DataFrame
244 The DataFrame to process.
245 multi_index_named : bool
247 Returns
248 -------
249 DataFrame
250 The processed DataFrame.
251 """
252 frame = self._do_date_conversions(frame.columns, frame)
253 frame = self._finalize_index(frame, multi_index_named)
254 frame = self._finalize_dtype(frame)
255 return frame
257 def _validate_usecols(self, usecols) -> None:
258 if lib.is_list_like(usecols) and not all(isinstance(x, str) for x in usecols):
259 raise ValueError(
260 "The pyarrow engine does not allow 'usecols' to be integer "
261 "column positions. Pass a list of string column names instead."
262 )
263 elif callable(usecols):
264 raise ValueError(
265 "The pyarrow engine does not allow 'usecols' to be a callable."
266 )
268 def read(self) -> DataFrame:
269 """
270 Reads the contents of a CSV file into a DataFrame and
271 processes it according to the kwargs passed in the
272 constructor.
274 Returns
275 -------
276 DataFrame
277 The DataFrame created from the CSV file.
278 """
279 pa = import_optional_dependency("pyarrow")
280 pyarrow_csv = import_optional_dependency("pyarrow.csv")
281 self._get_pyarrow_options()
282 convert_options = self._get_convert_options()
284 try:
285 table = pyarrow_csv.read_csv(
286 self.src,
287 read_options=pyarrow_csv.ReadOptions(**self.read_options),
288 parse_options=pyarrow_csv.ParseOptions(**self.parse_options),
289 convert_options=convert_options,
290 )
291 except pa.ArrowInvalid as e:
292 raise ParserError(e) from e
294 dtype_backend = self.kwds["dtype_backend"]
296 # Convert all pa.null() cols -> float64 (non nullable)
297 # else Int64 (nullable case, see below)
298 if dtype_backend is lib.no_default:
299 new_schema = table.schema
300 new_type = pa.float64()
301 for i, arrow_type in enumerate(table.schema.types):
302 if pa.types.is_null(arrow_type):
303 new_schema = new_schema.set(
304 i, new_schema.field(i).with_type(new_type)
305 )
307 table = table.cast(new_schema)
309 multi_index_named = self._adjust_column_names(table)
311 with warnings.catch_warnings():
312 warnings.filterwarnings(
313 "ignore",
314 "make_block is deprecated",
315 Pandas4Warning,
316 )
317 frame = arrow_table_to_pandas(
318 table,
319 dtype_backend=dtype_backend,
320 null_to_int64=True,
321 dtype=self.dtype,
322 names=self.names,
323 )
325 if self.header is None:
326 frame.columns = self.names
328 return self._finalize_pandas_output(frame, multi_index_named)