address comments

Author: sfc-gh-helmeleegy

src/snowflake/snowpark/session.py

@@ -763,10 +763,7 @@ def __init__(
                     _SNOWPARK_PANDAS_HYBRID_EXECUTION_ENABLED, AutoSwitchBackend().get()
                 )
             )
-            if pandas_hybrid_execution_enabled:
-                AutoSwitchBackend.enable()
-            else:
-                AutoSwitchBackend.disable()
+            AutoSwitchBackend.put(pandas_hybrid_execution_enabled)
 
         self._thread_store = create_thread_local(
             self._conn._thread_safe_session_enabled
@@ -1046,8 +1043,9 @@ def pandas_hybrid_execution_enabled(self) -> bool:
         This can significantly improve performance for operations that are more efficient in pandas than in Snowflake.
         """
         if not importlib.util.find_spec("modin"):
-            # If modin is not installed, always return False
-            return False
+            raise ImportError(
+                "The 'modin' package is required to enable this feature. Please install it first."
+            )
 
         from modin.config import AutoSwitchBackend
 
@@ -1232,16 +1230,14 @@ def dummy_row_pos_optimization_enabled(self, value: bool) -> None:
     def pandas_hybrid_execution_enabled(self, value: bool) -> None:
         """Set the value for pandas_hybrid_execution_enabled"""
         if not importlib.util.find_spec("modin"):
-            # If modin is not installed, treat this method as a no-op.
-            return
+            raise ImportError(
+                "The 'modin' package is required to enable this feature. Please install it first."
+            )
 
         from modin.config import AutoSwitchBackend
 
         if value in [True, False]:
-            if value:
-                AutoSwitchBackend.enable()
-            else:
-                AutoSwitchBackend.disable()
+            AutoSwitchBackend.put(value)
         else:
             raise ValueError(
                 "value for pandas_hybrid_execution_enabled must be True or False!"

tests/integ/modin/hybrid/test_data_movement.py

@@ -45,18 +45,6 @@ def test_unsupported_movement(session, pandas_df):
         assert move_from_result is NotImplemented
 
 
-@sql_count_checker(query_count=0)
-def test_unsupported_movement_using_hybrid_session_param(session, pandas_df):
-    with config_context(Backend="Snowflake"):
-        session.pandas_hybrid_execution_enabled = False
-        snow_df = pd.DataFrame(pandas_df)
-        mock_qc = MockQueryCompiler()
-        move_to_result = snow_df._query_compiler.move_to("UnsupportedBackend")
-        move_from_result = SnowflakeQueryCompiler.move_from(mock_qc)
-        assert move_to_result is NotImplemented
-        assert move_from_result is NotImplemented
-
-
 @pytest.mark.skipif(
     sys.version_info.minor >= 12,
     reason="snowflake-ml-python for efficient movement is not installed above python 3.12",
@@ -73,23 +61,6 @@ def test_move_to_ray(session, pandas_df):
         df_equals(result_df, snow_df)
 
 
-@pytest.mark.skipif(
-    sys.version_info.minor >= 12,
-    reason="snowflake-ml-python for efficient movement is not installed above python 3.12",
-)
-@sql_count_checker(query_count=9)
-def test_move_to_ray_using_hybrid_session_param(session, pandas_df):
-    with config_context(Backend="Snowflake"):
-        session.pandas_hybrid_execution_enabled = False
-        snow_df = pd.DataFrame(pandas_df)
-        assert snow_df.get_backend() == "Snowflake"
-        result = snow_df._query_compiler.move_to("Ray")
-        result_df = pd.DataFrame(query_compiler=result)
-        assert result_df.get_backend() == "Ray"
-        assert Backend.get() == "Snowflake"
-        df_equals(result_df, snow_df)
-
-
 @pytest.mark.skip(reason="SNOW-2276090")
 @sql_count_checker(query_count=4)
 def test_move_from_ray(session, pandas_df):

tests/integ/modin/hybrid/test_switch_operations.py

@@ -395,8 +395,9 @@ def test_tqdm_usage_during_snowflake_to_pandas_switch():
         ("Series", "transform", (lambda x: x * 2,)),  # declared in series_overrides
     ],
 )
+@pytest.mark.parametrize("use_session_param", [True, False])
 @sql_count_checker(query_count=1)
-def test_unimplemented_autoswitches(class_name, method_name, f_args):
+def test_unimplemented_autoswitches(class_name, method_name, f_args, use_session_param):
     # 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
@@ -405,6 +406,12 @@ def test_unimplemented_autoswitches(class_name, method_name, f_args):
     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):
+        if use_session_param:
+            from modin.config import AutoSwitchBackend
+
+            AutoSwitchBackend.enable()
+            pd.session.pandas_hybrid_execution_enabled = False
+            assert pd.session.pandas_hybrid_execution_enabled is False
         with pytest.raises(
             NotImplementedError, match="Snowpark pandas does not yet support the method"
         ):
@@ -425,49 +432,6 @@ def test_unimplemented_autoswitches(class_name, method_name, f_args):
             assert snow_result == pandas_result
 
 
-@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_using_hybrid_session_param(
-    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.
-    pd.session.pandas_hybrid_execution_enabled = 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.
-    pd.session.pandas_hybrid_execution_enabled = True
-    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
-
-
 @sql_count_checker(query_count=0)
 def test_to_datetime():
     assert Backend.get() == "Snowflake"
@@ -476,14 +440,15 @@ def test_to_datetime():
     assert isinstance(result, DatetimeIndex)
 
 
+@pytest.mark.parametrize("use_session_param", [True, False])
 @sql_count_checker(
     query_count=11,
     join_count=6,
     udtf_count=2,
     high_count_expected=True,
     high_count_reason="tests queries across different execution modes",
 )
-def test_query_count_no_switch(init_transaction_tables):
+def test_query_count_no_switch(init_transaction_tables, use_session_param):
     """
     Tests that when there is no switching behavior the query count is the
     same under hybrid mode and non-hybrid mode.
@@ -501,50 +466,23 @@ def inner_test(df_in):
     hybrid_len = None
     with pd.session.query_history() as query_history_orig:
         with config_context(AutoSwitchBackend=False, NativePandasMaxRows=10):
+            if use_session_param:
+                from modin.config import AutoSwitchBackend
+
+                AutoSwitchBackend.enable()
+                pd.session.pandas_hybrid_execution_enabled = False
+                assert pd.session.pandas_hybrid_execution_enabled is False
             df_result = inner_test(df_transactions)
             orig_len = len(df_result)
 
     with pd.session.query_history() as query_history_hybrid:
         with config_context(AutoSwitchBackend=True, NativePandasMaxRows=10):
-            df_result = inner_test(df_transactions)
-            hybrid_len = len(df_result)
-
-    assert orig_len == hybrid_len
-    assert len(query_history_orig.queries) == len(query_history_hybrid.queries)
-
-
-@sql_count_checker(
-    query_count=11,
-    join_count=6,
-    udtf_count=2,
-    high_count_expected=True,
-    high_count_reason="tests queries across different execution modes",
-)
-def test_query_count_no_switch_using_hybrid_session_param(init_transaction_tables):
-    """
-    Tests that when there is no switching behavior the query count is the
-    same under hybrid mode and non-hybrid mode.
-    """
-
-    def inner_test(df_in):
-        df_result = df_in[(df_in["REVENUE"] > 123) & (df_in["REVENUE"] < 200)]
-        df_result["REVENUE_DUPE"] = df_result["REVENUE"]
-        df_result["COUNT"] = df_result.groupby("DATE")["REVENUE"].transform("count")
-        return df_result
-
-    df_transactions = pd.read_snowflake("REVENUE_TRANSACTIONS")
-    inner_test(df_transactions)
-    orig_len = None
-    hybrid_len = None
-    with pd.session.query_history() as query_history_orig:
-        with config_context(NativePandasMaxRows=10):
-            pd.session.pandas_hybrid_execution_enabled = False
-            df_result = inner_test(df_transactions)
-            orig_len = len(df_result)
+            if use_session_param:
+                from modin.config import AutoSwitchBackend
 
-    with pd.session.query_history() as query_history_hybrid:
-        with config_context(NativePandasMaxRows=10):
-            pd.session.pandas_hybrid_execution_enabled = True
+                AutoSwitchBackend.disable()
+                pd.session.pandas_hybrid_execution_enabled = True
+                assert pd.session.pandas_hybrid_execution_enabled is True
             df_result = inner_test(df_transactions)
             hybrid_len = len(df_result)