pandas-dev/pandas · #70769

REF: build the DataFrame in ParserBase.read for the C and Python engines (GH#43135)

jbrockmendel · merged Oct 17, 20268 files · 127 + / 113 −
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         -------         DataFrame
pandas/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).