Set JOINED and MULTI_HOP properties from join_path length (cleanup)

dbt-labs · Apr 18, 2024 · ecbb5e9 · ecbb5e9
1 parent 23f976c
commit ecbb5e9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 22 deletions.
diff --git a/metricflow/model/semantics/linkable_spec_resolver.py b/metricflow/model/semantics/linkable_spec_resolver.py
@@ -858,13 +858,7 @@ def _get_joined_elements(self, measure_semantic_model: SemanticModel) -> Linkabl
                     )
                 )
         single_hop_elements = LinkableElementSet.merge_by_path_key(
-            [
-                self.create_linkable_element_set_from_join_path(
-                    join_path=join_path,
-                    with_properties=frozenset({LinkableElementProperty.JOINED}),
-                )
-                for join_path in join_paths
-            ]
+            [self.create_linkable_element_set_from_join_path(join_path) for join_path in join_paths]
         )
 
         # Create multi-hop elements. At each iteration, we generate the list of valid elements based on the current join
@@ -886,11 +880,7 @@ def _get_joined_elements(self, measure_semantic_model: SemanticModel) -> Linkabl
             multi_hop_elements = LinkableElementSet.merge_by_path_key(
                 (multi_hop_elements,)
                 + tuple(
-                    self.create_linkable_element_set_from_join_path(
-                        join_path=new_join_path,
-                        with_properties=frozenset({LinkableElementProperty.JOINED, LinkableElementProperty.MULTI_HOP}),
-                    )
-                    for new_join_path in new_join_paths
+                    self.create_linkable_element_set_from_join_path(new_join_path) for new_join_path in new_join_paths
                 )
             )
             join_paths = new_join_paths
@@ -1019,9 +1009,12 @@ def _find_next_possible_paths(
     def create_linkable_element_set_from_join_path(
         self,
         join_path: SemanticModelJoinPath,
-        with_properties: FrozenSet[LinkableElementProperty],
     ) -> LinkableElementSet:
         """Given the current path, generate the respective linkable elements from the last semantic model in the path."""
+        properties = frozenset({LinkableElementProperty.JOINED})
+        if len(join_path.path_elements) > 1:
+            properties = properties.union({LinkableElementProperty.MULTI_HOP})
+
         semantic_model = self._semantic_model_lookup.get_by_reference(join_path.last_semantic_model_reference)
         assert semantic_model
 
@@ -1038,7 +1031,7 @@ def create_linkable_element_set_from_join_path(
                         element_name=dimension.reference.element_name,
                         entity_links=join_path.entity_links,
                         join_path=join_path.path_elements,
-                        properties=with_properties,
+                        properties=properties,
                         time_granularity=None,
                         date_part=None,
                     )
@@ -1050,7 +1043,7 @@ def create_linkable_element_set_from_join_path(
                         dimension=dimension,
                         entity_links=join_path.entity_links,
                         join_path=join_path.path_elements,
-                        with_properties=with_properties,
+                        with_properties=properties,
                     )
                 )
             else:
@@ -1065,7 +1058,7 @@ def create_linkable_element_set_from_join_path(
                         element_name=entity.reference.element_name,
                         entity_links=join_path.entity_links,
                         join_path=join_path.path_elements,
-                        properties=with_properties.union({LinkableElementProperty.ENTITY}),
+                        properties=properties.union({LinkableElementProperty.ENTITY}),
                     )
                 )
 
@@ -1075,7 +1068,7 @@ def create_linkable_element_set_from_join_path(
                 entity_links=join_path.entity_links,
                 join_path=join_path.path_elements,
                 join_by_semantic_model=semantic_model.reference,
-                properties=with_properties.union({LinkableElementProperty.METRIC}),
+                properties=properties.union({LinkableElementProperty.METRIC}),
             )
             for metric in self._joinable_metrics_for_semantic_models.get(join_path.last_semantic_model_reference, set())
         ]

diff --git a/tests/model/semantics/test_linkable_spec_resolver.py b/tests/model/semantics/test_linkable_spec_resolver.py
@@ -15,6 +15,7 @@
 from metricflow.model.semantics.linkable_element_properties import LinkableElementProperty
 from metricflow.model.semantics.linkable_spec_resolver import (
     SemanticModelJoinPath,
+    SemanticModelJoinPathElement,
     ValidLinkableSpecResolver,
 )
 from metricflow.model.semantics.semantic_model_join_evaluator import MAX_JOIN_HOPS
@@ -145,7 +146,6 @@ def test_create_linkable_element_set_from_join_path(  # noqa: D103
                 semantic_model_reference=SemanticModelReference("listings_latest"),
                 join_on_entity=EntityReference("listing"),
             ),
-            with_properties=frozenset({LinkableElementProperty.JOINED}),
         ),
     )
 
@@ -160,11 +160,18 @@ def test_create_linkable_element_set_from_join_path_multi_hop(  # noqa: D103
         mf_test_configuration=mf_test_configuration,
         set_id="result0",
         linkable_element_set=simple_model_spec_resolver.create_linkable_element_set_from_join_path(
-            join_path=SemanticModelJoinPath.from_single_element(
-                semantic_model_reference=SemanticModelReference("listings_latest"),
-                join_on_entity=EntityReference("listing"),
+            SemanticModelJoinPath(
+                (
+                    SemanticModelJoinPathElement(
+                        semantic_model_reference=SemanticModelReference("bookings"),
+                        join_on_entity=EntityReference("guest"),
+                    ),
+                    SemanticModelJoinPathElement(
+                        semantic_model_reference=SemanticModelReference("listings_latest"),
+                        join_on_entity=EntityReference("listing"),
+                    ),
+                )
             ),
-            with_properties=frozenset({LinkableElementProperty.JOINED, LinkableElementProperty.MULTI_HOP}),
         ),
     )