pandas-dev/pandas · #66625

PERF: Arrow-backed groupby reductions avoid converting results through NumPy

fangchenli · merged Sep 19, 20263 files · 32 + / 32
doc/source/whatsnew/v3.1.0.rst1 + / 1
@@ -405,7 +405,7 @@ Deprecations Performance improvements ~~~~~~~~~~~~~~~~~~~~~~~~ - Performance improvement in :class:`.DataFrameGroupBy` aggregations (``sum``, ``mean``, ``min``, ``max``, ``prod``) when the grouping keys are already sorted (:issue:`65103`)-- Performance improvement in :class:`.DataFrameGroupBy` and :class:`.SeriesGroupBy` reductions for :class:`ArrowDtype` decimal columns (``sum``, ``prod``, ``min``, ``max``, ``mean``, ``var``) and for Arrow-backed string columns (``min``, ``max``), i.e. :class:`ArrowDtype` and :class:`StringDtype` with ``storage="pyarrow"``, by dispatching to PyArrow's native ``group_by`` instead of a slower fallback (:issue:`63416`)+- Performance improvement in :class:`.DataFrameGroupBy` and :class:`.SeriesGroupBy` reductions for :class:`ArrowDtype` decimal columns (``sum``, ``prod``, ``min``, ``max``, ``mean``, ``var``) and for Arrow-backed string columns (``min``, ``max``), i.e. :class:`ArrowDtype` and :class:`StringDtype` with ``storage="pyarrow"``, by dispatching to PyArrow's native ``group_by`` instead of a slower fallback and placing the results with PyArrow's ``take`` rather than converting them to NumPy and back (:issue:`63416`, :issue:`66625`) - Performance improvement in :class:`.DataFrameGroupBy` and :class:`.SeriesGroupBy` with ``sort=True``, as well as :func:`factorize` and :func:`merge` with ``sort=True``, when the keys are integers with many unique values (:issue:`66129`) - Performance improvement in casting integer and boolean dtypes to ``string[pyarrow]`` by using PyArrow's native cast instead of element-wise conversion (:issue:`56505`) - Performance improvement in :meth:`DataFrame.__getitem__` when selecting a
pandas/core/arrays/arrow/array.py22 + / 30
@@ -3471,38 +3471,30 @@ def _groupby_op_pyarrow(             )             result_values = pc.if_else(below_min_count, None, result_values) -        # Scatter results into output array ordered by group id.-        # Fallback to NumPy here due to the limitation of pc.scatter.-        # Another workaround is to use join + sort.-        # TODO: revisit this part when pc.scatter becomes more functionally complete.-        result_group_ids_np = result_group_ids.to_numpy(zero_copy_only=False).astype(-            np.int64, copy=False-        )-        result_values_np = result_values.to_numpy(zero_copy_only=False)+        # Place the results in group-id order: the inverse permutation takes+        # the row holding group i, and is null where group i had no rows.+        group_ids_np = result_group_ids.to_numpy(zero_copy_only=False)+        inverse = np.full(ngroups, -1, dtype=np.int64)+        inverse[group_ids_np] = np.arange(len(group_ids_np))+        indices = pa.array(inverse, mask=inverse < 0)++        if how in ["sum", "prod"] and pa.types.is_decimal(output_type):+            try:+                # take would carry an out-of-precision decimal through silently+                result_values.validate(full=True)+            except pa.ArrowInvalid:+                # needs more digits than the maximum precision, so let the+                # caller fall back to a type that can hold it+                return None -        default_py = default_value.as_py()-        try:-            if default_py is not None and min_count == 0:-                # Fill missing groups with identity element-                output_np = np.full(ngroups, default_py, dtype=result_values_np.dtype)-                output_np[result_group_ids_np] = result_values_np-                pa_result = pa.array(output_np, type=output_type)+        pa_result = pc.take(result_values, indices)+        if default_value.as_py() is not None and min_count == 0:+            if result_values.null_count == 0:+                # every null is a group with no rows+                pa_result = _safe_fill_null(pa_result, default_value)             else:-                # Fill missing groups with null-                output_np = np.empty(ngroups, dtype=result_values_np.dtype)-                null_mask = np.ones(ngroups, dtype=bool)-                output_np[result_group_ids_np] = result_values_np-                null_mask[result_group_ids_np] = False-                if result_values.null_count > 0:-                    result_nulls = pc.is_null(result_values).to_numpy()-                    null_mask[result_group_ids_np[result_nulls]] = True-                pa_result = pa.array(output_np, type=output_type, mask=null_mask)-        except pa.ArrowInvalid:-            # A group result does not fit the aggregated type, e.g. a decimal-            # sum or product needing more digits than the maximum precision.-            # Fall back so the result keeps the wider type it needs.-            return None-+                # keep the skipna=False nulls, fill only the empty groups+                pa_result = pc.if_else(pc.is_null(indices), default_value, pa_result)         return self._from_pyarrow_array(pa_result)      def _to_groupby_compatible(self) -> ExtensionArray:
pandas/tests/extension/test_arrow.py9 + / 1
@@ -4500,6 +4500,14 @@ def test_groupby_sum_prod_no_precision_overflow(self, agg_func, expected):                 [Decimal(3 * 10**37)] * 5,                 Decimal(15 * 10**37),             ),+            # same, but the input already has the maximum precision, so the+            # widening cast is a no-op and cannot report the overflow+            (+                "prod",+                pa.decimal128(38, 0),+                [Decimal(10**19), Decimal(12 * 10**18)],+                Decimal(12 * 10**37),+            ),         ],     )     def test_groupby_sum_prod_exceeds_max_precision(@@ -4561,7 +4569,7 @@ def test_groupby_decimal_ddof(self, how, ddof):         ids=["decimal", "ArrowDtype", "str[pyarrow]"],     )     def test_groupby_empty(self, dtype, how):-        # GH#63416 an empty input has no groups to scatter into+        # GH#63416 an empty input has no groups to place results into         ser = pd.Series([], dtype=dtype)         result = getattr(ser.groupby([]), how)()         assert len(result) == 0