parquet/arrow/array_reader/
primitive_array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::arrow::array_reader::{ArrayReader, read_records, skip_records};
6use crate::arrow::record_reader::RecordReader;
7use crate::arrow::schema::parquet_to_arrow_field;
8use crate::basic::Type as PhysicalType;
9use crate::column::page::PageIterator;
10use crate::data_type::{DataType, Int96};
11use crate::errors::Result;
12use crate::schema::types::ColumnDescPtr;
13use arrow_array::{
14    Array, ArrayRef, BooleanArray, Date64Array, Decimal64Array, Decimal128Array, Decimal256Array,
15    Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, PrimitiveArray,
16    UInt8Array, UInt16Array, builder::PrimitiveDictionaryBuilder, cast::AsArray, downcast_integer,
17    types::*,
18};
19use arrow_array::{
20    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
21    TimestampSecondArray, UInt32Array, UInt64Array,
22};
23use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer, i256};
24use arrow_schema::{DataType as ArrowType, TimeUnit};
25use std::any::Any;
26use std::sync::Arc;
27
28/// Provides conversion from `Vec<T>` to `Buffer`
29pub trait IntoBuffer {
30    fn into_buffer(self, target_type: &ArrowType) -> Buffer;
31}
32
33macro_rules! native_buffer {
34    ($($t:ty),*) => {
35        $(impl IntoBuffer for Vec<$t> {
36            fn into_buffer(self, _target_type: &ArrowType) -> Buffer {
37                Buffer::from_vec(self)
38            }
39        })*
40    };
41}
42native_buffer!(i8, i16, i32, i64, u8, u16, u32, u64, f32, f64);
43
44impl IntoBuffer for Vec<bool> {
45    fn into_buffer(self, _target_type: &ArrowType) -> Buffer {
46        BooleanBuffer::from(self.as_slice()).into_inner()
47    }
48}
49
50impl IntoBuffer for Vec<Int96> {
51    fn into_buffer(self, target_type: &ArrowType) -> Buffer {
52        let mut builder = Vec::with_capacity(self.len());
53        match target_type {
54            ArrowType::Timestamp(TimeUnit::Second, _) => {
55                builder.extend(self.iter().map(|x| x.to_seconds()));
56            }
57            ArrowType::Timestamp(TimeUnit::Millisecond, _) => {
58                builder.extend(self.iter().map(|x| x.to_millis()));
59            }
60            ArrowType::Timestamp(TimeUnit::Microsecond, _) => {
61                builder.extend(self.iter().map(|x| x.to_micros()));
62            }
63            ArrowType::Timestamp(TimeUnit::Nanosecond, _) => {
64                builder.extend(self.iter().map(|x| x.to_nanos()));
65            }
66            _ => unreachable!("Invalid target_type for Int96."),
67        }
68        Buffer::from_vec(builder)
69    }
70}
71
72/// Primitive array readers are leaves of array reader tree. They accept page iterator
73/// and read them into primitive arrays.
74pub struct PrimitiveArrayReader<T>
75where
76    T: DataType,
77    T::T: Copy + Default,
78    Vec<T::T>: IntoBuffer,
79{
80    data_type: ArrowType,
81    pages: Box<dyn PageIterator>,
82    def_levels_buffer: Option<Vec<i16>>,
83    rep_levels_buffer: Option<Vec<i16>>,
84    record_reader: RecordReader<T>,
85}
86
87impl<T> PrimitiveArrayReader<T>
88where
89    T: DataType,
90    T::T: Copy + Default,
91    Vec<T::T>: IntoBuffer,
92{
93    /// Construct primitive array reader.
94    ///
95    /// `padding_threshold` controls how null padding is applied. When
96    /// `None`, the reader pads all null positions (full padding). When
97    /// `Some(threshold)`, entries with `def < threshold` are excluded
98    /// from the value buffer (selective padding for list children).
99    pub fn new(
100        pages: Box<dyn PageIterator>,
101        column_desc: ColumnDescPtr,
102        arrow_type: Option<ArrowType>,
103        batch_size: usize,
104        padding_threshold: Option<i16>,
105    ) -> Result<Self> {
106        // Check if Arrow type is specified, else create it from Parquet type
107        let data_type = match arrow_type {
108            Some(t) => t,
109            None => parquet_to_arrow_field(column_desc.as_ref())?
110                .data_type()
111                .clone(),
112        };
113
114        let mut record_reader = RecordReader::<T>::new(column_desc, batch_size);
115        if let Some(threshold) = padding_threshold {
116            record_reader.set_padding_threshold(threshold);
117        }
118
119        Ok(Self {
120            data_type,
121            pages,
122            def_levels_buffer: None,
123            rep_levels_buffer: None,
124            record_reader,
125        })
126    }
127}
128
129/// Implementation of primitive array reader.
130impl<T> ArrayReader for PrimitiveArrayReader<T>
131where
132    T: DataType,
133    T::T: Copy + Default,
134    Vec<T::T>: IntoBuffer,
135{
136    fn as_any(&self) -> &dyn Any {
137        self
138    }
139
140    /// Returns data type of primitive array.
141    fn get_data_type(&self) -> &ArrowType {
142        &self.data_type
143    }
144
145    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
146        read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)
147    }
148
149    fn consume_batch(&mut self) -> Result<ArrayRef> {
150        let target_type = &self.data_type;
151
152        // Convert physical data to equivalent arrow type, and then perform
153        // coercion as needed
154        let len = self.record_reader.values_written();
155
156        let record_data = self
157            .record_reader
158            .consume_record_data()
159            .into_buffer(target_type);
160
161        let nulls = self
162            .record_reader
163            .consume_compact_bitmap()
164            .and_then(|b| NullBuffer::from_unsliced_buffer(b, len));
165
166        let array: ArrayRef = match T::get_physical_type() {
167            PhysicalType::BOOLEAN => Arc::new(BooleanArray::new(
168                BooleanBuffer::new(record_data, 0, len),
169                nulls,
170            )),
171            PhysicalType::INT32 => Arc::new(Int32Array::new(
172                ScalarBuffer::new(record_data, 0, len),
173                nulls,
174            )),
175            PhysicalType::INT64 => Arc::new(Int64Array::new(
176                ScalarBuffer::new(record_data, 0, len),
177                nulls,
178            )),
179            PhysicalType::FLOAT => Arc::new(Float32Array::new(
180                ScalarBuffer::new(record_data, 0, len),
181                nulls,
182            )),
183            PhysicalType::DOUBLE => Arc::new(Float64Array::new(
184                ScalarBuffer::new(record_data, 0, len),
185                nulls,
186            )),
187            PhysicalType::INT96 => Arc::new(Int64Array::new(
188                ScalarBuffer::new(record_data, 0, len),
189                nulls,
190            )),
191            PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
192                unreachable!("PrimitiveArrayReaders don't support complex physical types");
193            }
194        };
195
196        // Coerce the arrow type to the desired array type
197        let array = coerce_array(array, target_type)?;
198
199        // save definition and repetition buffers
200        self.def_levels_buffer = self.record_reader.consume_def_levels();
201        self.rep_levels_buffer = self.record_reader.consume_rep_levels();
202        self.record_reader.reset();
203        Ok(array)
204    }
205
206    fn skip_records(&mut self, num_records: usize) -> Result<usize> {
207        skip_records(&mut self.record_reader, self.pages.as_mut(), num_records)
208    }
209
210    fn get_def_levels(&self) -> Option<&[i16]> {
211        self.def_levels_buffer.as_deref()
212    }
213
214    fn get_rep_levels(&self) -> Option<&[i16]> {
215        self.rep_levels_buffer.as_deref()
216    }
217
218    fn max_def_level(&self) -> i16 {
219        self.record_reader.max_def_level()
220    }
221}
222
223/// Coerce the parquet physical type array to the target type
224///
225/// This should match the logic in schema::primitive::apply_hint
226fn coerce_array(array: ArrayRef, target_type: &ArrowType) -> Result<ArrayRef> {
227    if let ArrowType::Dictionary(key_type, value_type) = target_type {
228        let dictionary = pack_dictionary(key_type, array.as_ref())?;
229        let any_dictionary = dictionary.as_any_dictionary();
230
231        let coerced_values =
232            coerce_array(Arc::clone(any_dictionary.values()), value_type.as_ref())?;
233
234        return Ok(any_dictionary.with_values(coerced_values));
235    }
236
237    match array.data_type() {
238        ArrowType::Int32 => coerce_i32(array.as_primitive(), target_type),
239        ArrowType::Int64 => coerce_i64(array.as_primitive(), target_type),
240        ArrowType::Boolean | ArrowType::Float32 | ArrowType::Float64 => Ok(array),
241        _ => unreachable!("Cannot coerce array of type {}", array.data_type()),
242    }
243}
244
245fn coerce_i32(array: &Int32Array, target_type: &ArrowType) -> Result<ArrayRef> {
246    Ok(match target_type {
247        ArrowType::UInt8 => {
248            let array = array.unary(|i| i as u8) as UInt8Array;
249            Arc::new(array) as ArrayRef
250        }
251        ArrowType::Int8 => {
252            let array = array.unary(|i| i as i8) as Int8Array;
253            Arc::new(array) as ArrayRef
254        }
255        ArrowType::UInt16 => {
256            let array = array.unary(|i| i as u16) as UInt16Array;
257            Arc::new(array) as ArrayRef
258        }
259        ArrowType::Int16 => {
260            let array = array.unary(|i| i as i16) as Int16Array;
261            Arc::new(array) as ArrayRef
262        }
263        ArrowType::Int32 => Arc::new(array.clone()),
264        // follow C++ implementation and use overflow/reinterpret cast from  i32 to u32 which will map
265        // `i32::MIN..0` to `(i32::MAX as u32)..u32::MAX`
266        ArrowType::UInt32 => Arc::new(UInt32Array::new(
267            array.values().inner().clone().into(),
268            array.nulls().cloned(),
269        )) as ArrayRef,
270        ArrowType::Date32 => Arc::new(array.reinterpret_cast::<Date32Type>()) as _,
271        ArrowType::Date64 => {
272            let array: Date64Array = array.unary(|x| x as i64 * 86_400_000);
273            Arc::new(array) as ArrayRef
274        }
275        ArrowType::Time32(TimeUnit::Second) => {
276            Arc::new(array.reinterpret_cast::<Time32SecondType>()) as ArrayRef
277        }
278        ArrowType::Time32(TimeUnit::Millisecond) => {
279            Arc::new(array.reinterpret_cast::<Time32MillisecondType>()) as ArrayRef
280        }
281        ArrowType::Timestamp(time_unit, timezone) => match time_unit {
282            TimeUnit::Second => {
283                let array: TimestampSecondArray = array
284                    .unary(|x| x as i64)
285                    .with_timezone_opt(timezone.clone());
286                Arc::new(array) as _
287            }
288            TimeUnit::Millisecond => {
289                let array: TimestampMillisecondArray = array
290                    .unary(|x| x as i64)
291                    .with_timezone_opt(timezone.clone());
292                Arc::new(array) as _
293            }
294            TimeUnit::Microsecond => {
295                let array: TimestampMicrosecondArray = array
296                    .unary(|x| x as i64)
297                    .with_timezone_opt(timezone.clone());
298                Arc::new(array) as _
299            }
300            TimeUnit::Nanosecond => {
301                let array: TimestampNanosecondArray = array
302                    .unary(|x| x as i64)
303                    .with_timezone_opt(timezone.clone());
304                Arc::new(array) as _
305            }
306        },
307        ArrowType::Decimal32(p, s) => {
308            let array = array
309                .reinterpret_cast::<Decimal32Type>()
310                .with_precision_and_scale(*p, *s)?;
311            Arc::new(array) as ArrayRef
312        }
313        ArrowType::Decimal64(p, s) => {
314            let array: Decimal64Array =
315                array.unary(|i| i as i64).with_precision_and_scale(*p, *s)?;
316            Arc::new(array) as ArrayRef
317        }
318        ArrowType::Decimal128(p, s) => {
319            let array: Decimal128Array = array
320                .unary(|i| i as i128)
321                .with_precision_and_scale(*p, *s)?;
322            Arc::new(array) as ArrayRef
323        }
324        ArrowType::Decimal256(p, s) => {
325            let array: Decimal256Array = array
326                .unary(|i| i256::from_i128(i as i128))
327                .with_precision_and_scale(*p, *s)?;
328            Arc::new(array) as ArrayRef
329        }
330        _ => unreachable!("Cannot coerce i32 to {target_type}"),
331    })
332}
333
334fn coerce_i64(array: &Int64Array, target_type: &ArrowType) -> Result<ArrayRef> {
335    Ok(match target_type {
336        ArrowType::Int64 => Arc::new(array.clone()) as _,
337        // follow C++ implementation and use overflow/reinterpret cast from i64 to u64 which will map
338        // `i64::MIN..0` to `(i64::MAX as u64)..u64::MAX`
339        ArrowType::UInt64 => Arc::new(UInt64Array::new(
340            array.values().inner().clone().into(),
341            array.nulls().cloned(),
342        )) as ArrayRef,
343        ArrowType::Date64 => Arc::new(array.reinterpret_cast::<Date64Type>()) as _,
344        ArrowType::Time64(TimeUnit::Microsecond) => {
345            Arc::new(array.reinterpret_cast::<Time64MicrosecondType>()) as _
346        }
347        ArrowType::Time64(TimeUnit::Nanosecond) => {
348            Arc::new(array.reinterpret_cast::<Time64NanosecondType>()) as _
349        }
350        ArrowType::Duration(unit) => match unit {
351            TimeUnit::Second => Arc::new(array.reinterpret_cast::<DurationSecondType>()) as _,
352            TimeUnit::Millisecond => {
353                Arc::new(array.reinterpret_cast::<DurationMillisecondType>()) as _
354            }
355            TimeUnit::Microsecond => {
356                Arc::new(array.reinterpret_cast::<DurationMicrosecondType>()) as _
357            }
358            TimeUnit::Nanosecond => {
359                Arc::new(array.reinterpret_cast::<DurationNanosecondType>()) as _
360            }
361        },
362        ArrowType::Timestamp(time_unit, timezone) => match time_unit {
363            TimeUnit::Second => {
364                let array = array
365                    .reinterpret_cast::<TimestampSecondType>()
366                    .with_timezone_opt(timezone.clone());
367                Arc::new(array) as _
368            }
369            TimeUnit::Millisecond => {
370                let array = array
371                    .reinterpret_cast::<TimestampMillisecondType>()
372                    .with_timezone_opt(timezone.clone());
373                Arc::new(array) as _
374            }
375            TimeUnit::Microsecond => {
376                let array = array
377                    .reinterpret_cast::<TimestampMicrosecondType>()
378                    .with_timezone_opt(timezone.clone());
379                Arc::new(array) as _
380            }
381            TimeUnit::Nanosecond => {
382                let array = array
383                    .reinterpret_cast::<TimestampNanosecondType>()
384                    .with_timezone_opt(timezone.clone());
385                Arc::new(array) as _
386            }
387        },
388        ArrowType::Decimal64(p, s) => {
389            let array = array
390                .reinterpret_cast::<Decimal64Type>()
391                .with_precision_and_scale(*p, *s)?;
392            Arc::new(array) as _
393        }
394        ArrowType::Decimal128(p, s) => {
395            let array: Decimal128Array = array
396                .unary(|i| i as i128)
397                .with_precision_and_scale(*p, *s)?;
398            Arc::new(array) as _
399        }
400        ArrowType::Decimal256(p, s) => {
401            let array: Decimal256Array = array
402                .unary(|i| i256::from_i128(i as i128))
403                .with_precision_and_scale(*p, *s)?;
404            Arc::new(array) as _
405        }
406        _ => unreachable!("Cannot coerce i64 to {target_type}"),
407    })
408}
409
410macro_rules! pack_dictionary_helper {
411    ($t:ty, $values:ident) => {
412        match $values.data_type() {
413            ArrowType::Int32 => pack_dictionary_impl::<$t, Int32Type>($values.as_primitive()),
414            ArrowType::Int64 => pack_dictionary_impl::<$t, Int64Type>($values.as_primitive()),
415            ArrowType::Float32 => pack_dictionary_impl::<$t, Float32Type>($values.as_primitive()),
416            ArrowType::Float64 => pack_dictionary_impl::<$t, Float64Type>($values.as_primitive()),
417            _ => unreachable!("Invalid physical type"),
418        }
419    };
420}
421
422fn pack_dictionary(key: &ArrowType, values: &dyn Array) -> Result<ArrayRef> {
423    downcast_integer! {
424        key => (pack_dictionary_helper, values),
425        _ => unreachable!("Invalid key type"),
426    }
427}
428
429fn pack_dictionary_impl<K: ArrowDictionaryKeyType, V: ArrowPrimitiveType>(
430    values: &PrimitiveArray<V>,
431) -> Result<ArrayRef> {
432    let mut builder = PrimitiveDictionaryBuilder::<K, V>::with_capacity(1024, values.len());
433    builder.extend(values);
434    Ok(Arc::new(builder.finish()))
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440    use crate::arrow::array_reader::test_util::EmptyPageIterator;
441    use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE;
442    use crate::basic::Encoding;
443    use crate::column::page::Page;
444    use crate::data_type::{Int32Type, Int64Type};
445    use crate::schema::parser::parse_message_type;
446    use crate::schema::types::SchemaDescriptor;
447    use crate::util::InMemoryPageIterator;
448    use crate::util::test_common::rand_gen::make_pages;
449    use arrow::datatypes::ArrowPrimitiveType;
450    use arrow_array::{Array, Date32Array, PrimitiveArray};
451
452    use arrow::datatypes::DataType::{Date32, Decimal128};
453    use rand::distr::uniform::SampleUniform;
454    use std::collections::VecDeque;
455
456    #[expect(clippy::too_many_arguments)]
457    fn make_column_chunks<T: DataType>(
458        column_desc: ColumnDescPtr,
459        encoding: Encoding,
460        num_levels: usize,
461        min_value: T::T,
462        max_value: T::T,
463        def_levels: &mut Vec<i16>,
464        rep_levels: &mut Vec<i16>,
465        values: &mut Vec<T::T>,
466        page_lists: &mut Vec<Vec<Page>>,
467        use_v2: bool,
468        num_chunks: usize,
469    ) where
470        T::T: PartialOrd + SampleUniform + Copy,
471    {
472        for _i in 0..num_chunks {
473            let mut pages = VecDeque::new();
474            let mut data = Vec::new();
475            let mut page_def_levels = Vec::new();
476            let mut page_rep_levels = Vec::new();
477
478            make_pages::<T>(
479                column_desc.clone(),
480                encoding,
481                1,
482                num_levels,
483                min_value,
484                max_value,
485                &mut page_def_levels,
486                &mut page_rep_levels,
487                &mut data,
488                &mut pages,
489                use_v2,
490            );
491
492            def_levels.append(&mut page_def_levels);
493            rep_levels.append(&mut page_rep_levels);
494            values.append(&mut data);
495            page_lists.push(Vec::from(pages));
496        }
497    }
498
499    #[test]
500    fn test_primitive_array_reader_empty_pages() {
501        // Construct column schema
502        let message_type = "
503        message test_schema {
504          REQUIRED INT32 leaf;
505        }
506        ";
507
508        let schema = parse_message_type(message_type)
509            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
510            .unwrap();
511
512        let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
513            Box::<EmptyPageIterator>::default(),
514            schema.column(0),
515            None,
516            DEFAULT_BATCH_SIZE,
517            None,
518        )
519        .unwrap();
520
521        // expect no values to be read
522        let array = array_reader.next_batch(50).unwrap();
523        assert!(array.is_empty());
524    }
525
526    #[test]
527    fn test_primitive_array_reader_data() {
528        // Construct column schema
529        let message_type = "
530        message test_schema {
531          REQUIRED INT32 leaf;
532        }
533        ";
534
535        let schema = parse_message_type(message_type)
536            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
537            .unwrap();
538
539        let column_desc = schema.column(0);
540
541        // Construct page iterator
542        {
543            let mut data = Vec::new();
544            let mut page_lists = Vec::new();
545            make_column_chunks::<Int32Type>(
546                column_desc.clone(),
547                Encoding::PLAIN,
548                100,
549                1,
550                200,
551                &mut Vec::new(),
552                &mut Vec::new(),
553                &mut data,
554                &mut page_lists,
555                true,
556                2,
557            );
558            let page_iterator = InMemoryPageIterator::new(page_lists);
559
560            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
561                Box::new(page_iterator),
562                column_desc,
563                None,
564                DEFAULT_BATCH_SIZE,
565                None,
566            )
567            .unwrap();
568
569            // Read first 50 values, which are all from the first column chunk
570            let array = array_reader.next_batch(50).unwrap();
571            let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
572
573            assert_eq!(&Int32Array::from(data[0..50].to_vec()), array);
574
575            // Read next 100 values, the first 50 ones are from the first column chunk,
576            // and the last 50 ones are from the second column chunk
577            let array = array_reader.next_batch(100).unwrap();
578            let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
579
580            assert_eq!(&Int32Array::from(data[50..150].to_vec()), array);
581
582            // Try to read 100 values, however there are only 50 values
583            let array = array_reader.next_batch(100).unwrap();
584            let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
585
586            assert_eq!(&Int32Array::from(data[150..200].to_vec()), array);
587        }
588    }
589
590    macro_rules! test_primitive_array_reader_one_type {
591        (
592            $arrow_parquet_type:ty,
593            $physical_type:expr,
594            $converted_type_str:expr,
595            $result_arrow_type:ty,
596            $result_arrow_cast_type:ty,
597            $result_primitive_type:ty
598            $(, $timezone:expr)?
599        ) => {{
600            let message_type = format!(
601                "
602            message test_schema {{
603              REQUIRED {:?} leaf ({});
604          }}
605            ",
606                $physical_type, $converted_type_str
607            );
608            let schema = parse_message_type(&message_type)
609                .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
610                .unwrap();
611
612            let column_desc = schema.column(0);
613
614            // Construct page iterator
615            {
616                let mut data = Vec::new();
617                let mut page_lists = Vec::new();
618                make_column_chunks::<$arrow_parquet_type>(
619                    column_desc.clone(),
620                    Encoding::PLAIN,
621                    100,
622                    1,
623                    200,
624                    &mut Vec::new(),
625                    &mut Vec::new(),
626                    &mut data,
627                    &mut page_lists,
628                    true,
629                    2,
630                );
631                let page_iterator = InMemoryPageIterator::new(page_lists);
632                let mut array_reader = PrimitiveArrayReader::<$arrow_parquet_type>::new(
633                    Box::new(page_iterator),
634                    column_desc.clone(),
635                    None,
636                    DEFAULT_BATCH_SIZE,
637                    None,
638                )
639                .expect("Unable to get array reader");
640
641                let array = array_reader
642                    .next_batch(50)
643                    .expect("Unable to get batch from reader");
644
645                let result_data_type = <$result_arrow_type>::DATA_TYPE;
646                let array = array
647                    .as_any()
648                    .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
649                    .expect(
650                        format!(
651                            "Unable to downcast {:?} to {:?}",
652                            array.data_type(),
653                            result_data_type
654                        )
655                        .as_str(),
656                    )
657                    $(.clone().with_timezone($timezone))?
658                    ;
659
660                // create expected array as primitive, and cast to result type
661                #[expect(
662                    trivial_numeric_casts,
663                    reason = "the cast is a no-op when the source and result types are the same"
664                )]
665                let expected = PrimitiveArray::<$result_arrow_cast_type>::from(
666                    data[0..50]
667                        .iter()
668                        .map(|x| *x as $result_primitive_type)
669                        .collect::<Vec<$result_primitive_type>>(),
670                );
671                let expected = Arc::new(expected) as ArrayRef;
672                let expected = arrow::compute::cast(&expected, &result_data_type)
673                    .expect("Unable to cast expected array");
674                assert_eq!(expected.data_type(), &result_data_type);
675                let expected = expected
676                    .as_any()
677                    .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
678                    .expect(
679                        format!(
680                            "Unable to downcast expected {:?} to {:?}",
681                            expected.data_type(),
682                            result_data_type
683                        )
684                        .as_str(),
685                    )
686                    $(.clone().with_timezone($timezone))?
687                    ;
688                assert_eq!(expected, array);
689            }
690        }};
691    }
692
693    #[test]
694    fn test_primitive_array_reader_temporal_types() {
695        test_primitive_array_reader_one_type!(
696            crate::data_type::Int32Type,
697            PhysicalType::INT32,
698            "DATE",
699            arrow::datatypes::Date32Type,
700            arrow::datatypes::Int32Type,
701            i32
702        );
703        test_primitive_array_reader_one_type!(
704            crate::data_type::Int32Type,
705            PhysicalType::INT32,
706            "TIME_MILLIS",
707            arrow::datatypes::Time32MillisecondType,
708            arrow::datatypes::Int32Type,
709            i32
710        );
711        test_primitive_array_reader_one_type!(
712            crate::data_type::Int64Type,
713            PhysicalType::INT64,
714            "TIME_MICROS",
715            arrow::datatypes::Time64MicrosecondType,
716            arrow::datatypes::Int64Type,
717            i64
718        );
719        test_primitive_array_reader_one_type!(
720            crate::data_type::Int64Type,
721            PhysicalType::INT64,
722            "TIMESTAMP_MILLIS",
723            arrow::datatypes::TimestampMillisecondType,
724            arrow::datatypes::Int64Type,
725            i64,
726            "UTC"
727        );
728        test_primitive_array_reader_one_type!(
729            crate::data_type::Int64Type,
730            PhysicalType::INT64,
731            "TIMESTAMP_MICROS",
732            arrow::datatypes::TimestampMicrosecondType,
733            arrow::datatypes::Int64Type,
734            i64,
735            "UTC"
736        );
737    }
738
739    #[test]
740    fn test_primitive_array_reader_def_and_rep_levels() {
741        // Construct column schema
742        let message_type = "
743        message test_schema {
744            REPEATED Group test_mid {
745                OPTIONAL INT32 leaf;
746            }
747        }
748        ";
749
750        let schema = parse_message_type(message_type)
751            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
752            .unwrap();
753
754        let column_desc = schema.column(0);
755
756        // Construct page iterator
757        {
758            let mut def_levels = Vec::new();
759            let mut rep_levels = Vec::new();
760            let mut page_lists = Vec::new();
761            make_column_chunks::<Int32Type>(
762                column_desc.clone(),
763                Encoding::PLAIN,
764                100,
765                1,
766                200,
767                &mut def_levels,
768                &mut rep_levels,
769                &mut Vec::new(),
770                &mut page_lists,
771                true,
772                2,
773            );
774
775            let page_iterator = InMemoryPageIterator::new(page_lists);
776
777            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
778                Box::new(page_iterator),
779                column_desc,
780                None,
781                DEFAULT_BATCH_SIZE,
782                None,
783            )
784            .unwrap();
785
786            let mut accu_len: usize = 0;
787
788            // Read first 50 values, which are all from the first column chunk
789            let array = array_reader.next_batch(50).unwrap();
790            assert_eq!(
791                Some(&def_levels[accu_len..(accu_len + array.len())]),
792                array_reader.get_def_levels()
793            );
794            assert_eq!(
795                Some(&rep_levels[accu_len..(accu_len + array.len())]),
796                array_reader.get_rep_levels()
797            );
798            accu_len += array.len();
799
800            // Read next 100 values, the first 50 ones are from the first column chunk,
801            // and the last 50 ones are from the second column chunk
802            let array = array_reader.next_batch(100).unwrap();
803            assert_eq!(
804                Some(&def_levels[accu_len..(accu_len + array.len())]),
805                array_reader.get_def_levels()
806            );
807            assert_eq!(
808                Some(&rep_levels[accu_len..(accu_len + array.len())]),
809                array_reader.get_rep_levels()
810            );
811            accu_len += array.len();
812
813            // Try to read 100 values, however there are only 50 values
814            let array = array_reader.next_batch(100).unwrap();
815            assert_eq!(
816                Some(&def_levels[accu_len..(accu_len + array.len())]),
817                array_reader.get_def_levels()
818            );
819            assert_eq!(
820                Some(&rep_levels[accu_len..(accu_len + array.len())]),
821                array_reader.get_rep_levels()
822            );
823        }
824    }
825
826    #[test]
827    fn test_primitive_array_reader_decimal_types() {
828        // parquet `INT32` to decimal
829        let message_type = "
830            message test_schema {
831                REQUIRED INT32 decimal1 (DECIMAL(8,2));
832        }
833        ";
834        let schema = parse_message_type(message_type)
835            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
836            .unwrap();
837        let column_desc = schema.column(0);
838
839        // create the array reader
840        {
841            let mut data = Vec::new();
842            let mut page_lists = Vec::new();
843            make_column_chunks::<Int32Type>(
844                column_desc.clone(),
845                Encoding::PLAIN,
846                100,
847                -99999999,
848                99999999,
849                &mut Vec::new(),
850                &mut Vec::new(),
851                &mut data,
852                &mut page_lists,
853                true,
854                2,
855            );
856            let page_iterator = InMemoryPageIterator::new(page_lists);
857
858            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
859                Box::new(page_iterator),
860                column_desc,
861                None,
862                DEFAULT_BATCH_SIZE,
863                None,
864            )
865            .unwrap();
866
867            // read data from the reader
868            // the data type is decimal(8,2)
869            let array = array_reader.next_batch(50).unwrap();
870            assert_eq!(array.data_type(), &Decimal128(8, 2));
871            let array = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
872            let data_decimal_array = data[0..50]
873                .iter()
874                .copied()
875                .map(|v| Some(v as i128))
876                .collect::<Decimal128Array>()
877                .with_precision_and_scale(8, 2)
878                .unwrap();
879            assert_eq!(array, &data_decimal_array);
880
881            // not equal with different data type(precision and scale)
882            let data_decimal_array = data[0..50]
883                .iter()
884                .copied()
885                .map(|v| Some(v as i128))
886                .collect::<Decimal128Array>()
887                .with_precision_and_scale(9, 0)
888                .unwrap();
889            assert_ne!(array, &data_decimal_array)
890        }
891
892        // parquet `INT64` to decimal
893        let message_type = "
894            message test_schema {
895                REQUIRED INT64 decimal1 (DECIMAL(18,4));
896        }
897        ";
898        let schema = parse_message_type(message_type)
899            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
900            .unwrap();
901        let column_desc = schema.column(0);
902
903        // create the array reader
904        {
905            let mut data = Vec::new();
906            let mut page_lists = Vec::new();
907            make_column_chunks::<Int64Type>(
908                column_desc.clone(),
909                Encoding::PLAIN,
910                100,
911                -999999999999999999,
912                999999999999999999,
913                &mut Vec::new(),
914                &mut Vec::new(),
915                &mut data,
916                &mut page_lists,
917                true,
918                2,
919            );
920            let page_iterator = InMemoryPageIterator::new(page_lists);
921
922            let mut array_reader = PrimitiveArrayReader::<Int64Type>::new(
923                Box::new(page_iterator),
924                column_desc,
925                None,
926                DEFAULT_BATCH_SIZE,
927                None,
928            )
929            .unwrap();
930
931            // read data from the reader
932            // the data type is decimal(18,4)
933            let array = array_reader.next_batch(50).unwrap();
934            assert_eq!(array.data_type(), &Decimal128(18, 4));
935            let array = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
936            let data_decimal_array = data[0..50]
937                .iter()
938                .copied()
939                .map(|v| Some(v as i128))
940                .collect::<Decimal128Array>()
941                .with_precision_and_scale(18, 4)
942                .unwrap();
943            assert_eq!(array, &data_decimal_array);
944
945            // not equal with different data type(precision and scale)
946            let data_decimal_array = data[0..50]
947                .iter()
948                .copied()
949                .map(|v| Some(v as i128))
950                .collect::<Decimal128Array>()
951                .with_precision_and_scale(34, 0)
952                .unwrap();
953            assert_ne!(array, &data_decimal_array)
954        }
955    }
956
957    #[test]
958    fn test_primitive_array_reader_date32_type() {
959        // parquet `INT32` to date
960        let message_type = "
961            message test_schema {
962                REQUIRED INT32 date1 (DATE);
963        }
964        ";
965        let schema = parse_message_type(message_type)
966            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
967            .unwrap();
968        let column_desc = schema.column(0);
969
970        // create the array reader
971        {
972            let mut data = Vec::new();
973            let mut page_lists = Vec::new();
974            make_column_chunks::<Int32Type>(
975                column_desc.clone(),
976                Encoding::PLAIN,
977                100,
978                -99999999,
979                99999999,
980                &mut Vec::new(),
981                &mut Vec::new(),
982                &mut data,
983                &mut page_lists,
984                true,
985                2,
986            );
987            let page_iterator = InMemoryPageIterator::new(page_lists);
988
989            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
990                Box::new(page_iterator),
991                column_desc,
992                None,
993                DEFAULT_BATCH_SIZE,
994                None,
995            )
996            .unwrap();
997
998            // read data from the reader
999            // the data type is date
1000            let array = array_reader.next_batch(50).unwrap();
1001            assert_eq!(array.data_type(), &Date32);
1002            let array = array.as_any().downcast_ref::<Date32Array>().unwrap();
1003            let data_date_array = data[0..50]
1004                .iter()
1005                .copied()
1006                .map(Some)
1007                .collect::<Date32Array>();
1008            assert_eq!(array, &data_date_array);
1009        }
1010    }
1011}
parquet/arrow/array_reader/primitive_array.rs

parquet/arrow/array_reader/
primitive_array.rs
