parquet/record/

reader.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Contains implementation of record assembly and converting Parquet types into
//! [`Row`]s.

use std::{collections::HashMap, fmt, sync::Arc};

use crate::basic::{ConvertedType, Repetition};
use crate::errors::{ParquetError, Result};
use crate::file::reader::{FileReader, RowGroupReader};
use crate::record::{
    api::{Field, Row, make_list, make_map},
    triplet::TripletIter,
};
use crate::schema::types::{ColumnPath, SchemaDescPtr, SchemaDescriptor, Type, TypePtr};

/// Default batch size for a reader
const DEFAULT_BATCH_SIZE: usize = 1024;

/// Tree builder for `Reader` enum.
/// Serves as a container of options for building a reader tree and a builder, and
/// accessing a records iterator [`RowIter`].
pub struct TreeBuilder {
    // Batch size (>= 1) for triplet iterators
    batch_size: usize,
}

impl Default for TreeBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl TreeBuilder {
    /// Creates new tree builder with default parameters.
    pub fn new() -> Self {
        Self {
            batch_size: DEFAULT_BATCH_SIZE,
        }
    }

    /// Sets batch size for this tree builder.
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// Creates new root reader for provided schema and row group.
    pub fn build(
        &self,
        descr: SchemaDescPtr,
        row_group_reader: &dyn RowGroupReader,
    ) -> Result<Reader> {
        // Prepare lookup table of column path -> original column index
        // This allows to prune columns and map schema leaf nodes to the column readers
        let mut paths: HashMap<ColumnPath, usize> = HashMap::new();
        let row_group_metadata = row_group_reader.metadata();

        for col_index in 0..row_group_reader.num_columns() {
            let col_meta = row_group_metadata.column(col_index);
            let col_path = col_meta.column_path().clone();
            paths.insert(col_path, col_index);
        }

        // Build child readers for the message type
        let mut readers = Vec::new();
        let mut path = Vec::new();

        for field in descr.root_schema().get_fields() {
            let reader =
                self.reader_tree(field.clone(), &mut path, 0, 0, &paths, row_group_reader)?;
            readers.push(reader);
        }

        // Return group reader for message type,
        // it is always required with definition level 0
        Ok(Reader::GroupReader(None, 0, readers))
    }

    /// Creates iterator of `Row`s directly from schema descriptor and row group.
    pub fn as_iter(
        &self,
        descr: SchemaDescPtr,
        row_group_reader: &dyn RowGroupReader,
    ) -> Result<ReaderIter> {
        let num_records = row_group_reader.metadata().num_rows() as usize;
        ReaderIter::new(self.build(descr, row_group_reader)?, num_records)
    }

    /// Builds tree of readers for the current schema recursively.
    fn reader_tree(
        &self,
        field: TypePtr,
        path: &mut Vec<String>,
        mut curr_def_level: i16,
        mut curr_rep_level: i16,
        paths: &HashMap<ColumnPath, usize>,
        row_group_reader: &dyn RowGroupReader,
    ) -> Result<Reader> {
        assert!(field.get_basic_info().has_repetition());
        // Update current definition and repetition levels for this type
        let repetition = field.get_basic_info().repetition();
        match repetition {
            Repetition::OPTIONAL => {
                curr_def_level += 1;
            }
            Repetition::REPEATED => {
                curr_def_level += 1;
                curr_rep_level += 1;
            }
            _ => {}
        }

        path.push(String::from(field.name()));
        let reader = if field.is_primitive() {
            let col_path = ColumnPath::new(path.to_vec());
            let orig_index = *paths
                .get(&col_path)
                .ok_or(general_err!("Path {:?} not found", col_path))?;
            let col_descr = row_group_reader
                .metadata()
                .column(orig_index)
                .column_descr_ptr();
            let col_reader = row_group_reader.get_column_reader(orig_index)?;
            let column = TripletIter::new(col_descr, col_reader, self.batch_size);
            let reader = Reader::PrimitiveReader(field.clone(), Box::new(column));
            if repetition == Repetition::REPEATED {
                Reader::RepeatedReader(
                    field,
                    curr_def_level - 1,
                    curr_rep_level - 1,
                    Box::new(reader),
                )
            } else {
                reader
            }
        } else {
            match field.get_basic_info().converted_type() {
                // List types
                ConvertedType::LIST => {
                    assert_eq!(field.get_fields().len(), 1, "Invalid list type {field:?}");

                    let repeated_field = field.get_fields()[0].clone();
                    assert_eq!(
                        repeated_field.get_basic_info().repetition(),
                        Repetition::REPEATED,
                        "Invalid list type {field:?}"
                    );

                    if Reader::is_element_type(&repeated_field) {
                        // Support for backward compatible lists
                        let reader = self.reader_tree(
                            repeated_field,
                            path,
                            curr_def_level,
                            curr_rep_level,
                            paths,
                            row_group_reader,
                        )?;

                        Reader::RepeatedReader(
                            field,
                            curr_def_level,
                            curr_rep_level,
                            Box::new(reader),
                        )
                    } else {
                        let child_field = repeated_field.get_fields()[0].clone();

                        path.push(String::from(repeated_field.name()));

                        let reader = self.reader_tree(
                            child_field,
                            path,
                            curr_def_level + 1,
                            curr_rep_level + 1,
                            paths,
                            row_group_reader,
                        )?;

                        path.pop();

                        Reader::RepeatedReader(
                            field,
                            curr_def_level,
                            curr_rep_level,
                            Box::new(reader),
                        )
                    }
                }
                // Map types (key-value pairs)
                ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => {
                    assert_eq!(field.get_fields().len(), 1, "Invalid map type: {field:?}");
                    assert!(
                        !field.get_fields()[0].is_primitive(),
                        "Invalid map type: {field:?}"
                    );

                    let key_value_type = field.get_fields()[0].clone();
                    assert_eq!(
                        key_value_type.get_basic_info().repetition(),
                        Repetition::REPEATED,
                        "Invalid map type: {field:?}"
                    );
                    // Parquet spec allows no value. In that case treat as a list. #1642
                    if key_value_type.get_fields().len() != 1 {
                        // If not a list, then there can only be 2 fields in the struct
                        assert_eq!(
                            key_value_type.get_fields().len(),
                            2,
                            "Invalid map type: {field:?}"
                        );
                    }

                    path.push(String::from(key_value_type.name()));

                    let key_type = &key_value_type.get_fields()[0];
                    assert!(
                        key_type.is_primitive(),
                        "Map key type is expected to be a primitive type, but found {key_type:?}"
                    );
                    let key_reader = self.reader_tree(
                        key_type.clone(),
                        path,
                        curr_def_level + 1,
                        curr_rep_level + 1,
                        paths,
                        row_group_reader,
                    )?;

                    if key_value_type.get_fields().len() == 1 {
                        path.pop();
                        Reader::RepeatedReader(
                            field,
                            curr_def_level,
                            curr_rep_level,
                            Box::new(key_reader),
                        )
                    } else {
                        let value_type = &key_value_type.get_fields()[1];
                        let value_reader = self.reader_tree(
                            value_type.clone(),
                            path,
                            curr_def_level + 1,
                            curr_rep_level + 1,
                            paths,
                            row_group_reader,
                        )?;

                        path.pop();

                        Reader::KeyValueReader(
                            field,
                            curr_def_level,
                            curr_rep_level,
                            Box::new(key_reader),
                            Box::new(value_reader),
                        )
                    }
                }
                // A repeated field that is neither contained by a `LIST`- or
                // `MAP`-annotated group nor annotated by `LIST` or `MAP`
                // should be interpreted as a required list of required
                // elements where the element type is the type of the field.
                _ if repetition == Repetition::REPEATED => {
                    let required_field = Type::group_type_builder(field.name())
                        .with_repetition(Repetition::REQUIRED)
                        .with_converted_type(field.get_basic_info().converted_type())
                        .with_fields(field.get_fields().to_vec())
                        .build()?;

                    path.pop();

                    let reader = self.reader_tree(
                        Arc::new(required_field),
                        path,
                        curr_def_level,
                        curr_rep_level,
                        paths,
                        row_group_reader,
                    )?;

                    return Ok(Reader::RepeatedReader(
                        field,
                        curr_def_level - 1,
                        curr_rep_level - 1,
                        Box::new(reader),
                    ));
                }
                // Group types (structs)
                _ => {
                    let mut readers = Vec::new();
                    for child in field.get_fields() {
                        let reader = self.reader_tree(
                            child.clone(),
                            path,
                            curr_def_level,
                            curr_rep_level,
                            paths,
                            row_group_reader,
                        )?;
                        readers.push(reader);
                    }
                    Reader::GroupReader(Some(field), curr_def_level, readers)
                }
            }
        };
        path.pop();

        Ok(Reader::option(repetition, curr_def_level, reader))
    }
}

