mito2/sst/index/inverted_index/applier/builder/
regex_match.rs1use 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 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}