Skip to main content

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