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

datatypes/vectors/
helper.rs
