SNOW-2359402 Enabled autoswitching on some unsupported args (#3953)

Author: sfc-gh-jenzhang

CHANGELOG.md

@@ -26,6 +26,18 @@
 - Added support for mapping `np.percentile` with DataFrame and Series inputs to `Series.quantile`.
 - Added support for setting the `random_state` parameter to an integer when calling `DataFrame.sample` or `Series.sample`.
 
+#### Improvements
+
+- Enhanced autoswitching functionality from Snowflake to native Pandas for methods with unsupported argument combinations:
+  - `shift()` with `suffix` or non-integer `periods` parameters
+  - `sort_index()` with `axis=1` or `key` parameters
+  - `sort_values()` with `axis=1`
+  - `melt()` with `col_level` parameter
+  - `apply()` with `result_type` parameter for DataFrame
+  - `pivot_table()` with `sort=True`, non-string `index` list, non-string `columns` list, non-string `values` list, or `aggfunc` dict with non-string values
+  - `fillna()` with `downcast` parameter or using `limit` together with `value`
+  - `dropna()` with `axis=1`
+
 #### Bug Fixes
 
 - Fixed a bug in `DataFrameGroupBy.agg` where func is a list of tuples used to set the names of the output columns.

src/snowflake/snowpark/modin/plugin/_internal/pivot_utils.py

@@ -1866,3 +1866,58 @@ def generate_column_prefix_groupings(
         )
 
     return list(zip(margin_data_column_prefixes, margin_data_column_groupings))
+
+
+def check_pivot_table_unsupported_args(args: dict) -> Optional[str]:
+    """
+    Validate pivot_table arguments for unsupported conditions.
+
+    This helper function checks various argument combinations that are not yet
+    supported by Snowpark pandas pivot_table implementation.
+
+    Args:
+        args : dictionary of arguments passed to pivot_table
+
+    Returns:
+        Error message if an unsupported condition is found, None otherwise
+    """
+    # Check if index argument is a string or list of strings
+    index = args.get("index")
+    if (
+        index is not None
+        and not isinstance(index, str)
+        and not all(isinstance(v, str) for v in index)
+        and None not in index
+    ):
+        return "index argument should be a string or a list of strings"
+
+    # Check if columns argument is a string or list of strings
+    columns = args.get("columns")
+    if (
+        columns is not None
+        and not isinstance(columns, str)
+        and not all(isinstance(v, str) for v in columns)
+        and None not in columns
+    ):
+        return "columns argument should be a string or a list of strings"
+
+    # Check if values argument is a string or list of strings
+    values = args.get("values")
+    if (
+        values is not None
+        and not isinstance(values, str)
+        and not all(isinstance(v, str) for v in values)
+        and None not in values
+    ):
+        return "values argument should be a string or a list of strings"
+
+    # Check for dictionary aggfunc with non-string functions when index is None
+    aggfunc = args.get("aggfunc")
+    if (
+        isinstance(aggfunc, dict)
+        and any(not isinstance(af, str) for af in aggfunc.values())
+        and args.get("index") is None
+    ):
+        return "dictionary aggfunc with non-string aggregation functions is not yet supported for pivot_table when index is None"
+
+    return None

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

@@ -298,6 +298,7 @@
     OrderingColumn,
 )
 from snowflake.snowpark.modin.plugin._internal.pivot_utils import (
+    check_pivot_table_unsupported_args,
     expand_pivot_result_with_pivot_table_margins,
     expand_pivot_result_with_pivot_table_margins_no_groupby_columns,
     generate_pivot_aggregation_value_label_snowflake_quoted_identifier_mappings,
@@ -626,7 +627,7 @@ def get_unsupported_args_reason(
 
 
 def register_query_compiler_method_not_implemented(
-    api_cls_name: Optional[str],
+    api_cls_names: Union[list[Optional[str]], Optional[str]],
     method_name: str,
     unsupported_args: Optional["UnsupportedArgsRule"] = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
@@ -644,22 +645,30 @@ def register_query_compiler_method_not_implemented(
     without meaningful benefit.
 
     Args:
-        api_cls_name: Frontend class name (e.g., "BasePandasDataset", "Series", "DataFrame", "None").
+        api_cls_names: Frontend class names (e.g. "BasePandasDataset", "Series", "DataFrame", or None). It can be a list if multiple api_cls_names are needed.
         method_name: Method name to register.
         unsupported_args: UnsupportedArgsRule for args-based auto-switching.
                           If None, method is treated as completely unimplemented.
     """
-    reg_key = MethodKey(api_cls_name, method_name)
 
-    # register the method in the hybrid switch for unsupported args
-    if unsupported_args is None:
-        HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS.add(reg_key)
-    else:
-        HYBRID_SWITCH_FOR_UNSUPPORTED_ARGS[reg_key] = unsupported_args
+    if isinstance(api_cls_names, str) or api_cls_names is None:
+        api_cls_names = [api_cls_names]
+    assert (
+        api_cls_names
+    ), "api_cls_names must be a string (e.g., 'DataFrame', 'Series') or a list of strings (e.g., ['DataFrame', 'Series']) or None for top-level functions"
 
-    register_function_for_pre_op_switch(
-        class_name=api_cls_name, backend="Snowflake", method=method_name
-    )
+    for api_cls_name in api_cls_names:
+        reg_key = MethodKey(api_cls_name, method_name)
+
+        # register the method in the hybrid switch for unsupported args
+        if unsupported_args is None:
+            HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS.add(reg_key)
+        else:
+            HYBRID_SWITCH_FOR_UNSUPPORTED_ARGS[reg_key] = unsupported_args
+
+        register_function_for_pre_op_switch(
+            class_name=api_cls_name, backend="Snowflake", method=method_name
+        )
 
     def decorator(query_compiler_method: Callable[..., Any]) -> Callable[..., Any]:
         @functools.wraps(query_compiler_method)
@@ -1154,11 +1163,13 @@ def stay_cost(
 
                 return QCCoercionCost.COST_IMPOSSIBLE
 
-            return QCCoercionCost.COST_ZERO
-
         # Strongly discourage the use of these methods in snowflake
         if operation in HYBRID_ALL_EXPENSIVE_METHODS:
             return QCCoercionCost.COST_HIGH
+
+        if method_key in HYBRID_SWITCH_FOR_UNSUPPORTED_ARGS:
+            return QCCoercionCost.COST_ZERO
+
         return super().stay_cost(api_cls_name, operation, arguments)
 
     @classmethod
@@ -2615,6 +2626,22 @@ def _shift_index(self, periods: int, freq: Any) -> "SnowflakeQueryCompiler":  #
         # TODO: SNOW-1023324, implement shifting index only.
         ErrorMessage.not_implemented("shifting index values not yet supported.")
 
+    @register_query_compiler_method_not_implemented(
+        "BasePandasDataset",
+        "shift",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                (
+                    lambda args: args.get("suffix") is not None,
+                    "the 'suffix' parameter is not yet supported",
+                ),
+                (
+                    lambda args: not isinstance(args.get("periods"), int),
+                    "only int 'periods' is currently supported",
+                ),
+            ]
+        ),
+    )
     def shift(
         self,
         periods: Union[int, Sequence[int]] = 1,
@@ -4204,6 +4231,19 @@ def first_last_valid_index(
             )
         return None
 
+    @register_query_compiler_method_not_implemented(
+        "BasePandasDataset",
+        "sort_index",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("axis", 1),
+                (
+                    lambda args: args.get("key") is not None,
+                    "the 'key' parameter is not yet supported",
+                ),
+            ]
+        ),
+    )
     def sort_index(
         self,
         *,
@@ -4266,7 +4306,7 @@ def sort_index(
         1.0    c
         dtype: object
         """
-        if axis in (1, "index"):
+        if axis == 1:
             ErrorMessage.not_implemented(
                 "sort_index is not supported yet on axis=1 in Snowpark pandas."
             )
@@ -4290,8 +4330,17 @@ def sort_index(
             include_indexer=include_indexer,
         )
 
+    @register_query_compiler_method_not_implemented(
+        "BasePandasDataset",
+        "sort_values",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("axis", 1),
+            ]
+        ),
+    )
     def sort_columns_by_row_values(
-        self, rows: IndexLabel, ascending: bool = True, **kwargs: Any
+        self, rows: IndexLabel, ascending: bool = True, axis: int = 1, **kwargs: Any
     ) -> None:
         """
         Reorder the columns based on the lexicographic order of the given rows.
@@ -4301,6 +4350,9 @@ def sort_columns_by_row_values(
                 The row or rows to sort by.
             ascending : bool, default: True
                 Sort in ascending order (True) or descending order (False).
+            axis: Always set to 1. Required because the decorator compares frontend
+                method arguments during stay_cost computation (returning COST_IMPOSSIBLE)
+                but examines QC method arguments when calling the wrapped method.
             **kwargs : dict
                 Serves the compatibility purpose. Does not affect the result.
 
@@ -9085,6 +9137,18 @@ def cummax(
             ).frame
         )
 
+    @register_query_compiler_method_not_implemented(
+        None,
+        "melt",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                (
+                    lambda args: args.get("col_level") is not None,
+                    "col_level argument is not yet supported",
+                ),
+            ]
+        ),
+    )
     def melt(
         self,
         id_vars: list[str],
@@ -10142,6 +10206,18 @@ def align(
 
         return left_qc, right_qc
 
+    @register_query_compiler_method_not_implemented(
+        "DataFrame",
+        "apply",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                (
+                    lambda args: args.get("result_type") is not None,
+                    "the 'result_type' parameter is not yet supported",
+                ),
+            ]
+        ),
+    )
     def apply(
         self,
         func: Union[AggFuncType, UserDefinedFunction],
@@ -10814,6 +10890,20 @@ def pivot(
             sort=True,
         )
 
+    @register_query_compiler_method_not_implemented(
+        None,
+        "pivot_table",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("observed", True),
+                ("sort", False),
+                (
+                    lambda args: check_pivot_table_unsupported_args(args) is not None,
+                    check_pivot_table_unsupported_args,
+                ),
+            ]
+        ),
+    )
     def pivot_table(
         self,
         index: Any,
@@ -10889,7 +10979,7 @@ def pivot_table(
             index = [index]
 
         # TODO: SNOW-857485 Support for non-str and list of non-str for index/columns/values
-        if index and (
+        if index is not None and (
             not isinstance(index, str)
             and not all([isinstance(v, str) for v in index])
             and None not in index
@@ -10898,7 +10988,7 @@ def pivot_table(
                 f"Not implemented non-string of list of string {index}."
             )
 
-        if values and (
+        if values is not None and (
             not isinstance(values, str)
             and not all([isinstance(v, str) for v in values])
             and None not in values
@@ -10907,7 +10997,7 @@ def pivot_table(
                 f"Not implemented non-string of list of string {values}."
             )
 
-        if columns and (
+        if columns is not None and (
             not isinstance(columns, str)
             and not all([isinstance(v, str) for v in columns])
             and None not in columns
@@ -13021,6 +13111,23 @@ def _make_fill_expression_for_column_wise_fillna(
                 *columns_to_include,
             )
 
+    @register_query_compiler_method_not_implemented(
+        ["DataFrame", "Series"],
+        "fillna",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                (
+                    lambda kwargs: kwargs.get("value") is not None
+                    and kwargs.get("limit") is not None,
+                    "the 'limit' parameter with 'value' parameter is not yet supported",
+                ),
+                (
+                    lambda kwargs: kwargs.get("downcast") is not None,
+                    "the 'downcast' parameter is not yet supported",
+                ),
+            ]
+        ),
+    )
     def fillna(
         self,
         value: Optional[Union[Hashable, Mapping, "pd.DataFrame", "pd.Series"]] = None,
@@ -13266,6 +13373,15 @@ def fillna_expr(snowflake_quoted_id: str) -> SnowparkColumn:
             ).frame
         )
 
+    @register_query_compiler_method_not_implemented(
+        "DataFrame",
+        "dropna",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("axis", 1),
+            ]
+        ),
+    )
     def dropna(
         self,
         axis: int,
@@ -22121,7 +22237,7 @@ def _stack_helper(
         return qc
 
     @register_query_compiler_method_not_implemented(
-        api_cls_name="DataFrame",
+        api_cls_names="DataFrame",
         method_name="corr",
         unsupported_args=UnsupportedArgsRule(
             unsupported_conditions=[

tests/integ/modin/frame/test_apply.py

@@ -377,7 +377,7 @@ class TestNotImplemented:
     @sql_count_checker(query_count=0)
     def test_result_type(self, result_type):
         snow_df = pd.DataFrame([[1, 2], [3, 4]])
-        msg = "Snowpark pandas apply API doesn't yet support 'result_type' parameter"
+        msg = "Snowpark pandas apply does not yet support the parameter combination because the 'result_type' parameter is not yet supported."
         with pytest.raises(NotImplementedError, match=msg):
             snow_df.apply(lambda x: [1, 2], axis=1, result_type=result_type)
 

tests/integ/modin/frame/test_apply_axis_0.py

@@ -242,7 +242,7 @@ class TestNotImplemented:
     @sql_count_checker(query_count=0)
     def test_result_type(self, result_type):
         snow_df = pd.DataFrame([[1, 2], [3, 4]])
-        msg = "Snowpark pandas apply API doesn't yet support 'result_type' parameter"
+        msg = "Snowpark pandas apply does not yet support the parameter combination because the 'result_type' parameter is not yet supported."
         with pytest.raises(NotImplementedError, match=msg):
             snow_df.apply(lambda x: [1, 2], axis=0, result_type=result_type)