""" Generic data algorithms. This module is experimental at the moment and not intended for public consumption """ from __future__ import annotations import decimal import operator from typing import ( TYPE_CHECKING, Literal, TypeVar, cast, overload, ) import warnings import numpy as np from pandas._libs import ( algos, hashtable as htable, iNaT, lib, ) from pandas._libs.missing import NA from pandas.util._decorators import set_module from pandas.util._exceptions import find_stack_level from pandas.core.dtypes.cast import ( construct_1d_object_array_from_listlike, np_find_common_type, ) from pandas.core.dtypes.common import ( ensure_float64, ensure_object, ensure_platform_int, is_bool_dtype, is_complex_dtype, is_dict_like, is_dtype_equal, is_extension_array_dtype, is_float, is_float_dtype, is_integer, is_integer_dtype, is_list_like, is_object_dtype, is_signed_integer_dtype, needs_i8_conversion, ) from pandas.core.dtypes.concat import concat_compat from pandas.core.dtypes.dtypes import ( BaseMaskedDtype, CategoricalDtype, ExtensionDtype, NumpyEADtype, ) from pandas.core.dtypes.generic import ( ABCExtensionArray, ABCIndex, ABCMultiIndex, ABCNumpyExtensionArray, ABCSeries, ) from pandas.core.dtypes.missing import ( isna, na_value_for_dtype, ) from pandas.core.array_algos.take import take_nd from pandas.core.construction import ( array as pd_array, ensure_wrapped_if_datetimelike, extract_array, ) from pandas.core.indexers import validate_indices if TYPE_CHECKING: from pandas._typing import ( AnyArrayLike, ArrayLike, ArrayLikeT, AxisInt, DtypeObj, ListLike, NumpySorter, NumpyValueArrayLike, RankMethod, RankNaOption, TakeIndexer, npt, ) from pandas import ( Categorical, Index, Series, ) from pandas.core.arrays import ( BaseMaskedArray, ExtensionArray, ) T = TypeVar("T", bound=Index | Categorical | ExtensionArray) # --------------- # # dtype access # # --------------- # def _ensure_data(values: ArrayLike) -> np.ndarray: """ routine to ensure that our data is of the correct input dtype for lower-level routines This will coerce: - ints -> int64 - uint -> uint64 - bool -> uint8 - datetimelike -> i8 - datetime64tz -> i8 (in local tz) - categorical -> codes Parameters ---------- values : np.ndarray or ExtensionArray Returns ------- np.ndarray """ if not isinstance(values, ABCMultiIndex): # extract_array would raise values = extract_array(values, extract_numpy=True) if is_object_dtype(values.dtype): return ensure_object(np.asarray(values)) elif isinstance(values.dtype, BaseMaskedDtype): # i.e. BooleanArray, FloatingArray, IntegerArray values = cast("BaseMaskedArray", values) if not values._hasna: # No pd.NAs -> We can avoid an object-dtype cast (and copy) GH#41816 # recurse to avoid re-implementing logic for eg bool->uint8 return _ensure_data(values._data) return np.asarray(values) elif isinstance(values.dtype, CategoricalDtype): # NB: cases that go through here should NOT be using _reconstruct_data # on the back-end. values = cast("Categorical", values) return values.codes elif is_bool_dtype(values.dtype): if isinstance(values, np.ndarray): # i.e. actually dtype == np.dtype("bool") return np.asarray(values).view("uint8") else: # e.g. Sparse[bool, False], see test_mode_sparse_bool return np.asarray(values).astype("uint8", copy=False) elif is_integer_dtype(values.dtype): return np.asarray(values) elif is_float_dtype(values.dtype): # itemsize comes off the values, not the dtype: an ExtensionDtype # (e.g. Sparse[float64]) may not have one. # Note: checking `values.dtype == "float128"` raises on Windows and 32bit float_values = np.asarray(values) if float_values.dtype.itemsize in [2, 12, 16]: # we dont (yet) have float128 hashtable support return ensure_float64(float_values) return float_values elif is_complex_dtype(values.dtype): # NumpyExtensionArray needs to be unwrapped to the underlying ndarray # so the hashtable fused-type dispatch works (GH#54761) return np.asarray(values) # datetimelike elif needs_i8_conversion(values.dtype): npvalues = values.view("i8") npvalues = cast("np.ndarray", npvalues) return npvalues # we have failed, return object values = np.asarray(values, dtype=object) return ensure_object(values) def _reconstruct_data( values: ArrayLikeT, dtype: DtypeObj, original: AnyArrayLike ) -> ArrayLikeT: """ reverse of _ensure_data Parameters ---------- values : np.ndarray or ExtensionArray dtype : np.dtype or ExtensionDtype original : AnyArrayLike Returns ------- ExtensionArray or np.ndarray """ if isinstance(values, ABCExtensionArray) and values.dtype == dtype: # Catch DatetimeArray/TimedeltaArray return values if not isinstance(dtype, np.dtype): # i.e. ExtensionDtype; note we have ruled out above the possibility # that values.dtype == dtype cls = dtype.construct_array_type() # error: Incompatible return value type # (got "ExtensionArray", # expected "ndarray[tuple[Any, ...], dtype[Any]]") return cls._from_sequence(values, dtype=dtype) # type: ignore[return-value] # error: Incompatible return value type # (got "ndarray[tuple[Any, ...], dtype[Any]]", # expected "ExtensionArray") return values.astype(dtype, copy=False) # type: ignore[return-value] def _ensure_arraylike(values, func_name: str) -> ArrayLike: """ ensure that we are arraylike if not already """ if not isinstance( values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray, ABCNumpyExtensionArray), ): # GH#52986 if func_name != "isin-targets": # Make an exception for the comps argument in isin. raise TypeError( f"{func_name} requires a Series, Index, " f"ExtensionArray, np.ndarray or NumpyExtensionArray " f"got {type(values).__name__}." ) inferred = lib.infer_dtype(values, skipna=False) if inferred in ["mixed", "string", "mixed-integer"]: # "mixed-integer" to ensure we do not cast ["ss", 42] to str GH#22160 if isinstance(values, tuple): values = list(values) values = construct_1d_object_array_from_listlike(values) else: values = np.asarray(values) return values _hashtables = { "complex128": htable.Complex128HashTable, "complex64": htable.Complex64HashTable, "float64": htable.Float64HashTable, "float32": htable.Float32HashTable, "uint64": htable.UInt64HashTable, "uint32": htable.UInt32HashTable, "uint16": htable.UInt16HashTable, "uint8": htable.UInt8HashTable, "int64": htable.Int64HashTable, "int32": htable.Int32HashTable, "int16": htable.Int16HashTable, "int8": htable.Int8HashTable, "string": htable.StringHashTable, "object": htable.PyObjectHashTable, } def _get_hashtable_algo( values: np.ndarray, ) -> tuple[type[htable.HashTable], np.ndarray]: """ Parameters ---------- values : np.ndarray Returns ------- htable : HashTable subclass values : ndarray """ values = _ensure_data(values) ndtype = _check_object_for_strings(values) hashtable = _hashtables[ndtype] return hashtable, values def _check_object_for_strings(values: np.ndarray) -> str: """ Check if we can use string hashtable instead of object hashtable. Parameters ---------- values : ndarray Returns ------- str """ ndtype = values.dtype.name if ndtype == "object": # it's cheaper to use a String Hash Table than Object; we infer # including nulls because that is the only difference between # StringHashTable and ObjectHashtable if lib.is_string_array(values, skipna=False): ndtype = "string" return ndtype # --------------- # # top-level algos # # --------------- # @overload def unique(values: T) -> T: ... @overload def unique(values: np.ndarray | Series) -> np.ndarray: ... @set_module("pandas") def unique(values): """ Return unique values based on a hash table. Uniques are returned in order of appearance. This does NOT sort. Significantly faster than numpy.unique for long enough sequences. Includes NA values. Parameters ---------- values : 1d array-like The input array-like object containing values from which to extract unique values. Returns ------- Index, ExtensionArray or numpy.ndarray The return type depends on the type of the input: * :class:`Index` : when the input is an :class:`Index`, the result is an :class:`Index` (or a subclass such as :class:`DatetimeIndex`). * :class:`ExtensionArray` : when the input has an :class:`ExtensionDtype` (for example :class:`Categorical`, tz-aware ``datetime64``, :class:`IntervalArray`, a masked integer/boolean/float array, an Arrow-backed array, or :class:`NumpyExtensionArray`), the result is an :class:`ExtensionArray` of the same dtype. * :class:`numpy.ndarray` : for any other input (a NumPy-dtype :class:`Series`, a :class:`numpy.ndarray`, or any other 1-D array-like), the result is a :class:`numpy.ndarray`. See Also -------- Index.unique : Return unique values from an Index. Series.unique : Return unique values of Series object. Notes ----- When working with object-dtype arrays, boolean and integer values may not be distinguished since ``True == 1`` and ``False == 0`` in Python. >>> pd.unique(np.array([True, 1, False, 0], dtype=object)) array([True, False], dtype=object) Examples -------- >>> pd.unique(pd.Series([2, 1, 3, 3])) array([2, 1, 3]) >>> pd.unique(pd.Series([2] + [1] * 5)) array([2, 1]) >>> pd.unique(pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")])) array(['2016-01-01T00:00:00.000000'], dtype='datetime64[us]') >>> pd.unique( ... pd.Series( ... [ ... pd.Timestamp("20160101", tz="US/Eastern"), ... pd.Timestamp("20160101", tz="US/Eastern"), ... ], ... dtype="M8[ns, US/Eastern]", ... ) ... ) ['2016-01-01 00:00:00-05:00'] Length: 1, dtype: datetime64[ns, US/Eastern] >>> pd.unique( ... pd.Index( ... [ ... pd.Timestamp("20160101", tz="US/Eastern"), ... pd.Timestamp("20160101", tz="US/Eastern"), ... ], ... dtype="M8[ns, US/Eastern]", ... ) ... ) DatetimeIndex(['2016-01-01 00:00:00-05:00'], dtype='datetime64[ns, US/Eastern]', freq=None) >>> pd.unique(np.array(list("baabc"), dtype="O")) array(['b', 'a', 'c'], dtype=object) An unordered Categorical will return categories in the order of appearance. >>> pd.unique(pd.Series(pd.Categorical(list("baabc")))) ['b', 'a', 'c'] Categories (3, str): ['a', 'b', 'c'] >>> pd.unique(pd.Series(pd.Categorical(list("baabc"), categories=list("abc")))) ['b', 'a', 'c'] Categories (3, str): ['a', 'b', 'c'] An ordered Categorical preserves the category ordering. >>> pd.unique( ... pd.Series( ... pd.Categorical(list("baabc"), categories=list("abc"), ordered=True) ... ) ... ) ['b', 'a', 'c'] Categories (3, str): ['a' < 'b' < 'c'] An array of tuples >>> pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values) array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object) A NumpyExtensionArray of complex >>> pd.unique(pd.array([1 + 1j, 2, 3])) [(1+1j), (2+0j), (3+0j)] Length: 3, dtype: complex128 """ return unique_with_mask(values) def nunique_ints(values: ArrayLike) -> int: """ Return the number of unique values for integer array-likes. Significantly faster than pandas.unique for long enough sequences. No checks are done to ensure input is integral. Parameters ---------- values : 1d array-like Returns ------- int : The number of unique values in ``values`` """ if len(values) == 0: return 0 values = _ensure_data(values) # bincount requires intp result = (np.bincount(values.ravel().astype("intp")) != 0).sum() return result def unique_with_mask(values, mask: npt.NDArray[np.bool_] | None = None): """See algorithms.unique for docs. Takes a mask for masked arrays.""" values = _ensure_arraylike(values, func_name="unique") if isinstance(values.dtype, ExtensionDtype): # Dispatch to extension dtype's unique. return values.unique() if isinstance(values, ABCIndex): # Dispatch to Index's unique. return values.unique() original = values hashtable, values = _get_hashtable_algo(values) table = hashtable(len(values)) if mask is None: uniques = table.unique(values) uniques = _reconstruct_data(uniques, original.dtype, original) return uniques else: uniques, mask = table.unique(values, mask=mask) uniques = _reconstruct_data(uniques, original.dtype, original) assert mask is not None # for mypy return uniques, mask.astype("bool") unique1d = unique _MINIMUM_COMP_ARR_LEN = 1_000_000 # Integers with magnitude <= 2**53 are exactly representable as float64; above # this a 64-bit int cast to float64 may collide with a distinct value. _FLOAT64_INT_EXACT_MAX = 2**53 def _may_lose_precision_as_float64(arr: np.ndarray) -> bool: """ Whether casting ``arr`` to float64 could change an integer's identity. Only 64-bit integers can carry a magnitude that float64 cannot represent exactly; narrower widths always round-trip. """ if arr.dtype.kind not in "iu" or arr.dtype.itemsize < 8 or arr.size == 0: return False return ( int(arr.min()) < -_FLOAT64_INT_EXACT_MAX or int(arr.max()) > _FLOAT64_INT_EXACT_MAX ) def isin(comps: ListLike, values: ListLike) -> npt.NDArray[np.bool_]: """ Compute the isin boolean array. Parameters ---------- comps : list-like values : list-like Returns ------- ndarray[bool] Same length as `comps`. """ if not is_list_like(comps): raise TypeError( "only list-like objects are allowed to be passed " f"to isin(), you passed a `{type(comps).__name__}`" ) if not is_list_like(values): raise TypeError( "only list-like objects are allowed to be passed " f"to isin(), you passed a `{type(values).__name__}`" ) if isinstance(values, (set, frozenset)) and len(values) > 0: # GH#25507: for a set of values, membership can be tested directly # via the set, avoiding an O(len(values)) materialization that # otherwise dominates when comps is much smaller than values. # Restrict to integer/bool comps (i.e. dtypes that cannot contain # NaN), since Python set membership would mis-handle the case where # both sides contain NaN values that are not identical. if isinstance(comps, (ABCSeries, ABCIndex)): comps_arr = comps._values else: comps_arr = comps if ( isinstance(comps_arr, np.ndarray) and comps_arr.ndim == 1 and comps_arr.dtype.kind in "iub" # Only worth it when comps is smaller than values: this path does # an O(len(comps)) Python-level scan, which is much slower per # element than the vectorized htable path below. It wins only by # skipping the O(len(values)) materialization of the set, so it # must not fire for the common large-comps/small-set case. and comps_arr.size < len(values) # A 64-bit comp with magnitude >= 2**53 can match a float in the # set lossily under the materialized path below (via a float64 # cast) but never under exact Python equality, so the result would # depend on which path is taken. Only use this path when comps # values are strictly inside float64's exact-int range, where the # two agree. and ( comps_arr.dtype.itemsize < 8 or comps_arr.size == 0 or ( -_FLOAT64_INT_EXACT_MAX < int(comps_arr.min()) and int(comps_arr.max()) < _FLOAT64_INT_EXACT_MAX ) ) ): return np.fromiter( (item in values for item in comps_arr.tolist()), dtype=bool, count=comps_arr.size, ) if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)): orig_values = list(values) values = _ensure_arraylike(orig_values, func_name="isin-targets") if ( len(values) > 0 and values.dtype.kind in "iufcb" and not is_signed_integer_dtype(comps) and not is_dtype_equal(values, comps) ): # GH#46485: casting a 64-bit integer to float64 loses precision # for magnitudes > 2**53, so a numeric comparison can give wrong # results. Recast to object only when such a lossy cast is actually # at risk; otherwise keep the numeric ndarray so the downstream # htable dispatch can use the fast numeric path. # TODO: Share with _find_common_type_compat comps_dtype = getattr(comps, "dtype", None) # values came from _ensure_arraylike's numeric branch, so its # dtype is an np.dtype. values_dtype = cast("np.dtype", values.dtype) if not isinstance(comps_dtype, np.dtype): needs_object = True else: common = np_find_common_type(values_dtype, comps_dtype) if common.kind not in "iufcb": needs_object = True elif common.kind not in "fc": # integer/bool common type -> no lossy float cast # (complex128 components are float64, so "c" is as lossy # as "f") needs_object = False elif comps_dtype.kind in "iu" and comps_dtype.itemsize == 8: # comps is the (potentially large) caller and a 64-bit # integer that would be cast down to float64; too costly to # inspect element-wise, so fall back to object to be safe. needs_object = True elif values_dtype.kind in "iu" and values_dtype.itemsize == 8: # values is the (small) targets list -> cheap to check # exactly. Only object-cast if a target falls outside # float64's exact integer range; otherwise the numeric # cast is lossless and the fast path stays correct. # values is an ndarray here (it came from the numeric # branch of _ensure_arraylike above). values_arr = cast("np.ndarray", values) needs_object = _may_lose_precision_as_float64(values_arr) else: # float/complex values_dtype means _ensure_arraylike may # already have rounded large ints in a mixed int/float # list, so exactness cannot be checked from the array; # recover it from orig_values via the object cast. Smaller # ints/bools cast to float64 exactly. needs_object = values_dtype.kind in "fc" if needs_object: values = construct_1d_object_array_from_listlike(orig_values) elif isinstance(values, ABCMultiIndex): # Avoid raising in extract_array values = np.array(values) else: values = extract_array(values, extract_numpy=True, extract_range=True) comps_array = _ensure_arraylike(comps, func_name="isin") comps_array = extract_array(comps_array, extract_numpy=True) if not isinstance(comps_array, np.ndarray): # i.e. Extension Array return comps_array.isin(values) elif needs_i8_conversion(comps_array.dtype): # Dispatch to DatetimeLikeArrayMixin.isin return pd_array(comps_array).isin(values) elif needs_i8_conversion(values.dtype) and not is_object_dtype(comps_array.dtype): # e.g. comps_array are integers and values are datetime64s return np.zeros(comps_array.shape, dtype=bool) # TODO: not quite right ... Sparse/Categorical elif needs_i8_conversion(values.dtype): return isin(comps_array, values.astype(object)) elif isinstance(values.dtype, ExtensionDtype): return isin(np.asarray(comps_array), np.asarray(values)) # GH#59609: uint64 mixed with int64 has float64 as its common dtype, which # is lossy above 2**53. Targets outside the other side's range can never # match, so dropping them lets the rest be cast losslessly. if comps_array.dtype == np.uint64 and values.dtype == np.int64: values = values[values >= 0].astype(np.uint64) elif comps_array.dtype == np.int64 and values.dtype == np.uint64: values = values[values <= np.iinfo(np.int64).max].astype(np.int64) # GH#46485, GH#59609: whatever remains may still have a float64/complex128 # common dtype, which cannot represent every 64-bit integer; comparing # through it would report a match between distinct values above 2**53. common = np_find_common_type(values.dtype, comps_array.dtype) lossy_common = common.kind in "fc" and ( _may_lose_precision_as_float64(values) or _may_lose_precision_as_float64(comps_array) ) # GH16012 # Ensure np.isin doesn't get object types or it *may* throw an exception # Albeit hashmap has O(1) look-up (vs. O(logn) in sorted array), # isin is faster for small sizes # GH60678 # Ensure values don't contain , otherwise it throws exception with np.in1d if ( len(comps_array) > _MINIMUM_COMP_ARR_LEN and len(values) <= 26 and comps_array.dtype != object # np.isin would do the same lossy promotion internally and not lossy_common and not any(v is NA for v in values) ): # If the values include nan we need to check for nan explicitly # since np.nan it not equal to np.nan if isna(values).any(): def f(c, v): return np.logical_or(np.isin(c, v).ravel(), np.isnan(c)) else: f = lambda a, b: np.isin(a, b).ravel() else: if lossy_common: # Compare as Python ints instead of through the lossy float. common = np.dtype(object) values = values.astype(common, copy=False) comps_array = comps_array.astype(common, copy=False) f = htable.ismember return f(comps_array, values) def factorize_array( values: np.ndarray, use_na_sentinel: bool = True, size_hint: int | None = None, na_value: object = None, mask: npt.NDArray[np.bool_] | None = None, ) -> tuple[npt.NDArray[np.intp], np.ndarray]: """ Factorize a numpy array to codes and uniques. This doesn't do any coercion of types or unboxing before factorization. Parameters ---------- values : ndarray use_na_sentinel : bool, default True If True, the sentinel -1 will be used for NaN values. If False, NaN values will be encoded as non-negative integers and will not drop the NaN from the uniques of the values. size_hint : int, optional Passed through to the hashtable's 'get_labels' method na_value : object, optional A value in `values` to consider missing. Note: only use this parameter when you know that you don't have any values pandas would consider missing in the array (NaN for float data, iNaT for datetimes, etc.). mask : ndarray[bool], optional If not None, the mask is used as indicator for missing values (True = missing, False = valid) instead of `na_value` or condition "val != val". Returns ------- codes : ndarray[np.intp] uniques : ndarray """ original = values if values.dtype.kind in "mM": # _get_hashtable_algo will cast dt64/td64 to i8 via _ensure_data, so we # need to do the same to na_value. We are assuming here that the passed # na_value is an appropriately-typed NaT. # e.g. test_where_datetimelike_categorical na_value = iNaT hash_klass, values = _get_hashtable_algo(values) table = hash_klass(size_hint or len(values)) uniques, codes = table.factorize( values, na_sentinel=-1, na_value=na_value, mask=mask, ignore_na=use_na_sentinel, ) # re-cast e.g. i8->dt64/td64, uint8->bool uniques = _reconstruct_data(uniques, original.dtype, original) codes = ensure_platform_int(codes) return codes, uniques def factorize_monotonic_codes( arr: np.ndarray, ascending: bool, sort: bool, ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: """ Factorize an array known to be monotonic and of length >= 2. Uses adjacent-element comparison instead of hash table construction. The caller is responsible for having verified monotonicity; monotonic arrays contain no NA values (NAs break monotonicity checks), so NA handling is not needed. Parameters ---------- arr : ndarray ascending : bool Whether `arr` is monotonic increasing (otherwise decreasing). sort : bool Whether the uniques should be in sorted order. Returns ------- codes : ndarray[np.intp] uniques_indexer : ndarray[np.intp] Positions in `arr` of the unique values, ordered to match `codes`. """ n = len(arr) # Find group transitions via adjacent-element comparison. transitions = np.empty(n, dtype=bool) transitions[0] = True transitions[1:] = arr[1:] != arr[:-1] # Build codes using repeat (faster than cumsum on bool). group_starts = np.flatnonzero(transitions) ngroups = len(group_starts) run_lengths = np.diff(group_starts, append=n) codes = np.repeat(np.arange(ngroups, dtype=np.intp), run_lengths) if sort and not ascending: # Reverse uniques to sorted order and remap codes. codes = ensure_platform_int(ngroups - 1 - codes) group_starts = group_starts[::-1] return codes, group_starts @set_module("pandas") def factorize( values, sort: bool = False, use_na_sentinel: bool = True, size_hint: int | None = None, ) -> tuple[np.ndarray, np.ndarray | Index]: """ Encode the object as an enumerated type or categorical variable. This method is useful for obtaining a numeric representation of an array when all that matters is identifying distinct values. `factorize` is available as both a top-level function :func:`pandas.factorize`, and as a method :meth:`Series.factorize` and :meth:`Index.factorize`. Parameters ---------- values : sequence A 1-D sequence. Sequences that aren't pandas objects are coerced to ndarrays before factorization. sort : bool, default False Sort `uniques` and shuffle `codes` to maintain the relationship. use_na_sentinel : bool, default True If True, the sentinel -1 will be used for NaN values. If False, NaN values will be encoded as non-negative integers and will not drop the NaN from the uniques of the values. size_hint : int, optional Hint to the hashtable sizer. Returns ------- codes : ndarray An integer ndarray that's an indexer into `uniques`. ``uniques.take(codes)`` will have the same values as `values`. uniques : ndarray, Index, or Categorical The unique valid values. When `values` is Categorical, `uniques` is a Categorical. When `values` is some other pandas object, an `Index` is returned. Otherwise, a 1-D ndarray is returned. .. note:: Even if there's a missing value in `values`, `uniques` will *not* contain an entry for it. See Also -------- cut : Discretize continuous-valued array. unique : Find the unique value in an array. Notes ----- Reference :ref:`the user guide ` for more examples. When working with object-dtype arrays, boolean and integer values may not be distinguished since ``True == 1`` and ``False == 0`` in Python. >>> codes, uniques = pd.factorize(np.array([True, 1, False, 0], dtype=object)) >>> codes array([0, 0, 1, 1]) >>> uniques array([True, False], dtype=object) Examples -------- These examples all show factorize as a top-level method like ``pd.factorize(values)``. The results are identical for methods like :meth:`Series.factorize`. >>> codes, uniques = pd.factorize(np.array(["b", "b", "a", "c", "b"], dtype="O")) >>> codes array([0, 0, 1, 2, 0]) >>> uniques array(['b', 'a', 'c'], dtype=object) With ``sort=True``, the `uniques` will be sorted, and `codes` will be shuffled so that the relationship is the maintained. >>> codes, uniques = pd.factorize( ... np.array(["b", "b", "a", "c", "b"], dtype="O"), sort=True ... ) >>> codes array([1, 1, 0, 2, 1]) >>> uniques array(['a', 'b', 'c'], dtype=object) When ``use_na_sentinel=True`` (the default), missing values are indicated in the `codes` with the sentinel value ``-1`` and missing values are not included in `uniques`. >>> codes, uniques = pd.factorize(np.array(["b", None, "a", "c", "b"], dtype="O")) >>> codes array([ 0, -1, 1, 2, 0]) >>> uniques array(['b', 'a', 'c'], dtype=object) Thus far, we've only factorized lists (which are internally coerced to NumPy arrays). When factorizing pandas objects, the type of `uniques` will differ. For Categoricals, a `Categorical` is returned. >>> cat = pd.Categorical(["a", "a", "c"], categories=["a", "b", "c"]) >>> codes, uniques = pd.factorize(cat) >>> codes array([0, 0, 1]) >>> uniques ['a', 'c'] Categories (3, str): ['a', 'b', 'c'] Notice that ``'b'`` is in ``uniques.categories``, despite not being present in ``cat.values``. For all other pandas objects, an Index of the appropriate type is returned. >>> cat = pd.Series(["a", "a", "c"]) >>> codes, uniques = pd.factorize(cat) >>> codes array([0, 0, 1]) >>> uniques Index(['a', 'c'], dtype='str') If NaN is in the values, and we want to include NaN in the uniques of the values, it can be achieved by setting ``use_na_sentinel=False``. >>> values = np.array([1, 2, 1, np.nan]) >>> codes, uniques = pd.factorize(values) # default: use_na_sentinel=True >>> codes array([ 0, 1, 0, -1]) >>> uniques array([1., 2.]) >>> codes, uniques = pd.factorize(values, use_na_sentinel=False) >>> codes array([0, 1, 0, 2]) >>> uniques array([ 1., 2., nan]) """ # Implementation notes: This method is responsible for 3 things # 1.) coercing data to array-like (ndarray, Index, extension array) # 2.) factorizing codes and uniques # 3.) Maybe boxing the uniques in an Index # # Step 2 is dispatched to extension types (like Categorical). They are # responsible only for factorization. All data coercion, sorting and boxing # should happen here. if isinstance(values, (ABCIndex, ABCSeries)): return values.factorize(sort=sort, use_na_sentinel=use_na_sentinel) values = _ensure_arraylike(values, func_name="factorize") original = values if not isinstance(values, np.ndarray): # i.e. ExtensionArray codes, uniques = values.factorize(use_na_sentinel=use_na_sentinel) else: values = np.asarray(values) # convert DTA/TDA/MultiIndex if not use_na_sentinel and values.dtype == object: # factorize can now handle differentiating various types of null values. # These can only occur when the array has object dtype. # However, for backwards compatibility we only use the null for the # provided dtype. This may be revisited in the future, see GH#48476. null_mask = isna(values) if null_mask.any(): na_value = na_value_for_dtype(values.dtype, compat=False) # Don't modify (potentially user-provided) array # error: Argument 2 to "where" has incompatible type "Scalar"; # expected "_Buffer | _SupportsArray[dtype[Any]] | ..." values = np.where(null_mask, na_value, values) # type: ignore[arg-type] codes, uniques = factorize_array( values, use_na_sentinel=use_na_sentinel, size_hint=size_hint, ) if sort and len(uniques) > 0: uniques, codes = safe_sort( uniques, codes, use_na_sentinel=use_na_sentinel, assume_unique=True, verify=False, ) uniques = _reconstruct_data(uniques, original.dtype, original) return codes, uniques def value_counts_internal( values, sort: bool = True, ascending: bool = False, normalize: bool = False, bins=None, dropna: bool = True, ) -> Series: from pandas import ( DatetimeIndex, Index, Series, TimedeltaIndex, ) index_name = getattr(values, "name", None) name = "proportion" if normalize else "count" if bins is not None: from pandas.core.reshape.tile import cut if isinstance(values, Series): values = values._values try: ii = cut(values, bins, include_lowest=True) except TypeError as err: raise TypeError("bins argument only works with numeric data.") from err # count, remove nulls (from the index), and but the bins result = ii.value_counts(dropna=dropna) result.name = name result = result[result.index.notna()] result.index = result.index.astype("interval") result = result.sort_index() # if we are dropna and we have NO values if dropna and (result._values == 0).all(): result = result.iloc[0:0] # normalizing is by len of all (regardless of dropna) normalize_denominator = len(ii) else: normalize_denominator = None if is_extension_array_dtype(values): # handle Categorical and sparse, result = Series(values, copy=False)._values.value_counts(dropna=dropna) result.name = name result.index.name = index_name elif isinstance(values, ABCMultiIndex): # GH49558 levels = list(range(values.nlevels)) result = ( Series(index=values, name=name) .groupby(level=levels, dropna=dropna) .size() ) result.index.names = values.names else: values = _ensure_arraylike(values, func_name="value_counts") keys, counts, _ = value_counts_arraylike(values, dropna) if keys.dtype == np.float16: keys = keys.astype(np.float32) # Starting in 3.0, we no longer perform dtype inference on the # Index object we construct here, xref GH#56161 idx = Index(keys, dtype=keys.dtype, name=index_name, copy=False) if ( not sort and isinstance(values, (DatetimeIndex, TimedeltaIndex)) and idx.equals(values) and values._inferred_freq_str is not None ): # Preserve freq of original index idx.freq = values._inferred_freq_str # type: ignore[attr-defined] result = Series(counts, index=idx, name=name, copy=False) if sort: result = result.sort_values(ascending=ascending, kind="stable") if normalize: if normalize_denominator is not None: result = result / normalize_denominator else: result = result / result.sum() return result # Called once from SparseArray, otherwise could be private def value_counts_arraylike( values: np.ndarray, dropna: bool, mask: npt.NDArray[np.bool_] | None = None ) -> tuple[ArrayLike, npt.NDArray[np.int64], int]: """ Parameters ---------- values : np.ndarray dropna : bool mask : np.ndarray[bool] or None, default None Returns ------- uniques : np.ndarray counts : np.ndarray[np.int64] """ original = values values = _ensure_data(values) keys, counts, na_counter = htable.value_count(values, dropna, mask=mask) if needs_i8_conversion(original.dtype): # datetime, timedelta, or period if dropna: mask = keys != iNaT keys, counts = keys[mask], counts[mask] res_keys = _reconstruct_data(keys, original.dtype, original) return res_keys, counts, na_counter def duplicated( values: ArrayLike, keep: Literal["first", "last", False] = "first", mask: npt.NDArray[np.bool_] | None = None, ) -> npt.NDArray[np.bool_]: """ Return boolean ndarray denoting duplicate values. Parameters ---------- values : np.ndarray or ExtensionArray Array over which to check for duplicate values. keep : {'first', 'last', False}, default 'first' - ``first`` : Mark duplicates as ``True`` except for the first occurrence. - ``last`` : Mark duplicates as ``True`` except for the last occurrence. - False : Mark all duplicates as ``True``. mask : ndarray[bool], optional array indicating which elements to exclude from checking Returns ------- duplicated : ndarray[bool] """ values = _ensure_data(values) return htable.duplicated(values, keep=keep, mask=mask) def mode( values: ArrayLike, dropna: bool = True, mask: npt.NDArray[np.bool_] | None = None ) -> tuple[np.ndarray, npt.NDArray[np.bool_]] | ExtensionArray: """ Returns the mode(s) of an array. Parameters ---------- values : array-like Array over which to check for duplicate values. dropna : bool, default True Don't consider counts of NaN/NaT. Returns ------- Union[Tuple[np.ndarray, npt.NDArray[np.bool_]], ExtensionArray] """ values = _ensure_arraylike(values, func_name="mode") original = values if needs_i8_conversion(values.dtype): # Got here with ndarray; dispatch to DatetimeArray/TimedeltaArray. values = ensure_wrapped_if_datetimelike(values) values = cast("ExtensionArray", values) return values._mode(dropna=dropna) values = _ensure_data(values) npresult, res_mask = htable.mode(values, dropna=dropna, mask=mask) if res_mask is None: res_mask = np.zeros(npresult.shape, dtype=np.bool_) else: return npresult, res_mask try: npresult = safe_sort(npresult) except TypeError as err: warnings.warn( f"Unable to sort modes: {err}", stacklevel=find_stack_level(), ) result = _reconstruct_data(npresult, original.dtype, original) return result, res_mask def rank( values: ArrayLike, axis: AxisInt = 0, method: RankMethod = "average", na_option: RankNaOption = "keep", ascending: bool = True, pct: bool = False, mask: npt.NDArray[np.bool_] | None = None, ) -> npt.NDArray[np.float64]: """ Rank the values along a given axis. Parameters ---------- values : np.ndarray or ExtensionArray Array whose values will be ranked. The number of dimensions in this array must not exceed 2. axis : int, default 0 Axis over which to perform rankings. method : {'average', 'min', 'max', 'first', 'dense'}, default 'average' The method by which tiebreaks are broken during the ranking. na_option : {'keep', 'top'}, default 'keep' The method by which NaNs are placed in the ranking. - ``keep``: rank each NaN value with a NaN ranking - ``top``: replace each NaN with either +/- inf so that they there are ranked at the top ascending : bool, default True Whether or not the elements should be ranked in ascending order. pct : bool, default False Whether or not to the display the returned rankings in integer form (e.g. 1, 2, 3) or in percentile form (e.g. 0.333..., 0.666..., 1). mask : bool ndarray, optional Boolean array indicating which elements to exclude from ranking. """ is_datetimelike = needs_i8_conversion(values.dtype) values = _ensure_data(values) if values.ndim == 1: ranks = algos.rank_1d( values, is_datetimelike=is_datetimelike, ties_method=method, ascending=ascending, na_option=na_option, pct=pct, mask=mask, ) elif values.ndim == 2: assert mask is None ranks = algos.rank_2d( values, axis=axis, is_datetimelike=is_datetimelike, ties_method=method, ascending=ascending, na_option=na_option, pct=pct, ) else: raise TypeError("Array with ndim > 2 is not supported.") return ranks def is_monotonic(values: ArrayLike) -> tuple[bool, bool, bool]: """ Determine whether values are monotonic increasing/decreasing. Parameters ---------- values : np.ndarray or ExtensionArray Returns ------- tuple[bool, bool, bool] (is_monotonic_increasing, is_monotonic_decreasing, is_strict_monotonic) Raises ------ TypeError If the dtype is not orderable by the underlying routine (e.g. complex). """ timelike = needs_i8_conversion(values.dtype) values = _ensure_data(values) inc, dec, strict = algos.is_monotonic(values, timelike=timelike) return bool(inc), bool(dec), bool(strict) # ---- # # take # # ---- # def occurrence_rank(codes: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: """ For each element, the number of earlier elements with the same code. Parameters ---------- codes : np.ndarray[intp] Returns ------- np.ndarray[intp] Examples -------- >>> occurrence_rank(np.array([0, 1, 0, 0, 1])) array([0, 0, 1, 2, 1]) """ order = np.argsort(codes, kind="stable") sorted_codes = codes[order] n = len(codes) is_start = np.empty(n, dtype=bool) is_start[:1] = True is_start[1:] = sorted_codes[1:] != sorted_codes[:-1] group_start = np.maximum.accumulate(np.where(is_start, np.arange(n), 0)) rank = np.empty(n, dtype=np.intp) rank[order] = np.arange(n) - group_start return rank @set_module("pandas.api.extensions") def take( arr, indices: TakeIndexer, axis: AxisInt = 0, allow_fill: bool = False, fill_value=None, ): """ Take elements from an array. Parameters ---------- arr : numpy.ndarray, ExtensionArray, Index, or Series Input array. indices : sequence of int or one-dimensional np.ndarray of int Indices to be taken. axis : int, default 0 The axis over which to select values. allow_fill : bool, default False How to handle negative values in `indices`. * False: negative values in `indices` indicate positional indices from the right (the default). This is similar to :func:`numpy.take`. * True: negative values in `indices` indicate missing values. These values are set to `fill_value`. Any other negative values raise a ``ValueError``. fill_value : any, optional Fill value to use for NA-indices when `allow_fill` is True. This may be ``None``, in which case the default NA value for the type (``self.dtype.na_value``) is used. For multi-dimensional `arr`, each *element* is filled with `fill_value`. Returns ------- ndarray or ExtensionArray Same type as the input. Raises ------ IndexError When `indices` is out of bounds for the array. ValueError When the indexer contains negative values other than ``-1`` and `allow_fill` is True. Notes ----- When `allow_fill` is False, `indices` may be whatever dimensionality is accepted by NumPy for `arr`. When `allow_fill` is True, `indices` should be 1-D. See Also -------- numpy.take : Take elements from an array along an axis. Examples -------- >>> import pandas as pd With the default ``allow_fill=False``, negative numbers indicate positional indices from the right. >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1]) array([10, 10, 30]) Setting ``allow_fill=True`` will place `fill_value` in those positions. >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True) array([10., 10., nan]) >>> pd.api.extensions.take( ... np.array([10, 20, 30]), [0, 0, -1], allow_fill=True, fill_value=-10 ... ) array([ 10, 10, -10]) """ if not isinstance( arr, (np.ndarray, ABCExtensionArray, ABCIndex, ABCSeries, ABCNumpyExtensionArray), ): # GH#52981 raise TypeError( "pd.api.extensions.take requires a numpy.ndarray, ExtensionArray, " f"Index, Series, or NumpyExtensionArray got {type(arr).__name__}." ) indices = ensure_platform_int(indices) if allow_fill: # Pandas style, -1 means NA validate_indices(indices, arr.shape[axis]) # error: Argument 1 to "take_nd" has incompatible type # "ndarray[Any, Any] | ExtensionArray | Index | Series"; expected # "ndarray[Any, Any]" result = take_nd( arr, # type: ignore[arg-type] indices, axis=axis, allow_fill=True, fill_value=fill_value, ) else: # NumPy style # error: Unexpected keyword argument "axis" for "take" of "ExtensionArray" result = arr.take(indices, axis=axis) # type: ignore[call-arg,assignment] return result # ------------ # # searchsorted # # ------------ # def searchsorted( arr: ArrayLike, value: NumpyValueArrayLike | ExtensionArray, side: Literal["left", "right"] = "left", sorter: NumpySorter | None = None, ) -> npt.NDArray[np.intp] | np.intp: """ Find indices where elements should be inserted to maintain order. Find the indices into a sorted array `arr` (a) such that, if the corresponding elements in `value` were inserted before the indices, the order of `arr` would be preserved. Assuming that `arr` is sorted: ====== ================================ `side` returned index `i` satisfies ====== ================================ left ``arr[i-1] < value <= self[i]`` right ``arr[i-1] <= value < self[i]`` ====== ================================ Parameters ---------- arr: np.ndarray, ExtensionArray, Series Input array. If `sorter` is None, then it must be sorted in ascending order, otherwise `sorter` must be an array of indices that sort it. value : array-like or scalar Values to insert into `arr`. side : {'left', 'right'}, optional If 'left', the index of the first suitable location found is given. If 'right', return the last such index. If there is no suitable index, return either 0 or N (where N is the length of `self`). sorter : 1-D array-like, optional Optional array of integer indices that sort array a into ascending order. They are typically the result of argsort. Returns ------- array of ints or int If value is array-like, array of insertion points. If value is scalar, a single integer. See Also -------- numpy.searchsorted : Similar method from NumPy. """ if sorter is not None: sorter = ensure_platform_int(sorter) if ( isinstance(arr, np.ndarray) and arr.dtype.kind in "iu" and (is_integer(value) or is_integer_dtype(value)) ): # if `arr` and `value` have different dtypes, `arr` would be # recast by numpy, causing a slow search. # Before searching below, we therefore try to give `value` the # same dtype as `arr`, while guarding against integer overflows. iinfo = np.iinfo(arr.dtype.type) value_arr = np.array([value]) if is_integer(value) else np.array(value) if (value_arr >= iinfo.min).all() and (value_arr <= iinfo.max).all(): # value within bounds, so no overflow, so can convert value dtype # to dtype of arr dtype = arr.dtype else: dtype = value_arr.dtype if is_integer(value): # We know that value is int value = cast("int", dtype.type(value)) else: value = pd_array(cast("ArrayLike", value), dtype=dtype) else: # E.g. if `arr` is an array with dtype='datetime64[ns]' # and `value` is a pd.Timestamp, we may need to convert value arr = ensure_wrapped_if_datetimelike(arr) # Argument 1 to "searchsorted" of "ndarray" has incompatible type # "Union[NumpyValueArrayLike, ExtensionArray]"; expected "NumpyValueArrayLike" return arr.searchsorted(value, side=side, sorter=sorter) # type: ignore[arg-type] # ---- # # diff # # ---- # _diff_special = {"float64", "float32", "int64", "int32", "int16", "int8"} def diff(arr, n: int | float | np.integer | np.floating, axis: AxisInt = 0): """ difference of n between self, analogous to s-s.shift(n) Parameters ---------- arr : ndarray or ExtensionArray n : int number of periods axis : {0, 1} axis to shift on Returns ------- shifted """ # added a check on the integer value of period # see https://github.com/pandas-dev/pandas/issues/56607 if not lib.is_integer(n): if not (is_float(n) and n.is_integer()): raise ValueError("periods must be an integer") n = int(n) na = np.nan dtype = arr.dtype is_bool = is_bool_dtype(dtype) if is_bool: op = operator.xor else: op = operator.sub if isinstance(dtype, NumpyEADtype): # NumpyExtensionArray cannot necessarily hold shifted versions of itself. arr = arr.to_numpy() dtype = arr.dtype if not isinstance(arr, np.ndarray): # i.e ExtensionArray if hasattr(arr, f"__{op.__name__}__"): if axis >= arr.ndim: raise ValueError(f"cannot diff {type(arr).__name__} on axis={axis}") return op(arr, arr.shift(n)) else: raise TypeError( f"{type(arr).__name__} has no 'diff' method. " "Convert to a suitable dtype prior to calling 'diff'." ) is_timedelta = False if arr.dtype.kind in "mM": dtype = np.int64 arr = arr.view("i8") na = iNaT is_timedelta = True elif is_bool: # We have to cast in order to be able to hold np.nan dtype = np.object_ elif dtype.kind in "iu": # We have to cast in order to be able to hold np.nan # int8, int16 are incompatible with float64, # see https://github.com/cython/cython/issues/2646 if arr.dtype.name in ["int8", "int16"]: dtype = np.float32 else: dtype = np.float64 orig_ndim = arr.ndim if orig_ndim == 1: # reshape so we can always use algos.diff_2d arr = arr.reshape(-1, 1) # TODO: require axis == 0 dtype = np.dtype(dtype) out_arr = np.empty(arr.shape, dtype=dtype) na_indexer = [slice(None)] * 2 na_indexer[axis] = slice(None, n) if n >= 0 else slice(n, None) out_arr[tuple(na_indexer)] = na if arr.dtype.name in _diff_special: # TODO: can diff_2d dtype specialization troubles be fixed by defining # out_arr inside diff_2d? algos.diff_2d(arr, out_arr, int(n), axis, datetimelike=is_timedelta) else: # To keep mypy happy, _res_indexer is a list while res_indexer is # a tuple, ditto for lag_indexer. _res_indexer = [slice(None)] * 2 _res_indexer[axis] = slice(n, None) if n >= 0 else slice(None, n) res_indexer = tuple(_res_indexer) _lag_indexer = [slice(None)] * 2 _lag_indexer[axis] = slice(None, -n) if n > 0 else slice(-n, None) lag_indexer = tuple(_lag_indexer) out_arr[res_indexer] = op(arr[res_indexer], arr[lag_indexer]) if is_timedelta: out_arr = out_arr.view("timedelta64[ns]") if orig_ndim == 1: out_arr = out_arr[:, 0] return out_arr # -------------------------------------------------------------------- # Helper functions # Minimum number of values before safe_sort's counting-sort path is worth it. # Below this, argsort's n*log(n) is cheaper than the linear counting work. _MIN_COUNTING_SORT_LEN = 5_000 # Note: safe_sort is in algorithms.py instead of sorting.py because it is # low-dependency, is used in this module, and used private methods from # this module. def safe_sort( values: Index | ArrayLike, codes: npt.NDArray[np.intp] | None = None, use_na_sentinel: bool = True, assume_unique: bool = False, verify: bool = True, ) -> AnyArrayLike | tuple[AnyArrayLike, np.ndarray]: """ Sort ``values`` and reorder corresponding ``codes``. ``values`` should be unique if ``codes`` is not None. Safe for use with mixed types (int, str), orders ints before strs. Parameters ---------- values : list-like Sequence; must be unique if ``codes`` is not None. codes : np.ndarray[intp] or None, default None Indices to ``values``. All out of bound indices are treated as "not found" and will be masked with ``-1``. use_na_sentinel : bool, default True If True, the sentinel -1 will be used for NaN values. If False, NaN values will be encoded as non-negative integers and will not drop the NaN from the uniques of the values. assume_unique : bool, default False When True, ``values`` are assumed to be unique, which can speed up the calculation. Ignored when ``codes`` is None. verify : bool, default True Check if codes are out of bound for the values and put out of bound codes equal to ``-1``. If ``verify=False``, it is assumed there are no out of bound codes. Ignored when ``codes`` is None. Returns ------- ordered : AnyArrayLike Sorted ``values`` new_codes : ndarray Reordered ``codes``; returned when ``codes`` is not None. Raises ------ TypeError * If ``values`` is not list-like or if ``codes`` is neither None nor list-like * If ``values`` cannot be sorted ValueError * If ``codes`` is not None and ``values`` contain duplicates. """ if not isinstance(values, (np.ndarray, ABCExtensionArray, ABCIndex)): raise TypeError( "Only np.ndarray, ExtensionArray, and Index objects are allowed to " "be passed to safe_sort as values" ) sorter = None ordered: AnyArrayLike # For integer values spanning a modest range, the codes remapping below # can rank each value within that range (a counting sort) instead of # argsorting. use_counting = False vmin = rng_size = 0 if ( codes is not None and isinstance(values, np.ndarray) and values.dtype.kind in "iu" and len(values) >= _MIN_COUNTING_SORT_LEN ): vmin = int(values.min()) rng_size = int(values.max()) - vmin + 1 use_counting = rng_size <= 4 * len(values) if use_counting: ordered = np.sort(cast("np.ndarray", values)) elif ( not isinstance(values.dtype, ExtensionDtype) and lib.infer_dtype(values, skipna=False) == "mixed-integer" ): ordered = _sort_mixed(values) else: try: if ( codes is None and isinstance(values, np.ndarray) and values.dtype.kind in "iufb" ): # we only need the sorted values here, and a direct sort # is cheaper than argsort + take ordered = np.sort(values) else: sorter = values.argsort() ordered = values.take(sorter) except (TypeError, decimal.InvalidOperation): # Previous sorters failed or were not applicable, try `_sort_mixed` # which would work, but which fails for special case of 1d arrays # with tuples. if values.size and isinstance(values[0], tuple): # error: Argument 1 to "_sort_tuples" has incompatible type # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected # "ndarray[Any, Any]" ordered = _sort_tuples(values) # type: ignore[arg-type] else: ordered = _sort_mixed(values) # codes: if codes is None: return ordered if not is_list_like(codes): raise TypeError( "Only list-like objects or None are allowed to " "be passed to safe_sort as codes" ) codes = ensure_platform_int(np.asarray(codes)) # ranks[i] gives the position of values[i] in `ordered` if use_counting: arr = cast("np.ndarray", values) if arr.dtype.kind == "i": # go through int64 so differences don't overflow narrower signed # dtypes; int64 wraparound is exact since the true differences # are within rng_size shifted = arr.astype(np.int64, copy=False) - vmin else: # unsigned: differences always fit the unsigned dtype shifted = arr - vmin present = np.zeros(rng_size, dtype=bool) present[shifted] = True counts = present.cumsum(dtype=np.intp) # the counting pass gives the uniqueness check for free if not assume_unique and counts[-1] != len(values): raise ValueError("values should be unique if codes is not None") ranks = counts[shifted] ranks -= 1 else: if not assume_unique and not len(unique(values)) == len(values): raise ValueError("values should be unique if codes is not None") if sorter is not None: # sorter is a permutation, so scatter is a faster equivalent to # `ranks = sorter.argsort()` ranks = np.empty(len(sorter), dtype=np.intp) ranks[sorter] = np.arange(len(sorter), dtype=np.intp) else: # mixed types # error: Argument 1 to "_get_hashtable_algo" has incompatible type # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected # "ndarray[Any, Any]" hash_klass, values = _get_hashtable_algo(values) # type: ignore[arg-type] t = hash_klass(len(values)) # error: Argument 1 to "map_locations" of "HashTable" has incompatible # type "ExtensionArray | ndarray[Any, Any] | Index | Series"; # expected "ndarray" t.map_locations(ordered) # type: ignore[arg-type] ranks = ensure_platform_int(t.lookup(values)) if use_na_sentinel: # take_nd is faster, but only works for na_sentinels of -1 if verify: mask = (codes < -len(values)) | (codes >= len(values)) codes[mask] = -1 new_codes = take_nd(ranks, codes, fill_value=-1) else: # Out of bound indices will be masked with `-1` next, so we # may deal with them here without performance loss using `mode='wrap'` new_codes = ranks.take(codes, mode="wrap") return ordered, ensure_platform_int(new_codes) def _sort_mixed(values) -> AnyArrayLike: """order ints before strings before nulls in 1d arrays""" str_pos = np.array([isinstance(x, str) for x in values], dtype=bool) null_pos = np.array([isna(x) for x in values], dtype=bool) num_pos = ~str_pos & ~null_pos str_argsort = np.argsort(values[str_pos]) num_argsort = np.argsort(values[num_pos]) # convert boolean arrays to positional indices, then order by underlying values str_locs = str_pos.nonzero()[0].take(str_argsort) num_locs = num_pos.nonzero()[0].take(num_argsort) null_locs = null_pos.nonzero()[0] locs = np.concatenate([num_locs, str_locs, null_locs]) return values.take(locs) def _sort_tuples(values: np.ndarray) -> np.ndarray: """ Convert array of tuples (1d) to array of arrays (2d). We need to keep the columns separately as they contain different types and nans (can't use `np.sort` as it may fail when str and nan are mixed in a column as types cannot be compared). """ from pandas.core.internals.construction import to_arrays from pandas.core.sorting import lexsort_indexer arrays, _ = to_arrays(values, None) indexer = lexsort_indexer(arrays, orders=True) return values[indexer] def union_with_duplicates( lvals: ArrayLike | Index, rvals: ArrayLike | Index ) -> ArrayLike | Index: """ Extracts the union from lvals and rvals with respect to duplicates and nans in both arrays. Parameters ---------- lvals: np.ndarray or ExtensionArray left values which is ordered in front. rvals: np.ndarray or ExtensionArray right values ordered after lvals. Returns ------- np.ndarray or ExtensionArray Containing the unsorted union of both arrays. Notes ----- Caller is responsible for ensuring lvals.dtype == rvals.dtype. """ from pandas import Series l_count = value_counts_internal(lvals, dropna=False) r_count = value_counts_internal(rvals, dropna=False) l_count, r_count = l_count.align(r_count, fill_value=0) final_count = np.maximum(l_count.values, r_count.values) final_count = Series(final_count, index=l_count.index, dtype="int", copy=False) if isinstance(lvals, ABCMultiIndex) and isinstance(rvals, ABCMultiIndex): unique_vals = lvals.append(rvals).unique() else: if isinstance(lvals, ABCIndex): lvals = lvals._values if isinstance(rvals, ABCIndex): rvals = rvals._values # error: List item 0 has incompatible type "Union[ExtensionArray, # ndarray[Any, Any], Index]"; expected "Union[ExtensionArray, # ndarray[Any, Any]]" combined = concat_compat([lvals, rvals]) # type: ignore[list-item] unique_vals = unique(combined) unique_vals = ensure_wrapped_if_datetimelike(unique_vals) repeats = final_count.reindex(unique_vals).values return np.repeat(unique_vals, repeats) def map_array( arr: ArrayLike, mapper, na_action: Literal["ignore"] | None = None, ) -> np.ndarray | ExtensionArray | Index: """ Map values using an input mapping or function. Parameters ---------- mapper : function, dict, or Series Mapping correspondence. na_action : {None, 'ignore'}, default None If 'ignore', propagate NA values, without passing them to the mapping correspondence. Returns ------- Union[ndarray, Index, ExtensionArray] The output of the mapping function applied to the array. If the function returns a tuple with more than one element a MultiIndex will be returned. """ from pandas import Index if na_action not in (None, "ignore"): msg = f"na_action must either be 'ignore' or None, {na_action} was passed" raise ValueError(msg) # we can fastpath dict/Series to an efficient map # as we know that we are not going to have to yield # python types if is_dict_like(mapper): if isinstance(mapper, dict) and hasattr(mapper, "__missing__"): # If a dictionary subclass defines a default value method, # convert mapper to a lookup function (GH #15999). dict_with_default = mapper mapper = lambda x: dict_with_default[ np.nan if isinstance(x, float) and np.isnan(x) else x ] else: # Dictionary does not have a default. Thus it's safe to # convert to a Series for efficiency. # we specify the keys here to handle the # possibility that they are tuples # The return value of mapping with an empty mapper is # expected to be pd.Series(np.nan, ...). As np.nan is # of dtype float64 the return value of this method should # be float64 as well from pandas import Series if len(mapper) == 0: mapper = Series(mapper, dtype=np.float64) elif isinstance(mapper, dict): mapper = Series( mapper.values(), index=Index(mapper.keys(), tupleize_cols=False) ) else: mapper = Series(mapper) if isinstance(mapper, ABCSeries): if na_action == "ignore": mapper = mapper[mapper.index.notna()] # Since values were input this means we came from either # a dict or a series and mapper should be an index indexer = mapper.index.get_indexer(arr) new_values = take_nd(mapper._values, indexer) return new_values if not len(arr): return arr.copy() # we must convert to python types values = arr.astype(object, copy=False) if na_action is None: return lib.map_infer(values, mapper) else: return lib.map_infer_mask(values, mapper, mask=isna(values).view(np.uint8))