Skip to main content

parquet/schema/
types.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Contains structs and methods to build Parquet schema and schema descriptors.
19
20use std::vec::IntoIter;
21use std::{collections::HashMap, fmt, sync::Arc};
22
23use crate::file::metadata::HeapSize;
24use crate::file::metadata::thrift::SchemaElement;
25
26use crate::basic::{
27    ColumnOrder, ConvertedType, IntType, LogicalType, Repetition, SortOrder, TimeType, TimeUnit,
28    Type as PhysicalType,
29};
30use crate::errors::{ParquetError, Result};
31
32// ----------------------------------------------------------------------
33// Parquet Type definitions
34
35/// Type alias for `Arc<Type>`.
36pub type TypePtr = Arc<Type>;
37/// Type alias for `Arc<SchemaDescriptor>`.
38pub type SchemaDescPtr = Arc<SchemaDescriptor>;
39/// Type alias for `Arc<ColumnDescriptor>`.
40pub type ColumnDescPtr = Arc<ColumnDescriptor>;
41
42/// Representation of a Parquet type.
43///
44/// Used to describe primitive leaf fields and structs, including top-level schema.
45///
46/// Note that the top-level schema is represented using [`Type::GroupType`] whose
47/// repetition is `None`.
48#[derive(Clone, Debug, PartialEq)]
49pub enum Type {
50    /// Represents a primitive leaf field.
51    PrimitiveType {
52        /// Basic information about the type.
53        basic_info: BasicTypeInfo,
54        /// Physical type of this primitive type.
55        physical_type: PhysicalType,
56        /// Length of this type.
57        type_length: i32,
58        /// Scale of this type.
59        scale: i32,
60        /// Precision of this type.
61        precision: i32,
62    },
63    /// Represents a group of fields (similar to struct).
64    GroupType {
65        /// Basic information about the type.
66        basic_info: BasicTypeInfo,
67        /// Fields of this group type.
68        fields: Vec<TypePtr>,
69    },
70}
71
72impl HeapSize for Type {
73    fn heap_size(&self) -> usize {
74        match self {
75            Type::PrimitiveType { basic_info, .. } => basic_info.heap_size(),
76            Type::GroupType { basic_info, fields } => basic_info.heap_size() + fields.heap_size(),
77        }
78    }
79}
80
81impl Type {
82    /// Creates primitive type builder with provided field name and physical type.
83    pub fn primitive_type_builder(
84        name: &str,
85        physical_type: PhysicalType,
86    ) -> PrimitiveTypeBuilder<'_> {
87        PrimitiveTypeBuilder::new(name, physical_type)
88    }
89
90    /// Creates group type builder with provided column name.
91    pub fn group_type_builder(name: &str) -> GroupTypeBuilder<'_> {
92        GroupTypeBuilder::new(name)
93    }
94
95    /// Returns [`BasicTypeInfo`] information about the type.
96    pub fn get_basic_info(&self) -> &BasicTypeInfo {
97        match *self {
98            Type::PrimitiveType { ref basic_info, .. } => basic_info,
99            Type::GroupType { ref basic_info, .. } => basic_info,
100        }
101    }
102
103    /// Returns this type's field name.
104    pub fn name(&self) -> &str {
105        self.get_basic_info().name()
106    }
107
108    /// Gets the fields from this group type.
109    /// Note that this will panic if called on a non-group type.
110    // TODO: should we return `&[&Type]` here?
111    pub fn get_fields(&self) -> &[TypePtr] {
112        match *self {
113            Type::GroupType { ref fields, .. } => &fields[..],
114            _ => panic!("Cannot call get_fields() on a non-group type"),
115        }
116    }
117
118    /// Gets physical type of this primitive type.
119    /// Note that this will panic if called on a non-primitive type.
120    pub fn get_physical_type(&self) -> PhysicalType {
121        match *self {
122            Type::PrimitiveType {
123                basic_info: _,
124                physical_type,
125                ..
126            } => physical_type,
127            _ => panic!("Cannot call get_physical_type() on a non-primitive type"),
128        }
129    }
130
131    /// Gets precision of this primitive type.
132    /// Note that this will panic if called on a non-primitive type.
133    pub fn get_precision(&self) -> i32 {
134        match *self {
135            Type::PrimitiveType { precision, .. } => precision,
136            _ => panic!("Cannot call get_precision() on non-primitive type"),
137        }
138    }
139
140    /// Gets scale of this primitive type.
141    /// Note that this will panic if called on a non-primitive type.
142    pub fn get_scale(&self) -> i32 {
143        match *self {
144            Type::PrimitiveType { scale, .. } => scale,
145            _ => panic!("Cannot call get_scale() on non-primitive type"),
146        }
147    }
148
149    /// Checks if `sub_type` schema is part of current schema.
150    /// This method can be used to check if projected columns are part of the root schema.
151    pub fn check_contains(&self, sub_type: &Type) -> bool {
152        // Names match, and repetitions match or not set for both
153        let basic_match = self.get_basic_info().name() == sub_type.get_basic_info().name()
154            && (self.is_schema() && sub_type.is_schema()
155                || !self.is_schema()
156                    && !sub_type.is_schema()
157                    && self.get_basic_info().repetition()
158                        == sub_type.get_basic_info().repetition());
159
160        match *self {
161            Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => {
162                self.get_physical_type() == sub_type.get_physical_type()
163            }
164            Type::GroupType { .. } if basic_match && sub_type.is_group() => {
165                // build hashmap of name -> TypePtr
166                let mut field_map = HashMap::new();
167                for field in self.get_fields() {
168                    field_map.insert(field.name(), field);
169                }
170
171                for field in sub_type.get_fields() {
172                    if !field_map
173                        .get(field.name())
174                        .map(|tpe| tpe.check_contains(field))
175                        .unwrap_or(false)
176                    {
177                        return false;
178                    }
179                }
180                true
181            }
182            _ => false,
183        }
184    }
185
186    /// Returns `true` if this type is a primitive type, `false` otherwise.
187    pub fn is_primitive(&self) -> bool {
188        matches!(*self, Type::PrimitiveType { .. })
189    }
190
191    /// Returns `true` if this type is a group type, `false` otherwise.
192    pub fn is_group(&self) -> bool {
193        matches!(*self, Type::GroupType { .. })
194    }
195
196    /// Returns `true` if this type is the top-level schema type (message type).
197    pub fn is_schema(&self) -> bool {
198        match *self {
199            Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(),
200            _ => false,
201        }
202    }
203
204    /// Returns `true` if this type is repeated or optional.
205    /// If this type doesn't have repetition defined, we treat it as required.
206    pub fn is_optional(&self) -> bool {
207        self.get_basic_info().has_repetition()
208            && self.get_basic_info().repetition() != Repetition::REQUIRED
209    }
210
211    /// Returns `true` if this type is annotated as a list.
212    pub(crate) fn is_list(&self) -> bool {
213        if self.is_group() {
214            let basic_info = self.get_basic_info();
215            if let Some(logical_type) = basic_info.logical_type_ref() {
216                return logical_type == &LogicalType::List;
217            }
218            return basic_info.converted_type() == ConvertedType::LIST;
219        }
220        false
221    }
222
223    /// Returns `true` if this type is a group with a single child field that is `repeated`.
224    pub(crate) fn has_single_repeated_child(&self) -> bool {
225        if self.is_group() {
226            let children = self.get_fields();
227            return children.len() == 1
228                && children[0].get_basic_info().has_repetition()
229                && children[0].get_basic_info().repetition() == Repetition::REPEATED;
230        }
231        false
232    }
233}
234
235/// A builder for primitive types. All attributes are optional
236/// except the name and physical type.
237/// Note that if not specified explicitly, `Repetition::OPTIONAL` is used.
238pub struct PrimitiveTypeBuilder<'a> {
239    name: &'a str,
240    repetition: Repetition,
241    physical_type: PhysicalType,
242    converted_type: ConvertedType,
243    logical_type: Option<LogicalType>,
244    length: i32,
245    precision: i32,
246    scale: i32,
247    id: Option<i32>,
248}
249
250impl<'a> PrimitiveTypeBuilder<'a> {
251    /// Creates new primitive type builder with provided field name and physical type.
252    pub fn new(name: &'a str, physical_type: PhysicalType) -> Self {
253        Self {
254            name,
255            repetition: Repetition::OPTIONAL,
256            physical_type,
257            converted_type: ConvertedType::NONE,
258            logical_type: None,
259            length: -1,
260            precision: -1,
261            scale: -1,
262            id: None,
263        }
264    }
265
266    /// Sets [`Repetition`] for this field and returns itself.
267    pub fn with_repetition(self, repetition: Repetition) -> Self {
268        Self { repetition, ..self }
269    }
270
271    /// Sets [`ConvertedType`] for this field and returns itself.
272    pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
273        Self {
274            converted_type,
275            ..self
276        }
277    }
278
279    /// Sets [`LogicalType`] for this field and returns itself.
280    /// If only the logical type is populated for a primitive type, the converted type
281    /// will be automatically populated, and can thus be omitted.
282    pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
283        Self {
284            logical_type,
285            ..self
286        }
287    }
288
289    /// Sets type length and returns itself.
290    /// This is only applied to FIXED_LEN_BYTE_ARRAY and INT96 (INTERVAL) types, because
291    /// they maintain fixed size underlying byte array.
292    /// By default, value is `0`.
293    pub fn with_length(self, length: i32) -> Self {
294        Self { length, ..self }
295    }
296
297    /// Sets precision for Parquet DECIMAL physical type and returns itself.
298    /// By default, it equals to `0` and used only for decimal context.
299    pub fn with_precision(self, precision: i32) -> Self {
300        Self { precision, ..self }
301    }
302
303    /// Sets scale for Parquet DECIMAL physical type and returns itself.
304    /// By default, it equals to `0` and used only for decimal context.
305    pub fn with_scale(self, scale: i32) -> Self {
306        Self { scale, ..self }
307    }
308
309    /// Sets optional field id and returns itself.
310    pub fn with_id(self, id: Option<i32>) -> Self {
311        Self { id, ..self }
312    }
313
314    /// Creates a new `PrimitiveType` instance from the collected attributes.
315    /// Returns `Err` in case of any building conditions are not met.
316    pub fn build(self) -> Result<Type> {
317        let sort_order = ColumnOrder::column_order_for_type(
318            self.logical_type.as_ref(),
319            self.converted_type,
320            self.physical_type,
321        )
322        .sort_order();
323        let mut basic_info = BasicTypeInfo {
324            name: String::from(self.name),
325            repetition: Some(self.repetition),
326            converted_type: self.converted_type,
327            logical_type: self.logical_type.clone(),
328            id: self.id,
329            sort_order,
330        };
331
332        // Check length before logical type, since it is used for logical type validation.
333        if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 {
334            return Err(general_err!(
335                "Invalid FIXED_LEN_BYTE_ARRAY length: {} for field '{}'",
336                self.length,
337                self.name
338            ));
339        }
340
341        if let Some(logical_type) = &self.logical_type {
342            // If a converted type is populated, check that it is consistent with
343            // its logical type
344            if self.converted_type != ConvertedType::NONE {
345                if ConvertedType::from(self.logical_type.clone()) != self.converted_type {
346                    return Err(general_err!(
347                        "Logical type {:?} is incompatible with converted type {} for field '{}'",
348                        logical_type,
349                        self.converted_type,
350                        self.name
351                    ));
352                }
353            } else {
354                // Populate the converted type for backwards compatibility
355                basic_info.converted_type = self.logical_type.clone().into();
356            }
357            // Check that logical type and physical type are compatible
358            match (logical_type, self.physical_type) {
359                (LogicalType::Map | LogicalType::List, _) => {
360                    return Err(general_err!(
361                        "{:?} cannot be applied to a primitive type for field '{}'",
362                        logical_type,
363                        self.name
364                    ));
365                }
366                (LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {}
367                (LogicalType::Decimal(decimal), _) => {
368                    // Check that scale and precision are consistent with legacy values
369                    if decimal.scale != self.scale {
370                        return Err(general_err!(
371                            "DECIMAL logical type scale {} must match self.scale {} for field '{}'",
372                            decimal.scale,
373                            self.scale,
374                            self.name
375                        ));
376                    }
377                    if decimal.precision != self.precision {
378                        return Err(general_err!(
379                            "DECIMAL logical type precision {} must match self.precision {} for field '{}'",
380                            decimal.precision,
381                            self.precision,
382                            self.name
383                        ));
384                    }
385                    self.check_decimal_precision_scale()?;
386                }
387                (LogicalType::Date, PhysicalType::INT32) => {}
388                (
389                    LogicalType::Time(TimeType {
390                        unit: TimeUnit::MILLIS,
391                        ..
392                    }),
393                    PhysicalType::INT32,
394                ) => {}
395                (LogicalType::Time(time), PhysicalType::INT64) => {
396                    if time.unit == TimeUnit::MILLIS {
397                        return Err(general_err!(
398                            "Cannot use millisecond unit on INT64 type for field '{}'",
399                            self.name
400                        ));
401                    }
402                }
403                (LogicalType::Timestamp(_), PhysicalType::INT64) => {}
404                (LogicalType::Integer(int), PhysicalType::INT32) if int.bit_width <= 32 => {}
405                (LogicalType::Integer(int), PhysicalType::INT64) if int.bit_width == 64 => {}
406                // Null type
407                (LogicalType::Unknown, _) => {}
408                (LogicalType::String, PhysicalType::BYTE_ARRAY) => {}
409                (LogicalType::Json, PhysicalType::BYTE_ARRAY) => {}
410                (LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {}
411                (LogicalType::Geometry(_), PhysicalType::BYTE_ARRAY) => {}
412                (LogicalType::Geography(_), PhysicalType::BYTE_ARRAY) => {}
413                (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {}
414                (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
415                    return Err(general_err!(
416                        "UUID cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(16) field",
417                        self.name
418                    ));
419                }
420                (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 2 => {}
421                (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
422                    return Err(general_err!(
423                        "FLOAT16 cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(2) field",
424                        self.name
425                    ));
426                }
427                // unknown logical type means just use physical type
428                (LogicalType::_Unknown { .. }, _) => {}
429                (a, b) => {
430                    return Err(general_err!(
431                        "Cannot annotate {:?} from {} for field '{}'",
432                        a,
433                        b,
434                        self.name
435                    ));
436                }
437            }
438        }
439
440        match self.converted_type {
441            ConvertedType::NONE => {}
442            ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => {
443                if self.physical_type != PhysicalType::BYTE_ARRAY {
444                    return Err(general_err!(
445                        "{} cannot annotate field '{}' because it is not a BYTE_ARRAY field",
446                        self.converted_type,
447                        self.name
448                    ));
449                }
450            }
451            ConvertedType::DECIMAL => {
452                self.check_decimal_precision_scale()?;
453            }
454            ConvertedType::DATE
455            | ConvertedType::TIME_MILLIS
456            | ConvertedType::UINT_8
457            | ConvertedType::UINT_16
458            | ConvertedType::UINT_32
459            | ConvertedType::INT_8
460            | ConvertedType::INT_16
461            | ConvertedType::INT_32 => {
462                if self.physical_type != PhysicalType::INT32 {
463                    return Err(general_err!(
464                        "{} cannot annotate field '{}' because it is not a INT32 field",
465                        self.converted_type,
466                        self.name
467                    ));
468                }
469            }
470            ConvertedType::TIME_MICROS
471            | ConvertedType::TIMESTAMP_MILLIS
472            | ConvertedType::TIMESTAMP_MICROS
473            | ConvertedType::UINT_64
474            | ConvertedType::INT_64 => {
475                if self.physical_type != PhysicalType::INT64 {
476                    return Err(general_err!(
477                        "{} cannot annotate field '{}' because it is not a INT64 field",
478                        self.converted_type,
479                        self.name
480                    ));
481                }
482            }
483            ConvertedType::INTERVAL => {
484                if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY || self.length != 12 {
485                    return Err(general_err!(
486                        "INTERVAL cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(12) field",
487                        self.name
488                    ));
489                }
490            }
491            ConvertedType::ENUM => {
492                if self.physical_type != PhysicalType::BYTE_ARRAY {
493                    return Err(general_err!(
494                        "ENUM cannot annotate field '{}' because it is not a BYTE_ARRAY field",
495                        self.name
496                    ));
497                }
498            }
499            _ => {
500                return Err(general_err!(
501                    "{} cannot be applied to primitive field '{}'",
502                    self.converted_type,
503                    self.name
504                ));
505            }
506        }
507
508        Ok(Type::PrimitiveType {
509            basic_info,
510            physical_type: self.physical_type,
511            type_length: self.length,
512            scale: self.scale,
513            precision: self.precision,
514        })
515    }
516
517    #[inline]
518    fn check_decimal_precision_scale(&self) -> Result<()> {
519        match self.physical_type {
520            PhysicalType::INT32
521            | PhysicalType::INT64
522            | PhysicalType::BYTE_ARRAY
523            | PhysicalType::FIXED_LEN_BYTE_ARRAY => (),
524            _ => {
525                return Err(general_err!(
526                    "DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
527                ));
528            }
529        }
530
531        // Precision is required and must be a non-zero positive integer.
532        if self.precision < 1 {
533            return Err(general_err!(
534                "Invalid DECIMAL precision: {}",
535                self.precision
536            ));
537        }
538
539        // Scale must be zero or a positive integer less than the precision.
540        if self.scale < 0 {
541            return Err(general_err!("Invalid DECIMAL scale: {}", self.scale));
542        }
543
544        if self.scale > self.precision {
545            return Err(general_err!(
546                "Invalid DECIMAL: scale ({}) cannot be greater than precision \
547             ({})",
548                self.scale,
549                self.precision
550            ));
551        }
552
553        // Check precision and scale based on physical type limitations.
554        match self.physical_type {
555            PhysicalType::INT32 => {
556                if self.precision > 9 {
557                    return Err(general_err!(
558                        "Cannot represent INT32 as DECIMAL with precision {}",
559                        self.precision
560                    ));
561                }
562            }
563            PhysicalType::INT64 => {
564                if self.precision > 18 {
565                    return Err(general_err!(
566                        "Cannot represent INT64 as DECIMAL with precision {}",
567                        self.precision
568                    ));
569                }
570            }
571            PhysicalType::FIXED_LEN_BYTE_ARRAY => {
572                let length = self
573                    .length
574                    .checked_mul(8)
575                    .ok_or(general_err!("Invalid length {} for Decimal", self.length))?;
576                let max_precision = (2f64.powi(length - 1) - 1f64).log10().floor() as i32;
577
578                if self.precision > max_precision {
579                    return Err(general_err!(
580                        "Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \
581                        precision {}. The max precision can only be {}",
582                        self.length,
583                        self.precision,
584                        max_precision
585                    ));
586                }
587            }
588            _ => (), // For BYTE_ARRAY precision is not limited
589        }
590
591        Ok(())
592    }
593}
594
595/// A builder for group types. All attributes are optional except the name.
596/// Note that if not specified explicitly, `None` is used as the repetition of the group,
597/// which means it is a root (message) type.
598pub struct GroupTypeBuilder<'a> {
599    name: &'a str,
600    repetition: Option<Repetition>,
601    converted_type: ConvertedType,
602    logical_type: Option<LogicalType>,
603    fields: Vec<TypePtr>,
604    id: Option<i32>,
605}
606
607impl<'a> GroupTypeBuilder<'a> {
608    /// Creates new group type builder with provided field name.
609    pub fn new(name: &'a str) -> Self {
610        Self {
611            name,
612            repetition: None,
613            converted_type: ConvertedType::NONE,
614            logical_type: None,
615            fields: Vec::new(),
616            id: None,
617        }
618    }
619
620    /// Sets [`Repetition`] for this field and returns itself.
621    pub fn with_repetition(mut self, repetition: Repetition) -> Self {
622        self.repetition = Some(repetition);
623        self
624    }
625
626    /// Sets [`ConvertedType`] for this field and returns itself.
627    pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
628        Self {
629            converted_type,
630            ..self
631        }
632    }
633
634    /// Sets [`LogicalType`] for this field and returns itself.
635    pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
636        Self {
637            logical_type,
638            ..self
639        }
640    }
641
642    /// Sets a list of fields that should be child nodes of this field.
643    /// Returns updated self.
644    pub fn with_fields(self, fields: Vec<TypePtr>) -> Self {
645        Self { fields, ..self }
646    }
647
648    /// Sets optional field id and returns itself.
649    pub fn with_id(self, id: Option<i32>) -> Self {
650        Self { id, ..self }
651    }
652
653    /// Creates a new `GroupType` instance from the gathered attributes.
654    pub fn build(self) -> Result<Type> {
655        let mut basic_info = BasicTypeInfo {
656            name: String::from(self.name),
657            repetition: self.repetition,
658            converted_type: self.converted_type,
659            logical_type: self.logical_type.clone(),
660            id: self.id,
661            sort_order: SortOrder::UNDEFINED,
662        };
663        // Populate the converted type if only the logical type is populated
664        if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE {
665            basic_info.converted_type = self.logical_type.into();
666        }
667        Ok(Type::GroupType {
668            basic_info,
669            fields: self.fields,
670        })
671    }
672}
673
674/// Basic type info. This contains information such as the name of the type,
675/// the repetition level, the logical type and the kind of the type (group, primitive).
676#[derive(Clone, Debug, PartialEq, Eq)]
677pub struct BasicTypeInfo {
678    name: String,
679    repetition: Option<Repetition>,
680    converted_type: ConvertedType,
681    logical_type: Option<LogicalType>,
682    id: Option<i32>,
683    sort_order: SortOrder,
684}
685
686impl HeapSize for BasicTypeInfo {
687    fn heap_size(&self) -> usize {
688        // no heap allocations in any other subfield
689        self.name.heap_size()
690    }
691}
692
693impl BasicTypeInfo {
694    /// Returns field name.
695    pub fn name(&self) -> &str {
696        &self.name
697    }
698
699    /// Returns `true` if type has repetition field set, `false` otherwise.
700    /// This is mostly applied to group type, because primitive type always has
701    /// repetition set.
702    pub fn has_repetition(&self) -> bool {
703        self.repetition.is_some()
704    }
705
706    /// Returns [`Repetition`] value for the type.
707    pub fn repetition(&self) -> Repetition {
708        assert!(self.repetition.is_some());
709        self.repetition.unwrap()
710    }
711
712    /// Returns [`ConvertedType`] value for the type.
713    pub fn converted_type(&self) -> ConvertedType {
714        self.converted_type
715    }
716
717    /// Return a reference to the [`LogicalType`] value for the type.
718    pub fn logical_type_ref(&self) -> Option<&LogicalType> {
719        self.logical_type.as_ref()
720    }
721
722    /// Returns `true` if id is set, `false` otherwise.
723    pub fn has_id(&self) -> bool {
724        self.id.is_some()
725    }
726
727    /// Returns id value for the type.
728    pub fn id(&self) -> i32 {
729        assert!(self.id.is_some());
730        self.id.unwrap()
731    }
732
733    /// Returns [`SortOrder`] for the type.
734    pub fn sort_order(&self) -> SortOrder {
735        self.sort_order
736    }
737}
738
739// ----------------------------------------------------------------------
740// Parquet descriptor definitions
741
742/// Represents the location of a column in a Parquet schema
743///
744/// # Example: refer to column named `'my_column'`
745/// ```
746/// # use parquet::schema::types::ColumnPath;
747/// let column_path = ColumnPath::from("my_column");
748/// ```
749///
750/// # Example: refer to column named `c` in a nested struct `{a: {b: {c: ...}}}`
751/// ```
752/// # use parquet::schema::types::ColumnPath;
753/// // form path 'a.b.c'
754/// let column_path = ColumnPath::from(vec![
755///   String::from("a"),
756///   String::from("b"),
757///   String::from("c")
758/// ]);
759/// ```
760#[derive(Clone, PartialEq, Debug, Eq, Hash)]
761pub struct ColumnPath {
762    parts: Vec<String>,
763}
764
765impl HeapSize for ColumnPath {
766    fn heap_size(&self) -> usize {
767        self.parts.heap_size()
768    }
769}
770
771impl ColumnPath {
772    /// Creates new column path from vector of field names.
773    pub fn new(parts: Vec<String>) -> Self {
774        ColumnPath { parts }
775    }
776
777    /// Returns string representation of this column path.
778    /// ```rust
779    /// use parquet::schema::types::ColumnPath;
780    ///
781    /// let path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c".to_string()]);
782    /// assert_eq!(&path.string(), "a.b.c");
783    /// ```
784    pub fn string(&self) -> String {
785        self.parts.join(".")
786    }
787
788    /// Appends more components to end of column path.
789    /// ```rust
790    /// use parquet::schema::types::ColumnPath;
791    ///
792    /// let mut path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c"
793    /// .to_string()]);
794    /// assert_eq!(&path.string(), "a.b.c");
795    ///
796    /// path.append(vec!["d".to_string(), "e".to_string()]);
797    /// assert_eq!(&path.string(), "a.b.c.d.e");
798    /// ```
799    pub fn append(&mut self, mut tail: Vec<String>) {
800        self.parts.append(&mut tail);
801    }
802
803    /// Returns a slice of path components.
804    pub fn parts(&self) -> &[String] {
805        &self.parts
806    }
807}
808
809impl fmt::Display for ColumnPath {
810    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
811        write!(f, "{:?}", self.string())
812    }
813}
814
815impl From<Vec<String>> for ColumnPath {
816    fn from(parts: Vec<String>) -> Self {
817        ColumnPath { parts }
818    }
819}
820
821impl From<&str> for ColumnPath {
822    fn from(single_path: &str) -> Self {
823        let s = String::from(single_path);
824        ColumnPath::from(s)
825    }
826}
827
828impl From<String> for ColumnPath {
829    fn from(single_path: String) -> Self {
830        let v = vec![single_path];
831        ColumnPath { parts: v }
832    }
833}
834
835impl AsRef<[String]> for ColumnPath {
836    fn as_ref(&self) -> &[String] {
837        &self.parts
838    }
839}
840
841/// Physical type for leaf-level primitive columns.
842///
843/// Also includes the maximum definition and repetition levels required to
844/// re-assemble nested data.
845#[derive(Debug, PartialEq)]
846pub struct ColumnDescriptor {
847    /// The "leaf" primitive type of this column
848    primitive_type: TypePtr,
849
850    /// The maximum definition level for this column
851    max_def_level: i16,
852
853    /// The maximum repetition level for this column
854    max_rep_level: i16,
855
856    /// The definition level at the nearest REPEATED ancestor, or 0 if none.
857    repeated_ancestor_def_level: i16,
858
859    /// The path of this column. For instance, "a.b.c.d".
860    path: ColumnPath,
861}
862
863impl HeapSize for ColumnDescriptor {
864    fn heap_size(&self) -> usize {
865        // Don't include the heap size of primitive_type, this is already
866        // accounted for via SchemaDescriptor::schema
867        self.path.heap_size()
868    }
869}
870
871impl ColumnDescriptor {
872    /// Creates new descriptor for leaf-level column.
873    pub fn new(
874        primitive_type: TypePtr,
875        max_def_level: i16,
876        max_rep_level: i16,
877        path: ColumnPath,
878    ) -> Self {
879        Self::new_with_repeated_ancestor(primitive_type, max_def_level, max_rep_level, path, 0)
880    }
881
882    pub(crate) fn new_with_repeated_ancestor(
883        primitive_type: TypePtr,
884        max_def_level: i16,
885        max_rep_level: i16,
886        path: ColumnPath,
887        repeated_ancestor_def_level: i16,
888    ) -> Self {
889        Self {
890            primitive_type,
891            max_def_level,
892            max_rep_level,
893            repeated_ancestor_def_level,
894            path,
895        }
896    }
897
898    /// Returns maximum definition level for this column.
899    #[inline]
900    pub fn max_def_level(&self) -> i16 {
901        self.max_def_level
902    }
903
904    /// Returns maximum repetition level for this column.
905    #[inline]
906    pub fn max_rep_level(&self) -> i16 {
907        self.max_rep_level
908    }
909
910    /// Returns the definition level at the nearest REPEATED ancestor, or 0 if none.
911    #[inline]
912    pub fn repeated_ancestor_def_level(&self) -> i16 {
913        self.repeated_ancestor_def_level
914    }
915
916    /// Returns [`ColumnPath`] for this column.
917    pub fn path(&self) -> &ColumnPath {
918        &self.path
919    }
920
921    /// Returns self type [`Type`] for this leaf column.
922    pub fn self_type(&self) -> &Type {
923        self.primitive_type.as_ref()
924    }
925
926    /// Returns self type [`TypePtr`]  for this leaf
927    /// column.
928    pub fn self_type_ptr(&self) -> TypePtr {
929        self.primitive_type.clone()
930    }
931
932    /// Returns [`BasicTypeInfo`] information for this leaf column.
933    pub fn get_basic_info(&self) -> &BasicTypeInfo {
934        self.primitive_type.get_basic_info()
935    }
936
937    /// Returns column name.
938    pub fn name(&self) -> &str {
939        self.primitive_type.name()
940    }
941
942    /// Returns [`ConvertedType`] for this column.
943    pub fn converted_type(&self) -> ConvertedType {
944        self.primitive_type.get_basic_info().converted_type()
945    }
946
947    /// Returns a reference to the [`LogicalType`] for this column.
948    pub fn logical_type_ref(&self) -> Option<&LogicalType> {
949        self.primitive_type.get_basic_info().logical_type_ref()
950    }
951
952    /// Returns physical type for this column.
953    /// Note that it will panic if called on a non-primitive type.
954    pub fn physical_type(&self) -> PhysicalType {
955        match self.primitive_type.as_ref() {
956            Type::PrimitiveType { physical_type, .. } => *physical_type,
957            _ => panic!("Expected primitive type!"),
958        }
959    }
960
961    /// Returns type length for this column.
962    /// Note that it will panic if called on a non-primitive type.
963    pub fn type_length(&self) -> i32 {
964        match self.primitive_type.as_ref() {
965            Type::PrimitiveType { type_length, .. } => *type_length,
966            _ => panic!("Expected primitive type!"),
967        }
968    }
969
970    /// Returns type precision for this column.
971    /// Note that it will panic if called on a non-primitive type.
972    pub fn type_precision(&self) -> i32 {
973        match self.primitive_type.as_ref() {
974            Type::PrimitiveType { precision, .. } => *precision,
975            _ => panic!("Expected primitive type!"),
976        }
977    }
978
979    /// Returns type scale for this column.
980    /// Note that it will panic if called on a non-primitive type.
981    pub fn type_scale(&self) -> i32 {
982        match self.primitive_type.as_ref() {
983            Type::PrimitiveType { scale, .. } => *scale,
984            _ => panic!("Expected primitive type!"),
985        }
986    }
987
988    /// Returns the sort order for this column as currently defined for the logical or
989    /// physical type.
990    ///
991    /// Returns `SortOrder::UNDEFINED` for non-primitive types.
992    pub fn sort_order(&self) -> SortOrder {
993        self.primitive_type.get_basic_info().sort_order()
994    }
995}
996
997/// Schema of a Parquet file.
998///
999/// Encapsulates the file's schema ([`Type`]) and [`ColumnDescriptor`]s for
1000/// each primitive (leaf) column.
1001///
1002/// # Example
1003/// ```
1004/// # use std::sync::Arc;
1005/// use parquet::schema::types::{SchemaDescriptor, Type};
1006/// use parquet::basic; // note there are two `Type`s that are different
1007/// // Schema for a table with two columns: "a" (int64) and "b" (int32, stored as a date)
1008/// let descriptor = SchemaDescriptor::new(
1009///   Arc::new(
1010///     Type::group_type_builder("my_schema")
1011///       .with_fields(vec![
1012///         Arc::new(
1013///          Type::primitive_type_builder("a", basic::Type::INT64)
1014///           .build().unwrap()
1015///         ),
1016///         Arc::new(
1017///          Type::primitive_type_builder("b", basic::Type::INT32)
1018///           .with_converted_type(basic::ConvertedType::DATE)
1019///           .with_logical_type(Some(basic::LogicalType::Date))
1020///           .build().unwrap()
1021///         ),
1022///      ])
1023///      .build().unwrap()
1024///   )
1025/// );
1026/// ```
1027#[derive(PartialEq, Clone)]
1028pub struct SchemaDescriptor {
1029    /// The top-level logical schema (the "message" type).
1030    ///
1031    /// This must be a [`Type::GroupType`] where each field is a root
1032    /// column type in the schema.
1033    schema: TypePtr,
1034
1035    /// The descriptors for the physical type of each leaf column in this schema
1036    ///
1037    /// Constructed from `schema` in DFS order.
1038    leaves: Vec<ColumnDescPtr>,
1039
1040    /// Mapping from a leaf column's index to the root column index that it
1041    /// comes from.
1042    ///
1043    /// For instance: the leaf `a.b.c.d` would have a link back to `a`:
1044    /// ```text
1045    /// -- a  <-----+
1046    /// -- -- b     |
1047    /// -- -- -- c  |
1048    /// -- -- -- -- d
1049    /// ```
1050    leaf_to_base: Vec<usize>,
1051}
1052
1053impl fmt::Debug for SchemaDescriptor {
1054    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1055        // Skip leaves and leaf_to_base as they only a cache information already found in `schema`
1056        f.debug_struct("SchemaDescriptor")
1057            .field("schema", &self.schema)
1058            .finish()
1059    }
1060}
1061
1062// Need to implement HeapSize in this module as the fields are private
1063impl HeapSize for SchemaDescriptor {
1064    fn heap_size(&self) -> usize {
1065        self.schema.heap_size() + self.leaves.heap_size() + self.leaf_to_base.heap_size()
1066    }
1067}
1068
1069impl SchemaDescriptor {
1070    /// Creates new schema descriptor from Parquet schema.
1071    pub fn new(tp: TypePtr) -> Self {
1072        const INIT_SCHEMA_DEPTH: usize = 16;
1073        assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
1074        // unwrap should be safe since we just asserted tp is a group
1075        let n_leaves = num_leaves(&tp).unwrap();
1076        let mut leaves = Vec::with_capacity(n_leaves);
1077        let mut leaf_to_base = Vec::with_capacity(n_leaves);
1078        let mut path = Vec::with_capacity(INIT_SCHEMA_DEPTH);
1079        for (root_idx, f) in tp.get_fields().iter().enumerate() {
1080            path.clear();
1081            build_tree(
1082                f,
1083                root_idx,
1084                0,
1085                0,
1086                0,
1087                &mut leaves,
1088                &mut leaf_to_base,
1089                &mut path,
1090            );
1091        }
1092
1093        Self {
1094            schema: tp,
1095            leaves,
1096            leaf_to_base,
1097        }
1098    }
1099
1100    /// Returns [`ColumnDescriptor`] for a field position.
1101    pub fn column(&self, i: usize) -> ColumnDescPtr {
1102        assert!(
1103            i < self.leaves.len(),
1104            "Index out of bound: {} not in [0, {})",
1105            i,
1106            self.leaves.len()
1107        );
1108        self.leaves[i].clone()
1109    }
1110
1111    /// Returns slice of [`ColumnDescriptor`].
1112    pub fn columns(&self) -> &[ColumnDescPtr] {
1113        &self.leaves
1114    }
1115
1116    /// Returns number of leaf-level columns.
1117    pub fn num_columns(&self) -> usize {
1118        self.leaves.len()
1119    }
1120
1121    /// Returns column root [`Type`] for a leaf position.
1122    pub fn get_column_root(&self, i: usize) -> &Type {
1123        let result = self.column_root_of(i);
1124        result.as_ref()
1125    }
1126
1127    /// Returns column root [`Type`] pointer for a leaf position.
1128    pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
1129        let result = self.column_root_of(i);
1130        result.clone()
1131    }
1132
1133    /// Returns the index of the root column for a field position
1134    pub fn get_column_root_idx(&self, leaf: usize) -> usize {
1135        assert!(
1136            leaf < self.leaves.len(),
1137            "Index out of bound: {} not in [0, {})",
1138            leaf,
1139            self.leaves.len()
1140        );
1141
1142        *self
1143            .leaf_to_base
1144            .get(leaf)
1145            .unwrap_or_else(|| panic!("Expected a value for index {leaf} but found None"))
1146    }
1147
1148    fn column_root_of(&self, i: usize) -> &TypePtr {
1149        &self.schema.get_fields()[self.get_column_root_idx(i)]
1150    }
1151
1152    /// Returns schema as [`Type`].
1153    pub fn root_schema(&self) -> &Type {
1154        self.schema.as_ref()
1155    }
1156
1157    /// Returns schema as [`TypePtr`] for cheap cloning.
1158    pub fn root_schema_ptr(&self) -> TypePtr {
1159        self.schema.clone()
1160    }
1161
1162    /// Returns schema name.
1163    pub fn name(&self) -> &str {
1164        self.schema.name()
1165    }
1166}
1167
1168// walk tree and count nodes
1169pub(crate) fn num_nodes(tp: &TypePtr) -> Result<usize> {
1170    if !tp.is_group() {
1171        return Err(general_err!("Root schema must be Group type"));
1172    }
1173    let mut n_nodes = 1usize; // count root
1174    for f in tp.get_fields().iter() {
1175        count_nodes(f, &mut n_nodes);
1176    }
1177    Ok(n_nodes)
1178}
1179
1180pub(crate) fn count_nodes(tp: &TypePtr, n_nodes: &mut usize) {
1181    *n_nodes += 1;
1182    if let Type::GroupType { fields, .. } = tp.as_ref() {
1183        for f in fields {
1184            count_nodes(f, n_nodes);
1185        }
1186    }
1187}
1188
1189// do a quick walk of the tree to get proper sizing for SchemaDescriptor arrays
1190fn num_leaves(tp: &TypePtr) -> Result<usize> {
1191    if !tp.is_group() {
1192        return Err(general_err!("Root schema must be Group type"));
1193    }
1194    let mut n_leaves = 0usize;
1195    for f in tp.get_fields().iter() {
1196        count_leaves(f, &mut n_leaves);
1197    }
1198    Ok(n_leaves)
1199}
1200
1201fn count_leaves(tp: &TypePtr, n_leaves: &mut usize) {
1202    match tp.as_ref() {
1203        Type::PrimitiveType { .. } => *n_leaves += 1,
1204        Type::GroupType { fields, .. } => {
1205            for f in fields {
1206                count_leaves(f, n_leaves);
1207            }
1208        }
1209    }
1210}
1211
1212#[allow(clippy::too_many_arguments)]
1213fn build_tree<'a>(
1214    tp: &'a TypePtr,
1215    root_idx: usize,
1216    mut max_rep_level: i16,
1217    mut max_def_level: i16,
1218    mut repeated_ancestor_def_level: i16,
1219    leaves: &mut Vec<ColumnDescPtr>,
1220    leaf_to_base: &mut Vec<usize>,
1221    path_so_far: &mut Vec<&'a str>,
1222) {
1223    assert!(tp.get_basic_info().has_repetition());
1224
1225    path_so_far.push(tp.name());
1226    match tp.get_basic_info().repetition() {
1227        Repetition::OPTIONAL => {
1228            max_def_level += 1;
1229        }
1230        Repetition::REPEATED => {
1231            max_def_level += 1;
1232            max_rep_level += 1;
1233            repeated_ancestor_def_level = max_def_level;
1234        }
1235        _ => {}
1236    }
1237
1238    match tp.as_ref() {
1239        Type::PrimitiveType { .. } => {
1240            let mut path: Vec<String> = vec![];
1241            path.extend(path_so_far.iter().copied().map(String::from));
1242            let desc = ColumnDescriptor::new_with_repeated_ancestor(
1243                tp.clone(),
1244                max_def_level,
1245                max_rep_level,
1246                ColumnPath::new(path),
1247                repeated_ancestor_def_level,
1248            );
1249            leaves.push(Arc::new(desc));
1250            leaf_to_base.push(root_idx);
1251        }
1252        Type::GroupType { fields, .. } => {
1253            for f in fields {
1254                build_tree(
1255                    f,
1256                    root_idx,
1257                    max_rep_level,
1258                    max_def_level,
1259                    repeated_ancestor_def_level,
1260                    leaves,
1261                    leaf_to_base,
1262                    path_so_far,
1263                );
1264                path_so_far.pop();
1265            }
1266        }
1267    }
1268}
1269
1270/// Checks if the logical type is valid.
1271fn check_logical_type(logical_type: &Option<LogicalType>) -> Result<()> {
1272    if let Some(LogicalType::Integer(IntType { bit_width, .. })) = logical_type
1273        && *bit_width != 8
1274        && *bit_width != 16
1275        && *bit_width != 32
1276        && *bit_width != 64
1277    {
1278        return Err(general_err!(
1279            "Bit width must be 8, 16, 32, or 64 for Integer logical type"
1280        ));
1281    }
1282    Ok(())
1283}
1284
1285// convert thrift decoded array of `SchemaElement` into this crate's representation of
1286// parquet types. this function consumes `elements`.
1287pub(crate) fn parquet_schema_from_array<'a>(elements: Vec<SchemaElement<'a>>) -> Result<TypePtr> {
1288    let mut index = 0;
1289    let num_elements = elements.len();
1290    let mut schema_nodes = Vec::with_capacity(1); // there should only be one element when done
1291
1292    // turn into iterator so we can take ownership of elements of the vector
1293    let mut elements = elements.into_iter();
1294
1295    while index < num_elements {
1296        let t = schema_from_array_helper(&mut elements, num_elements, index)?;
1297        index = t.0;
1298        schema_nodes.push(t.1);
1299    }
1300    if schema_nodes.len() != 1 {
1301        return Err(general_err!(
1302            "Expected exactly one root node, but found {}",
1303            schema_nodes.len()
1304        ));
1305    }
1306
1307    if !schema_nodes[0].is_group() {
1308        return Err(general_err!("Expected root node to be a group type"));
1309    }
1310
1311    Ok(schema_nodes.remove(0))
1312}
1313
1314// recursive helper function for schema conversion
1315fn schema_from_array_helper<'a>(
1316    elements: &mut IntoIter<SchemaElement<'a>>,
1317    num_elements: usize,
1318    index: usize,
1319) -> Result<(usize, TypePtr)> {
1320    // Whether or not the current node is root (message type).
1321    // There is only one message type node in the schema tree.
1322    let is_root_node = index == 0;
1323
1324    if index >= num_elements {
1325        return Err(general_err!(
1326            "Index out of bound, index = {}, len = {}",
1327            index,
1328            num_elements
1329        ));
1330    }
1331    let element = elements.next().expect("schema vector should not be empty");
1332
1333    // Check for empty schema
1334    if let (true, None | Some(0)) = (is_root_node, element.num_children) {
1335        let builder = Type::group_type_builder(element.name);
1336        return Ok((index + 1, Arc::new(builder.build().unwrap())));
1337    }
1338
1339    let converted_type = element.converted_type.unwrap_or(ConvertedType::NONE);
1340
1341    // LogicalType is prefered to ConvertedType, but both may be present.
1342    let logical_type = element.logical_type;
1343
1344    check_logical_type(&logical_type)?;
1345
1346    let field_id = element.field_id;
1347    match element.num_children {
1348        // From parquet-format:
1349        //   The children count is used to construct the nested relationship.
1350        //   This field is not set when the element is a primitive type
1351        // Sometimes parquet-cpp sets num_children field to 0 for primitive types, so we
1352        // have to handle this case too.
1353        None | Some(0) => {
1354            // primitive type
1355            if element.repetition_type.is_none() {
1356                return Err(general_err!(
1357                    "Repetition level must be defined for a primitive type"
1358                ));
1359            }
1360            let repetition = element.repetition_type.unwrap();
1361            if let Some(physical_type) = element.r#type {
1362                let length = element.type_length.unwrap_or(-1);
1363                let scale = element.scale.unwrap_or(-1);
1364                let precision = element.precision.unwrap_or(-1);
1365                let name = element.name;
1366                let builder = Type::primitive_type_builder(name, physical_type)
1367                    .with_repetition(repetition)
1368                    .with_converted_type(converted_type)
1369                    .with_logical_type(logical_type)
1370                    .with_length(length)
1371                    .with_precision(precision)
1372                    .with_scale(scale)
1373                    .with_id(field_id);
1374                Ok((index + 1, Arc::new(builder.build()?)))
1375            } else {
1376                let mut builder = Type::group_type_builder(element.name)
1377                    .with_converted_type(converted_type)
1378                    .with_logical_type(logical_type)
1379                    .with_id(field_id);
1380                if !is_root_node {
1381                    // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or
1382                    // REPEATED for root node.
1383                    //
1384                    // We only set repetition for group types that are not top-level message
1385                    // type. According to parquet-format:
1386                    //   Root of the schema does not have a repetition_type.
1387                    //   All other types must have one.
1388                    builder = builder.with_repetition(repetition);
1389                }
1390                Ok((index + 1, Arc::new(builder.build().unwrap())))
1391            }
1392        }
1393        Some(n) => {
1394            let repetition = element.repetition_type;
1395
1396            let mut fields = Vec::with_capacity(usize::try_from(n)?);
1397            let mut next_index = index + 1;
1398            for _ in 0..n {
1399                let child_result = schema_from_array_helper(elements, num_elements, next_index)?;
1400                next_index = child_result.0;
1401                fields.push(child_result.1);
1402            }
1403
1404            let mut builder = Type::group_type_builder(element.name)
1405                .with_converted_type(converted_type)
1406                .with_logical_type(logical_type)
1407                .with_fields(fields)
1408                .with_id(field_id);
1409
1410            // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or
1411            // REPEATED for root node.
1412            //
1413            // We only set repetition for group types that are not top-level message
1414            // type. According to parquet-format:
1415            //   Root of the schema does not have a repetition_type.
1416            //   All other types must have one.
1417            if !is_root_node {
1418                let Some(rep) = repetition else {
1419                    return Err(general_err!(
1420                        "Repetition level must be defined for non-root types"
1421                    ));
1422                };
1423                builder = builder.with_repetition(rep);
1424            }
1425            Ok((next_index, Arc::new(builder.build()?)))
1426        }
1427    }
1428}
1429
1430#[cfg(test)]
1431mod tests {
1432    use super::*;
1433
1434    use crate::{
1435        file::metadata::thrift::tests::{buf_to_schema_list, roundtrip_schema, schema_to_buf},
1436        schema::parser::parse_message_type,
1437    };
1438
1439    // TODO: add tests for v2 types
1440
1441    #[test]
1442    fn test_primitive_type() {
1443        let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1444            .with_logical_type(Some(LogicalType::integer(32, true)))
1445            .with_id(Some(0))
1446            .build();
1447        assert!(result.is_ok());
1448
1449        if let Ok(tp) = result {
1450            assert!(tp.is_primitive());
1451            assert!(!tp.is_group());
1452            let basic_info = tp.get_basic_info();
1453            assert_eq!(basic_info.repetition(), Repetition::OPTIONAL);
1454            assert_eq!(
1455                basic_info.logical_type_ref(),
1456                Some(&LogicalType::integer(32, true))
1457            );
1458            assert_eq!(basic_info.converted_type(), ConvertedType::INT_32);
1459            assert_eq!(basic_info.id(), 0);
1460            match tp {
1461                Type::PrimitiveType { physical_type, .. } => {
1462                    assert_eq!(physical_type, PhysicalType::INT32);
1463                }
1464                _ => panic!(),
1465            }
1466        }
1467
1468        // Test illegal inputs with logical type
1469        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1470            .with_repetition(Repetition::REPEATED)
1471            .with_logical_type(Some(LogicalType::integer(8, true)))
1472            .build();
1473        assert!(result.is_err());
1474        if let Err(e) = result {
1475            assert_eq!(
1476                format!("{e}"),
1477                "Parquet error: Cannot annotate Integer(IntType { bit_width: 8, is_signed: true }) from INT64 for field 'foo'"
1478            );
1479        }
1480
1481        // Test illegal inputs with converted type
1482        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1483            .with_repetition(Repetition::REPEATED)
1484            .with_converted_type(ConvertedType::BSON)
1485            .build();
1486        assert!(result.is_err());
1487        if let Err(e) = result {
1488            assert_eq!(
1489                format!("{e}"),
1490                "Parquet error: BSON cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
1491            );
1492        }
1493
1494        result = Type::primitive_type_builder("foo", PhysicalType::INT96)
1495            .with_repetition(Repetition::REQUIRED)
1496            .with_converted_type(ConvertedType::DECIMAL)
1497            .with_precision(-1)
1498            .with_scale(-1)
1499            .build();
1500        assert!(result.is_err());
1501        if let Err(e) = result {
1502            assert_eq!(
1503                format!("{e}"),
1504                "Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
1505            );
1506        }
1507
1508        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1509            .with_repetition(Repetition::REQUIRED)
1510            .with_logical_type(Some(LogicalType::decimal(32, 12)))
1511            .with_precision(-1)
1512            .with_scale(-1)
1513            .build();
1514        assert!(result.is_err());
1515        if let Err(e) = result {
1516            assert_eq!(
1517                format!("{e}"),
1518                "Parquet error: DECIMAL logical type scale 32 must match self.scale -1 for field 'foo'"
1519            );
1520        }
1521
1522        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1523            .with_repetition(Repetition::REQUIRED)
1524            .with_converted_type(ConvertedType::DECIMAL)
1525            .with_precision(-1)
1526            .with_scale(-1)
1527            .build();
1528        assert!(result.is_err());
1529        if let Err(e) = result {
1530            assert_eq!(
1531                format!("{e}"),
1532                "Parquet error: Invalid DECIMAL precision: -1"
1533            );
1534        }
1535
1536        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1537            .with_repetition(Repetition::REQUIRED)
1538            .with_converted_type(ConvertedType::DECIMAL)
1539            .with_precision(0)
1540            .with_scale(-1)
1541            .build();
1542        assert!(result.is_err());
1543        if let Err(e) = result {
1544            assert_eq!(
1545                format!("{e}"),
1546                "Parquet error: Invalid DECIMAL precision: 0"
1547            );
1548        }
1549
1550        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1551            .with_repetition(Repetition::REQUIRED)
1552            .with_converted_type(ConvertedType::DECIMAL)
1553            .with_precision(1)
1554            .with_scale(-1)
1555            .build();
1556        assert!(result.is_err());
1557        if let Err(e) = result {
1558            assert_eq!(format!("{e}"), "Parquet error: Invalid DECIMAL scale: -1");
1559        }
1560
1561        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1562            .with_repetition(Repetition::REQUIRED)
1563            .with_converted_type(ConvertedType::DECIMAL)
1564            .with_precision(1)
1565            .with_scale(2)
1566            .build();
1567        assert!(result.is_err());
1568        if let Err(e) = result {
1569            assert_eq!(
1570                format!("{e}"),
1571                "Parquet error: Invalid DECIMAL: scale (2) cannot be greater than precision (1)"
1572            );
1573        }
1574
1575        // It is OK if precision == scale
1576        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1577            .with_repetition(Repetition::REQUIRED)
1578            .with_converted_type(ConvertedType::DECIMAL)
1579            .with_precision(1)
1580            .with_scale(1)
1581            .build();
1582        assert!(result.is_ok());
1583
1584        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1585            .with_repetition(Repetition::REQUIRED)
1586            .with_converted_type(ConvertedType::DECIMAL)
1587            .with_precision(18)
1588            .with_scale(2)
1589            .build();
1590        assert!(result.is_err());
1591        if let Err(e) = result {
1592            assert_eq!(
1593                format!("{e}"),
1594                "Parquet error: Cannot represent INT32 as DECIMAL with precision 18"
1595            );
1596        }
1597
1598        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1599            .with_repetition(Repetition::REQUIRED)
1600            .with_converted_type(ConvertedType::DECIMAL)
1601            .with_precision(32)
1602            .with_scale(2)
1603            .build();
1604        assert!(result.is_err());
1605        if let Err(e) = result {
1606            assert_eq!(
1607                format!("{e}"),
1608                "Parquet error: Cannot represent INT64 as DECIMAL with precision 32"
1609            );
1610        }
1611
1612        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1613            .with_repetition(Repetition::REQUIRED)
1614            .with_converted_type(ConvertedType::DECIMAL)
1615            .with_length(5)
1616            .with_precision(12)
1617            .with_scale(2)
1618            .build();
1619        assert!(result.is_err());
1620        if let Err(e) = result {
1621            assert_eq!(
1622                format!("{e}"),
1623                "Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11"
1624            );
1625        }
1626
1627        result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1628            .with_repetition(Repetition::REQUIRED)
1629            .with_converted_type(ConvertedType::UINT_8)
1630            .build();
1631        assert!(result.is_err());
1632        if let Err(e) = result {
1633            assert_eq!(
1634                format!("{e}"),
1635                "Parquet error: UINT_8 cannot annotate field 'foo' because it is not a INT32 field"
1636            );
1637        }
1638
1639        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1640            .with_repetition(Repetition::REQUIRED)
1641            .with_converted_type(ConvertedType::TIME_MICROS)
1642            .build();
1643        assert!(result.is_err());
1644        if let Err(e) = result {
1645            assert_eq!(
1646                format!("{e}"),
1647                "Parquet error: TIME_MICROS cannot annotate field 'foo' because it is not a INT64 field"
1648            );
1649        }
1650
1651        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1652            .with_repetition(Repetition::REQUIRED)
1653            .with_converted_type(ConvertedType::INTERVAL)
1654            .build();
1655        assert!(result.is_err());
1656        if let Err(e) = result {
1657            assert_eq!(
1658                format!("{e}"),
1659                "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
1660            );
1661        }
1662
1663        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1664            .with_repetition(Repetition::REQUIRED)
1665            .with_converted_type(ConvertedType::INTERVAL)
1666            .with_length(1)
1667            .build();
1668        assert!(result.is_err());
1669        if let Err(e) = result {
1670            assert_eq!(
1671                format!("{e}"),
1672                "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
1673            );
1674        }
1675
1676        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1677            .with_repetition(Repetition::REQUIRED)
1678            .with_converted_type(ConvertedType::ENUM)
1679            .build();
1680        assert!(result.is_err());
1681        if let Err(e) = result {
1682            assert_eq!(
1683                format!("{e}"),
1684                "Parquet error: ENUM cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
1685            );
1686        }
1687
1688        result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1689            .with_repetition(Repetition::REQUIRED)
1690            .with_converted_type(ConvertedType::MAP)
1691            .build();
1692        assert!(result.is_err());
1693        if let Err(e) = result {
1694            assert_eq!(
1695                format!("{e}"),
1696                "Parquet error: MAP cannot be applied to primitive field 'foo'"
1697            );
1698        }
1699
1700        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1701            .with_repetition(Repetition::REQUIRED)
1702            .with_converted_type(ConvertedType::DECIMAL)
1703            .with_length(-1)
1704            .build();
1705        assert!(result.is_err());
1706        if let Err(e) = result {
1707            assert_eq!(
1708                format!("{e}"),
1709                "Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1 for field 'foo'"
1710            );
1711        }
1712
1713        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1714            .with_repetition(Repetition::REQUIRED)
1715            .with_logical_type(Some(LogicalType::Float16))
1716            .with_length(2)
1717            .build();
1718        assert!(result.is_ok());
1719
1720        // Can't be other than FIXED_LEN_BYTE_ARRAY for physical type
1721        result = Type::primitive_type_builder("foo", PhysicalType::FLOAT)
1722            .with_repetition(Repetition::REQUIRED)
1723            .with_logical_type(Some(LogicalType::Float16))
1724            .with_length(2)
1725            .build();
1726        assert!(result.is_err());
1727        if let Err(e) = result {
1728            assert_eq!(
1729                format!("{e}"),
1730                "Parquet error: Cannot annotate Float16 from FLOAT for field 'foo'"
1731            );
1732        }
1733
1734        // Must have length 2
1735        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1736            .with_repetition(Repetition::REQUIRED)
1737            .with_logical_type(Some(LogicalType::Float16))
1738            .with_length(4)
1739            .build();
1740        assert!(result.is_err());
1741        if let Err(e) = result {
1742            assert_eq!(
1743                format!("{e}"),
1744                "Parquet error: FLOAT16 cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(2) field"
1745            );
1746        }
1747
1748        // Must have length 16
1749        result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1750            .with_repetition(Repetition::REQUIRED)
1751            .with_logical_type(Some(LogicalType::Uuid))
1752            .with_length(15)
1753            .build();
1754        assert!(result.is_err());
1755        if let Err(e) = result {
1756            assert_eq!(
1757                format!("{e}"),
1758                "Parquet error: UUID cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(16) field"
1759            );
1760        }
1761
1762        // test unknown logical types are ok
1763        result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1764            .with_logical_type(Some(LogicalType::_Unknown { field_id: 100 }))
1765            .build();
1766        assert!(result.is_ok());
1767    }
1768
1769    #[test]
1770    fn test_group_type() {
1771        let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
1772            .with_converted_type(ConvertedType::INT_32)
1773            .with_id(Some(0))
1774            .build();
1775        assert!(f1.is_ok());
1776        let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
1777            .with_converted_type(ConvertedType::UTF8)
1778            .with_id(Some(1))
1779            .build();
1780        assert!(f2.is_ok());
1781
1782        let fields = vec![Arc::new(f1.unwrap()), Arc::new(f2.unwrap())];
1783
1784        let result = Type::group_type_builder("foo")
1785            .with_repetition(Repetition::REPEATED)
1786            .with_logical_type(Some(LogicalType::List))
1787            .with_fields(fields)
1788            .with_id(Some(1))
1789            .build();
1790        assert!(result.is_ok());
1791
1792        let tp = result.unwrap();
1793        let basic_info = tp.get_basic_info();
1794        assert!(tp.is_group());
1795        assert!(!tp.is_primitive());
1796        assert_eq!(basic_info.repetition(), Repetition::REPEATED);
1797        assert_eq!(basic_info.logical_type_ref(), Some(&LogicalType::List));
1798        assert_eq!(basic_info.converted_type(), ConvertedType::LIST);
1799        assert_eq!(basic_info.id(), 1);
1800        assert_eq!(tp.get_fields().len(), 2);
1801        assert_eq!(tp.get_fields()[0].name(), "f1");
1802        assert_eq!(tp.get_fields()[1].name(), "f2");
1803    }
1804
1805    #[test]
1806    fn test_column_descriptor() {
1807        let result = test_column_descriptor_helper();
1808        assert!(
1809            result.is_ok(),
1810            "Expected result to be OK but got err:\n {}",
1811            result.unwrap_err()
1812        );
1813    }
1814
1815    fn test_column_descriptor_helper() -> Result<()> {
1816        let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY)
1817            .with_converted_type(ConvertedType::UTF8)
1818            .build()?;
1819
1820        let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name"));
1821
1822        assert_eq!(descr.path(), &ColumnPath::from("name"));
1823        assert_eq!(descr.converted_type(), ConvertedType::UTF8);
1824        assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY);
1825        assert_eq!(descr.max_def_level(), 4);
1826        assert_eq!(descr.max_rep_level(), 1);
1827        assert_eq!(descr.name(), "name");
1828        assert_eq!(descr.type_length(), -1);
1829        assert_eq!(descr.type_precision(), -1);
1830        assert_eq!(descr.type_scale(), -1);
1831
1832        Ok(())
1833    }
1834
1835    #[test]
1836    fn test_schema_descriptor() {
1837        let result = test_schema_descriptor_helper();
1838        assert!(
1839            result.is_ok(),
1840            "Expected result to be OK but got err:\n {}",
1841            result.unwrap_err()
1842        );
1843    }
1844
1845    // A helper fn to avoid handling the results from type creation
1846    fn test_schema_descriptor_helper() -> Result<()> {
1847        let mut fields = vec![];
1848
1849        let inta = Type::primitive_type_builder("a", PhysicalType::INT32)
1850            .with_repetition(Repetition::REQUIRED)
1851            .with_converted_type(ConvertedType::INT_32)
1852            .build()?;
1853        fields.push(Arc::new(inta));
1854        let intb = Type::primitive_type_builder("b", PhysicalType::INT64)
1855            .with_converted_type(ConvertedType::INT_64)
1856            .build()?;
1857        fields.push(Arc::new(intb));
1858        let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY)
1859            .with_repetition(Repetition::REPEATED)
1860            .with_converted_type(ConvertedType::UTF8)
1861            .build()?;
1862        fields.push(Arc::new(intc));
1863
1864        // 3-level list encoding
1865        let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64)
1866            .with_repetition(Repetition::REQUIRED)
1867            .with_converted_type(ConvertedType::INT_64)
1868            .build()?;
1869        let item2 = Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?;
1870        let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32)
1871            .with_repetition(Repetition::REPEATED)
1872            .with_converted_type(ConvertedType::INT_32)
1873            .build()?;
1874        let list = Type::group_type_builder("records")
1875            .with_repetition(Repetition::REPEATED)
1876            .with_converted_type(ConvertedType::LIST)
1877            .with_fields(vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)])
1878            .build()?;
1879        let bag = Type::group_type_builder("bag")
1880            .with_repetition(Repetition::OPTIONAL)
1881            .with_fields(vec![Arc::new(list)])
1882            .build()?;
1883        fields.push(Arc::new(bag));
1884
1885        let schema = Type::group_type_builder("schema")
1886            .with_repetition(Repetition::REPEATED)
1887            .with_fields(fields)
1888            .build()?;
1889        let descr = SchemaDescriptor::new(Arc::new(schema));
1890
1891        let nleaves = 6;
1892        assert_eq!(descr.num_columns(), nleaves);
1893
1894        //                             mdef mrep
1895        // required int32 a            0    0
1896        // optional int64 b            1    0
1897        // repeated byte_array c       1    1
1898        // optional group bag          1    0
1899        //   repeated group records    2    1
1900        //     required int64 item1    2    1
1901        //     optional boolean item2  3    1
1902        //     repeated int32 item3    3    2
1903        let ex_max_def_levels = [0, 1, 1, 2, 3, 3];
1904        let ex_max_rep_levels = [0, 0, 1, 1, 1, 2];
1905
1906        for i in 0..nleaves {
1907            let col = descr.column(i);
1908            assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{i}");
1909            assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{i}");
1910        }
1911
1912        assert_eq!(descr.column(0).path().string(), "a");
1913        assert_eq!(descr.column(1).path().string(), "b");
1914        assert_eq!(descr.column(2).path().string(), "c");
1915        assert_eq!(descr.column(3).path().string(), "bag.records.item1");
1916        assert_eq!(descr.column(4).path().string(), "bag.records.item2");
1917        assert_eq!(descr.column(5).path().string(), "bag.records.item3");
1918
1919        assert_eq!(descr.get_column_root(0).name(), "a");
1920        assert_eq!(descr.get_column_root(3).name(), "bag");
1921        assert_eq!(descr.get_column_root(4).name(), "bag");
1922        assert_eq!(descr.get_column_root(5).name(), "bag");
1923
1924        Ok(())
1925    }
1926
1927    #[test]
1928    fn test_schema_build_tree_def_rep_levels() {
1929        let message_type = "
1930    message spark_schema {
1931      REQUIRED INT32 a;
1932      OPTIONAL group b {
1933        OPTIONAL INT32 _1;
1934        OPTIONAL INT32 _2;
1935      }
1936      OPTIONAL group c (LIST) {
1937        REPEATED group list {
1938          OPTIONAL INT32 element;
1939        }
1940      }
1941    }
1942    ";
1943        let schema = parse_message_type(message_type).expect("should parse schema");
1944        let descr = SchemaDescriptor::new(Arc::new(schema));
1945        // required int32 a
1946        assert_eq!(descr.column(0).max_def_level(), 0);
1947        assert_eq!(descr.column(0).max_rep_level(), 0);
1948        // optional int32 b._1
1949        assert_eq!(descr.column(1).max_def_level(), 2);
1950        assert_eq!(descr.column(1).max_rep_level(), 0);
1951        // optional int32 b._2
1952        assert_eq!(descr.column(2).max_def_level(), 2);
1953        assert_eq!(descr.column(2).max_rep_level(), 0);
1954        // repeated optional int32 c.list.element
1955        assert_eq!(descr.column(3).max_def_level(), 3);
1956        assert_eq!(descr.column(3).max_rep_level(), 1);
1957    }
1958
1959    #[test]
1960    fn test_schema_build_tree_repeated_ancestor_def_level() {
1961        // Flat columns: no REPEATED ancestor → repeated_ancestor_def_level = 0
1962        let message_type = "
1963    message m {
1964      REQUIRED INT32 a;
1965      OPTIONAL INT32 b;
1966      OPTIONAL group s {
1967        OPTIONAL INT32 x;
1968      }
1969    }
1970    ";
1971        let schema = parse_message_type(message_type).expect("should parse schema");
1972        let descr = SchemaDescriptor::new(Arc::new(schema));
1973        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 0); // a
1974        assert_eq!(descr.column(1).repeated_ancestor_def_level(), 0); // b
1975        assert_eq!(descr.column(2).repeated_ancestor_def_level(), 0); // s.x
1976
1977        // Standard list: OPTIONAL outer, REPEATED group, OPTIONAL element
1978        // repeated_ancestor_def_level is the def_level at the REPEATED group (= 2)
1979        let message_type = "
1980    message m {
1981      OPTIONAL group c (LIST) {
1982        REPEATED group list {
1983          OPTIONAL INT32 element;
1984        }
1985      }
1986    }
1987    ";
1988        let schema = parse_message_type(message_type).expect("should parse schema");
1989        let descr = SchemaDescriptor::new(Arc::new(schema));
1990        // c(optional)=1, list(repeated)=2, element(optional)=3
1991        assert_eq!(descr.column(0).max_def_level(), 3);
1992        assert_eq!(descr.column(0).max_rep_level(), 1);
1993        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2);
1994
1995        // Required list: REQUIRED outer, REPEATED group, REQUIRED element
1996        // No OPTIONAL nodes between REPEATED and leaf, so repeated_ancestor_def_level == max_def_level
1997        let message_type = "
1998    message m {
1999      REQUIRED group c (LIST) {
2000        REPEATED group list {
2001          REQUIRED INT32 element;
2002        }
2003      }
2004    }
2005    ";
2006        let schema = parse_message_type(message_type).expect("should parse schema");
2007        let descr = SchemaDescriptor::new(Arc::new(schema));
2008        // list(repeated)=1, element(required)=1
2009        assert_eq!(descr.column(0).max_def_level(), 1);
2010        assert_eq!(descr.column(0).max_rep_level(), 1);
2011        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 1);
2012
2013        // Nested lists: innermost REPEATED wins
2014        let message_type = "
2015    message m {
2016      OPTIONAL group outer (LIST) {
2017        REPEATED group list {
2018          OPTIONAL group inner (LIST) {
2019            REPEATED group list2 {
2020              OPTIONAL INT32 element;
2021            }
2022          }
2023        }
2024      }
2025    }
2026    ";
2027        let schema = parse_message_type(message_type).expect("should parse schema");
2028        let descr = SchemaDescriptor::new(Arc::new(schema));
2029        // outer(opt)=1, list(rep)=2, inner(opt)=3, list2(rep)=4, element(opt)=5
2030        assert_eq!(descr.column(0).max_def_level(), 5);
2031        assert_eq!(descr.column(0).max_rep_level(), 2);
2032        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 4);
2033
2034        // Struct inside list: all sibling leaves share the same repeated_ancestor_def_level
2035        let message_type = "
2036    message m {
2037      OPTIONAL group bag (LIST) {
2038        REPEATED group list {
2039          REQUIRED group item {
2040            OPTIONAL INT32 x;
2041            REQUIRED INT32 y;
2042          }
2043        }
2044      }
2045    }
2046    ";
2047        let schema = parse_message_type(message_type).expect("should parse schema");
2048        let descr = SchemaDescriptor::new(Arc::new(schema));
2049        // bag(opt)=1, list(rep)=2, item(req)=2, x(opt)=3
2050        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); // bag.list.item.x
2051        // bag(opt)=1, list(rep)=2, item(req)=2, y(req)=2
2052        assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); // bag.list.item.y
2053
2054        // Map type: key (required) and value (optional) under the same REPEATED group
2055        let message_type = "
2056    message m {
2057      OPTIONAL group my_map (MAP) {
2058        REPEATED group key_value {
2059          REQUIRED BYTE_ARRAY key (UTF8);
2060          OPTIONAL INT32 value;
2061        }
2062      }
2063    }
2064    ";
2065        let schema = parse_message_type(message_type).expect("should parse schema");
2066        let descr = SchemaDescriptor::new(Arc::new(schema));
2067        // my_map(opt)=1, key_value(rep)=2, key(req)=2
2068        assert_eq!(descr.column(0).max_def_level(), 2);
2069        assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); // key: max_def == repeated_ancestor
2070        // my_map(opt)=1, key_value(rep)=2, value(opt)=3
2071        assert_eq!(descr.column(1).max_def_level(), 3);
2072        assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); // value: max_def > repeated_ancestor
2073    }
2074
2075    #[test]
2076    #[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")]
2077    fn test_get_physical_type_panic() {
2078        let list = Type::group_type_builder("records")
2079            .with_repetition(Repetition::REPEATED)
2080            .build()
2081            .unwrap();
2082        list.get_physical_type();
2083    }
2084
2085    #[test]
2086    fn test_get_physical_type_primitive() {
2087        let f = Type::primitive_type_builder("f", PhysicalType::INT64)
2088            .build()
2089            .unwrap();
2090        assert_eq!(f.get_physical_type(), PhysicalType::INT64);
2091
2092        let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY)
2093            .build()
2094            .unwrap();
2095        assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY);
2096    }
2097
2098    #[test]
2099    fn test_check_contains_primitive_primitive() {
2100        // OK
2101        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2102            .build()
2103            .unwrap();
2104        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2105            .build()
2106            .unwrap();
2107        assert!(f1.check_contains(&f2));
2108
2109        // OK: different logical type does not affect check_contains
2110        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2111            .with_converted_type(ConvertedType::UINT_8)
2112            .build()
2113            .unwrap();
2114        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2115            .with_converted_type(ConvertedType::UINT_16)
2116            .build()
2117            .unwrap();
2118        assert!(f1.check_contains(&f2));
2119
2120        // KO: different name
2121        let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
2122            .build()
2123            .unwrap();
2124        let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
2125            .build()
2126            .unwrap();
2127        assert!(!f1.check_contains(&f2));
2128
2129        // KO: different type
2130        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2131            .build()
2132            .unwrap();
2133        let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
2134            .build()
2135            .unwrap();
2136        assert!(!f1.check_contains(&f2));
2137
2138        // KO: different repetition
2139        let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2140            .with_repetition(Repetition::REQUIRED)
2141            .build()
2142            .unwrap();
2143        let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2144            .with_repetition(Repetition::OPTIONAL)
2145            .build()
2146            .unwrap();
2147        assert!(!f1.check_contains(&f2));
2148    }
2149
2150    // function to create a new group type for testing
2151    fn test_new_group_type(name: &str, repetition: Repetition, types: Vec<Type>) -> Type {
2152        Type::group_type_builder(name)
2153            .with_repetition(repetition)
2154            .with_fields(types.into_iter().map(Arc::new).collect())
2155            .build()
2156            .unwrap()
2157    }
2158
2159    #[test]
2160    fn test_check_contains_group_group() {
2161        // OK: should match okay with empty fields
2162        let f1 = Type::group_type_builder("f").build().unwrap();
2163        let f2 = Type::group_type_builder("f").build().unwrap();
2164        assert!(f1.check_contains(&f2));
2165        assert!(!f1.is_optional());
2166
2167        // OK: fields match
2168        let f1 = test_new_group_type(
2169            "f",
2170            Repetition::REPEATED,
2171            vec![
2172                Type::primitive_type_builder("f1", PhysicalType::INT32)
2173                    .build()
2174                    .unwrap(),
2175                Type::primitive_type_builder("f2", PhysicalType::INT64)
2176                    .build()
2177                    .unwrap(),
2178            ],
2179        );
2180        let f2 = test_new_group_type(
2181            "f",
2182            Repetition::REPEATED,
2183            vec![
2184                Type::primitive_type_builder("f1", PhysicalType::INT32)
2185                    .build()
2186                    .unwrap(),
2187                Type::primitive_type_builder("f2", PhysicalType::INT64)
2188                    .build()
2189                    .unwrap(),
2190            ],
2191        );
2192        assert!(f1.check_contains(&f2));
2193
2194        // OK: subset of fields
2195        let f1 = test_new_group_type(
2196            "f",
2197            Repetition::REPEATED,
2198            vec![
2199                Type::primitive_type_builder("f1", PhysicalType::INT32)
2200                    .build()
2201                    .unwrap(),
2202                Type::primitive_type_builder("f2", PhysicalType::INT64)
2203                    .build()
2204                    .unwrap(),
2205            ],
2206        );
2207        let f2 = test_new_group_type(
2208            "f",
2209            Repetition::REPEATED,
2210            vec![
2211                Type::primitive_type_builder("f2", PhysicalType::INT64)
2212                    .build()
2213                    .unwrap(),
2214            ],
2215        );
2216        assert!(f1.check_contains(&f2));
2217
2218        // KO: different name
2219        let f1 = Type::group_type_builder("f1").build().unwrap();
2220        let f2 = Type::group_type_builder("f2").build().unwrap();
2221        assert!(!f1.check_contains(&f2));
2222
2223        // KO: different repetition
2224        let f1 = Type::group_type_builder("f")
2225            .with_repetition(Repetition::OPTIONAL)
2226            .build()
2227            .unwrap();
2228        let f2 = Type::group_type_builder("f")
2229            .with_repetition(Repetition::REPEATED)
2230            .build()
2231            .unwrap();
2232        assert!(!f1.check_contains(&f2));
2233
2234        // KO: different fields
2235        let f1 = test_new_group_type(
2236            "f",
2237            Repetition::REPEATED,
2238            vec![
2239                Type::primitive_type_builder("f1", PhysicalType::INT32)
2240                    .build()
2241                    .unwrap(),
2242                Type::primitive_type_builder("f2", PhysicalType::INT64)
2243                    .build()
2244                    .unwrap(),
2245            ],
2246        );
2247        let f2 = test_new_group_type(
2248            "f",
2249            Repetition::REPEATED,
2250            vec![
2251                Type::primitive_type_builder("f1", PhysicalType::INT32)
2252                    .build()
2253                    .unwrap(),
2254                Type::primitive_type_builder("f2", PhysicalType::BOOLEAN)
2255                    .build()
2256                    .unwrap(),
2257            ],
2258        );
2259        assert!(!f1.check_contains(&f2));
2260
2261        // KO: different fields
2262        let f1 = test_new_group_type(
2263            "f",
2264            Repetition::REPEATED,
2265            vec![
2266                Type::primitive_type_builder("f1", PhysicalType::INT32)
2267                    .build()
2268                    .unwrap(),
2269                Type::primitive_type_builder("f2", PhysicalType::INT64)
2270                    .build()
2271                    .unwrap(),
2272            ],
2273        );
2274        let f2 = test_new_group_type(
2275            "f",
2276            Repetition::REPEATED,
2277            vec![
2278                Type::primitive_type_builder("f3", PhysicalType::INT32)
2279                    .build()
2280                    .unwrap(),
2281            ],
2282        );
2283        assert!(!f1.check_contains(&f2));
2284    }
2285
2286    #[test]
2287    fn test_check_contains_group_primitive() {
2288        // KO: should not match
2289        let f1 = Type::group_type_builder("f").build().unwrap();
2290        let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
2291            .build()
2292            .unwrap();
2293        assert!(!f1.check_contains(&f2));
2294        assert!(!f2.check_contains(&f1));
2295
2296        // KO: should not match when primitive field is part of group type
2297        let f1 = test_new_group_type(
2298            "f",
2299            Repetition::REPEATED,
2300            vec![
2301                Type::primitive_type_builder("f1", PhysicalType::INT32)
2302                    .build()
2303                    .unwrap(),
2304            ],
2305        );
2306        let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32)
2307            .build()
2308            .unwrap();
2309        assert!(!f1.check_contains(&f2));
2310        assert!(!f2.check_contains(&f1));
2311
2312        // OK: match nested types
2313        let f1 = test_new_group_type(
2314            "a",
2315            Repetition::REPEATED,
2316            vec![
2317                test_new_group_type(
2318                    "b",
2319                    Repetition::REPEATED,
2320                    vec![
2321                        Type::primitive_type_builder("c", PhysicalType::INT32)
2322                            .build()
2323                            .unwrap(),
2324                    ],
2325                ),
2326                Type::primitive_type_builder("d", PhysicalType::INT64)
2327                    .build()
2328                    .unwrap(),
2329                Type::primitive_type_builder("e", PhysicalType::BOOLEAN)
2330                    .build()
2331                    .unwrap(),
2332            ],
2333        );
2334        let f2 = test_new_group_type(
2335            "a",
2336            Repetition::REPEATED,
2337            vec![test_new_group_type(
2338                "b",
2339                Repetition::REPEATED,
2340                vec![
2341                    Type::primitive_type_builder("c", PhysicalType::INT32)
2342                        .build()
2343                        .unwrap(),
2344                ],
2345            )],
2346        );
2347        assert!(f1.check_contains(&f2)); // should match
2348        assert!(!f2.check_contains(&f1)); // should fail
2349    }
2350
2351    #[test]
2352    fn test_schema_type_thrift_conversion_err() {
2353        let schema = Type::primitive_type_builder("col", PhysicalType::INT32)
2354            .build()
2355            .unwrap();
2356        let schema = Arc::new(schema);
2357        let thrift_schema = schema_to_buf(&schema);
2358        assert!(thrift_schema.is_err());
2359        if let Err(e) = thrift_schema {
2360            assert_eq!(
2361                format!("{e}"),
2362                "Parquet error: Root schema must be Group type"
2363            );
2364        }
2365    }
2366
2367    #[test]
2368    fn test_schema_type_thrift_conversion() {
2369        let message_type = "
2370    message conversions {
2371      REQUIRED INT64 id;
2372      OPTIONAL FIXED_LEN_BYTE_ARRAY (2) f16 (FLOAT16);
2373      OPTIONAL group int_array_Array (LIST) {
2374        REPEATED group list {
2375          OPTIONAL group element (LIST) {
2376            REPEATED group list {
2377              OPTIONAL INT32 element;
2378            }
2379          }
2380        }
2381      }
2382      OPTIONAL group int_map (MAP) {
2383        REPEATED group map (MAP_KEY_VALUE) {
2384          REQUIRED BYTE_ARRAY key (UTF8);
2385          OPTIONAL INT32 value;
2386        }
2387      }
2388      OPTIONAL group int_Map_Array (LIST) {
2389        REPEATED group list {
2390          OPTIONAL group g (MAP) {
2391            REPEATED group map (MAP_KEY_VALUE) {
2392              REQUIRED BYTE_ARRAY key (UTF8);
2393              OPTIONAL group value {
2394                OPTIONAL group H {
2395                  OPTIONAL group i (LIST) {
2396                    REPEATED group list {
2397                      OPTIONAL DOUBLE element;
2398                    }
2399                  }
2400                }
2401              }
2402            }
2403          }
2404        }
2405      }
2406      OPTIONAL group nested_struct {
2407        OPTIONAL INT32 A;
2408        OPTIONAL group b (LIST) {
2409          REPEATED group list {
2410            REQUIRED FIXED_LEN_BYTE_ARRAY (16) element;
2411          }
2412        }
2413      }
2414    }
2415    ";
2416        let expected_schema = parse_message_type(message_type).unwrap();
2417        let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap();
2418        assert_eq!(result_schema, Arc::new(expected_schema));
2419    }
2420
2421    #[test]
2422    fn test_schema_type_thrift_conversion_decimal() {
2423        let message_type = "
2424    message decimals {
2425      OPTIONAL INT32 field0;
2426      OPTIONAL INT64 field1 (DECIMAL (18, 2));
2427      OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18));
2428      OPTIONAL BYTE_ARRAY field3 (DECIMAL (9));
2429    }
2430    ";
2431        let expected_schema = parse_message_type(message_type).unwrap();
2432        let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap();
2433        assert_eq!(result_schema, Arc::new(expected_schema));
2434    }
2435
2436    // Tests schema conversion from thrift, when num_children is set to Some(0) for a
2437    // primitive type.
2438    #[test]
2439    fn test_schema_from_thrift_with_num_children_set() {
2440        // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
2441        let message_type = "
2442    message schema {
2443      OPTIONAL BYTE_ARRAY id (UTF8);
2444      OPTIONAL BYTE_ARRAY name (UTF8);
2445      OPTIONAL BYTE_ARRAY message (UTF8);
2446      OPTIONAL INT32 type (UINT_8);
2447      OPTIONAL INT64 author_time (TIMESTAMP_MILLIS);
2448      OPTIONAL INT64 __index_level_0__;
2449    }
2450    ";
2451
2452        let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2453        let mut buf = schema_to_buf(&expected_schema).unwrap();
2454        let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2455
2456        // Change all of None to Some(0)
2457        for elem in &mut thrift_schema[..] {
2458            if elem.num_children.is_none() {
2459                elem.num_children = Some(0);
2460            }
2461        }
2462
2463        let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2464        assert_eq!(result_schema, expected_schema);
2465    }
2466
2467    // Sometimes parquet-cpp sets repetition level for the root node, which is against
2468    // the format definition, but we need to handle it by setting it back to None.
2469    #[test]
2470    fn test_schema_from_thrift_root_has_repetition() {
2471        // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
2472        let message_type = "
2473    message schema {
2474      OPTIONAL BYTE_ARRAY a (UTF8);
2475      OPTIONAL INT32 b (UINT_8);
2476    }
2477    ";
2478
2479        let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2480        let mut buf = schema_to_buf(&expected_schema).unwrap();
2481        let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2482        thrift_schema[0].repetition_type = Some(Repetition::REQUIRED);
2483
2484        let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2485        assert_eq!(result_schema, expected_schema);
2486    }
2487
2488    #[test]
2489    fn test_schema_from_thrift_group_has_no_child() {
2490        let message_type = "message schema {}";
2491
2492        let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2493        let mut buf = schema_to_buf(&expected_schema).unwrap();
2494        let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2495        thrift_schema[0].repetition_type = Some(Repetition::REQUIRED);
2496
2497        let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2498        assert_eq!(result_schema, expected_schema);
2499    }
2500
2501    #[test]
2502    fn test_parquet_schema_from_array_rejects_negative_num_children() {
2503        let elements = vec![SchemaElement {
2504            r#type: None,
2505            type_length: None,
2506            repetition_type: Some(Repetition::REQUIRED),
2507            name: "schema",
2508            num_children: Some(-1),
2509            converted_type: None,
2510            scale: None,
2511            precision: None,
2512            field_id: None,
2513            logical_type: None,
2514        }];
2515        let result = parquet_schema_from_array(elements);
2516        assert!(result.unwrap_err().to_string().contains("Integer overflow"));
2517    }
2518}