Migrate existing use cases to DateutilTimePeriodAdjuster.

dbt-labs · May 31, 2024 · b1e19e4 · b1e19e4
1 parent 0c5dd7a
commit b1e19e4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 27 deletions.
diff --git a/metricflow-semantics/metricflow_semantics/query/query_parser.py b/metricflow-semantics/metricflow_semantics/query/query_parser.py
@@ -5,7 +5,6 @@
 from dataclasses import dataclass
 from typing import List, Optional, Sequence, Tuple, Union
 
-import pandas as pd
 from dbt_semantic_interfaces.implementations.filters.where_filter import (
     PydanticWhereFilter,
     PydanticWhereFilterIntersection,
@@ -64,12 +63,7 @@
     TimeDimensionSpec,
 )
 from metricflow_semantics.specs.spec_set import group_specs_by_type
-from metricflow_semantics.time.pandas_adjuster import (
-    adjust_to_end_of_period,
-    adjust_to_start_of_period,
-    is_period_end,
-    is_period_start,
-)
+from metricflow_semantics.time.dateutil_adjuster import DateutilTimePeriodAdjuster
 
 logger = logging.getLogger(__name__)
 
@@ -95,6 +89,7 @@ def __init__(  # noqa: D107
             DunderNamingScheme(),
         )
         self._where_filter_pattern_factory = where_filter_pattern_factory
+        self._time_period_adjuster = DateutilTimePeriodAdjuster()
 
     def parse_and_validate_saved_query(
         self,
@@ -191,23 +186,10 @@ def _adjust_time_constraint(
 
         e.g. [2020-01-15, 2020-2-15] with MONTH granularity -> [2020-01-01, 2020-02-29]
         """
-        constraint_start = time_constraint.start_time
-        constraint_end = time_constraint.end_time
-
-        start_ts = pd.Timestamp(time_constraint.start_time)
-        if not is_period_start(metric_time_granularity, start_ts):
-            constraint_start = adjust_to_start_of_period(metric_time_granularity, start_ts).to_pydatetime()
-
-        end_ts = pd.Timestamp(time_constraint.end_time)
-        if not is_period_end(metric_time_granularity, end_ts):
-            constraint_end = adjust_to_end_of_period(metric_time_granularity, end_ts).to_pydatetime()
-
-        if constraint_start < TimeRangeConstraint.ALL_TIME_BEGIN():
-            constraint_start = TimeRangeConstraint.ALL_TIME_BEGIN()
-        if constraint_end > TimeRangeConstraint.ALL_TIME_END():
-            constraint_end = TimeRangeConstraint.ALL_TIME_END()
-
-        return TimeRangeConstraint(start_time=constraint_start, end_time=constraint_end)
+        return self._time_period_adjuster.expand_time_constraint_to_fill_granularity(
+            time_constraint=time_constraint,
+            granularity=metric_time_granularity,
+        )
 
     def _parse_order_by_names(
         self,

diff --git a/metricflow/dataflow/builder/dataflow_plan_builder.py b/metricflow/dataflow/builder/dataflow_plan_builder.py
@@ -52,7 +52,7 @@
 from metricflow_semantics.specs.spec_set import InstanceSpecSet, group_specs_by_type
 from metricflow_semantics.specs.where_filter_transform import WhereSpecFactory
 from metricflow_semantics.sql.sql_join_type import SqlJoinType
-from metricflow_semantics.time.pandas_adjuster import PandasTimePeriodAdjuster
+from metricflow_semantics.time.dateutil_adjuster import DateutilTimePeriodAdjuster
 
 from metricflow.dataflow.builder.node_data_set import DataflowPlanNodeOutputDataSetResolver
 from metricflow.dataflow.builder.node_evaluator import (
@@ -135,7 +135,7 @@ def __init__(  # noqa: D107
         self._column_association_resolver = column_association_resolver
         self._node_data_set_resolver = node_output_resolver
         self._source_node_builder = source_node_builder
-        self._time_period_adjuster = PandasTimePeriodAdjuster()
+        self._time_period_adjuster = DateutilTimePeriodAdjuster()
 
     def build_plan(
         self,
@@ -1298,7 +1298,7 @@ def _build_aggregated_measure_from_measure_source_node(
                     predicate_pushdown_state.time_range_constraint, granularity, count
                 )
             )
-            logger.info(f"Adjusted time range constraint {cumulative_metric_adjusted_time_constraint}")
+            logger.info(f"Adjusted time range constraint to: {cumulative_metric_adjusted_time_constraint}")
 
         required_linkable_specs, extraneous_linkable_specs = self.__get_required_and_extraneous_linkable_specs(
             queried_linkable_specs=queried_linkable_specs,