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)
8from pandas._libs import (
9 NaT,
10 lib,
11)
12from pandas.errors import InvalidIndexError
14from pandas.core.dtypes.cast import find_common_type
16from pandas.core.algorithms import safe_sort
17from pandas.core.indexes.base import (
18 Index,
19 _new_Index,
20 ensure_index,
21 ensure_index_from_sequences,
22 get_unanimous_names,
23 maybe_sequence_to_range,
24)
25from pandas.core.indexes.category import CategoricalIndex
26from pandas.core.indexes.datetimes import DatetimeIndex
27from pandas.core.indexes.interval import IntervalIndex
28from pandas.core.indexes.multi import MultiIndex
29from pandas.core.indexes.period import PeriodIndex
30from pandas.core.indexes.range import RangeIndex
31from pandas.core.indexes.timedeltas import TimedeltaIndex
33if TYPE_CHECKING:
34 import numpy as np
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) -> tuple[Index, bool]:
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 The combined index.
91 bool
92 Whether all indexes are equal.
93 """
94 obs_idxes = [obj._get_axis(axis) for obj in objs]
95 return _get_combined_index(
96 obs_idxes,
97 intersect=intersect,
98 sort=sort,
99 )
102def _get_distinct_objs(objs: list[Index]) -> list[Index]:
103 """
104 Return a list with distinct elements of "objs" (different ids).
105 Preserves order.
106 """
107 ids = set()
108 res = []
109 for obj in objs:
110 if obj._id not in ids:
111 ids.add(obj._id)
112 res.append(obj)
113 return res
116def _get_combined_index(
117 indexes: list[Index],
118 intersect: bool = False,
119 sort: bool | lib.NoDefault = False,
120) -> tuple[Index, bool]:
121 """
122 Return the union or intersection of indexes.
124 Parameters
125 ----------
126 indexes : list of Index or list objects
127 When intersect=True, do not accept list of lists.
128 intersect : bool, default False
129 If True, calculate the intersection between indexes. Otherwise,
130 calculate the union.
131 sort : bool, default False
132 Whether the result index should come out sorted or not. NoDefault
133 used for deprecation of GH#57335
135 Returns
136 -------
137 Index
138 The combined index.
139 bool
140 Whether all indexes are equal.
141 """
142 # TODO: handle index names!
143 indexes = _get_distinct_objs(indexes)
144 if len(indexes) == 0:
145 index: Index = default_index(0)
146 all_equal = True
147 elif len(indexes) == 1:
148 index = indexes[0]
149 all_equal = True
150 elif intersect:
151 index = indexes[0]
152 for other in indexes[1:]:
153 index = index.intersection(other)
154 all_equal = False
155 else:
156 index, all_equal = union_indexes(
157 indexes,
158 sort=sort if sort is lib.no_default else False,
159 )
160 index = ensure_index(index)
162 if sort and sort is not lib.no_default:
163 index = safe_sort_index(index)
164 all_equal = False
165 return index, all_equal
168def safe_sort_index(index: Index) -> Index:
169 """
170 Returns the sorted index
172 We keep the dtypes and the name attributes.
174 Parameters
175 ----------
176 index : an Index
178 Returns
179 -------
180 Index
181 """
182 if index.is_monotonic_increasing:
183 return index
185 try:
186 array_sorted = safe_sort(index)
187 except TypeError:
188 pass
189 else:
190 if isinstance(array_sorted, Index):
191 return array_sorted
193 array_sorted = cast("np.ndarray", array_sorted)
194 if isinstance(index, MultiIndex):
195 index = MultiIndex.from_tuples(array_sorted, names=index.names)
196 else:
197 index = Index(array_sorted, name=index.name, dtype=index.dtype)
199 return index
202def union_indexes(
203 indexes,
204 sort: bool | lib.NoDefault = True,
205) -> tuple[Index, bool]:
206 """
207 Return the union of indexes.
209 The behavior of sort and names is not consistent.
211 Parameters
212 ----------
213 indexes : list of Index or list objects
214 sort : bool, default True
215 Whether the result index should come out sorted or not. NoDefault
216 used for deprecation of GH#57335.
218 Returns
219 -------
220 Index
221 The union index.
222 bool
223 Whether all indexes are equal.
224 """
226 if len(indexes) == 0:
227 raise AssertionError("Must have at least 1 Index to union")
228 if len(indexes) == 1:
229 result = indexes[0]
230 if isinstance(result, list):
231 if not sort or sort is lib.no_default:
232 result = Index(result)
233 else:
234 result = Index(sorted(result))
235 return result, True
237 indexes, kind = _sanitize_and_check(indexes)
239 if kind == "special":
240 result = indexes[0]
242 num_dtis = 0
243 num_dti_tzs = 0
244 for idx in indexes:
245 if isinstance(idx, DatetimeIndex):
246 num_dtis += 1
247 if idx.tz is not None:
248 num_dti_tzs += 1
249 if num_dti_tzs not in [0, num_dtis]:
250 # TODO: this behavior is not tested (so may not be desired),
251 # but is kept in order to keep behavior the same when
252 # deprecating union_many
253 # test_frame_from_dict_with_mixed_indexes
254 raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")
256 if num_dtis == len(indexes):
257 if sort is lib.no_default:
258 sort = True
259 result = indexes[0]
261 elif num_dtis > 1:
262 # If we have mixed timezones, our casting behavior may depend on
263 # the order of indexes, which we don't want.
264 sort = False
266 # TODO: what about Categorical[dt64]?
267 # test_frame_from_dict_with_mixed_indexes
268 indexes = [x.astype(object, copy=False) for x in indexes]
269 result = indexes[0]
271 for other in indexes[1:]:
272 result = result.union(other, sort=None if sort else False)
273 return result, False
275 elif kind == "array":
276 all_equal = all_indexes_same(indexes)
277 if not all_equal:
278 dtype = find_common_type([idx.dtype for idx in indexes])
279 inds = [ind.astype(dtype, copy=False) for ind in indexes]
280 index = inds[0].unique()
281 other = inds[1].append(inds[2:])
282 diff = other[index.get_indexer_for(other) == -1]
283 if len(diff):
284 index = index.append(diff.unique())
285 if sort:
286 index = index.sort_values()
287 else:
288 index = indexes[0]
290 name = get_unanimous_names(*indexes)[0]
291 if name != index.name:
292 index = index.rename(name)
293 return index, all_equal
295 elif kind == "list":
296 dtypes = [idx.dtype for idx in indexes if isinstance(idx, Index)]
297 if dtypes:
298 dtype = find_common_type(dtypes)
299 else:
300 dtype = None
301 all_lists = (idx.tolist() if isinstance(idx, Index) else idx for idx in indexes)
302 result = Index(
303 lib.fast_unique_multiple_list_gen(all_lists, sort=bool(sort)),
304 dtype=dtype,
305 )
306 return result, False
307 else:
308 raise ValueError(f"{kind=} must be 'special', 'array' or 'list'.")
311def _sanitize_and_check(indexes):
312 """
313 Verify the type of indexes and convert lists to Index.
315 Cases:
317 - [list, list, ...]: Return ([list, list, ...], 'list')
318 - [list, Index, ...]: Return _sanitize_and_check([Index, Index, ...])
319 Lists are sorted and converted to Index.
320 - [Index, Index, ...]: Return ([Index, Index, ...], TYPE)
321 TYPE = 'special' if at least one special type, 'array' otherwise.
323 Parameters
324 ----------
325 indexes : list of Index or list objects
327 Returns
328 -------
329 sanitized_indexes : list of Index or list objects
330 type : {'list', 'array', 'special'}
331 """
332 kinds = {type(index) for index in indexes}
334 if list in kinds:
335 if len(kinds) > 1:
336 indexes = [
337 Index(list(x)) if not isinstance(x, Index) else x for x in indexes
338 ]
339 kinds -= {list}
340 else:
341 return indexes, "list"
343 if len(kinds) > 1 or Index not in kinds:
344 return indexes, "special"
345 else:
346 return indexes, "array"
349def all_indexes_same(indexes) -> bool:
350 """
351 Determine if all indexes contain the same elements.
353 Parameters
354 ----------
355 indexes : iterable of Index objects
357 Returns
358 -------
359 bool
360 True if all indexes contain the same elements, False otherwise.
361 """
362 itr = iter(indexes)
363 first = next(itr)
364 return all(first.equals(index) for index in itr)
367def default_index(n: int) -> RangeIndex:
368 rng = range(n)
369 return RangeIndex._simple_new(rng, name=None)