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 algos as libalgos,
13 internals as libinternals,
14 lib,
15)
16from pandas._libs.missing import NA
17from pandas.util._decorators import cache_readonly
19from pandas.core.dtypes.cast import (
20 ensure_dtype_can_hold_na,
21 find_common_type,
22)
23from pandas.core.dtypes.common import (
24 is_1d_only_ea_dtype,
25 needs_i8_conversion,
26)
27from pandas.core.dtypes.concat import concat_compat
28from pandas.core.dtypes.dtypes import ExtensionDtype
29from pandas.core.dtypes.missing import is_valid_na_for_dtype
31from pandas.core.construction import ensure_wrapped_if_datetimelike
32from pandas.core.internals.blocks import (
33 ensure_block_shape,
34 new_block_2d,
35)
36from pandas.core.internals.managers import (
37 BlockManager,
38 make_na_array,
39)
41if TYPE_CHECKING:
42 from collections.abc import (
43 Generator,
44 Sequence,
45 )
47 from pandas._typing import (
48 ArrayLike,
49 AxisInt,
50 DtypeObj,
51 Shape,
52 )
54 from pandas import Index
55 from pandas.core.internals.blocks import (
56 Block,
57 BlockPlacement,
58 )
61def concatenate_managers(
62 mgrs_indexers, axes: list[Index], concat_axis: AxisInt, copy: bool
63) -> BlockManager:
64 """
65 Concatenate block managers into one.
67 Parameters
68 ----------
69 mgrs_indexers : list of (BlockManager, {axis: indexer,...}) tuples
70 axes : list of Index
71 concat_axis : int
72 copy : bool
74 Returns
75 -------
76 BlockManager
77 """
79 needs_copy = copy and concat_axis == 0
81 # Assertions disabled for performance
82 # for tup in mgrs_indexers:
83 # # caller is responsible for ensuring this
84 # indexers = tup[1]
85 # assert concat_axis not in indexers
87 if concat_axis == 0:
88 mgrs = _maybe_reindex_columns_na_proxy(axes, mgrs_indexers, needs_copy)
89 return mgrs[0].concat_horizontal(mgrs, axes)
91 if len(mgrs_indexers) > 0 and mgrs_indexers[0][0].nblocks > 0:
92 first_dtype = mgrs_indexers[0][0].blocks[0].dtype
93 if first_dtype in [np.float64, np.float32]:
94 # TODO: support more dtypes here. This will be simpler once
95 # JoinUnit.is_na behavior is deprecated.
96 # (update 2024-04-13 that deprecation has been enforced)
97 if (
98 all(_is_homogeneous_mgr(mgr, first_dtype) for mgr, _ in mgrs_indexers)
99 and len(mgrs_indexers) > 1
100 ):
101 # Fastpath!
102 # Length restriction is just to avoid having to worry about 'copy'
103 shape = tuple(len(x) for x in axes)
104 nb = _concat_homogeneous_fastpath(mgrs_indexers, shape, first_dtype)
105 return BlockManager((nb,), axes)
107 mgrs = _maybe_reindex_columns_na_proxy(axes, mgrs_indexers, needs_copy)
109 if len(mgrs) == 1:
110 mgr = mgrs[0]
111 out = mgr.copy(deep=False)
112 out.axes = axes
113 return out
115 blocks = []
116 values: ArrayLike
118 for placement, join_units in _get_combined_plan(mgrs):
119 unit = join_units[0]
120 blk = unit.block
122 if _is_uniform_join_units(join_units):
123 vals = [ju.block.values for ju in join_units]
125 if not blk.is_extension:
126 # _is_uniform_join_units ensures a single dtype, so
127 # we can use np.concatenate, which is more performant
128 # than concat_compat
129 # error: Argument 1 to "concatenate" has incompatible type
130 # "List[Union[ndarray[Any, Any], ExtensionArray]]";
131 # expected "Union[_SupportsArray[dtype[Any]],
132 # _NestedSequence[_SupportsArray[dtype[Any]]]]"
133 values = np.concatenate(vals, axis=1) # type: ignore[arg-type]
134 elif is_1d_only_ea_dtype(blk.dtype):
135 # TODO(EA2D): special-casing not needed with 2D EAs
136 values = concat_compat(vals, axis=0, ea_compat_axis=True)
137 values = ensure_block_shape(values, ndim=2)
138 else:
139 values = concat_compat(vals, axis=1)
141 values = ensure_wrapped_if_datetimelike(values)
143 fastpath = blk.values.dtype == values.dtype
144 else:
145 values = _concatenate_join_units(join_units, copy=copy)
146 fastpath = False
148 if fastpath:
149 b = blk.make_block_same_class(values, placement=placement)
150 else:
151 b = new_block_2d(values, placement=placement)
153 blocks.append(b)
155 return BlockManager(tuple(blocks), axes)
158def _maybe_reindex_columns_na_proxy(
159 axes: list[Index],
160 mgrs_indexers: list[tuple[BlockManager, dict[int, np.ndarray]]],
161 needs_copy: bool,
162) -> list[BlockManager]:
163 """
164 Reindex along columns so that all of the BlockManagers being concatenated
165 have matching columns.
167 Columns added in this reindexing have dtype=np.void, indicating they
168 should be ignored when choosing a column's final dtype.
169 """
170 new_mgrs = []
172 for mgr, indexers in mgrs_indexers:
173 # For axis=0 (i.e. columns) we use_na_proxy and only_slice, so this
174 # is a cheap reindexing.
175 for i, indexer in indexers.items():
176 mgr = mgr.reindex_indexer(
177 axes[i],
178 indexer,
179 axis=i,
180 only_slice=True, # only relevant for i==0
181 allow_dups=True,
182 use_na_proxy=True, # only relevant for i==0
183 )
184 if needs_copy and not indexers:
185 mgr = mgr.copy(deep=True)
187 new_mgrs.append(mgr)
188 return new_mgrs
191def _is_homogeneous_mgr(mgr: BlockManager, first_dtype: DtypeObj) -> bool:
192 """
193 Check if this Manager can be treated as a single ndarray.
194 """
195 if mgr.nblocks != 1:
196 return False
197 blk = mgr.blocks[0]
198 if not (blk.mgr_locs.is_slice_like and blk.mgr_locs.as_slice.step == 1):
199 return False
201 return blk.dtype == first_dtype
204def _concat_homogeneous_fastpath(
205 mgrs_indexers, shape: Shape, first_dtype: np.dtype
206) -> Block:
207 """
208 With single-Block managers with homogeneous dtypes (that can already hold nan),
209 we avoid [...]
210 """
211 # assumes
212 # all(_is_homogeneous_mgr(mgr, first_dtype) for mgr, _ in in mgrs_indexers)
214 if all(not indexers for _, indexers in mgrs_indexers):
215 # https://github.com/pandas-dev/pandas/pull/52685#issuecomment-1523287739
216 arrs = [mgr.blocks[0].values.T for mgr, _ in mgrs_indexers]
217 arr = np.concatenate(arrs).T
218 bp = libinternals.BlockPlacement(slice(shape[0]))
219 nb = new_block_2d(arr, bp)
220 return nb
222 arr = np.empty(shape, dtype=first_dtype)
224 if first_dtype == np.float64:
225 take_func = libalgos.take_2d_axis0_float64_float64
226 else:
227 take_func = libalgos.take_2d_axis0_float32_float32
229 start = 0
230 for mgr, indexers in mgrs_indexers:
231 mgr_len = mgr.shape[1]
232 end = start + mgr_len
234 if 0 in indexers:
235 take_func(
236 mgr.blocks[0].values,
237 indexers[0],
238 arr[:, start:end],
239 )
240 else:
241 # No reindexing necessary, we can copy values directly
242 arr[:, start:end] = mgr.blocks[0].values
244 start += mgr_len
246 bp = libinternals.BlockPlacement(slice(shape[0]))
247 nb = new_block_2d(arr, bp)
248 return nb
251def _get_combined_plan(
252 mgrs: list[BlockManager],
253) -> Generator[tuple[BlockPlacement, list[JoinUnit]]]:
254 max_len = mgrs[0].shape[0]
256 blknos_list = [mgr.blknos for mgr in mgrs]
257 pairs = libinternals.get_concat_blkno_indexers(blknos_list)
258 for blknos, bp in pairs:
259 # assert bp.is_slice_like
260 # assert len(bp) > 0
262 units_for_bp = []
263 for k, mgr in enumerate(mgrs):
264 blkno = blknos[k]
266 nb = _get_block_for_concat_plan(mgr, bp, blkno, max_len=max_len)
267 unit = JoinUnit(nb)
268 units_for_bp.append(unit)
270 yield bp, units_for_bp
273def _get_block_for_concat_plan(
274 mgr: BlockManager, bp: BlockPlacement, blkno: int, *, max_len: int
275) -> Block:
276 blk = mgr.blocks[blkno]
277 # Assertions disabled for performance:
278 # assert bp.is_slice_like
279 # assert blkno != -1
280 # assert (mgr.blknos[bp] == blkno).all()
282 if len(bp) == len(blk.mgr_locs) and (
283 blk.mgr_locs.is_slice_like and blk.mgr_locs.as_slice.step == 1
284 ):
285 nb = blk
286 else:
287 ax0_blk_indexer = mgr.blklocs[bp.indexer]
289 slc = lib.maybe_indices_to_slice(ax0_blk_indexer, max_len)
290 # TODO: in all extant test cases 2023-04-08 we have a slice here.
291 # Will this always be the case?
292 if isinstance(slc, slice):
293 nb = blk.slice_block_columns(slc)
294 else:
295 nb = blk.take_block_columns(slc)
297 # assert nb.shape == (len(bp), mgr.shape[1])
298 return nb
301class JoinUnit:
302 def __init__(self, block: Block) -> None:
303 self.block = block
305 def __repr__(self) -> str:
306 return f"{type(self).__name__}({self.block!r})"
308 def _is_valid_na_for(self, dtype: DtypeObj) -> bool:
309 """
310 Check that we are all-NA of a type/dtype that is compatible with this dtype.
311 Augments `self.is_na` with an additional check of the type of NA values.
312 """
313 if not self.is_na:
314 return False
316 blk = self.block
317 if blk.dtype.kind == "V":
318 return True
320 if blk.dtype == object:
321 values = blk.values
322 return all(is_valid_na_for_dtype(x, dtype) for x in values.ravel(order="K"))
324 na_value = blk.fill_value
325 if na_value is NaT and blk.dtype != dtype:
326 # e.g. we are dt64 and other is td64
327 # fill_values match but we should not cast blk.values to dtype
328 # TODO: this will need updating if we ever have non-nano dt64/td64
329 return False
331 if na_value is NA and needs_i8_conversion(dtype):
332 # FIXME: kludge; test_append_empty_frame_with_timedelta64ns_nat
333 # e.g. blk.dtype == "Int64" and dtype is td64, we dont want
334 # to consider these as matching
335 return False
337 # TODO: better to use can_hold_element?
338 return is_valid_na_for_dtype(na_value, dtype)
340 @cache_readonly
341 def is_na(self) -> bool:
342 blk = self.block
343 if blk.dtype.kind == "V":
344 return True
345 return False
347 def get_reindexed_values(self, empty_dtype: DtypeObj, upcasted_na) -> ArrayLike:
348 values: ArrayLike
350 if upcasted_na is None and self.block.dtype.kind != "V":
351 # No upcasting is necessary
352 return self.block.values
353 else:
354 fill_value = upcasted_na
356 if self._is_valid_na_for(empty_dtype):
357 # note: always holds when self.block.dtype.kind == "V"
358 blk_dtype = self.block.dtype
360 if blk_dtype == np.dtype("object"):
361 # we want to avoid filling with np.nan if we are
362 # using None; we already know that we are all
363 # nulls
364 values = cast(np.ndarray, self.block.values)
365 if values.size and values[0, 0] is None:
366 fill_value = None
368 return make_na_array(empty_dtype, self.block.shape, fill_value)
370 return self.block.values
373def _concatenate_join_units(join_units: list[JoinUnit], copy: bool) -> ArrayLike:
374 """
375 Concatenate values from several join units along axis=1.
376 """
377 empty_dtype = _get_empty_dtype(join_units)
379 has_none_blocks = any(unit.block.dtype.kind == "V" for unit in join_units)
380 upcasted_na = _dtype_to_na_value(empty_dtype, has_none_blocks)
382 to_concat = [
383 ju.get_reindexed_values(empty_dtype=empty_dtype, upcasted_na=upcasted_na)
384 for ju in join_units
385 ]
387 if any(is_1d_only_ea_dtype(t.dtype) for t in to_concat):
388 # TODO(EA2D): special case not needed if all EAs used HybridBlocks
390 # error: No overload variant of "__getitem__" of "ExtensionArray" matches
391 # argument type "Tuple[int, slice]"
392 to_concat = [
393 t if is_1d_only_ea_dtype(t.dtype) else t[0, :] # type: ignore[call-overload]
394 for t in to_concat
395 ]
396 concat_values = concat_compat(to_concat, axis=0, ea_compat_axis=True)
397 concat_values = ensure_block_shape(concat_values, 2)
399 else:
400 concat_values = concat_compat(to_concat, axis=1)
402 return concat_values
405def _dtype_to_na_value(dtype: DtypeObj, has_none_blocks: bool):
406 """
407 Find the NA value to go with this dtype.
408 """
409 if isinstance(dtype, ExtensionDtype):
410 return dtype.na_value
411 elif dtype.kind in "mM":
412 return dtype.type("NaT", np.datetime_data(dtype)[0])
413 elif dtype.kind in "fc":
414 return dtype.type("NaN")
415 elif dtype.kind == "b":
416 # different from missing.na_value_for_dtype
417 return None
418 elif dtype.kind in "iu":
419 if not has_none_blocks:
420 # different from missing.na_value_for_dtype
421 return None
422 return np.nan
423 elif dtype.kind == "O":
424 return np.nan
425 raise NotImplementedError
428def _get_empty_dtype(join_units: Sequence[JoinUnit]) -> DtypeObj:
429 """
430 Return dtype and N/A values to use when concatenating specified units.
432 Returned N/A value may be None which means there was no casting involved.
434 Returns
435 -------
436 dtype
437 """
438 if lib.dtypes_all_equal([ju.block.dtype for ju in join_units]):
439 empty_dtype = join_units[0].block.dtype
440 return empty_dtype
442 has_none_blocks = any(unit.block.dtype.kind == "V" for unit in join_units)
444 dtypes = [unit.block.dtype for unit in join_units if not unit.is_na]
446 dtype = find_common_type(dtypes)
447 if has_none_blocks:
448 dtype = ensure_dtype_can_hold_na(dtype)
450 return dtype
453def _is_uniform_join_units(join_units: list[JoinUnit]) -> bool:
454 """
455 Check if the join units consist of blocks of uniform type that can
456 be concatenated using Block.concat_same_type instead of the generic
457 _concatenate_join_units (which uses `concat_compat`).
459 """
460 first = join_units[0].block
461 if first.dtype.kind == "V":
462 return False
463 return (
464 # exclude cases where a) ju.block is None or b) we have e.g. Int64+int64
465 all(type(ju.block) is type(first) for ju in join_units)
466 and
467 # e.g. DatetimeLikeBlock can be dt64 or td64, but these are not uniform
468 all(
469 ju.block.dtype == first.dtype
470 # GH#42092 we only want the dtype_equal check for non-numeric blocks
471 # (for now, may change but that would need a deprecation)
472 or ju.block.dtype.kind in "iub"
473 for ju in join_units
474 )
475 and
476 # no blocks that would get missing values (can lead to type upcasts)
477 # unless we're an extension dtype.
478 all(not ju.is_na or ju.block.is_extension for ju in join_units)
479 )