Switch query parameter classes onto ExpandedTimeGranularity

metricflow-semantics/metricflow_semantics/naming/dunder_scheme.py

@@ -18,6 +18,7 @@
     ParameterSetField,
 )
 from metricflow_semantics.specs.spec_set import InstanceSpecSet, InstanceSpecSetTransform, group_spec_by_type
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 
 class DunderNamingScheme(QueryItemNamingScheme):
@@ -81,7 +82,8 @@ def spec_pattern(self, input_str: str, semantic_manifest_lookup: SemanticManifes
         # At this point, len(input_str_parts) >= 2
         for granularity in TimeGranularity:
             if input_str_parts[-1] == granularity.value:
-                time_grain = granularity
+                # TODO: [custom granularity] add support for custom granularity names here
+                time_grain = ExpandedTimeGranularity.from_time_granularity(granularity)
 
         # Has a time grain specified.
         if time_grain is not None:

metricflow-semantics/metricflow_semantics/naming/object_builder_scheme.py

@@ -23,6 +23,7 @@
 )
 from metricflow_semantics.specs.patterns.spec_pattern import SpecPattern
 from metricflow_semantics.specs.patterns.typed_patterns import DimensionPattern, TimeDimensionPattern
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 logger = logging.getLogger(__name__)
 
@@ -80,14 +81,18 @@ def spec_pattern(self, input_str: str, semantic_manifest_lookup: SemanticManifes
                 ParameterSetField.DATE_PART,
             ]
 
+            time_granularity: Optional[ExpandedTimeGranularity] = None
             if time_dimension_call_parameter_set.time_granularity is not None:
                 fields_to_compare.append(ParameterSetField.TIME_GRANULARITY)
