SNOW-2444072: Add support for groupby.apply in faster pandas (2nd attempt) (#3962)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -33,6 +33,7 @@
 #### Improvements
 
 - Add support for the following in faster pandas:
+  - `groupby.apply`
   - `groupby.nunique`
   - `groupby.size`
   - `concat`

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

@@ -4898,6 +4898,52 @@ def convert_func_to_agg_func_info(
         return query_compiler if as_index else query_compiler.reset_index(drop=drop)
 
     def groupby_apply(
+        self,
+        by: Any,
+        agg_func: AggFuncType,
+        axis: int,
+        groupby_kwargs: dict[str, Any],
+        agg_args: Any,
+        agg_kwargs: dict[str, Any],
+        series_groupby: bool,
+        include_groups: bool,
+        force_single_group: bool = False,
+        force_list_like_to_series: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_apply_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_apply_internal(
+                    by=by,
+                    agg_func=agg_func,
+                    axis=axis,
+                    groupby_kwargs=groupby_kwargs,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                    series_groupby=series_groupby,
+                    include_groups=include_groups,
+                    force_single_group=force_single_group,
+                    force_list_like_to_series=force_list_like_to_series,
+                )
+            )
+        qc = self._groupby_apply_internal(
+            by=by,
+            agg_func=agg_func,
+            axis=axis,
+            groupby_kwargs=groupby_kwargs,
+            agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+            series_groupby=series_groupby,
+            include_groups=include_groups,
+            force_single_group=force_single_group,
+            force_list_like_to_series=force_list_like_to_series,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _groupby_apply_internal(
         self,
         by: Any,
         agg_func: Callable,

tests/integ/modin/test_faster_pandas.py

@@ -504,6 +504,41 @@ def test_groupby_agg(session, func):
         assert_frame_equal(snow_result4, native_result4, check_dtype=False)
 
 
+@sql_count_checker(query_count=9, join_count=1, udtf_count=1)
+def test_groupby_apply(session):
+    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 = df.groupby("A").apply(lambda x: x + 1)
+
+        # 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 = native_df.groupby("A").apply(lambda x: x + 1)
+
+        # compare results
+        assert_frame_equal(
+            snow_result, native_result, check_dtype=False, check_index_type=False
+        )
+
+
 @sql_count_checker(query_count=5)
 def test_iloc_head(session):
     with session_parameter_override(