parquet/schema/

types.rs

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

//! Contains structs and methods to build Parquet schema and schema descriptors.

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

use crate::file::metadata::HeapSize;
use crate::format::SchemaElement;

use crate::basic::{
    ColumnOrder, ConvertedType, LogicalType, Repetition, SortOrder, TimeUnit, Type as PhysicalType,
};
use crate::errors::{ParquetError, Result};

// ----------------------------------------------------------------------
// Parquet Type definitions

/// Type alias for `Arc<Type>`.
pub type TypePtr = Arc<Type>;
/// Type alias for `Arc<SchemaDescriptor>`.
pub type SchemaDescPtr = Arc<SchemaDescriptor>;
/// Type alias for `Arc<ColumnDescriptor>`.
pub type ColumnDescPtr = Arc<ColumnDescriptor>;

/// Representation of a Parquet type.
///
/// Used to describe primitive leaf fields and structs, including top-level schema.
///
/// Note that the top-level schema is represented using [`Type::GroupType`] whose
/// repetition is `None`.
#[derive(Clone, Debug, PartialEq)]
pub enum Type {
    /// Represents a primitive leaf field.
    PrimitiveType {
        /// Basic information about the type.
        basic_info: BasicTypeInfo,
        /// Physical type of this primitive type.
        physical_type: PhysicalType,
        /// Length of this type.
        type_length: i32,
        /// Scale of this type.
        scale: i32,
        /// Precision of this type.
        precision: i32,
    },
    /// Represents a group of fields (similar to struct).
    GroupType {
        /// Basic information about the type.
        basic_info: BasicTypeInfo,
        /// Fields of this group type.
        fields: Vec<TypePtr>,
    },
}

impl HeapSize for Type {
    fn heap_size(&self) -> usize {
        match self {
            Type::PrimitiveType { basic_info, .. } => basic_info.heap_size(),
            Type::GroupType { basic_info, fields } => basic_info.heap_size() + fields.heap_size(),
        }
    }
}

impl Type {
    /// Creates primitive type builder with provided field name and physical type.
    pub fn primitive_type_builder(name: &str, physical_type: PhysicalType) -> PrimitiveTypeBuilder {
        PrimitiveTypeBuilder::new(name, physical_type)
    }

    /// Creates group type builder with provided column name.
    pub fn group_type_builder(name: &str) -> GroupTypeBuilder {
        GroupTypeBuilder::new(name)
    }

    /// Returns [`BasicTypeInfo`] information about the type.
    pub fn get_basic_info(&self) -> &BasicTypeInfo {
        match *self {
            Type::PrimitiveType { ref basic_info, .. } => basic_info,
            Type::GroupType { ref basic_info, .. } => basic_info,
        }
    }

    /// Returns this type's field name.
    pub fn name(&self) -> &str {
        self.get_basic_info().name()
    }

    /// Gets the fields from this group type.
    /// Note that this will panic if called on a non-group type.
    // TODO: should we return `&[&Type]` here?
    pub fn get_fields(&self) -> &[TypePtr] {
        match *self {
            Type::GroupType { ref fields, .. } => &fields[..],
            _ => panic!("Cannot call get_fields() on a non-group type"),
        }
    }

    /// Gets physical type of this primitive type.
    /// Note that this will panic if called on a non-primitive type.
    pub fn get_physical_type(&self) -> PhysicalType {
        match *self {
            Type::PrimitiveType {
                basic_info: _,
                physical_type,
                ..
            } => physical_type,
            _ => panic!("Cannot call get_physical_type() on a non-primitive type"),
        }
    }

    /// Gets precision of this primitive type.
    /// Note that this will panic if called on a non-primitive type.
    pub fn get_precision(&self) -> i32 {
        match *self {
            Type::PrimitiveType { precision, .. } => precision,
            _ => panic!("Cannot call get_precision() on non-primitive type"),
        }
    }

    /// Gets scale of this primitive type.
    /// Note that this will panic if called on a non-primitive type.
    pub fn get_scale(&self) -> i32 {
        match *self {
            Type::PrimitiveType { scale, .. } => scale,
            _ => panic!("Cannot call get_scale() on non-primitive type"),
        }
    }

    /// Checks if `sub_type` schema is part of current schema.
    /// This method can be used to check if projected columns are part of the root schema.
    pub fn check_contains(&self, sub_type: &Type) -> bool {
        // Names match, and repetitions match or not set for both
        let basic_match = self.get_basic_info().name() == sub_type.get_basic_info().name()
            && (self.is_schema() && sub_type.is_schema()
                || !self.is_schema()
                    && !sub_type.is_schema()
                    && self.get_basic_info().repetition()
                        == sub_type.get_basic_info().repetition());

        match *self {
            Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => {
                self.get_physical_type() == sub_type.get_physical_type()
            }
            Type::GroupType { .. } if basic_match && sub_type.is_group() => {
                // build hashmap of name -> TypePtr
                let mut field_map = HashMap::new();
                for field in self.get_fields() {
                    field_map.insert(field.name(), field);
                }

                for field in sub_type.get_fields() {
                    if !field_map
                        .get(field.name())
                        .map(|tpe| tpe.check_contains(field))
                        .unwrap_or(false)
                    {
                        return false;
                    }
                }
                true
            }
            _ => false,
        }
    }

    /// Returns `true` if this type is a primitive type, `false` otherwise.
    pub fn is_primitive(&self) -> bool {
        matches!(*self, Type::PrimitiveType { .. })
    }

    /// Returns `true` if this type is a group type, `false` otherwise.
    pub fn is_group(&self) -> bool {
        matches!(*self, Type::GroupType { .. })
    }

    /// Returns `true` if this type is the top-level schema type (message type).
    pub fn is_schema(&self) -> bool {
        match *self {
            Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(),
            _ => false,
        }
    }

    /// Returns `true` if this type is repeated or optional.
    /// If this type doesn't have repetition defined, we treat it as required.
    pub fn is_optional(&self) -> bool {
        self.get_basic_info().has_repetition()
            && self.get_basic_info().repetition() != Repetition::REQUIRED
    }
}

/// A builder for primitive types. All attributes are optional
/// except the name and physical type.
/// Note that if not specified explicitly, `Repetition::OPTIONAL` is used.
pub struct PrimitiveTypeBuilder<'a> {
    name: &'a str,
    repetition: Repetition,
    physical_type: PhysicalType,
    converted_type: ConvertedType,
    logical_type: Option<LogicalType>,
    length: i32,
    precision: i32,
    scale: i32,
    id: Option<i32>,
}

impl<'a> PrimitiveTypeBuilder<'a> {
    /// Creates new primitive type builder with provided field name and physical type.
    pub fn new(name: &'a str, physical_type: PhysicalType) -> Self {
        Self {
            name,
            repetition: Repetition::OPTIONAL,
            physical_type,
            converted_type: ConvertedType::NONE,
            logical_type: None,
            length: -1,
            precision: -1,
            scale: -1,
            id: None,
        }
    }

    /// Sets [`Repetition`] for this field and returns itself.
    pub fn with_repetition(self, repetition: Repetition) -> Self {
        Self { repetition, ..self }
    }