+                time_granularity = ExpandedTimeGranularity.from_time_granularity(
+                    time_dimension_call_parameter_set.time_granularity
+                )
 
             return TimeDimensionPattern(
                 EntityLinkPatternParameterSet.from_parameters(
                     element_name=time_dimension_call_parameter_set.time_dimension_reference.element_name,
                     entity_links=time_dimension_call_parameter_set.entity_path,
-                    time_granularity=time_dimension_call_parameter_set.time_granularity,
+                    time_granularity=time_granularity,
                     date_part=time_dimension_call_parameter_set.date_part,
                     fields_to_compare=tuple(fields_to_compare),
                 )

metricflow-semantics/metricflow_semantics/protocols/query_parameter.py

@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING, Optional, Protocol, Union, runtime_checkable
 
-from dbt_semantic_interfaces.type_enums import TimeGranularity
 from dbt_semantic_interfaces.type_enums.date_part import DatePart
 
 if TYPE_CHECKING:
@@ -52,8 +51,12 @@ def name(self) -> str:
         raise NotImplementedError
 
     @property
-    def grain(self) -> Optional[TimeGranularity]:
-        """The time granularity."""
+    def grain(self) -> Optional[str]:
+        """The name of the time granularity.
+
+        This may be the name of a custom granularity or the string value of an entry in the standard
+        TimeGranularity enum.
+        """
         raise NotImplementedError
 
     @property

metricflow-semantics/metricflow_semantics/specs/patterns/entity_link_pattern.py

@@ -6,14 +6,14 @@
 from typing import Any, List, Optional, Sequence, Tuple
 
 from dbt_semantic_interfaces.references import EntityReference
-from dbt_semantic_interfaces.type_enums import TimeGranularity
 from dbt_semantic_interfaces.type_enums.date_part import DatePart
 from more_itertools import is_sorted
 from typing_extensions import override
 
 from metricflow_semantics.specs.instance_spec import InstanceSpec, LinkableInstanceSpec
 from metricflow_semantics.specs.patterns.spec_pattern import SpecPattern
 from metricflow_semantics.specs.spec_set import group_specs_by_type
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +50,7 @@ class EntityLinkPatternParameterSet:
     # The entities used for joining semantic models.
     entity_links: Optional[Tuple[EntityReference, ...]] = None
     # Properties of time dimensions to match.
-    time_granularity: Optional[TimeGranularity] = None
+    time_granularity: Optional[ExpandedTimeGranularity] = None
     date_part: Optional[DatePart] = None
     metric_subquery_entity_links: Optional[Tuple[EntityReference, ...]] = None
 
@@ -59,7 +59,7 @@ def from_parameters(  # noqa: D102
         fields_to_compare: Sequence[ParameterSetField],
         element_name: Optional[str] = None,
         entity_links: Optional[Sequence[EntityReference]] = None,
-        time_granularity: Optional[TimeGranularity] = None,
+        time_granularity: Optional[ExpandedTimeGranularity] = None,
         date_part: Optional[DatePart] = None,
         metric_subquery_entity_links: Optional[Tuple[EntityReference, ...]] = None,
     ) -> EntityLinkPatternParameterSet:
@@ -112,22 +112,6 @@ def _match_entity_links(self, candidate_specs: Sequence[LinkableInstanceSpec]) -
         shortest_entity_link_length = min(len(matching_spec.entity_links) for matching_spec in matching_specs)
         return tuple(spec for spec in matching_specs if len(spec.entity_links) == shortest_entity_link_length)
 
-    def _match_time_granularities(
-        self, candidate_specs: Sequence[LinkableInstanceSpec]
-    ) -> Sequence[LinkableInstanceSpec]:
-        """Do a partial match on time granularities.
-
-        TODO: [custom granularity] Support custom granularities properly. This requires us to allow these pattern classes
-        to take in ExpandedTimeGranularity types, which should be viable. Once that is done, this method can be removed.
-        """
-        matching_specs: Sequence[LinkableInstanceSpec] = tuple(
-            candidate_spec
-            for candidate_spec in group_specs_by_type(candidate_specs).time_dimension_specs
-            if candidate_spec.time_granularity.base_granularity == self.parameter_set.time_granularity
-        )
-
-        return matching_specs
-
     @override
     def match(self, candidate_specs: Sequence[InstanceSpec]) -> Sequence[LinkableInstanceSpec]:
         filtered_candidate_specs = group_specs_by_type(candidate_specs).linkable_specs
@@ -136,13 +120,10 @@ def match(self, candidate_specs: Sequence[InstanceSpec]) -> Sequence[LinkableIns
         # Entity links could be a partial match, so it's handled separately.
         if ParameterSetField.ENTITY_LINKS in self.parameter_set.fields_to_compare:
             filtered_candidate_specs = self._match_entity_links(filtered_candidate_specs)
-        # Time granularities are special, so they are also handled separately.
-        if ParameterSetField.TIME_GRANULARITY in self.parameter_set.fields_to_compare:
-            filtered_candidate_specs = self._match_time_granularities(filtered_candidate_specs)
 
         other_keys_to_check = set(
             field_to_compare.value for field_to_compare in self.parameter_set.fields_to_compare
-        ).difference({ParameterSetField.ENTITY_LINKS.value, ParameterSetField.TIME_GRANULARITY.value})
+        ).difference({ParameterSetField.ENTITY_LINKS.value})
 
         matching_specs: List[LinkableInstanceSpec] = []
         parameter_set_values = tuple(getattr(self.parameter_set, key_to_check) for key_to_check in other_keys_to_check)

metricflow-semantics/metricflow_semantics/specs/patterns/typed_patterns.py

@@ -20,6 +20,7 @@
     ParameterSetField,
 )
 from metricflow_semantics.specs.spec_set import group_specs_by_type
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 
 @dataclass(frozen=True)
@@ -78,15 +79,19 @@ def from_call_parameter_set(
             ParameterSetField.DATE_PART,
         ]
 
+        time_granularity = None
         if time_dimension_call_parameter_set.time_granularity is not None:
             fields_to_compare.append(ParameterSetField.TIME_GRANULARITY)
