SNOW-2296598: Estimate row count only when Hybrid is enabled (#3708)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -51,6 +51,8 @@
 - Set the default transfer limit in hybrid execution for data leaving Snowflake to 100k, which can be overridden with the SnowflakePandasTransferThreshold environment variable. This configuration is appropriate for scenarios with two available engines, "Pandas" and "Snowflake" on relational workloads.
 - Improve import error message by adding '--upgrade' to 'pip install "snowflake-snowpark-python[modin]"' in the error message.
 - Reduce the telemetry messages from the modin client by pre-aggregating into 5 second windows and only keeping a narrow band of metrics which are useful for tracking hybrid execution and native pandas performance.
+- Set the initial row count only when hybrid execution is enabled. This reduces the number of queries issued for many workloads.
+
 
 #### Dependency Updates
 

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

@@ -76,7 +76,7 @@ def upper_bound(
         elif operation == DataFrameOperation.UNION_ALL:
             other: OrderedDataFrame = args["other"]
             other_bound = other.row_count_upper_bound or other.row_count
-            if other_bound is None:
+            if other_bound is None or current is None:
                 # Cannot estimate row count: other DataFrame has no row count information
                 return None
             return current + other_bound

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

@@ -498,12 +498,20 @@ def create_initial_ordered_dataframe(
         row_position_snowflake_quoted_identifier = (
             ordered_dataframe.row_position_snowflake_quoted_identifier
         )
-    # Set the materialized row count
-    materialized_row_count = ordered_dataframe._dataframe_ref.snowpark_dataframe.count(
-        statement_params=get_default_snowpark_pandas_statement_params(), _emit_ast=False
-    )
-    ordered_dataframe.row_count = materialized_row_count
-    ordered_dataframe.row_count_upper_bound = materialized_row_count
+
+    from modin.config import AutoSwitchBackend
+
+    if AutoSwitchBackend.get():
+        # Set the materialized row count
+        materialized_row_count = (
+            initial_ordered_dataframe._dataframe_ref.snowpark_dataframe.count(
+                statement_params=get_default_snowpark_pandas_statement_params(),
+                _emit_ast=False,
+            )
+        )
+        ordered_dataframe.row_count = materialized_row_count
+        ordered_dataframe.row_count_upper_bound = materialized_row_count
+
     return ordered_dataframe, row_position_snowflake_quoted_identifier
 
 

tests/integ/modin/frame/test_aggregate.py

@@ -812,7 +812,7 @@ def test_sum_min_count(min_count, axis):
     )
 
 
-@sql_count_checker(query_count=3, union_count=4)
+@sql_count_checker(query_count=2, union_count=4)
 def test_agg_valid_variant_col(session, test_table_name):
     pandas_df = native_pd.DataFrame(
         {

tests/integ/modin/frame/test_create_or_replace_dynamic_table.py

@@ -14,7 +14,7 @@
 from tests.utils import Utils
 
 
-@sql_count_checker(query_count=7)
+@sql_count_checker(query_count=5)
 def test_create_or_replace_dynamic_table_enforce_ordering_raises(session) -> None:
     try:
         # create table
@@ -48,7 +48,7 @@ def test_create_or_replace_dynamic_table_enforce_ordering_raises(session) -> Non
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=5)
 def test_create_or_replace_dynamic_table_no_enforce_ordering(session) -> None:
     try:
         # create table
@@ -84,7 +84,7 @@ def test_create_or_replace_dynamic_table_no_enforce_ordering(session) -> None:
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=5)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_dynamic_table_multiple_sessions_no_enforce_ordering(
     session,
     db_parameters,
@@ -131,7 +131,7 @@ def test_create_or_replace_dynamic_table_multiple_sessions_no_enforce_ordering(
 
 @pytest.mark.parametrize("index", [True, False])
 @pytest.mark.parametrize("index_labels", [None, ["my_index"]])
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_dynamic_table_index(session, index, index_labels):
     try:
         # create table
@@ -177,7 +177,7 @@ def test_create_or_replace_dynamic_table_index(session, index, index_labels):
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_dynamic_table_multiindex(session):
     try:
         # create table

tests/integ/modin/frame/test_create_or_replace_view.py

@@ -42,7 +42,7 @@ def test_create_or_replace_view_basic(session, native_pandas_df_basic) -> None:
         Utils.drop_view(session, view_name)
 
 
-@sql_count_checker(query_count=8)
+@sql_count_checker(query_count=6)
 def test_create_or_replace_view_multiple_sessions_enforce_ordering_raises(
     session,
     db_parameters,
@@ -85,7 +85,7 @@ def test_create_or_replace_view_multiple_sessions_enforce_ordering_raises(
         pd.session = session
 
 
-@sql_count_checker(query_count=5)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_view_multiple_sessions_no_enforce_ordering(
     session,
     db_parameters,
@@ -127,7 +127,7 @@ def test_create_or_replace_view_multiple_sessions_no_enforce_ordering(
 
 @pytest.mark.parametrize("index", [True, False])
 @pytest.mark.parametrize("index_labels", [None, ["my_index"]])
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_view_index(session, index, index_labels):
     try:
         # create table
@@ -167,7 +167,7 @@ def test_create_or_replace_view_index(session, index, index_labels):
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_create_or_replace_view_multiindex(session):
     try:
         # create table

tests/integ/modin/frame/test_describe.py

@@ -349,7 +349,7 @@ def helper(df):
 
 
 @sql_count_checker(
-    query_count=4,
+    query_count=3,
     union_count=8,
 )
 # SNOW-1320296 - pd.concat SQL Compilation ambigious __row_position__ issue

tests/integ/modin/frame/test_to_dynamic_table.py

@@ -27,7 +27,7 @@ def to_dynamic_table(request):
     return request.param
 
 
-@sql_count_checker(query_count=5)
+@sql_count_checker(query_count=4)
 def test_to_dynamic_table_enforce_ordering_raises(session, to_dynamic_table) -> None:
     try:
         # create table
@@ -62,7 +62,7 @@ def test_to_dynamic_table_enforce_ordering_raises(session, to_dynamic_table) ->
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=5)
 def test_to_dynamic_table_no_enforce_ordering(session, to_dynamic_table) -> None:
     try:
         # create table
@@ -98,7 +98,7 @@ def test_to_dynamic_table_no_enforce_ordering(session, to_dynamic_table) -> None
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=5)
+@sql_count_checker(query_count=4)
 def test_to_dynamic_table_multiple_sessions_no_enforce_ordering(
     session,
     db_parameters,
@@ -153,7 +153,7 @@ def test_to_dynamic_table_multiple_sessions_no_enforce_ordering(
         (False, ["my_index"], []),
     ],
 )
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_to_dynamic_table_index(
     session, index, index_labels, expected_index_columns, to_dynamic_table
 ):
@@ -204,7 +204,7 @@ def test_to_dynamic_table_index(
         Utils.drop_table(session, table_name)
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_to_dynamic_table_multiindex(session, to_dynamic_table):
     try:
         # create table

tests/integ/modin/frame/test_to_iceberg.py

@@ -41,7 +41,7 @@ def to_iceberg(request):
     return request.param
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=5)
 def test_to_iceberg(session, native_pandas_df_basic, to_iceberg):
     if not iceberg_supported(session, local_testing_mode=False):
         pytest.skip("Test requires iceberg support.")

tests/integ/modin/frame/test_to_snowflake.py

@@ -17,7 +17,7 @@
 @pytest.mark.parametrize("index", [True, False])
 @pytest.mark.parametrize("index_labels", [None, ["my_index"]])
 # one extra query to convert index to native pandas when creating the snowpark pandas dataframe
-@sql_count_checker(query_count=3)
+@sql_count_checker(query_count=2)
 def test_to_snowflake_index(test_table_name, index, index_labels):
     df = pd.DataFrame(
         {"a": [1, 2, 3], "b": [4, 5, 6]}, index=pd.Index([2, 3, 4], name="index")
@@ -38,7 +38,7 @@ def test_to_snowflake_index(test_table_name, index, index_labels):
     verify_columns(test_table_name, expected_columns)
 
 
-@sql_count_checker(query_count=2)
+@sql_count_checker(query_count=1)
 def test_to_snowflake_multiindex(test_table_name):
     index = native_pd.MultiIndex.from_arrays(
         [[1, 1, 2, 2], ["red", "blue", "red", "blue"]], names=("number", "color")
@@ -94,7 +94,7 @@ def test_to_snowflake_if_exists(session, test_table_name):
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
 
     # Verify new table is created
-    with SqlCounter(query_count=3):
+    with SqlCounter(query_count=2):
         df.to_snowflake(test_table_name, if_exists="fail", index=False)
         verify_columns(test_table_name, ["a", "b"])
 
@@ -110,19 +110,19 @@ def test_to_snowflake_if_exists(session, test_table_name):
 
     # Verify existing table is replaced with new data
     df = pd.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]})
-    with SqlCounter(query_count=3):
+    with SqlCounter(query_count=2):
         df.to_snowflake(test_table_name, if_exists="replace", index=False)
         verify_columns(test_table_name, ["a", "c"])
         verify_num_rows(session, test_table_name, 3)
 
     # Verify data is appended to existing table
-    with SqlCounter(query_count=4):
+    with SqlCounter(query_count=3):
         df.to_snowflake(test_table_name, if_exists="append", index=False)
         verify_columns(test_table_name, ["a", "c"])
         verify_num_rows(session, test_table_name, 6)
 
     # Verify pd.to_snowflake operates the same
-    with SqlCounter(query_count=4):
+    with SqlCounter(query_count=3):
         pd.to_snowflake(df, test_table_name, if_exists="append", index=False)
         verify_columns(test_table_name, ["a", "c"])
         verify_num_rows(session, test_table_name, 9)
@@ -134,7 +134,7 @@ def test_to_snowflake_if_exists(session, test_table_name):
 
 
 @pytest.mark.parametrize("index_label", VALID_PANDAS_LABELS)
-@sql_count_checker(query_count=2)
+@sql_count_checker(query_count=1)
 def test_to_snowflake_index_labels(index_label, test_table_name):
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
     df.to_snowflake(
@@ -144,7 +144,7 @@ def test_to_snowflake_index_labels(index_label, test_table_name):
 
 
 @pytest.mark.parametrize("col_name", VALID_PANDAS_LABELS)
-@sql_count_checker(query_count=2)
+@sql_count_checker(query_count=1)
 def test_to_snowflake_column_names_from_panadas(col_name, test_table_name):
     df = pd.DataFrame({col_name: [1, 2, 3], "b": [4, 5, 6]})
     df.to_snowflake(test_table_name, if_exists="replace", index=False)
@@ -156,7 +156,7 @@ def test_to_snowflake_column_names_from_panadas(col_name, test_table_name):
 def test_column_names_with_read_snowflake_and_to_snowflake(
     col_name, if_exists, session
 ):
-    with SqlCounter(query_count=7 if if_exists == "append" else 6):
+    with SqlCounter(query_count=6 if if_exists == "append" else 5):
         # Create a table
         session.sql(f"create or replace table t1 ({col_name} int)").collect()
         session.sql("insert into t1 values (1), (2), (3)").collect()
@@ -173,7 +173,7 @@ def test_column_names_with_read_snowflake_and_to_snowflake(
         assert len(data) == (6 if if_exists == "append" else 3)
 
 
-@sql_count_checker(query_count=2)
+@sql_count_checker(query_count=1)
 def test_to_snowflake_column_with_quotes(session, test_table_name):
     df = pd.DataFrame({'a"b': [1, 2, 3], 'a""b': [4, 5, 6]})
     df.to_snowflake(test_table_name, if_exists="replace", index=False)
@@ -183,28 +183,28 @@ def test_to_snowflake_column_with_quotes(session, test_table_name):
 # one extra query to convert index to native pandas when creating the snowpark pandas dataframe
 def test_to_snowflake_index_label_none(test_table_name):
     # no index
-    with SqlCounter(query_count=2):
+    with SqlCounter(query_count=1):
         df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
         df.to_snowflake(test_table_name, if_exists="replace")
         verify_columns(test_table_name, ["index", "a", "b"])
 
     # named index
-    with SqlCounter(query_count=3):
+    with SqlCounter(query_count=2):
         df = pd.DataFrame(
             {"a": [1, 2, 3], "b": [4, 5, 6]}, index=pd.Index([2, 3, 4], name="index")
         )
         df.to_snowflake(test_table_name, if_exists="replace", index_label=[None])
         verify_columns(test_table_name, ["index", "a", "b"])
 
     # nameless index
-    with SqlCounter(query_count=3):
+    with SqlCounter(query_count=2):
         df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=pd.Index([2, 3, 4]))
         df.to_snowflake(test_table_name, if_exists="replace", index_label=[None])
         verify_columns(test_table_name, ["index", "a", "b"])
 
 
 # one extra query to convert index to native pandas when creating the snowpark pandas dataframe
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=4)
 def test_to_snowflake_index_label_none_data_column_conflict(test_table_name):
     df = pd.DataFrame({"index": [1, 2, 3], "a": [4, 5, 6]})
     df.to_snowflake(test_table_name, if_exists="replace")
@@ -260,7 +260,7 @@ def verify_num_rows(session, table_name: str, expected: int) -> None:
     assert actual == expected
 
 
-@sql_count_checker(query_count=2)
+@sql_count_checker(query_count=1)
 def test_timedelta_to_snowflake_with_read_snowflake(test_table_name, caplog):
     with caplog.at_level(logging.WARNING):
         df = pd.DataFrame(