SNOW-2189056: Support some snowflake extensions on other backends. (#3683)

Make sure that we support the function and method variants of `to_snowflake()`, `to_snowpark()`, `to_iceberg()`, `read_snowflake()`, and `to_pandas()` on the pandas and Ray backends. Prior to this commit, we only supported some of these functions and methods on the Pandas backend.

Fixes SNOW-2189056

Signed-off-by: sfc-gh-mvashishtha <mahesh.vashishtha@snowflake.com>
Co-authored-by: Hazem Elmeleegy <hazem.elmeleegy@snowflake.com>

Author: sfc-gh-mvashishtha

CHANGELOG.md

@@ -27,13 +27,26 @@
 ### Snowpark pandas API Updates
 
 #### New Features
+- Completed support for `pd.read_snowflake()`, `pd.to_iceberg()`,
+  `pd.to_pandas()`, `pd.to_snowpark()`, `pd.to_snowflake()`,
+  `DataFrame.to_iceberg()`, `DataFrame.to_pandas()`, `DataFrame.to_snowpark()`,
+  `DataFrame.to_snowflake()`, `Series.to_iceberg()`, `Series.to_pandas()`,
+  `Series.to_snowpark()`, and `Series.to_snowflake()` on the "Pandas" and "Ray"
+  backends. Previously, only some of these functions and methods were supported
+  on the Pandas backend.
 
 #### Improvements
 - Set the default transfer limit in hybrid execution for data leaving Snowflake to 100k, which can be overridden with the SnowflakePandasTransferThreshold environment variable. This configuration is appropriate for scenarios with two available engines, "Pandas" and "Snowflake" on relational workloads.
 - Improve import error message by adding '--upgrade' to 'pip install "snowflake-snowpark-python[modin]"' in the error message.
 
 #### Dependency Updates
+
 #### Bug Fixes
+- Raised `NotImplementedError` instead of `AttributeError` on attempting to call
+  Snowflake extension functions/methods `to_dynamic_table()`, `cache_result()`,
+  `to_view()`, `create_or_replace_dynamic_table()`, and
+  `create_or_replace_view()` on dataframes or series using the pandas or ray 
+  backends.
 
 ## 1.37.0 (2025-08-18)
 

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

@@ -21,7 +21,10 @@
 from snowflake.snowpark._internal.type_utils import ColumnOrName
 from snowflake.snowpark.async_job import AsyncJob
 from snowflake.snowpark.dataframe import DataFrame as SnowparkDataFrame
-from snowflake.snowpark.modin.plugin.extensions.utils import add_cache_result_docstring
+from snowflake.snowpark.modin.plugin.extensions.utils import (
+    add_cache_result_docstring,
+    register_non_snowflake_accessors,
+)
 from snowflake.snowpark.modin.plugin.utils.warning_message import (
     materialization_warning,
 )
@@ -30,17 +33,8 @@
 register_dataframe_accessor = functools.partial(
     _register_dataframe_accessor, backend="Snowflake"
 )
-_register_dataframe_accessor(name="to_pandas", backend="Pandas")(
-    pd.DataFrame._to_pandas
-)
-_register_dataframe_accessor(name="to_snowflake", backend="Pandas")(
-    lambda self, *args, **kwargs: self.move_to("Snowflake").to_snowflake(
-        *args, **kwargs
-    )
-)
-_register_dataframe_accessor(name="to_snowpark", backend="Pandas")(
-    lambda self, *args, **kwargs: self.move_to("Snowflake").to_snowpark(*args, **kwargs)
-)
+
+register_non_snowflake_accessors(_register_dataframe_accessor, "DataFrame")
 
 
 # Snowflake specific dataframe methods

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

@@ -6,6 +6,7 @@
 File containing top-level APIs defined in Snowpark pandas but not the Modin API layer
 under the `pd` namespace, such as `pd.read_snowflake`.
 """
+from functools import wraps
 from typing import Any, Iterable, List, Literal, Optional, Union
 
 from modin.pandas import DataFrame, Series
@@ -16,7 +17,7 @@
 from snowflake.snowpark._internal.type_utils import ColumnOrName
 from snowflake.snowpark.async_job import AsyncJob
 from snowflake.snowpark.row import Row
-from .general_overrides import register_pd_accessor
+from .general_overrides import register_pd_accessor as register_snowflake_accessor
 from pandas._typing import IndexLabel
 import pandas as native_pd
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
@@ -36,16 +37,32 @@
 )
 
 
-register_pd_accessor("Index")(Index)
-register_pd_accessor("DatetimeIndex")(DatetimeIndex)
-register_pd_accessor("TimedeltaIndex")(TimedeltaIndex)
+register_snowflake_accessor("Index")(Index)
+register_snowflake_accessor("DatetimeIndex")(DatetimeIndex)
+register_snowflake_accessor("TimedeltaIndex")(TimedeltaIndex)
 
 
 def _snowpark_pandas_obj_check(obj: Union[DataFrame, Series]):
     if not isinstance(obj, (DataFrame, Series)):
         raise TypeError("obj must be a Snowpark pandas DataFrame or Series")
 
 
+def _check_obj_and_set_backend_to_snowflake(
+    obj: Any,
+) -> Union[Series, DataFrame]:
+    """
+    Check if the object is a Snowpark pandas object and set the backend to Snowflake.
+
+    Args:
+        obj: The object to be checked and moved to Snowflake backend.
+
+    Returns:
+        The Series or DataFrame on the Snowflake backend.
+    """
+    _snowpark_pandas_obj_check(obj)
+    return obj.set_backend("Snowflake") if obj.get_backend() != "Snowflake" else obj
+
+
 # Use a template string so that we can share it between the read_snowflake
 # functions on the Snowflake and Pandas backends. We can't use the exact same
 # docstring because each doctest creates and inserts to a temp table, and also
@@ -397,7 +414,7 @@ def _snowpark_pandas_obj_check(obj: Union[DataFrame, Series]):
 """
 
 
-@register_pd_accessor("read_snowflake")
+@register_snowflake_accessor("read_snowflake")
 @doc(
     _READ_SNOWFLAKE_DOC,
     table_name="RESULT_0",
@@ -441,7 +458,26 @@ def _read_snowflake_pandas_backend(
     return df.set_backend("Pandas")
 
 
-@register_pd_accessor("to_snowflake")
+@_register_pd_accessor("read_snowflake", backend="Ray")
+@doc(
+    _READ_SNOWFLAKE_DOC,
+    table_name="RESULT_2",
+    stored_procedure_name="MULTIPLY_COL_BY_VALUE_2",
+)
+def _read_snowflake_ray_backend(
+    name_or_query, index_col=None, columns=None, enforce_ordering=False
+) -> pd.DataFrame:
+    with config_context(Backend="Snowflake"):
+        df = pd.read_snowflake(
+            name_or_query,
+            index_col=index_col,
+            columns=columns,
+            enforce_ordering=enforce_ordering,
+        )
+    return df.set_backend("Ray")
+
+
+@_register_pd_accessor("to_snowflake")
 def to_snowflake(
     obj: Union[DataFrame, Series],
     name: Union[str, Iterable[str]],
@@ -477,14 +513,12 @@ def to_snowflake(
         - :func:`Series.to_snowflake <modin.pandas.Series.to_snowflake>`
         - :func:`read_snowflake <modin.pandas.read_snowflake>`
     """
-    _snowpark_pandas_obj_check(obj)
-
-    return obj._query_compiler.to_snowflake(
+    return _check_obj_and_set_backend_to_snowflake(obj)._query_compiler.to_snowflake(
         name, if_exists, index, index_label, table_type
     )
 
 
-@register_pd_accessor("to_snowpark")
+@_register_pd_accessor("to_snowpark")
 def to_snowpark(
     obj: Union[DataFrame, Series],
     index: bool = True,
@@ -628,12 +662,12 @@ def to_snowpark(
         ------------
         <BLANKLINE>
     """
-    _snowpark_pandas_obj_check(obj)
-
-    return obj._query_compiler.to_snowpark(index, index_label)
+    return _check_obj_and_set_backend_to_snowflake(obj)._query_compiler.to_snowpark(
+        index, index_label
+    )
 
 
-@register_pd_accessor("to_pandas")
+@register_snowflake_accessor("to_pandas")
 @materialization_warning
 def to_pandas(
     obj: Union[DataFrame, Series],
@@ -678,7 +712,7 @@ def to_pandas(
     return obj.to_pandas(statement_params=statement_params, *kwargs)
 
 
-@register_pd_accessor("to_view")
+@register_snowflake_accessor("to_view")
 def to_view(
     obj: Union[DataFrame, Series],
     name: Union[str, Iterable[str]],
@@ -709,7 +743,6 @@ def to_view(
             then the index names are used. A sequence should be given if the DataFrame uses MultiIndex.
     """
     _snowpark_pandas_obj_check(obj)
-
     return obj.to_view(
         name=name,
         comment=comment,
@@ -718,7 +751,7 @@ def to_view(
     )
 
 
-@register_pd_accessor("to_dynamic_table")
+@register_snowflake_accessor("to_dynamic_table")
 def to_dynamic_table(
     obj: Union[DataFrame, Series],
     name: Union[str, Iterable[str]],
@@ -791,7 +824,6 @@ def to_dynamic_table(
         for more details on refresh mode.
     """
     _snowpark_pandas_obj_check(obj)
-
     return obj.to_dynamic_table(
         name=name,
         warehouse=warehouse,
@@ -810,7 +842,7 @@ def to_dynamic_table(
     )
 
 
-@register_pd_accessor("to_iceberg")
+@_register_pd_accessor("to_iceberg")
 def to_iceberg(
     obj: Union[DataFrame, Series],
     table_name: Union[str, Iterable[str]],
@@ -905,8 +937,7 @@ def to_iceberg(
         ... }
         >>> pd.to_iceberg(df.to_snowpark_pandas(), "my_table", iceberg_config=iceberg_config, mode="overwrite") # doctest: +SKIP
     """
-    _snowpark_pandas_obj_check(obj)
-    return obj._query_compiler.to_iceberg(
+    return _check_obj_and_set_backend_to_snowflake(obj)._query_compiler.to_iceberg(
         table_name=table_name,
         iceberg_config=iceberg_config,
         mode=mode,
@@ -924,7 +955,34 @@ def to_iceberg(
     )
 
 
-@register_pd_accessor("explain_switch")
+def _make_unimplemented_extension(name: str, to_wrap: callable):
+    """
+    Make an extension for an unimplemented function.
+
+    Args:
+        name: The name of the function.
+        to_wrap: The function to wrap.
+
+    Returns:
+        A function that raises NotImplementedError.
+    """
+
+    @wraps(to_wrap)
+    def _unimplemented_extension(obj, *args, **kwargs):
+        _snowpark_pandas_obj_check(obj)
+        # Let the object take care of raising the NotImplementedError.
+        return getattr(obj, name)(*args, **kwargs)
+
+    return _unimplemented_extension
+
+
+for function in (to_dynamic_table, to_view):
+    _register_pd_accessor(name=function.__name__, backend=None)(
+        _make_unimplemented_extension(name=function.__name__, to_wrap=function)
+    )
+
+
+@register_snowflake_accessor("explain_switch")
 def explain_switch(simple=True) -> Union[native_pd.DataFrame, None]:
     """
     Shows a log of all backend switching decisions made by Snowpark pandas.

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

@@ -21,7 +21,10 @@
 from snowflake.snowpark._internal.type_utils import ColumnOrName
 from snowflake.snowpark.async_job import AsyncJob
 from snowflake.snowpark.dataframe import DataFrame as SnowparkDataFrame
-from snowflake.snowpark.modin.plugin.extensions.utils import add_cache_result_docstring
+from snowflake.snowpark.modin.plugin.extensions.utils import (
+    add_cache_result_docstring,
+    register_non_snowflake_accessors,
+)
 from snowflake.snowpark.modin.plugin.utils.warning_message import (
     materialization_warning,
 )
@@ -31,15 +34,8 @@
 register_series_accessor = functools.partial(
     _register_series_accessor, backend="Snowflake"
 )
-_register_series_accessor(name="to_pandas", backend="Pandas")(pd.Series._to_pandas)
-_register_series_accessor("to_snowflake", backend="Pandas")(
-    lambda self, *args, **kwargs: self.move_to("Snowflake").to_snowflake(
-        *args, **kwargs
-    )
-)
-_register_series_accessor("to_snowpark", backend="Pandas")(
-    lambda self, *args, **kwargs: self.move_to("Snowflake").to_snowpark(*args, **kwargs)
-)
+
+register_non_snowflake_accessors(_register_series_accessor, "Series")
 
 
 @register_series_accessor("_set_axis_name")

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

@@ -9,7 +9,7 @@
 
 from collections.abc import Hashable, Sequence
 from types import BuiltinFunctionType
-from typing import Any, Callable
+from typing import Any, Callable, NoReturn
 
 import modin.pandas as pd
 import numpy as np
@@ -580,3 +580,82 @@ def try_convert_index_to_native(index_like: Any) -> Any:
     if isinstance(index_like, Index):
         index_like = index_like.to_pandas()
     return index_like
+
+
+def _make_non_snowflake_accessor(method: str) -> Callable:
+    """
+    Create a method that moves the object to Snowflake backend and calls the specified method.
+
+    Args:
+        method: The method name to call on the Snowflake-backend object.
+
+    Returns:
+        A method that moves the object to Snowflake and calls the method.
+    """
+    return lambda self, *args, **kwargs: getattr(self.move_to("Snowflake"), method)(
+        *args, **kwargs
+    )
+
+
+def _make_non_snowflake_not_implemented_accessor(
+    backend: str, method: str, object_type: str
+) -> Callable[..., NoReturn]:
+    """
+    Create a function that raises NotImplementedError for unsupported backend/method combinations.
+
+    Args:
+        backend: The backend name (e.g., "Pandas", "Ray").
+        method: The method name that is not supported.
+        object_type: The object type ("DataFrame" or "Series").
+
+    Returns:
+        A function that raises NotImplementedError when called.
+    """
+
+    def raise_error(self, *args, **kwargs) -> NoReturn:
+        raise NotImplementedError(
+            f"Modin supports the method {object_type}.{method} on the Snowflake backend, but not on the backend {backend}."
+        )
+
+    return raise_error
+
+
+def register_non_snowflake_accessors(
+    register_accessor_func: Callable, object_type: str
+) -> None:
+    """
+    Register accessors for non-Snowflake backends.
+
+    Args:
+        register_accessor_func: The function that registers the accessor.
+        object_type: The name of the object type.
+    """
+    for backend in ("Pandas", "Ray"):
+        register_accessor_func(name="to_pandas", backend=backend)(
+            pd.DataFrame._to_pandas
+            if object_type == "DataFrame"
+            else pd.Series._to_pandas
+        )
+
+        # Register methods that move to Snowflake backend
+        # TODO(SNOW-2288765): Improve performance of to_snowflake() by writing
+        # to a Snowflake table directly instead of going through Snowpark
+        # pandas.
+        for method in ("to_snowflake", "to_snowpark", "to_iceberg"):
+            register_accessor_func(name=method, backend=backend)(
+                _make_non_snowflake_accessor(method=method)
+            )
+
+        # Register methods that are not implemented on non-Snowflake backends
+        for method in (
+            "create_or_replace_view",
+            "create_or_replace_dynamic_table",
+            "to_view",
+            "to_dynamic_table",
+            "cache_result",
+        ):
+            register_accessor_func(name=method, backend=backend)(
+                _make_non_snowflake_not_implemented_accessor(
+                    backend=backend, method=method, object_type=object_type
+                )
+            )