From 2a18b9cc6b641871d6a0b2be366622c9bffd2fec Mon Sep 17 00:00:00 2001
From: Paul Yang <github@ardorful.com>
Date: Mon, 30 Sep 2024 13:54:40 -0700
Subject: [PATCH] /* PR_START p--short-term-perf 13 */ Cache frequently-used
 methods in `MetricLookup`.

---
 .../model/semantics/metric_lookup.py          | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py b/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
index e39e9063ec..02e7ea3a85 100644
--- a/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
+++ b/metricflow-semantics/metricflow_semantics/model/semantics/metric_lookup.py
@@ -54,6 +54,14 @@ def __init__(
             max_entity_links=MAX_JOIN_HOPS,
         )
 
+        # Cache for `get_min_queryable_time_granularity()`
+        self._metric_reference_to_min_metric_time_grain: Dict[MetricReference, TimeGranularity] = {}
+
+        # Cache for `get_valid_agg_time_dimensions_for_metric()`.
+        self._metric_reference_to_valid_agg_time_dimension_specs: Dict[
+            MetricReference, Sequence[TimeDimensionSpec]
+        ] = {}
+
     @functools.lru_cache
     def linkable_elements_for_measure(
         self,
@@ -183,6 +191,18 @@ def get_valid_agg_time_dimensions_for_metric(
         self, metric_reference: MetricReference
     ) -> Sequence[TimeDimensionSpec]:
         """Get the agg time dimension specs that can be used in place of metric time for this metric, if applicable."""
+        result = self._metric_reference_to_valid_agg_time_dimension_specs.get(metric_reference)
+        if result is not None:
+            return result
+
+        result = self._get_valid_agg_time_dimensions_for_metric(metric_reference)
+        self._metric_reference_to_valid_agg_time_dimension_specs[metric_reference] = result
+
+        return result
+
+    def _get_valid_agg_time_dimensions_for_metric(
+        self, metric_reference: MetricReference
+    ) -> Sequence[TimeDimensionSpec]:
         agg_time_dimension_specs = self._get_agg_time_dimension_specs_for_metric(metric_reference)
         distinct_agg_time_dimension_identifiers = set(
             [(spec.reference, spec.entity_links) for spec in agg_time_dimension_specs]
@@ -204,6 +224,15 @@ def get_min_queryable_time_granularity(self, metric_reference: MetricReference)
 
         Maps to the largest granularity defined for any of the metric's agg_time_dimensions.
         """
+        result = self._metric_reference_to_min_metric_time_grain.get(metric_reference)
+        if result is not None:
+            return result
+
+        result = self._get_min_queryable_time_granularity(metric_reference)
+        self._metric_reference_to_min_metric_time_grain[metric_reference] = result
+        return result
+
+    def _get_min_queryable_time_granularity(self, metric_reference: MetricReference) -> TimeGranularity:
         agg_time_dimension_specs = self._get_agg_time_dimension_specs_for_metric(metric_reference)
         assert (
             agg_time_dimension_specs