From d18e478f395f09bec0cf4cf6cc1f2986f015f162 Mon Sep 17 00:00:00 2001 From: Paul Yang Date: Mon, 28 Oct 2024 17:54:40 -0700 Subject: [PATCH] /* PR_START p--short-term-perf 29 */ Move static methods to `SemanticModelHelper`. --- .../model/semantics/linkable_spec_resolver.py | 5 +- .../model/semantics/semantic_model_helper.py | 92 +++++++++++++++++++ .../model/semantics/semantic_model_lookup.py | 87 ++---------------- metricflow/dataset/convert_semantic_model.py | 4 +- metricflow/engine/metricflow_engine.py | 8 +- 5 files changed, 108 insertions(+), 88 deletions(-) create mode 100644 metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py b/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py index e9b142b4c0..3dd0468a54 100644 --- a/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py +++ b/metricflow-semantics/metricflow_semantics/model/semantics/linkable_spec_resolver.py @@ -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 @@ -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 @@ -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: diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py new file mode 100644 index 0000000000..268ead0a11 --- /dev/null +++ b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_helper.py @@ -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})" + ) diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py index a3c3ef8388..8fef598e2c 100644 --- a/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py +++ b/metricflow-semantics/metricflow_semantics/model/semantics/semantic_model_lookup.py @@ -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 @@ -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. @@ -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, @@ -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 ) @@ -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. " @@ -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: @@ -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.", diff --git a/metricflow/dataset/convert_semantic_model.py b/metricflow/dataset/convert_semantic_model.py index 1abee6ddaa..4141606704 100644 --- a/metricflow/dataset/convert_semantic_model.py +++ b/metricflow/dataset/convert_semantic_model.py @@ -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 @@ -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 diff --git a/metricflow/engine/metricflow_engine.py b/metricflow/engine/metricflow_engine.py index 83099dd28f..0d684ce0ac 100644 --- a/metricflow/engine/metricflow_engine.py +++ b/metricflow/engine/metricflow_engine.py @@ -20,7 +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_lookup import SemanticModelLookup +from metricflow_semantics.model.semantics.semantic_model_helper import SemanticModelHelper from metricflow_semantics.naming.linkable_spec_name import StructuredLinkableSpecName from metricflow_semantics.protocols.query_parameter import GroupByParameter, MetricQueryParameter, OrderByQueryParameter from metricflow_semantics.query.query_exceptions import InvalidQueryException @@ -653,7 +653,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, ), @@ -671,7 +671,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),), @@ -706,7 +706,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), )