Add dimension_lookup to ModelObjectLookup

metricflow-semantics/metricflow_semantics/experimental/dsi/dimension_lookup.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from functools import cached_property
+from typing import Iterable
+
+from dbt_semantic_interfaces.protocols import Dimension
+from dbt_semantic_interfaces.type_enums import DimensionType
+from typing_extensions import override
+
+from metricflow_semantics.mf_logging.attribute_pretty_format import AttributeMapping, AttributePrettyFormattable
+
+logger = logging.getLogger(__name__)
+
+
+class DimensionLookup(AttributePrettyFormattable):
+    """A lookup for entities within a semantic model."""
+
+    def __init__(self, dimensions: Iterable[Dimension]) -> None:  # noqa: D107
+        self._dimensions = tuple(dimensions)
+
+    @cached_property
+    def dimension_name_to_dimension(self) -> Mapping[str, Dimension]:  # noqa: D102
+        return {dimension.name: dimension for dimension in self._dimensions}
+
+    @cached_property
+    def dimension_name_to_type(self) -> Mapping[str, DimensionType]:  # noqa: D102
+        return {
+            dimension_name: dimension.type for dimension_name, dimension in self.dimension_name_to_dimension.items()
+        }
+
+    @cached_property
+    @override
+    def _attribute_mapping(self) -> AttributeMapping:
+        return dict(
+            **super()._attribute_mapping,
+            **{attribute_name: getattr(self, attribute_name) for attribute_name in ("dimension_name_to_type",)},
+        )

metricflow-semantics/metricflow_semantics/experimental/dsi/entity_lookup.py

@@ -10,7 +10,7 @@
 from dbt_semantic_interfaces.type_enums import EntityType
 from typing_extensions import override
 
-from metricflow_semantics.experimental.ordered_set import FrozenOrderedSet, MutableOrderedSet
+from metricflow_semantics.experimental.ordered_set import MutableOrderedSet, OrderedSet
 from metricflow_semantics.mf_logging.attribute_pretty_format import AttributeMapping, AttributePrettyFormattable
 
 logger = logging.getLogger(__name__)
@@ -21,17 +21,18 @@ class EntityLookup(AttributePrettyFormattable):
 
     def __init__(self, entities: Iterable[Entity]) -> None:  # noqa: D107
         self._entity_name_to_entity: Mapping[str, Entity] = {entity.name: entity for entity in entities}
+        self._entities = tuple(entities)
 
     @cached_property
     def entity_name_to_type(self) -> Mapping[str, EntityType]:  # noqa: D102
-        return {entity_name: entity.type for entity_name, entity in self._entity_name_to_entity.items()}
+        return {entity.name: entity.type for entity in self._entities}
 
     @cached_property
-    def entity_type_to_names(self) -> Mapping[EntityType, FrozenOrderedSet[str]]:  # noqa: D102
+    def entity_type_to_names(self) -> Mapping[EntityType, OrderedSet[str]]:  # noqa: D102
         entity_type_to_names: dict[EntityType, MutableOrderedSet[str]] = defaultdict(MutableOrderedSet)
-        for entity_name, entity_type in self.entity_name_to_type.items():
-            entity_type_to_names[entity_type].add(entity_name)
-        return {entity_type: entity_names.as_frozen() for entity_type, entity_names in entity_type_to_names.items()}
+        for entity in self._entities:
+            entity_type_to_names[entity.type].add(entity.name)
+        return {entity_type: entity_names for entity_type, entity_names in entity_type_to_names.items()}
 
     @cached_property
     @override

metricflow-semantics/metricflow_semantics/experimental/dsi/model_object_lookup.py

@@ -6,12 +6,13 @@
 
 from dbt_semantic_interfaces.enum_extension import assert_values_exhausted
 from dbt_semantic_interfaces.protocols import SemanticModel
-from dbt_semantic_interfaces.type_enums import DimensionType, TimeGranularity
+from dbt_semantic_interfaces.type_enums import DimensionType, EntityType, TimeGranularity
 from typing_extensions import override
 
-from metricflow_semantics.collection_helpers.syntactic_sugar import mf_first_non_none_or_raise
+from metricflow_semantics.collection_helpers.syntactic_sugar import mf_first_item, mf_first_non_none_or_raise
+from metricflow_semantics.experimental.dsi.dimension_lookup import DimensionLookup
 from metricflow_semantics.experimental.dsi.entity_lookup import EntityLookup
