Handle time constraints for dimension only queries

dbt-labs · Dec 6, 2023 · 2c630da · 2c630da
1 parent 2aaee50
commit 2c630da
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 18 deletions.
diff --git a/metricflow/dataflow/builder/dataflow_plan_builder.py b/metricflow/dataflow/builder/dataflow_plan_builder.py
@@ -353,32 +353,29 @@ def build_plan_for_distinct_values(self, query_spec: MetricFlowQuerySpec) -> Dat
         required_linkable_specs, _ = self.__get_required_and_extraneous_linkable_specs(
             queried_linkable_specs=query_spec.linkable_specs, where_constraint=query_spec.where_constraint
         )
-        dataflow_recipe = self._find_dataflow_recipe(linkable_spec_set=required_linkable_specs)
+        dataflow_recipe = self._find_dataflow_recipe(
+            linkable_spec_set=required_linkable_specs, time_range_constraint=query_spec.time_range_constraint
+        )
         if not dataflow_recipe:
             raise UnableToSatisfyQueryError(f"Recipe not found for linkable specs: {required_linkable_specs}")
 
-        joined_node: Optional[JoinToBaseOutputNode] = None
+        output_node = dataflow_recipe.source_node
         if dataflow_recipe.join_targets:
-            joined_node = JoinToBaseOutputNode(
-                left_node=dataflow_recipe.source_node, join_targets=dataflow_recipe.join_targets
-            )
+            output_node = JoinToBaseOutputNode(left_node=output_node, join_targets=dataflow_recipe.join_targets)
 
-        where_constraint_node: Optional[WhereConstraintNode] = None
         if query_spec.where_constraint:
-            where_constraint_node = WhereConstraintNode(
-                parent_node=joined_node or dataflow_recipe.source_node, where_constraint=query_spec.where_constraint
+            output_node = WhereConstraintNode(parent_node=output_node, where_constraint=query_spec.where_constraint)
+        if query_spec.time_range_constraint:
+            output_node = ConstrainTimeRangeNode(
+                parent_node=output_node, time_range_constraint=query_spec.time_range_constraint
             )
 
-        distinct_values_node = FilterElementsNode(
-            parent_node=where_constraint_node or joined_node or dataflow_recipe.source_node,
-            include_specs=query_spec.linkable_specs.as_spec_set,
-            distinct=True,
+        output_node = FilterElementsNode(
+            parent_node=output_node, include_specs=query_spec.linkable_specs.as_spec_set, distinct=True
         )
 
         sink_node = self.build_sink_node(
-            parent_node=distinct_values_node,
-            order_by_specs=query_spec.order_by_specs,
-            limit=query_spec.limit,
+            parent_node=output_node, order_by_specs=query_spec.order_by_specs, limit=query_spec.limit
         )
 
         plan_id = IdGeneratorRegistry.for_class(DataflowPlanBuilder).create_id(DATAFLOW_PLAN_PREFIX)

diff --git a/metricflow/test/integration/test_cases/itest_dimensions.yaml b/metricflow/test/integration/test_cases/itest_dimensions.yaml
@@ -229,7 +229,6 @@ integration_test:
   description: Query metric_time with other dimensions and filters
   model: SIMPLE_MODEL
   group_bys: ["metric_time", "metric_time__month", "user__home_state_latest", "listing__is_lux_latest"]
-  time_constraint: ["2020-01-03", "2020-01-03"]
   check_query: |
     SELECT
       {{ render_date_trunc("ts.ds", TimeGranularity.DAY) }} AS metric_time__day
@@ -248,6 +247,26 @@ integration_test:
       , {{ render_date_trunc("ts.ds", TimeGranularity.MONTH) }}
       , l.is_lux
       , u.home_state_latest
+---
+integration_test:
+  name: query_dimensions_with_time_constraint
+  description: Query multiple dimensions without metrics
+  model: SIMPLE_MODEL
+  group_bys: ["user__home_state_latest", "listing__country_latest", "metric_time__day"]
+  time_constraint: ["2020-01-03", "2020-01-03"]
+  check_query: |
+    SELECT
+      {{ render_date_trunc("t.ds", TimeGranularity.DAY) }} AS metric_time__day
+      , l.country AS listing__country_latest
+      , u.home_state_latest AS user__home_state_latest
+    FROM {{ source_schema }}.dim_listings_latest l
+    CROSS JOIN {{ source_schema }}.mf_time_spine t
+    FULL OUTER JOIN {{ source_schema }}.dim_users_latest u ON l.user_id = u.user_id
+    WHERE {{ render_date_trunc("t.ds", TimeGranularity.DAY) }} BETWEEN '2020-01-03' AND '2020-01-03'
+    GROUP BY
+      {{ render_date_trunc("t.ds", TimeGranularity.DAY) }}
+      , l.country
+      , u.home_state_latest
 
 # case where where filter not in selections, for derived metric plan that doesn't use normal one
 # TODO: add test with date part

diff --git a/metricflow/test/integration/test_configured_cases.py b/metricflow/test/integration/test_configured_cases.py
@@ -207,7 +207,6 @@ def filter_not_supported_features(
 @pytest.mark.parametrize(
     "name",
     CONFIGURED_INTEGRATION_TESTS_REPOSITORY.all_test_case_names,
-    # ["itest_dimensions.yaml/metric_time_with_other_dimensions"],
     ids=lambda name: f"name={name}",
 )
 def test_case(
@@ -302,7 +301,6 @@ def test_case(
         )
     )
 
-    # assert 0, query_result.sql
     actual = query_result.result_df
 
     expected = sql_client.query(