datatypes/vectors/
helper.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
15//! Vector helper functions, inspired by databend Series mod
16
17use std::any::Any;
18use std::sync::Arc;
19
20use arrow::array::{Array, ArrayRef, StringArray};
21use arrow::compute;
22use arrow::compute::kernels::comparison;
23use arrow::datatypes::{
24    DataType as ArrowDataType, Int8Type, Int16Type, Int32Type, Int64Type, TimeUnit, UInt8Type,
25    UInt16Type, UInt32Type, UInt64Type,
26};
27use arrow_array::{DictionaryArray, StructArray};
28use arrow_schema::IntervalUnit;
29use datafusion_common::ScalarValue;
30use snafu::{OptionExt, ResultExt};
31
32use crate::data_type::ConcreteDataType;
33use crate::error::{self, ConvertScalarToArrowArraySnafu, Result};
34use crate::prelude::DataType;
35use crate::scalars::Scalar;
36use crate::types::StructType;
37use crate::value::Value;
38use crate::vectors::struct_vector::StructVector;
39use crate::vectors::{
40    BinaryVector, BooleanVector, DateVector, Decimal128Vector, DictionaryVector,
41    DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector,
42    DurationSecondVector, Float32Vector, Float64Vector, Int8Vector, Int16Vector, Int32Vector,
43    Int64Vector, IntervalDayTimeVector, IntervalMonthDayNanoVector, IntervalYearMonthVector,
44    ListVector, NullVector, StringVector, TimeMicrosecondVector, TimeMillisecondVector,
45    TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, TimestampMillisecondVector,
46    TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt16Vector, UInt32Vector,
47    UInt64Vector, Vector, VectorRef,
48};
49
50/// Helper functions for `Vector`.
51pub struct Helper;
52
53impl Helper {
54    /// Get a pointer to the underlying data of this vectors.
55    /// Can be useful for fast comparisons.
56    /// # Safety
57    /// Assumes that the `vector` is  T.
58    pub unsafe fn static_cast<T: Any>(vector: &VectorRef) -> &T {
59        let object = vector.as_ref();
60        debug_assert!(object.as_any().is::<T>());
61        unsafe { &*(object as *const dyn Vector as *const T) }
62    }
63
64    pub fn check_get_scalar<T: Scalar>(vector: &VectorRef) -> Result<&<T as Scalar>::VectorType> {
65        vector
66            .as_any()
67            .downcast_ref::<<T as Scalar>::VectorType>()
68            .with_context(|| error::UnknownVectorSnafu {
69                msg: format!(
70                    "downcast vector error, vector type: {:?}, expected vector: {:?}",
71                    vector.vector_type_name(),
72                    std::any::type_name::<T>(),
73                ),
74            })
75    }
76
77    pub fn check_get<T: 'static + Vector>(vector: &VectorRef) -> Result<&T> {
78        vector
79            .as_any()
80            .downcast_ref::<T>()
81            .with_context(|| error::UnknownVectorSnafu {
82                msg: format!(
83                    "downcast vector error, vector type: {:?}, expected vector: {:?}",
84                    vector.vector_type_name(),
85                    std::any::type_name::<T>(),
86                ),
87            })
88    }
89
90    pub fn check_get_scalar_vector<T: Scalar>(
91        vector: &VectorRef,
92    ) -> Result<&<T as Scalar>::VectorType> {
93        vector
94            .as_any()
95            .downcast_ref::<<T as Scalar>::VectorType>()
96            .with_context(|| error::UnknownVectorSnafu {
97                msg: format!(
98                    "downcast vector error, vector type: {:?}, expected vector: {:?}",
99                    vector.vector_type_name(),
100                    std::any::type_name::<T>(),
101                ),
102            })
103    }
104
105    /// Materializes an Arrow scalar as a vector of the given length.
106    ///
107    /// With a type hint, casts the array to the requested representation and returns
108    /// an error if the vector cannot preserve that concrete type. Without a hint,
109    /// normalizes the scalar to a supported GreptimeDB representation.
110    pub fn try_from_scalar_value(
111        value: ScalarValue,
112        length: usize,
113        type_hint: Option<&ConcreteDataType>,
114    ) -> Result<VectorRef> {
115        if let Some(data_type) = type_hint {
116            let mut array = value
117                .to_array_of_size(length)
118                .context(ConvertScalarToArrowArraySnafu)?;
119            let arrow_type = data_type.as_arrow_type();
120            if array.data_type() != &arrow_type {
121                array = compute::cast(&array, &arrow_type).context(error::ArrowComputeSnafu)?;
122            }
123            let vector = Self::try_into_vector(array)?;
124            if &vector.data_type() != data_type {
125                return error::CastTypeSnafu {
126                    msg: format!(
127                        "Scalar materialization produced {:?}, expected {data_type:?}",
128                        vector.data_type()
129                    ),
130                }
131                .fail();
132            }
133            return Ok(vector);
134        }
135
136        let value = match value {
137            // GreptimeDB doesn't support Float16 vectors.
138            ScalarValue::Float16(v) => ScalarValue::Float32(v.map(f32::from)),
139            ScalarValue::LargeUtf8(v) => ScalarValue::Utf8(v),
140            ScalarValue::LargeBinary(v) | ScalarValue::FixedSizeBinary(_, v) => {
141                ScalarValue::Binary(v)
142            }
143            // Timezones are not supported by GreptimeDB vectors.
144            ScalarValue::TimestampSecond(v, _) => ScalarValue::TimestampSecond(v, None),
145            ScalarValue::TimestampMillisecond(v, _) => ScalarValue::TimestampMillisecond(v, None),
146            ScalarValue::TimestampMicrosecond(v, _) => ScalarValue::TimestampMicrosecond(v, None),
147            ScalarValue::TimestampNanosecond(v, _) => ScalarValue::TimestampNanosecond(v, None),
148            value @ (ScalarValue::Decimal32(_, _, _)
149            | ScalarValue::Decimal64(_, _, _)
150            | ScalarValue::Decimal256(_, _, _)
151            | ScalarValue::FixedSizeList(_)
152            | ScalarValue::LargeList(_)
153            | ScalarValue::ListView(_)
154            | ScalarValue::LargeListView(_)
155            | ScalarValue::Dictionary(_, _)
156            | ScalarValue::Union(_, _, _)
157            | ScalarValue::Utf8View(_)
158            | ScalarValue::BinaryView(_)
159            | ScalarValue::Map(_)
160            | ScalarValue::Date64(_)
161            | ScalarValue::RunEndEncoded(_, _, _)) => {
162                return error::ConversionSnafu {
163                    from: format!("Unsupported scalar value: {value}"),
164                }
165                .fail();
166            }
167            value => value,
168        };
169
170        let array = value
171            .to_array_of_size(length)
172            .context(ConvertScalarToArrowArraySnafu)?;
173        Self::try_into_vector(array)
174    }
175
176    /// Try to cast an arrow array into vector
177    ///
178    /// # Panics
179    /// Panic if given arrow data type is not supported.
180    pub fn try_into_vector(array: impl AsRef<dyn Array>) -> Result<VectorRef> {
181        Ok(match array.as_ref().data_type() {
182            ArrowDataType::Null => Arc::new(NullVector::try_from_arrow_array(array)?),
183            ArrowDataType::Boolean => Arc::new(BooleanVector::try_from_arrow_array(array)?),
184            ArrowDataType::Binary | ArrowDataType::BinaryView => {
185                Arc::new(BinaryVector::try_from_arrow_array(array)?)
186            }
187            ArrowDataType::LargeBinary | ArrowDataType::FixedSizeBinary(_) => {
188                let array = arrow::compute::cast(array.as_ref(), &ArrowDataType::Binary)
189                    .context(crate::error::ArrowComputeSnafu)?;
190                Arc::new(BinaryVector::try_from_arrow_array(array)?)
191            }
192            ArrowDataType::Int8 => Arc::new(Int8Vector::try_from_arrow_array(array)?),
193            ArrowDataType::Int16 => Arc::new(Int16Vector::try_from_arrow_array(array)?),
194            ArrowDataType::Int32 => Arc::new(Int32Vector::try_from_arrow_array(array)?),
195            ArrowDataType::Int64 => Arc::new(Int64Vector::try_from_arrow_array(array)?),
196            ArrowDataType::UInt8 => Arc::new(UInt8Vector::try_from_arrow_array(array)?),
197            ArrowDataType::UInt16 => Arc::new(UInt16Vector::try_from_arrow_array(array)?),
198            ArrowDataType::UInt32 => Arc::new(UInt32Vector::try_from_arrow_array(array)?),
199            ArrowDataType::UInt64 => Arc::new(UInt64Vector::try_from_arrow_array(array)?),
200            ArrowDataType::Float32 => Arc::new(Float32Vector::try_from_arrow_array(array)?),
201            ArrowDataType::Float64 => Arc::new(Float64Vector::try_from_arrow_array(array)?),
202            ArrowDataType::Utf8 | ArrowDataType::LargeUtf8 | ArrowDataType::Utf8View => {
203                Arc::new(StringVector::try_from_arrow_array(array)?)
204            }
205            ArrowDataType::Date32 => Arc::new(DateVector::try_from_arrow_array(array)?),
206            ArrowDataType::List(_) => Arc::new(ListVector::try_from_arrow_array(array)?),
207            ArrowDataType::Timestamp(unit, _) => match unit {
208                TimeUnit::Second => Arc::new(TimestampSecondVector::try_from_arrow_array(array)?),
209                TimeUnit::Millisecond => {
210                    Arc::new(TimestampMillisecondVector::try_from_arrow_array(array)?)
211                }
212                TimeUnit::Microsecond => {
213                    Arc::new(TimestampMicrosecondVector::try_from_arrow_array(array)?)
214                }
215                TimeUnit::Nanosecond => {
216                    Arc::new(TimestampNanosecondVector::try_from_arrow_array(array)?)
217                }
218            },
219            ArrowDataType::Time32(unit) => match unit {
220                TimeUnit::Second => Arc::new(TimeSecondVector::try_from_arrow_array(array)?),
221                TimeUnit::Millisecond => {
222                    Arc::new(TimeMillisecondVector::try_from_arrow_array(array)?)
223                }
224                // Arrow use time32 for second/millisecond.
225                _ => unreachable!(
226                    "unexpected arrow array datatype: {:?}",
227                    array.as_ref().data_type()
228                ),
229            },
230            ArrowDataType::Time64(unit) => match unit {
231                TimeUnit::Microsecond => {
232                    Arc::new(TimeMicrosecondVector::try_from_arrow_array(array)?)
233                }
234                TimeUnit::Nanosecond => {
235                    Arc::new(TimeNanosecondVector::try_from_arrow_array(array)?)
236                }
237                // Arrow use time64 for microsecond/nanosecond.
238                _ => unreachable!(
239                    "unexpected arrow array datatype: {:?}",
240                    array.as_ref().data_type()
241                ),
242            },
243            ArrowDataType::Interval(unit) => match unit {
244                IntervalUnit::YearMonth => {
245                    Arc::new(IntervalYearMonthVector::try_from_arrow_array(array)?)
246                }
247                IntervalUnit::DayTime => {
248                    Arc::new(IntervalDayTimeVector::try_from_arrow_array(array)?)
249                }
250                IntervalUnit::MonthDayNano => {
251                    Arc::new(IntervalMonthDayNanoVector::try_from_arrow_array(array)?)
252                }
253            },
254            ArrowDataType::Duration(unit) => match unit {
255                TimeUnit::Second => Arc::new(DurationSecondVector::try_from_arrow_array(array)?),
256                TimeUnit::Millisecond => {
257                    Arc::new(DurationMillisecondVector::try_from_arrow_array(array)?)
258                }
259                TimeUnit::Microsecond => {
260                    Arc::new(DurationMicrosecondVector::try_from_arrow_array(array)?)
261                }
262                TimeUnit::Nanosecond => {
263                    Arc::new(DurationNanosecondVector::try_from_arrow_array(array)?)
264                }
265            },
266            ArrowDataType::Decimal128(_, _) => {
267                Arc::new(Decimal128Vector::try_from_arrow_array(array)?)
268            }
269            ArrowDataType::Dictionary(key, value) => {
270                macro_rules! handle_dictionary_key_type {
271                    ($key_type:ident) => {{
272                        let array = array
273                            .as_ref()
274                            .as_any()
275                            .downcast_ref::<DictionaryArray<$key_type>>()
276                            .unwrap(); // Safety: the type is guarded by match arm condition
277                        Arc::new(DictionaryVector::new(
278                            array.clone(),
279                            ConcreteDataType::try_from(value.as_ref())?,
280                        )?)
281                    }};
282                }
283
284                match key.as_ref() {
285                    ArrowDataType::Int8 => handle_dictionary_key_type!(Int8Type),
286                    ArrowDataType::Int16 => handle_dictionary_key_type!(Int16Type),
287                    ArrowDataType::Int32 => handle_dictionary_key_type!(Int32Type),
288                    ArrowDataType::Int64 => handle_dictionary_key_type!(Int64Type),
289                    ArrowDataType::UInt8 => handle_dictionary_key_type!(UInt8Type),
290                    ArrowDataType::UInt16 => handle_dictionary_key_type!(UInt16Type),
291                    ArrowDataType::UInt32 => handle_dictionary_key_type!(UInt32Type),
292                    ArrowDataType::UInt64 => handle_dictionary_key_type!(UInt64Type),
293                    _ => {
294                        return error::UnsupportedArrowTypeSnafu {
295                            arrow_type: array.as_ref().data_type().clone(),
296                        }
297                        .fail();
298                    }
299                }
300            }
301
302            ArrowDataType::Struct(fields) => {
303                let array = array
304                    .as_ref()
305                    .as_any()
306                    .downcast_ref::<StructArray>()
307                    .unwrap();
308                Arc::new(StructVector::try_new(
309                    StructType::from(fields),
310                    array.clone(),
311                )?)
312            }
313            ArrowDataType::Float16
314            | ArrowDataType::LargeList(_)
315            | ArrowDataType::FixedSizeList(_, _)
316            | ArrowDataType::Union(_, _)
317            | ArrowDataType::Decimal256(_, _)
318            | ArrowDataType::Map(_, _)
319            | ArrowDataType::RunEndEncoded(_, _)
320            | ArrowDataType::ListView(_)
321            | ArrowDataType::LargeListView(_)
322            | ArrowDataType::Date64
323            | ArrowDataType::Decimal32(_, _)
324            | ArrowDataType::Decimal64(_, _) => {
325                return error::UnsupportedArrowTypeSnafu {
326                    arrow_type: array.as_ref().data_type().clone(),
327                }
328                .fail();
329            }
330        })
331    }
332
333    /// Try to cast an vec of values into vector, fail if type is not the same across all values.
334    pub fn try_from_row_into_vector(row: &[Value], dt: &ConcreteDataType) -> Result<VectorRef> {
335        let mut builder = dt.create_mutable_vector(row.len());
336        for val in row {
337            builder.try_push_value_ref(&val.as_value_ref())?;
338        }
339        let vector = builder.to_vector();
340        Ok(vector)
341    }
342
343    /// Try to cast slice of `arrays` to vectors.
344    pub fn try_into_vectors(arrays: &[ArrayRef]) -> Result<Vec<VectorRef>> {
345        arrays.iter().map(Self::try_into_vector).collect()
346    }
347
348    /// Perform SQL like operation on `names` and a scalar `s`.
349    pub fn like_utf8(names: Vec<String>, s: &str) -> Result<VectorRef> {
350        let array = StringArray::from(names);
351
352        let s = StringArray::new_scalar(s);
353        let filter = comparison::like(&array, &s).context(error::ArrowComputeSnafu)?;
354
355        let result = compute::filter(&array, &filter).context(error::ArrowComputeSnafu)?;
356        Helper::try_into_vector(result)
357    }
358
359    pub fn like_utf8_filter(names: Vec<String>, s: &str) -> Result<(VectorRef, BooleanVector)> {
360        let array = StringArray::from(names);
361        let s = StringArray::new_scalar(s);
362        let filter = comparison::like(&array, &s).context(error::ArrowComputeSnafu)?;
363        let result = compute::filter(&array, &filter).context(error::ArrowComputeSnafu)?;
364        let vector = Helper::try_into_vector(result)?;
365
366        Ok((vector, BooleanVector::from(filter)))
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use arrow::array::{
373        ArrayRef, BooleanArray, Date32Array, Float32Array, Float64Array, Int8Array, Int16Array,
374        Int32Array, Int64Array, LargeBinaryArray, ListArray, NullArray, Time32MillisecondArray,
375        Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
376        TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
377        TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
378    };
379    use arrow::buffer::Buffer;
380    use arrow::datatypes::{Int32Type, IntervalMonthDayNano};
381    use arrow_array::{
382        BinaryArray, DictionaryArray, FixedSizeBinaryArray, LargeStringArray, StructArray,
383    };
384    use arrow_schema::{DataType, Field, Fields};
385    use common_decimal::Decimal128;
386    use common_time::time::Time;
387    use common_time::timestamp::TimeUnit;
388    use common_time::{Date, Duration};
389    use datafusion_common::scalar::ScalarStructBuilder;
390
391    use super::*;
392    use crate::value::Value;
393    use crate::vectors::ConcreteDataType;
394
395    #[test]
396    fn test_try_into_vectors() {
397        let arrays: Vec<ArrayRef> = vec![
398            Arc::new(Int32Array::from(vec![1])),
399            Arc::new(Int32Array::from(vec![2])),
400            Arc::new(Int32Array::from(vec![3])),
401        ];
402        let vectors = Helper::try_into_vectors(&arrays).unwrap();
403        vectors.iter().for_each(|v| assert_eq!(1, v.len()));
404        assert_eq!(Value::Int32(1), vectors[0].get(0));
405        assert_eq!(Value::Int32(2), vectors[1].get(0));
406        assert_eq!(Value::Int32(3), vectors[2].get(0));
407    }
408
409    #[test]
410    fn test_try_into_date_vector() {
411        let vector = DateVector::from(vec![Some(1), Some(2), None]);
412        let arrow_array = vector.to_arrow_array();
413        assert_eq!(&ArrowDataType::Date32, arrow_array.data_type());
414        let vector_converted = Helper::try_into_vector(arrow_array).unwrap();
415        assert_eq!(vector.len(), vector_converted.len());
416        for i in 0..vector_converted.len() {
417            assert_eq!(vector.get(i), vector_converted.get(i));
418        }
419    }
420
421    #[test]
422    fn test_try_from_scalar_date_value() {
423        let vector = Helper::try_from_scalar_value(ScalarValue::Date32(Some(42)), 3, None).unwrap();
424        assert_eq!(ConcreteDataType::date_datatype(), vector.data_type());
425        assert_eq!(3, vector.len());
426        for i in 0..vector.len() {
427            assert_eq!(Value::Date(Date::new(42)), vector.get(i));
428        }
429    }
430
431    #[test]
432    fn test_try_from_scalar_duration_value() {
433        let vector =
434            Helper::try_from_scalar_value(ScalarValue::DurationSecond(Some(42)), 3, None).unwrap();
435        assert_eq!(
436            ConcreteDataType::duration_second_datatype(),
437            vector.data_type()
438        );
439        assert_eq!(3, vector.len());
440        for i in 0..vector.len() {
441            assert_eq!(
442                Value::Duration(Duration::new(42, TimeUnit::Second)),
443                vector.get(i)
444            );
445        }
446    }
447
448    #[test]
449    fn test_try_from_scalar_decimal128_value() {
450        let vector =
451            Helper::try_from_scalar_value(ScalarValue::Decimal128(Some(42), 3, 1), 3, None)
452                .unwrap();
453        assert_eq!(
454            ConcreteDataType::decimal128_datatype(3, 1),
455            vector.data_type()
456        );
457        assert_eq!(3, vector.len());
458        for i in 0..vector.len() {
459            assert_eq!(Value::Decimal128(Decimal128::new(42, 3, 1)), vector.get(i));
460        }
461    }
462
463    #[test]
464    fn test_try_from_list_value() {
465        let value = ScalarValue::List(ScalarValue::new_list(
466            &[ScalarValue::Int32(Some(1)), ScalarValue::Int32(Some(2))],
467            &ArrowDataType::Int32,
468            true,
469        ));
470        let vector = Helper::try_from_scalar_value(value, 3, None).unwrap();
471        assert_eq!(
472            ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())),
473            vector.data_type()
474        );
475        assert_eq!(3, vector.len());
476        for i in 0..vector.len() {
477            let v = vector.get(i);
478            let items = v.as_list().unwrap().unwrap().items();
479            assert_eq!(vec![Value::Int32(1), Value::Int32(2)], items);
480        }
481    }
482
483    #[test]
484    fn test_try_from_scalar_value_materializes_values() {
485        let vector = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 4, None).unwrap();
486        assert_eq!(ConcreteDataType::int32_datatype(), vector.data_type());
487        assert_eq!(4, vector.len());
488        assert_eq!(0, vector.null_count());
489        for i in 0..vector.len() {
490            assert_eq!(Value::Int32(42), vector.get(i));
491        }
492
493        let empty = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 0, None).unwrap();
494        assert_eq!(ConcreteDataType::int32_datatype(), empty.data_type());
495        assert!(empty.is_empty());
496
497        let nulls = Helper::try_from_scalar_value(ScalarValue::Int32(None), 3, None).unwrap();
498        assert_eq!(3, nulls.len());
499        assert_eq!(3, nulls.null_count());
500        for i in 0..nulls.len() {
501            assert_eq!(Value::Null, nulls.get(i));
502        }
503    }
504
505    #[test]
506    fn test_try_from_scalar_struct_value() {
507        let fields = Fields::from(vec![
508            Field::new("id", ArrowDataType::Int32, false),
509            Field::new("name", ArrowDataType::Utf8, true),
510        ]);
511        let value = ScalarValue::Struct(Arc::new(StructArray::new(
512            fields.clone(),
513            vec![
514                ScalarValue::Int32(Some(7)).to_array().unwrap(),
515                ScalarValue::Utf8(Some("greptime".to_string()))
516                    .to_array()
517                    .unwrap(),
518            ],
519            None,
520        )));
521
522        let vector = Helper::try_from_scalar_value(value, 3, None).unwrap();
523        assert_eq!(
524            ConcreteDataType::struct_datatype(StructType::from(&fields)),
525            vector.data_type()
526        );
527        assert_eq!(3, vector.len());
528        for i in 0..vector.len() {
529            let Value::Struct(value) = vector.get(i) else {
530                panic!("expected struct value");
531            };
532            assert_eq!(
533                &[Value::Int32(7), Value::String("greptime".into())],
534                value.items()
535            );
536        }
537
538        let null = ScalarStructBuilder::new_null(fields);
539        let vector = Helper::try_from_scalar_value(null, 2, None).unwrap();
540        assert_eq!(2, vector.len());
541        assert_eq!(2, vector.null_count());
542        assert_eq!(Value::Null, vector.get(0));
543        assert_eq!(Value::Null, vector.get(1));
544    }
545
546    #[test]
547    fn test_try_from_scalar_value_normalizes_arrow_types() {
548        let string = Helper::try_from_scalar_value(
549            ScalarValue::LargeUtf8(Some("greptime".to_string())),
550            2,
551            None,
552        )
553        .unwrap();
554        assert_eq!(ConcreteDataType::string_datatype(), string.data_type());
555        assert_eq!(&ArrowDataType::Utf8, string.to_arrow_array().data_type());
556
557        let binary = Helper::try_from_scalar_value(
558            ScalarValue::FixedSizeBinary(2, Some(vec![1, 2])),
559            2,
560            None,
561        )
562        .unwrap();
563        assert_eq!(ConcreteDataType::binary_datatype(), binary.data_type());
564        assert_eq!(&ArrowDataType::Binary, binary.to_arrow_array().data_type());
565
566        let timestamp = Helper::try_from_scalar_value(
567            ScalarValue::TimestampMillisecond(Some(42), Some("UTC".into())),
568            2,
569            None,
570        )
571        .unwrap();
572        assert_eq!(
573            &ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
574            timestamp.to_arrow_array().data_type()
575        );
576    }
577
578    #[test]
579    fn test_like_utf8() {
580        fn assert_vector(expected: Vec<&str>, actual: &VectorRef) {
581            let actual = actual.as_any().downcast_ref::<StringVector>().unwrap();
582            assert_eq!(*actual, StringVector::from(expected));
583        }
584
585        let names: Vec<String> = vec!["greptime", "hello", "public", "world"]
586            .into_iter()
587            .map(|x| x.to_string())
588            .collect();
589
590        let ret = Helper::like_utf8(names.clone(), "%ll%").unwrap();
591        assert_vector(vec!["hello"], &ret);
592
593        let ret = Helper::like_utf8(names.clone(), "%time").unwrap();
594        assert_vector(vec!["greptime"], &ret);
595
596        let ret = Helper::like_utf8(names.clone(), "%ld").unwrap();
597        assert_vector(vec!["world"], &ret);
598
599        let ret = Helper::like_utf8(names, "%").unwrap();
600        assert_vector(vec!["greptime", "hello", "public", "world"], &ret);
601    }
602
603    #[test]
604    fn test_like_utf8_filter() {
605        fn assert_vector(expected: Vec<&str>, actual: &VectorRef) {
606            let actual = actual.as_any().downcast_ref::<StringVector>().unwrap();
607            assert_eq!(*actual, StringVector::from(expected));
608        }
609
610        fn assert_filter(array: Vec<String>, s: &str, expected_filter: &BooleanVector) {
611            let array = StringArray::from(array);
612            let s = StringArray::new_scalar(s);
613            let actual_filter = comparison::like(&array, &s).unwrap();
614            assert_eq!(BooleanVector::from(actual_filter), *expected_filter);
615        }
616
617        let names: Vec<String> = vec!["greptime", "timeseries", "cloud", "database"]
618            .into_iter()
619            .map(|x| x.to_string())
620            .collect();
621
622        let (table, filter) = Helper::like_utf8_filter(names.clone(), "%ti%").unwrap();
623        assert_vector(vec!["greptime", "timeseries"], &table);
624        assert_filter(names.clone(), "%ti%", &filter);
625
626        let (tables, filter) = Helper::like_utf8_filter(names.clone(), "%lou").unwrap();
627        assert_vector(vec![], &tables);
628        assert_filter(names.clone(), "%lou", &filter);
629
630        let (tables, filter) = Helper::like_utf8_filter(names.clone(), "%d%").unwrap();
631        assert_vector(vec!["cloud", "database"], &tables);
632        assert_filter(names.clone(), "%d%", &filter);
633    }
634
635    fn check_try_into_vector(array: impl Array + 'static) {
636        let array: ArrayRef = Arc::new(array);
637        let vector = Helper::try_into_vector(array.clone()).unwrap();
638        assert_eq!(&array, &vector.to_arrow_array());
639    }
640
641    #[test]
642    fn test_try_into_vector() {
643        check_try_into_vector(NullArray::new(2));
644        check_try_into_vector(BooleanArray::from(vec![true, false]));
645        check_try_into_vector(Int8Array::from(vec![1, 2, 3]));
646        check_try_into_vector(Int16Array::from(vec![1, 2, 3]));
647        check_try_into_vector(Int32Array::from(vec![1, 2, 3]));
648        check_try_into_vector(Int64Array::from(vec![1, 2, 3]));
649        check_try_into_vector(UInt8Array::from(vec![1, 2, 3]));
650        check_try_into_vector(UInt16Array::from(vec![1, 2, 3]));
651        check_try_into_vector(UInt32Array::from(vec![1, 2, 3]));
652        check_try_into_vector(UInt64Array::from(vec![1, 2, 3]));
653        check_try_into_vector(Float32Array::from(vec![1.0, 2.0, 3.0]));
654        check_try_into_vector(Float64Array::from(vec![1.0, 2.0, 3.0]));
655        check_try_into_vector(StringArray::from(vec!["hello", "world"]));
656        check_try_into_vector(Date32Array::from(vec![1, 2, 3]));
657        let data = vec![None, Some(vec![Some(6), Some(7)])];
658        let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
659        check_try_into_vector(list_array);
660        check_try_into_vector(TimestampSecondArray::from(vec![1, 2, 3]));
661        check_try_into_vector(TimestampMillisecondArray::from(vec![1, 2, 3]));
662        check_try_into_vector(TimestampMicrosecondArray::from(vec![1, 2, 3]));
663        check_try_into_vector(TimestampNanosecondArray::from(vec![1, 2, 3]));
664        check_try_into_vector(Time32SecondArray::from(vec![1, 2, 3]));
665        check_try_into_vector(Time32MillisecondArray::from(vec![1, 2, 3]));
666        check_try_into_vector(Time64MicrosecondArray::from(vec![1, 2, 3]));
667        check_try_into_vector(Time64NanosecondArray::from(vec![1, 2, 3]));
668
669        // Test dictionary arrays with different key types
670        let values = StringArray::from_iter_values(["a", "b", "c"]);
671
672        // Test Int8 keys
673        let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
674        let array: ArrayRef =
675            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
676        Helper::try_into_vector(array).unwrap();
677
678        // Test Int16 keys
679        let keys = Int16Array::from_iter_values([0, 0, 1, 2]);
680        let array: ArrayRef =
681            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
682        Helper::try_into_vector(array).unwrap();
683
684        // Test Int32 keys
685        let keys = Int32Array::from_iter_values([0, 0, 1, 2]);
686        let array: ArrayRef =
687            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
688        Helper::try_into_vector(array).unwrap();
689
690        // Test Int64 keys
691        let keys = Int64Array::from_iter_values([0, 0, 1, 2]);
692        let array: ArrayRef =
693            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
694        Helper::try_into_vector(array).unwrap();
695
696        // Test UInt8 keys
697        let keys = UInt8Array::from_iter_values([0, 0, 1, 2]);
698        let array: ArrayRef =
699            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
700        Helper::try_into_vector(array).unwrap();
701
702        // Test UInt16 keys
703        let keys = UInt16Array::from_iter_values([0, 0, 1, 2]);
704        let array: ArrayRef =
705            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
706        Helper::try_into_vector(array).unwrap();
707
708        // Test UInt32 keys
709        let keys = UInt32Array::from_iter_values([0, 0, 1, 2]);
710        let array: ArrayRef =
711            Arc::new(DictionaryArray::try_new(keys, Arc::new(values.clone())).unwrap());
712        Helper::try_into_vector(array).unwrap();
713
714        // Test UInt64 keys
715        let keys = UInt64Array::from_iter_values([0, 0, 1, 2]);
716        let array: ArrayRef = Arc::new(DictionaryArray::try_new(keys, Arc::new(values)).unwrap());
717        Helper::try_into_vector(array).unwrap();
718    }
719
720    #[test]
721    fn test_try_binary_array_into_vector() {
722        let input_vec: Vec<&[u8]> = vec!["hello".as_bytes(), "world".as_bytes()];
723        let assertion_vector = BinaryVector::from(input_vec.clone());
724
725        let input_arrays: Vec<ArrayRef> = vec![
726            Arc::new(LargeBinaryArray::from(input_vec.clone())) as ArrayRef,
727            Arc::new(BinaryArray::from(input_vec.clone())) as ArrayRef,
728            Arc::new(FixedSizeBinaryArray::new(
729                5,
730                Buffer::from_vec("helloworld".as_bytes().to_vec()),
731                None,
732            )) as ArrayRef,
733        ];
734
735        for input_array in input_arrays {
736            let vector = Helper::try_into_vector(input_array).unwrap();
737
738            assert_eq!(2, vector.len());
739            assert_eq!(0, vector.null_count());
740
741            let output_arrow_array: ArrayRef = vector.to_arrow_array();
742            assert_eq!(&DataType::Binary, output_arrow_array.data_type());
743            assert_eq!(&assertion_vector.to_arrow_array(), &output_arrow_array);
744        }
745    }
746
747    #[test]
748    fn test_large_string_array_into_vector() {
749        let input_vec = vec!["a", "b"];
750        let assertion_array = LargeStringArray::from(input_vec.clone());
751
752        let large_string_array: ArrayRef = Arc::new(LargeStringArray::from(input_vec));
753        let vector = Helper::try_into_vector(large_string_array).unwrap();
754        assert_eq!(2, vector.len());
755        assert_eq!(0, vector.null_count());
756
757        let output_arrow_array: LargeStringArray = vector
758            .to_arrow_array()
759            .as_any()
760            .downcast_ref::<LargeStringArray>()
761            .unwrap()
762            .clone();
763        assert_eq!(&assertion_array, &output_arrow_array);
764    }
765
766    #[test]
767    fn test_try_from_scalar_time_value() {
768        let vector =
769            Helper::try_from_scalar_value(ScalarValue::Time32Second(Some(42)), 3, None).unwrap();
770        assert_eq!(ConcreteDataType::time_second_datatype(), vector.data_type());
771        assert_eq!(3, vector.len());
772        for i in 0..vector.len() {
773            assert_eq!(Value::Time(Time::new_second(42)), vector.get(i));
774        }
775    }
776
777    #[test]
778    fn test_try_from_scalar_interval_value() {
779        let vector = Helper::try_from_scalar_value(
780            ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(1, 1, 2000))),
781            3,
782            None,
783        )
784        .unwrap();
785
786        assert_eq!(
787            ConcreteDataType::interval_month_day_nano_datatype(),
788            vector.data_type()
789        );
790        assert_eq!(3, vector.len());
791        for i in 0..vector.len() {
792            assert_eq!(
793                Value::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 1, 2000).into()),
794                vector.get(i)
795            );
796        }
797    }
798
799    fn check_try_from_row_to_vector(row: Vec<Value>, dt: &ConcreteDataType) {
800        let vector = Helper::try_from_row_into_vector(&row, dt).unwrap();
801        for (i, item) in row.iter().enumerate().take(vector.len()) {
802            assert_eq!(*item, vector.get(i));
803        }
804    }
805
806    fn check_into_and_from(array: impl Array + 'static) {
807        let array: ArrayRef = Arc::new(array);
808        let vector = Helper::try_into_vector(array.clone()).unwrap();
809        assert_eq!(&array, &vector.to_arrow_array());
810        let row: Vec<Value> = (0..array.len()).map(|i| vector.get(i)).collect();
811        let dt = vector.data_type();
812        check_try_from_row_to_vector(row, &dt);
813    }
814
815    #[test]
816    fn test_try_from_row_to_vector() {
817        check_into_and_from(NullArray::new(2));
818        check_into_and_from(BooleanArray::from(vec![true, false]));
819        check_into_and_from(Int8Array::from(vec![1, 2, 3]));
820        check_into_and_from(Int16Array::from(vec![1, 2, 3]));
821        check_into_and_from(Int32Array::from(vec![1, 2, 3]));
822        check_into_and_from(Int64Array::from(vec![1, 2, 3]));
823        check_into_and_from(UInt8Array::from(vec![1, 2, 3]));
824        check_into_and_from(UInt16Array::from(vec![1, 2, 3]));
825        check_into_and_from(UInt32Array::from(vec![1, 2, 3]));
826        check_into_and_from(UInt64Array::from(vec![1, 2, 3]));
827        check_into_and_from(Float32Array::from(vec![1.0, 2.0, 3.0]));
828        check_into_and_from(Float64Array::from(vec![1.0, 2.0, 3.0]));
829        check_into_and_from(StringArray::from(vec!["hello", "world"]));
830        check_into_and_from(Date32Array::from(vec![1, 2, 3]));
831
832        check_into_and_from(TimestampSecondArray::from(vec![1, 2, 3]));
833        check_into_and_from(TimestampMillisecondArray::from(vec![1, 2, 3]));
834        check_into_and_from(TimestampMicrosecondArray::from(vec![1, 2, 3]));
835        check_into_and_from(TimestampNanosecondArray::from(vec![1, 2, 3]));
836        check_into_and_from(Time32SecondArray::from(vec![1, 2, 3]));
837        check_into_and_from(Time32MillisecondArray::from(vec![1, 2, 3]));
838        check_into_and_from(Time64MicrosecondArray::from(vec![1, 2, 3]));
839        check_into_and_from(Time64NanosecondArray::from(vec![1, 2, 3]));
840    }
841}
datatypes/vectors/helper.rs

datatypes/vectors/
helper.rs
