timeseries_table_datafusion/
test_utils.rs

1use chrono::{DateTime, Utc};
2use datafusion::logical_expr::{Expr, Operator};
3
4use crate::ts_table_provider::{IntervalTruth, ParsedTz, TimePred, UnifiedInterval};
5
6/// Simplified, public view of a compiled time predicate for tests.
7#[derive(Debug, PartialEq)]
8pub enum CompiledTimePred {
9    /// Predicate always true.
10    True,
11    /// Predicate always false.
12    False,
13    /// Predicate does not reference the timestamp column.
14    NonTime,
15    /// Predicate references ts but cannot be evaluated for pruning.
16    Unknown,
17    /// A direct comparison against a timestamp literal.
18    Cmp {
19        /// Comparison operator.
20        op: Operator,
21        /// Literal timestamp in UTC.
22        ts: DateTime<Utc>,
23    },
24    /// Any other form (e.g., And/Or/Not trees).
25    Other,
26}
27
28/// Result of evaluating a compiled time predicate against a segment.
29#[derive(Debug, PartialEq, Eq)]
30pub enum CompiledIntervalTruth {
31    /// Predicate always true over the segment interval.
32    AlwaysTrue,
33    /// Predicate always false over the segment interval.
34    AlwaysFalse,
35    /// Predicate may be true for some values in the segment interval.
36    MaybeTrue,
37}
38
39/// Compile a time predicate and map it into `CompiledTimePred`.
40pub fn compile_time_pred_for_tests(expr: &Expr, ts_col: &str) -> CompiledTimePred {
41    compile_time_pred_for_tests_with_tz(expr, ts_col, None)
42}
43
44/// Compile a time predicate with an explicit timezone.
45pub(crate) fn compile_time_pred_for_tests_with_tz(
46    expr: &Expr,
47    ts_col: &str,
48    tz: Option<&ParsedTz>,
49) -> CompiledTimePred {
50    match crate::ts_table_provider::compile_time_pred(expr, ts_col, tz) {
51        TimePred::True => CompiledTimePred::True,
52        TimePred::False => CompiledTimePred::False,
53        TimePred::NonTime => CompiledTimePred::NonTime,
54        TimePred::Unknown => CompiledTimePred::Unknown,
55        TimePred::Cmp { op, ts } => CompiledTimePred::Cmp { op, ts },
56        _ => CompiledTimePred::Other,
57    }
58}
59
60/// Evaluate a time predicate against a segment interval using production logic.
61pub fn eval_time_pred_on_segment_for_tests(
62    expr: &Expr,
63    ts_col: &str,
64    seg_min: DateTime<Utc>,
65    seg_max: DateTime<Utc>,
66) -> CompiledIntervalTruth {
67    eval_time_pred_on_segment_for_tests_with_tz(expr, ts_col, seg_min, seg_max, None)
68}
69
70/// Evaluate a time predicate against a segment interval with an explicit timezone.
71pub(crate) fn eval_time_pred_on_segment_for_tests_with_tz(
72    expr: &Expr,
73    ts_col: &str,
74    seg_min: DateTime<Utc>,
75    seg_max: DateTime<Utc>,
76    tz: Option<&ParsedTz>,
77) -> CompiledIntervalTruth {
78    let pred = crate::ts_table_provider::compile_time_pred(expr, ts_col, tz);
79    match crate::ts_table_provider::eval_time_pred_on_segment(&pred, seg_min, seg_max) {
80        IntervalTruth::AlwaysTrue => CompiledIntervalTruth::AlwaysTrue,
81        IntervalTruth::AlwaysFalse => CompiledIntervalTruth::AlwaysFalse,
82        IntervalTruth::MaybeTrue => CompiledIntervalTruth::MaybeTrue,
83    }
84}
85
86/// Evaluate a time predicate against a segment interval using UTC boundaries.
87pub fn eval_time_pred_on_segment_for_tests_utc(
88    expr: &Expr,
89    ts_col: &str,
90    seg_min: DateTime<Utc>,
91    seg_max: DateTime<Utc>,
92) -> CompiledIntervalTruth {
93    eval_time_pred_on_segment_for_tests_with_tz(
94        expr,
95        ts_col,
96        seg_min,
97        seg_max,
98        Some(&ParsedTz::Utc),
99    )
100}
101
102/// Interval value for tests, split into calendar and sub-day components.
103pub struct TestInterval {
104    /// Calendar months.
105    pub months: i32,
106    /// Calendar days.
107    pub days: i32,
108    /// Sub-day nanos.
109    pub nanos: i64,
110}
111
112/// Apply a test interval to a timestamp using the production logic.
113pub fn add_interval_for_tests(
114    dt: DateTime<Utc>,
115    interval: TestInterval,
116    sign: i32,
117) -> Option<DateTime<Utc>> {
118    crate::ts_table_provider::add_interval(
119        dt,
120        UnifiedInterval {
121            months: interval.months,
122            days: interval.days,
123            nanos: interval.nanos,
124        },
125        sign,
126    )
127}
timeseries_table_datafusion/test_utils.rs

timeseries_table_datafusion/
test_utils.rs
