diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index 5447afc..08393e9 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -22,6 +22,7 @@ from functools import reduce import operator +from leanframe.core.window import Rolling from leanframe.core.dtypes import convert_ibis_to_pandas from leanframe.core.indexing import ( Index, @@ -126,6 +127,19 @@ def dtypes(self) -> pd.Series: types = [convert_ibis_to_pandas(t) for t in self._data.schema().types] return pd.Series(types, index=names, name="dtypes") + def rolling(self, window: int, min_periods: int | None = None) -> Rolling: + """Provide rolling window calculations. + + Args: + window: Size of the moving window. This is the number of observations + used for calculating the statistic. + min_periods: Minimum number of observations in window required to have + a value; otherwise, result is null. Defaults to window size. + + Returns: + A Rolling object. + """ + return Rolling(self, window=window, min_periods=min_periods) def __getitem__(self, key: str) -> DataFrame: """Get a column. diff --git a/leanframe/core/window.py b/leanframe/core/window.py new file mode 100644 index 0000000..0e5efb2 --- /dev/null +++ b/leanframe/core/window.py @@ -0,0 +1,101 @@ +# Copyright 2025 Google LLC, LeanFrame Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Windowed operations on DataFrames.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import ibis + +if TYPE_CHECKING: + from leanframe.core.frame import DataFrame + + +class Rolling: + """Windowed operations on DataFrames.""" + + def __init__( + self, + obj: DataFrame, + window: int, + min_periods: int | None = None, + ): + """Initialize a Rolling object. + + Args: + obj: The DataFrame to apply rolling operations to. + window: The size of the moving window. + min_periods: Minimum number of observations in window required to + have a value. Defaults to window size. + """ + self._obj = obj + self._window = window + self._min_periods = min_periods if min_periods is not None else window + + def _apply_aggregation(self, op: str, **kwargs: Any) -> DataFrame: + from leanframe.core.frame import DataFrame + + t = self._obj._data + + w = ibis.window(preceding=self._window - 1, following=0) + + exprs = [] + for c in t.columns: + col = t[c] + + if op == "count": + agg = col.count() + elif op == "sum": + agg = col.sum() + elif op == "mean": + agg = col.mean() + elif op == "std": + agg = col.std(how="sample") + else: + raise NotImplementedError(f"Unsupported operation: {op}") + + agg_over = agg.over(w) + + if op == "count": + # For count, min_periods applies to the TOTAL number of rows in the window (including nulls) + # We can trick Ibis into counting all rows by coalescing the column. + valid_count = ibis.coalesce(col, 0).count().over(w) + else: + # For others, min_periods applies to the number of non-null observations + valid_count = col.count().over(w) + + final_expr = ibis.ifelse( + valid_count >= self._min_periods, agg_over, ibis.null() + ).name(c) + exprs.append(final_expr) + + return DataFrame(t.select(exprs)) + + def count(self) -> DataFrame: + """Calculate the rolling count.""" + return self._apply_aggregation("count") + + def sum(self) -> DataFrame: + """Calculate the rolling sum.""" + return self._apply_aggregation("sum") + + def mean(self) -> DataFrame: + """Calculate the rolling mean.""" + return self._apply_aggregation("mean") + + def std(self) -> DataFrame: + """Calculate the rolling standard deviation.""" + return self._apply_aggregation("std") diff --git a/tests/unit/test_window.py b/tests/unit/test_window.py new file mode 100644 index 0000000..bd29ba3 --- /dev/null +++ b/tests/unit/test_window.py @@ -0,0 +1,42 @@ +import numpy as np +import pandas as pd +import pandas.testing as pdt +import ibis + +from leanframe.core.frame import DataFrame + + +def test_dataframe_rolling_aggregations(): + pdf = pd.DataFrame( + { + "a": [1.0, 2.0, 3.0, 4.0, 5.0], + "b": [10.0, np.nan, 30.0, 40.0, 50.0], + } + ) + ldf = DataFrame(ibis.memtable(pdf)) + + # Test sum with default min_periods (min_periods=window) + p_sum = pdf.rolling(window=3).sum() + l_sum = ldf.rolling(window=3).sum().to_pandas() + pdt.assert_frame_equal(p_sum, l_sum, check_dtype=False) + + # Test sum with min_periods=1 + p_sum_min = pdf.rolling(window=3, min_periods=1).sum() + l_sum_min = ldf.rolling(window=3, min_periods=1).sum().to_pandas() + pdt.assert_frame_equal(p_sum_min, l_sum_min, check_dtype=False) + + # Test count + p_count = pdf.rolling(window=2).count() + # Pandas count returns float if there's nan, int if not. Ibis returns int. Let pandas handle the check_dtype=False + l_count = ldf.rolling(window=2).count().to_pandas() + pdt.assert_frame_equal(p_count, l_count, check_dtype=False) + + # Test mean + p_mean = pdf.rolling(window=3, min_periods=2).mean() + l_mean = ldf.rolling(window=3, min_periods=2).mean().to_pandas() + pdt.assert_frame_equal(p_mean, l_mean, check_dtype=False) + + # Test std + p_std = pdf.rolling(window=3, min_periods=2).std() + l_std = ldf.rolling(window=3, min_periods=2).std().to_pandas() + pdt.assert_frame_equal(p_std, l_std, check_dtype=False)