common_recordbatch/filter/
timestamp.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Timestamp-unit adaptation for predicates evaluated against historical SSTs.
16
17use common_time::timestamp::div_mod_units;
18use datafusion::logical_expr::Operator;
19use datafusion_common::ScalarValue;
20use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar};
21use datatypes::arrow::datatypes::TimeUnit;
22
23use crate::filter::predicate::SimplePredicate;
24
25/// Builds a tz-naive timestamp literal for `value` in `unit`.
26pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue {
27    match unit {
28        TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
29        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
30        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
31        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
32    }
33}
34
35impl SimplePredicate {
36    pub(crate) fn cast_timestamp_unit(&self, target_unit: TimeUnit) -> Option<Self> {
37        // Constants represent scan outcomes here. As in the public Matched/Pruned
38        // result, UNKNOWN can be collapsed to false under these AND/OR-only trees.
39        match self {
40            Self::And(left, right) | Self::Or(left, right) => {
41                let left = left.cast_timestamp_unit(target_unit)?;
42                let right = right.cast_timestamp_unit(target_unit)?;
43                let and = matches!(self, Self::And(..));
44                Some(match (left, right) {
45                    (Self::Constant(Some(false)), _) | (_, Self::Constant(Some(false))) if and => {
46                        Self::Constant(Some(false))
47                    }
48                    (Self::Constant(Some(true)), _) | (_, Self::Constant(Some(true))) if !and => {
49                        Self::Constant(Some(true))
50                    }
51                    (Self::Constant(Some(true)), other) | (other, Self::Constant(Some(true)))
52                        if and =>
53                    {
54                        other
55                    }
56                    (Self::Constant(Some(false)), other) | (other, Self::Constant(Some(false)))
57                        if !and =>
58                    {
59                        other
60                    }
61                    (left, right) => {
62                        if and {
63                            Self::And(Box::new(left), Box::new(right))
64                        } else {
65                            Self::Or(Box::new(left), Box::new(right))
66                        }
67                    }
68                })
69            }
70            Self::IsNull { .. } => Some(self.clone()),
71            Self::Constant(value) => Some(Self::Constant(Some(value.unwrap_or(false)))),
72            Self::OrEqChain { literals } => {
73                let mut converted = Vec::with_capacity(literals.len());
74                for literal in literals {
75                    let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?;
76                    let (value, unit) = timestamp_scalar_parts(&scalar)?;
77                    let Some(value) = value else {
78                        continue;
79                    };
80                    let cast = div_mod_units(value, unit.into(), target_unit.into())?;
81                    if cast.remainder == 0 {
82                        converted.push(timestamp_scalar(cast.quotient, target_unit)?);
83                    }
84                }
85                if converted.is_empty() {
86                    return Some(Self::Constant(Some(false)));
87                }
88                Some(Self::OrEqChain {
89                    literals: converted,
90                })
91            }
92            Self::Comparison { literal, op, .. } => {
93                let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?;
94                let (value, unit) = timestamp_scalar_parts(&scalar)?;
95                let Some(value) = value else {
96                    return Some(Self::Constant(Some(false)));
97                };
98                if unit == target_unit {
99                    return Some(self.clone());
100                }
101                let cast = div_mod_units(value, unit.into(), target_unit.into())?;
102                let divisible = cast.remainder == 0;
103                let literal = timestamp_scalar(cast.quotient, target_unit)?;
104                let filter = |op| Self::Comparison {
105                    literal: literal.clone(),
106                    op,
107                    regex: None,
108                    regex_negative: false,
109                };
110                Some(match op {
111                    Operator::Eq if divisible => filter(Operator::Eq),
112                    Operator::Eq => Self::Constant(Some(false)),
113                    Operator::NotEq if divisible => filter(Operator::NotEq),
114                    Operator::NotEq => Self::Constant(Some(true)),
115                    // The quotient is floor(L), including negative timestamps.
116                    // v > L is v > quotient whether or not L is representable.
117                    Operator::Gt => filter(Operator::Gt),
118                    Operator::GtEq if divisible => filter(Operator::GtEq),
119                    // Between target values, v >= L is v > quotient.
120                    Operator::GtEq => filter(Operator::Gt),
121                    Operator::Lt if divisible => filter(Operator::Lt),
122                    // Between target values, v < L is v <= quotient.
123                    Operator::Lt => filter(Operator::LtEq),
124                    Operator::LtEq => filter(Operator::LtEq),
125                    _ => return None,
126                })
127            }
128        }
129    }
130}
131
132/// Extracts the value and unit from a tz-naive timestamp scalar.
133fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option<i64>, TimeUnit)> {
134    let (value, unit, timezone) = match scalar {
135        ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz),
136        ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz),
137        ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz),
138        ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz),
139        _ => return None,
140    };
141    // A timezone-aware literal doesn't compare against a tz-naive column;
142    // leave it to the caller instead of guessing the intended semantics.
143    (timezone.is_none()).then_some((value, unit))
144}
145
146/// Builds a one-element scalar array holding `value` in `unit`.
147fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option<Scalar<ArrayRef>> {
148    timestamp_scalar_value(value, unit).to_scalar().ok()
149}
150
151#[cfg(test)]
152mod tests {
153    use std::sync::Arc;
154
155    use datafusion::execution::context::ExecutionProps;
156    use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
157    use datafusion::logical_expr::{Expr, col, lit};
158    use datafusion::physical_expr::create_physical_expr;
159    use datafusion_common::arrow::array::ArrayRef;
160    use datafusion_common::arrow::buffer::BooleanBuffer;
161    use datafusion_common::cast::as_boolean_array;
162    use datafusion_common::{DFSchema, ScalarValue};
163    use datatypes::arrow::array::{
164        RecordBatch, TimestampMillisecondArray, TimestampNanosecondArray,
165    };
166    use datatypes::arrow::datatypes::{Field, Schema};
167    use datatypes::data_type::ConcreteDataType;
168
169    use crate::filter::{SimpleFilterEvaluator, TimestampUnitCast, boolean_array_to_scan_mask};
170
171    #[test]
172    fn compound_timestamp_filters_preserve_unit_conversion() {
173        use datatypes::arrow::array::TimestampMicrosecondArray;
174
175        let null = || lit(ScalarValue::TimestampMicrosecond(None, None));
176        let predicates = [
177            col("ts")
178                .gt(lit(ts_us(-1500)))
179                .and(col("ts").lt(lit(ts_us(500)))),
180            col("ts")
181                .eq(lit(ts_us(500)))
182                .or(col("ts").gt_eq(lit(ts_us(1500)))),
183            col("ts")
184                .is_not_null()
185                .and(col("ts").not_eq(lit(ts_us(500)))),
186            col("ts").eq(null()).or(col("ts").eq(lit(ts_us(-1000)))),
187            col("ts")
188                .eq(lit(ts_us(500)))
189                .and(col("ts").gt(lit(ts_us(0)))),
190            col("ts")
191                .not_eq(lit(ts_us(500)))
192                .or(col("ts").lt(lit(ts_us(0)))),
193        ];
194        let native: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![
195            -2000, -1000, 0, 1000, 2000,
196        ]));
197        let old: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![-2, -1, 0, 1, 2]));
198        let schema = Arc::new(Schema::new(vec![Field::new(
199            "ts",
200            native.data_type().clone(),
201            false,
202        )]));
203        let batch = RecordBatch::try_new(schema.clone(), vec![native]).unwrap();
204        for expr in predicates {
205            let physical = create_physical_expr(
206                &expr,
207                &DFSchema::try_from(schema.as_ref().clone()).unwrap(),
208                &ExecutionProps::new(),
209                &PhysicalPlanningContext::default(),
210            )
211            .unwrap();
212            let result = physical
213                .evaluate(&batch)
214                .unwrap()
215                .into_array(batch.num_rows())
216                .unwrap();
217            let expected = boolean_array_to_scan_mask(as_boolean_array(&result).unwrap())
218                .values()
219                .clone();
220            let actual = match cast_to_ms(&expr).unwrap() {
221                TimestampUnitCast::Matched => BooleanBuffer::new_set(old.len()),
222                TimestampUnitCast::Pruned => BooleanBuffer::new_unset(old.len()),
223                TimestampUnitCast::Filter(filter) => filter.evaluate_array(&old).unwrap(),
224            };
225            assert_eq!(actual, expected, "{expr}");
226        }
227    }
228
229    fn ts_us(v: i64) -> ScalarValue {
230        ScalarValue::TimestampMicrosecond(Some(v), None)
231    }
232
233    fn cast_to_ms(expr: &Expr) -> Option<TimestampUnitCast> {
234        SimpleFilterEvaluator::try_new(expr)
235            .unwrap()
236            .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype())
237    }
238
239    /// Evaluates a cast `Filter` outcome against millisecond values and
240    /// returns the mask.
241    fn eval_ms_mask(cast: Option<TimestampUnitCast>, values: &[i64]) -> Vec<bool> {
242        let filter = match cast.expect("cast must apply") {
243            TimestampUnitCast::Filter(filter) => filter,
244            other => panic!("expected Filter outcome, got {other:?}"),
245        };
246        let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef;
247        filter.evaluate_array(&array).unwrap().iter().collect()
248    }
249
250    #[test]
251    fn cast_timestamp_unit_converts_representable_literal() {
252        // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms.
253        assert_eq!(
254            vec![false, true, false],
255            eval_ms_mask(
256                cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))),
257                &[6_999, 7_000, 7_001]
258            )
259        );
260        // != complements =.
261        assert_eq!(
262            vec![true, false, true],
263            eval_ms_mask(
264                cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))),
265                &[6_999, 7_000, 7_001]
266            )
267        );
268    }
269
270    #[test]
271    fn cast_timestamp_unit_prunes_non_representable_equality() {
272        // 7_000_500us is not a whole millisecond: no ms row can equal it,
273        // and every ms row satisfies `!=`.
274        assert!(matches!(
275            cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))),
276            Some(TimestampUnitCast::Pruned)
277        ));
278        assert!(matches!(
279            cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))),
280            Some(TimestampUnitCast::Matched)
281        ));
282    }
283
284    #[test]
285    fn cast_timestamp_unit_strengthens_inequalities() {
286        // 2_500_500us is strictly between 2500ms and 2501ms: the cast must
287        // not round the literal to a boundary the original predicate
288        // excludes (>= 2500ms would wrongly match 2500ms).
289        let values = vec![2_499, 2_500, 2_501];
290        assert_eq!(
291            vec![false, false, true],
292            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values)
293        );
294        assert_eq!(
295            vec![false, false, true],
296            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values)
297        );
298        assert_eq!(
299            vec![true, true, false],
300            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values)
301        );
302        assert_eq!(
303            vec![true, true, false],
304            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values)
305        );
306
307        // Representable boundary keeps the original operator.
308        let values = vec![2_499, 2_500, 2_501];
309        assert_eq!(
310            vec![false, false, true],
311            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values)
312        );
313        assert_eq!(
314            vec![false, true, true],
315            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values)
316        );
317        assert_eq!(
318            vec![true, false, false],
319            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values)
320        );
321        assert_eq!(
322            vec![true, true, false],
323            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values)
324        );
325    }
326
327    #[test]
328    fn cast_timestamp_unit_handles_negative_instants() {
329        // -2_500_500us floors to -2501ms with remainder 500us: only rows
330        // with instant >= -2.5005ms may match >= / >.
331        let values = vec![-2_502, -2_501, -2_500];
332        assert_eq!(
333            vec![false, false, true],
334            eval_ms_mask(
335                cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))),
336                &values
337            )
338        );
339        assert_eq!(
340            vec![false, false, true],
341            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values)
342        );
343        assert_eq!(
344            vec![true, true, false],
345            eval_ms_mask(
346                cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))),
347                &values
348            )
349        );
350        // Exactly representable negative literal: -2_500_000us == -2500ms.
351        assert_eq!(
352            vec![false, false, true],
353            eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values)
354        );
355    }
356
357    #[test]
358    fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() {
359        // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match.
360        let expr = col("ts")
361            .eq(lit(ts_us(7_000_500)))
362            .or(col("ts").eq(lit(ts_us(7_000_000))));
363        assert_eq!(
364            vec![false, true],
365            eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000])
366        );
367
368        // A chain where no literal is representable matches nothing.
369        let expr = col("ts")
370            .eq(lit(ts_us(7_000_500)))
371            .or(col("ts").eq(lit(ts_us(6_000_500))));
372        assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned)));
373    }
374
375    #[test]
376    fn cast_timestamp_unit_null_literal_prunes() {
377        // A NULL literal never compares true.
378        assert!(matches!(
379            cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))),
380            Some(TimestampUnitCast::Pruned)
381        ));
382    }
383
384    #[test]
385    fn cast_timestamp_unit_same_unit_returns_filter_directly() {
386        // A literal already in the target unit is returned unchanged.
387        let expr = col("ts").eq(lit(ts_us(7_000_000)));
388        let filter = SimpleFilterEvaluator::try_new(&expr).unwrap();
389        let cast = filter
390            .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype())
391            .unwrap();
392        let TimestampUnitCast::Filter(f) = cast else {
393            panic!("expected Filter, got {cast:?}")
394        };
395        assert!(filter.is_eq() && f.is_eq());
396        assert_eq!(filter.column_name(), f.column_name());
397        assert_eq!(filter.literal_value(), f.literal_value());
398    }
399
400    #[test]
401    fn cast_timestamp_unit_rejects_non_timestamps() {
402        // Non-timestamp target type.
403        assert!(
404            SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1))))
405                .unwrap()
406                .cast_timestamp_unit(&ConcreteDataType::int64_datatype())
407                .is_none()
408        );
409        // Non-timestamp literal against a timestamp target.
410        assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none());
411    }
412
413    #[test]
414    fn cast_timestamp_unit_to_finer_unit() {
415        // 5 seconds against a nanosecond column becomes 5_000_000_000ns.
416        let filter = SimpleFilterEvaluator::try_new(
417            &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))),
418        )
419        .unwrap()
420        .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype())
421        .and_then(|cast| match cast {
422            TimestampUnitCast::Filter(filter) => Some(filter),
423            _ => None,
424        })
425        .unwrap();
426        let array = Arc::new(TimestampNanosecondArray::from(vec![
427            4_999_999_999,
428            5_000_000_000,
429        ])) as ArrayRef;
430        assert_eq!(
431            vec![false, true],
432            filter
433                .evaluate_array(&array)
434                .unwrap()
435                .iter()
436                .collect::<Vec<_>>()
437        );
438    }
439}
common_recordbatch/filter/timestamp.rs

common_recordbatch/filter/
timestamp.rs
