common_recordbatch/filter/
predicate.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//! Predicate recognition, execution, and precompiled regular expressions.
16
17use datafusion::logical_expr::{Expr, Operator};
18use datafusion_common::ScalarValue;
19use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar};
20use datafusion_common::arrow::buffer::BooleanBuffer;
21use datafusion_common::arrow::compute::kernels::cmp;
22use datafusion_common::cast::as_string_array;
23use datatypes::arrow::array::{
24    Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray,
25    StringArrayType,
26};
27use datatypes::arrow::datatypes::{DataType, UInt32Type};
28use datatypes::arrow::error::ArrowError;
29use datatypes::compute::or_kleene;
30use regex::Regex;
31use snafu::ResultExt;
32
33use crate::error::{ArrowComputeSnafu, Result, UnsupportedOperationSnafu};
34
35#[derive(Debug, Clone)]
36pub(crate) enum SimplePredicate {
37    Comparison {
38        literal: Scalar<ArrayRef>,
39        op: Operator,
40        regex: Option<Regex>,
41        regex_negative: bool,
42    },
43    OrEqChain {
44        literals: Vec<Scalar<ArrayRef>>,
45    },
46    IsNull {
47        negated: bool,
48    },
49    Constant(Option<bool>),
50    And(Box<Self>, Box<Self>),
51    Or(Box<Self>, Box<Self>),
52}
53
54impl SimplePredicate {
55    pub(crate) fn try_new(expr: &Expr, column_type: Option<&DataType>) -> Option<Self> {
56        match expr {
57            Expr::BinaryExpr(binary) if matches!(binary.op, Operator::And | Operator::Or) => {
58                let left = Self::try_new(&binary.left, column_type)?;
59                let right = Self::try_new(&binary.right, column_type)?;
60                if binary.op == Operator::And {
61                    return Some(Self::And(Box::new(left), Box::new(right)));
62                }
63                // Preserve the encoded-key IN-list fast path for pure equality disjunctions.
64                if let (Some(mut lhs), Some(rhs)) =
65                    (left.equality_literals(), right.equality_literals())
66                {
67                    lhs.extend(rhs);
68                    return Some(Self::OrEqChain { literals: lhs });
69                }
70                Some(Self::Or(Box::new(left), Box::new(right)))
71            }
72            Expr::BinaryExpr(binary) => {
73                let mut op = binary.op;
74                if !matches!(
75                    op,
76                    Operator::Eq
77                        | Operator::NotEq
78                        | Operator::Lt
79                        | Operator::LtEq
80                        | Operator::Gt
81                        | Operator::GtEq
82                        | Operator::RegexMatch
83                        | Operator::RegexIMatch
84                        | Operator::RegexNotMatch
85                        | Operator::RegexNotIMatch
86                ) {
87                    return None;
88                }
89                let literal = if let Expr::Literal(literal, _) = &*binary.right {
90                    Self::check_column(&binary.left, column_type)?;
91                    literal
92                } else if let Expr::Literal(literal, _) = &*binary.left {
93                    Self::check_column(&binary.right, column_type)?;
94                    op = op.swap()?;
95                    literal
96                } else {
97                    return None;
98                };
99                let (regex, regex_negative) = Self::maybe_build_regex(op, literal).ok()?;
100                Some(Self::Comparison {
101                    literal: literal.to_scalar().ok()?,
102                    op,
103                    regex,
104                    regex_negative,
105                })
106            }
107            Expr::IsNull(column) | Expr::IsNotNull(column) => {
108                Self::check_column(column, column_type)?;
109                Some(Self::IsNull {
110                    negated: matches!(expr, Expr::IsNotNull(_)),
111                })
112            }
113            Expr::Literal(ScalarValue::Boolean(value), _) => Some(Self::Constant(*value)),
114            _ => None,
115        }
116    }
117
118    fn check_column(expr: &Expr, column_type: Option<&DataType>) -> Option<()> {
119        match expr {
120            Expr::Column(_) => Some(()),
121            Expr::Cast(cast) if cast.field.data_type() == &DataType::Utf8 => {
122                let Expr::Column(_) = &*cast.expr else {
123                    return None;
124                };
125                match column_type? {
126                    DataType::Utf8 => Some(()),
127                    DataType::Dictionary(key, value)
128                        if **key == DataType::UInt32 && **value == DataType::Utf8 =>
129                    {
130                        Some(())
131                    }
132                    _ => None,
133                }
134            }
135            _ => None,
136        }
137    }
138
139    fn equality_literals(&self) -> Option<Vec<Scalar<ArrayRef>>> {
140        match self {
141            Self::Comparison {
142                literal,
143                op: Operator::Eq,
144                ..
145            } => Some(vec![literal.clone()]),
146            Self::OrEqChain { literals } => Some(literals.clone()),
147            _ => None,
148        }
149    }
150
151    pub(crate) fn evaluate(&self, input: &impl Datum, len: usize) -> Result<BooleanArray> {
152        match self {
153            Self::Comparison {
154                literal,
155                op,
156                regex,
157                regex_negative,
158            } => match op {
159                Operator::Eq => cmp::eq(input, literal),
160                Operator::NotEq => cmp::neq(input, literal),
161                Operator::Lt => cmp::lt(input, literal),
162                Operator::LtEq => cmp::lt_eq(input, literal),
163                Operator::Gt => cmp::gt(input, literal),
164                Operator::GtEq => cmp::gt_eq(input, literal),
165                Operator::RegexMatch
166                | Operator::RegexIMatch
167                | Operator::RegexNotMatch
168                | Operator::RegexNotIMatch => {
169                    Self::regex_match(input, regex.as_ref(), *regex_negative)
170                }
171                _ => {
172                    return UnsupportedOperationSnafu {
173                        reason: format!("{op:?}"),
174                    }
175                    .fail();
176                }
177            }
178            .context(ArrowComputeSnafu),
179            Self::OrEqChain { literals } => {
180                let mut result = BooleanArray::from(vec![false; len]);
181                for literal in literals {
182                    let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?;
183                    result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?;
184                }
185                Ok(result)
186            }
187            Self::IsNull { negated } => {
188                let values = match input.get().0.logical_nulls() {
189                    Some(nulls) if *negated => nulls.inner().clone(),
190                    Some(nulls) => !nulls.inner(),
191                    None if *negated => BooleanBuffer::new_set(len),
192                    None => BooleanBuffer::new_unset(len),
193                };
194                Ok(BooleanArray::new(values, None))
195            }
196            Self::Constant(value) => Ok(BooleanArray::from(vec![*value; len])),
197            Self::And(left, right) => datatypes::compute::and_kleene(
198                &left.evaluate(input, len)?,
199                &right.evaluate(input, len)?,
200            )
201            .context(ArrowComputeSnafu),
202            Self::Or(left, right) => {
203                or_kleene(&left.evaluate(input, len)?, &right.evaluate(input, len)?)
204                    .context(ArrowComputeSnafu)
205            }
206        }
207    }
208
209    /// Builds a regex pattern from a scalar value and operator.
210    /// Returns the `(regex, negative)` and if successful.
211    ///
212    /// Returns `Err` if
213    /// - the value is not a string
214    /// - the regex pattern is invalid
215    ///
216    /// The regex is `None` if
217    /// - the operator is not a regex operator
218    /// - the pattern is empty
219    fn maybe_build_regex(
220        operator: Operator,
221        value: &ScalarValue,
222    ) -> Result<(Option<Regex>, bool), ArrowError> {
223        let (ignore_case, negative) = match operator {
224            Operator::RegexMatch => (false, false),
225            Operator::RegexIMatch => (true, false),
226            Operator::RegexNotMatch => (false, true),
227            Operator::RegexNotIMatch => (true, true),
228            _ => return Ok((None, false)),
229        };
230        let flag = if ignore_case { Some("i") } else { None };
231        let regex = value
232            .try_as_str()
233            .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))?
234            .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?;
235        let pattern = match flag {
236            Some(flag) => format!("(?{flag}){regex}"),
237            None => regex.to_string(),
238        };
239        if pattern.is_empty() {
240            Ok((None, negative))
241        } else {
242            Regex::new(pattern.as_str())
243                .map_err(|e| {
244                    ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}"))
245                })
246                .map(|regex| (Some(regex), negative))
247        }
248    }
249
250    fn regex_match(
251        input: &impl Datum,
252        regex: Option<&Regex>,
253        negative: bool,
254    ) -> std::result::Result<BooleanArray, ArrowError> {
255        let array = input.get().0;
256
257        // Try to cast to StringArray first
258        if let Ok(string_array) = as_string_array(array) {
259            let mut result = regexp_is_match_scalar(string_array, regex)?;
260            if negative {
261                result = datatypes::compute::not(&result)?;
262            }
263            return Ok(result);
264        }
265
266        // Try to cast to StringDictionaryArray
267        if let Some(dict_array) = array.as_any().downcast_ref::<DictionaryArray<UInt32Type>>() {
268            let mut result = regexp_is_match_dictionary(dict_array, regex)?;
269            if negative {
270                result = datatypes::compute::not(&result)?;
271            }
272            return Ok(result);
273        }
274
275        Err(ArrowError::CastError(format!(
276            "Cannot cast {:?} to StringArray or StringDictionaryArray",
277            array.data_type()
278        )))
279    }
280}
281
282/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar())
283/// with pre-compiled regex.
284/// See <https://github.com/apache/arrow-rs/blob/54.2.0/arrow-string/src/regexp.rs#L204-L246> for the implementation details.
285pub fn regexp_is_match_scalar<'a, S>(
286    array: &'a S,
287    regex: Option<&Regex>,
288) -> Result<BooleanArray, ArrowError>
289where
290    &'a S: StringArrayType<'a>,
291{
292    let null_bit_buffer = array.nulls().map(|x| x.inner().sliced());
293    let mut result = BooleanBufferBuilder::new(array.len());
294
295    if let Some(re) = regex {
296        for i in 0..array.len() {
297            let value = array.value(i);
298            result.append(re.is_match(value));
299        }
300    } else {
301        result.append_n(array.len(), true);
302    }
303
304    let buffer = result.into();
305    let data = unsafe {
306        ArrayData::new_unchecked(
307            DataType::Boolean,
308            array.len(),
309            None,
310            null_bit_buffer,
311            0,
312            vec![buffer],
313            vec![],
314        )
315    };
316
317    Ok(BooleanArray::from(data))
318}
319
320/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray.
321/// Iterates through dictionary keys to get string values and applies regex matching.
322pub fn regexp_is_match_dictionary(
323    dict_array: &DictionaryArray<UInt32Type>,
324    regex: Option<&Regex>,
325) -> Result<BooleanArray, ArrowError> {
326    // Get the string values from the dictionary
327    let string_values = dict_array
328        .values()
329        .as_any()
330        .downcast_ref::<datatypes::arrow::array::StringArray>()
331        .ok_or_else(|| {
332            ArrowError::CastError("Dictionary values must be StringArray".to_string())
333        })?;
334
335    // Dictionary logical nulls include both null keys and keys whose dictionary value is null.
336    let logical_nulls = dict_array.logical_nulls();
337    let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced());
338    let mut result = BooleanBufferBuilder::new(dict_array.len());
339
340    if let Some(re) = regex {
341        let keys = dict_array.keys().values();
342        for i in 0..dict_array.len() {
343            if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) {
344                result.append(false);
345            } else {
346                let key = keys[i] as usize;
347                let string_value = string_values.value(key);
348                result.append(re.is_match(string_value));
349            }
350        }
351    } else {
352        result.append_n(dict_array.len(), true);
353    }
354
355    let buffer = result.into();
356    let data = unsafe {
357        ArrayData::new_unchecked(
358            DataType::Boolean,
359            dict_array.len(),
360            None,
361            null_bit_buffer,
362            0,
363            vec![buffer],
364            vec![],
365        )
366    };
367
368    Ok(BooleanArray::from(data))
369}
370
371#[cfg(test)]
372mod tests {
373    use std::sync::Arc;
374
375    use datafusion::execution::context::ExecutionProps;
376    use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
377    use datafusion::logical_expr::{BinaryExpr, Expr, Literal, Operator, col, lit};
378    use datafusion::physical_expr::create_physical_expr;
379    use datafusion_common::arrow::array::ArrayRef;
380    use datafusion_common::arrow::buffer::BooleanBuffer;
381    use datafusion_common::{Column, DFSchema, ScalarValue};
382    use datatypes::arrow::array::{Array, DictionaryArray, RecordBatch};
383    use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type};
384    use datatypes::value::Value;
385    use regex::Regex;
386
387    use crate::filter::predicate::{SimplePredicate, regexp_is_match_dictionary};
388    use crate::filter::{SimpleFilterEvaluator, batch_filter};
389
390    #[test]
391    fn unsupported_filter_op() {
392        // `+` is not supported
393        let expr = Expr::BinaryExpr(BinaryExpr {
394            left: Box::new(Expr::Column(Column::from_name("foo"))),
395            op: Operator::Plus,
396            right: Box::new(1.lit()),
397        });
398        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
399
400        // two literal is not supported
401        let expr = Expr::BinaryExpr(BinaryExpr {
402            left: Box::new(1.lit()),
403            op: Operator::Eq,
404            right: Box::new(1.lit()),
405        });
406        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
407
408        // two column is not supported
409        let expr = Expr::BinaryExpr(BinaryExpr {
410            left: Box::new(Expr::Column(Column::from_name("foo"))),
411            op: Operator::Eq,
412            right: Box::new(Expr::Column(Column::from_name("bar"))),
413        });
414        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
415
416        // compound expr is not supported
417        let expr = Expr::BinaryExpr(BinaryExpr {
418            left: Box::new(Expr::BinaryExpr(BinaryExpr {
419                left: Box::new(Expr::Column(Column::from_name("foo"))),
420                op: Operator::Eq,
421                right: Box::new(1.lit()),
422            })),
423            op: Operator::Eq,
424            right: Box::new(1.lit()),
425        });
426        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
427        for negated in [false, true] {
428            let membership = col("label").in_list(vec![lit(""), lit("lo")], negated);
429            assert!(SimpleFilterEvaluator::try_new(&membership).is_none());
430            assert!(
431                SimpleFilterEvaluator::try_new(&col("label").is_null().or(membership)).is_none()
432            );
433        }
434        for expr in [
435            col("other").is_null().or(col("label").not_eq(lit("lo"))),
436            col("a.label")
437                .is_null()
438                .or(col("b.label").not_eq(lit("lo"))),
439            col("label").is_null().or(col("label").eq(col("other"))),
440            col("label")
441                .is_null()
442                .or(Expr::Cast(datafusion::logical_expr::Cast::new(
443                    Box::new(col("label")),
444                    DataType::Int64,
445                ))
446                .eq(lit("lo"))),
447        ] {
448            assert!(SimpleFilterEvaluator::try_new(&expr).is_none(), "{expr}");
449        }
450    }
451
452    #[test]
453    fn string_cast_requires_matching_column_type() {
454        let cast = Expr::Cast(datafusion::logical_expr::Cast::new(
455            Box::new(col("label")),
456            DataType::Utf8,
457        ));
458        let expr = col("label").is_null().or(cast.not_eq(lit("lo")));
459        for data_type in [
460            DataType::Int64,
461            DataType::Timestamp(TimeUnit::Millisecond, None),
462            DataType::Binary,
463        ] {
464            assert!(
465                SimpleFilterEvaluator::try_new_with_column_type(&expr, &|_| Some(
466                    data_type.clone()
467                ))
468                .is_none()
469            );
470        }
471        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
472    }
473
474    #[tokio::test]
475    async fn single_column_predicates_match_datafusion() {
476        use datafusion::prelude::SessionContext;
477        use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
478        use datatypes::arrow::array::{Int32Array, StringArray, UInt32Array};
479
480        let values = vec![
481            None,
482            Some(""),
483            Some("lo"),
484            Some("eth0"),
485            Some("a\nb"),
486            Some("服务"),
487            None,
488        ];
489        let strings: ArrayRef = Arc::new(StringArray::from(values));
490        let dictionary: ArrayRef = Arc::new(DictionaryArray::<UInt32Type>::new(
491            UInt32Array::from(vec![
492                None,
493                Some(0),
494                Some(1),
495                Some(2),
496                Some(3),
497                Some(4),
498                Some(5),
499            ]),
500            Arc::new(StringArray::from(vec![
501                Some(""),
502                Some("lo"),
503                Some("eth0"),
504                Some("a\nb"),
505                Some("服务"),
506                None,
507            ])),
508        ));
509        let null = || lit(ScalarValue::Utf8(None));
510        let mut predicates = vec![
511            col("label").is_null(),
512            col("label").is_not_null(),
513            col("label")
514                .is_null()
515                .or(col("label").eq(lit("")))
516                .and(col("label").is_not_null()),
517            col("label")
518                .not_eq(lit("lo"))
519                .and(col("label").not_eq(lit("eth0"))),
520            col("label")
521                .eq(lit("lo"))
522                .or(col("label").eq(lit("eth0")))
523                .or(col("label").eq(lit(""))),
524            col("label").eq(null()).or(col("label").eq(lit("eth0"))),
525            col("label").gt(lit("a")).and(col("label").lt(lit("z"))),
526            lit("lo").not_eq(col("label")),
527        ];
528        for (op, literal) in [
529            (Operator::Eq, ""),
530            (Operator::NotEq, "lo"),
531            (Operator::RegexNotMatch, "^(?:lo|vfat)$"),
532            (Operator::RegexMatch, "^(?:|eth0)$"),
533        ] {
534            let comparison = Expr::BinaryExpr(BinaryExpr::new(
535                Box::new(col("label")),
536                op,
537                Box::new(lit(literal)),
538            ));
539            predicates.extend([
540                col("label").is_null().or(comparison.clone()),
541                comparison.or(col("label").is_null()),
542            ]);
543        }
544        for input in [strings, dictionary] {
545            let ctx = SessionContext::new();
546            let schema = Arc::new(Schema::new(vec![
547                Field::new("label", input.data_type().clone(), true),
548                Field::new("id", DataType::Int32, false),
549            ]));
550            ctx.register_batch(
551                "labels",
552                RecordBatch::try_new(
553                    schema,
554                    vec![
555                        input.clone(),
556                        Arc::new(Int32Array::from_iter_values(0..input.len() as i32)),
557                    ],
558                )
559                .unwrap(),
560            )
561            .unwrap();
562            for expr in &predicates {
563                let df = ctx
564                    .table("labels")
565                    .await
566                    .unwrap()
567                    .filter(expr.clone())
568                    .unwrap();
569                let plan = df.clone().into_optimized_plan().unwrap();
570                let expected = df
571                    .collect()
572                    .await
573                    .unwrap()
574                    .iter()
575                    .flat_map(|batch| {
576                        batch
577                            .column_by_name("id")
578                            .unwrap()
579                            .as_any()
580                            .downcast_ref::<Int32Array>()
581                            .unwrap()
582                            .values()
583                            .to_vec()
584                    })
585                    .collect::<Vec<_>>();
586                let evaluate = |expr: &Expr| {
587                    let evaluator =
588                        SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| {
589                            (name == "label").then(|| input.data_type().clone())
590                        })
591                        .unwrap_or_else(|| panic!("unsupported: {expr}"));
592                    let mask = evaluator.evaluate_array(&input).unwrap();
593                    for (i, matched) in mask.iter().enumerate() {
594                        let scalar = ScalarValue::try_from_array(&input, i).unwrap();
595                        assert_eq!(
596                            evaluator.evaluate_scalar(&scalar).unwrap(),
597                            matched,
598                            "{expr}, row={i}"
599                        );
600                    }
601                    mask
602                };
603                let selected = |mask: BooleanBuffer| {
604                    mask.iter()
605                        .enumerate()
606                        .filter_map(|(i, matched)| matched.then_some(i as i32))
607                        .collect::<Vec<_>>()
608                };
609                assert_eq!(selected(evaluate(expr)), expected, "{expr}");
610                let mut filters = Vec::new();
611                plan.apply(|plan| {
612                    if let datafusion::logical_expr::LogicalPlan::Filter(filter) = plan {
613                        filters.push(filter.predicate.clone());
614                    }
615                    Ok(TreeNodeRecursion::Continue)
616                })
617                .unwrap();
618                if !filters.is_empty() {
619                    let mut mask = BooleanBuffer::new_set(input.len());
620                    for filter in filters {
621                        mask = &mask & &evaluate(&filter);
622                    }
623                    assert_eq!(selected(mask), expected, "optimized: {plan}");
624                }
625            }
626        }
627    }
628
629    #[test]
630    fn supported_filter_op() {
631        // equal
632        let expr = Expr::BinaryExpr(BinaryExpr {
633            left: Box::new(Expr::Column(Column::from_name("foo"))),
634            op: Operator::Eq,
635            right: Box::new(1.lit()),
636        });
637        let _ = SimpleFilterEvaluator::try_new(&expr).unwrap();
638
639        // swap operands
640        let expr = Expr::BinaryExpr(BinaryExpr {
641            left: Box::new(1.lit()),
642            op: Operator::Lt,
643            right: Box::new(Expr::Column(Column::from_name("foo"))),
644        });
645        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
646        assert!(evaluator.is_gt());
647        assert_eq!(evaluator.column_name, "foo".to_string());
648    }
649
650    #[test]
651    fn run_on_array() {
652        let expr = Expr::BinaryExpr(BinaryExpr {
653            left: Box::new(Expr::Column(Column::from_name("foo"))),
654            op: Operator::Eq,
655            right: Box::new(1i64.lit()),
656        });
657        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
658
659        let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _;
660        let result = evaluator.evaluate_array(&input_1).unwrap();
661        assert_eq!(result, BooleanBuffer::from(vec![true, false, false]));
662
663        let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _;
664        let result = evaluator.evaluate_array(&input_2).unwrap();
665        assert_eq!(result, BooleanBuffer::from(vec![true, true, true]));
666
667        let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _;
668        let result = evaluator.evaluate_array(&input_3).unwrap();
669        assert_eq!(result, BooleanBuffer::from(vec![]));
670    }
671
672    #[test]
673    fn run_on_scalar() {
674        let expr = Expr::BinaryExpr(BinaryExpr {
675            left: Box::new(Expr::Column(Column::from_name("foo"))),
676            op: Operator::Lt,
677            right: Box::new(1i64.lit()),
678        });
679        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
680
681        let input_1 = ScalarValue::Int64(Some(1));
682        let result = evaluator.evaluate_scalar(&input_1).unwrap();
683        assert!(!result);
684
685        let input_2 = ScalarValue::Int64(Some(0));
686        let result = evaluator.evaluate_scalar(&input_2).unwrap();
687        assert!(result);
688
689        let input_3 = ScalarValue::Int64(None);
690        let result = evaluator.evaluate_scalar(&input_3).unwrap();
691        assert!(!result);
692    }
693
694    #[test]
695    fn test_complex_filter_expression() {
696        // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D'
697        let col_eq_b = col("col").eq(lit("B"));
698        let col_eq_c = col("col").eq(lit("C"));
699        let col_eq_d = col("col").eq(lit("D"));
700
701        // Build the OR chain
702        let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d);
703
704        // Check that SimpleFilterEvaluator can handle OR chain
705        let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap();
706        assert_eq!(or_evaluator.column_name, "col");
707        assert!(or_evaluator.is_or_eq_chain());
708        assert_eq!(
709            or_evaluator.literal_list_values().unwrap(),
710            vec![Value::from("B"), Value::from("C"), Value::from("D")]
711        );
712
713        // Create a schema and batch for testing
714        let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
715        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
716        let props = ExecutionProps::new();
717        let physical_expr = create_physical_expr(
718            &col_or_expr,
719            &df_schema,
720            &props,
721            &PhysicalPlanningContext::default(),
722        )
723        .unwrap();
724
725        // Create test data
726        let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![
727            "B", "C", "E", "B", "C", "D", "F",
728        ]));
729        let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap();
730        let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]);
731
732        // Filter the batch
733        let filtered_batch = batch_filter(&batch, &physical_expr).unwrap();
734
735        // Expected: rows with col in ("B", "C", "D")
736        // That would be rows 0, 1, 3, 4, 5
737        assert_eq!(filtered_batch.num_rows(), 5);
738
739        let col_filtered = filtered_batch
740            .column(0)
741            .as_any()
742            .downcast_ref::<datatypes::arrow::array::StringArray>()
743            .unwrap();
744        assert_eq!(col_filtered, &expected);
745    }
746
747    #[test]
748    fn test_maybe_build_regex() {
749        // Test case for RegexMatch (case sensitive, non-negative)
750        let (regex, negative) = SimplePredicate::maybe_build_regex(
751            Operator::RegexMatch,
752            &ScalarValue::Utf8(Some("a.*b".to_string())),
753        )
754        .unwrap();
755        assert!(regex.is_some());
756        assert!(!negative);
757        assert!(regex.unwrap().is_match("axxb"));
758
759        // Test case for RegexIMatch (case insensitive, non-negative)
760        let (regex, negative) = SimplePredicate::maybe_build_regex(
761            Operator::RegexIMatch,
762            &ScalarValue::Utf8(Some("a.*b".to_string())),
763        )
764        .unwrap();
765        assert!(regex.is_some());
766        assert!(!negative);
767        assert!(regex.unwrap().is_match("AxxB"));
768
769        // Test case for RegexNotMatch (case sensitive, negative)
770        let (regex, negative) = SimplePredicate::maybe_build_regex(
771            Operator::RegexNotMatch,
772            &ScalarValue::Utf8(Some("a.*b".to_string())),
773        )
774        .unwrap();
775        assert!(regex.is_some());
776        assert!(negative);
777
778        // Test case for RegexNotIMatch (case insensitive, negative)
779        let (regex, negative) = SimplePredicate::maybe_build_regex(
780            Operator::RegexNotIMatch,
781            &ScalarValue::Utf8(Some("a.*b".to_string())),
782        )
783        .unwrap();
784        assert!(regex.is_some());
785        assert!(negative);
786
787        // Test with empty regex pattern
788        let (regex, negative) = SimplePredicate::maybe_build_regex(
789            Operator::RegexMatch,
790            &ScalarValue::Utf8(Some("".to_string())),
791        )
792        .unwrap();
793        assert!(regex.is_none());
794        assert!(!negative);
795
796        // Test with non-regex operator
797        let (regex, negative) = SimplePredicate::maybe_build_regex(
798            Operator::Eq,
799            &ScalarValue::Utf8(Some("a.*b".to_string())),
800        )
801        .unwrap();
802        assert!(regex.is_none());
803        assert!(!negative);
804
805        // Test with invalid regex pattern
806        let result = SimplePredicate::maybe_build_regex(
807            Operator::RegexMatch,
808            &ScalarValue::Utf8(Some("a(b".to_string())),
809        );
810        assert!(result.is_err());
811
812        // Test with non-string value
813        let result = SimplePredicate::maybe_build_regex(
814            Operator::RegexMatch,
815            &ScalarValue::Int64(Some(123)),
816        );
817        assert!(result.is_err());
818
819        // Test with null value
820        let result =
821            SimplePredicate::maybe_build_regex(Operator::RegexMatch, &ScalarValue::Utf8(None));
822        assert!(result.is_err());
823    }
824
825    #[test]
826    fn test_regex_match_dictionary_array() {
827        use datatypes::arrow::array::StringDictionaryBuilder;
828
829        // Create a StringDictionaryArray
830        let mut builder = StringDictionaryBuilder::<UInt32Type>::new();
831        builder.append("apple").unwrap();
832        builder.append("banana").unwrap();
833        builder.append("apple").unwrap();
834        builder.append("cherry").unwrap();
835        let dict_array = builder.finish();
836
837        // Test regex that matches "apple"
838        let regex = regex::Regex::new(r"app.*").unwrap();
839        let result = regexp_is_match_dictionary(&dict_array, Some(&regex)).unwrap();
840
841        // Should match indices 0 and 2 (both "apple")
842        assert_eq!(result.len(), 4);
843        assert!(result.value(0)); // "apple"
844        assert!(!result.value(1)); // "banana"
845        assert!(result.value(2)); // "apple"
846        assert!(!result.value(3)); // "cherry"
847
848        // Test regex that matches "banana"
849        let regex2 = regex::Regex::new(r"ban.*").unwrap();
850        let result2 = regexp_is_match_dictionary(&dict_array, Some(&regex2)).unwrap();
851
852        assert!(!result2.value(0)); // "apple"
853        assert!(result2.value(1)); // "banana"
854        assert!(!result2.value(2)); // "apple"
855        assert!(!result2.value(3)); // "cherry"
856
857        // Test with no regex (should match all)
858        let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap();
859        assert!(result3.value(0));
860        assert!(result3.value(1));
861        assert!(result3.value(2));
862        assert!(result3.value(3));
863    }
864
865    #[test]
866    fn test_regex_scan_masks_preserve_sql_null_semantics() {
867        let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![
868            Some("api"),
869            Some("API"),
870            Some("db"),
871            None,
872        ])) as ArrayRef;
873        assert_regex_scan_masks(
874            &plain,
875            [
876                vec![true, false, false, false],
877                vec![true, true, false, false],
878                vec![false, true, true, false],
879                vec![false, false, true, false],
880            ],
881        );
882
883        let dictionary = DictionaryArray::new(
884            datatypes::arrow::array::UInt32Array::from(vec![
885                Some(0),
886                Some(1),
887                Some(2),
888                None,
889                Some(3),
890            ]),
891            Arc::new(datatypes::arrow::array::StringArray::from(vec![
892                Some("api"),
893                Some("API"),
894                Some("db"),
895                None,
896            ])),
897        );
898        let raw =
899            regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap();
900        assert!(raw.is_null(3)); // null dictionary key
901        assert!(raw.is_null(4)); // non-null key referencing a null dictionary value
902
903        let dictionary = Arc::new(dictionary) as ArrayRef;
904        assert_regex_scan_masks(
905            &dictionary,
906            [
907                vec![true, false, false, false, false],
908                vec![true, true, false, false, false],
909                vec![false, true, true, false, false],
910                vec![false, false, true, false, false],
911            ],
912        );
913
914        let negative = regex_evaluator(Operator::RegexNotMatch);
915        assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap());
916    }
917
918    fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec<bool>; 4]) {
919        for (op, expected) in [
920            Operator::RegexMatch,
921            Operator::RegexIMatch,
922            Operator::RegexNotMatch,
923            Operator::RegexNotIMatch,
924        ]
925        .into_iter()
926        .zip(expected)
927        {
928            assert_eq!(
929                BooleanBuffer::from(expected),
930                regex_evaluator(op).evaluate_array(input).unwrap(),
931                "{op:?}"
932            );
933        }
934    }
935
936    fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator {
937        let expr = Expr::BinaryExpr(BinaryExpr {
938            left: Box::new(Expr::Column(Column::from_name("host"))),
939            op,
940            right: Box::new("^api$".lit()),
941        });
942        SimpleFilterEvaluator::try_new(&expr).unwrap()
943    }
944}
common_recordbatch/filter/predicate.rs

common_recordbatch/filter/
predicate.rs
