SNOW-2333472: More hybrid benchmark fixes (#3771)

Author: sfc-gh-joshi

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

@@ -70,6 +70,7 @@
 from snowflake.snowpark.modin.plugin.compiler.snowflake_query_compiler import (
     HYBRID_SWITCH_FOR_UNIMPLEMENTED_METHODS,
 )
+from snowflake.snowpark.modin.plugin._internal.utils import new_snow_series
 from snowflake.snowpark.modin.plugin.extensions.utils import (
     ensure_index,
     extract_validate_and_try_convert_named_aggs_from_kwargs,
@@ -1166,10 +1167,11 @@ def _to_series_list(self, index: pd.Index) -> list[pd.Series]:
     # TODO: SNOW-1119855: Modin upgrade - modin.pandas.base.BasePandasDataset
     if isinstance(index, pd.MultiIndex):
         return [
-            pd.Series(index.get_level_values(level)) for level in range(index.nlevels)
+            new_snow_series(index.get_level_values(level))
+            for level in range(index.nlevels)
         ]
     elif isinstance(index, pd.Index):
-        return [pd.Series(index)]
+        return [new_snow_series(index)]
     else:
         raise Exception("invalid index: " + str(index))
 

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

@@ -77,6 +77,7 @@
     is_snowflake_agg_func,
 )
 from snowflake.snowpark.modin.plugin._internal.utils import (
+    new_snow_series,
     add_extra_columns_and_select_required_columns,
     assert_fields_are_none,
     convert_index_to_list_of_qcs,
@@ -1827,7 +1828,7 @@ def rename(
         )
 
     if isinstance(index, dict):
-        index = Series(index)
+        index = new_snow_series(index)
 
     new_qc = self._query_compiler.rename(
         index_renamer=index, columns_renamer=columns, level=level, errors=errors

tests/integ/modin/hybrid/test_switch_operations.py

@@ -608,3 +608,23 @@ def test_switch_then_iloc():
     # Setting should similarly not error
     df.iloc[1, 1] = 100
     assert df.iloc[1, 1] == 100
+
+
+@sql_count_checker(query_count=1, join_count=1)
+def test_rename():
+    # SNOW-2333472: Switching backends then performing a rename should be valid.
+    df = pd.DataFrame([[0] * 3] * 3)
+    assert df.get_backend() == "Pandas"
+    assert_snowpark_pandas_equal_to_pandas(
+        df.move_to("Snowflake").rename({0: "a", 1: "b", 2: "c"}),
+        # Perform to_pandas first due to modin issue 7667
+        df.to_pandas().rename({0: "a", 1: "b", 2: "c"}),
+    )
+
+
+@sql_count_checker(query_count=1, join_count=1)
+def test_set_index():
+    s = pd.Series([0]).move_to("Snowflake")
+    # SNOW-2333472: Switching backends then setting the index should be valid.
+    s.index = ["a"]
+    assert_snowpark_pandas_equal_to_pandas(s, native_pd.Series([0], index=["a"]))