feat: Add agg/aggregate methods to windows (#2288)

Author: TrevorBergeron

bigframes/core/window/rolling.py

@@ -15,19 +15,24 @@
 from __future__ importannotations
 
 importdatetime
-importtyping
+fromtypingimportLiteral, Mapping, Sequence, TYPE_CHECKING, Union
 
 importbigframes_vendored.pandas.core.window.rollingasvendored_pandas_rolling
 importnumpy
 importpandas
 
 frombigframesimportdtypes
+frombigframes.coreimportagg_expressions
 frombigframes.coreimportexpressionasex
-frombigframes.coreimportlog_adapter, ordering, window_spec
+frombigframes.coreimportlog_adapter, ordering, utils, window_spec
 importbigframes.core.blocksasblocks
 frombigframes.core.windowimportorderingaswindow_ordering
 importbigframes.operations.aggregationsasagg_ops
 
+ifTYPE_CHECKING:
+importbigframes.dataframeasdf
+importbigframes.seriesasseries
+
 
 @log_adapter.class_logger
 classWindow(vendored_pandas_rolling.Window):
@@ -37,7 +42,7 @@ def __init__(
 self,
 block: blocks.Block,
 window_spec: window_spec.WindowSpec,
-value_column_ids: typing.Sequence[str],
+value_column_ids: Sequence[str],
 drop_null_groups: bool=True,
 is_series: bool=False,
 skip_agg_column_id: str|None=None,
@@ -52,55 +57,106 @@ def __init__(
 self._skip_agg_column_id=skip_agg_column_id
 
 defcount(self):
-returnself._apply_aggregate(agg_ops.count_op)
+returnself._apply_aggregate_op(agg_ops.count_op)
 
 defsum(self):
-returnself._apply_aggregate(agg_ops.sum_op)
+returnself._apply_aggregate_op(agg_ops.sum_op)
 
 defmean(self):
-returnself._apply_aggregate(agg_ops.mean_op)
+returnself._apply_aggregate_op(agg_ops.mean_op)
 
 defvar(self):
-returnself._apply_aggregate(agg_ops.var_op)
+returnself._apply_aggregate_op(agg_ops.var_op)
 
 defstd(self):
-returnself._apply_aggregate(agg_ops.std_op)
+returnself._apply_aggregate_op(agg_ops.std_op)
 
 defmax(self):
-returnself._apply_aggregate(agg_ops.max_op)
+returnself._apply_aggregate_op(agg_ops.max_op)
 
 defmin(self):
-returnself._apply_aggregate(agg_ops.min_op)
+returnself._apply_aggregate_op(agg_ops.min_op)
 
-def_apply_aggregate(
-self,
-op: agg_ops.UnaryAggregateOp,
- ):
-agg_block=self._aggregate_block(op)
+defagg(self, func) ->Union[df.DataFrame, series.Series]:
+ifutils.is_dict_like(func):
+returnself._agg_dict(func)
+elifutils.is_list_like(func):
+returnself._agg_list(func)
+else:
+returnself._agg_func(func)
 
-ifself._is_series:
-frombigframes.seriesimportSeries
+aggregate=agg
+
+def_agg_func(self, func) ->df.DataFrame:
+ids, labels=self._aggregated_columns()
+aggregations= [agg(col_id, agg_ops.lookup_agg_func(func)[0]) forcol_idinids]
+returnself._apply_aggs(aggregations, labels)
+
+def_agg_dict(self, func: Mapping) ->df.DataFrame:
+aggregations: list[agg_expressions.Aggregation] = []
+column_labels= []
+function_labels= []
 
-returnSeries(agg_block)
+want_aggfunc_level=any(utils.is_list_like(aggs) foraggsinfunc.values())
+
+forlabel, funcs_for_idinfunc.items():
+col_id=self._block.label_to_col_id[label][-1] # get last matching column
+func_list= (
+funcs_for_idifutils.is_list_like(funcs_for_id) else [funcs_for_id]
+ )
+forfinfunc_list:
+f_op, f_label=agg_ops.lookup_agg_func(f)
+aggregations.append(agg(col_id, f_op))
+column_labels.append(label)
+function_labels.append(f_label)
+ifwant_aggfunc_level:
+result_labels: pandas.Index=utils.combine_indices(
+pandas.Index(column_labels),
+pandas.Index(function_labels),
+ )
 else:
-frombigframes.dataframeimportDataFrame
+result_labels=pandas.Index(column_labels)
 
-# Preserve column order.
-column_labels= [
-self._block.col_id_to_label[col_id] forcol_idinself._value_column_ids
- ]
-returnDataFrame(agg_block)._reindex_columns(column_labels)
+returnself._apply_aggs(aggregations, result_labels)
 
-def_aggregate_block(self, op: agg_ops.UnaryAggregateOp) ->blocks.Block:
-agg_col_ids= [
-col_id
-forcol_idinself._value_column_ids
-ifcol_id!=self._skip_agg_column_id
+def_agg_list(self, func: Sequence) ->df.DataFrame:
+ids, labels=self._aggregated_columns()
+aggregations= [
+agg(col_id, agg_ops.lookup_agg_func(f)[0]) forcol_idinidsforfinfunc
  ]
-block, result_ids=self._block.multi_apply_window_op(
-agg_col_ids,
-op,
-self._window_spec,
+
+ifself._is_series:
+# if series, no need to rebuild
+result_cols_idx=pandas.Index(
+ [agg_ops.lookup_agg_func(f)[1] forfinfunc]
+ )
+else:
+ifself._block.column_labels.nlevels>1:
+# Restructure MultiIndex for proper format: (idx1, idx2, func)
+# rather than ((idx1, idx2), func).
+column_labels= [
+tuple(label) + (agg_ops.lookup_agg_func(f)[1],)
+forlabelinlabels.to_frame(index=False).to_numpy()
+forfinfunc
+ ]
+else: # Single-level index
+column_labels= [
+ (label, agg_ops.lookup_agg_func(f)[1])
+forlabelinlabels
+forfinfunc
+ ]
+result_cols_idx=pandas.MultiIndex.from_tuples(
+column_labels, names=[*self._block.column_labels.names, None]
+ )
+returnself._apply_aggs(aggregations, result_cols_idx)
+
+def_apply_aggs(
+self, exprs: Sequence[agg_expressions.Aggregation], labels: pandas.Index
+ ):
+block, ids=self._block.apply_analytic(
+agg_exprs=exprs,
+window=self._window_spec,
+result_labels=labels,
 skip_null_groups=self._drop_null_groups,
  )
 
@@ -115,24 +171,50 @@ def _aggregate_block(self, op: agg_ops.UnaryAggregateOp) -> blocks.Block:
  )
 block=block.set_index(col_ids=index_ids)
 
-labels= [self._block.col_id_to_label[col] forcolinagg_col_ids]
 ifself._skip_agg_column_idisnotNone:
-result_ids= [self._skip_agg_column_id, *result_ids]
-labels.insert(0, self._block.col_id_to_label[self._skip_agg_column_id])
+block=block.select_columns([self._skip_agg_column_id, *ids])
+else:
+block=block.select_columns(ids).with_column_labels(labels)
+
+ifself._is_seriesand (len(block.value_columns) ==1):
+importbigframes.seriesasseries
+
+returnseries.Series(block)
+else:
+importbigframes.dataframeasdf
+
+returndf.DataFrame(block)
+
+def_apply_aggregate_op(
+self,
+op: agg_ops.UnaryAggregateOp,
+ ):
+ids, labels=self._aggregated_columns()
+aggregations= [agg(col_id, op) forcol_idinids]
+returnself._apply_aggs(aggregations, labels)
 
-returnblock.select_columns(result_ids).with_column_labels(labels)
+def_aggregated_columns(self) ->tuple[Sequence[str], pandas.Index]:
+agg_col_ids= [
+col_id
+forcol_idinself._value_column_ids
+ifcol_id!=self._skip_agg_column_id
+ ]
+labels: pandas.Index=pandas.Index(
+ [self._block.col_id_to_label[col] forcolinagg_col_ids]
+ )
+returnagg_col_ids, labels
 
 
 defcreate_range_window(
 block: blocks.Block,
 window: pandas.Timedelta|numpy.timedelta64|datetime.timedelta|str,
 *,
-value_column_ids: typing.Sequence[str] =tuple(),
+value_column_ids: Sequence[str] =tuple(),
 min_periods: int|None,
 on: str|None=None,
-closed: typing.Literal["right", "left", "both", "neither"],
+closed: Literal["right", "left", "both", "neither"],
 is_series: bool,
-grouping_keys: typing.Sequence[str] =tuple(),
+grouping_keys: Sequence[str] =tuple(),
 drop_null_groups: bool=True,
 ) ->Window:
 
@@ -184,3 +266,11 @@ def create_range_window(
 skip_agg_column_id=NoneifonisNoneelserolling_key_col_id,
 drop_null_groups=drop_null_groups,
  )
+
+
+defagg(input: str, op: agg_ops.AggregateOp) ->agg_expressions.Aggregation:
+ifisinstance(op, agg_ops.UnaryAggregateOp):
+returnagg_expressions.UnaryAggregation(op, ex.deref(input))
+else:
+assertisinstance(op, agg_ops.NullaryAggregateOp)
+returnagg_expressions.NullaryAggregation(op)

tests/system/small/test_window.py

@@ -228,6 +228,75 @@ def test_dataframe_window_agg_ops(scalars_dfs, windowing, agg_op):
 pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False)
 
 
+@pytest.mark.parametrize(
+ ("windowing"),
+ [
+pytest.param(lambdax: x.expanding(), id="expanding"),
+pytest.param(lambdax: x.rolling(3, min_periods=3), id="rolling"),
+pytest.param(
+lambdax: x.groupby(level=0).rolling(3, min_periods=3), id="rollinggroupby"
+ ),
+pytest.param(
+lambdax: x.groupby("int64_too").expanding(min_periods=2),
+id="expandinggroupby",
+ ),
+ ],
+)
+@pytest.mark.parametrize(
+ ("func"),
+ [
+pytest.param("sum", id="sum_by_name"),
+pytest.param(np.sum, id="sum_by_by_np"),
+pytest.param([np.sum, np.mean], id="list_of_funcs"),
+pytest.param(
+ {"int64_col": np.sum, "float64_col": "mean"}, id="dict_of_single_funcs"
+ ),
+pytest.param(
+ {"int64_col": np.sum, "float64_col": ["mean", np.max]},
+id="dict_of_lists_and_single_funcs",
+ ),
+ ],
+)
+deftest_dataframe_window_agg_func(scalars_dfs, windowing, func):
+bf_df, pd_df=scalars_dfs
+target_columns= ["int64_too", "float64_col", "bool_col", "int64_col"]
+index_column="bool_col"
+bf_df=bf_df[target_columns].set_index(index_column)
+pd_df=pd_df[target_columns].set_index(index_column)
+
+bf_result=windowing(bf_df).agg(func).to_pandas()
+
+pd_result=windowing(pd_df).agg(func)
+
+pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False)
+
+
+deftest_series_window_agg_single_func(scalars_dfs):
+bf_df, pd_df=scalars_dfs
+index_column="bool_col"
+bf_series=bf_df.set_index(index_column).int64_too
+pd_series=pd_df.set_index(index_column).int64_too
+
+bf_result=bf_series.expanding().agg("sum").to_pandas()
+
+pd_result=pd_series.expanding().agg("sum")
+
+pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False)
+
+
+deftest_series_window_agg_multi_func(scalars_dfs):
+bf_df, pd_df=scalars_dfs
+index_column="bool_col"
+bf_series=bf_df.set_index(index_column).int64_too
+pd_series=pd_df.set_index(index_column).int64_too
+
+bf_result=bf_series.expanding().agg(["sum", np.mean]).to_pandas()
+
+pd_result=pd_series.expanding().agg(["sum", np.mean])
+
+pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False)
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize(
 "window", # skipped numpy timedelta because Pandas does not support it.

third_party/bigframes_vendored/pandas/core/window/rolling.py

@@ -37,3 +37,52 @@ def max(self):
 defmin(self):
 """Calculate the weighted window minimum."""
 raiseNotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+defagg(self, func):
+"""
+ Aggregate using one or more operations over the specified axis.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+
+ >>> df = bpd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
+ >>> df
+ A B C
+ 0 1 4 7
+ 1 2 5 8
+ 2 3 6 9
+ <BLANKLINE>
+ [3 rows x 3 columns]
+
+ >>> df.rolling(2).sum()
+ A B C
+ 0 <NA> <NA> <NA>
+ 1 3 9 15
+ 2 5 11 17
+ <BLANKLINE>
+ [3 rows x 3 columns]
+
+ >>> df.rolling(2).agg({"A": "sum", "B": "min"})
+ A B
+ 0 <NA> <NA>
+ 1 3 4
+ 2 5 5
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ Args:
+ func (function, str, list or dict):
+ Function to use for aggregating the data.
+
+ Accepted combinations are:
+
+ - string function name
+ - list of function names, e.g. ``['sum', 'mean']``
+ - dict of axis labels -> function names or list of such.
+
+ Returns:
+ Series or DataFrame
+
+ """
+raiseNotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)