datatypes/vectors/operations/
replicate.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
15use crate::prelude::*;
16pub(crate) use crate::vectors::decimal::replicate_decimal128;
17pub(crate) use crate::vectors::null::replicate_null;
18pub(crate) use crate::vectors::primitive::replicate_primitive;
19use crate::vectors::{ListVector, ListVectorBuilder};
20
21pub(crate) fn replicate_scalar<C: ScalarVector>(c: &C, offsets: &[usize]) -> VectorRef {
22    assert_eq!(offsets.len(), c.len());
23
24    if offsets.is_empty() {
25        return c.slice(0, 0);
26    }
27    let mut builder = <<C as ScalarVector>::Builder>::with_capacity(c.len());
28
29    let mut previous_offset = 0;
30    for (i, offset) in offsets.iter().enumerate() {
31        let data = c.get_data(i);
32        for _ in previous_offset..*offset {
33            builder.push(data.clone());
34        }
35        previous_offset = *offset;
36    }
37    builder.to_vector()
38}
39
40pub(crate) fn replicate_list(c: &ListVector, offsets: &[usize]) -> VectorRef {
41    assert_eq!(offsets.len(), c.len());
42
43    if offsets.is_empty() {
44        return c.slice(0, 0);
45    }
46    let mut builder =
47        ListVectorBuilder::with_type_capacity(c.item_type(), *offsets.last().unwrap());
48
49    let mut previous_offset = 0;
50    for (i, offset) in offsets.iter().enumerate() {
51        let data = c.get_data(i);
52        for _ in previous_offset..*offset {
53            builder.push(data.clone());
54        }
55        previous_offset = *offset;
56    }
57    builder.to_vector()
58}
59
60#[cfg(test)]
61mod tests {
62    use std::sync::Arc;
63
64    use common_time::timestamp::TimeUnit;
65    use common_time::{Date, Timestamp};
66    use paste::paste;
67
68    use super::*;
69    use crate::value::{ListValue, ListValueRef};
70    use crate::vectors::constant::ConstantVector;
71    use crate::vectors::{
72        Decimal128Vector, Int32Vector, ListVectorBuilder, NullVector, StringVector, VectorOp,
73    };
74
75    #[test]
76    fn test_replicate_primitive() {
77        let v = Int32Vector::from_iterator(0..5);
78        let offsets = [0, 1, 2, 3, 4];
79
80        let v = v.replicate(&offsets);
81        assert_eq!(4, v.len());
82
83        for i in 0..4 {
84            assert_eq!(Value::Int32(i as i32 + 1), v.get(i));
85        }
86    }
87
88    #[test]
89    fn test_replicate_nullable_primitive() {
90        let v = Int32Vector::from(vec![None, Some(1), None, Some(2)]);
91        let offsets = [2, 4, 6, 8];
92        let v = v.replicate(&offsets);
93        assert_eq!(8, v.len());
94
95        let expect: VectorRef = Arc::new(Int32Vector::from(vec![
96            None,
97            None,
98            Some(1),
99            Some(1),
100            None,
101            None,
102            Some(2),
103            Some(2),
104        ]));
105        assert_eq!(expect, v);
106    }
107
108    #[test]
109    fn test_replicate_scalar() {
110        let v = StringVector::from_slice(&["0", "1", "2", "3"]);
111        let offsets = [1, 3, 5, 6];
112
113        let v = v.replicate(&offsets);
114        assert_eq!(6, v.len());
115
116        let expect: VectorRef = Arc::new(StringVector::from_slice(&["0", "1", "1", "2", "2", "3"]));
117        assert_eq!(expect, v);
118    }
119
120    #[test]
121    fn test_replicate_list() {
122        let item_type = Arc::new(ConcreteDataType::int32_datatype());
123        let first = ListValue::new(vec![Value::Int32(1), Value::Int32(2)], item_type.clone());
124        let second = ListValue::new(vec![Value::Int32(3)], item_type.clone());
125        let mut builder = ListVectorBuilder::with_type_capacity(item_type, 2);
126        builder.push(Some(ListValueRef::Ref { val: &first }));
127        builder.push(Some(ListValueRef::Ref { val: &second }));
128        let v = builder.finish();
129
130        let v = v.replicate(&[1, 3]);
131        let mut expect_builder =
132            ListVectorBuilder::with_type_capacity(Arc::new(ConcreteDataType::int32_datatype()), 3);
133        expect_builder.push(Some(ListValueRef::Ref { val: &first }));
134        expect_builder.push(Some(ListValueRef::Ref { val: &second }));
135        expect_builder.push(Some(ListValueRef::Ref { val: &second }));
136        let expect = expect_builder.to_vector();
137        assert_eq!(expect, v);
138    }
139
140    #[test]
141    fn test_replicate_constant() {
142        let v = Arc::new(StringVector::from_slice(&["hello"]));
143        let cv = ConstantVector::new(v.clone(), 2);
144        let offsets = [1, 4];
145
146        let cv = cv.replicate(&offsets);
147        assert_eq!(4, cv.len());
148
149        let expect: VectorRef = Arc::new(ConstantVector::new(v, 4));
150        assert_eq!(expect, cv);
151    }
152
153    #[test]
154    fn test_replicate_null() {
155        let v = NullVector::new(0);
156        let offsets = [];
157        let v = v.replicate(&offsets);
158        assert!(v.is_empty());
159
160        let v = NullVector::new(3);
161        let offsets = [1, 3, 5];
162
163        let v = v.replicate(&offsets);
164        assert_eq!(5, v.len());
165    }
166
167    macro_rules! impl_replicate_date_like_test {
168        ($VectorType: ident, $ValueType: ident, $method: ident) => {{
169            use $crate::vectors::$VectorType;
170
171            let v = $VectorType::from_iterator((0..5).map($ValueType::$method));
172            let offsets = [0, 1, 2, 3, 4];
173
174            let v = v.replicate(&offsets);
175            assert_eq!(4, v.len());
176
177            for i in 0..4 {
178                assert_eq!(
179                    Value::$ValueType($ValueType::$method((i as i32 + 1).into())),
180                    v.get(i)
181                );
182            }
183        }};
184    }
185
186    macro_rules! impl_replicate_timestamp_test {
187        ($unit: ident) => {{
188            paste!{
189                use $crate::vectors::[<Timestamp $unit Vector>];
190                use $crate::timestamp::[<Timestamp $unit>];
191                let v = [<Timestamp $unit Vector>]::from_iterator((0..5).map([<Timestamp $unit>]::from));
192                let offsets = [0, 1, 2, 3, 4];
193                let v = v.replicate(&offsets);
194                assert_eq!(4, v.len());
195                for i in 0..4 {
196                    assert_eq!(
197                        Value::Timestamp(Timestamp::new(i as i64 + 1, TimeUnit::$unit)),
198                        v.get(i)
199                    );
200                }
201            }
202        }};
203    }
204
205    #[test]
206    fn test_replicate_date_like() {
207        impl_replicate_date_like_test!(DateVector, Date, new);
208        impl_replicate_timestamp_test!(Second);
209        impl_replicate_timestamp_test!(Millisecond);
210        impl_replicate_timestamp_test!(Microsecond);
211        impl_replicate_timestamp_test!(Nanosecond);
212    }
213
214    #[test]
215    fn test_replicate_decimal() {
216        let data = vec![100];
217        // create a decimal vector
218        let v = Decimal128Vector::from_values(data.clone())
219            .with_precision_and_scale(10, 2)
220            .unwrap();
221        let offsets = [5];
222        let v = v.replicate(&offsets);
223        assert_eq!(5, v.len());
224
225        let expect: VectorRef = Arc::new(
226            Decimal128Vector::from_values(vec![100; 5])
227                .with_precision_and_scale(10, 2)
228                .unwrap(),
229        );
230        assert_eq!(expect, v);
231    }
232}
datatypes/vectors/operations/replicate.rs

datatypes/vectors/operations/
replicate.rs
