From 1878e548e55c4ef8a2a2bd73eff1e842985ad1d2 Mon Sep 17 00:00:00 2001
From: Paul Yang <github@ardorful.com>
Date: Tue, 1 Oct 2024 22:37:40 -0700
Subject: [PATCH] Cache methods in `ExpandedTimeGranularity` (#1436)

The methods in `ExpandedTimeGranularity` benefits from caching as they
are called in many cases where `TimeDimensionSpecs` are created. `*Spec`
objects are one of the most frequently created / used objects.
---
 .../specs/time_dimension_spec.py              | 30 ++++++++++++-------
 .../metricflow_semantics/time/granularity.py  | 26 ++++++++++------
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py b/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py
index dc6457c73a..9b42dcf2ed 100644
--- a/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py
+++ b/metricflow-semantics/metricflow_semantics/specs/time_dimension_spec.py
@@ -2,6 +2,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
+from functools import lru_cache
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from dbt_semantic_interfaces.naming.keywords import METRIC_TIME_ELEMENT_NAME
@@ -190,8 +191,18 @@ def comparison_key(self, exclude_fields: Sequence[TimeDimensionSpecField] = ())
             exclude_fields=exclude_fields,
         )
 
-    @staticmethod
+    @classmethod
+    @lru_cache
+    def _get_compatible_grain_and_date_part(cls) -> Sequence[Tuple[ExpandedTimeGranularity, DatePart]]:
+        items = []
+        for date_part in DatePart:
+            for compatible_granularity in date_part.compatible_granularities:
+                items.append((ExpandedTimeGranularity.from_time_granularity(compatible_granularity), date_part))
+        return items
+
+    @classmethod
     def generate_possible_specs_for_time_dimension(
+        cls,
         time_dimension_reference: TimeDimensionReference,
         entity_links: Tuple[EntityReference, ...],
         custom_granularities: Dict[str, ExpandedTimeGranularity],
@@ -210,16 +221,15 @@ def generate_possible_specs_for_time_dimension(
                     date_part=None,
                 )
             )
-        for date_part in DatePart:
-            for compatible_granularity in date_part.compatible_granularities:
-                time_dimension_specs.append(
-                    TimeDimensionSpec(
-                        element_name=time_dimension_reference.element_name,
-                        entity_links=entity_links,
-                        time_granularity=ExpandedTimeGranularity.from_time_granularity(compatible_granularity),
-                        date_part=date_part,
-                    )
+        for grain, date_part in cls._get_compatible_grain_and_date_part():
+            time_dimension_specs.append(
+                TimeDimensionSpec(
+                    element_name=time_dimension_reference.element_name,
+                    entity_links=entity_links,
+                    time_granularity=grain,
+                    date_part=date_part,
                 )
+            )
         return time_dimension_specs
 
     @property
diff --git a/metricflow-semantics/metricflow_semantics/time/granularity.py b/metricflow-semantics/metricflow_semantics/time/granularity.py
index cf036b869c..1903f46e59 100644
--- a/metricflow-semantics/metricflow_semantics/time/granularity.py
+++ b/metricflow-semantics/metricflow_semantics/time/granularity.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cached_property, lru_cache
+from typing import FrozenSet
 
 from dbt_semantic_interfaces.dataclass_serialization import SerializableDataclass
 from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
@@ -32,20 +34,26 @@ def __post_init__(self) -> None:
                 f"{self.base_granularity}."
             )
 
-    @property
+    @cached_property
     def is_custom_granularity(self) -> bool:  # noqa: D102
         return self.base_granularity.value != self.name
 
     @classmethod
+    @lru_cache
     def from_time_granularity(cls, granularity: TimeGranularity) -> ExpandedTimeGranularity:
-        """Factory method for creating an ExpandedTimeGranularity from a standard TimeGranularity enumeration value."""
+        """Factory method for creating an ExpandedTimeGranularity from a standard TimeGranularity enumeration value.
+
+        This should be appropriate to use with `@lru_cache` since the number of `TimeGranularity` is small.
+        """
         return ExpandedTimeGranularity(name=granularity.value, base_granularity=granularity)
 
-    @staticmethod
-    def is_standard_granularity_name(time_granularity_name: str) -> bool:
-        """Helper for checking if a given time granularity name is part of the standard TimeGranularity enumeration."""
-        for granularity in TimeGranularity:
-            if time_granularity_name == granularity.value:
-                return True
+    @classmethod
+    @lru_cache
+    def _standard_time_granularity_names(cls) -> FrozenSet:
+        """This should be appropriate to use with `@lru_cache` since the number of `TimeGranularity` is small."""
+        return frozenset(granularity.value for granularity in TimeGranularity)
 
-        return False
+    @classmethod
+    def is_standard_granularity_name(cls, time_granularity_name: str) -> bool:
+        """Helper for checking if a given time granularity name is part of the standard TimeGranularity enumeration."""
+        return time_granularity_name in cls._standard_time_granularity_names()