Skip to content

Commit

Permalink
SNOW-1541091: Add plan compiler component (#1946)
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-aalam authored Jul 21, 2024
1 parent dc16701 commit e14b78b
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 6 deletions.
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
10 changes: 4 additions & 6 deletions src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]]:
Expand Down
3 changes: 3 additions & 0 deletions src/snowflake/snowpark/_internal/compiler/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
#
# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
#
28 changes: 28 additions & 0 deletions src/snowflake/snowpark/_internal/compiler/plan_compiler.py
Original file line number Diff line number Diff line change
@@ -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,
}

0 comments on commit e14b78b

Please sign in to comment.