SNOW-2476998: Add support for expanding.min/max/count/sum/mean/std/var/sem in faster pandas (#3974)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -93,6 +93,14 @@
   - `rolling.var`
   - `rolling.sem`
   - `rolling.corr`
+  - `expanding.min`
+  - `expanding.max`
+  - `expanding.count`
+  - `expanding.sum`
+  - `expanding.mean`
+  - `expanding.std`
+  - `expanding.var`
+  - `expanding.sem`
 - Make faster pandas disabled by default (opt-in instead of opt-out).
 
 ## 1.42.0 (2025-10-28)

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -17767,6 +17767,31 @@ def expanding_count(
         fold_axis: Union[int, str],
         expanding_kwargs: dict,
         numeric_only: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_count_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_count_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    numeric_only=numeric_only,
+                )
+            )
+        qc = self._expanding_count_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            numeric_only=numeric_only,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_count_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        numeric_only: bool = False,
     ) -> "SnowflakeQueryCompiler":
         return self._window_agg(
             window_func=WindowFunction.EXPANDING,
@@ -17782,6 +17807,37 @@ def expanding_sum(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_sum_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_sum_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_sum_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_sum_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "expanding_sum", engine, engine_kwargs
@@ -17800,6 +17856,37 @@ def expanding_mean(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_mean_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_mean_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_mean_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_mean_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "expanding_mean", engine, engine_kwargs
@@ -17829,6 +17916,40 @@ def expanding_var(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_var_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_var_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    ddof=ddof,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_var_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            ddof=ddof,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_var_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        ddof: int = 1,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "rolling_var", engine, engine_kwargs
@@ -17848,6 +17969,40 @@ def expanding_std(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_std_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_std_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    ddof=ddof,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_std_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            ddof=ddof,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_std_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        ddof: int = 1,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "rolling_std", engine, engine_kwargs
@@ -17866,6 +18021,37 @@ def expanding_min(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_min_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_min_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_min_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_min_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "expanding_min", engine, engine_kwargs
@@ -17884,6 +18070,37 @@ def expanding_max(
         numeric_only: bool = False,
         engine: Optional[Literal["cython", "numba"]] = None,
         engine_kwargs: Optional[dict[str, bool]] = None,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_max_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_max_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    numeric_only=numeric_only,
+                    engine=engine,
+                    engine_kwargs=engine_kwargs,
+                )
+            )
+        qc = self._expanding_max_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            numeric_only=numeric_only,
+            engine=engine,
+            engine_kwargs=engine_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_max_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        numeric_only: bool = False,
+        engine: Optional[Literal["cython", "numba"]] = None,
+        engine_kwargs: Optional[dict[str, bool]] = None,
     ) -> "SnowflakeQueryCompiler":
         WarningMessage.warning_if_engine_args_is_set(
             "expanding_max", engine, engine_kwargs
@@ -17972,6 +18189,34 @@ def expanding_sem(
         expanding_kwargs: dict,
         ddof: int = 1,
         numeric_only: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _expanding_sem_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._expanding_sem_internal(
+                    fold_axis=fold_axis,
+                    expanding_kwargs=expanding_kwargs,
+                    ddof=ddof,
+                    numeric_only=numeric_only,
+                )
+            )
+        qc = self._expanding_sem_internal(
+            fold_axis=fold_axis,
+            expanding_kwargs=expanding_kwargs,
+            ddof=ddof,
+            numeric_only=numeric_only,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _expanding_sem_internal(
+        self,
+        fold_axis: Union[int, str],
+        expanding_kwargs: dict,
+        ddof: int = 1,
+        numeric_only: bool = False,
     ) -> "SnowflakeQueryCompiler":
         return self._window_agg(
             window_func=WindowFunction.EXPANDING,

tests/integ/modin/test_faster_pandas.py

@@ -439,6 +439,52 @@ def test_duplicated(session):
         assert_series_equal(snow_result, native_result)
 
 
+@pytest.mark.parametrize(
+    "func",
+    [
+        "min",
+        "max",
+        "count",
+        "sum",
+        "mean",
+        "std",
+        "var",
+        "sem",
+    ],
+)
+@sql_count_checker(query_count=3)
+def test_expanding(session, func):
+    with session_parameter_override(
+        session, "dummy_row_pos_optimization_enabled", True
+    ):
+        # create tables
+        table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+        session.create_dataframe(
+            native_pd.DataFrame([[1, 11], [2, 12], [3, 13]], columns=["A", "B"])
+        ).write.save_as_table(table_name, table_type="temp")
+
+        # create snow dataframes
+        df = pd.read_snowflake(table_name)
+        snow_result = getattr(df.expanding(), func)()
+
+        # verify that the input dataframe has a populated relaxed query compiler
+        assert df._query_compiler._relaxed_query_compiler is not None
+        assert df._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+        # verify that the output dataframe also has a populated relaxed query compiler
+        assert snow_result._query_compiler._relaxed_query_compiler is not None
+        assert (
+            snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode
+            is True
+        )
+
+        # create pandas dataframes
+        native_df = df.to_pandas()
+        native_result = getattr(native_df.expanding(), func)()
+
+        # compare results
+        assert_frame_equal(snow_result, native_result, check_dtype=False)
+
+
 @pytest.mark.parametrize(
     "func",
     [