From e14b78b18f66dba9fea35b0ee96de6c9481c56f8 Mon Sep 17 00:00:00 2001
From: Afroz Alam <afroz.alam@snowflake.com>
Date: Sun, 21 Jul 2024 00:00:54 +0000
Subject: [PATCH] SNOW-1541091: Add plan compiler component (#1946)

---
 setup.py                                      |  1 +
 .../_internal/analyzer/snowflake_plan.py      | 10 +++----
 .../snowpark/_internal/compiler/__init__.py   |  3 ++
 .../_internal/compiler/plan_compiler.py       | 28 +++++++++++++++++++
 4 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 src/snowflake/snowpark/_internal/compiler/__init__.py
 create mode 100644 src/snowflake/snowpark/_internal/compiler/plan_compiler.py

diff --git a/setup.py b/setup.py
index a363f4ad46a..9994ee58ced 100644
--- a/setup.py
+++ b/setup.py
@@ -87,6 +87,7 @@
         "snowflake.snowpark",
         "snowflake.snowpark._internal",
         "snowflake.snowpark._internal.analyzer",
+        "snowflake.snowpark._internal.compiler",
         "snowflake.snowpark.mock",
         "snowflake.snowpark.modin",
         "snowflake.snowpark.modin.config",
diff --git a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
index 9f9d39055d8..47f3a24b565 100644
--- a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
+++ b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
@@ -266,12 +266,10 @@ def execution_queries(self) -> Dict["PlanQueryType", List["Query"]]:
         -------
         A mapping between the PlanQueryType and the list of Queries corresponds to the type.
         """
-        # apply optimizations
-        final_plan = self.replace_repeated_subquery_with_cte()
-        return {
-            PlanQueryType.QUERIES: final_plan.queries,
-            PlanQueryType.POST_ACTIONS: final_plan.post_actions,
-        }
+        from snowflake.snowpark._internal.compiler.plan_compiler import PlanCompiler
+
+        compiler = PlanCompiler(self)
+        return compiler.compile()
 
     @property
     def children_plan_nodes(self) -> List[Union["Selectable", "SnowflakePlan"]]:
diff --git a/src/snowflake/snowpark/_internal/compiler/__init__.py b/src/snowflake/snowpark/_internal/compiler/__init__.py
new file mode 100644
index 00000000000..0fbef920926
--- /dev/null
+++ b/src/snowflake/snowpark/_internal/compiler/__init__.py
@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
+#
diff --git a/src/snowflake/snowpark/_internal/compiler/plan_compiler.py b/src/snowflake/snowpark/_internal/compiler/plan_compiler.py
new file mode 100644
index 00000000000..3d6702c1b5d
--- /dev/null
+++ b/src/snowflake/snowpark/_internal/compiler/plan_compiler.py
@@ -0,0 +1,28 @@
+#
+# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Dict, List
+
+from snowflake.snowpark._internal.analyzer.snowflake_plan import (
+    PlanQueryType,
+    Query,
+    SnowflakePlan,
+)
+
+
+class PlanCompiler:
+    """This class is common point of entry from compiling SnowflakePlan to list of queries and post actions.
+    We run pre-check and pre-process steps for optimization steps, and apply the activated optimizations.
+    """
+
+    def __init__(self, plan: SnowflakePlan) -> None:
+        self._plan = plan
+
+    def compile(self) -> Dict[PlanQueryType, List[Query]]:
+        # apply optimizations
+        final_plan = self._plan.replace_repeated_subquery_with_cte()
+        return {
+            PlanQueryType.QUERIES: final_plan.queries,
+            PlanQueryType.POST_ACTIONS: final_plan.post_actions,
+        }