vortex_fsst/compute/
like.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::BoolArray;
9use vortex_array::arrays::PrimitiveArray;
10use vortex_array::arrays::varbin::VarBinArraySlotsExt;
11use vortex_array::match_each_integer_ptype;
12use vortex_array::scalar_fn::fns::like::LikeKernel;
13use vortex_array::scalar_fn::fns::like::LikeOptions;
14use vortex_error::VortexResult;
15
16use crate::FSST;
17use crate::FSSTArrayExt;
18use crate::dfa::FsstMatcher;
19use crate::dfa::dfa_scan_to_bitbuf;
20
21impl LikeKernel for FSST {
22    fn like(
23        array: ArrayView<'_, Self>,
24        pattern: &ArrayRef,
25        options: LikeOptions,
26        ctx: &mut ExecutionCtx,
27    ) -> VortexResult<Option<ArrayRef>> {
28        let Some(pattern_scalar) = pattern.as_constant() else {
29            return Ok(None);
30        };
31
32        if options.case_insensitive {
33            return Ok(None);
34        }
35
36        let pattern_bytes: &[u8] = if let Some(s) = pattern_scalar.as_utf8_opt() {
37            let Some(v) = s.value() else {
38                return Ok(None);
39            };
40            v.as_ref()
41        } else if let Some(b) = pattern_scalar.as_binary_opt() {
42            let Some(v) = b.value() else {
43                return Ok(None);
44            };
45            v
46        } else {
47            return Ok(None);
48        };
49
50        let Some(matcher) =
51            FsstMatcher::try_new(array.symbols(), array.symbol_lengths(), pattern_bytes)?
52        else {
53            return Ok(None);
54        };
55
56        let negated = options.negated;
57        let codes = array.codes();
58        let offsets = codes.offsets().clone().execute::<PrimitiveArray>(ctx)?;
59        let all_bytes = codes.bytes();
60        let all_bytes = all_bytes.as_slice();
61        let n = codes.len();
62
63        let result = match_each_integer_ptype!(offsets.ptype(), |T| {
64            let off = offsets.as_slice::<T>();
65            dfa_scan_to_bitbuf(n, off, all_bytes, negated, |codes| matcher.matches(codes))
66        });
67
68        // FSST delegates validity to its codes array, so we can read it
69        // directly without cloning the entire FSSTArray into an ArrayRef.
70        let validity = array
71            .codes()
72            .validity()?
73            .union_nullability(pattern_scalar.dtype().nullability());
74
75        Ok(Some(BoolArray::new(result, validity).into_array()))
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use std::sync::LazyLock;
82
83    use vortex_array::Canonical;
84    use vortex_array::IntoArray;
85    use vortex_array::VortexSessionExecute;
86    use vortex_array::arrays::BoolArray;
87    use vortex_array::arrays::ConstantArray;
88    use vortex_array::arrays::VarBinArray;
89    use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
90    use vortex_array::assert_arrays_eq;
91    use vortex_array::dtype::DType;
92    use vortex_array::dtype::Nullability;
93    use vortex_array::scalar_fn::fns::like::Like;
94    use vortex_array::scalar_fn::fns::like::LikeKernel;
95    use vortex_array::scalar_fn::fns::like::LikeOptions;
96    use vortex_error::VortexResult;
97    use vortex_session::VortexSession;
98
99    use crate::FSST;
100    use crate::FSSTArray;
101    use crate::fsst_compress;
102    use crate::fsst_train_compressor;
103
104    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
105        let session = vortex_array::array_session();
106        crate::initialize(&session);
107        session
108    });
109
110    fn make_fsst(strings: &[Option<&str>], nullability: Nullability) -> FSSTArray {
111        let array =
112            VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)).into_array();
113        let mut ctx = SESSION.create_execution_ctx();
114        let compressor = fsst_train_compressor(&array, &mut ctx).unwrap();
115        fsst_compress(&array, &compressor, &mut ctx).unwrap()
116    }
117
118    fn run_like(array: FSSTArray, pattern: &str, opts: LikeOptions) -> VortexResult<BoolArray> {
119        let len = array.len();
120        let arr = array.into_array();
121        let pattern = ConstantArray::new(pattern, len).into_array();
122        let result = Like
123            .try_new_array(len, opts, [arr, pattern])?
124            .into_array()
125            .execute::<Canonical>(&mut SESSION.create_execution_ctx())?;
126        Ok(result.into_bool())
127    }
128
129    fn like(array: FSSTArray, pattern: &str) -> VortexResult<BoolArray> {
130        run_like(array, pattern, LikeOptions::default())
131    }
132
133    #[test]
134    fn test_like_prefix() -> VortexResult<()> {
135        let fsst = make_fsst(
136            &[
137                Some("http://example.com"),
138                Some("http://test.org"),
139                Some("ftp://files.net"),
140                Some("http://vortex.dev"),
141                Some("ssh://server.io"),
142            ],
143            Nullability::NonNullable,
144        );
145        let result = like(fsst, "http%")?;
146        assert_arrays_eq!(
147            &result,
148            &BoolArray::from_iter([true, true, false, true, false]),
149            &mut SESSION.create_execution_ctx()
150        );
151        Ok(())
152    }
153
154    #[test]
155    fn test_like_prefix_with_nulls() -> VortexResult<()> {
156        let fsst = make_fsst(
157            &[Some("hello"), None, Some("help"), None, Some("goodbye")],
158            Nullability::Nullable,
159        );
160        let result = like(fsst, "hel%")?; // spellchecker:disable-line
161        assert_arrays_eq!(
162            &result,
163            &BoolArray::from_iter([Some(true), None, Some(true), None, Some(false)]),
164            &mut SESSION.create_execution_ctx()
165        );
166        Ok(())
167    }
168
169    #[test]
170    fn test_like_contains() -> VortexResult<()> {
171        let fsst = make_fsst(
172            &[
173                Some("hello world"),
174                Some("say hello"),
175                Some("goodbye"),
176                Some("hellooo"),
177            ],
178            Nullability::NonNullable,
179        );
180        let result = like(fsst, "%hello%")?;
181        assert_arrays_eq!(
182            &result,
183            &BoolArray::from_iter([true, true, false, true]),
184            &mut SESSION.create_execution_ctx()
185        );
186        Ok(())
187    }
188
189    #[test]
190    fn test_like_contains_cross_symbol() -> VortexResult<()> {
191        let fsst = make_fsst(
192            &[
193                Some("the quick brown fox jumps over the lazy dog"),
194                Some("a short string"),
195                Some("the lazy dog sleeps"),
196                Some("no match"),
197            ],
198            Nullability::NonNullable,
199        );
200        let result = like(fsst, "%lazy dog%")?;
201        assert_arrays_eq!(
202            &result,
203            &BoolArray::from_iter([true, false, true, false]),
204            &mut SESSION.create_execution_ctx()
205        );
206        Ok(())
207    }
208
209    #[test]
210    fn test_not_like_contains() -> VortexResult<()> {
211        let fsst = make_fsst(
212            &[Some("foobar_sdf"), Some("sdf_start"), Some("nothing")],
213            Nullability::NonNullable,
214        );
215        let opts = LikeOptions {
216            negated: true,
217            case_insensitive: false,
218        };
219        let result = run_like(fsst, "%sdf%", opts)?;
220        assert_arrays_eq!(
221            &result,
222            &BoolArray::from_iter([false, false, true]),
223            &mut SESSION.create_execution_ctx()
224        );
225        Ok(())
226    }
227
228    #[test]
229    fn test_like_match_all() -> VortexResult<()> {
230        let fsst = make_fsst(
231            &[Some("abc"), Some(""), Some("xyz")],
232            Nullability::NonNullable,
233        );
234        let result = like(fsst, "%")?;
235        assert_arrays_eq!(
236            &result,
237            &BoolArray::from_iter([true, true, true]),
238            &mut SESSION.create_execution_ctx()
239        );
240        Ok(())
241    }
242
243    /// Call `LikeKernel::like` directly on the FSSTArray and verify it
244    /// returns `Some(...)` (i.e. the kernel handles it, rather than
245    /// returning `None` which would mean "fall back to decompress").
246    #[test]
247    fn test_like_prefix_kernel_handles() -> VortexResult<()> {
248        let fsst = make_fsst(
249            &[Some("http://a.com"), Some("ftp://b.com")],
250            Nullability::NonNullable,
251        );
252        let pattern = ConstantArray::new("http%", fsst.len()).into_array();
253        let mut ctx = SESSION.create_execution_ctx();
254
255        let fsst = fsst.as_view();
256        let result = <FSST as LikeKernel>::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?;
257        assert!(result.is_some(), "FSST LikeKernel should handle prefix%");
258        assert_arrays_eq!(
259            result.unwrap(),
260            BoolArray::from_iter([true, false]),
261            &mut ctx
262        );
263        Ok(())
264    }
265
266    /// Same direct-call check for the contains pattern `%needle%`.
267    #[test]
268    fn test_like_contains_kernel_handles() -> VortexResult<()> {
269        let fsst = make_fsst(
270            &[Some("hello world"), Some("goodbye")],
271            Nullability::NonNullable,
272        );
273        let pattern = ConstantArray::new("%world%", fsst.len()).into_array();
274        let mut ctx = SESSION.create_execution_ctx();
275
276        let fsst = fsst.as_view();
277        let result = <FSST as LikeKernel>::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?;
278        assert!(result.is_some(), "FSST LikeKernel should handle %needle%");
279        assert_arrays_eq!(
280            result.unwrap(),
281            BoolArray::from_iter([true, false]),
282            &mut ctx
283        );
284        Ok(())
285    }
286
287    /// Patterns we can't handle should return `None` (fall back).
288    #[test]
289    fn test_like_kernel_falls_back_for_complex_pattern() -> VortexResult<()> {
290        let fsst = make_fsst(&[Some("abc"), Some("def")], Nullability::NonNullable);
291        let mut ctx = SESSION.create_execution_ctx();
292
293        // Underscore wildcard -- not handled.
294        let pattern = ConstantArray::new("a_c", fsst.len()).into_array();
295        let fsst_v = fsst.as_view();
296        let result =
297            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
298        assert!(result.is_none(), "underscore pattern should fall back");
299
300        // Case-insensitive -- not handled.
301        let pattern = ConstantArray::new("abc%", fsst.len()).into_array();
302        let opts = LikeOptions {
303            negated: false,
304            case_insensitive: true,
305        };
306        let result = <FSST as LikeKernel>::like(fsst_v, &pattern, opts, &mut ctx)?;
307        assert!(result.is_none(), "ilike should fall back");
308
309        // Suffix patterns are still unsupported, even when the suffix is an escaped literal.
310        let pattern = ConstantArray::new(r"%\%", fsst.len()).into_array();
311        let result =
312            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
313        assert!(result.is_none(), "escaped suffix pattern should fall back");
314
315        Ok(())
316    }
317
318    #[test]
319    fn test_like_kernel_handles_escaped_prefix_and_contains() -> VortexResult<()> {
320        let fsst = make_fsst(
321            &[
322                Some("%front"),
323                Some("_front"),
324                Some("\\front"),
325                Some("middle%value"),
326                Some("middle_value"),
327                Some("middle\\value"),
328                Some("front"),
329            ],
330            Nullability::NonNullable,
331        );
332        let fsst_v = fsst.as_view();
333        let mut ctx = SESSION.create_execution_ctx();
334
335        let pattern = ConstantArray::new(r"\%%", fsst.len()).into_array();
336        let result =
337            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
338        assert!(result.is_some(), "escaped percent prefix should use FSST");
339        assert_arrays_eq!(
340            result.unwrap(),
341            BoolArray::from_iter([true, false, false, false, false, false, false]),
342            &mut ctx
343        );
344
345        let pattern = ConstantArray::new(r"\_%", fsst.len()).into_array();
346        let result =
347            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
348        assert!(
349            result.is_some(),
350            "escaped underscore prefix should use FSST"
351        );
352        assert_arrays_eq!(
353            result.unwrap(),
354            BoolArray::from_iter([false, true, false, false, false, false, false]),
355            &mut ctx
356        );
357
358        let pattern = ConstantArray::new(r"\\%", fsst.len()).into_array();
359        let result =
360            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
361        assert!(result.is_some(), "escaped backslash prefix should use FSST");
362        assert_arrays_eq!(
363            result.unwrap(),
364            BoolArray::from_iter([false, false, true, false, false, false, false]),
365            &mut ctx
366        );
367
368        let pattern = ConstantArray::new(r"%\%%", fsst.len()).into_array();
369        let result =
370            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
371        assert!(result.is_some(), "escaped percent contains should use FSST");
372        assert_arrays_eq!(
373            result.unwrap(),
374            BoolArray::from_iter([true, false, false, true, false, false, false]),
375            &mut ctx
376        );
377
378        let pattern = ConstantArray::new(r"%\_%", fsst.len()).into_array();
379        let result =
380            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
381        assert!(
382            result.is_some(),
383            "escaped underscore contains should use FSST"
384        );
385        assert_arrays_eq!(
386            result.unwrap(),
387            BoolArray::from_iter([false, true, false, false, true, false, false]),
388            &mut ctx
389        );
390
391        let pattern = ConstantArray::new(r"%\\%", fsst.len()).into_array();
392        let result =
393            <FSST as LikeKernel>::like(fsst_v, &pattern, LikeOptions::default(), &mut ctx)?;
394        assert!(
395            result.is_some(),
396            "escaped backslash contains should use FSST"
397        );
398        assert_arrays_eq!(
399            result.unwrap(),
400            BoolArray::from_iter([false, false, true, false, false, true, false]),
401            &mut ctx
402        );
403
404        Ok(())
405    }
406
407    #[test]
408    fn test_like_long_prefix_handled_by_flat_dfa() -> VortexResult<()> {
409        let fsst = make_fsst(
410            &[
411                Some("abcdefghijklmn-tail"),
412                Some("abcdefghijklmx-tail"),
413                Some("abcdefghijklmn"),
414            ],
415            Nullability::NonNullable,
416        );
417        let pattern = "abcdefghijklmn%";
418
419        let fsst = fsst.as_view();
420        let direct = <FSST as LikeKernel>::like(
421            fsst,
422            &ConstantArray::new(pattern, fsst.len()).into_array(),
423            LikeOptions::default(),
424            &mut SESSION.create_execution_ctx(),
425        )?;
426        assert!(
427            direct.is_some(),
428            "14-byte prefixes are now handled by the flat prefix DFA"
429        );
430        assert_arrays_eq!(
431            direct.unwrap(),
432            BoolArray::from_iter([true, false, true]),
433            &mut SESSION.create_execution_ctx()
434        );
435        Ok(())
436    }
437
438    #[test]
439    fn test_like_long_contains_falls_back_but_still_matches() -> VortexResult<()> {
440        let needle = "a".repeat(255);
441        let matching = format!("xx{needle}yy");
442        let non_matching = format!("xx{}byy", "a".repeat(254));
443        let exact = needle.clone();
444        let pattern = format!("%{needle}%");
445
446        let fsst = make_fsst(
447            &[Some(&matching), Some(&non_matching), Some(&exact)],
448            Nullability::NonNullable,
449        );
450
451        let fsst_v = fsst.as_view();
452        let direct = <FSST as LikeKernel>::like(
453            fsst_v,
454            &ConstantArray::new(pattern.as_str(), fsst.len()).into_array(),
455            LikeOptions::default(),
456            &mut SESSION.create_execution_ctx(),
457        )?;
458        assert!(
459            direct.is_none(),
460            "contains needles longer than 254 bytes exceed the DFA's u8 state space"
461        );
462
463        let result = like(fsst, &pattern)?;
464        assert_arrays_eq!(
465            &result,
466            &BoolArray::from_iter([true, false, true]),
467            &mut SESSION.create_execution_ctx()
468        );
469        Ok(())
470    }
471
472    #[test]
473    fn test_like_contains_len_254_kernel_handles() -> VortexResult<()> {
474        let needle = "a".repeat(254);
475        let matching = format!("xx{needle}yy");
476        let non_matching = format!("xx{}byy", "a".repeat(253));
477        let pattern = format!("%{needle}%");
478
479        let fsst = make_fsst(
480            &[Some(&matching), Some(&non_matching), Some(needle.as_str())],
481            Nullability::NonNullable,
482        );
483
484        let fsst = fsst.as_view();
485        let direct = <FSST as LikeKernel>::like(
486            fsst,
487            &ConstantArray::new(pattern.as_str(), fsst.len()).into_array(),
488            LikeOptions::default(),
489            &mut SESSION.create_execution_ctx(),
490        )?;
491        assert!(
492            direct.is_some(),
493            "254-byte contains needle should stay on the DFA path"
494        );
495        assert_arrays_eq!(
496            direct.unwrap(),
497            BoolArray::from_iter([true, false, true]),
498            &mut SESSION.create_execution_ctx()
499        );
500        Ok(())
501    }
502}
vortex_fsst/compute/like.rs

vortex_fsst/compute/
like.rs
