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)
20import numpy as np
22from pandas._libs import (
23 NaT,
24 lib,
25)
26import pandas._libs.groupby as libgroupby
27from pandas._typing import (
28 ArrayLike,
29 AxisInt,
30 NDFrameT,
31 Shape,
32 npt,
33)
34from pandas.errors import AbstractMethodError
35from pandas.util._decorators import cache_readonly
37from pandas.core.dtypes.cast import (
38 maybe_downcast_to_dtype,
39)
40from pandas.core.dtypes.common import (
41 ensure_float64,
42 ensure_int64,
43 ensure_platform_int,
44 ensure_uint64,
45 is_1d_only_ea_dtype,
46)
47from pandas.core.dtypes.missing import (
48 isna,
49 maybe_fill,
50)
52from pandas.core.arrays import Categorical
53from pandas.core.frame import DataFrame
54from pandas.core.groupby import grouper
55from pandas.core.indexes.api import (
56 CategoricalIndex,
57 Index,
58 MultiIndex,
59 ensure_index,
60)
61from pandas.core.series import Series
62from pandas.core.sorting import (
63 compress_group_index,
64 decons_obs_group_ids,
65 get_group_index,
66 get_group_index_sorter,
67 get_indexer_dict,
68)
70if TYPE_CHECKING:
71 from collections.abc import (
72 Callable,
73 Generator,
74 Hashable,
75 Iterator,
76 )
78 from pandas.core.generic import NDFrame
81def check_result_array(obj, dtype) -> None:
82 # Our operation is supposed to be an aggregation/reduction. If
83 # it returns an ndarray, this likely means an invalid operation has
84 # been passed. See test_apply_without_aggregation, test_agg_must_agg
85 if isinstance(obj, np.ndarray):
86 if dtype != object:
87 # If it is object dtype, the function can be a reduction/aggregation
88 # and still return an ndarray e.g. test_agg_over_numpy_arrays
89 raise ValueError("Must produce aggregated value")
92def extract_result(res):
93 """
94 Extract the result object, it might be a 0-dim ndarray
95 or a len-1 0-dim, or a scalar
96 """
97 if hasattr(res, "_values"):
98 # Preserve EA
99 res = res._values
100 if res.ndim == 1 and len(res) == 1:
101 # see test_agg_lambda_with_timezone, test_resampler_grouper.py::test_apply
102 res = res[0]
103 return res
106class WrappedCythonOp:
107 """
108 Dispatch logic for functions defined in _libs.groupby
110 Parameters
111 ----------
112 kind: str
113 Whether the operation is an aggregate or transform.
114 how: str
115 Operation name, e.g. "mean".
116 has_dropped_na: bool
117 True precisely when dropna=True and the grouper contains a null value.
118 """
120 # Functions for which we do _not_ attempt to cast the cython result
121 # back to the original dtype.
122 cast_blocklist = frozenset(
123 ["any", "all", "rank", "count", "size", "idxmin", "idxmax"]
124 )
126 def __init__(self, kind: str, how: str, has_dropped_na: bool) -> None:
127 self.kind = kind
128 self.how = how
129 self.has_dropped_na = has_dropped_na
131 _CYTHON_FUNCTIONS: dict[str, dict] = {
132 "aggregate": {
133 "any": functools.partial(libgroupby.group_any_all, val_test="any"),
134 "all": functools.partial(libgroupby.group_any_all, val_test="all"),
135 "sum": "group_sum",
136 "prod": "group_prod",
137 "idxmin": functools.partial(libgroupby.group_idxmin_idxmax, name="idxmin"),
138 "idxmax": functools.partial(libgroupby.group_idxmin_idxmax, name="idxmax"),
139 "min": "group_min",
140 "max": "group_max",
141 "mean": "group_mean",
142 "median": "group_median_float64",
143 "var": "group_var",
144 "std": functools.partial(libgroupby.group_var, name="std"),
145 "sem": functools.partial(libgroupby.group_var, name="sem"),
146 "skew": "group_skew",
147 "kurt": "group_kurt",
148 "first": "group_nth",
149 "last": "group_last",
150 "ohlc": "group_ohlc",
151 },
152 "transform": {
153 "cumprod": "group_cumprod",
154 "cumsum": "group_cumsum",
155 "cummin": "group_cummin",
156 "cummax": "group_cummax",
157 "rank": "group_rank",
158 },
159 }
161 _cython_arity = {"ohlc": 4} # OHLC
163 @classmethod
164 def get_kind_from_how(cls, how: str) -> str:
165 if how in cls._CYTHON_FUNCTIONS["aggregate"]:
166 return "aggregate"
167 return "transform"
169 # Note: we make this a classmethod and pass kind+how so that caching
170 # works at the class level and not the instance level
171 @classmethod
172 @functools.cache
173 def _get_cython_function(
174 cls, kind: str, how: str, dtype: np.dtype, is_numeric: bool
175 ):
176 dtype_str = dtype.name
177 ftype = cls._CYTHON_FUNCTIONS[kind][how]
179 # see if there is a fused-type version of function
180 # only valid for numeric
181 if callable(ftype):
182 f = ftype
183 else:
184 f = getattr(libgroupby, ftype)
185 if is_numeric:
186 return f
187 elif dtype == np.dtype(object):
188 if how in ["median", "cumprod"]:
189 # no fused types -> no __signatures__
190 raise NotImplementedError(
191 f"function is not implemented for this dtype: "
192 f"[how->{how},dtype->{dtype_str}]"
193 )
194 elif how in ["std", "sem", "idxmin", "idxmax"]:
195 # We have a partial object that does not have __signatures__
196 return f
197 elif how in ["skew", "kurt"]:
198 # _get_cython_vals will convert to float64
199 pass
200 elif "object" not in f.__signatures__:
201 # raise NotImplementedError here rather than TypeError later
202 raise NotImplementedError(
203 f"function is not implemented for this dtype: "
204 f"[how->{how},dtype->{dtype_str}]"
205 )
206 return f
207 else:
208 raise NotImplementedError(
209 "This should not be reached. Please report a bug at "
210 "github.com/pandas-dev/pandas/",
211 dtype,
212 )
214 def _get_cython_vals(self, values: np.ndarray) -> np.ndarray:
215 """
216 Cast numeric dtypes to float64 for functions that only support that.
218 Parameters
219 ----------
220 values : np.ndarray
222 Returns
223 -------
224 values : np.ndarray
225 """
226 how = self.how
228 if how in ["median", "std", "sem", "skew", "kurt"]:
229 # median only has a float64 implementation
230 # We should only get here with is_numeric, as non-numeric cases
231 # should raise in _get_cython_function
232 values = ensure_float64(values)
234 elif values.dtype.kind in "iu":
235 if how in ["var", "mean"] or (
236 self.kind == "transform" and self.has_dropped_na
237 ):
238 # has_dropped_na check need for test_null_group_str_transformer
239 # result may still include NaN, so we have to cast
240 values = ensure_float64(values)
242 elif how in ["sum", "ohlc", "prod", "cumsum", "cumprod"]:
243 # Avoid overflow during group op
244 if values.dtype.kind == "i":
245 values = ensure_int64(values)
246 else:
247 values = ensure_uint64(values)
249 return values
251 def _get_output_shape(self, ngroups: int, values: np.ndarray) -> Shape:
252 how = self.how
253 kind = self.kind
255 arity = self._cython_arity.get(how, 1)
257 out_shape: Shape
258 if how == "ohlc":
259 out_shape = (ngroups, arity)
260 elif arity > 1:
261 raise NotImplementedError(
262 "arity of more than 1 is not supported for the 'how' argument"
263 )
264 elif kind == "transform":
265 out_shape = values.shape
266 else:
267 out_shape = (ngroups, *values.shape[1:])
268 return out_shape
270 def _get_out_dtype(self, dtype: np.dtype) -> np.dtype:
271 how = self.how
273 if how == "rank":
274 out_dtype = "float64"
275 elif how in ["idxmin", "idxmax"]:
276 # The Cython implementation only produces the row number; we'll take
277 # from the index using this in post processing
278 out_dtype = "intp"
279 elif dtype.kind in "iufcb":
280 out_dtype = f"{dtype.kind}{dtype.itemsize}"
281 else:
282 out_dtype = "object"
283 return np.dtype(out_dtype)
285 def _get_result_dtype(self, dtype: np.dtype) -> np.dtype:
286 """
287 Get the desired dtype of a result based on the
288 input dtype and how it was computed.
290 Parameters
291 ----------
292 dtype : np.dtype
294 Returns
295 -------
296 np.dtype
297 The desired dtype of the result.
298 """
299 how = self.how
301 if how in ["sum", "cumsum", "sum", "prod", "cumprod"]:
302 if dtype == np.dtype(bool):
303 return np.dtype(np.int64)
304 elif how in ["mean", "median", "var", "std", "sem"]:
305 if dtype.kind in "fc":
306 return dtype
307 elif dtype.kind in "iub":
308 return np.dtype(np.float64)
309 return dtype
311 @final
312 def _cython_op_ndim_compat(
313 self,
314 values: np.ndarray,
315 *,
316 min_count: int,
317 ngroups: int,
318 comp_ids: np.ndarray,
319 mask: npt.NDArray[np.bool_] | None = None,
320 result_mask: npt.NDArray[np.bool_] | None = None,
321 initial: Any = 0,
322 **kwargs,
323 ) -> np.ndarray:
324 if values.ndim == 1:
325 # expand to 2d, dispatch, then squeeze if appropriate
326 values2d = values[None, :]
327 if mask is not None:
328 mask = mask[None, :]
329 if result_mask is not None:
330 result_mask = result_mask[None, :]
331 res = self._call_cython_op(
332 values2d,
333 min_count=min_count,
334 ngroups=ngroups,
335 comp_ids=comp_ids,
336 mask=mask,
337 result_mask=result_mask,
338 initial=initial,
339 **kwargs,
340 )
341 if res.shape[0] == 1:
342 return res[0]
344 # otherwise we have OHLC
345 return res.T
347 return self._call_cython_op(
348 values,
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 )
358 @final
359 def _call_cython_op(
360 self,
361 values: np.ndarray, # np.ndarray[ndim=2]
362 *,
363 min_count: int,
364 ngroups: int,
365 comp_ids: np.ndarray,
366 mask: npt.NDArray[np.bool_] | None,
367 result_mask: npt.NDArray[np.bool_] | None,
368 initial: Any = 0,
369 **kwargs,
370 ) -> np.ndarray: # np.ndarray[ndim=2]
371 orig_values = values
373 dtype = values.dtype
374 is_numeric = dtype.kind in "iufcb"
376 is_datetimelike = dtype.kind in "mM"
378 if self.how in ["any", "all"]:
379 if mask is None:
380 mask = isna(values)
382 if is_datetimelike:
383 values = values.view("int64")
384 is_numeric = True
385 elif dtype.kind == "b":
386 values = values.view("uint8")
387 if values.dtype == "float16":
388 values = values.astype(np.float32)
390 if self.how in ["any", "all"]:
391 if dtype == object:
392 if kwargs["skipna"]:
393 # GH#37501: don't raise on pd.NA when skipna=True
394 if mask is not None and mask.any():
395 # mask on original values computed separately
396 values = values.copy()
397 values[mask] = True
398 values = values.astype(bool, copy=False).view(np.int8)
399 is_numeric = True
401 values = values.T
402 if mask is not None:
403 mask = mask.T
404 if result_mask is not None:
405 result_mask = result_mask.T
407 out_shape = self._get_output_shape(ngroups, values)
408 func = self._get_cython_function(self.kind, self.how, values.dtype, is_numeric)
409 values = self._get_cython_vals(values)
410 out_dtype = self._get_out_dtype(values.dtype)
412 result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
413 if self.kind == "aggregate":
414 counts = np.zeros(ngroups, dtype=np.int64)
415 if self.how in [
416 "idxmin",
417 "idxmax",
418 "min",
419 "max",
420 "mean",
421 "last",
422 "first",
423 "sum",
424 "median",
425 ]:
426 if self.how == "sum":
427 # pass in through kwargs only for sum (other functions don't have
428 # the keyword)
429 kwargs["initial"] = initial
430 func(
431 out=result,
432 counts=counts,
433 values=values,
434 labels=comp_ids,
435 min_count=min_count,
436 mask=mask,
437 result_mask=result_mask,
438 is_datetimelike=is_datetimelike,
439 **kwargs,
440 )
441 elif self.how in ["sem", "std", "var", "ohlc", "prod"]:
442 if self.how in ["std", "sem"]:
443 kwargs["is_datetimelike"] = is_datetimelike
444 func(
445 result,
446 counts,
447 values,
448 comp_ids,
449 min_count=min_count,
450 mask=mask,
451 result_mask=result_mask,
452 **kwargs,
453 )
454 elif self.how in ["any", "all"]:
455 func(
456 out=result,
457 values=values,
458 labels=comp_ids,
459 mask=mask,
460 result_mask=result_mask,
461 **kwargs,
462 )
463 result = result.astype(bool, copy=False)
464 elif self.how in ["skew", "kurt"]:
465 func(
466 out=result,
467 counts=counts,
468 values=values,
469 labels=comp_ids,
470 mask=mask,
471 result_mask=result_mask,
472 **kwargs,
473 )
474 if dtype == object:
475 result = result.astype(object)
477 else:
478 raise NotImplementedError(f"{self.how} is not implemented")
479 else:
480 # TODO: min_count
481 if self.how != "rank":
482 # TODO: should rank take result_mask?
483 kwargs["result_mask"] = result_mask
484 func(
485 out=result,
486 values=values,
487 labels=comp_ids,
488 ngroups=ngroups,
489 is_datetimelike=is_datetimelike,
490 mask=mask,
491 **kwargs,
492 )
494 if self.kind == "aggregate" and self.how not in ["idxmin", "idxmax"]:
495 # i.e. counts is defined. Locations where count<min_count
496 # need to have the result set to np.nan, which may require casting,
497 # see GH#40767. For idxmin/idxmax is handled specially via post-processing
498 if result.dtype.kind in "iu" and not is_datetimelike:
499 # if the op keeps the int dtypes, we have to use 0
500 cutoff = max(0 if self.how in ["sum", "prod"] else 1, min_count)
501 empty_groups = counts < cutoff
502 if empty_groups.any():
503 if result_mask is not None:
504 assert result_mask[empty_groups].all()
505 else:
506 # Note: this conversion could be lossy, see GH#40767
507 result = result.astype("float64")
508 result[empty_groups] = np.nan
510 result = result.T
512 if self.how not in self.cast_blocklist:
513 # e.g. if we are int64 and need to restore to datetime64/timedelta64
514 # "rank" is the only member of cast_blocklist we get here
515 # Casting only needed for float16, bool, datetimelike,
516 # and self.how in ["sum", "prod", "ohlc", "cumprod"]
517 res_dtype = self._get_result_dtype(orig_values.dtype)
518 op_result = maybe_downcast_to_dtype(result, res_dtype)
519 else:
520 op_result = result
522 return op_result
524 @final
525 def _validate_axis(self, axis: AxisInt, values: ArrayLike) -> None:
526 if values.ndim > 2:
527 raise NotImplementedError("number of dimensions is currently limited to 2")
528 if values.ndim == 2:
529 assert axis == 1, axis
530 elif not is_1d_only_ea_dtype(values.dtype):
531 # Note: it is *not* the case that axis is always 0 for 1-dim values,
532 # as we can have 1D ExtensionArrays that we need to treat as 2D
533 assert axis == 0
535 @final
536 def cython_operation(
537 self,
538 *,
539 values: ArrayLike,
540 axis: AxisInt,
541 min_count: int = -1,
542 comp_ids: np.ndarray,
543 ngroups: int,
544 **kwargs,
545 ) -> ArrayLike:
546 """
547 Call our cython function, with appropriate pre- and post- processing.
548 """
549 self._validate_axis(axis, values)
551 if not isinstance(values, np.ndarray):
552 # i.e. ExtensionArray
553 return values._groupby_op(
554 how=self.how,
555 has_dropped_na=self.has_dropped_na,
556 min_count=min_count,
557 ngroups=ngroups,
558 ids=comp_ids,
559 **kwargs,
560 )
562 return self._cython_op_ndim_compat(
563 values,
564 min_count=min_count,
565 ngroups=ngroups,
566 comp_ids=comp_ids,
567 mask=None,
568 **kwargs,
569 )
572class BaseGrouper:
573 """
574 This is an internal Grouper class, which actually holds
575 the generated groups
577 Parameters
578 ----------
579 axis : Index
580 groupings : Sequence[Grouping]
581 all the grouping instances to handle in this grouper
582 for example for grouper list to groupby, need to pass the list
583 sort : bool, default True
584 whether this grouper will give sorted result or not
586 """
588 axis: Index
590 def __init__(
591 self,
592 axis: Index,
593 groupings: list[grouper.Grouping],
594 sort: bool = True,
595 dropna: bool = True,
596 ) -> None:
597 assert isinstance(axis, Index), axis
599 self.axis = axis
600 self._groupings = groupings
601 self._sort = sort
602 self.dropna = dropna
604 @property
605 def groupings(self) -> list[grouper.Grouping]:
606 return self._groupings
608 def __iter__(self) -> Iterator[Hashable]:
609 return iter(self.indices)
611 @property
612 def nkeys(self) -> int:
613 return len(self.groupings)
615 def get_iterator(self, data: NDFrameT) -> Iterator[tuple[Hashable, NDFrameT]]:
616 """
617 Groupby iterator
619 Returns
620 -------
621 Generator yielding sequence of (name, subsetted object)
622 for each group
623 """
624 splitter = self._get_splitter(data)
625 # TODO: Would be more efficient to skip unobserved for transforms
626 keys = self.result_index
627 yield from zip(keys, splitter, strict=True)
629 @final
630 def _get_splitter(self, data: NDFrame) -> DataSplitter:
631 """
632 Returns
633 -------
634 Generator yielding subsetted objects
635 """
636 if isinstance(data, Series):
637 klass: type[DataSplitter] = SeriesSplitter
638 else:
639 # i.e. DataFrame
640 klass = FrameSplitter
642 return klass(
643 data,
644 self.ngroups,
645 sorted_ids=self._sorted_ids,
646 sort_idx=self.result_ilocs,
647 )
649 @cache_readonly
650 def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]:
651 """dict {group name -> group indices}"""
652 if len(self.groupings) == 1 and isinstance(self.result_index, CategoricalIndex):
653 # This shows unused categories in indices GH#38642
654 result = self.groupings[0].indices
655 else:
656 codes_list = [ping.codes for ping in self.groupings]
657 result = get_indexer_dict(codes_list, self.levels)
658 if not self.dropna:
659 has_mi = isinstance(self.result_index, MultiIndex)
660 if not has_mi and self.result_index.hasnans:
661 result = {
662 np.nan if isna(key) else key: value for key, value in result.items()
663 }
664 elif has_mi:
665 # MultiIndex has no efficient way to tell if there are NAs
666 result = {
667 # error: "Hashable" has no attribute "__iter__" (not iterable)
668 tuple(np.nan if isna(comp) else comp for comp in key): value # type: ignore[attr-defined]
669 for key, value in result.items()
670 }
672 return result
674 @final
675 @cache_readonly
676 def result_ilocs(self) -> npt.NDArray[np.intp]:
677 """
678 Get the original integer locations of result_index in the input.
679 """
680 # Original indices are where group_index would go via sorting.
681 # But when dropna is true, we need to remove null values while accounting for
682 # any gaps that then occur because of them.
683 ids = self.ids
685 if self.has_dropped_na:
686 mask = np.where(ids >= 0)
687 # Count how many gaps are caused by previous null values for each position
688 null_gaps = np.cumsum(ids == -1)[mask]
689 ids = ids[mask]
691 result = get_group_index_sorter(ids, self.ngroups)
693 if self.has_dropped_na:
694 # Shift by the number of prior null gaps
695 result += np.take(null_gaps, result)
697 return result
699 @property
700 def codes(self) -> list[npt.NDArray[np.signedinteger]]:
701 return [ping.codes for ping in self.groupings]
703 @property
704 def levels(self) -> list[Index]:
705 if len(self.groupings) > 1:
706 # mypy doesn't know result_index must be a MultiIndex
707 return list(self.result_index.levels) # type: ignore[attr-defined]
708 else:
709 return [self.result_index]
711 @property
712 def names(self) -> list[Hashable]:
713 return [ping.name for ping in self.groupings]
715 @final
716 def size(self) -> Series:
717 """
718 Compute group sizes.
719 """
720 ids = self.ids
721 ngroups = self.ngroups
722 out: np.ndarray | list
723 if ngroups:
724 out = np.bincount(ids[ids != -1], minlength=ngroups)
725 else:
726 out = []
727 return Series(out, index=self.result_index, dtype="int64", copy=False)
729 @cache_readonly
730 def groups(self) -> dict[Hashable, Index]:
731 """dict {group name -> group labels}"""
732 if len(self.groupings) == 1:
733 return self.groupings[0].groups
734 result_index, ids = self.result_index_and_ids
735 values = result_index._values
736 categories = Categorical.from_codes(ids, categories=range(len(result_index)))
737 result = {
738 # mypy is not aware that group has to be an integer
739 values[group]: self.axis.take(axis_ilocs) # type: ignore[call-overload]
740 for group, axis_ilocs in categories._reverse_indexer().items()
741 }
742 return result
744 @final
745 @cache_readonly
746 def is_monotonic(self) -> bool:
747 # return if my group orderings are monotonic
748 return Index(self.ids, copy=False).is_monotonic_increasing
750 @final
751 @cache_readonly
752 def has_dropped_na(self) -> bool:
753 """
754 Whether grouper has null value(s) that are dropped.
755 """
756 return bool((self.ids < 0).any())
758 @cache_readonly
759 def codes_info(self) -> npt.NDArray[np.intp]:
760 # return the codes of items in original grouped axis
761 return self.ids
763 @final
764 @cache_readonly
765 def ngroups(self) -> int:
766 return len(self.result_index)
768 @property
769 def result_index(self) -> Index:
770 return self.result_index_and_ids[0]
772 @property
773 def ids(self) -> npt.NDArray[np.intp]:
774 return self.result_index_and_ids[1]
776 @cache_readonly
777 def result_index_and_ids(self) -> tuple[Index, npt.NDArray[np.intp]]:
778 levels = [
779 Index._with_infer(ping.uniques, copy=False) for ping in self.groupings
780 ]
781 obs = [
782 ping._observed or not ping._passed_categorical for ping in self.groupings
783 ]
784 sorts = [ping._sort for ping in self.groupings]
785 # When passed a categorical grouping, keep all categories
786 for k, (ping, level) in enumerate(zip(self.groupings, levels, strict=True)):
787 if ping._passed_categorical:
788 levels[k] = level.set_categories(ping._orig_cats)
790 if len(self.groupings) == 1:
791 result_index = levels[0]
792 result_index.name = self.names[0]
793 ids = ensure_platform_int(self.codes[0])
794 elif all(obs):
795 result_index, ids = self._ob_index_and_ids(
796 levels, self.codes, self.names, sorts
797 )
798 elif not any(obs):
799 result_index, ids = self._unob_index_and_ids(levels, self.codes, self.names)
800 else:
801 # Combine unobserved and observed parts
802 names = self.names
803 codes = [ping.codes for ping in self.groupings]
804 ob_indices = [idx for idx, ob in enumerate(obs) if ob]
805 unob_indices = [idx for idx, ob in enumerate(obs) if not ob]
806 ob_index, ob_ids = self._ob_index_and_ids(
807 levels=[levels[idx] for idx in ob_indices],
808 codes=[codes[idx] for idx in ob_indices],
809 names=[names[idx] for idx in ob_indices],
810 sorts=[sorts[idx] for idx in ob_indices],
811 )
812 unob_index, unob_ids = self._unob_index_and_ids(
813 levels=[levels[idx] for idx in unob_indices],
814 codes=[codes[idx] for idx in unob_indices],
815 names=[names[idx] for idx in unob_indices],
816 )
818 result_index_codes = np.concatenate(
819 [
820 np.tile(unob_index.codes, len(ob_index)),
821 np.repeat(ob_index.codes, len(unob_index), axis=1),
822 ],
823 axis=0,
824 )
825 _, index = np.unique(unob_indices + ob_indices, return_index=True)
826 result_index = MultiIndex(
827 levels=list(unob_index.levels) + list(ob_index.levels),
828 codes=result_index_codes,
829 names=list(unob_index.names) + list(ob_index.names),
830 ).reorder_levels(index)
832 # The sum here will get -1 values wrong when dropna=True;
833 # we will fix at the end.
834 ids = len(unob_index) * ob_ids + unob_ids
836 if any(sorts):
837 # Sort result_index and recode ids using the new order
838 n_levels = len(sorts)
839 drop_levels = [
840 n_levels - idx
841 for idx, sort in enumerate(reversed(sorts), 1)
842 if not sort
843 ]
844 if len(drop_levels) > 0:
845 sorter = result_index._drop_level_numbers(drop_levels).argsort()
846 else:
847 sorter = result_index.argsort()
848 result_index = result_index.take(sorter)
849 _, index = np.unique(sorter, return_index=True)
850 ids = ensure_platform_int(ids)
851 ids = index.take(ids)
852 else:
853 # Recode ids and reorder result_index with observed groups up front,
854 # unobserved at the end
855 ids, uniques = compress_group_index(ids, sort=False)
856 ids = ensure_platform_int(ids)
857 taker = np.concatenate(
858 [uniques, np.delete(np.arange(len(result_index)), uniques)]
859 )
860 result_index = result_index.take(taker)
862 if self.dropna:
863 ids = np.where((ob_ids < 0) | (unob_ids < 0), -1, ids)
865 return result_index, ids
867 @property
868 def observed_grouper(self) -> BaseGrouper:
869 if all(ping._observed for ping in self.groupings):
870 return self
872 return self._observed_grouper
874 @cache_readonly
875 def _observed_grouper(self) -> BaseGrouper:
876 groupings = [ping.observed_grouping for ping in self.groupings]
877 grouper = BaseGrouper(self.axis, groupings, sort=self._sort, dropna=self.dropna)
878 return grouper
880 def _ob_index_and_ids(
881 self,
882 levels: list[Index],
883 codes: list[npt.NDArray[np.intp]],
884 names: list[Hashable],
885 sorts: list[bool],
886 ) -> tuple[MultiIndex, npt.NDArray[np.intp]]:
887 consistent_sorting = all(sorts[0] == sort for sort in sorts[1:])
888 sort_in_compress = sorts[0] if consistent_sorting else False
889 shape = tuple(len(level) for level in levels)
890 group_index = get_group_index(codes, shape, sort=True, xnull=True)
891 ob_ids, obs_group_ids = compress_group_index(group_index, sort=sort_in_compress)
892 ob_ids = ensure_platform_int(ob_ids)
893 ob_index_codes = decons_obs_group_ids(
894 ob_ids, obs_group_ids, shape, codes, xnull=True
895 )
896 ob_index = MultiIndex(
897 levels=levels,
898 codes=ob_index_codes,
899 names=names,
900 verify_integrity=False,
901 )
902 if not consistent_sorting and len(ob_index) > 0:
903 # Sort by the levels where the corresponding sort argument is True
904 n_levels = len(sorts)
905 drop_levels = [
906 n_levels - idx
907 for idx, sort in enumerate(reversed(sorts), 1)
908 if not sort
909 ]
910 if len(drop_levels) > 0:
911 sorter = ob_index._drop_level_numbers(drop_levels).argsort()
912 else:
913 sorter = ob_index.argsort()
914 ob_index = ob_index.take(sorter)
915 _, index = np.unique(sorter, return_index=True)
916 ob_ids = np.where(ob_ids == -1, -1, index.take(ob_ids))
917 ob_ids = ensure_platform_int(ob_ids)
918 return ob_index, ob_ids
920 def _unob_index_and_ids(
921 self,
922 levels: list[Index],
923 codes: list[npt.NDArray[np.intp]],
924 names: list[Hashable],
925 ) -> tuple[MultiIndex, npt.NDArray[np.intp]]:
926 shape = tuple(len(level) for level in levels)
927 unob_ids = get_group_index(codes, shape, sort=True, xnull=True)
928 unob_index = MultiIndex.from_product(levels, names=names)
929 unob_ids = ensure_platform_int(unob_ids)
930 return unob_index, unob_ids
932 @final
933 def get_group_levels(self) -> Generator[Index]:
934 # Note: only called from _insert_inaxis_grouper, which
935 # is only called for BaseGrouper, never for BinGrouper
936 result_index = self.result_index
937 if len(self.groupings) == 1:
938 yield result_index
939 else:
940 for level in range(result_index.nlevels - 1, -1, -1):
941 yield result_index.get_level_values(level)
943 # ------------------------------------------------------------
944 # Aggregation functions
946 @final
947 def _cython_operation(
948 self,
949 kind: str,
950 values,
951 how: str,
952 axis: AxisInt,
953 min_count: int = -1,
954 **kwargs,
955 ) -> ArrayLike:
956 """
957 Returns the values of a cython operation.
958 """
959 assert kind in ["transform", "aggregate"]
961 cy_op = WrappedCythonOp(kind=kind, how=how, has_dropped_na=self.has_dropped_na)
963 return cy_op.cython_operation(
964 values=values,
965 axis=axis,
966 min_count=min_count,
967 comp_ids=self.ids,
968 ngroups=self.ngroups,
969 **kwargs,
970 )
972 @final
973 def agg_series(
974 self, obj: Series, func: Callable, preserve_dtype: bool = False
975 ) -> ArrayLike:
976 """
977 Parameters
978 ----------
979 obj : Series
980 func : function taking a Series and returning a scalar-like
981 preserve_dtype : bool
982 Whether the aggregation is known to be dtype-preserving.
984 Returns
985 -------
986 np.ndarray or ExtensionArray
987 """
988 result = self._aggregate_series_pure_python(obj, func)
989 return obj.array._cast_pointwise_result(result)
991 @final
992 def _aggregate_series_pure_python(
993 self, obj: Series, func: Callable
994 ) -> npt.NDArray[np.object_]:
995 result = np.empty(self.ngroups, dtype="O")
996 initialized = False
998 splitter = self._get_splitter(obj)
1000 for i, group in enumerate(splitter):
1001 res = func(group)
1002 res = extract_result(res)
1004 if not initialized:
1005 # We only do this validation on the first iteration
1006 check_result_array(res, group.dtype)
1007 initialized = True
1009 result[i] = res
1011 return result
1013 @final
1014 def apply_groupwise(
1015 self, f: Callable, data: DataFrame | Series
1016 ) -> tuple[list, bool]:
1017 mutated = False
1018 splitter = self._get_splitter(data)
1019 group_keys = self.result_index
1020 result_values = []
1022 # This calls DataSplitter.__iter__
1023 zipped = zip(group_keys, splitter, strict=True)
1025 for key, group in zipped:
1026 # Pinning name is needed for
1027 # test_group_apply_once_per_group,
1028 # test_inconsistent_return_type, test_set_group_name,
1029 # test_group_name_available_in_inference_pass,
1030 # test_groupby_multi_timezone
1031 object.__setattr__(group, "name", key)
1033 # group might be modified
1034 group_axes = group.axes
1035 res = f(group)
1036 if not mutated and not _is_indexed_like(res, group_axes):
1037 mutated = True
1038 result_values.append(res)
1039 # getattr pattern for __name__ is needed for functools.partial objects
1040 if len(group_keys) == 0 and getattr(f, "__name__", None) in [
1041 "skew",
1042 "kurt",
1043 "sum",
1044 "prod",
1045 ]:
1046 # If group_keys is empty, then no function calls have been made,
1047 # so we will not have raised even if this is an invalid dtype.
1048 # So do one dummy call here to raise appropriate TypeError.
1049 f(data.iloc[:0])
1051 return result_values, mutated
1053 # ------------------------------------------------------------
1054 # Methods for sorting subsets of our GroupBy's object
1056 @final
1057 @cache_readonly
1058 def _sorted_ids(self) -> npt.NDArray[np.intp]:
1059 result = self.ids.take(self.result_ilocs)
1060 if getattr(self, "dropna", True):
1061 # BinGrouper has no dropna
1062 result = result[result >= 0]
1063 return result
1066class BinGrouper(BaseGrouper):
1067 """
1068 This is an internal Grouper class
1070 Parameters
1071 ----------
1072 bins : the split index of binlabels to group the item of axis
1073 binlabels : the label list
1074 indexer : np.ndarray[np.intp], optional
1075 the indexer created by Grouper
1076 some groupers (TimeGrouper) will sort its axis and its
1077 group_info is also sorted, so need the indexer to reorder
1079 Examples
1080 --------
1081 bins: [2, 4, 6, 8, 10]
1082 binlabels: DatetimeIndex(['2005-01-01', '2005-01-03',
1083 '2005-01-05', '2005-01-07', '2005-01-09'],
1084 dtype='datetime64[ns]', freq='2D')
1086 the group_info, which contains the label of each item in grouped
1087 axis, the index of label in label list, group number, is
1089 (array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4]), array([0, 1, 2, 3, 4]), 5)
1091 means that, the grouped axis has 10 items, can be grouped into 5
1092 labels, the first and second items belong to the first label, the
1093 third and forth items belong to the second label, and so on
1095 """
1097 bins: npt.NDArray[np.int64]
1098 binlabels: Index
1100 def __init__(
1101 self,
1102 bins,
1103 binlabels,
1104 indexer=None,
1105 ) -> None:
1106 self.bins = ensure_int64(bins)
1107 self.binlabels = ensure_index(binlabels)
1108 self.indexer = indexer
1110 # These lengths must match, otherwise we could call agg_series
1111 # with empty self.bins, which would raise later.
1112 assert len(self.binlabels) == len(self.bins)
1114 @cache_readonly
1115 def groups(self):
1116 """dict {group name -> group labels}"""
1117 # this is mainly for compat
1118 # GH 3881
1119 result = {
1120 key: value
1121 for key, value in zip(self.binlabels, self.bins, strict=True)
1122 if key is not NaT
1123 }
1124 return result
1126 @property
1127 def nkeys(self) -> int:
1128 # still matches len(self.groupings), but we can hard-code
1129 return 1
1131 @cache_readonly
1132 def codes_info(self) -> npt.NDArray[np.intp]:
1133 # return the codes of items in original grouped axis
1134 ids = self.ids
1135 if self.indexer is not None:
1136 sorter = np.lexsort((ids, self.indexer))
1137 ids = ids[sorter]
1138 return ids
1140 def get_iterator(self, data: NDFrame):
1141 """
1142 Groupby iterator
1144 Returns
1145 -------
1146 Generator yielding sequence of (name, subsetted object)
1147 for each group
1148 """
1149 slicer = lambda start, edge: data.iloc[start:edge]
1151 start: np.int64 | int = 0
1152 for edge, label in zip(self.bins, self.binlabels, strict=True):
1153 if label is not NaT:
1154 yield label, slicer(start, edge)
1155 start = edge
1157 if start < len(data):
1158 yield self.binlabels[-1], slicer(start, None)
1160 @cache_readonly
1161 def indices(self):
1162 indices = collections.defaultdict(list)
1164 i: np.int64 | int = 0
1165 for label, bin in zip(self.binlabels, self.bins, strict=True):
1166 if i < bin:
1167 if label is not NaT:
1168 indices[label] = list(range(i, bin))
1169 i = bin
1170 return indices
1172 @cache_readonly
1173 def codes(self) -> list[npt.NDArray[np.intp]]:
1174 return [self.ids]
1176 @cache_readonly
1177 def result_index_and_ids(self):
1178 result_index = self.binlabels
1179 if len(self.binlabels) != 0 and isna(self.binlabels[0]):
1180 result_index = result_index[1:]
1182 ngroups = len(result_index)
1183 rep = np.diff(np.r_[0, self.bins])
1185 rep = ensure_platform_int(rep)
1186 if ngroups == len(self.bins):
1187 ids = np.repeat(np.arange(ngroups), rep)
1188 else:
1189 ids = np.repeat(np.r_[-1, np.arange(ngroups)], rep)
1190 ids = ensure_platform_int(ids)
1192 return result_index, ids
1194 @property
1195 def levels(self) -> list[Index]:
1196 return [self.binlabels]
1198 @property
1199 def names(self) -> list[Hashable]:
1200 return [self.binlabels.name]
1202 @property
1203 def groupings(self) -> list[grouper.Grouping]:
1204 lev = self.binlabels
1205 codes = self.ids
1206 labels = lev.take(codes)
1207 ping = grouper.Grouping(
1208 labels, labels, in_axis=False, level=None, uniques=lev._values
1209 )
1210 return [ping]
1212 @property
1213 def observed_grouper(self) -> BinGrouper:
1214 return self
1217def _is_indexed_like(obj, axes) -> bool:
1218 if isinstance(obj, Series):
1219 if len(axes) > 1:
1220 return False
1221 return obj.index.equals(axes[0])
1222 elif isinstance(obj, DataFrame):
1223 return obj.index.equals(axes[0])
1225 return False
1228# ----------------------------------------------------------------------
1229# Splitting / application
1232class DataSplitter(Generic[NDFrameT]):
1233 def __init__(
1234 self,
1235 data: NDFrameT,
1236 ngroups: int,
1237 *,
1238 sort_idx: npt.NDArray[np.intp],
1239 sorted_ids: npt.NDArray[np.intp],
1240 ) -> None:
1241 self.data = data
1242 self.ngroups = ngroups
1244 self._slabels = sorted_ids
1245 self._sort_idx = sort_idx
1247 def __iter__(self) -> Iterator:
1248 if self.ngroups == 0:
1249 # we are inside a generator, rather than raise StopIteration
1250 # we merely return signal the end
1251 return
1253 starts, ends = lib.generate_slices(self._slabels, self.ngroups)
1254 sdata = self._sorted_data
1255 for start, end in zip(starts, ends, strict=True):
1256 yield self._chop(sdata, slice(start, end))
1258 @cache_readonly
1259 def _sorted_data(self) -> NDFrameT:
1260 return self.data.take(self._sort_idx, axis=0)
1262 def _chop(self, sdata, slice_obj: slice) -> NDFrame:
1263 raise AbstractMethodError(self)
1266class SeriesSplitter(DataSplitter):
1267 def _chop(self, sdata: Series, slice_obj: slice) -> Series:
1268 # fastpath equivalent to `sdata.iloc[slice_obj]`
1269 mgr = sdata._mgr.get_slice(slice_obj)
1270 ser = sdata._constructor_from_mgr(mgr, axes=mgr.axes)
1271 ser._name = sdata.name
1272 return ser.__finalize__(sdata, method="groupby")
1275class FrameSplitter(DataSplitter):
1276 def _chop(self, sdata: DataFrame, slice_obj: slice) -> DataFrame:
1277 # Fastpath equivalent to:
1278 # return sdata.iloc[slice_obj]
1279 mgr = sdata._mgr.get_slice(slice_obj, axis=1)
1280 df = sdata._constructor_from_mgr(mgr, axes=mgr.axes)
1281 return df.__finalize__(sdata, method="groupby")