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 TYPE_CHECKING
5import numpy as np
7from pandas._libs import (
8 algos as libalgos,
9 internals as libinternals,
10 lib,
11)
12from pandas.util._decorators import cache_readonly
14from pandas.core.dtypes.cast import (
15 ensure_dtype_can_hold_na,
16 find_common_type,
17)
18from pandas.core.dtypes.common import is_1d_only_ea_dtype
19from pandas.core.dtypes.concat import (
20 concat_compat,
21 union_categories_compat,
22)
23from pandas.core.dtypes.dtypes import (
24 CategoricalDtype,
25 ExtensionDtype,
26)
28from pandas.core.arrays import Categorical
29from pandas.core.construction import ensure_wrapped_if_datetimelike
30from pandas.core.internals.blocks import (
31 ensure_block_shape,
32 new_block_2d,
33)
34from pandas.core.internals.managers import (
35 BlockManager,
36 make_na_array,
37)
39if TYPE_CHECKING:
40 from collections.abc import (
41 Generator,
42 Sequence,
43 )
45 from pandas._typing import (
46 ArrayLike,
47 AxisInt,
48 DtypeObj,
49 Shape,
50 )
52 from pandas import Index
53 from pandas.core.internals.blocks import (
54 Block,
55 BlockPlacement,
56 )
59def concatenate_managers(
60 mgrs_indexers,
61 axes: list[Index],
62 concat_axis: AxisInt,
63 copy: bool,
64 union_categories: bool = False,
65) -> BlockManager:
66 """
67 Concatenate block managers into one.
69 Parameters
70 ----------
71 mgrs_indexers : list of (BlockManager, {axis: indexer,...}) tuples
72 axes : list of Index
73 concat_axis : int
74 copy : bool
75 union_categories : bool, default False
76 If True, union the categories of categorical blocks being concatenated
77 rather than falling back to a non-categorical dtype.
79 Returns
80 -------
81 BlockManager
82 """
84 needs_copy = copy and concat_axis == 0
86 # Assertions disabled for performance
87 # for tup in mgrs_indexers:
88 # # caller is responsible for ensuring this
89 # indexers = tup[1]
90 # assert concat_axis not in indexers
92 if concat_axis == 0:
93 mgrs = _maybe_reindex_columns_na_proxy(axes, mgrs_indexers, needs_copy)
94 return mgrs[0].concat_horizontal(mgrs, axes)
96 if len(mgrs_indexers) > 0 and mgrs_indexers[0][0].nblocks > 0:
97 first_dtype = mgrs_indexers[0][0].blocks[0].dtype
98 if first_dtype in [np.float64, np.float32]:
99 # TODO: support more dtypes here. This will be simpler once
100 # JoinUnit.is_na behavior is deprecated.
101 # (update 2024-04-13 that deprecation has been enforced)
102 if (
103 all(_is_homogeneous_mgr(mgr, first_dtype) for mgr, _ in mgrs_indexers)
104 and len(mgrs_indexers) > 1
105 ):
106 # Fastpath!
107 # Length restriction is just to avoid having to worry about 'copy'
108 shape = tuple(len(x) for x in axes)
109 nb = _concat_homogeneous_fastpath(mgrs_indexers, shape, first_dtype)
110 return BlockManager((nb,), axes)
112 mgrs = _maybe_reindex_columns_na_proxy(axes, mgrs_indexers, needs_copy)
114 if len(mgrs) == 1:
115 mgr = mgrs[0]
116 out = mgr.copy(deep=False)
117 out.axes = axes
118 return out
120 blocks = []
121 values: ArrayLike
123 for placement, join_units in _get_combined_plan(mgrs):
124 unit = join_units[0]
125 blk = unit.block
127 if _is_uniform_join_units(join_units):
128 vals = [ju.block.values for ju in join_units]
130 if not blk.is_extension:
131 # _is_uniform_join_units ensures a single dtype, so
132 # we can use np.concatenate, which is more performant
133 # than concat_compat
134 # error: Argument 1 to "concatenate" has incompatible type
135 # "List[Union[ndarray[Any, Any], ExtensionArray]]";
136 # expected "Union[_SupportsArray[dtype[Any]],
137 # _NestedSequence[_SupportsArray[dtype[Any]]]]"
138 values = np.concatenate(vals, axis=1) # type: ignore[arg-type]
139 elif is_1d_only_ea_dtype(blk.dtype):
140 # TODO(EA2D): special-casing not needed with 2D EAs
141 values = concat_compat(
142 vals,
143 axis=0,
144 ea_compat_axis=True,
145 union_categories=union_categories,
146 )
147 values = ensure_block_shape(values, ndim=2)
148 else:
149 values = concat_compat(vals, axis=1, union_categories=union_categories)
151 values = ensure_wrapped_if_datetimelike(values)
153 fastpath = blk.values.dtype == values.dtype
154 else:
155 values = _concatenate_join_units(
156 join_units, copy=copy, union_categories=union_categories
157 )
158 fastpath = False
160 if fastpath:
161 b = blk.make_block_same_class(values, placement=placement)
162 else:
163 b = new_block_2d(values, placement=placement)
165 blocks.append(b)
167 return BlockManager(tuple(blocks), axes)
170def _maybe_reindex_columns_na_proxy(
171 axes: list[Index],
172 mgrs_indexers: list[tuple[BlockManager, dict[int, np.ndarray]]],
173 needs_copy: bool,
174) -> list[BlockManager]:
175 """
176 Reindex along columns so that all of the BlockManagers being concatenated
177 have matching columns.
179 Columns added in this reindexing have dtype=np.void, indicating they
180 should be ignored when choosing a column's final dtype.
181 """
182 new_mgrs = []
184 for mgr, indexers in mgrs_indexers:
185 # For axis=0 (i.e. columns) we use_na_proxy and only_slice, so this
186 # is a cheap reindexing.
187 for i, indexer in indexers.items():
188 mgr = mgr.reindex_indexer(
189 axes[i],
190 indexer,
191 axis=i,
192 only_slice=True, # only relevant for i==0
193 allow_dups=True,
194 use_na_proxy=True, # only relevant for i==0
195 )
196 if needs_copy and not indexers:
197 mgr = mgr.copy(deep=True)
199 new_mgrs.append(mgr)
200 return new_mgrs
203def _is_homogeneous_mgr(mgr: BlockManager, first_dtype: DtypeObj) -> bool:
204 """
205 Check if this Manager can be treated as a single ndarray.
206 """
207 if mgr.nblocks != 1:
208 return False
209 blk = mgr.blocks[0]
210 if not (blk.mgr_locs.is_slice_like and blk.mgr_locs.as_slice.step == 1):
211 return False
213 return blk.dtype == first_dtype
216def _concat_homogeneous_fastpath(
217 mgrs_indexers, shape: Shape, first_dtype: np.dtype
218) -> Block:
219 """
220 With single-Block managers with homogeneous dtypes (that can already hold nan),
221 we avoid [...]
222 """
223 # assumes
224 # all(_is_homogeneous_mgr(mgr, first_dtype) for mgr, _ in mgrs_indexers)
226 if all(not indexers for _, indexers in mgrs_indexers):
227 # https://github.com/pandas-dev/pandas/pull/52685#issuecomment-1523287739
228 arrs = [mgr.blocks[0].values.T for mgr, _ in mgrs_indexers]
229 arr = np.concatenate(arrs).T
230 bp = libinternals.BlockPlacement(slice(shape[0]))
231 nb = new_block_2d(arr, bp)
232 return nb
234 arr = np.empty(shape, dtype=first_dtype)
236 if first_dtype == np.float64:
237 take_func = libalgos.take_2d_axis0_float64_float64
238 else:
239 take_func = libalgos.take_2d_axis0_float32_float32
241 start = 0
242 for mgr, indexers in mgrs_indexers:
243 mgr_len = mgr.shape[1]
244 end = start + mgr_len
246 if 0 in indexers:
247 take_func(
248 mgr.blocks[0].values,
249 indexers[0],
250 arr[:, start:end],
251 )
252 else:
253 # No reindexing necessary, we can copy values directly
254 arr[:, start:end] = mgr.blocks[0].values
256 start += mgr_len
258 bp = libinternals.BlockPlacement(slice(shape[0]))
259 nb = new_block_2d(arr, bp)
260 return nb
263def _get_combined_plan(
264 mgrs: list[BlockManager],
265) -> Generator[tuple[BlockPlacement, list[JoinUnit]]]:
266 max_len = mgrs[0].shape[0]
268 blknos_list = [mgr.blknos for mgr in mgrs]
269 pairs = libinternals.get_concat_blkno_indexers(blknos_list)
270 for blknos, bp in pairs:
271 # assert bp.is_slice_like
272 # assert len(bp) > 0
274 units_for_bp = []
275 for k, mgr in enumerate(mgrs):
276 blkno = blknos[k]
278 nb = _get_block_for_concat_plan(mgr, bp, blkno, max_len=max_len)
279 unit = JoinUnit(nb)
280 units_for_bp.append(unit)
282 yield bp, units_for_bp
285def _get_block_for_concat_plan(
286 mgr: BlockManager, bp: BlockPlacement, blkno: int, *, max_len: int
287) -> Block:
288 blk = mgr.blocks[blkno]
289 # Assertions disabled for performance:
290 # assert bp.is_slice_like
291 # assert blkno != -1
292 # assert (mgr.blknos[bp] == blkno).all()
294 if len(bp) == len(blk.mgr_locs) and (
295 blk.mgr_locs.is_slice_like and blk.mgr_locs.as_slice.step == 1
296 ):
297 nb = blk
298 else:
299 ax0_blk_indexer = mgr.blklocs[bp.indexer]
301 slc = lib.maybe_indices_to_slice(ax0_blk_indexer, max_len)
302 # TODO: in all extant test cases 2023-04-08 we have a slice here.
303 # Will this always be the case?
304 if isinstance(slc, slice):
305 nb = blk.slice_block_columns(slc)
306 else:
307 nb = blk.take_block_columns(slc)
309 # assert nb.shape == (len(bp), mgr.shape[1])
310 return nb
313class JoinUnit:
314 def __init__(self, block: Block) -> None:
315 self.block = block
317 def __repr__(self) -> str:
318 return f"{type(self).__name__}({self.block!r})"
320 @cache_readonly
321 def is_na(self) -> bool:
322 return self.block.dtype.kind == "V"
324 def get_reindexed_values(self, empty_dtype: DtypeObj, upcasted_na) -> ArrayLike:
325 if self.block.dtype.kind == "V":
326 return make_na_array(empty_dtype, self.block.shape, upcasted_na)
327 return self.block.values
330def _concatenate_join_units(
331 join_units: list[JoinUnit],
332 copy: bool,
333 union_categories: bool = False,
334) -> ArrayLike:
335 """
336 Concatenate values from several join units along axis=1.
337 """
338 empty_dtype = _get_empty_dtype(join_units, union_categories=union_categories)
340 has_none_blocks = any(unit.block.dtype.kind == "V" for unit in join_units)
341 upcasted_na = _dtype_to_na_value(empty_dtype, has_none_blocks)
343 to_concat = [
344 ju.get_reindexed_values(empty_dtype=empty_dtype, upcasted_na=upcasted_na)
345 for ju in join_units
346 ]
348 if any(is_1d_only_ea_dtype(t.dtype) for t in to_concat):
349 # TODO(EA2D): special case not needed if all EAs used HybridBlocks
351 # error: No overload variant of "__getitem__" of "ExtensionArray" matches
352 # argument type "Tuple[int, slice]"
353 to_concat = [
354 t if is_1d_only_ea_dtype(t.dtype) else t[0, :] # type: ignore[call-overload]
355 for t in to_concat
356 ]
357 concat_values = concat_compat(
358 to_concat, axis=0, ea_compat_axis=True, union_categories=union_categories
359 )
360 concat_values = ensure_block_shape(concat_values, 2)
362 else:
363 concat_values = concat_compat(
364 to_concat, axis=1, union_categories=union_categories
365 )
367 return concat_values
370def _dtype_to_na_value(dtype: DtypeObj, has_none_blocks: bool):
371 """
372 Find the NA value to go with this dtype.
373 """
374 if isinstance(dtype, ExtensionDtype):
375 return dtype.na_value
376 elif dtype.kind in "mM":
377 return dtype.type("NaT", np.datetime_data(dtype)[0])
378 elif dtype.kind in "fc":
379 return dtype.type("NaN")
380 elif dtype.kind == "b":
381 # different from missing.na_value_for_dtype
382 return None
383 elif dtype.kind in "iu":
384 if not has_none_blocks:
385 # different from missing.na_value_for_dtype
386 return None
387 return np.nan
388 elif dtype.kind == "O":
389 return np.nan
390 raise NotImplementedError
393def _get_empty_dtype(
394 join_units: Sequence[JoinUnit], union_categories: bool = False
395) -> DtypeObj:
396 """
397 Return dtype and N/A values to use when concatenating specified units.
399 Returned N/A value may be None which means there was no casting involved.
401 Returns
402 -------
403 dtype
404 """
405 if lib.dtypes_all_equal([ju.block.dtype for ju in join_units]):
406 empty_dtype = join_units[0].block.dtype
407 return empty_dtype
409 has_none_blocks = any(unit.block.dtype.kind == "V" for unit in join_units)
411 dtypes = [unit.block.dtype for unit in join_units if not unit.is_na]
413 if union_categories and has_none_blocks and dtypes:
414 if all(isinstance(dt, CategoricalDtype) for dt in dtypes):
415 # GH#14177 give the all-NA filler for missing columns the unioned
416 # categorical dtype so that concat_compat sees all-categorical
417 # inputs and unions instead of falling back to the common dtype.
418 empties = [Categorical([], dtype=dt) for dt in dtypes]
419 return union_categories_compat(empties).dtype
421 dtype = find_common_type(dtypes)
422 if has_none_blocks:
423 dtype = ensure_dtype_can_hold_na(dtype)
425 return dtype
428def _is_uniform_join_units(join_units: list[JoinUnit]) -> bool:
429 """
430 Check if the join units consist of blocks of uniform type that can
431 be concatenated using Block.concat_same_type instead of the generic
432 _concatenate_join_units (which uses `concat_compat`).
434 """
435 first = join_units[0].block
436 if first.dtype.kind == "V":
437 return False
438 return (
439 # exclude cases where a) ju.block is None or b) we have e.g. Int64+int64
440 all(type(ju.block) is type(first) for ju in join_units)
441 and
442 # e.g. DatetimeLikeBlock can be dt64 or td64, but these are not uniform
443 all(
444 ju.block.dtype == first.dtype
445 # GH#42092 we only want the dtype_equal check for non-numeric blocks
446 # (for now, may change but that would need a deprecation)
447 or ju.block.dtype.kind in "iub"
448 for ju in join_units
449 )
450 and
451 # no blocks that would get missing values (can lead to type upcasts)
452 # unless we're an extension dtype.
453 all(not ju.is_na or ju.block.is_extension for ju in join_units)
454 )