fix query counts and test failures

Author: sfc-gh-joshi

src/snowflake/snowpark/modin/plugin/_internal/frame.py

@@ -1144,6 +1144,8 @@ def update_snowflake_quoted_identifiers_with_expressions(
         self,
         quoted_identifier_to_column_map: dict[str, SnowparkColumn],
         snowpark_pandas_types: Optional[list[Optional[SnowparkPandasType]]] = None,
+        *,
+        new_index_column_pandas_labels: Optional[list[Hashable]] = None,
     ) -> UpdatedInternalFrameResult:
         """
         Points Snowflake quoted identifiers to column expression given by `quoted_identifier_to_column_map`.
@@ -1171,6 +1173,8 @@ def update_snowflake_quoted_identifiers_with_expressions(
                 must be index columns and data columns in the original internal frame.
             data_column_snowpark_pandas_types: The optional Snowpark pandas types for the new
                 expressions, in the order of the keys of quoted_identifier_to_column_map.
+            new_index_column_pandas_labels: The optional list of labels to be used as
+                index_column_pandas_labels for the result.
 
         Returns:
             UpdatedInternalFrameResult: A tuple containing the new InternalFrame with updated column references, and a mapping
@@ -1252,7 +1256,9 @@ def update_snowflake_quoted_identifiers_with_expressions(
                 data_column_pandas_labels=self.data_column_pandas_labels,
                 data_column_snowflake_quoted_identifiers=new_data_column_snowflake_quoted_identifiers,
                 data_column_pandas_index_names=self.data_column_pandas_index_names,
-                index_column_pandas_labels=self.index_column_pandas_labels,
+                index_column_pandas_labels=self.index_column_pandas_labels
+                if new_index_column_pandas_labels is None
+                else new_index_column_pandas_labels,
                 index_column_snowflake_quoted_identifiers=new_index_column_snowflake_quoted_identifiers,
                 data_column_types=[
                     new_type_mapping[k]

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

@@ -11888,7 +11888,9 @@ def transpose_single_row(self) -> "SnowflakeQueryCompiler":
                         frame.index_column_snowflake_quoted_identifiers[0]: pandas_lit(
                             frame.data_column_pandas_labels[0]
                         ),
-                    }
+                    },
+                    # Swap the name of the index/columns objects
+                    new_index_column_pandas_labels=frame.data_column_pandas_index_names,
                 )[0]
             ).set_columns([None])
 

tests/integ/modin/frame/test_aggregate.py

@@ -193,7 +193,9 @@ def test_string_sum_with_nulls():
     with pytest.raises(TypeError):
         pandas_df.sum(numeric_only=False)
     snow_result = snow_df.sum(numeric_only=False)
-    assert_series_equal(snow_result.to_pandas(), native_pd.Series(["ab"]))
+    assert_series_equal(
+        snow_result.to_pandas(), native_pd.Series(["ab"]), check_index_type=False
+    )
 
 
 class TestTimedelta:

tests/integ/modin/frame/test_iloc.py

@@ -1077,21 +1077,9 @@ def iloc_helper(df):
         else:
             return native_pd.Series([]) if axis == "row" else df.iloc[:, []]
 
-    def determine_query_count():
-        # Multiple queries because of squeeze() - in range is 2, out-of-bounds is 1.
-        if axis == "col":
-            num_queries = 1
-        else:
-            if not -8 < key < 7:  # key is out of bound
-                num_queries = 2
-            else:
-                num_queries = 1
-        return num_queries
-
-    query_count = determine_query_count()
     # test df with default index
     num_cols = 7
-    with SqlCounter(query_count=query_count):
+    with SqlCounter(query_count=1):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df,
             default_index_native_df,
@@ -1101,21 +1089,20 @@ def determine_query_count():
 
     # test df with non-default index
     num_cols = 6  # set_index() makes the number of columns 6
-    with SqlCounter(query_count=query_count):
+    with SqlCounter(query_count=1):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df.set_index("D"),
             default_index_native_df.set_index("D"),
             iloc_helper,
             test_attrs=False,
         )
 
-    query_count = determine_query_count()
     # test df with MultiIndex
     # Index dtype is different between Snowpark and native pandas if key produces empty df.
     num_cols = 7
     native_df = default_index_native_df.set_index(multiindex_native)
     snowpark_df = pd.DataFrame(native_df)
-    with SqlCounter(query_count=query_count):
+    with SqlCounter(query_count=1):
         eval_snowpark_pandas_result(
             snowpark_df,
             native_df,
@@ -1129,7 +1116,7 @@ def determine_query_count():
         native_df_with_multiindex_columns
     )
     in_range = True if (-8 < key < 7) else False
-    with SqlCounter(query_count=query_count):
+    with SqlCounter(query_count=1):
         if axis == "row" or in_range:  # series result
             eval_snowpark_pandas_result(
                 snowpark_df_with_multiindex_columns,
@@ -1151,7 +1138,7 @@ def determine_query_count():
     # test df with MultiIndex on both index and columns
     native_df = native_df_with_multiindex_columns.set_index(multiindex_native)
     snowpark_df = pd.DataFrame(native_df)
-    with SqlCounter(query_count=query_count):
+    with SqlCounter(query_count=1):
         if axis == "row" or in_range:  # series result
             eval_snowpark_pandas_result(
                 snowpark_df,

tests/integ/modin/frame/test_squeeze.py

@@ -31,11 +31,7 @@ def test_n_by_1(axis, dtype):
 
 @pytest.mark.parametrize("dtype", ["int", "timedelta64[ns]"])
 def test_1_by_n(axis, dtype):
-    if axis is None:
-        expected_query_count = 2
-    else:
-        expected_query_count = 1
-    with SqlCounter(query_count=expected_query_count):
+    with SqlCounter(query_count=1):
         eval_snowpark_pandas_result(
             *create_test_dfs({"a": [1], "b": [2], "c": [3]}, dtype=dtype),
             lambda df: df.squeeze(axis=axis),