datatypes/
vectors.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 std::any::Any;
16use std::fmt::Debug;
17use std::sync::Arc;
18
19use arrow::array::{Array, ArrayRef};
20use snafu::ensure;
21
22use crate::data_type::ConcreteDataType;
23use crate::error::{self, Result};
24use crate::serialize::Serializable;
25use crate::value::{Value, ValueRef};
26use crate::vectors::operations::VectorOp;
27
28mod binary;
29mod boolean;
30mod date;
31mod decimal;
32mod dictionary;
33mod duration;
34mod eq;
35mod helper;
36mod interval;
37pub mod json;
38mod list;
39mod null;
40pub(crate) mod operations;
41mod primitive;
42mod string;
43mod struct_vector;
44mod time;
45mod timestamp;
46mod validity;
47
48pub use binary::{BinaryVector, BinaryVectorBuilder};
49pub use boolean::{BooleanVector, BooleanVectorBuilder};
50pub use date::{DateVector, DateVectorBuilder};
51pub use decimal::{Decimal128Vector, Decimal128VectorBuilder};
52pub(crate) use dictionary::StringDictionaryVectorBuilder;
53pub use dictionary::{DictionaryIter, DictionaryVector};
54pub use duration::{
55    DurationMicrosecondVector, DurationMicrosecondVectorBuilder, DurationMillisecondVector,
56    DurationMillisecondVectorBuilder, DurationNanosecondVector, DurationNanosecondVectorBuilder,
57    DurationSecondVector, DurationSecondVectorBuilder,
58};
59pub use helper::Helper;
60pub use interval::{
61    IntervalDayTimeVector, IntervalDayTimeVectorBuilder, IntervalMonthDayNanoVector,
62    IntervalMonthDayNanoVectorBuilder, IntervalYearMonthVector, IntervalYearMonthVectorBuilder,
63};
64pub use list::{ListIter, ListVector, ListVectorBuilder};
65pub use null::{NullVector, NullVectorBuilder};
66pub use primitive::{
67    Float32Vector, Float32VectorBuilder, Float64Vector, Float64VectorBuilder, Int8Vector,
68    Int8VectorBuilder, Int16Vector, Int16VectorBuilder, Int32Vector, Int32VectorBuilder,
69    Int64Vector, Int64VectorBuilder, PrimitiveIter, PrimitiveVector, PrimitiveVectorBuilder,
70    UInt8Vector, UInt8VectorBuilder, UInt16Vector, UInt16VectorBuilder, UInt32Vector,
71    UInt32VectorBuilder, UInt64Vector, UInt64VectorBuilder,
72};
73pub use string::{StringVector, StringVectorBuilder};
74pub use struct_vector::{StructVector, StructVectorBuilder};
75pub use time::{
76    TimeMicrosecondVector, TimeMicrosecondVectorBuilder, TimeMillisecondVector,
77    TimeMillisecondVectorBuilder, TimeNanosecondVector, TimeNanosecondVectorBuilder,
78    TimeSecondVector, TimeSecondVectorBuilder,
79};
80pub use timestamp::{
81    TimestampMicrosecondVector, TimestampMicrosecondVectorBuilder, TimestampMillisecondVector,
82    TimestampMillisecondVectorBuilder, TimestampNanosecondVector, TimestampNanosecondVectorBuilder,
83    TimestampSecondVector, TimestampSecondVectorBuilder,
84};
85pub use validity::Validity;
86
87// TODO(yingwen): arrow 28.0 implements Clone for all arrays, we could upgrade to it and simplify
88// some codes in methods such as `to_arrow_array()` and `to_boxed_arrow_array()`.
89/// Vector of data values.
90pub trait Vector: Send + Sync + Serializable + Debug + VectorOp {
91    /// Returns the data type of the vector.
92    ///
93    /// This may require heap allocation.
94    fn data_type(&self) -> ConcreteDataType;
95
96    fn vector_type_name(&self) -> String;
97
98    /// Returns the vector as [Any](std::any::Any) so that it can be
99    /// downcast to a specific implementation.
100    fn as_any(&self) -> &dyn Any;
101
102    /// Returns number of elements in the vector.
103    fn len(&self) -> usize;
104
105    /// Returns whether the vector is empty.
106    fn is_empty(&self) -> bool {
107        self.len() == 0
108    }
109
110    /// Convert this vector to a new arrow [ArrayRef].
111    fn to_arrow_array(&self) -> ArrayRef;
112
113    /// Convert this vector to a new boxed arrow [Array].
114    fn to_boxed_arrow_array(&self) -> Box<dyn Array>;
115
116    /// Returns the validity of the Array.
117    fn validity(&self) -> Validity;
118
119    /// Returns the memory size of vector.
120    fn memory_size(&self) -> usize;
121
122    /// The number of null slots on this [`Vector`].
123    /// # Implementation
124    /// This is `O(1)`.
125    fn null_count(&self) -> usize;
126
127    /// Returns whether row is null.
128    fn is_null(&self, row: usize) -> bool;
129
130    /// If the vector only contains NULL.
131    fn only_null(&self) -> bool {
132        self.null_count() == self.len()
133    }
134
135    /// Slices the `Vector`, returning a new `VectorRef`.
136    ///
137    /// # Panics
138    /// This function panics if `offset + length > self.len()`.
139    fn slice(&self, offset: usize, length: usize) -> VectorRef;
140
141    /// Returns the clone of value at `index`.
142    ///
143    /// # Panics
144    /// Panic if `index` is out of bound.
145    fn get(&self, index: usize) -> Value;
146
147    /// Returns the clone of value at `index` or error if `index`
148    /// is out of bound.
149    fn try_get(&self, index: usize) -> Result<Value> {
150        ensure!(
151            index < self.len(),
152            error::BadArrayAccessSnafu {
153                index,
154                size: self.len()
155            }
156        );
157        Ok(self.get(index))
158    }
159
160    /// Returns the reference of value at `index`.
161    ///
162    /// # Panics
163    /// Panic if `index` is out of bound.
164    fn get_ref(&self, index: usize) -> ValueRef<'_>;
165}
166
167pub type VectorRef = Arc<dyn Vector>;
168
169/// Mutable vector that could be used to build an immutable vector.
170pub trait MutableVector: Send + Sync {
171    /// Returns the data type of the vector.
172    fn data_type(&self) -> ConcreteDataType;
173
174    /// Returns the length of the vector.
175    fn len(&self) -> usize;
176
177    /// Returns whether the vector is empty.
178    fn is_empty(&self) -> bool {
179        self.len() == 0
180    }
181
182    /// Convert to Any, to enable dynamic casting.
183    fn as_any(&self) -> &dyn Any;
184
185    /// Convert to mutable Any, to enable dynamic casting.
186    fn as_mut_any(&mut self) -> &mut dyn Any;
187
188    /// Convert `self` to an (immutable) [VectorRef] and reset `self`.
189    fn to_vector(&mut self) -> VectorRef;
190
191    /// Convert `self` to an (immutable) [VectorRef] and without resetting `self`.
192    fn to_vector_cloned(&self) -> VectorRef;
193
194    /// Try to push value ref to this mutable vector.
195    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()>;
196
197    /// Push value ref to this mutable vector.
198    ///
199    /// # Panics
200    /// Panics if error if data types mismatch.
201    fn push_value_ref(&mut self, value: &ValueRef) {
202        self.try_push_value_ref(value).unwrap_or_else(|_| {
203            panic!(
204                "expecting pushing value of datatype {:?}, actual {:?}",
205                self.data_type(),
206                value
207            );
208        });
209    }
210
211    /// Push null to this mutable vector.
212    fn push_null(&mut self);
213
214    /// Push nulls to this mutable vector.
215    fn push_nulls(&mut self, num_nulls: usize) {
216        for _ in 0..num_nulls {
217            self.push_null();
218        }
219    }
220
221    /// Extend this mutable vector by slice of `vector`.
222    ///
223    /// Returns error if data types mismatch.
224    ///
225    /// # Panics
226    /// Panics if `offset + length > vector.len()`.
227    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()>;
228}
229
230/// Helper to define `try_from_arrow_array(array: arrow::array::ArrayRef)` function.
231macro_rules! impl_try_from_arrow_array_for_vector {
232    ($Array: ident, $Vector: ident) => {
233        impl $Vector {
234            pub fn try_from_arrow_array(
235                array: impl AsRef<dyn arrow::array::Array>,
236            ) -> crate::error::Result<$Vector> {
237                use snafu::OptionExt;
238
239                let arrow_array = array
240                    .as_ref()
241                    .as_any()
242                    .downcast_ref::<$Array>()
243                    .with_context(|| crate::error::ConversionSnafu {
244                        from: std::format!("{:?}", array.as_ref().data_type()),
245                    })?
246                    .clone();
247
248                Ok($Vector::from(arrow_array))
249            }
250        }
251    };
252}
253
254macro_rules! impl_validity_for_vector {
255    ($array: expr) => {
256        Validity::from_array_data($array.to_data())
257    };
258}
259
260macro_rules! impl_get_for_vector {
261    ($array: expr, $index: ident) => {
262        if $array.is_valid($index) {
263            // Safety: The index have been checked by `is_valid()`.
264            unsafe { $array.value_unchecked($index).into() }
265        } else {
266            Value::Null
267        }
268    };
269}
270
271macro_rules! impl_get_ref_for_vector {
272    ($array: expr, $index: ident) => {
273        if $array.is_valid($index) {
274            // Safety: The index have been checked by `is_valid()`.
275            unsafe { $array.value_unchecked($index).into() }
276        } else {
277            ValueRef::Null
278        }
279    };
280}
281
282macro_rules! impl_extend_for_builder {
283    ($mutable_vector: expr, $vector: ident, $VectorType: ident, $offset: ident, $length: ident) => {{
284        use snafu::OptionExt;
285
286        let sliced_vector = $vector.slice($offset, $length);
287        let concrete_vector = sliced_vector
288            .as_any()
289            .downcast_ref::<$VectorType>()
290            .with_context(|| crate::error::CastTypeSnafu {
291                msg: format!(
292                    "Failed to cast vector from {} to {}",
293                    $vector.vector_type_name(),
294                    stringify!($VectorType)
295                ),
296            })?;
297        for value in concrete_vector.iter_data() {
298            $mutable_vector.push(value);
299        }
300        Ok(())
301    }};
302}
303
304pub(crate) use impl_extend_for_builder;
305pub(crate) use impl_get_for_vector;
306pub(crate) use impl_get_ref_for_vector;
307pub(crate) use impl_try_from_arrow_array_for_vector;
308pub(crate) use impl_validity_for_vector;
309
310#[cfg(test)]
311pub mod tests {
312    use arrow::array::{Array, Int32Array, UInt8Array};
313    use paste::paste;
314    use serde_json;
315
316    use super::*;
317    use crate::data_type::DataType;
318    use crate::prelude::ScalarVectorBuilder;
319    use crate::types::{Int32Type, LogicalPrimitiveType};
320    use crate::vectors::helper::Helper;
321
322    #[test]
323    fn test_df_columns_to_vector() {
324        let df_column: Arc<dyn Array> = Arc::new(Int32Array::from(vec![1, 2, 3]));
325        let vector = Helper::try_into_vector(df_column).unwrap();
326        assert_eq!(
327            Int32Type::build_data_type().as_arrow_type(),
328            vector.data_type().as_arrow_type()
329        );
330    }
331
332    #[test]
333    fn test_serialize_i32_vector() {
334        let df_column: Arc<dyn Array> = Arc::new(Int32Array::from(vec![1, 2, 3]));
335        let json_value = Helper::try_into_vector(df_column)
336            .unwrap()
337            .serialize_to_json()
338            .unwrap();
339        assert_eq!("[1,2,3]", serde_json::to_string(&json_value).unwrap());
340    }
341
342    #[test]
343    fn test_serialize_i8_vector() {
344        let df_column: Arc<dyn Array> = Arc::new(UInt8Array::from(vec![1, 2, 3]));
345        let json_value = Helper::try_into_vector(df_column)
346            .unwrap()
347            .serialize_to_json()
348            .unwrap();
349        assert_eq!("[1,2,3]", serde_json::to_string(&json_value).unwrap());
350    }
351
352    #[test]
353    fn test_mutable_vector_data_type() {
354        macro_rules! mutable_primitive_data_type_eq_with_lower {
355            ($($type: ident),*) => {
356                $(
357                    paste! {
358                        let mutable_vector = [<$type VectorBuilder>]::with_capacity(1024);
359                        assert_eq!(mutable_vector.data_type(), ConcreteDataType::[<$type:lower _datatype>]());
360                    }
361                )*
362            };
363        }
364
365        macro_rules! mutable_time_data_type_eq_with_snake {
366            ($($type: ident),*) => {
367                $(
368                    paste! {
369                        let mutable_vector = [<$type VectorBuilder>]::with_capacity(1024);
370                        assert_eq!(mutable_vector.data_type(), ConcreteDataType::[<$type:snake _datatype>]());
371                    }
372                )*
373            };
374        }
375        // Test Primitive types
376        mutable_primitive_data_type_eq_with_lower!(
377            Boolean, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, Float32, Float64,
378            Date, Binary, String
379        );
380
381        // Test types about time
382        mutable_time_data_type_eq_with_snake!(
383            TimeSecond,
384            TimeMillisecond,
385            TimeMicrosecond,
386            TimeNanosecond,
387            TimestampSecond,
388            TimestampMillisecond,
389            TimestampMicrosecond,
390            TimestampNanosecond,
391            DurationSecond,
392            DurationMillisecond,
393            DurationMicrosecond,
394            DurationNanosecond,
395            IntervalYearMonth,
396            IntervalDayTime,
397            IntervalMonthDayNano
398        );
399
400        // Null type
401        let builder = NullVectorBuilder::default();
402        assert_eq!(builder.data_type(), ConcreteDataType::null_datatype());
403
404        // Decimal128 type
405        let builder = Decimal128VectorBuilder::with_capacity(1024);
406        assert_eq!(
407            builder.data_type(),
408            ConcreteDataType::decimal128_datatype(38, 10)
409        );
410
411        let builder = Decimal128VectorBuilder::with_capacity(1024)
412            .with_precision_and_scale(3, 2)
413            .unwrap();
414        assert_eq!(
415            builder.data_type(),
416            ConcreteDataType::decimal128_datatype(3, 2)
417        );
418    }
419
420    #[test]
421    #[should_panic(expected = "Must use ListVectorBuilder::with_type_capacity()")]
422    fn test_mutable_vector_list_data_type() {
423        let item_type = Arc::new(ConcreteDataType::int32_datatype());
424        // List type
425        let builder = ListVectorBuilder::with_type_capacity(item_type.clone(), 1024);
426        assert_eq!(
427            builder.data_type(),
428            ConcreteDataType::list_datatype(item_type)
429        );
430
431        // Panic with_capacity
432        let _ = ListVectorBuilder::with_capacity(1024);
433    }
434
435    #[test]
436    fn test_mutable_vector_to_vector_cloned() {
437        // create a string vector builder
438        let mut builder = ConcreteDataType::string_datatype().create_mutable_vector(1024);
439        builder.push_value_ref(&ValueRef::String("hello"));
440        builder.push_value_ref(&ValueRef::String("world"));
441        builder.push_value_ref(&ValueRef::String("!"));
442
443        // use MutableVector trait to_vector_cloned won't reset builder
444        let vector = builder.to_vector_cloned();
445        assert_eq!(vector.len(), 3);
446        assert_eq!(builder.len(), 3);
447    }
448}
datatypes/vectors.rs

datatypes/
vectors.rs
