REF: implement Block.reduce for DataFrame._reduce (#35867)

Author: jbrockmendel

pandas/core/frame.py

@@ -8647,13 +8647,11 @@ def blk_func(values):
  return op(values, axis=1, skipna=skipna, **kwds)
 
  # After possibly _get_data and transposing, we are now in the
- # simple case where we can use BlockManager._reduce
+ # simple case where we can use BlockManager.reduce
  res = df._mgr.reduce(blk_func)
- assert isinstance(res, dict)
- if len(res):
- assert len(res) == max(list(res.keys())) + 1, res.keys()
- out = df._constructor_sliced(res, index=range(len(res)), dtype=out_dtype)
- out.index = df.columns
+ out = df._constructor(res,).iloc[0].rename(None)
+ if out_dtype is not None:
+ out = out.astype(out_dtype)
  if axis == 0 and is_object_dtype(out.dtype):
  out[:] = coerce_to_dtypes(out.values, df.dtypes)
  return out

pandas/core/internals/blocks.py

@@ -346,6 +346,21 @@ def apply(self, func, **kwargs) -> List["Block"]:
 
  return self._split_op_result(result)
 
+ def reduce(self, func) -> List["Block"]:
+ # We will apply the function and reshape the result into a single-row
+ # Block with the same mgr_locs; squeezing will be done at a higher level
+ assert self.ndim == 2
+
+ result = func(self.values)
+ if np.ndim(result) == 0:
+ # TODO(EA2D): special case not needed with 2D EAs
+ res_values = np.array([[result]])
+ else:
+ res_values = result.reshape(-1, 1)
+
+ nb = self.make_block(res_values)
+ return [nb]
+
  def _split_op_result(self, result) -> List["Block"]:
  # See also: split_and_operate
  if is_extension_array_dtype(result) and result.ndim > 1:

pandas/core/internals/managers.py

@@ -330,31 +330,18 @@ def _verify_integrity(self) -> None:
  f"tot_items: {tot_items}"
  )
 
- def reduce(self, func):
+ def reduce(self: T, func) -> T:
  # If 2D, we assume that we're operating column-wise
- if self.ndim == 1:
- # we'll be returning a scalar
- blk = self.blocks[0]
- return func(blk.values)
+ assert self.ndim == 2
 
- res = {}
+ res_blocks = []
  for blk in self.blocks:
- bres = func(blk.values)
-
- if np.ndim(bres) == 0:
- # EA
- assert blk.shape[0] == 1
- new_res = zip(blk.mgr_locs.as_array, [bres])
- else:
- assert bres.ndim == 1, bres.shape
- assert blk.shape[0] == len(bres), (blk.shape, bres.shape)
- new_res = zip(blk.mgr_locs.as_array, bres)
-
- nr = dict(new_res)
- assert not any(key in res for key in nr)
- res.update(nr)
+ nbs = blk.reduce(func)
+ res_blocks.extend(nbs)
 
- return res
+ index = Index([0]) # placeholder
+ new_mgr = BlockManager.from_blocks(res_blocks, [self.items, index])
+ return new_mgr
 
  def operate_blockwise(self, other: "BlockManager", array_op) -> "BlockManager":
  """