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

vortex_array/arrays/dict/compute/
like.rs
