parquet/record/
triplet.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::basic::Type as PhysicalType;
6use crate::column::reader::{ColumnReader, ColumnReaderImpl, get_typed_column_reader};
7use crate::data_type::*;
8use crate::errors::{ParquetError, Result};
9use crate::record::api::Field;
10use crate::schema::types::ColumnDescPtr;
11
12/// Macro to generate simple functions that cover all types of triplet iterator.
13/// $func is a function of a typed triplet iterator and $token is a either {`ref`} or
14/// {`ref`, `mut`}
15macro_rules! triplet_enum_func {
16  ($self:ident, $func:ident, $( $token:tt ),*) => ({
17    match *$self {
18      TripletIter::BoolTripletIter($($token)* typed) => typed.$func(),
19      TripletIter::Int32TripletIter($($token)* typed) => typed.$func(),
20      TripletIter::Int64TripletIter($($token)* typed) => typed.$func(),
21      TripletIter::Int96TripletIter($($token)* typed) => typed.$func(),
22      TripletIter::FloatTripletIter($($token)* typed) => typed.$func(),
23      TripletIter::DoubleTripletIter($($token)* typed) => typed.$func(),
24      TripletIter::ByteArrayTripletIter($($token)* typed) => typed.$func(),
25      TripletIter::FixedLenByteArrayTripletIter($($token)* typed) => typed.$func()
26    }
27  });
28}
29
30/// High level API wrapper on column reader.
31/// Provides per-element access for each primitive column.
32#[expect(clippy::enum_variant_names)]
33pub enum TripletIter {
34    BoolTripletIter(TypedTripletIter<BoolType>),
35    Int32TripletIter(TypedTripletIter<Int32Type>),
36    Int64TripletIter(TypedTripletIter<Int64Type>),
37    Int96TripletIter(TypedTripletIter<Int96Type>),
38    FloatTripletIter(TypedTripletIter<FloatType>),
39    DoubleTripletIter(TypedTripletIter<DoubleType>),
40    ByteArrayTripletIter(TypedTripletIter<ByteArrayType>),
41    FixedLenByteArrayTripletIter(TypedTripletIter<FixedLenByteArrayType>),
42}
43
44impl TripletIter {
45    /// Creates new triplet for column reader
46    pub fn new(descr: ColumnDescPtr, reader: ColumnReader, batch_size: usize) -> Self {
47        match descr.physical_type() {
48            PhysicalType::BOOLEAN => {
49                TripletIter::BoolTripletIter(TypedTripletIter::new(descr, batch_size, reader))
50            }
51            PhysicalType::INT32 => {
52                TripletIter::Int32TripletIter(TypedTripletIter::new(descr, batch_size, reader))
53            }
54            PhysicalType::INT64 => {
55                TripletIter::Int64TripletIter(TypedTripletIter::new(descr, batch_size, reader))
56            }
57            PhysicalType::INT96 => {
58                TripletIter::Int96TripletIter(TypedTripletIter::new(descr, batch_size, reader))
59            }
60            PhysicalType::FLOAT => {
61                TripletIter::FloatTripletIter(TypedTripletIter::new(descr, batch_size, reader))
62            }
63            PhysicalType::DOUBLE => {
64                TripletIter::DoubleTripletIter(TypedTripletIter::new(descr, batch_size, reader))
65            }
66            PhysicalType::BYTE_ARRAY => {
67                TripletIter::ByteArrayTripletIter(TypedTripletIter::new(descr, batch_size, reader))
68            }
69            PhysicalType::FIXED_LEN_BYTE_ARRAY => TripletIter::FixedLenByteArrayTripletIter(
70                TypedTripletIter::new(descr, batch_size, reader),
71            ),
72        }
73    }
74
75    /// Invokes underlying typed triplet iterator to buffer current value.
76    /// Should be called once - either before `is_null` or `current_value`.
77    #[inline]
78    pub fn read_next(&mut self) -> Result<bool> {
79        triplet_enum_func!(self, read_next, ref, mut)
80    }
81
82    /// Provides check on values/levels left without invoking the underlying typed triplet
83    /// iterator.
84    /// Returns true if more values/levels exist, false otherwise.
85    /// It is always in sync with `read_next` method.
86    #[inline]
87    pub fn has_next(&self) -> bool {
88        triplet_enum_func!(self, has_next, ref)
89    }
90
91    /// Returns current definition level for a leaf triplet iterator
92    #[inline]
93    pub fn current_def_level(&self) -> i16 {
94        triplet_enum_func!(self, current_def_level, ref)
95    }
96
97    /// Returns max definition level for a leaf triplet iterator
98    #[inline]
99    pub fn max_def_level(&self) -> i16 {
100        triplet_enum_func!(self, max_def_level, ref)
101    }
102
103    /// Returns current repetition level for a leaf triplet iterator
104    #[inline]
105    pub fn current_rep_level(&self) -> i16 {
106        triplet_enum_func!(self, current_rep_level, ref)
107    }
108
109    /// Returns max repetition level for a leaf triplet iterator
110    #[inline]
111    pub fn max_rep_level(&self) -> i16 {
112        triplet_enum_func!(self, max_rep_level, ref)
113    }
114
115    /// Returns true, if current value is null.
116    /// Based on the fact that for non-null value current definition level
117    /// equals to max definition level.
118    #[inline]
119    pub fn is_null(&self) -> bool {
120        self.current_def_level() < self.max_def_level()
121    }
122
123    /// Updates non-null value for current row.
124    pub fn current_value(&self) -> Result<Field> {
125        if self.is_null() {
126            return Ok(Field::Null);
127        }
128        let field = match *self {
129            TripletIter::BoolTripletIter(ref typed) => {
130                Field::convert_bool(typed.column_descr(), *typed.current_value())
131            }
132            TripletIter::Int32TripletIter(ref typed) => {
133                Field::convert_int32(typed.column_descr(), *typed.current_value())
134            }
135            TripletIter::Int64TripletIter(ref typed) => {
136                Field::convert_int64(typed.column_descr(), *typed.current_value())
137            }
138            TripletIter::Int96TripletIter(ref typed) => {
139                Field::convert_int96(typed.column_descr(), *typed.current_value())
140            }
141            TripletIter::FloatTripletIter(ref typed) => {
142                Field::convert_float(typed.column_descr(), *typed.current_value())
143            }
144            TripletIter::DoubleTripletIter(ref typed) => {
145                Field::convert_double(typed.column_descr(), *typed.current_value())
146            }
147            TripletIter::ByteArrayTripletIter(ref typed) => {
148                Field::convert_byte_array(typed.column_descr(), typed.current_value().clone())?
149            }
150            TripletIter::FixedLenByteArrayTripletIter(ref typed) => Field::convert_byte_array(
151                typed.column_descr(),
152                typed.current_value().clone().into(),
153            )?,
154        };
155        Ok(field)
156    }
157}
158
159/// Internal typed triplet iterator as a wrapper for column reader
160/// (primitive leaf column), provides per-element access.
161pub struct TypedTripletIter<T: DataType> {
162    reader: ColumnReaderImpl<T>,
163    column_descr: ColumnDescPtr,
164    batch_size: usize,
165    // type properties
166    max_def_level: i16,
167    max_rep_level: i16,
168    // values and levels
169    values: Vec<T::T>,
170    def_levels: Option<Vec<i16>>,
171    rep_levels: Option<Vec<i16>>,
172    // current index for the triplet (value, def, rep)
173    curr_triplet_index: usize,
174    // how many triplets are left before we need to buffer
175    triplets_left: usize,
176    // helper flag to quickly check if we have more values/levels to read
177    has_next: bool,
178}
179
180impl<T: DataType> TypedTripletIter<T> {
181    /// Creates new typed triplet iterator based on provided column reader.
182    /// Use batch size to specify the amount of values to buffer from column reader.
183    fn new(descr: ColumnDescPtr, batch_size: usize, column_reader: ColumnReader) -> Self {
184        assert!(
185            batch_size > 0,
186            "Expected positive batch size, found: {batch_size}"
187        );
188
189        let max_def_level = descr.max_def_level();
190        let max_rep_level = descr.max_rep_level();
191
192        let def_levels = if max_def_level == 0 {
193            None
194        } else {
195            Some(vec![0; batch_size])
196        };
197        let rep_levels = if max_rep_level == 0 {
198            None
199        } else {
200            Some(vec![0; batch_size])
201        };
202
203        Self {
204            reader: get_typed_column_reader(column_reader),
205            column_descr: descr,
206            batch_size,
207            max_def_level,
208            max_rep_level,
209            values: vec![T::T::default(); batch_size],
210            def_levels,
211            rep_levels,
212            curr_triplet_index: 0,
213            triplets_left: 0,
214            has_next: false,
215        }
216    }
217
218    /// Returns column descriptor reference for the current typed triplet iterator.
219    #[inline]
220    pub fn column_descr(&self) -> &ColumnDescPtr {
221        &self.column_descr
222    }
223
224    /// Returns maximum definition level for the triplet iterator (leaf column).
225    #[inline]
226    fn max_def_level(&self) -> i16 {
227        self.max_def_level
228    }
229
230    /// Returns maximum repetition level for the triplet iterator (leaf column).
231    #[inline]
232    fn max_rep_level(&self) -> i16 {
233        self.max_rep_level
234    }
235
236    /// Returns current value.
237    /// Method does not advance the iterator, therefore can be called multiple times.
238    #[inline]
239    fn current_value(&self) -> &T::T {
240        assert!(
241            self.current_def_level() == self.max_def_level(),
242            "Cannot extract value, max definition level: {}, current level: {}",
243            self.max_def_level(),
244            self.current_def_level()
245        );
246        &self.values[self.curr_triplet_index]
247    }
248
249    /// Returns current definition level.
250    /// If field is required, then maximum definition level is returned.
251    #[inline]
252    fn current_def_level(&self) -> i16 {
253        if !self.has_next {
254            return 0;
255        }
256        match self.def_levels {
257            Some(ref vec) => vec[self.curr_triplet_index],
258            None => self.max_def_level,
259        }
260    }
261
262    /// Returns current repetition level.
263    /// If field is required, then maximum repetition level is returned.
264    #[inline]
265    fn current_rep_level(&self) -> i16 {
266        if !self.has_next {
267            return 0;
268        }
269        match self.rep_levels {
270            Some(ref vec) => vec[self.curr_triplet_index],
271            None => self.max_rep_level,
272        }
273    }
274
275    /// Quick check if iterator has more values/levels to read.
276    /// It is updated as a result of `read_next` method, so they are synchronized.
277    #[inline]
278    fn has_next(&self) -> bool {
279        self.has_next
280    }
281
282    /// Advances to the next triplet.
283    /// Returns true, if there are more records to read, false there are no records left.
284    fn read_next(&mut self) -> Result<bool> {
285        self.curr_triplet_index += 1;
286
287        // A loop is required to handle the case of a batch size of 1, as in such a case
288        // on reaching the end of a record, read_records will return `Ok((1, 0, 0))`
289        // and therefore not advance `self.triplets_left`
290        while self.curr_triplet_index >= self.triplets_left {
291            let (records_read, values_read, levels_read) = {
292                self.values.clear();
293                if let Some(x) = &mut self.def_levels {
294                    x.clear()
295                }
296                if let Some(x) = &mut self.rep_levels {
297                    x.clear()
298                }
299
300                // Buffer triplets
301                self.reader.read_records(
302                    self.batch_size,
303                    self.def_levels.as_mut(),
304                    self.rep_levels.as_mut(),
305                    &mut self.values,
306                )?
307            };
308
309            // No more values or levels to read
310            if records_read == 0 && values_read == 0 && levels_read == 0 {
311                self.curr_triplet_index = 0;
312                self.has_next = false;
313                return Ok(false);
314            }
315
316            // We never read values more than levels
317            if levels_read == 0 || values_read == levels_read {
318                // There are no definition levels to read, column is required
319                // or definition levels match values, so it does not require spacing
320                self.curr_triplet_index = 0;
321                self.triplets_left = values_read;
322            } else if values_read < levels_read {
323                // Add spacing for triplets.
324                // The idea is setting values for positions in def_levels when current
325                // definition level equals to maximum definition level.
326                // Values and levels are guaranteed to line up, because of
327                // the column reader method.
328
329                // Note: if values_read == 0, then spacing will not be triggered
330                let mut idx = values_read;
331                let def_levels = self.def_levels.as_ref().unwrap();
332                self.values.resize(levels_read, T::T::default());
333                for i in 0..levels_read {
334                    if def_levels[levels_read - i - 1] == self.max_def_level {
335                        idx -= 1; // This is done to avoid usize becoming a negative value
336                        self.values.swap(levels_read - i - 1, idx);
337                    }
338                }
339                self.curr_triplet_index = 0;
340                self.triplets_left = levels_read;
341            } else {
342                return Err(general_err!(
343                    "Spacing of values/levels is wrong, values_read: {}, levels_read: {}",
344                    values_read,
345                    levels_read
346                ));
347            }
348        }
349
350        self.has_next = true;
351        Ok(true)
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    use crate::file::reader::{FileReader, SerializedFileReader};
360    use crate::schema::types::ColumnPath;
361    use crate::util::test_common::file_util::get_test_file;
362
363    #[test]
364    #[should_panic(expected = "Expected positive batch size, found: 0")]
365    fn test_triplet_zero_batch_size() {
366        let column_path = ColumnPath::from(vec!["b_struct".to_string(), "b_c_int".to_string()]);
367        test_column_in_file("nulls.snappy.parquet", 0, &column_path, &[], &[], &[]);
368    }
369
370    #[test]
371    fn test_triplet_null_column() {
372        let path = vec!["b_struct", "b_c_int"];
373        let values = vec![];
374        let def_levels = vec![1, 1, 1, 1, 1, 1, 1, 1];
375        let rep_levels = vec![0, 0, 0, 0, 0, 0, 0, 0];
376        test_triplet_iter(
377            "nulls.snappy.parquet",
378            path,
379            &values,
380            &def_levels,
381            &rep_levels,
382        );
383    }
384
385    #[test]
386    #[cfg_attr(miri, ignore)] // Takes too long
387    fn test_triplet_required_column() {
388        let path = vec!["ID"];
389        let values = vec![Field::Long(8)];
390        let def_levels = vec![0];
391        let rep_levels = vec![0];
392        test_triplet_iter(
393            "nonnullable.impala.parquet",
394            path,
395            &values,
396            &def_levels,
397            &rep_levels,
398        );
399    }
400
401    #[test]
402    #[cfg_attr(miri, ignore)] // Takes too long
403    fn test_triplet_optional_column() {
404        let path = vec!["nested_struct", "A"];
405        let values = vec![Field::Int(1), Field::Int(7)];
406        let def_levels = vec![2, 1, 1, 1, 1, 0, 2];
407        let rep_levels = vec![0, 0, 0, 0, 0, 0, 0];
408        test_triplet_iter(
409            "nullable.impala.parquet",
410            path,
411            &values,
412            &def_levels,
413            &rep_levels,
414        );
415    }
416
417    #[test]
418    #[cfg_attr(miri, ignore)] // Takes too long
419    fn test_triplet_optional_list_column() {
420        let path = vec!["a", "list", "element", "list", "element", "list", "element"];
421        let values = vec![
422            Field::Str("a".to_string()),
423            Field::Str("b".to_string()),
424            Field::Str("c".to_string()),
425            Field::Str("d".to_string()),
426            Field::Str("a".to_string()),
427            Field::Str("b".to_string()),
428            Field::Str("c".to_string()),
429            Field::Str("d".to_string()),
430            Field::Str("e".to_string()),
431            Field::Str("a".to_string()),
432            Field::Str("b".to_string()),
433            Field::Str("c".to_string()),
434            Field::Str("d".to_string()),
435            Field::Str("e".to_string()),
436            Field::Str("f".to_string()),
437        ];
438        let def_levels = vec![7, 7, 7, 4, 7, 7, 7, 7, 7, 4, 7, 7, 7, 7, 7, 7, 4, 7];
439        let rep_levels = vec![0, 3, 2, 1, 2, 0, 3, 2, 3, 1, 2, 0, 3, 2, 3, 2, 1, 2];
440        test_triplet_iter(
441            "nested_lists.snappy.parquet",
442            path,
443            &values,
444            &def_levels,
445            &rep_levels,
446        );
447    }
448
449    #[test]
450    #[cfg_attr(miri, ignore)] // Takes too long
451    fn test_triplet_optional_map_column() {
452        let path = vec!["a", "key_value", "value", "key_value", "key"];
453        let values = vec![
454            Field::Int(1),
455            Field::Int(2),
456            Field::Int(1),
457            Field::Int(1),
458            Field::Int(3),
459            Field::Int(4),
460            Field::Int(5),
461        ];
462        let def_levels = vec![4, 4, 4, 2, 3, 4, 4, 4, 4];
463        let rep_levels = vec![0, 2, 0, 0, 0, 0, 0, 2, 2];
464        test_triplet_iter(
465            "nested_maps.snappy.parquet",
466            path,
467            &values,
468            &def_levels,
469            &rep_levels,
470        );
471    }
472
473    // Check triplet iterator across different batch sizes
474    fn test_triplet_iter(
475        file_name: &str,
476        column_path: Vec<&str>,
477        expected_values: &[Field],
478        expected_def_levels: &[i16],
479        expected_rep_levels: &[i16],
480    ) {
481        // Convert path into column path
482        let path: Vec<String> = column_path.iter().map(|x| x.to_string()).collect();
483        let column_path = ColumnPath::from(path);
484
485        let batch_sizes = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 128, 256];
486        for batch_size in batch_sizes {
487            test_column_in_file(
488                file_name,
489                batch_size,
490                &column_path,
491                expected_values,
492                expected_def_levels,
493                expected_rep_levels,
494            );
495        }
496    }
497
498    // Check values of a selected column in a file
499    fn test_column_in_file(
500        file_name: &str,
501        batch_size: usize,
502        column_path: &ColumnPath,
503        expected_values: &[Field],
504        expected_def_levels: &[i16],
505        expected_rep_levels: &[i16],
506    ) {
507        let file = get_test_file(file_name);
508        let file_reader = SerializedFileReader::new(file).unwrap();
509        let metadata = file_reader.metadata();
510        // Get schema descriptor
511        let file_metadata = metadata.file_metadata();
512        let schema = file_metadata.schema_descr();
513        // Get first row group
514        let row_group_reader = file_reader.get_row_group(0).unwrap();
515
516        for i in 0..schema.num_columns() {
517            let descr = schema.column(i);
518            if descr.path() == column_path {
519                let reader = row_group_reader.get_column_reader(i).unwrap();
520                test_triplet_column(
521                    descr,
522                    reader,
523                    batch_size,
524                    expected_values,
525                    expected_def_levels,
526                    expected_rep_levels,
527                );
528            }
529        }
530    }
531
532    // Check values for individual triplet iterator
533    fn test_triplet_column(
534        descr: ColumnDescPtr,
535        reader: ColumnReader,
536        batch_size: usize,
537        expected_values: &[Field],
538        expected_def_levels: &[i16],
539        expected_rep_levels: &[i16],
540    ) {
541        let mut iter = TripletIter::new(descr.clone(), reader, batch_size);
542        let mut values: Vec<Field> = Vec::new();
543        let mut def_levels: Vec<i16> = Vec::new();
544        let mut rep_levels: Vec<i16> = Vec::new();
545
546        assert_eq!(iter.max_def_level(), descr.max_def_level());
547        assert_eq!(iter.max_rep_level(), descr.max_rep_level());
548
549        while matches!(iter.read_next(), Ok(true)) {
550            assert!(iter.has_next());
551            if !iter.is_null() {
552                values.push(iter.current_value().unwrap());
553            }
554            def_levels.push(iter.current_def_level());
555            rep_levels.push(iter.current_rep_level());
556        }
557
558        assert_eq!(values, expected_values);
559        assert_eq!(def_levels, expected_def_levels);
560        assert_eq!(rep_levels, expected_rep_levels);
561    }
562
563    fn open_triplet_iter(file_name: &str, path: &[&str], batch_size: usize) -> TripletIter {
564        let column_path = ColumnPath::from(path.iter().map(|x| x.to_string()).collect::<Vec<_>>());
565        let file = get_test_file(file_name);
566        let file_reader = SerializedFileReader::new(file).unwrap();
567        let metadata = file_reader.metadata();
568        let schema = metadata.file_metadata().schema_descr();
569        let row_group_reader = file_reader.get_row_group(0).unwrap();
570        for i in 0..schema.num_columns() {
571            let descr = schema.column(i);
572            if descr.path() == &column_path {
573                let reader = row_group_reader.get_column_reader(i).unwrap();
574                return TripletIter::new(descr.clone(), reader, batch_size);
575            }
576        }
577        panic!("Column {column_path:?} not found in {file_name}");
578    }
579
580    #[test]
581    fn test_current_def_level_safe_after_exhaustion() {
582        let mut iter = open_triplet_iter("nulls.snappy.parquet", &["b_struct", "b_c_int"], 256);
583        while matches!(iter.read_next(), Ok(true)) {}
584        assert!(!iter.has_next());
585        assert_eq!(iter.current_def_level(), 0);
586    }
587
588    #[test]
589    fn test_current_rep_level_safe_after_exhaustion() {
590        let mut iter = open_triplet_iter(
591            "nested_lists.snappy.parquet",
592            &["a", "list", "element", "list", "element", "list", "element"],
593            256,
594        );
595        while matches!(iter.read_next(), Ok(true)) {}
596        assert!(!iter.has_next());
597        assert_eq!(iter.current_rep_level(), 0);
598    }
599}
parquet/record/triplet.rs

parquet/record/
triplet.rs
