From 3b2cc9fc5aa9e49ce298f69bb63ba3b1d9c42edd Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 26 Nov 2021 13:54:17 -0800 Subject: [PATCH 1/2] API: don't downcast ints in Block.where --- pandas/core/internals/blocks.py | 39 ++++++++++++++------------------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 550bc4ac56d4b..53b72ad6e1663 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -40,7 +40,6 @@ can_hold_element, find_common_type, infer_dtype_from, - maybe_downcast_numeric, maybe_downcast_to_dtype, maybe_upcast, soft_convert_objects, @@ -1215,31 +1214,25 @@ def where(self, other, cond) -> list[Block]: else: # By the time we get here, we should have all Series/Index # args extracted to ndarray + if ( + is_list_like(other) + and not isinstance(other, np.ndarray) + and len(other) == self.shape[-1] + ): + # If we don't do this broadcasting here, then expressions.where + # will broadcast a 1D other to be row-like instead of + # column-like. + other = np.array(other).reshape(values.shape) + # If lengths don't match (or len(other)==1), we will raise + # inside expressions.where, see test_series_where + + # Note: expressions.where may upcast. result = expressions.where(~icond, values, other) - if self._can_hold_na or self.ndim == 1: - - if transpose: - result = result.T - - return [self.make_block(result)] - - # might need to separate out blocks - cond = ~icond - axis = cond.ndim - 1 - cond = cond.swapaxes(axis, 0) - mask = cond.all(axis=1) - - result_blocks: list[Block] = [] - for m in [mask, ~mask]: - if m.any(): - result = cast(np.ndarray, result) # EABlock overrides where - taken = result.take(m.nonzero()[0], axis=axis) - r = maybe_downcast_numeric(taken, self.dtype) - nb = self.make_block(r.T, placement=self._mgr_locs[m]) - result_blocks.append(nb) + if transpose: + result = result.T - return result_blocks + return [self.make_block(result)] def _unstack( self, From 7695edbf2fc2d618feb5ea5d9a266be26ba0eb30 Mon Sep 17 00:00:00 2001 From: Brock Date: Mon, 29 Nov 2021 21:35:42 -0800 Subject: [PATCH 2/2] restore casting, so just the expressions left --- pandas/core/internals/blocks.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 22d3ea5e066d9..662a58a54caf4 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -40,6 +40,7 @@ can_hold_element, find_common_type, infer_dtype_from, + maybe_downcast_numeric, maybe_downcast_to_dtype, maybe_upcast, soft_convert_objects, @@ -1229,10 +1230,29 @@ def where(self, other, cond) -> list[Block]: # Note: expressions.where may upcast. result = expressions.where(~icond, values, other) - if transpose: - result = result.T + if self._can_hold_na or self.ndim == 1: + + if transpose: + result = result.T + + return [self.make_block(result)] + + # might need to separate out blocks + cond = ~icond + axis = cond.ndim - 1 + cond = cond.swapaxes(axis, 0) + mask = cond.all(axis=1) + + result_blocks: list[Block] = [] + for m in [mask, ~mask]: + if m.any(): + result = cast(np.ndarray, result) # EABlock overrides where + taken = result.take(m.nonzero()[0], axis=axis) + r = maybe_downcast_numeric(taken, self.dtype) + nb = self.make_block(r.T, placement=self._mgr_locs[m]) + result_blocks.append(nb) - return [self.make_block(result)] + return result_blocks def _unstack( self,