    /// Sets [`ConvertedType`] for this field and returns itself.
    pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
        Self {
            converted_type,
            ..self
        }
    }

    /// Sets [`LogicalType`] for this field and returns itself.
    /// If only the logical type is populated for a primitive type, the converted type
    /// will be automatically populated, and can thus be omitted.
    pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
        Self {
            logical_type,
            ..self
        }
    }

    /// Sets type length and returns itself.
    /// This is only applied to FIXED_LEN_BYTE_ARRAY and INT96 (INTERVAL) types, because
    /// they maintain fixed size underlying byte array.
    /// By default, value is `0`.
    pub fn with_length(self, length: i32) -> Self {
        Self { length, ..self }
    }

    /// Sets precision for Parquet DECIMAL physical type and returns itself.
    /// By default, it equals to `0` and used only for decimal context.
    pub fn with_precision(self, precision: i32) -> Self {
        Self { precision, ..self }
    }

    /// Sets scale for Parquet DECIMAL physical type and returns itself.
    /// By default, it equals to `0` and used only for decimal context.
    pub fn with_scale(self, scale: i32) -> Self {
        Self { scale, ..self }
    }

    /// Sets optional field id and returns itself.
    pub fn with_id(self, id: Option<i32>) -> Self {
        Self { id, ..self }
    }

    /// Creates a new `PrimitiveType` instance from the collected attributes.
    /// Returns `Err` in case of any building conditions are not met.
    pub fn build(self) -> Result<Type> {
        let mut basic_info = BasicTypeInfo {
            name: String::from(self.name),
            repetition: Some(self.repetition),
            converted_type: self.converted_type,
            logical_type: self.logical_type.clone(),
            id: self.id,
        };

        // Check length before logical type, since it is used for logical type validation.
        if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 {
            return Err(general_err!(
                "Invalid FIXED_LEN_BYTE_ARRAY length: {} for field '{}'",
                self.length,
                self.name
            ));
        }

        if let Some(logical_type) = &self.logical_type {
            // If a converted type is populated, check that it is consistent with
            // its logical type
            if self.converted_type != ConvertedType::NONE {
                if ConvertedType::from(self.logical_type.clone()) != self.converted_type {
                    return Err(general_err!(
                        "Logical type {:?} is incompatible with converted type {} for field '{}'",
                        logical_type,
                        self.converted_type,
                        self.name
                    ));
                }
            } else {
                // Populate the converted type for backwards compatibility
                basic_info.converted_type = self.logical_type.clone().into();
            }
            // Check that logical type and physical type are compatible
            match (logical_type, self.physical_type) {
                (LogicalType::Map, _) | (LogicalType::List, _) => {
                    return Err(general_err!(
                        "{:?} cannot be applied to a primitive type for field '{}'",
                        logical_type,
                        self.name
                    ));
                }
                (LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {}
                (LogicalType::Decimal { scale, precision }, _) => {
                    // Check that scale and precision are consistent with legacy values
                    if *scale != self.scale {
                        return Err(general_err!(
                            "DECIMAL logical type scale {} must match self.scale {} for field '{}'",
                            scale,
                            self.scale,
                            self.name
                        ));
                    }
                    if *precision != self.precision {
                        return Err(general_err!(
                            "DECIMAL logical type precision {} must match self.precision {} for field '{}'",
                            precision,
                            self.precision,
                            self.name
                        ));
                    }
                    self.check_decimal_precision_scale()?;
                }
                (LogicalType::Date, PhysicalType::INT32) => {}
                (
                    LogicalType::Time {
                        unit: TimeUnit::MILLIS(_),
                        ..
                    },
                    PhysicalType::INT32,
                ) => {}
                (LogicalType::Time { unit, .. }, PhysicalType::INT64) => {
                    if *unit == TimeUnit::MILLIS(Default::default()) {
                        return Err(general_err!(
                            "Cannot use millisecond unit on INT64 type for field '{}'",
                            self.name
                        ));
                    }
                }
                (LogicalType::Timestamp { .. }, PhysicalType::INT64) => {}
                (LogicalType::Integer { bit_width, .. }, PhysicalType::INT32)
                    if *bit_width <= 32 => {}
                (LogicalType::Integer { bit_width, .. }, PhysicalType::INT64)
                    if *bit_width == 64 => {}
                // Null type
                (LogicalType::Unknown, PhysicalType::INT32) => {}
                (LogicalType::String, PhysicalType::BYTE_ARRAY) => {}
                (LogicalType::Json, PhysicalType::BYTE_ARRAY) => {}
                (LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {}
                (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {}
                (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
                    return Err(general_err!(
                        "UUID cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(16) field",
                        self.name
                    ))
                }
                (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY)
                    if self.length == 2 => {}
                (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
                    return Err(general_err!(
                        "FLOAT16 cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(2) field",
                        self.name
                    ))
                }
                (a, b) => {
                    return Err(general_err!(
                        "Cannot annotate {:?} from {} for field '{}'",
                        a,
                        b,
                        self.name
                    ))
                }
            }
        }

        match self.converted_type {
            ConvertedType::NONE => {}
            ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => {
                if self.physical_type != PhysicalType::BYTE_ARRAY {
                    return Err(general_err!(
                        "{} cannot annotate field '{}' because it is not a BYTE_ARRAY field",
                        self.converted_type,
                        self.name
                    ));
                }
            }
            ConvertedType::DECIMAL => {
                self.check_decimal_precision_scale()?;
            }
            ConvertedType::DATE
            | ConvertedType::TIME_MILLIS
            | ConvertedType::UINT_8
            | ConvertedType::UINT_16
            | ConvertedType::UINT_32
            | ConvertedType::INT_8
            | ConvertedType::INT_16
            | ConvertedType::INT_32 => {
                if self.physical_type != PhysicalType::INT32 {
                    return Err(general_err!(
                        "{} cannot annotate field '{}' because it is not a INT32 field",
                        self.converted_type,
                        self.name
                    ));
                }
            }
            ConvertedType::TIME_MICROS
            | ConvertedType::TIMESTAMP_MILLIS
            | ConvertedType::TIMESTAMP_MICROS
            | ConvertedType::UINT_64
            | ConvertedType::INT_64 => {
                if self.physical_type != PhysicalType::INT64 {
                    return Err(general_err!(
                        "{} cannot annotate field '{}' because it is not a INT64 field",
                        self.converted_type,
                        self.name
                    ));
                }
            }
            ConvertedType::INTERVAL => {
                if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY || self.length != 12 {
                    return Err(general_err!(
                        "INTERVAL cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(12) field",
                        self.name
                    ));
                }
            }
            ConvertedType::ENUM => {
                if self.physical_type != PhysicalType::BYTE_ARRAY {
                    return Err(general_err!(
                        "ENUM cannot annotate field '{}' because it is not a BYTE_ARRAY field",
                        self.name
                    ));
                }
            }
            _ => {
                return Err(general_err!(
                    "{} cannot be applied to primitive field '{}'",
                    self.converted_type,
                    self.name
                ));
            }
        }

        Ok(Type::PrimitiveType {
            basic_info,
            physical_type: self.physical_type,
            type_length: self.length,
            scale: self.scale,
            precision: self.precision,
        })
    }

    #[inline]
    fn check_decimal_precision_scale(&self) -> Result<()> {
        match self.physical_type {
            PhysicalType::INT32
            | PhysicalType::INT64
            | PhysicalType::BYTE_ARRAY
            | PhysicalType::FIXED_LEN_BYTE_ARRAY => (),
            _ => {
                return Err(general_err!(
                    "DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
                ));
            }
        }

        // Precision is required and must be a non-zero positive integer.
        if self.precision < 1 {
            return Err(general_err!(
                "Invalid DECIMAL precision: {}",
                self.precision
            ));
        }

        // Scale must be zero or a positive integer less than the precision.
        if self.scale < 0 {
            return Err(general_err!("Invalid DECIMAL scale: {}", self.scale));
        }

        if self.scale > self.precision {
            return Err(general_err!(
                "Invalid DECIMAL: scale ({}) cannot be greater than precision \
             ({})",
                self.scale,
                self.precision
            ));
        }

        // Check precision and scale based on physical type limitations.
        match self.physical_type {
            PhysicalType::INT32 => {
                if self.precision > 9 {
                    return Err(general_err!(
                        "Cannot represent INT32 as DECIMAL with precision {}",
                        self.precision
                    ));
                }
            }
            PhysicalType::INT64 => {
                if self.precision > 18 {
                    return Err(general_err!(
                        "Cannot represent INT64 as DECIMAL with precision {}",
                        self.precision
                    ));
                }
            }
            PhysicalType::FIXED_LEN_BYTE_ARRAY => {
                let max_precision = (2f64.powi(8 * self.length - 1) - 1f64).log10().floor() as i32;

                if self.precision > max_precision {
                    return Err(general_err!(
                        "Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \
                        precision {}. The max precision can only be {}",
                        self.length,
                        self.precision,
                        max_precision
                    ));
                }
            }
            _ => (), // For BYTE_ARRAY precision is not limited
        }

        Ok(())
    }
}

/// A builder for group types. All attributes are optional except the name.
/// Note that if not specified explicitly, `None` is used as the repetition of the group,
/// which means it is a root (message) type.
pub struct GroupTypeBuilder<'a> {
    name: &'a str,
    repetition: Option<Repetition>,
    converted_type: ConvertedType,
    logical_type: Option<LogicalType>,
    fields: Vec<TypePtr>,
    id: Option<i32>,
}

impl<'a> GroupTypeBuilder<'a> {
    /// Creates new group type builder with provided field name.
    pub fn new(name: &'a str) -> Self {
        Self {
            name,
            repetition: None,
            converted_type: ConvertedType::NONE,
            logical_type: None,
            fields: Vec::new(),
            id: None,
        }
    }

    /// Sets [`Repetition`] for this field and returns itself.
    pub fn with_repetition(mut self, repetition: Repetition) -> Self {
        self.repetition = Some(repetition);
        self
    }

    /// Sets [`ConvertedType`] for this field and returns itself.
    pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
        Self {
            converted_type,
            ..self
        }
    }

    /// Sets [`LogicalType`] for this field and returns itself.
    pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
        Self {
            logical_type,
            ..self
        }
    }

    /// Sets a list of fields that should be child nodes of this field.
    /// Returns updated self.
    pub fn with_fields(self, fields: Vec<TypePtr>) -> Self {
        Self { fields, ..self }
    }

    /// Sets optional field id and returns itself.
    pub fn with_id(self, id: Option<i32>) -> Self {
        Self { id, ..self }
    }

    /// Creates a new `GroupType` instance from the gathered attributes.
    pub fn build(self) -> Result<Type> {
        let mut basic_info = BasicTypeInfo {
            name: String::from(self.name),
            repetition: self.repetition,
            converted_type: self.converted_type,
            logical_type: self.logical_type.clone(),
            id: self.id,
        };
        // Populate the converted type if only the logical type is populated
        if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE {
            basic_info.converted_type = self.logical_type.into();
        }
        Ok(Type::GroupType {
            basic_info,
            fields: self.fields,
        })
    }
}

