Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Provide classes to perform the groupby aggregate operations.
4These are not exposed to the user and provide implementations of the grouping
5operations, primarily in cython. These classes (BaseGrouper and BinGrouper)
6are contained *in* the SeriesGroupBy and DataFrameGroupBy objects.
7"""
9from __future__ import annotations
11import collections
12import functools
13from typing import (
14 TYPE_CHECKING,
15 Any,
16 Generic,
17 final,
18)
19import warnings
21import numpy as np
23from pandas._libs import (
24 NaT,
25 lib,
26)
27import pandas._libs.groupby as libgroupby
28from pandas._typing import (
29 ArrayLike,
30 AxisInt,
31 NDFrameT,
32 Shape,
33 npt,
34)
35from pandas.errors import (
36 AbstractMethodError,
37 Pandas4Warning,
38)
39from pandas.util._decorators import cache_readonly
40from pandas.util._exceptions import find_stack_level
42from pandas.core.dtypes.cast import (
43 maybe_downcast_to_dtype,
44)
45from pandas.core.dtypes.common import (
46 ensure_float64,
47 ensure_int64,
48 ensure_platform_int,
49 ensure_uint64,
50 is_1d_only_ea_dtype,
51)
52from pandas.core.dtypes.missing import (
53 isna,
54 maybe_fill,
55)
57from pandas.core.arrays import Categorical
58from pandas.core.frame import DataFrame
59from pandas.core.groupby import grouper
60from pandas.core.indexes.api import (
61 CategoricalIndex,
62 Index,
63 MultiIndex,
64 ensure_index,
65)
66from pandas.core.series import Series
67from pandas.core.sorting import (
68 compress_group_index,
69 decons_obs_group_ids,
70 get_group_index,
71 get_group_index_sorter,
72)
74if TYPE_CHECKING:
75 from collections.abc import (
76 Callable,
77 Generator,
78 Hashable,
79 Iterator,
80 )
82 from pandas.core.generic import NDFrame
85def check_result_array(obj, dtype) -> None:
86 # Our operation is supposed to be an aggregation/reduction. If
87 # it returns an ndarray, this likely means an invalid operation has
88 # been passed. See test_apply_without_aggregation, test_agg_must_agg
89 if isinstance(obj, np.ndarray):
90 if dtype != object:
91 # If it is object dtype, the function can be a reduction/aggregation
92 # and still return an ndarray e.g. test_agg_over_numpy_arrays
93 raise ValueError("Must produce aggregated value")
96def extract_result(res):
97 """
98 Extract the result object, it might be a 0-dim ndarray
99 or a len-1 0-dim, or a scalar
100 """
101 if hasattr(res, "_values"):
102 # Preserve EA
103 res = res._values
104 if res.ndim == 1 and len(res) == 1:
105 warnings.warn(
106 "Converting a Series or array of length 1 into a scalar is "
107 "deprecated and will be removed in a future version. If you wish "
108 "to preserve the current behavior, have ``func`` return scalars.",
109 category=Pandas4Warning,
110 stacklevel=find_stack_level(),
111 )
112 # see test_agg_lambda_with_timezone, test_resampler_grouper.py::test_apply
113 res = res[0]
114 return res
117_REDUCEAT_UFUNCS: dict[str, np.ufunc] = {
118 "sum": np.add,
119 "prod": np.multiply,
120 "min": np.minimum,
121 "max": np.maximum,
122 "mean": np.add,
123}
126class WrappedCythonOp:
127 """
128 Dispatch logic for functions defined in _libs.groupby
130 Parameters
131 ----------
132 how: str
133 Operation name, e.g. "mean".
134 has_dropped_na: bool
135 True precisely when dropna=True and the grouper contains a null value.
136 """
138 # Functions for which we do _not_ attempt to cast the cython result
139 # back to the original dtype.
140 cast_blocklist = frozenset(
141 ["any", "all", "rank", "count", "size", "idxmin", "idxmax"]
142 )
144 def __init__(self, how: str, has_dropped_na: bool) -> None:
145 self.kind = (
146 "aggregate" if how in self._CYTHON_FUNCTIONS["aggregate"] else "transform"
147 )
148 self.how = how
149 self.has_dropped_na = has_dropped_na
151 _CYTHON_FUNCTIONS: dict[str, dict] = {
152 "aggregate": {
153 "any": functools.partial(libgroupby.group_any_all, val_test="any"),
154 "all": functools.partial(libgroupby.group_any_all, val_test="all"),
155 "sum": "group_sum",
156 "prod": "group_prod",
157 "idxmin": functools.partial(libgroupby.group_idxmin_idxmax, name="idxmin"),
158 "idxmax": functools.partial(libgroupby.group_idxmin_idxmax, name="idxmax"),
159 "min": "group_min",
160 "max": "group_max",
161 "mean": "group_mean",
162 "median": "group_median_float64",
163 "var": "group_var",
164 "std": functools.partial(libgroupby.group_var, name="std"),
165 "sem": functools.partial(libgroupby.group_var, name="sem"),
166 "skew": "group_skew",
167 "kurt": "group_kurt",
168 "first": "group_nth",
169 "last": "group_last",
170 "ohlc": "group_ohlc",
171 },
172 "transform": {
173 "cumprod": "group_cumprod",
174 "cumsum": "group_cumsum",
175 "cummin": "group_cummin",
176 "cummax": "group_cummax",
177 "rank": "group_rank",
178 },
179 }
181 _cython_arity = {"ohlc": 4} # OHLC
183 # Note: we make this a classmethod and pass kind+how so that caching
184 # works at the class level and not the instance level
185 @classmethod
186 @functools.cache
187 def _get_cython_function(
188 cls, kind: str, how: str, dtype: np.dtype, is_numeric: bool
189 ):
190 dtype_str = dtype.name
191 ftype = cls._CYTHON_FUNCTIONS[kind][how]
193 # see if there is a fused-type version of function
194 # only valid for numeric
195 if callable(ftype):
196 f = ftype
197 else:
198 f = getattr(libgroupby, ftype)
199 if is_numeric:
200 return f
201 elif dtype == np.dtype(object):
202 if how in ["median", "cumprod"]:
203 # no fused types -> no __signatures__
204 raise NotImplementedError(
205 f"function is not implemented for this dtype: "
206 f"[how->{how},dtype->{dtype_str}]"
207 )
208 elif how in ["std", "sem", "idxmin", "idxmax"]:
209 # We have a partial object that does not have __signatures__
210 return f
211 elif how in ["skew", "kurt"]:
212 # _get_cython_vals will convert to float64
213 pass
214 elif "object" not in f.__signatures__:
215 # raise NotImplementedError here rather than TypeError later
216 raise NotImplementedError(
217 f"function is not implemented for this dtype: "
218 f"[how->{how},dtype->{dtype_str}]"
219 )
220 return f
221 else:
222 raise NotImplementedError(
223 "This should not be reached. Please report a bug at "
224 "github.com/pandas-dev/pandas/",
225 dtype,
226 )
228 def _get_cython_vals(
229 self, values: np.ndarray, uses_mask: bool = False
230 ) -> np.ndarray:
231 """
232 Cast numeric dtypes to float64 for functions that only support that.
234 Parameters
235 ----------
236 values : np.ndarray
238 Returns
239 -------
240 values : np.ndarray
241 """
242 how = self.how
244 if how in ["median", "std", "sem", "skew", "kurt"]:
245 # median only has a float64 implementation
246 # We should only get here with is_numeric, as non-numeric cases
247 # should raise in _get_cython_function
248 values = ensure_float64(values)
250 elif values.dtype.kind in "iu":
251 if how in ["var", "mean"] or (
252 self.kind == "transform" and self.has_dropped_na and not uses_mask
253 ):
254 # has_dropped_na check need for test_null_group_str_transformer
255 # result may still include NaN, so we have to cast
256 values = ensure_float64(values)
258 elif how in ["sum", "ohlc", "prod", "cumsum", "cumprod"]:
259 # Avoid overflow during group op
260 if values.dtype.kind == "i":
261 values = ensure_int64(values)
262 else:
263 values = ensure_uint64(values)
265 return values
267 def _get_output_shape(self, ngroups: int, values: np.ndarray) -> Shape:
268 how = self.how
269 kind = self.kind
271 arity = self._cython_arity.get(how, 1)
273 out_shape: Shape
274 if how == "ohlc":
275 out_shape = (ngroups, arity)
276 elif arity > 1:
277 raise NotImplementedError(
278 "arity of more than 1 is not supported for the 'how' argument"
279 )
280 elif kind == "transform":
281 out_shape = values.shape
282 else:
283 out_shape = (ngroups, *values.shape[1:])
284 return out_shape
286 def _get_out_dtype(self, dtype: np.dtype) -> np.dtype:
287 how = self.how
289 if how == "rank":
290 out_dtype = "float64"
291 elif how in ["idxmin", "idxmax"]:
292 # The Cython implementation only produces the row number; we'll take
293 # from the index using this in post processing
294 out_dtype = "intp"
295 elif dtype.kind in "iufcb":
296 out_dtype = f"{dtype.kind}{dtype.itemsize}"
297 else:
298 out_dtype = "object"
299 return np.dtype(out_dtype)
301 def _get_result_dtype(self, dtype: np.dtype) -> np.dtype:
302 """
303 Get the desired dtype of a result based on the
304 input dtype and how it was computed.
306 Parameters
307 ----------
308 dtype : np.dtype
310 Returns
311 -------
312 np.dtype
313 The desired dtype of the result.
314 """
315 how = self.how
317 if how in ["sum", "cumsum", "sum", "prod", "cumprod"]:
318 if dtype == np.dtype(bool):
319 return np.dtype(np.int64)
320 elif how in ["mean", "median", "var", "std", "sem"]:
321 if dtype.kind in "fc":
322 return dtype
323 elif dtype.kind in "iub":
324 return np.dtype(np.float64)
325 return dtype
327 @final
328 def _cython_op_ndim_compat(
329 self,
330 values: np.ndarray,
331 *,
332 min_count: int,
333 ngroups: int,
334 comp_ids: np.ndarray,
335 mask: npt.NDArray[np.bool_] | None = None,
336 result_mask: npt.NDArray[np.bool_] | None = None,
337 initial: Any = 0,
338 **kwargs,
339 ) -> np.ndarray:
340 if values.ndim == 1:
341 # expand to 2d, dispatch, then squeeze if appropriate
342 values2d = values[None, :]
343 if mask is not None:
344 mask = mask[None, :]
345 if result_mask is not None:
346 result_mask = result_mask[None, :]
347 res = self._call_cython_op(
348 values2d,
349 min_count=min_count,
350 ngroups=ngroups,
351 comp_ids=comp_ids,
352 mask=mask,
353 result_mask=result_mask,
354 initial=initial,
355 **kwargs,
356 )
357 if res.shape[0] == 1:
358 return res[0]
360 # otherwise we have OHLC
361 return res.T
363 return self._call_cython_op(
364 values,
365 min_count=min_count,
366 ngroups=ngroups,
367 comp_ids=comp_ids,
368 mask=mask,
369 result_mask=result_mask,
370 initial=initial,
371 **kwargs,
372 )
374 @final
375 def _call_cython_op(
376 self,
377 values: np.ndarray, # np.ndarray[ndim=2]
378 *,
379 min_count: int,
380 ngroups: int,
381 comp_ids: np.ndarray,
382 mask: npt.NDArray[np.bool_] | None,
383 result_mask: npt.NDArray[np.bool_] | None,
384 initial: Any = 0,
385 **kwargs,
386 ) -> np.ndarray: # np.ndarray[ndim=2]
387 orig_values = values
389 dtype = values.dtype
390 is_numeric = dtype.kind in "iufcb"
392 is_datetimelike = dtype.kind in "mM"
394 if self.how in ["any", "all"]:
395 if mask is None:
396 mask = isna(values)
398 if is_datetimelike:
399 values = values.view("int64")
400 is_numeric = True
401 elif dtype.kind == "b":
402 values = values.view("uint8")
403 if values.dtype == "float16":
404 values = values.astype(np.float32)
406 if self.how in ["any", "all"]:
407 if dtype == object:
408 if kwargs["skipna"]:
409 # GH#37501: don't raise on pd.NA when skipna=True
410 if mask is not None and mask.any():
411 # mask on original values computed separately
412 values = values.copy()
413 values[mask] = True
414 values = values.astype(bool, copy=False).view(np.int8)
415 is_numeric = True
417 values = values.T
418 if mask is not None:
419 mask = mask.T
420 if result_mask is not None:
421 result_mask = result_mask.T
423 out_shape = self._get_output_shape(ngroups, values)
424 func = self._get_cython_function(self.kind, self.how, values.dtype, is_numeric)
425 values = self._get_cython_vals(values, uses_mask=mask is not None)
426 out_dtype = self._get_out_dtype(values.dtype)
428 result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
429 if self.kind == "aggregate":
430 counts = np.zeros(ngroups, dtype=np.int64)
431 if self.how in [
432 "idxmin",
433 "idxmax",
434 "min",
435 "max",
436 "mean",
437 "last",
438 "first",
439 "sum",
440 "median",
441 ]:
442 if self.how == "sum":
443 # pass in through kwargs only for sum (other functions don't have
444 # the keyword)
445 kwargs["initial"] = initial
446 func(
447 out=result,
448 counts=counts,
449 values=values,
450 labels=comp_ids,
451 min_count=min_count,
452 mask=mask,
453 result_mask=result_mask,
454 is_datetimelike=is_datetimelike,
455 **kwargs,
456 )
457 elif self.how in ["sem", "std", "var", "ohlc", "prod"]:
458 if self.how in ["std", "sem"]:
459 kwargs["is_datetimelike"] = is_datetimelike
460 func(
461 result,
462 counts,
463 values,
464 comp_ids,
465 min_count=min_count,
466 mask=mask,
467 result_mask=result_mask,
468 **kwargs,
469 )
470 elif self.how in ["any", "all"]:
471 func(
472 out=result,
473 values=values,
474 labels=comp_ids,
475 mask=mask,
476 result_mask=result_mask,
477 **kwargs,
478 )
479 result = result.astype(bool, copy=False)
480 elif self.how in ["skew", "kurt"]:
481 func(
482 out=result,
483 counts=counts,
484 values=values,
485 labels=comp_ids,
486 mask=mask,
487 result_mask=result_mask,
488 **kwargs,
489 )
490 if dtype == object:
491 result = result.astype(object)
493 else:
494 raise NotImplementedError(f"{self.how} is not implemented")
495 else:
496 # TODO: min_count
497 if self.how != "rank":
498 # TODO: should rank take result_mask?
499 kwargs["result_mask"] = result_mask
500 func(
501 out=result,
502 values=values,
503 labels=comp_ids,
504 ngroups=ngroups,
505 is_datetimelike=is_datetimelike,
506 mask=mask,
507 **kwargs,
508 )
510 if self.kind == "aggregate" and self.how not in ["idxmin", "idxmax"]:
511 # i.e. counts is defined. Locations where count<min_count
512 # need to have the result set to np.nan, which may require casting,
513 # see GH#40767. For idxmin/idxmax is handled specially via post-processing
514 if result.dtype.kind in "iu" and not is_datetimelike:
515 # if the op keeps the int dtypes, we have to use 0
516 cutoff = max(0 if self.how in ["sum", "prod"] else 1, min_count)
517 empty_groups = counts < cutoff
518 if empty_groups.any():
519 if result_mask is not None:
520 assert result_mask[empty_groups].all()
521 else:
522 # Note: this conversion could be lossy, see GH#40767
523 result = result.astype("float64")
524 result[empty_groups] = np.nan
526 result = result.T
528 if self.how not in self.cast_blocklist:
529 # e.g. if we are int64 and need to restore to datetime64/timedelta64
530 # "rank" is the only member of cast_blocklist we get here
531 # Casting only needed for float16, bool, datetimelike,
532 # and self.how in ["sum", "prod", "ohlc", "cumprod"]
533 res_dtype = self._get_result_dtype(orig_values.dtype)
534 op_result = maybe_downcast_to_dtype(result, res_dtype)
535 else:
536 op_result = result
538 return op_result
540 @final
541 def _validate_axis(self, axis: AxisInt, values: ArrayLike) -> None:
542 if values.ndim > 2:
543 raise NotImplementedError("number of dimensions is currently limited to 2")
544 if values.ndim == 2:
545 assert axis == 1, axis
546 elif not is_1d_only_ea_dtype(values.dtype):
547 # Note: it is *not* the case that axis is always 0 for 1-dim values,
548 # as we can have 1D ExtensionArrays that we need to treat as 2D
549 assert axis == 0
551 @final
552 def cython_operation(
553 self,
554 *,
555 values: ArrayLike,
556 axis: AxisInt,
557 min_count: int = -1,
558 comp_ids: np.ndarray,
559 ngroups: int,
560 **kwargs,
561 ) -> ArrayLike:
562 """
563 Call our cython function, with appropriate pre- and post- processing.
564 """
565 self._validate_axis(axis, values)
567 if not isinstance(values, np.ndarray):
568 # i.e. ExtensionArray
569 return values._groupby_op(
570 how=self.how,
571 has_dropped_na=self.has_dropped_na,
572 min_count=min_count,
573 ngroups=ngroups,
574 ids=comp_ids,
575 **kwargs,
576 )
578 return self._cython_op_ndim_compat(
579 values,
580 min_count=min_count,
581 ngroups=ngroups,
582 comp_ids=comp_ids,
583 mask=None,
584 **kwargs,
585 )
588class BaseGrouper:
589 """
590 This is an internal Grouper class, which actually holds
591 the generated groups
593 Parameters
594 ----------
595 axis : Index
596 groupings : Sequence[Grouping]
597 all the grouping instances to handle in this grouper
598 for example for grouper list to groupby, need to pass the list
599 sort : bool, default True
600 whether this grouper will give sorted result or not
602 """
604 axis: Index
606 def __init__(
607 self,
608 axis: Index,
609 groupings: list[grouper.Grouping],
610 sort: bool = True,
611 dropna: bool = True,
612 ) -> None:
613 assert isinstance(axis, Index), axis
615 self.axis = axis
616 self._groupings = groupings
617 self._sort = sort
618 self.dropna = dropna
619 self._is_resample = False
621 @property
622 def groupings(self) -> list[grouper.Grouping]:
623 return self._groupings
625 def __iter__(self) -> Iterator[Hashable]:
626 return iter(self.indices)
628 @property
629 def nkeys(self) -> int:
630 return len(self.groupings)
632 def get_iterator(self, data: NDFrameT) -> Iterator[tuple[Hashable, NDFrameT]]:
633 """
634 Groupby iterator
636 Returns
637 -------
638 Generator yielding sequence of (name, subsetted object)
639 for each group
640 """
641 splitter = self._get_splitter(data)
642 # TODO: Would be more efficient to skip unobserved for transforms
643 keys = self.result_index
644 yield from zip(keys, splitter, strict=True)
646 @final
647 def _get_splitter(self, data: NDFrame) -> DataSplitter:
648 """
649 Returns
650 -------
651 Generator yielding subsetted objects
652 """
653 if isinstance(data, Series):
654 klass: type[DataSplitter] = SeriesSplitter
655 else:
656 # i.e. DataFrame
657 klass = FrameSplitter
659 return klass(
660 data,
661 self.ngroups,
662 sorted_ids=self._sorted_ids,
663 sort_idx=self.result_ilocs,
664 )
666 @cache_readonly
667 def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]:
668 """dict {group name -> group indices}"""
669 if len(self.groupings) == 1 and isinstance(self.result_index, CategoricalIndex):
670 # This shows unused categories in indices GH#38642
671 result = self.groupings[0].indices
672 else:
673 result_index, ids = self.result_index_and_ids
674 values = result_index._values
675 categories = Categorical.from_codes(
676 ids, categories=range(len(result_index))
677 )
678 result = {
679 # mypy is not aware that group has to be an integer
680 values[group]: axis_ilocs # type: ignore[call-overload]
681 for group, axis_ilocs in categories._reverse_indexer().items()
682 if len(axis_ilocs)
683 }
684 if not self.dropna:
685 has_mi = isinstance(self.result_index, MultiIndex)
686 if not has_mi and self.result_index.hasnans:
687 result = {
688 np.nan if isna(key) else key: value for key, value in result.items()
689 }
690 elif has_mi:
691 # MultiIndex has no efficient way to tell if there are NAs
692 result = {
693 # error: "Hashable" has no attribute "__iter__" (not iterable)
694 tuple(np.nan if isna(comp) else comp for comp in key): value # type: ignore[attr-defined]
695 for key, value in result.items()
696 }
698 return result
700 @final
701 @cache_readonly
702 def result_ilocs(self) -> npt.NDArray[np.intp]:
703 """
704 Get the original integer locations of result_index in the input.
705 """
706 # Original indices are where group_index would go via sorting.
707 # But when dropna is true, we need to remove null values while accounting for
708 # any gaps that then occur because of them.
709 ids = self.ids
711 if self.has_dropped_na:
712 mask = np.where(ids >= 0)
713 # Count how many gaps are caused by previous null values for each position
714 null_gaps = np.cumsum(ids == -1)[mask]
715 ids = ids[mask]
717 result = get_group_index_sorter(ids, self.ngroups)
719 if self.has_dropped_na:
720 # Shift by the number of prior null gaps
721 result += np.take(null_gaps, result)
723 return result
725 @property
726 def codes(self) -> list[npt.NDArray[np.signedinteger]]:
727 return [ping.codes for ping in self.groupings]
729 @property
730 def levels(self) -> list[Index]:
731 if len(self.groupings) > 1:
732 # mypy doesn't know result_index must be a MultiIndex
733 return list(self.result_index.levels) # type: ignore[attr-defined]
734 else:
735 return [self.result_index]
737 @property
738 def names(self) -> list[Hashable]:
739 return [ping.name for ping in self.groupings]
741 @final
742 def size(self) -> Series:
743 """
744 Compute group sizes.
745 """
746 ids = self.ids
747 ngroups = self.ngroups
748 out: np.ndarray | list
749 if ngroups:
750 if self.has_dropped_na:
751 out = np.bincount(ids + 1, minlength=ngroups + 1)[1:]
752 else:
753 out = np.bincount(ids, minlength=ngroups)
754 else:
755 out = []
756 return Series(out, index=self.result_index, dtype="int64", copy=False)
758 @cache_readonly
759 def groups(self) -> dict[Hashable, Index]:
760 """dict {group name -> group labels}"""
761 if len(self.groupings) == 1:
762 return self.groupings[0].groups
763 result_index, ids = self.result_index_and_ids
764 values = result_index._values
765 categories = Categorical.from_codes(ids, categories=range(len(result_index)))
766 result = {
767 # mypy is not aware that group has to be an integer
768 values[group]: self.axis.take(axis_ilocs) # type: ignore[call-overload]
769 for group, axis_ilocs in categories._reverse_indexer().items()
770 }
771 return result
773 @final
774 @cache_readonly
775 def is_monotonic(self) -> bool:
776 # return if my group orderings are monotonic
777 return Index(self.ids, copy=False).is_monotonic_increasing
779 @final
780 @cache_readonly
781 def has_dropped_na(self) -> bool:
782 """
783 Whether grouper has null value(s) that are dropped.
784 """
785 return bool((self.ids < 0).any())
787 @final
788 @cache_readonly
789 def _reduceat_segment_info(
790 self,
791 ) -> tuple[int, npt.NDArray[np.intp], npt.NDArray[np.intp], np.ndarray] | None:
792 """
793 Precompute segment boundaries for _reduceat_aggregate.
795 Returns (na_count, group_starts, segment_ids, segment_sizes),
796 or None when comp_ids is empty after dropping NA rows.
797 """
798 comp_ids = self.ids
800 if self.has_dropped_na:
801 na_count = int((comp_ids < 0).sum())
802 comp_ids = comp_ids[na_count:]
803 else:
804 na_count = 0
806 if comp_ids.size == 0:
807 return None
809 group_starts = np.flatnonzero(np.r_[True, comp_ids[1:] != comp_ids[:-1]])
810 segment_ids = comp_ids[group_starts]
811 segment_sizes = np.diff(group_starts, append=comp_ids.size)
813 return na_count, group_starts, segment_ids, segment_sizes
815 @cache_readonly
816 def codes_info(self) -> npt.NDArray[np.intp]:
817 # return the codes of items in original grouped axis
818 return self.ids
820 @final
821 @cache_readonly
822 def ngroups(self) -> int:
823 return len(self.result_index)
825 @property
826 def result_index(self) -> Index:
827 return self.result_index_and_ids[0]
829 @property
830 def ids(self) -> npt.NDArray[np.intp]:
831 return self.result_index_and_ids[1]
833 @cache_readonly
834 def result_index_and_ids(self) -> tuple[Index, npt.NDArray[np.intp]]:
835 levels = [
836 Index._with_infer(ping.uniques, copy=False) for ping in self.groupings
837 ]
838 obs = [
839 ping._observed or not ping._passed_categorical for ping in self.groupings
840 ]
841 sorts = [ping._sort for ping in self.groupings]
842 # When passed a categorical grouping, keep all categories
843 for k, (ping, level) in enumerate(zip(self.groupings, levels, strict=True)):
844 if ping._passed_categorical:
845 levels[k] = level.set_categories(ping._orig_cats) # type: ignore[attr-defined]
847 if len(self.groupings) == 1:
848 result_index = levels[0]
849 result_index.name = self.names[0]
850 ids = ensure_platform_int(self.codes[0])
851 elif all(obs):
852 result_index, ids = self._ob_index_and_ids(
853 levels, self.codes, self.names, sorts
854 )
855 elif not any(obs):
856 result_index, ids = self._unob_index_and_ids(levels, self.codes, self.names)
857 else:
858 # Combine unobserved and observed parts
859 names = self.names
860 codes = [ping.codes for ping in self.groupings]
861 ob_indices = [idx for idx, ob in enumerate(obs) if ob]
862 unob_indices = [idx for idx, ob in enumerate(obs) if not ob]
863 ob_index, ob_ids = self._ob_index_and_ids(
864 levels=[levels[idx] for idx in ob_indices],
865 codes=[codes[idx] for idx in ob_indices],
866 names=[names[idx] for idx in ob_indices],
867 sorts=[sorts[idx] for idx in ob_indices],
868 )
869 unob_index, unob_ids = self._unob_index_and_ids(
870 levels=[levels[idx] for idx in unob_indices],
871 codes=[codes[idx] for idx in unob_indices],
872 names=[names[idx] for idx in unob_indices],
873 )
875 result_index_codes = np.concatenate(
876 [
877 np.tile(unob_index.codes, len(ob_index)),
878 np.repeat(ob_index.codes, len(unob_index), axis=1),
879 ],
880 axis=0,
881 )
882 _, index = np.unique(unob_indices + ob_indices, return_index=True)
883 result_index = MultiIndex(
884 levels=list(unob_index.levels) + list(ob_index.levels),
885 codes=result_index_codes,
886 names=list(unob_index.names) + list(ob_index.names),
887 ).reorder_levels(index)
889 # A dropped NA key is -1 in ob_ids/unob_ids and needs to come through
890 # as -1 in `ids`; the sum can be non-negative, so test the operands.
891 ids = len(unob_index) * ob_ids + unob_ids
892 if self.dropna:
893 ids = np.where((ob_ids < 0) | (unob_ids < 0), -1, ids)
895 if any(sorts):
896 # Sort result_index and recode ids using the new order
897 n_levels = len(sorts)
898 drop_levels = [
899 n_levels - idx
900 for idx, sort in enumerate(reversed(sorts), 1)
901 if not sort
902 ]
903 if len(drop_levels) > 0:
904 sorter = result_index._drop_level_numbers(drop_levels).argsort()
905 else:
906 sorter = result_index.argsort()
907 result_index = result_index.take(sorter)
908 _, index = np.unique(sorter, return_index=True)
909 # ids is -1 for the dropped NA keys; the sentinel at the end
910 # of recode maps those back to -1 instead of to a position.
911 recode = np.empty(len(index) + 1, dtype=np.intp)
912 recode[:-1] = index
913 recode[-1] = -1
914 ids = recode.take(ids)
915 else:
916 # Recode ids and reorder result_index with observed groups up front,
917 # unobserved at the end
918 ids, uniques = compress_group_index(ids, sort=False)
919 ids = ensure_platform_int(ids)
920 taker = np.concatenate(
921 [uniques, np.delete(np.arange(len(result_index)), uniques)]
922 )
923 result_index = result_index.take(taker)
925 return result_index, ids
927 @property
928 def observed_grouper(self) -> BaseGrouper:
929 if all(ping._observed for ping in self.groupings):
930 return self
932 return self._observed_grouper
934 @cache_readonly
935 def _observed_grouper(self) -> BaseGrouper:
936 groupings = [ping.observed_grouping for ping in self.groupings]
937 grouper = BaseGrouper(self.axis, groupings, sort=self._sort, dropna=self.dropna)
938 return grouper
940 def _ob_index_and_ids(
941 self,
942 levels: list[Index],
943 codes: list[npt.NDArray[np.intp]],
944 names: list[Hashable],
945 sorts: list[bool],
946 ) -> tuple[MultiIndex, npt.NDArray[np.intp]]:
947 consistent_sorting = all(sorts[0] == sort for sort in sorts[1:])
948 sort_in_compress = sorts[0] if consistent_sorting else False
949 shape = tuple(len(level) for level in levels)
950 group_index = get_group_index(codes, shape, sort=True, xnull=True)
951 ob_ids, obs_group_ids = compress_group_index(group_index, sort=sort_in_compress)
952 ob_ids = ensure_platform_int(ob_ids)
953 ob_index_codes = decons_obs_group_ids(
954 ob_ids, obs_group_ids, shape, codes, xnull=True
955 )
956 ob_index = MultiIndex(
957 levels=levels,
958 codes=ob_index_codes,
959 names=names,
960 verify_integrity=False,
961 )
962 if not consistent_sorting and len(ob_index) > 0:
963 # Sort by the levels where the corresponding sort argument is True
964 n_levels = len(sorts)
965 drop_levels = [
966 n_levels - idx
967 for idx, sort in enumerate(reversed(sorts), 1)
968 if not sort
969 ]
970 if len(drop_levels) > 0:
971 sorter = ob_index._drop_level_numbers(drop_levels).argsort()
972 else:
973 sorter = ob_index.argsort()
974 ob_index = ob_index.take(sorter)
975 _, index = np.unique(sorter, return_index=True)
976 ob_ids = np.where(ob_ids == -1, -1, index.take(ob_ids))
977 ob_ids = ensure_platform_int(ob_ids)
978 return ob_index, ob_ids
980 def _unob_index_and_ids(
981 self,
982 levels: list[Index],
983 codes: list[npt.NDArray[np.intp]],
984 names: list[Hashable],
985 ) -> tuple[MultiIndex, npt.NDArray[np.intp]]:
986 shape = tuple(len(level) for level in levels)
987 unob_ids = get_group_index(codes, shape, sort=True, xnull=True)
988 unob_index = MultiIndex.from_product(levels, names=names)
989 unob_ids = ensure_platform_int(unob_ids)
990 return unob_index, unob_ids
992 @final
993 def get_group_levels(self) -> Generator[Index]:
994 # Note: only called from _insert_inaxis_grouper, which
995 # is only called for BaseGrouper, never for BinGrouper
996 result_index = self.result_index
997 if len(self.groupings) == 1:
998 yield result_index
999 else:
1000 for level in range(result_index.nlevels - 1, -1, -1):
1001 yield result_index.get_level_values(level)
1003 # ------------------------------------------------------------
1004 # Aggregation functions
1006 @final
1007 def _cython_operation(
1008 self,
1009 values,
1010 how: str,
1011 axis: AxisInt,
1012 min_count: int = -1,
1013 **kwargs,
1014 ) -> ArrayLike:
1015 """
1016 Returns the values of a cython operation.
1017 """
1018 if (
1019 how in _REDUCEAT_UFUNCS
1020 and isinstance(values, np.ndarray)
1021 and values.dtype.kind in "iufb"
1022 and self.is_monotonic
1023 ):
1024 result = self._reduceat_aggregate(values, how, axis, min_count)
1025 if result is not None:
1026 return result
1028 cy_op = WrappedCythonOp(how=how, has_dropped_na=self.has_dropped_na)
1030 return cy_op.cython_operation(
1031 values=values,
1032 axis=axis,
1033 min_count=min_count,
1034 comp_ids=self.ids,
1035 ngroups=self.ngroups,
1036 **kwargs,
1037 )
1039 @final
1040 def _reduceat_aggregate(
1041 self,
1042 values: np.ndarray,
1043 how: str,
1044 axis: AxisInt,
1045 min_count: int,
1046 ) -> np.ndarray | None:
1047 """
1048 Use ufunc.reduceat for sorted group ids. Returns None to fall back.
1049 """
1050 info = self._reduceat_segment_info
1051 if info is None:
1052 return None
1054 na_count, group_starts, segment_ids, segment_sizes = info
1055 ngroups = self.ngroups
1056 orig_dtype = values.dtype
1058 # reduceat propagates NaN, but the Cython kernels skip it
1059 if orig_dtype.kind == "f" and np.isnan(values).any():
1060 return None
1062 # dropped-NA rows (comp_ids < 0) are at the front when monotonic
1063 if na_count > 0:
1064 if axis == 0:
1065 values = values[na_count:]
1066 else:
1067 values = values[:, na_count:]
1069 # avoid overflow, matching _get_cython_vals
1070 work_dtype: np.dtype
1071 if how == "mean":
1072 work_dtype = np.dtype(np.float64)
1073 elif how in ("sum", "prod") and orig_dtype.kind in "ib":
1074 work_dtype = np.dtype(np.int64)
1075 elif how in ("sum", "prod") and orig_dtype.kind == "u":
1076 work_dtype = np.dtype(np.uint64)
1077 else:
1078 work_dtype = orig_dtype
1080 work_values = values.astype(work_dtype, copy=False)
1082 ufunc = _REDUCEAT_UFUNCS[how]
1083 seg = ufunc.reduceat(work_values, group_starts, axis=axis)
1085 if how == "mean":
1086 sizes = segment_sizes.astype(np.float64)
1087 if seg.ndim == 1:
1088 seg = seg / sizes
1089 else:
1090 shape = [1] * seg.ndim
1091 shape[axis] = len(sizes)
1092 seg = seg / sizes.reshape(shape)
1094 # inline _get_result_dtype for the ops we handle
1095 res_dtype: np.dtype
1096 if how == "mean":
1097 res_dtype = np.dtype(np.float64)
1098 elif how in ("sum", "prod") and orig_dtype.kind == "b":
1099 res_dtype = np.dtype(np.int64)
1100 else:
1101 res_dtype = orig_dtype
1103 # ngroups may exceed present segments
1104 out_shape = list(values.shape)
1105 out_shape[axis] = ngroups
1107 if how == "sum":
1108 result = np.zeros(out_shape, dtype=work_dtype)
1109 elif how == "prod":
1110 result = np.ones(out_shape, dtype=work_dtype)
1111 elif work_dtype.kind == "f":
1112 # mean, or float min/max — NaN for empty groups
1113 result = np.full(out_shape, np.nan, dtype=work_dtype)
1114 # int/bool min/max: fill with dtype extreme; empty groups
1115 # will be converted to NaN below
1116 elif work_dtype.kind == "b":
1117 result = np.full(out_shape, how == "min", dtype=work_dtype)
1118 else:
1119 iinfo = np.iinfo(work_dtype)
1120 fill = iinfo.max if how == "min" else iinfo.min
1121 result = np.full(out_shape, fill, dtype=work_dtype)
1123 indexer: list[Any] = [slice(None)] * len(out_shape)
1124 indexer[axis] = segment_ids
1125 result[tuple(indexer)] = seg
1127 # see _call_cython_op for the min_count logic
1128 if how in ("sum", "prod"):
1129 cutoff = max(0, min_count)
1130 else:
1131 cutoff = max(1, min_count)
1133 if cutoff > 0:
1134 group_counts = np.zeros(ngroups, dtype=np.int64)
1135 group_counts[segment_ids] = segment_sizes
1136 empty_mask = group_counts < cutoff
1137 if empty_mask.any():
1138 if result.dtype.kind in "iub":
1139 result = result.astype(np.float64)
1140 indexer2: list[Any] = [slice(None)] * result.ndim
1141 indexer2[axis] = empty_mask
1142 result[tuple(indexer2)] = np.nan
1144 result = maybe_downcast_to_dtype(result, res_dtype)
1146 return result
1148 @final
1149 def agg_series(self, obj: Series, func: Callable) -> ArrayLike:
1150 """
1151 Parameters
1152 ----------
1153 obj : Series
1154 func : function taking a Series and returning a scalar-like
1156 Returns
1157 -------
1158 np.ndarray or ExtensionArray
1159 """
1160 result = self._aggregate_series_pure_python(obj, func)
1161 return obj.array._cast_pointwise_result(result)
1163 @final
1164 def _aggregate_series_pure_python(
1165 self, obj: Series, func: Callable
1166 ) -> npt.NDArray[np.object_]:
1167 result = np.empty(self.ngroups, dtype="O")
1168 initialized = False
1170 splitter = self._get_splitter(obj)
1172 for i, group in enumerate(splitter):
1173 res = func(group)
1174 res = extract_result(res)
1176 if self._is_resample and not initialized:
1177 # We only do this validation on the first iteration
1178 check_result_array(res, group.dtype)
1179 initialized = True
1181 result[i] = res
1183 return result
1185 @final
1186 def apply_groupwise(
1187 self, f: Callable, data: DataFrame | Series
1188 ) -> tuple[list, bool]:
1189 mutated = False
1190 splitter = self._get_splitter(data)
1191 group_keys = self.result_index
1192 result_values = []
1194 # This calls DataSplitter.__iter__
1195 zipped = zip(group_keys, splitter, strict=True)
1197 for key, group in zipped:
1198 # Pinning name is needed for
1199 # test_group_apply_once_per_group,
1200 # test_inconsistent_return_type, test_set_group_name,
1201 # test_group_name_available_in_inference_pass,
1202 # test_groupby_multi_timezone
1203 object.__setattr__(group, "name", key)
1205 # group might be modified
1206 group_axes = group.axes
1207 res = f(group)
1208 if not mutated and not _is_indexed_like(res, group_axes):
1209 mutated = True
1210 result_values.append(res)
1211 # getattr pattern for __name__ is needed for functools.partial objects
1212 if len(group_keys) == 0 and getattr(f, "__name__", None) in [
1213 "skew",
1214 "kurt",
1215 "sum",
1216 "prod",
1217 ]:
1218 # If group_keys is empty, then no function calls have been made,
1219 # so we will not have raised even if this is an invalid dtype.
1220 # So do one dummy call here to raise appropriate TypeError.
1221 f(data.iloc[:0])
1223 return result_values, mutated
1225 # ------------------------------------------------------------
1226 # Methods for sorting subsets of our GroupBy's object
1228 @final
1229 @cache_readonly
1230 def _sorted_ids(self) -> npt.NDArray[np.intp]:
1231 result = self.ids.take(self.result_ilocs)
1232 if getattr(self, "dropna", True):
1233 # BinGrouper has no dropna
1234 result = result[result >= 0]
1235 return result
1238class BinGrouper(BaseGrouper):
1239 """
1240 This is an internal Grouper class
1242 Parameters
1243 ----------
1244 bins : the split index of binlabels to group the item of axis
1245 binlabels : the label list
1246 indexer : np.ndarray[np.intp], optional
1247 the indexer created by Grouper
1248 some groupers (TimeGrouper) will sort its axis and its
1249 group_info is also sorted, so need the indexer to reorder
1251 Examples
1252 --------
1253 bins: [2, 4, 6, 8, 10]
1254 binlabels: DatetimeIndex(['2005-01-01', '2005-01-03',
1255 '2005-01-05', '2005-01-07', '2005-01-09'],
1256 dtype='datetime64[ns]', freq='2D')
1258 the group_info, which contains the label of each item in grouped
1259 axis, the index of label in label list, group number, is
1261 (array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4]), array([0, 1, 2, 3, 4]), 5)
1263 means that, the grouped axis has 10 items, can be grouped into 5
1264 labels, the first and second items belong to the first label, the
1265 third and forth items belong to the second label, and so on
1267 """
1269 bins: npt.NDArray[np.int64]
1270 binlabels: Index
1272 def __init__(
1273 self,
1274 bins,
1275 binlabels,
1276 indexer=None,
1277 ) -> None:
1278 self.bins = ensure_int64(bins)
1279 self.binlabels = ensure_index(binlabels)
1280 self.indexer = indexer
1282 # These lengths must match, otherwise we could call agg_series
1283 # with empty self.bins, which would raise later.
1284 assert len(self.binlabels) == len(self.bins)
1286 @cache_readonly
1287 def groups(self):
1288 """dict {group name -> group labels}"""
1289 # this is mainly for compat
1290 # GH 3881
1291 result = {
1292 key: value
1293 for key, value in zip(self.binlabels, self.bins, strict=True)
1294 if key is not NaT
1295 }
1296 return result
1298 @property
1299 def nkeys(self) -> int:
1300 # still matches len(self.groupings), but we can hard-code
1301 return 1
1303 @cache_readonly
1304 def codes_info(self) -> npt.NDArray[np.intp]:
1305 # return the codes of items in original grouped axis
1306 ids = self.ids
1307 if self.indexer is not None:
1308 sorter = np.lexsort((ids, self.indexer))
1309 ids = ids[sorter]
1310 return ids
1312 def get_iterator(self, data: NDFrame):
1313 """
1314 Groupby iterator
1316 Returns
1317 -------
1318 Generator yielding sequence of (name, subsetted object)
1319 for each group
1320 """
1321 slicer = lambda start, edge: data.iloc[start:edge]
1323 start: np.int64 | int = 0
1324 for edge, label in zip(self.bins, self.binlabels, strict=True):
1325 if label is not NaT:
1326 yield label, slicer(start, edge)
1327 start = edge
1329 if start < len(data):
1330 yield self.binlabels[-1], slicer(start, None)
1332 @cache_readonly
1333 def indices(self):
1334 indices = collections.defaultdict(list)
1336 i: np.int64 | int = 0
1337 for label, bin in zip(self.binlabels, self.bins, strict=True):
1338 if i < bin:
1339 if label is not NaT:
1340 indices[label] = list(range(i, bin))
1341 i = bin
1342 return indices
1344 @cache_readonly
1345 def codes(self) -> list[npt.NDArray[np.intp]]:
1346 return [self.ids]
1348 @cache_readonly
1349 def result_index_and_ids(self):
1350 result_index = self.binlabels
1351 if len(self.binlabels) != 0 and isna(self.binlabels[0]):
1352 result_index = result_index[1:]
1354 ngroups = len(result_index)
1355 rep = np.diff(np.r_[0, self.bins])
1357 rep = ensure_platform_int(rep)
1358 if ngroups == len(self.bins):
1359 ids = np.repeat(np.arange(ngroups), rep)
1360 else:
1361 ids = np.repeat(np.r_[-1, np.arange(ngroups)], rep)
1362 ids = ensure_platform_int(ids)
1364 return result_index, ids
1366 @property
1367 def levels(self) -> list[Index]:
1368 return [self.binlabels]
1370 @property
1371 def names(self) -> list[Hashable]:
1372 return [self.binlabels.name]
1374 @property
1375 def groupings(self) -> list[grouper.Grouping]:
1376 lev = self.binlabels
1377 codes = self.ids
1378 labels = lev.take(codes)
1379 ping = grouper.Grouping(
1380 labels, labels, in_axis=False, level=None, uniques=lev._values
1381 )
1382 return [ping]
1384 @property
1385 def observed_grouper(self) -> BinGrouper:
1386 return self
1389def _is_indexed_like(obj, axes) -> bool:
1390 if isinstance(obj, Series):
1391 if len(axes) > 1:
1392 return False
1393 return obj.index.equals(axes[0])
1394 elif isinstance(obj, DataFrame):
1395 return obj.index.equals(axes[0])
1397 return False
1400# ----------------------------------------------------------------------
1401# Splitting / application
1404class DataSplitter(Generic[NDFrameT]):
1405 def __init__(
1406 self,
1407 data: NDFrameT,
1408 ngroups: int,
1409 *,
1410 sort_idx: npt.NDArray[np.intp],
1411 sorted_ids: npt.NDArray[np.intp],
1412 ) -> None:
1413 self.data = data
1414 self.ngroups = ngroups
1416 self._slabels = sorted_ids
1417 self._sort_idx = sort_idx
1419 def __iter__(self) -> Iterator:
1420 if self.ngroups == 0:
1421 # we are inside a generator, rather than raise StopIteration
1422 # we merely return signal the end
1423 return
1425 starts, ends = lib.generate_slices(self._slabels, self.ngroups)
1426 sdata = self._sorted_data
1427 # __finalize__ is a no-op for an exact Series/DataFrame with empty
1428 # attrs and default flags; skip it in that case. Safe to compute the
1429 # gate once because sdata is splitter-owned and its type/attrs/flags
1430 # can't be mutated through the yielded chunks.
1431 needs_finalize = (
1432 type(sdata) is not self._sorted_cls
1433 or bool(sdata.attrs)
1434 or not sdata._flags._allows_duplicate_labels
1435 )
1436 for start, end in zip(starts, ends, strict=True):
1437 yield self._chop(sdata, slice(start, end), needs_finalize)
1439 @cache_readonly
1440 def _sorted_data(self) -> NDFrameT:
1441 return self.data.take(self._sort_idx, axis=0)
1443 _sorted_cls: type[NDFrame]
1445 def _chop(self, sdata, slice_obj: slice, needs_finalize: bool) -> NDFrame:
1446 raise AbstractMethodError(self)
1449class SeriesSplitter(DataSplitter):
1450 _sorted_cls = Series
1452 def _chop(self, sdata: Series, slice_obj: slice, needs_finalize: bool) -> Series:
1453 # fastpath equivalent to `sdata.iloc[slice_obj]`
1454 mgr = sdata._mgr.get_slice(slice_obj)
1455 ser = sdata._constructor_from_mgr(mgr, axes=mgr.axes)
1456 # Use object.__setattr__ to bypass NDFrame.__setattr__ overhead
1457 object.__setattr__(ser, "_name", sdata.name)
1458 if needs_finalize:
1459 return ser.__finalize__(sdata, method="groupby")
1460 return ser
1463class FrameSplitter(DataSplitter):
1464 _sorted_cls = DataFrame
1466 def _chop(
1467 self, sdata: DataFrame, slice_obj: slice, needs_finalize: bool
1468 ) -> DataFrame:
1469 # Fastpath equivalent to:
1470 # return sdata.iloc[slice_obj]
1471 mgr = sdata._mgr.get_slice(slice_obj, axis=1)
1472 df = sdata._constructor_from_mgr(mgr, axes=mgr.axes)
1473 if needs_finalize:
1474 return df.__finalize__(sdata, method="groupby")
1475 return df