SNOW-2022995: Enable automatic backend switching for unimplemented pandas methods (#3512)

When AutoSwitchBackend (hybrid execution PrPr) is enabled, the client will automatically switch to the pandas backend when a method registered via a register_*_not_implemented annotation is called.

Author: sfc-gh-joshi

CHANGELOG.md

@@ -47,7 +47,8 @@ local ingestion. By default, local ingestion uses multithreading. Multiprocessin
 - Support `pd.read_snowflake` when the global modin backend is `Pandas`.
 
 #### Improvements
-- Add a data type guard to the cost functions for hybrid execution which checks for data type compatibility
+- Add a data type guard to the cost functions for hybrid execution mode (PrPr) which checks for data type compatibility.
+- Added automatic switching to the pandas backend in hybrid execution mode (PrPr) for many methods that are not directly implemented in Snowpark pandas.
 
 #### Dependency Updates
 - Added tqdm and ipywidgets as dependencies so that progress bars appear when switching between modin backends.

docs/source/modin/hybrid_execution.rst

@@ -19,6 +19,12 @@ Example Pre-Operation Switchpoints:
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 apply, iterrows, itertuples, items, plot, quantile, __init__, plot, quantile, T, read_csv, read_json, concat, merge 
 
+Many methods that are not yet implemented in Snowpark pandas are also registered as
+pre-operation switch points, and will automatically move data to local pandas for execution when
+called. This includes most methods that are ordinarily completely unsupported by Snowpark pandas,
+and have `N` in their implemented status in the :doc:`DataFrame <supported/dataframe_supported>` and
+:doc:`Series <supported/series_supported>` supported API lists.
+
 Post-Operation Switchpoints:
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 read_snowflake, value_counts, tail, var, std, sum, sem, max, min, mean, agg, aggregate, count, nunique, cummax, cummin, cumprod, cumsum

src/snowflake/snowpark/modin/plugin/__init__.py

@@ -185,6 +185,9 @@
     AutoSwitchBackend().disable()
 
     # Hybrid Mode Registration
+    # In hybrid execution mode, the client will automatically switch backends when a
+    # wholly-unimplemented method is called. Those switch points are registered separately in
+    # extensions files via the register_*_not_implemented family of methods.
     pre_op_switch_points: list[dict[str, Union[str, None]]] = [
         {"class_name": "DataFrame", "method": "__init__"},
         {"class_name": "Series", "method": "__init__"},

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

@@ -18,7 +18,18 @@
 from datetime import timedelta, tzinfo
 from functools import reduce
 from types import MappingProxyType
-from typing import Any, Callable, List, Literal, Optional, TypeVar, Union, get_args
+from typing import (
+    Any,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    get_args,
+    Set,
+    Tuple,
+)
 
 import modin.pandas as pd
 from modin.pandas import Series, DataFrame
@@ -481,6 +492,9 @@
 HYBRID_ALL_EXPENSIVE_METHODS = (
     HYBRID_HIGH_OVERHEAD_METHODS + HYBRID_ITERATIVE_STYLE_METHODS
 )
+# Set of (class name, method name) tuples for methods that are wholly unimplemented by
+# Snowpark pandas. This list is populated by the register_*_not_implemented decorators.
+HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS: Set[Tuple[str, str]] = set()
 
 T = TypeVar("T", bound=Callable[..., Any])
 
@@ -858,7 +872,10 @@ def stay_cost(
         operation: str,
         arguments: MappingProxyType[str, Any],
     ) -> Optional[int]:
-        if self._is_in_memory_init(api_cls_name, operation, arguments):
+        if (
+            self._is_in_memory_init(api_cls_name, operation, arguments)
+            or (api_cls_name, operation) in HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS
+        ):
             return QCCoercionCost.COST_IMPOSSIBLE
         # Strongly discourage the use of these methods in snowflake
         if operation in HYBRID_ALL_EXPENSIVE_METHODS:
@@ -891,9 +908,11 @@ def move_to_me_cost(
             None if the cost cannot be determined.
         """
         # in-memory intialization should not move to Snowflake
-        if cls._is_in_memory_init(api_cls_name, operation, arguments):
-            return QCCoercionCost.COST_IMPOSSIBLE
-        if not cls._are_dtypes_compatible_with_snowflake(other_qc):
+        if (
+            cls._is_in_memory_init(api_cls_name, operation, arguments)
+            or not cls._are_dtypes_compatible_with_snowflake(other_qc)
+            or (api_cls_name, operation) in HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS
+        ):
             return QCCoercionCost.COST_IMPOSSIBLE
         # Strongly discourage the use of these methods in snowflake
         if operation in HYBRID_ALL_EXPENSIVE_METHODS:

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

@@ -68,6 +68,9 @@
     MODIN_IS_AT_LEAST_0_34_0,
 )
 from snowflake.snowpark.modin.plugin._typing import ListLike
+from snowflake.snowpark.modin.plugin.compiler.snowflake_query_compiler import (
+    HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS,
+)
 from snowflake.snowpark.modin.plugin.extensions.utils import (
     ensure_index,
     extract_validate_and_try_convert_named_aggs_from_kwargs,
@@ -91,14 +94,23 @@
 
 if MODIN_IS_AT_LEAST_0_33_0:
     from modin.pandas.api.extensions import register_base_accessor
+    from modin.core.storage_formats.pandas.query_compiler_caster import (
+        register_function_for_pre_op_switch,
+    )
 
     register_base_override = functools.partial(
         register_base_accessor, backend="Snowflake"
     )
 
     def register_base_not_implemented():
         def decorator(base_method: Any):
-            return register_base_override(name=base_method.__name__)(
+            name = base_method.__name__
+            HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS.add(("BasePandasDataset", name))
+            if MODIN_IS_AT_LEAST_0_33_0:
+                register_function_for_pre_op_switch(
+                    class_name="BasePandasDataset", backend="Snowflake", method=name
+                )
+            return register_base_override(name=name)(
                 base_not_implemented()(base_method)
             )
 

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

@@ -85,6 +85,7 @@
 from snowflake.snowpark.modin.plugin._typing import ListLike
 from snowflake.snowpark.modin.plugin.compiler.snowflake_query_compiler import (
     SnowflakeQueryCompiler,
+    HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS,
 )
 from snowflake.snowpark.modin.plugin.extensions.index import Index
 from snowflake.snowpark.modin.plugin.extensions.snow_partition_iterator import (
@@ -123,6 +124,9 @@
     from modin.pandas.api.extensions import (
         register_dataframe_accessor as _register_dataframe_accessor,
     )
+    from modin.core.storage_formats.pandas.query_compiler_caster import (
+        register_function_for_pre_op_switch,
+    )
 
     register_dataframe_accessor = functools.partial(
         _register_dataframe_accessor, backend="Snowflake"
@@ -138,7 +142,13 @@
 def register_dataframe_not_implemented():
     def decorator(base_method: Any):
         func = dataframe_not_implemented()(base_method)
-        register_dataframe_accessor(base_method.__name__)(func)
+        name = base_method.__name__
+        HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS.add(("DataFrame", name))
+        if MODIN_IS_AT_LEAST_0_33_0:
+            register_function_for_pre_op_switch(
+                class_name="DataFrame", backend="Snowflake", method=name
+            )
+        register_dataframe_accessor(name)(func)
         return func
 
     return decorator

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

@@ -50,6 +50,9 @@
     error_checking_for_init,
     MODIN_IS_AT_LEAST_0_33_0,
 )
+from snowflake.snowpark.modin.plugin.compiler.snowflake_query_compiler import (
+    HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS,
+)
 from snowflake.snowpark.modin.plugin._typing import DropKeep, ListLike
 from snowflake.snowpark.modin.plugin.extensions.snow_partition_iterator import (
     SnowparkPandasRowPartitionIterator,
@@ -80,6 +83,9 @@
     from modin.pandas.api.extensions import (
         register_series_accessor as _register_series_accessor,
     )
+    from modin.core.storage_formats.pandas.query_compiler_caster import (
+        register_function_for_pre_op_switch,
+    )
 
     register_series_accessor = functools.partial(
         _register_series_accessor, backend="Snowflake"
@@ -96,6 +102,11 @@ def decorator(base_method: Any):
             if isinstance(base_method, property)
             else base_method.__name__
         )
+        HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS.add(("Series", name))
+        if MODIN_IS_AT_LEAST_0_33_0:
+            register_function_for_pre_op_switch(
+                class_name="Series", backend="Snowflake", method=name
+            )
         register_series_accessor(name)(func)
         return func
 

tests/integ/modin/hybrid/test_switch_operations.py

@@ -8,8 +8,10 @@
 from unittest.mock import patch
 import tqdm.auto
 
+import pandas as native_pd
 import numpy as np
 from numpy.testing import assert_array_equal
+from modin.config import context as config_context
 import modin.pandas as pd
 import snowflake.snowpark.modin.plugin  # noqa: F401
 from snowflake.snowpark.modin.plugin._internal.utils import (
@@ -265,3 +267,42 @@ def test_tqdm_usage_during_snowflake_to_pandas_switch():
         df.set_backend("Pandas")
 
     mock_trange.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "class_name, method_name, f_args",
+    [
+        ("DataFrame", "to_json", ()),  # declared in base_overrides
+        ("Series", "to_json", ()),  # declared in base_overrides
+        ("DataFrame", "dot", ([6],)),  # declared in dataframe_overrides
+        ("Series", "transform", (lambda x: x * 2,)),  # declared in series_overrides
+    ],
+)
+@sql_count_checker(query_count=1)
+def test_unimplemented_autoswitches(class_name, method_name, f_args):
+    # Unimplemented methods declared via register_*_not_implemented should automatically
+    # default to local pandas execution.
+    # This test needs to be modified if any of the APIs in question are ever natively implemented
+    # for Snowpark pandas.
+    data = [1, 2, 3]
+    method = getattr(getattr(pd, class_name)(data).move_to("Snowflake"), method_name)
+    # Attempting to call the method without switching should raise.
+    with config_context(AutoSwitchBackend=False):
+        with pytest.raises(
+            NotImplementedError, match="Snowpark pandas does not yet support the method"
+        ):
+            method(*f_args)
+    # Attempting to call the method while switching is enabled should work fine.
+    snow_result = method(*f_args)
+    pandas_result = getattr(getattr(native_pd, class_name)(data), method_name)(*f_args)
+    if isinstance(snow_result, (pd.DataFrame, pd.Series)):
+        assert snow_result.get_backend() == "Pandas"
+        assert_array_equal(snow_result.to_numpy(), pandas_result.to_numpy())
+    else:
+        # Series.to_json will output an extraneous level for the __reduced__ column, but that's OK
+        # since we don't officially support the method.
+        # See modin bug: https://github.com/modin-project/modin/issues/7624
+        if class_name == "Series" and method_name == "to_json":
+            assert snow_result == '{"__reduced__":{"0":1,"1":2,"2":3}}'
+        else:
+            assert snow_result == pandas_result