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

@@ -6,7 +6,10 @@
 
 from snowflake.snowpark._internal.analyzer.analyzer import Analyzer
 from snowflake.snowpark._internal.analyzer.expression import Attribute
-from snowflake.snowpark._internal.analyzer.select_statement import Selectable
+from snowflake.snowpark._internal.analyzer.select_statement import (
+    SelectSnowflakePlan,
+    Selectable,
+)
 from snowflake.snowpark._internal.analyzer.snowflake_plan import (
     PlanQueryType,
     Query,
@@ -66,6 +69,16 @@ def __init__(
         # between the CTE definition is satisfied.
         self.resolved_with_query_block: Dict[str, Query] = {}
 
+    def to_selectable(self, plan: LogicalPlan) -> Selectable:
+        """Given a LogicalPlan, convert it to a Selectable."""
+        if isinstance(plan, Selectable):
+            return plan
+
+        snowflake_plan = self.resolve(plan)
+        selectable = SelectSnowflakePlan(snowflake_plan, analyzer=self)
+        selectable._is_valid_for_replacement = snowflake_plan._is_valid_for_replacement
+        return selectable
+
     def generate_queries(
         self, logical_plans: List[LogicalPlan]
     ) -> Dict[PlanQueryType, List[Query]]:

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

@@ -118,14 +118,6 @@ def replace_child(
     based on the parent node type.
     """
 
-    def to_selectable(plan: LogicalPlan, query_generator: QueryGenerator) -> Selectable:
-        """Given a LogicalPlan, convert it to a Selectable."""
-        if isinstance(plan, Selectable):
-            return plan
-
-        snowflake_plan = query_generator.resolve(plan)
-        return SelectSnowflakePlan(snowflake_plan, analyzer=query_generator)
-
     if not parent._is_valid_for_replacement:
         raise ValueError(f"parent node {parent} is not valid for replacement.")
 
@@ -143,13 +135,13 @@ def to_selectable(plan: LogicalPlan, query_generator: QueryGenerator) -> Selecta
         replace_child(parent.source_plan, old_child, new_child, query_generator)
 
     elif isinstance(parent, SelectStatement):
-        parent.from_ = to_selectable(new_child, query_generator)
+        parent.from_ = query_generator.to_selectable(new_child)
         # once the subquery is updated, set _merge_projection_complexity_with_subquery to False to
         # disable the projection complexity merge
         parent._merge_projection_complexity_with_subquery = False
 
     elif isinstance(parent, SetStatement):
-        new_child_as_selectable = to_selectable(new_child, query_generator)
+        new_child_as_selectable = query_generator.to_selectable(new_child)
         parent._nodes = [
             node if node != old_child else new_child_as_selectable
             for node in parent._nodes
@@ -421,6 +413,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
@@ -429,7 +424,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")
 
@@ -439,7 +434,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"

tests/unit/compiler/test_replace_child_and_update_node.py

@@ -3,6 +3,7 @@
 #
 
 import copy
+from functools import partial
 from unittest import mock
 
 import pytest
@@ -67,6 +68,7 @@ def mock_snowflake_plan() -> SnowflakePlan:
     with_query_block = WithQueryBlock(name="TEST_CTE", child=LogicalPlan())
     fake_snowflake_plan.referenced_ctes = {with_query_block: 1}
     fake_snowflake_plan._cumulative_node_complexity = {}
+    fake_snowflake_plan._is_valid_for_replacement = True
     return fake_snowflake_plan
 
 
@@ -82,6 +84,9 @@ def mock_resolve(x):
     fake_query_generator = mock.create_autospec(QueryGenerator)
     fake_query_generator.resolve.side_effect = mock_resolve
     fake_query_generator.session = mock_session
+    fake_query_generator.to_selectable = partial(
+        QueryGenerator.to_selectable, fake_query_generator
+    )
     return fake_query_generator