Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/sorting.py: 12%
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 Hashable,
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 level is not None:
94 _, indexer = target.sortlevel(
95 level,
96 ascending=ascending,
97 sort_remaining=sort_remaining,
98 na_position=na_position,
99 )
100 elif (np.all(ascending) and target.is_monotonic_increasing) or (
101 not np.any(ascending) and target.is_monotonic_decreasing
102 ):
103 # Check monotonic-ness before sort an index (GH 11080)
104 return None
105 elif isinstance(target, ABCMultiIndex):
106 codes = [lev.codes for lev in target._get_codes_for_sorting()]
107 indexer = lexsort_indexer(
108 codes, orders=ascending, na_position=na_position, codes_given=True
109 )
110 else:
111 # ascending can only be a Sequence for MultiIndex
112 indexer = nargsort(
113 target,
114 kind=kind,
115 ascending=cast(bool, ascending),
116 na_position=na_position,
117 )
118 return indexer
121def get_group_index(
122 labels, shape: Shape, sort: bool, xnull: bool
123) -> npt.NDArray[np.int64]:
124 """
125 For the particular label_list, gets the offsets into the hypothetical list
126 representing the totally ordered cartesian product of all possible label
127 combinations, *as long as* this space fits within int64 bounds;
128 otherwise, though group indices identify unique combinations of
129 labels, they cannot be deconstructed.
130 - If `sort`, rank of returned ids preserve lexical ranks of labels.
131 i.e. returned id's can be used to do lexical sort on labels;
132 - If `xnull` nulls (-1 labels) are passed through.
134 Parameters
135 ----------
136 labels : sequence of arrays
137 Integers identifying levels at each location
138 shape : tuple[int, ...]
139 Number of unique levels at each location
140 sort : bool
141 If the ranks of returned ids should match lexical ranks of labels
142 xnull : bool
143 If true nulls are excluded. i.e. -1 values in the labels are
144 passed through.
146 Returns
147 -------
148 An array of type int64 where two elements are equal if their corresponding
149 labels are equal at all location.
151 Notes
152 -----
153 The length of `labels` and `shape` must be identical.
154 """
156 def _int64_cut_off(shape) -> int:
157 acc = 1
158 for i, mul in enumerate(shape):
159 acc *= int(mul)
160 if not acc < lib.i8max:
161 return i
162 return len(shape)
164 def maybe_lift(lab, size: int) -> tuple[np.ndarray, int]:
165 # promote nan values (assigned -1 label in lab array)
166 # so that all output values are non-negative
167 return (lab + 1, size + 1) if (lab == -1).any() else (lab, size)
169 labels = [ensure_int64(x) for x in labels]
170 lshape = list(shape)
171 if not xnull:
172 for i, (lab, size) in enumerate(zip(labels, shape, strict=True)):
173 labels[i], lshape[i] = maybe_lift(lab, size)
175 # Iteratively process all the labels in chunks sized so less
176 # than lib.i8max unique int ids will be required for each chunk
177 while True:
178 # how many levels can be done without overflow:
179 nlev = _int64_cut_off(lshape)
181 # compute flat ids for the first `nlev` levels
182 stride = np.prod(lshape[1:nlev], dtype="i8")
183 out = stride * labels[0].astype("i8", subok=False, copy=False)
185 for i in range(1, nlev):
186 if lshape[i] == 0:
187 stride = np.int64(0)
188 else:
189 stride //= lshape[i]
190 out += labels[i] * stride
192 if xnull: # exclude nulls
193 mask = labels[0] == -1
194 for lab in labels[1:nlev]:
195 mask |= lab == -1
196 out[mask] = -1
198 if nlev == len(lshape): # all levels done!
199 break
201 # compress what has been done so far in order to avoid overflow
202 # to retain lexical ranks, obs_ids should be sorted
203 comp_ids, obs_ids = compress_group_index(out, sort=sort)
205 labels = [comp_ids, *labels[nlev:]]
206 lshape = [len(obs_ids), *lshape[nlev:]]
208 return out
211def get_compressed_ids(
212 labels, sizes: Shape
213) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.int64]]:
214 """
215 Group_index is offsets into cartesian product of all possible labels. This
216 space can be huge, so this function compresses it, by computing offsets
217 (comp_ids) into the list of unique labels (obs_group_ids).
219 Parameters
220 ----------
221 labels : list of label arrays
222 sizes : tuple[int] of size of the levels
224 Returns
225 -------
226 np.ndarray[np.intp]
227 comp_ids
228 np.ndarray[np.int64]
229 obs_group_ids
230 """
231 ids = get_group_index(labels, sizes, sort=True, xnull=False)
232 return compress_group_index(ids, sort=True)
235def is_int64_overflow_possible(shape: Shape) -> bool:
236 the_prod = 1
237 for x in shape:
238 the_prod *= int(x)
240 return the_prod >= lib.i8max
243def _decons_group_index(
244 comp_labels: npt.NDArray[np.intp], shape: Shape
245) -> list[npt.NDArray[np.intp]]:
246 # reconstruct labels
247 if is_int64_overflow_possible(shape):
248 # at some point group indices are factorized,
249 # and may not be deconstructed here! wrong path!
250 raise ValueError("cannot deconstruct factorized group indices!")
252 label_list = []
253 factor = 1
254 y = np.array(0)
255 x = comp_labels
256 for i in reversed(range(len(shape))):
257 labels = (x - y) % (factor * shape[i]) // factor
258 np.putmask(labels, comp_labels < 0, -1)
259 label_list.append(labels)
260 y = labels * factor
261 factor *= shape[i]
262 return label_list[::-1]
265def decons_obs_group_ids(
266 comp_ids: npt.NDArray[np.intp],
267 obs_ids: npt.NDArray[np.intp],
268 shape: Shape,
269 labels: Sequence[npt.NDArray[np.signedinteger]],
270 xnull: bool,
271) -> list[npt.NDArray[np.intp]]:
272 """
273 Reconstruct labels from observed group ids.
275 Parameters
276 ----------
277 comp_ids : np.ndarray[np.intp]
278 obs_ids: np.ndarray[np.intp]
279 shape : tuple[int]
280 labels : Sequence[np.ndarray[np.signedinteger]]
281 xnull : bool
282 If nulls are excluded; i.e. -1 labels are passed through.
283 """
284 if not xnull:
285 lift = np.fromiter(((a == -1).any() for a in labels), dtype=np.intp)
286 arr_shape = np.asarray(shape, dtype=np.intp) + lift
287 shape = tuple(arr_shape)
289 if not is_int64_overflow_possible(shape):
290 # obs ids are deconstructable! take the fast route!
291 out = _decons_group_index(obs_ids, shape)
292 return (
293 out
294 if xnull or not lift.any()
295 else [x - y for x, y in zip(out, lift, strict=True)]
296 )
298 indexer = unique_label_indices(comp_ids)
299 return [lab[indexer].astype(np.intp, subok=False, copy=True) for lab in labels]
302def lexsort_indexer(
303 keys: Sequence[ArrayLike | Index | Series],
304 orders=None,
305 na_position: str = "last",
306 key: Callable | None = None,
307 codes_given: bool = False,
308) -> npt.NDArray[np.intp]:
309 """
310 Performs lexical sorting on a set of keys
312 Parameters
313 ----------
314 keys : Sequence[ArrayLike | Index | Series]
315 Sequence of arrays to be sorted by the indexer
316 Sequence[Series] is only if key is not None.
317 orders : bool or list of booleans, optional
318 Determines the sorting order for each element in keys. If a list,
319 it must be the same length as keys. This determines whether the
320 corresponding element in keys should be sorted in ascending
321 (True) or descending (False) order. if bool, applied to all
322 elements as above. if None, defaults to True.
323 na_position : {'first', 'last'}, default 'last'
324 Determines placement of NA elements in the sorted list ("last" or "first")
325 key : Callable, optional
326 Callable key function applied to every element in keys before sorting
327 codes_given: bool, False
328 Avoid categorical materialization if codes are already provided.
330 Returns
331 -------
332 np.ndarray[np.intp]
333 """
334 from pandas.core.arrays import Categorical
336 if na_position not in ["last", "first"]:
337 raise ValueError(f"invalid na_position: {na_position}")
339 if isinstance(orders, bool):
340 orders = itertools.repeat(orders, len(keys))
341 elif orders is None:
342 orders = itertools.repeat(True, len(keys))
343 else:
344 orders = reversed(orders)
346 labels = []
348 for k, order in zip(reversed(keys), orders, strict=True):
349 k = ensure_key_mapped(k, key)
350 if codes_given:
351 codes = cast(np.ndarray, k)
352 n = codes.max() + 1 if len(codes) else 0
353 else:
354 cat = Categorical(k, ordered=True)
355 codes = cat.codes
356 n = len(cat.categories)
358 mask = codes == -1
360 if na_position == "last" and mask.any():
361 codes = np.where(mask, n, codes)
363 # not order means descending
364 if not order:
365 codes = np.where(mask, codes, n - codes - 1)
367 labels.append(codes)
369 return np.lexsort(labels)
372def nargsort(
373 items: ArrayLike | Index | Series,
374 kind: SortKind = "quicksort",
375 ascending: bool = True,
376 na_position: str = "last",
377 key: Callable | None = None,
378 mask: npt.NDArray[np.bool_] | None = None,
379) -> npt.NDArray[np.intp]:
380 """
381 Intended to be a drop-in replacement for np.argsort which handles NaNs.
383 Adds ascending, na_position, and key parameters.
385 (GH #6399, #5231, #27237)
387 Parameters
388 ----------
389 items : np.ndarray, ExtensionArray, Index, or Series
390 kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, default 'quicksort'
391 ascending : bool, default True
392 na_position : {'first', 'last'}, default 'last'
393 key : Optional[Callable], default None
394 mask : Optional[np.ndarray[bool]], default None
395 Passed when called by ExtensionArray.argsort.
397 Returns
398 -------
399 np.ndarray[np.intp]
400 """
402 if key is not None:
403 # see TestDataFrameSortKey, TestRangeIndex::test_sort_values_key
404 items = ensure_key_mapped(items, key)
405 return nargsort(
406 items,
407 kind=kind,
408 ascending=ascending,
409 na_position=na_position,
410 key=None,
411 mask=mask,
412 )
414 if isinstance(items, ABCRangeIndex):
415 return items.argsort(ascending=ascending)
416 elif not isinstance(items, ABCMultiIndex):
417 items = extract_array(items)
418 else:
419 raise TypeError(
420 "nargsort does not support MultiIndex. Use index.sort_values instead."
421 )
423 if mask is None:
424 mask = np.asarray(isna(items))
426 if not isinstance(items, np.ndarray):
427 # i.e. ExtensionArray
428 return items.argsort(
429 ascending=ascending,
430 kind=kind,
431 na_position=na_position,
432 )
434 idx = np.arange(len(items))
435 non_nans = items[~mask]
436 non_nan_idx = idx[~mask]
438 nan_idx = np.nonzero(mask)[0]
439 if not ascending:
440 non_nans = non_nans[::-1]
441 non_nan_idx = non_nan_idx[::-1]
442 indexer = non_nan_idx[non_nans.argsort(kind=kind)]
443 if not ascending:
444 indexer = indexer[::-1]
445 # Finally, place the NaNs at the end or the beginning according to
446 # na_position
447 if na_position == "last":
448 indexer = np.concatenate([indexer, nan_idx])
449 elif na_position == "first":
450 indexer = np.concatenate([nan_idx, indexer])
451 else:
452 raise ValueError(f"invalid na_position: {na_position}")
453 return ensure_platform_int(indexer)
456def nargminmax(values: ExtensionArray, method: str, axis: AxisInt = 0):
457 """
458 Implementation of np.argmin/argmax but for ExtensionArray and which
459 handles missing values.
461 Parameters
462 ----------
463 values : ExtensionArray
464 method : {"argmax", "argmin"}
465 axis : int, default 0
467 Returns
468 -------
469 int
470 """
471 assert method in {"argmax", "argmin"}
472 func = np.argmax if method == "argmax" else np.argmin
474 mask = np.asarray(isna(values))
475 arr_values = values._values_for_argsort()
477 if arr_values.ndim > 1:
478 if mask.any():
479 if axis == 1:
480 zipped = zip(arr_values, mask, strict=True)
481 else:
482 zipped = zip(arr_values.T, mask.T, strict=True)
483 return np.array([_nanargminmax(v, m, func) for v, m in zipped])
484 return func(arr_values, axis=axis)
486 return _nanargminmax(arr_values, mask, func)
489def _nanargminmax(values: np.ndarray, mask: npt.NDArray[np.bool_], func) -> int:
490 """
491 See nanargminmax.__doc__.
492 """
493 idx = np.arange(values.shape[0])
494 non_nans = values[~mask]
495 non_nan_idx = idx[~mask]
497 return non_nan_idx[func(non_nans)]
500def _ensure_key_mapped_multiindex(
501 index: MultiIndex, key: Callable, level=None
502) -> MultiIndex:
503 """
504 Returns a new MultiIndex in which key has been applied
505 to all levels specified in level (or all levels if level
506 is None). Used for key sorting for MultiIndex.
508 Parameters
509 ----------
510 index : MultiIndex
511 Index to which to apply the key function on the
512 specified levels.
513 key : Callable
514 Function that takes an Index and returns an Index of
515 the same shape. This key is applied to each level
516 separately. The name of the level can be used to
517 distinguish different levels for application.
518 level : list-like, int or str, default None
519 Level or list of levels to apply the key function to.
520 If None, key function is applied to all levels. Other
521 levels are left unchanged.
523 Returns
524 -------
525 labels : MultiIndex
526 Resulting MultiIndex with modified levels.
527 """
529 if level is not None:
530 if isinstance(level, (str, int)):
531 level_iter = [level]
532 else:
533 level_iter = level
535 sort_levels: range | set = {index._get_level_number(lev) for lev in level_iter}
536 else:
537 sort_levels = range(index.nlevels)
539 mapped = [
540 (
541 ensure_key_mapped(index._get_level_values(level), key)
542 if level in sort_levels
543 else index._get_level_values(level)
544 )
545 for level in range(index.nlevels)
546 ]
548 return type(index).from_arrays(mapped)
551def ensure_key_mapped(
552 values: ArrayLike | Index | Series, key: Callable | None, levels=None
553) -> ArrayLike | Index | Series:
554 """
555 Applies a callable key function to the values function and checks
556 that the resulting value has the same shape. Can be called on Index
557 subclasses, Series, DataFrames, or ndarrays.
559 Parameters
560 ----------
561 values : Series, DataFrame, Index subclass, or ndarray
562 key : Optional[Callable], key to be called on the values array
563 levels : Optional[List], if values is a MultiIndex, list of levels to
564 apply the key to.
565 """
566 from pandas.core.indexes.api import Index
568 if not key:
569 return values
571 if isinstance(values, ABCMultiIndex):
572 return _ensure_key_mapped_multiindex(values, key, level=levels)
574 result = key(values.copy())
575 if len(result) != len(values):
576 raise ValueError(
577 "User-provided `key` function must not change the shape of the array."
578 )
580 try:
581 if isinstance(
582 values, Index
583 ): # convert to a new Index subclass, not necessarily the same
584 result = Index(result, tupleize_cols=False)
585 else:
586 # try to revert to original type otherwise
587 type_of_values = type(values)
588 # error: Too many arguments for "ExtensionArray"
589 result = type_of_values(result) # type: ignore[call-arg]
590 except TypeError as err:
591 raise TypeError(
592 f"User-provided `key` function returned an invalid type {type(result)} \
593 which could not be converted to {type(values)}."
594 ) from err
596 return result
599def get_indexer_dict(
600 label_list: list[np.ndarray], keys: list[Index]
601) -> dict[Hashable, npt.NDArray[np.intp]]:
602 """
603 Returns
604 -------
605 dict:
606 Labels mapped to indexers.
607 """
608 shape = tuple(len(x) for x in keys)
610 group_index = get_group_index(label_list, shape, sort=True, xnull=True)
611 if np.all(group_index == -1):
612 # Short-circuit, lib.indices_fast will return the same
613 return {}
614 ngroups = (
615 ((group_index.size and group_index.max()) + 1)
616 if is_int64_overflow_possible(shape)
617 else np.prod(shape, dtype="i8")
618 )
620 sorter = get_group_index_sorter(group_index, ngroups)
622 sorted_labels = [lab.take(sorter) for lab in label_list]
623 group_index = group_index.take(sorter)
625 return lib.indices_fast(sorter, group_index, keys, sorted_labels)
628# ----------------------------------------------------------------------
629# sorting levels...cleverly?
632def get_group_index_sorter(
633 group_index: npt.NDArray[np.intp], ngroups: int | None = None
634) -> npt.NDArray[np.intp]:
635 """
636 algos.groupsort_indexer implements `counting sort` and it is at least
637 O(ngroups), where
638 ngroups = prod(shape)
639 shape = map(len, keys)
640 that is, linear in the number of combinations (cartesian product) of unique
641 values of groupby keys. This can be huge when doing multi-key groupby.
642 np.argsort(kind='mergesort') is O(count x log(count)) where count is the
643 length of the data-frame;
644 Both algorithms are `stable` sort and that is necessary for correctness of
645 groupby operations. e.g. consider:
646 df.groupby(key)[col].transform('first')
648 Parameters
649 ----------
650 group_index : np.ndarray[np.intp]
651 signed integer dtype
652 ngroups : int or None, default None
654 Returns
655 -------
656 np.ndarray[np.intp]
657 """
658 if ngroups is None:
659 ngroups = 1 + group_index.max()
660 count = len(group_index)
661 alpha = 0.0 # taking complexities literally; there may be
662 beta = 1.0 # some room for fine-tuning these parameters
663 do_groupsort = count > 0 and ((alpha + beta * ngroups) < (count * np.log(count)))
664 if do_groupsort:
665 sorter, _ = algos.groupsort_indexer(
666 ensure_platform_int(group_index),
667 ngroups,
668 )
669 # sorter _should_ already be intp, but mypy is not yet able to verify
670 else:
671 sorter = group_index.argsort(kind="mergesort")
672 return ensure_platform_int(sorter)
675def compress_group_index(
676 group_index: npt.NDArray[np.int64], sort: bool = True
677) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
678 """
679 Group_index is offsets into cartesian product of all possible labels. This
680 space can be huge, so this function compresses it, by computing offsets
681 (comp_ids) into the list of unique labels (obs_group_ids).
682 """
683 if len(group_index) and np.all(group_index[1:] >= group_index[:-1]):
684 # GH 53806: fast path for sorted group_index
685 unique_mask = np.concatenate(
686 [group_index[:1] > -1, group_index[1:] != group_index[:-1]]
687 )
688 comp_ids = unique_mask.cumsum()
689 comp_ids -= 1
690 obs_group_ids = group_index[unique_mask]
691 else:
692 size_hint = len(group_index)
693 table = hashtable.Int64HashTable(size_hint)
695 group_index = ensure_int64(group_index)
697 # note, group labels come out ascending (ie, 1,2,3 etc)
698 comp_ids, obs_group_ids = table.get_labels_groupby(group_index)
700 if sort and len(obs_group_ids) > 0:
701 obs_group_ids, comp_ids = _reorder_by_uniques(obs_group_ids, comp_ids)
703 return ensure_int64(comp_ids), ensure_int64(obs_group_ids)
706def _reorder_by_uniques(
707 uniques: npt.NDArray[np.int64], labels: npt.NDArray[np.intp]
708) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.intp]]:
709 """
710 Parameters
711 ----------
712 uniques : np.ndarray[np.int64]
713 labels : np.ndarray[np.intp]
715 Returns
716 -------
717 np.ndarray[np.int64]
718 np.ndarray[np.intp]
719 """
720 # sorter is index where elements ought to go
721 sorter = uniques.argsort()
723 # reverse_indexer is where elements came from
724 reverse_indexer = np.empty(len(sorter), dtype=np.intp)
725 reverse_indexer.put(sorter, np.arange(len(sorter)))
727 mask = labels < 0
729 # move labels to right locations (ie, unsort ascending labels)
730 labels = reverse_indexer.take(labels)
731 np.putmask(labels, mask, -1)
733 # sort observed ids
734 uniques = uniques.take(sorter)
736 return uniques, labels