Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from typing import (
4 TYPE_CHECKING,
5 cast,
6)
8import numpy as np
10from pandas._libs import (
11 NaT,
12 lib,
13)
14from pandas.errors import InvalidIndexError
16from pandas.core.dtypes.cast import find_common_type
18from pandas.core.algorithms import safe_sort
19from pandas.core.indexes.base import (
20 Index,
21 _new_Index,
22 ensure_index,
23 ensure_index_from_sequences,
24 get_unanimous_names,
25 maybe_sequence_to_range,
26)
27from pandas.core.indexes.category import CategoricalIndex
28from pandas.core.indexes.datetimes import DatetimeIndex
29from pandas.core.indexes.interval import IntervalIndex
30from pandas.core.indexes.multi import MultiIndex
31from pandas.core.indexes.period import PeriodIndex
32from pandas.core.indexes.range import RangeIndex
33from pandas.core.indexes.timedeltas import TimedeltaIndex
35if TYPE_CHECKING:
36 from pandas._typing import Axis
39__all__ = [
40 "CategoricalIndex",
41 "DatetimeIndex",
42 "Index",
43 "IntervalIndex",
44 "InvalidIndexError",
45 "MultiIndex",
46 "NaT",
47 "PeriodIndex",
48 "RangeIndex",
49 "TimedeltaIndex",
50 "_new_Index",
51 "all_indexes_same",
52 "default_index",
53 "ensure_index",
54 "ensure_index_from_sequences",
55 "get_objs_combined_axis",
56 "get_unanimous_names",
57 "maybe_sequence_to_range",
58 "safe_sort_index",
59 "union_indexes",
60]
63def get_objs_combined_axis(
64 objs,
65 intersect: bool = False,
66 axis: Axis = 0,
67 sort: bool | lib.NoDefault = True,
68) -> Index:
69 """
70 Extract combined index: return intersection or union (depending on the
71 value of "intersect") of indexes on given axis, or None if all objects
72 lack indexes (e.g. they are numpy arrays).
74 Parameters
75 ----------
76 objs : list
77 Series or DataFrame objects, may be mix of the two.
78 intersect : bool, default False
79 If True, calculate the intersection between indexes. Otherwise,
80 calculate the union.
81 axis : {0 or 'index', 1 or 'outer'}, default 0
82 The axis to extract indexes from.
83 sort : bool, default True
84 Whether the result index should come out sorted or not. NoDefault
85 use for deprecation in GH#57335.
87 Returns
88 -------
89 Index
90 """
91 obs_idxes = [obj._get_axis(axis) for obj in objs]
92 return _get_combined_index(obs_idxes, intersect=intersect, sort=sort)
95def _get_distinct_objs(objs: list[Index]) -> list[Index]:
96 """
97 Return a list with distinct elements of "objs" (different ids).
98 Preserves order.
99 """
100 ids: set[int] = set()
101 res = []
102 for obj in objs:
103 if id(obj) not in ids:
104 ids.add(id(obj))
105 res.append(obj)
106 return res
109def _get_combined_index(
110 indexes: list[Index],
111 intersect: bool = False,
112 sort: bool | lib.NoDefault = False,
113) -> Index:
114 """
115 Return the union or intersection of indexes.
117 Parameters
118 ----------
119 indexes : list of Index or list objects
120 When intersect=True, do not accept list of lists.
121 intersect : bool, default False
122 If True, calculate the intersection between indexes. Otherwise,
123 calculate the union.
124 sort : bool, default False
125 Whether the result index should come out sorted or not. NoDefault
126 used for deprecation of GH#57335
128 Returns
129 -------
130 Index
131 """
132 # TODO: handle index names!
133 indexes = _get_distinct_objs(indexes)
134 if len(indexes) == 0:
135 index: Index = default_index(0)
136 elif len(indexes) == 1:
137 index = indexes[0]
138 elif intersect:
139 index = indexes[0]
140 for other in indexes[1:]:
141 index = index.intersection(other)
142 else:
143 index = union_indexes(indexes, sort=sort if sort is lib.no_default else False)
144 index = ensure_index(index)
146 if sort and sort is not lib.no_default:
147 index = safe_sort_index(index)
148 return index
151def safe_sort_index(index: Index) -> Index:
152 """
153 Returns the sorted index
155 We keep the dtypes and the name attributes.
157 Parameters
158 ----------
159 index : an Index
161 Returns
162 -------
163 Index
164 """
165 if index.is_monotonic_increasing:
166 return index
168 try:
169 array_sorted = safe_sort(index)
170 except TypeError:
171 pass
172 else:
173 if isinstance(array_sorted, Index):
174 return array_sorted
176 array_sorted = cast(np.ndarray, array_sorted)
177 if isinstance(index, MultiIndex):
178 index = MultiIndex.from_tuples(array_sorted, names=index.names)
179 else:
180 index = Index(array_sorted, name=index.name, dtype=index.dtype)
182 return index
185def union_indexes(indexes, sort: bool | lib.NoDefault = True) -> Index:
186 """
187 Return the union of indexes.
189 The behavior of sort and names is not consistent.
191 Parameters
192 ----------
193 indexes : list of Index or list objects
194 sort : bool, default True
195 Whether the result index should come out sorted or not. NoDefault
196 used for deprecation of GH#57335.
198 Returns
199 -------
200 Index
201 """
202 if len(indexes) == 0:
203 raise AssertionError("Must have at least 1 Index to union")
204 if len(indexes) == 1:
205 result = indexes[0]
206 if isinstance(result, list):
207 if not sort or sort is lib.no_default:
208 result = Index(result)
209 else:
210 result = Index(sorted(result))
211 return result
213 indexes, kind = _sanitize_and_check(indexes)
215 if kind == "special":
216 result = indexes[0]
218 num_dtis = 0
219 num_dti_tzs = 0
220 for idx in indexes:
221 if isinstance(idx, DatetimeIndex):
222 num_dtis += 1
223 if idx.tz is not None:
224 num_dti_tzs += 1
225 if num_dti_tzs not in [0, num_dtis]:
226 # TODO: this behavior is not tested (so may not be desired),
227 # but is kept in order to keep behavior the same when
228 # deprecating union_many
229 # test_frame_from_dict_with_mixed_indexes
230 raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")
232 if num_dtis == len(indexes):
233 if sort is lib.no_default:
234 sort = True
235 result = indexes[0]
237 elif num_dtis > 1:
238 # If we have mixed timezones, our casting behavior may depend on
239 # the order of indexes, which we don't want.
240 sort = False
242 # TODO: what about Categorical[dt64]?
243 # test_frame_from_dict_with_mixed_indexes
244 indexes = [x.astype(object, copy=False) for x in indexes]
245 result = indexes[0]
247 for other in indexes[1:]:
248 result = result.union(other, sort=None if sort else False)
249 return result
251 elif kind == "array":
252 if not all_indexes_same(indexes):
253 dtype = find_common_type([idx.dtype for idx in indexes])
254 inds = [ind.astype(dtype, copy=False) for ind in indexes]
255 index = inds[0].unique()
256 other = inds[1].append(inds[2:])
257 diff = other[index.get_indexer_for(other) == -1]
258 if len(diff):
259 index = index.append(diff.unique())
260 if sort:
261 index = index.sort_values()
262 else:
263 index = indexes[0]
265 name = get_unanimous_names(*indexes)[0]
266 if name != index.name:
267 index = index.rename(name)
268 return index
269 elif kind == "list":
270 dtypes = [idx.dtype for idx in indexes if isinstance(idx, Index)]
271 if dtypes:
272 dtype = find_common_type(dtypes)
273 else:
274 dtype = None
275 all_lists = (idx.tolist() if isinstance(idx, Index) else idx for idx in indexes)
276 return Index(
277 lib.fast_unique_multiple_list_gen(all_lists, sort=bool(sort)),
278 dtype=dtype,
279 )
280 else:
281 raise ValueError(f"{kind=} must be 'special', 'array' or 'list'.")
284def _sanitize_and_check(indexes):
285 """
286 Verify the type of indexes and convert lists to Index.
288 Cases:
290 - [list, list, ...]: Return ([list, list, ...], 'list')
291 - [list, Index, ...]: Return _sanitize_and_check([Index, Index, ...])
292 Lists are sorted and converted to Index.
293 - [Index, Index, ...]: Return ([Index, Index, ...], TYPE)
294 TYPE = 'special' if at least one special type, 'array' otherwise.
296 Parameters
297 ----------
298 indexes : list of Index or list objects
300 Returns
301 -------
302 sanitized_indexes : list of Index or list objects
303 type : {'list', 'array', 'special'}
304 """
305 kinds = {type(index) for index in indexes}
307 if list in kinds:
308 if len(kinds) > 1:
309 indexes = [
310 Index(list(x)) if not isinstance(x, Index) else x for x in indexes
311 ]
312 kinds -= {list}
313 else:
314 return indexes, "list"
316 if len(kinds) > 1 or Index not in kinds:
317 return indexes, "special"
318 else:
319 return indexes, "array"
322def all_indexes_same(indexes) -> bool:
323 """
324 Determine if all indexes contain the same elements.
326 Parameters
327 ----------
328 indexes : iterable of Index objects
330 Returns
331 -------
332 bool
333 True if all indexes contain the same elements, False otherwise.
334 """
335 itr = iter(indexes)
336 first = next(itr)
337 return all(first.equals(index) for index in itr)
340def default_index(n: int) -> RangeIndex:
341 rng = range(n)
342 return RangeIndex._simple_new(rng, name=None)