catalog/system_schema/
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
15use arrow::array::StringArray;
16use arrow::compute::kernels::comparison;
17use datafusion::common::ScalarValue;
18use datafusion::logical_expr::expr::Like;
19use datafusion::logical_expr::{Expr, Operator};
20use datatypes::value::Value;
21use store_api::storage::ScanRequest;
22
23type ColumnName = String;
24/// Predicate to filter `information_schema` tables stream,
25/// we only support these simple predicates currently.
26/// TODO(dennis): supports more predicate types.
27#[derive(Clone, PartialEq, Eq, Debug)]
28pub(crate) enum Predicate {
29    Eq(ColumnName, Value),
30    Like(ColumnName, String, bool),
31    NotEq(ColumnName, Value),
32    InList(ColumnName, Vec<Value>),
33    And(Box<Predicate>, Box<Predicate>),
34    Or(Box<Predicate>, Box<Predicate>),
35    Not(Box<Predicate>),
36}
37
38impl Predicate {
39    /// Evaluate the predicate with the row, returns:
40    /// - `None` when the predicate can't evaluate with the row.
41    /// - `Some(true)` when the predicate is satisfied,
42    /// - `Some(false)` when the predicate is not satisfied,
43    fn eval(&self, row: &[(&str, &Value)]) -> Option<bool> {
44        match self {
45            Predicate::Eq(c, v) => {
46                for (column, value) in row {
47                    if c != column {
48                        continue;
49                    }
50                    return Some(v == *value);
51                }
52            }
53            Predicate::Like(c, pattern, case_insensitive) => {
54                for (column, value) in row {
55                    if c != column {
56                        continue;
57                    }
58
59                    let Value::String(bs) = value else {
60                        continue;
61                    };
62
63                    return like_utf8(bs.as_utf8(), pattern, case_insensitive);
64                }
65            }
66            Predicate::NotEq(c, v) => {
67                for (column, value) in row {
68                    if c != column {
69                        continue;
70                    }
71                    return Some(v != *value);
72                }
73            }
74            Predicate::InList(c, values) => {
75                for (column, value) in row {
76                    if c != column {
77                        continue;
78                    }
79                    return Some(values.iter().any(|v| v == *value));
80                }
81            }
82            Predicate::And(left, right) => {
83                let left = left.eval(row);
84
85                // short-circuit
86                if matches!(left, Some(false)) {
87                    return Some(false);
88                }
89
90                return match (left, right.eval(row)) {
91                    (Some(left), Some(right)) => Some(left && right),
92                    (None, Some(false)) => Some(false),
93                    _ => None,
94                };
95            }
96            Predicate::Or(left, right) => {
97                let left = left.eval(row);
98
99                // short-circuit
100                if matches!(left, Some(true)) {
101                    return Some(true);
102                }
103
104                return match (left, right.eval(row)) {
105                    (Some(left), Some(right)) => Some(left || right),
106                    (None, Some(true)) => Some(true),
107                    _ => None,
108                };
109            }
110            Predicate::Not(p) => {
111                return Some(!p.eval(row)?);
112            }
113        }
114
115        // Can't evaluate predicate with the row
116        None
117    }
118
119    /// Try to create a predicate from datafusion [`Expr`], return None if fails.
120    fn from_expr(expr: Expr) -> Option<Predicate> {
121        match expr {
122            // NOT expr
123            Expr::Not(expr) => Some(Predicate::Not(Box::new(Self::from_expr(*expr)?))),
124            // expr LIKE pattern
125            Expr::Like(Like {
126                negated,
127                expr,
128                pattern,
129                case_insensitive,
130                ..
131            }) if is_column(&expr) && is_string_literal(&pattern) => {
132                // Safety: ensured by gurad
133                let Expr::Column(c) = *expr else {
134                    unreachable!();
135                };
136                let Expr::Literal(ScalarValue::Utf8(Some(pattern)), _) = *pattern else {
137                    unreachable!();
138                };
139
140                let p = Predicate::Like(c.name, pattern, case_insensitive);
141
142                if negated {
143                    Some(Predicate::Not(Box::new(p)))
144                } else {
145                    Some(p)
146                }
147            }
148            // left OP right
149            Expr::BinaryExpr(bin) => match (*bin.left, bin.op, *bin.right) {
150                // left == right
151                (Expr::Literal(scalar, _), Operator::Eq, Expr::Column(c))
152                | (Expr::Column(c), Operator::Eq, Expr::Literal(scalar, _)) => {
153                    let Ok(v) = Value::try_from(scalar) else {
154                        return None;
155                    };
156
157                    Some(Predicate::Eq(c.name, v))
158                }
159                // left != right
160                (Expr::Literal(scalar, _), Operator::NotEq, Expr::Column(c))
161                | (Expr::Column(c), Operator::NotEq, Expr::Literal(scalar, _)) => {
162                    let Ok(v) = Value::try_from(scalar) else {
163                        return None;
164                    };
165
166                    Some(Predicate::NotEq(c.name, v))
167                }
168                // left AND right
169                (left, Operator::And, right) => {
170                    let left = Self::from_expr(left)?;
171                    let right = Self::from_expr(right)?;
172
173                    Some(Predicate::And(Box::new(left), Box::new(right)))
174                }
175                // left OR right
176                (left, Operator::Or, right) => {
177                    let left = Self::from_expr(left)?;
178                    let right = Self::from_expr(right)?;
179
180                    Some(Predicate::Or(Box::new(left), Box::new(right)))
181                }
182                _ => None,
183            },
184            // [NOT] IN (LIST)
185            Expr::InList(list) => {
186                match (*list.expr, list.list, list.negated) {
187                    // column [NOT] IN (v1, v2, v3, ...)
188                    (Expr::Column(c), list, negated) if is_all_scalars(&list) => {
189                        let mut values = Vec::with_capacity(list.len());
190                        for scalar in list {
191                            // Safety: checked by `is_all_scalars`
192                            let Expr::Literal(scalar, _) = scalar else {
193                                unreachable!();
194                            };
195
196                            let Ok(value) = Value::try_from(scalar) else {
197                                return None;
198                            };
199
200                            values.push(value);
201                        }
202
203                        let predicate = Predicate::InList(c.name, values);
204
205                        if negated {
206                            Some(Predicate::Not(Box::new(predicate)))
207                        } else {
208                            Some(predicate)
209                        }
210                    }
211                    _ => None,
212                }
213            }
214            _ => None,
215        }
216    }
217
218    /// Collects the values that constrain `column` to be equal to them at this
219    /// predicate's top-level conjunctions.
220    ///
221    /// Constraints under `Or`/`Not` are ignored because they don't guarantee
222    /// that `column` equals a single value.
223    fn collect_eq_values<'a>(&'a self, column: &str, values: &mut Vec<&'a Value>) {
224        match self {
225            Predicate::Eq(c, v) if c == column => values.push(v),
226            Predicate::And(left, right) => {
227                left.collect_eq_values(column, values);
228                right.collect_eq_values(column, values);
229            }
230            _ => {}
231        }
232    }
233}
234
235/// Perform SQL left LIKE right, return `None` if fail to evaluate.
236/// - `s` the target string
237/// - `pattern` the pattern just like '%abc'
238/// - `case_insensitive` whether to perform case-insensitive like or not.
239fn like_utf8(s: &str, pattern: &str, case_insensitive: &bool) -> Option<bool> {
240    let array = StringArray::from(vec![s]);
241    let patterns = StringArray::new_scalar(pattern);
242
243    let Ok(booleans) = (if *case_insensitive {
244        comparison::ilike(&array, &patterns)
245    } else {
246        comparison::like(&array, &patterns)
247    }) else {
248        return None;
249    };
250
251    // Safety: at least one value in result
252    Some(booleans.value(0))
253}
254
255fn is_string_literal(expr: &Expr) -> bool {
256    matches!(expr, Expr::Literal(ScalarValue::Utf8(Some(_)), _))
257}
258
259fn is_column(expr: &Expr) -> bool {
260    matches!(expr, Expr::Column(_))
261}
262
263/// A list of predicate
264pub struct Predicates {
265    predicates: Vec<Predicate>,
266}
267
268impl Predicates {
269    /// Try its best to create predicates from [`ScanRequest`].
270    pub fn from_scan_request(request: &Option<ScanRequest>) -> Predicates {
271        if let Some(request) = request {
272            let mut predicates = Vec::with_capacity(request.filters.len());
273
274            for filter in &request.filters {
275                if let Some(predicate) = Predicate::from_expr(filter.clone()) {
276                    predicates.push(predicate);
277                }
278            }
279
280            Self { predicates }
281        } else {
282            Self {
283                predicates: Vec::new(),
284            }
285        }
286    }
287
288    /// Returns the single equality value that constrains `column` in the
289    /// top-level conjunctions of the predicates.
290    ///
291    /// Returns `None` if there is no equality constraint on `column`, or if
292    /// the predicates constrain it with multiple distinct values.
293    pub fn find_eq_value(&self, column: &str) -> Option<&Value> {
294        let mut values = Vec::new();
295        for predicate in &self.predicates {
296            predicate.collect_eq_values(column, &mut values);
297        }
298
299        let (first, rest) = values.split_first()?;
300        if rest.iter().all(|v| v == first) {
301            Some(*first)
302        } else {
303            None
304        }
305    }
306
307    /// Evaluate the predicates with the row.
308    /// returns true when all the predicates are satisfied or can't be evaluated.
309    pub fn eval(&self, row: &[(&str, &Value)]) -> bool {
310        // fast path
311        if self.predicates.is_empty() {
312            return true;
313        }
314
315        self.predicates
316            .iter()
317            .filter_map(|p| p.eval(row))
318            .all(|b| b)
319    }
320}
321
322/// Returns true when the values are all [`DfExpr::Literal`].
323fn is_all_scalars(list: &[Expr]) -> bool {
324    list.iter().all(|v| matches!(v, Expr::Literal(_, _)))
325}
326
327#[cfg(test)]
328mod tests {
329    use datafusion::common::Column;
330    use datafusion::logical_expr::expr::InList;
331    use datafusion::logical_expr::{BinaryExpr, Literal};
332
333    use super::*;
334
335    #[test]
336    fn test_predicate_eval() {
337        let a_col = "a".to_string();
338        let b_col = "b".to_string();
339        let a_value = Value::from("a_value");
340        let b_value = Value::from("b_value");
341        let wrong_value = Value::from("wrong_value");
342
343        let a_row = [(a_col.as_str(), &a_value)];
344        let b_row = [("b", &wrong_value)];
345        let wrong_row = [(a_col.as_str(), &wrong_value)];
346
347        // Predicate::Eq
348        let p = Predicate::Eq(a_col.clone(), a_value.clone());
349        assert!(p.eval(&a_row).unwrap());
350        assert!(p.eval(&b_row).is_none());
351        assert!(!p.eval(&wrong_row).unwrap());
352
353        // Predicate::NotEq
354        let p = Predicate::NotEq(a_col.clone(), a_value.clone());
355        assert!(!p.eval(&a_row).unwrap());
356        assert!(p.eval(&b_row).is_none());
357        assert!(p.eval(&wrong_row).unwrap());
358
359        // Predicate::InList
360        let p = Predicate::InList(a_col.clone(), vec![a_value.clone(), b_value.clone()]);
361        assert!(p.eval(&a_row).unwrap());
362        assert!(p.eval(&b_row).is_none());
363        assert!(!p.eval(&wrong_row).unwrap());
364        assert!(p.eval(&[(&a_col, &b_value)]).unwrap());
365
366        let p1 = Predicate::Eq(a_col.clone(), a_value.clone());
367        let p2 = Predicate::Eq(b_col.clone(), b_value.clone());
368        let row = [(a_col.as_str(), &a_value), (b_col.as_str(), &b_value)];
369        let wrong_row = [(a_col.as_str(), &a_value), (b_col.as_str(), &wrong_value)];
370
371        //Predicate::And
372        let p = Predicate::And(Box::new(p1.clone()), Box::new(p2.clone()));
373        assert!(p.eval(&row).unwrap());
374        assert!(!p.eval(&wrong_row).unwrap());
375        assert!(p.eval(&[]).is_none());
376        assert!(p.eval(&[("c", &a_value)]).is_none());
377        assert!(
378            !p.eval(&[(a_col.as_str(), &b_value), (b_col.as_str(), &a_value)])
379                .unwrap()
380        );
381        assert!(
382            !p.eval(&[(a_col.as_str(), &b_value), (b_col.as_str(), &b_value)])
383                .unwrap()
384        );
385        assert!(
386            p.eval(&[(a_col.as_ref(), &a_value), ("c", &a_value)])
387                .is_none()
388        );
389        assert!(
390            !p.eval(&[(a_col.as_ref(), &b_value), ("c", &a_value)])
391                .unwrap()
392        );
393
394        //Predicate::Or
395        let p = Predicate::Or(Box::new(p1), Box::new(p2));
396        assert!(p.eval(&row).unwrap());
397        assert!(p.eval(&wrong_row).unwrap());
398        assert!(p.eval(&[]).is_none());
399        assert!(p.eval(&[("c", &a_value)]).is_none());
400        assert!(
401            !p.eval(&[(a_col.as_str(), &b_value), (b_col.as_str(), &a_value)])
402                .unwrap()
403        );
404        assert!(
405            p.eval(&[(a_col.as_str(), &b_value), (b_col.as_str(), &b_value)])
406                .unwrap()
407        );
408        assert!(
409            p.eval(&[(a_col.as_ref(), &a_value), ("c", &a_value)])
410                .unwrap()
411        );
412        assert!(
413            p.eval(&[(a_col.as_ref(), &b_value), ("c", &a_value)])
414                .is_none()
415        );
416    }
417
418    #[test]
419    fn test_predicate_like() {
420        // case insensitive
421        let expr = Expr::Like(Like {
422            negated: false,
423            expr: Box::new(column("a")),
424            pattern: Box::new("%abc".lit()),
425            case_insensitive: true,
426            escape_char: None,
427        });
428
429        let p = Predicate::from_expr(expr).unwrap();
430        assert!(
431            matches!(&p, Predicate::Like(c, pattern, case_insensitive) if
432                         c == "a"
433                         && pattern == "%abc"
434                         && *case_insensitive)
435        );
436
437        let match_row = [
438            ("a", &Value::from("hello AbC")),
439            ("b", &Value::from("b value")),
440        ];
441        let unmatch_row = [("a", &Value::from("bca")), ("b", &Value::from("b value"))];
442
443        assert!(p.eval(&match_row).unwrap());
444        assert!(!p.eval(&unmatch_row).unwrap());
445        assert!(p.eval(&[]).is_none());
446
447        // case sensitive
448        let expr = Expr::Like(Like {
449            negated: false,
450            expr: Box::new(column("a")),
451            pattern: Box::new("%abc".lit()),
452            case_insensitive: false,
453            escape_char: None,
454        });
455
456        let p = Predicate::from_expr(expr).unwrap();
457        assert!(
458            matches!(&p, Predicate::Like(c, pattern, case_insensitive) if
459                         c == "a"
460                         && pattern == "%abc"
461                         && !*case_insensitive)
462        );
463        assert!(!p.eval(&match_row).unwrap());
464        assert!(!p.eval(&unmatch_row).unwrap());
465        assert!(p.eval(&[]).is_none());
466
467        // not like
468        let expr = Expr::Like(Like {
469            negated: true,
470            expr: Box::new(column("a")),
471            pattern: Box::new("%abc".lit()),
472            case_insensitive: true,
473            escape_char: None,
474        });
475
476        let p = Predicate::from_expr(expr).unwrap();
477        assert!(!p.eval(&match_row).unwrap());
478        assert!(p.eval(&unmatch_row).unwrap());
479        assert!(p.eval(&[]).is_none());
480    }
481
482    fn column(name: &str) -> Expr {
483        Expr::Column(Column::from_name(name))
484    }
485
486    fn match_string_value(v: &Value, expected: &str) -> bool {
487        matches!(v, Value::String(bs) if bs.as_utf8() == expected)
488    }
489
490    fn match_string_values(vs: &[Value], expected: &[&str]) -> bool {
491        assert_eq!(vs.len(), expected.len());
492
493        let mut result = true;
494        for (i, v) in vs.iter().enumerate() {
495            result = result && match_string_value(v, expected[i]);
496        }
497
498        result
499    }
500
501    fn mock_exprs() -> (Expr, Expr) {
502        let expr1 = Expr::BinaryExpr(BinaryExpr {
503            left: Box::new(column("a")),
504            op: Operator::Eq,
505            right: Box::new("a_value".lit()),
506        });
507
508        let expr2 = Expr::BinaryExpr(BinaryExpr {
509            left: Box::new(column("b")),
510            op: Operator::NotEq,
511            right: Box::new("b_value".lit()),
512        });
513
514        (expr1, expr2)
515    }
516
517    #[test]
518    fn test_predicate_from_expr() {
519        let (expr1, expr2) = mock_exprs();
520
521        let p1 = Predicate::from_expr(expr1.clone()).unwrap();
522        assert!(matches!(&p1, Predicate::Eq(column, v) if column == "a"
523                         && match_string_value(v, "a_value")));
524
525        let p2 = Predicate::from_expr(expr2.clone()).unwrap();
526        assert!(matches!(&p2, Predicate::NotEq(column, v) if column == "b"
527                         && match_string_value(v, "b_value")));
528
529        let and_expr = Expr::BinaryExpr(BinaryExpr {
530            left: Box::new(expr1.clone()),
531            op: Operator::And,
532            right: Box::new(expr2.clone()),
533        });
534        let or_expr = Expr::BinaryExpr(BinaryExpr {
535            left: Box::new(expr1.clone()),
536            op: Operator::Or,
537            right: Box::new(expr2.clone()),
538        });
539        let not_expr = Expr::Not(Box::new(expr1.clone()));
540
541        let and_p = Predicate::from_expr(and_expr).unwrap();
542        assert!(matches!(and_p, Predicate::And(left, right) if *left == p1 && *right == p2));
543        let or_p = Predicate::from_expr(or_expr).unwrap();
544        assert!(matches!(or_p, Predicate::Or(left, right) if *left == p1 && *right == p2));
545        let not_p = Predicate::from_expr(not_expr).unwrap();
546        assert!(matches!(not_p, Predicate::Not(p) if *p == p1));
547
548        let inlist_expr = Expr::InList(InList {
549            expr: Box::new(column("a")),
550            list: vec!["a1".lit(), "a2".lit()],
551            negated: false,
552        });
553
554        let inlist_p = Predicate::from_expr(inlist_expr).unwrap();
555        assert!(matches!(&inlist_p, Predicate::InList(c, values) if c == "a"
556                         && match_string_values(values, &["a1", "a2"])));
557
558        let inlist_expr = Expr::InList(InList {
559            expr: Box::new(column("a")),
560            list: vec!["a1".lit(), "a2".lit()],
561            negated: true,
562        });
563        let inlist_p = Predicate::from_expr(inlist_expr).unwrap();
564        assert!(matches!(inlist_p, Predicate::Not(p) if
565                         matches!(&*p,
566                                  Predicate::InList(c, values) if c == "a"
567                                  && match_string_values(values, &["a1", "a2"]))));
568    }
569
570    #[test]
571    fn test_predicates_from_scan_request() {
572        let predicates = Predicates::from_scan_request(&None);
573        assert!(predicates.predicates.is_empty());
574
575        let (expr1, expr2) = mock_exprs();
576
577        let request = ScanRequest {
578            filters: vec![expr1, expr2],
579            ..Default::default()
580        };
581        let predicates = Predicates::from_scan_request(&Some(request));
582
583        assert_eq!(2, predicates.predicates.len());
584        assert!(
585            matches!(&predicates.predicates[0], Predicate::Eq(column, v) if column == "a"
586                     && match_string_value(v, "a_value"))
587        );
588        assert!(
589            matches!(&predicates.predicates[1], Predicate::NotEq(column, v) if column == "b"
590                     && match_string_value(v, "b_value"))
591        );
592    }
593
594    #[test]
595    fn test_find_eq_value() {
596        fn eq(name: &str, value: &str) -> Predicate {
597            Predicate::Eq(name.to_string(), Value::from(value))
598        }
599
600        fn and(left: Predicate, right: Predicate) -> Predicate {
601            Predicate::And(Box::new(left), Box::new(right))
602        }
603
604        fn or(left: Predicate, right: Predicate) -> Predicate {
605            Predicate::Or(Box::new(left), Box::new(right))
606        }
607
608        fn predicates(predicates: Vec<Predicate>) -> Predicates {
609            Predicates { predicates }
610        }
611
612        // No predicates, no equality value.
613        assert!(predicates(vec![]).find_eq_value("a").is_none());
614
615        // No equality on the column.
616        let p = predicates(vec![Predicate::NotEq("a".to_string(), Value::from("x"))]);
617        assert!(p.find_eq_value("a").is_none());
618        assert!(p.find_eq_value("b").is_none());
619
620        // Simple equality.
621        let p = predicates(vec![eq("a", "x")]);
622        let v = p.find_eq_value("a").unwrap();
623        assert!(match_string_value(v, "x"));
624        assert!(p.find_eq_value("b").is_none());
625
626        // Duplicated equality on the same value.
627        let p = predicates(vec![and(eq("a", "x"), eq("a", "x"))]);
628        let v = p.find_eq_value("a").unwrap();
629        assert!(match_string_value(v, "x"));
630
631        // Conflicting equalities reject the lookup.
632        let p = predicates(vec![and(eq("a", "x"), eq("a", "y"))]);
633        assert!(p.find_eq_value("a").is_none());
634
635        // Conflicting equalities in separate predicates reject the lookup too.
636        let p = predicates(vec![eq("a", "x"), eq("a", "y")]);
637        assert!(p.find_eq_value("a").is_none());
638
639        // Equality nested in deeper conjunctions.
640        let p = predicates(vec![and(and(eq("b", "y"), eq("a", "x")), eq("a", "x"))]);
641        let v = p.find_eq_value("a").unwrap();
642        assert!(match_string_value(v, "x"));
643
644        // Equality under a disjunction or a negation doesn't constrain the column.
645        let p = predicates(vec![or(eq("a", "x"), eq("a", "x"))]);
646        assert!(p.find_eq_value("a").is_none());
647        let p = predicates(vec![Predicate::Not(Box::new(eq("a", "x")))]);
648        assert!(p.find_eq_value("a").is_none());
649
650        // Equality value in a top-level conjunction composed of separate predicates.
651        let p = predicates(vec![
652            eq("a", "x"),
653            Predicate::Like("b".to_string(), "%y".to_string(), false),
654        ]);
655        let v = p.find_eq_value("a").unwrap();
656        assert!(match_string_value(v, "x"));
657    }
658
659    #[test]
660    fn test_predicates_eval_row() {
661        let wrong_row = [
662            ("a", &Value::from("a_value")),
663            ("b", &Value::from("b_value")),
664            ("c", &Value::from("c_value")),
665        ];
666        let row = [
667            ("a", &Value::from("a_value")),
668            ("b", &Value::from("not_b_value")),
669            ("c", &Value::from("c_value")),
670        ];
671        let c_row = [("c", &Value::from("c_value"))];
672
673        // test empty predicates, always returns true
674        let predicates = Predicates::from_scan_request(&None);
675        assert!(predicates.eval(&row));
676        assert!(predicates.eval(&wrong_row));
677        assert!(predicates.eval(&c_row));
678
679        let (expr1, expr2) = mock_exprs();
680        let request = ScanRequest {
681            filters: vec![expr1, expr2],
682            ..Default::default()
683        };
684        let predicates = Predicates::from_scan_request(&Some(request));
685        assert!(predicates.eval(&row));
686        assert!(!predicates.eval(&wrong_row));
687        assert!(predicates.eval(&c_row));
688    }
689}
catalog/system_schema/predicate.rs

catalog/system_schema/
predicate.rs
