query/optimizer/
const_normalization.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
15use std::sync::Arc;
16
17use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
18use datafusion::config::ConfigOptions;
19use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
20use datafusion_common::{DFSchemaRef, Result, ScalarValue};
21use datafusion_expr::expr::{Cast, InList, Like, TryCast};
22use datafusion_expr::{Between, BinaryExpr, Expr, ExprSchemable, LogicalPlan, Operator, lit};
23use datafusion_expr_common::casts::try_cast_literal_to_type;
24use datafusion_optimizer::analyzer::AnalyzerRule;
25
26/// ConstNormalizationRule rewrites castable constants against their
27/// non-constant comparison operand ahead of filter pushdown.
28#[derive(Debug)]
29pub struct ConstNormalizationRule;
30
31impl AnalyzerRule for ConstNormalizationRule {
32    fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result<LogicalPlan> {
33        plan.transform(|plan| match plan {
34            LogicalPlan::Filter(filter) => {
35                let schema = filter.input.schema().clone();
36                rewrite_plan_exprs(LogicalPlan::Filter(filter), schema)
37            }
38            LogicalPlan::TableScan(scan) => {
39                let schema = scan.projected_schema.clone();
40                rewrite_plan_exprs(LogicalPlan::TableScan(scan), schema)
41            }
42            _ => Ok(Transformed::no(plan)),
43        })
44        .map(|x| x.data)
45    }
46
47    fn name(&self) -> &str {
48        "ConstNormalizationRule"
49    }
50}
51
52fn rewrite_plan_exprs(plan: LogicalPlan, schema: DFSchemaRef) -> Result<Transformed<LogicalPlan>> {
53    let mut rewriter = ConstNormalizationRewriter { schema };
54    plan.map_expressions(|expr| expr.rewrite(&mut rewriter))
55}
56
57struct ConstNormalizationRewriter {
58    schema: DFSchemaRef,
59}
60
61impl TreeNodeRewriter for ConstNormalizationRewriter {
62    type Node = Expr;
63
64    fn f_down(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
65        let recursion = if matches!(
66            expr,
67            Expr::Exists(_) | Expr::InSubquery(_) | Expr::ScalarSubquery(_)
68        ) {
69            TreeNodeRecursion::Jump
70        } else {
71            TreeNodeRecursion::Continue
72        };
73
74        Ok(Transformed::new(expr, false, recursion))
75    }
76
77    fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
78        rewrite_expr_node(expr, &self.schema)
79    }
80}
81
82fn rewrite_expr_node(expr: Expr, schema: &DFSchemaRef) -> Result<Transformed<Expr>> {
83    // Borrow while checking eligibility: cloning unchanged subtrees here can overflow
84    // the stack even though DataFusion protects the surrounding rewrite traversal.
85    let rewritten = match &expr {
86        Expr::BinaryExpr(binary) => rewrite_binary_expr(binary, schema)?,
87        Expr::Between(between) => rewrite_between_expr(between, schema)?,
88        Expr::InList(in_list) => rewrite_in_list_expr(in_list, schema)?,
89        Expr::Like(like) => rewrite_like_expr(like, PatternMatchKind::Like, schema)?,
90        Expr::SimilarTo(like) => rewrite_like_expr(like, PatternMatchKind::SimilarTo, schema)?,
91        _ => None,
92    };
93    Ok(match rewritten {
94        Some(rewritten) => Transformed::yes(rewritten),
95        None => Transformed::no(expr),
96    })
97}
98
99fn rewrite_between_expr(between: &Between, schema: &DFSchemaRef) -> Result<Option<Expr>> {
100    let Between {
101        expr,
102        negated,
103        low,
104        high,
105    } = between;
106    let Some((target, constants)) =
107        extract_rewrite_operands(expr, [low.as_ref(), high.as_ref()], schema)?
108    else {
109        return Ok(None);
110    };
111
112    if let Some(mut constants) = target.normalize_constants(&constants) {
113        let high = constants
114            .pop()
115            .expect("between normalization expects high constant");
116        let low = constants
117            .pop()
118            .expect("between normalization expects low constant");
119        return Ok(Some(Expr::Between(Between {
120            expr: Box::new(target.expr.clone()),
121            negated: *negated,
122            low: Box::new(lit(low)),
123            high: Box::new(lit(high)),
124        })));
125    }
126
127    Ok((!negated)
128        .then(|| target.normalize_timestamp_between(&constants[0], &constants[1]))
129        .flatten())
130}
131
132fn rewrite_in_list_expr(in_list: &InList, schema: &DFSchemaRef) -> Result<Option<Expr>> {
133    let InList {
134        expr,
135        list,
136        negated,
137    } = in_list;
138    let Some((target, constants)) = extract_rewrite_operands(expr, list, schema)? else {
139        return Ok(None);
140    };
141
142    Ok(target.normalize_constants(&constants).map(|constants| {
143        target
144            .expr
145            .clone()
146            .in_list(constants.into_iter().map(lit).collect(), *negated)
147    }))
148}
149
150fn rewrite_like_expr(
151    like: &Like,
152    kind: PatternMatchKind,
153    schema: &DFSchemaRef,
154) -> Result<Option<Expr>> {
155    let Like {
156        negated,
157        expr,
158        pattern,
159        escape_char,
160        case_insensitive,
161    } = like;
162    let Some((target, constants)) = extract_rewrite_operands(expr, [pattern.as_ref()], schema)?
163    else {
164        return Ok(None);
165    };
166    let Some(mut constants) = target.normalize_constants(&constants) else {
167        return Ok(None);
168    };
169
170    let pattern = lit(constants
171        .pop()
172        .expect("pattern normalization expects one constant"));
173    let like = Like::new(
174        *negated,
175        Box::new(target.expr.clone()),
176        Box::new(pattern),
177        *escape_char,
178        *case_insensitive,
179    );
180    let rewritten = match kind {
181        PatternMatchKind::Like => Expr::Like(like),
182        PatternMatchKind::SimilarTo => Expr::SimilarTo(like),
183    };
184    Ok(Some(rewritten))
185}
186
187fn rewrite_binary_expr(binary: &BinaryExpr, schema: &DFSchemaRef) -> Result<Option<Expr>> {
188    if let Some(expr) = rewrite_dictionary_string_regex(binary, schema)? {
189        return Ok(Some(expr));
190    }
191
192    if !binary.op.supports_propagation() {
193        return Ok(None);
194    }
195
196    let BinaryExpr { left, op, right } = binary;
197    if let Some(expr) = rewrite_binary_side(left, *op, right, schema)? {
198        return Ok(Some(expr));
199    }
200
201    let Some(swapped_op) = op.swap() else {
202        return Ok(None);
203    };
204
205    rewrite_binary_side(right, swapped_op, left, schema)
206}
207
208/// Removes the string coercion/schema-reconciliation cast present before physical regex planning.
209///
210/// Keeping the dictionary input lets DataFusion's physical regex kernel evaluate scalar patterns
211/// against dictionary values instead of materializing the string column first.
212fn rewrite_dictionary_string_regex(
213    binary: &BinaryExpr,
214    schema: &DFSchemaRef,
215) -> Result<Option<Expr>> {
216    let BinaryExpr { left, op, right } = binary;
217    if !matches!(
218        op,
219        Operator::RegexMatch
220            | Operator::RegexIMatch
221            | Operator::RegexNotMatch
222            | Operator::RegexNotIMatch
223    ) || !matches!(right.as_literal(), Some(ScalarValue::Utf8(Some(_))))
224    {
225        return Ok(None);
226    }
227
228    let Some((CastInputKind::Cast, source, DataType::Utf8)) = extract_cast_input(left) else {
229        return Ok(None);
230    };
231    if !matches!(source, Expr::Column(_))
232        || !matches!(
233            source.get_type(schema)?,
234            DataType::Dictionary(key_type, value_type)
235                if key_type.as_ref() == &DataType::UInt32 && value_type.as_ref() == &DataType::Utf8
236        )
237    {
238        return Ok(None);
239    }
240
241    Ok(Some(Expr::BinaryExpr(BinaryExpr {
242        left: Box::new(source.clone()),
243        op: *op,
244        right: right.clone(),
245    })))
246}
247
248fn rewrite_binary_side(
249    target_expr: &Expr,
250    op: Operator,
251    constant_expr: &Expr,
252    schema: &DFSchemaRef,
253) -> Result<Option<Expr>> {
254    let Some((target, constants)) = extract_rewrite_operands(target_expr, [constant_expr], schema)?
255    else {
256        return Ok(None);
257    };
258
259    if let Some(mut constants) = target.normalize_constants(&constants) {
260        let constant = constants
261            .pop()
262            .expect("binary normalization expects one constant");
263        return Ok(Some(Expr::BinaryExpr(BinaryExpr {
264            left: Box::new(target.expr.clone()),
265            op,
266            right: Box::new(lit(constant)),
267        })));
268    }
269
270    Ok(target.normalize_timestamp_binary(op, &constants[0]))
271}
272
273fn extract_rewrite_operands<'a, 'b>(
274    target_expr: &'a Expr,
275    constant_exprs: impl IntoIterator<Item = &'b Expr>,
276    schema: &DFSchemaRef,
277) -> Result<Option<(NormalizationTarget<'a>, Vec<ScalarValue>)>> {
278    let Some(target) = extract_normalization_target(target_expr, schema)? else {
279        return Ok(None);
280    };
281
282    extract_constant_scalars(constant_exprs)
283        .map(|constants| constants.map(|constants| (target, constants)))
284}
285
286#[derive(Clone)]
287struct NormalizationTarget<'a> {
288    expr: &'a Expr,
289    data_type: DataType,
290    kind: NormalizationKind,
291}
292
293#[derive(Clone)]
294enum NormalizationKind {
295    /// The cast preserves every source value exactly, so literals can be cast directly.
296    Lossless,
297    /// The cast drops timestamp precision and must widen predicate bounds to preserve semantics.
298    TimestampDowncast {
299        source_unit: ArrowTimeUnit,
300        target_unit: ArrowTimeUnit,
301        timezone: Option<Arc<str>>,
302    },
303}
304
305impl NormalizationTarget<'_> {
306    /// Normalizes constants for rewrites that can preserve the original predicate with a direct
307    /// literal cast. Timestamp precision-changing casts are handled by timestamp-specific helpers.
308    fn normalize_constants(&self, constants: &[ScalarValue]) -> Option<Vec<ScalarValue>> {
309        constants
310            .iter()
311            .map(|constant| self.normalize_constant(constant))
312            .collect()
313    }
314
315    fn normalize_constant(&self, constant: &ScalarValue) -> Option<ScalarValue> {
316        match self.kind {
317            NormalizationKind::TimestampDowncast { .. } => None,
318            NormalizationKind::Lossless => try_cast_literal_to_type(constant, &self.data_type),
319        }
320    }
321
322    /// Rewrites predicates over timestamp downcasts into source-side half-open bounds.
323    fn normalize_timestamp_binary(&self, op: Operator, constant: &ScalarValue) -> Option<Expr> {
324        let NormalizationKind::TimestampDowncast {
325            source_unit,
326            target_unit,
327            timezone,
328        } = &self.kind
329        else {
330            return None;
331        };
332
333        let constant = constant
334            .cast_to(&DataType::Timestamp(*target_unit, timezone.clone()))
335            .ok()?;
336        let value = timestamp_scalar_value(&constant)?;
337        let bound = match op {
338            Operator::GtEq => lower_bound_for_ge(value, *source_unit, *target_unit)?,
339            Operator::Gt => lower_bound_for_ge(value.checked_add(1)?, *source_unit, *target_unit)?,
340            Operator::Lt => lower_bound_for_ge(value, *source_unit, *target_unit)?,
341            Operator::LtEq => {
342                lower_bound_for_ge(value.checked_add(1)?, *source_unit, *target_unit)?
343            }
344            _ => return None,
345        };
346
347        let normalized_op = match op {
348            Operator::GtEq | Operator::Gt => Operator::GtEq,
349            Operator::Lt | Operator::LtEq => Operator::Lt,
350            _ => return None,
351        };
352
353        Some(match normalized_op {
354            Operator::GtEq => self.expr.clone().gt_eq(lit(timestamp_scalar(
355                *source_unit,
356                timezone.clone(),
357                bound,
358            ))),
359            Operator::Lt => {
360                self.expr
361                    .clone()
362                    .lt(lit(timestamp_scalar(*source_unit, timezone.clone(), bound)))
363            }
364            _ => unreachable!("timestamp normalization only rewrites to >= or <"),
365        })
366    }
367
368    /// Rewrites `BETWEEN` over timestamp downcasts into an inclusive lower bound and exclusive
369    /// upper bound over the source timestamp unit.
370    fn normalize_timestamp_between(&self, low: &ScalarValue, high: &ScalarValue) -> Option<Expr> {
371        let NormalizationKind::TimestampDowncast {
372            source_unit,
373            target_unit,
374            timezone,
375        } = &self.kind
376        else {
377            return None;
378        };
379
380        let target_type = DataType::Timestamp(*target_unit, timezone.clone());
381        let low = low.cast_to(&target_type).ok()?;
382        let high = high.cast_to(&target_type).ok()?;
383        let low = timestamp_scalar_value(&low)?;
384        let high = timestamp_scalar_value(&high)?;
385
386        let lower = lower_bound_for_ge(low, *source_unit, *target_unit)?;
387        let upper = lower_bound_for_ge(high.checked_add(1)?, *source_unit, *target_unit)?;
388
389        Some(
390            self.expr
391                .clone()
392                .gt_eq(lit(timestamp_scalar(*source_unit, timezone.clone(), lower)))
393                .and(self.expr.clone().lt(lit(timestamp_scalar(
394                    *source_unit,
395                    timezone.clone(),
396                    upper,
397                )))),
398        )
399    }
400}
401
402/// Returns the non-constant side we should normalize against.
403///
404/// Plain expressions normalize literals to their own type. Cast expressions only participate when
405/// the cast is lossless or when timestamp downcasts can be rewritten as wider source-side bounds.
406fn extract_normalization_target<'a>(
407    expr: &'a Expr,
408    schema: &DFSchemaRef,
409) -> Result<Option<NormalizationTarget<'a>>> {
410    if extract_constant_scalar(expr)?.is_some() {
411        return Ok(None);
412    }
413
414    let Some((_, source_expr, target_type)) = extract_cast_input(expr) else {
415        return Ok(Some(NormalizationTarget {
416            expr,
417            data_type: expr.get_type(schema)?,
418            kind: NormalizationKind::Lossless,
419        }));
420    };
421
422    let data_type = source_expr.get_type(schema)?;
423    let Some(kind) = classify_normalization_kind(&data_type, target_type) else {
424        return Ok(None);
425    };
426
427    Ok(Some(NormalizationTarget {
428        expr: source_expr,
429        data_type,
430        kind,
431    }))
432}
433
434fn classify_normalization_kind(
435    source_type: &DataType,
436    target_type: &DataType,
437) -> Option<NormalizationKind> {
438    // Timestamp casts that change precision need boundary-aware rewrites. A finer target literal
439    // may not map exactly back to the coarser source unit, so the generic lossless path is only
440    // safe for timestamp casts that keep the same unit.
441    if is_lossless_cast(source_type, target_type) {
442        return Some(NormalizationKind::Lossless);
443    }
444
445    match (source_type, target_type) {
446        (
447            DataType::Timestamp(source_unit, source_tz),
448            DataType::Timestamp(target_unit, target_tz),
449        ) if source_tz == target_tz
450            && time_unit_rank(*source_unit) > time_unit_rank(*target_unit) =>
451        {
452            Some(NormalizationKind::TimestampDowncast {
453                source_unit: *source_unit,
454                target_unit: *target_unit,
455                timezone: source_tz.clone(),
456            })
457        }
458        _ => None,
459    }
460}
461
462/// Returns whether every value of `source_type` is representable in `target_type`.
463fn is_lossless_cast(source_type: &DataType, target_type: &DataType) -> bool {
464    match (source_type, target_type) {
465        (DataType::Int8, DataType::Int16 | DataType::Int32 | DataType::Int64)
466        | (DataType::Int16, DataType::Int32 | DataType::Int64)
467        | (DataType::Int32, DataType::Int64)
468        | (DataType::UInt8, DataType::UInt16 | DataType::UInt32 | DataType::UInt64)
469        | (DataType::UInt8, DataType::Int16 | DataType::Int32 | DataType::Int64)
470        | (DataType::UInt16, DataType::UInt32 | DataType::UInt64)
471        | (DataType::UInt16, DataType::Int32 | DataType::Int64)
472        | (DataType::UInt32, DataType::UInt64 | DataType::Int64)
473        | (DataType::Utf8, DataType::Utf8View | DataType::LargeUtf8) => true,
474        (
475            DataType::Timestamp(source_unit, source_tz),
476            DataType::Timestamp(target_unit, target_tz),
477        ) => source_tz == target_tz && source_unit == target_unit,
478        _ => false,
479    }
480}
481
482#[derive(Clone, Copy)]
483enum PatternMatchKind {
484    Like,
485    SimilarTo,
486}
487
488fn extract_constant_scalars<'a>(
489    exprs: impl IntoIterator<Item = &'a Expr>,
490) -> Result<Option<Vec<ScalarValue>>> {
491    let exprs = exprs.into_iter();
492    let mut values = Vec::with_capacity(exprs.size_hint().0);
493    for expr in exprs {
494        let Some(value) = extract_constant_scalar(expr)? else {
495            return Ok(None);
496        };
497        values.push(value);
498    }
499
500    Ok(Some(values))
501}
502
503/// Extracts a literal scalar from an expression, folding constant `CAST` and `TRY_CAST` nodes.
504fn extract_constant_scalar(expr: &Expr) -> Result<Option<ScalarValue>> {
505    if let Some(value) = expr.as_literal() {
506        return Ok(Some(value.clone()));
507    }
508
509    let Some((kind, expr, data_type)) = extract_cast_input(expr) else {
510        return Ok(None);
511    };
512
513    match kind {
514        CastInputKind::Cast => extract_constant_scalar(expr)?
515            .map(|value| value.cast_to(data_type))
516            .transpose(),
517        CastInputKind::TryCast => {
518            Ok(extract_constant_scalar(expr)?.and_then(|value| value.cast_to(data_type).ok()))
519        }
520    }
521}
522
523#[derive(Clone, Copy)]
524enum CastInputKind {
525    Cast,
526    TryCast,
527}
528
529/// Returns the input expression and target type for `CAST` and `TRY_CAST` expressions.
530fn extract_cast_input(expr: &Expr) -> Option<(CastInputKind, &Expr, &DataType)> {
531    match expr {
532        Expr::Cast(Cast { expr, field }) => {
533            Some((CastInputKind::Cast, expr.as_ref(), field.data_type()))
534        }
535        Expr::TryCast(TryCast { expr, field }) => {
536            Some((CastInputKind::TryCast, expr.as_ref(), field.data_type()))
537        }
538        _ => None,
539    }
540}
541
542fn time_unit_rank(unit: ArrowTimeUnit) -> usize {
543    match unit {
544        ArrowTimeUnit::Second => 0,
545        ArrowTimeUnit::Millisecond => 1,
546        ArrowTimeUnit::Microsecond => 2,
547        ArrowTimeUnit::Nanosecond => 3,
548    }
549}
550
551fn time_unit_scale(unit: ArrowTimeUnit) -> i64 {
552    match unit {
553        ArrowTimeUnit::Second => 1,
554        ArrowTimeUnit::Millisecond => 1_000,
555        ArrowTimeUnit::Microsecond => 1_000_000,
556        ArrowTimeUnit::Nanosecond => 1_000_000_000,
557    }
558}
559
560/// Returns the number of source-unit ticks in one target-unit tick for finer-to-coarser casts.
561fn finer_to_coarser_ratio(source_unit: ArrowTimeUnit, target_unit: ArrowTimeUnit) -> Option<i64> {
562    let source_scale = time_unit_scale(source_unit);
563    let target_scale = time_unit_scale(target_unit);
564    (source_scale >= target_scale).then_some(source_scale / target_scale)
565}
566
567/// Returns the smallest source-unit timestamp whose downcast is greater than or equal to
568/// `target_value`.
569///
570/// DataFusion timestamp downcasts truncate toward zero. For non-positive buckets that means the
571/// bucket starts before `target_value * ratio`, so `<= x` can be rewritten as `< lower_bound(x+1)`
572/// without dropping rows near zero or across negative boundaries.
573fn lower_bound_for_ge(
574    target_value: i64,
575    source_unit: ArrowTimeUnit,
576    target_unit: ArrowTimeUnit,
577) -> Option<i64> {
578    let ratio = finer_to_coarser_ratio(source_unit, target_unit)?;
579    let base = target_value.checked_mul(ratio)?;
580    if target_value <= 0 {
581        base.checked_sub(ratio - 1)
582    } else {
583        Some(base)
584    }
585}
586
587fn timestamp_scalar_value(value: &ScalarValue) -> Option<i64> {
588    match value {
589        ScalarValue::TimestampSecond(Some(value), _)
590        | ScalarValue::TimestampMillisecond(Some(value), _)
591        | ScalarValue::TimestampMicrosecond(Some(value), _)
592        | ScalarValue::TimestampNanosecond(Some(value), _) => Some(*value),
593        _ => None,
594    }
595}
596
597fn timestamp_scalar(unit: ArrowTimeUnit, timezone: Option<Arc<str>>, value: i64) -> ScalarValue {
598    match unit {
599        ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), timezone),
600        ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), timezone),
601        ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), timezone),
602        ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), timezone),
603    }
604}
605
606#[cfg(test)]
607mod tests {
608    use std::sync::Arc;
609
610    use arrow::array::{DictionaryArray, StringArray, UInt32Array};
611    use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
612    use async_trait::async_trait;
613    use common_time::Timestamp;
614    use common_time::range::TimestampRange;
615    use common_time::timestamp::TimeUnit;
616    use datafusion::catalog::Session;
617    use datafusion::config::ConfigOptions;
618    use datafusion::datasource::{MemTable, TableProvider, provider_as_source};
619    use datafusion::execution::SessionStateBuilder;
620    use datafusion::execution::context::SessionContext;
621    use datafusion::physical_plan::filter::FilterExec;
622    use datafusion::physical_plan::{ExecutionPlan, collect};
623    use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
624    use datafusion_common::arrow::datatypes::Field;
625    use datafusion_common::{DFSchema, ScalarValue, ToDFSchema};
626    use datafusion_expr::expr::{Between, BinaryExpr, Like};
627    use datafusion_expr::expr_fn::{cast, col, try_cast};
628    use datafusion_expr::{
629        Expr, LogicalPlan, LogicalPlanBuilder, Operator, TableProviderFilterPushDown, TableScan,
630        TableSource, TableType, lit,
631    };
632    use datafusion_optimizer::analyzer::AnalyzerRule;
633    use datafusion_optimizer::optimizer::{Optimizer, OptimizerContext};
634    use datafusion_optimizer::push_down_filter::PushDownFilter;
635    use datafusion_optimizer::simplify_expressions::SimplifyExpressions;
636    use table::predicate::build_time_range_predicate;
637
638    use super::{
639        ConstNormalizationRule, PatternMatchKind, lower_bound_for_ge,
640        rewrite_dictionary_string_regex, try_cast_literal_to_type,
641    };
642
643    #[test]
644    fn test_normalize_deep_predicate() {
645        std::thread::Builder::new()
646            .stack_size(2 * 1024 * 1024)
647            .spawn(|| {
648                let mut predicate = col("t.v").eq(lit(0_i64));
649                for value in 1..1024_i64 {
650                    predicate = predicate.or(col("t.v").eq(lit(value)));
651                }
652                let schema = test_schema(vec![Field::new("v", DataType::Int32, false)]);
653                // Construct the filter directly to isolate this rule from other planner passes.
654                let plan = analyze_plan(LogicalPlan::Filter(
655                    datafusion_expr::Filter::try_new(predicate, Arc::new(build_scan_plan(schema)))
656                        .unwrap(),
657                ));
658                let LogicalPlan::Filter(filter) = plan else {
659                    panic!("expected filter");
660                };
661                // Inspect and consume the chain without recursive formatting or equality.
662                let mut predicate = filter.predicate;
663                for value in (1..1024_i32).rev() {
664                    let Expr::BinaryExpr(BinaryExpr { left, op, right }) = predicate else {
665                        panic!("expected binary predicate");
666                    };
667                    assert_eq!(op, Operator::Or);
668                    assert_eq!(*right, col("t.v").eq(lit(value)));
669                    predicate = *left;
670                }
671                assert_eq!(predicate, col("t.v").eq(lit(0_i32)));
672            })
673            .unwrap()
674            .join()
675            .unwrap();
676    }
677
678    #[test]
679    fn test_normalize_direct_integer_cast_comparison() {
680        assert_filter_plan(
681            vec![Field::new("v", DataType::Int32, false)],
682            cast(col("v"), DataType::Int64).gt_eq(lit(42_i64)),
683            "Filter: t.v >= Int32(42)\n  TableScan: t",
684        );
685    }
686
687    #[test]
688    fn test_normalize_non_column_operand() {
689        assert_filter_plan(
690            vec![Field::new("v", DataType::Int32, false)],
691            cast(col("v") + lit(1_i32), DataType::Int64).gt_eq(lit(42_i64)),
692            "Filter: t.v + Int32(1) >= Int32(42)\n  TableScan: t",
693        );
694    }
695
696    #[test]
697    fn test_normalize_swapped_binary_comparison() {
698        assert_filter_plan(
699            vec![Field::new("v", DataType::Int16, false)],
700            lit(42_i64).lt_eq(cast(col("v"), DataType::Int64)),
701            "Filter: t.v >= Int16(42)\n  TableScan: t",
702        );
703    }
704
705    #[test]
706    fn test_normalize_try_cast_target() {
707        assert_filter_plan(
708            vec![Field::new("v", DataType::Int16, false)],
709            try_cast(col("v"), DataType::Int64).gt_eq(lit(42_i64)),
710            "Filter: t.v >= Int16(42)\n  TableScan: t",
711        );
712    }
713
714    #[test]
715    fn test_normalize_casted_constants() {
716        let fields = vec![Field::new("v", DataType::Int16, false)];
717        let cases = [
718            (
719                col("v").gt_eq(cast(lit(42_i8), DataType::Int64)),
720                "Filter: t.v >= Int16(42)\n  TableScan: t",
721            ),
722            (
723                col("v").in_list(
724                    vec![
725                        cast(lit(1_i8), DataType::Int64),
726                        try_cast(lit(2_i8), DataType::Int64),
727                    ],
728                    false,
729                ),
730                "Filter: t.v IN ([Int16(1), Int16(2)])\n  TableScan: t",
731            ),
732        ];
733
734        for (predicate, expected) in cases {
735            assert_filter_plan(fields.clone(), predicate, expected);
736        }
737    }
738
739    #[test]
740    fn test_normalize_plain_integer_literals() {
741        let fields = vec![Field::new("v", DataType::Int16, false)];
742        let cases = [
743            (
744                col("v").gt_eq(lit(42_i64)),
745                "Filter: t.v >= Int16(42)\n  TableScan: t",
746            ),
747            (
748                col("v").in_list(vec![lit(1_i64), lit(2_i64)], false),
749                "Filter: t.v IN ([Int16(1), Int16(2)])\n  TableScan: t",
750            ),
751            (
752                col("v").between(lit(3_i64), lit(5_i64)),
753                "Filter: t.v BETWEEN Int16(3) AND Int16(5)\n  TableScan: t",
754            ),
755        ];
756
757        for (predicate, expected) in cases {
758            assert_filter_plan(fields.clone(), predicate, expected);
759        }
760    }
761
762    #[test]
763    fn test_normalize_unsigned_to_signed_literals() {
764        let cases = [
765            (
766                vec![Field::new("v", DataType::UInt8, false)],
767                cast(col("v"), DataType::Int16).lt_eq(lit(255_i16)),
768                "Filter: t.v <= UInt8(255)\n  TableScan: t",
769            ),
770            (
771                vec![Field::new("v", DataType::UInt16, false)],
772                cast(col("v"), DataType::Int32).gt_eq(lit(42_i32)),
773                "Filter: t.v >= UInt16(42)\n  TableScan: t",
774            ),
775            (
776                vec![Field::new("v", DataType::UInt32, false)],
777                cast(col("v"), DataType::Int64).between(lit(3_i64), lit(5_i64)),
778                "Filter: t.v BETWEEN UInt32(3) AND UInt32(5)\n  TableScan: t",
779            ),
780        ];
781
782        for (fields, predicate, expected) in cases {
783            assert_filter_plan(fields, predicate, expected);
784        }
785    }
786
787    #[test]
788    fn test_normalize_in_list_and_between() {
789        let fields = vec![Field::new("v", DataType::Int16, false)];
790        let cases = [
791            (
792                cast(col("v"), DataType::Int64).in_list(vec![lit(1_i64), lit(2_i64)], false),
793                "Filter: t.v IN ([Int16(1), Int16(2)])\n  TableScan: t",
794            ),
795            (
796                cast(col("v"), DataType::Int64).between(lit(3_i64), lit(5_i64)),
797                "Filter: t.v BETWEEN Int16(3) AND Int16(5)\n  TableScan: t",
798            ),
799        ];
800
801        for (predicate, expected) in cases {
802            assert_filter_plan(fields.clone(), predicate, expected);
803        }
804    }
805
806    #[test]
807    fn test_keep_non_lossless_literal_unchanged() {
808        assert_filter_plan(
809            vec![Field::new("v", DataType::Int16, false)],
810            col("v").gt_eq(lit(100_000_i64)),
811            "Filter: t.v >= Int64(100000)\n  TableScan: t",
812        );
813    }
814
815    #[test]
816    fn test_normalize_scan_filters() {
817        let scan = build_scan_plan(test_schema(vec![Field::new("v", DataType::Int16, false)]));
818        let LogicalPlan::TableScan(scan) = scan else {
819            panic!("expected table scan");
820        };
821        let plan = LogicalPlan::TableScan(TableScan {
822            filters: vec![cast(col("v"), DataType::Int64).gt_eq(lit(42_i64))],
823            ..scan
824        });
825
826        let analyzed = analyze_plan(plan);
827
828        assert_eq!(
829            vec![col("v").gt_eq(lit(42_i16))],
830            extract_scan_filters(&analyzed)
831        );
832    }
833
834    #[test]
835    fn test_normalize_negated_between() {
836        assert_filter_plan(
837            vec![Field::new("v", DataType::Int16, false)],
838            Expr::Between(Between {
839                expr: Box::new(cast(col("v"), DataType::Int64)),
840                negated: true,
841                low: Box::new(lit(3_i64)),
842                high: Box::new(lit(5_i64)),
843            }),
844            "Filter: t.v NOT BETWEEN Int16(3) AND Int16(5)\n  TableScan: t",
845        );
846    }
847
848    #[test]
849    fn test_normalize_like_literal() {
850        assert_pattern_match_plan(
851            PatternMatchKind::Like,
852            ScalarValue::LargeUtf8(Some("api%".to_string())),
853            "Filter: t.s LIKE Utf8(\"api%\")\n  TableScan: t",
854        );
855    }
856
857    #[test]
858    fn test_normalize_similar_to_literal() {
859        assert_pattern_match_plan(
860            PatternMatchKind::SimilarTo,
861            ScalarValue::LargeUtf8(Some("api.*".to_string())),
862            "Filter: t.s SIMILAR TO Utf8(\"api.*\")\n  TableScan: t",
863        );
864    }
865
866    #[tokio::test]
867    async fn test_dictionary_regex_filter_keeps_dictionary_input() {
868        let dictionary_type =
869            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
870        let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
871            "host",
872            dictionary_type.clone(),
873            true,
874        )]));
875        let host = DictionaryArray::new(
876            UInt32Array::from(vec![Some(0), Some(1), Some(2), None, Some(3)]),
877            Arc::new(StringArray::from(vec![
878                Some("api"),
879                Some("API"),
880                Some("db"),
881                None,
882            ])),
883        );
884        let batch = datafusion::arrow::record_batch::RecordBatch::try_new(
885            schema.clone(),
886            vec![Arc::new(host)],
887        )
888        .unwrap();
889
890        for (op, expected_rows) in [
891            (Operator::RegexMatch, 1),
892            (Operator::RegexIMatch, 2),
893            (Operator::RegexNotMatch, 2),
894            (Operator::RegexNotIMatch, 1),
895        ] {
896            let table = MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap();
897            let predicate = Expr::BinaryExpr(BinaryExpr {
898                // This string coercion/schema-reconciliation cast is present before physical
899                // regex planning and would bypass DataFusion's dictionary-aware scalar regex
900                // kernel.
901                left: Box::new(cast(col("host"), DataType::Utf8)),
902                op,
903                right: Box::new(lit("^api$")),
904            });
905            let plan = LogicalPlanBuilder::scan("t", provider_as_source(Arc::new(table)), None)
906                .unwrap()
907                .filter(predicate)
908                .unwrap()
909                .build()
910                .unwrap();
911            let analyzed = analyze_plan(plan);
912
913            let LogicalPlan::Filter(filter) = &analyzed else {
914                panic!("expected filter plan");
915            };
916            let Expr::BinaryExpr(BinaryExpr { left, .. }) = &filter.predicate else {
917                panic!("expected regex binary predicate");
918            };
919            assert!(matches!(left.as_ref(), Expr::Column(_)));
920
921            let session_state = SessionStateBuilder::new().with_default_features().build();
922            let physical_plan = DefaultPhysicalPlanner::default()
923                .create_physical_plan(&analyzed, &session_state)
924                .await
925                .unwrap();
926            let filter = physical_plan
927                .downcast_ref::<FilterExec>()
928                .expect("regex residual must remain a FilterExec");
929            assert!(matches!(
930                filter.schema().field(0).data_type(),
931                DataType::Dictionary(_, value_type) if value_type.as_ref() == &DataType::Utf8
932            ));
933            assert!(!format!("{:?}", filter.predicate()).contains("Cast"));
934
935            let batches = collect(physical_plan, SessionContext::new().task_ctx())
936                .await
937                .unwrap();
938            assert_eq!(
939                expected_rows,
940                batches.iter().map(|batch| batch.num_rows()).sum::<usize>()
941            );
942        }
943    }
944
945    #[test]
946    fn test_dictionary_regex_rewrite_requires_scalar_utf8_pattern() {
947        assert_filter_left_is_cast(
948            vec![
949                Field::new(
950                    "host",
951                    DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
952                    true,
953                ),
954                Field::new("pattern", DataType::Utf8, true),
955            ],
956            Expr::BinaryExpr(BinaryExpr {
957                left: Box::new(cast(col("host"), DataType::Utf8)),
958                op: Operator::RegexMatch,
959                right: Box::new(col("pattern")),
960            }),
961        );
962    }
963
964    #[test]
965    fn test_dictionary_regex_rewrite_excludes_non_regex_and_non_utf8_dictionary() {
966        assert_filter_left_is_cast(
967            vec![Field::new(
968                "host",
969                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
970                true,
971            )],
972            Expr::BinaryExpr(BinaryExpr {
973                left: Box::new(cast(col("host"), DataType::Utf8)),
974                op: Operator::Eq,
975                right: Box::new(lit("api")),
976            }),
977        );
978        assert_filter_left_is_cast(
979            vec![Field::new(
980                "host",
981                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::LargeUtf8)),
982                true,
983            )],
984            Expr::BinaryExpr(BinaryExpr {
985                left: Box::new(cast(col("host"), DataType::Utf8)),
986                op: Operator::RegexMatch,
987                right: Box::new(lit("^api$")),
988            }),
989        );
990    }
991
992    #[test]
993    fn test_dictionary_regex_rewrite_requires_exact_contract() {
994        let dictionary_utf8 = || {
995            vec![Field::new(
996                "host",
997                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
998                true,
999            )]
1000        };
1001
1002        for op in [
1003            Operator::RegexMatch,
1004            Operator::RegexIMatch,
1005            Operator::RegexNotMatch,
1006            Operator::RegexNotIMatch,
1007        ] {
1008            let rewritten = rewrite_dictionary_regex(
1009                dictionary_utf8(),
1010                cast(col("host"), DataType::Utf8),
1011                lit("^api$"),
1012                op,
1013            );
1014            assert!(matches!(
1015                rewritten,
1016                Some(Expr::BinaryExpr(BinaryExpr { left, .. })) if matches!(left.as_ref(), Expr::Column(_))
1017            ));
1018        }
1019
1020        for left in [
1021            try_cast(col("host"), DataType::Utf8),
1022            cast(cast(col("host"), DataType::Utf8), DataType::Utf8),
1023        ] {
1024            assert!(
1025                rewrite_dictionary_regex(
1026                    dictionary_utf8(),
1027                    left,
1028                    lit("^api$"),
1029                    Operator::RegexMatch,
1030                )
1031                .is_none()
1032            );
1033        }
1034        assert!(
1035            rewrite_dictionary_regex(
1036                dictionary_utf8(),
1037                cast(col("host"), DataType::Utf8),
1038                lit(ScalarValue::Utf8(None)),
1039                Operator::RegexMatch,
1040            )
1041            .is_none()
1042        );
1043        assert!(
1044            rewrite_dictionary_regex(
1045                vec![Field::new(
1046                    "host",
1047                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1048                    true,
1049                )],
1050                cast(col("host"), DataType::Utf8),
1051                lit("^api$"),
1052                Operator::RegexMatch,
1053            )
1054            .is_none()
1055        );
1056    }
1057
1058    #[test]
1059    fn test_normalize_direct_timestamp_filter() {
1060        assert_timestamp_pushdown(
1061            vec![
1062                Field::new(
1063                    "ts",
1064                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1065                    false,
1066                ),
1067                Field::new("tag", DataType::Utf8, true),
1068            ],
1069            ts_cast_to_ms()
1070                .gt_eq(ts_ms_literal(-299_999))
1071                .and(ts_cast_to_ms().lt_eq(ts_ms_literal(10_000)))
1072                .and(col("tag").eq(lit("api"))),
1073            "Filter: t.ts >= TimestampNanosecond(-299999999999, None) AND t.ts < TimestampNanosecond(10001000000, None) AND t.tag = Utf8(\"api\")\n  TableScan: t",
1074            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(-299999999999, None), t.ts < TimestampNanosecond(10001000000, None), t.tag = Utf8(\"api\")]",
1075            TimestampRange::new_inclusive(
1076                Some(Timestamp::new_nanosecond(-299_999_999_999)),
1077                Some(Timestamp::new_nanosecond(10_000_999_999)),
1078            ),
1079        );
1080    }
1081
1082    #[test]
1083    fn test_normalize_timestamp_between_filter() {
1084        assert_timestamp_pushdown(
1085            vec![Field::new(
1086                "ts",
1087                DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1088                false,
1089            )],
1090            ts_cast_to_ms().between(ts_ms_literal(-299_999), ts_ms_literal(10_000)),
1091            "Filter: t.ts >= TimestampNanosecond(-299999999999, None) AND t.ts < TimestampNanosecond(10001000000, None)\n  TableScan: t",
1092            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(-299999999999, None), t.ts < TimestampNanosecond(10001000000, None)]",
1093            TimestampRange::new_inclusive(
1094                Some(Timestamp::new_nanosecond(-299_999_999_999)),
1095                Some(Timestamp::new_nanosecond(10_000_999_999)),
1096            ),
1097        );
1098    }
1099
1100    #[test]
1101    fn test_normalize_strict_timestamp_filter() {
1102        assert_timestamp_pushdown(
1103            vec![Field::new(
1104                "ts",
1105                DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1106                false,
1107            )],
1108            ts_cast_to_ms()
1109                .gt(ts_ms_literal(10_000))
1110                .and(ts_cast_to_ms().lt(ts_ms_literal(20_000))),
1111            "Filter: t.ts >= TimestampNanosecond(10001000000, None) AND t.ts < TimestampNanosecond(20000000000, None)\n  TableScan: t",
1112            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(10001000000, None), t.ts < TimestampNanosecond(20000000000, None)]",
1113            TimestampRange::new_inclusive(
1114                Some(Timestamp::new_nanosecond(10_001_000_000)),
1115                Some(Timestamp::new_nanosecond(19_999_999_999)),
1116            ),
1117        );
1118    }
1119
1120    #[test]
1121    fn test_normalize_zero_boundary_timestamp_filter() {
1122        let fields = vec![Field::new(
1123            "ts",
1124            DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1125            false,
1126        )];
1127
1128        assert_timestamp_pushdown(
1129            fields.clone(),
1130            ts_cast_to_ms().gt_eq(ts_ms_literal(0)),
1131            "Filter: t.ts >= TimestampNanosecond(-999999, None)\n  TableScan: t",
1132            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(-999999, None)]",
1133            TimestampRange::from_start(Timestamp::new_nanosecond(-999_999)),
1134        );
1135
1136        assert_timestamp_pushdown(
1137            fields.clone(),
1138            ts_cast_to_ms().lt(ts_ms_literal(0)),
1139            "Filter: t.ts < TimestampNanosecond(-999999, None)\n  TableScan: t",
1140            "TableScan: t, full_filters=[t.ts < TimestampNanosecond(-999999, None)]",
1141            TimestampRange::until_end(Timestamp::new_nanosecond(-999_999), false),
1142        );
1143
1144        assert_timestamp_pushdown(
1145            fields,
1146            ts_cast_to_ms().between(ts_ms_literal(0), ts_ms_literal(0)),
1147            "Filter: t.ts >= TimestampNanosecond(-999999, None) AND t.ts < TimestampNanosecond(1000000, None)\n  TableScan: t",
1148            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(-999999, None), t.ts < TimestampNanosecond(1000000, None)]",
1149            TimestampRange::new_inclusive(
1150                Some(Timestamp::new_nanosecond(-999_999)),
1151                Some(Timestamp::new_nanosecond(999_999)),
1152            ),
1153        );
1154    }
1155
1156    #[test]
1157    fn test_timestamp_downcast_contract_matches_datafusion_casts() {
1158        let cases = [
1159            (-1_000_001, -1),
1160            (-1_000_000, -1),
1161            (-999_999, 0),
1162            (-1, 0),
1163            (0, 0),
1164            (999_999, 0),
1165            (1_000_000, 1),
1166        ];
1167
1168        for (source, expected) in cases {
1169            let casted = try_cast_literal_to_type(
1170                &ScalarValue::TimestampNanosecond(Some(source), None),
1171                &DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1172            )
1173            .unwrap();
1174            assert_eq!(
1175                ScalarValue::TimestampMillisecond(Some(expected), None),
1176                casted
1177            );
1178        }
1179
1180        assert_eq!(
1181            Some(-1_999_999),
1182            lower_bound_for_ge(-1, ArrowTimeUnit::Nanosecond, ArrowTimeUnit::Millisecond)
1183        );
1184        assert_eq!(
1185            Some(-999_999),
1186            lower_bound_for_ge(0, ArrowTimeUnit::Nanosecond, ArrowTimeUnit::Millisecond)
1187        );
1188        assert_eq!(
1189            Some(1_000_000),
1190            lower_bound_for_ge(1, ArrowTimeUnit::Nanosecond, ArrowTimeUnit::Millisecond)
1191        );
1192    }
1193
1194    #[test]
1195    fn test_normalize_plain_timestamp_literals() {
1196        assert_timestamp_pushdown(
1197            vec![Field::new(
1198                "ts",
1199                DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1200                false,
1201            )],
1202            col("ts")
1203                .gt_eq(ts_ms_literal(-299_999))
1204                .and(col("ts").lt_eq(ts_ms_literal(10_000))),
1205            "Filter: t.ts >= TimestampNanosecond(-299999000000, None) AND t.ts <= TimestampNanosecond(10000000000, None)\n  TableScan: t",
1206            "TableScan: t, full_filters=[t.ts >= TimestampNanosecond(-299999000000, None), t.ts <= TimestampNanosecond(10000000000, None)]",
1207            TimestampRange::new_inclusive(
1208                Some(Timestamp::new_nanosecond(-299_999_000_000)),
1209                Some(Timestamp::new_nanosecond(10_000_000_000)),
1210            ),
1211        );
1212    }
1213
1214    #[test]
1215    fn test_keep_timestamp_upcast_filter_unchanged() {
1216        assert_filter_plan(
1217            vec![Field::new(
1218                "ts",
1219                DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1220                false,
1221            )],
1222            cast(
1223                col("ts"),
1224                DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1225            )
1226            .gt_eq(lit(ScalarValue::TimestampNanosecond(Some(1), None))),
1227            "Filter: CAST(t.ts AS Timestamp(ns)) >= TimestampNanosecond(1, None)\n  TableScan: t",
1228        );
1229    }
1230
1231    #[test]
1232    fn test_const_normalization_vs_datafusion_cast_preimage_overlap() {
1233        struct Case {
1234            name: &'static str,
1235            fields: Vec<Field>,
1236            predicate: Expr,
1237            expected_greptime: &'static str,
1238            expected_datafusion: &'static str,
1239        }
1240
1241        let cases = [
1242            Case {
1243                name: "integer widening binary",
1244                fields: vec![Field::new("v", DataType::Int16, false)],
1245                predicate: cast(col("v"), DataType::Int64).gt_eq(lit(42_i64)),
1246                expected_greptime: "Filter: t.v >= Int16(42)\n  TableScan: t",
1247                expected_datafusion: "Filter: t.v >= Int16(42)\n  TableScan: t",
1248            },
1249            Case {
1250                name: "swapped integer comparison",
1251                fields: vec![Field::new("v", DataType::Int16, false)],
1252                predicate: lit(42_i64).lt_eq(cast(col("v"), DataType::Int64)),
1253                expected_greptime: "Filter: t.v >= Int16(42)\n  TableScan: t",
1254                expected_datafusion: "Filter: t.v >= Int16(42)\n  TableScan: t",
1255            },
1256            Case {
1257                name: "try_cast integer widening binary",
1258                fields: vec![Field::new("v", DataType::Int16, false)],
1259                predicate: try_cast(col("v"), DataType::Int64).gt_eq(lit(42_i64)),
1260                expected_greptime: "Filter: t.v >= Int16(42)\n  TableScan: t",
1261                expected_datafusion: "Filter: t.v >= Int16(42)\n  TableScan: t",
1262            },
1263            Case {
1264                name: "exact in-list",
1265                fields: vec![Field::new("v", DataType::Int16, false)],
1266                predicate: cast(col("v"), DataType::Int64)
1267                    .in_list(vec![lit(1_i64), lit(2_i64)], false),
1268                expected_greptime: "Filter: t.v IN ([Int16(1), Int16(2)])\n  TableScan: t",
1269                expected_datafusion: "Filter: t.v = Int16(1) OR t.v = Int16(2)\n  TableScan: t",
1270            },
1271            Case {
1272                name: "integer between",
1273                fields: vec![Field::new("v", DataType::Int16, false)],
1274                predicate: cast(col("v"), DataType::Int64).between(lit(3_i64), lit(5_i64)),
1275                expected_greptime: "Filter: t.v BETWEEN Int16(3) AND Int16(5)\n  TableScan: t",
1276                expected_datafusion: "Filter: t.v >= Int16(3) AND t.v <= Int16(5)\n  TableScan: t",
1277            },
1278            Case {
1279                name: "not between",
1280                fields: vec![Field::new("v", DataType::Int16, false)],
1281                predicate: Expr::Between(Between {
1282                    expr: Box::new(cast(col("v"), DataType::Int64)),
1283                    negated: true,
1284                    low: Box::new(lit(3_i64)),
1285                    high: Box::new(lit(5_i64)),
1286                }),
1287                expected_greptime: "Filter: t.v NOT BETWEEN Int16(3) AND Int16(5)\n  TableScan: t",
1288                expected_datafusion: "Filter: t.v < Int16(3) OR t.v > Int16(5)\n  TableScan: t",
1289            },
1290            Case {
1291                name: "plain literal",
1292                fields: vec![Field::new("v", DataType::Int16, false)],
1293                predicate: col("v").gt_eq(lit(42_i64)),
1294                expected_greptime: "Filter: t.v >= Int16(42)\n  TableScan: t",
1295                expected_datafusion: "Filter: t.v >= Int64(42)\n  TableScan: t",
1296            },
1297            Case {
1298                name: "casted constant",
1299                fields: vec![Field::new("v", DataType::Int16, false)],
1300                predicate: col("v").gt_eq(cast(lit(42_i8), DataType::Int64)),
1301                expected_greptime: "Filter: t.v >= Int16(42)\n  TableScan: t",
1302                expected_datafusion: "Filter: t.v >= Int64(42)\n  TableScan: t",
1303            },
1304            Case {
1305                name: "timestamp downcast equality",
1306                fields: vec![Field::new(
1307                    "ts_ns",
1308                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1309                    false,
1310                )],
1311                predicate: cast(
1312                    col("ts_ns"),
1313                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1314                )
1315                .eq(ts_ms_literal(5000)),
1316                expected_greptime: "Filter: CAST(t.ts_ns AS Timestamp(ms)) = TimestampMillisecond(5000, None)\n  TableScan: t",
1317                expected_datafusion: "Filter: t.ts_ns >= TimestampNanosecond(5000000000, None) AND t.ts_ns < TimestampNanosecond(5001000000, None)\n  TableScan: t",
1318            },
1319            Case {
1320                name: "timestamp widening exact",
1321                fields: vec![Field::new(
1322                    "ts_ms",
1323                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1324                    false,
1325                )],
1326                predicate: cast(
1327                    col("ts_ms"),
1328                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1329                )
1330                .eq(lit(ScalarValue::TimestampNanosecond(
1331                    Some(5_000_000_000),
1332                    None,
1333                ))),
1334                expected_greptime: "Filter: CAST(t.ts_ms AS Timestamp(ns)) = TimestampNanosecond(5000000000, None)\n  TableScan: t",
1335                expected_datafusion: "Filter: t.ts_ms = TimestampMillisecond(5000, None)\n  TableScan: t",
1336            },
1337            Case {
1338                name: "timestamp widening try_cast exact",
1339                fields: vec![Field::new(
1340                    "ts_ms",
1341                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1342                    false,
1343                )],
1344                predicate: try_cast(
1345                    col("ts_ms"),
1346                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1347                )
1348                .eq(lit(ScalarValue::TimestampNanosecond(
1349                    Some(5_000_000_000),
1350                    None,
1351                ))),
1352                expected_greptime: "Filter: TRY_CAST(t.ts_ms AS Timestamp(ns)) = TimestampNanosecond(5000000000, None)\n  TableScan: t",
1353                expected_datafusion: "Filter: t.ts_ms = TimestampMillisecond(5000, None)\n  TableScan: t",
1354            },
1355        ];
1356
1357        for case in cases {
1358            let greptime =
1359                greptime_const_normalized_filter(case.fields.clone(), case.predicate.clone());
1360            let datafusion = datafusion_simplified_filter(case.fields, case.predicate);
1361            assert_eq!(case.expected_greptime, greptime, "{} greptime", case.name);
1362            assert_eq!(
1363                case.expected_datafusion, datafusion,
1364                "{} datafusion",
1365                case.name
1366            );
1367        }
1368    }
1369
1370    fn assert_pattern_match_plan(kind: PatternMatchKind, pattern: ScalarValue, expected: &str) {
1371        let predicate = match kind {
1372            PatternMatchKind::Like => Expr::Like(Like::new(
1373                false,
1374                Box::new(cast(col("s"), DataType::LargeUtf8)),
1375                Box::new(lit(pattern)),
1376                None,
1377                false,
1378            )),
1379            PatternMatchKind::SimilarTo => Expr::SimilarTo(Like::new(
1380                false,
1381                Box::new(cast(col("s"), DataType::LargeUtf8)),
1382                Box::new(lit(pattern)),
1383                None,
1384                false,
1385            )),
1386        };
1387
1388        assert_filter_plan(
1389            vec![Field::new("s", DataType::Utf8, false)],
1390            predicate,
1391            expected,
1392        );
1393    }
1394
1395    fn assert_filter_plan(fields: Vec<Field>, predicate: Expr, expected: &str) {
1396        assert_eq!(expected, analyze_filter(fields, predicate).to_string());
1397    }
1398
1399    fn assert_filter_left_is_cast(fields: Vec<Field>, predicate: Expr) {
1400        let analyzed = analyze_filter(fields, predicate);
1401        let LogicalPlan::Filter(filter) = analyzed else {
1402            panic!("expected filter plan");
1403        };
1404        let Expr::BinaryExpr(BinaryExpr { left, .. }) = filter.predicate else {
1405            panic!("expected binary predicate");
1406        };
1407        assert!(matches!(left.as_ref(), Expr::Cast(_)));
1408    }
1409
1410    fn rewrite_dictionary_regex(
1411        fields: Vec<Field>,
1412        left: Expr,
1413        right: Expr,
1414        op: Operator,
1415    ) -> Option<Expr> {
1416        rewrite_dictionary_string_regex(
1417            &BinaryExpr {
1418                left: Box::new(left),
1419                op,
1420                right: Box::new(right),
1421            },
1422            &test_schema(fields),
1423        )
1424        .unwrap()
1425    }
1426
1427    fn assert_timestamp_pushdown(
1428        fields: Vec<Field>,
1429        predicate: Expr,
1430        expected_analyzed: &str,
1431        expected_pushed: &str,
1432        expected_range: TimestampRange,
1433    ) {
1434        let analyzed = analyze_filter(fields, predicate);
1435        assert_eq!(expected_analyzed, analyzed.to_string());
1436
1437        let pushed = push_down_filters(analyzed);
1438        assert_eq!(expected_pushed, pushed.to_string());
1439
1440        let range =
1441            build_time_range_predicate("ts", TimeUnit::Nanosecond, &extract_scan_filters(&pushed));
1442        assert_eq!(expected_range, range);
1443    }
1444
1445    fn analyze_filter(fields: Vec<Field>, predicate: Expr) -> LogicalPlan {
1446        analyze_plan(build_filter_plan(test_schema(fields), predicate))
1447    }
1448
1449    fn greptime_const_normalized_filter(fields: Vec<Field>, predicate: Expr) -> String {
1450        analyze_filter(fields, predicate).to_string()
1451    }
1452
1453    fn datafusion_simplified_filter(fields: Vec<Field>, predicate: Expr) -> String {
1454        let plan = build_filter_plan(test_schema(fields), predicate);
1455        Optimizer::with_rules(vec![Arc::new(SimplifyExpressions::new())])
1456            .optimize(plan, &OptimizerContext::new(), |_, _| {})
1457            .unwrap()
1458            .to_string()
1459    }
1460
1461    fn analyze_plan(plan: LogicalPlan) -> LogicalPlan {
1462        ConstNormalizationRule
1463            .analyze(plan, &ConfigOptions::default())
1464            .unwrap()
1465    }
1466
1467    fn build_filter_plan(schema: Arc<DFSchema>, predicate: Expr) -> LogicalPlan {
1468        LogicalPlanBuilder::scan("t", test_source(schema), None)
1469            .unwrap()
1470            .filter(predicate)
1471            .unwrap()
1472            .build()
1473            .unwrap()
1474    }
1475
1476    fn build_scan_plan(schema: Arc<DFSchema>) -> LogicalPlan {
1477        LogicalPlanBuilder::scan("t", test_source(schema), None)
1478            .unwrap()
1479            .build()
1480            .unwrap()
1481    }
1482
1483    fn push_down_filters(plan: LogicalPlan) -> LogicalPlan {
1484        Optimizer::with_rules(vec![Arc::new(PushDownFilter::new())])
1485            .optimize(plan, &OptimizerContext::new(), |_, _| {})
1486            .unwrap()
1487    }
1488
1489    fn ts_cast_to_ms() -> Expr {
1490        cast(
1491            col("ts"),
1492            DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1493        )
1494    }
1495
1496    fn ts_ms_literal(value: i64) -> Expr {
1497        lit(ScalarValue::TimestampMillisecond(Some(value), None))
1498    }
1499
1500    fn extract_scan_filters(plan: &LogicalPlan) -> Vec<Expr> {
1501        match plan {
1502            LogicalPlan::TableScan(scan) => scan.filters.clone(),
1503            _ => plan
1504                .inputs()
1505                .into_iter()
1506                .flat_map(extract_scan_filters)
1507                .collect(),
1508        }
1509    }
1510
1511    fn test_schema(fields: Vec<Field>) -> Arc<DFSchema> {
1512        arrow_schema::Schema::new(fields).to_dfschema_ref().unwrap()
1513    }
1514
1515    fn test_source(schema: Arc<DFSchema>) -> Arc<dyn TableSource> {
1516        let table = ExactPushdownProvider {
1517            schema: Arc::new(schema.as_ref().as_arrow().clone()),
1518        };
1519        provider_as_source(Arc::new(table))
1520    }
1521
1522    #[derive(Debug)]
1523    struct ExactPushdownProvider {
1524        schema: arrow_schema::SchemaRef,
1525    }
1526
1527    #[async_trait]
1528    impl TableProvider for ExactPushdownProvider {
1529        fn schema(&self) -> arrow_schema::SchemaRef {
1530            self.schema.clone()
1531        }
1532
1533        fn table_type(&self) -> TableType {
1534            TableType::Base
1535        }
1536
1537        async fn scan(
1538            &self,
1539            _state: &dyn Session,
1540            _projection: Option<&Vec<usize>>,
1541            _filters: &[Expr],
1542            _limit: Option<usize>,
1543        ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
1544            unreachable!("scan should not be called in const_normalization tests")
1545        }
1546
1547        fn supports_filters_pushdown(
1548            &self,
1549            filters: &[&Expr],
1550        ) -> datafusion::error::Result<Vec<TableProviderFilterPushDown>> {
1551            Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
1552        }
1553    }
1554}
query/optimizer/const_normalization.rs

query/optimizer/
const_normalization.rs
