1use 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
32pub type TypePtr = Arc<Type>;
37pub type SchemaDescPtr = Arc<SchemaDescriptor>;
39pub type ColumnDescPtr = Arc<ColumnDescriptor>;
41
42#[derive(Clone, Debug, PartialEq)]
49pub enum Type {
50 PrimitiveType {
52 basic_info: BasicTypeInfo,
54 physical_type: PhysicalType,
56 type_length: i32,
58 scale: i32,
60 precision: i32,
62 },
63 GroupType {
65 basic_info: BasicTypeInfo,
67 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 pub fn primitive_type_builder(
84 name: &str,
85 physical_type: PhysicalType,
86 ) -> PrimitiveTypeBuilder<'_> {
87 PrimitiveTypeBuilder::new(name, physical_type)
88 }
89
90 pub fn group_type_builder(name: &str) -> GroupTypeBuilder<'_> {
92 GroupTypeBuilder::new(name)
93 }
94
95 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 pub fn name(&self) -> &str {
105 self.get_basic_info().name()
106 }
107
108 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 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 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 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 pub fn check_contains(&self, sub_type: &Type) -> bool {
152 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 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 pub fn is_primitive(&self) -> bool {
188 matches!(*self, Type::PrimitiveType { .. })
189 }
190
191 pub fn is_group(&self) -> bool {
193 matches!(*self, Type::GroupType { .. })
194 }
195
196 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 pub fn is_optional(&self) -> bool {
207 self.get_basic_info().has_repetition()
208 && self.get_basic_info().repetition() != Repetition::REQUIRED
209 }
210
211 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 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
235pub 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 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 pub fn with_repetition(self, repetition: Repetition) -> Self {
268 Self { repetition, ..self }
269 }
270
271 pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
273 Self {
274 converted_type,
275 ..self
276 }
277 }
278
279 pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
283 Self {
284 logical_type,
285 ..self
286 }
287 }
288
289 pub fn with_length(self, length: i32) -> Self {
294 Self { length, ..self }
295 }
296
297 pub fn with_precision(self, precision: i32) -> Self {
300 Self { precision, ..self }
301 }
302
303 pub fn with_scale(self, scale: i32) -> Self {
306 Self { scale, ..self }
307 }
308
309 pub fn with_id(self, id: Option<i32>) -> Self {
311 Self { id, ..self }
312 }
313
314 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 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 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 basic_info.converted_type = self.logical_type.clone().into();
356 }
357 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 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 (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 (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 if self.precision < 1 {
533 return Err(general_err!(
534 "Invalid DECIMAL precision: {}",
535 self.precision
536 ));
537 }
538
539 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 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 _ => (), }
590
591 Ok(())
592 }
593}
594
595pub 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 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 pub fn with_repetition(mut self, repetition: Repetition) -> Self {
622 self.repetition = Some(repetition);
623 self
624 }
625
626 pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
628 Self {
629 converted_type,
630 ..self
631 }
632 }
633
634 pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
636 Self {
637 logical_type,
638 ..self
639 }
640 }
641
642 pub fn with_fields(self, fields: Vec<TypePtr>) -> Self {
645 Self { fields, ..self }
646 }
647
648 pub fn with_id(self, id: Option<i32>) -> Self {
650 Self { id, ..self }
651 }
652
653 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 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#[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 self.name.heap_size()
690 }
691}
692
693impl BasicTypeInfo {
694 pub fn name(&self) -> &str {
696 &self.name
697 }
698
699 pub fn has_repetition(&self) -> bool {
703 self.repetition.is_some()
704 }
705
706 pub fn repetition(&self) -> Repetition {
708 assert!(self.repetition.is_some());
709 self.repetition.unwrap()
710 }
711
712 pub fn converted_type(&self) -> ConvertedType {
714 self.converted_type
715 }
716
717 pub fn logical_type_ref(&self) -> Option<&LogicalType> {
719 self.logical_type.as_ref()
720 }
721
722 pub fn has_id(&self) -> bool {
724 self.id.is_some()
725 }
726
727 pub fn id(&self) -> i32 {
729 assert!(self.id.is_some());
730 self.id.unwrap()
731 }
732
733 pub fn sort_order(&self) -> SortOrder {
735 self.sort_order
736 }
737}
738
739#[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 pub fn new(parts: Vec<String>) -> Self {
774 ColumnPath { parts }
775 }
776
777 pub fn string(&self) -> String {
785 self.parts.join(".")
786 }
787
788 pub fn append(&mut self, mut tail: Vec<String>) {
800 self.parts.append(&mut tail);
801 }
802
803 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#[derive(Debug, PartialEq)]
846pub struct ColumnDescriptor {
847 primitive_type: TypePtr,
849
850 max_def_level: i16,
852
853 max_rep_level: i16,
855
856 repeated_ancestor_def_level: i16,
858
859 path: ColumnPath,
861}
862
863impl HeapSize for ColumnDescriptor {
864 fn heap_size(&self) -> usize {
865 self.path.heap_size()
868 }
869}
870
871impl ColumnDescriptor {
872 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 #[inline]
900 pub fn max_def_level(&self) -> i16 {
901 self.max_def_level
902 }
903
904 #[inline]
906 pub fn max_rep_level(&self) -> i16 {
907 self.max_rep_level
908 }
909
910 #[inline]
912 pub fn repeated_ancestor_def_level(&self) -> i16 {
913 self.repeated_ancestor_def_level
914 }
915
916 pub fn path(&self) -> &ColumnPath {
918 &self.path
919 }
920
921 pub fn self_type(&self) -> &Type {
923 self.primitive_type.as_ref()
924 }
925
926 pub fn self_type_ptr(&self) -> TypePtr {
929 self.primitive_type.clone()
930 }
931
932 pub fn get_basic_info(&self) -> &BasicTypeInfo {
934 self.primitive_type.get_basic_info()
935 }
936
937 pub fn name(&self) -> &str {
939 self.primitive_type.name()
940 }
941
942 pub fn converted_type(&self) -> ConvertedType {
944 self.primitive_type.get_basic_info().converted_type()
945 }
946
947 pub fn logical_type_ref(&self) -> Option<&LogicalType> {
949 self.primitive_type.get_basic_info().logical_type_ref()
950 }
951
952 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 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 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 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 pub fn sort_order(&self) -> SortOrder {
993 self.primitive_type.get_basic_info().sort_order()
994 }
995}
996
997#[derive(PartialEq, Clone)]
1028pub struct SchemaDescriptor {
1029 schema: TypePtr,
1034
1035 leaves: Vec<ColumnDescPtr>,
1039
1040 leaf_to_base: Vec<usize>,
1051}
1052
1053impl fmt::Debug for SchemaDescriptor {
1054 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1055 f.debug_struct("SchemaDescriptor")
1057 .field("schema", &self.schema)
1058 .finish()
1059 }
1060}
1061
1062impl 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 pub fn new(tp: TypePtr) -> Self {
1072 const INIT_SCHEMA_DEPTH: usize = 16;
1073 assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
1074 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 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 pub fn columns(&self) -> &[ColumnDescPtr] {
1113 &self.leaves
1114 }
1115
1116 pub fn num_columns(&self) -> usize {
1118 self.leaves.len()
1119 }
1120
1121 pub fn get_column_root(&self, i: usize) -> &Type {
1123 let result = self.column_root_of(i);
1124 result.as_ref()
1125 }
1126
1127 pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
1129 let result = self.column_root_of(i);
1130 result.clone()
1131 }
1132
1133 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 pub fn root_schema(&self) -> &Type {
1154 self.schema.as_ref()
1155 }
1156
1157 pub fn root_schema_ptr(&self) -> TypePtr {
1159 self.schema.clone()
1160 }
1161
1162 pub fn name(&self) -> &str {
1164 self.schema.name()
1165 }
1166}
1167
1168pub(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; 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
1189fn 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
1270fn 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
1285pub(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); 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
1314fn schema_from_array_helper<'a>(
1316 elements: &mut IntoIter<SchemaElement<'a>>,
1317 num_elements: usize,
1318 index: usize,
1319) -> Result<(usize, TypePtr)> {
1320 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 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 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 None | Some(0) => {
1354 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 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 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 #[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 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 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 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 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 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 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 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 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 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 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 assert_eq!(descr.column(0).max_def_level(), 0);
1947 assert_eq!(descr.column(0).max_rep_level(), 0);
1948 assert_eq!(descr.column(1).max_def_level(), 2);
1950 assert_eq!(descr.column(1).max_rep_level(), 0);
1951 assert_eq!(descr.column(2).max_def_level(), 2);
1953 assert_eq!(descr.column(2).max_rep_level(), 0);
1954 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 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); assert_eq!(descr.column(1).repeated_ancestor_def_level(), 0); assert_eq!(descr.column(2).repeated_ancestor_def_level(), 0); 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 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 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 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 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 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 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 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); 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 assert_eq!(descr.column(0).max_def_level(), 2);
2069 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); assert_eq!(descr.column(1).max_def_level(), 3);
2072 assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); }
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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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)); assert!(!f2.check_contains(&f1)); }
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 #[test]
2439 fn test_schema_from_thrift_with_num_children_set() {
2440 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 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 #[test]
2470 fn test_schema_from_thrift_root_has_repetition() {
2471 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}