common_recordbatch/
filter.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//! Util record batch stream wrapper that can perform precise filter.
16
17use std::sync::Arc;
18
19use common_time::timestamp::div_mod_units;
20use datafusion::error::Result as DfResult;
21use datafusion::logical_expr::{Expr, Literal, Operator};
22use datafusion::physical_plan::PhysicalExpr;
23use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar};
24use datafusion_common::arrow::buffer::BooleanBuffer;
25use datafusion_common::arrow::compute::kernels::cmp;
26use datafusion_common::cast::{as_boolean_array, as_null_array, as_string_array};
27use datafusion_common::{DataFusionError, ScalarValue, internal_err};
28use datatypes::arrow::array::{
29    Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray,
30    RecordBatch, StringArrayType,
31};
32use datatypes::arrow::compute::filter_record_batch;
33use datatypes::arrow::datatypes::{DataType, TimeUnit, UInt32Type};
34use datatypes::arrow::error::ArrowError;
35use datatypes::compute::or_kleene;
36use datatypes::data_type::{ConcreteDataType, DataType as _};
37use datatypes::value::Value;
38use datatypes::vectors::VectorRef;
39use regex::Regex;
40use snafu::ResultExt;
41
42use crate::error::{ArrowComputeSnafu, Result, ToArrowScalarSnafu, UnsupportedOperationSnafu};
43
44/// An inplace expr evaluator for simple filter. Only support
45/// - `col` `op` `literal`
46/// - `literal` `op` `col`
47///
48/// And the `op` is one of `=`, `!=`, `>`, `>=`, `<`, `<=`,
49/// or regex operators: `~`, `~*`, `!~`, `!~*`.
50///
51/// This struct contains normalized predicate expr. In the form of
52/// `col` `op` `literal` where the `col` is provided from input.
53#[derive(Debug, Clone)]
54pub struct SimpleFilterEvaluator {
55    /// Name of the referenced column.
56    column_name: String,
57    /// The literal value.
58    literal: Scalar<ArrayRef>,
59    /// The operator.
60    op: Operator,
61    /// Only used when the operator is `Or`-chain.
62    literal_list: Vec<Scalar<ArrayRef>>,
63    /// Pre-compiled regex.
64    /// Only used when the operator is regex operators.
65    /// If the regex is empty, it is also `None`.
66    regex: Option<Regex>,
67    /// Whether the regex is negative.
68    regex_negative: bool,
69}
70
71impl SimpleFilterEvaluator {
72    pub fn new<T: Literal>(column_name: String, lit: T, op: Operator) -> Option<Self> {
73        match op {
74            Operator::Eq
75            | Operator::NotEq
76            | Operator::Lt
77            | Operator::LtEq
78            | Operator::Gt
79            | Operator::GtEq => {}
80            _ => return None,
81        }
82
83        let Expr::Literal(val, _) = lit.lit() else {
84            return None;
85        };
86
87        Some(Self {
88            column_name,
89            literal: val.to_scalar().ok()?,
90            op,
91            literal_list: vec![],
92            regex: None,
93            regex_negative: false,
94        })
95    }
96
97    pub fn try_new(predicate: &Expr) -> Option<Self> {
98        match predicate {
99            Expr::BinaryExpr(binary) => {
100                // check if the expr is in the supported form
101                match binary.op {
102                    Operator::Eq
103                    | Operator::NotEq
104                    | Operator::Lt
105                    | Operator::LtEq
106                    | Operator::Gt
107                    | Operator::GtEq
108                    | Operator::RegexMatch
109                    | Operator::RegexIMatch
110                    | Operator::RegexNotMatch
111                    | Operator::RegexNotIMatch => {}
112                    Operator::Or => {
113                        let lhs = Self::try_new(&binary.left)?;
114                        let rhs = Self::try_new(&binary.right)?;
115                        if lhs.column_name != rhs.column_name
116                            || !matches!(lhs.op, Operator::Eq | Operator::Or)
117                            || !matches!(rhs.op, Operator::Eq | Operator::Or)
118                        {
119                            return None;
120                        }
121                        let mut list = vec![];
122                        let placeholder_literal = lhs.literal.clone();
123                        // above check guarantees the op is either `Eq` or `Or`
124                        if matches!(lhs.op, Operator::Or) {
125                            list.extend(lhs.literal_list);
126                        } else {
127                            list.push(lhs.literal);
128                        }
129                        if matches!(rhs.op, Operator::Or) {
130                            list.extend(rhs.literal_list);
131                        } else {
132                            list.push(rhs.literal);
133                        }
134                        return Some(Self {
135                            column_name: lhs.column_name,
136                            literal: placeholder_literal,
137                            op: Operator::Or,
138                            literal_list: list,
139                            regex: None,
140                            regex_negative: false,
141                        });
142                    }
143                    _ => return None,
144                }
145
146                // swap the expr if it is in the form of `literal` `op` `col`
147                let mut op = binary.op;
148                let (lhs, rhs) = match (&*binary.left, &*binary.right) {
149                    (Expr::Column(col), Expr::Literal(lit, _)) => (col, lit),
150                    (Expr::Literal(lit, _), Expr::Column(col)) => {
151                        // safety: The previous check ensures the operator is able to swap.
152                        op = op.swap().unwrap();
153                        (col, lit)
154                    }
155                    _ => return None,
156                };
157
158                let (regex, regex_negative) = Self::maybe_build_regex(op, rhs).ok()?;
159                let literal = rhs.to_scalar().ok()?;
160                Some(Self {
161                    column_name: lhs.name.clone(),
162                    literal,
163                    op,
164                    literal_list: vec![],
165                    regex,
166                    regex_negative,
167                })
168            }
169            _ => None,
170        }
171    }
172
173    /// Get the name of the referenced column.
174    pub fn column_name(&self) -> &str {
175        &self.column_name
176    }
177
178    pub fn is_eq(&self) -> bool {
179        matches!(self.op, Operator::Eq)
180    }
181
182    pub fn is_not_eq(&self) -> bool {
183        matches!(self.op, Operator::NotEq)
184    }
185
186    pub fn is_lt(&self) -> bool {
187        matches!(self.op, Operator::Lt)
188    }
189
190    pub fn is_lt_eq(&self) -> bool {
191        matches!(self.op, Operator::LtEq)
192    }
193
194    pub fn is_gt(&self) -> bool {
195        matches!(self.op, Operator::Gt)
196    }
197
198    pub fn is_gt_eq(&self) -> bool {
199        matches!(self.op, Operator::GtEq)
200    }
201
202    /// Returns true if this filter represents an `OR` chain of equality comparisons, e.g.
203    /// `col = lit1 OR col = lit2 ...`.
204    pub fn is_or_eq_chain(&self) -> bool {
205        matches!(self.op, Operator::Or)
206    }
207
208    /// Returns the literal as a [`Value`]. It returns `None` if the literal can't be converted.
209    pub fn literal_value(&self) -> Option<Value> {
210        let array = self.literal.get().0;
211        let scalar = ScalarValue::try_from_array(array, 0).ok()?;
212        Value::try_from(scalar).ok()
213    }
214
215    /// Returns the literal list as a list of [`Value`]s. It returns `None` if any literal can't be
216    /// converted.
217    pub fn literal_list_values(&self) -> Option<Vec<Value>> {
218        self.literal_list
219            .iter()
220            .map(|scalar| {
221                let array = scalar.get().0;
222                let scalar = ScalarValue::try_from_array(array, 0).ok()?;
223                Value::try_from(scalar).ok()
224            })
225            .collect()
226    }
227
228    pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result<bool> {
229        let input = input
230            .to_scalar()
231            .with_context(|_| ToArrowScalarSnafu { v: input.clone() })?;
232        let result = self.evaluate_datum(&input, 1)?;
233        Ok(result.value(0))
234    }
235
236    pub fn evaluate_array(&self, input: &ArrayRef) -> Result<BooleanBuffer> {
237        self.evaluate_datum(input, input.len())
238    }
239
240    pub fn evaluate_vector(&self, input: &VectorRef) -> Result<BooleanBuffer> {
241        self.evaluate_datum(&input.to_arrow_array(), input.len())
242    }
243
244    fn evaluate_datum(&self, input: &impl Datum, input_len: usize) -> Result<BooleanBuffer> {
245        let result = match self.op {
246            Operator::Eq => cmp::eq(input, &self.literal),
247            Operator::NotEq => cmp::neq(input, &self.literal),
248            Operator::Lt => cmp::lt(input, &self.literal),
249            Operator::LtEq => cmp::lt_eq(input, &self.literal),
250            Operator::Gt => cmp::gt(input, &self.literal),
251            Operator::GtEq => cmp::gt_eq(input, &self.literal),
252            Operator::RegexMatch => self.regex_match(input),
253            Operator::RegexIMatch => self.regex_match(input),
254            Operator::RegexNotMatch => self.regex_match(input),
255            Operator::RegexNotIMatch => self.regex_match(input),
256            Operator::Or => {
257                // OR operator stands for OR-chained EQs (or INLIST in other words)
258                let mut result: BooleanArray = vec![false; input_len].into();
259                for literal in &self.literal_list {
260                    let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?;
261                    result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?;
262                }
263                Ok(result)
264            }
265            _ => {
266                return UnsupportedOperationSnafu {
267                    reason: format!("{:?}", self.op),
268                }
269                .fail();
270            }
271        };
272        result
273            .context(ArrowComputeSnafu)
274            .map(|array| boolean_array_to_scan_mask(&array).values().clone())
275    }
276
277    /// Builds a regex pattern from a scalar value and operator.
278    /// Returns the `(regex, negative)` and if successful.
279    ///
280    /// Returns `Err` if
281    /// - the value is not a string
282    /// - the regex pattern is invalid
283    ///
284    /// The regex is `None` if
285    /// - the operator is not a regex operator
286    /// - the pattern is empty
287    fn maybe_build_regex(
288        operator: Operator,
289        value: &ScalarValue,
290    ) -> Result<(Option<Regex>, bool), ArrowError> {
291        let (ignore_case, negative) = match operator {
292            Operator::RegexMatch => (false, false),
293            Operator::RegexIMatch => (true, false),
294            Operator::RegexNotMatch => (false, true),
295            Operator::RegexNotIMatch => (true, true),
296            _ => return Ok((None, false)),
297        };
298        let flag = if ignore_case { Some("i") } else { None };
299        let regex = value
300            .try_as_str()
301            .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))?
302            .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?;
303        let pattern = match flag {
304            Some(flag) => format!("(?{flag}){regex}"),
305            None => regex.to_string(),
306        };
307        if pattern.is_empty() {
308            Ok((None, negative))
309        } else {
310            Regex::new(pattern.as_str())
311                .map_err(|e| {
312                    ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}"))
313                })
314                .map(|regex| (Some(regex), negative))
315        }
316    }
317
318    fn regex_match(&self, input: &impl Datum) -> std::result::Result<BooleanArray, ArrowError> {
319        let array = input.get().0;
320
321        // Try to cast to StringArray first
322        if let Ok(string_array) = as_string_array(array) {
323            let mut result = regexp_is_match_scalar(string_array, self.regex.as_ref())?;
324            if self.regex_negative {
325                result = datatypes::compute::not(&result)?;
326            }
327            return Ok(result);
328        }
329
330        // Try to cast to StringDictionaryArray
331        if let Some(dict_array) = array.as_any().downcast_ref::<DictionaryArray<UInt32Type>>() {
332            let mut result = regexp_is_match_dictionary(dict_array, self.regex.as_ref())?;
333            if self.regex_negative {
334                result = datatypes::compute::not(&result)?;
335            }
336            return Ok(result);
337        }
338
339        Err(ArrowError::CastError(format!(
340            "Cannot cast {:?} to StringArray or StringDictionaryArray",
341            array.data_type()
342        )))
343    }
344
345    /// Casts the filter's timestamp literal into the unit of the `target`
346    /// timestamp type, so it can be evaluated against a column stored in that
347    /// unit (e.g. an old-unit SST after the time index unit was widened).
348    ///
349    /// When the literal is not representable in `target`'s unit (e.g.
350    /// `ts = 7_000_500us` against a millisecond column), the outcome keeps the
351    /// row set instead of rounding the literal: `=` prunes / `!=` matches, and
352    /// inequalities strengthen the operator (e.g. `>= 2_500_500us` becomes
353    /// `> 2500ms`) so the excluded boundary row stays excluded.
354    ///
355    /// Returns `None` when the filter doesn't compare against a tz-naive
356    /// timestamp literal, or `target` is not a timestamp type.
357    pub fn cast_timestamp_unit(&self, target: &ConcreteDataType) -> Option<TimestampUnitCast> {
358        let target_unit = match target.as_arrow_type() {
359            DataType::Timestamp(unit, _) => unit,
360            _ => return None,
361        };
362
363        // An `OR` chain of equalities (an IN list): convert each literal
364        // with `=` semantics; literals not representable in the target unit
365        // cannot match any row and simply drop out of the chain.
366        if self.op == Operator::Or {
367            let mut literal_list = Vec::with_capacity(self.literal_list.len());
368            for literal in &self.literal_list {
369                let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?;
370                let Some((value, unit)) = timestamp_scalar_parts(&scalar) else {
371                    continue;
372                };
373                let Some(value) = value else {
374                    continue;
375                };
376                let cast = div_mod_units(value, unit.into(), target_unit.into())?;
377                if cast.remainder == 0 {
378                    literal_list.push(timestamp_scalar(cast.quotient, target_unit)?);
379                }
380            }
381            if literal_list.is_empty() {
382                return Some(TimestampUnitCast::Pruned);
383            }
384            let literal = literal_list[0].clone();
385            return Some(TimestampUnitCast::Filter(Self {
386                column_name: self.column_name.clone(),
387                literal,
388                op: Operator::Or,
389                literal_list,
390                regex: None,
391                regex_negative: false,
392            }));
393        }
394
395        let scalar = ScalarValue::try_from_array(self.literal.get().0, 0).ok()?;
396        let (value, unit) = timestamp_scalar_parts(&scalar)?;
397        let Some(value) = value else {
398            // A NULL literal never compares true (arrow comparison
399            // semantics: null results are filtered out).
400            return Some(TimestampUnitCast::Pruned);
401        };
402        if unit == target_unit {
403            return Some(TimestampUnitCast::Filter(self.clone()));
404        }
405        let cast = div_mod_units(value, unit.into(), target_unit.into())?;
406        let divisible = cast.remainder == 0;
407        let literal = timestamp_scalar(cast.quotient, target_unit)?;
408        let filter = |op: Operator| Self {
409            column_name: self.column_name.clone(),
410            literal: literal.clone(),
411            op,
412            literal_list: vec![],
413            regex: None,
414            regex_negative: false,
415        };
416
417        Some(match self.op {
418            Operator::Eq if divisible => TimestampUnitCast::Filter(filter(Operator::Eq)),
419            Operator::Eq => TimestampUnitCast::Pruned,
420            Operator::NotEq if divisible => TimestampUnitCast::Filter(filter(Operator::NotEq)),
421            Operator::NotEq => TimestampUnitCast::Matched,
422            // v > L holds exactly for v > quotient, whether or not L is
423            // representable in the target unit.
424            Operator::Gt => TimestampUnitCast::Filter(filter(Operator::Gt)),
425            Operator::GtEq if divisible => TimestampUnitCast::Filter(filter(Operator::GtEq)),
426            // L strictly between two target values: v >= L is v > quotient.
427            Operator::GtEq => TimestampUnitCast::Filter(filter(Operator::Gt)),
428            Operator::Lt if divisible => TimestampUnitCast::Filter(filter(Operator::Lt)),
429            // L strictly between two target values: v < L is v <= quotient.
430            Operator::Lt => TimestampUnitCast::Filter(filter(Operator::LtEq)),
431            // v <= L holds exactly for v <= quotient.
432            Operator::LtEq => TimestampUnitCast::Filter(filter(Operator::LtEq)),
433            // Regex predicates don't apply to timestamps.
434            _ => return None,
435        })
436    }
437}
438
439/// The result of casting a [`SimpleFilterEvaluator`] to a different timestamp
440/// unit, preserving the predicate's semantics on rows stored in that unit.
441/// See [`SimpleFilterEvaluator::cast_timestamp_unit`].
442#[derive(Debug, Clone)]
443pub enum TimestampUnitCast {
444    /// The cast filter; evaluates column values stored in the target unit.
445    Filter(SimpleFilterEvaluator),
446    /// No value in the target unit satisfies the filter.
447    Pruned,
448    /// Every value in the target unit satisfies the filter.
449    Matched,
450}
451
452/// Extracts the value and unit from a tz-naive timestamp scalar.
453fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option<i64>, TimeUnit)> {
454    let (value, unit, timezone) = match scalar {
455        ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz),
456        ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz),
457        ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz),
458        ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz),
459        _ => return None,
460    };
461    // A timezone-aware literal doesn't compare against a tz-naive column;
462    // leave it to the caller instead of guessing the intended semantics.
463    (timezone.is_none()).then_some((value, unit))
464}
465
466/// Builds a tz-naive timestamp literal for `value` in `unit`.
467pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue {
468    match unit {
469        TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
470        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
471        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
472        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
473    }
474}
475
476/// Builds a one-element scalar array holding `value` in `unit`.
477fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option<Scalar<ArrayRef>> {
478    timestamp_scalar_value(value, unit).to_scalar().ok()
479}
480
481/// Evaluate the predicate on the input [RecordBatch], and return a new [RecordBatch].
482/// Copy from datafusion::physical_plan::src::filter.rs
483pub fn batch_filter(
484    batch: &RecordBatch,
485    predicate: &Arc<dyn PhysicalExpr>,
486) -> DfResult<RecordBatch> {
487    predicate
488        .evaluate(batch)
489        .and_then(|v| v.into_array(batch.num_rows()))
490        .and_then(|array| {
491            let filter_array = match as_boolean_array(&array) {
492                Ok(boolean_array) => Ok(boolean_array.clone()),
493                Err(_) => {
494                    let Ok(null_array) = as_null_array(&array) else {
495                        return internal_err!(
496                            "Cannot create filter_array from non-boolean predicates"
497                        );
498                    };
499
500                    // if the predicate is null, then the result is also null
501                    Ok::<BooleanArray, DataFusionError>(BooleanArray::new_null(null_array.len()))
502                }
503            }?;
504            Ok(filter_record_batch(
505                batch,
506                &boolean_array_to_scan_mask(&filter_array),
507            )?)
508        })
509}
510
511/// Converts nullable SQL predicate values to a scan mask, where `NULL` is `false`.
512fn boolean_array_to_scan_mask(array: &BooleanArray) -> BooleanArray {
513    if array.null_count() == 0 {
514        return array.clone();
515    }
516
517    let mut values = BooleanBufferBuilder::new(array.len());
518    for index in 0..array.len() {
519        values.append(array.is_valid(index) && array.value(index));
520    }
521    BooleanArray::new(values.into(), None)
522}
523
524/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar())
525/// with pre-compiled regex.
526/// See <https://github.com/apache/arrow-rs/blob/54.2.0/arrow-string/src/regexp.rs#L204-L246> for the implementation details.
527pub fn regexp_is_match_scalar<'a, S>(
528    array: &'a S,
529    regex: Option<&Regex>,
530) -> Result<BooleanArray, ArrowError>
531where
532    &'a S: StringArrayType<'a>,
533{
534    let null_bit_buffer = array.nulls().map(|x| x.inner().sliced());
535    let mut result = BooleanBufferBuilder::new(array.len());
536
537    if let Some(re) = regex {
538        for i in 0..array.len() {
539            let value = array.value(i);
540            result.append(re.is_match(value));
541        }
542    } else {
543        result.append_n(array.len(), true);
544    }
545
546    let buffer = result.into();
547    let data = unsafe {
548        ArrayData::new_unchecked(
549            DataType::Boolean,
550            array.len(),
551            None,
552            null_bit_buffer,
553            0,
554            vec![buffer],
555            vec![],
556        )
557    };
558
559    Ok(BooleanArray::from(data))
560}
561
562/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray.
563/// Iterates through dictionary keys to get string values and applies regex matching.
564pub fn regexp_is_match_dictionary(
565    dict_array: &DictionaryArray<UInt32Type>,
566    regex: Option<&Regex>,
567) -> Result<BooleanArray, ArrowError> {
568    // Get the string values from the dictionary
569    let string_values = dict_array
570        .values()
571        .as_any()
572        .downcast_ref::<datatypes::arrow::array::StringArray>()
573        .ok_or_else(|| {
574            ArrowError::CastError("Dictionary values must be StringArray".to_string())
575        })?;
576
577    // Dictionary logical nulls include both null keys and keys whose dictionary value is null.
578    let logical_nulls = dict_array.logical_nulls();
579    let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced());
580    let mut result = BooleanBufferBuilder::new(dict_array.len());
581
582    if let Some(re) = regex {
583        let keys = dict_array.keys().values();
584        for i in 0..dict_array.len() {
585            if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) {
586                result.append(false);
587            } else {
588                let key = keys[i] as usize;
589                let string_value = string_values.value(key);
590                result.append(re.is_match(string_value));
591            }
592        }
593    } else {
594        result.append_n(dict_array.len(), true);
595    }
596
597    let buffer = result.into();
598    let data = unsafe {
599        ArrayData::new_unchecked(
600            DataType::Boolean,
601            dict_array.len(),
602            None,
603            null_bit_buffer,
604            0,
605            vec![buffer],
606            vec![],
607        )
608    };
609
610    Ok(BooleanArray::from(data))
611}
612
613#[cfg(test)]
614mod test {
615
616    use std::sync::Arc;
617
618    use datafusion::execution::context::ExecutionProps;
619    use datafusion::logical_expr::{BinaryExpr, col, lit};
620    use datafusion::physical_expr::create_physical_expr;
621    use datafusion_common::{Column, DFSchema};
622    use datatypes::arrow::array::{TimestampMillisecondArray, TimestampNanosecondArray};
623    use datatypes::arrow::datatypes::{DataType, Field, Schema};
624
625    use super::*;
626
627    #[test]
628    fn unsupported_filter_op() {
629        // `+` is not supported
630        let expr = Expr::BinaryExpr(BinaryExpr {
631            left: Box::new(Expr::Column(Column::from_name("foo"))),
632            op: Operator::Plus,
633            right: Box::new(1.lit()),
634        });
635        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
636
637        // two literal is not supported
638        let expr = Expr::BinaryExpr(BinaryExpr {
639            left: Box::new(1.lit()),
640            op: Operator::Eq,
641            right: Box::new(1.lit()),
642        });
643        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
644
645        // two column is not supported
646        let expr = Expr::BinaryExpr(BinaryExpr {
647            left: Box::new(Expr::Column(Column::from_name("foo"))),
648            op: Operator::Eq,
649            right: Box::new(Expr::Column(Column::from_name("bar"))),
650        });
651        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
652
653        // compound expr is not supported
654        let expr = Expr::BinaryExpr(BinaryExpr {
655            left: Box::new(Expr::BinaryExpr(BinaryExpr {
656                left: Box::new(Expr::Column(Column::from_name("foo"))),
657                op: Operator::Eq,
658                right: Box::new(1.lit()),
659            })),
660            op: Operator::Eq,
661            right: Box::new(1.lit()),
662        });
663        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
664    }
665
666    #[test]
667    fn supported_filter_op() {
668        // equal
669        let expr = Expr::BinaryExpr(BinaryExpr {
670            left: Box::new(Expr::Column(Column::from_name("foo"))),
671            op: Operator::Eq,
672            right: Box::new(1.lit()),
673        });
674        let _ = SimpleFilterEvaluator::try_new(&expr).unwrap();
675
676        // swap operands
677        let expr = Expr::BinaryExpr(BinaryExpr {
678            left: Box::new(1.lit()),
679            op: Operator::Lt,
680            right: Box::new(Expr::Column(Column::from_name("foo"))),
681        });
682        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
683        assert_eq!(evaluator.op, Operator::Gt);
684        assert_eq!(evaluator.column_name, "foo".to_string());
685    }
686
687    #[test]
688    fn run_on_array() {
689        let expr = Expr::BinaryExpr(BinaryExpr {
690            left: Box::new(Expr::Column(Column::from_name("foo"))),
691            op: Operator::Eq,
692            right: Box::new(1i64.lit()),
693        });
694        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
695
696        let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _;
697        let result = evaluator.evaluate_array(&input_1).unwrap();
698        assert_eq!(result, BooleanBuffer::from(vec![true, false, false]));
699
700        let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _;
701        let result = evaluator.evaluate_array(&input_2).unwrap();
702        assert_eq!(result, BooleanBuffer::from(vec![true, true, true]));
703
704        let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _;
705        let result = evaluator.evaluate_array(&input_3).unwrap();
706        assert_eq!(result, BooleanBuffer::from(vec![]));
707    }
708
709    #[test]
710    fn run_on_scalar() {
711        let expr = Expr::BinaryExpr(BinaryExpr {
712            left: Box::new(Expr::Column(Column::from_name("foo"))),
713            op: Operator::Lt,
714            right: Box::new(1i64.lit()),
715        });
716        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
717
718        let input_1 = ScalarValue::Int64(Some(1));
719        let result = evaluator.evaluate_scalar(&input_1).unwrap();
720        assert!(!result);
721
722        let input_2 = ScalarValue::Int64(Some(0));
723        let result = evaluator.evaluate_scalar(&input_2).unwrap();
724        assert!(result);
725
726        let input_3 = ScalarValue::Int64(None);
727        let result = evaluator.evaluate_scalar(&input_3).unwrap();
728        assert!(!result);
729    }
730
731    #[test]
732    fn batch_filter_test() {
733        let expr = col("ts").gt(lit(123456u64));
734        let schema = Schema::new(vec![
735            Field::new("a", DataType::Int32, true),
736            Field::new("ts", DataType::UInt64, false),
737        ]);
738        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
739        let props = ExecutionProps::new();
740        let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
741        let batch = RecordBatch::try_new(
742            Arc::new(schema),
743            vec![
744                Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])),
745                Arc::new(datatypes::arrow::array::UInt64Array::from(vec![
746                    123456, 123457, 123458,
747                ])),
748            ],
749        )
750        .unwrap();
751        let new_batch = batch_filter(&batch, &physical_expr).unwrap();
752        assert_eq!(new_batch.num_rows(), 2);
753        let first_column_values = new_batch
754            .column(0)
755            .as_any()
756            .downcast_ref::<datatypes::arrow::array::Int32Array>()
757            .unwrap();
758        let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]);
759        assert_eq!(first_column_values, &expected);
760    }
761
762    #[test]
763    fn test_complex_filter_expression() {
764        // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D'
765        let col_eq_b = col("col").eq(lit("B"));
766        let col_eq_c = col("col").eq(lit("C"));
767        let col_eq_d = col("col").eq(lit("D"));
768
769        // Build the OR chain
770        let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d);
771
772        // Check that SimpleFilterEvaluator can handle OR chain
773        let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap();
774        assert_eq!(or_evaluator.column_name, "col");
775        assert_eq!(or_evaluator.op, Operator::Or);
776        assert_eq!(or_evaluator.literal_list.len(), 3);
777        assert_eq!(
778            format!("{:?}", or_evaluator.literal_list),
779            "[Scalar(StringArray\n[\n  \"B\",\n]), Scalar(StringArray\n[\n  \"C\",\n]), Scalar(StringArray\n[\n  \"D\",\n])]"
780        );
781
782        // Create a schema and batch for testing
783        let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
784        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
785        let props = ExecutionProps::new();
786        let physical_expr = create_physical_expr(&col_or_expr, &df_schema, &props).unwrap();
787
788        // Create test data
789        let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![
790            "B", "C", "E", "B", "C", "D", "F",
791        ]));
792        let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap();
793        let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]);
794
795        // Filter the batch
796        let filtered_batch = batch_filter(&batch, &physical_expr).unwrap();
797
798        // Expected: rows with col in ("B", "C", "D")
799        // That would be rows 0, 1, 3, 4, 5
800        assert_eq!(filtered_batch.num_rows(), 5);
801
802        let col_filtered = filtered_batch
803            .column(0)
804            .as_any()
805            .downcast_ref::<datatypes::arrow::array::StringArray>()
806            .unwrap();
807        assert_eq!(col_filtered, &expected);
808    }
809
810    #[test]
811    fn test_maybe_build_regex() {
812        // Test case for RegexMatch (case sensitive, non-negative)
813        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
814            Operator::RegexMatch,
815            &ScalarValue::Utf8(Some("a.*b".to_string())),
816        )
817        .unwrap();
818        assert!(regex.is_some());
819        assert!(!negative);
820        assert!(regex.unwrap().is_match("axxb"));
821
822        // Test case for RegexIMatch (case insensitive, non-negative)
823        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
824            Operator::RegexIMatch,
825            &ScalarValue::Utf8(Some("a.*b".to_string())),
826        )
827        .unwrap();
828        assert!(regex.is_some());
829        assert!(!negative);
830        assert!(regex.unwrap().is_match("AxxB"));
831
832        // Test case for RegexNotMatch (case sensitive, negative)
833        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
834            Operator::RegexNotMatch,
835            &ScalarValue::Utf8(Some("a.*b".to_string())),
836        )
837        .unwrap();
838        assert!(regex.is_some());
839        assert!(negative);
840
841        // Test case for RegexNotIMatch (case insensitive, negative)
842        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
843            Operator::RegexNotIMatch,
844            &ScalarValue::Utf8(Some("a.*b".to_string())),
845        )
846        .unwrap();
847        assert!(regex.is_some());
848        assert!(negative);
849
850        // Test with empty regex pattern
851        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
852            Operator::RegexMatch,
853            &ScalarValue::Utf8(Some("".to_string())),
854        )
855        .unwrap();
856        assert!(regex.is_none());
857        assert!(!negative);
858
859        // Test with non-regex operator
860        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
861            Operator::Eq,
862            &ScalarValue::Utf8(Some("a.*b".to_string())),
863        )
864        .unwrap();
865        assert!(regex.is_none());
866        assert!(!negative);
867
868        // Test with invalid regex pattern
869        let result = SimpleFilterEvaluator::maybe_build_regex(
870            Operator::RegexMatch,
871            &ScalarValue::Utf8(Some("a(b".to_string())),
872        );
873        assert!(result.is_err());
874
875        // Test with non-string value
876        let result = SimpleFilterEvaluator::maybe_build_regex(
877            Operator::RegexMatch,
878            &ScalarValue::Int64(Some(123)),
879        );
880        assert!(result.is_err());
881
882        // Test with null value
883        let result = SimpleFilterEvaluator::maybe_build_regex(
884            Operator::RegexMatch,
885            &ScalarValue::Utf8(None),
886        );
887        assert!(result.is_err());
888    }
889
890    #[test]
891    fn test_regex_match_dictionary_array() {
892        use datatypes::arrow::array::StringDictionaryBuilder;
893
894        // Create a StringDictionaryArray
895        let mut builder = StringDictionaryBuilder::<UInt32Type>::new();
896        builder.append("apple").unwrap();
897        builder.append("banana").unwrap();
898        builder.append("apple").unwrap();
899        builder.append("cherry").unwrap();
900        let dict_array = builder.finish();
901
902        // Test regex that matches "apple"
903        let regex = regex::Regex::new(r"app.*").unwrap();
904        let result = regexp_is_match_dictionary(&dict_array, Some(&regex)).unwrap();
905
906        // Should match indices 0 and 2 (both "apple")
907        assert_eq!(result.len(), 4);
908        assert!(result.value(0)); // "apple"
909        assert!(!result.value(1)); // "banana"
910        assert!(result.value(2)); // "apple"
911        assert!(!result.value(3)); // "cherry"
912
913        // Test regex that matches "banana"
914        let regex2 = regex::Regex::new(r"ban.*").unwrap();
915        let result2 = regexp_is_match_dictionary(&dict_array, Some(&regex2)).unwrap();
916
917        assert!(!result2.value(0)); // "apple"
918        assert!(result2.value(1)); // "banana"
919        assert!(!result2.value(2)); // "apple"
920        assert!(!result2.value(3)); // "cherry"
921
922        // Test with no regex (should match all)
923        let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap();
924        assert!(result3.value(0));
925        assert!(result3.value(1));
926        assert!(result3.value(2));
927        assert!(result3.value(3));
928    }
929
930    #[test]
931    fn test_regex_scan_masks_preserve_sql_null_semantics() {
932        let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![
933            Some("api"),
934            Some("API"),
935            Some("db"),
936            None,
937        ])) as ArrayRef;
938        assert_regex_scan_masks(
939            &plain,
940            [
941                vec![true, false, false, false],
942                vec![true, true, false, false],
943                vec![false, true, true, false],
944                vec![false, false, true, false],
945            ],
946        );
947
948        let dictionary = DictionaryArray::new(
949            datatypes::arrow::array::UInt32Array::from(vec![
950                Some(0),
951                Some(1),
952                Some(2),
953                None,
954                Some(3),
955            ]),
956            Arc::new(datatypes::arrow::array::StringArray::from(vec![
957                Some("api"),
958                Some("API"),
959                Some("db"),
960                None,
961            ])),
962        );
963        let raw =
964            regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap();
965        assert!(raw.is_null(3)); // null dictionary key
966        assert!(raw.is_null(4)); // non-null key referencing a null dictionary value
967
968        let dictionary = Arc::new(dictionary) as ArrayRef;
969        assert_regex_scan_masks(
970            &dictionary,
971            [
972                vec![true, false, false, false, false],
973                vec![true, true, false, false, false],
974                vec![false, true, true, false, false],
975                vec![false, false, true, false, false],
976            ],
977        );
978
979        let negative = regex_evaluator(Operator::RegexNotMatch);
980        assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap());
981    }
982
983    #[test]
984    fn test_nullable_boolean_predicate_becomes_scan_mask() {
985        let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]);
986        assert_eq!(
987            BooleanArray::from(vec![true, false, false]),
988            boolean_array_to_scan_mask(&predicate)
989        );
990    }
991
992    fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec<bool>; 4]) {
993        for (op, expected) in [
994            Operator::RegexMatch,
995            Operator::RegexIMatch,
996            Operator::RegexNotMatch,
997            Operator::RegexNotIMatch,
998        ]
999        .into_iter()
1000        .zip(expected)
1001        {
1002            assert_eq!(
1003                BooleanBuffer::from(expected),
1004                regex_evaluator(op).evaluate_array(input).unwrap(),
1005                "{op:?}"
1006            );
1007        }
1008    }
1009
1010    fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator {
1011        let expr = Expr::BinaryExpr(BinaryExpr {
1012            left: Box::new(Expr::Column(Column::from_name("host"))),
1013            op,
1014            right: Box::new("^api$".lit()),
1015        });
1016        SimpleFilterEvaluator::try_new(&expr).unwrap()
1017    }
1018
1019    fn ts_us(v: i64) -> ScalarValue {
1020        ScalarValue::TimestampMicrosecond(Some(v), None)
1021    }
1022
1023    fn cast_to_ms(expr: &Expr) -> Option<TimestampUnitCast> {
1024        SimpleFilterEvaluator::try_new(expr)
1025            .unwrap()
1026            .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype())
1027    }
1028
1029    /// Evaluates a cast `Filter` outcome against millisecond values and
1030    /// returns the mask.
1031    fn eval_ms_mask(cast: Option<TimestampUnitCast>, values: &[i64]) -> Vec<bool> {
1032        let filter = match cast.expect("cast must apply") {
1033            TimestampUnitCast::Filter(filter) => filter,
1034            other => panic!("expected Filter outcome, got {other:?}"),
1035        };
1036        let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef;
1037        filter.evaluate_array(&array).unwrap().iter().collect()
1038    }
1039
1040    #[test]
1041    fn cast_timestamp_unit_converts_representable_literal() {
1042        // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms.
1043        assert_eq!(
1044            vec![false, true, false],
1045            eval_ms_mask(
1046                cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))),
1047                &[6_999, 7_000, 7_001]
1048            )
1049        );
1050        // != complements =.
1051        assert_eq!(
1052            vec![true, false, true],
1053            eval_ms_mask(
1054                cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))),
1055                &[6_999, 7_000, 7_001]
1056            )
1057        );
1058    }
1059
1060    #[test]
1061    fn cast_timestamp_unit_prunes_non_representable_equality() {
1062        // 7_000_500us is not a whole millisecond: no ms row can equal it,
1063        // and every ms row satisfies `!=`.
1064        assert!(matches!(
1065            cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))),
1066            Some(TimestampUnitCast::Pruned)
1067        ));
1068        assert!(matches!(
1069            cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))),
1070            Some(TimestampUnitCast::Matched)
1071        ));
1072    }
1073
1074    #[test]
1075    fn cast_timestamp_unit_strengthens_inequalities() {
1076        // 2_500_500us is strictly between 2500ms and 2501ms: the cast must
1077        // not round the literal to a boundary the original predicate
1078        // excludes (>= 2500ms would wrongly match 2500ms).
1079        let values = vec![2_499, 2_500, 2_501];
1080        assert_eq!(
1081            vec![false, false, true],
1082            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values)
1083        );
1084        assert_eq!(
1085            vec![false, false, true],
1086            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values)
1087        );
1088        assert_eq!(
1089            vec![true, true, false],
1090            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values)
1091        );
1092        assert_eq!(
1093            vec![true, true, false],
1094            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values)
1095        );
1096
1097        // Representable boundary keeps the original operator.
1098        let values = vec![2_499, 2_500, 2_501];
1099        assert_eq!(
1100            vec![false, false, true],
1101            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values)
1102        );
1103        assert_eq!(
1104            vec![false, true, true],
1105            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values)
1106        );
1107        assert_eq!(
1108            vec![true, false, false],
1109            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values)
1110        );
1111        assert_eq!(
1112            vec![true, true, false],
1113            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values)
1114        );
1115    }
1116
1117    #[test]
1118    fn cast_timestamp_unit_handles_negative_instants() {
1119        // -2_500_500us floors to -2501ms with remainder 500us: only rows
1120        // with instant >= -2.5005ms may match >= / >.
1121        let values = vec![-2_502, -2_501, -2_500];
1122        assert_eq!(
1123            vec![false, false, true],
1124            eval_ms_mask(
1125                cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))),
1126                &values
1127            )
1128        );
1129        assert_eq!(
1130            vec![false, false, true],
1131            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values)
1132        );
1133        assert_eq!(
1134            vec![true, true, false],
1135            eval_ms_mask(
1136                cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))),
1137                &values
1138            )
1139        );
1140        // Exactly representable negative literal: -2_500_000us == -2500ms.
1141        assert_eq!(
1142            vec![false, false, true],
1143            eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values)
1144        );
1145    }
1146
1147    #[test]
1148    fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() {
1149        // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match.
1150        let expr = col("ts")
1151            .eq(lit(ts_us(7_000_500)))
1152            .or(col("ts").eq(lit(ts_us(7_000_000))));
1153        assert_eq!(
1154            vec![false, true],
1155            eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000])
1156        );
1157
1158        // A chain where no literal is representable matches nothing.
1159        let expr = col("ts")
1160            .eq(lit(ts_us(7_000_500)))
1161            .or(col("ts").eq(lit(ts_us(6_000_500))));
1162        assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned)));
1163    }
1164
1165    #[test]
1166    fn cast_timestamp_unit_null_literal_prunes() {
1167        // A NULL literal never compares true.
1168        assert!(matches!(
1169            cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))),
1170            Some(TimestampUnitCast::Pruned)
1171        ));
1172    }
1173
1174    #[test]
1175    fn cast_timestamp_unit_same_unit_returns_filter_directly() {
1176        // A literal already in the target unit is returned unchanged.
1177        let expr = col("ts").eq(lit(ts_us(7_000_000)));
1178        let filter = SimpleFilterEvaluator::try_new(&expr).unwrap();
1179        let cast = filter
1180            .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype())
1181            .unwrap();
1182        let TimestampUnitCast::Filter(f) = cast else {
1183            panic!("expected Filter, got {cast:?}")
1184        };
1185        assert_eq!(filter.op, f.op);
1186        assert_eq!(filter.column_name(), f.column_name());
1187        assert_eq!(
1188            ts_us(7_000_000),
1189            ScalarValue::try_from_array(f.literal.get().0, 0).unwrap()
1190        );
1191    }
1192
1193    #[test]
1194    fn cast_timestamp_unit_rejects_non_timestamps() {
1195        // Non-timestamp target type.
1196        assert!(
1197            SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1))))
1198                .unwrap()
1199                .cast_timestamp_unit(&ConcreteDataType::int64_datatype())
1200                .is_none()
1201        );
1202        // Non-timestamp literal against a timestamp target.
1203        assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none());
1204    }
1205
1206    #[test]
1207    fn cast_timestamp_unit_to_finer_unit() {
1208        // 5 seconds against a nanosecond column becomes 5_000_000_000ns.
1209        let filter = SimpleFilterEvaluator::try_new(
1210            &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))),
1211        )
1212        .unwrap()
1213        .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype())
1214        .and_then(|cast| match cast {
1215            TimestampUnitCast::Filter(filter) => Some(filter),
1216            _ => None,
1217        })
1218        .unwrap();
1219        let array = Arc::new(TimestampNanosecondArray::from(vec![
1220            4_999_999_999,
1221            5_000_000_000,
1222        ])) as ArrayRef;
1223        assert_eq!(
1224            vec![false, true],
1225            filter
1226                .evaluate_array(&array)
1227                .unwrap()
1228                .iter()
1229                .collect::<Vec<_>>()
1230        );
1231    }
1232}
common_recordbatch/filter.rs

common_recordbatch/
filter.rs
