use rstest::rstest;
use std::collections::HashMap;
use crate::incremental::{compute_incremental_ranges_for_pipeline, TimeRange};
use hamelin_lib::tree::{
ast::{expression::TruncUnit, identifier::Identifier},
builder::{at_hour, call, eq, field_ref, hours, pipeline, string, PipelineBuilder},
};
use super::helpers::{
build_pipeline, stale_ranges, time_range, time_range_exclusive, timestamp_field,
};
#[rstest]
#[case::snap_to_hour_boundaries(
pipeline()
.from(|f| f.table_reference("events"))
.where_cmd(eq(field_ref("severity"), string("high"))),
stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
TruncUnit::Hour,
time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"), // query_range: inclusive
time_range_exclusive("2024-01-01 14:00:00", "2024-01-01 17:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_to_day_boundaries(
pipeline()
.from(|f| f.table_reference("events"))
.where_cmd(eq(field_ref("severity"), string("high"))),
stale_ranges("events", "2024-01-01 14:30:00", "2024-01-03 16:45:00"),
TruncUnit::Day,
time_range("2024-01-01 00:00:00", "2024-01-04 00:00:00"), // query_range: inclusive
time_range_exclusive("2024-01-01 00:00:00", "2024-01-04 00:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_hourly_agg_to_daily_partitions(
pipeline()
.from(|f| f.table_reference("events"))
.agg(|a| a
.named_aggregate("count", call("count"))
.named_group("timestamp", at_hour(field_ref("timestamp")))
),
stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
TruncUnit::Day,
time_range("2024-01-01 00:00:00", "2024-01-02 01:00:00"), // query_range: inclusive
time_range_exclusive("2024-01-01 00:00:00", "2024-01-02 00:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_with_window_lookback(
pipeline()
.from(|f| f.table_reference("events"))
.window(|w| w
.named_field("sum_amount", call("sum").arg(field_ref("amount")))
.within(hours(-1))
),
stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
TruncUnit::Hour,
time_range("2024-01-01 13:00:00", "2024-01-01 17:00:00"), // query_range: inclusive
time_range_exclusive("2024-01-01 14:00:00", "2024-01-01 17:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_to_month_boundaries(
pipeline()
.from(|f| f.table_reference("events"))
.where_cmd(eq(field_ref("severity"), string("high"))),
stale_ranges("events", "2024-01-15 10:00:00", "2024-02-20 14:00:00"),
TruncUnit::Month,
time_range("2024-01-01 00:00:00", "2024-03-01 00:00:00"), // query_range: inclusive
time_range_exclusive("2024-01-01 00:00:00", "2024-03-01 00:00:00") // replace_range: snapped, exclusive
)]
#[case::no_snap_when_none(
pipeline()
.from(|f| f.table_reference("events"))
.where_cmd(eq(field_ref("severity"), string("high"))),
stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
TruncUnit::Second, // We'll use this as a placeholder, but pass None
time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"), // query_range: unchanged
time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00") // replace_range: unchanged (no snap)
)]
fn test_partition_snapping(
#[case] pipeline_builder: PipelineBuilder,
#[case] stale_ranges_map: HashMap<Identifier, Option<TimeRange>>,
#[case] partition_unit: TruncUnit,
#[case] expected_query_range: crate::incremental::TimeRange,
#[case] expected_replace_range: crate::incremental::TimeRange,
) {
let pipeline = build_pipeline(pipeline_builder);
let partition_unit_opt = if partition_unit == TruncUnit::Second {
None
} else {
Some((partition_unit, 1))
};
let result = compute_incremental_ranges_for_pipeline(
&pipeline,
stale_ranges_map,
×tamp_field(),
partition_unit_opt,
false,
);
assert!(result.is_ok(), "Expected success, got error: {:?}", result);
let ranges = result.unwrap();
assert_eq!(
ranges.query_range, expected_query_range,
"Query range mismatch"
);
assert_eq!(
ranges.replace_range, expected_replace_range,
"Replace range mismatch"
);
}