From a26656edc28fb18b023b6cb5b2a19229a7e66858 Mon Sep 17 00:00:00 2001 From: Afroz Alam Date: Tue, 22 Oct 2024 16:06:51 -0700 Subject: [PATCH] address comments --- .../snowpark/_internal/compiler/large_query_breakdown.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py b/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py index 767cefcd9c..e2ca9fef02 100644 --- 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