// 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.

use std::any::Any;
use std::cmp::{max, min};
use std::collections::{HashMap, HashSet};
use std::marker::PhantomData;
use std::mem::size_of;
use std::result::Result::Ok;
use std::sync::Arc;
use std::vec::Vec;

use arrow::array::{
    new_empty_array, Array, ArrayData, ArrayDataBuilder, ArrayRef, BinaryArray,
    BinaryBuilder, BooleanArray, BooleanBufferBuilder, BooleanBuilder, DecimalBuilder,
    FixedSizeBinaryArray, FixedSizeBinaryBuilder, GenericListArray, Int16BufferBuilder,
    Int32Array, Int64Array, MapArray, OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder,
    StringArray, StringBuilder, StructArray,
};
use arrow::buffer::{Buffer, MutableBuffer};
use arrow::datatypes::{
    ArrowPrimitiveType, BooleanType as ArrowBooleanType, DataType as ArrowType,
    Date32Type as ArrowDate32Type, Date64Type as ArrowDate64Type,
    DurationMicrosecondType as ArrowDurationMicrosecondType,
    DurationMillisecondType as ArrowDurationMillisecondType,
    DurationNanosecondType as ArrowDurationNanosecondType,
    DurationSecondType as ArrowDurationSecondType, Field,
    Float32Type as ArrowFloat32Type, Float64Type as ArrowFloat64Type,
    Int16Type as ArrowInt16Type, Int32Type as ArrowInt32Type,
    Int64Type as ArrowInt64Type, Int8Type as ArrowInt8Type, IntervalUnit, Schema,
    Time32MillisecondType as ArrowTime32MillisecondType,
    Time32SecondType as ArrowTime32SecondType,
    Time64MicrosecondType as ArrowTime64MicrosecondType,
    Time64NanosecondType as ArrowTime64NanosecondType, TimeUnit as ArrowTimeUnit,
    TimestampMicrosecondType as ArrowTimestampMicrosecondType,
    TimestampMillisecondType as ArrowTimestampMillisecondType,
    TimestampNanosecondType as ArrowTimestampNanosecondType,
    TimestampSecondType as ArrowTimestampSecondType, ToByteSlice,
    UInt16Type as ArrowUInt16Type, UInt32Type as ArrowUInt32Type,
    UInt64Type as ArrowUInt64Type, UInt8Type as ArrowUInt8Type,
};
use arrow::util::bit_util;

use crate::arrow::converter::{
    Converter, DecimalArrayConverter, DecimalConverter, FixedLenBinaryConverter,
    FixedSizeArrayConverter, Int96ArrayConverter, Int96Converter,
    IntervalDayTimeArrayConverter, IntervalDayTimeConverter,
    IntervalYearMonthArrayConverter, IntervalYearMonthConverter,
};
use crate::arrow::record_reader::buffer::{ScalarValue, ValuesBuffer};
use crate::arrow::record_reader::{GenericRecordReader, RecordReader};
use crate::arrow::schema::parquet_to_arrow_field;
use crate::basic::{ConvertedType, Repetition, Type as PhysicalType};
use crate::column::page::PageIterator;
use crate::column::reader::decoder::ColumnValueDecoder;
use crate::column::reader::ColumnReaderImpl;
use crate::data_type::{
    BoolType, DataType, DoubleType, FixedLenByteArrayType, FloatType, Int32Type,
    Int64Type, Int96Type,
};
use crate::errors::{ParquetError, ParquetError::ArrowError, Result};
use crate::file::reader::{FilePageIterator, FileReader};
use crate::schema::types::{
    ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, Type, TypePtr,
};
use crate::schema::visitor::TypeVisitor;

mod byte_array;
mod byte_array_dictionary;
mod dictionary_buffer;
mod offset_buffer;

#[cfg(test)]
mod test_util;

pub use byte_array::make_byte_array_reader;
pub use byte_array_dictionary::make_byte_array_dictionary_reader;

/// Array reader reads parquet data into arrow array.
pub trait ArrayReader {
    fn as_any(&self) -> &dyn Any;

    /// Returns the arrow type of this array reader.
    fn get_data_type(&self) -> &ArrowType;

    /// Reads at most `batch_size` records into an arrow array and return it.
    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef>;

    /// If this array has a non-zero definition level, i.e. has a nullable parent
    /// array, returns the definition levels of data from the last call of `next_batch`
    ///
    /// Otherwise returns None
    ///
    /// This is used by parent [`ArrayReader`] to compute their null bitmaps
    fn get_def_levels(&self) -> Option<&[i16]>;

    /// If this array has a non-zero repetition level, i.e. has a repeated parent
    /// array, returns the repetition levels of data from the last call of `next_batch`
    ///
    /// Otherwise returns None
    ///
    /// This is used by parent [`ArrayReader`] to compute their array offsets
    fn get_rep_levels(&self) -> Option<&[i16]>;
}

/// Uses `record_reader` to read up to `batch_size` records from `pages`
///
/// Returns the number of records read, which can be less than batch_size if
/// pages is exhausted.
fn read_records<V, CV>(
    record_reader: &mut GenericRecordReader<V, CV>,
    pages: &mut dyn PageIterator,
    batch_size: usize,
) -> Result<usize>
where
    V: ValuesBuffer + Default,
    CV: ColumnValueDecoder<Slice = V::Slice>,
{
    let mut records_read = 0usize;
    while records_read < batch_size {
        let records_to_read = batch_size - records_read;

        let records_read_once = record_reader.read_records(records_to_read)?;
        records_read += records_read_once;

        // Record reader exhausted
        if records_read_once < records_to_read {
            if let Some(page_reader) = pages.next() {
                // Read from new page reader (i.e. column chunk)
                record_reader.set_page_reader(page_reader?)?;
            } else {
                // Page reader also exhausted
                break;
            }
        }
    }
    Ok(records_read)
}

/// A NullArrayReader reads Parquet columns stored as null int32s with an Arrow
/// NullArray type.
pub struct NullArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    data_type: ArrowType,
    pages: Box<dyn PageIterator>,
    def_levels_buffer: Option<Buffer>,
    rep_levels_buffer: Option<Buffer>,
    column_desc: ColumnDescPtr,
    record_reader: RecordReader<T>,
    _type_marker: PhantomData<T>,
}

impl<T> NullArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    /// Construct null array reader.
    pub fn new(pages: Box<dyn PageIterator>, column_desc: ColumnDescPtr) -> Result<Self> {
        let record_reader = RecordReader::<T>::new(column_desc.clone());

        Ok(Self {
            data_type: ArrowType::Null,
            pages,
            def_levels_buffer: None,
            rep_levels_buffer: None,
            column_desc,
            record_reader,
            _type_marker: PhantomData,
        })
    }
}

/// Implementation of primitive array reader.
impl<T> ArrayReader for NullArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns data type of primitive array.
    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    /// Reads at most `batch_size` records into array.
    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        let records_read =
            read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)?;

        // convert to arrays
        let array = arrow::array::NullArray::new(records_read);

        // save definition and repetition buffers
        self.def_levels_buffer = self.record_reader.consume_def_levels()?;
        self.rep_levels_buffer = self.record_reader.consume_rep_levels()?;
        self.record_reader.reset();
        Ok(Arc::new(array))
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_levels_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_levels_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }
}

/// Primitive array readers are leaves of array reader tree. They accept page iterator
/// and read them into primitive arrays.
pub struct PrimitiveArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    data_type: ArrowType,
    pages: Box<dyn PageIterator>,
    def_levels_buffer: Option<Buffer>,
    rep_levels_buffer: Option<Buffer>,
    column_desc: ColumnDescPtr,
    record_reader: RecordReader<T>,
}

impl<T> PrimitiveArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    /// Construct primitive array reader.
    pub fn new(
        pages: Box<dyn PageIterator>,
        column_desc: ColumnDescPtr,
        arrow_type: Option<ArrowType>,
    ) -> Result<Self> {
        Self::new_with_options(pages, column_desc, arrow_type, false)
    }

    /// Construct primitive array reader with ability to only compute null mask and not
    /// buffer level data
    pub fn new_with_options(
        pages: Box<dyn PageIterator>,
        column_desc: ColumnDescPtr,
        arrow_type: Option<ArrowType>,
        null_mask_only: bool,
    ) -> Result<Self> {
        // Check if Arrow type is specified, else create it from Parquet type
        let data_type = match arrow_type {
            Some(t) => t,
            None => parquet_to_arrow_field(column_desc.as_ref())?
                .data_type()
                .clone(),
        };

        let record_reader =
            RecordReader::<T>::new_with_options(column_desc.clone(), null_mask_only);

        Ok(Self {
            data_type,
            pages,
            def_levels_buffer: None,
            rep_levels_buffer: None,
            column_desc,
            record_reader,
        })
    }
}

