testing how smart cursor is

Author: sfc-gh-joshi

tests/integ/modin/binary/test_binary_op.py

@@ -2094,7 +2094,7 @@ def test_binary_add_dataframe_and_series_axis0(df, s):
     snow_s = pd.Series(s)
 
     # DataFrame <op> Series
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         ans = df.add(s, axis=0)
         snow_ans = snow_df.add(snow_s, axis=0)
 
@@ -2160,7 +2160,7 @@ def test_binary_op_between_dataframe_and_series_axis0(opname, df, s):
     snow_s = pd.Series(s)
 
     # DataFrame <op> Series
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         ans = getattr(df, opname)(s, axis=0)
         snow_ans = getattr(snow_df, opname)(snow_s, axis=0)
 
@@ -2385,12 +2385,12 @@ def test_binary_sub_dataframe_and_dataframe(df1, df2):
     snow_df1 = pd.DataFrame(df1)
     snow_df2 = pd.DataFrame(df2)
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             (snow_df1, snow_df2), (df1, df2), lambda t: t[0] - t[1]
         )
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             (snow_df2, snow_df1), (df2, df1), lambda t: t[0] - t[1]
         )
@@ -2402,14 +2402,14 @@ def test_binary_sub_dataframe_and_dataframe_with_fill_value(df1, df2, fill_value
     snow_df1 = pd.DataFrame(df1)
     snow_df2 = pd.DataFrame(df2)
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             (snow_df1, snow_df2),
             (df1, df2),
             lambda t: t[0].sub(t[1], fill_value=fill_value),
         )
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             (snow_df2, snow_df1),
             (df2, df1),

tests/integ/modin/frame/test_apply.py

@@ -93,7 +93,12 @@ def test_axis_1_basic_types_without_type_hints(data, func, return_type):
     native_df = native_pd.DataFrame(data, columns=["A", "b"])
     snow_df = pd.DataFrame(data, columns=["A", "b"])
     # np.min is mapped to sql builtin function.
-    with SqlCounter(query_count=1 if func == np.min else 5):
+    query_count = 1 if func == np.min else 5
+    join_count = 0 if func == np.min else 2
+    udtf_count = 0 if func == np.min else 1
+    with SqlCounter(
+        query_count=query_count, join_count=join_count, udtf_count=udtf_count
+    ):
         eval_snowpark_pandas_result(snow_df, native_df, lambda x: x.apply(func, axis=1))
 
 
@@ -107,7 +112,7 @@ def test_axis_1_basic_types_with_type_hints(data, func, return_type):
     snow_df = pd.DataFrame(data, columns=["A", "b"])
     func_with_type_hint = create_func_with_return_type_hint(func, return_type)
     #  Invoking a single UDF typically requires 3 queries (package management, code upload, UDF registration) upfront.
-    with SqlCounter(query_count=4, join_count=0, udtf_count=0):
+    with SqlCounter(query_count=4, join_count=0, udtf_count=0, udf_count=1):
         eval_snowpark_pandas_result(
             snow_df, native_df, lambda x: x.apply(func_with_type_hint, axis=1)
         )
@@ -144,7 +149,7 @@ def foo(row) -> str:
 
     snow_df = pd.DataFrame(df)
     #  Invoking a single UDF typically requires 3 queries (package management, code upload, UDF registration) upfront.
-    with SqlCounter(query_count=4, join_count=0, udtf_count=0):
+    with SqlCounter(query_count=4, join_count=0, udtf_count=0, udf_count=1):
         eval_snowpark_pandas_result(snow_df, df, lambda x: x.apply(foo, axis=1))
 
 
@@ -640,6 +645,8 @@ def test_basic_dataframe_transform(data, apply_func, expected_query_count):
         snow_df = pd.DataFrame(data)
         with SqlCounter(
             query_count=expected_query_count,
+            join_count=5,
+            udtf_count=3,
             high_count_expected=expected_query_count > 8,
             high_count_reason=msg,
         ):