/// Reader tree for record assembly
pub enum Reader {
    /// Primitive reader with type information and triplet iterator
    PrimitiveReader(TypePtr, Box<TripletIter>),
    /// Optional reader with definition level of a parent and a reader
    OptionReader(i16, Box<Reader>),
    /// Group (struct) reader with type information, definition level and list of child
    /// readers. When it represents message type, type information is None
    GroupReader(Option<TypePtr>, i16, Vec<Reader>),
    /// Reader for repeated values, e.g. lists, contains type information, definition
    /// level, repetition level and a child reader
    RepeatedReader(TypePtr, i16, i16, Box<Reader>),
    /// Reader of key-value pairs, e.g. maps, contains type information, definition
    /// level, repetition level, child reader for keys and child reader for values
    KeyValueReader(TypePtr, i16, i16, Box<Reader>, Box<Reader>),
}

impl Reader {
    /// Wraps reader in option reader based on repetition.
    fn option(repetition: Repetition, def_level: i16, reader: Reader) -> Self {
        if repetition == Repetition::OPTIONAL {
            Reader::OptionReader(def_level - 1, Box::new(reader))
        } else {
            reader
        }
    }

    /// Returns true if repeated type is an element type for the list.
    /// Used to determine legacy list types.
    /// This method is copied from Spark Parquet reader and is based on the reference:
    /// <https://github.com/apache/parquet-format/blob/master/LogicalTypes.md>
    ///   #backward-compatibility-rules
    fn is_element_type(repeated_type: &Type) -> bool {
        // For legacy 2-level list types whose element type is a 2-level list
        //
        //    // ARRAY<ARRAY<INT>> (nullable list, non-null elements)
        //    optional group my_list (LIST) {
        //      repeated group array (LIST) {
        //        repeated int32 array;
        //      };
        //    }
        //
        if repeated_type.is_list() || repeated_type.has_single_repeated_child() {
            return false;
        }

        // For legacy 2-level list types with primitive element type, e.g.:
        //
        //    // ARRAY<INT> (nullable list, non-null elements)
        //    optional group my_list (LIST) {
        //      repeated int32 element;
        //    }
        //
        repeated_type.is_primitive() ||
    // For legacy 2-level list types whose element type is a group type with 2 or more
    // fields, e.g.:
    //
    //    // ARRAY<STRUCT<str: STRING, num: INT>> (nullable list, non-null elements)
    //    optional group my_list (LIST) {
    //      repeated group element {
    //        required binary str (UTF8);
    //        required int32 num;
    //      };
    //    }
    //
    repeated_type.is_group() && repeated_type.get_fields().len() > 1 ||
    // For legacy 2-level list types generated by parquet-avro (Parquet version < 1.6.0),
    // e.g.:
    //
    //    // ARRAY<STRUCT<str: STRING>> (nullable list, non-null elements)
    //    optional group my_list (LIST) {
    //      repeated group array {
    //        required binary str (UTF8);
    //      };
    //    }
    //
    repeated_type.name() == "array" ||
    // For Parquet data generated by parquet-thrift, e.g.:
    //
    //    // ARRAY<STRUCT<str: STRING>> (nullable list, non-null elements)
    //    optional group my_list (LIST) {
    //      repeated group my_list_tuple {
    //        required binary str (UTF8);
    //      };
    //    }
    //
    repeated_type.name().ends_with("_tuple")
    }

    /// Reads current record as `Row` from the reader tree.
    /// Automatically advances all necessary readers.
    /// This must be called on the root level reader (i.e., for Message type).
    /// Otherwise, it will panic.
    fn read(&mut self) -> Result<Row> {
        match *self {
            Reader::GroupReader(_, _, ref mut readers) => {
                let mut fields = Vec::new();
                for reader in readers {
                    fields.push((String::from(reader.field_name()), reader.read_field()?));
                }
                Ok(Row::new(fields))
            }
            _ => panic!("Cannot call read() on {self}"),
        }
    }