/// Basic type info. This contains information such as the name of the type,
/// the repetition level, the logical type and the kind of the type (group, primitive).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BasicTypeInfo {
    name: String,
    repetition: Option<Repetition>,
    converted_type: ConvertedType,
    logical_type: Option<LogicalType>,
    id: Option<i32>,
}

impl HeapSize for BasicTypeInfo {
    fn heap_size(&self) -> usize {
        // no heap allocations in any other subfield
        self.name.heap_size()
    }
}

impl BasicTypeInfo {
    /// Returns field name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns `true` if type has repetition field set, `false` otherwise.
    /// This is mostly applied to group type, because primitive type always has
    /// repetition set.
    pub fn has_repetition(&self) -> bool {
        self.repetition.is_some()
    }

    /// Returns [`Repetition`] value for the type.
    pub fn repetition(&self) -> Repetition {
        assert!(self.repetition.is_some());
        self.repetition.unwrap()
    }

    /// Returns [`ConvertedType`] value for the type.
    pub fn converted_type(&self) -> ConvertedType {
        self.converted_type
    }

    /// Returns [`LogicalType`] value for the type.
    pub fn logical_type(&self) -> Option<LogicalType> {
        // Unlike ConvertedType, LogicalType cannot implement Copy, thus we clone it
        self.logical_type.clone()
    }

    /// Returns `true` if id is set, `false` otherwise.
    pub fn has_id(&self) -> bool {
        self.id.is_some()
    }

    /// Returns id value for the type.
    pub fn id(&self) -> i32 {
        assert!(self.id.is_some());
        self.id.unwrap()
    }
}

// ----------------------------------------------------------------------
// Parquet descriptor definitions

/// Represents the location of a column in a Parquet schema
///
/// # Example: refer to column named `'my_column'`
/// ```
/// # use parquet::schema::types::ColumnPath;
/// let column_path = ColumnPath::from("my_column");
/// ```
///
/// # Example: refer to column named `c` in a nested struct `{a: {b: {c: ...}}}`
/// ```
/// # use parquet::schema::types::ColumnPath;
/// // form path 'a.b.c'
/// let column_path = ColumnPath::from(vec![
///   String::from("a"),
///   String::from("b"),
///   String::from("c")
/// ]);
/// ```
#[derive(Clone, PartialEq, Debug, Eq, Hash)]
pub struct ColumnPath {
    parts: Vec<String>,
}

impl HeapSize for ColumnPath {
    fn heap_size(&self) -> usize {
        self.parts.heap_size()
    }
}

impl ColumnPath {
    /// Creates new column path from vector of field names.
    pub fn new(parts: Vec<String>) -> Self {
        ColumnPath { parts }
    }

    /// Returns string representation of this column path.
    /// ```rust
    /// use parquet::schema::types::ColumnPath;
    ///
    /// let path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c".to_string()]);
    /// assert_eq!(&path.string(), "a.b.c");
    /// ```
    pub fn string(&self) -> String {
        self.parts.join(".")
    }

    /// Appends more components to end of column path.
    /// ```rust
    /// use parquet::schema::types::ColumnPath;
    ///
    /// let mut path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c"
    /// .to_string()]);
    /// assert_eq!(&path.string(), "a.b.c");
    ///
    /// path.append(vec!["d".to_string(), "e".to_string()]);
    /// assert_eq!(&path.string(), "a.b.c.d.e");
    /// ```
    pub fn append(&mut self, mut tail: Vec<String>) {
        self.parts.append(&mut tail);
    }

    /// Returns a slice of path components.
    pub fn parts(&self) -> &[String] {
        &self.parts
    }
}

impl fmt::Display for ColumnPath {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.string())
    }
}

