SNOW-1830648: Fix nest query complexity adjustment (#2691)

src/snowflake/snowpark/_internal/analyzer/select_statement.py

@@ -997,14 +997,15 @@ def projection_complexities(self) -> List[Dict[PlanNodeCategory, int]]:
                     #       - {PlanNodeCategory.COLUMN: 1} + col2_complexity
                     #       - {PlanNodeCategory.COLUMN: 1} + col1_complexity
                     dependent_columns = proj.dependent_column_names_with_duplication()
-                    projection_complexity = proj.cumulative_node_complexity
+                    projection_complexity = proj.cumulative_node_complexity.copy()
                     for dependent_column in dependent_columns:
                         dependent_column_complexity = (
                             subquery_projection_name_complexity_map[dependent_column]
                         )
-                        projection_complexity[PlanNodeCategory.COLUMN] -= 1
                         projection_complexity = sum_node_complexities(
-                            projection_complexity, dependent_column_complexity
+                            projection_complexity,
+                            dependent_column_complexity,
+                            {PlanNodeCategory.COLUMN: -1},
                         )
 
                     self._projection_complexities.append(projection_complexity)

tests/integ/test_large_query_breakdown.py

@@ -18,6 +18,9 @@
     DEFAULT_COMPLEXITY_SCORE_UPPER_BOUND,
     Session,
 )
+from tests.integ.test_deepcopy import (
+    create_df_with_deep_nested_with_column_dependencies,
+)
 from tests.integ.utils.sql_counter import SqlCounter, sql_count_checker
 from tests.utils import IS_IN_STORED_PROC, Utils
 
@@ -604,6 +607,31 @@ def test_optimization_skipped_with_views_and_dynamic_tables(session, caplog):
         Utils.drop_table(session, source_table)
 
 
+def test_large_query_breakdown_with_nested_select(session):
+    if not session.sql_simplifier_enabled:
+        pytest.skip(
+            "the nested select optimization is only enabled with sql simplifier"
+        )
+
+    temp_table_name = Utils.random_table_name()
+    final_df = create_df_with_deep_nested_with_column_dependencies(
+        session, temp_table_name, 8
+    )
+
+    with SqlCounter(query_count=1, describe_count=0):
+        queries = final_df.queries
+    assert len(queries["queries"]) == 3
+    assert queries["queries"][0].startswith("CREATE  SCOPED TEMPORARY  TABLE")
+    assert queries["queries"][1].startswith("CREATE  SCOPED TEMPORARY  TABLE")
+
+    assert len(queries["post_actions"]) == 2
+    assert queries["post_actions"][0].startswith("DROP  TABLE  If  EXISTS")
+    assert queries["post_actions"][1].startswith("DROP  TABLE  If  EXISTS")
+
+    with SqlCounter(query_count=7, describe_count=0):
+        check_result_with_and_without_breakdown(session, final_df)
+
+
 @pytest.mark.skipif(
     IS_IN_STORED_PROC, reason="cannot create a new session in stored procedure"
 )

tests/integ/test_multithreading.py

@@ -688,7 +688,7 @@ def change_config_value(session_):
     run=False,
 )
 def test_large_query_breakdown_with_cte(threadsafe_session):
-    bounds = (300, 600) if threadsafe_session.sql_simplifier_enabled else (50, 70)
+    bounds = (300, 520) if threadsafe_session.sql_simplifier_enabled else (50, 70)
     try:
         original_query_compilation_stage_enabled = (
             threadsafe_session._query_compilation_stage_enabled

tests/integ/test_nested_select_plan_analysis.py

@@ -495,7 +495,7 @@ def test_deep_nested_with_columns(session):
                 PlanNodeCategory.FUNCTION: 60,
             },
             complexity_after_merge={
-                PlanNodeCategory.COLUMN: 389,
+                PlanNodeCategory.COLUMN: 532,
                 PlanNodeCategory.CASE_WHEN: 40,
                 PlanNodeCategory.LOW_IMPACT: 80,
                 PlanNodeCategory.LITERAL: 200,