/// Implementation of primitive array reader.
impl<T> ArrayReader for PrimitiveArrayReader<T>
where
    T: DataType,
    T::T: ScalarValue,
{
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns data type of primitive array.
    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    /// Reads at most `batch_size` records into array.
    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)?;

        let target_type = self.get_data_type().clone();
        let arrow_data_type = match T::get_physical_type() {
            PhysicalType::BOOLEAN => ArrowBooleanType::DATA_TYPE,
            PhysicalType::INT32 => {
                match target_type {
                    ArrowType::UInt32 => {
                        // follow C++ implementation and use overflow/reinterpret cast from  i32 to u32 which will map
                        // `i32::MIN..0` to `(i32::MAX as u32)..u32::MAX`
                        ArrowUInt32Type::DATA_TYPE
                    }
                    _ => ArrowInt32Type::DATA_TYPE,
                }
            }
            PhysicalType::INT64 => {
                match target_type {
                    ArrowType::UInt64 => {
                        // follow C++ implementation and use overflow/reinterpret cast from  i64 to u64 which will map
                        // `i64::MIN..0` to `(i64::MAX as u64)..u64::MAX`
                        ArrowUInt64Type::DATA_TYPE
                    }
                    _ => ArrowInt64Type::DATA_TYPE,
                }
            }
            PhysicalType::FLOAT => ArrowFloat32Type::DATA_TYPE,
            PhysicalType::DOUBLE => ArrowFloat64Type::DATA_TYPE,
            PhysicalType::INT96
            | PhysicalType::BYTE_ARRAY
            | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
                unreachable!(
                    "PrimitiveArrayReaders don't support complex physical types"
                );
            }
        };

        // Convert to arrays by using the Parquet physical type.
        // The physical types are then cast to Arrow types if necessary

        let mut record_data = self.record_reader.consume_record_data()?;

        if T::get_physical_type() == PhysicalType::BOOLEAN {
            let mut boolean_buffer = BooleanBufferBuilder::new(record_data.len());

            for e in record_data.as_slice() {
                boolean_buffer.append(*e > 0);
            }
            record_data = boolean_buffer.finish();
        }

        let mut array_data = ArrayDataBuilder::new(arrow_data_type)
            .len(self.record_reader.num_values())
            .add_buffer(record_data);

        if let Some(b) = self.record_reader.consume_bitmap_buffer()? {
            array_data = array_data.null_bit_buffer(b);
        }

        let array_data = unsafe { array_data.build_unchecked() };
        let array = match T::get_physical_type() {
            PhysicalType::BOOLEAN => Arc::new(BooleanArray::from(array_data)) as ArrayRef,
            PhysicalType::INT32 => {
                Arc::new(PrimitiveArray::<ArrowInt32Type>::from(array_data)) as ArrayRef
            }
            PhysicalType::INT64 => {
                Arc::new(PrimitiveArray::<ArrowInt64Type>::from(array_data)) as ArrayRef
            }
            PhysicalType::FLOAT => {
                Arc::new(PrimitiveArray::<ArrowFloat32Type>::from(array_data)) as ArrayRef
            }
            PhysicalType::DOUBLE => {
                Arc::new(PrimitiveArray::<ArrowFloat64Type>::from(array_data)) as ArrayRef
            }
            PhysicalType::INT96
            | PhysicalType::BYTE_ARRAY
            | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
                unreachable!(
                    "PrimitiveArrayReaders don't support complex physical types"
                );
            }
        };

        // cast to Arrow type
        // We make a strong assumption here that the casts should be infallible.
        // If the cast fails because of incompatible datatypes, then there might
        // be a bigger problem with how Arrow schemas are converted to Parquet.
        //
        // As there is not always a 1:1 mapping between Arrow and Parquet, there
        // are datatypes which we must convert explicitly.
        // These are:
        // - date64: we should cast int32 to date32, then date32 to date64.
        let array = match target_type {
            ArrowType::Date64 => {
                // this is cheap as it internally reinterprets the data
                let a = arrow::compute::cast(&array, &ArrowType::Date32)?;
                arrow::compute::cast(&a, &target_type)?
            }
            ArrowType::Decimal(p, s) => {
                let mut builder = DecimalBuilder::new(array.len(), p, s);
                match array.data_type() {
                    ArrowType::Int32 => {
                        let values = array.as_any().downcast_ref::<Int32Array>().unwrap();
                        for maybe_value in values.iter() {
                            match maybe_value {
                                Some(value) => builder.append_value(value as i128)?,
                                None => builder.append_null()?,
                            }
                        }
                    }
                    ArrowType::Int64 => {
                        let values = array.as_any().downcast_ref::<Int64Array>().unwrap();
                        for maybe_value in values.iter() {
                            match maybe_value {
                                Some(value) => builder.append_value(value as i128)?,
                                None => builder.append_null()?,
                            }
                        }
                    }
                    _ => {
                        return Err(ArrowError(format!(
                            "Cannot convert {:?} to decimal",
                            array.data_type()
                        )))
                    }
                }
                Arc::new(builder.finish()) as ArrayRef
            }
            _ => arrow::compute::cast(&array, &target_type)?,
        };

        // save definition and repetition buffers
        self.def_levels_buffer = self.record_reader.consume_def_levels()?;
        self.rep_levels_buffer = self.record_reader.consume_rep_levels()?;
        self.record_reader.reset();
        Ok(array)
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_levels_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_levels_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }
}

/// Primitive array readers are leaves of array reader tree. They accept page iterator
/// and read them into primitive arrays.
pub struct ComplexObjectArrayReader<T, C>
where
    T: DataType,
    C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
{
    data_type: ArrowType,
    pages: Box<dyn PageIterator>,
    def_levels_buffer: Option<Vec<i16>>,
    rep_levels_buffer: Option<Vec<i16>>,
    column_desc: ColumnDescPtr,
    column_reader: Option<ColumnReaderImpl<T>>,
    converter: C,
    _parquet_type_marker: PhantomData<T>,
    _converter_marker: PhantomData<C>,
}

impl<T, C> ArrayReader for ComplexObjectArrayReader<T, C>
where
    T: DataType,
    C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
{
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        // Try to initialize column reader
        if self.column_reader.is_none() {
            self.next_column_reader()?;
        }

        let mut data_buffer: Vec<T::T> = Vec::with_capacity(batch_size);
        data_buffer.resize_with(batch_size, T::T::default);

        let mut def_levels_buffer = if self.column_desc.max_def_level() > 0 {
            let mut buf: Vec<i16> = Vec::with_capacity(batch_size);
            buf.resize_with(batch_size, || 0);
            Some(buf)
        } else {
            None
        };

        let mut rep_levels_buffer = if self.column_desc.max_rep_level() > 0 {
            let mut buf: Vec<i16> = Vec::with_capacity(batch_size);
            buf.resize_with(batch_size, || 0);
            Some(buf)
        } else {
            None
        };

        let mut num_read = 0;

        while self.column_reader.is_some() && num_read < batch_size {
            let num_to_read = batch_size - num_read;
            let cur_data_buf = &mut data_buffer[num_read..];
            let cur_def_levels_buf =
                def_levels_buffer.as_mut().map(|b| &mut b[num_read..]);
            let cur_rep_levels_buf =
                rep_levels_buffer.as_mut().map(|b| &mut b[num_read..]);
            let (data_read, levels_read) =
                self.column_reader.as_mut().unwrap().read_batch(
                    num_to_read,
                    cur_def_levels_buf,
                    cur_rep_levels_buf,
                    cur_data_buf,
                )?;

            // Fill space
            if levels_read > data_read {
                def_levels_buffer.iter().for_each(|def_levels_buffer| {
                    let (mut level_pos, mut data_pos) = (levels_read, data_read);
                    while level_pos > 0 && data_pos > 0 {
                        if def_levels_buffer[num_read + level_pos - 1]
                            == self.column_desc.max_def_level()
                        {
                            cur_data_buf.swap(level_pos - 1, data_pos - 1);
                            level_pos -= 1;
                            data_pos -= 1;
                        } else {
                            level_pos -= 1;
                        }
                    }
                });
            }

            let values_read = max(levels_read, data_read);
            num_read += values_read;
            // current page exhausted && page iterator exhausted
            if values_read < num_to_read && !self.next_column_reader()? {
                break;
            }
        }

        data_buffer.truncate(num_read);
        def_levels_buffer
            .iter_mut()
            .for_each(|buf| buf.truncate(num_read));
        rep_levels_buffer
            .iter_mut()
            .for_each(|buf| buf.truncate(num_read));

        self.def_levels_buffer = def_levels_buffer;
        self.rep_levels_buffer = rep_levels_buffer;

        let data: Vec<Option<T::T>> = if self.def_levels_buffer.is_some() {
            data_buffer
                .into_iter()
                .zip(self.def_levels_buffer.as_ref().unwrap().iter())
                .map(|(t, def_level)| {
                    if *def_level == self.column_desc.max_def_level() {
                        Some(t)
                    } else {
                        None
                    }
                })
                .collect()
        } else {
            data_buffer.into_iter().map(Some).collect()
        };

        let mut array = self.converter.convert(data)?;

        if let ArrowType::Dictionary(_, _) = self.data_type {
            array = arrow::compute::cast(&array, &self.data_type)?;
        }

        Ok(array)
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_levels_buffer.as_deref()
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_levels_buffer.as_deref()
    }
}

impl<T, C> ComplexObjectArrayReader<T, C>
where
    T: DataType,
    C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
{
    pub fn new(
        pages: Box<dyn PageIterator>,
        column_desc: ColumnDescPtr,
        converter: C,
        arrow_type: Option<ArrowType>,
    ) -> Result<Self> {
        let data_type = match arrow_type {
            Some(t) => t,
            None => parquet_to_arrow_field(column_desc.as_ref())?
                .data_type()
                .clone(),
        };

        Ok(Self {
            data_type,
            pages,
            def_levels_buffer: None,
            rep_levels_buffer: None,
            column_desc,
            column_reader: None,
            converter,
            _parquet_type_marker: PhantomData,
            _converter_marker: PhantomData,
        })
    }

    fn next_column_reader(&mut self) -> Result<bool> {
        Ok(match self.pages.next() {
            Some(page) => {
                self.column_reader =
                    Some(ColumnReaderImpl::<T>::new(self.column_desc.clone(), page?));
                true
            }
            None => false,
        })
    }
}

/// Implementation of list array reader.
pub struct ListArrayReader<OffsetSize: OffsetSizeTrait> {
    item_reader: Box<dyn ArrayReader>,
    data_type: ArrowType,
    item_type: ArrowType,
    list_def_level: i16,
    list_rep_level: i16,
    list_empty_def_level: i16,
    list_null_def_level: i16,
    def_level_buffer: Option<Buffer>,
    rep_level_buffer: Option<Buffer>,
    _marker: PhantomData<OffsetSize>,
}

impl<OffsetSize: OffsetSizeTrait> ListArrayReader<OffsetSize> {
    /// Construct list array reader.
    pub fn new(
        item_reader: Box<dyn ArrayReader>,
        data_type: ArrowType,
        item_type: ArrowType,
        def_level: i16,
        rep_level: i16,
        list_null_def_level: i16,
        list_empty_def_level: i16,
    ) -> Self {
        Self {
            item_reader,
            data_type,
            item_type,
            list_def_level: def_level,
            list_rep_level: rep_level,
            list_null_def_level,
            list_empty_def_level,
            def_level_buffer: None,
            rep_level_buffer: None,
            _marker: PhantomData,
        }
    }
}

macro_rules! remove_primitive_array_indices {
    ($arr: expr, $item_type:ty, $indices:expr) => {{
        let array_data = match $arr.as_any().downcast_ref::<PrimitiveArray<$item_type>>() {
            Some(a) => a,
            _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
        };
        let mut builder = PrimitiveBuilder::<$item_type>::new($arr.len());
        for i in 0..array_data.len() {
            if !$indices.contains(&i) {
                if array_data.is_null(i) {
                    builder.append_null()?;
                } else {
                    builder.append_value(array_data.value(i))?;
                }
            }
        }
        Ok(Arc::new(builder.finish()))
    }};
}