impl From<Vec<String>> for ColumnPath {
    fn from(parts: Vec<String>) -> Self {
        ColumnPath { parts }
    }
}

impl From<&str> for ColumnPath {
    fn from(single_path: &str) -> Self {
        let s = String::from(single_path);
        ColumnPath::from(s)
    }
}

impl From<String> for ColumnPath {
    fn from(single_path: String) -> Self {
        let v = vec![single_path];
        ColumnPath { parts: v }
    }
}

impl AsRef<[String]> for ColumnPath {
    fn as_ref(&self) -> &[String] {
        &self.parts
    }
}

/// Physical type for leaf-level primitive columns.
///
/// Also includes the maximum definition and repetition levels required to
/// re-assemble nested data.
#[derive(Debug, PartialEq)]
pub struct ColumnDescriptor {
    /// The "leaf" primitive type of this column
    primitive_type: TypePtr,

    /// The maximum definition level for this column
    max_def_level: i16,

    /// The maximum repetition level for this column
    max_rep_level: i16,

    /// The path of this column. For instance, "a.b.c.d".
    path: ColumnPath,
}

impl HeapSize for ColumnDescriptor {
    fn heap_size(&self) -> usize {
        self.primitive_type.heap_size() + self.path.heap_size()
    }
}

impl ColumnDescriptor {
    /// Creates new descriptor for leaf-level column.
    pub fn new(
        primitive_type: TypePtr,
        max_def_level: i16,
        max_rep_level: i16,
        path: ColumnPath,
    ) -> Self {
        Self {
            primitive_type,
            max_def_level,
            max_rep_level,
            path,
        }
    }

    /// Returns maximum definition level for this column.
    #[inline]
    pub fn max_def_level(&self) -> i16 {
        self.max_def_level
    }

    /// Returns maximum repetition level for this column.
    #[inline]
    pub fn max_rep_level(&self) -> i16 {
        self.max_rep_level
    }

    /// Returns [`ColumnPath`] for this column.
    pub fn path(&self) -> &ColumnPath {
        &self.path
    }

    /// Returns self type [`Type`] for this leaf column.
    pub fn self_type(&self) -> &Type {
        self.primitive_type.as_ref()
    }

    /// Returns self type [`TypePtr`]  for this leaf
    /// column.
    pub fn self_type_ptr(&self) -> TypePtr {
        self.primitive_type.clone()
    }

    /// Returns column name.
    pub fn name(&self) -> &str {
        self.primitive_type.name()
    }

    /// Returns [`ConvertedType`] for this column.
    pub fn converted_type(&self) -> ConvertedType {
        self.primitive_type.get_basic_info().converted_type()
    }

    /// Returns [`LogicalType`] for this column.
    pub fn logical_type(&self) -> Option<LogicalType> {
        self.primitive_type.get_basic_info().logical_type()
    }

    /// Returns physical type for this column.
    /// Note that it will panic if called on a non-primitive type.
    pub fn physical_type(&self) -> PhysicalType {
        match self.primitive_type.as_ref() {
            Type::PrimitiveType { physical_type, .. } => *physical_type,
            _ => panic!("Expected primitive type!"),
        }
    }

    /// Returns type length for this column.
    /// Note that it will panic if called on a non-primitive type.
    pub fn type_length(&self) -> i32 {
        match self.primitive_type.as_ref() {
            Type::PrimitiveType { type_length, .. } => *type_length,
            _ => panic!("Expected primitive type!"),
        }
    }

    /// Returns type precision for this column.
    /// Note that it will panic if called on a non-primitive type.
    pub fn type_precision(&self) -> i32 {
        match self.primitive_type.as_ref() {
            Type::PrimitiveType { precision, .. } => *precision,
            _ => panic!("Expected primitive type!"),
        }
    }

    /// Returns type scale for this column.
    /// Note that it will panic if called on a non-primitive type.
    pub fn type_scale(&self) -> i32 {
        match self.primitive_type.as_ref() {
            Type::PrimitiveType { scale, .. } => *scale,
            _ => panic!("Expected primitive type!"),
        }
    }

    /// Returns the sort order for this column
    pub fn sort_order(&self) -> SortOrder {
        ColumnOrder::get_sort_order(
            self.logical_type(),
            self.converted_type(),
            self.physical_type(),
        )
    }
}

/// Schema of a Parquet file.
///
/// Encapsulates the file's schema ([`Type`]) and [`ColumnDescriptor`]s for
/// each primitive (leaf) column.
#[derive(PartialEq)]
pub struct SchemaDescriptor {
    /// The top-level logical schema (the "message" type).
    ///
    /// This must be a [`Type::GroupType`] where each field is a root
    /// column type in the schema.
    schema: TypePtr,

    /// The descriptors for the physical type of each leaf column in this schema
    ///
    /// Constructed from `schema` in DFS order.
    leaves: Vec<ColumnDescPtr>,

    /// Mapping from a leaf column's index to the root column index that it
    /// comes from.
    ///
    /// For instance: the leaf `a.b.c.d` would have a link back to `a`:
    /// ```text
    /// -- a  <-----+
    /// -- -- b     |
    /// -- -- -- c  |
    /// -- -- -- -- d
    /// ```
    leaf_to_base: Vec<usize>,
}

impl fmt::Debug for SchemaDescriptor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Skip leaves and leaf_to_base as they only a cache information already found in `schema`
        f.debug_struct("SchemaDescriptor")
            .field("schema", &self.schema)
            .finish()
    }
}

// Need to implement HeapSize in this module as the fields are private
impl HeapSize for SchemaDescriptor {
    fn heap_size(&self) -> usize {
        self.schema.heap_size() + self.leaves.heap_size() + self.leaf_to_base.heap_size()
    }
}

impl SchemaDescriptor {
    /// Creates new schema descriptor from Parquet schema.
    pub fn new(tp: TypePtr) -> Self {
        assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
        let mut leaves = vec![];
        let mut leaf_to_base = Vec::new();
        for (root_idx, f) in tp.get_fields().iter().enumerate() {
            let mut path = vec![];
            build_tree(f, root_idx, 0, 0, &mut leaves, &mut leaf_to_base, &mut path);
        }

        Self {
            schema: tp,
            leaves,
            leaf_to_base,
        }
    }

    /// Returns [`ColumnDescriptor`] for a field position.
    pub fn column(&self, i: usize) -> ColumnDescPtr {
        assert!(
            i < self.leaves.len(),
            "Index out of bound: {} not in [0, {})",
            i,
            self.leaves.len()
        );
        self.leaves[i].clone()
    }

    /// Returns slice of [`ColumnDescriptor`].
    pub fn columns(&self) -> &[ColumnDescPtr] {
        &self.leaves
    }

    /// Returns number of leaf-level columns.
    pub fn num_columns(&self) -> usize {
        self.leaves.len()
    }

    /// Returns column root [`Type`] for a leaf position.
    pub fn get_column_root(&self, i: usize) -> &Type {
        let result = self.column_root_of(i);
        result.as_ref()
    }

