SNOW-2676991: Add support for to_snowflake in faster pandas (#3988)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -157,6 +157,7 @@
   - `groupby.any`
   - `groupby.all`
   - `groupby.unique`
+  - `to_snowflake`
 - Make faster pandas disabled by default (opt-in instead of opt-out).
 - Improve performance of `drop_duplicates` by avoiding joins when `keep!=False` in faster pandas.
 

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

@@ -836,9 +836,8 @@ def _raise_not_implemented_error_for_timedelta(
                 method = method_frame.f_code.co_name  # type: ignore[union-attr]
                 ErrorMessage.not_implemented_for_timedelta(method)
 
-    def _warn_lost_snowpark_pandas_type(self) -> None:
+    def _warn_lost_snowpark_pandas_type(self, method: str) -> None:
         """Warn Snowpark pandas type can be lost in current operation."""
-        method = inspect.currentframe().f_back.f_back.f_code.co_name  # type: ignore[union-attr]
         snowpark_pandas_types = [
             type(t).__name__
             for t in set(
@@ -2284,7 +2283,35 @@ def to_snowflake(
         index_label: Optional[IndexLabel] = None,
         table_type: Literal["", "temp", "temporary", "transient"] = "",
     ) -> None:
-        self._warn_lost_snowpark_pandas_type()
+        """
+        Wrapper around _to_snowflake_internal to be supported in faster pandas.
+        """
+        if self._relaxed_query_compiler is not None and not index:
+            self._relaxed_query_compiler._to_snowflake_internal(
+                name=name,
+                if_exists=if_exists,
+                index=index,
+                index_label=index_label,
+                table_type=table_type,
+            )
+        else:
+            self._to_snowflake_internal(
+                name=name,
+                if_exists=if_exists,
+                index=index,
+                index_label=index_label,
+                table_type=table_type,
+            )
+
+    def _to_snowflake_internal(
+        self,
+        name: Union[str, Iterable[str]],
+        if_exists: Optional[Literal["fail", "replace", "append"]] = "fail",
+        index: bool = True,
+        index_label: Optional[IndexLabel] = None,
+        table_type: Literal["", "temp", "temporary", "transient"] = "",
+    ) -> None:
+        self._warn_lost_snowpark_pandas_type("to_snowflake")
         handle_if_exists_for_to_snowflake(if_exists=if_exists, name=name)
 
         if if_exists == "fail":
@@ -2322,7 +2349,7 @@ def to_snowpark(
 
         For details, please see comment in _to_snowpark_dataframe_of_pandas_dataframe.
         """
-        self._warn_lost_snowpark_pandas_type()
+        self._warn_lost_snowpark_pandas_type("to_snowpark")
 
         return self._to_snowpark_dataframe_from_snowpark_pandas_dataframe(
             index, index_label

tests/integ/modin/test_faster_pandas.py

@@ -1752,6 +1752,40 @@ def test_series_to_datetime(session):
         assert_series_equal(snow_result, native_result)
 
 
+@sql_count_checker(query_count=4)
+def test_to_snowflake(session):
+    with session_parameter_override(
+        session, "dummy_row_pos_optimization_enabled", True
+    ):
+        # create tables
+        table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+        session.create_dataframe(
+            native_pd.DataFrame([[2, 12], [1, 11], [3, 13]], columns=["A", "B"])
+        ).write.save_as_table(table_name, table_type="temp")
+
+        # create snow dataframes
+        df = pd.read_snowflake(table_name)
+        df.to_snowflake(table_name, if_exists="replace", index=False)
+        snow_result = pd.read_snowflake(table_name).sort_values(by="A")
+
+        # verify that the input dataframe has a populated relaxed query compiler
+        assert df._query_compiler._relaxed_query_compiler is not None
+        assert df._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+        # verify that the output dataframe also has a populated relaxed query compiler
+        assert snow_result._query_compiler._relaxed_query_compiler is not None
+        assert (
+            snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode
+            is True
+        )
+
+        # create pandas dataframes
+        native_df = df.to_pandas()
+        native_result = native_df.sort_values(by="A")
+
+        # compare results
+        assert_frame_equal(snow_result, native_result)
+
+
 @sql_count_checker(query_count=0)
 def test_dummy_row_pos_optimization_enabled_on_session(db_parameters):
     with Session.builder.configs(db_parameters).create() as new_session: