Fix output result and result min_periods handling

pandas/core/window/numba_.py

@@ -189,6 +189,7 @@ def generate_numba_table_func(
  """
  nopython, nogil, parallel = get_jit_arguments(engine_kwargs, kwargs)
 
+ # TODO: change this key
  cache_key = (func, "rolling_apply")
  if cache_key in NUMBA_FUNC_CACHE:
  return NUMBA_FUNC_CACHE[cache_key]
@@ -206,15 +207,17 @@ def roll_table(
  ):
  # TODO: consider axis argument, len should be replaced with axis aware result
  result = np.empty(values.shape)
+ min_periods_mask = np.empty(values.shape)
  for i in loop_range(len(result)):
  start = begin[i]
  stop = end[i]
- window = values[start:stop, :]
- count_nan = np.sum(np.isnan(window))
+ window = values[start:stop]
+ count_nan = np.sum(np.isnan(window), axis=0)
  sub_result = numba_func(window, *args)
  nan_mask = len(window) - count_nan >= minimum_periods
- sub_result[~nan_mask] = np.nan
+ min_periods_mask[i, :] = nan_mask
  result[i, :] = sub_result
+ result = np.where(min_periods_mask, result, np.nan)
  return result
 
  return roll_table

pandas/core/window/rolling.py

@@ -66,7 +66,10 @@
  GroupbyIndexer,
  VariableWindowIndexer,
 )
-from pandas.core.window.numba_ import generate_numba_apply_func
+from pandas.core.window.numba_ import (
+ generate_numba_apply_func,
+ generate_numba_table_func,
+)
 
 if TYPE_CHECKING:
  from pandas import DataFrame, Series
@@ -418,10 +421,17 @@ def _apply_tablewise(
  if self._selected_obj.ndim == 1:
  raise ValueError("method='table' not applicable for Series objects.")
  obj = self._create_data(self._selected_obj)
- values = self._prep_values(obj)
- # Need to ensure we wrap this correctly, will return ndarray
+ values = self._prep_values(obj.to_numpy())
  result = homogeneous_func(values)
- return result
+ out = obj._constructor(result, index=obj.index, columns=obj.columns)
+
+ if out.shape[1] == 0 and obj.shape[1] > 0:
+ raise DataError("No numeric types to aggregate")
+ elif out.shape[1] == 0:
+ return obj.astype("float64")
+
+ self._insert_on_column(out, obj)
+ return out
 
  def _apply(
  self,
@@ -482,7 +492,10 @@ def calc(x):
 
  return result
 
- return self._apply_blockwise(homogeneous_func, name)
+ if self.method == "column":
+ return self._apply_blockwise(homogeneous_func, name)
+ else:
+ return self._apply_tablewise(homogeneous_func, name)
 
  def aggregate(self, func, *args, **kwargs):
  result, how = aggregate(self, func, *args, **kwargs)
@@ -1326,8 +1339,16 @@ def apply(
  if maybe_use_numba(engine):
  if raw is False:
  raise ValueError("raw must be `True` when using the numba engine")
- apply_func = generate_numba_apply_func(args, kwargs, func, engine_kwargs)
- numba_cache_key = (func, "rolling_apply")
+ if self.method == "column":
+ apply_func = generate_numba_apply_func(
+ args, kwargs, func, engine_kwargs
+ )
+ numba_cache_key = (func, "rolling_apply_column")
+ else:
+ apply_func = generate_numba_table_func(
+ args, kwargs, func, engine_kwargs
+ )
+ numba_cache_key = (func, "rolling_apply_table")
  elif engine in ("cython", None):
  if engine_kwargs is not None:
  raise ValueError("cython engine does not accept engine_kwargs")