From cb8a4cc4429bde16f0da87e0fba16a66701c5f93 Mon Sep 17 00:00:00 2001 From: serramatutu Date: Thu, 12 Dec 2024 18:28:31 +0100 Subject: [PATCH] Add metric alias step in dataflow plan builder This commit makes the `DataflowPlanBuilder` add an extra `AliasSpecsNode` at the end of the plan in case any of the input metric specs has an alias. --- .../dataflow/builder/dataflow_plan_builder.py | 18 ++++++- .../builder/test_dataflow_plan_builder.py | 32 ++++++++++++ .../test_metric_with_alias_plan__dfp_0.xml | 52 +++++++++++++++++++ 3 files changed, 100 insertions(+), 2 deletions(-) create mode 100644 tests_metricflow/snapshots/test_dataflow_plan_builder.py/DataflowPlan/test_metric_with_alias_plan__dfp_0.xml diff --git a/metricflow/dataflow/builder/dataflow_plan_builder.py b/metricflow/dataflow/builder/dataflow_plan_builder.py index 348ba5b4e8..0d46126053 100644 --- a/metricflow/dataflow/builder/dataflow_plan_builder.py +++ b/metricflow/dataflow/builder/dataflow_plan_builder.py @@ -161,7 +161,6 @@ def _build_query_output_node( len(metric_spec.filter_spec_set.all_filter_specs) > 0 or metric_spec.offset_to_grain is not None or metric_spec.offset_window is not None - or metric_spec.alias is not None ): raise ValueError( f"The metric specs in the query spec should not contain any metric modifiers. Got: {metric_spec}" @@ -213,6 +212,7 @@ def _build_plan( sink_node = DataflowPlanBuilder.build_sink_node( parent_node=metrics_output_node, + metric_specs=query_spec.metric_specs, order_by_specs=query_spec.order_by_specs, output_sql_table=output_sql_table, limit=query_spec.limit, @@ -866,7 +866,10 @@ def _build_plan_for_distinct_values( output_node = MinMaxNode.create(parent_node=output_node) sink_node = self.build_sink_node( - parent_node=output_node, order_by_specs=query_spec.order_by_specs, limit=query_spec.limit + parent_node=output_node, + metric_specs=query_spec.metric_specs, + order_by_specs=query_spec.order_by_specs, + limit=query_spec.limit, ) plan = DataflowPlan(sink_nodes=[sink_node]) @@ -875,6 +878,7 @@ def _build_plan_for_distinct_values( @staticmethod def build_sink_node( parent_node: DataflowPlanNode, + metric_specs: Sequence[MetricSpec], order_by_specs: Sequence[OrderBySpec], output_sql_table: Optional[SqlTable] = None, limit: Optional[int] = None, @@ -893,6 +897,16 @@ def build_sink_node( parent_node=pre_result_node or parent_node, include_specs=output_selection_specs ) + alias_specs = tuple( + SpecToAlias(MetricSpec(metric.element_name), MetricSpec(metric.alias)) + for metric in metric_specs + if metric.alias is not None + ) + if len(alias_specs) > 0: + pre_result_node = AliasSpecsNode.create( + parent_node=pre_result_node or parent_node, change_specs=alias_specs + ) + write_result_node: DataflowPlanNode if not output_sql_table: write_result_node = WriteToResultDataTableNode.create(parent_node=pre_result_node or parent_node) diff --git a/tests_metricflow/dataflow/builder/test_dataflow_plan_builder.py b/tests_metricflow/dataflow/builder/test_dataflow_plan_builder.py index 3a965e1c6f..3f2b006b01 100644 --- a/tests_metricflow/dataflow/builder/test_dataflow_plan_builder.py +++ b/tests_metricflow/dataflow/builder/test_dataflow_plan_builder.py @@ -1403,3 +1403,35 @@ def test_derived_cumulative_metric_with_non_default_grain( mf_test_configuration=mf_test_configuration, dag_graph=dataflow_plan, ) + + +def test_metric_with_alias_plan( + request: FixtureRequest, + mf_test_configuration: MetricFlowTestConfiguration, + dataflow_plan_builder: DataflowPlanBuilder, +) -> None: + """Tests a simple plan getting a metric and a local dimension.""" + dataflow_plan = dataflow_plan_builder.build_plan( + MetricFlowQuerySpec( + metric_specs=(MetricSpec(element_name="bookings", alias="bookings_alias"),), + dimension_specs=( + DimensionSpec( + element_name="is_instant", + entity_links=(EntityReference("booking"),), + ), + ), + ) + ) + + assert_plan_snapshot_text_equal( + request=request, + mf_test_configuration=mf_test_configuration, + plan=dataflow_plan, + plan_snapshot_text=dataflow_plan.structure_text(), + ) + + display_graph_if_requested( + request=request, + mf_test_configuration=mf_test_configuration, + dag_graph=dataflow_plan, + ) diff --git a/tests_metricflow/snapshots/test_dataflow_plan_builder.py/DataflowPlan/test_metric_with_alias_plan__dfp_0.xml b/tests_metricflow/snapshots/test_dataflow_plan_builder.py/DataflowPlan/test_metric_with_alias_plan__dfp_0.xml new file mode 100644 index 0000000000..bea2564cd4 --- /dev/null +++ b/tests_metricflow/snapshots/test_dataflow_plan_builder.py/DataflowPlan/test_metric_with_alias_plan__dfp_0.xml @@ -0,0 +1,52 @@ +test_name: test_metric_with_alias_plan +test_filename: test_dataflow_plan_builder.py +docstring: + Tests a simple plan getting a metric and a local dimension. +--- + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +