Skip to content

Commit

Permalink
/* PR_START p--short-term-perf 29 */ Move static methods to `Semantic…
Browse files Browse the repository at this point in the history
…ModelHelper`.
  • Loading branch information
plypaul committed Oct 30, 2024
1 parent c611b44 commit 9537a88
Show file tree
Hide file tree
Showing 5 changed files with 108 additions and 87 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
SemanticModelToMetricSubqueryJoinPath,
)
from metricflow_semantics.model.semantics.linkable_element_set import LinkableElementSet
from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper
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
Expand Down Expand Up @@ -358,7 +359,7 @@ def _get_elements_in_semantic_model(self, semantic_model: SemanticModel) -> Link
properties=entity_properties,
)
)
for entity_link in self._semantic_model_lookup.entity_links_for_local_elements(semantic_model):
for entity_link in SemanticModelHelper.entity_links_for_local_elements(semantic_model):
# Avoid creating "booking_id__booking_id"
if entity_link == entity.reference:
continue
Expand All @@ -378,7 +379,7 @@ def _get_elements_in_semantic_model(self, semantic_model: SemanticModel) -> Link
if semantic_model_is_scd:
dimension_properties = dimension_properties.union({LinkableElementProperty.SCD_HOP})

for entity_link in self._semantic_model_lookup.entity_links_for_local_elements(semantic_model):
for entity_link in SemanticModelHelper.entity_links_for_local_elements(semantic_model):
for dimension in semantic_model.dimensions:
dimension_type = dimension.type
if dimension_type is DimensionType.CATEGORICAL:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
from __future__ import annotations

from typing import Optional, Sequence

from dbt_semantic_interfaces.protocols import Dimension
from dbt_semantic_interfaces.protocols.entity import Entity
from dbt_semantic_interfaces.protocols.measure import Measure
from dbt_semantic_interfaces.protocols.semantic_model import SemanticModel
from dbt_semantic_interfaces.references import (
EntityReference,
LinkableElementReference,
MeasureReference,
)
from dbt_semantic_interfaces.type_enums import EntityType


class SemanticModelHelper:
"""Static helper methods for retrieving items from a semantic model."""

@staticmethod
def get_entity_from_semantic_model(
semantic_model: SemanticModel, entity_reference: LinkableElementReference
) -> Entity:
"""Get entity from semantic model."""
for entity in semantic_model.entities:
if entity.reference == entity_reference:
return entity

raise ValueError(
f"No entity with name ({entity_reference}) in semantic_model with name ({semantic_model.name})"
)

@staticmethod
def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
"""Return the primary entity for dimensions in the model."""
primary_entity_reference = semantic_model.primary_entity_reference

entities_with_type_primary = tuple(
entity for entity in semantic_model.entities if entity.type == EntityType.PRIMARY
)

# This should be caught by the validation, but adding a sanity check.
assert len(entities_with_type_primary) <= 1, f"Found > 1 primary entity in {semantic_model}"
if primary_entity_reference is not None:
assert len(entities_with_type_primary) == 0, (
f"The primary_entity field was set to {primary_entity_reference}, but there are non-zero entities with "
f"type {EntityType.PRIMARY} in {semantic_model}"
)
return primary_entity_reference

if len(entities_with_type_primary) > 0:
return entities_with_type_primary[0].reference

return None

@staticmethod
def entity_links_for_local_elements(semantic_model: SemanticModel) -> Sequence[EntityReference]:
"""Return the entity prefix that can be used to access dimensions defined in the semantic model."""
primary_entity_reference = semantic_model.primary_entity_reference

possible_entity_links = set()
if primary_entity_reference is not None:
possible_entity_links.add(primary_entity_reference)

for entity in semantic_model.entities:
if entity.is_linkable_entity_type:
possible_entity_links.add(entity.reference)

return sorted(possible_entity_links, key=lambda entity_reference: entity_reference.element_name)

@staticmethod
def get_measure_from_semantic_model(semantic_model: SemanticModel, measure_reference: MeasureReference) -> Measure:
"""Get measure from semantic model."""
for measure in semantic_model.measures:
if measure.reference == measure_reference:
return measure

raise ValueError(
f"No dimension with name ({measure_reference.element_name}) in semantic_model with name ({semantic_model.name})"
)

@staticmethod
def get_dimension_from_semantic_model(
semantic_model: SemanticModel, dimension_reference: LinkableElementReference
) -> Dimension:
"""Get dimension from semantic model."""
for dim in semantic_model.dimensions:
if dim.reference == dimension_reference:
return dim
raise ValueError(
f"No dimension with name ({dimension_reference}) in semantic_model with name ({semantic_model.name})"
)
Original file line number Diff line number Diff line change
Expand Up @@ -11,18 +11,18 @@
from dbt_semantic_interfaces.references import (
DimensionReference,
EntityReference,
LinkableElementReference,
MeasureReference,
SemanticModelElementReference,
SemanticModelReference,
TimeDimensionReference,
)
from dbt_semantic_interfaces.type_enums import AggregationType, DimensionType, EntityType, TimeGranularity
from dbt_semantic_interfaces.type_enums import AggregationType, DimensionType, TimeGranularity

from metricflow_semantics.errors.error_classes import InvalidSemanticModelError
from metricflow_semantics.mf_logging.lazy_formattable import LazyFormat
from metricflow_semantics.mf_logging.pretty_print import mf_pformat
from metricflow_semantics.model.semantics.element_group import ElementGrouper
from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper
from metricflow_semantics.model.spec_converters import MeasureConverter
from metricflow_semantics.naming.linkable_spec_name import StructuredLinkableSpecName
from metricflow_semantics.specs.dimension_spec import DimensionSpec
Expand Down Expand Up @@ -74,18 +74,6 @@ def get_dimension_references(self) -> Sequence[DimensionReference]:
"""Retrieve all dimension references from the collection of semantic models."""
return tuple(self._dimension_index.keys())

@staticmethod
def get_dimension_from_semantic_model(
semantic_model: SemanticModel, dimension_reference: LinkableElementReference
) -> Dimension:
"""Get dimension from semantic model."""
for dim in semantic_model.dimensions:
if dim.reference == dimension_reference:
return dim
raise ValueError(
f"No dimension with name ({dimension_reference}) in semantic_model with name ({semantic_model.name})"
)

def get_dimension(self, dimension_reference: DimensionReference) -> Dimension:
"""Retrieves a full dimension object by name."""
# If the reference passed is a TimeDimensionReference, convert to DimensionReference.
Expand All @@ -97,7 +85,7 @@ def get_dimension(self, dimension_reference: DimensionReference) -> Dimension:
f"Could not find dimension with name '{dimension_reference.element_name}' in configured semantic models"
)

return SemanticModelLookup.get_dimension_from_semantic_model(
return SemanticModelHelper.get_dimension_from_semantic_model(
# Dimension object should match across semantic models, so just use the first semantic model.
semantic_model=semantic_models[0],
dimension_reference=dimension_reference,
Expand All @@ -120,23 +108,12 @@ def non_additive_dimension_specs_by_measure(self) -> Dict[MeasureReference, NonA
"""
return self._measure_non_additive_dimension_specs

@staticmethod
def get_measure_from_semantic_model(semantic_model: SemanticModel, measure_reference: MeasureReference) -> Measure:
"""Get measure from semantic model."""
for measure in semantic_model.measures:
if measure.reference == measure_reference:
return measure

raise ValueError(
f"No dimension with name ({measure_reference.element_name}) in semantic_model with name ({semantic_model.name})"
)

def get_measure(self, measure_reference: MeasureReference) -> Measure:
"""Retrieve the measure model object associated with the measure reference."""
if measure_reference not in self._measure_index:
raise ValueError(f"Could not find measure with name ({measure_reference}) in configured semantic models")

return SemanticModelLookup.get_measure_from_semantic_model(
return SemanticModelHelper.get_measure_from_semantic_model(
semantic_model=self.get_semantic_model_for_measure(measure_reference), measure_reference=measure_reference
)

Expand Down Expand Up @@ -276,7 +253,8 @@ def get_primary_entity_else_error(self, semantic_model: SemanticModel) -> Entity
also assume there must be a primary entity because measures are required to have an `agg_time_dimension`
defined in the same semantic model.
"""
primary_entity = SemanticModelLookup.resolved_primary_entity(semantic_model)
# TODO: Move me.
primary_entity = SemanticModelHelper.resolved_primary_entity(semantic_model)
if primary_entity is None:
raise RuntimeError(
f"The semantic model should have a primary entity since there are dimensions, but it does not. "
Expand All @@ -301,57 +279,6 @@ def get_semantic_models_for_dimension(self, dimension_reference: DimensionRefere
"""Return all semantic models associated with a dimension reference."""
return set(self._dimension_index.get(dimension_reference, []))

@staticmethod
def get_entity_from_semantic_model(
semantic_model: SemanticModel, entity_reference: LinkableElementReference
) -> Entity:
"""Get entity from semantic model."""
for entity in semantic_model.entities:
if entity.reference == entity_reference:
return entity

raise ValueError(
f"No entity with name ({entity_reference}) in semantic_model with name ({semantic_model.name})"
)

@staticmethod
def resolved_primary_entity(semantic_model: SemanticModel) -> Optional[EntityReference]:
"""Return the primary entity for dimensions in the model."""
primary_entity_reference = semantic_model.primary_entity_reference

entities_with_type_primary = tuple(
entity for entity in semantic_model.entities if entity.type == EntityType.PRIMARY
)

# This should be caught by the validation, but adding a sanity check.
assert len(entities_with_type_primary) <= 1, f"Found > 1 primary entity in {semantic_model}"
if primary_entity_reference is not None:
assert len(entities_with_type_primary) == 0, (
f"The primary_entity field was set to {primary_entity_reference}, but there are non-zero entities with "
f"type {EntityType.PRIMARY} in {semantic_model}"
)
return primary_entity_reference

if len(entities_with_type_primary) > 0:
return entities_with_type_primary[0].reference

return None

@staticmethod
def entity_links_for_local_elements(semantic_model: SemanticModel) -> Sequence[EntityReference]:
"""Return the entity prefix that can be used to access dimensions defined in the semantic model."""
primary_entity_reference = semantic_model.primary_entity_reference

possible_entity_links = set()
if primary_entity_reference is not None:
possible_entity_links.add(primary_entity_reference)

for entity in semantic_model.entities:
if entity.is_linkable_entity_type:
possible_entity_links.add(entity.reference)

return sorted(possible_entity_links, key=lambda entity_reference: entity_reference.element_name)

def get_element_spec_for_name(self, element_name: str) -> LinkableInstanceSpec:
"""Returns the spec for the given name of a linkable element (dimension or entity)."""
if TimeDimensionReference(element_name=element_name) in self._dimension_ref_to_spec:
Expand Down Expand Up @@ -382,7 +309,7 @@ def _get_agg_time_dimension_specs_for_measure(
# A measure's agg_time_dimension is required to be in the same semantic model as the measure,
# so we can assume the same semantic model for both measure and dimension.
semantic_model = self.get_semantic_model_for_measure(measure_reference)
entity_link = self.resolved_primary_entity(semantic_model)
entity_link = SemanticModelHelper.resolved_primary_entity(semantic_model)
assert entity_link is not None, (
f"Expected semantic model {semantic_model} to have a primary entity since it has a "
"measure requiring an agg_time_dimension, but found none.",
Expand Down
4 changes: 2 additions & 2 deletions metricflow/dataset/convert_semantic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
MeasureInstance,
TimeDimensionInstance,
)
from metricflow_semantics.model.semantics.semantic_model_lookup import SemanticModelLookup
from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper
from metricflow_semantics.model.spec_converters import MeasureConverter
from metricflow_semantics.specs.column_assoc import ColumnAssociationResolver
from metricflow_semantics.specs.dimension_spec import DimensionSpec
Expand Down Expand Up @@ -441,7 +441,7 @@ def create_sql_source_data_set(self, semantic_model: SemanticModel) -> SemanticM
(),
]

for entity_link in SemanticModelLookup.entity_links_for_local_elements(semantic_model):
for entity_link in SemanticModelHelper.entity_links_for_local_elements(semantic_model):
possible_entity_links.append((entity_link,))

# Handle dimensions
Expand Down
7 changes: 4 additions & 3 deletions metricflow/engine/metricflow_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from metricflow_semantics.model.semantic_manifest_lookup import SemanticManifestLookup
from metricflow_semantics.model.semantics.element_filter import LinkableElementFilter
from metricflow_semantics.model.semantics.linkable_element import LinkableDimension
from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper
from metricflow_semantics.model.semantics.semantic_model_lookup import SemanticModelLookup
from metricflow_semantics.naming.linkable_spec_name import StructuredLinkableSpecName
from metricflow_semantics.protocols.query_parameter import GroupByParameter, MetricQueryParameter, OrderByQueryParameter
Expand Down Expand Up @@ -653,7 +654,7 @@ def simple_dimensions_for_metrics( # noqa: D102
assert semantic_model
dimensions.append(
Dimension.from_pydantic(
pydantic_dimension=SemanticModelLookup.get_dimension_from_semantic_model(
pydantic_dimension=SemanticModelHelper.get_dimension_from_semantic_model(
semantic_model=semantic_model,
dimension_reference=linkable_dimension.reference,
),
Expand All @@ -671,7 +672,7 @@ def list_dimensions(self) -> List[Dimension]: # noqa: D102
for semantic_model in semantic_model_lookup.get_semantic_models_for_dimension(dimension_reference):
dimensions.append(
Dimension.from_pydantic(
pydantic_dimension=semantic_model_lookup.get_dimension_from_semantic_model(
pydantic_dimension=SemanticModelHelper.get_dimension_from_semantic_model(
semantic_model=semantic_model, dimension_reference=dimension_reference
),
entity_links=(semantic_model_lookup.get_primary_entity_else_error(semantic_model),),
Expand Down Expand Up @@ -706,7 +707,7 @@ def entities_for_metrics(self, metric_names: List[str]) -> List[Entity]: # noqa
assert semantic_model
entities.append(
Entity.from_pydantic(
pydantic_entity=SemanticModelLookup.get_entity_from_semantic_model(
pydantic_entity=SemanticModelHelper.get_entity_from_semantic_model(
semantic_model=semantic_model,
entity_reference=EntityReference(element_name=linkable_entity.element_name),
)
Expand Down

0 comments on commit 9537a88

Please sign in to comment.