Skip to content

Commit

Permalink
Move _make_time_range_comparison_expr.
Browse files Browse the repository at this point in the history
  • Loading branch information
plypaul committed Nov 12, 2024
1 parent 7344c25 commit 687300c
Showing 1 changed file with 37 additions and 35 deletions.
72 changes: 37 additions & 35 deletions metricflow/plan_conversion/dataflow_to_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,39 +210,6 @@ def convert_to_sql_query_plan(
)


def _make_time_range_comparison_expr(
table_alias: str, column_alias: str, time_range_constraint: TimeRangeConstraint
) -> SqlExpressionNode:
"""Build an expression like "ds BETWEEN CAST('2020-01-01' AS TIMESTAMP) AND CAST('2020-01-02' AS TIMESTAMP).
If the constraint uses day or larger grain, only render to the date level. Otherwise, render to the timestamp level.
"""

def strip_time_from_dt(ts: dt.datetime) -> dt.datetime:
date_obj = ts.date()
return dt.datetime(date_obj.year, date_obj.month, date_obj.day)

constraint_uses_day_or_larger_grain = True
for constraint_input in (time_range_constraint.start_time, time_range_constraint.end_time):
if strip_time_from_dt(constraint_input) != constraint_input:
constraint_uses_day_or_larger_grain = False
break

time_format_to_render = ISO8601_PYTHON_FORMAT if constraint_uses_day_or_larger_grain else ISO8601_PYTHON_TS_FORMAT

return SqlBetweenExpression.create(
column_arg=SqlColumnReferenceExpression.create(
SqlColumnReference(table_alias=table_alias, column_name=column_alias)
),
start_expr=SqlStringLiteralExpression.create(
literal_value=time_range_constraint.start_time.strftime(time_format_to_render),
),
end_expr=SqlStringLiteralExpression.create(
literal_value=time_range_constraint.end_time.strftime(time_format_to_render),
),
)


class DataflowNodeToSqlSubqueryVisitor(DataflowPlanNodeVisitor[SqlDataSet]):
"""Generates a SQL query plan by converting a node's parents to sub-queries.
Expand Down Expand Up @@ -360,7 +327,7 @@ def _make_time_spine_data_set(
from_source_alias=time_spine_table_alias,
group_bys=select_columns if apply_group_by else (),
where=(
_make_time_range_comparison_expr(
DataflowNodeToSqlSubqueryVisitor._make_time_range_comparison_expr(
table_alias=time_spine_table_alias,
column_alias=time_spine_source.base_column,
time_range_constraint=time_range_constraint,
Expand Down Expand Up @@ -1125,7 +1092,7 @@ def visit_constrain_time_range_node(self, node: ConstrainTimeRangeNode) -> SqlDa
time_dimension_instance_for_metric_time = time_dimension_instances_for_metric_time[0]

# Build an expression like "ds >= CAST('2020-01-01' AS TIMESTAMP) AND ds <= CAST('2020-01-02' AS TIMESTAMP)"
constrain_metric_time_column_condition = _make_time_range_comparison_expr(
constrain_metric_time_column_condition = DataflowNodeToSqlSubqueryVisitor._make_time_range_comparison_expr(
table_alias=from_data_set_alias,
column_alias=time_dimension_instance_for_metric_time.associated_column.column_name,
time_range_constraint=node.time_range_constraint,
Expand Down Expand Up @@ -1956,3 +1923,38 @@ def visit_window_reaggregation_node(self, node: WindowReaggregationNode) -> SqlD
group_bys=outer_query_select_columns,
),
)

@staticmethod
def _make_time_range_comparison_expr(
table_alias: str, column_alias: str, time_range_constraint: TimeRangeConstraint
) -> SqlExpressionNode:
"""Build an expression like "ds BETWEEN CAST('2020-01-01' AS TIMESTAMP) AND CAST('2020-01-02' AS TIMESTAMP).
If the constraint uses day or larger grain, only render to the date level. Otherwise, render to the timestamp level.
"""

def strip_time_from_dt(ts: dt.datetime) -> dt.datetime:
date_obj = ts.date()
return dt.datetime(date_obj.year, date_obj.month, date_obj.day)

constraint_uses_day_or_larger_grain = True
for constraint_input in (time_range_constraint.start_time, time_range_constraint.end_time):
if strip_time_from_dt(constraint_input) != constraint_input:
constraint_uses_day_or_larger_grain = False
break

time_format_to_render = (
ISO8601_PYTHON_FORMAT if constraint_uses_day_or_larger_grain else ISO8601_PYTHON_TS_FORMAT
)

return SqlBetweenExpression.create(
column_arg=SqlColumnReferenceExpression.create(
SqlColumnReference(table_alias=table_alias, column_name=column_alias)
),
start_expr=SqlStringLiteralExpression.create(
literal_value=time_range_constraint.start_time.strftime(time_format_to_render),
),
end_expr=SqlStringLiteralExpression.create(
literal_value=time_range_constraint.end_time.strftime(time_format_to_render),
),
)

0 comments on commit 687300c

Please sign in to comment.