SNOW-2444072: Add support for groupby.apply in faster pandas

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -152,6 +152,9 @@
   - `groupby.median`
   - `groupby.std`
   - `groupby.var`
+  - `groupby.nunique`
+  - `groupby.size`
+  - `groupby.apply`
   - `drop_duplicates`
 - Reuse row count from the relaxed query compiler in `get_axis_len`.
 

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

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

@@ -392,6 +392,37 @@ 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):
+    # 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):
     # create tables