Handle time constraint adjustment for sub-daily time granularity

dbt-labs · Aug 13, 2024 · 9edd1a6 · 9edd1a6
1 parent 30d469d
commit 9edd1a6
Show file tree

Hide file tree

Showing 3 changed files with 5,891 additions and 5,507 deletions.
diff --git a/metricflow-semantics/metricflow_semantics/time/dateutil_adjuster.py b/metricflow-semantics/metricflow_semantics/time/dateutil_adjuster.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import datetime
+from copy import deepcopy
 from typing import Optional
 
 import dateutil.relativedelta
@@ -22,13 +23,12 @@ class DateutilTimePeriodAdjuster(TimePeriodAdjuster):
 
     def _relative_delta_for_window(self, time_granularity: TimeGranularity, count: int) -> relativedelta:
         """Relative-delta to cover time windows specified at different grains."""
-        if time_granularity is TimeGranularity.NANOSECOND:
-            # TODO: figure out a workaround when enabling time constraints
-            raise ValueError("`relativedelta` does not support nanoseconds.")
-        elif time_granularity is TimeGranularity.MICROSECOND:
-            return relativedelta(microseconds=count)
-        elif time_granularity is TimeGranularity.MILLISECOND:
-            return relativedelta(microseconds=count * 1000)
+        if (
+            time_granularity is TimeGranularity.NANOSECOND
+            or time_granularity is TimeGranularity.MICROSECOND
+            or time_granularity is TimeGranularity.MILLISECOND
+        ):
+            raise ValueError("Time constraints only support SECOND or larger time granularities.")
         elif time_granularity is TimeGranularity.SECOND:
             return relativedelta(seconds=count)
         elif time_granularity is TimeGranularity.MINUTE:
@@ -66,18 +66,35 @@ def expand_time_constraint_to_fill_granularity(
     def adjust_to_start_of_period(
         self, time_granularity: TimeGranularity, date_to_adjust: datetime.datetime
     ) -> datetime.datetime:
-        # TODO: update these options once time constraints support a full timestamp
+        """Takes a date and adjusts it to the start of the given granularity period.
+
+        Used to adjust time constraints when the time dimension they will be applied to has a larger granularity
+        than the time constraint requested.
+
+        Args:
+            time_granularity: The granularity of the time dimension that will be constrained.
+            date_to_adjust: The time constraint start date.
+        """
+        date_to_adjust = deepcopy(date_to_adjust)
+        # Strip time data that is irrelevant to avoid rendering an unnecessarily complex time constraint.
+        date_to_adjust = date_to_adjust.replace(microsecond=0)
+        if time_granularity.to_int() >= TimeGranularity.DAY.to_int():
+            date_to_adjust = date_to_adjust.replace(hour=0, minute=0, second=0)
+
         if (
             time_granularity is TimeGranularity.NANOSECOND
             or time_granularity is TimeGranularity.MICROSECOND
             or time_granularity is TimeGranularity.MILLISECOND
-            or time_granularity is TimeGranularity.SECOND
-            or time_granularity is TimeGranularity.MINUTE
-            or time_granularity is TimeGranularity.HOUR
-            or time_granularity is TimeGranularity.DAY
         ):
+            raise ValueError("Time constraints only support SECOND and larger time granularities.")
+        elif time_granularity is TimeGranularity.SECOND:
+            return date_to_adjust
+        elif time_granularity is TimeGranularity.MINUTE:
+            return date_to_adjust + relativedelta(second=0, microsecond=0)
+        elif time_granularity is TimeGranularity.HOUR:
+            return date_to_adjust + relativedelta(minute=0, second=0, microsecond=0)
+        elif time_granularity is TimeGranularity.DAY:
             return date_to_adjust
-
         elif time_granularity is TimeGranularity.WEEK:
             return date_to_adjust + relativedelta(weekday=dateutil.relativedelta.MO(-1))
         elif time_granularity is TimeGranularity.MONTH:
@@ -100,18 +117,35 @@ def adjust_to_start_of_period(
     def adjust_to_end_of_period(
         self, time_granularity: TimeGranularity, date_to_adjust: datetime.datetime
     ) -> datetime.datetime:
-        # TODO: update these options once time constraints support a full timestamp
+        """Takes a date and adjusts it to the end of the given granularity period.
+
+        Used to adjust time constraints when the time dimension they will be applied to has a larger granularity
+        than the time constraint requested.
+
+        Args:
+            time_granularity: The granularity of the time dimension that will be constrained.
+            date_to_adjust: The time constraint end date.
+        """
+        date_to_adjust = deepcopy(date_to_adjust)
+        # Strip time data that is irrelevant to avoid rendering an unnecessarily complex time constraint.
+        date_to_adjust = date_to_adjust.replace(microsecond=0)
+        if time_granularity.to_int() >= TimeGranularity.DAY.to_int():
+            date_to_adjust = date_to_adjust.replace(hour=0, minute=0, second=0)
+
         if (
             time_granularity is TimeGranularity.NANOSECOND
             or time_granularity is TimeGranularity.MICROSECOND
             or time_granularity is TimeGranularity.MILLISECOND
-            or time_granularity is TimeGranularity.SECOND
-            or time_granularity is TimeGranularity.MINUTE
-            or time_granularity is TimeGranularity.HOUR
-            or time_granularity is TimeGranularity.DAY
         ):
+            raise ValueError("Time constraints only support SECOND and larger time granularities.")
+        elif time_granularity is TimeGranularity.SECOND:
+            return date_to_adjust
+        elif time_granularity is TimeGranularity.MINUTE:
+            return date_to_adjust + relativedelta(second=59, microsecond=999999)
+        elif time_granularity is TimeGranularity.HOUR:
+            return date_to_adjust + relativedelta(minute=59, second=59, microsecond=999999)
+        elif time_granularity is TimeGranularity.DAY:
             return date_to_adjust
-
         elif time_granularity is TimeGranularity.WEEK:
             return date_to_adjust + relativedelta(weekday=dateutil.relativedelta.SU(1))
         elif time_granularity is TimeGranularity.MONTH: