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

datatypes/
vectors.rs
