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

datatypes/vectors/operations/
replicate.rs