macro_rules! remove_array_indices_custom_builder {
    ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr) => {{
        let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
            Some(a) => a,
            _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
        };
        let mut builder = $item_builder::new(array_data.len());

        for i in 0..array_data.len() {
            if !$indices.contains(&i) {
                if array_data.is_null(i) {
                    builder.append_null()?;
                } else {
                    builder.append_value(array_data.value(i))?;
                }
            }
        }
        Ok(Arc::new(builder.finish()))
    }};
}

macro_rules! remove_fixed_size_binary_array_indices {
    ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr, $len:expr) => {{
        let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
            Some(a) => a,
            _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
        };
        let mut builder = FixedSizeBinaryBuilder::new(array_data.len(), $len);
        for i in 0..array_data.len() {
            if !$indices.contains(&i) {
                if array_data.is_null(i) {
                    builder.append_null()?;
                } else {
                    builder.append_value(array_data.value(i))?;
                }
            }
        }
        Ok(Arc::new(builder.finish()))
    }};
}

fn remove_indices(
    arr: ArrayRef,
    item_type: ArrowType,
    indices: Vec<usize>,
) -> Result<ArrayRef> {
    match item_type {
        ArrowType::UInt8 => remove_primitive_array_indices!(arr, ArrowUInt8Type, indices),
        ArrowType::UInt16 => {
            remove_primitive_array_indices!(arr, ArrowUInt16Type, indices)
        }
        ArrowType::UInt32 => {
            remove_primitive_array_indices!(arr, ArrowUInt32Type, indices)
        }
        ArrowType::UInt64 => {
            remove_primitive_array_indices!(arr, ArrowUInt64Type, indices)
        }
        ArrowType::Int8 => remove_primitive_array_indices!(arr, ArrowInt8Type, indices),
        ArrowType::Int16 => remove_primitive_array_indices!(arr, ArrowInt16Type, indices),
        ArrowType::Int32 => remove_primitive_array_indices!(arr, ArrowInt32Type, indices),
        ArrowType::Int64 => remove_primitive_array_indices!(arr, ArrowInt64Type, indices),
        ArrowType::Float32 => {
            remove_primitive_array_indices!(arr, ArrowFloat32Type, indices)
        }
        ArrowType::Float64 => {
            remove_primitive_array_indices!(arr, ArrowFloat64Type, indices)
        }
        ArrowType::Boolean => {
            remove_array_indices_custom_builder!(
                arr,
                BooleanArray,
                BooleanBuilder,
                indices
            )
        }
        ArrowType::Date32 => {
            remove_primitive_array_indices!(arr, ArrowDate32Type, indices)
        }
        ArrowType::Date64 => {
            remove_primitive_array_indices!(arr, ArrowDate64Type, indices)
        }
        ArrowType::Time32(ArrowTimeUnit::Second) => {
            remove_primitive_array_indices!(arr, ArrowTime32SecondType, indices)
        }
        ArrowType::Time32(ArrowTimeUnit::Millisecond) => {
            remove_primitive_array_indices!(arr, ArrowTime32MillisecondType, indices)
        }
        ArrowType::Time64(ArrowTimeUnit::Microsecond) => {
            remove_primitive_array_indices!(arr, ArrowTime64MicrosecondType, indices)
        }
        ArrowType::Time64(ArrowTimeUnit::Nanosecond) => {
            remove_primitive_array_indices!(arr, ArrowTime64NanosecondType, indices)
        }
        ArrowType::Duration(ArrowTimeUnit::Second) => {
            remove_primitive_array_indices!(arr, ArrowDurationSecondType, indices)
        }
        ArrowType::Duration(ArrowTimeUnit::Millisecond) => {
            remove_primitive_array_indices!(arr, ArrowDurationMillisecondType, indices)
        }
        ArrowType::Duration(ArrowTimeUnit::Microsecond) => {
            remove_primitive_array_indices!(arr, ArrowDurationMicrosecondType, indices)
        }
        ArrowType::Duration(ArrowTimeUnit::Nanosecond) => {
            remove_primitive_array_indices!(arr, ArrowDurationNanosecondType, indices)
        }
        ArrowType::Timestamp(ArrowTimeUnit::Second, _) => {
            remove_primitive_array_indices!(arr, ArrowTimestampSecondType, indices)
        }
        ArrowType::Timestamp(ArrowTimeUnit::Millisecond, _) => {
            remove_primitive_array_indices!(arr, ArrowTimestampMillisecondType, indices)
        }
        ArrowType::Timestamp(ArrowTimeUnit::Microsecond, _) => {
            remove_primitive_array_indices!(arr, ArrowTimestampMicrosecondType, indices)
        }
        ArrowType::Timestamp(ArrowTimeUnit::Nanosecond, _) => {
            remove_primitive_array_indices!(arr, ArrowTimestampNanosecondType, indices)
        }
        ArrowType::Utf8 => {
            remove_array_indices_custom_builder!(arr, StringArray, StringBuilder, indices)
        }
        ArrowType::Binary => {
            remove_array_indices_custom_builder!(arr, BinaryArray, BinaryBuilder, indices)
        }
        ArrowType::FixedSizeBinary(size) => remove_fixed_size_binary_array_indices!(
            arr,
            FixedSizeBinaryArray,
            FixedSizeBinaryBuilder,
            indices,
            size
        ),
        ArrowType::Struct(fields) => {
            let struct_array = arr
                .as_any()
                .downcast_ref::<StructArray>()
                .expect("Array should be a struct");

            // Recursively call remove indices on each of the structs fields
            let new_columns = fields
                .into_iter()
                .zip(struct_array.columns())
                .map(|(field, column)| {
                    let dt = field.data_type().clone();
                    Ok((field, remove_indices(column.clone(), dt, indices.clone())?))
                })
                .collect::<Result<Vec<_>>>()?;

            if arr.data().null_count() == 0 {
                // No nulls, nothing to do.
                Ok(Arc::new(StructArray::from(new_columns)))
            } else {
                // Construct a new validity buffer by removing `indices` from the original validity
                // map.
                let mut valid = BooleanBufferBuilder::new(arr.len() - indices.len());
                for idx in 0..arr.len() {
                    if !indices.contains(&idx) {
                        valid.append(!arr.is_null(idx));
                    }
                }
                Ok(Arc::new(StructArray::from((new_columns, valid.finish()))))
            }
        }
        _ => Err(ParquetError::General(format!(
            "ListArray of type List({:?}) is not supported by array_reader",
            item_type
        ))),
    }
}

/// Implementation of ListArrayReader. Nested lists and lists of structs are not yet supported.
impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns data type.
    /// This must be a List.
    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        let next_batch_array = self.item_reader.next_batch(batch_size)?;
        let item_type = self.item_reader.get_data_type().clone();

        if next_batch_array.len() == 0 {
            return Ok(new_empty_array(&self.data_type));
        }
        let def_levels = self
            .item_reader
            .get_def_levels()
            .ok_or_else(|| ArrowError("item_reader def levels are None.".to_string()))?;
        let rep_levels = self
            .item_reader
            .get_rep_levels()
            .ok_or_else(|| ArrowError("item_reader rep levels are None.".to_string()))?;

        if !((def_levels.len() == rep_levels.len())
            && (rep_levels.len() == next_batch_array.len()))
        {
            return Err(ArrowError(
                "Expected item_reader def_levels and rep_levels to be same length as batch".to_string(),
            ));
        }

        // List definitions can be encoded as 4 values:
        // - n + 0: the list slot is null
        // - n + 1: the list slot is not null, but is empty (i.e. [])
        // - n + 2: the list slot is not null, but its child is empty (i.e. [ null ])
        // - n + 3: the list slot is not null, and its child is not empty
        // Where n is the max definition level of the list's parent.
        // If a Parquet schema's only leaf is the list, then n = 0.

        // If the list index is at empty definition, the child slot is null
        let null_list_indices: Vec<usize> = def_levels
            .iter()
            .enumerate()
            .filter_map(|(index, def)| {
                if *def <= self.list_empty_def_level {
                    Some(index)
                } else {
                    None
                }
            })
            .collect();
        let batch_values = match null_list_indices.len() {
            0 => next_batch_array.clone(),
            _ => remove_indices(next_batch_array.clone(), item_type, null_list_indices)?,
        };

        // first item in each list has rep_level = 0, subsequent items have rep_level = 1
        let mut offsets: Vec<OffsetSize> = Vec::new();
        let mut cur_offset = OffsetSize::zero();
        def_levels.iter().zip(rep_levels).for_each(|(d, r)| {
            if *r == 0 || d == &self.list_empty_def_level {
                offsets.push(cur_offset);
            }
            if d > &self.list_empty_def_level {
                cur_offset += OffsetSize::one();
            }
        });
        offsets.push(cur_offset);

        let num_bytes = bit_util::ceil(offsets.len(), 8);
        // TODO: A useful optimization is to use the null count to fill with
        // 0 or null, to reduce individual bits set in a loop.
        // To favour dense data, set every slot to true, then unset
        let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
        let null_slice = null_buf.as_slice_mut();
        let mut list_index = 0;
        for i in 0..rep_levels.len() {
            // If the level is lower than empty, then the slot is null.
            // When a list is non-nullable, its empty level = null level,
            // so this automatically factors that in.
            if rep_levels[i] == 0 && def_levels[i] < self.list_empty_def_level {
                bit_util::unset_bit(null_slice, list_index);
            }
            if rep_levels[i] == 0 {
                list_index += 1;
            }
        }
        let value_offsets = Buffer::from(&offsets.to_byte_slice());

        let list_data = ArrayData::builder(self.get_data_type().clone())
            .len(offsets.len() - 1)
            .add_buffer(value_offsets)
            .add_child_data(batch_values.data().clone())
            .null_bit_buffer(null_buf.into())
            .offset(next_batch_array.offset());

        let list_data = unsafe { list_data.build_unchecked() };

        let result_array = GenericListArray::<OffsetSize>::from(list_data);
        Ok(Arc::new(result_array))
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }
}

/// Implementation of a map array reader.
pub struct MapArrayReader {
    key_reader: Box<dyn ArrayReader>,
    value_reader: Box<dyn ArrayReader>,
    data_type: ArrowType,
    map_def_level: i16,
    map_rep_level: i16,
    def_level_buffer: Option<Buffer>,
    rep_level_buffer: Option<Buffer>,
}

impl MapArrayReader {
    pub fn new(
        key_reader: Box<dyn ArrayReader>,
        value_reader: Box<dyn ArrayReader>,
        data_type: ArrowType,
        def_level: i16,
        rep_level: i16,
    ) -> Self {
        Self {
            key_reader,
            value_reader,
            data_type,
            map_def_level: rep_level,
            map_rep_level: def_level,
            def_level_buffer: None,
            rep_level_buffer: None,
        }
    }
}

