1"""
2Generic data algorithms. This module is experimental at the moment and not
3intended for public consumption
4"""
5
6from __future__ import annotations
7
8import decimal
9import operator
10from typing import (
11 TYPE_CHECKING,
12 Literal,
13 TypeVar,
14 cast,
15 overload,
16)
17import warnings
18
19import numpy as np
20
21from pandas._libs import (
22 algos,
23 hashtable as htable,
24 iNaT,
25 lib,
26)
27from pandas._libs.missing import NA
28from pandas._typing import (
29 AnyArrayLike,
30 ArrayLike,
31 ArrayLikeT,
32 AxisInt,
33 DtypeObj,
34 TakeIndexer,
35 npt,
36)
37from pandas.util._decorators import set_module
38from pandas.util._exceptions import find_stack_level
39
40from pandas.core.dtypes.cast import (
41 construct_1d_object_array_from_listlike,
42 np_find_common_type,
43)
44from pandas.core.dtypes.common import (
45 ensure_float64,
46 ensure_object,
47 ensure_platform_int,
48 is_bool_dtype,
49 is_complex_dtype,
50 is_dict_like,
51 is_dtype_equal,
52 is_extension_array_dtype,
53 is_float,
54 is_float_dtype,
55 is_integer,
56 is_integer_dtype,
57 is_list_like,
58 is_object_dtype,
59 is_signed_integer_dtype,
60 needs_i8_conversion,
61)
62from pandas.core.dtypes.concat import concat_compat
63from pandas.core.dtypes.dtypes import (
64 BaseMaskedDtype,
65 CategoricalDtype,
66 ExtensionDtype,
67 NumpyEADtype,
68)
69from pandas.core.dtypes.generic import (
70 ABCDatetimeArray,
71 ABCExtensionArray,
72 ABCIndex,
73 ABCMultiIndex,
74 ABCNumpyExtensionArray,
75 ABCSeries,
76 ABCTimedeltaArray,
77)
78from pandas.core.dtypes.missing import (
79 isna,
80 na_value_for_dtype,
81)
82
83from pandas.core.array_algos.take import take_nd
84from pandas.core.construction import (
85 array as pd_array,
86 ensure_wrapped_if_datetimelike,
87 extract_array,
88)
89from pandas.core.indexers import validate_indices
90
91if TYPE_CHECKING:
92 from pandas._typing import (
93 ListLike,
94 NumpySorter,
95 NumpyValueArrayLike,
96 )
97
98 from pandas import (
99 Categorical,
100 Index,
101 Series,
102 )
103 from pandas.core.arrays import (
104 BaseMaskedArray,
105 ExtensionArray,
106 )
107
108 T = TypeVar("T", bound=Index | Categorical | ExtensionArray)
109
110
111# --------------- #
112# dtype access #
113# --------------- #
114def _ensure_data(values: ArrayLike) -> np.ndarray:
115 """
116 routine to ensure that our data is of the correct
117 input dtype for lower-level routines
118
119 This will coerce:
120 - ints -> int64
121 - uint -> uint64
122 - bool -> uint8
123 - datetimelike -> i8
124 - datetime64tz -> i8 (in local tz)
125 - categorical -> codes
126
127 Parameters
128 ----------
129 values : np.ndarray or ExtensionArray
130
131 Returns
132 -------
133 np.ndarray
134 """
135
136 if not isinstance(values, ABCMultiIndex):
137 # extract_array would raise
138 values = extract_array(values, extract_numpy=True)
139
140 if is_object_dtype(values.dtype):
141 return ensure_object(np.asarray(values))
142
143 elif isinstance(values.dtype, BaseMaskedDtype):
144 # i.e. BooleanArray, FloatingArray, IntegerArray
145 values = cast("BaseMaskedArray", values)
146 if not values._hasna:
147 # No pd.NAs -> We can avoid an object-dtype cast (and copy) GH#41816
148 # recurse to avoid re-implementing logic for eg bool->uint8
149 return _ensure_data(values._data)
150 return np.asarray(values)
151
152 elif isinstance(values.dtype, CategoricalDtype):
153 # NB: cases that go through here should NOT be using _reconstruct_data
154 # on the back-end.
155 values = cast("Categorical", values)
156 return values.codes
157
158 elif is_bool_dtype(values.dtype):
159 if isinstance(values, np.ndarray):
160 # i.e. actually dtype == np.dtype("bool")
161 return np.asarray(values).view("uint8")
162 else:
163 # e.g. Sparse[bool, False] # TODO: no test cases get here
164 return np.asarray(values).astype("uint8", copy=False)
165
166 elif is_integer_dtype(values.dtype):
167 return np.asarray(values)
168
169 elif is_float_dtype(values.dtype):
170 # Note: checking `values.dtype == "float128"` raises on Windows and 32bit
171 # error: Item "ExtensionDtype" of "Union[Any, ExtensionDtype, dtype[Any]]"
172 # has no attribute "itemsize"
173 if values.dtype.itemsize in [2, 12, 16]: # type: ignore[union-attr]
174 # we dont (yet) have float128 hashtable support
175 return ensure_float64(values)
176 return np.asarray(values)
177
178 elif is_complex_dtype(values.dtype):
179 return cast(np.ndarray, values)
180
181 # datetimelike
182 elif needs_i8_conversion(values.dtype):
183 npvalues = values.view("i8")
184 npvalues = cast(np.ndarray, npvalues)
185 return npvalues
186
187 # we have failed, return object
188 values = np.asarray(values, dtype=object)
189 return ensure_object(values)
190
191
192def _reconstruct_data(
193 values: ArrayLikeT, dtype: DtypeObj, original: AnyArrayLike
194) -> ArrayLikeT:
195 """
196 reverse of _ensure_data
197
198 Parameters
199 ----------
200 values : np.ndarray or ExtensionArray
201 dtype : np.dtype or ExtensionDtype
202 original : AnyArrayLike
203
204 Returns
205 -------
206 ExtensionArray or np.ndarray
207 """
208 if isinstance(values, ABCExtensionArray) and values.dtype == dtype:
209 # Catch DatetimeArray/TimedeltaArray
210 return values
211
212 if not isinstance(dtype, np.dtype):
213 # i.e. ExtensionDtype; note we have ruled out above the possibility
214 # that values.dtype == dtype
215 cls = dtype.construct_array_type()
216
217 # error: Incompatible return value type
218 # (got "ExtensionArray",
219 # expected "ndarray[tuple[Any, ...], dtype[Any]]")
220 return cls._from_sequence(values, dtype=dtype) # type: ignore[return-value]
221
222 # error: Incompatible return value type
223 # (got "ndarray[tuple[Any, ...], dtype[Any]]",
224 # expected "ExtensionArray")
225 return values.astype(dtype, copy=False) # type: ignore[return-value]
226
227
228def _ensure_arraylike(values, func_name: str) -> ArrayLike:
229 """
230 ensure that we are arraylike if not already
231 """
232 if not isinstance(
233 values,
234 (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray, ABCNumpyExtensionArray),
235 ):
236 # GH#52986
237 if func_name != "isin-targets":
238 # Make an exception for the comps argument in isin.
239 raise TypeError(
240 f"{func_name} requires a Series, Index, "
241 f"ExtensionArray, np.ndarray or NumpyExtensionArray "
242 f"got {type(values).__name__}."
243 )
244
245 inferred = lib.infer_dtype(values, skipna=False)
246 if inferred in ["mixed", "string", "mixed-integer"]:
247 # "mixed-integer" to ensure we do not cast ["ss", 42] to str GH#22160
248 if isinstance(values, tuple):
249 values = list(values)
250 values = construct_1d_object_array_from_listlike(values)
251 else:
252 values = np.asarray(values)
253 return values
254
255
256_hashtables = {
257 "complex128": htable.Complex128HashTable,
258 "complex64": htable.Complex64HashTable,
259 "float64": htable.Float64HashTable,
260 "float32": htable.Float32HashTable,
261 "uint64": htable.UInt64HashTable,
262 "uint32": htable.UInt32HashTable,
263 "uint16": htable.UInt16HashTable,
264 "uint8": htable.UInt8HashTable,
265 "int64": htable.Int64HashTable,
266 "int32": htable.Int32HashTable,
267 "int16": htable.Int16HashTable,
268 "int8": htable.Int8HashTable,
269 "string": htable.StringHashTable,
270 "object": htable.PyObjectHashTable,
271}
272
273
274def _get_hashtable_algo(
275 values: np.ndarray,
276) -> tuple[type[htable.HashTable], np.ndarray]:
277 """
278 Parameters
279 ----------
280 values : np.ndarray
281
282 Returns
283 -------
284 htable : HashTable subclass
285 values : ndarray
286 """
287 values = _ensure_data(values)
288
289 ndtype = _check_object_for_strings(values)
290 hashtable = _hashtables[ndtype]
291 return hashtable, values
292
293
294def _check_object_for_strings(values: np.ndarray) -> str:
295 """
296 Check if we can use string hashtable instead of object hashtable.
297
298 Parameters
299 ----------
300 values : ndarray
301
302 Returns
303 -------
304 str
305 """
306 ndtype = values.dtype.name
307 if ndtype == "object":
308 # it's cheaper to use a String Hash Table than Object; we infer
309 # including nulls because that is the only difference between
310 # StringHashTable and ObjectHashtable
311 if lib.is_string_array(values, skipna=False):
312 ndtype = "string"
313 return ndtype
314
315
316# --------------- #
317# top-level algos #
318# --------------- #
319
320
321@overload
322def unique(values: T) -> T: ...
323@overload
324def unique(values: np.ndarray | Series) -> np.ndarray: ...
325
326
327@set_module("pandas")
328def unique(values):
329 """
330 Return unique values based on a hash table.
331
332 Uniques are returned in order of appearance. This does NOT sort.
333
334 Significantly faster than numpy.unique for long enough sequences.
335 Includes NA values.
336
337 Parameters
338 ----------
339 values : 1d array-like
340 The input array-like object containing values from which to extract
341 unique values.
342
343 Returns
344 -------
345 numpy.ndarray, ExtensionArray or NumpyExtensionArray
346
347 The return can be:
348
349 * Index : when the input is an Index
350 * Categorical : when the input is a Categorical dtype
351 * ndarray : when the input is a Series/ndarray
352
353 Return numpy.ndarray, ExtensionArray or NumpyExtensionArray.
354
355 See Also
356 --------
357 Index.unique : Return unique values from an Index.
358 Series.unique : Return unique values of Series object.
359
360 Examples
361 --------
362 >>> pd.unique(pd.Series([2, 1, 3, 3]))
363 array([2, 1, 3])
364
365 >>> pd.unique(pd.Series([2] + [1] * 5))
366 array([2, 1])
367
368 >>> pd.unique(pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")]))
369 array(['2016-01-01T00:00:00.000000'], dtype='datetime64[us]')
370
371 >>> pd.unique(
372 ... pd.Series(
373 ... [
374 ... pd.Timestamp("20160101", tz="US/Eastern"),
375 ... pd.Timestamp("20160101", tz="US/Eastern"),
376 ... ],
377 ... dtype="M8[ns, US/Eastern]",
378 ... )
379 ... )
380 <DatetimeArray>
381 ['2016-01-01 00:00:00-05:00']
382 Length: 1, dtype: datetime64[ns, US/Eastern]
383
384 >>> pd.unique(
385 ... pd.Index(
386 ... [
387 ... pd.Timestamp("20160101", tz="US/Eastern"),
388 ... pd.Timestamp("20160101", tz="US/Eastern"),
389 ... ],
390 ... dtype="M8[ns, US/Eastern]",
391 ... )
392 ... )
393 DatetimeIndex(['2016-01-01 00:00:00-05:00'],
394 dtype='datetime64[ns, US/Eastern]',
395 freq=None)
396
397 >>> pd.unique(np.array(list("baabc"), dtype="O"))
398 array(['b', 'a', 'c'], dtype=object)
399
400 An unordered Categorical will return categories in the
401 order of appearance.
402
403 >>> pd.unique(pd.Series(pd.Categorical(list("baabc"))))
404 ['b', 'a', 'c']
405 Categories (3, str): ['a', 'b', 'c']
406
407 >>> pd.unique(pd.Series(pd.Categorical(list("baabc"), categories=list("abc"))))
408 ['b', 'a', 'c']
409 Categories (3, str): ['a', 'b', 'c']
410
411 An ordered Categorical preserves the category ordering.
412
413 >>> pd.unique(
414 ... pd.Series(
415 ... pd.Categorical(list("baabc"), categories=list("abc"), ordered=True)
416 ... )
417 ... )
418 ['b', 'a', 'c']
419 Categories (3, str): ['a' < 'b' < 'c']
420
421 An array of tuples
422
423 >>> pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)
424 array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object)
425
426 A NumpyExtensionArray of complex
427
428 >>> pd.unique(pd.array([1 + 1j, 2, 3]))
429 <NumpyExtensionArray>
430 [(1+1j), (2+0j), (3+0j)]
431 Length: 3, dtype: complex128
432 """
433 return unique_with_mask(values)
434
435
436def nunique_ints(values: ArrayLike) -> int:
437 """
438 Return the number of unique values for integer array-likes.
439
440 Significantly faster than pandas.unique for long enough sequences.
441 No checks are done to ensure input is integral.
442
443 Parameters
444 ----------
445 values : 1d array-like
446
447 Returns
448 -------
449 int : The number of unique values in ``values``
450 """
451 if len(values) == 0:
452 return 0
453 values = _ensure_data(values)
454 # bincount requires intp
455 result = (np.bincount(values.ravel().astype("intp")) != 0).sum()
456 return result
457
458
459def unique_with_mask(values, mask: npt.NDArray[np.bool_] | None = None):
460 """See algorithms.unique for docs. Takes a mask for masked arrays."""
461 values = _ensure_arraylike(values, func_name="unique")
462
463 if isinstance(values.dtype, ExtensionDtype):
464 # Dispatch to extension dtype's unique.
465 return values.unique()
466
467 if isinstance(values, ABCIndex):
468 # Dispatch to Index's unique.
469 return values.unique()
470
471 original = values
472 hashtable, values = _get_hashtable_algo(values)
473
474 table = hashtable(len(values))
475 if mask is None:
476 uniques = table.unique(values)
477 uniques = _reconstruct_data(uniques, original.dtype, original)
478 return uniques
479
480 else:
481 uniques, mask = table.unique(values, mask=mask)
482 uniques = _reconstruct_data(uniques, original.dtype, original)
483 assert mask is not None # for mypy
484 return uniques, mask.astype("bool")
485
486
487unique1d = unique
488
489
490_MINIMUM_COMP_ARR_LEN = 1_000_000
491
492
493def isin(comps: ListLike, values: ListLike) -> npt.NDArray[np.bool_]:
494 """
495 Compute the isin boolean array.
496
497 Parameters
498 ----------
499 comps : list-like
500 values : list-like
501
502 Returns
503 -------
504 ndarray[bool]
505 Same length as `comps`.
506 """
507 if not is_list_like(comps):
508 raise TypeError(
509 "only list-like objects are allowed to be passed "
510 f"to isin(), you passed a `{type(comps).__name__}`"
511 )
512 if not is_list_like(values):
513 raise TypeError(
514 "only list-like objects are allowed to be passed "
515 f"to isin(), you passed a `{type(values).__name__}`"
516 )
517
518 if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)):
519 orig_values = list(values)
520 values = _ensure_arraylike(orig_values, func_name="isin-targets")
521
522 if (
523 len(values) > 0
524 and values.dtype.kind in "iufcb"
525 and not is_signed_integer_dtype(comps)
526 and not is_dtype_equal(values, comps)
527 ):
528 # GH#46485 Use object to avoid upcast to float64 later
529 # TODO: Share with _find_common_type_compat
530 values = construct_1d_object_array_from_listlike(orig_values)
531
532 elif isinstance(values, ABCMultiIndex):
533 # Avoid raising in extract_array
534 values = np.array(values)
535 else:
536 values = extract_array(values, extract_numpy=True, extract_range=True)
537
538 comps_array = _ensure_arraylike(comps, func_name="isin")
539 comps_array = extract_array(comps_array, extract_numpy=True)
540 if not isinstance(comps_array, np.ndarray):
541 # i.e. Extension Array
542 return comps_array.isin(values)
543
544 elif needs_i8_conversion(comps_array.dtype):
545 # Dispatch to DatetimeLikeArrayMixin.isin
546 return pd_array(comps_array).isin(values)
547 elif needs_i8_conversion(values.dtype) and not is_object_dtype(comps_array.dtype):
548 # e.g. comps_array are integers and values are datetime64s
549 return np.zeros(comps_array.shape, dtype=bool)
550 # TODO: not quite right ... Sparse/Categorical
551 elif needs_i8_conversion(values.dtype):
552 return isin(comps_array, values.astype(object))
553
554 elif isinstance(values.dtype, ExtensionDtype):
555 return isin(np.asarray(comps_array), np.asarray(values))
556
557 # GH16012
558 # Ensure np.isin doesn't get object types or it *may* throw an exception
559 # Albeit hashmap has O(1) look-up (vs. O(logn) in sorted array),
560 # isin is faster for small sizes
561
562 # GH60678
563 # Ensure values don't contain <NA>, otherwise it throws exception with np.in1d
564
565 if (
566 len(comps_array) > _MINIMUM_COMP_ARR_LEN
567 and len(values) <= 26
568 and comps_array.dtype != object
569 and not any(v is NA for v in values)
570 ):
571 # If the values include nan we need to check for nan explicitly
572 # since np.nan it not equal to np.nan
573 if isna(values).any():
574
575 def f(c, v):
576 return np.logical_or(np.isin(c, v).ravel(), np.isnan(c))
577
578 else:
579 f = lambda a, b: np.isin(a, b).ravel()
580
581 else:
582 common = np_find_common_type(values.dtype, comps_array.dtype)
583 values = values.astype(common, copy=False)
584 comps_array = comps_array.astype(common, copy=False)
585 f = htable.ismember
586
587 return f(comps_array, values)
588
589
590def factorize_array(
591 values: np.ndarray,
592 use_na_sentinel: bool = True,
593 size_hint: int | None = None,
594 na_value: object = None,
595 mask: npt.NDArray[np.bool_] | None = None,
596) -> tuple[npt.NDArray[np.intp], np.ndarray]:
597 """
598 Factorize a numpy array to codes and uniques.
599
600 This doesn't do any coercion of types or unboxing before factorization.
601
602 Parameters
603 ----------
604 values : ndarray
605 use_na_sentinel : bool, default True
606 If True, the sentinel -1 will be used for NaN values. If False,
607 NaN values will be encoded as non-negative integers and will not drop the
608 NaN from the uniques of the values.
609 size_hint : int, optional
610 Passed through to the hashtable's 'get_labels' method
611 na_value : object, optional
612 A value in `values` to consider missing. Note: only use this
613 parameter when you know that you don't have any values pandas would
614 consider missing in the array (NaN for float data, iNaT for
615 datetimes, etc.).
616 mask : ndarray[bool], optional
617 If not None, the mask is used as indicator for missing values
618 (True = missing, False = valid) instead of `na_value` or
619 condition "val != val".
620
621 Returns
622 -------
623 codes : ndarray[np.intp]
624 uniques : ndarray
625 """
626 original = values
627 if values.dtype.kind in "mM":
628 # _get_hashtable_algo will cast dt64/td64 to i8 via _ensure_data, so we
629 # need to do the same to na_value. We are assuming here that the passed
630 # na_value is an appropriately-typed NaT.
631 # e.g. test_where_datetimelike_categorical
632 na_value = iNaT
633
634 hash_klass, values = _get_hashtable_algo(values)
635
636 table = hash_klass(size_hint or len(values))
637 uniques, codes = table.factorize(
638 values,
639 na_sentinel=-1,
640 na_value=na_value,
641 mask=mask,
642 ignore_na=use_na_sentinel,
643 )
644
645 # re-cast e.g. i8->dt64/td64, uint8->bool
646 uniques = _reconstruct_data(uniques, original.dtype, original)
647
648 codes = ensure_platform_int(codes)
649 return codes, uniques
650
651
652@set_module("pandas")
653def factorize(
654 values,
655 sort: bool = False,
656 use_na_sentinel: bool = True,
657 size_hint: int | None = None,
658) -> tuple[np.ndarray, np.ndarray | Index]:
659 """
660 Encode the object as an enumerated type or categorical variable.
661
662 This method is useful for obtaining a numeric representation of an
663 array when all that matters is identifying distinct values. `factorize`
664 is available as both a top-level function :func:`pandas.factorize`,
665 and as a method :meth:`Series.factorize` and :meth:`Index.factorize`.
666
667 Parameters
668 ----------
669 values : sequence
670 A 1-D sequence. Sequences that aren't pandas objects are
671 coerced to ndarrays before factorization.
672 sort : bool, default False
673 Sort `uniques` and shuffle `codes` to maintain the
674 relationship.
675 use_na_sentinel : bool, default True
676 If True, the sentinel -1 will be used for NaN values. If False,
677 NaN values will be encoded as non-negative integers and will not drop the
678 NaN from the uniques of the values.
679 size_hint : int, optional
680 Hint to the hashtable sizer.
681
682 Returns
683 -------
684 codes : ndarray
685 An integer ndarray that's an indexer into `uniques`.
686 ``uniques.take(codes)`` will have the same values as `values`.
687 uniques : ndarray, Index, or Categorical
688 The unique valid values. When `values` is Categorical, `uniques`
689 is a Categorical. When `values` is some other pandas object, an
690 `Index` is returned. Otherwise, a 1-D ndarray is returned.
691
692 .. note::
693
694 Even if there's a missing value in `values`, `uniques` will
695 *not* contain an entry for it.
696
697 See Also
698 --------
699 cut : Discretize continuous-valued array.
700 unique : Find the unique value in an array.
701
702 Notes
703 -----
704 Reference :ref:`the user guide <reshaping.factorize>` for more examples.
705
706 Examples
707 --------
708 These examples all show factorize as a top-level method like
709 ``pd.factorize(values)``. The results are identical for methods like
710 :meth:`Series.factorize`.
711
712 >>> codes, uniques = pd.factorize(np.array(["b", "b", "a", "c", "b"], dtype="O"))
713 >>> codes
714 array([0, 0, 1, 2, 0])
715 >>> uniques
716 array(['b', 'a', 'c'], dtype=object)
717
718 With ``sort=True``, the `uniques` will be sorted, and `codes` will be
719 shuffled so that the relationship is the maintained.
720
721 >>> codes, uniques = pd.factorize(
722 ... np.array(["b", "b", "a", "c", "b"], dtype="O"), sort=True
723 ... )
724 >>> codes
725 array([1, 1, 0, 2, 1])
726 >>> uniques
727 array(['a', 'b', 'c'], dtype=object)
728
729 When ``use_na_sentinel=True`` (the default), missing values are indicated in
730 the `codes` with the sentinel value ``-1`` and missing values are not
731 included in `uniques`.
732
733 >>> codes, uniques = pd.factorize(np.array(["b", None, "a", "c", "b"], dtype="O"))
734 >>> codes
735 array([ 0, -1, 1, 2, 0])
736 >>> uniques
737 array(['b', 'a', 'c'], dtype=object)
738
739 Thus far, we've only factorized lists (which are internally coerced to
740 NumPy arrays). When factorizing pandas objects, the type of `uniques`
741 will differ. For Categoricals, a `Categorical` is returned.
742
743 >>> cat = pd.Categorical(["a", "a", "c"], categories=["a", "b", "c"])
744 >>> codes, uniques = pd.factorize(cat)
745 >>> codes
746 array([0, 0, 1])
747 >>> uniques
748 ['a', 'c']
749 Categories (3, str): ['a', 'b', 'c']
750
751 Notice that ``'b'`` is in ``uniques.categories``, despite not being
752 present in ``cat.values``.
753
754 For all other pandas objects, an Index of the appropriate type is
755 returned.
756
757 >>> cat = pd.Series(["a", "a", "c"])
758 >>> codes, uniques = pd.factorize(cat)
759 >>> codes
760 array([0, 0, 1])
761 >>> uniques
762 Index(['a', 'c'], dtype='str')
763
764 If NaN is in the values, and we want to include NaN in the uniques of the
765 values, it can be achieved by setting ``use_na_sentinel=False``.
766
767 >>> values = np.array([1, 2, 1, np.nan])
768 >>> codes, uniques = pd.factorize(values) # default: use_na_sentinel=True
769 >>> codes
770 array([ 0, 1, 0, -1])
771 >>> uniques
772 array([1., 2.])
773
774 >>> codes, uniques = pd.factorize(values, use_na_sentinel=False)
775 >>> codes
776 array([0, 1, 0, 2])
777 >>> uniques
778 array([ 1., 2., nan])
779 """
780 # Implementation notes: This method is responsible for 3 things
781 # 1.) coercing data to array-like (ndarray, Index, extension array)
782 # 2.) factorizing codes and uniques
783 # 3.) Maybe boxing the uniques in an Index
784 #
785 # Step 2 is dispatched to extension types (like Categorical). They are
786 # responsible only for factorization. All data coercion, sorting and boxing
787 # should happen here.
788 if isinstance(values, (ABCIndex, ABCSeries)):
789 return values.factorize(sort=sort, use_na_sentinel=use_na_sentinel)
790
791 values = _ensure_arraylike(values, func_name="factorize")
792 original = values
793
794 if (
795 isinstance(values, (ABCDatetimeArray, ABCTimedeltaArray))
796 and values.freq is not None
797 ):
798 # The presence of 'freq' means we can fast-path sorting and know there
799 # aren't NAs
800 codes, uniques = values.factorize(sort=sort)
801 return codes, uniques
802
803 elif not isinstance(values, np.ndarray):
804 # i.e. ExtensionArray
805 codes, uniques = values.factorize(use_na_sentinel=use_na_sentinel)
806
807 else:
808 values = np.asarray(values) # convert DTA/TDA/MultiIndex
809
810 if not use_na_sentinel and values.dtype == object:
811 # factorize can now handle differentiating various types of null values.
812 # These can only occur when the array has object dtype.
813 # However, for backwards compatibility we only use the null for the
814 # provided dtype. This may be revisited in the future, see GH#48476.
815 null_mask = isna(values)
816 if null_mask.any():
817 na_value = na_value_for_dtype(values.dtype, compat=False)
818 # Don't modify (potentially user-provided) array
819 values = np.where(null_mask, na_value, values)
820
821 codes, uniques = factorize_array(
822 values,
823 use_na_sentinel=use_na_sentinel,
824 size_hint=size_hint,
825 )
826
827 if sort and len(uniques) > 0:
828 uniques, codes = safe_sort(
829 uniques,
830 codes,
831 use_na_sentinel=use_na_sentinel,
832 assume_unique=True,
833 verify=False,
834 )
835
836 uniques = _reconstruct_data(uniques, original.dtype, original)
837
838 return codes, uniques
839
840
841def value_counts_internal(
842 values,
843 sort: bool = True,
844 ascending: bool = False,
845 normalize: bool = False,
846 bins=None,
847 dropna: bool = True,
848) -> Series:
849 from pandas import (
850 DatetimeIndex,
851 Index,
852 Series,
853 TimedeltaIndex,
854 )
855
856 index_name = getattr(values, "name", None)
857 name = "proportion" if normalize else "count"
858
859 if bins is not None:
860 from pandas.core.reshape.tile import cut
861
862 if isinstance(values, Series):
863 values = values._values
864
865 try:
866 ii = cut(values, bins, include_lowest=True)
867 except TypeError as err:
868 raise TypeError("bins argument only works with numeric data.") from err
869
870 # count, remove nulls (from the index), and but the bins
871 result = ii.value_counts(dropna=dropna)
872 result.name = name
873 result = result[result.index.notna()]
874 result.index = result.index.astype("interval")
875 result = result.sort_index()
876
877 # if we are dropna and we have NO values
878 if dropna and (result._values == 0).all():
879 result = result.iloc[0:0]
880
881 # normalizing is by len of all (regardless of dropna)
882 normalize_denominator = len(ii)
883
884 else:
885 normalize_denominator = None
886 if is_extension_array_dtype(values):
887 # handle Categorical and sparse,
888 result = Series(values, copy=False)._values.value_counts(dropna=dropna)
889 result.name = name
890 result.index.name = index_name
891
892 elif isinstance(values, ABCMultiIndex):
893 # GH49558
894 levels = list(range(values.nlevels))
895 result = (
896 Series(index=values, name=name)
897 .groupby(level=levels, dropna=dropna)
898 .size()
899 )
900 result.index.names = values.names
901
902 else:
903 values = _ensure_arraylike(values, func_name="value_counts")
904 keys, counts, _ = value_counts_arraylike(values, dropna)
905 if keys.dtype == np.float16:
906 keys = keys.astype(np.float32)
907
908 # Starting in 3.0, we no longer perform dtype inference on the
909 # Index object we construct here, xref GH#56161
910 idx = Index(keys, dtype=keys.dtype, name=index_name, copy=False)
911
912 if (
913 not sort
914 and isinstance(values, (DatetimeIndex, TimedeltaIndex))
915 and idx.equals(values)
916 and values.inferred_freq is not None
917 ):
918 # Preserve freq of original index
919 idx.freq = values.inferred_freq # type: ignore[attr-defined]
920
921 result = Series(counts, index=idx, name=name, copy=False)
922
923 if sort:
924 result = result.sort_values(ascending=ascending, kind="stable")
925
926 if normalize:
927 if normalize_denominator is not None:
928 result = result / normalize_denominator
929 else:
930 result = result / result.sum()
931
932 return result
933
934
935# Called once from SparseArray, otherwise could be private
936def value_counts_arraylike(
937 values: np.ndarray, dropna: bool, mask: npt.NDArray[np.bool_] | None = None
938) -> tuple[ArrayLike, npt.NDArray[np.int64], int]:
939 """
940 Parameters
941 ----------
942 values : np.ndarray
943 dropna : bool
944 mask : np.ndarray[bool] or None, default None
945
946 Returns
947 -------
948 uniques : np.ndarray
949 counts : np.ndarray[np.int64]
950 """
951 original = values
952 values = _ensure_data(values)
953
954 keys, counts, na_counter = htable.value_count(values, dropna, mask=mask)
955
956 if needs_i8_conversion(original.dtype):
957 # datetime, timedelta, or period
958
959 if dropna:
960 mask = keys != iNaT
961 keys, counts = keys[mask], counts[mask]
962
963 res_keys = _reconstruct_data(keys, original.dtype, original)
964 return res_keys, counts, na_counter
965
966
967def duplicated(
968 values: ArrayLike,
969 keep: Literal["first", "last", False] = "first",
970 mask: npt.NDArray[np.bool_] | None = None,
971) -> npt.NDArray[np.bool_]:
972 """
973 Return boolean ndarray denoting duplicate values.
974
975 Parameters
976 ----------
977 values : np.ndarray or ExtensionArray
978 Array over which to check for duplicate values.
979 keep : {'first', 'last', False}, default 'first'
980 - ``first`` : Mark duplicates as ``True`` except for the first
981 occurrence.
982 - ``last`` : Mark duplicates as ``True`` except for the last
983 occurrence.
984 - False : Mark all duplicates as ``True``.
985 mask : ndarray[bool], optional
986 array indicating which elements to exclude from checking
987
988 Returns
989 -------
990 duplicated : ndarray[bool]
991 """
992 values = _ensure_data(values)
993 return htable.duplicated(values, keep=keep, mask=mask)
994
995
996def mode(
997 values: ArrayLike, dropna: bool = True, mask: npt.NDArray[np.bool_] | None = None
998) -> tuple[np.ndarray, npt.NDArray[np.bool_]] | ExtensionArray:
999 """
1000 Returns the mode(s) of an array.
1001
1002 Parameters
1003 ----------
1004 values : array-like
1005 Array over which to check for duplicate values.
1006 dropna : bool, default True
1007 Don't consider counts of NaN/NaT.
1008
1009 Returns
1010 -------
1011 Union[Tuple[np.ndarray, npt.NDArray[np.bool_]], ExtensionArray]
1012 """
1013 values = _ensure_arraylike(values, func_name="mode")
1014 original = values
1015
1016 if needs_i8_conversion(values.dtype):
1017 # Got here with ndarray; dispatch to DatetimeArray/TimedeltaArray.
1018 values = ensure_wrapped_if_datetimelike(values)
1019 values = cast("ExtensionArray", values)
1020 return values._mode(dropna=dropna)
1021
1022 values = _ensure_data(values)
1023
1024 npresult, res_mask = htable.mode(values, dropna=dropna, mask=mask)
1025 if res_mask is None:
1026 res_mask = np.zeros(npresult.shape, dtype=np.bool_)
1027 else:
1028 return npresult, res_mask
1029
1030 try:
1031 npresult = safe_sort(npresult)
1032 except TypeError as err:
1033 warnings.warn(
1034 f"Unable to sort modes: {err}",
1035 stacklevel=find_stack_level(),
1036 )
1037
1038 result = _reconstruct_data(npresult, original.dtype, original)
1039 return result, res_mask
1040
1041
1042def rank(
1043 values: ArrayLike,
1044 axis: AxisInt = 0,
1045 method: str = "average",
1046 na_option: str = "keep",
1047 ascending: bool = True,
1048 pct: bool = False,
1049 mask: npt.NDArray[np.bool_] | None = None,
1050) -> npt.NDArray[np.float64]:
1051 """
1052 Rank the values along a given axis.
1053
1054 Parameters
1055 ----------
1056 values : np.ndarray or ExtensionArray
1057 Array whose values will be ranked. The number of dimensions in this
1058 array must not exceed 2.
1059 axis : int, default 0
1060 Axis over which to perform rankings.
1061 method : {'average', 'min', 'max', 'first', 'dense'}, default 'average'
1062 The method by which tiebreaks are broken during the ranking.
1063 na_option : {'keep', 'top'}, default 'keep'
1064 The method by which NaNs are placed in the ranking.
1065 - ``keep``: rank each NaN value with a NaN ranking
1066 - ``top``: replace each NaN with either +/- inf so that they
1067 there are ranked at the top
1068 ascending : bool, default True
1069 Whether or not the elements should be ranked in ascending order.
1070 pct : bool, default False
1071 Whether or not to the display the returned rankings in integer form
1072 (e.g. 1, 2, 3) or in percentile form (e.g. 0.333..., 0.666..., 1).
1073 mask : bool ndarray, optional
1074 Boolean array indicating which elements to exclude from ranking.
1075 """
1076 is_datetimelike = needs_i8_conversion(values.dtype)
1077 values = _ensure_data(values)
1078
1079 if values.ndim == 1:
1080 ranks = algos.rank_1d(
1081 values,
1082 is_datetimelike=is_datetimelike,
1083 ties_method=method,
1084 ascending=ascending,
1085 na_option=na_option,
1086 pct=pct,
1087 mask=mask,
1088 )
1089 elif values.ndim == 2:
1090 assert mask is None
1091 ranks = algos.rank_2d(
1092 values,
1093 axis=axis,
1094 is_datetimelike=is_datetimelike,
1095 ties_method=method,
1096 ascending=ascending,
1097 na_option=na_option,
1098 pct=pct,
1099 )
1100 else:
1101 raise TypeError("Array with ndim > 2 are not supported.")
1102
1103 return ranks
1104
1105
1106# ---- #
1107# take #
1108# ---- #
1109
1110
1111@set_module("pandas.api.extensions")
1112def take(
1113 arr,
1114 indices: TakeIndexer,
1115 axis: AxisInt = 0,
1116 allow_fill: bool = False,
1117 fill_value=None,
1118):
1119 """
1120 Take elements from an array.
1121
1122 Parameters
1123 ----------
1124 arr : numpy.ndarray, ExtensionArray, Index, or Series
1125 Input array.
1126 indices : sequence of int or one-dimensional np.ndarray of int
1127 Indices to be taken.
1128 axis : int, default 0
1129 The axis over which to select values.
1130 allow_fill : bool, default False
1131 How to handle negative values in `indices`.
1132
1133 * False: negative values in `indices` indicate positional indices
1134 from the right (the default). This is similar to :func:`numpy.take`.
1135
1136 * True: negative values in `indices` indicate
1137 missing values. These values are set to `fill_value`. Any other
1138 negative values raise a ``ValueError``.
1139
1140 fill_value : any, optional
1141 Fill value to use for NA-indices when `allow_fill` is True.
1142 This may be ``None``, in which case the default NA value for
1143 the type (``self.dtype.na_value``) is used.
1144
1145 For multi-dimensional `arr`, each *element* is filled with
1146 `fill_value`.
1147
1148 Returns
1149 -------
1150 ndarray or ExtensionArray
1151 Same type as the input.
1152
1153 Raises
1154 ------
1155 IndexError
1156 When `indices` is out of bounds for the array.
1157 ValueError
1158 When the indexer contains negative values other than ``-1``
1159 and `allow_fill` is True.
1160
1161 Notes
1162 -----
1163 When `allow_fill` is False, `indices` may be whatever dimensionality
1164 is accepted by NumPy for `arr`.
1165
1166 When `allow_fill` is True, `indices` should be 1-D.
1167
1168 See Also
1169 --------
1170 numpy.take : Take elements from an array along an axis.
1171
1172 Examples
1173 --------
1174 >>> import pandas as pd
1175
1176 With the default ``allow_fill=False``, negative numbers indicate
1177 positional indices from the right.
1178
1179 >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1])
1180 array([10, 10, 30])
1181
1182 Setting ``allow_fill=True`` will place `fill_value` in those positions.
1183
1184 >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True)
1185 array([10., 10., nan])
1186
1187 >>> pd.api.extensions.take(
1188 ... np.array([10, 20, 30]), [0, 0, -1], allow_fill=True, fill_value=-10
1189 ... )
1190 array([ 10, 10, -10])
1191 """
1192 if not isinstance(
1193 arr,
1194 (np.ndarray, ABCExtensionArray, ABCIndex, ABCSeries, ABCNumpyExtensionArray),
1195 ):
1196 # GH#52981
1197 raise TypeError(
1198 "pd.api.extensions.take requires a numpy.ndarray, ExtensionArray, "
1199 f"Index, Series, or NumpyExtensionArray got {type(arr).__name__}."
1200 )
1201
1202 indices = ensure_platform_int(indices)
1203
1204 if allow_fill:
1205 # Pandas style, -1 means NA
1206 validate_indices(indices, arr.shape[axis])
1207 # error: Argument 1 to "take_nd" has incompatible type
1208 # "ndarray[Any, Any] | ExtensionArray | Index | Series"; expected
1209 # "ndarray[Any, Any]"
1210 result = take_nd(
1211 arr, # type: ignore[arg-type]
1212 indices,
1213 axis=axis,
1214 allow_fill=True,
1215 fill_value=fill_value,
1216 )
1217 else:
1218 # NumPy style
1219 # error: Unexpected keyword argument "axis" for "take" of "ExtensionArray"
1220 result = arr.take(indices, axis=axis) # type: ignore[call-arg,assignment]
1221 return result
1222
1223
1224# ------------ #
1225# searchsorted #
1226# ------------ #
1227
1228
1229def searchsorted(
1230 arr: ArrayLike,
1231 value: NumpyValueArrayLike | ExtensionArray,
1232 side: Literal["left", "right"] = "left",
1233 sorter: NumpySorter | None = None,
1234) -> npt.NDArray[np.intp] | np.intp:
1235 """
1236 Find indices where elements should be inserted to maintain order.
1237
1238 Find the indices into a sorted array `arr` (a) such that, if the
1239 corresponding elements in `value` were inserted before the indices,
1240 the order of `arr` would be preserved.
1241
1242 Assuming that `arr` is sorted:
1243
1244 ====== ================================
1245 `side` returned index `i` satisfies
1246 ====== ================================
1247 left ``arr[i-1] < value <= self[i]``
1248 right ``arr[i-1] <= value < self[i]``
1249 ====== ================================
1250
1251 Parameters
1252 ----------
1253 arr: np.ndarray, ExtensionArray, Series
1254 Input array. If `sorter` is None, then it must be sorted in
1255 ascending order, otherwise `sorter` must be an array of indices
1256 that sort it.
1257 value : array-like or scalar
1258 Values to insert into `arr`.
1259 side : {'left', 'right'}, optional
1260 If 'left', the index of the first suitable location found is given.
1261 If 'right', return the last such index. If there is no suitable
1262 index, return either 0 or N (where N is the length of `self`).
1263 sorter : 1-D array-like, optional
1264 Optional array of integer indices that sort array a into ascending
1265 order. They are typically the result of argsort.
1266
1267 Returns
1268 -------
1269 array of ints or int
1270 If value is array-like, array of insertion points.
1271 If value is scalar, a single integer.
1272
1273 See Also
1274 --------
1275 numpy.searchsorted : Similar method from NumPy.
1276 """
1277 if sorter is not None:
1278 sorter = ensure_platform_int(sorter)
1279
1280 if (
1281 isinstance(arr, np.ndarray)
1282 and arr.dtype.kind in "iu"
1283 and (is_integer(value) or is_integer_dtype(value))
1284 ):
1285 # if `arr` and `value` have different dtypes, `arr` would be
1286 # recast by numpy, causing a slow search.
1287 # Before searching below, we therefore try to give `value` the
1288 # same dtype as `arr`, while guarding against integer overflows.
1289 iinfo = np.iinfo(arr.dtype.type)
1290 value_arr = np.array([value]) if is_integer(value) else np.array(value)
1291 if (value_arr >= iinfo.min).all() and (value_arr <= iinfo.max).all():
1292 # value within bounds, so no overflow, so can convert value dtype
1293 # to dtype of arr
1294 dtype = arr.dtype
1295 else:
1296 dtype = value_arr.dtype
1297
1298 if is_integer(value):
1299 # We know that value is int
1300 value = cast(int, dtype.type(value))
1301 else:
1302 value = pd_array(cast(ArrayLike, value), dtype=dtype)
1303 else:
1304 # E.g. if `arr` is an array with dtype='datetime64[ns]'
1305 # and `value` is a pd.Timestamp, we may need to convert value
1306 arr = ensure_wrapped_if_datetimelike(arr)
1307
1308 # Argument 1 to "searchsorted" of "ndarray" has incompatible type
1309 # "Union[NumpyValueArrayLike, ExtensionArray]"; expected "NumpyValueArrayLike"
1310 return arr.searchsorted(value, side=side, sorter=sorter) # type: ignore[arg-type]
1311
1312
1313# ---- #
1314# diff #
1315# ---- #
1316
1317_diff_special = {"float64", "float32", "int64", "int32", "int16", "int8"}
1318
1319
1320def diff(arr, n: int | float | np.integer | np.floating, axis: AxisInt = 0):
1321 """
1322 difference of n between self,
1323 analogous to s-s.shift(n)
1324
1325 Parameters
1326 ----------
1327 arr : ndarray or ExtensionArray
1328 n : int
1329 number of periods
1330 axis : {0, 1}
1331 axis to shift on
1332 stacklevel : int, default 3
1333 The stacklevel for the lost dtype warning.
1334
1335 Returns
1336 -------
1337 shifted
1338 """
1339
1340 # added a check on the integer value of period
1341 # see https://github.com/pandas-dev/pandas/issues/56607
1342 if not lib.is_integer(n):
1343 if not (is_float(n) and n.is_integer()):
1344 raise ValueError("periods must be an integer")
1345 n = int(n)
1346 na = np.nan
1347 dtype = arr.dtype
1348
1349 is_bool = is_bool_dtype(dtype)
1350 if is_bool:
1351 op = operator.xor
1352 else:
1353 op = operator.sub
1354
1355 if isinstance(dtype, NumpyEADtype):
1356 # NumpyExtensionArray cannot necessarily hold shifted versions of itself.
1357 arr = arr.to_numpy()
1358 dtype = arr.dtype
1359
1360 if not isinstance(arr, np.ndarray):
1361 # i.e ExtensionArray
1362 if hasattr(arr, f"__{op.__name__}__"):
1363 if axis != 0:
1364 raise ValueError(f"cannot diff {type(arr).__name__} on axis={axis}")
1365 return op(arr, arr.shift(n))
1366 else:
1367 raise TypeError(
1368 f"{type(arr).__name__} has no 'diff' method. "
1369 "Convert to a suitable dtype prior to calling 'diff'."
1370 )
1371
1372 is_timedelta = False
1373 if arr.dtype.kind in "mM":
1374 dtype = np.int64
1375 arr = arr.view("i8")
1376 na = iNaT
1377 is_timedelta = True
1378
1379 elif is_bool:
1380 # We have to cast in order to be able to hold np.nan
1381 dtype = np.object_
1382
1383 elif dtype.kind in "iu":
1384 # We have to cast in order to be able to hold np.nan
1385
1386 # int8, int16 are incompatible with float64,
1387 # see https://github.com/cython/cython/issues/2646
1388 if arr.dtype.name in ["int8", "int16"]:
1389 dtype = np.float32
1390 else:
1391 dtype = np.float64
1392
1393 orig_ndim = arr.ndim
1394 if orig_ndim == 1:
1395 # reshape so we can always use algos.diff_2d
1396 arr = arr.reshape(-1, 1)
1397 # TODO: require axis == 0
1398
1399 dtype = np.dtype(dtype)
1400 out_arr = np.empty(arr.shape, dtype=dtype)
1401
1402 na_indexer = [slice(None)] * 2
1403 na_indexer[axis] = slice(None, n) if n >= 0 else slice(n, None)
1404 out_arr[tuple(na_indexer)] = na
1405
1406 if arr.dtype.name in _diff_special:
1407 # TODO: can diff_2d dtype specialization troubles be fixed by defining
1408 # out_arr inside diff_2d?
1409 algos.diff_2d(arr, out_arr, int(n), axis, datetimelike=is_timedelta)
1410 else:
1411 # To keep mypy happy, _res_indexer is a list while res_indexer is
1412 # a tuple, ditto for lag_indexer.
1413 _res_indexer = [slice(None)] * 2
1414 _res_indexer[axis] = slice(n, None) if n >= 0 else slice(None, n)
1415 res_indexer = tuple(_res_indexer)
1416
1417 _lag_indexer = [slice(None)] * 2
1418 _lag_indexer[axis] = slice(None, -n) if n > 0 else slice(-n, None)
1419 lag_indexer = tuple(_lag_indexer)
1420
1421 out_arr[res_indexer] = op(arr[res_indexer], arr[lag_indexer])
1422
1423 if is_timedelta:
1424 out_arr = out_arr.view("timedelta64[ns]")
1425
1426 if orig_ndim == 1:
1427 out_arr = out_arr[:, 0]
1428 return out_arr
1429
1430
1431# --------------------------------------------------------------------
1432# Helper functions
1433
1434
1435# Note: safe_sort is in algorithms.py instead of sorting.py because it is
1436# low-dependency, is used in this module, and used private methods from
1437# this module.
1438def safe_sort(
1439 values: Index | ArrayLike,
1440 codes: npt.NDArray[np.intp] | None = None,
1441 use_na_sentinel: bool = True,
1442 assume_unique: bool = False,
1443 verify: bool = True,
1444) -> AnyArrayLike | tuple[AnyArrayLike, np.ndarray]:
1445 """
1446 Sort ``values`` and reorder corresponding ``codes``.
1447
1448 ``values`` should be unique if ``codes`` is not None.
1449 Safe for use with mixed types (int, str), orders ints before strs.
1450
1451 Parameters
1452 ----------
1453 values : list-like
1454 Sequence; must be unique if ``codes`` is not None.
1455 codes : np.ndarray[intp] or None, default None
1456 Indices to ``values``. All out of bound indices are treated as
1457 "not found" and will be masked with ``-1``.
1458 use_na_sentinel : bool, default True
1459 If True, the sentinel -1 will be used for NaN values. If False,
1460 NaN values will be encoded as non-negative integers and will not drop the
1461 NaN from the uniques of the values.
1462 assume_unique : bool, default False
1463 When True, ``values`` are assumed to be unique, which can speed up
1464 the calculation. Ignored when ``codes`` is None.
1465 verify : bool, default True
1466 Check if codes are out of bound for the values and put out of bound
1467 codes equal to ``-1``. If ``verify=False``, it is assumed there
1468 are no out of bound codes. Ignored when ``codes`` is None.
1469
1470 Returns
1471 -------
1472 ordered : AnyArrayLike
1473 Sorted ``values``
1474 new_codes : ndarray
1475 Reordered ``codes``; returned when ``codes`` is not None.
1476
1477 Raises
1478 ------
1479 TypeError
1480 * If ``values`` is not list-like or if ``codes`` is neither None
1481 nor list-like
1482 * If ``values`` cannot be sorted
1483 ValueError
1484 * If ``codes`` is not None and ``values`` contain duplicates.
1485 """
1486 if not isinstance(values, (np.ndarray, ABCExtensionArray, ABCIndex)):
1487 raise TypeError(
1488 "Only np.ndarray, ExtensionArray, and Index objects are allowed to "
1489 "be passed to safe_sort as values"
1490 )
1491
1492 sorter = None
1493 ordered: AnyArrayLike
1494
1495 if (
1496 not isinstance(values.dtype, ExtensionDtype)
1497 and lib.infer_dtype(values, skipna=False) == "mixed-integer"
1498 ):
1499 ordered = _sort_mixed(values)
1500 else:
1501 try:
1502 sorter = values.argsort()
1503 ordered = values.take(sorter)
1504 except (TypeError, decimal.InvalidOperation):
1505 # Previous sorters failed or were not applicable, try `_sort_mixed`
1506 # which would work, but which fails for special case of 1d arrays
1507 # with tuples.
1508 if values.size and isinstance(values[0], tuple):
1509 # error: Argument 1 to "_sort_tuples" has incompatible type
1510 # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected
1511 # "ndarray[Any, Any]"
1512 ordered = _sort_tuples(values) # type: ignore[arg-type]
1513 else:
1514 ordered = _sort_mixed(values)
1515
1516 # codes:
1517
1518 if codes is None:
1519 return ordered
1520
1521 if not is_list_like(codes):
1522 raise TypeError(
1523 "Only list-like objects or None are allowed to "
1524 "be passed to safe_sort as codes"
1525 )
1526 codes = ensure_platform_int(np.asarray(codes))
1527
1528 if not assume_unique and not len(unique(values)) == len(values):
1529 raise ValueError("values should be unique if codes is not None")
1530
1531 if sorter is None:
1532 # mixed types
1533 # error: Argument 1 to "_get_hashtable_algo" has incompatible type
1534 # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected
1535 # "ndarray[Any, Any]"
1536 hash_klass, values = _get_hashtable_algo(values) # type: ignore[arg-type]
1537 t = hash_klass(len(values))
1538 t.map_locations(values)
1539 # error: Argument 1 to "lookup" of "HashTable" has incompatible type
1540 # "ExtensionArray | ndarray[Any, Any] | Index | Series"; expected "ndarray"
1541 sorter = ensure_platform_int(t.lookup(ordered)) # type: ignore[arg-type]
1542
1543 if use_na_sentinel:
1544 # take_nd is faster, but only works for na_sentinels of -1
1545 order2 = sorter.argsort()
1546 if verify:
1547 mask = (codes < -len(values)) | (codes >= len(values))
1548 codes[mask] = -1
1549 new_codes = take_nd(order2, codes, fill_value=-1)
1550 else:
1551 reverse_indexer = np.empty(len(sorter), dtype=int)
1552 reverse_indexer.put(sorter, np.arange(len(sorter)))
1553 # Out of bound indices will be masked with `-1` next, so we
1554 # may deal with them here without performance loss using `mode='wrap'`
1555 new_codes = reverse_indexer.take(codes, mode="wrap")
1556
1557 return ordered, ensure_platform_int(new_codes)
1558
1559
1560def _sort_mixed(values) -> AnyArrayLike:
1561 """order ints before strings before nulls in 1d arrays"""
1562 str_pos = np.array([isinstance(x, str) for x in values], dtype=bool)
1563 null_pos = np.array([isna(x) for x in values], dtype=bool)
1564 num_pos = ~str_pos & ~null_pos
1565 str_argsort = np.argsort(values[str_pos])
1566 num_argsort = np.argsort(values[num_pos])
1567 # convert boolean arrays to positional indices, then order by underlying values
1568 str_locs = str_pos.nonzero()[0].take(str_argsort)
1569 num_locs = num_pos.nonzero()[0].take(num_argsort)
1570 null_locs = null_pos.nonzero()[0]
1571 locs = np.concatenate([num_locs, str_locs, null_locs])
1572 return values.take(locs)
1573
1574
1575def _sort_tuples(values: np.ndarray) -> np.ndarray:
1576 """
1577 Convert array of tuples (1d) to array of arrays (2d).
1578 We need to keep the columns separately as they contain different types and
1579 nans (can't use `np.sort` as it may fail when str and nan are mixed in a
1580 column as types cannot be compared).
1581 """
1582 from pandas.core.internals.construction import to_arrays
1583 from pandas.core.sorting import lexsort_indexer
1584
1585 arrays, _ = to_arrays(values, None)
1586 indexer = lexsort_indexer(arrays, orders=True)
1587 return values[indexer]
1588
1589
1590def union_with_duplicates(
1591 lvals: ArrayLike | Index, rvals: ArrayLike | Index
1592) -> ArrayLike | Index:
1593 """
1594 Extracts the union from lvals and rvals with respect to duplicates and nans in
1595 both arrays.
1596
1597 Parameters
1598 ----------
1599 lvals: np.ndarray or ExtensionArray
1600 left values which is ordered in front.
1601 rvals: np.ndarray or ExtensionArray
1602 right values ordered after lvals.
1603
1604 Returns
1605 -------
1606 np.ndarray or ExtensionArray
1607 Containing the unsorted union of both arrays.
1608
1609 Notes
1610 -----
1611 Caller is responsible for ensuring lvals.dtype == rvals.dtype.
1612 """
1613 from pandas import Series
1614
1615 l_count = value_counts_internal(lvals, dropna=False)
1616 r_count = value_counts_internal(rvals, dropna=False)
1617 l_count, r_count = l_count.align(r_count, fill_value=0)
1618 final_count = np.maximum(l_count.values, r_count.values)
1619 final_count = Series(final_count, index=l_count.index, dtype="int", copy=False)
1620 if isinstance(lvals, ABCMultiIndex) and isinstance(rvals, ABCMultiIndex):
1621 unique_vals = lvals.append(rvals).unique()
1622 else:
1623 if isinstance(lvals, ABCIndex):
1624 lvals = lvals._values
1625 if isinstance(rvals, ABCIndex):
1626 rvals = rvals._values
1627 # error: List item 0 has incompatible type "Union[ExtensionArray,
1628 # ndarray[Any, Any], Index]"; expected "Union[ExtensionArray,
1629 # ndarray[Any, Any]]"
1630 combined = concat_compat([lvals, rvals]) # type: ignore[list-item]
1631 unique_vals = unique(combined)
1632 unique_vals = ensure_wrapped_if_datetimelike(unique_vals)
1633 repeats = final_count.reindex(unique_vals).values
1634 return np.repeat(unique_vals, repeats)
1635
1636
1637def map_array(
1638 arr: ArrayLike,
1639 mapper,
1640 na_action: Literal["ignore"] | None = None,
1641) -> np.ndarray | ExtensionArray | Index:
1642 """
1643 Map values using an input mapping or function.
1644
1645 Parameters
1646 ----------
1647 mapper : function, dict, or Series
1648 Mapping correspondence.
1649 na_action : {None, 'ignore'}, default None
1650 If 'ignore', propagate NA values, without passing them to the
1651 mapping correspondence.
1652
1653 Returns
1654 -------
1655 Union[ndarray, Index, ExtensionArray]
1656 The output of the mapping function applied to the array.
1657 If the function returns a tuple with more than one element
1658 a MultiIndex will be returned.
1659 """
1660 from pandas import Index
1661
1662 if na_action not in (None, "ignore"):
1663 msg = f"na_action must either be 'ignore' or None, {na_action} was passed"
1664 raise ValueError(msg)
1665
1666 # we can fastpath dict/Series to an efficient map
1667 # as we know that we are not going to have to yield
1668 # python types
1669 if is_dict_like(mapper):
1670 if isinstance(mapper, dict) and hasattr(mapper, "__missing__"):
1671 # If a dictionary subclass defines a default value method,
1672 # convert mapper to a lookup function (GH #15999).
1673 dict_with_default = mapper
1674 mapper = lambda x: dict_with_default[
1675 np.nan if isinstance(x, float) and np.isnan(x) else x
1676 ]
1677 else:
1678 # Dictionary does not have a default. Thus it's safe to
1679 # convert to a Series for efficiency.
1680 # we specify the keys here to handle the
1681 # possibility that they are tuples
1682
1683 # The return value of mapping with an empty mapper is
1684 # expected to be pd.Series(np.nan, ...). As np.nan is
1685 # of dtype float64 the return value of this method should
1686 # be float64 as well
1687 from pandas import Series
1688
1689 if len(mapper) == 0:
1690 mapper = Series(mapper, dtype=np.float64)
1691 elif isinstance(mapper, dict):
1692 mapper = Series(
1693 mapper.values(), index=Index(mapper.keys(), tupleize_cols=False)
1694 )
1695 else:
1696 mapper = Series(mapper)
1697
1698 if isinstance(mapper, ABCSeries):
1699 if na_action == "ignore":
1700 mapper = mapper[mapper.index.notna()]
1701
1702 # Since values were input this means we came from either
1703 # a dict or a series and mapper should be an index
1704 indexer = mapper.index.get_indexer(arr)
1705 new_values = take_nd(mapper._values, indexer)
1706
1707 return new_values
1708
1709 if not len(arr):
1710 return arr.copy()
1711
1712 # we must convert to python types
1713 values = arr.astype(object, copy=False)
1714 if na_action is None:
1715 return lib.map_infer(values, mapper)
1716 else:
1717 return lib.map_infer_mask(values, mapper, mask=isna(values).view(np.uint8))