risingwave_common/array/
mod.rs

1// Copyright 2022 RisingWave Labs
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//! `Array` defines all in-memory representations of vectorized execution framework.
16
17pub mod arrow;
18mod bool_array;
19pub mod bytes_array;
20mod chrono_array;
21mod data_chunk;
22pub mod data_chunk_iter;
23mod decimal_array;
24pub mod error;
25pub mod interval_array;
26mod iterator;
27mod jsonb_array;
28pub mod list_array;
29mod map_array;
30mod num256_array;
31mod primitive_array;
32mod proto_reader;
33pub mod stream_chunk;
34pub mod stream_chunk_builder;
35mod stream_chunk_iter;
36pub mod stream_record;
37pub mod struct_array;
38mod utf8_array;
39mod variant_array;
40mod vector_array;
41
42use std::convert::From;
43use std::hash::{Hash, Hasher};
44use std::sync::Arc;
45
46pub use bool_array::{BoolArray, BoolArrayBuilder};
47pub use bytes_array::*;
48pub use chrono_array::{
49    DateArray, DateArrayBuilder, TimeArray, TimeArrayBuilder, TimestampArray,
50    TimestampArrayBuilder, TimestamptzArray, TimestamptzArrayBuilder,
51};
52pub use data_chunk::{DataChunk, DataChunkTestExt};
53pub use data_chunk_iter::RowRef;
54pub use decimal_array::{DecimalArray, DecimalArrayBuilder};
55pub use interval_array::{IntervalArray, IntervalArrayBuilder};
56pub use iterator::ArrayIterator;
57pub use jsonb_array::{JsonbArray, JsonbArrayBuilder};
58pub use list_array::{ListArray, ListArrayBuilder, ListRef, ListValue, ListWrite, ListWriter};
59pub use map_array::{MapArray, MapArrayBuilder, MapRef, MapValue};
60use paste::paste;
61pub use primitive_array::{PrimitiveArray, PrimitiveArrayBuilder, PrimitiveArrayItemType};
62use risingwave_common_estimate_size::EstimateSize;
63use risingwave_pb::data::PbArray;
64pub use stream_chunk::{Op, StreamChunk, StreamChunkTestExt};
65pub use stream_chunk_builder::StreamChunkBuilder;
66pub use struct_array::{StructArray, StructArrayBuilder, StructRef, StructValue};
67pub use utf8_array::*;
68pub use variant_array::{VariantArray, VariantArrayBuilder};
69pub use vector_array::{
70    Finite32, VECTOR_AS_LIST_TYPE, VECTOR_DISTANCE_TYPE, VECTOR_ITEM_TYPE, VectorArray,
71    VectorArrayBuilder, VectorDistanceType, VectorItemType, VectorRef, VectorVal,
72};
73
74pub use self::error::ArrayError;
75pub use crate::array::num256_array::{Int256Array, Int256ArrayBuilder};
76use crate::bitmap::Bitmap;
77use crate::types::*;
78use crate::{dispatch_array_builder_variants, dispatch_array_variants, for_all_variants};
79pub type ArrayResult<T> = Result<T, ArrayError>;
80
81pub type I64Array = PrimitiveArray<i64>;
82pub type I32Array = PrimitiveArray<i32>;
83pub type I16Array = PrimitiveArray<i16>;
84pub type F64Array = PrimitiveArray<F64>;
85pub type F32Array = PrimitiveArray<F32>;
86pub type SerialArray = PrimitiveArray<Serial>;
87
88pub type I64ArrayBuilder = PrimitiveArrayBuilder<i64>;
89pub type I32ArrayBuilder = PrimitiveArrayBuilder<i32>;
90pub type I16ArrayBuilder = PrimitiveArrayBuilder<i16>;
91pub type F64ArrayBuilder = PrimitiveArrayBuilder<F64>;
92pub type F32ArrayBuilder = PrimitiveArrayBuilder<F32>;
93pub type SerialArrayBuilder = PrimitiveArrayBuilder<Serial>;
94
95// alias for expr macros
96pub type ArrayImplBuilder = ArrayBuilderImpl;
97
98/// The hash source for `None` values when hashing an item.
99pub(crate) const NULL_VAL_FOR_HASH: u32 = 0xfffffff0;
100
101/// A trait over all array builders.
102///
103/// `ArrayBuilder` is a trait over all builders. You could build an array with
104/// `append` with the help of `ArrayBuilder` trait. The `append` function always
105/// accepts reference to an element if it is not primitive. e.g. for `PrimitiveArray`,
106/// you could do `builder.append(Some(1))`. For `Utf8Array`, you must do
107/// `builder.append(Some("xxx"))`. Note that you don't need to construct a `String`.
108///
109/// The associated type `ArrayType` is the type of the corresponding array. It is the
110/// return type of `finish`.
111pub trait ArrayBuilder: Send + Sync + Sized + 'static {
112    /// Corresponding `Array` of this builder, which is reciprocal to `ArrayBuilder`.
113    type ArrayType: Array<Builder = Self>;
114
115    /// Create a new builder with `capacity`.
116    /// TODO: remove this function from the trait. Let it be methods of each concrete builders.
117    fn new(capacity: usize) -> Self;
118
119    /// # Panics
120    /// Panics if `meta`'s type mismatches with the array type.
121    fn with_type(capacity: usize, ty: DataType) -> Self;
122
123    /// Append a value multiple times.
124    ///
125    /// This should be more efficient than calling `append` multiple times.
126    fn append_n(&mut self, n: usize, value: Option<<Self::ArrayType as Array>::RefItem<'_>>);
127
128    /// Append a value to builder.
129    fn append(&mut self, value: Option<<Self::ArrayType as Array>::RefItem<'_>>) {
130        self.append_n(1, value);
131    }
132
133    /// Append an owned value to builder.
134    fn append_owned(&mut self, value: Option<<Self::ArrayType as Array>::OwnedItem>) {
135        let value = value.as_ref().map(|s| s.as_scalar_ref());
136        self.append(value)
137    }
138
139    fn append_null(&mut self) {
140        self.append(None)
141    }
142
143    /// Append an array to builder.
144    fn append_array(&mut self, other: &Self::ArrayType);
145
146    /// Pop an element from the builder.
147    ///
148    /// It's used in `rollback` in source parser.
149    ///
150    /// # Returns
151    ///
152    /// Returns `None` if there is no elements in the builder.
153    fn pop(&mut self) -> Option<()>;
154
155    /// Append an element in another array into builder.
156    fn append_array_element(&mut self, other: &Self::ArrayType, idx: usize) {
157        self.append(other.value_at(idx));
158    }
159
160    /// Return the number of elements in the builder.
161    fn len(&self) -> usize;
162
163    /// Return `true` if the array has a length of 0.
164    fn is_empty(&self) -> bool {
165        self.len() == 0
166    }
167
168    /// Finish build and return a new array.
169    fn finish(self) -> Self::ArrayType;
170}
171
172/// A trait over all array.
173///
174/// `Array` must be built with an `ArrayBuilder`. The array trait provides several
175/// unified interface on an array, like `len`, `value_at` and `iter`.
176///
177/// The `Builder` associated type is the builder for this array.
178///
179/// The `Iter` associated type is the iterator of this array. And the `RefItem` is
180/// the item you could retrieve from this array.
181/// For example, `PrimitiveArray` could return an `Option<u32>`, and `Utf8Array` will
182/// return an `Option<&str>`.
183///
184/// In some cases, we will need to store owned data. For example, when aggregating min
185/// and max, we need to store current maximum in the aggregator. In this case, we
186/// could use `A::OwnedItem` in aggregator struct.
187pub trait Array:
188    std::fmt::Debug + Send + Sync + Sized + 'static + Into<ArrayImpl> + EstimateSize
189{
190    /// A reference to item in array, as well as return type of `value_at`, which is
191    /// reciprocal to `Self::OwnedItem`.
192    type RefItem<'a>: ScalarRef<'a, ScalarType = Self::OwnedItem>
193    where
194        Self: 'a;
195
196    /// Owned type of item in array, which is reciprocal to `Self::RefItem`.
197    type OwnedItem: Clone
198        + std::fmt::Debug
199        + EstimateSize
200        + for<'a> Scalar<ScalarRefType<'a> = Self::RefItem<'a>>;
201
202    /// Corresponding builder of this array, which is reciprocal to `Array`.
203    type Builder: ArrayBuilder<ArrayType = Self>;
204
205    /// Retrieve a reference to value regardless of whether it is null
206    /// without checking the index boundary.
207    ///
208    /// The returned value for NULL values is the default value.
209    ///
210    /// # Safety
211    ///
212    /// Index must be within the bounds.
213    unsafe fn raw_value_at_unchecked(&self, idx: usize) -> Self::RefItem<'_>;
214
215    /// Retrieve a reference to value.
216    #[inline]
217    fn value_at(&self, idx: usize) -> Option<Self::RefItem<'_>> {
218        if !self.is_null(idx) {
219            // Safety: the above `is_null` check ensures that the index is valid.
220            Some(unsafe { self.raw_value_at_unchecked(idx) })
221        } else {
222            None
223        }
224    }
225
226    /// # Safety
227    ///
228    /// Retrieve a reference to value without checking the index boundary.
229    #[inline]
230    unsafe fn value_at_unchecked(&self, idx: usize) -> Option<Self::RefItem<'_>> {
231        unsafe {
232            if !self.is_null_unchecked(idx) {
233                Some(self.raw_value_at_unchecked(idx))
234            } else {
235                None
236            }
237        }
238    }
239
240    /// Number of items of array.
241    fn len(&self) -> usize;
242
243    /// Get iterator of current array.
244    fn iter(&self) -> ArrayIterator<'_, Self> {
245        ArrayIterator::new(self)
246    }
247
248    /// Get raw iterator of current array.
249    ///
250    /// The raw iterator simply iterates values without checking the null bitmap.
251    /// The returned value for NULL values is undefined.
252    fn raw_iter(&self) -> impl ExactSizeIterator<Item = Self::RefItem<'_>> {
253        (0..self.len()).map(|i| unsafe { self.raw_value_at_unchecked(i) })
254    }
255
256    /// Serialize to protobuf
257    fn to_protobuf(&self) -> PbArray;
258
259    /// Get the null `Bitmap` from `Array`.
260    fn null_bitmap(&self) -> &Bitmap;
261
262    /// Get the owned null `Bitmap` from `Array`.
263    fn into_null_bitmap(self) -> Bitmap;
264
265    /// Check if an element is `null` or not.
266    fn is_null(&self, idx: usize) -> bool {
267        !self.null_bitmap().is_set(idx)
268    }
269
270    /// # Safety
271    ///
272    /// The unchecked version of `is_null`, ignore index out of bound check. It is
273    /// the caller's responsibility to ensure the index is valid.
274    unsafe fn is_null_unchecked(&self, idx: usize) -> bool {
275        unsafe { !self.null_bitmap().is_set_unchecked(idx) }
276    }
277
278    fn set_bitmap(&mut self, bitmap: Bitmap);
279
280    /// Feed the value at `idx` into the given [`Hasher`].
281    #[inline(always)]
282    fn hash_at<H: Hasher>(&self, idx: usize, state: &mut H) {
283        // We use a default implementation for all arrays for now, as retrieving the reference
284        // should be lightweight.
285        if let Some(value) = self.value_at(idx) {
286            value.hash_scalar(state);
287        } else {
288            NULL_VAL_FOR_HASH.hash(state);
289        }
290    }
291
292    fn hash_vec<H: Hasher>(&self, hashers: &mut [H], vis: &Bitmap) {
293        assert_eq!(hashers.len(), self.len());
294        for idx in vis.iter_ones() {
295            self.hash_at(idx, &mut hashers[idx]);
296        }
297    }
298
299    fn is_empty(&self) -> bool {
300        self.len() == 0
301    }
302
303    fn create_builder(&self, capacity: usize) -> Self::Builder {
304        Self::Builder::with_type(capacity, self.data_type())
305    }
306
307    fn data_type(&self) -> DataType;
308
309    /// Converts the array into an [`ArrayRef`].
310    fn into_ref(self) -> ArrayRef {
311        Arc::new(self.into())
312    }
313}
314
315/// Implement `compact_vis` on array, which removes element according to `visibility`.
316#[easy_ext::ext(ArrayCompactVisExt)]
317impl<A: Array> A {
318    /// Select some elements from `Array` based on `visibility` bitmap.
319    /// `cardinality` is only used to decide capacity of the new `Array`.
320    pub fn compact_vis(&self, visibility: &Bitmap, cardinality: usize) -> Self {
321        let mut builder = A::Builder::with_type(cardinality, self.data_type());
322        for idx in visibility.iter_ones() {
323            // SAFETY(value_at_unchecked): the idx is always in bound.
324            unsafe {
325                builder.append(self.value_at_unchecked(idx));
326            }
327        }
328        builder.finish()
329    }
330}
331
332/// Define `ArrayImpl` with macro.
333macro_rules! array_impl_enum {
334    ( $( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
335        /// `ArrayImpl` embeds all possible array in `array` module.
336        #[derive(Debug, Clone, EstimateSize)]
337        pub enum ArrayImpl {
338            $( $variant_name($array) ),*
339        }
340    };
341}
342
343for_all_variants! { array_impl_enum }
344
345// We cannot put the From implementations in impl_convert,
346// because then we can't prove for all `T: PrimitiveArrayItemType`,
347// it's implemented.
348
349impl<T: PrimitiveArrayItemType> From<PrimitiveArray<T>> for ArrayImpl {
350    fn from(arr: PrimitiveArray<T>) -> Self {
351        T::erase_array_type(arr)
352    }
353}
354
355impl From<Int256Array> for ArrayImpl {
356    fn from(arr: Int256Array) -> Self {
357        Self::Int256(arr)
358    }
359}
360
361impl From<BoolArray> for ArrayImpl {
362    fn from(arr: BoolArray) -> Self {
363        Self::Bool(arr)
364    }
365}
366
367impl From<Utf8Array> for ArrayImpl {
368    fn from(arr: Utf8Array) -> Self {
369        Self::Utf8(arr)
370    }
371}
372
373impl From<JsonbArray> for ArrayImpl {
374    fn from(arr: JsonbArray) -> Self {
375        Self::Jsonb(arr)
376    }
377}
378
379impl From<VariantArray> for ArrayImpl {
380    fn from(arr: VariantArray) -> Self {
381        Self::Variant(arr)
382    }
383}
384
385impl From<StructArray> for ArrayImpl {
386    fn from(arr: StructArray) -> Self {
387        Self::Struct(arr)
388    }
389}
390
391impl From<ListArray> for ArrayImpl {
392    fn from(arr: ListArray) -> Self {
393        Self::List(arr)
394    }
395}
396
397impl From<VectorArray> for ArrayImpl {
398    fn from(arr: VectorArray) -> Self {
399        Self::Vector(arr)
400    }
401}
402
403impl From<BytesArray> for ArrayImpl {
404    fn from(arr: BytesArray) -> Self {
405        Self::Bytea(arr)
406    }
407}
408
409impl From<MapArray> for ArrayImpl {
410    fn from(arr: MapArray) -> Self {
411        Self::Map(arr)
412    }
413}
414
415/// `impl_convert` implements several conversions for `Array` and `ArrayBuilder`.
416/// * `ArrayImpl -> &Array` with `impl.as_int16()`.
417/// * `ArrayImpl -> Array` with `impl.into_int16()`.
418/// * `&ArrayImpl -> &Array` with `From` trait.
419/// * `ArrayImpl -> Array` with `From` trait.
420/// * `ArrayBuilder -> ArrayBuilderImpl` with `From` trait.
421macro_rules! impl_convert {
422    ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
423        $(
424            paste! {
425                impl ArrayImpl {
426                    /// # Panics
427                    ///
428                    /// Panics if type mismatches.
429                    pub fn [<as_ $suffix_name>](&self) -> &$array {
430                        match self {
431                            Self::$variant_name(array) => array,
432                            other_array => panic!("cannot convert ArrayImpl::{} to concrete type {}", other_array.get_ident(), stringify!($variant_name))
433                        }
434                    }
435
436                    /// # Panics
437                    ///
438                    /// Panics if type mismatches.
439                    pub fn [<into_ $suffix_name>](self) -> $array {
440                        match self {
441                            Self::$variant_name(array) => array,
442                            other_array => panic!("cannot convert ArrayImpl::{} to concrete type {}", other_array.get_ident(), stringify!($variant_name))
443                        }
444                    }
445                }
446
447                // FIXME: panic in From here is not proper.
448                impl <'a> From<&'a ArrayImpl> for &'a $array {
449                    fn from(array: &'a ArrayImpl) -> Self {
450                        match array {
451                            ArrayImpl::$variant_name(inner) => inner,
452                            other_array => panic!("cannot convert ArrayImpl::{} to concrete type {}", other_array.get_ident(), stringify!($variant_name))
453                        }
454                    }
455                }
456
457                impl From<ArrayImpl> for $array {
458                    fn from(array: ArrayImpl) -> Self {
459                        match array {
460                            ArrayImpl::$variant_name(inner) => inner,
461                            other_array => panic!("cannot convert ArrayImpl::{} to concrete type {}", other_array.get_ident(), stringify!($variant_name))
462                        }
463                    }
464                }
465
466                impl From<$builder> for ArrayBuilderImpl {
467                    fn from(builder: $builder) -> Self {
468                        Self::$variant_name(builder)
469                    }
470                }
471            }
472        )*
473    };
474}
475
476for_all_variants! { impl_convert }
477
478/// Define `ArrayImplBuilder` with macro.
479macro_rules! array_builder_impl_enum {
480    ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
481        /// `ArrayBuilderImpl` embeds all possible array in `array` module.
482        #[derive(Debug, Clone, EstimateSize)]
483        pub enum ArrayBuilderImpl {
484            $( $variant_name($builder) ),*
485        }
486    };
487}
488
489for_all_variants! { array_builder_impl_enum }
490
491/// Implements all `ArrayBuilder` functions with `for_all_variant`.
492impl ArrayBuilderImpl {
493    pub fn with_type(capacity: usize, ty: DataType) -> Self {
494        ty.create_array_builder(capacity)
495    }
496
497    pub fn append_array(&mut self, other: &ArrayImpl) {
498        dispatch_array_builder_variants!(self, inner, { inner.append_array(other.into()) })
499    }
500
501    pub fn append_null(&mut self) {
502        dispatch_array_builder_variants!(self, inner, { inner.append(None) })
503    }
504
505    pub fn append_n_null(&mut self, n: usize) {
506        dispatch_array_builder_variants!(self, inner, { inner.append_n(n, None) })
507    }
508
509    /// Append a [`Datum`] or [`DatumRef`] multiple times,
510    /// panicking if the datum's type does not match the array builder's type.
511    pub fn append_n(&mut self, n: usize, datum: impl ToDatumRef) {
512        match datum.to_datum_ref() {
513            None => dispatch_array_builder_variants!(self, inner, { inner.append_n(n, None) }),
514
515            Some(scalar_ref) => {
516                dispatch_array_builder_variants!(self, inner, [I = VARIANT_NAME], {
517                    inner.append_n(
518                        n,
519                        Some(scalar_ref.try_into().unwrap_or_else(|_| {
520                            panic!(
521                                "type mismatch, array builder type: {}, scalar type: {}",
522                                I,
523                                scalar_ref.get_ident()
524                            )
525                        })),
526                    )
527                })
528            }
529        }
530    }
531
532    /// Append a [`Datum`] or [`DatumRef`], return error while type not match.
533    pub fn append(&mut self, datum: impl ToDatumRef) {
534        self.append_n(1, datum);
535    }
536
537    pub fn append_array_element(&mut self, other: &ArrayImpl, idx: usize) {
538        dispatch_array_builder_variants!(self, inner, {
539            inner.append_array_element(other.into(), idx)
540        })
541    }
542
543    pub fn pop(&mut self) -> Option<()> {
544        dispatch_array_builder_variants!(self, inner, { inner.pop() })
545    }
546
547    pub fn finish(self) -> ArrayImpl {
548        dispatch_array_builder_variants!(self, inner, { inner.finish().into() })
549    }
550
551    pub fn get_ident(&self) -> &'static str {
552        dispatch_array_builder_variants!(self, [I = VARIANT_NAME], { I })
553    }
554
555    pub fn len(&self) -> usize {
556        dispatch_array_builder_variants!(self, inner, { inner.len() })
557    }
558
559    pub fn is_empty(&self) -> bool {
560        self.len() == 0
561    }
562}
563
564impl ArrayImpl {
565    /// Number of items in array.
566    pub fn len(&self) -> usize {
567        dispatch_array_variants!(self, inner, { inner.len() })
568    }
569
570    pub fn is_empty(&self) -> bool {
571        self.len() == 0
572    }
573
574    /// Get the null `Bitmap` of the array.
575    pub fn null_bitmap(&self) -> &Bitmap {
576        dispatch_array_variants!(self, inner, { inner.null_bitmap() })
577    }
578
579    pub fn into_null_bitmap(self) -> Bitmap {
580        dispatch_array_variants!(self, inner, { inner.into_null_bitmap() })
581    }
582
583    pub fn to_protobuf(&self) -> PbArray {
584        dispatch_array_variants!(self, inner, { inner.to_protobuf() })
585    }
586
587    pub fn hash_at<H: Hasher>(&self, idx: usize, state: &mut H) {
588        dispatch_array_variants!(self, inner, { inner.hash_at(idx, state) })
589    }
590
591    pub fn hash_vec<H: Hasher>(&self, hashers: &mut [H], vis: &Bitmap) {
592        dispatch_array_variants!(self, inner, { inner.hash_vec(hashers, vis) })
593    }
594
595    /// Select some elements from `Array` based on `visibility` bitmap.
596    pub fn compact_vis(&self, visibility: &Bitmap, cardinality: usize) -> Self {
597        dispatch_array_variants!(self, inner, {
598            inner.compact_vis(visibility, cardinality).into()
599        })
600    }
601
602    pub fn get_ident(&self) -> &'static str {
603        dispatch_array_variants!(self, [I = VARIANT_NAME], { I })
604    }
605
606    /// Get the enum-wrapped `Datum` out of the `Array`.
607    pub fn datum_at(&self, idx: usize) -> Datum {
608        self.value_at(idx).to_owned_datum()
609    }
610
611    /// If the array only have one single element, convert it to `Datum`.
612    pub fn to_datum(&self) -> Datum {
613        assert_eq!(self.len(), 1);
614        self.datum_at(0)
615    }
616
617    /// Get the enum-wrapped `ScalarRefImpl` out of the `Array`.
618    pub fn value_at(&self, idx: usize) -> DatumRef<'_> {
619        dispatch_array_variants!(self, inner, {
620            inner.value_at(idx).map(ScalarRefImpl::from)
621        })
622    }
623
624    /// # Safety
625    ///
626    /// This function is unsafe because it does not check the validity of `idx`. It is caller's
627    /// responsibility to ensure the validity of `idx`.
628    ///
629    /// Unsafe version of getting the enum-wrapped `ScalarRefImpl` out of the `Array`.
630    pub unsafe fn value_at_unchecked(&self, idx: usize) -> DatumRef<'_> {
631        unsafe {
632            dispatch_array_variants!(self, inner, {
633                inner.value_at_unchecked(idx).map(ScalarRefImpl::from)
634            })
635        }
636    }
637
638    pub fn set_bitmap(&mut self, bitmap: Bitmap) {
639        dispatch_array_variants!(self, inner, { inner.set_bitmap(bitmap) })
640    }
641
642    pub fn create_builder(&self, capacity: usize) -> ArrayBuilderImpl {
643        dispatch_array_variants!(self, inner, { inner.create_builder(capacity).into() })
644    }
645
646    /// Returns the `DataType` of this array.
647    pub fn data_type(&self) -> DataType {
648        dispatch_array_variants!(self, inner, { inner.data_type() })
649    }
650
651    pub fn into_ref(self) -> ArrayRef {
652        Arc::new(self)
653    }
654
655    pub fn iter(&self) -> impl DoubleEndedIterator<Item = DatumRef<'_>> + ExactSizeIterator {
656        (0..self.len()).map(|i| self.value_at(i))
657    }
658}
659
660pub type ArrayRef = Arc<ArrayImpl>;
661
662impl PartialEq for ArrayImpl {
663    fn eq(&self, other: &Self) -> bool {
664        self.iter().eq(other.iter())
665    }
666}
667
668impl Eq for ArrayImpl {}
669
670#[cfg(test)]
671mod tests {
672
673    use super::*;
674    use crate::util::iter_util::ZipEqFast;
675
676    fn filter<'a, A, F>(data: &'a A, pred: F) -> ArrayResult<A>
677    where
678        A: Array + 'a,
679        F: Fn(Option<A::RefItem<'a>>) -> bool,
680    {
681        let mut builder = A::Builder::with_type(data.len(), data.data_type());
682        for i in 0..data.len() {
683            if pred(data.value_at(i)) {
684                builder.append(data.value_at(i));
685            }
686        }
687        Ok(builder.finish())
688    }
689
690    #[test]
691    fn test_filter() {
692        let mut builder = PrimitiveArrayBuilder::<i32>::new(0);
693        for i in 0..=60 {
694            builder.append(Some(i));
695        }
696        let array = filter(&builder.finish(), |x| x.unwrap_or(0) >= 60).unwrap();
697        assert_eq!(array.iter().collect::<Vec<Option<i32>>>(), vec![Some(60)]);
698    }
699
700    use num_traits::ops::checked::CheckedAdd;
701
702    fn vec_add<T1, T2, T3>(
703        a: &PrimitiveArray<T1>,
704        b: &PrimitiveArray<T2>,
705    ) -> ArrayResult<PrimitiveArray<T3>>
706    where
707        T1: PrimitiveArrayItemType,
708        T2: PrimitiveArrayItemType,
709        T3: PrimitiveArrayItemType + CheckedAdd + From<T1> + From<T2>,
710    {
711        let mut builder = PrimitiveArrayBuilder::<T3>::new(a.len());
712        for (a, b) in a.iter().zip_eq_fast(b.iter()) {
713            let item = match (a, b) {
714                (Some(a), Some(b)) => Some(T3::from(a) + T3::from(b)),
715                _ => None,
716            };
717            builder.append(item);
718        }
719        Ok(builder.finish())
720    }
721
722    #[test]
723    fn test_vectorized_add() {
724        let mut builder = PrimitiveArrayBuilder::<i32>::new(0);
725        for i in 0..=60 {
726            builder.append(Some(i));
727        }
728        let array1 = builder.finish();
729
730        let mut builder = PrimitiveArrayBuilder::<i16>::new(0);
731        for i in 0..=60 {
732            builder.append(Some(i as i16));
733        }
734        let array2 = builder.finish();
735
736        let final_array = vec_add(&array1, &array2).unwrap() as PrimitiveArray<i64>;
737
738        assert_eq!(final_array.len(), array1.len());
739        for (idx, data) in final_array.iter().enumerate() {
740            assert_eq!(data, Some(idx as i64 * 2));
741        }
742    }
743}
744
745#[cfg(test)]
746mod test_util {
747    use std::hash::{BuildHasher, Hasher};
748
749    use super::Array;
750    use crate::bitmap::Bitmap;
751    use crate::util::iter_util::ZipEqFast;
752
753    pub fn hash_finish<H: Hasher>(hashers: &[H]) -> Vec<u64> {
754        hashers
755            .iter()
756            .map(|hasher| hasher.finish())
757            .collect::<Vec<u64>>()
758    }
759
760    pub fn test_hash<H: BuildHasher, A: Array>(arrs: Vec<A>, expects: Vec<u64>, hasher_builder: H) {
761        let len = expects.len();
762        let mut states_scalar = Vec::with_capacity(len);
763        states_scalar.resize_with(len, || hasher_builder.build_hasher());
764        let mut states_vec = Vec::with_capacity(len);
765        states_vec.resize_with(len, || hasher_builder.build_hasher());
766
767        arrs.iter().for_each(|arr| {
768            for (i, state) in states_scalar.iter_mut().enumerate() {
769                arr.hash_at(i, state)
770            }
771        });
772        let vis = Bitmap::ones(len);
773        arrs.iter()
774            .for_each(|arr| arr.hash_vec(&mut states_vec[..], &vis));
775        itertools::cons_tuples(
776            expects
777                .iter()
778                .zip_eq_fast(hash_finish(&states_scalar[..]))
779                .zip_eq_fast(hash_finish(&states_vec[..])),
780        )
781        .all(|(a, b, c)| *a == b && b == c);
782    }
783}
risingwave_common/array/mod.rs

risingwave_common/array/
mod.rs