impl ArrayReader for MapArrayReader {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        let key_array = self.key_reader.next_batch(batch_size)?;
        let value_array = self.value_reader.next_batch(batch_size)?;

        // Check that key and value have the same lengths
        let key_length = key_array.len();
        if key_length != value_array.len() {
            return Err(general_err!(
                "Map key and value should have the same lengths."
            ));
        }

        let def_levels = self
            .key_reader
            .get_def_levels()
            .ok_or_else(|| ArrowError("item_reader def levels are None.".to_string()))?;
        let rep_levels = self
            .key_reader
            .get_rep_levels()
            .ok_or_else(|| ArrowError("item_reader rep levels are None.".to_string()))?;

        if !((def_levels.len() == rep_levels.len()) && (rep_levels.len() == key_length)) {
            return Err(ArrowError(
                "Expected item_reader def_levels and rep_levels to be same length as batch".to_string(),
            ));
        }

        let entry_data_type = if let ArrowType::Map(field, _) = &self.data_type {
            field.data_type().clone()
        } else {
            return Err(ArrowError("Expected a map arrow type".to_string()));
        };

        let entry_data = ArrayDataBuilder::new(entry_data_type)
            .len(key_length)
            .add_child_data(key_array.data().clone())
            .add_child_data(value_array.data().clone());
        let entry_data = unsafe { entry_data.build_unchecked() };

        let entry_len = rep_levels.iter().filter(|level| **level == 0).count();

        // first item in each list has rep_level = 0, subsequent items have rep_level = 1
        let mut offsets: Vec<i32> = Vec::new();
        let mut cur_offset = 0;
        def_levels.iter().zip(rep_levels).for_each(|(d, r)| {
            if *r == 0 || d == &self.map_def_level {
                offsets.push(cur_offset);
            }
            if d > &self.map_def_level {
                cur_offset += 1;
            }
        });
        offsets.push(cur_offset);

        let num_bytes = bit_util::ceil(offsets.len(), 8);
        // TODO: A useful optimization is to use the null count to fill with
        // 0 or null, to reduce individual bits set in a loop.
        // To favour dense data, set every slot to true, then unset
        let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
        let null_slice = null_buf.as_slice_mut();
        let mut list_index = 0;
        for i in 0..rep_levels.len() {
            // If the level is lower than empty, then the slot is null.
            // When a list is non-nullable, its empty level = null level,
            // so this automatically factors that in.
            if rep_levels[i] == 0 && def_levels[i] < self.map_def_level {
                // should be empty list
                bit_util::unset_bit(null_slice, list_index);
            }
            if rep_levels[i] == 0 {
                list_index += 1;
            }
        }
        let value_offsets = Buffer::from(&offsets.to_byte_slice());

        // Now we can build array data
        let array_data = ArrayDataBuilder::new(self.data_type.clone())
            .len(entry_len)
            .add_buffer(value_offsets)
            .null_bit_buffer(null_buf.into())
            .add_child_data(entry_data);

        let array_data = unsafe { array_data.build_unchecked() };

        Ok(Arc::new(MapArray::from(array_data)))
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }
}

/// Implementation of struct array reader.
pub struct StructArrayReader {
    children: Vec<Box<dyn ArrayReader>>,
    data_type: ArrowType,
    struct_def_level: i16,
    struct_rep_level: i16,
    def_level_buffer: Option<Buffer>,
    rep_level_buffer: Option<Buffer>,
}

impl StructArrayReader {
    /// Construct struct array reader.
    pub fn new(
        data_type: ArrowType,
        children: Vec<Box<dyn ArrayReader>>,
        def_level: i16,
        rep_level: i16,
    ) -> Self {
        Self {
            data_type,
            children,
            struct_def_level: def_level,
            struct_rep_level: rep_level,
            def_level_buffer: None,
            rep_level_buffer: None,
        }
    }
}

impl ArrayReader for StructArrayReader {
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns data type.
    /// This must be a struct.
    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    /// Read `batch_size` struct records.
    ///
    /// Definition levels of struct array is calculated as following:
    /// ```ignore
    /// def_levels[i] = min(child1_def_levels[i], child2_def_levels[i], ...,
    /// childn_def_levels[i]);
    /// ```
    ///
    /// Repetition levels of struct array is calculated as following:
    /// ```ignore
    /// rep_levels[i] = child1_rep_levels[i];
    /// ```
    ///
    /// The null bitmap of struct array is calculated from def_levels:
    /// ```ignore
    /// null_bitmap[i] = (def_levels[i] >= self.def_level);
    /// ```
    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
        if self.children.is_empty() {
            self.def_level_buffer = None;
            self.rep_level_buffer = None;
            return Ok(Arc::new(StructArray::from(Vec::new())));
        }

        let children_array = self
            .children
            .iter_mut()
            .map(|reader| reader.next_batch(batch_size))
            .try_fold(
                Vec::new(),
                |mut result, child_array| -> Result<Vec<ArrayRef>> {
                    result.push(child_array?);
                    Ok(result)
                },
            )?;

        // check that array child data has same size
        let children_array_len =
            children_array.first().map(|arr| arr.len()).ok_or_else(|| {
                general_err!("Struct array reader should have at least one child!")
            })?;

        let all_children_len_eq = children_array
            .iter()
            .all(|arr| arr.len() == children_array_len);
        if !all_children_len_eq {
            return Err(general_err!("Not all children array length are the same!"));
        }

        // Now we can build array data
        let mut array_data_builder = ArrayDataBuilder::new(self.data_type.clone())
            .len(children_array_len)
            .child_data(
                children_array
                    .iter()
                    .map(|x| x.data().clone())
                    .collect::<Vec<ArrayData>>(),
            );

        if self.struct_def_level != 0 {
            // calculate struct def level data
            let buffer_size = children_array_len * size_of::<i16>();
            let mut def_level_data_buffer = MutableBuffer::new(buffer_size);
            def_level_data_buffer.resize(buffer_size, 0);

            // Safety: the buffer is always treated as `u16` in the code below
            let def_level_data = unsafe { def_level_data_buffer.typed_data_mut() };

            def_level_data
                .iter_mut()
                .for_each(|v| *v = self.struct_def_level);

            for child in &self.children {
                if let Some(current_child_def_levels) = child.get_def_levels() {
                    if current_child_def_levels.len() != children_array_len {
                        return Err(general_err!("Child array length are not equal!"));
                    } else {
                        for i in 0..children_array_len {
                            def_level_data[i] =
                                min(def_level_data[i], current_child_def_levels[i]);
                        }
                    }
                }
            }

            // calculate bitmap for current array
            let mut bitmap_builder = BooleanBufferBuilder::new(children_array_len);
            for def_level in def_level_data {
                let not_null = *def_level >= self.struct_def_level;
                bitmap_builder.append(not_null);
            }

            array_data_builder =
                array_data_builder.null_bit_buffer(bitmap_builder.finish());

            self.def_level_buffer = Some(def_level_data_buffer.into());
        }

        let array_data = unsafe { array_data_builder.build_unchecked() };

        if self.struct_rep_level != 0 {
            // calculate struct rep level data, since struct doesn't add to repetition
            // levels, here we just need to keep repetition levels of first array
            // TODO: Verify that all children array reader has same repetition levels
            let rep_level_data = self
                .children
                .first()
                .ok_or_else(|| {
                    general_err!("Struct array reader should have at least one child!")
                })?
                .get_rep_levels()
                .map(|data| -> Result<Buffer> {
                    let mut buffer = Int16BufferBuilder::new(children_array_len);
                    buffer.append_slice(data);
                    Ok(buffer.finish())
                })
                .transpose()?;

            self.rep_level_buffer = rep_level_data;
        }
        Ok(Arc::new(StructArray::from(array_data)))
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.def_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.rep_level_buffer
            .as_ref()
            .map(|buf| unsafe { buf.typed_data() })
    }
}

/// Create array reader from parquet schema, column indices, and parquet file reader.
pub fn build_array_reader<T>(
    parquet_schema: SchemaDescPtr,
    arrow_schema: Schema,
    column_indices: T,
    file_reader: Arc<dyn FileReader>,
) -> Result<Box<dyn ArrayReader>>
where
    T: IntoIterator<Item = usize>,
{
    let mut leaves = HashMap::<*const Type, usize>::new();

    let mut filtered_root_names = HashSet::<String>::new();

    for c in column_indices {
        let column = parquet_schema.column(c).self_type() as *const Type;

        leaves.insert(column, c);

        let root = parquet_schema.get_column_root_ptr(c);
        filtered_root_names.insert(root.name().to_string());
    }

    if leaves.is_empty() {
        return Err(general_err!("Can't build array reader without columns!"));
    }

    // Only pass root fields that take part in the projection
    // to avoid traversal of columns that are not read.
    // TODO: also prune unread parts of the tree in child structures
    let filtered_root_fields = parquet_schema
        .root_schema()
        .get_fields()
        .iter()
        .filter(|field| filtered_root_names.contains(field.name()))
        .cloned()
        .collect::<Vec<_>>();

    let proj = Type::GroupType {
        basic_info: parquet_schema.root_schema().get_basic_info().clone(),
        fields: filtered_root_fields,
    };

    ArrayReaderBuilder::new(
        Arc::new(proj),
        Arc::new(arrow_schema),
        Arc::new(leaves),
        file_reader,
    )
    .build_array_reader()
}

/// Used to build array reader.
struct ArrayReaderBuilder {
    root_schema: TypePtr,
    arrow_schema: Arc<Schema>,
    // Key: columns that need to be included in final array builder
    // Value: column index in schema
    columns_included: Arc<HashMap<*const Type, usize>>,
    file_reader: Arc<dyn FileReader>,
}

/// Used in type visitor.
#[derive(Clone)]
struct ArrayReaderBuilderContext {
    def_level: i16,
    rep_level: i16,
    path: ColumnPath,
}

impl Default for ArrayReaderBuilderContext {
    fn default() -> Self {
        Self {
            def_level: 0i16,
            rep_level: 0i16,
            path: ColumnPath::new(Vec::new()),
        }
    }
}