    /// Returns column root [`Type`] pointer for a leaf position.
    pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
        let result = self.column_root_of(i);
        result.clone()
    }

    /// Returns the index of the root column for a field position
    pub fn get_column_root_idx(&self, leaf: usize) -> usize {
        assert!(
            leaf < self.leaves.len(),
            "Index out of bound: {} not in [0, {})",
            leaf,
            self.leaves.len()
        );

        *self
            .leaf_to_base
            .get(leaf)
            .unwrap_or_else(|| panic!("Expected a value for index {leaf} but found None"))
    }

    fn column_root_of(&self, i: usize) -> &TypePtr {
        &self.schema.get_fields()[self.get_column_root_idx(i)]
    }

    /// Returns schema as [`Type`].
    pub fn root_schema(&self) -> &Type {
        self.schema.as_ref()
    }

    /// Returns schema as [`TypePtr`] for cheap cloning.
    pub fn root_schema_ptr(&self) -> TypePtr {
        self.schema.clone()
    }

    /// Returns schema name.
    pub fn name(&self) -> &str {
        self.schema.name()
    }
}

fn build_tree<'a>(
    tp: &'a TypePtr,
    root_idx: usize,
    mut max_rep_level: i16,
    mut max_def_level: i16,
    leaves: &mut Vec<ColumnDescPtr>,
    leaf_to_base: &mut Vec<usize>,
    path_so_far: &mut Vec<&'a str>,
) {
    assert!(tp.get_basic_info().has_repetition());

    path_so_far.push(tp.name());
    match tp.get_basic_info().repetition() {
        Repetition::OPTIONAL => {
            max_def_level += 1;
        }
        Repetition::REPEATED => {
            max_def_level += 1;
            max_rep_level += 1;
        }
        _ => {}
    }

    match tp.as_ref() {
        Type::PrimitiveType { .. } => {
            let mut path: Vec<String> = vec![];
            path.extend(path_so_far.iter().copied().map(String::from));
            leaves.push(Arc::new(ColumnDescriptor::new(
                tp.clone(),
                max_def_level,
                max_rep_level,
                ColumnPath::new(path),
            )));
            leaf_to_base.push(root_idx);
        }
        Type::GroupType { ref fields, .. } => {
            for f in fields {
                build_tree(
                    f,
                    root_idx,
                    max_rep_level,
                    max_def_level,
                    leaves,
                    leaf_to_base,
                    path_so_far,
                );
                path_so_far.pop();
            }
        }
    }
}

/// Method to convert from Thrift.
pub fn from_thrift(elements: &[SchemaElement]) -> Result<TypePtr> {
    let mut index = 0;
    let mut schema_nodes = Vec::new();
    while index < elements.len() {
        let t = from_thrift_helper(elements, index)?;
        index = t.0;
        schema_nodes.push(t.1);
    }
    if schema_nodes.len() != 1 {
        return Err(general_err!(
            "Expected exactly one root node, but found {}",
            schema_nodes.len()
        ));
    }

    Ok(schema_nodes.remove(0))
}

/// Constructs a new Type from the `elements`, starting at index `index`.
/// The first result is the starting index for the next Type after this one. If it is
/// equal to `elements.len()`, then this Type is the last one.
/// The second result is the result Type.
fn from_thrift_helper(elements: &[SchemaElement], index: usize) -> Result<(usize, TypePtr)> {
    // Whether or not the current node is root (message type).
    // There is only one message type node in the schema tree.
    let is_root_node = index == 0;

    if index > elements.len() {
        return Err(general_err!(
            "Index out of bound, index = {}, len = {}",
            index,
            elements.len()
        ));
    }
    let element = &elements[index];
    let converted_type = ConvertedType::try_from(element.converted_type)?;
    // LogicalType is only present in v2 Parquet files. ConvertedType is always
    // populated, regardless of the version of the file (v1 or v2).
    let logical_type = element
        .logical_type
        .as_ref()
        .map(|value| LogicalType::from(value.clone()));
    let field_id = elements[index].field_id;
    match elements[index].num_children {
        // From parquet-format:
        //   The children count is used to construct the nested relationship.
        //   This field is not set when the element is a primitive type
        // Sometimes parquet-cpp sets num_children field to 0 for primitive types, so we
        // have to handle this case too.
        None | Some(0) => {
            // primitive type
            if elements[index].repetition_type.is_none() {
                return Err(general_err!(
                    "Repetition level must be defined for a primitive type"
                ));
            }
            let repetition = Repetition::try_from(elements[index].repetition_type.unwrap())?;
            if let Some(type_) = elements[index].type_ {
                let physical_type = PhysicalType::try_from(type_)?;
                let length = elements[index].type_length.unwrap_or(-1);
                let scale = elements[index].scale.unwrap_or(-1);
                let precision = elements[index].precision.unwrap_or(-1);
                let name = &elements[index].name;
                let builder = Type::primitive_type_builder(name, physical_type)
                    .with_repetition(repetition)
                    .with_converted_type(converted_type)
                    .with_logical_type(logical_type)
                    .with_length(length)
                    .with_precision(precision)
                    .with_scale(scale)
                    .with_id(field_id);
                Ok((index + 1, Arc::new(builder.build()?)))
            } else {
                let mut builder = Type::group_type_builder(&elements[index].name)
                    .with_converted_type(converted_type)
                    .with_logical_type(logical_type)
                    .with_id(field_id);
                if !is_root_node {
                    // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or
                    // REPEATED for root node.
                    //
                    // We only set repetition for group types that are not top-level message
                    // type. According to parquet-format:
                    //   Root of the schema does not have a repetition_type.
                    //   All other types must have one.
                    builder = builder.with_repetition(repetition);
                }
                Ok((index + 1, Arc::new(builder.build().unwrap())))
            }
        }
        Some(n) => {
            let repetition = elements[index]
                .repetition_type
                .map(Repetition::try_from)
                .transpose()?;

            let mut fields = vec![];
            let mut next_index = index + 1;
            for _ in 0..n {
                let child_result = from_thrift_helper(elements, next_index)?;
                next_index = child_result.0;
                fields.push(child_result.1);
            }

            let mut builder = Type::group_type_builder(&elements[index].name)
                .with_converted_type(converted_type)
                .with_logical_type(logical_type)
                .with_fields(fields)
                .with_id(field_id);
            if let Some(rep) = repetition {
                // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or
                // REPEATED for root node.
                //
                // We only set repetition for group types that are not top-level message
                // type. According to parquet-format:
                //   Root of the schema does not have a repetition_type.
                //   All other types must have one.
                if !is_root_node {
                    builder = builder.with_repetition(rep);
                }
            }
            Ok((next_index, Arc::new(builder.build().unwrap())))
        }
    }
}

/// Method to convert to Thrift.
pub fn to_thrift(schema: &Type) -> Result<Vec<SchemaElement>> {
    if !schema.is_group() {
        return Err(general_err!("Root schema must be Group type"));
    }
    let mut elements: Vec<SchemaElement> = Vec::new();
    to_thrift_helper(schema, &mut elements);
    Ok(elements)
}

