SNOW-2679277: Add support for groupby.get_group/resample/rolling in faster pandas (#3991)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -161,6 +161,9 @@
   - `groupby.any`
   - `groupby.all`
   - `groupby.unique`
+  - `groupby.get_group`
+  - `groupby.rolling`
+  - `groupby.resample`
   - `to_snowflake`
   - `to_snowpark`
 - Make faster pandas disabled by default (opt-in instead of opt-out).

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

@@ -5975,6 +5975,40 @@ def groupby_resample(
         agg_args: Any,
         agg_kwargs: dict[str, Any],
     ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_resample_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._groupby_resample_internal(
+                    resample_kwargs=resample_kwargs,
+                    resample_method=resample_method,
+                    groupby_kwargs=groupby_kwargs,
+                    is_series=is_series,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                )
+            )
+        qc = self._groupby_resample_internal(
+            resample_kwargs=resample_kwargs,
+            resample_method=resample_method,
+            groupby_kwargs=groupby_kwargs,
+            is_series=is_series,
+            agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _groupby_resample_internal(
+        self,
+        resample_kwargs: dict[str, Any],
+        resample_method: AggFuncType,
+        groupby_kwargs: dict[str, Any],
+        is_series: bool,
+        agg_args: Any,
+        agg_kwargs: dict[str, Any],
+    ) -> "SnowflakeQueryCompiler":
 
         validate_groupby_resample_supported_by_snowflake(resample_kwargs)
         level = groupby_kwargs.get("level", None)
@@ -6128,6 +6162,40 @@ def groupby_rolling(
         is_series: bool,
         agg_args: Any,
         agg_kwargs: dict[str, Any],
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_rolling_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._groupby_rolling_internal(
+                    rolling_kwargs=rolling_kwargs,
+                    rolling_method=rolling_method,
+                    groupby_kwargs=groupby_kwargs,
+                    is_series=is_series,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                )
+            )
+        qc = self._groupby_rolling_internal(
+            rolling_kwargs=rolling_kwargs,
+            rolling_method=rolling_method,
+            groupby_kwargs=groupby_kwargs,
+            is_series=is_series,
+            agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _groupby_rolling_internal(
+        self,
+        rolling_kwargs: dict[str, Any],
+        rolling_method: AggFuncType,
+        groupby_kwargs: dict[str, Any],
+        is_series: bool,
+        agg_args: Any,
+        agg_kwargs: dict[str, Any],
     ) -> "SnowflakeQueryCompiler":
         """
         Return a rolling grouper, providing rolling functionality per group.
@@ -6548,6 +6616,43 @@ def groupby_get_group(
         agg_kwargs: dict[str, Any],
         drop: bool = False,
         **kwargs: dict[str, Any],
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_get_group_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._groupby_get_group_internal(
+                    by=by,
+                    axis=axis,
+                    groupby_kwargs=groupby_kwargs,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                    drop=drop,
+                    **kwargs,
+                )
+            )
+        qc = self._groupby_get_group_internal(
+            by=by,
+            axis=axis,
+            groupby_kwargs=groupby_kwargs,
+            agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+            drop=drop,
+            **kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _groupby_get_group_internal(
+        self,
+        by: Any,
+        axis: int,
+        groupby_kwargs: dict[str, Any],
+        agg_args: tuple[Any],
+        agg_kwargs: dict[str, Any],
+        drop: bool = False,
+        **kwargs: dict[str, Any],
     ) -> "SnowflakeQueryCompiler":
         """
         Get all rows that match a given group name in the `by` column.