SNOW-1819531: propagate referenced ctes to all root nodes

src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py

@@ -477,6 +477,14 @@ def _is_relaxed_pipeline_breaker(self, node: LogicalPlan) -> bool:
         if isinstance(node, SelectStatement):
             return True
 
+        if isinstance(node, SnowflakePlan):
+            return node.source_plan is not None and self._is_relaxed_pipeline_breaker(
+                node.source_plan
+            )
+
+        if isinstance(node, SelectSnowflakePlan):
+            return self._is_relaxed_pipeline_breaker(node.snowflake_plan)
+
         return False
 
     def _is_node_pipeline_breaker(self, node: LogicalPlan) -> bool:

src/snowflake/snowpark/_internal/compiler/query_generator.py

@@ -169,6 +169,7 @@ def do_resolve_with_resolved_children(
                 iceberg_config=logical_plan.iceberg_config,
                 table_exists=logical_plan.table_exists,
             )
+            resolved_plan.referenced_ctes = resolved_child.referenced_ctes
 
         elif isinstance(
             logical_plan,
@@ -197,6 +198,7 @@ def do_resolve_with_resolved_children(
             resolved_plan = super().do_resolve_with_resolved_children(
                 logical_plan, resolved_children, df_aliased_col_name_to_real_col_name
             )
+            resolved_plan.referenced_ctes = resolved_child.referenced_ctes
 
         elif isinstance(logical_plan, Selectable):
             # overwrite the Selectable resolving to make sure we are triggering

src/snowflake/snowpark/_internal/compiler/utils.py

@@ -34,7 +34,10 @@
     TableMerge,
     TableUpdate,
 )
-from snowflake.snowpark._internal.analyzer.unary_plan_node import UnaryNode
+from snowflake.snowpark._internal.analyzer.unary_plan_node import (
+    CreateViewCommand,
+    UnaryNode,
+)
 from snowflake.snowpark._internal.compiler.query_generator import (
     QueryGenerator,
     SnowflakeCreateTablePlanInfo,
@@ -121,7 +124,9 @@ def to_selectable(plan: LogicalPlan, query_generator: QueryGenerator) -> Selecta
             return plan
 
         snowflake_plan = query_generator.resolve(plan)
-        return SelectSnowflakePlan(snowflake_plan, analyzer=query_generator)
+        selectable = SelectSnowflakePlan(snowflake_plan, analyzer=query_generator)
+        selectable._is_valid_for_replacement = snowflake_plan._is_valid_for_replacement
+        return selectable
 
     if not parent._is_valid_for_replacement:
         raise ValueError(f"parent node {parent} is not valid for replacement.")
@@ -306,7 +311,28 @@ def get_snowflake_plan_queries(
 
     plan_queries = plan.queries
     post_action_queries = plan.post_actions
-    if len(plan.referenced_ctes) > 0:
+    # If the plan has referenced ctes, we need to add the cte definition before
+    # the final query. This is done for all source plan except for the following
+    # cases:
+    # - SnowflakeCreateTable
+    # - CreateViewCommand
+    # - TableUpdate
+    # - TableDelete
+    # - TableMerge
+    # - CopyIntoLocationNode
+    # because the generated_queries by QueryGenerator for these nodes already include the cte
+    # definition. Adding the cte definition before the query again will cause a syntax error.
+    if len(plan.referenced_ctes) > 0 and not isinstance(
+        plan.source_plan,
+        (
+            SnowflakeCreateTable,
+            CreateViewCommand,
+            TableUpdate,
+            TableDelete,
+            TableMerge,
+            CopyIntoLocationNode,
+        ),
+    ):
         # make a copy of the original query to avoid any update to the
         # original query object
         plan_queries = copy.deepcopy(plan.queries)
@@ -397,6 +423,9 @@ def get_name(node: Optional[LogicalPlan]) -> str:
             name = f"{name} :: ({'| '.join(properties)})"
 
         score = get_complexity_score(node)
+        num_ref_ctes = "nil"
+        if isinstance(node, (SnowflakePlan, Selectable)):
+            num_ref_ctes = len(node.referenced_ctes)
         sql_text = ""
         if isinstance(node, Selectable):
             sql_text = node.sql_query
@@ -405,7 +434,7 @@ def get_name(node: Optional[LogicalPlan]) -> str:
         sql_size = len(sql_text)
         sql_preview = sql_text[:50]
 
-        return f"{name=}\n{score=}, {sql_size=}\n{sql_preview=}"
+        return f"{name=}\n{score=}, {num_ref_ctes=}, {sql_size=}\n{sql_preview=}"
 
     g = graphviz.Graph(format="png")
 
@@ -415,7 +444,8 @@ def get_name(node: Optional[LogicalPlan]) -> str:
         next_level = []
         for node in curr_level:
             node_id = hex(id(node))
-            g.node(node_id, get_stat(node))
+            color = "lightblue" if node._is_valid_for_replacement else "red"
+            g.node(node_id, get_stat(node), color=color)
             if isinstance(node, (Selectable, SnowflakePlan)):
                 children = node.children_plan_nodes
             else:

tests/integ/scala/test_datatype_suite.py

@@ -600,11 +600,14 @@ def test_iceberg_nested_fields(
         Utils.drop_table(structured_type_session, transformed_table_name)
 
 
-@pytest.mark.skip(
-    reason="SNOW-1819531: Error in _contains_external_cte_ref when analyzing lqb"
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="local testing does not fully support structured types yet.",
+    run=False,
 )
+@pytest.mark.parametrize("cte_enabled", [True, False])
 def test_struct_dtype_iceberg_lqb(
-    structured_type_session, local_testing_mode, structured_type_support
+    structured_type_session, local_testing_mode, structured_type_support, cte_enabled
 ):
     if not (
         structured_type_support
@@ -641,12 +644,14 @@ def test_struct_dtype_iceberg_lqb(
     is_query_compilation_stage_enabled = (
         structured_type_session._query_compilation_stage_enabled
     )
+    is_cte_optimization_enabled = structured_type_session._cte_optimization_enabled
     is_large_query_breakdown_enabled = (
         structured_type_session._large_query_breakdown_enabled
     )
     original_bounds = structured_type_session._large_query_breakdown_complexity_bounds
     try:
         structured_type_session._query_compilation_stage_enabled = True
+        structured_type_session._cte_optimization_enabled = cte_enabled
         structured_type_session._large_query_breakdown_enabled = True
         structured_type_session._large_query_breakdown_complexity_bounds = (300, 600)
 
@@ -695,9 +700,14 @@ def test_struct_dtype_iceberg_lqb(
         )
 
         queries = union_df.queries
-        # assert that the queries are broken down into 2 queries and 1 post action
-        assert len(queries["queries"]) == 2, queries["queries"]
-        assert len(queries["post_actions"]) == 1
+        if cte_enabled:
+            # when CTE is enabled, WithQueryBlock makes pipeline breaker nodes ineligible
+            assert len(queries["queries"]) == 1
+            assert len(queries["post_actions"]) == 0
+        else:
+            # assert that the queries are broken down into 2 queries and 1 post action
+            assert len(queries["queries"]) == 2, queries["queries"]
+            assert len(queries["post_actions"]) == 1
         final_df = structured_type_session.table(write_table)
 
         # assert that
@@ -707,6 +717,7 @@ def test_struct_dtype_iceberg_lqb(
         structured_type_session._query_compilation_stage_enabled = (
             is_query_compilation_stage_enabled
         )
+        structured_type_session._cte_optimization_enabled = is_cte_optimization_enabled
         structured_type_session._large_query_breakdown_enabled = (
             is_large_query_breakdown_enabled
         )

tests/unit/compiler/test_large_query_breakdown.py

@@ -135,3 +135,42 @@ def test_pipeline_breaker_node(mock_session, mock_analyzer, node_generator, expe
         large_query_breakdown._is_node_pipeline_breaker(select_snowflake_plan)
         is expected
     ), "SelectSnowflakePlan node is not detected as a pipeline breaker node"
+
+
+@pytest.mark.parametrize(
+    "node_generator,expected",
+    [
+        (
+            lambda x: SelectStatement(
+                from_=empty_selectable, order_by=[empty_expression], analyzer=x
+            ),
+            True,
+        ),
+    ],
+)
+def test_relaxed_pipeline_breaker_node(
+    mock_session, mock_analyzer, node_generator, expected
+):
+    large_query_breakdown = LargeQueryBreakdown(
+        mock_session,
+        mock_analyzer,
+        [],
+        mock_session.large_query_breakdown_complexity_bounds,
+    )
+    node = node_generator(mock_analyzer)
+
+    assert (
+        large_query_breakdown._is_relaxed_pipeline_breaker(node) is expected
+    ), f"Node {type(node)} is not detected as a pipeline breaker node"
+
+    resolved_node = mock_analyzer.resolve(node)
+    assert isinstance(resolved_node, SnowflakePlan)
+    assert (
+        large_query_breakdown._is_relaxed_pipeline_breaker(resolved_node) is expected
+    ), f"Resolved node of {type(node)} is not detected as a pipeline breaker node"
+
+    select_snowflake_plan = SelectSnowflakePlan(resolved_node, analyzer=mock_analyzer)
+    assert (
+        large_query_breakdown._is_relaxed_pipeline_breaker(select_snowflake_plan)
+        is expected
+    ), "SelectSnowflakePlan node is not detected as a pipeline breaker node"