Add ExpandedTimeGranularity struct to linkable elements (#1385)

With the impending introduction of custom granularities,
which are named granularities accessed via a lookup keyed
by a date/time value with a specified base grain, our reliance
on raw TimeGranularity enumeration elements will no longer
be sufficient for managing dataflow plans or queries.

In order to be able to reference custom granularities in
our TimeDimensionSpec and associated ElementPathKey and
LinkableElement subclasses we need to make space for both
the granularity name and a base grain value.

This change adds an ExpandedTimeGranularity struct to
encapsulate both of these pieces of information in a single
object value, and switches the LinkableDimension time granularity
property to use it in place of the bare TimeGranularity enum value.

This is just the first step. Later changes will migrate most
internal TimeGranularity-typed properties to the new ExpandedTimeGranularity
construct.

Once we have a proof of concept for this working we will decide
whether this construct belongs in dbt-semantic-interfaces or if we can
keep this internal to MetricFlow. The main reason to expand this
is for more natural support for time windows and offsets, but
we may be able to simply allow strings in the protocol spec and
convert them to the struct types inside of MetricFlow.

metricflow-semantics/metricflow_semantics/model/semantics/linkable_element.py

@@ -19,12 +19,12 @@
     SemanticModelReference,
 )
 from dbt_semantic_interfaces.type_enums.date_part import DatePart
-from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
 from typing_extensions import override
 
 from metricflow_semantics.assert_one_arg import assert_exactly_one_arg_set
 from metricflow_semantics.model.linkable_element_property import LinkableElementProperty
 from metricflow_semantics.model.semantic_model_derivation import SemanticModelDerivation
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 from metricflow_semantics.workarounds.reference import sorted_semantic_model_references
 
 logger = logging.getLogger(__name__)
@@ -64,7 +64,7 @@ class ElementPathKey:
     element_name: str
     element_type: LinkableElementType
     entity_links: Tuple[EntityReference, ...]
-    time_granularity: Optional[TimeGranularity] = None
+    time_granularity: Optional[ExpandedTimeGranularity] = None
     date_part: Optional[DatePart] = None
     metric_subquery_entity_links: Tuple[EntityReference, ...] = ()
 
@@ -148,7 +148,7 @@ class LinkableDimension(LinkableElement, SerializableDataclass):
     dimension_type: DimensionType
     entity_links: Tuple[EntityReference, ...]
     join_path: SemanticModelJoinPath
-    time_granularity: Optional[TimeGranularity]
+    time_granularity: Optional[ExpandedTimeGranularity]
     date_part: Optional[DatePart]
 
     @staticmethod
@@ -159,7 +159,7 @@ def create(  # noqa: D102
         dimension_type: DimensionType,
         entity_links: Tuple[EntityReference, ...],
         join_path: SemanticModelJoinPath,
-        time_granularity: Optional[TimeGranularity],
+        time_granularity: Optional[ExpandedTimeGranularity],
         date_part: Optional[DatePart],
     ) -> LinkableDimension:
         return LinkableDimension(

metricflow-semantics/metricflow_semantics/model/semantics/linkable_element_set.py

@@ -355,10 +355,12 @@ 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,
-                time_granularity=path_key.time_granularity,
+                time_granularity=path_key.time_granularity.base_granularity,
                 date_part=path_key.date_part,
             )
         elif path_key.element_type is LinkableElementType.ENTITY:

metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py

@@ -40,6 +40,7 @@
 from metricflow_semantics.model.semantics.linkable_element_set import LinkableElementSet
 from metricflow_semantics.model.semantics.semantic_model_join_evaluator import SemanticModelJoinEvaluator
 from metricflow_semantics.specs.time_dimension_spec import DEFAULT_TIME_GRANULARITY
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 from metricflow_semantics.time.time_spine_source import TimeSpineSource
 
 if TYPE_CHECKING:
@@ -76,7 +77,7 @@ def _generate_linkable_time_dimensions(
                 dimension_type=DimensionType.TIME,
                 entity_links=entity_links,
                 join_path=join_path,
-                time_granularity=time_granularity,
+                time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
                 date_part=None,
                 properties=tuple(sorted(properties)),
             )
@@ -92,7 +93,7 @@ def _generate_linkable_time_dimensions(
                         dimension_type=DimensionType.TIME,
                         entity_links=entity_links,
                         join_path=join_path,
-                        time_granularity=time_granularity,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
                         date_part=date_part,
                         properties=frozenset(properties),
                     )
@@ -496,7 +497,7 @@ def _get_metric_time_elements(self, measure_reference: Optional[MeasureReference
                     element_name=MetricFlowReservedKeywords.METRIC_TIME.value,
                     element_type=LinkableElementType.TIME_DIMENSION,
                     entity_links=(),
-                    time_granularity=time_granularity,
+                    time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
                     date_part=date_part,
                 )
                 path_key_to_linkable_dimensions[path_key].append(
@@ -524,7 +525,7 @@ def _get_metric_time_elements(self, measure_reference: Optional[MeasureReference
                                 }
                             )
                         ),
-                        time_granularity=time_granularity,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(time_granularity),
                         date_part=date_part,
                     )
                 )

metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py

@@ -14,6 +14,7 @@
 from metricflow_semantics.naming.linkable_spec_name import StructuredLinkableSpecName
 from metricflow_semantics.specs.dimension_spec import DimensionSpec
 from metricflow_semantics.specs.instance_spec import InstanceSpecVisitor
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 from metricflow_semantics.visitor import VisitorOutputT
 
 
@@ -141,7 +142,7 @@ def element_path_key(self) -> ElementPathKey:
             element_name=self.element_name,
             element_type=LinkableElementType.TIME_DIMENSION,
             entity_links=self.entity_links,
-            time_granularity=self.time_granularity,
+            time_granularity=ExpandedTimeGranularity.from_time_granularity(self.time_granularity),
             date_part=self.date_part,
         )
 

metricflow-semantics/metricflow_semantics/time/granularity.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from dbt_semantic_interfaces.dataclass_serialization import SerializableDataclass
+from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
+
+
+@dataclass(frozen=True)
+class ExpandedTimeGranularity(SerializableDataclass):
+    """Dataclass container for custom granularity extensions to the base TimeGranularity enumeration.
+
+    This includes the granularity name, which is either the custom granularity or the TimeGranularity string value,
+    and an associated base time granularity value which we use as a pointer to the base grain used to look up the
+    time spine. This will allow for some level of comparison between custom granularities.
+
+    Note: this assumes that any base TimeGranularity value will derive the name from the TimeGranularity. It might be
+    worth adding validation to ensure that is always the case, meaning that no `name` value can be a value in the
+    TimeGranularity enumeration.
+    """
+
+    name: str
+    base_granularity: TimeGranularity
+
+    @property
+    def is_custom_granularity(self) -> bool:  # noqa: D102
+        return self.base_granularity.value != self.name
+
+    @classmethod
+    def from_time_granularity(cls, granularity: TimeGranularity) -> ExpandedTimeGranularity:
+        """Factory method for creating an ExpandedTimeGranularity from a standard TimeGranularity enumeration value."""
+        return ExpandedTimeGranularity(name=granularity.value, base_granularity=granularity)

metricflow-semantics/tests_metricflow_semantics/model/semantics/test_linkable_element_set.py

@@ -40,6 +40,7 @@
     ParameterSetField,
 )
 from metricflow_semantics.specs.time_dimension_spec import TimeDimensionSpec
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 from more_itertools import bucket
 
 AMBIGUOUS_NAME = "ambiguous"
@@ -105,7 +106,7 @@
     defined_in_semantic_model=_base_semantic_model,
     join_path=SemanticModelJoinPath(left_semantic_model_reference=_measure_semantic_model),
     properties=frozenset([LinkableElementProperty.LOCAL_LINKED]),
-    time_granularity=TimeGranularity.DAY,
+    time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
     date_part=None,
 )
 # Resolves to the same local linked name name as _ambiguous_entity
@@ -598,7 +599,7 @@ def linkable_set() -> LinkableElementSet:  # noqa: D103
                 element_name="time_dimension_element",
                 entity_links=(entity_1,),
                 element_type=LinkableElementType.TIME_DIMENSION,
-                time_granularity=TimeGranularity.DAY,
+                time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
             ): (
                 LinkableDimension.create(
                     defined_in_semantic_model=SemanticModelReference("time_dimension_source"),
@@ -615,7 +616,7 @@ def linkable_set() -> LinkableElementSet:  # noqa: D103
                         ),
                     ),
                     properties=frozenset(),
-                    time_granularity=TimeGranularity.DAY,
+                    time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
                     date_part=None,
                 ),
             ),

metricflow-semantics/tests_metricflow_semantics/model/test_where_filter_spec.py

@@ -56,6 +56,7 @@
 from metricflow_semantics.specs.time_dimension_spec import TimeDimensionSpec
 from metricflow_semantics.specs.where_filter.where_filter_spec import WhereFilterSpec
 from metricflow_semantics.specs.where_filter.where_filter_transform import WhereSpecFactory
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 from tests_metricflow_semantics.specs.conftest import EXAMPLE_FILTER_LOCATION
 
@@ -167,7 +168,7 @@ def test_dimension_in_filter_with_grain(  # noqa: D103
                         element_name="created_at",
                         element_type=LinkableElementType.TIME_DIMENSION,
                         entity_links=(EntityReference("listing"),),
-                        time_granularity=TimeGranularity.WEEK,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.WEEK),
                         date_part=None,
                     ): (
                         LinkableDimension.create(
@@ -179,7 +180,7 @@ def test_dimension_in_filter_with_grain(  # noqa: D103
                                 left_semantic_model_reference=SemanticModelReference("bookings_source"),
                             ),
                             properties=frozenset(),
-                            time_granularity=TimeGranularity.WEEK,
+                            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.WEEK),
                             date_part=None,
                         ),
                     )