    /// Reads current record as `Field` from the reader tree.
    /// Automatically advances all necessary readers.
    fn read_field(&mut self) -> Result<Field> {
        let field = match *self {
            Reader::PrimitiveReader(_, ref mut column) => {
                if !column.has_next() {
                    return Err(general_err!("Unexpected end of column data"));
                }
                let value = column.current_value()?;
                column.read_next()?;
                value
            }
            Reader::OptionReader(def_level, ref mut reader) => {
                if !reader.has_next() {
                    return Err(general_err!("Unexpected end of column data"));
                }
                if reader.current_def_level() > def_level {
                    reader.read_field()?
                } else {
                    reader.advance_columns()?;
                    Field::Null
                }
            }
            Reader::GroupReader(_, def_level, ref mut readers) => {
                let mut fields = Vec::new();
                for reader in readers {
                    if reader.repetition() != Repetition::OPTIONAL
                        || reader.current_def_level() > def_level
                    {
                        fields.push((String::from(reader.field_name()), reader.read_field()?));
                    } else {
                        reader.advance_columns()?;
                        fields.push((String::from(reader.field_name()), Field::Null));
                    }
                }
                let row = Row::new(fields);
                Field::Group(row)
            }
            Reader::RepeatedReader(_, def_level, rep_level, ref mut reader) => {
                if !reader.has_next() {
                    return Err(general_err!("Unexpected end of column data"));
                }
                let mut elements = Vec::new();
                loop {
                    if reader.current_def_level() > def_level {
                        elements.push(reader.read_field()?);
                    } else {
                        reader.advance_columns()?;
                        // If the current definition level is equal to the definition
                        // level of this repeated type, then the
                        // result is an empty list and the repetition level
                        // will always be <= rl.
                        break;
                    }

                    // This covers case when we are out of repetition levels and should
                    // close the group, or there are no values left to
                    // buffer.
                    if !reader.has_next() || reader.current_rep_level() <= rep_level {
                        break;
                    }
                }
                Field::ListInternal(make_list(elements))
            }
            Reader::KeyValueReader(_, def_level, rep_level, ref mut keys, ref mut values) => {
                if !keys.has_next() {
                    return Err(general_err!("Unexpected end of column data"));
                }
                let mut pairs = Vec::new();
                loop {
                    if keys.current_def_level() > def_level {
                        pairs.push((keys.read_field()?, values.read_field()?));
                    } else {
                        keys.advance_columns()?;
                        values.advance_columns()?;
                        // If the current definition level is equal to the definition
                        // level of this repeated type, then the
                        // result is an empty list and the repetition level
                        // will always be <= rl.
                        break;
                    }

                    // This covers case when we are out of repetition levels and should
                    // close the group, or there are no values left to
                    // buffer.
                    if !keys.has_next() || keys.current_rep_level() <= rep_level {
                        break;
                    }
                }

                Field::MapInternal(make_map(pairs))
            }
        };
        Ok(field)
    }

    /// Returns field name for the current reader.
    fn field_name(&self) -> &str {
        match *self {
            Reader::PrimitiveReader(ref field, _) => field.name(),
            Reader::OptionReader(_, ref reader) => reader.field_name(),
            Reader::GroupReader(ref opt, ..) => match opt {
                Some(field) => field.name(),
                None => panic!("Field is None for group reader"),
            },
            Reader::RepeatedReader(ref field, ..) => field.name(),
            Reader::KeyValueReader(ref field, ..) => field.name(),
        }
    }

    /// Returns repetition for the current reader.
    fn repetition(&self) -> Repetition {
        match *self {
            Reader::PrimitiveReader(ref field, _) => field.get_basic_info().repetition(),
            Reader::OptionReader(_, ref reader) => reader.repetition(),
            Reader::GroupReader(ref opt, ..) => match opt {
                Some(field) => field.get_basic_info().repetition(),
                None => panic!("Field is None for group reader"),
            },
            Reader::RepeatedReader(ref field, ..) => field.get_basic_info().repetition(),
            Reader::KeyValueReader(ref field, ..) => field.get_basic_info().repetition(),
        }
    }

    /// Returns true, if current reader has more values, false otherwise.
    /// Method does not advance internal iterator.
    fn has_next(&self) -> bool {
        match *self {
            Reader::PrimitiveReader(_, ref column) => column.has_next(),
            Reader::OptionReader(_, ref reader) => reader.has_next(),
            Reader::GroupReader(_, _, ref readers) => readers.first().unwrap().has_next(),
            Reader::RepeatedReader(_, _, _, ref reader) => reader.has_next(),
            Reader::KeyValueReader(_, _, _, ref keys, _) => keys.has_next(),
        }
    }

    /// Returns current definition level,
    /// Method does not advance internal iterator.
    fn current_def_level(&self) -> i16 {
        match *self {
            Reader::PrimitiveReader(_, ref column) => column.current_def_level(),
            Reader::OptionReader(_, ref reader) => reader.current_def_level(),
            Reader::GroupReader(_, _, ref readers) => match readers.first() {
                Some(reader) => reader.current_def_level(),
                None => panic!("Current definition level: empty group reader"),
            },
            Reader::RepeatedReader(_, _, _, ref reader) => reader.current_def_level(),
            Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_def_level(),
        }
    }

    /// Returns current repetition level.
    /// Method does not advance internal iterator.
    fn current_rep_level(&self) -> i16 {
        match *self {
            Reader::PrimitiveReader(_, ref column) => column.current_rep_level(),
            Reader::OptionReader(_, ref reader) => reader.current_rep_level(),
            Reader::GroupReader(_, _, ref readers) => match readers.first() {
                Some(reader) => reader.current_rep_level(),
                None => panic!("Current repetition level: empty group reader"),
            },
            Reader::RepeatedReader(_, _, _, ref reader) => reader.current_rep_level(),
            Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_rep_level(),
        }
    }

    /// Advances leaf columns for the current reader.
    fn advance_columns(&mut self) -> Result<()> {
        match *self {
            Reader::PrimitiveReader(_, ref mut column) => column.read_next().map(|_| ()),
            Reader::OptionReader(_, ref mut reader) => reader.advance_columns(),
            Reader::GroupReader(_, _, ref mut readers) => {
                for reader in readers {
                    reader.advance_columns()?;
                }
                Ok(())
            }
            Reader::RepeatedReader(_, _, _, ref mut reader) => reader.advance_columns(),
            Reader::KeyValueReader(_, _, _, ref mut keys, ref mut values) => {
                keys.advance_columns()?;
                values.advance_columns()
            }
        }
    }
}

impl fmt::Display for Reader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = match self {
            Reader::PrimitiveReader(..) => "PrimitiveReader",
            Reader::OptionReader(..) => "OptionReader",
            Reader::GroupReader(..) => "GroupReader",
            Reader::RepeatedReader(..) => "RepeatedReader",
            Reader::KeyValueReader(..) => "KeyValueReader",
        };
        write!(f, "{s}")
    }
}

// ----------------------------------------------------------------------
// Row iterators

/// The enum Either with variants That represents a reference and a box of
/// [`FileReader`].
enum Either<'a> {
    Left(&'a dyn FileReader),
    Right(Box<dyn FileReader>),
}

impl Either<'_> {
    fn reader(&self) -> &dyn FileReader {
        match *self {
            Either::Left(r) => r,
            Either::Right(ref r) => &**r,
        }
    }
}

/// Access parquet data as an iterator of [`Row`]
///
/// # Caveats
///
/// Parquet stores data in a columnar fashion using [Dremel] encoding, and is therefore highly
/// optimised for reading data by column, not row. As a consequence applications concerned with
/// performance should prefer the columnar arrow or [ColumnReader] APIs.
///
/// Additionally the current implementation does not correctly handle repeated fields ([#2394]),
/// and workloads looking to handle such schema should use the other APIs.
///
/// [#2394]: https://github.com/apache/arrow-rs/issues/2394
/// [ColumnReader]: crate::file::reader::RowGroupReader::get_column_reader
/// [Dremel]: https://research.google/pubs/pub36632/
pub struct RowIter<'a> {
    descr: SchemaDescPtr,
    tree_builder: TreeBuilder,
    file_reader: Option<Either<'a>>,
    current_row_group: usize,
    num_row_groups: usize,
    row_iter: Option<ReaderIter>,
}

