From 6fcf71fd43b8b891b3f36d93a3c6e4f86b72f261 Mon Sep 17 00:00:00 2001
From: Jianzhun Du <jianzhun.du@snowflake.com>
Date: Tue, 5 Mar 2024 10:02:51 -0800
Subject: [PATCH] copy

---
 .../_internal/analyzer/snowflake_plan.py      | 38 +++++++++++++------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
index ed7734b7d5c..e2661cb9ac0 100644
--- a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
+++ b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
@@ -320,18 +320,32 @@ def output_dict(self) -> Dict[str, Any]:
         return self._output_dict
 
     def __copy__(self) -> "SnowflakePlan":
-        return SnowflakePlan(
-            copy.deepcopy(self.queries) if self.queries else [],
-            self.schema_query,
-            copy.deepcopy(self.post_actions) if self.post_actions else None,
-            dict(self.expr_to_alias) if self.expr_to_alias else None,
-            self.source_plan,
-            self.is_ddl_on_temp_object,
-            copy.deepcopy(self.api_calls) if self.api_calls else None,
-            self.df_aliased_col_name_to_real_col_name,
-            session=self.session,
-            placeholder_query=self.placeholder_query,
-        )
+        if self.session._cte_optimization_enabled:
+            return SnowflakePlan(
+                copy.deepcopy(self.queries) if self.queries else [],
+                self.schema_query,
+                copy.deepcopy(self.post_actions) if self.post_actions else None,
+                dict(self.expr_to_alias) if self.expr_to_alias else None,
+                self.source_plan,
+                self.is_ddl_on_temp_object,
+                copy.deepcopy(self.api_calls) if self.api_calls else None,
+                self.df_aliased_col_name_to_real_col_name,
+                session=self.session,
+                placeholder_query=self.placeholder_query,
+            )
+        else:
+            return SnowflakePlan(
+                self.queries.copy() if self.queries else [],
+                self.schema_query,
+                self.post_actions.copy() if self.post_actions else None,
+                dict(self.expr_to_alias) if self.expr_to_alias else None,
+                self.source_plan,
+                self.is_ddl_on_temp_object,
+                self.api_calls.copy() if self.api_calls else None,
+                self.df_aliased_col_name_to_real_col_name,
+                session=self.session,
+                placeholder_query=self.placeholder_query,
+            )
 
     def add_aliases(self, to_add: Dict) -> None:
         self.expr_to_alias = {**self.expr_to_alias, **to_add}