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] {
115 match *self {
116 Type::GroupType { ref fields, .. } => &fields[..],
117 Type::PrimitiveType { .. } => panic!("Cannot call get_fields() on a non-group type"),
118 }
119 }
120
121 pub fn get_physical_type(&self) -> PhysicalType {
127 match *self {
128 Type::PrimitiveType {
129 basic_info: _,
130 physical_type,
131 ..
132 } => physical_type,
133 Type::GroupType { .. } => {
134 panic!("Cannot call get_physical_type() on a non-primitive type")
135 }
136 }
137 }
138
139 pub fn get_precision(&self) -> i32 {
145 match *self {
146 Type::PrimitiveType { precision, .. } => precision,
147 Type::GroupType { .. } => panic!("Cannot call get_precision() on non-primitive type"),
148 }
149 }
150
151 pub fn get_scale(&self) -> i32 {
157 match *self {
158 Type::PrimitiveType { scale, .. } => scale,
159 Type::GroupType { .. } => panic!("Cannot call get_scale() on non-primitive type"),
160 }
161 }
162
163 pub fn check_contains(&self, sub_type: &Type) -> bool {
166 let basic_match = self.get_basic_info().name() == sub_type.get_basic_info().name()
168 && (self.is_schema() && sub_type.is_schema()
169 || !self.is_schema()
170 && !sub_type.is_schema()
171 && self.get_basic_info().repetition()
172 == sub_type.get_basic_info().repetition());
173
174 match *self {
175 Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => {
176 self.get_physical_type() == sub_type.get_physical_type()
177 }
178 Type::GroupType { .. } if basic_match && sub_type.is_group() => {
179 let mut field_map = HashMap::new();
181 for field in self.get_fields() {
182 field_map.insert(field.name(), field);
183 }
184
185 for field in sub_type.get_fields() {
186 if !field_map
187 .get(field.name())
188 .map(|tpe| tpe.check_contains(field))
189 .unwrap_or(false)
190 {
191 return false;
192 }
193 }
194 true
195 }
196 _ => false,
197 }
198 }
199
200 pub fn is_primitive(&self) -> bool {
202 matches!(*self, Type::PrimitiveType { .. })
203 }
204
205 pub fn is_group(&self) -> bool {
207 matches!(*self, Type::GroupType { .. })
208 }
209
210 pub fn is_schema(&self) -> bool {
212 match *self {
213 Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(),
214 Type::PrimitiveType { .. } => false,
215 }
216 }
217
218 pub fn is_optional(&self) -> bool {
221 self.get_basic_info().has_repetition()
222 && self.get_basic_info().repetition() != Repetition::REQUIRED
223 }
224
225 pub(crate) fn is_list(&self) -> bool {
227 if self.is_group() {
228 let basic_info = self.get_basic_info();
229 if let Some(logical_type) = basic_info.logical_type_ref() {
230 return logical_type == &LogicalType::List;
231 }
232 return basic_info.converted_type() == ConvertedType::LIST;
233 }
234 false
235 }
236
237 pub(crate) fn has_single_repeated_child(&self) -> bool {
239 if self.is_group() {
240 let children = self.get_fields();
241 return children.len() == 1
242 && children[0].get_basic_info().has_repetition()
243 && children[0].get_basic_info().repetition() == Repetition::REPEATED;
244 }
245 false
246 }
247}
248
249pub struct PrimitiveTypeBuilder<'a> {
253 name: &'a str,
254 repetition: Repetition,
255 physical_type: PhysicalType,
256 converted_type: ConvertedType,
257 logical_type: Option<LogicalType>,
258 length: i32,
259 precision: i32,
260 scale: i32,
261 id: Option<i32>,
262}
263
264impl<'a> PrimitiveTypeBuilder<'a> {
265 pub fn new(name: &'a str, physical_type: PhysicalType) -> Self {
267 Self {
268 name,
269 repetition: Repetition::OPTIONAL,
270 physical_type,
271 converted_type: ConvertedType::NONE,
272 logical_type: None,
273 length: -1,
274 precision: -1,
275 scale: -1,
276 id: None,
277 }
278 }
279
280 pub fn with_repetition(self, repetition: Repetition) -> Self {
282 Self { repetition, ..self }
283 }
284
285 pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
287 Self {
288 converted_type,
289 ..self
290 }
291 }
292
293 pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
297 Self {
298 logical_type,
299 ..self
300 }
301 }
302
303 pub fn with_length(self, length: i32) -> Self {
308 Self { length, ..self }
309 }
310
311 pub fn with_precision(self, precision: i32) -> Self {
314 Self { precision, ..self }
315 }
316
317 pub fn with_scale(self, scale: i32) -> Self {
320 Self { scale, ..self }
321 }
322
323 pub fn with_id(self, id: Option<i32>) -> Self {
325 Self { id, ..self }
326 }
327
328 pub fn build(self) -> Result<Type> {
331 let sort_order = ColumnOrder::column_order_for_type(
332 self.logical_type.as_ref(),
333 self.converted_type,
334 self.physical_type,
335 )
336 .sort_order();
337 let mut basic_info = BasicTypeInfo {
338 name: String::from(self.name),
339 repetition: Some(self.repetition),
340 converted_type: self.converted_type,
341 logical_type: self.logical_type.clone(),
342 id: self.id,
343 sort_order,
344 };
345
346 if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 {
348 return Err(general_err!(
349 "Invalid FIXED_LEN_BYTE_ARRAY length: {} for field '{}'",
350 self.length,
351 self.name
352 ));
353 }
354
355 if let Some(logical_type) = &self.logical_type {
356 if self.converted_type != ConvertedType::NONE {
359 if ConvertedType::from(self.logical_type.clone()) != self.converted_type {
360 return Err(general_err!(
361 "Logical type {:?} is incompatible with converted type {} for field '{}'",
362 logical_type,
363 self.converted_type,
364 self.name
365 ));
366 }
367 } else {
368 basic_info.converted_type = self.logical_type.clone().into();
370 }
371 match (logical_type, self.physical_type) {
373 (LogicalType::Map | LogicalType::List | LogicalType::File, _) => {
374 return Err(general_err!(
375 "{:?} cannot be applied to a primitive type for field '{}'",
376 logical_type,
377 self.name
378 ));
379 }
380 (LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {}
381 (LogicalType::Decimal(decimal), _) => {
382 if decimal.scale != self.scale {
384 return Err(general_err!(
385 "DECIMAL logical type scale {} must match self.scale {} for field '{}'",
386 decimal.scale,
387 self.scale,
388 self.name
389 ));
390 }
391 if decimal.precision != self.precision {
392 return Err(general_err!(
393 "DECIMAL logical type precision {} must match self.precision {} for field '{}'",
394 decimal.precision,
395 self.precision,
396 self.name
397 ));
398 }
399 self.check_decimal_precision_scale()?;
400 }
401 (LogicalType::Date, PhysicalType::INT32) => {}
402 (
403 LogicalType::Time(TimeType {
404 unit: TimeUnit::MILLIS,
405 ..
406 }),
407 PhysicalType::INT32,
408 ) => {}
409 (LogicalType::Time(time), PhysicalType::INT64) => {
410 if time.unit == TimeUnit::MILLIS {
411 return Err(general_err!(
412 "Cannot use millisecond unit on INT64 type for field '{}'",
413 self.name
414 ));
415 }
416 }
417 (LogicalType::Timestamp(_), PhysicalType::INT64) => {}
418 (LogicalType::Integer(int), PhysicalType::INT32) if int.bit_width <= 32 => {}
419 (LogicalType::Integer(int), PhysicalType::INT64) if int.bit_width == 64 => {}
420 (LogicalType::Unknown, _) => {}
422 (LogicalType::String, PhysicalType::BYTE_ARRAY) => {}
423 (LogicalType::Json, PhysicalType::BYTE_ARRAY) => {}
424 (LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {}
425 (LogicalType::Geometry(_), PhysicalType::BYTE_ARRAY) => {}
426 (LogicalType::Geography(_), PhysicalType::BYTE_ARRAY) => {}
427 (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {}
428 (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
429 return Err(general_err!(
430 "UUID cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(16) field",
431 self.name
432 ));
433 }
434 (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 2 => {}
435 (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) => {
436 return Err(general_err!(
437 "FLOAT16 cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(2) field",
438 self.name
439 ));
440 }
441 (LogicalType::_Unknown { .. }, _) => {}
443 (a, b) => {
444 return Err(general_err!(
445 "Cannot annotate {:?} from {} for field '{}'",
446 a,
447 b,
448 self.name
449 ));
450 }
451 }
452 }
453
454 match self.converted_type {
455 ConvertedType::NONE => {}
456 ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => {
457 if self.physical_type != PhysicalType::BYTE_ARRAY {
458 return Err(general_err!(
459 "{} cannot annotate field '{}' because it is not a BYTE_ARRAY field",
460 self.converted_type,
461 self.name
462 ));
463 }
464 }
465 ConvertedType::DECIMAL => {
466 self.check_decimal_precision_scale()?;
467 }
468 ConvertedType::DATE
469 | ConvertedType::TIME_MILLIS
470 | ConvertedType::UINT_8
471 | ConvertedType::UINT_16
472 | ConvertedType::UINT_32
473 | ConvertedType::INT_8
474 | ConvertedType::INT_16
475 | ConvertedType::INT_32 => {
476 if self.physical_type != PhysicalType::INT32 {
477 return Err(general_err!(
478 "{} cannot annotate field '{}' because it is not a INT32 field",
479 self.converted_type,
480 self.name
481 ));
482 }
483 }
484 ConvertedType::TIME_MICROS
485 | ConvertedType::TIMESTAMP_MILLIS
486 | ConvertedType::TIMESTAMP_MICROS
487 | ConvertedType::UINT_64
488 | ConvertedType::INT_64 => {
489 if self.physical_type != PhysicalType::INT64 {
490 return Err(general_err!(
491 "{} cannot annotate field '{}' because it is not a INT64 field",
492 self.converted_type,
493 self.name
494 ));
495 }
496 }
497 ConvertedType::INTERVAL => {
498 if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY || self.length != 12 {
499 return Err(general_err!(
500 "INTERVAL cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(12) field",
501 self.name
502 ));
503 }
504 }
505 ConvertedType::ENUM => {
506 if self.physical_type != PhysicalType::BYTE_ARRAY {
507 return Err(general_err!(
508 "ENUM cannot annotate field '{}' because it is not a BYTE_ARRAY field",
509 self.name
510 ));
511 }
512 }
513 _ => {
514 return Err(general_err!(
515 "{} cannot be applied to primitive field '{}'",
516 self.converted_type,
517 self.name
518 ));
519 }
520 }
521
522 Ok(Type::PrimitiveType {
523 basic_info,
524 physical_type: self.physical_type,
525 type_length: self.length,
526 scale: self.scale,
527 precision: self.precision,
528 })
529 }
530
531 #[inline]
532 fn check_decimal_precision_scale(&self) -> Result<()> {
533 match self.physical_type {
534 PhysicalType::INT32
535 | PhysicalType::INT64
536 | PhysicalType::BYTE_ARRAY
537 | PhysicalType::FIXED_LEN_BYTE_ARRAY => (),
538 _ => {
539 return Err(general_err!(
540 "DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
541 ));
542 }
543 }
544
545 if self.precision < 1 {
547 return Err(general_err!(
548 "Invalid DECIMAL precision: {}",
549 self.precision
550 ));
551 }
552
553 if self.scale < 0 {
555 return Err(general_err!("Invalid DECIMAL scale: {}", self.scale));
556 }
557
558 if self.scale > self.precision {
559 return Err(general_err!(
560 "Invalid DECIMAL: scale ({}) cannot be greater than precision \
561 ({})",
562 self.scale,
563 self.precision
564 ));
565 }
566
567 match self.physical_type {
569 PhysicalType::INT32 => {
570 if self.precision > 9 {
571 return Err(general_err!(
572 "Cannot represent INT32 as DECIMAL with precision {}",
573 self.precision
574 ));
575 }
576 }
577 PhysicalType::INT64 => {
578 if self.precision > 18 {
579 return Err(general_err!(
580 "Cannot represent INT64 as DECIMAL with precision {}",
581 self.precision
582 ));
583 }
584 }
585 PhysicalType::FIXED_LEN_BYTE_ARRAY => {
586 let length = self
587 .length
588 .checked_mul(8)
589 .ok_or(general_err!("Invalid length {} for Decimal", self.length))?;
590 let max_precision = (2f64.powi(length - 1) - 1f64).log10().floor() as i32;
591
592 if self.precision > max_precision {
593 return Err(general_err!(
594 "Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \
595 precision {}. The max precision can only be {}",
596 self.length,
597 self.precision,
598 max_precision
599 ));
600 }
601 }
602 _ => (), }
604
605 Ok(())
606 }
607}
608
609pub struct GroupTypeBuilder<'a> {
613 name: &'a str,
614 repetition: Option<Repetition>,
615 converted_type: ConvertedType,
616 logical_type: Option<LogicalType>,
617 fields: Vec<TypePtr>,
618 id: Option<i32>,
619}
620
621impl<'a> GroupTypeBuilder<'a> {
622 pub fn new(name: &'a str) -> Self {
624 Self {
625 name,
626 repetition: None,
627 converted_type: ConvertedType::NONE,
628 logical_type: None,
629 fields: Vec::new(),
630 id: None,
631 }
632 }
633
634 pub fn with_repetition(mut self, repetition: Repetition) -> Self {
636 self.repetition = Some(repetition);
637 self
638 }
639
640 pub fn with_converted_type(self, converted_type: ConvertedType) -> Self {
642 Self {
643 converted_type,
644 ..self
645 }
646 }
647
648 pub fn with_logical_type(self, logical_type: Option<LogicalType>) -> Self {
650 Self {
651 logical_type,
652 ..self
653 }
654 }
655
656 pub fn with_fields(self, fields: Vec<TypePtr>) -> Self {
659 Self { fields, ..self }
660 }
661
662 pub fn with_id(self, id: Option<i32>) -> Self {
664 Self { id, ..self }
665 }
666
667 pub fn build(self) -> Result<Type> {
669 if matches!(&self.logical_type, Some(LogicalType::File)) {
670 validate_file_type_fields(self.name, &self.fields)?;
671 }
672 let mut basic_info = BasicTypeInfo {
673 name: String::from(self.name),
674 repetition: self.repetition,
675 converted_type: self.converted_type,
676 logical_type: self.logical_type.clone(),
677 id: self.id,
678 sort_order: SortOrder::UNDEFINED,
679 };
680 if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE {
682 basic_info.converted_type = self.logical_type.into();
683 }
684 Ok(Type::GroupType {
685 basic_info,
686 fields: self.fields,
687 })
688 }
689}
690
691fn validate_file_type_fields(name: &str, fields: &[TypePtr]) -> Result<()> {
711 const VALID_FIELDS: &[(&str, PhysicalType, Option<LogicalType>)] = &[
713 ("uri", PhysicalType::BYTE_ARRAY, Some(LogicalType::String)),
714 ("offset", PhysicalType::INT64, None),
715 ("size", PhysicalType::INT64, None),
716 (
717 "content_type",
718 PhysicalType::BYTE_ARRAY,
719 Some(LogicalType::String),
720 ),
721 (
722 "checksum",
723 PhysicalType::BYTE_ARRAY,
724 Some(LogicalType::String),
725 ),
726 ("inline", PhysicalType::BYTE_ARRAY, None),
727 ];
728
729 for field in fields {
730 let field_name = field.get_basic_info().name();
731 let Some((_, expected_physical, expected_logical)) =
732 VALID_FIELDS.iter().find(|(n, _, _)| *n == field_name)
733 else {
734 return Err(general_err!(
735 "FILE type group '{}' contains unrecognized field '{}'. \
736 Valid fields are: uri, offset, size, content_type, checksum, inline",
737 name,
738 field_name
739 ));
740 };
741
742 let is_optional = field.get_basic_info().has_repetition()
744 && field.get_basic_info().repetition() == Repetition::OPTIONAL;
745 if !is_optional {
746 return Err(general_err!(
747 "FILE type field '{}' must be OPTIONAL in group '{}'",
748 field_name,
749 name
750 ));
751 }
752
753 if field.is_group() {
755 return Err(general_err!(
756 "FILE type field '{}' in group '{}' must be a primitive type",
757 field_name,
758 name
759 ));
760 }
761 if field.get_physical_type() != *expected_physical {
762 return Err(general_err!(
763 "FILE type field '{}' in group '{}' must have physical type {:?}",
764 field_name,
765 name,
766 expected_physical
767 ));
768 }
769 if field.get_basic_info().logical_type_ref() != expected_logical.as_ref() {
770 return Err(general_err!(
771 "FILE type field '{}' in group '{}' must have logical type {:?}",
772 field_name,
773 name,
774 expected_logical
775 ));
776 }
777 }
778 Ok(())
779}
780
781#[derive(Clone, Debug, PartialEq, Eq)]
784pub struct BasicTypeInfo {
785 name: String,
786 repetition: Option<Repetition>,
787 converted_type: ConvertedType,
788 logical_type: Option<LogicalType>,
789 id: Option<i32>,
790 sort_order: SortOrder,
791}
792
793impl HeapSize for BasicTypeInfo {
794 fn heap_size(&self) -> usize {
795 self.name.heap_size()
797 }
798}
799
800impl BasicTypeInfo {
801 pub fn name(&self) -> &str {
803 &self.name
804 }
805
806 pub fn has_repetition(&self) -> bool {
810 self.repetition.is_some()
811 }
812
813 pub fn repetition(&self) -> Repetition {
819 assert!(self.repetition.is_some());
820 self.repetition.unwrap()
821 }
822
823 pub fn converted_type(&self) -> ConvertedType {
825 self.converted_type
826 }
827
828 pub fn logical_type_ref(&self) -> Option<&LogicalType> {
830 self.logical_type.as_ref()
831 }
832
833 pub fn has_id(&self) -> bool {
835 self.id.is_some()
836 }
837
838 pub fn id(&self) -> i32 {
844 assert!(self.id.is_some());
845 self.id.unwrap()
846 }
847
848 pub fn sort_order(&self) -> SortOrder {
850 self.sort_order
851 }
852}
853
854#[derive(Clone, PartialEq, Debug, Eq, Hash)]
876pub struct ColumnPath {
877 parts: Vec<String>,
878}
879
880impl HeapSize for ColumnPath {
881 fn heap_size(&self) -> usize {
882 self.parts.heap_size()
883 }
884}
885
886impl ColumnPath {
887 pub fn new(parts: Vec<String>) -> Self {
889 ColumnPath { parts }
890 }
891
892 pub fn string(&self) -> String {
900 self.parts.join(".")
901 }
902
903 pub fn append(&mut self, mut tail: Vec<String>) {
915 self.parts.append(&mut tail);
916 }
917
918 pub fn parts(&self) -> &[String] {
920 &self.parts
921 }
922}
923
924impl fmt::Display for ColumnPath {
925 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
926 write!(f, "{:?}", self.string())
927 }
928}
929
930impl From<Vec<String>> for ColumnPath {
931 fn from(parts: Vec<String>) -> Self {
932 ColumnPath { parts }
933 }
934}
935
936impl From<&str> for ColumnPath {
937 fn from(single_path: &str) -> Self {
938 let s = String::from(single_path);
939 ColumnPath::from(s)
940 }
941}
942
943impl From<String> for ColumnPath {
944 fn from(single_path: String) -> Self {
945 let v = vec![single_path];
946 ColumnPath { parts: v }
947 }
948}
949
950impl AsRef<[String]> for ColumnPath {
951 fn as_ref(&self) -> &[String] {
952 &self.parts
953 }
954}
955
956#[derive(Debug, PartialEq)]
961pub struct ColumnDescriptor {
962 primitive_type: TypePtr,
964
965 max_def_level: i16,
967
968 max_rep_level: i16,
970
971 repeated_ancestor_def_level: i16,
973
974 path: ColumnPath,
976}
977
978impl HeapSize for ColumnDescriptor {
979 fn heap_size(&self) -> usize {
980 self.path.heap_size()
983 }
984}
985
986impl ColumnDescriptor {
987 pub fn new(
989 primitive_type: TypePtr,
990 max_def_level: i16,
991 max_rep_level: i16,
992 path: ColumnPath,
993 ) -> Self {
994 Self::new_with_repeated_ancestor(primitive_type, max_def_level, max_rep_level, path, 0)
995 }
996
997 pub(crate) fn new_with_repeated_ancestor(
998 primitive_type: TypePtr,
999 max_def_level: i16,
1000 max_rep_level: i16,
1001 path: ColumnPath,
1002 repeated_ancestor_def_level: i16,
1003 ) -> Self {
1004 Self {
1005 primitive_type,
1006 max_def_level,
1007 max_rep_level,
1008 repeated_ancestor_def_level,
1009 path,
1010 }
1011 }
1012
1013 #[inline]
1015 pub fn max_def_level(&self) -> i16 {
1016 self.max_def_level
1017 }
1018
1019 #[inline]
1021 pub fn max_rep_level(&self) -> i16 {
1022 self.max_rep_level
1023 }
1024
1025 #[inline]
1027 pub fn repeated_ancestor_def_level(&self) -> i16 {
1028 self.repeated_ancestor_def_level
1029 }
1030
1031 pub fn path(&self) -> &ColumnPath {
1033 &self.path
1034 }
1035
1036 pub fn self_type(&self) -> &Type {
1038 self.primitive_type.as_ref()
1039 }
1040
1041 pub fn self_type_ptr(&self) -> TypePtr {
1044 self.primitive_type.clone()
1045 }
1046
1047 pub fn get_basic_info(&self) -> &BasicTypeInfo {
1049 self.primitive_type.get_basic_info()
1050 }
1051
1052 pub fn name(&self) -> &str {
1054 self.primitive_type.name()
1055 }
1056
1057 pub fn converted_type(&self) -> ConvertedType {
1059 self.primitive_type.get_basic_info().converted_type()
1060 }
1061
1062 pub fn logical_type_ref(&self) -> Option<&LogicalType> {
1064 self.primitive_type.get_basic_info().logical_type_ref()
1065 }
1066
1067 pub fn physical_type(&self) -> PhysicalType {
1073 match self.primitive_type.as_ref() {
1074 Type::PrimitiveType { physical_type, .. } => *physical_type,
1075 Type::GroupType { .. } => panic!("Expected primitive type!"),
1076 }
1077 }
1078
1079 pub fn type_length(&self) -> i32 {
1085 match self.primitive_type.as_ref() {
1086 Type::PrimitiveType { type_length, .. } => *type_length,
1087 Type::GroupType { .. } => panic!("Expected primitive type!"),
1088 }
1089 }
1090
1091 pub fn type_precision(&self) -> i32 {
1097 match self.primitive_type.as_ref() {
1098 Type::PrimitiveType { precision, .. } => *precision,
1099 Type::GroupType { .. } => panic!("Expected primitive type!"),
1100 }
1101 }
1102
1103 pub fn type_scale(&self) -> i32 {
1109 match self.primitive_type.as_ref() {
1110 Type::PrimitiveType { scale, .. } => *scale,
1111 Type::GroupType { .. } => panic!("Expected primitive type!"),
1112 }
1113 }
1114
1115 pub fn sort_order(&self) -> SortOrder {
1120 self.primitive_type.get_basic_info().sort_order()
1121 }
1122}
1123
1124#[derive(PartialEq, Clone)]
1155pub struct SchemaDescriptor {
1156 schema: TypePtr,
1161
1162 leaves: Vec<ColumnDescPtr>,
1166
1167 leaf_to_base: Vec<usize>,
1178}
1179
1180impl fmt::Debug for SchemaDescriptor {
1181 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1182 f.debug_struct("SchemaDescriptor")
1184 .field("schema", &self.schema)
1185 .finish()
1186 }
1187}
1188
1189impl HeapSize for SchemaDescriptor {
1191 fn heap_size(&self) -> usize {
1192 self.schema.heap_size() + self.leaves.heap_size() + self.leaf_to_base.heap_size()
1193 }
1194}
1195
1196impl SchemaDescriptor {
1197 pub fn new(tp: TypePtr) -> Self {
1199 const INIT_SCHEMA_DEPTH: usize = 16;
1200 assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
1201 let n_leaves = num_leaves(&tp).unwrap();
1203 let mut leaves = Vec::with_capacity(n_leaves);
1204 let mut leaf_to_base = Vec::with_capacity(n_leaves);
1205 let mut path = Vec::with_capacity(INIT_SCHEMA_DEPTH);
1206 for (root_idx, f) in tp.get_fields().iter().enumerate() {
1207 path.clear();
1208 build_tree(
1209 f,
1210 root_idx,
1211 0,
1212 0,
1213 0,
1214 &mut leaves,
1215 &mut leaf_to_base,
1216 &mut path,
1217 );
1218 }
1219
1220 Self {
1221 schema: tp,
1222 leaves,
1223 leaf_to_base,
1224 }
1225 }
1226
1227 pub fn column(&self, i: usize) -> ColumnDescPtr {
1233 assert!(
1234 i < self.leaves.len(),
1235 "Index out of bound: {} not in [0, {})",
1236 i,
1237 self.leaves.len()
1238 );
1239 self.leaves[i].clone()
1240 }
1241
1242 pub fn columns(&self) -> &[ColumnDescPtr] {
1244 &self.leaves
1245 }
1246
1247 pub fn num_columns(&self) -> usize {
1249 self.leaves.len()
1250 }
1251
1252 pub fn get_column_root(&self, i: usize) -> &Type {
1254 let result = self.column_root_of(i);
1255 result.as_ref()
1256 }
1257
1258 pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
1264 let result = self.column_root_of(i);
1265 result.clone()
1266 }
1267
1268 pub fn get_column_root_idx(&self, leaf: usize) -> usize {
1274 assert!(
1275 leaf < self.leaves.len(),
1276 "Index out of bound: {} not in [0, {})",
1277 leaf,
1278 self.leaves.len()
1279 );
1280
1281 *self
1282 .leaf_to_base
1283 .get(leaf)
1284 .unwrap_or_else(|| panic!("Expected a value for index {leaf} but found None"))
1285 }
1286
1287 fn column_root_of(&self, i: usize) -> &TypePtr {
1288 &self.schema.get_fields()[self.get_column_root_idx(i)]
1289 }
1290
1291 pub fn root_schema(&self) -> &Type {
1293 self.schema.as_ref()
1294 }
1295
1296 pub fn root_schema_ptr(&self) -> TypePtr {
1298 self.schema.clone()
1299 }
1300
1301 pub fn name(&self) -> &str {
1303 self.schema.name()
1304 }
1305}
1306
1307pub(crate) fn num_nodes(tp: &TypePtr) -> Result<usize> {
1309 if !tp.is_group() {
1310 return Err(general_err!("Root schema must be Group type"));
1311 }
1312 let mut n_nodes = 1usize; for f in tp.get_fields() {
1314 count_nodes(f, &mut n_nodes);
1315 }
1316 Ok(n_nodes)
1317}
1318
1319pub(crate) fn count_nodes(tp: &TypePtr, n_nodes: &mut usize) {
1320 *n_nodes += 1;
1321 if let Type::GroupType { fields, .. } = tp.as_ref() {
1322 for f in fields {
1323 count_nodes(f, n_nodes);
1324 }
1325 }
1326}
1327
1328fn num_leaves(tp: &TypePtr) -> Result<usize> {
1330 if !tp.is_group() {
1331 return Err(general_err!("Root schema must be Group type"));
1332 }
1333 let mut n_leaves = 0usize;
1334 for f in tp.get_fields() {
1335 count_leaves(f, &mut n_leaves);
1336 }
1337 Ok(n_leaves)
1338}
1339
1340fn count_leaves(tp: &TypePtr, n_leaves: &mut usize) {
1341 match tp.as_ref() {
1342 Type::PrimitiveType { .. } => *n_leaves += 1,
1343 Type::GroupType { fields, .. } => {
1344 for f in fields {
1345 count_leaves(f, n_leaves);
1346 }
1347 }
1348 }
1349}
1350
1351#[expect(clippy::too_many_arguments)]
1352fn build_tree<'a>(
1353 tp: &'a TypePtr,
1354 root_idx: usize,
1355 mut max_rep_level: i16,
1356 mut max_def_level: i16,
1357 mut repeated_ancestor_def_level: i16,
1358 leaves: &mut Vec<ColumnDescPtr>,
1359 leaf_to_base: &mut Vec<usize>,
1360 path_so_far: &mut Vec<&'a str>,
1361) {
1362 assert!(tp.get_basic_info().has_repetition());
1363
1364 path_so_far.push(tp.name());
1365 match tp.get_basic_info().repetition() {
1366 Repetition::OPTIONAL => {
1367 max_def_level += 1;
1368 }
1369 Repetition::REPEATED => {
1370 max_def_level += 1;
1371 max_rep_level += 1;
1372 repeated_ancestor_def_level = max_def_level;
1373 }
1374 Repetition::REQUIRED => {}
1375 }
1376
1377 match tp.as_ref() {
1378 Type::PrimitiveType { .. } => {
1379 let mut path: Vec<String> = vec![];
1380 path.extend(path_so_far.iter().copied().map(String::from));
1381 let desc = ColumnDescriptor::new_with_repeated_ancestor(
1382 tp.clone(),
1383 max_def_level,
1384 max_rep_level,
1385 ColumnPath::new(path),
1386 repeated_ancestor_def_level,
1387 );
1388 leaves.push(Arc::new(desc));
1389 leaf_to_base.push(root_idx);
1390 }
1391 Type::GroupType { fields, .. } => {
1392 for f in fields {
1393 build_tree(
1394 f,
1395 root_idx,
1396 max_rep_level,
1397 max_def_level,
1398 repeated_ancestor_def_level,
1399 leaves,
1400 leaf_to_base,
1401 path_so_far,
1402 );
1403 path_so_far.pop();
1404 }
1405 }
1406 }
1407}
1408
1409fn check_logical_type(logical_type: Option<&LogicalType>) -> Result<()> {
1411 if let Some(LogicalType::Integer(IntType { bit_width, .. })) = logical_type
1412 && *bit_width != 8
1413 && *bit_width != 16
1414 && *bit_width != 32
1415 && *bit_width != 64
1416 {
1417 return Err(general_err!(
1418 "Bit width must be 8, 16, 32, or 64 for Integer logical type"
1419 ));
1420 }
1421 Ok(())
1422}
1423
1424pub(crate) fn parquet_schema_from_array(elements: Vec<SchemaElement<'_>>) -> Result<TypePtr> {
1427 let mut index = 0;
1428 let num_elements = elements.len();
1429 let mut schema_nodes = Vec::with_capacity(1); let mut elements = elements.into_iter();
1433
1434 while index < num_elements {
1435 let t = schema_from_array_helper(&mut elements, num_elements, index)?;
1436 index = t.0;
1437 schema_nodes.push(t.1);
1438 }
1439 if schema_nodes.len() != 1 {
1440 return Err(general_err!(
1441 "Expected exactly one root node, but found {}",
1442 schema_nodes.len()
1443 ));
1444 }
1445
1446 if !schema_nodes[0].is_group() {
1447 return Err(general_err!("Expected root node to be a group type"));
1448 }
1449
1450 Ok(schema_nodes.remove(0))
1451}
1452
1453fn schema_from_array_helper(
1455 elements: &mut IntoIter<SchemaElement<'_>>,
1456 num_elements: usize,
1457 index: usize,
1458) -> Result<(usize, TypePtr)> {
1459 let is_root_node = index == 0;
1462
1463 if index >= num_elements {
1464 return Err(general_err!(
1465 "Index out of bound, index = {}, len = {}",
1466 index,
1467 num_elements
1468 ));
1469 }
1470 let element = elements.next().expect("schema vector should not be empty");
1471
1472 if let (true, None | Some(0)) = (is_root_node, element.num_children) {
1474 let builder = Type::group_type_builder(element.name);
1475 return Ok((index + 1, Arc::new(builder.build().unwrap())));
1476 }
1477
1478 let converted_type = element.converted_type.unwrap_or(ConvertedType::NONE);
1479
1480 let logical_type = element.logical_type;
1482
1483 check_logical_type(logical_type.as_ref())?;
1484
1485 let field_id = element.field_id;
1486 match element.num_children {
1487 None | Some(0) => {
1493 if element.repetition_type.is_none() {
1495 return Err(general_err!(
1496 "Repetition level must be defined for a primitive type"
1497 ));
1498 }
1499 let repetition = element.repetition_type.unwrap();
1500 if let Some(physical_type) = element.r#type {
1501 let length = element.type_length.unwrap_or(-1);
1502 let scale = element.scale.unwrap_or(-1);
1503 let precision = element.precision.unwrap_or(-1);
1504 let name = element.name;
1505 let builder = Type::primitive_type_builder(name, physical_type)
1506 .with_repetition(repetition)
1507 .with_converted_type(converted_type)
1508 .with_logical_type(logical_type)
1509 .with_length(length)
1510 .with_precision(precision)
1511 .with_scale(scale)
1512 .with_id(field_id);
1513 Ok((index + 1, Arc::new(builder.build()?)))
1514 } else {
1515 let mut builder = Type::group_type_builder(element.name)
1516 .with_converted_type(converted_type)
1517 .with_logical_type(logical_type)
1518 .with_id(field_id);
1519 if !is_root_node {
1520 builder = builder.with_repetition(repetition);
1528 }
1529 Ok((index + 1, Arc::new(builder.build().unwrap())))
1530 }
1531 }
1532 Some(n) => {
1533 let repetition = element.repetition_type;
1534
1535 let mut fields = Vec::with_capacity(usize::try_from(n)?);
1536 let mut next_index = index + 1;
1537 for _ in 0..n {
1538 let child_result = schema_from_array_helper(elements, num_elements, next_index)?;
1539 next_index = child_result.0;
1540 fields.push(child_result.1);
1541 }
1542
1543 let mut builder = Type::group_type_builder(element.name)
1544 .with_converted_type(converted_type)
1545 .with_logical_type(logical_type)
1546 .with_fields(fields)
1547 .with_id(field_id);
1548
1549 if !is_root_node {
1557 let Some(rep) = repetition else {
1558 return Err(general_err!(
1559 "Repetition level must be defined for non-root types"
1560 ));
1561 };
1562 builder = builder.with_repetition(rep);
1563 }
1564 Ok((next_index, Arc::new(builder.build()?)))
1565 }
1566 }
1567}
1568
1569#[cfg(test)]
1570mod tests {
1571 use super::*;
1572
1573 use crate::{
1574 file::metadata::thrift::tests::{buf_to_schema_list, roundtrip_schema, schema_to_buf},
1575 schema::parser::parse_message_type,
1576 };
1577
1578 #[test]
1581 fn test_primitive_type() {
1582 let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1583 .with_logical_type(Some(LogicalType::integer(32, true)))
1584 .with_id(Some(0))
1585 .build();
1586 assert!(result.is_ok());
1587
1588 if let Ok(tp) = result {
1589 assert!(tp.is_primitive());
1590 assert!(!tp.is_group());
1591 let basic_info = tp.get_basic_info();
1592 assert_eq!(basic_info.repetition(), Repetition::OPTIONAL);
1593 assert_eq!(
1594 basic_info.logical_type_ref(),
1595 Some(&LogicalType::integer(32, true))
1596 );
1597 assert_eq!(basic_info.converted_type(), ConvertedType::INT_32);
1598 assert_eq!(basic_info.id(), 0);
1599 match tp {
1600 Type::PrimitiveType { physical_type, .. } => {
1601 assert_eq!(physical_type, PhysicalType::INT32);
1602 }
1603 Type::GroupType { .. } => panic!(),
1604 }
1605 }
1606
1607 result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1609 .with_repetition(Repetition::REPEATED)
1610 .with_logical_type(Some(LogicalType::integer(8, true)))
1611 .build();
1612 assert!(result.is_err());
1613 if let Err(e) = result {
1614 assert_eq!(
1615 format!("{e}"),
1616 "Parquet error: Cannot annotate Integer(IntType { bit_width: 8, is_signed: true }) from INT64 for field 'foo'"
1617 );
1618 }
1619
1620 result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1622 .with_repetition(Repetition::REPEATED)
1623 .with_converted_type(ConvertedType::BSON)
1624 .build();
1625 assert!(result.is_err());
1626 if let Err(e) = result {
1627 assert_eq!(
1628 format!("{e}"),
1629 "Parquet error: BSON cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
1630 );
1631 }
1632
1633 result = Type::primitive_type_builder("foo", PhysicalType::INT96)
1634 .with_repetition(Repetition::REQUIRED)
1635 .with_converted_type(ConvertedType::DECIMAL)
1636 .with_precision(-1)
1637 .with_scale(-1)
1638 .build();
1639 assert!(result.is_err());
1640 if let Err(e) = result {
1641 assert_eq!(
1642 format!("{e}"),
1643 "Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
1644 );
1645 }
1646
1647 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1648 .with_repetition(Repetition::REQUIRED)
1649 .with_logical_type(Some(LogicalType::decimal(32, 12)))
1650 .with_precision(-1)
1651 .with_scale(-1)
1652 .build();
1653 assert!(result.is_err());
1654 if let Err(e) = result {
1655 assert_eq!(
1656 format!("{e}"),
1657 "Parquet error: DECIMAL logical type scale 32 must match self.scale -1 for field 'foo'"
1658 );
1659 }
1660
1661 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1662 .with_repetition(Repetition::REQUIRED)
1663 .with_converted_type(ConvertedType::DECIMAL)
1664 .with_precision(-1)
1665 .with_scale(-1)
1666 .build();
1667 assert!(result.is_err());
1668 if let Err(e) = result {
1669 assert_eq!(
1670 format!("{e}"),
1671 "Parquet error: Invalid DECIMAL precision: -1"
1672 );
1673 }
1674
1675 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1676 .with_repetition(Repetition::REQUIRED)
1677 .with_converted_type(ConvertedType::DECIMAL)
1678 .with_precision(0)
1679 .with_scale(-1)
1680 .build();
1681 assert!(result.is_err());
1682 if let Err(e) = result {
1683 assert_eq!(
1684 format!("{e}"),
1685 "Parquet error: Invalid DECIMAL precision: 0"
1686 );
1687 }
1688
1689 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1690 .with_repetition(Repetition::REQUIRED)
1691 .with_converted_type(ConvertedType::DECIMAL)
1692 .with_precision(1)
1693 .with_scale(-1)
1694 .build();
1695 assert!(result.is_err());
1696 if let Err(e) = result {
1697 assert_eq!(format!("{e}"), "Parquet error: Invalid DECIMAL scale: -1");
1698 }
1699
1700 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1701 .with_repetition(Repetition::REQUIRED)
1702 .with_converted_type(ConvertedType::DECIMAL)
1703 .with_precision(1)
1704 .with_scale(2)
1705 .build();
1706 assert!(result.is_err());
1707 if let Err(e) = result {
1708 assert_eq!(
1709 format!("{e}"),
1710 "Parquet error: Invalid DECIMAL: scale (2) cannot be greater than precision (1)"
1711 );
1712 }
1713
1714 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1716 .with_repetition(Repetition::REQUIRED)
1717 .with_converted_type(ConvertedType::DECIMAL)
1718 .with_precision(1)
1719 .with_scale(1)
1720 .build();
1721 assert!(result.is_ok());
1722
1723 result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1724 .with_repetition(Repetition::REQUIRED)
1725 .with_converted_type(ConvertedType::DECIMAL)
1726 .with_precision(18)
1727 .with_scale(2)
1728 .build();
1729 assert!(result.is_err());
1730 if let Err(e) = result {
1731 assert_eq!(
1732 format!("{e}"),
1733 "Parquet error: Cannot represent INT32 as DECIMAL with precision 18"
1734 );
1735 }
1736
1737 result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1738 .with_repetition(Repetition::REQUIRED)
1739 .with_converted_type(ConvertedType::DECIMAL)
1740 .with_precision(32)
1741 .with_scale(2)
1742 .build();
1743 assert!(result.is_err());
1744 if let Err(e) = result {
1745 assert_eq!(
1746 format!("{e}"),
1747 "Parquet error: Cannot represent INT64 as DECIMAL with precision 32"
1748 );
1749 }
1750
1751 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1752 .with_repetition(Repetition::REQUIRED)
1753 .with_converted_type(ConvertedType::DECIMAL)
1754 .with_length(5)
1755 .with_precision(12)
1756 .with_scale(2)
1757 .build();
1758 assert!(result.is_err());
1759 if let Err(e) = result {
1760 assert_eq!(
1761 format!("{e}"),
1762 "Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11"
1763 );
1764 }
1765
1766 result = Type::primitive_type_builder("foo", PhysicalType::INT64)
1767 .with_repetition(Repetition::REQUIRED)
1768 .with_converted_type(ConvertedType::UINT_8)
1769 .build();
1770 assert!(result.is_err());
1771 if let Err(e) = result {
1772 assert_eq!(
1773 format!("{e}"),
1774 "Parquet error: UINT_8 cannot annotate field 'foo' because it is not a INT32 field"
1775 );
1776 }
1777
1778 result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1779 .with_repetition(Repetition::REQUIRED)
1780 .with_converted_type(ConvertedType::TIME_MICROS)
1781 .build();
1782 assert!(result.is_err());
1783 if let Err(e) = result {
1784 assert_eq!(
1785 format!("{e}"),
1786 "Parquet error: TIME_MICROS cannot annotate field 'foo' because it is not a INT64 field"
1787 );
1788 }
1789
1790 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1791 .with_repetition(Repetition::REQUIRED)
1792 .with_converted_type(ConvertedType::INTERVAL)
1793 .build();
1794 assert!(result.is_err());
1795 if let Err(e) = result {
1796 assert_eq!(
1797 format!("{e}"),
1798 "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
1799 );
1800 }
1801
1802 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1803 .with_repetition(Repetition::REQUIRED)
1804 .with_converted_type(ConvertedType::INTERVAL)
1805 .with_length(1)
1806 .build();
1807 assert!(result.is_err());
1808 if let Err(e) = result {
1809 assert_eq!(
1810 format!("{e}"),
1811 "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field"
1812 );
1813 }
1814
1815 result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1816 .with_repetition(Repetition::REQUIRED)
1817 .with_converted_type(ConvertedType::ENUM)
1818 .build();
1819 assert!(result.is_err());
1820 if let Err(e) = result {
1821 assert_eq!(
1822 format!("{e}"),
1823 "Parquet error: ENUM cannot annotate field 'foo' because it is not a BYTE_ARRAY field"
1824 );
1825 }
1826
1827 result = Type::primitive_type_builder("foo", PhysicalType::INT32)
1828 .with_repetition(Repetition::REQUIRED)
1829 .with_converted_type(ConvertedType::MAP)
1830 .build();
1831 assert!(result.is_err());
1832 if let Err(e) = result {
1833 assert_eq!(
1834 format!("{e}"),
1835 "Parquet error: MAP cannot be applied to primitive field 'foo'"
1836 );
1837 }
1838
1839 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1840 .with_repetition(Repetition::REQUIRED)
1841 .with_converted_type(ConvertedType::DECIMAL)
1842 .with_length(-1)
1843 .build();
1844 assert!(result.is_err());
1845 if let Err(e) = result {
1846 assert_eq!(
1847 format!("{e}"),
1848 "Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1 for field 'foo'"
1849 );
1850 }
1851
1852 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1853 .with_repetition(Repetition::REQUIRED)
1854 .with_logical_type(Some(LogicalType::Float16))
1855 .with_length(2)
1856 .build();
1857 assert!(result.is_ok());
1858
1859 result = Type::primitive_type_builder("foo", PhysicalType::FLOAT)
1861 .with_repetition(Repetition::REQUIRED)
1862 .with_logical_type(Some(LogicalType::Float16))
1863 .with_length(2)
1864 .build();
1865 assert!(result.is_err());
1866 if let Err(e) = result {
1867 assert_eq!(
1868 format!("{e}"),
1869 "Parquet error: Cannot annotate Float16 from FLOAT for field 'foo'"
1870 );
1871 }
1872
1873 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1875 .with_repetition(Repetition::REQUIRED)
1876 .with_logical_type(Some(LogicalType::Float16))
1877 .with_length(4)
1878 .build();
1879 assert!(result.is_err());
1880 if let Err(e) = result {
1881 assert_eq!(
1882 format!("{e}"),
1883 "Parquet error: FLOAT16 cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(2) field"
1884 );
1885 }
1886
1887 result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
1889 .with_repetition(Repetition::REQUIRED)
1890 .with_logical_type(Some(LogicalType::Uuid))
1891 .with_length(15)
1892 .build();
1893 assert!(result.is_err());
1894 if let Err(e) = result {
1895 assert_eq!(
1896 format!("{e}"),
1897 "Parquet error: UUID cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(16) field"
1898 );
1899 }
1900
1901 result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
1903 .with_logical_type(Some(LogicalType::_Unknown { field_id: 100 }))
1904 .build();
1905 assert!(result.is_ok());
1906 }
1907
1908 #[test]
1909 fn test_group_type() {
1910 let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
1911 .with_converted_type(ConvertedType::INT_32)
1912 .with_id(Some(0))
1913 .build();
1914 assert!(f1.is_ok());
1915 let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
1916 .with_converted_type(ConvertedType::UTF8)
1917 .with_id(Some(1))
1918 .build();
1919 assert!(f2.is_ok());
1920
1921 let fields = vec![Arc::new(f1.unwrap()), Arc::new(f2.unwrap())];
1922
1923 let result = Type::group_type_builder("foo")
1924 .with_repetition(Repetition::REPEATED)
1925 .with_logical_type(Some(LogicalType::List))
1926 .with_fields(fields)
1927 .with_id(Some(1))
1928 .build();
1929 assert!(result.is_ok());
1930
1931 let tp = result.unwrap();
1932 let basic_info = tp.get_basic_info();
1933 assert!(tp.is_group());
1934 assert!(!tp.is_primitive());
1935 assert_eq!(basic_info.repetition(), Repetition::REPEATED);
1936 assert_eq!(basic_info.logical_type_ref(), Some(&LogicalType::List));
1937 assert_eq!(basic_info.converted_type(), ConvertedType::LIST);
1938 assert_eq!(basic_info.id(), 1);
1939 assert_eq!(tp.get_fields().len(), 2);
1940 assert_eq!(tp.get_fields()[0].name(), "f1");
1941 assert_eq!(tp.get_fields()[1].name(), "f2");
1942 }
1943
1944 #[test]
1945 fn test_column_descriptor() {
1946 let result = test_column_descriptor_helper();
1947 assert!(
1948 result.is_ok(),
1949 "Expected result to be OK but got err:\n {}",
1950 result.unwrap_err()
1951 );
1952 }
1953
1954 fn test_column_descriptor_helper() -> Result<()> {
1955 let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY)
1956 .with_converted_type(ConvertedType::UTF8)
1957 .build()?;
1958
1959 let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name"));
1960
1961 assert_eq!(descr.path(), &ColumnPath::from("name"));
1962 assert_eq!(descr.converted_type(), ConvertedType::UTF8);
1963 assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY);
1964 assert_eq!(descr.max_def_level(), 4);
1965 assert_eq!(descr.max_rep_level(), 1);
1966 assert_eq!(descr.name(), "name");
1967 assert_eq!(descr.type_length(), -1);
1968 assert_eq!(descr.type_precision(), -1);
1969 assert_eq!(descr.type_scale(), -1);
1970
1971 Ok(())
1972 }
1973
1974 #[test]
1975 fn test_schema_descriptor() {
1976 let result = test_schema_descriptor_helper();
1977 assert!(
1978 result.is_ok(),
1979 "Expected result to be OK but got err:\n {}",
1980 result.unwrap_err()
1981 );
1982 }
1983
1984 fn test_schema_descriptor_helper() -> Result<()> {
1986 let mut fields = vec![];
1987
1988 let inta = Type::primitive_type_builder("a", PhysicalType::INT32)
1989 .with_repetition(Repetition::REQUIRED)
1990 .with_converted_type(ConvertedType::INT_32)
1991 .build()?;
1992 fields.push(Arc::new(inta));
1993 let intb = Type::primitive_type_builder("b", PhysicalType::INT64)
1994 .with_converted_type(ConvertedType::INT_64)
1995 .build()?;
1996 fields.push(Arc::new(intb));
1997 let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY)
1998 .with_repetition(Repetition::REPEATED)
1999 .with_converted_type(ConvertedType::UTF8)
2000 .build()?;
2001 fields.push(Arc::new(intc));
2002
2003 let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64)
2005 .with_repetition(Repetition::REQUIRED)
2006 .with_converted_type(ConvertedType::INT_64)
2007 .build()?;
2008 let item2 = Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?;
2009 let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32)
2010 .with_repetition(Repetition::REPEATED)
2011 .with_converted_type(ConvertedType::INT_32)
2012 .build()?;
2013 let list = Type::group_type_builder("records")
2014 .with_repetition(Repetition::REPEATED)
2015 .with_converted_type(ConvertedType::LIST)
2016 .with_fields(vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)])
2017 .build()?;
2018 let bag = Type::group_type_builder("bag")
2019 .with_repetition(Repetition::OPTIONAL)
2020 .with_fields(vec![Arc::new(list)])
2021 .build()?;
2022 fields.push(Arc::new(bag));
2023
2024 let schema = Type::group_type_builder("schema")
2025 .with_repetition(Repetition::REPEATED)
2026 .with_fields(fields)
2027 .build()?;
2028 let descr = SchemaDescriptor::new(Arc::new(schema));
2029
2030 let nleaves = 6;
2031 assert_eq!(descr.num_columns(), nleaves);
2032
2033 let ex_max_def_levels = [0, 1, 1, 2, 3, 3];
2043 let ex_max_rep_levels = [0, 0, 1, 1, 1, 2];
2044
2045 for i in 0..nleaves {
2046 let col = descr.column(i);
2047 assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{i}");
2048 assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{i}");
2049 }
2050
2051 assert_eq!(descr.column(0).path().string(), "a");
2052 assert_eq!(descr.column(1).path().string(), "b");
2053 assert_eq!(descr.column(2).path().string(), "c");
2054 assert_eq!(descr.column(3).path().string(), "bag.records.item1");
2055 assert_eq!(descr.column(4).path().string(), "bag.records.item2");
2056 assert_eq!(descr.column(5).path().string(), "bag.records.item3");
2057
2058 assert_eq!(descr.get_column_root(0).name(), "a");
2059 assert_eq!(descr.get_column_root(3).name(), "bag");
2060 assert_eq!(descr.get_column_root(4).name(), "bag");
2061 assert_eq!(descr.get_column_root(5).name(), "bag");
2062
2063 Ok(())
2064 }
2065
2066 #[test]
2067 fn test_schema_build_tree_def_rep_levels() {
2068 let message_type = "
2069 message spark_schema {
2070 REQUIRED INT32 a;
2071 OPTIONAL group b {
2072 OPTIONAL INT32 _1;
2073 OPTIONAL INT32 _2;
2074 }
2075 OPTIONAL group c (LIST) {
2076 REPEATED group list {
2077 OPTIONAL INT32 element;
2078 }
2079 }
2080 }
2081 ";
2082 let schema = parse_message_type(message_type).expect("should parse schema");
2083 let descr = SchemaDescriptor::new(Arc::new(schema));
2084 assert_eq!(descr.column(0).max_def_level(), 0);
2086 assert_eq!(descr.column(0).max_rep_level(), 0);
2087 assert_eq!(descr.column(1).max_def_level(), 2);
2089 assert_eq!(descr.column(1).max_rep_level(), 0);
2090 assert_eq!(descr.column(2).max_def_level(), 2);
2092 assert_eq!(descr.column(2).max_rep_level(), 0);
2093 assert_eq!(descr.column(3).max_def_level(), 3);
2095 assert_eq!(descr.column(3).max_rep_level(), 1);
2096 }
2097
2098 #[test]
2099 fn test_schema_build_tree_repeated_ancestor_def_level() {
2100 let message_type = "
2102 message m {
2103 REQUIRED INT32 a;
2104 OPTIONAL INT32 b;
2105 OPTIONAL group s {
2106 OPTIONAL INT32 x;
2107 }
2108 }
2109 ";
2110 let schema = parse_message_type(message_type).expect("should parse schema");
2111 let descr = SchemaDescriptor::new(Arc::new(schema));
2112 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 = "
2119 message m {
2120 OPTIONAL group c (LIST) {
2121 REPEATED group list {
2122 OPTIONAL INT32 element;
2123 }
2124 }
2125 }
2126 ";
2127 let schema = parse_message_type(message_type).expect("should parse schema");
2128 let descr = SchemaDescriptor::new(Arc::new(schema));
2129 assert_eq!(descr.column(0).max_def_level(), 3);
2131 assert_eq!(descr.column(0).max_rep_level(), 1);
2132 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2);
2133
2134 let message_type = "
2137 message m {
2138 REQUIRED group c (LIST) {
2139 REPEATED group list {
2140 REQUIRED INT32 element;
2141 }
2142 }
2143 }
2144 ";
2145 let schema = parse_message_type(message_type).expect("should parse schema");
2146 let descr = SchemaDescriptor::new(Arc::new(schema));
2147 assert_eq!(descr.column(0).max_def_level(), 1);
2149 assert_eq!(descr.column(0).max_rep_level(), 1);
2150 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 1);
2151
2152 let message_type = "
2154 message m {
2155 OPTIONAL group outer (LIST) {
2156 REPEATED group list {
2157 OPTIONAL group inner (LIST) {
2158 REPEATED group list2 {
2159 OPTIONAL INT32 element;
2160 }
2161 }
2162 }
2163 }
2164 }
2165 ";
2166 let schema = parse_message_type(message_type).expect("should parse schema");
2167 let descr = SchemaDescriptor::new(Arc::new(schema));
2168 assert_eq!(descr.column(0).max_def_level(), 5);
2170 assert_eq!(descr.column(0).max_rep_level(), 2);
2171 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 4);
2172
2173 let message_type = "
2175 message m {
2176 OPTIONAL group bag (LIST) {
2177 REPEATED group list {
2178 REQUIRED group item {
2179 OPTIONAL INT32 x;
2180 REQUIRED INT32 y;
2181 }
2182 }
2183 }
2184 }
2185 ";
2186 let schema = parse_message_type(message_type).expect("should parse schema");
2187 let descr = SchemaDescriptor::new(Arc::new(schema));
2188 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); let message_type = "
2195 message m {
2196 OPTIONAL group my_map (MAP) {
2197 REPEATED group key_value {
2198 REQUIRED BYTE_ARRAY key (UTF8);
2199 OPTIONAL INT32 value;
2200 }
2201 }
2202 }
2203 ";
2204 let schema = parse_message_type(message_type).expect("should parse schema");
2205 let descr = SchemaDescriptor::new(Arc::new(schema));
2206 assert_eq!(descr.column(0).max_def_level(), 2);
2208 assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); assert_eq!(descr.column(1).max_def_level(), 3);
2211 assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); }
2213
2214 #[test]
2215 #[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")]
2216 fn test_get_physical_type_panic() {
2217 let list = Type::group_type_builder("records")
2218 .with_repetition(Repetition::REPEATED)
2219 .build()
2220 .unwrap();
2221 list.get_physical_type();
2222 }
2223
2224 #[test]
2225 fn test_get_physical_type_primitive() {
2226 let f = Type::primitive_type_builder("f", PhysicalType::INT64)
2227 .build()
2228 .unwrap();
2229 assert_eq!(f.get_physical_type(), PhysicalType::INT64);
2230
2231 let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY)
2232 .build()
2233 .unwrap();
2234 assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY);
2235 }
2236
2237 #[test]
2238 fn test_check_contains_primitive_primitive() {
2239 let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2241 .build()
2242 .unwrap();
2243 let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2244 .build()
2245 .unwrap();
2246 assert!(f1.check_contains(&f2));
2247
2248 let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2250 .with_converted_type(ConvertedType::UINT_8)
2251 .build()
2252 .unwrap();
2253 let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2254 .with_converted_type(ConvertedType::UINT_16)
2255 .build()
2256 .unwrap();
2257 assert!(f1.check_contains(&f2));
2258
2259 let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
2261 .build()
2262 .unwrap();
2263 let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
2264 .build()
2265 .unwrap();
2266 assert!(!f1.check_contains(&f2));
2267
2268 let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2270 .build()
2271 .unwrap();
2272 let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
2273 .build()
2274 .unwrap();
2275 assert!(!f1.check_contains(&f2));
2276
2277 let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
2279 .with_repetition(Repetition::REQUIRED)
2280 .build()
2281 .unwrap();
2282 let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
2283 .with_repetition(Repetition::OPTIONAL)
2284 .build()
2285 .unwrap();
2286 assert!(!f1.check_contains(&f2));
2287 }
2288
2289 fn test_new_group_type(name: &str, repetition: Repetition, types: Vec<Type>) -> Type {
2291 Type::group_type_builder(name)
2292 .with_repetition(repetition)
2293 .with_fields(types.into_iter().map(Arc::new).collect())
2294 .build()
2295 .unwrap()
2296 }
2297
2298 #[test]
2299 fn test_check_contains_group_group() {
2300 let f1 = Type::group_type_builder("f").build().unwrap();
2302 let f2 = Type::group_type_builder("f").build().unwrap();
2303 assert!(f1.check_contains(&f2));
2304 assert!(!f1.is_optional());
2305
2306 let f1 = test_new_group_type(
2308 "f",
2309 Repetition::REPEATED,
2310 vec![
2311 Type::primitive_type_builder("f1", PhysicalType::INT32)
2312 .build()
2313 .unwrap(),
2314 Type::primitive_type_builder("f2", PhysicalType::INT64)
2315 .build()
2316 .unwrap(),
2317 ],
2318 );
2319 let f2 = test_new_group_type(
2320 "f",
2321 Repetition::REPEATED,
2322 vec![
2323 Type::primitive_type_builder("f1", PhysicalType::INT32)
2324 .build()
2325 .unwrap(),
2326 Type::primitive_type_builder("f2", PhysicalType::INT64)
2327 .build()
2328 .unwrap(),
2329 ],
2330 );
2331 assert!(f1.check_contains(&f2));
2332
2333 let f1 = test_new_group_type(
2335 "f",
2336 Repetition::REPEATED,
2337 vec![
2338 Type::primitive_type_builder("f1", PhysicalType::INT32)
2339 .build()
2340 .unwrap(),
2341 Type::primitive_type_builder("f2", PhysicalType::INT64)
2342 .build()
2343 .unwrap(),
2344 ],
2345 );
2346 let f2 = test_new_group_type(
2347 "f",
2348 Repetition::REPEATED,
2349 vec![
2350 Type::primitive_type_builder("f2", PhysicalType::INT64)
2351 .build()
2352 .unwrap(),
2353 ],
2354 );
2355 assert!(f1.check_contains(&f2));
2356
2357 let f1 = Type::group_type_builder("f1").build().unwrap();
2359 let f2 = Type::group_type_builder("f2").build().unwrap();
2360 assert!(!f1.check_contains(&f2));
2361
2362 let f1 = Type::group_type_builder("f")
2364 .with_repetition(Repetition::OPTIONAL)
2365 .build()
2366 .unwrap();
2367 let f2 = Type::group_type_builder("f")
2368 .with_repetition(Repetition::REPEATED)
2369 .build()
2370 .unwrap();
2371 assert!(!f1.check_contains(&f2));
2372
2373 let f1 = test_new_group_type(
2375 "f",
2376 Repetition::REPEATED,
2377 vec![
2378 Type::primitive_type_builder("f1", PhysicalType::INT32)
2379 .build()
2380 .unwrap(),
2381 Type::primitive_type_builder("f2", PhysicalType::INT64)
2382 .build()
2383 .unwrap(),
2384 ],
2385 );
2386 let f2 = test_new_group_type(
2387 "f",
2388 Repetition::REPEATED,
2389 vec![
2390 Type::primitive_type_builder("f1", PhysicalType::INT32)
2391 .build()
2392 .unwrap(),
2393 Type::primitive_type_builder("f2", PhysicalType::BOOLEAN)
2394 .build()
2395 .unwrap(),
2396 ],
2397 );
2398 assert!(!f1.check_contains(&f2));
2399
2400 let f1 = test_new_group_type(
2402 "f",
2403 Repetition::REPEATED,
2404 vec![
2405 Type::primitive_type_builder("f1", PhysicalType::INT32)
2406 .build()
2407 .unwrap(),
2408 Type::primitive_type_builder("f2", PhysicalType::INT64)
2409 .build()
2410 .unwrap(),
2411 ],
2412 );
2413 let f2 = test_new_group_type(
2414 "f",
2415 Repetition::REPEATED,
2416 vec![
2417 Type::primitive_type_builder("f3", PhysicalType::INT32)
2418 .build()
2419 .unwrap(),
2420 ],
2421 );
2422 assert!(!f1.check_contains(&f2));
2423 }
2424
2425 #[test]
2426 fn test_check_contains_group_primitive() {
2427 let f1 = Type::group_type_builder("f").build().unwrap();
2429 let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
2430 .build()
2431 .unwrap();
2432 assert!(!f1.check_contains(&f2));
2433 assert!(!f2.check_contains(&f1));
2434
2435 let f1 = test_new_group_type(
2437 "f",
2438 Repetition::REPEATED,
2439 vec![
2440 Type::primitive_type_builder("f1", PhysicalType::INT32)
2441 .build()
2442 .unwrap(),
2443 ],
2444 );
2445 let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32)
2446 .build()
2447 .unwrap();
2448 assert!(!f1.check_contains(&f2));
2449 assert!(!f2.check_contains(&f1));
2450
2451 let f1 = test_new_group_type(
2453 "a",
2454 Repetition::REPEATED,
2455 vec![
2456 test_new_group_type(
2457 "b",
2458 Repetition::REPEATED,
2459 vec![
2460 Type::primitive_type_builder("c", PhysicalType::INT32)
2461 .build()
2462 .unwrap(),
2463 ],
2464 ),
2465 Type::primitive_type_builder("d", PhysicalType::INT64)
2466 .build()
2467 .unwrap(),
2468 Type::primitive_type_builder("e", PhysicalType::BOOLEAN)
2469 .build()
2470 .unwrap(),
2471 ],
2472 );
2473 let f2 = test_new_group_type(
2474 "a",
2475 Repetition::REPEATED,
2476 vec![test_new_group_type(
2477 "b",
2478 Repetition::REPEATED,
2479 vec![
2480 Type::primitive_type_builder("c", PhysicalType::INT32)
2481 .build()
2482 .unwrap(),
2483 ],
2484 )],
2485 );
2486 assert!(f1.check_contains(&f2)); assert!(!f2.check_contains(&f1)); }
2489
2490 #[test]
2491 fn test_schema_type_thrift_conversion_err() {
2492 let schema = Type::primitive_type_builder("col", PhysicalType::INT32)
2493 .build()
2494 .unwrap();
2495 let schema = Arc::new(schema);
2496 let thrift_schema = schema_to_buf(&schema);
2497 assert!(thrift_schema.is_err());
2498 if let Err(e) = thrift_schema {
2499 assert_eq!(
2500 format!("{e}"),
2501 "Parquet error: Root schema must be Group type"
2502 );
2503 }
2504 }
2505
2506 #[test]
2507 fn test_schema_type_thrift_conversion() {
2508 let message_type = "
2509 message conversions {
2510 REQUIRED INT64 id;
2511 OPTIONAL FIXED_LEN_BYTE_ARRAY (2) f16 (FLOAT16);
2512 OPTIONAL group int_array_Array (LIST) {
2513 REPEATED group list {
2514 OPTIONAL group element (LIST) {
2515 REPEATED group list {
2516 OPTIONAL INT32 element;
2517 }
2518 }
2519 }
2520 }
2521 OPTIONAL group int_map (MAP) {
2522 REPEATED group map (MAP_KEY_VALUE) {
2523 REQUIRED BYTE_ARRAY key (UTF8);
2524 OPTIONAL INT32 value;
2525 }
2526 }
2527 OPTIONAL group int_Map_Array (LIST) {
2528 REPEATED group list {
2529 OPTIONAL group g (MAP) {
2530 REPEATED group map (MAP_KEY_VALUE) {
2531 REQUIRED BYTE_ARRAY key (UTF8);
2532 OPTIONAL group value {
2533 OPTIONAL group H {
2534 OPTIONAL group i (LIST) {
2535 REPEATED group list {
2536 OPTIONAL DOUBLE element;
2537 }
2538 }
2539 }
2540 }
2541 }
2542 }
2543 }
2544 }
2545 OPTIONAL group nested_struct {
2546 OPTIONAL INT32 A;
2547 OPTIONAL group b (LIST) {
2548 REPEATED group list {
2549 REQUIRED FIXED_LEN_BYTE_ARRAY (16) element;
2550 }
2551 }
2552 }
2553 }
2554 ";
2555 let expected_schema = parse_message_type(message_type).unwrap();
2556 let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap();
2557 assert_eq!(result_schema, Arc::new(expected_schema));
2558 }
2559
2560 #[test]
2561 fn test_schema_type_thrift_conversion_decimal() {
2562 let message_type = "
2563 message decimals {
2564 OPTIONAL INT32 field0;
2565 OPTIONAL INT64 field1 (DECIMAL (18, 2));
2566 OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18));
2567 OPTIONAL BYTE_ARRAY field3 (DECIMAL (9));
2568 }
2569 ";
2570 let expected_schema = parse_message_type(message_type).unwrap();
2571 let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap();
2572 assert_eq!(result_schema, Arc::new(expected_schema));
2573 }
2574
2575 #[test]
2578 fn test_schema_from_thrift_with_num_children_set() {
2579 let message_type = "
2581 message schema {
2582 OPTIONAL BYTE_ARRAY id (UTF8);
2583 OPTIONAL BYTE_ARRAY name (UTF8);
2584 OPTIONAL BYTE_ARRAY message (UTF8);
2585 OPTIONAL INT32 type (UINT_8);
2586 OPTIONAL INT64 author_time (TIMESTAMP_MILLIS);
2587 OPTIONAL INT64 __index_level_0__;
2588 }
2589 ";
2590
2591 let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2592 let mut buf = schema_to_buf(&expected_schema).unwrap();
2593 let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2594
2595 for elem in &mut thrift_schema[..] {
2597 if elem.num_children.is_none() {
2598 elem.num_children = Some(0);
2599 }
2600 }
2601
2602 let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2603 assert_eq!(result_schema, expected_schema);
2604 }
2605
2606 #[test]
2609 fn test_schema_from_thrift_root_has_repetition() {
2610 let message_type = "
2612 message schema {
2613 OPTIONAL BYTE_ARRAY a (UTF8);
2614 OPTIONAL INT32 b (UINT_8);
2615 }
2616 ";
2617
2618 let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2619 let mut buf = schema_to_buf(&expected_schema).unwrap();
2620 let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2621 thrift_schema[0].repetition_type = Some(Repetition::REQUIRED);
2622
2623 let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2624 assert_eq!(result_schema, expected_schema);
2625 }
2626
2627 #[test]
2628 fn test_schema_from_thrift_group_has_no_child() {
2629 let message_type = "message schema {}";
2630
2631 let expected_schema = Arc::new(parse_message_type(message_type).unwrap());
2632 let mut buf = schema_to_buf(&expected_schema).unwrap();
2633 let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap();
2634 thrift_schema[0].repetition_type = Some(Repetition::REQUIRED);
2635
2636 let result_schema = parquet_schema_from_array(thrift_schema).unwrap();
2637 assert_eq!(result_schema, expected_schema);
2638 }
2639
2640 fn file_field(name: &str, physical: PhysicalType, logical: Option<LogicalType>) -> TypePtr {
2642 Arc::new(
2643 Type::primitive_type_builder(name, physical)
2644 .with_repetition(Repetition::OPTIONAL)
2645 .with_logical_type(logical)
2646 .build()
2647 .unwrap(),
2648 )
2649 }
2650
2651 fn all_file_fields() -> Vec<TypePtr> {
2653 vec![
2654 file_field("uri", PhysicalType::BYTE_ARRAY, Some(LogicalType::String)),
2655 file_field("offset", PhysicalType::INT64, None),
2656 file_field("size", PhysicalType::INT64, None),
2657 file_field(
2658 "content_type",
2659 PhysicalType::BYTE_ARRAY,
2660 Some(LogicalType::String),
2661 ),
2662 file_field(
2663 "checksum",
2664 PhysicalType::BYTE_ARRAY,
2665 Some(LogicalType::String),
2666 ),
2667 file_field("inline", PhysicalType::BYTE_ARRAY, None),
2668 ]
2669 }
2670
2671 #[test]
2672 fn test_file_logical_type_roundtrip() {
2673 let file_group = Arc::new(
2674 Type::group_type_builder("f")
2675 .with_repetition(Repetition::REQUIRED)
2676 .with_logical_type(Some(LogicalType::File))
2677 .with_fields(all_file_fields())
2678 .build()
2679 .unwrap(),
2680 );
2681 let schema = Arc::new(
2682 Type::group_type_builder("example")
2683 .with_fields(vec![file_group])
2684 .build()
2685 .unwrap(),
2686 );
2687 let result = roundtrip_schema(schema.clone()).unwrap();
2688 assert_eq!(result, schema);
2689 assert_eq!(
2690 result.get_fields()[0].get_basic_info().logical_type_ref(),
2691 Some(&LogicalType::File)
2692 );
2693 }
2694
2695 #[test]
2696 fn test_file_logical_type_all_fields() {
2697 let result = Type::group_type_builder("file_field")
2698 .with_repetition(Repetition::REQUIRED)
2699 .with_logical_type(Some(LogicalType::File))
2700 .with_fields(all_file_fields())
2701 .build();
2702 assert!(result.is_ok());
2703 assert_eq!(result.unwrap().get_fields().len(), 6);
2704 }
2705
2706 #[test]
2707 fn test_file_logical_type_uri_only() {
2708 let result = Type::group_type_builder("file_field")
2710 .with_repetition(Repetition::REQUIRED)
2711 .with_logical_type(Some(LogicalType::File))
2712 .with_fields(vec![file_field(
2713 "uri",
2714 PhysicalType::BYTE_ARRAY,
2715 Some(LogicalType::String),
2716 )])
2717 .build();
2718 assert!(result.is_ok());
2719 assert_eq!(
2720 result.unwrap().get_basic_info().logical_type_ref(),
2721 Some(&LogicalType::File)
2722 );
2723 }
2724
2725 #[test]
2726 fn test_file_logical_type_inline_only() {
2727 let result = Type::group_type_builder("inline_file")
2729 .with_repetition(Repetition::REQUIRED)
2730 .with_logical_type(Some(LogicalType::File))
2731 .with_fields(vec![file_field("inline", PhysicalType::BYTE_ARRAY, None)])
2732 .build();
2733 assert!(result.is_ok());
2734 }
2735
2736 #[test]
2737 fn test_file_logical_type_empty_group_is_allowed() {
2738 let result = Type::group_type_builder("empty_file")
2740 .with_repetition(Repetition::REQUIRED)
2741 .with_logical_type(Some(LogicalType::File))
2742 .with_fields(vec![])
2743 .build();
2744 assert!(result.is_ok());
2745 }
2746
2747 #[test]
2748 fn test_file_logical_type_rejects_unrecognized_field() {
2749 let unknown_field = file_field("unknown_field", PhysicalType::BYTE_ARRAY, None);
2750 let result = Type::group_type_builder("bad_file")
2751 .with_repetition(Repetition::REQUIRED)
2752 .with_logical_type(Some(LogicalType::File))
2753 .with_fields(vec![
2754 file_field("uri", PhysicalType::BYTE_ARRAY, Some(LogicalType::String)),
2755 unknown_field,
2756 ])
2757 .build();
2758 assert_eq!(
2759 result.unwrap_err().to_string(),
2760 "Parquet error: FILE type group 'bad_file' contains unrecognized field \
2761 'unknown_field'. Valid fields are: uri, offset, size, content_type, \
2762 checksum, inline"
2763 );
2764 }
2765
2766 #[test]
2767 fn test_file_logical_type_requires_optional_fields() {
2768 let uri_field = Arc::new(
2770 Type::primitive_type_builder("uri", PhysicalType::BYTE_ARRAY)
2771 .with_repetition(Repetition::REQUIRED)
2772 .with_logical_type(Some(LogicalType::String))
2773 .build()
2774 .unwrap(),
2775 );
2776 let result = Type::group_type_builder("required_uri")
2777 .with_repetition(Repetition::REQUIRED)
2778 .with_logical_type(Some(LogicalType::File))
2779 .with_fields(vec![uri_field])
2780 .build();
2781 assert_eq!(
2782 result.unwrap_err().to_string(),
2783 "Parquet error: FILE type field 'uri' must be OPTIONAL in group 'required_uri'"
2784 );
2785 }
2786
2787 #[test]
2788 fn test_file_logical_type_rejects_wrong_physical_type() {
2789 let bad_size = file_field("size", PhysicalType::BYTE_ARRAY, None);
2791 let result = Type::group_type_builder("bad_size")
2792 .with_repetition(Repetition::REQUIRED)
2793 .with_logical_type(Some(LogicalType::File))
2794 .with_fields(vec![bad_size])
2795 .build();
2796 assert_eq!(
2797 result.unwrap_err().to_string(),
2798 "Parquet error: FILE type field 'size' in group 'bad_size' must have physical type INT64"
2799 );
2800 }
2801
2802 #[test]
2803 fn test_file_logical_type_rejects_wrong_logical_type() {
2804 let bad_uri = file_field("uri", PhysicalType::BYTE_ARRAY, None);
2806 let result = Type::group_type_builder("bad_uri")
2807 .with_repetition(Repetition::REQUIRED)
2808 .with_logical_type(Some(LogicalType::File))
2809 .with_fields(vec![bad_uri])
2810 .build();
2811 assert_eq!(
2812 result.unwrap_err().to_string(),
2813 "Parquet error: FILE type field 'uri' in group 'bad_uri' must have logical type \
2814 Some(String)"
2815 );
2816 }
2817
2818 #[test]
2819 fn test_file_logical_type_not_allowed_on_primitive() {
2820 let result = Type::primitive_type_builder("bad", PhysicalType::BYTE_ARRAY)
2821 .with_repetition(Repetition::REQUIRED)
2822 .with_logical_type(Some(LogicalType::File))
2823 .build();
2824 assert!(result.is_err());
2825 }
2826
2827 #[test]
2828 fn test_parquet_schema_from_array_rejects_negative_num_children() {
2829 let elements = vec![SchemaElement {
2830 r#type: None,
2831 type_length: None,
2832 repetition_type: Some(Repetition::REQUIRED),
2833 name: "schema",
2834 num_children: Some(-1),
2835 converted_type: None,
2836 scale: None,
2837 precision: None,
2838 field_id: None,
2839 logical_type: None,
2840 }];
2841 let result = parquet_schema_from_array(elements);
2842 assert!(result.unwrap_err().to_string().contains("Integer overflow"));
2843 }
2844}