impl<'a> RowIter<'a> {
    /// Creates a new iterator of [`Row`]s.
    fn new(
        file_reader: Option<Either<'a>>,
        row_iter: Option<ReaderIter>,
        descr: SchemaDescPtr,
    ) -> Self {
        let tree_builder = Self::tree_builder();
        let num_row_groups = match file_reader {
            Some(ref r) => r.reader().num_row_groups(),
            None => 0,
        };

        Self {
            descr,
            file_reader,
            tree_builder,
            num_row_groups,
            row_iter,
            current_row_group: 0,
        }
    }

    /// Creates iterator of [`Row`]s for all row groups in a
    /// file.
    pub fn from_file(proj: Option<Type>, reader: &'a dyn FileReader) -> Result<Self> {
        let either = Either::Left(reader);
        let descr =
            Self::get_proj_descr(proj, reader.metadata().file_metadata().schema_descr_ptr())?;

        Ok(Self::new(Some(either), None, descr))
    }

    /// Creates iterator of [`Row`]s for a specific row group.
    pub fn from_row_group(proj: Option<Type>, reader: &'a dyn RowGroupReader) -> Result<Self> {
        let descr = Self::get_proj_descr(proj, reader.metadata().schema_descr_ptr())?;
        let tree_builder = Self::tree_builder();
        let row_iter = tree_builder.as_iter(descr.clone(), reader)?;

        // For row group we need to set `current_row_group` >= `num_row_groups`, because
        // we only have one row group and can't buffer more.
        Ok(Self::new(None, Some(row_iter), descr))
    }

    /// Creates a iterator of [`Row`]s from a [`FileReader`] using the full file schema.
    pub fn from_file_into(reader: Box<dyn FileReader>) -> Self {
        let either = Either::Right(reader);
        let descr = either
            .reader()
            .metadata()
            .file_metadata()
            .schema_descr_ptr();

        Self::new(Some(either), None, descr)
    }

    /// Tries to create a iterator of [`Row`]s using projections.
    /// Returns a error if a file reader is not the source of this iterator.
    ///
    /// The Projected schema can be a subset of or equal to the file schema,
    /// when it is None, full file schema is assumed.
    pub fn project(self, proj: Option<Type>) -> Result<Self> {
        match self.file_reader {
            Some(ref either) => {
                let schema = either
                    .reader()
                    .metadata()
                    .file_metadata()
                    .schema_descr_ptr();
                let descr = Self::get_proj_descr(proj, schema)?;

                Ok(Self::new(self.file_reader, None, descr))
            }
            None => Err(general_err!("File reader is required to use projections")),
        }
    }

    /// Helper method to get schema descriptor for projected schema.
    /// If projection is None, then full schema is returned.
    #[inline]
    fn get_proj_descr(proj: Option<Type>, root_descr: SchemaDescPtr) -> Result<SchemaDescPtr> {
        match proj {
            Some(projection) => {
                // check if projection is part of file schema
                let root_schema = root_descr.root_schema();
                if !root_schema.check_contains(&projection) {
                    return Err(general_err!("Root schema does not contain projection"));
                }
                Ok(Arc::new(SchemaDescriptor::new(Arc::new(projection))))
            }
            None => Ok(root_descr),
        }
    }

    /// Sets batch size for this row iter.
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.tree_builder = self.tree_builder.with_batch_size(batch_size);
        self
    }

    /// Returns common tree builder, so the same settings are applied to both iterators
    /// from file reader and row group.
    #[inline]
    fn tree_builder() -> TreeBuilder {
        TreeBuilder::new()
    }
}

impl Iterator for RowIter<'_> {
    type Item = Result<Row>;

    fn next(&mut self) -> Option<Result<Row>> {
        let mut row = self.row_iter.as_mut().and_then(|iter| iter.next());

        while row.is_none() && self.current_row_group < self.num_row_groups {
            // We do not expect any failures when accessing a row group, and file reader
            // must be set for selecting next row group.
            if let Some(ref either) = self.file_reader {
                let file_reader = either.reader();
                let row_group_reader = &*file_reader
                    .get_row_group(self.current_row_group)
                    .expect("Row group is required to advance");

                match self
                    .tree_builder
                    .as_iter(self.descr.clone(), row_group_reader)
                {
                    Ok(mut iter) => {
                        row = iter.next();

                        self.current_row_group += 1;
                        self.row_iter = Some(iter);
                    }
                    Err(e) => return Some(Err(e)),
                }
            }
        }

        row
    }
}

/// Internal iterator of [`Row`]s for a reader.
pub struct ReaderIter {
    root_reader: Reader,
    records_left: usize,
}

impl ReaderIter {
    fn new(mut root_reader: Reader, num_records: usize) -> Result<Self> {
        // Prepare root reader by advancing all column vectors
        root_reader.advance_columns()?;
        Ok(Self {
            root_reader,
            records_left: num_records,
        })
    }
}

impl Iterator for ReaderIter {
    type Item = Result<Row>;

