Add custom granularities to LinkableSpecResolver

This includes a bit of a refactor to simplify the time dimension generation code.
dbt-labs · Sep 20, 2024 · 42c6514 · 42c6514
1 parent 7c9f56c
commit 42c6514
Show file tree

Hide file tree

Showing 9 changed files with 98 additions and 56 deletions.
diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/linkable_element_set.py b/metricflow-semantics/metricflow_semantics/model/semantics/linkable_element_set.py
@@ -355,8 +355,6 @@ def _path_key_to_spec(path_key: ElementPathKey) -> LinkableInstanceSpec:
             )
         elif path_key.element_type is LinkableElementType.TIME_DIMENSION:
             assert path_key.time_granularity is not None
-            # TODO: [custom granularity] Remove block against custom granularity values
-            assert not path_key.time_granularity.is_custom_granularity, "Custom granularities are not yet supported!"
             return TimeDimensionSpec(
                 element_name=path_key.element_name,
                 entity_links=path_key.entity_links,

diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py b/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py
@@ -50,58 +50,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _generate_linkable_time_dimensions(
-    semantic_model_origin: SemanticModelReference,
-    dimension: Dimension,
-    entity_links: Tuple[EntityReference, ...],
-    join_path: SemanticModelJoinPath,
-    with_properties: FrozenSet[LinkableElementProperty],
-) -> Sequence[LinkableDimension]:
-    """Generates different versions of the given dimension, but at other valid time granularities."""
-    linkable_dimensions = []
-
-    defined_time_granularity = (
-        dimension.type_params.time_granularity if dimension.type_params else DEFAULT_TIME_GRANULARITY
-    )
-    for time_granularity in TimeGranularity:
-        if time_granularity < defined_time_granularity:
-            continue
-        properties = set(with_properties)
-        if time_granularity != defined_time_granularity:
-            properties.add(LinkableElementProperty.DERIVED_TIME_GRANULARITY)
-
-        linkable_dimensions.append(
-            LinkableDimension.create(
-                defined_in_semantic_model=semantic_model_origin,
-                element_name=dimension.reference.element_name,
-                dimension_type=DimensionType.TIME,
-                entity_links=entity_links,
-                join_path=join_path,
-                time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
-                date_part=None,
-                properties=tuple(sorted(properties)),
-            )
-        )
-
-        # Add the time dimension aggregated to a different date part.
-        for date_part in DatePart:
-            if time_granularity.to_int() <= date_part.to_int():
-                linkable_dimensions.append(
-                    LinkableDimension.create(
-                        defined_in_semantic_model=semantic_model_origin,
-                        element_name=dimension.reference.element_name,
-                        dimension_type=DimensionType.TIME,
-                        entity_links=entity_links,
-                        join_path=join_path,
-                        time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
-                        date_part=date_part,
-                        properties=frozenset(properties),
-                    )
-                )
-
-    return linkable_dimensions
-
-
 class ValidLinkableSpecResolver:
     """Figures out what linkable specs are valid for a given metric.
 
@@ -127,6 +75,7 @@ def __init__(
         self._semantic_models = sorted(self._semantic_manifest.semantic_models, key=lambda x: x.name)
         self._join_evaluator = SemanticModelJoinEvaluator(semantic_model_lookup)
         self._time_spine_sources = TimeSpineSource.build_standard_time_spine_sources(self._semantic_manifest)
+        self._custom_granularities = TimeSpineSource.build_custom_granularities(list(self._time_spine_sources.values()))
 
         assert max_entity_links >= 0
         self._max_entity_links = max_entity_links
@@ -239,6 +188,54 @@ def __init__(
 
         logger.info(f"Building valid group-by-item indexes took: {time.time() - start_time:.2f}s")
 
+    def _generate_linkable_time_dimensions(
+        self,
+        semantic_model_origin: SemanticModelReference,
+        dimension: Dimension,
+        entity_links: Tuple[EntityReference, ...],
+        join_path: SemanticModelJoinPath,
+        with_properties: FrozenSet[LinkableElementProperty],
+    ) -> Sequence[LinkableDimension]:
+        """Generates different versions of the given time dimension with all compatible time granularities & date parts."""
+        defined_time_granularity = (
+            dimension.type_params.time_granularity if dimension.type_params else DEFAULT_TIME_GRANULARITY
+        )
+
+        def create(
+            time_granularity: ExpandedTimeGranularity, date_part: Optional[DatePart] = None
+        ) -> LinkableDimension:
+            properties = set(with_properties)
+            if time_granularity.is_custom_granularity or time_granularity.base_granularity != defined_time_granularity:
+                properties.add(LinkableElementProperty.DERIVED_TIME_GRANULARITY)
+            return LinkableDimension.create(
+                defined_in_semantic_model=semantic_model_origin,
+                element_name=dimension.reference.element_name,
+                dimension_type=DimensionType.TIME,
+                entity_links=entity_links,
+                join_path=join_path,
+                time_granularity=time_granularity,
+                date_part=date_part,
+                properties=tuple(sorted(properties)),
+            )
+
+        # Build linkable dimensions for all compatible standard granularities and date parts.
+        linkable_dimensions = []
+        for time_granularity in TimeGranularity:
+            if time_granularity.to_int() < defined_time_granularity.to_int():
+                continue
+            expanded_granularity = ExpandedTimeGranularity.from_time_granularity(time_granularity)
+            linkable_dimensions.append(create(expanded_granularity))
+            for date_part in DatePart:
+                if time_granularity.to_int() <= date_part.to_int():
+                    linkable_dimensions.append(create(time_granularity=expanded_granularity, date_part=date_part))
+
+        # Build linkable dimensions for all compatible custom granularities.
+        for custom_grain in self._custom_granularities.values():
+            if custom_grain.base_granularity.to_int() >= defined_time_granularity.to_int():
+                linkable_dimensions.append(create(custom_grain))
+
+        return linkable_dimensions
+
     def _metric_requires_metric_time(self, metric: Metric) -> bool:
         """Checks if the metric can only be queried with metric_time. Also checks input metrics.
 
@@ -378,7 +375,7 @@ def _get_elements_in_semantic_model(self, semantic_model: SemanticModel) -> Link
                     )
                 elif dimension_type is DimensionType.TIME:
                     linkable_dimensions.extend(
-                        _generate_linkable_time_dimensions(
+                        self._generate_linkable_time_dimensions(
                             semantic_model_origin=semantic_model.reference,
                             dimension=dimension,
                             entity_links=(entity_link,),
@@ -741,7 +738,7 @@ def create_linkable_element_set_from_join_path(
                 )
             elif dimension_type == DimensionType.TIME:
                 linkable_dimensions.extend(
-                    _generate_linkable_time_dimensions(
+                    self._generate_linkable_time_dimensions(
                         semantic_model_origin=semantic_model.reference,
                         dimension=dimension,
                         entity_links=join_path.entity_links,

diff --git a/...shots/test_linkable_spec_resolver.py/list/test_linkable_element_set_as_spec_set__set0.txt b/...shots/test_linkable_spec_resolver.py/list/test_linkable_element_set_as_spec_set__set0.txt
@@ -80,6 +80,7 @@
   'listing__created_at__extract_year',
   'listing__created_at__extract_year',
   'listing__created_at__extract_year',
+  'listing__created_at__martian_day',
   'listing__created_at__month',
   'listing__created_at__quarter',
   'listing__created_at__week',
@@ -104,6 +105,7 @@
   'listing__ds__extract_year',
   'listing__ds__extract_year',
   'listing__ds__extract_year',
+  'listing__ds__martian_day',
   'listing__ds__month',
   'listing__ds__quarter',
   'listing__ds__week',
@@ -183,6 +185,7 @@
   'user__archived_at__extract_year',
   'user__archived_at__extract_year',
   'user__archived_at__hour',
+  'user__archived_at__martian_day',
   'user__archived_at__month',
   'user__archived_at__quarter',
   'user__archived_at__week',
@@ -223,6 +226,7 @@
   'user__bio_added_ts__extract_year',
   'user__bio_added_ts__extract_year',
   'user__bio_added_ts__hour',
+  'user__bio_added_ts__martian_day',
   'user__bio_added_ts__minute',
   'user__bio_added_ts__month',
   'user__bio_added_ts__quarter',
@@ -247,6 +251,7 @@
   'user__created_at__extract_year',
   'user__created_at__extract_year',
   'user__created_at__extract_year',
+  'user__created_at__martian_day',
   'user__created_at__month',
   'user__created_at__quarter',
   'user__created_at__week',
@@ -268,6 +273,7 @@
   'user__ds__extract_year',
   'user__ds__extract_year',
   'user__ds__extract_year',
+  'user__ds__martian_day',
   'user__ds__month',
   'user__ds__quarter',
   'user__ds__week',
@@ -288,6 +294,7 @@
   'user__ds_latest__extract_year',
   'user__ds_latest__extract_year',
   'user__ds_latest__extract_year',
+  'user__ds_latest__martian_day',
   'user__ds_latest__month',
   'user__ds_latest__quarter',
   'user__ds_latest__week',
@@ -308,6 +315,7 @@
   'user__ds_partitioned__extract_year',
   'user__ds_partitioned__extract_year',
   'user__ds_partitioned__extract_year',
+  'user__ds_partitioned__martian_day',
   'user__ds_partitioned__month',
   'user__ds_partitioned__quarter',
   'user__ds_partitioned__week',
@@ -345,6 +353,7 @@
   'user__last_login_ts__extract_year',
   'user__last_login_ts__extract_year',
   'user__last_login_ts__hour',
+  'user__last_login_ts__martian_day',
   'user__last_login_ts__minute',
   'user__last_login_ts__month',
   'user__last_login_ts__quarter',
@@ -391,6 +400,7 @@
   'user__last_profile_edit_ts__extract_year',
   'user__last_profile_edit_ts__extract_year',
   'user__last_profile_edit_ts__hour',
+  'user__last_profile_edit_ts__martian_day',
   'user__last_profile_edit_ts__millisecond',
   'user__last_profile_edit_ts__minute',
   'user__last_profile_edit_ts__month',