From 0a9bbc7a3e067cb97eac656453ac9e7ef733157d Mon Sep 17 00:00:00 2001 From: Varnika Budati Date: Mon, 26 Aug 2024 18:07:39 -0700 Subject: [PATCH] SNOW-1458137 Tests to verify `set_index` and `reset_index` work as expected (#2138) 1. Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR. Fixes SNOW-1458137 2. Fill out the following pre-review checklist: - [x] I am adding a new automated test(s) to verify correctness of my new code - [ ] If this test skips Local Testing mode, I'm requesting review from @snowflakedb/local-testing - [ ] I am adding new logging messages - [ ] I am adding a new telemetry message - [ ] I am adding new credentials - [ ] I am adding a new dependency - [ ] If this is a new feature/behavior, I'm adding the Local Testing parity changes. 3. Please describe how your code solves the related issue. Verifying that `df.index = new_index` is implemented correctly. --- .../compiler/snowflake_query_compiler.py | 8 -- tests/integ/modin/frame/test_axis.py | 133 ++---------------- tests/integ/modin/frame/test_drop.py | 2 +- .../integ/modin/frame/test_drop_duplicates.py | 2 +- .../modin/frame/test_nlargest_nsmallest.py | 2 +- tests/integ/modin/frame/test_set_index.py | 102 +++----------- tests/integ/modin/index/test_index_methods.py | 28 ++++ tests/integ/modin/series/test_axis.py | 40 ++---- tests/integ/modin/test_concat.py | 4 +- tests/integ/modin/test_telemetry.py | 2 +- 10 files changed, 75 insertions(+), 248 deletions(-) diff --git a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py index d25cb2f7e2c..31904a58b65 100644 --- a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py +++ b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py @@ -1705,17 +1705,9 @@ def set_index( if not any(isinstance(k, SnowflakeQueryCompiler) for k in keys): return self.set_index_from_columns(keys, drop=drop, append=append) - self_num_rows = self.get_axis_len(axis=0) new_qc = self for key in keys: if isinstance(key, SnowflakeQueryCompiler): - assert ( - len(key._modin_frame.data_column_pandas_labels) == 1 - ), "need to be a series" - if key.get_axis_len(0) != self_num_rows: - raise ValueError( - f"Length mismatch: Expected {self_num_rows} rows, received array of length {key.get_axis_len(0)}" - ) new_qc = new_qc.set_index_from_series(key, append) else: new_qc = new_qc.set_index_from_columns([key], drop, append) diff --git a/tests/integ/modin/frame/test_axis.py b/tests/integ/modin/frame/test_axis.py index 28cf55dee40..a6a156a05fe 100644 --- a/tests/integ/modin/frame/test_axis.py +++ b/tests/integ/modin/frame/test_axis.py @@ -81,7 +81,7 @@ def test_index(test_df): @pytest.mark.parametrize("test_df", test_dfs) -@sql_count_checker(query_count=8, join_count=3) +@sql_count_checker(query_count=3, join_count=3) def test_set_and_assign_index(test_df): def assign_index(df, keys): df.index = keys @@ -289,7 +289,7 @@ def test_duplicate_labels_assignment(): native_pd.DataFrame({"A": [3.14, 1.414, 1.732], "B": [9.8, 1.0, 0]}), "rows", [None] * 3, - 5, + 3, 2, ], [ # Labels is a MultiIndex from tuples. @@ -299,14 +299,14 @@ def test_duplicate_labels_assignment(): [("r0", "rA", "rR"), ("r1", "rB", "rS"), ("r2", "rC", "rT")], names=["Courses", "Fee", "Random"], ), - 7, + 3, 6, ], [ native_pd.DataFrame({"A": ["foo", "bar", 3], "B": [4, "baz", 6]}), 0, {1: "c", 2: "b", 3: "a"}, - 5, + 3, 2, ], [ @@ -326,7 +326,7 @@ def test_duplicate_labels_assignment(): ), 0, ['"row 1"', "row 2"], - 5, + 3, 2, ], [ @@ -339,7 +339,7 @@ def test_duplicate_labels_assignment(): ), "rows", list(range(10)), - 5, + 3, 2, ], [ @@ -367,7 +367,7 @@ def test_duplicate_labels_assignment(): native_pd.MultiIndex.from_product( [["NJ", "CA"], ["temp", "precip"]], names=["number", "color"] ), - 6, + 3, 4, ], # Set columns. @@ -720,105 +720,6 @@ def test_duplicate_labels_assignment(): ] -# Invalid data which raises ValueError different from native pandas -# ----------------------------------------------------------------- -# Format: df, axis, and invalid labels. -# - This data cover the negative case for DataFrame.set_axis() with invalid labels. -# - Invalid labels here consist of: passing None, too many values, too few values, empty list, -# Index, and MultiIndex objects as invalid labels for row-like axis. -TEST_DATA_FOR_DF_SET_AXIS_RAISES_VALUE_ERROR_DIFF_ERROR_MSG = [ - # invalid row labels - [ - native_pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}), - "index", - ["a", "b", "c", "d"], - "Length mismatch: Expected 3 rows, received array of length 4", - ], - [ - native_pd.DataFrame( - { - "A": [None] * 3, - "B": [4, 5, 6], - "C": [4, 5, 6], - "D": [7, 8, 9], - "E": [-1, -2, -3], - } - ), - 0, - [99], # too few labels, - "Length mismatch: Expected 3 rows, received array of length 1", - ], - [ - native_pd.DataFrame( - { - "Artists": ["Monet", "Manet", "Gogh"], - "Museums": ["MoMA", "SoMA", "The High"], - } - ), - "rows", - [None], # too few labels - "Length mismatch: Expected 3 rows, received array of length 1", - ], - [ - native_pd.DataFrame( - {"A": ["a", "b", "c", "d", "e"], "B": ["e", "d", "c", "b", "a"]} - ), - "index", - native_pd.Index([11, 111, 1111, 11111, 111111, 11111111]), # too many labels - "Length mismatch: Expected 5 rows, received array of length 6", - ], - [ - native_pd.DataFrame({"foo": [None], "bar": [None], "baz": [None]}), - 0, - native_pd.Index([None] * 10), # too many labels - "Length mismatch: Expected 1 rows, received array of length 10", - ], - [ - native_pd.DataFrame( - {"echo": ["echo", "echo"], "not an echo": [". . .", ". . ."]} - ), - "rows", - native_pd.Index([]), # too few labels - "Length mismatch: Expected 2 rows, received array of length 0", - ], - [ # Labels is a MultiIndex from tuples. - native_pd.DataFrame({"A": [1], -2515 / 135: [4]}), - "index", - native_pd.MultiIndex.from_tuples( - [("r0", "rA", "rR"), ("r1", "rB", "rS"), ("r2", "rC", "rT")], - names=["Courses", "Fee", "Random"], - ), # too many labels - "Length mismatch: Expected 1 rows, received array of length 3", - ], - [ # Labels is a MultiIndex from a DataFrame. - native_pd.DataFrame( - { - "A": [1, 2, 3, 4], - "B": [4, 5, 6, 7], - "C": [7, 8, 9, 10], - "D": [10, 11, 12, 13], - } - ), - "rows", - native_pd.MultiIndex.from_frame( - native_pd.DataFrame( - [], - columns=["a", "b"], - ), - ), # too few labels - "Length mismatch: Expected 4 rows, received array of length 0", - ], - [ # Labels is a MultiIndex from a product. - native_pd.DataFrame({1: [1], 2: [2], 3: [3], 4: [4], 5: [5], 6: [6]}), - 0, - native_pd.MultiIndex.from_product( - [[0], ["green", "purple"]], names=["number", "color"] - ), # too many labels - "Length mismatch: Expected 1 rows, received array of length 2", - ], -] - - @pytest.mark.parametrize( "native_df, axis, labels, num_queries, num_joins", TEST_DATA_FOR_DF_SET_AXIS ) @@ -865,20 +766,6 @@ def test_set_axis_df_raises_value_error(native_df, axis, labels): ) -@pytest.mark.parametrize( - "native_df, axis, labels, error_msg", - TEST_DATA_FOR_DF_SET_AXIS_RAISES_VALUE_ERROR_DIFF_ERROR_MSG, -) -def test_set_axis_df_raises_value_error_diff_error_msg( - native_df, axis, labels, error_msg -): - # Should raise a ValueError if the labels for row-like axis are invalid. - # The error messages do not match native pandas. - with SqlCounter(query_count=2): - with pytest.raises(ValueError, match=error_msg): - pd.DataFrame(native_df).set_axis(labels, axis=axis) - - @pytest.mark.parametrize( "native_df, axis, labels, error_msg", TEST_DATA_FOR_DF_SET_AXIS_RAISES_TYPE_ERROR ) @@ -892,7 +779,7 @@ def test_set_axis_df_raises_type_error_diff_error_msg( pd.DataFrame(native_df).set_axis(labels, axis=axis) -@sql_count_checker(query_count=3, join_count=1) +@sql_count_checker(query_count=1, join_count=1) def test_df_set_axis_copy_true(caplog): # Test that warning is raised when copy argument is used. native_df = native_pd.DataFrame({"A": [1.25], "B": [3]}) @@ -933,11 +820,11 @@ def test_df_set_axis_with_quoted_index(): # check first that operation result is the same snow_df = pd.DataFrame(data) native_df = native_pd.DataFrame(data) - with SqlCounter(query_count=3): + with SqlCounter(query_count=1): eval_snowpark_pandas_result(snow_df, native_df, helper) # then, explicitly compare axes - with SqlCounter(query_count=1): + with SqlCounter(query_count=0): ans = helper(snow_df) native_ans = helper(native_df) diff --git a/tests/integ/modin/frame/test_drop.py b/tests/integ/modin/frame/test_drop.py index aee61fb69a4..e71999dd28d 100644 --- a/tests/integ/modin/frame/test_drop.py +++ b/tests/integ/modin/frame/test_drop.py @@ -108,7 +108,7 @@ def test_drop_duplicate_columns(native_df, labels): @pytest.mark.parametrize( "labels, expected_query_count, expected_join_count", - [([], 3, 1), (1, 4, 2), (2, 4, 2), ([1, 2], 5, 3)], + [([], 1, 1), (1, 2, 2), (2, 2, 2), ([1, 2], 3, 3)], ) def test_drop_duplicate_row_labels( native_df, labels, expected_query_count, expected_join_count diff --git a/tests/integ/modin/frame/test_drop_duplicates.py b/tests/integ/modin/frame/test_drop_duplicates.py index 35c4a8edb05..aff95263f4b 100644 --- a/tests/integ/modin/frame/test_drop_duplicates.py +++ b/tests/integ/modin/frame/test_drop_duplicates.py @@ -64,7 +64,7 @@ def test_drop_duplicates(subset, keep, ignore_index): query_count = 1 join_count = 2 if ignore_index is True: - query_count += 2 + query_count += 1 join_count += 3 with SqlCounter(query_count=query_count, join_count=join_count): assert_frame_equal( diff --git a/tests/integ/modin/frame/test_nlargest_nsmallest.py b/tests/integ/modin/frame/test_nlargest_nsmallest.py index 1a2b13db7a5..c32fb64a80e 100644 --- a/tests/integ/modin/frame/test_nlargest_nsmallest.py +++ b/tests/integ/modin/frame/test_nlargest_nsmallest.py @@ -54,7 +54,7 @@ def test_nlargest_nsmallest_large_n(snow_df, native_df, method): ) -@sql_count_checker(query_count=3) +@sql_count_checker(query_count=1) def test_nlargest_nsmallest_overlapping_index_name(snow_df, native_df, method): snow_df = snow_df.rename_axis("A") native_df = native_df.rename_axis("A") diff --git a/tests/integ/modin/frame/test_set_index.py b/tests/integ/modin/frame/test_set_index.py index ae035f0b3a4..7e2c08042cd 100644 --- a/tests/integ/modin/frame/test_set_index.py +++ b/tests/integ/modin/frame/test_set_index.py @@ -8,7 +8,7 @@ import snowflake.snowpark.modin.plugin # noqa: F401 from tests.integ.modin.sql_counter import SqlCounter, sql_count_checker -from tests.integ.modin.utils import eval_snowpark_pandas_result +from tests.integ.modin.utils import assert_frame_equal, eval_snowpark_pandas_result @pytest.fixture @@ -31,7 +31,7 @@ def drop(request): return request.param -@sql_count_checker(query_count=4, join_count=2) +@sql_count_checker(query_count=1, join_count=2) def test_set_index_multiindex(snow_df, native_df): index = native_pd.MultiIndex.from_tuples([(5, 4), (4, 5), (5, 5)]) eval_snowpark_pandas_result( @@ -80,20 +80,14 @@ def test_set_index_multiindex_columns(snow_df): ) -@sql_count_checker(query_count=2) -def test_set_index_negative(snow_df, native_df): +@sql_count_checker(query_count=1, join_count=1) +def test_set_index_different_index_length(snow_df): index = pd.Index([1, 2]) - native_index = native_pd.Index([1, 2]) - eval_snowpark_pandas_result( - snow_df, - native_df, - lambda df: df.set_index(native_index) - if isinstance(df, native_pd.DataFrame) - else df.set_index(index), - expect_exception=True, - expect_exception_match="Length mismatch: Expected 3 rows, received array of length 2", - expect_exception_type=ValueError, + actual_df = snow_df.set_index(index) + expected_df = native_pd.DataFrame( + {"a": [1, 2, 2], "b": [3, 4, 5], ("c", "d"): [0, 0, 1]}, index=[1, 2, np.nan] ) + assert_frame_equal(actual_df, expected_df) @sql_count_checker(query_count=1) @@ -121,17 +115,17 @@ def test_set_index_names(snow_df): # Verify name from input index is set. index = pd.Index([1, 2, 0]) index.names = ["iname"] - with SqlCounter(query_count=2): + with SqlCounter(query_count=0): assert snow_df.set_index(index).index.names == ["iname"] # Verify names from input multiindex are set. multi_index = native_pd.MultiIndex.from_arrays( [[1, 1, 2], [1, 2, 1]], names=["a", "b"] ) - with SqlCounter(query_count=3, join_count=2): + with SqlCounter(query_count=1, join_count=2): assert snow_df.set_index(multi_index).index.names == ["a", "b"] - with SqlCounter(query_count=6, join_count=4): + with SqlCounter(query_count=2, join_count=4): # Verify that [MultiIndex, MultiIndex] yields a MultiIndex rather # than a pair of tuples multi_index2 = multi_index.rename(["C", "D"]) @@ -227,9 +221,9 @@ def test_set_index_pass_single_array(obj_type, drop, append, native_df): expect_exception=True, ) else: - expected_query_count = 3 + expected_query_count = 1 if obj_type == pd.Series or obj_type == pd.Index: - expected_query_count = 4 + expected_query_count = 2 with SqlCounter(query_count=expected_query_count, join_count=1): eval_snowpark_pandas_result( snow_df, @@ -264,7 +258,7 @@ def test_set_index_pass_arrays(obj_type, drop, append, native_df): "a", key.to_pandas() if isinstance(key, (pd.Series, pd.Index)) else key, ] - with SqlCounter(query_count=3, join_count=1): + with SqlCounter(query_count=1, join_count=1): eval_snowpark_pandas_result( snow_df, native_df, @@ -312,7 +306,7 @@ def test_set_index_pass_arrays_duplicate(obj_type1, obj_type2, drop, append, nat obj_type2 = native_pd.Index native_keys = [obj_type1(array), obj_type2(array)] - with SqlCounter(query_count=4, join_count=2): + with SqlCounter(query_count=1, join_count=2): eval_snowpark_pandas_result( snow_df, native_df, @@ -324,7 +318,7 @@ def test_set_index_pass_arrays_duplicate(obj_type1, obj_type2, drop, append, nat ) -@sql_count_checker(query_count=4, join_count=2) +@sql_count_checker(query_count=1, join_count=2) def test_set_index_pass_multiindex(drop, append, native_df): snow_df = pd.DataFrame(native_df) index_data = [["one", "two", "three"], [9, 3, 7]] @@ -340,7 +334,7 @@ def test_set_index_pass_multiindex(drop, append, native_df): "keys, expected_query_count", [ (["a"], 3), - ([[1, 6, 6]], 5), + ([[1, 6, 6]], 3), ], ) def test_set_index_verify_integrity_negative(native_df, keys, expected_query_count): @@ -400,66 +394,6 @@ def test_set_index_raise_on_invalid_type_set_negative(keys, drop, append, native ) -@pytest.mark.parametrize( - "obj_type", - [ - pd.Series, - pd.Index, - np.array, - iter, - lambda x: native_pd.MultiIndex.from_arrays([x]), - ], - ids=["Series", "Index", "np.array", "iter", "MultiIndex"], -) -@pytest.mark.parametrize("length", [2, 6], ids=["too_short", "too_long"]) -def test_set_index_raise_on_len(length, obj_type, drop, append, native_df): - snow_df = pd.DataFrame(native_df) - values = np.random.randint(0, 10, (length,)) - key = obj_type(values) - if obj_type == pd.Series: - obj_type = native_pd.Series - elif obj_type == pd.Index: - obj_type = native_pd.Index - native_key = obj_type(values) - - msg = "Length mismatch: Expected 3 rows, received array of length.*" - # wrong length directly - # one extra query to create the series to set index - with SqlCounter(query_count=2): - eval_snowpark_pandas_result( - snow_df, - native_df, - lambda df: df.set_index( - key if isinstance(df, pd.DataFrame) else native_key, - drop=drop, - append=append, - ), - expect_exception=True, - expect_exception_type=ValueError, - expect_exception_match=msg, - ) - - # wrong length in list - expected_query_count = 1 - if obj_type == native_pd.Series: - expected_query_count = 0 - keys = ["a", key] - native_keys = ["a", native_key] - with SqlCounter(query_count=expected_query_count): - eval_snowpark_pandas_result( - snow_df, - native_df, - lambda df: df.set_index( - keys if isinstance(df, pd.DataFrame) else native_keys, - drop=drop, - append=append, - ), - expect_exception=True, - expect_exception_type=ValueError, - expect_exception_match=msg, - ) - - class TestSetIndexCustomLabelType: class CustomLabel: def __init__(self, name, color) -> None: @@ -560,7 +494,7 @@ def __str__(self) -> str: native_pd.Series([1, 2, 3, 4], name="num"), ], ) - @sql_count_checker(query_count=3, join_count=1) + @sql_count_checker(query_count=1, join_count=1) def test_set_index_with_index_series_name(self, sample): df = native_pd.DataFrame( { diff --git a/tests/integ/modin/index/test_index_methods.py b/tests/integ/modin/index/test_index_methods.py index ad829dcff5f..16407a5a015 100644 --- a/tests/integ/modin/index/test_index_methods.py +++ b/tests/integ/modin/index/test_index_methods.py @@ -22,6 +22,7 @@ assert_index_equal, assert_series_equal, assert_snowpark_pandas_equals_to_pandas_without_dtypecheck, + eval_snowpark_pandas_result, ) @@ -469,3 +470,30 @@ def test_index_is_type(native_index, func): snow_res = getattr(snow_index, func)() native_res = getattr(native_index, func)() assert snow_res == native_res + + +@pytest.mark.parametrize("obj_type", ["df", "series"]) +def test_df_series_set_index_and_reset_index(obj_type): + obj = {"A": [1, 2, 3], "B": [4, 5, 6]} + original_index = ["A", "B", "C"] + assert_equal = assert_frame_equal if obj_type == "df" else assert_series_equal + native_obj = ( + native_pd.DataFrame(obj, index=original_index) + if obj_type == "df" + else native_pd.Series(obj, index=original_index) + ) + snow_obj = pd.DataFrame(native_obj) if obj_type == "df" else pd.Series(native_obj) + + # Index object to change obj's index to. + native_idx = native_pd.Index([11, 22, 33]) + snow_idx = pd.Index(native_idx) + + # Test that df.index = new_index works with lazy index. + with SqlCounter(query_count=1): + native_obj.index = native_idx + snow_obj.index = snow_idx + assert_equal(snow_obj, native_obj) + + # Check if reset_index works with lazy index. + with SqlCounter(query_count=1): + eval_snowpark_pandas_result(snow_obj, native_obj, lambda df: df.reset_index()) diff --git a/tests/integ/modin/series/test_axis.py b/tests/integ/modin/series/test_axis.py index d099272d6e9..6b21bfe0c76 100644 --- a/tests/integ/modin/series/test_axis.py +++ b/tests/integ/modin/series/test_axis.py @@ -30,7 +30,7 @@ native_pd.Series({"A": [1, 2, 3], 5 / 6: [4, 5, 6]}), "index", [None] * 2, - 3, + 1, 1, ], [ @@ -44,7 +44,7 @@ ), "index", ["iccanobif", "serauqs", "semirp"], - 3, + 1, 1, ], [ @@ -58,7 +58,7 @@ ), "index", native_pd.Series(["iccanobif", "serauqs", "semirp"], name="reverse names"), - 3, + 1, 1, ], [ @@ -73,7 +73,7 @@ ), 0, native_pd.Index([99, 999, 9999, 99999, 999999]), - 3, + 1, 1, ], [ @@ -88,7 +88,7 @@ ), 0, native_pd.Index([99, 999, 9999, 99999, 999999], name="index with name"), - 3, + 1, 1, ], [ @@ -104,7 +104,7 @@ ), 0, native_pd.Index([99, 999, 9999, 99999, 999999], name="index with name"), - 3, + 1, 1, ], [ # Index is a MultiIndex from tuples. @@ -113,7 +113,7 @@ native_pd.MultiIndex.from_tuples( [("r0", "rA"), ("r1", "rB")], names=["Courses", "Fee"] ), - 4, + 1, 2, ], [ # Index is a MultiIndex from arrays. @@ -135,7 +135,7 @@ ], names=("tea", "steep time", "caffeine"), ), - 5, + 1, 3, ], [ # Index is a MultiIndex from a DataFrame. @@ -149,7 +149,7 @@ columns=["a", "b"], ), ), - 4, + 1, 2, ], [ # Index is a MultiIndex from a product. @@ -158,21 +158,21 @@ native_pd.MultiIndex.from_product( [[0, 1, 2], ["green", "purple"]], names=["number", "color"] ), - 4, + 1, 2, ], [ native_pd.Series({"A": ["foo", "bar", 3], "B": [4, "baz", 6]}), "index", {1: 1, 2: 2}, - 3, + 1, 1, ], [ native_pd.Series({"A": ["foo", "bar", 3], "B": [4, "baz", 6]}), "rows", {1, 2}, - 3, + 1, 1, ], ] @@ -431,20 +431,6 @@ def test_set_axis_series_copy(native_series, axis, labels, num_queries, num_join # Invalid input tests for Series.set_axis(). -@pytest.mark.parametrize( - "ser, axis, labels, error_msg", - TEST_DATA_FOR_SERIES_SET_AXIS_RAISES_VALUE_ERROR_DIFF_ERROR_MSG, -) -def test_set_axis_series_raises_value_error_diff_error_msg( - ser, axis, labels, error_msg -): - # Should raise a ValueError if length of labels passed in - # don't match the number of rows. - with SqlCounter(query_count=2): - with pytest.raises(ValueError, match=error_msg): - pd.Series(ser).set_axis(labels, axis=axis) - - @pytest.mark.parametrize( "ser, axis, labels", TEST_DATA_FOR_SERIES_SET_AXIS_RAISES_VALUE_ERROR, @@ -474,7 +460,7 @@ def test_set_axis_series_raises_type_error(ser, axis, labels, error_msg): pd.Series(ser).set_axis(labels, axis=axis) -@sql_count_checker(query_count=3, join_count=1) +@sql_count_checker(query_count=1, join_count=1) def test_series_set_axis_copy_true(caplog): # Test that warning is raised when copy argument is used. series = native_pd.Series([1.25]) diff --git a/tests/integ/modin/test_concat.py b/tests/integ/modin/test_concat.py index f3e149a37fc..7e11a3537af 100644 --- a/tests/integ/modin/test_concat.py +++ b/tests/integ/modin/test_concat.py @@ -656,9 +656,9 @@ def test_concat_keys_with_none(df1, df2, axis): "name1, name2", [("one", "two"), ("one", None), (None, "two"), (None, None)] ) def test_concat_with_keys_and_names(df1, df2, names, name1, name2, axis): - with SqlCounter(query_count=0 if name1 is None or axis == 1 else 2): + with SqlCounter(query_count=0): df1 = df1.rename_axis(name1, axis=axis) - with SqlCounter(query_count=0 if name2 is None or axis == 1 else 2): + with SqlCounter(query_count=0): df2 = df2.rename_axis(name2, axis=axis) # One extra query to convert index to native pandas when creating df diff --git a/tests/integ/modin/test_telemetry.py b/tests/integ/modin/test_telemetry.py index 887e67e59e0..9c24c6b6853 100644 --- a/tests/integ/modin/test_telemetry.py +++ b/tests/integ/modin/test_telemetry.py @@ -445,7 +445,7 @@ def test_telemetry_private_method(name, method, expected_query_count): assert data["api_calls"] == [{"name": f"DataFrame.DataFrame.{name}"}] -@sql_count_checker(query_count=2) +@sql_count_checker(query_count=0) def test_telemetry_property_index(): df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) df._query_compiler.snowpark_pandas_api_calls.clear()