SNOW-2643972: Add support for groupby properties (groupby.groups/indices) in faster pandas (#3984)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -144,6 +144,8 @@
   - `cumsum`
   - `cummin`
   - `cummax`
+  - `groupby.groups`
+  - `groupby.indices`
   - `groupby.first`
   - `groupby.last`
   - `groupby.rank`

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

@@ -823,14 +823,17 @@ def __init__(self, frame: InternalFrame) -> None:
     storage_format = property(lambda self: "Snowflake")
 
     def _raise_not_implemented_error_for_timedelta(
-        self, frame: InternalFrame = None
+        self, frame: InternalFrame = None, stack_depth: int = 2
     ) -> None:
         """Raise NotImplementedError for SnowflakeQueryCompiler methods which does not support timedelta yet."""
         if frame is None:
             frame = self._modin_frame
         for val in frame.snowflake_quoted_identifier_to_snowpark_pandas_type.values():
             if isinstance(val, TimedeltaType):
-                method = inspect.currentframe().f_back.f_back.f_code.co_name  # type: ignore[union-attr]
+                method_frame = inspect.currentframe()
+                for _ in range(stack_depth):
+                    method_frame = method_frame.f_back  # type: ignore[union-attr]
+                method = method_frame.f_code.co_name  # type: ignore[union-attr]
                 ErrorMessage.not_implemented_for_timedelta(method)
 
     def _warn_lost_snowpark_pandas_type(self) -> None:
@@ -5737,6 +5740,49 @@ def groupby_rank(
         na_option: Literal["keep", "top", "bottom"] = "keep",
         ascending: bool = True,
         pct: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_rank_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_rank_internal(
+                    by=by,
+                    groupby_kwargs=groupby_kwargs,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                    axis=axis,
+                    method=method,
+                    na_option=na_option,
+                    ascending=ascending,
+                    pct=pct,
+                )
+            )
+        qc = self._groupby_rank_internal(
+            by=by,
+            groupby_kwargs=groupby_kwargs,
+            agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+            axis=axis,
+            method=method,
+            na_option=na_option,
+            ascending=ascending,
+            pct=pct,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _groupby_rank_internal(
+        self,
+        by: Any,
+        groupby_kwargs: dict[str, Any],
+        agg_args: Any,
+        agg_kwargs: dict[str, Any],
+        axis: Axis = 0,
+        method: Literal["average", "min", "max", "first", "dense"] = "average",
+        na_option: Literal["keep", "top", "bottom"] = "keep",
+        ascending: bool = True,
+        pct: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Compute groupby with rank.
@@ -6624,6 +6670,27 @@ def groupby_groups(
         by: Any,
         axis: int,
         groupby_kwargs: dict[str, Any],
+    ) -> PrettyDict[Hashable, "pd.Index"]:
+        """
+        Wrapper around _groupby_groups_internal to be supported in faster pandas.
+        """
+        if self._relaxed_query_compiler is not None:
+            return self._relaxed_query_compiler._groupby_groups_internal(
+                by=by,
+                axis=axis,
+                groupby_kwargs=groupby_kwargs,
+            )
+        return self._groupby_groups_internal(
+            by=by,
+            axis=axis,
+            groupby_kwargs=groupby_kwargs,
+        )
+
+    def _groupby_groups_internal(
+        self,
+        by: Any,
+        axis: int,
+        groupby_kwargs: dict[str, Any],
     ) -> PrettyDict[Hashable, "pd.Index"]:
         """
         Get a PrettyDict mapping group keys to row labels.
@@ -6667,7 +6734,7 @@ def groupby_groups(
         4    5        2                     4                     5
         0    8        9                     0                     8
         """
-        self._raise_not_implemented_error_for_timedelta()
+        self._raise_not_implemented_error_for_timedelta(stack_depth=4)
 
         original_index_names = self.get_index_names()
         frame = self._modin_frame
@@ -6764,6 +6831,30 @@ def groupby_indices(
         axis: int,
         groupby_kwargs: dict[str, Any],
         values_as_np_array: bool = True,
+    ) -> dict[Hashable, np.ndarray]:
+        """
+        Wrapper around _groupby_indices_internal to be supported in faster pandas.
+        """
+        if self._relaxed_query_compiler is not None:
+            return self._relaxed_query_compiler._groupby_indices_internal(
+                by=by,
+                axis=axis,
+                groupby_kwargs=groupby_kwargs,
+                values_as_np_array=values_as_np_array,
+            )
+        return self._groupby_indices_internal(
+            by=by,
+            axis=axis,
+            groupby_kwargs=groupby_kwargs,
+            values_as_np_array=values_as_np_array,
+        )
+
+    def _groupby_indices_internal(
+        self,
+        by: Any,
+        axis: int,
+        groupby_kwargs: dict[str, Any],
+        values_as_np_array: bool = True,
     ) -> dict[Hashable, np.ndarray]:
         """
         Get a dict mapping group keys to row labels.

tests/integ/modin/test_faster_pandas.py

@@ -300,6 +300,12 @@ def test_groupby_no_param_functions(session, 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()
@@ -662,6 +668,36 @@ def test_groupby_apply(session):
         )
 
 
+@pytest.mark.parametrize("property_name", ["groups", "indices"])
+@sql_count_checker(query_count=3)
+def test_groupby_properties(session, property_name):
+    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([[2, 12], [2, 11], [3, 13]], columns=["A", "B"])
+        ).write.save_as_table(table_name, table_type="temp")
+
+        # create snow dataframes
+        df = pd.read_snowflake(table_name).sort_values("B", ignore_index=True)
+        snow_result = getattr(df.groupby("A"), property_name)
+
+        # 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
+
+        # create pandas dataframes
+        native_df = df.to_pandas()
+        native_result = getattr(native_df.groupby("A"), property_name)
+
+        # compare results
+        snow_result = {k: list(v) for k, v in snow_result.items()}
+        native_result = {k: list(v) for k, v in native_result.items()}
+        assert snow_result == native_result
+
+
 @sql_count_checker(query_count=5)
 def test_iloc_head(session):
     with session_parameter_override(