-from metricflow_semantics.experimental.metricflow_exception import InvalidManifestException
+from metricflow_semantics.experimental.metricflow_exception import InvalidManifestException, MetricFlowInternalError
 from metricflow_semantics.experimental.semantic_graph.model_id import SemanticModelId
 from metricflow_semantics.mf_logging.attribute_pretty_format import AttributeMapping, AttributePrettyFormattable
 from metricflow_semantics.mf_logging.lazy_formattable import LazyFormat
@@ -62,6 +63,27 @@ def time_dimension_name_to_grain(self) -> Mapping[str, TimeGranularity]:  # noqa
     def entity_lookup(self) -> EntityLookup:  # noqa: D102
         return EntityLookup(self._semantic_model.entities)
 
+    @cached_property
+    def dimension_lookup(self) -> DimensionLookup:  # noqa: D102
+        return DimensionLookup(self._semantic_model.dimensions)
+
+    @cached_property
+    def resolved_primary_entity_name(self) -> str:
+        """Return the name of the primary entity for this mode."""
+        if self._semantic_model.primary_entity is not None:
+            return self._semantic_model.primary_entity
+
+        primary_entity_set = self.entity_lookup.entity_type_to_names.get(EntityType.PRIMARY)
+        if primary_entity_set is None or len(primary_entity_set) != 1:
+            raise MetricFlowInternalError(
+                LazyFormat(
+                    "Did not find exactly 1 primary entity in the semantic model",
+                    primary_entity_set=primary_entity_set,
+                    semantic_model=self._semantic_model,
+                )
+            )
+        return mf_first_item(primary_entity_set)
+
     @cached_property
     @override
     def _attribute_mapping(self) -> AttributeMapping:

metricflow/engine/metricflow_engine.py

@@ -727,26 +727,26 @@ def list_dimensions(
         order_by: GroupByOrderByAttribute = GroupByOrderByAttribute.DUNDER_NAME,
     ) -> List[Dimension]:
         """Get full dimension object for all dimensions in the semantic manifest."""
-        dimensions: List[Dimension] = []
+        engine_model_dimensions: List[Dimension] = []
         if metric_names:
-            dimensions = self.simple_dimensions_for_metrics(metric_names=metric_names)
+            engine_model_dimensions = self.simple_dimensions_for_metrics(metric_names=metric_names)
         else:
-            semantic_model_lookup = self._semantic_manifest_lookup.semantic_model_lookup
-            for dimension_reference in semantic_model_lookup.get_dimension_references():
-                for semantic_model in semantic_model_lookup.get_semantic_models_for_dimension(dimension_reference):
-                    dimension = Dimension.from_pydantic(
-                        pydantic_dimension=SemanticModelHelper.get_dimension_from_semantic_model(
-                            semantic_model=semantic_model, dimension_reference=dimension_reference
-                        ),
-                        entity_links=(SemanticModelHelper.resolved_primary_entity(semantic_model),),
-                        semantic_model_reference=semantic_model.reference,
+            model_object_lookups = self._semantic_manifest_lookup.manifest_object_lookup.model_object_lookups
+
+            for model_object_lookup in model_object_lookups:
+                primary_entity_name = model_object_lookup.resolved_primary_entity_name
+                for pydantic_dimension in model_object_lookup.dimension_lookup.dimension_name_to_dimension.values():
+                    engine_model_dimension = Dimension.from_pydantic(
+                        pydantic_dimension=pydantic_dimension,
+                        entity_links=(EntityReference(primary_entity_name),),
+                        semantic_model_reference=model_object_lookup.semantic_model.reference,
                     )
-                    dimensions.append(dimension)
+                    engine_model_dimensions.append(engine_model_dimension)
 
-        def sort_dimensions(dimension: Dimension) -> Tuple[str, ...]:
-            if order_by == GroupByOrderByAttribute.DUNDER_NAME:
+        def _dimension_sort_key(dimension: Dimension) -> Tuple[str, ...]:
+            if order_by is GroupByOrderByAttribute.DUNDER_NAME:
                 return (dimension.dunder_name,)
-            elif order_by == GroupByOrderByAttribute.SEMANTIC_MODEL_NAME:
+            elif order_by is GroupByOrderByAttribute.SEMANTIC_MODEL_NAME:
                 return (
                     (
                         dimension.semantic_model_reference.semantic_model_name
@@ -758,7 +758,7 @@ def sort_dimensions(dimension: Dimension) -> Tuple[str, ...]:
             else:
                 assert_values_exhausted(order_by)
 
-        return sorted(set(dimensions), key=sort_dimensions)
+        return sorted(set(engine_model_dimensions), key=_dimension_sort_key)
 
     def entities_for_metrics(self, metric_names: List[str]) -> List[Entity]:  # noqa: D102
         group_by_item_set = self._semantic_manifest_lookup.metric_lookup.get_common_group_by_items(