/// Create array reader by visiting schema.
impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext>
    for ArrayReaderBuilder
{
    /// Build array reader for primitive type.
    fn visit_primitive(
        &mut self,
        cur_type: TypePtr,
        context: &'a ArrayReaderBuilderContext,
    ) -> Result<Option<Box<dyn ArrayReader>>> {
        if self.is_included(cur_type.as_ref()) {
            let mut new_context = context.clone();
            new_context.path.append(vec![cur_type.name().to_string()]);

            let null_mask_only = match cur_type.get_basic_info().repetition() {
                Repetition::REPEATED => {
                    new_context.def_level += 1;
                    new_context.rep_level += 1;
                    false
                }
                Repetition::OPTIONAL => {
                    new_context.def_level += 1;

                    // Can just compute null mask if no parent
                    context.def_level == 0 && context.rep_level == 0
                }
                _ => false,
            };

            let reader = self.build_for_primitive_type_inner(
                cur_type.clone(),
                &new_context,
                null_mask_only,
            )?;

            if cur_type.get_basic_info().repetition() == Repetition::REPEATED {
                Err(ArrowError(
                    "Reading repeated field is not supported yet!".to_string(),
                ))
            } else {
                Ok(Some(reader))
            }
        } else {
            Ok(None)
        }
    }

    /// Build array reader for struct type.
    fn visit_struct(
        &mut self,
        cur_type: Arc<Type>,
        context: &'a ArrayReaderBuilderContext,
    ) -> Result<Option<Box<dyn ArrayReader>>> {
        let mut new_context = context.clone();
        new_context.path.append(vec![cur_type.name().to_string()]);

        if cur_type.get_basic_info().has_repetition() {
            match cur_type.get_basic_info().repetition() {
                Repetition::REPEATED => {
                    new_context.def_level += 1;
                    new_context.rep_level += 1;
                }
                Repetition::OPTIONAL => {
                    new_context.def_level += 1;
                }
                _ => (),
            }
        }

        if let Some(reader) = self.build_for_struct_type_inner(&cur_type, &new_context)? {
            if cur_type.get_basic_info().has_repetition()
                && cur_type.get_basic_info().repetition() == Repetition::REPEATED
            {
                Err(ArrowError(
                    "Reading repeated field is not supported yet!".to_string(),
                ))
            } else {
                Ok(Some(reader))
            }
        } else {
            Ok(None)
        }
    }

    /// Build array reader for map type.
    fn visit_map(
        &mut self,
        map_type: Arc<Type>,
        context: &'a ArrayReaderBuilderContext,
    ) -> Result<Option<Box<dyn ArrayReader>>> {
        // Add map type to context
        let mut new_context = context.clone();
        new_context.path.append(vec![map_type.name().to_string()]);
        if let Repetition::OPTIONAL = map_type.get_basic_info().repetition() {
            new_context.def_level += 1;
        }

        // Add map entry (key_value) to context
        let map_key_value = map_type.get_fields().first().ok_or_else(|| {
            ArrowError("Map field must have a key_value entry".to_string())
        })?;
        new_context
            .path
            .append(vec![map_key_value.name().to_string()]);
        new_context.rep_level += 1;

        // Get key and value, and create context for each
        let map_key = map_key_value
            .get_fields()
            .first()
            .ok_or_else(|| ArrowError("Map entry must have a key".to_string()))?;
        let map_value = map_key_value
            .get_fields()
            .get(1)
            .ok_or_else(|| ArrowError("Map entry must have a value".to_string()))?;

        let key_reader = {
            let mut key_context = new_context.clone();
            key_context.def_level += 1;
            key_context.path.append(vec![map_key.name().to_string()]);
            self.dispatch(map_key.clone(), &key_context)?.unwrap()
        };
        let value_reader = {
            let mut value_context = new_context.clone();
            if let Repetition::OPTIONAL = map_value.get_basic_info().repetition() {
                value_context.def_level += 1;
            }
            self.dispatch(map_value.clone(), &value_context)?.unwrap()
        };

        let arrow_type = self
            .arrow_schema
            .field_with_name(map_type.name())
            .ok()
            .map(|f| f.data_type().to_owned())
            .unwrap_or_else(|| {
                ArrowType::Map(
                    Box::new(Field::new(
                        map_key_value.name(),
                        ArrowType::Struct(vec![
                            Field::new(
                                map_key.name(),
                                key_reader.get_data_type().clone(),
                                false,
                            ),
                            Field::new(
                                map_value.name(),
                                value_reader.get_data_type().clone(),
                                map_value.is_optional(),
                            ),
                        ]),
                        map_type.is_optional(),
                    )),
                    false,
                )
            });

        let key_array_reader: Box<dyn ArrayReader> = Box::new(MapArrayReader::new(
            key_reader,
            value_reader,
            arrow_type,
            new_context.def_level,
            new_context.rep_level,
        ));

        Ok(Some(key_array_reader))
    }

    /// Build array reader for list type.
    fn visit_list_with_item(
        &mut self,
        list_type: Arc<Type>,
        item_type: Arc<Type>,
        context: &'a ArrayReaderBuilderContext,
    ) -> Result<Option<Box<dyn ArrayReader>>> {
        let mut list_child = &list_type
            .get_fields()
            .first()
            .ok_or_else(|| ArrowError("List field must have a child.".to_string()))?
            .clone();
        let mut new_context = context.clone();

        new_context.path.append(vec![list_type.name().to_string()]);
        // We need to know at what definition a list or its child is null
        let list_null_def = new_context.def_level;
        let mut list_empty_def = new_context.def_level;

        // If the list's root is nullable
        if let Repetition::OPTIONAL = list_type.get_basic_info().repetition() {
            new_context.def_level += 1;
            // current level is nullable, increment to get level for empty list slot
            list_empty_def += 1;
        }

        match list_child.get_basic_info().repetition() {
            Repetition::REPEATED => {
                new_context.def_level += 1;
                new_context.rep_level += 1;
            }
            Repetition::OPTIONAL => {
                new_context.def_level += 1;
            }
            _ => (),
        }

        let item_reader = self
            .dispatch(item_type.clone(), &new_context)
            .unwrap()
            .unwrap();

        let item_reader_type = item_reader.get_data_type().clone();

        match item_reader_type {
            ArrowType::List(_)
            | ArrowType::FixedSizeList(_, _)
            | ArrowType::Dictionary(_, _) => Err(ArrowError(format!(
                "reading List({:?}) into arrow not supported yet",
                item_type
            ))),
            _ => {
                // a list is a group type with a single child. The list child's
                // name comes from the child's field name.
                // if the child's name is "list" and it has a child, then use this child
                if list_child.name() == "list" && !list_child.get_fields().is_empty() {
                    list_child = list_child.get_fields().first().unwrap();
                }
                let arrow_type = self
                    .arrow_schema
                    .field_with_name(list_type.name())
                    .ok()
                    .map(|f| f.data_type().to_owned())
                    .unwrap_or_else(|| {
                        ArrowType::List(Box::new(Field::new(
                            list_child.name(),
                            item_reader_type.clone(),
                            list_child.is_optional(),
                        )))
                    });

                let list_array_reader: Box<dyn ArrayReader> = match arrow_type {
                    ArrowType::List(_) => Box::new(ListArrayReader::<i32>::new(
                        item_reader,
                        arrow_type,
                        item_reader_type,
                        new_context.def_level,
                        new_context.rep_level,
                        list_null_def,
                        list_empty_def,
                    )),
                    ArrowType::LargeList(_) => Box::new(ListArrayReader::<i64>::new(
                        item_reader,
                        arrow_type,
                        item_reader_type,
                        new_context.def_level,
                        new_context.rep_level,
                        list_null_def,
                        list_empty_def,
                    )),

                    _ => {
                        return Err(ArrowError(format!(
                        "creating ListArrayReader with type {:?} should be unreachable",
                        arrow_type
                    )))
                    }
                };

                Ok(Some(list_array_reader))
            }
        }
    }
}