/// Constructs list of `SchemaElement` from the schema using depth-first traversal.
/// Here we assume that schema is always valid and starts with group type.
fn to_thrift_helper(schema: &Type, elements: &mut Vec<SchemaElement>) {
    match *schema {
        Type::PrimitiveType {
            ref basic_info,
            physical_type,
            type_length,
            scale,
            precision,
        } => {
            let element = SchemaElement {
                type_: Some(physical_type.into()),
                type_length: if type_length >= 0 {
                    Some(type_length)
                } else {
                    None
                },
                repetition_type: Some(basic_info.repetition().into()),
                name: basic_info.name().to_owned(),
                num_children: None,
                converted_type: basic_info.converted_type().into(),
                scale: if scale >= 0 { Some(scale) } else { None },
                precision: if precision >= 0 {
                    Some(precision)
                } else {
                    None
                },
                field_id: if basic_info.has_id() {
                    Some(basic_info.id())
                } else {
                    None
                },
                logical_type: basic_info.logical_type().map(|value| value.into()),
            };

            elements.push(element);
        }
        Type::GroupType {
            ref basic_info,
            ref fields,
        } => {
            let repetition = if basic_info.has_repetition() {
                Some(basic_info.repetition().into())
            } else {
                None
            };

            let element = SchemaElement {
                type_: None,
                type_length: None,
                repetition_type: repetition,
                name: basic_info.name().to_owned(),
                num_children: Some(fields.len() as i32),
                converted_type: basic_info.converted_type().into(),
                scale: None,
                precision: None,
                field_id: if basic_info.has_id() {
                    Some(basic_info.id())
                } else {
                    None
                },
                logical_type: basic_info.logical_type().map(|value| value.into()),
            };

            elements.push(element);

            // Add child elements for a group
            for field in fields {
                to_thrift_helper(field, elements);
            }
        }
    }
}

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

    use crate::schema::parser::parse_message_type;

    // TODO: add tests for v2 types

    #[test]
    fn test_primitive_type() {
        let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32)
            .with_logical_type(Some(LogicalType::Integer {
                bit_width: 32,
                is_signed: true,
            }))
            .with_id(Some(0))
            .build();
        assert!(result.is_ok());

        if let Ok(tp) = result {
            assert!(tp.is_primitive());
            assert!(!tp.is_group());
            let basic_info = tp.get_basic_info();
            assert_eq!(basic_info.repetition(), Repetition::OPTIONAL);
            assert_eq!(
                basic_info.logical_type(),
                Some(LogicalType::Integer {
                    bit_width: 32,
                    is_signed: true
                })
            );
            assert_eq!(basic_info.converted_type(), ConvertedType::INT_32);
            assert_eq!(basic_info.id(), 0);
            match tp {
                Type::PrimitiveType { physical_type, .. } => {
                    assert_eq!(physical_type, PhysicalType::INT32);
                }
                _ => panic!(),
            }
        }

        // Test illegal inputs with logical type
        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
            .with_repetition(Repetition::REPEATED)
            .with_logical_type(Some(LogicalType::Integer {
                is_signed: true,
                bit_width: 8,
            }))
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Cannot annotate Integer { bit_width: 8, is_signed: true } from INT64 for field 'foo'"
            );
        }

        // Test illegal inputs with converted type
        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
            .with_repetition(Repetition::REPEATED)
            .with_converted_type(ConvertedType::BSON)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: BSON cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT96)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(-1)
            .with_scale(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_logical_type(Some(LogicalType::Decimal {
                scale: 32,
                precision: 12,
            }))
            .with_precision(-1)
            .with_scale(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: DECIMAL logical type scale 32 must match self.scale -1 for field 'foo'"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(-1)
            .with_scale(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Invalid DECIMAL precision: -1"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(0)
            .with_scale(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Invalid DECIMAL precision: 0"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(1)
            .with_scale(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(format!("{e}"), "Parquet error: Invalid DECIMAL scale: -1");
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(1)
            .with_scale(2)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Invalid DECIMAL: scale (2) cannot be greater than precision (1)"
            );
        }

        // It is OK if precision == scale
        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(1)
            .with_scale(1)
            .build();
        assert!(result.is_ok());

        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(18)
            .with_scale(2)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Cannot represent INT32 as DECIMAL with precision 18"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_precision(32)
            .with_scale(2)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Cannot represent INT64 as DECIMAL with precision 32"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_length(5)
            .with_precision(12)
            .with_scale(2)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::UINT_8)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: UINT_8 cannot annotate field 'foo' because it is not a INT32 field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::TIME_MICROS)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: TIME_MICROS cannot annotate field 'foo' because it is not a INT64 field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::INTERVAL)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::INTERVAL)
            .with_length(1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::ENUM)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: ENUM cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::MAP)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: MAP cannot be applied to primitive field 'foo'"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::DECIMAL)
            .with_length(-1)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1 for field 'foo'"
            );
        }

        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_logical_type(Some(LogicalType::Float16))
            .with_length(2)
            .build();
        assert!(result.is_ok());

        // Can't be other than FIXED_LEN_BYTE_ARRAY for physical type
        result = Type::primitive_type_builder("foo", PhysicalType::FLOAT)
            .with_repetition(Repetition::REQUIRED)
            .with_logical_type(Some(LogicalType::Float16))
            .with_length(2)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Cannot annotate Float16 from FLOAT for field 'foo'"
            );
        }

        // Must have length 2
        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_logical_type(Some(LogicalType::Float16))
            .with_length(4)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: FLOAT16 cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(2) field"
            );
        }

        // Must have length 16
        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
            .with_repetition(Repetition::REQUIRED)
            .with_logical_type(Some(LogicalType::Uuid))
            .with_length(15)
            .build();
        assert!(result.is_err());
        if let Err(e) = result {
            assert_eq!(
                format!("{e}"),
                "Parquet error: UUID cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(16) field"
            );
        }
    }

    #[test]
    fn test_group_type() {
        let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
            .with_converted_type(ConvertedType::INT_32)
            .with_id(Some(0))
            .build();
        assert!(f1.is_ok());
        let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
            .with_converted_type(ConvertedType::UTF8)
            .with_id(Some(1))
            .build();
        assert!(f2.is_ok());

        let fields = vec![Arc::new(f1.unwrap()), Arc::new(f2.unwrap())];

        let result = Type::group_type_builder("foo")
            .with_repetition(Repetition::REPEATED)
            .with_logical_type(Some(LogicalType::List))
            .with_fields(fields)
            .with_id(Some(1))
            .build();
        assert!(result.is_ok());

        let tp = result.unwrap();
        let basic_info = tp.get_basic_info();
        assert!(tp.is_group());
        assert!(!tp.is_primitive());
        assert_eq!(basic_info.repetition(), Repetition::REPEATED);
        assert_eq!(basic_info.logical_type(), Some(LogicalType::List));
        assert_eq!(basic_info.converted_type(), ConvertedType::LIST);
        assert_eq!(basic_info.id(), 1);
        assert_eq!(tp.get_fields().len(), 2);
        assert_eq!(tp.get_fields()[0].name(), "f1");
        assert_eq!(tp.get_fields()[1].name(), "f2");
    }

    #[test]
    fn test_column_descriptor() {
        let result = test_column_descriptor_helper();
        assert!(
            result.is_ok(),
            "Expected result to be OK but got err:\n {}",
            result.unwrap_err()
        );
    }

    fn test_column_descriptor_helper() -> Result<()> {
        let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY)
            .with_converted_type(ConvertedType::UTF8)
            .build()?;

        let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name"));

        assert_eq!(descr.path(), &ColumnPath::from("name"));
        assert_eq!(descr.converted_type(), ConvertedType::UTF8);
        assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY);
        assert_eq!(descr.max_def_level(), 4);
        assert_eq!(descr.max_rep_level(), 1);
        assert_eq!(descr.name(), "name");
        assert_eq!(descr.type_length(), -1);
        assert_eq!(descr.type_precision(), -1);
        assert_eq!(descr.type_scale(), -1);

        Ok(())
    }

    #[test]
    fn test_schema_descriptor() {
        let result = test_schema_descriptor_helper();
        assert!(
            result.is_ok(),
            "Expected result to be OK but got err:\n {}",
            result.unwrap_err()
        );
    }

    // A helper fn to avoid handling the results from type creation
    fn test_schema_descriptor_helper() -> Result<()> {
        let mut fields = vec![];

        let inta = Type::primitive_type_builder("a", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::INT_32)
            .build()?;
        fields.push(Arc::new(inta));
        let intb = Type::primitive_type_builder("b", PhysicalType::INT64)
            .with_converted_type(ConvertedType::INT_64)
            .build()?;
        fields.push(Arc::new(intb));
        let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY)
            .with_repetition(Repetition::REPEATED)
            .with_converted_type(ConvertedType::UTF8)
            .build()?;
        fields.push(Arc::new(intc));

        // 3-level list encoding
        let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64)
            .with_repetition(Repetition::REQUIRED)
            .with_converted_type(ConvertedType::INT_64)
            .build()?;
        let item2 = Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?;
        let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32)
            .with_repetition(Repetition::REPEATED)
            .with_converted_type(ConvertedType::INT_32)
            .build()?;
        let list = Type::group_type_builder("records")
            .with_repetition(Repetition::REPEATED)
            .with_converted_type(ConvertedType::LIST)
            .with_fields(vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)])
            .build()?;
        let bag = Type::group_type_builder("bag")
            .with_repetition(Repetition::OPTIONAL)
            .with_fields(vec![Arc::new(list)])
            .build()?;
        fields.push(Arc::new(bag));

        let schema = Type::group_type_builder("schema")
            .with_repetition(Repetition::REPEATED)
            .with_fields(fields)
            .build()?;
        let descr = SchemaDescriptor::new(Arc::new(schema));

        let nleaves = 6;
        assert_eq!(descr.num_columns(), nleaves);

        //                             mdef mrep
        // required int32 a            0    0
        // optional int64 b            1    0
        // repeated byte_array c       1    1
        // optional group bag          1    0
        //   repeated group records    2    1
        //     required int64 item1    2    1
        //     optional boolean item2  3    1
        //     repeated int32 item3    3    2
        let ex_max_def_levels = [0, 1, 1, 2, 3, 3];
        let ex_max_rep_levels = [0, 0, 1, 1, 1, 2];

        for i in 0..nleaves {
            let col = descr.column(i);
            assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{i}");
            assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{i}");
        }

        assert_eq!(descr.column(0).path().string(), "a");
        assert_eq!(descr.column(1).path().string(), "b");
        assert_eq!(descr.column(2).path().string(), "c");
        assert_eq!(descr.column(3).path().string(), "bag.records.item1");
        assert_eq!(descr.column(4).path().string(), "bag.records.item2");
        assert_eq!(descr.column(5).path().string(), "bag.records.item3");

        assert_eq!(descr.get_column_root(0).name(), "a");
        assert_eq!(descr.get_column_root(3).name(), "bag");
        assert_eq!(descr.get_column_root(4).name(), "bag");
        assert_eq!(descr.get_column_root(5).name(), "bag");

        Ok(())
    }

    #[test]
    fn test_schema_build_tree_def_rep_levels() {
        let message_type = "
    message spark_schema {
      REQUIRED INT32 a;
      OPTIONAL group b {
        OPTIONAL INT32 _1;
        OPTIONAL INT32 _2;
      }
      OPTIONAL group c (LIST) {
        REPEATED group list {
          OPTIONAL INT32 element;
        }
      }
    }
    ";
        let schema = parse_message_type(message_type).expect("should parse schema");
        let descr = SchemaDescriptor::new(Arc::new(schema));
        // required int32 a
        assert_eq!(descr.column(0).max_def_level(), 0);
        assert_eq!(descr.column(0).max_rep_level(), 0);
        // optional int32 b._1
        assert_eq!(descr.column(1).max_def_level(), 2);
        assert_eq!(descr.column(1).max_rep_level(), 0);
        // optional int32 b._2
        assert_eq!(descr.column(2).max_def_level(), 2);
        assert_eq!(descr.column(2).max_rep_level(), 0);
        // repeated optional int32 c.list.element
        assert_eq!(descr.column(3).max_def_level(), 3);
        assert_eq!(descr.column(3).max_rep_level(), 1);
    }

    #[test]
    #[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")]
    fn test_get_physical_type_panic() {
        let list = Type::group_type_builder("records")
            .with_repetition(Repetition::REPEATED)
            .build()
            .unwrap();
        list.get_physical_type();
    }

    #[test]
    fn test_get_physical_type_primitive() {
        let f = Type::primitive_type_builder("f", PhysicalType::INT64)
            .build()
            .unwrap();
        assert_eq!(f.get_physical_type(), PhysicalType::INT64);

        let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY)
            .build()
            .unwrap();
        assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY);
    }

    #[test]
    fn test_check_contains_primitive_primitive() {
        // OK
        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .build()
            .unwrap();
        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .build()
            .unwrap();
        assert!(f1.check_contains(&f2));

        // OK: different logical type does not affect check_contains
        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .with_converted_type(ConvertedType::UINT_8)
            .build()
            .unwrap();
        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .with_converted_type(ConvertedType::UINT_16)
            .build()
            .unwrap();
        assert!(f1.check_contains(&f2));

        // KO: different name
        let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
            .build()
            .unwrap();
        let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));

        // KO: different type
        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .build()
            .unwrap();
        let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));

        // KO: different repetition
        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .with_repetition(Repetition::REQUIRED)
            .build()
            .unwrap();
        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
            .with_repetition(Repetition::OPTIONAL)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));
    }

    // function to create a new group type for testing
    fn test_new_group_type(name: &str, repetition: Repetition, types: Vec<Type>) -> Type {
        Type::group_type_builder(name)
            .with_repetition(repetition)
            .with_fields(types.into_iter().map(Arc::new).collect())
            .build()
            .unwrap()
    }

    #[test]
    fn test_check_contains_group_group() {
        // OK: should match okay with empty fields
        let f1 = Type::group_type_builder("f").build().unwrap();
        let f2 = Type::group_type_builder("f").build().unwrap();
        assert!(f1.check_contains(&f2));
        assert!(!f1.is_optional());

        // OK: fields match
        let f1 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::INT64)
                    .build()
                    .unwrap(),
            ],
        );
        let f2 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::INT64)
                    .build()
                    .unwrap(),
            ],
        );
        assert!(f1.check_contains(&f2));

        // OK: subset of fields
        let f1 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::INT64)
                    .build()
                    .unwrap(),
            ],
        );
        let f2 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![Type::primitive_type_builder("f2", PhysicalType::INT64)
                .build()
                .unwrap()],
        );
        assert!(f1.check_contains(&f2));

        // KO: different name
        let f1 = Type::group_type_builder("f1").build().unwrap();
        let f2 = Type::group_type_builder("f2").build().unwrap();
        assert!(!f1.check_contains(&f2));

        // KO: different repetition
        let f1 = Type::group_type_builder("f")
            .with_repetition(Repetition::OPTIONAL)
            .build()
            .unwrap();
        let f2 = Type::group_type_builder("f")
            .with_repetition(Repetition::REPEATED)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));

        // KO: different fields
        let f1 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::INT64)
                    .build()
                    .unwrap(),
            ],
        );
        let f2 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::BOOLEAN)
                    .build()
                    .unwrap(),
            ],
        );
        assert!(!f1.check_contains(&f2));

        // KO: different fields
        let f1 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![
                Type::primitive_type_builder("f1", PhysicalType::INT32)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("f2", PhysicalType::INT64)
                    .build()
                    .unwrap(),
            ],
        );
        let f2 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![Type::primitive_type_builder("f3", PhysicalType::INT32)
                .build()
                .unwrap()],
        );
        assert!(!f1.check_contains(&f2));
    }

    #[test]
    fn test_check_contains_group_primitive() {
        // KO: should not match
        let f1 = Type::group_type_builder("f").build().unwrap();
        let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));
        assert!(!f2.check_contains(&f1));

        // KO: should not match when primitive field is part of group type
        let f1 = test_new_group_type(
            "f",
            Repetition::REPEATED,
            vec![Type::primitive_type_builder("f1", PhysicalType::INT32)
                .build()
                .unwrap()],
        );
        let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32)
            .build()
            .unwrap();
        assert!(!f1.check_contains(&f2));
        assert!(!f2.check_contains(&f1));

        // OK: match nested types
        let f1 = test_new_group_type(
            "a",
            Repetition::REPEATED,
            vec![
                test_new_group_type(
                    "b",
                    Repetition::REPEATED,
                    vec![Type::primitive_type_builder("c", PhysicalType::INT32)
                        .build()
                        .unwrap()],
                ),
                Type::primitive_type_builder("d", PhysicalType::INT64)
                    .build()
                    .unwrap(),
                Type::primitive_type_builder("e", PhysicalType::BOOLEAN)
                    .build()
                    .unwrap(),
            ],
        );
        let f2 = test_new_group_type(
            "a",
            Repetition::REPEATED,
            vec![test_new_group_type(
                "b",
                Repetition::REPEATED,
                vec![Type::primitive_type_builder("c", PhysicalType::INT32)
                    .build()
                    .unwrap()],
            )],
        );
        assert!(f1.check_contains(&f2)); // should match
        assert!(!f2.check_contains(&f1)); // should fail
    }

    #[test]
    fn test_schema_type_thrift_conversion_err() {
        let schema = Type::primitive_type_builder("col", PhysicalType::INT32)
            .build()
            .unwrap();
        let thrift_schema = to_thrift(&schema);
        assert!(thrift_schema.is_err());
        if let Err(e) = thrift_schema {
            assert_eq!(
                format!("{e}"),
                "Parquet error: Root schema must be Group type"
            );
        }
    }

    #[test]
    fn test_schema_type_thrift_conversion() {
        let message_type = "
    message conversions {
      REQUIRED INT64 id;
      OPTIONAL FIXED_LEN_BYTE_ARRAY (2) f16 (FLOAT16);
      OPTIONAL group int_array_Array (LIST) {
        REPEATED group list {
          OPTIONAL group element (LIST) {
            REPEATED group list {
              OPTIONAL INT32 element;
            }
          }
        }
      }
      OPTIONAL group int_map (MAP) {
        REPEATED group map (MAP_KEY_VALUE) {
          REQUIRED BYTE_ARRAY key (UTF8);
          OPTIONAL INT32 value;
        }
      }
      OPTIONAL group int_Map_Array (LIST) {
        REPEATED group list {
          OPTIONAL group g (MAP) {
            REPEATED group map (MAP_KEY_VALUE) {
              REQUIRED BYTE_ARRAY key (UTF8);
              OPTIONAL group value {
                OPTIONAL group H {
                  OPTIONAL group i (LIST) {
                    REPEATED group list {
                      OPTIONAL DOUBLE element;
                    }
                  }
                }
              }
            }
          }
        }
      }
      OPTIONAL group nested_struct {
        OPTIONAL INT32 A;
        OPTIONAL group b (LIST) {
          REPEATED group list {
            REQUIRED FIXED_LEN_BYTE_ARRAY (16) element;
          }
        }
      }
    }
    ";
        let expected_schema = parse_message_type(message_type).unwrap();
        let thrift_schema = to_thrift(&expected_schema).unwrap();
        let result_schema = from_thrift(&thrift_schema).unwrap();
        assert_eq!(result_schema, Arc::new(expected_schema));
    }

    #[test]
    fn test_schema_type_thrift_conversion_decimal() {
        let message_type = "
    message decimals {
      OPTIONAL INT32 field0;
      OPTIONAL INT64 field1 (DECIMAL (18, 2));
      OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18));
      OPTIONAL BYTE_ARRAY field3 (DECIMAL (9));
    }
    ";
        let expected_schema = parse_message_type(message_type).unwrap();
        let thrift_schema = to_thrift(&expected_schema).unwrap();
        let result_schema = from_thrift(&thrift_schema).unwrap();
        assert_eq!(result_schema, Arc::new(expected_schema));
    }

    // Tests schema conversion from thrift, when num_children is set to Some(0) for a
    // primitive type.
    #[test]
    fn test_schema_from_thrift_with_num_children_set() {
        // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
        let message_type = "
    message schema {
      OPTIONAL BYTE_ARRAY id (UTF8);
      OPTIONAL BYTE_ARRAY name (UTF8);
      OPTIONAL BYTE_ARRAY message (UTF8);
      OPTIONAL INT32 type (UINT_8);
      OPTIONAL INT64 author_time (TIMESTAMP_MILLIS);
      OPTIONAL INT64 __index_level_0__;
    }
    ";

        let expected_schema = parse_message_type(message_type).unwrap();
        let mut thrift_schema = to_thrift(&expected_schema).unwrap();
        // Change all of None to Some(0)
        for elem in &mut thrift_schema[..] {
            if elem.num_children.is_none() {
                elem.num_children = Some(0);
            }
        }

        let result_schema = from_thrift(&thrift_schema).unwrap();
        assert_eq!(result_schema, Arc::new(expected_schema));
    }

    // Sometimes parquet-cpp sets repetition level for the root node, which is against
    // the format definition, but we need to handle it by setting it back to None.
    #[test]
    fn test_schema_from_thrift_root_has_repetition() {
        // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
        let message_type = "
    message schema {
      OPTIONAL BYTE_ARRAY a (UTF8);
      OPTIONAL INT32 b (UINT_8);
    }
    ";

        let expected_schema = parse_message_type(message_type).unwrap();
        let mut thrift_schema = to_thrift(&expected_schema).unwrap();
        thrift_schema[0].repetition_type = Some(Repetition::REQUIRED.into());

        let result_schema = from_thrift(&thrift_schema).unwrap();
        assert_eq!(result_schema, Arc::new(expected_schema));
    }

    #[test]
    fn test_schema_from_thrift_group_has_no_child() {
        let message_type = "message schema {}";

        let expected_schema = parse_message_type(message_type).unwrap();
        let mut thrift_schema = to_thrift(&expected_schema).unwrap();
        thrift_schema[0].repetition_type = Some(Repetition::REQUIRED.into());

        let result_schema = from_thrift(&thrift_schema).unwrap();
        assert_eq!(result_schema, Arc::new(expected_schema));
    }
}