Skip to content

Commit

Permalink
Add window_function attribute to TimeDimensionSpec
Browse files Browse the repository at this point in the history
This will allow us to track which specs have had a window function applied between DataflowPlan nodes
  • Loading branch information
courtneyholcomb committed Dec 21, 2024
1 parent 1d34121 commit 1b71261
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,11 @@ def visit_time_dimension_spec(self, time_dimension_spec: TimeDimensionSpec) -> C
if time_dimension_spec.aggregation_state
else ""
)
+ (
f"{DUNDER}{time_dimension_spec.window_function.value.lower()}"
if time_dimension_spec.window_function
else ""
)
)

def visit_entity_spec(self, entity_spec: EntitySpec) -> ColumnAssociation: # noqa: D102
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from metricflow_semantics.naming.linkable_spec_name import StructuredLinkableSpecName
from metricflow_semantics.specs.dimension_spec import DimensionSpec
from metricflow_semantics.specs.instance_spec import InstanceSpecVisitor
from metricflow_semantics.sql.sql_exprs import SqlWindowFunction
from metricflow_semantics.time.granularity import ExpandedTimeGranularity
from metricflow_semantics.visitor import VisitorOutputT

Expand Down Expand Up @@ -91,6 +92,8 @@ class TimeDimensionSpec(DimensionSpec): # noqa: D101
# Used for semi-additive joins. Some more thought is needed, but this may be useful in InstanceSpec.
aggregation_state: Optional[AggregationState] = None

window_function: Optional[SqlWindowFunction] = None

@property
def without_first_entity_link(self) -> TimeDimensionSpec: # noqa: D102
assert len(self.entity_links) > 0, f"Spec does not have any entity links: {self}"
Expand All @@ -99,6 +102,8 @@ def without_first_entity_link(self) -> TimeDimensionSpec: # noqa: D102
entity_links=self.entity_links[1:],
time_granularity=self.time_granularity,
date_part=self.date_part,
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

@property
Expand All @@ -108,6 +113,8 @@ def without_entity_links(self) -> TimeDimensionSpec: # noqa: D102
time_granularity=self.time_granularity,
date_part=self.date_part,
entity_links=(),
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

@property
Expand Down Expand Up @@ -153,6 +160,7 @@ def with_grain(self, time_granularity: ExpandedTimeGranularity) -> TimeDimension
time_granularity=time_granularity,
date_part=self.date_part,
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

def with_base_grain(self) -> TimeDimensionSpec: # noqa: D102
Expand All @@ -162,6 +170,7 @@ def with_base_grain(self) -> TimeDimensionSpec: # noqa: D102
time_granularity=ExpandedTimeGranularity.from_time_granularity(self.time_granularity.base_granularity),
date_part=self.date_part,
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

def with_grain_and_date_part( # noqa: D102
Expand All @@ -173,6 +182,7 @@ def with_grain_and_date_part( # noqa: D102
time_granularity=time_granularity,
date_part=date_part,
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

def with_aggregation_state(self, aggregation_state: AggregationState) -> TimeDimensionSpec: # noqa: D102
Expand All @@ -182,6 +192,17 @@ def with_aggregation_state(self, aggregation_state: AggregationState) -> TimeDim
time_granularity=self.time_granularity,
date_part=self.date_part,
aggregation_state=aggregation_state,
window_function=self.window_function,
)

def with_window_function(self, window_function: SqlWindowFunction) -> TimeDimensionSpec: # noqa: D102
return TimeDimensionSpec(
element_name=self.element_name,
entity_links=self.entity_links,
time_granularity=self.time_granularity,
date_part=self.date_part,
aggregation_state=self.aggregation_state,
window_function=window_function,
)

def comparison_key(self, exclude_fields: Sequence[TimeDimensionSpecField] = ()) -> TimeDimensionSpecComparisonKey:
Expand Down Expand Up @@ -243,6 +264,7 @@ def with_entity_prefix(self, entity_prefix: EntityReference) -> TimeDimensionSpe
time_granularity=self.time_granularity,
date_part=self.date_part,
aggregation_state=self.aggregation_state,
window_function=self.window_function,
)

@staticmethod
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def test_classes() -> None: # noqa: D103
time_granularity=ExpandedTimeGranularity(name='day', base_granularity=DAY),
date_part=None,
aggregation_state=None,
window_function=None,
)
"""
).rstrip()
Expand Down

0 comments on commit 1b71261

Please sign in to comment.