@@ -900,6 +907,7 @@ def test_apply_axis1_with_3rd_party_libraries_and_decorator(
 
     with SqlCounter(
         query_count=expected_query_count,
+        udf_count=1,
         high_count_expected=True,
         high_count_reason="Snowpark package upload requires many queries.",
     ):

tests/integ/modin/frame/test_apply_axis_0.py

@@ -88,6 +88,7 @@ def test_axis_0_basic_types_without_type_hints(data, func, return_type):
         query_count=1 if func == np.min else 11,
         join_count=0 if func == np.min else 2,
         udtf_count=0 if func == np.min else 2,
+        union_count=1,
         high_count_expected=func != np.min,
         high_count_reason="SNOW-1650644 & SNOW-1345395: Avoid extra caching and repeatedly creating same temp function",
     ):
@@ -108,6 +109,7 @@ def test_axis_0_basic_types_with_type_hints(data, func, return_type):
         query_count=11,
         join_count=2,
         udtf_count=2,
+        union_count=1,
         high_count_expected=True,
         high_count_reason="SNOW-1650644 & SNOW-1345395: Avoid extra caching and repeatedly creating same temp function",
     ):
@@ -161,6 +163,7 @@ def foo(row) -> str:
         query_count=11,
         join_count=2,
         udtf_count=2,
+        union_count=1,
         high_count_expected=True,
         high_count_reason="SNOW-1650644 & SNOW-1345395: Avoid extra caching and repeatedly creating same temp function",
     ):

tests/integ/modin/frame/test_iloc.py

@@ -166,7 +166,12 @@ def test_df_row_return_dataframe(
     default_index_snowpark_pandas_df,
     default_index_native_df,
 ):
-    with SqlCounter(query_count=expected_query_count, join_count=expected_join_count):
+    union_count = 1 if isinstance(key, (list, slice)) else 0
+    with SqlCounter(
+        query_count=expected_query_count,
+        join_count=expected_join_count,
+        union_count=union_count,
+    ):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df,
             default_index_native_df,
@@ -406,7 +411,7 @@ def test_df_iloc_get_empty_key(
     default_index_snowpark_pandas_df,
     default_index_native_df,
 ):
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=2):
         eval_snowpark_pandas_result(
             empty_snowpark_pandas_df,
             native_pd.DataFrame(),
@@ -416,7 +421,7 @@ def test_df_iloc_get_empty_key(
             # from native pandas (Snowpark pandas gives "empty" vs. native pandas "integer")
             check_column_type=False,
         )
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=2):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df,
             default_index_native_df,
@@ -1079,7 +1084,7 @@ def iloc_helper(df):
 
     # test df with default index
     num_cols = 7
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, union_count=1):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df,
             default_index_native_df,
@@ -1089,7 +1094,7 @@ def iloc_helper(df):
 
     # test df with non-default index
     num_cols = 6  # set_index() makes the number of columns 6
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, union_count=1):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df.set_index("D"),
             default_index_native_df.set_index("D"),
@@ -1102,7 +1107,7 @@ def iloc_helper(df):
     num_cols = 7
     native_df = default_index_native_df.set_index(multiindex_native)
     snowpark_df = pd.DataFrame(native_df)
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, union_count=1):
         eval_snowpark_pandas_result(
             snowpark_df,
             native_df,
@@ -1116,7 +1121,7 @@ def iloc_helper(df):
         native_df_with_multiindex_columns
     )
     in_range = True if (-8 < key < 7) else False
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, union_count=1):
         if axis == "row" or in_range:  # series result
             eval_snowpark_pandas_result(
                 snowpark_df_with_multiindex_columns,
@@ -1138,7 +1143,7 @@ def iloc_helper(df):
     # 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=1):
+    with SqlCounter(query_count=1, union_count=1):
         if axis == "row" or in_range:  # series result
             eval_snowpark_pandas_result(
                 snowpark_df,

tests/integ/modin/frame/test_isin.py

@@ -38,7 +38,8 @@ def _test_isin_with_snowflake_logic(df: Any, values, query_count=0):  # noqa: E3
         # set expected query counts to 0 if native pandas is used.
         query_count = 0
 
-    with SqlCounter(query_count=query_count):
+    join_count = 1 if isinstance(df, pd.DataFrame) else 0
+    with SqlCounter(query_count=query_count, join_count=join_count):
         ans = df.isin(values)
 
         # Following code is to emulate Snowflake behavior:

tests/integ/modin/frame/test_loc.py

@@ -152,6 +152,7 @@ def test_df_loc_get_tuple_key(
 
     with SqlCounter(
         query_count=query_count,
+        join_count=2,
     ):
         eval_snowpark_pandas_result(
             str_index_snowpark_pandas_df,
@@ -314,15 +315,15 @@ def test_df_loc_get_empty_key(
     default_index_native_df,
 ):
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             empty_snowpark_pandas_df,
             native_pd.DataFrame(),
             lambda df: df.loc[key],
             comparator=assert_snowpark_pandas_equal_to_pandas,
             check_column_type=False,
         )
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             default_index_snowpark_pandas_df,
             default_index_native_df,
@@ -1913,6 +1914,7 @@ def _test_df_loc_with_1k_shape(
     with SqlCounter(
         query_count=query_count,
         join_count=1,
+        union_count=1,
         high_count_expected=high_count_expected,
         high_count_reason=high_count_reason,
     ):
@@ -2575,7 +2577,7 @@ def test_empty_df_loc_set_scalar():
     with pytest.raises(ValueError, match="cannot set a frame with no defined columns"):
         native_df.loc[0] = 1
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         snow_df.loc[0] = 1
         assert_snowpark_pandas_equal_to_pandas(
             snow_df,
@@ -2603,23 +2605,23 @@ def col_loc(df):
     native_df = native_pd.DataFrame(index=[0, 1, 2])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with row scalar on empty DataFrame with non-empty index.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             snow_df, native_df, row_loc, inplace=True, check_column_type=False
         )
 
     native_df = native_pd.DataFrame(index=[0, 1, 2])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with column scalar on empty DataFrame with non-empty index.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             snow_df, native_df, col_loc, inplace=True, check_column_type=False
         )
 
     native_df = native_pd.DataFrame(columns=["A", "B", "C"])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with row scalar on empty DataFrame with non-empty columns.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             snow_df,
             native_df,
@@ -2630,7 +2632,7 @@ def col_loc(df):
     native_df = native_pd.DataFrame(columns=["A", "B", "C"])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with column scalar on empty DataFrame with non-empty columns.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         col_loc(snow_df)
         assert_snowpark_pandas_equal_to_pandas(
             snow_df,
@@ -2642,7 +2644,7 @@ def col_loc(df):
     native_df = native_pd.DataFrame(index=[0, 1, 2], columns=["A", "B", "C"])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with row scalar on empty DataFrame with non-empty index and columns.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             snow_df,
             native_df,
@@ -2653,7 +2655,7 @@ def col_loc(df):
     native_df = native_pd.DataFrame(index=[0, 1, 2], columns=["A", "B", "C"])
     snow_df = pd.DataFrame(native_df)
     # Check `loc` with column scalar on empty DataFrame with non-empty index and columns.
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, join_count=1):
         eval_snowpark_pandas_result(
             snow_df,
             native_df,
@@ -3959,7 +3961,12 @@ def test_df_loc_get_with_timedelta(key, query_count, join_count):
     idx = ["1 days", "2 days", "3 days"]
     native_df = native_pd.DataFrame(data, index=native_pd.to_timedelta(idx))
     snow_df = pd.DataFrame(data, index=pd.to_timedelta(idx))
-    with SqlCounter(query_count=query_count, join_count=join_count):
+    union_count = (
+        1 if isinstance(key, (list, slice)) and not isinstance(key, str) else 0
+    )
+    with SqlCounter(
+        query_count=query_count, join_count=join_count, union_count=union_count
+    ):
         eval_snowpark_pandas_result(snow_df, native_df, lambda df: df.loc[key])
 
 

tests/integ/modin/frame/test_nunique.py

@@ -42,7 +42,7 @@ def test_dataframe_nunique(axes_slices, dropna):
         ]
     )
 
-    with SqlCounter(query_count=1):
+    with SqlCounter(query_count=1, union_count=1 if axes_slices[0] == 0 else 0):
         eval_snowpark_pandas_result(
             df,
             native_df,

tests/integ/modin/frame/test_squeeze.py

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