impl<'a> ArrayReaderBuilder {
    /// Construct array reader builder.
    fn new(
        root_schema: TypePtr,
        arrow_schema: Arc<Schema>,
        columns_included: Arc<HashMap<*const Type, usize>>,
        file_reader: Arc<dyn FileReader>,
    ) -> Self {
        Self {
            root_schema,
            arrow_schema,
            columns_included,
            file_reader,
        }
    }

    /// Main entry point.
    fn build_array_reader(&mut self) -> Result<Box<dyn ArrayReader>> {
        let context = ArrayReaderBuilderContext::default();

        self.visit_struct(self.root_schema.clone(), &context)
            .and_then(|reader_opt| {
                reader_opt.ok_or_else(|| general_err!("Failed to build array reader!"))
            })
    }

    // Utility functions

    /// Check whether one column in included in this array reader builder.
    fn is_included(&self, t: &Type) -> bool {
        self.columns_included.contains_key(&(t as *const Type))
    }

    /// Creates primitive array reader for each primitive type.
    fn build_for_primitive_type_inner(
        &self,
        cur_type: TypePtr,
        context: &'a ArrayReaderBuilderContext,
        null_mask_only: bool,
    ) -> Result<Box<dyn ArrayReader>> {
        let column_desc = Arc::new(ColumnDescriptor::new(
            cur_type.clone(),
            context.def_level,
            context.rep_level,
            context.path.clone(),
        ));
        let page_iterator = Box::new(FilePageIterator::new(
            self.columns_included[&(cur_type.as_ref() as *const Type)],
            self.file_reader.clone(),
        )?);

        let arrow_type: Option<ArrowType> = self
            .get_arrow_field(&cur_type, context)
            .map(|f| f.data_type().clone());

        match cur_type.get_physical_type() {
            PhysicalType::BOOLEAN => Ok(Box::new(
                PrimitiveArrayReader::<BoolType>::new_with_options(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                )?,
            )),
            PhysicalType::INT32 => {
                if let Some(ArrowType::Null) = arrow_type {
                    Ok(Box::new(NullArrayReader::<Int32Type>::new(
                        page_iterator,
                        column_desc,
                    )?))
                } else {
                    Ok(Box::new(
                        PrimitiveArrayReader::<Int32Type>::new_with_options(
                            page_iterator,
                            column_desc,
                            arrow_type,
                            null_mask_only,
                        )?,
                    ))
                }
            }
            PhysicalType::INT64 => Ok(Box::new(
                PrimitiveArrayReader::<Int64Type>::new_with_options(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                )?,
            )),
            PhysicalType::INT96 => {
                // get the optional timezone information from arrow type
                let timezone = arrow_type
                    .as_ref()
                    .map(|data_type| {
                        if let ArrowType::Timestamp(_, tz) = data_type {
                            tz.clone()
                        } else {
                            None
                        }
                    })
                    .flatten();
                let converter = Int96Converter::new(Int96ArrayConverter { timezone });
                Ok(Box::new(ComplexObjectArrayReader::<
                    Int96Type,
                    Int96Converter,
                >::new(
                    page_iterator,
                    column_desc,
                    converter,
                    arrow_type,
                )?))
            }
            PhysicalType::FLOAT => Ok(Box::new(
                PrimitiveArrayReader::<FloatType>::new_with_options(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                )?,
            )),
            PhysicalType::DOUBLE => Ok(Box::new(
                PrimitiveArrayReader::<DoubleType>::new_with_options(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                )?,
            )),
            PhysicalType::BYTE_ARRAY => match arrow_type {
                Some(ArrowType::Dictionary(_, _)) => make_byte_array_dictionary_reader(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                ),
                _ => make_byte_array_reader(
                    page_iterator,
                    column_desc,
                    arrow_type,
                    null_mask_only,
                ),
            },
            PhysicalType::FIXED_LEN_BYTE_ARRAY
                if cur_type.get_basic_info().converted_type()
                    == ConvertedType::DECIMAL =>
            {
                let converter = DecimalConverter::new(DecimalArrayConverter::new(
                    cur_type.get_precision(),
                    cur_type.get_scale(),
                ));
                Ok(Box::new(ComplexObjectArrayReader::<
                    FixedLenByteArrayType,
                    DecimalConverter,
                >::new(
                    page_iterator,
                    column_desc,
                    converter,
                    arrow_type,
                )?))
            }
            PhysicalType::FIXED_LEN_BYTE_ARRAY => {
                let byte_width = match *cur_type {
                    Type::PrimitiveType {
                        ref type_length, ..
                    } => *type_length,
                    _ => {
                        return Err(ArrowError(
                            "Expected a physical type, not a group type".to_string(),
                        ))
                    }
                };
                if cur_type.get_basic_info().converted_type() == ConvertedType::INTERVAL {
                    if byte_width != 12 {
                        return Err(ArrowError(format!(
                            "Parquet interval type should have length of 12, found {}",
                            byte_width
                        )));
                    }
                    match arrow_type {
                        Some(ArrowType::Interval(IntervalUnit::DayTime)) => {
                            let converter = IntervalDayTimeConverter::new(
                                IntervalDayTimeArrayConverter {},
                            );
                            Ok(Box::new(ComplexObjectArrayReader::<
                                FixedLenByteArrayType,
                                IntervalDayTimeConverter,
                            >::new(
                                page_iterator,
                                column_desc,
                                converter,
                                arrow_type,
                            )?))
                        }
                        Some(ArrowType::Interval(IntervalUnit::YearMonth)) => {
                            let converter = IntervalYearMonthConverter::new(
                                IntervalYearMonthArrayConverter {},
                            );
                            Ok(Box::new(ComplexObjectArrayReader::<
                                FixedLenByteArrayType,
                                IntervalYearMonthConverter,
                            >::new(
                                page_iterator,
                                column_desc,
                                converter,
                                arrow_type,
                            )?))
                        }
                        Some(t) => Err(ArrowError(format!(
                            "Cannot write a Parquet interval to {:?}",
                            t
                        ))),
                        None => {
                            // we do not support an interval not matched to an Arrow type,
                            // because we risk data loss as we won't know which of the 12 bytes
                            // are or should be populated
                            Err(ArrowError(
                                "Cannot write a Parquet interval with no Arrow type specified.
                                There is a risk of data loss as Arrow either supports YearMonth or
                                DayTime precision. Without the Arrow type, we cannot infer the type.
                                ".to_string()
                            ))
                        }
                    }
                } else {
                    let converter = FixedLenBinaryConverter::new(
                        FixedSizeArrayConverter::new(byte_width),
                    );
                    Ok(Box::new(ComplexObjectArrayReader::<
                        FixedLenByteArrayType,
                        FixedLenBinaryConverter,
                    >::new(
                        page_iterator,
                        column_desc,
                        converter,
                        arrow_type,
                    )?))
                }
            }
        }
    }

    /// Constructs struct array reader without considering repetition.
    fn build_for_struct_type_inner(
        &mut self,
        cur_type: &Type,
        context: &'a ArrayReaderBuilderContext,
    ) -> Result<Option<Box<dyn ArrayReader>>> {
        let mut fields = Vec::with_capacity(cur_type.get_fields().len());
        let mut children_reader = Vec::with_capacity(cur_type.get_fields().len());

        for child in cur_type.get_fields() {
            let mut struct_context = context.clone();
            if let Some(child_reader) = self.dispatch(child.clone(), context)? {
                // TODO: this results in calling get_arrow_field twice, it could be reused
                // from child_reader above, by making child_reader carry its `Field`
                struct_context.path.append(vec![child.name().to_string()]);
                let field = match self.get_arrow_field(child, &struct_context) {
                    Some(f) => f.clone(),
                    _ => Field::new(
                        child.name(),
                        child_reader.get_data_type().clone(),
                        child.is_optional(),
                    ),
                };
                fields.push(field);
                children_reader.push(child_reader);
            }
        }

        if !fields.is_empty() {
            let arrow_type = ArrowType::Struct(fields);
            Ok(Some(Box::new(StructArrayReader::new(
                arrow_type,
                children_reader,
                context.def_level,
                context.rep_level,
            ))))
        } else {
            Ok(None)
        }
    }

    fn get_arrow_field(
        &self,
        cur_type: &Type,
        context: &'a ArrayReaderBuilderContext,
    ) -> Option<&Field> {
        let parts: Vec<&str> = context
            .path
            .parts()
            .iter()
            .map(|x| -> &str { x })
            .collect::<Vec<&str>>();

        // If the parts length is one it'll have the top level "schema" type. If
        // it's two then it'll be a top-level type that we can get from the arrow
        // schema directly.
        if parts.len() <= 2 {
            self.arrow_schema.field_with_name(cur_type.name()).ok()
        } else {
            // If it's greater than two then we need to traverse the type path
            // until we find the actual field we're looking for.
            let mut field: Option<&Field> = None;

            for (i, part) in parts.iter().enumerate().skip(1) {
                if i == 1 {
                    field = self.arrow_schema.field_with_name(part).ok();
                } else if let Some(f) = field {
                    if let ArrowType::Struct(fields) = f.data_type() {
                        field = fields.iter().find(|f| f.name() == part)
                    } else {
                        field = None
                    }
                } else {
                    field = None
                }
            }
            field
        }
    }
}

#[cfg(test)]
mod tests {
    use std::any::Any;
    use std::collections::VecDeque;
    use std::sync::Arc;

    use rand::distributions::uniform::SampleUniform;
    use rand::{thread_rng, Rng};

    use arrow::array::{
        Array, ArrayRef, LargeListArray, ListArray, PrimitiveArray, StringArray,
        StructArray,
    };
    use arrow::datatypes::{
        ArrowPrimitiveType, DataType as ArrowType, Date32Type as ArrowDate32, Field,
        Int32Type as ArrowInt32, Int64Type as ArrowInt64,
        Time32MillisecondType as ArrowTime32MillisecondArray,
        Time64MicrosecondType as ArrowTime64MicrosecondArray,
        TimestampMicrosecondType as ArrowTimestampMicrosecondType,
        TimestampMillisecondType as ArrowTimestampMillisecondType,
    };

    use crate::arrow::converter::{Utf8ArrayConverter, Utf8Converter};
    use crate::arrow::schema::parquet_to_arrow_schema;
    use crate::basic::{Encoding, Type as PhysicalType};
    use crate::column::page::{Page, PageReader};
    use crate::data_type::{ByteArray, ByteArrayType, DataType, Int32Type, Int64Type};
    use crate::errors::Result;
    use crate::file::reader::{FileReader, SerializedFileReader};
    use crate::schema::parser::parse_message_type;
    use crate::schema::types::{ColumnDescPtr, SchemaDescriptor};
    use crate::util::test_common::page_util::{
        DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,
    };
    use crate::util::test_common::{get_test_file, make_pages};

    use super::*;

    fn make_column_chunks<T: DataType>(
        column_desc: ColumnDescPtr,
        encoding: Encoding,
        num_levels: usize,
        min_value: T::T,
        max_value: T::T,
        def_levels: &mut Vec<i16>,
        rep_levels: &mut Vec<i16>,
        values: &mut Vec<T::T>,
        page_lists: &mut Vec<Vec<Page>>,
        use_v2: bool,
        num_chunks: usize,
    ) where
        T::T: PartialOrd + SampleUniform + Copy,
    {
        for _i in 0..num_chunks {
            let mut pages = VecDeque::new();
            let mut data = Vec::new();
            let mut page_def_levels = Vec::new();
            let mut page_rep_levels = Vec::new();

            make_pages::<T>(
                column_desc.clone(),
                encoding,
                1,
                num_levels,
                min_value,
                max_value,
                &mut page_def_levels,
                &mut page_rep_levels,
                &mut data,
                &mut pages,
                use_v2,
            );

            def_levels.append(&mut page_def_levels);
            rep_levels.append(&mut page_rep_levels);
            values.append(&mut data);
            page_lists.push(Vec::from(pages));
        }
    }

    #[test]
    fn test_primitive_array_reader_empty_pages() {
        // Construct column schema
        let message_type = "
        message test_schema {
          REQUIRED INT32 leaf;
        }
        ";

        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();

        let column_desc = schema.column(0);
        let page_iterator = EmptyPageIterator::new(schema);

        let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
            Box::new(page_iterator),
            column_desc,
            None,
        )
        .unwrap();

