// 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`](crate::record::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::{make_list, make_map, make_row, Field, Row},
    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 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,
    ) -> 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
        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,
    ) -> 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,
        mut path: &mut Vec<String>,
        mut curr_def_level: i16,
        mut curr_rep_level: i16,
        paths: &HashMap<ColumnPath, usize>,
        row_group_reader: &dyn RowGroupReader,
    ) -> 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).unwrap();
            let col_descr = row_group_reader
                .metadata()
                .column(orig_index)
                .column_descr_ptr();
            let col_reader = row_group_reader.get_column_reader(orig_index).unwrap();
            let column = TripletIter::new(col_descr, col_reader, self.batch_size);
            Reader::PrimitiveReader(field, column)
        } 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,
                            &mut 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,
                            &mut 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
                    );
                    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(),
                        &mut path,
                        curr_def_level + 1,
                        curr_rep_level + 1,
                        paths,
                        row_group_reader,
                    );

                    let value_type = &key_value_type.get_fields()[1];
                    let value_reader = self.reader_tree(
                        value_type.clone(),
                        &mut 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(&mut Vec::from(field.get_fields()))
                        .build()
                        .unwrap();

                    path.pop();

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

                    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(),
                            &mut path,
                            curr_def_level,
                            curr_rep_level,
                            paths,
                            row_group_reader,
                        );
                        readers.push(reader);
                    }
                    Reader::GroupReader(Some(field), curr_def_level, readers)
                }
            }
        };
        path.pop();

        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, 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 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) -> 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()));
                }
                make_row(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) -> Field {
        match *self {
            Reader::PrimitiveReader(_, ref mut column) => {
                let value = column.current_value();
                column.read_next().unwrap();
                value
            }
            Reader::OptionReader(def_level, ref mut reader) => {
                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 = make_row(fields);
                Field::Group(row)
            }
            Reader::RepeatedReader(_, def_level, rep_level, ref mut reader) => {
                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,
            ) => {
                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))
            }
        }
    }

    /// 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(ref 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(ref 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) {
        match *self {
            Reader::PrimitiveReader(_, ref mut column) => {
                column.read_next().unwrap();
            }
            Reader::OptionReader(_, ref mut reader) => {
                reader.advance_columns();
            }
            Reader::GroupReader(_, _, ref mut readers) => {
                for reader in readers {
                    reader.advance_columns();
                }
            }
            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 represet a reference and a box of
/// [`FileReader`](crate::file::reader::FileReader).
enum Either<'a> {
    Left(&'a dyn FileReader),
    Right(Box<dyn FileReader>),
}

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

/// Iterator of [`Row`](crate::record::Row)s.
/// It is used either for a single row group to iterate over data in that row group, or
/// an entire file with auto buffering of all row groups.
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`](crate::record::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`](crate::record::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`](crate::record::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`](crate::record::Row)s from a
    /// [`FileReader`](crate::file::reader::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`](crate::record::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),
        }
    }

    /// 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<'a> Iterator for RowIter<'a> {
    type Item = Row;

    fn next(&mut self) -> Option<Row> {
        let mut row = None;
        if let Some(ref mut iter) = self.row_iter {
            row = 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");

                let mut iter = self
                    .tree_builder
                    .as_iter(self.descr.clone(), row_group_reader);

                row = iter.next();

                self.current_row_group += 1;
                self.row_iter = Some(iter);
            }
        }

        row
    }
}

/// Internal iterator of [`Row`](crate::record::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) -> Self {
        // Prepare root reader by advancing all column vectors
        root_reader.advance_columns();
        Self {
            root_reader,
            records_left: num_records,
        }
    }
}

impl Iterator for ReaderIter {
    type Item = Row;

    fn next(&mut self) -> Option<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::errors::{ParquetError, Result};
    use crate::file::reader::{FileReader, SerializedFileReader};
    use crate::record::api::{Field, Row, RowAccessor, RowFormatter};
    use crate::schema::parser::parse_message_type;
    use crate::util::test_common::{get_test_file, get_test_path};
    use std::convert::TryFrom;

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

    macro_rules! row {
        () => {
            {
                let result = Vec::new();
                make_row(result)
            }
        };
        ( $( $e:expr ), + ) => {
            {
                let mut result = Vec::new();
                $(
                    result.push($e);
                )*
                    make_row(result)
            }
        }
    }

    macro_rules! list {
        () => {
            {
                let result = Vec::new();
                Field::ListInternal(make_list(result))
            }
        };
        ( $( $e:expr ), + ) => {
            {
                let mut result = Vec::new();
                $(
                    result.push($e);
                )*
                    Field::ListInternal(make_list(result))
            }
        }
    }

    macro_rules! map {
        () => {
            {
                let result = Vec::new();
                Field::MapInternal(make_map(result))
            }
        };
        ( $( $e:expr ), + ) => {
            {
                let mut result = Vec::new();
                $(
                    result.push($e);
                )*
                    Field::MapInternal(make_map(result))
            }
        }
    }

    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_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!(res.is_err());
        assert_eq!(
            res.unwrap_err(),
            general_err!("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!(res.is_err());
        assert_eq!(
            res.unwrap_err(),
            general_err!("Root schema does not contain projection")
        );
    }

    #[test]
    #[should_panic(expected = "Invalid map type")]
    fn test_file_reader_rows_invalid_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 vec = vec![path]
            .iter()
            .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
            .flat_map(|r| RowIter::from_file_into(Box::new(r)))
            .flat_map(|r| r.get_int(0))
            .collect::<Vec<_>>();

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

    #[test]
    fn test_file_reader_iter_projection() {
        let path = get_test_path("alltypes_plain.parquet");
        let values = vec![path]
            .iter()
            .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
            .flat_map(|r| {
                let schema = "message schema { OPTIONAL INT32 id; }";
                let proj = parse_message_type(&schema).ok();

                RowIter::from_file_into(Box::new(r)).project(proj).unwrap()
            })
            .map(|r| format!("id:{}", r.fmt(0)))
            .collect::<Vec<_>>()
            .join(", ");

        assert_eq!(values, "id:4, id:5, id:6, id:7, id:2, id:3, id:0, id: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!(res.is_err());
        assert_eq!(
            res.err().unwrap(),
            general_err!("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);
    }

    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.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.collect())
    }
}