diff --git a/src/snowflake/snowpark/modin/plugin/extensions/base_overrides.py b/src/snowflake/snowpark/modin/plugin/extensions/base_overrides.py index 2ee1e0a0db..3f65c2913d 100644 --- a/src/snowflake/snowpark/modin/plugin/extensions/base_overrides.py +++ b/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)) diff --git a/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py b/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py index d4151d2bc3..732a79cfa2 100644 --- a/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py +++ b/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 diff --git a/tests/integ/modin/hybrid/test_switch_operations.py b/tests/integ/modin/hybrid/test_switch_operations.py index 1c56966845..e469c41cb0 100644 --- a/tests/integ/modin/hybrid/test_switch_operations.py +++ b/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"]))