SNOW-2441444: Add support for groupby.nunique/size in faster pandas (#3923)

Author: sfc-gh-helmeleegy

CHANGELOG.md

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

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

@@ -6078,6 +6078,43 @@ def groupby_size(
         agg_kwargs: dict[str, Any],
         drop: bool = False,
         **kwargs: dict[str, Any],
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_size_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_size_internal(
+                    by=by,
+                    axis=axis,
+                    groupby_kwargs=groupby_kwargs,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                    drop=drop,
+                    **kwargs,
+                )
+            )
+        qc = self._groupby_size_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_size_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":
         """
         compute groupby with size.
@@ -6495,6 +6532,43 @@ def groupby_nunique(
         agg_kwargs: dict[str, Any],
         drop: bool = False,
         **kwargs: Any,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _groupby_nunique_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_nunique_internal(
+                    by=by,
+                    axis=axis,
+                    groupby_kwargs=groupby_kwargs,
+                    agg_args=agg_args,
+                    agg_kwargs=agg_kwargs,
+                    drop=drop,
+                    **kwargs,
+                )
+            )
+        qc = self._groupby_nunique_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_nunique_internal(
+        self,
+        by: Any,
+        axis: int,
+        groupby_kwargs: dict[str, Any],
+        agg_args: Any,
+        agg_kwargs: dict[str, Any],
+        drop: bool = False,
+        **kwargs: Any,
     ) -> "SnowflakeQueryCompiler":
         # We have to override the Modin version of this function because our groupby frontend passes the
         # ignored numeric_only argument to this query compiler method, and BaseQueryCompiler

tests/integ/modin/test_faster_pandas.py

@@ -348,6 +348,8 @@ def test_duplicated(session):
         "median",
         "std",
         "var",
+        "nunique",
+        "size",
     ],
 )
 @sql_count_checker(query_count=6)
@@ -386,7 +388,10 @@ def test_groupby_agg(session, func):
     native_result4 = native_df.groupby("A")["B"].agg([func])
 
     # compare results
-    assert_frame_equal(snow_result1, native_result1, check_dtype=False)
+    if func == "size":
+        assert_series_equal(snow_result1, native_result1, check_dtype=False)
+    else:
+        assert_frame_equal(snow_result1, native_result1, check_dtype=False)
     assert_frame_equal(snow_result2, native_result2, check_dtype=False)
     assert_series_equal(snow_result3, native_result3, check_dtype=False)
     assert_frame_equal(snow_result4, native_result4, check_dtype=False)