feat: implement DataFrame.rolling with count, sum, mean, std

Co-authored-by: tswast <247555+tswast@users.noreply.github.com>

Author: google-labs-jules[bot]

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.
 

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")

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)