SNOW-2359402 Enabled autoswitching on some unsupported args

CHANGELOG.md

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

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: List[Optional[str]],
     method_name: str,
     unsupported_args: Optional["UnsupportedArgsRule"] = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
@@ -644,22 +644,24 @@ 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", None]). This is a list because some methods are implemented for both DataFrames and Series.
         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
+    for api_cls_name in api_cls_names:
+        reg_key = MethodKey(api_cls_name, method_name)
 
-    register_function_for_pre_op_switch(
-        class_name=api_cls_name, backend="Snowflake", method=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)
@@ -2591,6 +2593,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,
@@ -4178,6 +4196,22 @@ def first_last_valid_index(
             )
         return None
 
+    @register_query_compiler_method_not_implemented(
+        ["BasePandasDataset"],
+        "sort_index",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                (
+                    lambda args: args.get("axis") in (1, "index"),
+                    "axis=1 is not yet supported",
+                ),
+                (
+                    lambda args: args.get("key") is not None,
+                    "the 'key' parameter is not yet supported",
+                ),
+            ]
+        ),
+    )
     def sort_index(
         self,
         *,
@@ -4264,8 +4298,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.
@@ -6997,7 +7040,7 @@ def groupby_pct_change(
         )
 
     @register_query_compiler_method_not_implemented(
-        None,
+        [None],
         "get_dummies",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -8545,7 +8588,7 @@ def concat(
         return qc
 
     @register_query_compiler_method_not_implemented(
-        "BasePandasDataset",
+        ["BasePandasDataset"],
         "cumsum",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -8583,7 +8626,7 @@ def cumsum(
         )
 
     @register_query_compiler_method_not_implemented(
-        "BasePandasDataset",
+        ["BasePandasDataset"],
         "cummin",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -8621,7 +8664,7 @@ def cummin(
         )
 
     @register_query_compiler_method_not_implemented(
-        "BasePandasDataset",
+        ["BasePandasDataset"],
         "cummax",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -8658,6 +8701,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],
@@ -9705,6 +9760,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],
@@ -10377,6 +10444,63 @@ def pivot(
             sort=True,
         )
 
+    @register_query_compiler_method_not_implemented(
+        [None],
+        "pivot_table",
+        UnsupportedArgsRule(
+            unsupported_conditions=[
+                ("sort", False),
+                (
+                    lambda args: (
+                        args.get("index")
+                        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")
+                        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")
+                        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,
@@ -12541,6 +12665,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,
@@ -12786,6 +12927,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,
@@ -15281,7 +15431,7 @@ def _quantiles_single_col(
         return SnowflakeQueryCompiler(internal_frame)
 
     @register_query_compiler_method_not_implemented(
-        "BasePandasDataset",
+        ["BasePandasDataset"],
         "skew",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -17362,7 +17512,7 @@ def create_lazy_type_functions(
         return SnowflakeQueryCompiler(new_frame)
 
     @register_query_compiler_method_not_implemented(
-        "BasePandasDataset",
+        ["BasePandasDataset"],
         "round",
         UnsupportedArgsRule(
             unsupported_conditions=[
@@ -21239,7 +21389,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=[

src/snowflake/snowpark/modin/plugin/extensions/base_overrides.py

@@ -1623,6 +1623,7 @@ def sort_values(
             na_position=na_position,
             ignore_index=ignore_index,
             key=key,
+            axis=axis,
         )
     return self._create_or_update_from_compiler(result, inplace)