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::physical_planning_context::PhysicalPlanningContext;
620    use datafusion::logical_expr::{BinaryExpr, col, lit};
621    use datafusion::physical_expr::create_physical_expr;
622    use datafusion_common::{Column, DFSchema};
623    use datatypes::arrow::array::{TimestampMillisecondArray, TimestampNanosecondArray};
624    use datatypes::arrow::datatypes::{DataType, Field, Schema};
625
626    use super::*;
627
628    #[test]
629    fn unsupported_filter_op() {
630        // `+` is not supported
631        let expr = Expr::BinaryExpr(BinaryExpr {
632            left: Box::new(Expr::Column(Column::from_name("foo"))),
633            op: Operator::Plus,
634            right: Box::new(1.lit()),
635        });
636        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
637
638        // two literal is not supported
639        let expr = Expr::BinaryExpr(BinaryExpr {
640            left: Box::new(1.lit()),
641            op: Operator::Eq,
642            right: Box::new(1.lit()),
643        });
644        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
645
646        // two column is not supported
647        let expr = Expr::BinaryExpr(BinaryExpr {
648            left: Box::new(Expr::Column(Column::from_name("foo"))),
649            op: Operator::Eq,
650            right: Box::new(Expr::Column(Column::from_name("bar"))),
651        });
652        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
653
654        // compound expr is not supported
655        let expr = Expr::BinaryExpr(BinaryExpr {
656            left: Box::new(Expr::BinaryExpr(BinaryExpr {
657                left: Box::new(Expr::Column(Column::from_name("foo"))),
658                op: Operator::Eq,
659                right: Box::new(1.lit()),
660            })),
661            op: Operator::Eq,
662            right: Box::new(1.lit()),
663        });
664        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
665    }
666
667    #[test]
668    fn supported_filter_op() {
669        // equal
670        let expr = Expr::BinaryExpr(BinaryExpr {
671            left: Box::new(Expr::Column(Column::from_name("foo"))),
672            op: Operator::Eq,
673            right: Box::new(1.lit()),
674        });
675        let _ = SimpleFilterEvaluator::try_new(&expr).unwrap();
676
677        // swap operands
678        let expr = Expr::BinaryExpr(BinaryExpr {
679            left: Box::new(1.lit()),
680            op: Operator::Lt,
681            right: Box::new(Expr::Column(Column::from_name("foo"))),
682        });
683        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
684        assert_eq!(evaluator.op, Operator::Gt);
685        assert_eq!(evaluator.column_name, "foo".to_string());
686    }
687
688    #[test]
689    fn run_on_array() {
690        let expr = Expr::BinaryExpr(BinaryExpr {
691            left: Box::new(Expr::Column(Column::from_name("foo"))),
692            op: Operator::Eq,
693            right: Box::new(1i64.lit()),
694        });
695        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
696
697        let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _;
698        let result = evaluator.evaluate_array(&input_1).unwrap();
699        assert_eq!(result, BooleanBuffer::from(vec![true, false, false]));
700
701        let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _;
702        let result = evaluator.evaluate_array(&input_2).unwrap();
703        assert_eq!(result, BooleanBuffer::from(vec![true, true, true]));
704
705        let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _;
706        let result = evaluator.evaluate_array(&input_3).unwrap();
707        assert_eq!(result, BooleanBuffer::from(vec![]));
708    }
709
710    #[test]
711    fn run_on_scalar() {
712        let expr = Expr::BinaryExpr(BinaryExpr {
713            left: Box::new(Expr::Column(Column::from_name("foo"))),
714            op: Operator::Lt,
715            right: Box::new(1i64.lit()),
716        });
717        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
718
719        let input_1 = ScalarValue::Int64(Some(1));
720        let result = evaluator.evaluate_scalar(&input_1).unwrap();
721        assert!(!result);
722
723        let input_2 = ScalarValue::Int64(Some(0));
724        let result = evaluator.evaluate_scalar(&input_2).unwrap();
725        assert!(result);
726
727        let input_3 = ScalarValue::Int64(None);
728        let result = evaluator.evaluate_scalar(&input_3).unwrap();
729        assert!(!result);
730    }
731
732    #[test]
733    fn batch_filter_test() {
734        let expr = col("ts").gt(lit(123456u64));
735        let schema = Schema::new(vec![
736            Field::new("a", DataType::Int32, true),
737            Field::new("ts", DataType::UInt64, false),
738        ]);
739        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
740        let props = ExecutionProps::new();
741        let physical_expr = create_physical_expr(
742            &expr,
743            &df_schema,
744            &props,
745            &PhysicalPlanningContext::default(),
746        )
747        .unwrap();
748        let batch = RecordBatch::try_new(
749            Arc::new(schema),
750            vec![
751                Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])),
752                Arc::new(datatypes::arrow::array::UInt64Array::from(vec![
753                    123456, 123457, 123458,
754                ])),
755            ],
756        )
757        .unwrap();
758        let new_batch = batch_filter(&batch, &physical_expr).unwrap();
759        assert_eq!(new_batch.num_rows(), 2);
760        let first_column_values = new_batch
761            .column(0)
762            .as_any()
763            .downcast_ref::<datatypes::arrow::array::Int32Array>()
764            .unwrap();
765        let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]);
766        assert_eq!(first_column_values, &expected);
767    }
768
769    #[test]
770    fn test_complex_filter_expression() {
771        // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D'
772        let col_eq_b = col("col").eq(lit("B"));
773        let col_eq_c = col("col").eq(lit("C"));
774        let col_eq_d = col("col").eq(lit("D"));
775
776        // Build the OR chain
777        let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d);
778
779        // Check that SimpleFilterEvaluator can handle OR chain
780        let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap();
781        assert_eq!(or_evaluator.column_name, "col");
782        assert_eq!(or_evaluator.op, Operator::Or);
783        assert_eq!(or_evaluator.literal_list.len(), 3);
784        assert_eq!(
785            format!("{:?}", or_evaluator.literal_list),
786            "[Scalar(StringArray\n[\n  \"B\",\n]), Scalar(StringArray\n[\n  \"C\",\n]), Scalar(StringArray\n[\n  \"D\",\n])]"
787        );
788
789        // Create a schema and batch for testing
790        let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
791        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
792        let props = ExecutionProps::new();
793        let physical_expr = create_physical_expr(
794            &col_or_expr,
795            &df_schema,
796            &props,
797            &PhysicalPlanningContext::default(),
798        )
799        .unwrap();
800
801        // Create test data
802        let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![
803            "B", "C", "E", "B", "C", "D", "F",
804        ]));
805        let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap();
806        let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]);
807
808        // Filter the batch
809        let filtered_batch = batch_filter(&batch, &physical_expr).unwrap();
810
811        // Expected: rows with col in ("B", "C", "D")
812        // That would be rows 0, 1, 3, 4, 5
813        assert_eq!(filtered_batch.num_rows(), 5);
814
815        let col_filtered = filtered_batch
816            .column(0)
817            .as_any()
818            .downcast_ref::<datatypes::arrow::array::StringArray>()
819            .unwrap();
820        assert_eq!(col_filtered, &expected);
821    }
822
823    #[test]
824    fn test_maybe_build_regex() {
825        // Test case for RegexMatch (case sensitive, non-negative)
826        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
827            Operator::RegexMatch,
828            &ScalarValue::Utf8(Some("a.*b".to_string())),
829        )
830        .unwrap();
831        assert!(regex.is_some());
832        assert!(!negative);
833        assert!(regex.unwrap().is_match("axxb"));
834
835        // Test case for RegexIMatch (case insensitive, non-negative)
836        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
837            Operator::RegexIMatch,
838            &ScalarValue::Utf8(Some("a.*b".to_string())),
839        )
840        .unwrap();
841        assert!(regex.is_some());
842        assert!(!negative);
843        assert!(regex.unwrap().is_match("AxxB"));
844
845        // Test case for RegexNotMatch (case sensitive, negative)
846        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
847            Operator::RegexNotMatch,
848            &ScalarValue::Utf8(Some("a.*b".to_string())),
849        )
850        .unwrap();
851        assert!(regex.is_some());
852        assert!(negative);
853
854        // Test case for RegexNotIMatch (case insensitive, negative)
855        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
856            Operator::RegexNotIMatch,
857            &ScalarValue::Utf8(Some("a.*b".to_string())),
858        )
859        .unwrap();
860        assert!(regex.is_some());
861        assert!(negative);
862
863        // Test with empty regex pattern
864        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
865            Operator::RegexMatch,
866            &ScalarValue::Utf8(Some("".to_string())),
867        )
868        .unwrap();
869        assert!(regex.is_none());
870        assert!(!negative);
871
872        // Test with non-regex operator
873        let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex(
874            Operator::Eq,
875            &ScalarValue::Utf8(Some("a.*b".to_string())),
876        )
877        .unwrap();
878        assert!(regex.is_none());
879        assert!(!negative);
880
881        // Test with invalid regex pattern
882        let result = SimpleFilterEvaluator::maybe_build_regex(
883            Operator::RegexMatch,
884            &ScalarValue::Utf8(Some("a(b".to_string())),
885        );
886        assert!(result.is_err());
887
888        // Test with non-string value
889        let result = SimpleFilterEvaluator::maybe_build_regex(
890            Operator::RegexMatch,
891            &ScalarValue::Int64(Some(123)),
892        );
893        assert!(result.is_err());
894
895        // Test with null value
896        let result = SimpleFilterEvaluator::maybe_build_regex(
897            Operator::RegexMatch,
898            &ScalarValue::Utf8(None),
899        );
900        assert!(result.is_err());
901    }
902
903    #[test]
904    fn test_regex_match_dictionary_array() {
905        use datatypes::arrow::array::StringDictionaryBuilder;
906
907        // Create a StringDictionaryArray
908        let mut builder = StringDictionaryBuilder::<UInt32Type>::new();
909        builder.append("apple").unwrap();
910        builder.append("banana").unwrap();
911        builder.append("apple").unwrap();
912        builder.append("cherry").unwrap();
913        let dict_array = builder.finish();
914
915        // Test regex that matches "apple"
916        let regex = regex::Regex::new(r"app.*").unwrap();
917        let result = regexp_is_match_dictionary(&dict_array, Some(&regex)).unwrap();
918
919        // Should match indices 0 and 2 (both "apple")
920        assert_eq!(result.len(), 4);
921        assert!(result.value(0)); // "apple"
922        assert!(!result.value(1)); // "banana"
923        assert!(result.value(2)); // "apple"
924        assert!(!result.value(3)); // "cherry"
925
926        // Test regex that matches "banana"
927        let regex2 = regex::Regex::new(r"ban.*").unwrap();
928        let result2 = regexp_is_match_dictionary(&dict_array, Some(&regex2)).unwrap();
929
930        assert!(!result2.value(0)); // "apple"
931        assert!(result2.value(1)); // "banana"
932        assert!(!result2.value(2)); // "apple"
933        assert!(!result2.value(3)); // "cherry"
934
935        // Test with no regex (should match all)
936        let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap();
937        assert!(result3.value(0));
938        assert!(result3.value(1));
939        assert!(result3.value(2));
940        assert!(result3.value(3));
941    }
942
943    #[test]
944    fn test_regex_scan_masks_preserve_sql_null_semantics() {
945        let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![
946            Some("api"),
947            Some("API"),
948            Some("db"),
949            None,
950        ])) as ArrayRef;
951        assert_regex_scan_masks(
952            &plain,
953            [
954                vec![true, false, false, false],
955                vec![true, true, false, false],
956                vec![false, true, true, false],
957                vec![false, false, true, false],
958            ],
959        );
960
961        let dictionary = DictionaryArray::new(
962            datatypes::arrow::array::UInt32Array::from(vec![
963                Some(0),
964                Some(1),
965                Some(2),
966                None,
967                Some(3),
968            ]),
969            Arc::new(datatypes::arrow::array::StringArray::from(vec![
970                Some("api"),
971                Some("API"),
972                Some("db"),
973                None,
974            ])),
975        );
976        let raw =
977            regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap();
978        assert!(raw.is_null(3)); // null dictionary key
979        assert!(raw.is_null(4)); // non-null key referencing a null dictionary value
980
981        let dictionary = Arc::new(dictionary) as ArrayRef;
982        assert_regex_scan_masks(
983            &dictionary,
984            [
985                vec![true, false, false, false, false],
986                vec![true, true, false, false, false],
987                vec![false, true, true, false, false],
988                vec![false, false, true, false, false],
989            ],
990        );
991
992        let negative = regex_evaluator(Operator::RegexNotMatch);
993        assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap());
994    }
995
996    #[test]
997    fn test_nullable_boolean_predicate_becomes_scan_mask() {
998        let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]);
999        assert_eq!(
1000            BooleanArray::from(vec![true, false, false]),
1001            boolean_array_to_scan_mask(&predicate)
1002        );
1003    }
1004
1005    fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec<bool>; 4]) {
1006        for (op, expected) in [
1007            Operator::RegexMatch,
1008            Operator::RegexIMatch,
1009            Operator::RegexNotMatch,
1010            Operator::RegexNotIMatch,
1011        ]
1012        .into_iter()
1013        .zip(expected)
1014        {
1015            assert_eq!(
1016                BooleanBuffer::from(expected),
1017                regex_evaluator(op).evaluate_array(input).unwrap(),
1018                "{op:?}"
1019            );
1020        }
1021    }
1022
1023    fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator {
1024        let expr = Expr::BinaryExpr(BinaryExpr {
1025            left: Box::new(Expr::Column(Column::from_name("host"))),
1026            op,
1027            right: Box::new("^api$".lit()),
1028        });
1029        SimpleFilterEvaluator::try_new(&expr).unwrap()
1030    }
1031
1032    fn ts_us(v: i64) -> ScalarValue {
1033        ScalarValue::TimestampMicrosecond(Some(v), None)
1034    }
1035
1036    fn cast_to_ms(expr: &Expr) -> Option<TimestampUnitCast> {
1037        SimpleFilterEvaluator::try_new(expr)
1038            .unwrap()
1039            .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype())
1040    }
1041
1042    /// Evaluates a cast `Filter` outcome against millisecond values and
1043    /// returns the mask.
1044    fn eval_ms_mask(cast: Option<TimestampUnitCast>, values: &[i64]) -> Vec<bool> {
1045        let filter = match cast.expect("cast must apply") {
1046            TimestampUnitCast::Filter(filter) => filter,
1047            other => panic!("expected Filter outcome, got {other:?}"),
1048        };
1049        let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef;
1050        filter.evaluate_array(&array).unwrap().iter().collect()
1051    }
1052
1053    #[test]
1054    fn cast_timestamp_unit_converts_representable_literal() {
1055        // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms.
1056        assert_eq!(
1057            vec![false, true, false],
1058            eval_ms_mask(
1059                cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))),
1060                &[6_999, 7_000, 7_001]
1061            )
1062        );
1063        // != complements =.
1064        assert_eq!(
1065            vec![true, false, true],
1066            eval_ms_mask(
1067                cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))),
1068                &[6_999, 7_000, 7_001]
1069            )
1070        );
1071    }
1072
1073    #[test]
1074    fn cast_timestamp_unit_prunes_non_representable_equality() {
1075        // 7_000_500us is not a whole millisecond: no ms row can equal it,
1076        // and every ms row satisfies `!=`.
1077        assert!(matches!(
1078            cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))),
1079            Some(TimestampUnitCast::Pruned)
1080        ));
1081        assert!(matches!(
1082            cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))),
1083            Some(TimestampUnitCast::Matched)
1084        ));
1085    }
1086
1087    #[test]
1088    fn cast_timestamp_unit_strengthens_inequalities() {
1089        // 2_500_500us is strictly between 2500ms and 2501ms: the cast must
1090        // not round the literal to a boundary the original predicate
1091        // excludes (>= 2500ms would wrongly match 2500ms).
1092        let values = vec![2_499, 2_500, 2_501];
1093        assert_eq!(
1094            vec![false, false, true],
1095            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values)
1096        );
1097        assert_eq!(
1098            vec![false, false, true],
1099            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values)
1100        );
1101        assert_eq!(
1102            vec![true, true, false],
1103            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values)
1104        );
1105        assert_eq!(
1106            vec![true, true, false],
1107            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values)
1108        );
1109
1110        // Representable boundary keeps the original operator.
1111        let values = vec![2_499, 2_500, 2_501];
1112        assert_eq!(
1113            vec![false, false, true],
1114            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values)
1115        );
1116        assert_eq!(
1117            vec![false, true, true],
1118            eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values)
1119        );
1120        assert_eq!(
1121            vec![true, false, false],
1122            eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values)
1123        );
1124        assert_eq!(
1125            vec![true, true, false],
1126            eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values)
1127        );
1128    }
1129
1130    #[test]
1131    fn cast_timestamp_unit_handles_negative_instants() {
1132        // -2_500_500us floors to -2501ms with remainder 500us: only rows
1133        // with instant >= -2.5005ms may match >= / >.
1134        let values = vec![-2_502, -2_501, -2_500];
1135        assert_eq!(
1136            vec![false, false, true],
1137            eval_ms_mask(
1138                cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))),
1139                &values
1140            )
1141        );
1142        assert_eq!(
1143            vec![false, false, true],
1144            eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values)
1145        );
1146        assert_eq!(
1147            vec![true, true, false],
1148            eval_ms_mask(
1149                cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))),
1150                &values
1151            )
1152        );
1153        // Exactly representable negative literal: -2_500_000us == -2500ms.
1154        assert_eq!(
1155            vec![false, false, true],
1156            eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values)
1157        );
1158    }
1159
1160    #[test]
1161    fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() {
1162        // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match.
1163        let expr = col("ts")
1164            .eq(lit(ts_us(7_000_500)))
1165            .or(col("ts").eq(lit(ts_us(7_000_000))));
1166        assert_eq!(
1167            vec![false, true],
1168            eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000])
1169        );
1170
1171        // A chain where no literal is representable matches nothing.
1172        let expr = col("ts")
1173            .eq(lit(ts_us(7_000_500)))
1174            .or(col("ts").eq(lit(ts_us(6_000_500))));
1175        assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned)));
1176    }
1177
1178    #[test]
1179    fn cast_timestamp_unit_null_literal_prunes() {
1180        // A NULL literal never compares true.
1181        assert!(matches!(
1182            cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))),
1183            Some(TimestampUnitCast::Pruned)
1184        ));
1185    }
1186
1187    #[test]
1188    fn cast_timestamp_unit_same_unit_returns_filter_directly() {
1189        // A literal already in the target unit is returned unchanged.
1190        let expr = col("ts").eq(lit(ts_us(7_000_000)));
1191        let filter = SimpleFilterEvaluator::try_new(&expr).unwrap();
1192        let cast = filter
1193            .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype())
1194            .unwrap();
1195        let TimestampUnitCast::Filter(f) = cast else {
1196            panic!("expected Filter, got {cast:?}")
1197        };
1198        assert_eq!(filter.op, f.op);
1199        assert_eq!(filter.column_name(), f.column_name());
1200        assert_eq!(
1201            ts_us(7_000_000),
1202            ScalarValue::try_from_array(f.literal.get().0, 0).unwrap()
1203        );
1204    }
1205
1206    #[test]
1207    fn cast_timestamp_unit_rejects_non_timestamps() {
1208        // Non-timestamp target type.
1209        assert!(
1210            SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1))))
1211                .unwrap()
1212                .cast_timestamp_unit(&ConcreteDataType::int64_datatype())
1213                .is_none()
1214        );
1215        // Non-timestamp literal against a timestamp target.
1216        assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none());
1217    }
1218
1219    #[test]
1220    fn cast_timestamp_unit_to_finer_unit() {
1221        // 5 seconds against a nanosecond column becomes 5_000_000_000ns.
1222        let filter = SimpleFilterEvaluator::try_new(
1223            &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))),
1224        )
1225        .unwrap()
1226        .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype())
1227        .and_then(|cast| match cast {
1228            TimestampUnitCast::Filter(filter) => Some(filter),
1229            _ => None,
1230        })
1231        .unwrap();
1232        let array = Arc::new(TimestampNanosecondArray::from(vec![
1233            4_999_999_999,
1234            5_000_000_000,
1235        ])) as ArrayRef;
1236        assert_eq!(
1237            vec![false, true],
1238            filter
1239                .evaluate_array(&array)
1240                .unwrap()
1241                .iter()
1242                .collect::<Vec<_>>()
1243        );
1244    }
1245}
common_recordbatch/filter.rs

common_recordbatch/
filter.rs