@@ -231,19 +232,19 @@ def test_time_dimension_in_filter(  # noqa: D103
                         element_name="created_at",
                         element_type=LinkableElementType.TIME_DIMENSION,
                         entity_links=(EntityReference("listing"),),
-                        time_granularity=TimeGranularity.MONTH,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.MONTH),
                         date_part=None,
                     ): (
                         LinkableDimension.create(
                             defined_in_semantic_model=SemanticModelReference("listings_source"),
-                            dimension_type=DimensionType.CATEGORICAL,
+                            dimension_type=DimensionType.TIME,
                             element_name="created_at",
                             entity_links=(EntityReference("listing"),),
                             join_path=SemanticModelJoinPath(
                                 left_semantic_model_reference=SemanticModelReference("bookings_source"),
                             ),
                             properties=frozenset(),
-                            time_granularity=TimeGranularity.MONTH,
+                            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.MONTH),
                             date_part=None,
                         ),
                     )
@@ -295,19 +296,19 @@ def test_time_dimension_with_grain_in_name(  # noqa: D103
                         element_name="created_at",
                         element_type=LinkableElementType.TIME_DIMENSION,
                         entity_links=(EntityReference("listing"),),
-                        time_granularity=TimeGranularity.MONTH,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.MONTH),
                         date_part=None,
                     ): (
                         LinkableDimension.create(
                             defined_in_semantic_model=SemanticModelReference("listings_source"),
-                            dimension_type=DimensionType.CATEGORICAL,
+                            dimension_type=DimensionType.TIME,
                             element_name="created_at",
                             entity_links=(EntityReference("listing"),),
                             join_path=SemanticModelJoinPath(
                                 left_semantic_model_reference=SemanticModelReference("bookings_source"),
                             ),
                             properties=frozenset(),
-                            time_granularity=TimeGranularity.MONTH,
+                            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.MONTH),
                             date_part=None,
                         ),
                     )
@@ -360,7 +361,7 @@ def test_date_part_in_filter(  # noqa: D103
                         element_name="metric_time",
                         element_type=LinkableElementType.TIME_DIMENSION,
                         entity_links=(),
-                        time_granularity=TimeGranularity.DAY,
+                        time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
                         date_part=DatePart.YEAR,
                     ): (
                         LinkableDimension.create(
@@ -372,7 +373,7 @@ def test_date_part_in_filter(  # noqa: D103
                                 left_semantic_model_reference=SemanticModelReference("bookings_source"),
                             ),
                             properties=frozenset(),
-                            time_granularity=TimeGranularity.DAY,
+                            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
                             date_part=DatePart.YEAR,
                         ),
                     )
@@ -428,7 +429,7 @@ def resolved_spec_lookup() -> FilterSpecResolutionLookUp:
                             element_name="metric_time",
                             element_type=LinkableElementType.TIME_DIMENSION,
                             entity_links=(),
-                            time_granularity=TimeGranularity.WEEK,
+                            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.WEEK),
                             date_part=DatePart.YEAR,
                         ): (
                             LinkableDimension.create(
@@ -440,7 +441,7 @@ def resolved_spec_lookup() -> FilterSpecResolutionLookUp:
                                     left_semantic_model_reference=SemanticModelReference("bookings_source"),
                                 ),
                                 properties=frozenset(),
-                                time_granularity=TimeGranularity.WEEK,
+                                time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.WEEK),
                                 date_part=DatePart.YEAR,
                             ),
                         )
@@ -550,8 +551,8 @@ def test_entity_in_filter(  # noqa: D103
                         element_name="user",
                         element_type=LinkableElementType.ENTITY,
                         entity_links=(EntityReference("listing"),),
-                        time_granularity=TimeGranularity.DAY,
-                        date_part=DatePart.YEAR,
+                        time_granularity=None,
+                        date_part=None,
                     ): (
                         LinkableEntity.create(
                             defined_in_semantic_model=SemanticModelReference("bookings"),
@@ -659,9 +660,9 @@ def get_spec(dimension: str) -> WhereFilterSpec:
                             path_key_to_linkable_dimensions={
                                 ElementPathKey(
                                     element_name=METRIC_TIME_ELEMENT_NAME,
-                                    element_type=LinkableElementType.DIMENSION,
+                                    element_type=LinkableElementType.TIME_DIMENSION,
                                     entity_links=(EntityReference("listing"),),
-                                    time_granularity=TimeGranularity.DAY,
+                                    time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.DAY),
                                     date_part=DatePart.YEAR,
                                 ): (
                                     LinkableDimension.create(
@@ -673,7 +674,9 @@ def get_spec(dimension: str) -> WhereFilterSpec:
                                             left_semantic_model_reference=SemanticModelReference("bookings_source"),
                                         ),
                                         properties=frozenset(),
-                                        time_granularity=TimeGranularity.WEEK,
+                                        time_granularity=ExpandedTimeGranularity.from_time_granularity(
+                                            TimeGranularity.WEEK
+                                        ),
                                         date_part=DatePart.YEAR,
                                     ),
                                 )