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

vortex_fsst/compute/
like.rs
