Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/sorting.py: 11%
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"""miscellaneous sorting / groupby utilities"""
3from __future__ import annotations
5import itertools
6from typing import (
7 TYPE_CHECKING,
8 cast,
9)
11import numpy as np
13from pandas._libs import (
14 algos,
15 hashtable,
16 lib,
17)
18from pandas._libs.hashtable import unique_label_indices
20from pandas.core.dtypes.common import (
21 ensure_int64,
22 ensure_platform_int,
23)
24from pandas.core.dtypes.generic import (
25 ABCMultiIndex,
26 ABCRangeIndex,
27)
28from pandas.core.dtypes.missing import isna
30from pandas.core.construction import extract_array
32if TYPE_CHECKING:
33 from collections.abc import (
34 Callable,
35 Iterable,
36 Sequence,
37 )
39 from pandas._typing import (
40 ArrayLike,
41 AxisInt,
42 IndexKeyFunc,
43 Level,
44 NaPosition,
45 Shape,
46 SortKind,
47 npt,
48 )
50 from pandas import (
51 MultiIndex,
52 Series,
53 )
54 from pandas.core.arrays import ExtensionArray
55 from pandas.core.indexes.base import Index
58def get_indexer_indexer(
59 target: Index,
60 level: Level | list[Level] | None,
61 ascending: list[bool] | bool,
62 kind: SortKind,
63 na_position: NaPosition,
64 sort_remaining: bool,
65 key: IndexKeyFunc,
66) -> npt.NDArray[np.intp] | None:
67 """
68 Helper method that return the indexer according to input parameters for
69 the sort_index method of DataFrame and Series.
71 Parameters
72 ----------
73 target : Index
74 level : int or level name or list of ints or list of level names
75 ascending : bool or list of bools, default True
76 kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}
77 na_position : {'first', 'last'}
78 sort_remaining : bool
79 key : callable, optional
81 Returns
82 -------
83 Optional[ndarray[intp]]
84 The indexer for the new index.
85 """
87 # error: Incompatible types in assignment (expression has type
88 # "Union[ExtensionArray, ndarray[Any, Any], Index, Series]", variable has
89 # type "Index")
90 target = ensure_key_mapped(target, key, levels=level) # type: ignore[assignment]
91 target = target._sort_levels_monotonic()
93 if (np.all(ascending) and target.is_monotonic_increasing) or (
94 not np.any(ascending) and target.is_monotonic_decreasing
95 ):
96 # GH#11080, GH#64883: on a non-MultiIndex, `level` is a no-op,
97 # so short-circuit even when level is specified.
98 if level is None or not isinstance(target, ABCMultiIndex):
99 return None
101 if level is not None:
102 _, indexer = target.sortlevel(
103 level, # type: ignore[arg-type]
104 ascending=ascending,
105 sort_remaining=sort_remaining,
106 na_position=na_position,
107 )
108 elif isinstance(target, ABCMultiIndex):
109 codes = [lev.codes for lev in target._get_codes_for_sorting()]
110 indexer = lexsort_indexer(
111 codes, orders=ascending, na_position=na_position, codes_given=True
112 )
113 else:
114 # ascending can only be a Sequence for MultiIndex
115 indexer = nargsort(
116 target,
117 kind=kind,
118 ascending=cast("bool", ascending),
119 na_position=na_position,
120 )
121 return indexer
124def get_group_index(
125 labels: Sequence[npt.NDArray[np.signedinteger]],
126 shape: Shape,
127 sort: bool,
128 xnull: bool,
129) -> npt.NDArray[np.int64]:
130 """
131 For the particular label_list, gets the offsets into the hypothetical list
132 representing the totally ordered cartesian product of all possible label
133 combinations, *as long as* this space fits within int64 bounds;
134 otherwise, though group indices identify unique combinations of
135 labels, they cannot be deconstructed.
136 - If `sort`, rank of returned ids preserve lexical ranks of labels.
137 i.e. returned id's can be used to do lexical sort on labels;
138 - If `xnull` nulls (-1 labels) are passed through.
140 Parameters
141 ----------
142 labels : sequence of arrays
143 Integers identifying levels at each location
144 shape : tuple[int, ...]
145 Number of unique levels at each location
146 sort : bool
147 If the ranks of returned ids should match lexical ranks of labels
148 xnull : bool
149 If true nulls are excluded. i.e. -1 values in the labels are
150 passed through.
152 Returns
153 -------
154 An array of type int64 where two elements are equal if their corresponding
155 labels are equal at all location.
157 Notes
158 -----
159 The length of `labels` and `shape` must be identical.
160 """
162 def _int64_cut_off(shape: list[int]) -> int:
163 acc = 1
164 for i, mul in enumerate(shape):
165 acc *= int(mul)
166 if not acc < lib.i8max:
167 return i
168 return len(shape)
170 def maybe_lift(lab: np.ndarray, size: int) -> tuple[np.ndarray, int]:
171 # promote nan values (assigned -1 label in lab array)
172 # so that all output values are non-negative
173 return (lab + 1, size + 1) if lib.has_sentinel(lab, -1) else (lab, size)
175 labels = [ensure_int64(x) for x in labels]
176 lshape = list(shape)
177 if not xnull:
178 for i, (lab, size) in enumerate(zip(labels, shape, strict=True)):
179 labels[i], lshape[i] = maybe_lift(lab, size)
181 # Iteratively process all the labels in chunks sized so less
182 # than lib.i8max unique int ids will be required for each chunk
183 while True:
184 # how many levels can be done without overflow:
185 nlev = _int64_cut_off(lshape)
187 # compute flat ids for the first `nlev` levels
188 stride = np.prod(lshape[1:nlev], dtype="i8")
189 out = stride * labels[0].astype("i8", subok=False, copy=False)
191 for i in range(1, nlev):
192 if lshape[i] == 0:
193 stride = np.int64(0)
194 else:
195 stride //= lshape[i]
196 out += labels[i] * stride
198 if xnull: # exclude nulls
199 mask = labels[0] == -1
200 for lab in labels[1:nlev]:
201 mask |= lab == -1
202 out[mask] = -1
204 if nlev == len(lshape): # all levels done!
205 break
207 # compress what has been done so far in order to avoid overflow
208 # to retain lexical ranks, obs_ids should be sorted
209 comp_ids, obs_ids = compress_group_index(out, sort=sort)
211 labels = [comp_ids, *labels[nlev:]]
212 lshape = [len(obs_ids), *lshape[nlev:]]
214 return out
217def get_compressed_ids(
218 labels: Sequence[npt.NDArray[np.signedinteger]], sizes: Shape
219) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.int64]]:
220 """
221 Group_index is offsets into cartesian product of all possible labels. This
222 space can be huge, so this function compresses it, by computing offsets
223 (comp_ids) into the list of unique labels (obs_group_ids).
225 Parameters
226 ----------
227 labels : list of label arrays
228 sizes : tuple[int] of size of the levels
230 Returns
231 -------
232 np.ndarray[np.intp]
233 comp_ids
234 np.ndarray[np.int64]
235 obs_group_ids
236 """
237 ids = get_group_index(labels, sizes, sort=True, xnull=False)
238 return compress_group_index(ids, sort=True)
241def is_int64_overflow_possible(shape: Shape) -> bool:
242 the_prod = 1
243 for x in shape:
244 the_prod *= int(x)
246 return the_prod >= lib.i8max
249def _decons_group_index(
250 comp_labels: npt.NDArray[np.intp], shape: Shape
251) -> list[npt.NDArray[np.intp]]:
252 # reconstruct labels
253 if is_int64_overflow_possible(shape):
254 # at some point group indices are factorized,
255 # and may not be deconstructed here! wrong path!
256 raise ValueError("cannot deconstruct factorized group indices!")
258 label_list = []
259 factor = 1
260 y = np.array(0)
261 x = comp_labels
262 for i in reversed(range(len(shape))):
263 labels = (x - y) % (factor * shape[i]) // factor
264 np.putmask(labels, comp_labels < 0, -1)
265 label_list.append(labels)
266 y = labels * factor
267 factor *= shape[i]
268 return label_list[::-1]
271def decons_obs_group_ids(
272 comp_ids: npt.NDArray[np.intp],
273 obs_ids: npt.NDArray[np.intp],
274 shape: Shape,
275 labels: Sequence[npt.NDArray[np.signedinteger]],
276 xnull: bool,
277) -> list[npt.NDArray[np.intp]]:
278 """
279 Reconstruct labels from observed group ids.
281 Parameters
282 ----------
283 comp_ids : np.ndarray[np.intp]
284 obs_ids: np.ndarray[np.intp]
285 shape : tuple[int]
286 labels : Sequence[np.ndarray[np.signedinteger]]
287 xnull : bool
288 If nulls are excluded; i.e. -1 labels are passed through.
289 """
290 if not xnull:
291 lift = np.fromiter((lib.has_sentinel(a, -1) for a in labels), dtype=np.intp)
292 arr_shape = np.asarray(shape, dtype=np.intp) + lift
293 shape = tuple(arr_shape)
295 if not is_int64_overflow_possible(shape):
296 # obs ids are deconstructable! take the fast route!
297 out = _decons_group_index(obs_ids, shape)
298 return (
299 out
300 if xnull or not lift.any()
301 else [x - y for x, y in zip(out, lift, strict=True)]
302 )
304 indexer = unique_label_indices(comp_ids)
305 return [lab[indexer].astype(np.intp, subok=False, copy=True) for lab in labels]
308def lexsort_indexer(
309 keys: Sequence[ArrayLike | Index | Series],
310 orders: bool | Sequence[bool] | None = None,
311 na_position: str = "last",
312 key: Callable | None = None,
313 codes_given: bool = False,
314) -> npt.NDArray[np.intp]:
315 """
316 Performs lexical sorting on a set of keys
318 Parameters
319 ----------
320 keys : Sequence[ArrayLike | Index | Series]
321 Sequence of arrays to be sorted by the indexer
322 Sequence[Series] is only if key is not None.
323 orders : bool or list of booleans, optional
324 Determines the sorting order for each element in keys. If a list,
325 it must be the same length as keys. This determines whether the
326 corresponding element in keys should be sorted in ascending
327 (True) or descending (False) order. if bool, applied to all
328 elements as above. if None, defaults to True.
329 na_position : {'first', 'last'}, default 'last'
330 Determines placement of NA elements in the sorted list ("last" or "first")
331 key : Callable, optional
332 Callable key function applied to every element in keys before sorting
333 codes_given: bool, False
334 Avoid categorical materialization if codes are already provided.
336 Returns
337 -------
338 np.ndarray[np.intp]
339 """
340 from pandas.core.arrays import Categorical
342 if na_position not in ["last", "first"]:
343 raise ValueError(f"invalid na_position: {na_position}")
345 orders_iter: Iterable[bool]
346 if isinstance(orders, bool):
347 orders_iter = itertools.repeat(orders, len(keys))
348 elif orders is None:
349 orders_iter = itertools.repeat(True, len(keys))
350 else:
351 orders_iter = reversed(orders)
353 labels = []
355 for k, order in zip(reversed(keys), orders_iter, strict=True):
356 k = ensure_key_mapped(k, key)
357 if codes_given:
358 codes = cast("np.ndarray", k)
359 n = codes.max() + 1 if len(codes) else 0
360 else:
361 # Fast path for numeric numpy arrays: skip Categorical
362 # conversion and pass values directly to np.lexsort.
363 arr = extract_array(k, extract_numpy=True)
364 if isinstance(arr, np.ndarray) and arr.dtype.kind in "fiub":
365 if arr.dtype.kind == "f":
366 # For float dtypes, np.lexsort sorts NaN to the end.
367 mask = np.isnan(arr)
368 has_na = mask.any()
369 if not order:
370 # Descending: negate values. NaN stays NaN
371 # and still sorts last.
372 arr = -arr
373 if na_position == "first" and has_na:
374 arr[mask] = -np.inf
375 elif na_position == "first" and has_na:
376 arr = arr.copy()
377 arr[mask] = -np.inf
378 elif not order:
379 # int/uint/bool: no NaN possible, use bitwise NOT
380 # for descending to avoid overflow with negation.
381 arr = ~arr
382 labels.append(arr)
383 continue
385 cat = Categorical(k, ordered=True)
386 codes = cat.codes
387 n = len(cat.categories)
389 mask = codes == -1
391 if na_position == "last" and mask.any():
392 codes = np.where(mask, n, codes)
394 # not order means descending
395 if not order:
396 codes = np.where(mask, codes, n - codes - 1)
398 labels.append(codes)
400 return np.lexsort(labels)
403def nargsort(
404 items: ArrayLike | Index | Series,
405 kind: SortKind = "quicksort",
406 ascending: bool = True,
407 na_position: str = "last",
408 key: Callable | None = None,
409 mask: npt.NDArray[np.bool_] | None = None,
410) -> npt.NDArray[np.intp]:
411 """
412 Intended to be a drop-in replacement for np.argsort which handles NaNs.
414 Adds ascending, na_position, and key parameters.
416 (GH #6399, #5231, #27237)
418 Parameters
419 ----------
420 items : np.ndarray, ExtensionArray, Index, or Series
421 kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, default 'quicksort'
422 ascending : bool, default True
423 na_position : {'first', 'last'}, default 'last'
424 key : Optional[Callable], default None
425 mask : Optional[np.ndarray[bool]], default None
426 Passed when called by ExtensionArray.argsort.
428 Returns
429 -------
430 np.ndarray[np.intp]
431 """
433 if key is not None:
434 # see TestDataFrameSortKey, TestRangeIndex::test_sort_values_key
435 items = ensure_key_mapped(items, key)
436 return nargsort(
437 items,
438 kind=kind,
439 ascending=ascending,
440 na_position=na_position,
441 key=None,
442 mask=mask,
443 )
445 if isinstance(items, ABCRangeIndex):
446 return items.argsort(ascending=ascending)
447 elif not isinstance(items, ABCMultiIndex):
448 items = extract_array(items)
449 else:
450 raise TypeError(
451 "nargsort does not support MultiIndex. Use index.sort_values instead."
452 )
454 if mask is None:
455 mask = np.asarray(isna(items))
457 if not isinstance(items, np.ndarray):
458 # i.e. ExtensionArray
459 return items.argsort(
460 ascending=ascending,
461 kind=kind,
462 na_position=na_position,
463 )
465 idx = np.arange(len(items))
466 non_nans = items[~mask]
467 non_nan_idx = idx[~mask]
469 nan_idx = np.nonzero(mask)[0]
470 if not ascending:
471 non_nans = non_nans[::-1]
472 non_nan_idx = non_nan_idx[::-1]
473 indexer = non_nan_idx[non_nans.argsort(kind=kind)]
474 if not ascending:
475 indexer = indexer[::-1]
476 # Finally, place the NaNs at the end or the beginning according to
477 # na_position
478 if na_position == "last":
479 indexer = np.concatenate([indexer, nan_idx])
480 elif na_position == "first":
481 indexer = np.concatenate([nan_idx, indexer])
482 else:
483 raise ValueError(f"invalid na_position: {na_position}")
484 return ensure_platform_int(indexer)
487def nargminmax(
488 values: ExtensionArray, method: str, axis: AxisInt = 0
489) -> int | np.ndarray:
490 """
491 Implementation of np.argmin/argmax but for ExtensionArray and which
492 handles missing values.
494 Parameters
495 ----------
496 values : ExtensionArray
497 method : {"argmax", "argmin"}
498 axis : int, default 0
500 Returns
501 -------
502 int
503 """
504 assert method in {"argmax", "argmin"}
505 func = np.argmax if method == "argmax" else np.argmin
507 mask = np.asarray(isna(values))
508 arr_values = values._values_for_argsort()
510 if arr_values.ndim > 1:
511 if mask.any():
512 if axis == 1:
513 zipped = zip(arr_values, mask, strict=True)
514 else:
515 zipped = zip(arr_values.T, mask.T, strict=True)
516 return np.array([_nanargminmax(v, m, func) for v, m in zipped])
517 return func(arr_values, axis=axis)
519 return _nanargminmax(arr_values, mask, func)
522def _nanargminmax(
523 values: np.ndarray, mask: npt.NDArray[np.bool_], func: Callable
524) -> int:
525 """
526 See nanargminmax.__doc__.
527 """
528 if mask.size and mask.all():
529 # func would raise "empty sequence" for an array that is not empty;
530 # match the message nanops raises for the numpy-backed dtypes
531 raise ValueError("Encountered all NA values")
533 idx = np.arange(values.shape[0])
534 non_nans = values[~mask]
535 non_nan_idx = idx[~mask]
537 return non_nan_idx[func(non_nans)]
540def _ensure_key_mapped_multiindex(
541 index: MultiIndex,
542 key: Callable,
543 level: Level | list[Level] | None = None,
544) -> MultiIndex:
545 """
546 Returns a new MultiIndex in which key has been applied
547 to all levels specified in level (or all levels if level
548 is None). Used for key sorting for MultiIndex.
550 Parameters
551 ----------
552 index : MultiIndex
553 Index to which to apply the key function on the
554 specified levels.
555 key : Callable
556 Function that takes an Index and returns an Index of
557 the same shape. This key is applied to each level
558 separately. The name of the level can be used to
559 distinguish different levels for application.
560 level : list-like, int or str, default None
561 Level or list of levels to apply the key function to.
562 If None, key function is applied to all levels. Other
563 levels are left unchanged.
565 Returns
566 -------
567 labels : MultiIndex
568 Resulting MultiIndex with modified levels.
569 """
571 if level is not None:
572 if isinstance(level, (str, int)):
573 level_iter: list[Level] = [level]
574 else:
575 level_iter = cast("list[Level]", level)
577 sort_levels: range | set = {index._get_level_number(lev) for lev in level_iter}
578 else:
579 sort_levels = range(index.nlevels)
581 mapped = [
582 (
583 ensure_key_mapped(index._get_level_values(level), key)
584 if level in sort_levels
585 else index._get_level_values(level)
586 )
587 for level in range(index.nlevels)
588 ]
590 # key may return an array without a name, keep the level names so
591 # that levels can still be referred to by name
592 return type(index).from_arrays(mapped, names=index.names)
595def ensure_key_mapped(
596 values: ArrayLike | Index | Series,
597 key: Callable | None,
598 levels: Level | list[Level] | None = None,
599) -> ArrayLike | Index | Series:
600 """
601 Applies a callable key function to the values function and checks
602 that the resulting value has the same shape. Can be called on Index
603 subclasses, Series, DataFrames, or ndarrays.
605 Parameters
606 ----------
607 values : Series, DataFrame, Index subclass, or ndarray
608 key : Optional[Callable], key to be called on the values array
609 levels : Optional[List], if values is a MultiIndex, list of levels to
610 apply the key to.
611 """
612 from pandas.core.indexes.api import Index
614 if not key:
615 return values
617 if isinstance(values, ABCMultiIndex):
618 return _ensure_key_mapped_multiindex(values, key, level=levels)
620 result = key(values.copy())
621 if len(result) != len(values):
622 raise ValueError(
623 "User-provided `key` function must not change the shape of the array."
624 )
626 try:
627 if isinstance(
628 values, Index
629 ): # convert to a new Index subclass, not necessarily the same
630 result = Index(result, tupleize_cols=False)
631 else:
632 # try to revert to original type otherwise
633 type_of_values = type(values)
634 # error: Too many arguments for "ExtensionArray"
635 result = type_of_values(result) # type: ignore[call-arg]
636 except TypeError as err:
637 raise TypeError(
638 f"User-provided `key` function returned an invalid type {type(result)} \
639 which could not be converted to {type(values)}."
640 ) from err
642 return result
645# ----------------------------------------------------------------------
646# sorting levels...cleverly?
649def get_group_index_sorter(
650 group_index: npt.NDArray[np.intp], ngroups: int | None = None
651) -> npt.NDArray[np.intp]:
652 """
653 algos.groupsort_indexer implements `counting sort` and it is at least
654 O(ngroups), where
655 ngroups = prod(shape)
656 shape = map(len, keys)
657 that is, linear in the number of combinations (cartesian product) of unique
658 values of groupby keys. This can be huge when doing multi-key groupby.
659 np.argsort(kind='mergesort') is O(count x log(count)) where count is the
660 length of the data-frame;
661 Both algorithms are `stable` sort and that is necessary for correctness of
662 groupby operations. e.g. consider:
663 df.groupby(key)[col].transform('first')
665 Parameters
666 ----------
667 group_index : np.ndarray[np.intp]
668 signed integer dtype
669 ngroups : int or None, default None
671 Returns
672 -------
673 np.ndarray[np.intp]
674 """
675 if ngroups is None:
676 ngroups = 1 + group_index.max() # type: ignore[assignment]
677 count = len(group_index)
678 alpha = 0.0 # taking complexities literally; there may be
679 beta = 1.0 # some room for fine-tuning these parameters
680 do_groupsort = count > 0 and ((alpha + beta * ngroups) < (count * np.log(count))) # type: ignore[operator]
681 if do_groupsort:
682 sorter, _ = algos.groupsort_indexer(
683 ensure_platform_int(group_index),
684 ngroups, # type: ignore[arg-type]
685 )
686 # sorter _should_ already be intp, but mypy is not yet able to verify
687 else:
688 sorter = group_index.argsort(kind="mergesort")
689 return ensure_platform_int(sorter)
692def compress_group_index(
693 group_index: npt.NDArray[np.int64], sort: bool = True
694) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
695 """
696 Group_index is offsets into cartesian product of all possible labels. This
697 space can be huge, so this function compresses it, by computing offsets
698 (comp_ids) into the list of unique labels (obs_group_ids).
699 """
700 if len(group_index) and np.all(group_index[1:] >= group_index[:-1]):
701 # GH 53806: fast path for sorted group_index
702 unique_mask = np.concatenate(
703 [group_index[:1] > -1, group_index[1:] != group_index[:-1]]
704 )
705 comp_ids = unique_mask.cumsum()
706 comp_ids -= 1
707 obs_group_ids = group_index[unique_mask]
708 else:
709 size_hint = len(group_index)
710 table = hashtable.Int64HashTable(size_hint)
712 group_index = ensure_int64(group_index)
714 # note, group labels come out ascending (ie, 1,2,3 etc)
715 comp_ids, obs_group_ids = table.get_labels_groupby(group_index)
717 if sort and len(obs_group_ids) > 0:
718 obs_group_ids, comp_ids = _reorder_by_uniques(obs_group_ids, comp_ids)
720 return ensure_int64(comp_ids), ensure_int64(obs_group_ids)
723def _reorder_by_uniques(
724 uniques: npt.NDArray[np.int64], labels: npt.NDArray[np.intp]
725) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.intp]]:
726 """
727 Parameters
728 ----------
729 uniques : np.ndarray[np.int64]
730 labels : np.ndarray[np.intp]
732 Returns
733 -------
734 np.ndarray[np.int64]
735 np.ndarray[np.intp]
736 """
737 # sorter is index where elements ought to go
738 sorter = uniques.argsort()
740 # reverse_indexer is where elements came from
741 reverse_indexer = np.empty(len(sorter), dtype=np.intp)
742 reverse_indexer.put(sorter, np.arange(len(sorter)))
744 mask = labels < 0
746 # move labels to right locations (ie, unsort ascending labels)
747 labels = reverse_indexer.take(labels)
748 np.putmask(labels, mask, -1)
750 # sort observed ids
751 uniques = uniques.take(sorter)
753 return uniques, labels