SNOW-2454948: Revert 2 modin commits to fix CI. (#3945)

This commit reverts 2 commits that broke tests, b505e92 and 226b9d9. 

b505e92 broke several tests. 226b9d9 broke tests/integ/modin/hybrid/test_switch_operations.py::test_query_count_no_switch[True] and tests/integ/modin/hybrid/test_switch_operations.py::test_query_count_no_switch[False].

GitHub allowed us to merge the commits because I think I never saved the branch protection rule requiring modin-ubuntu-AWS tests to pass before merging to main. I'm struggling with the branch protection rules so I am still trying to fix them.

Author: sfc-gh-mvashishtha

CHANGELOG.md

@@ -95,16 +95,6 @@
   - `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`
@@ -166,7 +156,6 @@
   - `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

@@ -626,7 +626,7 @@ def get_unsupported_args_reason(
 
 
 def register_query_compiler_method_not_implemented(
-    api_cls_names: Union[list[Optional[str]], Optional[str]],
+    api_cls_name: Optional[str],
     method_name: str,
     unsupported_args: Optional["UnsupportedArgsRule"] = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
@@ -644,30 +644,22 @@ def register_query_compiler_method_not_implemented(
     without meaningful benefit.
 
     Args:
-        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.
+        api_cls_name: Frontend class name (e.g., "BasePandasDataset", "Series", "DataFrame", "None").
         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)
 
-    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"
-
-    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 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
-        )
+    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)
@@ -2613,22 +2605,6 @@ 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,
@@ -4216,19 +4192,6 @@ 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,
         *,
@@ -4291,7 +4254,7 @@ def sort_index(
         1.0    c
         dtype: object
         """
-        if axis == 1:
+        if axis in (1, "index"):
             ErrorMessage.not_implemented(
                 "sort_index is not supported yet on axis=1 in Snowpark pandas."
             )
@@ -4315,17 +4278,8 @@ 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, axis: int = 1, **kwargs: Any
+        self, rows: IndexLabel, ascending: bool = True, **kwargs: Any
     ) -> None:
         """
         Reorder the columns based on the lexicographic order of the given rows.
@@ -4335,9 +4289,6 @@ 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.
 
@@ -4935,52 +4886,6 @@ 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,
@@ -9056,18 +8961,6 @@ 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],
@@ -10115,18 +10008,6 @@ 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],
@@ -10799,63 +10680,6 @@ 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,
@@ -10931,7 +10755,7 @@ def pivot_table(
             index = [index]
 
         # TODO: SNOW-857485 Support for non-str and list of non-str for index/columns/values
-        if index is not None and (
+        if index and (
             not isinstance(index, str)
             and not all([isinstance(v, str) for v in index])
             and None not in index
@@ -10940,7 +10764,7 @@ def pivot_table(
                 f"Not implemented non-string of list of string {index}."
             )
 
-        if values is not None and (
+        if values and (
             not isinstance(values, str)
             and not all([isinstance(v, str) for v in values])
             and None not in values
@@ -10949,7 +10773,7 @@ def pivot_table(
                 f"Not implemented non-string of list of string {values}."
             )
 
-        if columns is not None and (
+        if columns and (
             not isinstance(columns, str)
             and not all([isinstance(v, str) for v in columns])
             and None not in columns
@@ -13050,23 +12874,6 @@ 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,
@@ -13312,15 +13119,6 @@ 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,
@@ -22018,7 +21816,7 @@ def _stack_helper(
         return qc
 
     @register_query_compiler_method_not_implemented(
-        api_cls_names="DataFrame",
+        api_cls_name="DataFrame",
         method_name="corr",
         unsupported_args=UnsupportedArgsRule(
             unsupported_conditions=[