/* PR_START p--cte 17 */ Fix SQL optimization level not getting passe…

…d from the request.
dbt-labs · Nov 13, 2024 · 600d5f9 · 600d5f9
1 parent be37249
commit 600d5f9
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 15 deletions.
diff --git a/metricflow/engine/metricflow_engine.py b/metricflow/engine/metricflow_engine.py
@@ -409,11 +409,6 @@ def __init__(
             column_association_resolver=self._column_association_resolver,
             semantic_manifest_lookup=self._semantic_manifest_lookup,
         )
-        self._to_execution_plan_converter = DataflowToExecutionPlanConverter(
-            sql_plan_converter=self._to_sql_query_plan_converter,
-            sql_plan_renderer=self._sql_client.sql_query_plan_renderer,
-            sql_client=sql_client,
-        )
         self._executor = SequentialPlanExecutor()
 
         self._query_parser = query_parser or MetricFlowQueryParser(
@@ -539,7 +534,13 @@ def _create_execution_plan(self, mf_query_request: MetricFlowQueryRequest) -> Me
             )
 
         logger.info(LazyFormat("Building execution plan"))
-        convert_to_execution_plan_result = self._to_execution_plan_converter.convert_to_execution_plan(dataflow_plan)
+        _to_execution_plan_converter = DataflowToExecutionPlanConverter(
+            sql_plan_converter=self._to_sql_query_plan_converter,
+            sql_plan_renderer=self._sql_client.sql_query_plan_renderer,
+            sql_client=self._sql_client,
+            sql_optimization_level=mf_query_request.sql_optimization_level,
+        )
+        convert_to_execution_plan_result = _to_execution_plan_converter.convert_to_execution_plan(dataflow_plan)
         return MetricFlowExplainResult(
             query_spec=query_spec,
             dataflow_plan=dataflow_plan,

diff --git a/metricflow/execution/dataflow_to_execution.py b/metricflow/execution/dataflow_to_execution.py
@@ -40,6 +40,7 @@
 from metricflow.plan_conversion.convert_to_sql_plan import ConvertToSqlPlanResult
 from metricflow.plan_conversion.dataflow_to_sql import DataflowToSqlQueryPlanConverter
 from metricflow.protocols.sql_client import SqlClient
+from metricflow.sql.optimizer.optimization_levels import SqlQueryOptimizationLevel
 from metricflow.sql.render.sql_plan_renderer import SqlPlanRenderResult, SqlQueryPlanRenderer
 
 logger = logging.getLogger(__name__)
@@ -53,22 +54,26 @@ def __init__(
         sql_plan_converter: DataflowToSqlQueryPlanConverter,
         sql_plan_renderer: SqlQueryPlanRenderer,
         sql_client: SqlClient,
+        sql_optimization_level: SqlQueryOptimizationLevel,
     ) -> None:
         """Constructor.
 
         Args:
             sql_plan_converter: Converts a dataflow plan node to a SQL query plan
             sql_plan_renderer: Converts a SQL query plan to SQL text
             sql_client: The client to use for running queries.
+            sql_optimization_level: The optimization level to use for generating the SQL.
         """
         self._sql_plan_converter = sql_plan_converter
         self._sql_plan_renderer = sql_plan_renderer
         self._sql_client = sql_client
+        self._optimization_level = sql_optimization_level
 
     def _convert_to_sql_plan(self, node: DataflowPlanNode) -> ConvertToSqlPlanResult:
         logger.debug(LazyFormat(lambda: f"Generating SQL query plan from {node.node_id}"))
         result = self._sql_plan_converter.convert_to_sql_query_plan(
             sql_engine_type=self._sql_client.sql_engine_type,
+            optimization_level=self._optimization_level,
             dataflow_plan_node=node,
         )
         logger.debug(LazyFormat(lambda: f"Generated SQL query plan is:\n{result.sql_plan.structure_text()}"))

diff --git a/tests_metricflow/plan_conversion/test_dataflow_to_execution.py b/tests_metricflow/plan_conversion/test_dataflow_to_execution.py
@@ -15,6 +15,7 @@
 from metricflow.execution.dataflow_to_execution import DataflowToExecutionPlanConverter
 from metricflow.plan_conversion.dataflow_to_sql import DataflowToSqlQueryPlanConverter
 from metricflow.protocols.sql_client import SqlClient
+from metricflow.sql.optimizer.optimization_levels import SqlQueryOptimizationLevel
 from metricflow.sql.render.sql_plan_renderer import DefaultSqlQueryPlanRenderer
 from tests_metricflow.snapshot_utils import assert_execution_plan_text_equal
 
@@ -30,6 +31,7 @@ def make_execution_plan_converter(  # noqa: D103
         ),
         sql_plan_renderer=DefaultSqlQueryPlanRenderer(),
         sql_client=sql_client,
+        sql_optimization_level=SqlQueryOptimizationLevel.O4,
     )
 
 
@@ -172,17 +174,10 @@ def test_multihop_joined_plan(
         )
     )
 
-    to_execution_plan_converter = DataflowToExecutionPlanConverter(
-        sql_plan_converter=DataflowToSqlQueryPlanConverter(
-            column_association_resolver=DunderColumnAssociationResolver(
-                partitioned_multi_hop_join_semantic_manifest_lookup
-            ),
-            semantic_manifest_lookup=partitioned_multi_hop_join_semantic_manifest_lookup,
-        ),
-        sql_plan_renderer=DefaultSqlQueryPlanRenderer(),
+    to_execution_plan_converter = make_execution_plan_converter(
+        semantic_manifest_lookup=partitioned_multi_hop_join_semantic_manifest_lookup,
         sql_client=sql_client,
     )
-
     execution_plan = to_execution_plan_converter.convert_to_execution_plan(dataflow_plan).execution_plan
 
     assert_execution_plan_text_equal(