datatypes/vectors/operations/
take.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15macro_rules! take_indices {
16    ($vector: expr, $VectorType: ty, $indices: ident) => {{
17        use std::sync::Arc;
18
19        use arrow::compute;
20        use snafu::ResultExt;
21
22        let arrow_array = $vector.as_arrow();
23        let taken = compute::take(arrow_array, $indices.as_arrow(), None)
24            .context(crate::error::ArrowComputeSnafu)?;
25        Ok(Arc::new(<$VectorType>::try_from_arrow_array(taken)?))
26    }};
27}
28
29pub(crate) use take_indices;
30
31#[cfg(test)]
32mod tests {
33    use std::sync::Arc;
34
35    use arrow::array::{PrimitiveArray, UInt32Array};
36    use common_time::Date;
37
38    use crate::prelude::VectorRef;
39    use crate::scalars::ScalarVector;
40    use crate::timestamp::{
41        TimestampMicrosecond, TimestampMillisecond, TimestampNanosecond, TimestampSecond,
42    };
43    use crate::types::{LogicalPrimitiveType, WrapperType};
44    use crate::vectors::operations::VectorOp;
45    use crate::vectors::{
46        BooleanVector, ConstantVector, Int32Vector, NullVector, PrimitiveVector, StringVector,
47        UInt32Vector,
48    };
49
50    fn check_take_primitive<T>(
51        input: Vec<Option<T::Native>>,
52        indices: Vec<Option<u32>>,
53        expect: Vec<Option<T::Native>>,
54    ) where
55        T: LogicalPrimitiveType,
56        PrimitiveArray<T::ArrowPrimitive>: From<Vec<Option<T::Native>>>,
57    {
58        let v = PrimitiveVector::<T>::new(PrimitiveArray::<T::ArrowPrimitive>::from(input));
59        let indices = UInt32Vector::new(UInt32Array::from(indices));
60        let output = v.take(&indices).unwrap();
61
62        let expected: VectorRef = Arc::new(PrimitiveVector::<T>::new(PrimitiveArray::<
63            T::ArrowPrimitive,
64        >::from(expect)));
65        assert_eq!(expected, output);
66    }
67
68    macro_rules! take_time_like_test {
69        ($VectorType: ident, $ValueType: ident, $method: ident) => {{
70            use $crate::vectors::{$VectorType, VectorRef};
71
72            let v = $VectorType::from_iterator((0..5).map($ValueType::$method));
73            let indices = UInt32Vector::from_slice(&[3, 0, 1, 4]);
74            let out = v.take(&indices).unwrap();
75
76            let expect: VectorRef = Arc::new($VectorType::from_iterator(
77                [3, 0, 1, 4].into_iter().map($ValueType::$method),
78            ));
79            assert_eq!(expect, out);
80        }};
81    }
82
83    #[test]
84    fn test_take_primitive() {
85        // nullable int32
86        check_take_primitive::<crate::types::Int32Type>(
87            vec![Some(1), None, Some(3), Some(4), Some(-5)],
88            vec![Some(3), None, Some(0), Some(1), Some(4)],
89            vec![Some(4), None, Some(1), None, Some(-5)],
90        );
91
92        // nullable float32
93        check_take_primitive::<crate::types::Float32Type>(
94            vec![Some(3.24), None, Some(1.34), Some(4.13), Some(5.13)],
95            vec![Some(3), None, Some(0), Some(1), Some(4)],
96            vec![Some(4.13), None, Some(3.24), None, Some(5.13)],
97        );
98
99        // nullable uint32
100        check_take_primitive::<crate::types::UInt32Type>(
101            vec![Some(0), None, Some(2), Some(3), Some(4)],
102            vec![Some(4), None, Some(2), Some(1), Some(3)],
103            vec![Some(4), None, Some(2), None, Some(3)],
104        );
105
106        // test date like type
107        take_time_like_test!(DateVector, Date, new);
108        take_time_like_test!(TimestampSecondVector, TimestampSecond, from_native);
109        take_time_like_test!(
110            TimestampMillisecondVector,
111            TimestampMillisecond,
112            from_native
113        );
114        take_time_like_test!(
115            TimestampMicrosecondVector,
116            TimestampMicrosecond,
117            from_native
118        );
119        take_time_like_test!(TimestampNanosecondVector, TimestampNanosecond, from_native);
120    }
121
122    fn check_take_constant(expect_length: usize, input_length: usize, indices: &[u32]) {
123        let v = ConstantVector::new(Arc::new(Int32Vector::from_slice([111])), input_length);
124        let indices = UInt32Vector::from_slice(indices);
125        let out = v.take(&indices).unwrap();
126
127        assert!(out.is_const());
128        assert_eq!(expect_length, out.len());
129    }
130
131    #[test]
132    fn test_take_constant() {
133        check_take_constant(2, 5, &[3, 4]);
134        check_take_constant(3, 10, &[1, 2, 3]);
135        check_take_constant(4, 10, &[1, 5, 3, 6]);
136        check_take_constant(5, 10, &[1, 9, 8, 7, 3]);
137    }
138
139    #[test]
140    #[should_panic]
141    fn test_take_constant_out_of_index() {
142        check_take_constant(2, 5, &[3, 5]);
143    }
144
145    #[test]
146    #[should_panic]
147    fn test_take_out_of_index() {
148        let v = Int32Vector::from_slice([1, 2, 3, 4, 5]);
149        let indies = UInt32Vector::from_slice([1, 5, 6]);
150        let _ = v.take(&indies);
151    }
152
153    #[test]
154    fn test_take_null() {
155        let v = NullVector::new(5);
156        let indices = UInt32Vector::from_slice([1, 3, 2]);
157        let out = v.take(&indices).unwrap();
158
159        let expect: VectorRef = Arc::new(NullVector::new(3));
160        assert_eq!(expect, out);
161    }
162
163    #[test]
164    fn test_take_scalar() {
165        let v = StringVector::from_slice(&["0", "1", "2", "3"]);
166        let indices = UInt32Vector::from_slice([1, 3, 2]);
167        let out = v.take(&indices).unwrap();
168
169        let expect: VectorRef = Arc::new(StringVector::from_slice(&["1", "3", "2"]));
170        assert_eq!(expect, out);
171    }
172
173    #[test]
174    fn test_take_bool() {
175        let v = BooleanVector::from_slice(&[false, true, false, true, false, false, true]);
176        let indices = UInt32Vector::from_slice([1, 3, 5, 6]);
177        let out = v.take(&indices).unwrap();
178        let expected: VectorRef = Arc::new(BooleanVector::from_slice(&[true, true, false, true]));
179        assert_eq!(out, expected);
180
181        let v = BooleanVector::from(vec![
182            Some(true),
183            None,
184            Some(false),
185            Some(true),
186            Some(false),
187            Some(false),
188            Some(true),
189            None,
190        ]);
191        let indices = UInt32Vector::from(vec![Some(1), None, Some(3), Some(5), Some(6)]);
192        let out = v.take(&indices).unwrap();
193        let expected: VectorRef = Arc::new(BooleanVector::from(vec![
194            None,
195            None,
196            Some(true),
197            Some(false),
198            Some(true),
199        ]));
200        assert_eq!(out, expected);
201    }
202}
datatypes/vectors/operations/take.rs

datatypes/vectors/operations/
take.rs
