Skip to main content

mito2/sst/index/inverted_index/applier/builder/
regex_match.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 datafusion_expr::Expr as DfExpr;
16use index::inverted_index::search::predicate::{Predicate, RegexMatchPredicate};
17
18use crate::error::Result;
19use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
20
21impl InvertedIndexApplierBuilder<'_> {
22    /// Collects a regex match expression in the form of `column ~ pattern`.
23    pub(crate) fn collect_regex_match(&mut self, column: &DfExpr, pattern: &DfExpr) -> Result<()> {
24        let Some((target, data_type)) = self.expr_to_index_target(column)? else {
25            return Ok(());
26        };
27        if !data_type.is_string() {
28            return Ok(());
29        }
30        let Some(pattern) = Self::nonnull_lit(pattern)
31            .and_then(|value| value.try_as_str())
32            .flatten()
33        else {
34            return Ok(());
35        };
36
37        let predicate = Predicate::RegexMatch(RegexMatchPredicate {
38            pattern: pattern.to_owned(),
39        });
40        self.add_predicate(target, predicate);
41        Ok(())
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use std::collections::HashSet;
48
49    use store_api::region_request::PathType;
50
51    use super::*;
52    use crate::error::Error;
53    use crate::sst::index::inverted_index::applier::builder::tests::{
54        field_column, int64_lit, nonexistent_column, string_lit, tag_column, test_object_store,
55        test_region_metadata,
56    };
57    use crate::sst::index::puffin_manager::PuffinManagerFactory;
58
59    #[test]
60    fn test_regex_match_basic() {
61        let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_regex_match_basic_");
62        let metadata = test_region_metadata();
63        let mut builder = InvertedIndexApplierBuilder::new(
64            "test".to_string(),
65            PathType::Bare,
66            test_object_store(),
67            &metadata,
68            HashSet::from_iter([1, 2, 3]),
69            factory,
70        );
71
72        builder
73            .collect_regex_match(&tag_column(), &string_lit("abc"))
74            .unwrap();
75
76        let predicates = builder
77            .output
78            .get(&index::target::IndexTarget::ColumnId(1))
79            .unwrap();
80        assert_eq!(predicates.len(), 1);
81        assert_eq!(
82            predicates[0],
83            Predicate::RegexMatch(RegexMatchPredicate {
84                pattern: "abc".to_string()
85            })
86        );
87    }
88
89    #[test]
90    fn test_regex_match_field_column() {
91        let (_d, factory) =
92            PuffinManagerFactory::new_for_test_block("test_regex_match_field_column_");
93        let metadata = test_region_metadata();
94        let mut builder = InvertedIndexApplierBuilder::new(
95            "test".to_string(),
96            PathType::Bare,
97            test_object_store(),
98            &metadata,
99            HashSet::from_iter([1, 2, 3]),
100            factory,
101        );
102
103        builder
104            .collect_regex_match(&field_column(), &string_lit("abc"))
105            .unwrap();
106
107        let predicates = builder
108            .output
109            .get(&index::target::IndexTarget::ColumnId(3))
110            .unwrap();
111        assert_eq!(predicates.len(), 1);
112        assert_eq!(
113            predicates[0],
114            Predicate::RegexMatch(RegexMatchPredicate {
115                pattern: "abc".to_string()
116            })
117        );
118    }
119
120    #[test]
121    fn test_regex_match_type_mismatch() {
122        let (_d, factory) =
123            PuffinManagerFactory::new_for_test_block("test_regex_match_type_mismatch_");
124        let metadata = test_region_metadata();
125        let mut builder = InvertedIndexApplierBuilder::new(
126            "test".to_string(),
127            PathType::Bare,
128            test_object_store(),
129            &metadata,
130            HashSet::from_iter([1, 2, 3]),
131            factory,
132        );
133
134        builder
135            .collect_regex_match(&tag_column(), &int64_lit(123))
136            .unwrap();
137
138        assert!(builder.output.is_empty());
139    }
140
141    #[test]
142    fn test_regex_match_type_nonexist_column() {
143        let (_d, factory) =
144            PuffinManagerFactory::new_for_test_block("test_regex_match_type_nonexist_column_");
145        let metadata = test_region_metadata();
146        let mut builder = InvertedIndexApplierBuilder::new(
147            "test".to_string(),
148            PathType::Bare,
149            test_object_store(),
150            &metadata,
151            HashSet::from_iter([1, 2, 3]),
152            factory,
153        );
154
155        let res = builder.collect_regex_match(&nonexistent_column(), &string_lit("abc"));
156        assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
157        assert!(builder.output.is_empty());
158    }
159}