From 7469801200e9b4201107ceb9bc3f727836e0a790 Mon Sep 17 00:00:00 2001
From: Paul Yang <github@ardorful.com>
Date: Tue, 1 Oct 2024 17:12:50 -0700
Subject: [PATCH] /* PR_START p--short-term-perf 18 */ Improve cache for
 `MetricLookup.linkable_elements_for_measure()`.

---
 .../model/semantics/metric_lookup.py          | 38 +++++++++++--------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py b/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
index 00a4def8a9..ce5dc150e2 100644
--- a/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
+++ b/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
@@ -2,8 +2,7 @@
 
 import functools
 import logging
-import time
-from typing import Dict, Optional, Sequence, Set
+from typing import Dict, Optional, Sequence, Set, Tuple
 
 from dbt_semantic_interfaces.enum_extension import assert_values_exhausted
 from dbt_semantic_interfaces.protocols.metric import Metric, MetricInputMeasure, MetricType
@@ -12,7 +11,7 @@
 from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
 
 from metricflow_semantics.errors.error_classes import DuplicateMetricError, MetricNotFoundError, NonExistentMeasureError
-from metricflow_semantics.mf_logging.lazy_formattable import LazyFormat
+from metricflow_semantics.model.linkable_element_property import LinkableElementProperty
 from metricflow_semantics.model.semantics.element_filter import LinkableElementFilter
 from metricflow_semantics.model.semantics.linkable_element_set import LinkableElementSet
 from metricflow_semantics.model.semantics.linkable_spec_resolver import (
@@ -62,25 +61,34 @@ def __init__(
             MetricReference, Sequence[TimeDimensionSpec]
         ] = {}
 
-    @functools.lru_cache
+        # Cache for `linkable_elements_for_measure()`.
+        self._linkable_element_set_for_measure_cache: Dict[
+            Tuple[MeasureReference, LinkableElementFilter], LinkableElementSet
+        ] = {}
+
     def linkable_elements_for_measure(
         self,
         measure_reference: MeasureReference,
         element_filter: LinkableElementFilter = LinkableElementFilter(),
     ) -> LinkableElementSet:
         """Return the set of linkable elements reachable from a given measure."""
-        start_time = time.time()
-        linkable_element_set = self._linkable_spec_resolver.get_linkable_element_set_for_measure(
-            measure_reference=measure_reference,
-            element_filter=element_filter,
-        )
-        logger.debug(
-            LazyFormat(
-                lambda: f"Getting valid linkable elements for measure '{measure_reference.element_name}' took: {time.time() - start_time:.2f}s"
-            )
-        )
+        # Don't cache the result when group-by-metrics are selected as there can be many of them and may significantly
+        # increase memory usage.
+        if (
+            LinkableElementProperty.METRIC in element_filter.with_any_of
+            and LinkableElementProperty.METRIC not in element_filter.without_any_of
+        ):
+            return self._linkable_spec_resolver.get_linkable_element_set_for_measure(measure_reference, element_filter)
+
+        cache_key = (measure_reference, element_filter)
+        result = self._linkable_element_set_for_measure_cache.get(cache_key)
+        if result is not None:
+            return result
+
+        result = self._linkable_spec_resolver.get_linkable_element_set_for_measure(measure_reference, element_filter)
+        self._linkable_element_set_for_measure_cache[cache_key] = result
 
-        return linkable_element_set
+        return result
 
     @functools.lru_cache
     def linkable_elements_for_no_metrics_query(