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

vortex_fsst/compute/
like.rs