    fn next(&mut self) -> Option<Result<Row>> {
        if self.records_left > 0 {
            self.records_left -= 1;
            Some(self.root_reader.read())
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::data_type::Int64Type;
    use crate::file::reader::SerializedFileReader;
    use crate::file::writer::SerializedFileWriter;
    use crate::record::api::RowAccessor;
    use crate::schema::parser::parse_message_type;
    use crate::util::test_common::file_util::{get_test_file, get_test_path};
    use bytes::Bytes;

    // Convenient macros to assemble row, list, map, and group.

    macro_rules! row {
        ($($e:tt)*) => {
            {
                Row::new(vec![$($e)*])
            }
        }
    }

    macro_rules! list {
        ($($e:tt)*) => {
            {
                Field::ListInternal(make_list(vec![$($e)*]))
            }
        }
    }

    macro_rules! map {
        ($($e:tt)*) => {
            {
                Field::MapInternal(make_map(vec![$($e)*]))
            }
        }
    }

    macro_rules! group {
        ( $( $e:expr ), * ) => {
            {
                Field::Group(row!($( $e ), *))
            }
        }
    }

    #[test]
    fn test_file_reader_rows_nulls() {
        let rows = test_file_reader_rows("nulls.snappy.parquet", None).unwrap();
        let expected_rows = vec![
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
            row![(
                "b_struct".to_string(),
                group![("b_c_int".to_string(), Field::Null)]
            )],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_file_reader_rows_nonnullable() {
        let rows = test_file_reader_rows("nonnullable.impala.parquet", None).unwrap();
        let expected_rows = vec![row![
            ("ID".to_string(), Field::Long(8)),
            ("Int_Array".to_string(), list![Field::Int(-1)]),
            (
                "int_array_array".to_string(),
                list![list![Field::Int(-1), Field::Int(-2)], list![]]
            ),
            (
                "Int_Map".to_string(),
                map![(Field::Str("k1".to_string()), Field::Int(-1))]
            ),
            (
                "int_map_array".to_string(),
                list![
                    map![],
                    map![(Field::Str("k1".to_string()), Field::Int(1))],
                    map![],
                    map![]
                ]
            ),
            (
                "nested_Struct".to_string(),
                group![
                    ("a".to_string(), Field::Int(-1)),
                    ("B".to_string(), list![Field::Int(-1)]),
                    (
                        "c".to_string(),
                        group![(
                            "D".to_string(),
                            list![list![group![
                                ("e".to_string(), Field::Int(-1)),
                                ("f".to_string(), Field::Str("nonnullable".to_string()))
                            ]]]
                        )]
                    ),
                    ("G".to_string(), map![])
                ]
            )
        ]];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_file_reader_rows_nullable() {
        let rows = test_file_reader_rows("nullable.impala.parquet", None).unwrap();
        let expected_rows = vec![
            row![
                ("id".to_string(), Field::Long(1)),
                (
                    "int_array".to_string(),
                    list![Field::Int(1), Field::Int(2), Field::Int(3)]
                ),
                (
                    "int_array_Array".to_string(),
                    list![
                        list![Field::Int(1), Field::Int(2)],
                        list![Field::Int(3), Field::Int(4)]
                    ]
                ),
                (
                    "int_map".to_string(),
                    map![
                        (Field::Str("k1".to_string()), Field::Int(1)),
                        (Field::Str("k2".to_string()), Field::Int(100))
                    ]
                ),
                (
                    "int_Map_Array".to_string(),
                    list![map![(Field::Str("k1".to_string()), Field::Int(1))]]
                ),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Int(1)),
                        ("b".to_string(), list![Field::Int(1)]),
                        (
                            "C".to_string(),
                            group![(
                                "d".to_string(),
                                list![
                                    list![
                                        group![
                                            ("E".to_string(), Field::Int(10)),
                                            ("F".to_string(), Field::Str("aaa".to_string()))
                                        ],
                                        group![
                                            ("E".to_string(), Field::Int(-10)),
                                            ("F".to_string(), Field::Str("bbb".to_string()))
                                        ]
                                    ],
                                    list![group![
                                        ("E".to_string(), Field::Int(11)),
                                        ("F".to_string(), Field::Str("c".to_string()))
                                    ]]
                                ]
                            )]
                        ),
                        (
                            "g".to_string(),
                            map![(
                                Field::Str("foo".to_string()),
                                group![(
                                    "H".to_string(),
                                    group![("i".to_string(), list![Field::Double(1.1)])]
                                )]
                            )]
                        )
                    ]
                )
            ],
            row![
                ("id".to_string(), Field::Long(2)),
                (
                    "int_array".to_string(),
                    list![
                        Field::Null,
                        Field::Int(1),
                        Field::Int(2),
                        Field::Null,
                        Field::Int(3),
                        Field::Null
                    ]
                ),
                (
                    "int_array_Array".to_string(),
                    list![
                        list![Field::Null, Field::Int(1), Field::Int(2), Field::Null],
                        list![Field::Int(3), Field::Null, Field::Int(4)],
                        list![],
                        Field::Null
                    ]
                ),
                (
                    "int_map".to_string(),
                    map![
                        (Field::Str("k1".to_string()), Field::Int(2)),
                        (Field::Str("k2".to_string()), Field::Null)
                    ]
                ),
                (
                    "int_Map_Array".to_string(),
                    list![
                        map![
                            (Field::Str("k3".to_string()), Field::Null),
                            (Field::Str("k1".to_string()), Field::Int(1))
                        ],
                        Field::Null,
                        map![]
                    ]
                ),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Null),
                        ("b".to_string(), list![Field::Null]),
                        (
                            "C".to_string(),
                            group![(
                                "d".to_string(),
                                list![
                                    list![
                                        group![
                                            ("E".to_string(), Field::Null),
                                            ("F".to_string(), Field::Null)
                                        ],
                                        group![
                                            ("E".to_string(), Field::Int(10)),
                                            ("F".to_string(), Field::Str("aaa".to_string()))
                                        ],
                                        group![
                                            ("E".to_string(), Field::Null),
                                            ("F".to_string(), Field::Null)
                                        ],
                                        group![
                                            ("E".to_string(), Field::Int(-10)),
                                            ("F".to_string(), Field::Str("bbb".to_string()))
                                        ],
                                        group![
                                            ("E".to_string(), Field::Null),
                                            ("F".to_string(), Field::Null)
                                        ]
                                    ],
                                    list![
                                        group![
                                            ("E".to_string(), Field::Int(11)),
                                            ("F".to_string(), Field::Str("c".to_string()))
                                        ],
                                        Field::Null
                                    ],
                                    list![],
                                    Field::Null
                                ]
                            )]
                        ),
                        (
                            "g".to_string(),
                            map![
                                (
                                    Field::Str("g1".to_string()),
                                    group![(
                                        "H".to_string(),
                                        group![(
                                            "i".to_string(),
                                            list![Field::Double(2.2), Field::Null]
                                        )]
                                    )]
                                ),
                                (
                                    Field::Str("g2".to_string()),
                                    group![("H".to_string(), group![("i".to_string(), list![])])]
                                ),
                                (Field::Str("g3".to_string()), Field::Null),
                                (
                                    Field::Str("g4".to_string()),
                                    group![(
                                        "H".to_string(),
                                        group![("i".to_string(), Field::Null)]
                                    )]
                                ),
                                (
                                    Field::Str("g5".to_string()),
                                    group![("H".to_string(), Field::Null)]
                                )
                            ]
                        )
                    ]
                )
            ],
            row![
                ("id".to_string(), Field::Long(3)),
                ("int_array".to_string(), list![]),
                ("int_array_Array".to_string(), list![Field::Null]),
                ("int_map".to_string(), map![]),
                ("int_Map_Array".to_string(), list![Field::Null, Field::Null]),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Null),
                        ("b".to_string(), Field::Null),
                        ("C".to_string(), group![("d".to_string(), list![])]),
                        ("g".to_string(), map![])
                    ]
                )
            ],
            row![
                ("id".to_string(), Field::Long(4)),
                ("int_array".to_string(), Field::Null),
                ("int_array_Array".to_string(), list![]),
                ("int_map".to_string(), map![]),
                ("int_Map_Array".to_string(), list![]),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Null),
                        ("b".to_string(), Field::Null),
                        ("C".to_string(), group![("d".to_string(), Field::Null)]),
                        ("g".to_string(), Field::Null)
                    ]
                )
            ],
            row![
                ("id".to_string(), Field::Long(5)),
                ("int_array".to_string(), Field::Null),
                ("int_array_Array".to_string(), Field::Null),
                ("int_map".to_string(), map![]),
                ("int_Map_Array".to_string(), Field::Null),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Null),
                        ("b".to_string(), Field::Null),
                        ("C".to_string(), Field::Null),
                        (
                            "g".to_string(),
                            map![(
                                Field::Str("foo".to_string()),
                                group![(
                                    "H".to_string(),
                                    group![(
                                        "i".to_string(),
                                        list![Field::Double(2.2), Field::Double(3.3)]
                                    )]
                                )]
                            )]
                        )
                    ]
                )
            ],
            row![
                ("id".to_string(), Field::Long(6)),
                ("int_array".to_string(), Field::Null),
                ("int_array_Array".to_string(), Field::Null),
                ("int_map".to_string(), Field::Null),
                ("int_Map_Array".to_string(), Field::Null),
                ("nested_struct".to_string(), Field::Null)
            ],
            row![
                ("id".to_string(), Field::Long(7)),
                ("int_array".to_string(), Field::Null),
                (
                    "int_array_Array".to_string(),
                    list![Field::Null, list![Field::Int(5), Field::Int(6)]]
                ),
                (
                    "int_map".to_string(),
                    map![
                        (Field::Str("k1".to_string()), Field::Null),
                        (Field::Str("k3".to_string()), Field::Null)
                    ]
                ),
                ("int_Map_Array".to_string(), Field::Null),
                (
                    "nested_struct".to_string(),
                    group![
                        ("A".to_string(), Field::Int(7)),
                        (
                            "b".to_string(),
                            list![Field::Int(2), Field::Int(3), Field::Null]
                        ),
                        (
                            "C".to_string(),
                            group![(
                                "d".to_string(),
                                list![list![], list![Field::Null], Field::Null]
                            )]
                        ),
                        ("g".to_string(), Field::Null)
                    ]
                )
            ],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_file_reader_rows_projection() {
        let schema = "
      message spark_schema {
        REQUIRED DOUBLE c;
        REQUIRED INT32 b;
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        let rows = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
        let expected_rows = vec![
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
            row![
                ("c".to_string(), Field::Double(1.0)),
                ("b".to_string(), Field::Int(1))
            ],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_iter_columns_in_row() {
        let r = row![
            ("c".to_string(), Field::Double(1.0)),
            ("b".to_string(), Field::Int(1))
        ];
        let mut result = Vec::new();
        for (name, record) in r.get_column_iter() {
            result.push((name, record));
        }
        assert_eq!(
            vec![
                (&"c".to_string(), &Field::Double(1.0)),
                (&"b".to_string(), &Field::Int(1))
            ],
            result
        );
    }

    #[test]
    fn test_into_columns_in_row() {
        let r = row![
            ("a".to_string(), Field::Str("My string".to_owned())),
            ("b".to_string(), Field::Int(1))
        ];
        assert_eq!(
            r.into_columns(),
            vec![
                ("a".to_string(), Field::Str("My string".to_owned())),
                ("b".to_string(), Field::Int(1)),
            ]
        );
    }

    #[test]
    fn test_file_reader_rows_projection_map() {
        let schema = "
      message spark_schema {
        OPTIONAL group a (MAP) {
          REPEATED group key_value {
            REQUIRED BYTE_ARRAY key (UTF8);
            OPTIONAL group value (MAP) {
              REPEATED group key_value {
                REQUIRED INT32 key;
                REQUIRED BOOLEAN value;
              }
            }
          }
        }
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        let rows = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
        let expected_rows = vec![
            row![(
                "a".to_string(),
                map![(
                    Field::Str("a".to_string()),
                    map![
                        (Field::Int(1), Field::Bool(true)),
                        (Field::Int(2), Field::Bool(false))
                    ]
                )]
            )],
            row![(
                "a".to_string(),
                map![(
                    Field::Str("b".to_string()),
                    map![(Field::Int(1), Field::Bool(true))]
                )]
            )],
            row![(
                "a".to_string(),
                map![(Field::Str("c".to_string()), Field::Null)]
            )],
            row![("a".to_string(), map![(Field::Str("d".to_string()), map![])])],
            row![(
                "a".to_string(),
                map![(
                    Field::Str("e".to_string()),
                    map![(Field::Int(1), Field::Bool(true))]
                )]
            )],
            row![(
                "a".to_string(),
                map![(
                    Field::Str("f".to_string()),
                    map![
                        (Field::Int(3), Field::Bool(true)),
                        (Field::Int(4), Field::Bool(false)),
                        (Field::Int(5), Field::Bool(true))
                    ]
                )]
            )],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_file_reader_rows_projection_list() {
        let schema = "
      message spark_schema {
        OPTIONAL group a (LIST) {
          REPEATED group list {
            OPTIONAL group element (LIST) {
              REPEATED group list {
                OPTIONAL group element (LIST) {
                  REPEATED group list {
                    OPTIONAL BYTE_ARRAY element (UTF8);
                  }
                }
              }
            }
          }
        }
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        let rows = test_file_reader_rows("nested_lists.snappy.parquet", Some(schema)).unwrap();
        let expected_rows = vec![
            row![(
                "a".to_string(),
                list![
                    list![
                        list![Field::Str("a".to_string()), Field::Str("b".to_string())],
                        list![Field::Str("c".to_string())]
                    ],
                    list![Field::Null, list![Field::Str("d".to_string())]]
                ]
            )],
            row![(
                "a".to_string(),
                list![
                    list![
                        list![Field::Str("a".to_string()), Field::Str("b".to_string())],
                        list![Field::Str("c".to_string()), Field::Str("d".to_string())]
                    ],
                    list![Field::Null, list![Field::Str("e".to_string())]]
                ]
            )],
            row![(
                "a".to_string(),
                list![
                    list![
                        list![Field::Str("a".to_string()), Field::Str("b".to_string())],
                        list![Field::Str("c".to_string()), Field::Str("d".to_string())],
                        list![Field::Str("e".to_string())]
                    ],
                    list![Field::Null, list![Field::Str("f".to_string())]]
                ]
            )],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_file_reader_rows_invalid_projection() {
        let schema = "
      message spark_schema {
        REQUIRED INT32 key;
        REQUIRED BOOLEAN value;
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        let res = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema));
        assert_eq!(
            res.unwrap_err().to_string(),
            "Parquet error: Root schema does not contain projection"
        );
    }

    #[test]
    fn test_row_group_rows_invalid_projection() {
        let schema = "
      message spark_schema {
        REQUIRED INT32 key;
        REQUIRED BOOLEAN value;
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        let res = test_row_group_rows("nested_maps.snappy.parquet", Some(schema));
        assert_eq!(
            res.unwrap_err().to_string(),
            "Parquet error: Root schema does not contain projection"
        );
    }

    #[test]
    fn test_file_reader_rows_nested_map_type() {
        let schema = "
      message spark_schema {
        OPTIONAL group a (MAP) {
          REPEATED group key_value {
            REQUIRED BYTE_ARRAY key (UTF8);
            OPTIONAL group value (MAP) {
              REPEATED group key_value {
                REQUIRED INT32 key;
              }
            }
          }
        }
      }
    ";
        let schema = parse_message_type(schema).unwrap();
        test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
    }

    #[test]
    fn test_file_reader_iter() {
        let path = get_test_path("alltypes_plain.parquet");
        let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
        let iter = RowIter::from_file_into(Box::new(reader));

        let values: Vec<_> = iter.flat_map(|r| r.unwrap().get_int(0)).collect();
        assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]);
    }

    #[test]
    fn test_file_reader_iter_projection() {
        let path = get_test_path("alltypes_plain.parquet");
        let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
        let schema = "message schema { OPTIONAL INT32 id; }";
        let proj = parse_message_type(schema).ok();

        let iter = RowIter::from_file_into(Box::new(reader))
            .project(proj)
            .unwrap();
        let values: Vec<_> = iter.flat_map(|r| r.unwrap().get_int(0)).collect();

        assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]);
    }

    #[test]
    fn test_file_reader_iter_projection_err() {
        let schema = "
      message spark_schema {
        REQUIRED INT32 key;
        REQUIRED BOOLEAN value;
      }
    ";
        let proj = parse_message_type(schema).ok();
        let path = get_test_path("nested_maps.snappy.parquet");
        let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
        let res = RowIter::from_file_into(Box::new(reader)).project(proj);

        assert_eq!(
            res.err().unwrap().to_string(),
            "Parquet error: Root schema does not contain projection"
        );
    }

    #[test]
    fn test_tree_reader_handle_repeated_fields_with_no_annotation() {
        // Array field `phoneNumbers` does not contain LIST annotation.
        // We parse it as struct with `phone` repeated field as array.
        let rows = test_file_reader_rows("repeated_no_annotation.parquet", None).unwrap();
        let expected_rows = vec![
            row![
                ("id".to_string(), Field::Int(1)),
                ("phoneNumbers".to_string(), Field::Null)
            ],
            row![
                ("id".to_string(), Field::Int(2)),
                ("phoneNumbers".to_string(), Field::Null)
            ],
            row![
                ("id".to_string(), Field::Int(3)),
                (
                    "phoneNumbers".to_string(),
                    group![("phone".to_string(), list![])]
                )
            ],
            row![
                ("id".to_string(), Field::Int(4)),
                (
                    "phoneNumbers".to_string(),
                    group![(
                        "phone".to_string(),
                        list![group![
                            ("number".to_string(), Field::Long(5555555555)),
                            ("kind".to_string(), Field::Null)
                        ]]
                    )]
                )
            ],
            row![
                ("id".to_string(), Field::Int(5)),
                (
                    "phoneNumbers".to_string(),
                    group![(
                        "phone".to_string(),
                        list![group![
                            ("number".to_string(), Field::Long(1111111111)),
                            ("kind".to_string(), Field::Str("home".to_string()))
                        ]]
                    )]
                )
            ],
            row![
                ("id".to_string(), Field::Int(6)),
                (
                    "phoneNumbers".to_string(),
                    group![(
                        "phone".to_string(),
                        list![
                            group![
                                ("number".to_string(), Field::Long(1111111111)),
                                ("kind".to_string(), Field::Str("home".to_string()))
                            ],
                            group![
                                ("number".to_string(), Field::Long(2222222222)),
                                ("kind".to_string(), Field::Null)
                            ],
                            group![
                                ("number".to_string(), Field::Long(3333333333)),
                                ("kind".to_string(), Field::Str("mobile".to_string()))
                            ]
                        ]
                    )]
                )
            ],
        ];

        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_tree_reader_handle_nested_repeated_fields_with_no_annotation() {
        // Create schema
        let schema = Arc::new(
            parse_message_type(
                "
            message schema {
                REPEATED group level1 {
                    REPEATED group level2 {
                        REQUIRED group level3 {
                            REQUIRED INT64 value3;
                        }
                    }
                    REQUIRED INT64 value1;
                }
            }",
            )
            .unwrap(),
        );

        // Write Parquet file to buffer
        let mut buffer: Vec<u8> = Vec::new();
        let mut file_writer =
            SerializedFileWriter::new(&mut buffer, schema, Default::default()).unwrap();
        let mut row_group_writer = file_writer.next_row_group().unwrap();

        // Write column level1.level2.level3.value3
        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
        column_writer
            .typed::<Int64Type>()
            .write_batch(&[30, 31, 32], Some(&[2, 2, 2]), Some(&[0, 0, 0]))
            .unwrap();
        column_writer.close().unwrap();

        // Write column level1.value1
        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
        column_writer
            .typed::<Int64Type>()
            .write_batch(&[10, 11, 12], Some(&[1, 1, 1]), Some(&[0, 0, 0]))
            .unwrap();
        column_writer.close().unwrap();

        // Finalize Parquet file
        row_group_writer.close().unwrap();
        file_writer.close().unwrap();
        assert_eq!(&buffer[0..4], b"PAR1");

        // Read Parquet file from buffer
        let file_reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
        let rows: Vec<_> = file_reader
            .get_row_iter(None)
            .unwrap()
            .map(|row| row.unwrap())
            .collect();

        let expected_rows = vec![
            row![(
                "level1".to_string(),
                list![group![
                    (
                        "level2".to_string(),
                        list![group![(
                            "level3".to_string(),
                            group![("value3".to_string(), Field::Long(30))]
                        )]]
                    ),
                    ("value1".to_string(), Field::Long(10))
                ]]
            )],
            row![(
                "level1".to_string(),
                list![group![
                    (
                        "level2".to_string(),
                        list![group![(
                            "level3".to_string(),
                            group![("value3".to_string(), Field::Long(31))]
                        )]]
                    ),
                    ("value1".to_string(), Field::Long(11))
                ]]
            )],
            row![(
                "level1".to_string(),
                list![group![
                    (
                        "level2".to_string(),
                        list![group![(
                            "level3".to_string(),
                            group![("value3".to_string(), Field::Long(32))]
                        )]]
                    ),
                    ("value1".to_string(), Field::Long(12))
                ]]
            )],
        ];

        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_tree_reader_handle_primitive_repeated_fields_with_no_annotation() {
        // In this test the REPEATED fields are primitives
        let rows = test_file_reader_rows("repeated_primitive_no_list.parquet", None).unwrap();
        let expected_rows = vec![
            row![
                (
                    "Int32_list".to_string(),
                    Field::ListInternal(make_list([0, 1, 2, 3].map(Field::Int).to_vec()))
                ),
                (
                    "String_list".to_string(),
                    Field::ListInternal(make_list(
                        ["foo", "zero", "one", "two"]
                            .map(|s| Field::Str(s.to_string()))
                            .to_vec()
                    ))
                ),
                (
                    "group_of_lists".to_string(),
                    group![
                        (
                            "Int32_list_in_group".to_string(),
                            Field::ListInternal(make_list([0, 1, 2, 3].map(Field::Int).to_vec()))
                        ),
                        (
                            "String_list_in_group".to_string(),
                            Field::ListInternal(make_list(
                                ["foo", "zero", "one", "two"]
                                    .map(|s| Field::Str(s.to_string()))
                                    .to_vec()
                            ))
                        )
                    ]
                )
            ],
            row![
                (
                    "Int32_list".to_string(),
                    Field::ListInternal(make_list(vec![]))
                ),
                (
                    "String_list".to_string(),
                    Field::ListInternal(make_list(
                        ["three"].map(|s| Field::Str(s.to_string())).to_vec()
                    ))
                ),
                (
                    "group_of_lists".to_string(),
                    group![
                        (
                            "Int32_list_in_group".to_string(),
                            Field::ListInternal(make_list(vec![]))
                        ),
                        (
                            "String_list_in_group".to_string(),
                            Field::ListInternal(make_list(
                                ["three"].map(|s| Field::Str(s.to_string())).to_vec()
                            ))
                        )
                    ]
                )
            ],
            row![
                (
                    "Int32_list".to_string(),
                    Field::ListInternal(make_list(vec![Field::Int(4)]))
                ),
                (
                    "String_list".to_string(),
                    Field::ListInternal(make_list(
                        ["four"].map(|s| Field::Str(s.to_string())).to_vec()
                    ))
                ),
                (
                    "group_of_lists".to_string(),
                    group![
                        (
                            "Int32_list_in_group".to_string(),
                            Field::ListInternal(make_list(vec![Field::Int(4)]))
                        ),
                        (
                            "String_list_in_group".to_string(),
                            Field::ListInternal(make_list(
                                ["four"].map(|s| Field::Str(s.to_string())).to_vec()
                            ))
                        )
                    ]
                )
            ],
            row![
                (
                    "Int32_list".to_string(),
                    Field::ListInternal(make_list([5, 6, 7, 8].map(Field::Int).to_vec()))
                ),
                (
                    "String_list".to_string(),
                    Field::ListInternal(make_list(
                        ["five", "six", "seven", "eight"]
                            .map(|s| Field::Str(s.to_string()))
                            .to_vec()
                    ))
                ),
                (
                    "group_of_lists".to_string(),
                    group![
                        (
                            "Int32_list_in_group".to_string(),
                            Field::ListInternal(make_list([5, 6, 7, 8].map(Field::Int).to_vec()))
                        ),
                        (
                            "String_list_in_group".to_string(),
                            Field::ListInternal(make_list(
                                ["five", "six", "seven", "eight"]
                                    .map(|s| Field::Str(s.to_string()))
                                    .to_vec()
                            ))
                        )
                    ]
                )
            ],
        ];
        assert_eq!(rows, expected_rows);
    }

    #[test]
    fn test_map_no_value() {
        // File schema:
        // message schema {
        //   required group my_map (MAP) {
        //     repeated group key_value {
        //       required int32 key;
        //       optional int32 value;
        //     }
        //   }
        //   required group my_map_no_v (MAP) {
        //     repeated group key_value {
        //       required int32 key;
        //     }
        //   }
        //   required group my_list (LIST) {
        //     repeated group list {
        //       required int32 element;
        //     }
        //   }
        // }
        let rows = test_file_reader_rows("map_no_value.parquet", None).unwrap();

        // the my_map_no_v and my_list columns should be equivalent lists by this point
        for row in rows {
            let cols = row.into_columns();
            assert_eq!(cols[1].1, cols[2].1);
        }
    }

    fn test_file_reader_rows(file_name: &str, schema: Option<Type>) -> Result<Vec<Row>> {
        let file = get_test_file(file_name);
        let file_reader: Box<dyn FileReader> = Box::new(SerializedFileReader::new(file)?);
        let iter = file_reader.get_row_iter(schema)?;
        Ok(iter.map(|row| row.unwrap()).collect())
    }

    fn test_row_group_rows(file_name: &str, schema: Option<Type>) -> Result<Vec<Row>> {
        let file = get_test_file(file_name);
        let file_reader: Box<dyn FileReader> = Box::new(SerializedFileReader::new(file)?);
        // Check the first row group only, because files will contain only single row
        // group
        let row_group_reader = file_reader.get_row_group(0).unwrap();
        let iter = row_group_reader.get_row_iter(schema)?;
        Ok(iter.map(|row| row.unwrap()).collect())
    }

    #[test]
    fn test_read_old_nested_list() {
        let rows = test_file_reader_rows("old_list_structure.parquet", None).unwrap();
        let expected_rows = vec![row![(
            "a".to_string(),
            Field::ListInternal(make_list(
                [
                    make_list([1, 2].map(Field::Int).to_vec()),
                    make_list([3, 4].map(Field::Int).to_vec())
                ]
                .map(Field::ListInternal)
                .to_vec()
            ))
        ),]];
        assert_eq!(rows, expected_rows);
    }

    fn assert_err_on_overcount(file_name: &str, proj_schema: Option<Type>) {
        let file = get_test_file(file_name);
        let file_reader = SerializedFileReader::new(file).unwrap();
        let metadata = file_reader.metadata();
        let row_group_reader = file_reader.get_row_group(0).unwrap();
        let actual_rows = row_group_reader.metadata().num_rows() as usize;

        let descr = match proj_schema {
            Some(schema) => Arc::new(SchemaDescriptor::new(Arc::new(schema))),
            None => metadata.file_metadata().schema_descr_ptr(),
        };
        let reader = TreeBuilder::new().build(descr, &*row_group_reader).unwrap();
        let iter = ReaderIter::new(reader, actual_rows + 1).unwrap();

        let rows: Vec<Result<Row>> = iter.collect();
        assert_eq!(rows.len(), actual_rows + 1);
        for row in &rows[..actual_rows] {
            assert!(row.is_ok(), "Expected Ok row, got: {:?}", row);
        }
        let err = rows[actual_rows].as_ref().unwrap_err();
        assert!(
            err.to_string().contains("Unexpected end of column data"),
            "Unexpected error message: {}",
            err
        );
    }

    #[test]
    fn test_reader_iter_returns_error_when_num_records_exceeds_data() {
        assert_err_on_overcount("nulls.snappy.parquet", None);
    }

    #[test]
    fn test_reader_iter_returns_error_for_repeated_field_when_num_records_exceeds_data() {
        assert_err_on_overcount("repeated_primitive_no_list.parquet", None);
    }

    #[test]
    fn test_reader_iter_returns_error_for_map_field_when_num_records_exceeds_data() {
        let schema = parse_message_type(
            "message schema {
               REQUIRED group my_map (MAP) {
                 REPEATED group key_value {
                   REQUIRED INT32 key;
                   OPTIONAL INT32 value;
                 }
               }
             }",
        )
        .unwrap();
        assert_err_on_overcount("map_no_value.parquet", Some(schema));
    }
}