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

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -123,6 +123,8 @@
   - `cumsum`
   - `cummin`
   - `cummax`
+  - `groupby.groups`
+  - `groupby.indices`
 - Make faster pandas disabled by default (opt-in instead of opt-out).
 - Improve performance of `drop_duplicates` by avoiding joins when `keep!=False` in faster pandas.
 

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

@@ -6513,6 +6513,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.
@@ -6653,6 +6674,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

@@ -619,6 +619,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(