1"""
2Utility functions related to concat.
3"""
4
5from __future__ import annotations
6
7from typing import (
8 TYPE_CHECKING,
9 cast,
10)
11
12import numpy as np
13
14from pandas._libs import lib
15from pandas.util._decorators import set_module
16
17from pandas.core.dtypes.astype import astype_array
18from pandas.core.dtypes.cast import (
19 common_dtype_categorical_compat,
20 find_common_type,
21 np_find_common_type,
22)
23from pandas.core.dtypes.dtypes import CategoricalDtype
24from pandas.core.dtypes.generic import (
25 ABCCategoricalIndex,
26 ABCSeries,
27)
28
29if TYPE_CHECKING:
30 from collections.abc import Sequence
31
32 from pandas._typing import (
33 ArrayLike,
34 AxisInt,
35 DtypeObj,
36 )
37
38 from pandas import (
39 CategoricalIndex,
40 Series,
41 )
42 from pandas.core.arrays import (
43 Categorical,
44 ExtensionArray,
45 )
46
47
48def _is_nonempty(x: ArrayLike, axis: AxisInt) -> bool:
49 # filter empty arrays
50 # 1-d dtypes always are included here
51 if x.ndim <= axis:
52 return True
53 return x.shape[axis] > 0
54
55
56def concat_compat(
57 to_concat: Sequence[ArrayLike],
58 axis: AxisInt = 0,
59 ea_compat_axis: bool = False,
60 union_categories: bool = False,
61) -> ArrayLike:
62 """
63 provide concatenation of an array of arrays each of which is a single
64 'normalized' dtypes (in that for example, if it's object, then it is a
65 non-datetimelike and provide a combined dtype for the resulting array that
66 preserves the overall dtype if possible)
67
68 Parameters
69 ----------
70 to_concat : sequence of arrays
71 axis : axis to provide concatenation
72 ea_compat_axis : bool, default False
73 For ExtensionArray compat, behave as if axis == 1 when determining
74 whether to drop empty arrays.
75 union_categories : bool, default False
76 If True and all arrays are Categorical, union the categories
77 instead of casting to the underlying dtype.
78
79 Returns
80 -------
81 a single array, preserving the combined dtypes
82 """
83 if (
84 union_categories
85 and len(to_concat)
86 and all(isinstance(x.dtype, CategoricalDtype) for x in to_concat)
87 ):
88 return union_categories_compat(cast("Sequence[Categorical]", to_concat))
89
90 if len(to_concat) and lib.dtypes_all_equal([obj.dtype for obj in to_concat]):
91 # fastpath!
92 obj = to_concat[0]
93 if isinstance(obj, np.ndarray):
94 to_concat_arrs = cast("Sequence[np.ndarray]", to_concat)
95 return np.concatenate(to_concat_arrs, axis=axis)
96
97 to_concat_eas = cast("Sequence[ExtensionArray]", to_concat)
98 if ea_compat_axis:
99 # We have 1D objects, that don't support axis keyword
100 return obj._concat_same_type(to_concat_eas)
101 elif axis == 0:
102 return obj._concat_same_type(to_concat_eas)
103 else:
104 # e.g. DatetimeArray
105 # NB: We are assuming here that ensure_wrapped_if_arraylike has
106 # been called where relevant.
107 return obj._concat_same_type(
108 # error: Unexpected keyword argument "axis" for "_concat_same_type"
109 # of "ExtensionArray"
110 to_concat_eas,
111 axis=axis, # type: ignore[call-arg]
112 )
113
114 # If all arrays are empty, there's nothing to convert, just short-cut to
115 # the concatenation, #3121.
116 #
117 # Creating an empty array directly is tempting, but the winnings would be
118 # marginal given that it would still require shape & dtype calculation and
119 # np.concatenate which has them both implemented is compiled.
120 non_empties = [x for x in to_concat if _is_nonempty(x, axis)]
121
122 any_ea, kinds, target_dtype = _get_result_dtype(to_concat, non_empties)
123
124 if target_dtype is not None:
125 to_concat = [astype_array(arr, target_dtype, copy=False) for arr in to_concat]
126
127 if not isinstance(to_concat[0], np.ndarray):
128 # i.e. isinstance(to_concat[0], ExtensionArray)
129 to_concat_eas = cast("Sequence[ExtensionArray]", to_concat)
130 cls = type(to_concat[0])
131 # GH#53640: eg. for datetime array, axis=1 but 0 is default
132 # However, class method `_concat_same_type()` for some classes
133 # may not support the `axis` keyword
134 if ea_compat_axis or axis == 0:
135 return cls._concat_same_type(to_concat_eas)
136 else:
137 return cls._concat_same_type(
138 to_concat_eas,
139 axis=axis, # type: ignore[call-arg]
140 )
141 else:
142 to_concat_arrs = cast("Sequence[np.ndarray]", to_concat)
143 result = np.concatenate(to_concat_arrs, axis=axis)
144
145 if not any_ea and "b" in kinds and result.dtype.kind in "iuf":
146 # GH#39817 cast to object instead of casting bools to numeric
147 result = result.astype(object, copy=False)
148 return result
149
150
151def union_categories_compat(to_union: Sequence[Categorical]) -> Categorical:
152 """
153 union_categoricals for concat(union_categories=True).
154
155 Unlike union_categoricals, categories with differing dtypes are cast to a
156 common dtype instead of raising, so that an all-categorical concatenation
157 always returns a Categorical. Orderedness is preserved only if every input
158 shares the same dtype after this cast.
159 """
160 from pandas import Categorical
161 from pandas.core.arrays.categorical import recode_for_categories
162
163 if not lib.dtypes_all_equal([x.categories.dtype for x in to_union]):
164 # GH#14177 e.g. int64 vs float64 categories. Empty categories carry no
165 # dtype information, so keep them out of the common dtype.
166 dtypes = [x.categories.dtype for x in to_union if len(x.categories)]
167 common = find_common_type(dtypes or [x.categories.dtype for x in to_union])
168
169 recast = []
170 for obj in to_union:
171 cats = obj.categories.astype(common)
172 codes = obj._codes
173 if not cats.is_unique:
174 # the cast collapsed distinct categories, e.g. int64 values
175 # that are not exactly representable as float64
176 uniques = cats.unique()
177 codes = recode_for_categories(codes, cats, uniques)
178 cats = uniques
179 recast.append(
180 Categorical._simple_new(codes, CategoricalDtype(cats, obj.ordered))
181 )
182 to_union = recast
183
184 ignore_order = not lib.dtypes_all_equal([x.dtype for x in to_union])
185 return union_categoricals(to_union, ignore_order=ignore_order)
186
187
188def _get_result_dtype(
189 to_concat: Sequence[ArrayLike], non_empties: Sequence[ArrayLike]
190) -> tuple[bool, set[str], DtypeObj | None]:
191 target_dtype = None
192
193 dtypes = {obj.dtype for obj in to_concat}
194 kinds = {obj.dtype.kind for obj in to_concat}
195
196 any_ea = any(not isinstance(x, np.ndarray) for x in to_concat)
197 if any_ea:
198 # i.e. any ExtensionArrays
199
200 # we ignore axis here, as internally concatting with EAs is always
201 # for axis=0
202 if len(dtypes) != 1:
203 target_dtype = find_common_type([x.dtype for x in to_concat])
204 target_dtype = common_dtype_categorical_compat(to_concat, target_dtype)
205
206 elif not len(non_empties):
207 # we have all empties, but may need to coerce the result dtype to
208 # object if we have non-numeric type operands (numpy would otherwise
209 # cast this to float)
210 if len(kinds) != 1:
211 if not len(kinds - {"i", "u", "f"}) or not len(kinds - {"b", "i", "u"}):
212 # let numpy coerce
213 pass
214 else:
215 # coerce to object
216 target_dtype = np.dtype(object)
217 kinds = {"o"}
218 elif "b" in kinds and len(kinds) > 1:
219 # GH#21108, GH#45101
220 target_dtype = np.dtype(object)
221 kinds = {"o"}
222 else:
223 # error: Argument 1 to "np_find_common_type" has incompatible type
224 # "*Set[Union[ExtensionDtype, Any]]"; expected "dtype[Any]"
225 target_dtype = np_find_common_type(*dtypes) # type: ignore[arg-type]
226
227 return any_ea, kinds, target_dtype
228
229
230@set_module("pandas.api.types")
231def union_categoricals(
232 to_union: Sequence[CategoricalIndex | Series | Categorical],
233 sort_categories: bool = False,
234 ignore_order: bool = False,
235) -> Categorical:
236 """
237 Combine list-like of Categorical-like, unioning categories.
238
239 All categories must have the same dtype.
240
241 Parameters
242 ----------
243 to_union : list-like
244 Categorical, CategoricalIndex, or Series with dtype='category'.
245 sort_categories : bool, default False
246 If true, resulting categories will be lexsorted, otherwise
247 they will be ordered as they appear in the data.
248 ignore_order : bool, default False
249 If true, the ordered attribute of the Categoricals will be ignored.
250 Results in an unordered categorical.
251
252 Returns
253 -------
254 Categorical
255 The union of categories being combined.
256
257 Raises
258 ------
259 TypeError
260 - all inputs do not have the same dtype
261 - all inputs do not have the same ordered property
262 - all inputs are ordered and their categories are not identical
263 - sort_categories=True and Categoricals are ordered
264 ValueError
265 Empty list of categoricals passed
266
267 See Also
268 --------
269 CategoricalDtype : Type for categorical data with the categories and orderedness.
270 Categorical : Represent a categorical variable in classic R / S-plus fashion.
271
272 Notes
273 -----
274 To learn more about categories, see `link
275 <https://pandas.pydata.org/pandas-docs/stable/user_guide/categorical.html#unioning>`__
276
277 Examples
278 --------
279 If you want to combine categoricals that do not necessarily have
280 the same categories, `union_categoricals` will combine a list-like
281 of categoricals. The new categories will be the union of the
282 categories being combined.
283
284 >>> a = pd.Categorical(["b", "c"])
285 >>> b = pd.Categorical(["a", "b"])
286 >>> pd.api.types.union_categoricals([a, b])
287 ['b', 'c', 'a', 'b']
288 Categories (3, str): ['b', 'c', 'a']
289
290 By default, the resulting categories will be ordered as they appear
291 in the `categories` of the data. If you want the categories to be
292 lexsorted, use `sort_categories=True` argument.
293
294 >>> pd.api.types.union_categoricals([a, b], sort_categories=True)
295 ['b', 'c', 'a', 'b']
296 Categories (3, str): ['a', 'b', 'c']
297
298 `union_categoricals` also works with the case of combining two
299 categoricals of the same categories and order information (e.g. what
300 you could also `append` for).
301
302 >>> a = pd.Categorical(["a", "b"], ordered=True)
303 >>> b = pd.Categorical(["a", "b", "a"], ordered=True)
304 >>> pd.api.types.union_categoricals([a, b])
305 ['a', 'b', 'a', 'b', 'a']
306 Categories (2, str): ['a' < 'b']
307
308 Raises `TypeError` because the categories are ordered and not identical.
309
310 >>> a = pd.Categorical(["a", "b"], ordered=True)
311 >>> b = pd.Categorical(["a", "b", "c"], ordered=True)
312 >>> pd.api.types.union_categoricals([a, b])
313 Traceback (most recent call last):
314 ...
315 TypeError: to union ordered Categoricals, all categories must be the same
316
317 Ordered categoricals with different categories or orderings can be
318 combined by using the `ignore_order=True` argument.
319
320 >>> a = pd.Categorical(["a", "b", "c"], ordered=True)
321 >>> b = pd.Categorical(["c", "b", "a"], ordered=True)
322 >>> pd.api.types.union_categoricals([a, b], ignore_order=True)
323 ['a', 'b', 'c', 'c', 'b', 'a']
324 Categories (3, str): ['a', 'b', 'c']
325
326 `union_categoricals` also works with a `CategoricalIndex`, or `Series`
327 containing categorical data, but note that the resulting array will
328 always be a plain `Categorical`
329
330 >>> a = pd.Series(["b", "c"], dtype="category")
331 >>> b = pd.Series(["a", "b"], dtype="category")
332 >>> pd.api.types.union_categoricals([a, b])
333 ['b', 'c', 'a', 'b']
334 Categories (3, str): ['b', 'c', 'a']
335 """
336 from pandas import Categorical
337 from pandas.core.arrays.categorical import recode_for_categories
338
339 if len(to_union) == 0:
340 raise ValueError("No Categoricals to union")
341
342 def _maybe_unwrap(
343 x: CategoricalIndex | Series | Categorical,
344 ) -> Categorical:
345 if isinstance(x, (ABCCategoricalIndex, ABCSeries)):
346 return x._values
347 elif isinstance(x, Categorical):
348 return x
349 else:
350 raise TypeError("all components to combine must be Categorical")
351
352 to_union = cast("list[Categorical]", [_maybe_unwrap(x) for x in to_union])
353 first = to_union[0]
354
355 if not lib.dtypes_all_equal([obj.categories.dtype for obj in to_union]):
356 raise TypeError("dtype of categories must be the same")
357
358 ordered: bool | None = False
359 if all(first._categories_match_up_to_permutation(other) for other in to_union[1:]):
360 # identical categories - fastpath
361 categories = first.categories
362 ordered = first.ordered
363
364 all_codes = [first._encode_with_my_categories(x)._codes for x in to_union]
365 new_codes = np.concatenate(all_codes)
366
367 if sort_categories and not ignore_order and ordered:
368 raise TypeError("Cannot use sort_categories=True with ordered Categoricals")
369
370 if sort_categories and not categories.is_monotonic_increasing:
371 categories = categories.sort_values()
372 indexer = categories.get_indexer(first.categories)
373
374 from pandas.core.algorithms import take_nd
375
376 new_codes = take_nd(indexer, new_codes, fill_value=-1)
377 elif ignore_order or all(not c.ordered for c in to_union):
378 # different categories - union and recode
379 cats = first.categories.append([c.categories for c in to_union[1:]])
380 categories = cats.unique()
381 if sort_categories:
382 categories = categories.sort_values()
383
384 all_codes = [
385 recode_for_categories(c.codes, c.categories, categories, copy=False)
386 for c in to_union
387 ]
388 new_codes = np.concatenate(all_codes)
389 else:
390 # ordered - to show a proper error message
391 if all(c.ordered for c in to_union):
392 msg = "to union ordered Categoricals, all categories must be the same"
393 raise TypeError(msg)
394 raise TypeError("Categorical.ordered must be the same")
395
396 if ignore_order:
397 ordered = False
398
399 dtype = CategoricalDtype(categories=categories, ordered=ordered)
400 return Categorical._simple_new(new_codes, dtype=dtype)