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

vortex_fsst/compute/
like.rs
