vortex_array/scalar_fn/fns/like/
mod.rs1mod kernel;
5
6use std::fmt::Display;
7use std::fmt::Formatter;
8
9pub use kernel::*;
10use prost::Message;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_err;
14use vortex_proto::expr as pb;
15use vortex_session::VortexSession;
16
17use crate::Array;
18use crate::ArrayRef;
19use crate::arrow::Datum;
20use crate::arrow::from_arrow_array_with_len;
21use crate::dtype::DType;
22use crate::expr::Expression;
23use crate::expr::StatsCatalog;
24use crate::expr::and;
25use crate::expr::gt;
26use crate::expr::gt_eq;
27use crate::expr::lit;
28use crate::expr::lt;
29use crate::expr::or;
30use crate::scalar::StringLike;
31use crate::scalar_fn::Arity;
32use crate::scalar_fn::ChildName;
33use crate::scalar_fn::ExecutionArgs;
34use crate::scalar_fn::ScalarFnId;
35use crate::scalar_fn::ScalarFnVTable;
36use crate::scalar_fn::fns::literal::Literal;
37
38#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
40pub struct LikeOptions {
41 pub negated: bool,
42 pub case_insensitive: bool,
43}
44
45impl Display for LikeOptions {
46 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
47 if self.negated {
48 write!(f, "NOT ")?;
49 }
50 if self.case_insensitive {
51 write!(f, "ILIKE")
52 } else {
53 write!(f, "LIKE")
54 }
55 }
56}
57
58#[derive(Clone)]
60pub struct Like;
61
62impl ScalarFnVTable for Like {
63 type Options = LikeOptions;
64
65 fn id(&self) -> ScalarFnId {
66 ScalarFnId::from("vortex.like")
67 }
68
69 fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
70 Ok(Some(
71 pb::LikeOpts {
72 negated: instance.negated,
73 case_insensitive: instance.case_insensitive,
74 }
75 .encode_to_vec(),
76 ))
77 }
78
79 fn deserialize(
80 &self,
81 _metadata: &[u8],
82 _session: &VortexSession,
83 ) -> VortexResult<Self::Options> {
84 let opts = pb::LikeOpts::decode(_metadata)?;
85 Ok(LikeOptions {
86 negated: opts.negated,
87 case_insensitive: opts.case_insensitive,
88 })
89 }
90
91 fn arity(&self, _options: &Self::Options) -> Arity {
92 Arity::Exact(2)
93 }
94
95 fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
96 match child_idx {
97 0 => ChildName::from("child"),
98 1 => ChildName::from("pattern"),
99 _ => unreachable!("Invalid child index {} for Like expression", child_idx),
100 }
101 }
102
103 fn fmt_sql(
104 &self,
105 options: &Self::Options,
106 expr: &Expression,
107 f: &mut Formatter<'_>,
108 ) -> std::fmt::Result {
109 expr.child(0).fmt_sql(f)?;
110 if options.negated {
111 write!(f, " not")?;
112 }
113 if options.case_insensitive {
114 write!(f, " ilike ")?;
115 } else {
116 write!(f, " like ")?;
117 }
118 expr.child(1).fmt_sql(f)
119 }
120
121 fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
122 let input = &arg_dtypes[0];
123 let pattern = &arg_dtypes[1];
124
125 if !input.is_utf8() {
126 vortex_bail!("LIKE expression requires UTF8 input dtype, got {}", input);
127 }
128 if !pattern.is_utf8() {
129 vortex_bail!(
130 "LIKE expression requires UTF8 pattern dtype, got {}",
131 pattern
132 );
133 }
134
135 Ok(DType::Bool(
136 (input.is_nullable() || pattern.is_nullable()).into(),
137 ))
138 }
139
140 fn execute(&self, options: &Self::Options, args: ExecutionArgs) -> VortexResult<ArrayRef> {
141 let [child, pattern]: [ArrayRef; _] = args
142 .inputs
143 .try_into()
144 .map_err(|_| vortex_err!("Wrong argument count"))?;
145
146 arrow_like(&child, &pattern, *options)
147 }
148
149 fn validity(
150 &self,
151 _options: &Self::Options,
152 expression: &Expression,
153 ) -> VortexResult<Option<Expression>> {
154 tracing::warn!("Computing validity for LIKE expression");
155 let child_validity = expression.child(0).validity()?;
156 let pattern_validity = expression.child(1).validity()?;
157 Ok(Some(and(child_validity, pattern_validity)))
158 }
159
160 fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {
161 false
162 }
163
164 fn stat_falsification(
165 &self,
166 like_opts: &LikeOptions,
167 expr: &Expression,
168 catalog: &dyn StatsCatalog,
169 ) -> Option<Expression> {
170 if like_opts.negated || like_opts.case_insensitive {
174 return None;
175 }
176
177 let pat = expr.child(1).as_::<Literal>();
179
180 let pat_str = pat.as_utf8().value()?;
182
183 let src = expr.child(0).clone();
184 let src_min = src.stat_min(catalog)?;
185 let src_max = src.stat_max(catalog)?;
186
187 match LikeVariant::from_str(pat_str)? {
188 LikeVariant::Exact(text) => {
189 Some(or(gt(src_min, lit(text)), lt(src_max, lit(text))))
191 }
192 LikeVariant::Prefix(prefix) => {
193 let succ = prefix.to_string().increment().ok()?;
195
196 Some(or(gt_eq(src_min, lit(succ)), lt(src_max, lit(prefix))))
197 }
198 }
199 }
200}
201
202pub(crate) fn arrow_like(
204 array: &dyn Array,
205 pattern: &dyn Array,
206 options: LikeOptions,
207) -> VortexResult<ArrayRef> {
208 let nullable = array.dtype().is_nullable() | pattern.dtype().is_nullable();
209 let len = array.len();
210 assert_eq!(
211 array.len(),
212 pattern.len(),
213 "Arrow Like: length mismatch for {}",
214 array.encoding_id()
215 );
216
217 let lhs = Datum::try_new(array)?;
219 let rhs = Datum::try_new_with_target_datatype(pattern, lhs.data_type())?;
220
221 let result = match (options.negated, options.case_insensitive) {
222 (false, false) => arrow_string::like::like(&lhs, &rhs)?,
223 (true, false) => arrow_string::like::nlike(&lhs, &rhs)?,
224 (false, true) => arrow_string::like::ilike(&lhs, &rhs)?,
225 (true, true) => arrow_string::like::nilike(&lhs, &rhs)?,
226 };
227
228 from_arrow_array_with_len(&result, len, nullable)
229}
230
231#[derive(Debug, PartialEq)]
233enum LikeVariant<'a> {
234 Exact(&'a str),
235 Prefix(&'a str),
236}
237
238impl<'a> LikeVariant<'a> {
239 fn from_str(string: &str) -> Option<LikeVariant<'_>> {
241 let Some(wildcard_pos) = string.find(['%', '_']) else {
242 return Some(LikeVariant::Exact(string));
243 };
244
245 if wildcard_pos == 0 {
247 return None;
248 }
249
250 let prefix = &string[..wildcard_pos];
251 Some(LikeVariant::Prefix(prefix))
252 }
253}
254
255#[cfg(test)]
256mod tests {
257 use crate::arrays::BoolArray;
258 use crate::assert_arrays_eq;
259 use crate::dtype::DType;
260 use crate::dtype::Nullability;
261 use crate::expr::col;
262 use crate::expr::get_item;
263 use crate::expr::ilike;
264 use crate::expr::like;
265 use crate::expr::lit;
266 use crate::expr::not;
267 use crate::expr::not_ilike;
268 use crate::expr::not_like;
269 use crate::expr::pruning::pruning_expr::TrackingStatsCatalog;
270 use crate::expr::root;
271 use crate::scalar_fn::fns::like::LikeVariant;
272
273 #[test]
274 fn invert_booleans() {
275 let not_expr = not(root());
276 let bools = BoolArray::from_iter([false, true, false, false, true, true]);
277 assert_arrays_eq!(
278 bools.to_array().apply(¬_expr).unwrap(),
279 BoolArray::from_iter([true, false, true, true, false, false])
280 );
281 }
282
283 #[test]
284 fn dtype() {
285 let dtype = DType::Utf8(Nullability::NonNullable);
286 let like_expr = like(root(), lit("%test%"));
287 assert_eq!(
288 like_expr.return_dtype(&dtype).unwrap(),
289 DType::Bool(Nullability::NonNullable)
290 );
291 }
292
293 #[test]
294 fn test_display() {
295 let expr = like(get_item("name", root()), lit("%john%"));
296 assert_eq!(expr.to_string(), "$.name like \"%john%\"");
297
298 let expr2 = not_ilike(root(), lit("test*"));
299 assert_eq!(expr2.to_string(), "$ not ilike \"test*\"");
300 }
301
302 #[test]
303 fn test_like_variant() {
304 assert_eq!(
306 LikeVariant::from_str("simple"),
307 Some(LikeVariant::Exact("simple"))
308 );
309 assert_eq!(
310 LikeVariant::from_str("prefix%"),
311 Some(LikeVariant::Prefix("prefix"))
312 );
313 assert_eq!(
314 LikeVariant::from_str("first%rest_stuff"),
315 Some(LikeVariant::Prefix("first"))
316 );
317
318 assert_eq!(LikeVariant::from_str("%suffix"), None);
320 assert_eq!(LikeVariant::from_str("_pattern"), None);
321 }
322
323 #[test]
324 fn test_like_pushdown() {
325 let catalog = TrackingStatsCatalog::default();
328
329 let pruning_expr = like(col("a"), lit("prefix%"))
330 .stat_falsification(&catalog)
331 .expect("LIKE stat falsification");
332
333 insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "prefiy") or ($.a_max < "prefix"))"#);
334
335 let pruning_expr = like(col("a"), lit("pref%ix%"))
337 .stat_falsification(&catalog)
338 .expect("LIKE stat falsification");
339 insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "preg") or ($.a_max < "pref"))"#);
340
341 let pruning_expr = like(col("a"), lit("pref_ix_"))
342 .stat_falsification(&catalog)
343 .expect("LIKE stat falsification");
344 insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "preg") or ($.a_max < "pref"))"#);
345
346 let pruning_expr = like(col("a"), lit("exactly"))
348 .stat_falsification(&catalog)
349 .expect("LIKE stat falsification");
350 insta::assert_snapshot!(pruning_expr, @r#"(($.a_min > "exactly") or ($.a_max < "exactly"))"#);
351
352 let pruning_expr = like(col("a"), lit("%suffix")).stat_falsification(&catalog);
354 assert_eq!(pruning_expr, None);
355
356 assert_eq!(
358 None,
359 not_like(col("a"), lit("a")).stat_falsification(&catalog)
360 );
361 assert_eq!(None, ilike(col("a"), lit("a")).stat_falsification(&catalog));
362 }
363}