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, Int32Vector, NullVector, PrimitiveVector, StringVector, UInt32Vector,
47    };
48
49    fn check_take_primitive<T>(
50        input: Vec<Option<T::Native>>,
51        indices: Vec<Option<u32>>,
52        expect: Vec<Option<T::Native>>,
53    ) where
54        T: LogicalPrimitiveType,
55        PrimitiveArray<T::ArrowPrimitive>: From<Vec<Option<T::Native>>>,
56    {
57        let v = PrimitiveVector::<T>::new(PrimitiveArray::<T::ArrowPrimitive>::from(input));
58        let indices = UInt32Vector::new(UInt32Array::from(indices));
59        let output = v.take(&indices).unwrap();
60
61        let expected: VectorRef = Arc::new(PrimitiveVector::<T>::new(PrimitiveArray::<
62            T::ArrowPrimitive,
63        >::from(expect)));
64        assert_eq!(expected, output);
65    }
66
67    macro_rules! take_time_like_test {
68        ($VectorType: ident, $ValueType: ident, $method: ident) => {{
69            use $crate::vectors::{VectorRef, $VectorType};
70
71            let v = $VectorType::from_iterator((0..5).map($ValueType::$method));
72            let indices = UInt32Vector::from_slice(&[3, 0, 1, 4]);
73            let out = v.take(&indices).unwrap();
74
75            let expect: VectorRef = Arc::new($VectorType::from_iterator(
76                [3, 0, 1, 4].into_iter().map($ValueType::$method),
77            ));
78            assert_eq!(expect, out);
79        }};
80    }
81
82    #[test]
83    fn test_take_primitive() {
84        // nullable int32
85        check_take_primitive::<crate::types::Int32Type>(
86            vec![Some(1), None, Some(3), Some(4), Some(-5)],
87            vec![Some(3), None, Some(0), Some(1), Some(4)],
88            vec![Some(4), None, Some(1), None, Some(-5)],
89        );
90
91        // nullable float32
92        check_take_primitive::<crate::types::Float32Type>(
93            vec![Some(3.24), None, Some(1.34), Some(4.13), Some(5.13)],
94            vec![Some(3), None, Some(0), Some(1), Some(4)],
95            vec![Some(4.13), None, Some(3.24), None, Some(5.13)],
96        );
97
98        // nullable uint32
99        check_take_primitive::<crate::types::UInt32Type>(
100            vec![Some(0), None, Some(2), Some(3), Some(4)],
101            vec![Some(4), None, Some(2), Some(1), Some(3)],
102            vec![Some(4), None, Some(2), None, Some(3)],
103        );
104
105        // test date like type
106        take_time_like_test!(DateVector, Date, new);
107        take_time_like_test!(TimestampSecondVector, TimestampSecond, from_native);
108        take_time_like_test!(
109            TimestampMillisecondVector,
110            TimestampMillisecond,
111            from_native
112        );
113        take_time_like_test!(
114            TimestampMicrosecondVector,
115            TimestampMicrosecond,
116            from_native
117        );
118        take_time_like_test!(TimestampNanosecondVector, TimestampNanosecond, from_native);
119    }
120
121    #[test]
122    #[should_panic]
123    fn test_take_out_of_index() {
124        let v = Int32Vector::from_slice([1, 2, 3, 4, 5]);
125        let indies = UInt32Vector::from_slice([1, 5, 6]);
126        let _ = v.take(&indies);
127    }
128
129    #[test]
130    fn test_take_null() {
131        let v = NullVector::new(5);
132        let indices = UInt32Vector::from_slice([1, 3, 2]);
133        let out = v.take(&indices).unwrap();
134
135        let expect: VectorRef = Arc::new(NullVector::new(3));
136        assert_eq!(expect, out);
137    }
138
139    #[test]
140    fn test_take_scalar() {
141        let v = StringVector::from_slice(&["0", "1", "2", "3"]);
142        let indices = UInt32Vector::from_slice([1, 3, 2]);
143        let out = v.take(&indices).unwrap();
144
145        let expect: VectorRef = Arc::new(StringVector::from_slice(&["1", "3", "2"]));
146        assert_eq!(expect, out);
147    }
148
149    #[test]
150    fn test_take_bool() {
151        let v = BooleanVector::from_slice(&[false, true, false, true, false, false, true]);
152        let indices = UInt32Vector::from_slice([1, 3, 5, 6]);
153        let out = v.take(&indices).unwrap();
154        let expected: VectorRef = Arc::new(BooleanVector::from_slice(&[true, true, false, true]));
155        assert_eq!(out, expected);
156
157        let v = BooleanVector::from(vec![
158            Some(true),
159            None,
160            Some(false),
161            Some(true),
162            Some(false),
163            Some(false),
164            Some(true),
165            None,
166        ]);
167        let indices = UInt32Vector::from(vec![Some(1), None, Some(3), Some(5), Some(6)]);
168        let out = v.take(&indices).unwrap();
169        let expected: VectorRef = Arc::new(BooleanVector::from(vec![
170            None,
171            None,
172            Some(true),
173            Some(false),
174            Some(true),
175        ]));
176        assert_eq!(out, expected);
177    }
178}
datatypes/vectors/operations/take.rs

datatypes/vectors/operations/
take.rs
