SNOW-1541091: Add plan compiler component (#1946)

Author: sfc-gh-aalam

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",

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"]]:

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

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
+#

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,
+        }