address comments

snowflakedb · Oct 22, 2024 · a26656e · a26656e
1 parent 3b797ce
commit a26656e
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py b/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py
@@ -401,9 +401,13 @@ def _contains_external_cte_ref(self, node: TreeNode, root: TreeNode) -> bool:
         if isinstance(node, SnowflakePlan) and isinstance(
             node.source_plan, WithQueryBlock
         ):
-            return False
+            ignore_with_query_block = node.source_plan
+        else:
+            ignore_with_query_block = None
 
         for with_query_block, node_count in node.referenced_ctes.items():
+            if with_query_block is ignore_with_query_block:
+                continue
             root_count = root.referenced_ctes[with_query_block]
             if node_count != root_count:
                 return True