vortex_array/arrays/dict/compute/
like.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use super::Dict;
7use super::DictArray;
8use crate::ArrayRef;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::ConstantArray;
12use crate::arrays::dict::DictArrayExt;
13use crate::arrays::dict::DictArraySlotsExt;
14use crate::optimizer::ArrayOptimizer;
15use crate::scalar_fn::fns::like::Like;
16use crate::scalar_fn::fns::like::LikeOptions;
17use crate::scalar_fn::fns::like::LikeReduce;
18
19impl LikeReduce for Dict {
20    fn like(
21        array: ArrayView<'_, Dict>,
22        pattern: &ArrayRef,
23        options: LikeOptions,
24    ) -> VortexResult<Option<ArrayRef>> {
25        // If we have more values than codes, it is faster to canonicalize first.
26        if array.values().len() > array.codes().len() {
27            return Ok(None);
28        }
29        if let Some(pattern) = pattern.as_constant() {
30            let pattern = ConstantArray::new(pattern, array.values().len()).into_array();
31
32            let values = Like::try_new(array.values().clone(), pattern, options)?
33                .into_array()
34                .optimize()?;
35
36            // SAFETY: LIKE preserves the len of the values, so codes are still pointing at
37            //  valid positions.
38            // Preserve all_values_referenced since codes are unchanged.
39            unsafe {
40                Ok(Some(
41                    DictArray::new_unchecked(array.codes().clone(), values)
42                        .set_all_values_referenced(array.has_all_values_referenced())
43                        .into_array(),
44                ))
45            }
46        } else {
47            Ok(None)
48        }
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use vortex_buffer::buffer;
55    use vortex_error::VortexResult;
56
57    use crate::IntoArray;
58    use crate::VortexSessionExecute;
59    use crate::array_session;
60    use crate::arrays::BoolArray;
61    use crate::arrays::DictArray;
62    use crate::arrays::VarBinArray;
63    use crate::arrays::dict::compute::like::ConstantArray;
64    use crate::assert_arrays_eq;
65    use crate::optimizer::ArrayOptimizer;
66    use crate::scalar_fn::fns::like::Like;
67    use crate::scalar_fn::fns::like::LikeOptions;
68
69    #[test]
70    fn like_reduce_dict() -> VortexResult<()> {
71        let mut ctx = array_session().create_execution_ctx();
72        let dict = DictArray::try_new(
73            buffer![0u8, 1, 0, 2].into_array(),
74            VarBinArray::from(vec!["hello", "world", "help"]).into_array(),
75        )?
76        .into_array();
77
78        let pattern = ConstantArray::new("hello%", 4).into_array();
79        let result = Like::try_new(dict, pattern, LikeOptions::default())?
80            .into_array()
81            .optimize()?;
82
83        assert_arrays_eq!(
84            result,
85            BoolArray::from_iter([true, false, true, false]),
86            &mut ctx
87        );
88        Ok(())
89    }
90}
vortex_array/arrays/dict/compute/like.rs

vortex_array/arrays/dict/compute/
like.rs
