PERF: GroupBy.quantile (#51722)

Author: jbrockmendel

doc/source/whatsnew/v2.1.0.rst

@@ -176,6 +176,7 @@ Other enhancements
 - Performance improvement in :func:`concat` with homogeneous ``np.float64`` or ``np.float32`` dtypes (:issue:`52685`)
 - Performance improvement in :meth:`DataFrame.filter` when ``items`` is given (:issue:`52941`)
 - Reductions :meth:`Series.argmax`, :meth:`Series.argmin`, :meth:`Series.idxmax`, :meth:`Series.idxmin`, :meth:`Index.argmax`, :meth:`Index.argmin`, :meth:`DataFrame.idxmax`, :meth:`DataFrame.idxmin` are now supported for object-dtype objects (:issue:`4279`, :issue:`18021`, :issue:`40685`, :issue:`43697`)
+- Performance improvement in :meth:`GroupBy.quantile` (:issue:`51722`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/_libs/groupby.pyi

@@ -121,10 +121,12 @@ def group_quantile(
     values: np.ndarray,  # ndarray[numeric, ndim=1]
     labels: npt.NDArray[np.intp],
     mask: npt.NDArray[np.uint8],
-    sort_indexer: npt.NDArray[np.intp],  # const
     qs: npt.NDArray[np.float64],  # const
+    starts: npt.NDArray[np.int64],
+    ends: npt.NDArray[np.int64],
     interpolation: Literal["linear", "lower", "higher", "nearest", "midpoint"],
-    result_mask: np.ndarray | None = ...,
+    result_mask: np.ndarray | None,
+    is_datetimelike: bool,
 ) -> None: ...
 def group_last(
     out: np.ndarray,  # rank_t[:, ::1]

pandas/_libs/groupby.pyx

@@ -1177,11 +1177,13 @@ def group_quantile(
     ndarray[float64_t, ndim=2] out,
     ndarray[numeric_t, ndim=1] values,
     ndarray[intp_t] labels,
-    ndarray[uint8_t] mask,
-    const intp_t[:] sort_indexer,
+    const uint8_t[:] mask,
     const float64_t[:] qs,
+    ndarray[int64_t] starts,
+    ndarray[int64_t] ends,
     str interpolation,
-    uint8_t[:, ::1] result_mask=None,
+    uint8_t[:, ::1] result_mask,
+    bint is_datetimelike,
 ) -> None:
     """
     Calculate the quantile per group.
@@ -1194,27 +1196,38 @@ def group_quantile(
         Array containing the values to apply the function against.
     labels : ndarray[np.intp]
         Array containing the unique group labels.
-    sort_indexer : ndarray[np.intp]
-        Indices describing sort order by values and labels.
     qs : ndarray[float64_t]
         The quantile values to search for.
+    starts : ndarray[int64]
+        Positions at which each group begins.
+    ends : ndarray[int64]
+        Positions at which each group ends.
     interpolation : {'linear', 'lower', 'highest', 'nearest', 'midpoint'}
+    result_mask : ndarray[bool, ndim=2] or None
+    is_datetimelike : bool
+        Whether int64 values represent datetime64-like values.
 
     Notes
     -----
     Rather than explicitly returning a value, this function modifies the
     provided `out` parameter.
     """
     cdef:
-        Py_ssize_t i, N=len(labels), ngroups, grp_sz, non_na_sz, k, nqs
-        Py_ssize_t grp_start=0, idx=0
-        intp_t lab
+        Py_ssize_t i, N=len(labels), ngroups, non_na_sz, k, nqs
+        Py_ssize_t idx=0
+        Py_ssize_t grp_size
         InterpolationEnumType interp
         float64_t q_val, q_idx, frac, val, next_val
-        int64_t[::1] counts, non_na_counts
         bint uses_result_mask = result_mask is not None
+        Py_ssize_t start, end
+        ndarray[numeric_t] grp
+        intp_t[::1] sort_indexer
+        const uint8_t[:] sub_mask
 
     assert values.shape[0] == N
+    assert starts is not None
+    assert ends is not None
+    assert len(starts) == len(ends)
 
     if any(not (0 <= q <= 1) for q in qs):
         wrong = [x for x in qs if not (0 <= x <= 1)][0]
@@ -1233,64 +1246,65 @@ def group_quantile(
 
     nqs = len(qs)
     ngroups = len(out)
-    counts = np.zeros(ngroups, dtype=np.int64)
-    non_na_counts = np.zeros(ngroups, dtype=np.int64)
-
-    # First figure out the size of every group
-    with nogil:
-        for i in range(N):
-            lab = labels[i]
-            if lab == -1:  # NA group label
-                continue
 
-            counts[lab] += 1
-            if not mask[i]:
-                non_na_counts[lab] += 1
+    # TODO: get cnp.PyArray_ArgSort to work with nogil so we can restore the rest
+    #  of this function as being `with nogil:`
+    for i in range(ngroups):
+        start = starts[i]
+        end = ends[i]
+
+        grp = values[start:end]
+
+        # Figure out how many group elements there are
+        sub_mask = mask[start:end]
+        grp_size = sub_mask.size
+        non_na_sz = 0
+        for k in range(grp_size):
+            if sub_mask[k] == 0:
+                non_na_sz += 1
+
+        # equiv: sort_indexer = grp.argsort()
+        if is_datetimelike:
+            # We need the argsort to put NaTs at the end, not the beginning
+            sort_indexer = cnp.PyArray_ArgSort(grp.view("M8[ns]"), 0, cnp.NPY_QUICKSORT)
+        else:
+            sort_indexer = cnp.PyArray_ArgSort(grp, 0, cnp.NPY_QUICKSORT)
 
-    with nogil:
-        for i in range(ngroups):
-            # Figure out how many group elements there are
-            grp_sz = counts[i]
-            non_na_sz = non_na_counts[i]
-
-            if non_na_sz == 0:
-                for k in range(nqs):
-                    if uses_result_mask:
-                        result_mask[i, k] = 1
-                    else:
-                        out[i, k] = NaN
-            else:
-                for k in range(nqs):
-                    q_val = qs[k]
+        if non_na_sz == 0:
+            for k in range(nqs):
+                if uses_result_mask:
+                    result_mask[i, k] = 1
+                else:
+                    out[i, k] = NaN
+        else:
+            for k in range(nqs):
+                q_val = qs[k]
 
-                    # Calculate where to retrieve the desired value
-                    # Casting to int will intentionally truncate result
-                    idx = grp_start + <int64_t>(q_val * <float64_t>(non_na_sz - 1))
+                # Calculate where to retrieve the desired value
+                # Casting to int will intentionally truncate result
+                idx = <int64_t>(q_val * <float64_t>(non_na_sz - 1))
 
-                    val = values[sort_indexer[idx]]
-                    # If requested quantile falls evenly on a particular index
-                    # then write that index's value out. Otherwise interpolate
-                    q_idx = q_val * (non_na_sz - 1)
-                    frac = q_idx % 1
+                val = grp[sort_indexer[idx]]
+                # If requested quantile falls evenly on a particular index
+                # then write that index's value out. Otherwise interpolate
+                q_idx = q_val * (non_na_sz - 1)
+                frac = q_idx % 1
 
-                    if frac == 0.0 or interp == INTERPOLATION_LOWER:
-                        out[i, k] = val
-                    else:
-                        next_val = values[sort_indexer[idx + 1]]
-                        if interp == INTERPOLATION_LINEAR:
-                            out[i, k] = val + (next_val - val) * frac
-                        elif interp == INTERPOLATION_HIGHER:
+                if frac == 0.0 or interp == INTERPOLATION_LOWER:
+                    out[i, k] = val
+                else:
+                    next_val = grp[sort_indexer[idx + 1]]
+                    if interp == INTERPOLATION_LINEAR:
+                        out[i, k] = val + (next_val - val) * frac
+                    elif interp == INTERPOLATION_HIGHER:
+                        out[i, k] = next_val
+                    elif interp == INTERPOLATION_MIDPOINT:
+                        out[i, k] = (val + next_val) / 2.0
+                    elif interp == INTERPOLATION_NEAREST:
+                        if frac > .5 or (frac == .5 and q_val > .5):  # Always OK?
                             out[i, k] = next_val
-                        elif interp == INTERPOLATION_MIDPOINT:
-                            out[i, k] = (val + next_val) / 2.0
-                        elif interp == INTERPOLATION_NEAREST:
-                            if frac > .5 or (frac == .5 and q_val > .5):  # Always OK?
-                                out[i, k] = next_val
-                            else:
-                                out[i, k] = val
-
-            # Increment the index reference in sorted_arr for the next group
-            grp_start += grp_sz
+                        else:
+                            out[i, k] = val
 
 
 # ----------------------------------------------------------------------

pandas/core/groupby/groupby.py

@@ -4226,6 +4226,16 @@ def quantile(
         a    2.0
         b    3.0
         """
+        mgr = self._get_data_to_aggregate(numeric_only=numeric_only, name="quantile")
+        obj = self._wrap_agged_manager(mgr)
+        if self.axis == 1:
+            splitter = self.grouper._get_splitter(obj.T, axis=self.axis)
+            sdata = splitter._sorted_data.T
+        else:
+            splitter = self.grouper._get_splitter(obj, axis=self.axis)
+            sdata = splitter._sorted_data
+
+        starts, ends = lib.generate_slices(splitter._slabels, splitter.ngroups)
 
         def pre_processor(vals: ArrayLike) -> tuple[np.ndarray, DtypeObj | None]:
             if is_object_dtype(vals.dtype):
@@ -4323,24 +4333,24 @@ def post_processor(
 
             return vals
 
-        orig_scalar = is_scalar(q)
-        if orig_scalar:
+        qs = np.array(q, dtype=np.float64)
+        pass_qs: np.ndarray | None = qs
+        if is_scalar(q):
             qs = np.array([q], dtype=np.float64)
-        else:
-            qs = np.array(q, dtype=np.float64)
+            pass_qs = None
+
         ids, _, ngroups = self.grouper.group_info
         nqs = len(qs)
 
         func = partial(
-            libgroupby.group_quantile, labels=ids, qs=qs, interpolation=interpolation
+            libgroupby.group_quantile,
+            labels=ids,
+            qs=qs,
+            interpolation=interpolation,
+            starts=starts,
+            ends=ends,
         )
 
-        # Put '-1' (NaN) labels as the last group so it does not interfere
-        # with the calculations. Note: length check avoids failure on empty
-        # labels. In that case, the value doesn't matter
-        na_label_for_sorting = ids.max() + 1 if len(ids) > 0 else 0
-        labels_for_lexsort = np.where(ids == -1, na_label_for_sorting, ids)
-
         def blk_func(values: ArrayLike) -> ArrayLike:
             orig_vals = values
             if isinstance(values, BaseMaskedArray):
@@ -4357,53 +4367,44 @@ def blk_func(values: ArrayLike) -> ArrayLike:
             ncols = 1
             if vals.ndim == 2:
                 ncols = vals.shape[0]
-                shaped_labels = np.broadcast_to(
-                    labels_for_lexsort, (ncols, len(labels_for_lexsort))
-                )
-            else:
-                shaped_labels = labels_for_lexsort
 
             out = np.empty((ncols, ngroups, nqs), dtype=np.float64)
 
-            # Get an index of values sorted by values and then labels
-            order = (vals, shaped_labels)
-            sort_arr = np.lexsort(order).astype(np.intp, copy=False)
-
             if is_datetimelike:
-                # This casting needs to happen after the lexsort in order
-                #  to ensure that NaTs are placed at the end and not the front
-                vals = vals.view("i8").astype(np.float64)
+                vals = vals.view("i8")
 
             if vals.ndim == 1:
-                # Ea is always 1d
+                # EA is always 1d
                 func(
                     out[0],
                     values=vals,
                     mask=mask,
-                    sort_indexer=sort_arr,
                     result_mask=result_mask,
+                    is_datetimelike=is_datetimelike,
                 )
             else:
                 for i in range(ncols):
-                    func(out[i], values=vals[i], mask=mask[i], sort_indexer=sort_arr[i])
+                    func(
+                        out[i],
+                        values=vals[i],
+                        mask=mask[i],
+                        result_mask=None,
+                        is_datetimelike=is_datetimelike,
+                    )
 
             if vals.ndim == 1:
                 out = out.ravel("K")
                 if result_mask is not None:
                     result_mask = result_mask.ravel("K")
             else:
                 out = out.reshape(ncols, ngroups * nqs)
+
             return post_processor(out, inference, result_mask, orig_vals)
 
-        data = self._get_data_to_aggregate(numeric_only=numeric_only, name="quantile")
-        res_mgr = data.grouped_reduce(blk_func)
+        res_mgr = sdata._mgr.grouped_reduce(blk_func)
 
         res = self._wrap_agged_manager(res_mgr)
-
-        if orig_scalar:
-            # Avoid expensive MultiIndex construction
-            return self._wrap_aggregated_output(res)
-        return self._wrap_aggregated_output(res, qs=qs)
+        return self._wrap_aggregated_output(res, qs=pass_qs)
 
     @final
     @Substitution(name="groupby")