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