+            time_granularity = ExpandedTimeGranularity.from_time_granularity(
+                time_dimension_call_parameter_set.time_granularity
+            )
 
         return TimeDimensionPattern(
             parameter_set=EntityLinkPatternParameterSet.from_parameters(
                 fields_to_compare=tuple(fields_to_compare),
                 element_name=time_dimension_call_parameter_set.time_dimension_reference.element_name,
                 entity_links=time_dimension_call_parameter_set.entity_path,
-                time_granularity=time_dimension_call_parameter_set.time_granularity,
+                time_granularity=time_granularity,
                 date_part=time_dimension_call_parameter_set.date_part,
             )
         )

metricflow-semantics/metricflow_semantics/specs/query_param_implementations.py

@@ -31,6 +31,7 @@
     EntityLinkPatternParameterSet,
     ParameterSetField,
 )
+from metricflow_semantics.time.granularity import ExpandedTimeGranularity
 
 
 @dataclass(frozen=True)
@@ -41,7 +42,7 @@ def _implements_protocol(self) -> TimeDimensionQueryParameter:
         return self
 
     name: str
-    grain: Optional[TimeGranularity] = None
+    grain: Optional[str] = None
     date_part: Optional[DatePart] = None
 
     def query_resolver_input(  # noqa: D102
@@ -54,9 +55,17 @@ def query_resolver_input(  # noqa: D102
             ParameterSetField.ENTITY_LINKS,
             ParameterSetField.DATE_PART,
         ]
+        time_granularity = None
         if self.grain is not None:
             fields_to_compare.append(ParameterSetField.TIME_GRANULARITY)
-
+            # TODO: [custom granularity] support custom granularity lookups
+            assert ExpandedTimeGranularity.is_standard_granularity_name(self.grain), (
+                f"We got a non-standard granularity name, `{self.grain}`, but we have not yet "
+                "implemented support for custom granularities!"
+            )
+            time_granularity = ExpandedTimeGranularity.from_time_granularity(TimeGranularity(self.grain))
+
+        # TODO: assert that the name does not include a time granularity marker
         name_structure = StructuredLinkableSpecName.from_name(self.name.lower())
 
         return ResolverInputForGroupByItem(
@@ -67,7 +76,7 @@ def query_resolver_input(  # noqa: D102
                     fields_to_compare=tuple(fields_to_compare),
                     element_name=name_structure.element_name,
                     entity_links=tuple(EntityReference(link_name) for link_name in name_structure.entity_link_names),
-                    time_granularity=self.grain,
+                    time_granularity=time_granularity,
                     date_part=self.date_part,
                 )
             ),

metricflow-semantics/tests_metricflow_semantics/query/test_query_parser.py

@@ -560,7 +560,7 @@ def test_date_part_parsing(
         query_parser.parse_and_validate_query(
             metric_names=["revenue"],
             group_by=(
-                TimeDimensionParameter(name="metric_time", grain=TimeGranularity.YEAR, date_part=DatePart.MONTH),
+                TimeDimensionParameter(name="metric_time", grain=TimeGranularity.YEAR.value, date_part=DatePart.MONTH),
             ),
         )
 

metricflow-semantics/tests_metricflow_semantics/specs/patterns/test_entity_link_pattern.py

@@ -145,7 +145,7 @@ def test_time_dimension_match(specs: Sequence[LinkableInstanceSpec]) -> None:  #
         EntityLinkPatternParameterSet.from_parameters(
             element_name=METRIC_TIME_ELEMENT_NAME,
             entity_links=(),
-            time_granularity=TimeGranularity.WEEK,
+            time_granularity=ExpandedTimeGranularity.from_time_granularity(TimeGranularity.WEEK),
             date_part=None,
             fields_to_compare=(
                 ParameterSetField.ELEMENT_NAME,

tests_metricflow/integration/test_configured_cases.py

@@ -261,8 +261,6 @@ def test_case(
         if date_part or grain:
             if date_part:
                 kwargs["date_part"] = DatePart(date_part)
-            if grain:
-                kwargs["grain"] = TimeGranularity(grain)
             group_by.append(TimeDimensionParameter(**kwargs))
         else:
             group_by.append(DimensionOrEntityParameter(**kwargs))