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

sfc-gh-jenzhang · web-flow · commit b505e92ac90a · 2025-10-22T16:35:22.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -95,6 +95,16 @@
   - `skew()` with `axis=1` or `numeric_only=False` parameters
   - `round()` with `decimals` parameter as a Series
   - `corr()` with `method!=pearson` parameter
+  - `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`
+
+
 - Set `cte_optimization_enabled` to True for all Snowpark pandas sessions.
 - Add support for the following in faster pandas:
   - `isin`
diff --git a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
@@ -626,7 +626,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 +644,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)
@@ -2605,6 +2613,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,
@@ -4192,6 +4216,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,
         *,
@@ -4254,7 +4291,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."
             )
@@ -4278,8 +4315,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.
@@ -4289,6 +4335,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.
 
@@ -9007,6 +9056,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],
@@ -10054,6 +10115,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],
@@ -10726,6 +10799,63 @@ def pivot(
             sort=True,
         )
 
+    @register_query_compiler_method_not_implemented(
+        None,
+        "pivot_table",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("sort", False),
+                (
+                    lambda args: (
+                        args.get("index") is not None
+                        and (
+                            not isinstance(args.get("index"), str)
+                            and not all([isinstance(v, str) for v in args.get("index")])
+                            and None not in args.get("index")
+                        )
+                    ),
+                    "non-string of list of string index is not yet supported for pivot_table",
+                ),
+                (
+                    lambda args: (
+                        args.get("columns") is not None
+                        and (
+                            not isinstance(args.get("columns"), str)
+                            and not all(
+                                [isinstance(v, str) for v in args.get("columns")]
+                            )
+                            and None not in args.get("columns")
+                        )
+                    ),
+                    "non-string of list of string columns is not yet supported for pivot_table",
+                ),
+                (
+                    lambda args: (
+                        args.get("values") is not None
+                        and (
+                            not isinstance(args.get("values"), str)
+                            and not all(
+                                [isinstance(v, str) for v in args.get("values")]
+                            )
+                            and None not in args.get("values")
+                        )
+                    ),
+                    "non-string of list of string values is not yet supported for pivot_table",
+                ),
+                (
+                    lambda args: (
+                        isinstance(args.get("aggfunc"), dict)
+                        and any(
+                            not isinstance(af, str)
+                            for af in args.get("aggfunc").values()
+                        )
+                        and args.get("index") is None
+                    ),
+                    "dictionary aggfunc with non-string aggregation functions is not yet supported for pivot_table with margins or when index is None",
+                ),
+            ]
+        ),
+    )
     def pivot_table(
         self,
         index: Any,
@@ -10801,7 +10931,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
@@ -10810,7 +10940,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
@@ -10819,7 +10949,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
@@ -12920,6 +13050,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,
@@ -13165,6 +13312,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,
@@ -21862,7 +22018,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=[
diff --git a/tests/integ/modin/hybrid/test_switch_operations.py b/tests/integ/modin/hybrid/test_switch_operations.py
