pandas-dev/pandas · #70769
REF: build the DataFrame in ParserBase.read for the C and Python engines (GH#43135)
doc/source/whatsnew/v3.2.0.rst1 + / 0 −
@@ -126,6 +126,7 @@ Other deprecations Performance improvements ~~~~~~~~~~~~~~~~~~~~~~~~ - Performance improvement in :func:`read_csv` with ``engine="c"`` across the board, in tokenizing and in the machinery shared by all column types; large uncompressed local files are additionally read in parallel across multiple CPU cores by default, see :ref:`above <whatsnew_310.enhancements.read_csv>` (:issue:`64347`, :issue:`64515`, :issue:`65350`, :issue:`66271`, :issue:`66272`, :issue:`66273`, :issue:`66274`, :issue:`66276`, :issue:`66798`, :issue:`68314`, :issue:`69632`, :issue:`69748`, :issue:`70140`)+- Performance improvement in :func:`read_csv` with a dict ``dtype`` on wide files or with a small ``chunksize`` (:issue:`70769`) - Performance improvement in :meth:`DataFrame.skew` and :meth:`DataFrame.kurt` with types different from ``float64`` (:issue:`66662`) - Performance improvement in :meth:`DataFrameGroupBy.pct_change` and :meth:`SeriesGroupBy.pct_change` (:issue:`70972`) - Performance improvement in :meth:`PeriodIndex.strftime` with a custom ``date_format`` (:issue:`70106`)pandas/io/parsers/arrow_parser_wrapper.py6 + / 1 −
@@ -560,12 +560,17 @@ def _validate_usecols(self, usecols) -> None: "The pyarrow engine does not allow 'usecols' to be a callable." ) - def read(self) -> DataFrame:+ def read(self, nrows: int | None = None, row_offset: int = 0) -> DataFrame: """ Reads the contents of a CSV file into a DataFrame and processes it according to the kwargs passed in the constructor. + Parameters+ ----------+ nrows, row_offset+ Ignored; the pyarrow engine always reads the whole file.+ Returns ------- DataFramepandas/io/parsers/base_parser.py71 + / 1 −
@@ -24,6 +24,7 @@ from pandas._libs.parsers import STR_NA_VALUES from pandas.compat._optional import import_optional_dependency from pandas.errors import (+ AbstractMethodError, ParserError, ParserWarning, )@@ -59,6 +60,7 @@ from pandas.core.indexes.api import ( Index, MultiIndex,+ RangeIndex, default_index, ensure_index_from_sequences, )@@ -76,6 +78,7 @@ ) from pandas._typing import (+ AnyArrayLike, ArrayLike, Dtype, DtypeArg,@@ -127,7 +130,9 @@ def __init__(self, kwds) -> None: self.na_filter = kwds.get("na_filter", False) self.keep_default_na = kwds.get("keep_default_na", True) - self.dtype = copy(kwds.get("dtype", None))+ # self.dtype may gain mangled duplicate-column keys; read() uses the original+ self._orig_dtype = kwds.get("dtype", None)+ self.dtype = copy(self._orig_dtype) self.converters = kwds.get("converters") self.dtype_backend = kwds.get("dtype_backend") @@ -181,6 +186,40 @@ def __init__(self, kwds) -> None: def close(self) -> None: pass + def _read_arrays(+ self, nrows: int | None = None+ ) -> tuple[+ Index | None,+ Sequence[Hashable] | MultiIndex,+ Mapping[Hashable, AnyArrayLike],+ ]:+ raise AbstractMethodError(self)++ def read(self, nrows: int | None = None, row_offset: int = 0) -> DataFrame:+ """+ Read up to ``nrows`` rows into a DataFrame.++ ``row_offset`` is the start of the default index, so successive+ chunks continue the row numbering.+ """+ index, columns, col_dict = self._read_arrays(nrows)++ if index is None and col_dict:+ # Any column is actually fine:+ new_rows = len(next(iter(col_dict.values())))+ index = RangeIndex(row_offset, row_offset + new_rows)++ col_dict = wrap_object_columns(col_dict, self._orig_dtype, index)++ return DataFrame(+ col_dict,+ # error: Argument "columns" to "DataFrame" has incompatible type+ # "Sequence[Hashable] | MultiIndex"; expected "... | None"+ columns=columns, # type: ignore[arg-type]+ index=index,+ copy=False,+ )+ @final def _warn_parser(self, msg: str) -> None: if self._warning_sink is not None:@@ -800,6 +839,37 @@ def _get_empty_meta( return index, columns, col_dict +def wrap_object_columns(+ col_dict: Mapping[Hashable, AnyArrayLike],+ dtype_arg: DtypeArg | None,+ index: Index | None,+) -> Mapping[Hashable, AnyArrayLike]:+ """+ Wrap the columns ``dtype_arg`` requests as object/str in a Series of that dtype.++ Otherwise the DataFrame constructor would infer them as str, see GH#56047.+ """+ if dtype_arg is None:+ return col_dict+ if isinstance(dtype_arg, dict):+ requested = {key: dtype_arg.get(key) for key in col_dict}+ elif pandas_dtype(dtype_arg) in (np.str_, np.object_):+ requested = dict.fromkeys(col_dict, dtype_arg)+ else:+ return col_dict+ return {+ key: (+ Series(values, index=index, dtype=requested[key], copy=False)+ if requested[key] is not None+ and pandas_dtype(requested[key]) in (np.str_, np.object_)+ # GH#57512 don't undo a parse_dates conversion+ and values.dtype.kind != "M"+ else values+ )+ for key, values in col_dict.items()+ }++ def date_converter( date_col, col: Hashable,pandas/io/parsers/c_parser_wrapper.py3 + / 3 −
@@ -68,7 +68,7 @@ class CParserWrapper(ParserBase): low_memory: bool _reader: parsers.TextReader _exhausted: bool- # When False, read() leaves deferred string columns as raw pending+ # When False, _read_arrays() leaves deferred string columns as raw pending # handles for the caller to materialize (one chunked ExtensionArray per # column, e.g. post-gather in a parallel read). wrap_deferred: bool = True@@ -106,7 +106,7 @@ def __init__(self, src: ReadCsvBuffer[str], **kwds) -> None: import_optional_dependency("pyarrow") self._reader = parsers.TextReader(src, **kwds) # Let the pyarrow string fast path return raw pending-column handles;- # read() wraps them into one ExtensionArray per column at the end.+ # _read_arrays() wraps them into one ExtensionArray per column at the end. self._reader.defer_pa_wrap = True self.unnamed_cols = self._reader.unnamed_cols@@ -278,7 +278,7 @@ def _low_memory_column_labels( ) return {position: named.get(position) for position in positions} - def read(+ def _read_arrays( self, nrows: int | None = None, ) -> tuple[pandas/io/parsers/python_parser.py3 + / 16 −
@@ -276,15 +276,15 @@ def _read(): return reader - def read(- self, rows: int | None = None+ def _read_arrays(+ self, nrows: int | None = None ) -> tuple[ Index | None, Sequence[Hashable] | MultiIndex, Mapping[Hashable, ArrayLike | Series], ]: try:- content = self._get_lines(rows)+ content = self._get_lines(nrows) except StopIteration: if self._first_chunk: content = []@@ -356,19 +356,6 @@ def _exclude_implicit_index( name: alldata[i + offset] for i, name in enumerate(names) if i < len_alldata }, names - # legacy- def get_chunk(- self, size: int | None = None- ) -> tuple[- Index | None,- Sequence[Hashable] | MultiIndex,- Mapping[Hashable, ArrayLike | Series],- ]:- if size is None:- # error: "PythonParser" has no attribute "chunksize"- size = self.chunksize # type: ignore[attr-defined]- return self.read(rows=size)- def _convert_data( self, data: Mapping[Hashable, np.ndarray],pandas/io/parsers/readers.py18 + / 92 −
@@ -9,7 +9,6 @@ import codecs from collections import ( abc,- defaultdict, deque, ) from concurrent.futures import ThreadPoolExecutor@@ -95,6 +94,7 @@ evaluate_callable_usecols, is_index_col, parser_defaults,+ wrap_object_columns, ) from pandas.io.parsers.c_parser_wrapper import ( CParserWrapper,@@ -1105,7 +1105,7 @@ def _worker() -> None: try: # Raw column arrays, not DataFrames: per-chunk frame assembly # holds the GIL and the concat would redo the block machinery.- _, chunk_columns, col_dict = reader._engine.read()+ _, chunk_columns, col_dict = reader._engine._read_arrays() except Exception: chunk_failed.set() return@@ -1133,8 +1133,8 @@ def _piece_ndarray(piece) -> np.ndarray | None: return None dtype_arg = kwds.get("dtype")- # Mirror TextFileReader.read: a dict or str/object dtype needs a- # Series-per-column wrap, so it skips direct block assembly.+ # A dict or str/object dtype may need object columns wrapped (see+ # wrap_object_columns), so it skips direct block assembly. needs_series_wrap = isinstance(dtype_arg, dict) or ( dtype_arg is not None and pandas_dtype(dtype_arg) in (np.str_, np.object_) )@@ -1309,25 +1309,12 @@ def _piece_ndarray(piece) -> np.ndarray | None: if needs_series_wrap: data = _concatenate_chunks(chunk_dicts, columns, warn_mixed=False)- if isinstance(dtype_arg, dict):- dtype: defaultdict[Hashable, Any] = defaultdict(lambda: None)- dtype.update(dtype_arg)- else:- dtype = defaultdict(lambda: dtype_arg)- new_col_dict = {- k: Series(- v,- index=index,- dtype=(- dtype[k]- if pandas_dtype(dtype[k]) in (np.str_, np.object_)- else None- ),- copy=False,- )- for k, v in data.items()- }- return DataFrame(new_col_dict, columns=columns, index=index, copy=False)+ return DataFrame(+ wrap_object_columns(data, dtype_arg, index),+ columns=columns,+ index=index,+ copy=False,+ ) if leftover_pos: leftover_names = [col_list[pos] for pos in leftover_pos]@@ -3205,76 +3192,15 @@ def read(self, nrows: int | None = None) -> DataFrame: >>> with pd.read_csv("data.csv", iterator=True) as reader: # doctest: +SKIP ... df = reader.read() """- if self.engine == "pyarrow":- try:- # error: "ParserBase" has no attribute "read"- df = self._engine.read() # type: ignore[attr-defined]- except Exception:- self.close()- raise- else:+ if self.engine != "pyarrow":+ # the pyarrow engine ignores nrows nrows = validate_integer("nrows", nrows)- try:- # error: "ParserBase" has no attribute "read"- (- index,- columns,- col_dict,- ) = self._engine.read( # type: ignore[attr-defined]- nrows- )- except Exception:- self.close()- raise-- if index is None:- if col_dict:- # Any column is actually fine:- new_rows = len(next(iter(col_dict.values())))- index = RangeIndex(self._currow, self._currow + new_rows)- else:- new_rows = 0- else:- new_rows = len(index)-- if hasattr(self, "orig_options"):- dtype_arg = self.orig_options.get("dtype", None)- else:- dtype_arg = None-- if isinstance(dtype_arg, dict):- dtype = defaultdict(lambda: None) # type: ignore[var-annotated]- dtype.update(dtype_arg)- elif dtype_arg is not None and pandas_dtype(dtype_arg) in (- np.str_,- np.object_,- ):- dtype = defaultdict(lambda: dtype_arg)- else:- dtype = None-- if dtype is not None:- new_col_dict = {}- for k, v in col_dict.items():- # GH#57512 don't undo a parse_dates conversion- d = (- dtype[k]- if pandas_dtype(dtype[k]) in (np.str_, np.object_)- and v.dtype.kind != "M"- else None- )- new_col_dict[k] = Series(v, index=index, dtype=d, copy=False)- else:- new_col_dict = col_dict-- df = DataFrame(- new_col_dict,- columns=columns,- index=index,- copy=False,- )-- self._currow += new_rows+ try:+ df = self._engine.read(nrows, row_offset=self._currow)+ except Exception:+ self.close()+ raise+ self._currow += len(df) return df def get_chunk(self, size: int | None = None) -> DataFrame:pandas/tests/io/parser/dtypes/test_dtypes_basic.py12 + / 0 −
@@ -22,6 +22,8 @@ import pandas._testing as tm from pandas.core.arrays import IntegerArray +from pandas.io.parsers.base_parser import wrap_object_columns+ xfail_pyarrow = pytest.mark.usefixtures("pyarrow_xfail") @@ -1091,3 +1093,13 @@ def test_unsafe_integer_dtype_raises_for_non_int64_source(all_parsers, data, dty parser = all_parsers with pytest.raises(ValueError, match="cannot safely convert passed user dtype"): parser.read_csv(StringIO(data), dtype={"x": dtype})+++def test_wrap_object_columns_only_wraps_object():+ # wrapping every column would be correct but slow on wide frames+ ints = np.array([1, 2])+ strings = np.array(["x", "y"], dtype=object)+ dtype = {"a": "int64", "b": object}+ result = wrap_object_columns({"a": ints, "b": strings}, dtype, pd.RangeIndex(2))+ assert result["a"] is ints+ tm.assert_series_equal(result["b"], pd.Series(strings, dtype=object))pandas/tests/io/parser/test_parallel_read.py13 + / 0 −
@@ -1031,6 +1031,19 @@ def test_parallel_dup_names_dict_dtype_matches_serial(tmp_path, monkeypatch): assert result.dtypes["a"] == result.dtypes["a.1"] +def test_parallel_object_dtype_dict_matches_serial(tmp_path, monkeypatch):+ # an object dtype must survive DataFrame construction, see GH#56047+ raw = b"a,b,c\n" + b"".join(f"x{i},y{i},{i}\n".encode() for i in range(500))+ path = tmp_path / "obj.csv"+ path.write_bytes(raw)++ dtype = {"a": object, "c": "float32"}+ result = _read_forced_parallel(path, monkeypatch, dtype=dtype)+ expected = pd.read_csv(io.BytesIO(raw), dtype=dtype)+ tm.assert_frame_equal(result, expected)+ assert result.dtypes["a"] == object++ def test_parallel_comment_before_header_matches_serial(tmp_path, monkeypatch): # A full-line comment above the header shifts the preamble byte offset; # the header row must not be ingested as data (GH#64347).