        // expect no values to be read
        let array = array_reader.next_batch(50).unwrap();
        assert!(array.is_empty());
    }

    #[test]
    fn test_primitive_array_reader_data() {
        // Construct column schema
        let message_type = "
        message test_schema {
          REQUIRED INT32 leaf;
        }
        ";

        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();

        let column_desc = schema.column(0);

        // Construct page iterator
        {
            let mut data = Vec::new();
            let mut page_lists = Vec::new();
            make_column_chunks::<Int32Type>(
                column_desc.clone(),
                Encoding::PLAIN,
                100,
                1,
                200,
                &mut Vec::new(),
                &mut Vec::new(),
                &mut data,
                &mut page_lists,
                true,
                2,
            );
            let page_iterator =
                InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);

            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
                Box::new(page_iterator),
                column_desc,
                None,
            )
            .unwrap();

            // Read first 50 values, which are all from the first column chunk
            let array = array_reader.next_batch(50).unwrap();
            let array = array
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap();

            assert_eq!(
                &PrimitiveArray::<ArrowInt32>::from(data[0..50].to_vec()),
                array
            );

            // Read next 100 values, the first 50 ones are from the first column chunk,
            // and the last 50 ones are from the second column chunk
            let array = array_reader.next_batch(100).unwrap();
            let array = array
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap();

            assert_eq!(
                &PrimitiveArray::<ArrowInt32>::from(data[50..150].to_vec()),
                array
            );

            // Try to read 100 values, however there are only 50 values
            let array = array_reader.next_batch(100).unwrap();
            let array = array
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap();

            assert_eq!(
                &PrimitiveArray::<ArrowInt32>::from(data[150..200].to_vec()),
                array
            );
        }
    }

    macro_rules! test_primitive_array_reader_one_type {
        ($arrow_parquet_type:ty, $physical_type:expr, $converted_type_str:expr, $result_arrow_type:ty, $result_arrow_cast_type:ty, $result_primitive_type:ty) => {{
            let message_type = format!(
                "
            message test_schema {{
              REQUIRED {:?} leaf ({});
          }}
            ",
                $physical_type, $converted_type_str
            );
            let schema = parse_message_type(&message_type)
                .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
                .unwrap();

            let column_desc = schema.column(0);

            // Construct page iterator
            {
                let mut data = Vec::new();
                let mut page_lists = Vec::new();
                make_column_chunks::<$arrow_parquet_type>(
                    column_desc.clone(),
                    Encoding::PLAIN,
                    100,
                    1,
                    200,
                    &mut Vec::new(),
                    &mut Vec::new(),
                    &mut data,
                    &mut page_lists,
                    true,
                    2,
                );
                let page_iterator = InMemoryPageIterator::new(
                    schema.clone(),
                    column_desc.clone(),
                    page_lists,
                );
                let mut array_reader = PrimitiveArrayReader::<$arrow_parquet_type>::new(
                    Box::new(page_iterator),
                    column_desc.clone(),
                    None,
                )
                .expect("Unable to get array reader");

                let array = array_reader
                    .next_batch(50)
                    .expect("Unable to get batch from reader");

                let result_data_type = <$result_arrow_type>::DATA_TYPE;
                let array = array
                    .as_any()
                    .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
                    .expect(
                        format!(
                            "Unable to downcast {:?} to {:?}",
                            array.data_type(),
                            result_data_type
                        )
                        .as_str(),
                    );

                // create expected array as primitive, and cast to result type
                let expected = PrimitiveArray::<$result_arrow_cast_type>::from(
                    data[0..50]
                        .iter()
                        .map(|x| *x as $result_primitive_type)
                        .collect::<Vec<$result_primitive_type>>(),
                );
                let expected = Arc::new(expected) as ArrayRef;
                let expected = arrow::compute::cast(&expected, &result_data_type)
                    .expect("Unable to cast expected array");
                assert_eq!(expected.data_type(), &result_data_type);
                let expected = expected
                    .as_any()
                    .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
                    .expect(
                        format!(
                            "Unable to downcast expected {:?} to {:?}",
                            expected.data_type(),
                            result_data_type
                        )
                        .as_str(),
                    );
                assert_eq!(expected, array);
            }
        }};
    }

    #[test]
    fn test_primitive_array_reader_temporal_types() {
        test_primitive_array_reader_one_type!(
            Int32Type,
            PhysicalType::INT32,
            "DATE",
            ArrowDate32,
            ArrowInt32,
            i32
        );
        test_primitive_array_reader_one_type!(
            Int32Type,
            PhysicalType::INT32,
            "TIME_MILLIS",
            ArrowTime32MillisecondArray,
            ArrowInt32,
            i32
        );
        test_primitive_array_reader_one_type!(
            Int64Type,
            PhysicalType::INT64,
            "TIME_MICROS",
            ArrowTime64MicrosecondArray,
            ArrowInt64,
            i64
        );
        test_primitive_array_reader_one_type!(
            Int64Type,
            PhysicalType::INT64,
            "TIMESTAMP_MILLIS",
            ArrowTimestampMillisecondType,
            ArrowInt64,
            i64
        );
        test_primitive_array_reader_one_type!(
            Int64Type,
            PhysicalType::INT64,
            "TIMESTAMP_MICROS",
            ArrowTimestampMicrosecondType,
            ArrowInt64,
            i64
        );
    }

    #[test]
    fn test_primitive_array_reader_def_and_rep_levels() {
        // Construct column schema
        let message_type = "
        message test_schema {
            REPEATED Group test_mid {
                OPTIONAL INT32 leaf;
            }
        }
        ";

        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();

        let column_desc = schema.column(0);

        // Construct page iterator
        {
            let mut def_levels = Vec::new();
            let mut rep_levels = Vec::new();
            let mut page_lists = Vec::new();
            make_column_chunks::<Int32Type>(
                column_desc.clone(),
                Encoding::PLAIN,
                100,
                1,
                200,
                &mut def_levels,
                &mut rep_levels,
                &mut Vec::new(),
                &mut page_lists,
                true,
                2,
            );

            let page_iterator =
                InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);

            let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
                Box::new(page_iterator),
                column_desc,
                None,
            )
            .unwrap();

            let mut accu_len: usize = 0;

            // Read first 50 values, which are all from the first column chunk
            let array = array_reader.next_batch(50).unwrap();
            assert_eq!(
                Some(&def_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_def_levels()
            );
            assert_eq!(
                Some(&rep_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_rep_levels()
            );
            accu_len += array.len();

            // Read next 100 values, the first 50 ones are from the first column chunk,
            // and the last 50 ones are from the second column chunk
            let array = array_reader.next_batch(100).unwrap();
            assert_eq!(
                Some(&def_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_def_levels()
            );
            assert_eq!(
                Some(&rep_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_rep_levels()
            );
            accu_len += array.len();

            // Try to read 100 values, however there are only 50 values
            let array = array_reader.next_batch(100).unwrap();
            assert_eq!(
                Some(&def_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_def_levels()
            );
            assert_eq!(
                Some(&rep_levels[accu_len..(accu_len + array.len())]),
                array_reader.get_rep_levels()
            );
        }
    }

    #[test]
    fn test_complex_array_reader_no_pages() {
        let message_type = "
        message test_schema {
            REPEATED Group test_mid {
                OPTIONAL BYTE_ARRAY leaf (UTF8);
            }
        }
        ";
        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();
        let column_desc = schema.column(0);
        let pages: Vec<Vec<Page>> = Vec::new();
        let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);

        let converter = Utf8Converter::new(Utf8ArrayConverter {});
        let mut array_reader =
            ComplexObjectArrayReader::<ByteArrayType, Utf8Converter>::new(
                Box::new(page_iterator),
                column_desc,
                converter,
                None,
            )
            .unwrap();

        let values_per_page = 100; // this value is arbitrary in this test - the result should always be an array of 0 length
        let array = array_reader.next_batch(values_per_page).unwrap();
        assert_eq!(array.len(), 0);
    }

    #[test]
    fn test_complex_array_reader_def_and_rep_levels() {
        // Construct column schema
        let message_type = "
        message test_schema {
            REPEATED Group test_mid {
                OPTIONAL BYTE_ARRAY leaf (UTF8);
            }
        }
        ";
        let num_pages = 2;
        let values_per_page = 100;
        let str_base = "Hello World";

        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();

        let max_def_level = schema.column(0).max_def_level();
        let max_rep_level = schema.column(0).max_rep_level();

        assert_eq!(max_def_level, 2);
        assert_eq!(max_rep_level, 1);

        let mut rng = thread_rng();
        let column_desc = schema.column(0);
        let mut pages: Vec<Vec<Page>> = Vec::new();

        let mut rep_levels = Vec::with_capacity(num_pages * values_per_page);
        let mut def_levels = Vec::with_capacity(num_pages * values_per_page);
        let mut all_values = Vec::with_capacity(num_pages * values_per_page);

        for i in 0..num_pages {
            let mut values = Vec::with_capacity(values_per_page);

            for _ in 0..values_per_page {
                let def_level = rng.gen_range(0..max_def_level + 1);
                let rep_level = rng.gen_range(0..max_rep_level + 1);
                if def_level == max_def_level {
                    let len = rng.gen_range(1..str_base.len());
                    let slice = &str_base[..len];
                    values.push(ByteArray::from(slice));
                    all_values.push(Some(slice.to_string()));
                } else {
                    all_values.push(None)
                }
                rep_levels.push(rep_level);
                def_levels.push(def_level)
            }

            let range = i * values_per_page..(i + 1) * values_per_page;
            let mut pb =
                DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);

            pb.add_rep_levels(max_rep_level, &rep_levels.as_slice()[range.clone()]);
            pb.add_def_levels(max_def_level, &def_levels.as_slice()[range]);
            pb.add_values::<ByteArrayType>(Encoding::PLAIN, values.as_slice());

            let data_page = pb.consume();
            pages.push(vec![data_page]);
        }

        let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);

        let converter = Utf8Converter::new(Utf8ArrayConverter {});
        let mut array_reader =
            ComplexObjectArrayReader::<ByteArrayType, Utf8Converter>::new(
                Box::new(page_iterator),
                column_desc,
                converter,
                None,
            )
            .unwrap();

        let mut accu_len: usize = 0;

        let array = array_reader.next_batch(values_per_page / 2).unwrap();
        assert_eq!(array.len(), values_per_page / 2);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
        accu_len += array.len();

        // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
        // and the last values_per_page/2 ones are from the second column chunk
        let array = array_reader.next_batch(values_per_page).unwrap();
        assert_eq!(array.len(), values_per_page);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
        for i in 0..array.len() {
            if array.is_valid(i) {
                assert_eq!(
                    all_values[i + accu_len].as_ref().unwrap().as_str(),
                    strings.value(i)
                )
            } else {
                assert_eq!(all_values[i + accu_len], None)
            }
        }
        accu_len += array.len();

        // Try to read values_per_page values, however there are only values_per_page/2 values
        let array = array_reader.next_batch(values_per_page).unwrap();
        assert_eq!(array.len(), values_per_page / 2);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
    }

    #[test]
    fn test_complex_array_reader_dict_enc_string() {
        use crate::encodings::encoding::{DictEncoder, Encoder};
        use crate::util::memory::MemTracker;
        // Construct column schema
        let message_type = "
        message test_schema {
            REPEATED Group test_mid {
                OPTIONAL BYTE_ARRAY leaf (UTF8);
            }
        }
        ";
        let num_pages = 2;
        let values_per_page = 100;
        let str_base = "Hello World";

        let schema = parse_message_type(message_type)
            .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
            .unwrap();
        let column_desc = schema.column(0);
        let max_def_level = column_desc.max_def_level();
        let max_rep_level = column_desc.max_rep_level();

        assert_eq!(max_def_level, 2);
        assert_eq!(max_rep_level, 1);

        let mut rng = thread_rng();
        let mut pages: Vec<Vec<Page>> = Vec::new();

        let mut rep_levels = Vec::with_capacity(num_pages * values_per_page);
        let mut def_levels = Vec::with_capacity(num_pages * values_per_page);
        let mut all_values = Vec::with_capacity(num_pages * values_per_page);

        for i in 0..num_pages {
            let mem_tracker = Arc::new(MemTracker::new());
            let mut dict_encoder =
                DictEncoder::<ByteArrayType>::new(column_desc.clone(), mem_tracker);
            // add data page
            let mut values = Vec::with_capacity(values_per_page);

            for _ in 0..values_per_page {
                let def_level = rng.gen_range(0..max_def_level + 1);
                let rep_level = rng.gen_range(0..max_rep_level + 1);
                if def_level == max_def_level {
                    let len = rng.gen_range(1..str_base.len());
                    let slice = &str_base[..len];
                    values.push(ByteArray::from(slice));
                    all_values.push(Some(slice.to_string()));
                } else {
                    all_values.push(None)
                }
                rep_levels.push(rep_level);
                def_levels.push(def_level)
            }

            let range = i * values_per_page..(i + 1) * values_per_page;
            let mut pb =
                DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
            pb.add_rep_levels(max_rep_level, &rep_levels.as_slice()[range.clone()]);
            pb.add_def_levels(max_def_level, &def_levels.as_slice()[range]);
            let _ = dict_encoder.put(&values);
            let indices = dict_encoder
                .write_indices()
                .expect("write_indices() should be OK");
            pb.add_indices(indices);
            let data_page = pb.consume();
            // for each page log num_values vs actual values in page
            // println!("page num_values: {}, values.len(): {}", data_page.num_values(), values.len());
            // add dictionary page
            let dict = dict_encoder
                .write_dict()
                .expect("write_dict() should be OK");
            let dict_page = Page::DictionaryPage {
                buf: dict,
                num_values: dict_encoder.num_entries() as u32,
                encoding: Encoding::RLE_DICTIONARY,
                is_sorted: false,
            };
            pages.push(vec![dict_page, data_page]);
        }

        let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);
        let converter = Utf8Converter::new(Utf8ArrayConverter {});
        let mut array_reader =
            ComplexObjectArrayReader::<ByteArrayType, Utf8Converter>::new(
                Box::new(page_iterator),
                column_desc,
                converter,
                None,
            )
            .unwrap();

        let mut accu_len: usize = 0;

        // println!("---------- reading a batch of {} values ----------", values_per_page / 2);
        let array = array_reader.next_batch(values_per_page / 2).unwrap();
        assert_eq!(array.len(), values_per_page / 2);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
        accu_len += array.len();

        // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
        // and the last values_per_page/2 ones are from the second column chunk
        // println!("---------- reading a batch of {} values ----------", values_per_page);
        let array = array_reader.next_batch(values_per_page).unwrap();
        assert_eq!(array.len(), values_per_page);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
        for i in 0..array.len() {
            if array.is_valid(i) {
                assert_eq!(
                    all_values[i + accu_len].as_ref().unwrap().as_str(),
                    strings.value(i)
                )
            } else {
                assert_eq!(all_values[i + accu_len], None)
            }
        }
        accu_len += array.len();

        // Try to read values_per_page values, however there are only values_per_page/2 values
        // println!("---------- reading a batch of {} values ----------", values_per_page);
        let array = array_reader.next_batch(values_per_page).unwrap();
        assert_eq!(array.len(), values_per_page / 2);
        assert_eq!(
            Some(&def_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_def_levels()
        );
        assert_eq!(
            Some(&rep_levels[accu_len..(accu_len + array.len())]),
            array_reader.get_rep_levels()
        );
    }

    /// Array reader for test.
    struct InMemoryArrayReader {
        data_type: ArrowType,
        array: ArrayRef,
        def_levels: Option<Vec<i16>>,
        rep_levels: Option<Vec<i16>>,
    }

    impl InMemoryArrayReader {
        pub fn new(
            data_type: ArrowType,
            array: ArrayRef,
            def_levels: Option<Vec<i16>>,
            rep_levels: Option<Vec<i16>>,
        ) -> Self {
            Self {
                data_type,
                array,
                def_levels,
                rep_levels,
            }
        }
    }

    impl ArrayReader for InMemoryArrayReader {
        fn as_any(&self) -> &dyn Any {
            self
        }

        fn get_data_type(&self) -> &ArrowType {
            &self.data_type
        }

        fn next_batch(&mut self, _batch_size: usize) -> Result<ArrayRef> {
            Ok(self.array.clone())
        }

        fn get_def_levels(&self) -> Option<&[i16]> {
            self.def_levels.as_deref()
        }

        fn get_rep_levels(&self) -> Option<&[i16]> {
            self.rep_levels.as_deref()
        }
    }

    /// Iterator for testing reading empty columns
    struct EmptyPageIterator {
        schema: SchemaDescPtr,
    }

    impl EmptyPageIterator {
        fn new(schema: SchemaDescPtr) -> Self {
            EmptyPageIterator { schema }
        }
    }

    impl Iterator for EmptyPageIterator {
        type Item = Result<Box<dyn PageReader>>;

        fn next(&mut self) -> Option<Self::Item> {
            None
        }
    }

    impl PageIterator for EmptyPageIterator {
        fn schema(&mut self) -> Result<SchemaDescPtr> {
            Ok(self.schema.clone())
        }

        fn column_schema(&mut self) -> Result<ColumnDescPtr> {
            Ok(self.schema.column(0))
        }
    }

    #[test]
    fn test_struct_array_reader() {
        let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));
        let array_reader_1 = InMemoryArrayReader::new(
            ArrowType::Int32,
            array_1.clone(),
            Some(vec![0, 1, 2, 3, 1]),
            Some(vec![1, 1, 1, 1, 1]),
        );

        let array_2 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![5, 4, 3, 2, 1]));
        let array_reader_2 = InMemoryArrayReader::new(
            ArrowType::Int32,
            array_2.clone(),
            Some(vec![0, 1, 3, 1, 2]),
            Some(vec![1, 1, 1, 1, 1]),
        );

        let struct_type = ArrowType::Struct(vec![
            Field::new("f1", array_1.data_type().clone(), true),
            Field::new("f2", array_2.data_type().clone(), true),
        ]);

        let mut struct_array_reader = StructArrayReader::new(
            struct_type,
            vec![Box::new(array_reader_1), Box::new(array_reader_2)],
            1,
            1,
        );

        let struct_array = struct_array_reader.next_batch(5).unwrap();
        let struct_array = struct_array.as_any().downcast_ref::<StructArray>().unwrap();

        assert_eq!(5, struct_array.len());
        assert_eq!(
            vec![true, false, false, false, false],
            (0..5)
                .map(|idx| struct_array.data_ref().is_null(idx))
                .collect::<Vec<bool>>()
        );
        assert_eq!(
            Some(vec![0, 1, 1, 1, 1].as_slice()),
            struct_array_reader.get_def_levels()
        );
        assert_eq!(
            Some(vec![1, 1, 1, 1, 1].as_slice()),
            struct_array_reader.get_rep_levels()
        );
    }

    #[test]
    fn test_create_array_reader() {
        let file = get_test_file("nulls.snappy.parquet");
        let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());

        let file_metadata = file_reader.metadata().file_metadata();
        let arrow_schema = parquet_to_arrow_schema(
            file_metadata.schema_descr(),
            file_metadata.key_value_metadata(),
        )
        .unwrap();

        let array_reader = build_array_reader(
            file_reader.metadata().file_metadata().schema_descr_ptr(),
            arrow_schema,
            vec![0usize].into_iter(),
            file_reader,
        )
        .unwrap();

        // Create arrow types
        let arrow_type = ArrowType::Struct(vec![Field::new(
            "b_struct",
            ArrowType::Struct(vec![Field::new("b_c_int", ArrowType::Int32, true)]),
            true,
        )]);

        assert_eq!(array_reader.get_data_type(), &arrow_type);
    }

    #[test]
    fn test_list_array_reader() {
        // [[1, null, 2], null, [3, 4]]
        let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
            Some(1),
            None,
            Some(2),
            None,
            Some(3),
            Some(4),
        ]));
        let item_array_reader = InMemoryArrayReader::new(
            ArrowType::Int32,
            array,
            Some(vec![3, 2, 3, 0, 3, 3]),
            Some(vec![0, 1, 1, 0, 0, 1]),
        );

        let mut list_array_reader = ListArrayReader::<i32>::new(
            Box::new(item_array_reader),
            ArrowType::List(Box::new(Field::new("item", ArrowType::Int32, true))),
            ArrowType::Int32,
            1,
            1,
            0,
            1,
        );

        let next_batch = list_array_reader.next_batch(1024).unwrap();
        let list_array = next_batch.as_any().downcast_ref::<ListArray>().unwrap();

        assert_eq!(3, list_array.len());
        // This passes as I expect
        assert_eq!(1, list_array.null_count());

        assert_eq!(
            list_array
                .value(0)
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap(),
            &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
        );

        assert!(list_array.is_null(1));

        assert_eq!(
            list_array
                .value(2)
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap(),
            &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
        );
    }

    #[test]
    fn test_large_list_array_reader() {
        // [[1, null, 2], null, [3, 4]]
        let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
            Some(1),
            None,
            Some(2),
            None,
            Some(3),
            Some(4),
        ]));
        let item_array_reader = InMemoryArrayReader::new(
            ArrowType::Int32,
            array,
            Some(vec![3, 2, 3, 0, 3, 3]),
            Some(vec![0, 1, 1, 0, 0, 1]),
        );

        let mut list_array_reader = ListArrayReader::<i64>::new(
            Box::new(item_array_reader),
            ArrowType::LargeList(Box::new(Field::new("item", ArrowType::Int32, true))),
            ArrowType::Int32,
            1,
            1,
            0,
            1,
        );

        let next_batch = list_array_reader.next_batch(1024).unwrap();
        let list_array = next_batch
            .as_any()
            .downcast_ref::<LargeListArray>()
            .unwrap();

        assert_eq!(3, list_array.len());

        assert_eq!(
            list_array
                .value(0)
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap(),
            &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
        );

        assert!(list_array.is_null(1));

        assert_eq!(
            list_array
                .value(2)
                .as_any()
                .downcast_ref::<PrimitiveArray<ArrowInt32>>()
                .unwrap(),
            &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
        );
    }
}
