pandas-dev/pandas · #66625
PERF: Arrow-backed groupby reductions avoid converting results through NumPy
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 apandas/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