1use base64::Engine;
21use base64::prelude::BASE64_STANDARD;
22use std::collections::HashMap;
23use std::sync::Arc;
24
25use arrow_ipc::writer;
26use arrow_schema::{DataType, Field, FieldRef, Fields, Schema, TimeUnit};
27
28use crate::basic::{
29 ConvertedType, LogicalType, Repetition, TimeUnit as ParquetTimeUnit, Type as PhysicalType,
30};
31use crate::errors::{ParquetError, Result};
32use crate::file::{metadata::KeyValue, properties::WriterProperties};
33use crate::schema::types::{ColumnDescriptor, SchemaDescriptor, Type};
34
35mod complex;
36mod extension;
37mod primitive;
38pub mod virtual_type;
39
40use super::PARQUET_FIELD_ID_META_KEY;
41use crate::arrow::ProjectionMask;
42use crate::arrow::schema::extension::{
43 has_extension_type, logical_type_for_binary, logical_type_for_binary_view,
44 logical_type_for_fixed_size_binary, logical_type_for_string, logical_type_for_struct,
45 try_add_extension_type,
46};
47pub(crate) use complex::{ParquetField, ParquetFieldType, VirtualColumnType};
48
49pub fn parquet_to_arrow_schema(
54 parquet_schema: &SchemaDescriptor,
55 key_value_metadata: Option<&Vec<KeyValue>>,
56) -> Result<Schema> {
57 parquet_to_arrow_schema_by_columns(parquet_schema, ProjectionMask::all(), key_value_metadata)
58}
59
60pub fn parquet_to_arrow_schema_by_columns(
63 parquet_schema: &SchemaDescriptor,
64 mask: ProjectionMask,
65 key_value_metadata: Option<&Vec<KeyValue>>,
66) -> Result<Schema> {
67 Ok(parquet_to_arrow_schema_and_fields(parquet_schema, mask, key_value_metadata, &[])?.0)
68}
69
70pub(crate) fn parquet_to_arrow_schema_and_fields(
76 parquet_schema: &SchemaDescriptor,
77 mask: ProjectionMask,
78 key_value_metadata: Option<&Vec<KeyValue>>,
79 virtual_columns: &[FieldRef],
80) -> Result<(Schema, Option<ParquetField>)> {
81 let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
82 let maybe_schema = metadata
83 .remove(super::ARROW_SCHEMA_META_KEY)
84 .map(|value| get_arrow_schema_from_metadata(&value))
85 .transpose()?;
86
87 if let Some(arrow_schema) = &maybe_schema {
89 arrow_schema.metadata().iter().for_each(|(k, v)| {
90 metadata.entry(k.clone()).or_insert_with(|| v.clone());
91 });
92 }
93
94 let hint = maybe_schema.as_ref().map(|s| s.fields());
95 let field_levels =
96 parquet_to_arrow_field_levels_with_virtual(parquet_schema, mask, hint, virtual_columns)?;
97 let schema = Schema::new_with_metadata(field_levels.fields, metadata);
98 Ok((schema, field_levels.levels))
99}
100
101#[derive(Debug, Clone)]
109pub struct FieldLevels {
110 pub(crate) fields: Fields,
111 pub(crate) levels: Option<ParquetField>,
112}
113
114pub fn parquet_to_arrow_field_levels(
135 schema: &SchemaDescriptor,
136 mask: ProjectionMask,
137 hint: Option<&Fields>,
138) -> Result<FieldLevels> {
139 parquet_to_arrow_field_levels_with_virtual(schema, mask, hint, &[])
140}
141
142pub fn parquet_to_arrow_field_levels_with_virtual(
172 schema: &SchemaDescriptor,
173 mask: ProjectionMask,
174 hint: Option<&Fields>,
175 virtual_columns: &[FieldRef],
176) -> Result<FieldLevels> {
177 for field in virtual_columns {
179 if !virtual_type::is_virtual_column(field) {
180 return Err(ParquetError::General(format!(
181 "Field '{}' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'",
182 field.name()
183 )));
184 }
185 }
186
187 let mut parquet_field = match complex::convert_schema(schema, mask, hint)? {
189 Some(field) => field,
190 None if virtual_columns.is_empty() => {
191 return Ok(FieldLevels {
192 fields: Fields::empty(),
193 levels: None,
194 });
195 }
196 None => {
197 ParquetField {
199 rep_level: 0,
200 def_level: 0,
201 nullable: false,
202 arrow_type: DataType::Struct(Fields::empty()),
203 field_type: ParquetFieldType::Group {
204 children: Vec::new(),
205 },
206 }
207 }
208 };
209
210 if !virtual_columns.is_empty() {
212 match &mut parquet_field.field_type {
213 ParquetFieldType::Group { children } => {
214 let DataType::Struct(ref mut fields) = parquet_field.arrow_type else {
216 unreachable!("Root field must be a struct");
217 };
218
219 let mut fields_vec: Vec<FieldRef> = fields.iter().cloned().collect();
221
222 for virtual_column in virtual_columns {
224 assert_eq!(
226 parquet_field.rep_level, 0,
227 "Virtual columns can only be added at rep level 0"
228 );
229 assert_eq!(
230 parquet_field.def_level, 0,
231 "Virtual columns can only be added at def level 0"
232 );
233
234 fields_vec.push(virtual_column.clone());
235 let virtual_parquet_field = complex::convert_virtual_field(
236 virtual_column,
237 parquet_field.rep_level,
238 parquet_field.def_level,
239 )?;
240 children.push(virtual_parquet_field);
241 }
242
243 parquet_field.arrow_type = DataType::Struct(Fields::from(fields_vec));
245 }
246 _ => unreachable!("Root field must be a group"),
247 }
248 }
249
250 match &parquet_field.arrow_type {
251 DataType::Struct(fields) => Ok(FieldLevels {
252 fields: fields.clone(),
253 levels: Some(parquet_field),
254 }),
255 _ => unreachable!(),
256 }
257}
258
259fn get_arrow_schema_from_metadata(encoded_meta: &str) -> Result<Schema> {
261 let decoded = BASE64_STANDARD.decode(encoded_meta);
262 match decoded {
263 Ok(bytes) => {
264 let slice = if bytes.len() > 8 && bytes[0..4] == [255u8; 4] {
265 &bytes[8..]
266 } else {
267 bytes.as_slice()
268 };
269 match arrow_ipc::root_as_message(slice) {
270 Ok(message) => {
271 let schema = message
272 .header_as_schema()
273 .ok_or_else(|| arrow_err!("the message is not Arrow Schema"))?;
274 arrow_ipc::convert::try_fb_to_schema(schema).map_err(Into::into)
275 }
276 Err(err) => {
277 Err(arrow_err!(
279 "Unable to get root as message stored in {}: {:?}",
280 super::ARROW_SCHEMA_META_KEY,
281 err
282 ))
283 }
284 }
285 }
286 Err(err) => {
287 Err(arrow_err!(
289 "Unable to decode the encoded schema stored in {}, {:?}",
290 super::ARROW_SCHEMA_META_KEY,
291 err
292 ))
293 }
294 }
295}
296
297pub fn encode_arrow_schema(schema: &Schema) -> String {
299 let options = writer::IpcWriteOptions::default();
300 let mut dictionary_tracker = writer::DictionaryTracker::new(true);
301 let data_gen = writer::IpcDataGenerator::default();
302 let mut serialized_schema =
303 data_gen.schema_to_bytes_with_dictionary_tracker(schema, &mut dictionary_tracker, &options);
304
305 let schema_len = serialized_schema.ipc_message.len();
308 let mut len_prefix_schema = Vec::with_capacity(schema_len + 8);
309 len_prefix_schema.append(&mut vec![255u8, 255, 255, 255]);
310 len_prefix_schema.append((schema_len as u32).to_le_bytes().to_vec().as_mut());
311 len_prefix_schema.append(&mut serialized_schema.ipc_message);
312
313 BASE64_STANDARD.encode(&len_prefix_schema)
314}
315
316pub(crate) fn validate_map_key_type(data_type: &DataType) -> Result<()> {
319 match data_type {
320 DataType::Map(entries, _) => {
321 let DataType::Struct(fields) = entries.data_type() else {
322 return Err(arrow_err!(
323 "DataType::Map should contain a struct field child"
324 ));
325 };
326 if fields.len() != 2 {
327 return Err(arrow_err!("DataType::Map entries must contain two fields"));
328 }
329 let key = &fields[0];
330 let mut key_type = key.data_type();
331 let mut nullable = key.is_nullable();
332 loop {
333 match key_type {
334 DataType::Dictionary(_, value) => key_type = value,
335 DataType::RunEndEncoded(_, value) => {
336 nullable |= value.is_nullable();
337 key_type = value.data_type();
338 }
339 _ => break,
340 }
341 }
342 if nullable {
343 return Err(arrow_err!(
344 "Map key field '{}' must be non-nullable, including its encoded values",
345 key.name()
346 ));
347 }
348 validate_map_key_type(entries.data_type())?;
349 }
350 DataType::Struct(fields) => {
351 for field in fields {
352 validate_map_key_type(field.data_type())?;
353 }
354 }
355 DataType::List(field)
356 | DataType::LargeList(field)
357 | DataType::FixedSizeList(field, _)
358 | DataType::ListView(field)
359 | DataType::LargeListView(field)
360 | DataType::RunEndEncoded(_, field) => validate_map_key_type(field.data_type())?,
361 DataType::Dictionary(_, value) => validate_map_key_type(value)?,
362 _ => {}
363 }
364 Ok(())
365}
366
367fn normalize_metadata_field(field: &Field) -> Field {
373 match field.data_type() {
374 DataType::RunEndEncoded(_, value_field) => {
375 let value = normalize_metadata_field(value_field);
376 field
377 .clone()
378 .with_data_type(value.data_type().clone())
379 .with_nullable(field.is_nullable() || value.is_nullable())
380 }
381 DataType::Dictionary(key, value_type) => {
382 let value_field = Field::new("", value_type.as_ref().clone(), false);
385 let value = normalize_metadata_field(&value_field);
386 let value_type = value.data_type();
387 let data_type = if value_type.is_nested() {
390 value_type.clone()
391 } else {
392 DataType::Dictionary(key.clone(), Box::new(value_type.clone()))
393 };
394 field
395 .clone()
396 .with_data_type(data_type)
397 .with_nullable(field.is_nullable() || value.is_nullable())
398 }
399 _ => field
400 .clone()
401 .with_data_type(normalize_metadata_children(field.data_type())),
402 }
403}
404
405fn normalize_metadata_children(dt: &DataType) -> DataType {
406 match dt {
407 DataType::Struct(fields) => DataType::Struct(
408 fields
409 .iter()
410 .map(|f| Arc::new(normalize_metadata_field(f)))
411 .collect(),
412 ),
413 DataType::List(f) => DataType::List(Arc::new(normalize_metadata_field(f))),
414 DataType::LargeList(f) => DataType::LargeList(Arc::new(normalize_metadata_field(f))),
415 DataType::FixedSizeList(f, n) => {
416 DataType::FixedSizeList(Arc::new(normalize_metadata_field(f)), *n)
417 }
418 DataType::ListView(f) => DataType::ListView(Arc::new(normalize_metadata_field(f))),
419 DataType::LargeListView(f) => {
420 DataType::LargeListView(Arc::new(normalize_metadata_field(f)))
421 }
422 DataType::Map(f, sorted) => DataType::Map(Arc::new(normalize_metadata_field(f)), *sorted),
423 _ => dt.clone(),
424 }
425}
426
427fn needs_metadata_normalization(dt: &DataType) -> bool {
429 match dt {
430 DataType::RunEndEncoded(_, _) => true,
431 DataType::Struct(fields) => fields
432 .iter()
433 .any(|f| needs_metadata_normalization(f.data_type())),
434 DataType::List(f)
435 | DataType::LargeList(f)
436 | DataType::FixedSizeList(f, _)
437 | DataType::ListView(f)
438 | DataType::LargeListView(f)
439 | DataType::Map(f, _) => needs_metadata_normalization(f.data_type()),
440 DataType::Dictionary(_, value) => value.is_nested() || needs_metadata_normalization(value),
441 _ => false,
442 }
443}
444
445pub fn add_encoded_arrow_schema_to_metadata(schema: &Schema, props: &mut WriterProperties) {
452 let needs_normalization = schema
453 .fields()
454 .iter()
455 .any(|f| needs_metadata_normalization(f.data_type()));
456 let metadata_schema;
457 let schema = if needs_normalization {
458 let fields: Vec<Field> = schema
459 .fields()
460 .iter()
461 .map(|f| normalize_metadata_field(f))
462 .collect();
463 metadata_schema = Schema::new_with_metadata(fields, schema.metadata().clone());
464 &metadata_schema
465 } else {
466 schema
467 };
468 let encoded = encode_arrow_schema(schema);
469
470 let schema_kv = KeyValue {
471 key: super::ARROW_SCHEMA_META_KEY.to_string(),
472 value: Some(encoded),
473 };
474
475 let meta = props
476 .key_value_metadata
477 .get_or_insert_with(Default::default);
478
479 let schema_meta = meta
481 .iter()
482 .enumerate()
483 .find(|(_, kv)| kv.key.as_str() == super::ARROW_SCHEMA_META_KEY);
484 match schema_meta {
485 Some((i, _)) => {
486 meta.remove(i);
487 meta.push(schema_kv);
488 }
489 None => {
490 meta.push(schema_kv);
491 }
492 }
493}
494
495#[derive(Debug)]
540pub struct ArrowSchemaConverter<'a> {
541 schema_root: &'a str,
543 coerce_types: bool,
547}
548
549impl Default for ArrowSchemaConverter<'_> {
550 fn default() -> Self {
551 Self::new()
552 }
553}
554
555impl<'a> ArrowSchemaConverter<'a> {
556 pub fn new() -> Self {
558 Self {
559 schema_root: "arrow_schema",
560 coerce_types: false,
561 }
562 }
563
564 pub fn with_coerce_types(mut self, coerce_types: bool) -> Self {
595 self.coerce_types = coerce_types;
596 self
597 }
598
599 pub fn schema_root(mut self, schema_root: &'a str) -> Self {
601 self.schema_root = schema_root;
602 self
603 }
604
605 pub fn convert(&self, schema: &Schema) -> Result<SchemaDescriptor> {
609 for field in schema.fields() {
610 validate_map_key_type(field.data_type())?;
611 }
612 let fields = schema
613 .fields()
614 .iter()
615 .map(|field| arrow_to_parquet_type(field, self.coerce_types).map(Arc::new))
616 .collect::<Result<_>>()?;
617 let group = Type::group_type_builder(self.schema_root)
618 .with_fields(fields)
619 .build()?;
620 Ok(SchemaDescriptor::new(Arc::new(group)))
621 }
622}
623
624fn parse_key_value_metadata(
625 key_value_metadata: Option<&Vec<KeyValue>>,
626) -> Option<HashMap<String, String>> {
627 match key_value_metadata {
628 Some(key_values) => {
629 let map: HashMap<String, String> = key_values
630 .iter()
631 .filter_map(|kv| {
632 kv.value
633 .as_ref()
634 .map(|value| (kv.key.clone(), value.clone()))
635 })
636 .collect();
637
638 if map.is_empty() { None } else { Some(map) }
639 }
640 None => None,
641 }
642}
643
644pub fn parquet_to_arrow_field(parquet_column: &ColumnDescriptor) -> Result<Field> {
646 let field = complex::convert_type(&parquet_column.self_type_ptr())?;
647 let mut ret = Field::new(parquet_column.name(), field.arrow_type, field.nullable);
648
649 let parquet_type = parquet_column.self_type();
650 let basic_info = parquet_type.get_basic_info();
651
652 let mut hash_map_size = 0;
653 if basic_info.has_id() {
654 hash_map_size += 1;
655 }
656 if has_extension_type(parquet_type) {
657 hash_map_size += 1;
658 }
659 if hash_map_size == 0 {
660 return Ok(ret);
661 }
662 ret.set_metadata(HashMap::with_capacity(hash_map_size));
663 if basic_info.has_id() {
664 ret.metadata_mut().insert(
665 PARQUET_FIELD_ID_META_KEY.to_string(),
666 basic_info.id().to_string(),
667 );
668 }
669 try_add_extension_type(ret, parquet_column.self_type())
670}
671
672pub fn decimal_length_from_precision(precision: u8) -> usize {
673 (((10.0_f64.powi(precision as i32) + 1.0).log2() + 1.0) / 8.0).ceil() as usize
681}
682
683fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result<Type> {
685 const PARQUET_LIST_ELEMENT_NAME: &str = "element";
686 const PARQUET_MAP_STRUCT_NAME: &str = "key_value";
687 const PARQUET_KEY_FIELD_NAME: &str = "key";
688 const PARQUET_VALUE_FIELD_NAME: &str = "value";
689
690 let name = field.name().as_str();
691 let repetition = if field.is_nullable() {
692 Repetition::OPTIONAL
693 } else {
694 Repetition::REQUIRED
695 };
696 let id = field_id(field);
697 match field.data_type() {
699 DataType::Null => Type::primitive_type_builder(name, PhysicalType::INT32)
700 .with_logical_type(Some(LogicalType::Unknown))
701 .with_repetition(repetition)
702 .with_id(id)
703 .build(),
704 DataType::Boolean => Type::primitive_type_builder(name, PhysicalType::BOOLEAN)
705 .with_repetition(repetition)
706 .with_id(id)
707 .build(),
708 DataType::Int8 => Type::primitive_type_builder(name, PhysicalType::INT32)
709 .with_logical_type(Some(LogicalType::integer(8, true)))
710 .with_repetition(repetition)
711 .with_id(id)
712 .build(),
713 DataType::Int16 => Type::primitive_type_builder(name, PhysicalType::INT32)
714 .with_logical_type(Some(LogicalType::integer(16, true)))
715 .with_repetition(repetition)
716 .with_id(id)
717 .build(),
718 DataType::Int32 => Type::primitive_type_builder(name, PhysicalType::INT32)
719 .with_repetition(repetition)
720 .with_id(id)
721 .build(),
722 DataType::Int64 => Type::primitive_type_builder(name, PhysicalType::INT64)
723 .with_repetition(repetition)
724 .with_id(id)
725 .build(),
726 DataType::UInt8 => Type::primitive_type_builder(name, PhysicalType::INT32)
727 .with_logical_type(Some(LogicalType::integer(8, false)))
728 .with_repetition(repetition)
729 .with_id(id)
730 .build(),
731 DataType::UInt16 => Type::primitive_type_builder(name, PhysicalType::INT32)
732 .with_logical_type(Some(LogicalType::integer(16, false)))
733 .with_repetition(repetition)
734 .with_id(id)
735 .build(),
736 DataType::UInt32 => Type::primitive_type_builder(name, PhysicalType::INT32)
737 .with_logical_type(Some(LogicalType::integer(32, false)))
738 .with_repetition(repetition)
739 .with_id(id)
740 .build(),
741 DataType::UInt64 => Type::primitive_type_builder(name, PhysicalType::INT64)
742 .with_logical_type(Some(LogicalType::integer(64, false)))
743 .with_repetition(repetition)
744 .with_id(id)
745 .build(),
746 DataType::Float16 => Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
747 .with_repetition(repetition)
748 .with_id(id)
749 .with_logical_type(Some(LogicalType::Float16))
750 .with_length(2)
751 .build(),
752 DataType::Float32 => Type::primitive_type_builder(name, PhysicalType::FLOAT)
753 .with_repetition(repetition)
754 .with_id(id)
755 .build(),
756 DataType::Float64 => Type::primitive_type_builder(name, PhysicalType::DOUBLE)
757 .with_repetition(repetition)
758 .with_id(id)
759 .build(),
760 DataType::Timestamp(TimeUnit::Second, _) => {
761 Type::primitive_type_builder(name, PhysicalType::INT64)
763 .with_repetition(repetition)
764 .with_id(id)
765 .build()
766 }
767 DataType::Timestamp(time_unit, tz) => {
768 Type::primitive_type_builder(name, PhysicalType::INT64)
769 .with_logical_type(Some(LogicalType::timestamp(
770 matches!(tz, Some(z) if !z.as_ref().is_empty()),
772 match time_unit {
773 TimeUnit::Second => unreachable!(),
774 TimeUnit::Millisecond => ParquetTimeUnit::MILLIS,
775 TimeUnit::Microsecond => ParquetTimeUnit::MICROS,
776 TimeUnit::Nanosecond => ParquetTimeUnit::NANOS,
777 },
778 )))
779 .with_repetition(repetition)
780 .with_id(id)
781 .build()
782 }
783 DataType::Date32 => Type::primitive_type_builder(name, PhysicalType::INT32)
784 .with_logical_type(Some(LogicalType::Date))
785 .with_repetition(repetition)
786 .with_id(id)
787 .build(),
788 DataType::Date64 => {
789 if coerce_types {
790 Type::primitive_type_builder(name, PhysicalType::INT32)
791 .with_logical_type(Some(LogicalType::Date))
792 .with_repetition(repetition)
793 .with_id(id)
794 .build()
795 } else {
796 Type::primitive_type_builder(name, PhysicalType::INT64)
797 .with_repetition(repetition)
798 .with_id(id)
799 .build()
800 }
801 }
802 DataType::Time32(TimeUnit::Second) => {
803 Type::primitive_type_builder(name, PhysicalType::INT32)
805 .with_repetition(repetition)
806 .with_id(id)
807 .build()
808 }
809 DataType::Time32(unit) => Type::primitive_type_builder(name, PhysicalType::INT32)
810 .with_logical_type(Some(LogicalType::time(
811 field.metadata().contains_key("adjusted_to_utc"),
812 match unit {
813 TimeUnit::Millisecond => ParquetTimeUnit::MILLIS,
814 u => unreachable!("Invalid unit for Time32: {:?}", u),
815 },
816 )))
817 .with_repetition(repetition)
818 .with_id(id)
819 .build(),
820 DataType::Time64(unit) => Type::primitive_type_builder(name, PhysicalType::INT64)
821 .with_logical_type(Some(LogicalType::time(
822 field.metadata().contains_key("adjusted_to_utc"),
823 match unit {
824 TimeUnit::Microsecond => ParquetTimeUnit::MICROS,
825 TimeUnit::Nanosecond => ParquetTimeUnit::NANOS,
826 u => unreachable!("Invalid unit for Time64: {:?}", u),
827 },
828 )))
829 .with_repetition(repetition)
830 .with_id(id)
831 .build(),
832 DataType::Duration(_) => Type::primitive_type_builder(name, PhysicalType::INT64)
833 .with_repetition(repetition)
834 .with_id(id)
835 .build(),
836 DataType::Interval(_) => {
837 Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
838 .with_converted_type(ConvertedType::INTERVAL)
839 .with_repetition(repetition)
840 .with_id(id)
841 .with_length(12)
842 .build()
843 }
844 DataType::Binary | DataType::LargeBinary => {
845 Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
846 .with_repetition(repetition)
847 .with_id(id)
848 .with_logical_type(logical_type_for_binary(field))
849 .build()
850 }
851 DataType::FixedSizeBinary(length) => {
852 Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
853 .with_repetition(repetition)
854 .with_id(id)
855 .with_length(*length)
856 .with_logical_type(logical_type_for_fixed_size_binary(field))
857 .build()
858 }
859 DataType::BinaryView => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
860 .with_repetition(repetition)
861 .with_id(id)
862 .with_logical_type(logical_type_for_binary_view(field))
863 .build(),
864 DataType::Decimal32(precision, scale)
865 | DataType::Decimal64(precision, scale)
866 | DataType::Decimal128(precision, scale)
867 | DataType::Decimal256(precision, scale) => {
868 let (physical_type, length) = if *precision >= 1 && *precision <= 9 {
871 (PhysicalType::INT32, -1)
872 } else if *precision <= 18 {
873 (PhysicalType::INT64, -1)
874 } else {
875 (
876 PhysicalType::FIXED_LEN_BYTE_ARRAY,
877 decimal_length_from_precision(*precision) as i32,
878 )
879 };
880 Type::primitive_type_builder(name, physical_type)
881 .with_repetition(repetition)
882 .with_id(id)
883 .with_length(length)
884 .with_logical_type(Some(LogicalType::decimal(*scale as i32, *precision as i32)))
885 .with_precision(*precision as i32)
886 .with_scale(*scale as i32)
887 .build()
888 }
889 DataType::Utf8 | DataType::LargeUtf8 => {
890 Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
891 .with_logical_type(logical_type_for_string(field))
892 .with_repetition(repetition)
893 .with_id(id)
894 .build()
895 }
896 DataType::Utf8View => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
897 .with_logical_type(logical_type_for_string(field))
898 .with_repetition(repetition)
899 .with_id(id)
900 .build(),
901 DataType::List(f)
902 | DataType::FixedSizeList(f, _)
903 | DataType::LargeList(f)
904 | DataType::ListView(f)
905 | DataType::LargeListView(f) => {
906 let field_ref = if coerce_types && f.name() != PARQUET_LIST_ELEMENT_NAME {
907 let ff = f.as_ref().clone().with_name(PARQUET_LIST_ELEMENT_NAME);
909 Arc::new(arrow_to_parquet_type(&ff, coerce_types)?)
910 } else {
911 Arc::new(arrow_to_parquet_type(f, coerce_types)?)
912 };
913
914 Type::group_type_builder(name)
915 .with_fields(vec![Arc::new(
916 Type::group_type_builder("list")
917 .with_fields(vec![field_ref])
918 .with_repetition(Repetition::REPEATED)
919 .build()?,
920 )])
921 .with_logical_type(Some(LogicalType::List))
922 .with_repetition(repetition)
923 .with_id(id)
924 .build()
925 }
926 DataType::Struct(fields) => {
927 if fields.is_empty() {
928 return Err(arrow_err!("Parquet does not support writing empty structs",));
929 }
930 let fields = fields
932 .iter()
933 .map(|f| arrow_to_parquet_type(f, coerce_types).map(Arc::new))
934 .collect::<Result<_>>()?;
935 Type::group_type_builder(name)
936 .with_fields(fields)
937 .with_repetition(repetition)
938 .with_id(id)
939 .with_logical_type(logical_type_for_struct(field))
940 .build()
941 }
942 DataType::Map(field, _) => {
943 if let DataType::Struct(struct_fields) = field.data_type() {
944 let map_struct_name = if coerce_types {
946 PARQUET_MAP_STRUCT_NAME
947 } else {
948 field.name()
949 };
950
951 let fix_map_field = |name: &str, fld: &Arc<Field>| -> Result<Arc<Type>> {
953 if coerce_types && fld.name() != name {
954 let f = fld.as_ref().clone().with_name(name);
955 Ok(Arc::new(arrow_to_parquet_type(&f, coerce_types)?))
956 } else {
957 Ok(Arc::new(arrow_to_parquet_type(fld, coerce_types)?))
958 }
959 };
960 let key_field = fix_map_field(PARQUET_KEY_FIELD_NAME, &struct_fields[0])?;
961 let val_field = fix_map_field(PARQUET_VALUE_FIELD_NAME, &struct_fields[1])?;
962
963 Type::group_type_builder(name)
964 .with_fields(vec![Arc::new(
965 Type::group_type_builder(map_struct_name)
966 .with_fields(vec![key_field, val_field])
967 .with_repetition(Repetition::REPEATED)
968 .build()?,
969 )])
970 .with_logical_type(Some(LogicalType::Map))
971 .with_repetition(repetition)
972 .with_id(id)
973 .build()
974 } else {
975 Err(arrow_err!(
976 "DataType::Map should contain a struct field child",
977 ))
978 }
979 }
980 DataType::Union(_, _) => unimplemented!("See ARROW-8817."),
981 DataType::Dictionary(_, value) => {
982 let dict_field = field.clone().with_data_type(value.as_ref().clone());
984 arrow_to_parquet_type(&dict_field, coerce_types)
985 }
986 DataType::RunEndEncoded(_, value) => {
987 let value_field = field
992 .clone()
993 .with_data_type(value.data_type().clone())
994 .with_nullable(field.is_nullable() || value.is_nullable());
995 arrow_to_parquet_type(&value_field, coerce_types)
996 }
997 }
998}
999
1000fn field_id(field: &Field) -> Option<i32> {
1001 let value = field.metadata().get(super::PARQUET_FIELD_ID_META_KEY)?;
1002 value.parse().ok() }
1004
1005#[cfg(test)]
1006mod tests {
1007 use super::*;
1008
1009 use std::{collections::HashMap, sync::Arc};
1010
1011 use crate::arrow::PARQUET_FIELD_ID_META_KEY;
1012 use crate::file::metadata::KeyValue;
1013 use crate::file::reader::FileReader;
1014 use crate::{
1015 arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder},
1016 schema::{parser::parse_message_type, types::SchemaDescriptor},
1017 };
1018 use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit};
1019
1020 #[test]
1021 fn test_flat_primitives() {
1022 let message_type = "
1023 message test_schema {
1024 REQUIRED BOOLEAN boolean;
1025 REQUIRED INT32 int8 (INT_8);
1026 REQUIRED INT32 int16 (INT_16);
1027 REQUIRED INT32 uint8 (INTEGER(8,false));
1028 REQUIRED INT32 uint16 (INTEGER(16,false));
1029 REQUIRED INT32 int32;
1030 REQUIRED INT64 int64;
1031 OPTIONAL DOUBLE double;
1032 OPTIONAL FLOAT float;
1033 OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16);
1034 OPTIONAL BINARY string (UTF8);
1035 OPTIONAL BINARY string_2 (STRING);
1036 OPTIONAL BINARY json (JSON);
1037 }
1038 ";
1039 let parquet_group_type = parse_message_type(message_type).unwrap();
1040
1041 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1042 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1043
1044 let arrow_fields = Fields::from(vec![
1045 Field::new("boolean", DataType::Boolean, false),
1046 Field::new("int8", DataType::Int8, false),
1047 Field::new("int16", DataType::Int16, false),
1048 Field::new("uint8", DataType::UInt8, false),
1049 Field::new("uint16", DataType::UInt16, false),
1050 Field::new("int32", DataType::Int32, false),
1051 Field::new("int64", DataType::Int64, false),
1052 Field::new("double", DataType::Float64, true),
1053 Field::new("float", DataType::Float32, true),
1054 Field::new("float16", DataType::Float16, true),
1055 Field::new("string", DataType::Utf8, true),
1056 Field::new("string_2", DataType::Utf8, true),
1057 json_field(),
1058 ]);
1059
1060 assert_eq!(&arrow_fields, converted_arrow_schema.fields());
1061 }
1062
1063 fn json_field() -> Field {
1066 #[cfg(feature = "arrow_canonical_extension_types")]
1067 {
1068 Field::new("json", DataType::Utf8, true)
1069 .with_extension_type(arrow_schema::extension::Json::default())
1070 }
1071 #[cfg(not(feature = "arrow_canonical_extension_types"))]
1072 {
1073 Field::new("json", DataType::Utf8, true)
1074 }
1075 }
1076
1077 #[test]
1078 fn test_decimal_fields() {
1079 let message_type = "
1080 message test_schema {
1081 REQUIRED INT32 decimal1 (DECIMAL(4,2));
1082 REQUIRED INT64 decimal2 (DECIMAL(12,2));
1083 REQUIRED FIXED_LEN_BYTE_ARRAY (16) decimal3 (DECIMAL(30,2));
1084 REQUIRED BYTE_ARRAY decimal4 (DECIMAL(33,2));
1085 REQUIRED BYTE_ARRAY decimal5 (DECIMAL(38,2));
1086 REQUIRED FIXED_LEN_BYTE_ARRAY (17) decimal6 (DECIMAL(39,2));
1087 REQUIRED BYTE_ARRAY decimal7 (DECIMAL(39,2));
1088 }
1089 ";
1090
1091 let parquet_group_type = parse_message_type(message_type).unwrap();
1092
1093 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1094 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1095
1096 let arrow_fields = Fields::from(vec![
1097 Field::new("decimal1", DataType::Decimal128(4, 2), false),
1098 Field::new("decimal2", DataType::Decimal128(12, 2), false),
1099 Field::new("decimal3", DataType::Decimal128(30, 2), false),
1100 Field::new("decimal4", DataType::Decimal128(33, 2), false),
1101 Field::new("decimal5", DataType::Decimal128(38, 2), false),
1102 Field::new("decimal6", DataType::Decimal256(39, 2), false),
1103 Field::new("decimal7", DataType::Decimal256(39, 2), false),
1104 ]);
1105 assert_eq!(&arrow_fields, converted_arrow_schema.fields());
1106 }
1107
1108 #[test]
1109 fn test_decimal_precision_physical_types() {
1110 let cases = [
1111 (1, PhysicalType::INT32, None),
1112 (9, PhysicalType::INT32, None),
1113 (10, PhysicalType::INT64, None),
1114 (18, PhysicalType::INT64, None),
1115 (19, PhysicalType::FIXED_LEN_BYTE_ARRAY, Some(9)),
1116 (38, PhysicalType::FIXED_LEN_BYTE_ARRAY, Some(16)),
1117 ];
1118
1119 for (precision, expected_type, expected_length) in cases {
1120 let arrow_schema = Schema::new(vec![Field::new(
1121 "decimal",
1122 DataType::Decimal128(precision, 0),
1123 false,
1124 )]);
1125 let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema).unwrap();
1126 let column = parquet_schema.column(0);
1127
1128 assert_eq!(
1129 column.physical_type(),
1130 expected_type,
1131 "precision {precision}"
1132 );
1133 if let Some(expected_length) = expected_length {
1134 assert_eq!(
1135 column.type_length(),
1136 expected_length,
1137 "precision {precision}"
1138 );
1139 }
1140 }
1141 }
1142
1143 #[test]
1144 fn test_byte_array_fields() {
1145 let message_type = "
1146 message test_schema {
1147 REQUIRED BYTE_ARRAY binary;
1148 REQUIRED FIXED_LEN_BYTE_ARRAY (20) fixed_binary;
1149 }
1150 ";
1151
1152 let parquet_group_type = parse_message_type(message_type).unwrap();
1153
1154 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1155 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1156
1157 let arrow_fields = Fields::from(vec![
1158 Field::new("binary", DataType::Binary, false),
1159 Field::new("fixed_binary", DataType::FixedSizeBinary(20), false),
1160 ]);
1161 assert_eq!(&arrow_fields, converted_arrow_schema.fields());
1162 }
1163
1164 #[test]
1165 fn test_duplicate_fields() {
1166 let message_type = "
1167 message test_schema {
1168 REQUIRED BOOLEAN boolean;
1169 REQUIRED INT32 int8 (INT_8);
1170 }
1171 ";
1172
1173 let parquet_group_type = parse_message_type(message_type).unwrap();
1174
1175 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1176 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1177
1178 let arrow_fields = Fields::from(vec![
1179 Field::new("boolean", DataType::Boolean, false),
1180 Field::new("int8", DataType::Int8, false),
1181 ]);
1182 assert_eq!(&arrow_fields, converted_arrow_schema.fields());
1183
1184 let converted_arrow_schema =
1185 parquet_to_arrow_schema_by_columns(&parquet_schema, ProjectionMask::all(), None)
1186 .unwrap();
1187 assert_eq!(&arrow_fields, converted_arrow_schema.fields());
1188 }
1189
1190 #[test]
1191 fn test_parquet_lists() {
1192 let mut arrow_fields = Vec::new();
1193
1194 let message_type = "
1196 message test_schema {
1197 REQUIRED GROUP my_list (LIST) {
1198 REPEATED GROUP list {
1199 OPTIONAL BINARY element (UTF8);
1200 }
1201 }
1202 OPTIONAL GROUP my_list (LIST) {
1203 REPEATED GROUP list {
1204 REQUIRED BINARY element (UTF8);
1205 }
1206 }
1207 OPTIONAL GROUP array_of_arrays (LIST) {
1208 REPEATED GROUP list {
1209 REQUIRED GROUP element (LIST) {
1210 REPEATED GROUP list {
1211 REQUIRED INT32 element;
1212 }
1213 }
1214 }
1215 }
1216 OPTIONAL GROUP my_list (LIST) {
1217 REPEATED GROUP element {
1218 REQUIRED BINARY str (UTF8);
1219 }
1220 }
1221 OPTIONAL GROUP my_list (LIST) {
1222 REPEATED INT32 element;
1223 }
1224 OPTIONAL GROUP my_list (LIST) {
1225 REPEATED GROUP element {
1226 REQUIRED BINARY str (UTF8);
1227 REQUIRED INT32 num;
1228 }
1229 }
1230 OPTIONAL GROUP my_list (LIST) {
1231 REPEATED GROUP array {
1232 REQUIRED BINARY str (UTF8);
1233 }
1234
1235 }
1236 OPTIONAL GROUP my_list (LIST) {
1237 REPEATED GROUP my_list_tuple {
1238 REQUIRED BINARY str (UTF8);
1239 }
1240 }
1241 REPEATED INT32 name;
1242 }
1243 ";
1244
1245 {
1252 arrow_fields.push(Field::new_list(
1253 "my_list",
1254 Field::new("element", DataType::Utf8, true),
1255 false,
1256 ));
1257 }
1258
1259 {
1266 arrow_fields.push(Field::new_list(
1267 "my_list",
1268 Field::new("element", DataType::Utf8, false),
1269 true,
1270 ));
1271 }
1272
1273 {
1286 let arrow_inner_list = Field::new("element", DataType::Int32, false);
1287 arrow_fields.push(Field::new_list(
1288 "array_of_arrays",
1289 Field::new_list("element", arrow_inner_list, false),
1290 true,
1291 ));
1292 }
1293
1294 {
1301 arrow_fields.push(Field::new_list(
1302 "my_list",
1303 Field::new("str", DataType::Utf8, false),
1304 true,
1305 ));
1306 }
1307
1308 {
1313 arrow_fields.push(Field::new_list(
1314 "my_list",
1315 Field::new("element", DataType::Int32, false),
1316 true,
1317 ));
1318 }
1319
1320 {
1328 let fields = vec![
1329 Field::new("str", DataType::Utf8, false),
1330 Field::new("num", DataType::Int32, false),
1331 ];
1332 arrow_fields.push(Field::new_list(
1333 "my_list",
1334 Field::new_struct("element", fields, false),
1335 true,
1336 ));
1337 }
1338
1339 {
1347 let fields = vec![Field::new("str", DataType::Utf8, false)];
1348 arrow_fields.push(Field::new_list(
1349 "my_list",
1350 Field::new_struct("array", fields, false),
1351 true,
1352 ));
1353 }
1354
1355 {
1363 let fields = vec![Field::new("str", DataType::Utf8, false)];
1364 arrow_fields.push(Field::new_list(
1365 "my_list",
1366 Field::new_struct("my_list_tuple", fields, false),
1367 true,
1368 ));
1369 }
1370
1371 {
1374 arrow_fields.push(Field::new_list(
1375 "name",
1376 Field::new("name", DataType::Int32, false),
1377 false,
1378 ));
1379 }
1380
1381 let parquet_group_type = parse_message_type(message_type).unwrap();
1382
1383 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1384 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1385 let converted_fields = converted_arrow_schema.fields();
1386
1387 assert_eq!(arrow_fields.len(), converted_fields.len());
1388 for i in 0..arrow_fields.len() {
1389 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref(), "{i}");
1390 }
1391 }
1392
1393 #[test]
1394 fn test_parquet_list_nullable() {
1395 let mut arrow_fields = Vec::new();
1396
1397 let message_type = "
1398 message test_schema {
1399 REQUIRED GROUP my_list1 (LIST) {
1400 REPEATED GROUP list {
1401 OPTIONAL BINARY element (UTF8);
1402 }
1403 }
1404 OPTIONAL GROUP my_list2 (LIST) {
1405 REPEATED GROUP list {
1406 REQUIRED BINARY element (UTF8);
1407 }
1408 }
1409 REQUIRED GROUP my_list3 (LIST) {
1410 REPEATED GROUP list {
1411 REQUIRED BINARY element (UTF8);
1412 }
1413 }
1414 }
1415 ";
1416
1417 {
1424 arrow_fields.push(Field::new_list(
1425 "my_list1",
1426 Field::new("element", DataType::Utf8, true),
1427 false,
1428 ));
1429 }
1430
1431 {
1438 arrow_fields.push(Field::new_list(
1439 "my_list2",
1440 Field::new("element", DataType::Utf8, false),
1441 true,
1442 ));
1443 }
1444
1445 {
1452 arrow_fields.push(Field::new_list(
1453 "my_list3",
1454 Field::new("element", DataType::Utf8, false),
1455 false,
1456 ));
1457 }
1458
1459 let parquet_group_type = parse_message_type(message_type).unwrap();
1460
1461 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1462 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1463 let converted_fields = converted_arrow_schema.fields();
1464
1465 assert_eq!(arrow_fields.len(), converted_fields.len());
1466 for i in 0..arrow_fields.len() {
1467 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1468 }
1469 }
1470
1471 #[test]
1472 fn test_parquet_maps() {
1473 let mut arrow_fields = Vec::new();
1474
1475 let message_type = "
1477 message test_schema {
1478 REQUIRED group my_map1 (MAP) {
1479 REPEATED group key_value {
1480 REQUIRED binary key (UTF8);
1481 OPTIONAL int32 value;
1482 }
1483 }
1484 OPTIONAL group my_map2 (MAP) {
1485 REPEATED group map {
1486 REQUIRED binary str (UTF8);
1487 REQUIRED int32 num;
1488 }
1489 }
1490 OPTIONAL group my_map3 (MAP_KEY_VALUE) {
1491 REPEATED group map {
1492 REQUIRED binary key (UTF8);
1493 OPTIONAL int32 value;
1494 }
1495 }
1496 REQUIRED group my_map4 (MAP) {
1497 REPEATED group map {
1498 OPTIONAL binary key (UTF8);
1499 REQUIRED int32 value;
1500 }
1501 }
1502 }
1503 ";
1504
1505 {
1513 arrow_fields.push(Field::new_map(
1514 "my_map1",
1515 "key_value",
1516 Field::new("key", DataType::Utf8, false),
1517 Field::new("value", DataType::Int32, true),
1518 false,
1519 false,
1520 ));
1521 }
1522
1523 {
1531 arrow_fields.push(Field::new_map(
1532 "my_map2",
1533 "map",
1534 Field::new("str", DataType::Utf8, false),
1535 Field::new("num", DataType::Int32, false),
1536 false,
1537 true,
1538 ));
1539 }
1540
1541 {
1549 arrow_fields.push(Field::new_map(
1550 "my_map3",
1551 "map",
1552 Field::new("key", DataType::Utf8, false),
1553 Field::new("value", DataType::Int32, true),
1554 false,
1555 true,
1556 ));
1557 }
1558
1559 {
1567 arrow_fields.push(Field::new_map(
1568 "my_map4",
1569 "map",
1570 Field::new("key", DataType::Utf8, false), Field::new("value", DataType::Int32, false),
1572 false,
1573 false,
1574 ));
1575 }
1576
1577 let parquet_group_type = parse_message_type(message_type).unwrap();
1578
1579 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1580 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1581 let converted_fields = converted_arrow_schema.fields();
1582
1583 assert_eq!(arrow_fields.len(), converted_fields.len());
1584 for i in 0..arrow_fields.len() {
1585 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1586 }
1587 }
1588
1589 #[test]
1590 fn test_nested_schema() {
1591 let mut arrow_fields = Vec::new();
1592 {
1593 let group1_fields = Fields::from(vec![
1594 Field::new("leaf1", DataType::Boolean, false),
1595 Field::new("leaf2", DataType::Int32, false),
1596 ]);
1597 let group1_struct = Field::new("group1", DataType::Struct(group1_fields), false);
1598 arrow_fields.push(group1_struct);
1599
1600 let leaf3_field = Field::new("leaf3", DataType::Int64, false);
1601 arrow_fields.push(leaf3_field);
1602 }
1603
1604 let message_type = "
1605 message test_schema {
1606 REQUIRED GROUP group1 {
1607 REQUIRED BOOLEAN leaf1;
1608 REQUIRED INT32 leaf2;
1609 }
1610 REQUIRED INT64 leaf3;
1611 }
1612 ";
1613 let parquet_group_type = parse_message_type(message_type).unwrap();
1614
1615 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1616 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1617 let converted_fields = converted_arrow_schema.fields();
1618
1619 assert_eq!(arrow_fields.len(), converted_fields.len());
1620 for i in 0..arrow_fields.len() {
1621 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1622 }
1623 }
1624
1625 #[test]
1626 fn test_nested_schema_partial() {
1627 let mut arrow_fields = Vec::new();
1628 {
1629 let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)].into();
1630 let group1 = Field::new("group1", DataType::Struct(group1_fields), false);
1631 arrow_fields.push(group1);
1632
1633 let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)].into();
1634 let group2 = Field::new("group2", DataType::Struct(group2_fields), false);
1635 arrow_fields.push(group2);
1636
1637 arrow_fields.push(Field::new("leaf5", DataType::Int64, false));
1638 }
1639
1640 let message_type = "
1641 message test_schema {
1642 REQUIRED GROUP group1 {
1643 REQUIRED INT64 leaf1;
1644 REQUIRED INT64 leaf2;
1645 }
1646 REQUIRED GROUP group2 {
1647 REQUIRED INT64 leaf3;
1648 REQUIRED INT64 leaf4;
1649 }
1650 REQUIRED INT64 leaf5;
1651 }
1652 ";
1653 let parquet_group_type = parse_message_type(message_type).unwrap();
1654
1655 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1665 let mask = ProjectionMask::leaves(&parquet_schema, [3, 0, 4, 4]);
1666 let converted_arrow_schema =
1667 parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap();
1668 let converted_fields = converted_arrow_schema.fields();
1669
1670 assert_eq!(arrow_fields.len(), converted_fields.len());
1671 for i in 0..arrow_fields.len() {
1672 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1673 }
1674 }
1675
1676 #[test]
1677 fn test_nested_schema_partial_ordering() {
1678 let mut arrow_fields = Vec::new();
1679 {
1680 let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)].into();
1681 let group1 = Field::new("group1", DataType::Struct(group1_fields), false);
1682 arrow_fields.push(group1);
1683
1684 let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)].into();
1685 let group2 = Field::new("group2", DataType::Struct(group2_fields), false);
1686 arrow_fields.push(group2);
1687
1688 arrow_fields.push(Field::new("leaf5", DataType::Int64, false));
1689 }
1690
1691 let message_type = "
1692 message test_schema {
1693 REQUIRED GROUP group1 {
1694 REQUIRED INT64 leaf1;
1695 REQUIRED INT64 leaf2;
1696 }
1697 REQUIRED GROUP group2 {
1698 REQUIRED INT64 leaf3;
1699 REQUIRED INT64 leaf4;
1700 }
1701 REQUIRED INT64 leaf5;
1702 }
1703 ";
1704 let parquet_group_type = parse_message_type(message_type).unwrap();
1705
1706 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1716 let mask = ProjectionMask::leaves(&parquet_schema, [3, 0, 4]);
1717 let converted_arrow_schema =
1718 parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap();
1719 let converted_fields = converted_arrow_schema.fields();
1720
1721 assert_eq!(arrow_fields.len(), converted_fields.len());
1722 for i in 0..arrow_fields.len() {
1723 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1724 }
1725
1726 let mask =
1727 ProjectionMask::columns(&parquet_schema, ["group2.leaf4", "group1.leaf1", "leaf5"]);
1728 let converted_arrow_schema =
1729 parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap();
1730 let converted_fields = converted_arrow_schema.fields();
1731
1732 assert_eq!(arrow_fields.len(), converted_fields.len());
1733 for i in 0..arrow_fields.len() {
1734 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1735 }
1736 }
1737
1738 #[test]
1739 fn test_repeated_nested_schema() {
1740 let mut arrow_fields = Vec::new();
1741 {
1742 arrow_fields.push(Field::new("leaf1", DataType::Int32, true));
1743
1744 let inner_group_list = Field::new_list(
1745 "innerGroup",
1746 Field::new_struct(
1747 "innerGroup",
1748 vec![Field::new("leaf3", DataType::Int32, true)],
1749 false,
1750 ),
1751 false,
1752 );
1753
1754 let outer_group_list = Field::new_list(
1755 "outerGroup",
1756 Field::new_struct(
1757 "outerGroup",
1758 vec![Field::new("leaf2", DataType::Int32, true), inner_group_list],
1759 false,
1760 ),
1761 false,
1762 );
1763 arrow_fields.push(outer_group_list);
1764 }
1765
1766 let message_type = "
1767 message test_schema {
1768 OPTIONAL INT32 leaf1;
1769 REPEATED GROUP outerGroup {
1770 OPTIONAL INT32 leaf2;
1771 REPEATED GROUP innerGroup {
1772 OPTIONAL INT32 leaf3;
1773 }
1774 }
1775 }
1776 ";
1777 let parquet_group_type = parse_message_type(message_type).unwrap();
1778
1779 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1780 let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap();
1781 let converted_fields = converted_arrow_schema.fields();
1782
1783 assert_eq!(arrow_fields.len(), converted_fields.len());
1784 for i in 0..arrow_fields.len() {
1785 assert_eq!(&arrow_fields[i], converted_fields[i].as_ref());
1786 }
1787 }
1788
1789 #[test]
1790 fn test_column_desc_to_field() {
1791 let message_type = "
1792 message test_schema {
1793 REQUIRED BOOLEAN boolean;
1794 REQUIRED INT32 int8 (INT_8);
1795 REQUIRED INT32 uint8 (INTEGER(8,false));
1796 REQUIRED INT32 int16 (INT_16);
1797 REQUIRED INT32 uint16 (INTEGER(16,false));
1798 REQUIRED INT32 int32;
1799 REQUIRED INT64 int64;
1800 OPTIONAL DOUBLE double;
1801 OPTIONAL FLOAT float;
1802 OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16);
1803 OPTIONAL BINARY string (UTF8);
1804 REPEATED BOOLEAN bools;
1805 OPTIONAL INT32 date (DATE);
1806 OPTIONAL INT32 time_milli (TIME_MILLIS);
1807 OPTIONAL INT64 time_micro (TIME_MICROS);
1808 OPTIONAL INT64 time_nano (TIME(NANOS,false));
1809 OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS);
1810 REQUIRED INT64 ts_micro (TIMESTAMP_MICROS);
1811 REQUIRED INT64 ts_nano (TIMESTAMP(NANOS,true));
1812 REPEATED INT32 int_list;
1813 REPEATED BINARY byte_list;
1814 REPEATED BINARY string_list (UTF8);
1815 REQUIRED INT32 decimal_int32 (DECIMAL(8,2));
1816 REQUIRED INT64 decimal_int64 (DECIMAL(16,2));
1817 REQUIRED FIXED_LEN_BYTE_ARRAY (13) decimal_fix_length (DECIMAL(30,2));
1818 }
1819 ";
1820 let parquet_group_type = parse_message_type(message_type).unwrap();
1821
1822 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1823 let converted_arrow_fields = parquet_schema
1824 .columns()
1825 .iter()
1826 .map(|c| parquet_to_arrow_field(c).unwrap())
1827 .collect::<Vec<Field>>();
1828
1829 let arrow_fields = vec![
1830 Field::new("boolean", DataType::Boolean, false),
1831 Field::new("int8", DataType::Int8, false),
1832 Field::new("uint8", DataType::UInt8, false),
1833 Field::new("int16", DataType::Int16, false),
1834 Field::new("uint16", DataType::UInt16, false),
1835 Field::new("int32", DataType::Int32, false),
1836 Field::new("int64", DataType::Int64, false),
1837 Field::new("double", DataType::Float64, true),
1838 Field::new("float", DataType::Float32, true),
1839 Field::new("float16", DataType::Float16, true),
1840 Field::new("string", DataType::Utf8, true),
1841 Field::new_list(
1842 "bools",
1843 Field::new("bools", DataType::Boolean, false),
1844 false,
1845 ),
1846 Field::new("date", DataType::Date32, true),
1847 Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true),
1848 Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true),
1849 Field::new("time_nano", DataType::Time64(TimeUnit::Nanosecond), true),
1850 Field::new(
1851 "ts_milli",
1852 DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
1853 true,
1854 ),
1855 Field::new(
1856 "ts_micro",
1857 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
1858 false,
1859 ),
1860 Field::new(
1861 "ts_nano",
1862 DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
1863 false,
1864 ),
1865 Field::new_list(
1866 "int_list",
1867 Field::new("int_list", DataType::Int32, false),
1868 false,
1869 ),
1870 Field::new_list(
1871 "byte_list",
1872 Field::new("byte_list", DataType::Binary, false),
1873 false,
1874 ),
1875 Field::new_list(
1876 "string_list",
1877 Field::new("string_list", DataType::Utf8, false),
1878 false,
1879 ),
1880 Field::new("decimal_int32", DataType::Decimal128(8, 2), false),
1881 Field::new("decimal_int64", DataType::Decimal128(16, 2), false),
1882 Field::new("decimal_fix_length", DataType::Decimal128(30, 2), false),
1883 ];
1884
1885 assert_eq!(arrow_fields, converted_arrow_fields);
1886 }
1887
1888 #[test]
1889 fn test_coerced_map_list() {
1890 let arrow_fields = vec![
1892 Field::new_list(
1893 "my_list",
1894 Field::new("item", DataType::Boolean, true),
1895 false,
1896 ),
1897 Field::new_map(
1898 "my_map",
1899 "my_entries",
1900 Field::new("my_keys", DataType::Utf8, false),
1901 Field::new("my_values", DataType::Int32, true),
1902 false,
1903 true,
1904 ),
1905 ];
1906 let arrow_schema = Schema::new(arrow_fields);
1907
1908 let message_type = "
1910 message parquet_schema {
1911 REQUIRED GROUP my_list (LIST) {
1912 REPEATED GROUP list {
1913 OPTIONAL BOOLEAN element;
1914 }
1915 }
1916 OPTIONAL GROUP my_map (MAP) {
1917 REPEATED GROUP key_value {
1918 REQUIRED BINARY key (STRING);
1919 OPTIONAL INT32 value;
1920 }
1921 }
1922 }
1923 ";
1924 let parquet_group_type = parse_message_type(message_type).unwrap();
1925 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1926 let converted_arrow_schema = ArrowSchemaConverter::new()
1927 .with_coerce_types(true)
1928 .convert(&arrow_schema)
1929 .unwrap();
1930 assert_eq!(
1931 parquet_schema.columns().len(),
1932 converted_arrow_schema.columns().len()
1933 );
1934
1935 let message_type = "
1937 message parquet_schema {
1938 REQUIRED GROUP my_list (LIST) {
1939 REPEATED GROUP list {
1940 OPTIONAL BOOLEAN item;
1941 }
1942 }
1943 OPTIONAL GROUP my_map (MAP) {
1944 REPEATED GROUP my_entries {
1945 REQUIRED BINARY my_keys (STRING);
1946 OPTIONAL INT32 my_values;
1947 }
1948 }
1949 }
1950 ";
1951 let parquet_group_type = parse_message_type(message_type).unwrap();
1952 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
1953 let converted_arrow_schema = ArrowSchemaConverter::new()
1954 .with_coerce_types(false)
1955 .convert(&arrow_schema)
1956 .unwrap();
1957 assert_eq!(
1958 parquet_schema.columns().len(),
1959 converted_arrow_schema.columns().len()
1960 );
1961 }
1962
1963 #[test]
1964 fn test_field_to_column_desc() {
1965 let message_type = "
1966 message arrow_schema {
1967 REQUIRED BOOLEAN boolean;
1968 REQUIRED INT32 int8 (INT_8);
1969 REQUIRED INT32 int16 (INTEGER(16,true));
1970 REQUIRED INT32 int32;
1971 REQUIRED INT64 int64;
1972 OPTIONAL DOUBLE double;
1973 OPTIONAL FLOAT float;
1974 OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16);
1975 OPTIONAL BINARY string (STRING);
1976 OPTIONAL GROUP bools (LIST) {
1977 REPEATED GROUP list {
1978 OPTIONAL BOOLEAN element;
1979 }
1980 }
1981 REQUIRED GROUP bools_non_null (LIST) {
1982 REPEATED GROUP list {
1983 REQUIRED BOOLEAN element;
1984 }
1985 }
1986 OPTIONAL INT32 date (DATE);
1987 OPTIONAL INT32 time_milli (TIME(MILLIS,false));
1988 OPTIONAL INT32 time_milli_utc (TIME(MILLIS,true));
1989 OPTIONAL INT64 time_micro (TIME_MICROS);
1990 OPTIONAL INT64 time_micro_utc (TIME(MICROS, true));
1991 OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS);
1992 REQUIRED INT64 ts_micro (TIMESTAMP(MICROS,false));
1993 REQUIRED INT64 ts_seconds;
1994 REQUIRED INT64 ts_micro_utc (TIMESTAMP(MICROS, true));
1995 REQUIRED INT64 ts_millis_zero_offset (TIMESTAMP(MILLIS, true));
1996 REQUIRED INT64 ts_millis_zero_negative_offset (TIMESTAMP(MILLIS, true));
1997 REQUIRED INT64 ts_micro_non_utc (TIMESTAMP(MICROS, true));
1998 REQUIRED GROUP struct {
1999 REQUIRED BOOLEAN bools;
2000 REQUIRED INT32 uint32 (INTEGER(32,false));
2001 REQUIRED GROUP int32 (LIST) {
2002 REPEATED GROUP list {
2003 OPTIONAL INT32 element;
2004 }
2005 }
2006 }
2007 REQUIRED BINARY dictionary_strings (STRING);
2008 REQUIRED INT32 decimal_int32 (DECIMAL(8,2));
2009 REQUIRED INT64 decimal_int64 (DECIMAL(16,2));
2010 REQUIRED FIXED_LEN_BYTE_ARRAY (13) decimal_fix_length (DECIMAL(30,2));
2011 REQUIRED FIXED_LEN_BYTE_ARRAY (16) decimal128 (DECIMAL(38,2));
2012 REQUIRED FIXED_LEN_BYTE_ARRAY (17) decimal256 (DECIMAL(39,2));
2013 }
2014 ";
2015 let parquet_group_type = parse_message_type(message_type).unwrap();
2016
2017 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
2018
2019 let arrow_fields = vec![
2020 Field::new("boolean", DataType::Boolean, false),
2021 Field::new("int8", DataType::Int8, false),
2022 Field::new("int16", DataType::Int16, false),
2023 Field::new("int32", DataType::Int32, false),
2024 Field::new("int64", DataType::Int64, false),
2025 Field::new("double", DataType::Float64, true),
2026 Field::new("float", DataType::Float32, true),
2027 Field::new("float16", DataType::Float16, true),
2028 Field::new("string", DataType::Utf8, true),
2029 Field::new_list(
2030 "bools",
2031 Field::new("element", DataType::Boolean, true),
2032 true,
2033 ),
2034 Field::new_list(
2035 "bools_non_null",
2036 Field::new("element", DataType::Boolean, false),
2037 false,
2038 ),
2039 Field::new("date", DataType::Date32, true),
2040 Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true),
2041 Field::new(
2042 "time_milli_utc",
2043 DataType::Time32(TimeUnit::Millisecond),
2044 true,
2045 )
2046 .with_metadata(HashMap::from_iter(vec![(
2047 "adjusted_to_utc".to_string(),
2048 String::new(),
2049 )])),
2050 Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true),
2051 Field::new(
2052 "time_micro_utc",
2053 DataType::Time64(TimeUnit::Microsecond),
2054 true,
2055 )
2056 .with_metadata(HashMap::from_iter(vec![(
2057 "adjusted_to_utc".to_string(),
2058 String::new(),
2059 )])),
2060 Field::new(
2061 "ts_milli",
2062 DataType::Timestamp(TimeUnit::Millisecond, None),
2063 true,
2064 ),
2065 Field::new(
2066 "ts_micro",
2067 DataType::Timestamp(TimeUnit::Microsecond, None),
2068 false,
2069 ),
2070 Field::new(
2071 "ts_seconds",
2072 DataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
2073 false,
2074 ),
2075 Field::new(
2076 "ts_micro_utc",
2077 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
2078 false,
2079 ),
2080 Field::new(
2081 "ts_millis_zero_offset",
2082 DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
2083 false,
2084 ),
2085 Field::new(
2086 "ts_millis_zero_negative_offset",
2087 DataType::Timestamp(TimeUnit::Millisecond, Some("-00:00".into())),
2088 false,
2089 ),
2090 Field::new(
2091 "ts_micro_non_utc",
2092 DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())),
2093 false,
2094 ),
2095 Field::new_struct(
2096 "struct",
2097 vec![
2098 Field::new("bools", DataType::Boolean, false),
2099 Field::new("uint32", DataType::UInt32, false),
2100 Field::new_list("int32", Field::new("element", DataType::Int32, true), false),
2101 ],
2102 false,
2103 ),
2104 Field::new_dictionary("dictionary_strings", DataType::Int32, DataType::Utf8, false),
2105 Field::new("decimal_int32", DataType::Decimal128(8, 2), false),
2106 Field::new("decimal_int64", DataType::Decimal128(16, 2), false),
2107 Field::new("decimal_fix_length", DataType::Decimal128(30, 2), false),
2108 Field::new("decimal128", DataType::Decimal128(38, 2), false),
2109 Field::new("decimal256", DataType::Decimal256(39, 2), false),
2110 ];
2111 let arrow_schema = Schema::new(arrow_fields);
2112 let converted_arrow_schema = ArrowSchemaConverter::new().convert(&arrow_schema).unwrap();
2113
2114 assert_eq!(
2115 parquet_schema.columns().len(),
2116 converted_arrow_schema.columns().len()
2117 );
2118 parquet_schema
2119 .columns()
2120 .iter()
2121 .zip(converted_arrow_schema.columns())
2122 .for_each(|(a, b)| {
2123 match a.logical_type_ref() {
2128 Some(_) => {
2129 assert_eq!(a, b)
2130 }
2131 None => {
2132 assert_eq!(a.name(), b.name());
2133 assert_eq!(a.physical_type(), b.physical_type());
2134 assert_eq!(a.converted_type(), b.converted_type());
2135 }
2136 }
2137 });
2138 }
2139
2140 #[test]
2141 #[should_panic(expected = "Parquet does not support writing empty structs")]
2142 fn test_empty_struct_field() {
2143 let arrow_fields = vec![Field::new(
2144 "struct",
2145 DataType::Struct(Fields::empty()),
2146 false,
2147 )];
2148 let arrow_schema = Schema::new(arrow_fields);
2149 let converted_arrow_schema = ArrowSchemaConverter::new()
2150 .with_coerce_types(true)
2151 .convert(&arrow_schema);
2152
2153 converted_arrow_schema.unwrap();
2154 }
2155
2156 #[test]
2157 fn test_run_end_schema_conversion_recurses_into_value_type() {
2158 fn run_end_type(value_type: DataType, nullable_values: bool) -> DataType {
2159 DataType::RunEndEncoded(
2160 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2161 Arc::new(Field::new("values", value_type, nullable_values)),
2162 )
2163 }
2164
2165 let leaf_schema = Schema::new(vec![Field::new(
2168 "ree",
2169 run_end_type(DataType::Int32, true),
2170 false,
2171 )]);
2172 let converted = ArrowSchemaConverter::new().convert(&leaf_schema).unwrap();
2173 assert_eq!(converted.columns().len(), 1);
2174 assert_eq!(converted.column(0).name(), "ree");
2175 assert_eq!(converted.column(0).physical_type(), PhysicalType::INT32);
2176 assert_eq!(converted.column(0).max_def_level(), 1);
2177
2178 let list_schema = Schema::new(vec![Field::new(
2181 "ree_list",
2182 run_end_type(
2183 DataType::List(Arc::new(Field::new("element", DataType::Int32, true))),
2184 true,
2185 ),
2186 true,
2187 )]);
2188 let converted = ArrowSchemaConverter::new().convert(&list_schema).unwrap();
2189 assert_eq!(converted.columns().len(), 1);
2190 assert_eq!(converted.column(0).physical_type(), PhysicalType::INT32);
2191 assert!(
2192 converted.column(0).max_rep_level() >= 1,
2193 "list element should be repeated"
2194 );
2195
2196 let dictionary_schema = Schema::new(vec![Field::new(
2197 "ree_dictionary",
2198 run_end_type(
2199 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2200 true,
2201 ),
2202 true,
2203 )]);
2204 let converted = ArrowSchemaConverter::new()
2205 .convert(&dictionary_schema)
2206 .unwrap();
2207 assert_eq!(converted.columns().len(), 1);
2208 assert_eq!(
2209 converted.column(0).physical_type(),
2210 PhysicalType::BYTE_ARRAY
2211 );
2212 }
2213
2214 fn map_type_with_key(key: Field) -> DataType {
2215 DataType::Map(
2216 Arc::new(Field::new(
2217 "entries",
2218 DataType::Struct(vec![key, Field::new("value", DataType::Int32, true)].into()),
2219 false,
2220 )),
2221 false,
2222 )
2223 }
2224
2225 #[test]
2226 fn test_map_key_schema_rejects_logically_nullable_keys() {
2227 let ree = DataType::RunEndEncoded(
2228 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2229 Arc::new(Field::new("values", DataType::Int32, true)),
2230 );
2231 let dictionary_ree = DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree.clone()));
2232 let nested_ree = DataType::RunEndEncoded(
2233 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2234 Arc::new(Field::new("values", dictionary_ree.clone(), false)),
2235 );
2236 for key in [
2237 Field::new("key", DataType::Int32, true),
2238 Field::new("key", ree, false),
2239 Field::new("key", dictionary_ree, false),
2240 Field::new("key", nested_ree, false),
2241 ] {
2242 let map = map_type_with_key(key);
2243 for data_type in [
2244 map.clone(),
2245 DataType::Struct(vec![Field::new("nested", map.clone(), true)].into()),
2246 DataType::List(Arc::new(Field::new("item", map.clone(), true))),
2247 DataType::Dictionary(Box::new(DataType::Int8), Box::new(map.clone())),
2248 DataType::RunEndEncoded(
2249 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2250 Arc::new(Field::new("values", map, true)),
2251 ),
2252 ] {
2253 let schema = Schema::new(vec![Field::new("map", data_type, true)]);
2254 for coerce_types in [false, true] {
2255 let err = ArrowSchemaConverter::new()
2256 .with_coerce_types(coerce_types)
2257 .convert(&schema)
2258 .unwrap_err();
2259 assert!(
2260 err.to_string()
2261 .contains("Map key field 'key' must be non-nullable"),
2262 "{err}"
2263 );
2264 }
2265 }
2266 }
2267 }
2268
2269 #[test]
2270 fn test_map_key_schema_preserves_required_keys_and_nullable_members() {
2271 for key_type in [
2272 DataType::Int32,
2273 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2274 DataType::RunEndEncoded(
2275 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2276 Arc::new(Field::new("values", DataType::Int32, false)),
2277 ),
2278 DataType::Struct(vec![Field::new("member", DataType::Int32, true)].into()),
2279 ] {
2280 let schema = Schema::new(vec![Field::new(
2281 "map",
2282 map_type_with_key(Field::new("key", key_type, false)),
2283 true,
2284 )]);
2285 let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
2286 let key = &parquet.root_schema().get_fields()[0].get_fields()[0].get_fields()[0];
2287 assert_eq!(key.get_basic_info().repetition(), Repetition::REQUIRED);
2288
2289 let mut props = WriterProperties::builder().build();
2290 add_encoded_arrow_schema_to_metadata(&schema, &mut props);
2291 let encoded = props
2292 .key_value_metadata()
2293 .unwrap()
2294 .iter()
2295 .find(|kv| kv.key == crate::arrow::ARROW_SCHEMA_META_KEY)
2296 .unwrap()
2297 .value
2298 .as_deref()
2299 .unwrap();
2300 let decoded = get_arrow_schema_from_metadata(encoded).unwrap();
2301 let DataType::Map(entries, _) = decoded.field(0).data_type() else {
2302 unreachable!()
2303 };
2304 let DataType::Struct(fields) = entries.data_type() else {
2305 unreachable!()
2306 };
2307 assert!(!fields[0].is_nullable());
2308 assert!(fields[1].is_nullable());
2309 if let DataType::Struct(members) = fields[0].data_type() {
2310 assert!(members[0].is_nullable());
2311 }
2312 }
2313 }
2314
2315 #[test]
2316 fn test_arrow_metadata_normalization_merges_ree_value_nullability() {
2317 let run_ends = Arc::new(Field::new("run_ends", DataType::Int32, false));
2318 let inner = Arc::new(Field::new("inner", DataType::Int32, true));
2319 let nested = Arc::new(Field::new(
2320 "nested",
2321 DataType::RunEndEncoded(Arc::clone(&run_ends), inner),
2322 false,
2323 ));
2324 let field = Field::new("outer", DataType::RunEndEncoded(run_ends, nested), false);
2325
2326 let normalized = normalize_metadata_field(&field);
2327 assert_eq!(normalized.data_type(), &DataType::Int32);
2328 assert!(normalized.is_nullable());
2329 }
2330
2331 #[test]
2332 fn test_arrow_metadata_normalization_descends_dictionary_values() {
2333 let ree = DataType::RunEndEncoded(
2334 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2335 Arc::new(Field::new("values", DataType::Int64, true)),
2336 );
2337 let field = Field::new(
2338 "dictionary_ree",
2339 DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
2340 false,
2341 );
2342
2343 assert!(needs_metadata_normalization(field.data_type()));
2344 let normalized = normalize_metadata_field(&field);
2345 assert_eq!(
2346 normalized.data_type(),
2347 &DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64))
2348 );
2349 assert!(normalized.is_nullable());
2350 assert!(!needs_metadata_normalization(normalized.data_type()));
2351
2352 let mut props = WriterProperties::builder().build();
2353 add_encoded_arrow_schema_to_metadata(&Schema::new(vec![field]), &mut props);
2354 let encoded = props
2355 .key_value_metadata()
2356 .unwrap()
2357 .iter()
2358 .find(|kv| kv.key == crate::arrow::ARROW_SCHEMA_META_KEY)
2359 .unwrap()
2360 .value
2361 .as_deref()
2362 .unwrap();
2363 let decoded = get_arrow_schema_from_metadata(encoded).unwrap();
2364 assert_eq!(decoded.field(0), &normalized);
2365 }
2366
2367 #[test]
2368 fn test_arrow_metadata_normalization_preserves_view_dictionary_siblings() {
2369 let dictionary = |dt| DataType::Dictionary(Box::new(DataType::Int8), Box::new(dt));
2370 let ree = DataType::RunEndEncoded(
2371 Arc::new(Field::new(
2372 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2373 DataType::Int32,
2374 false,
2375 )),
2376 Arc::new(Field::new(
2377 Field::REE_VALUES_FIELD_DEFAULT_NAME,
2378 DataType::Int32,
2379 true,
2380 )),
2381 );
2382 let fields = vec![
2383 Field::new(
2385 "nested",
2386 dictionary(DataType::Struct(
2387 vec![Field::new("child", ree, false)].into(),
2388 )),
2389 false,
2390 ),
2391 Field::new("view", dictionary(DataType::Utf8View), true),
2394 Field::new("ordinary", dictionary(DataType::Utf8), true),
2395 Field::new("binary_view", dictionary(DataType::BinaryView), true),
2396 ];
2397 let schema = Schema::new(fields);
2398 assert!(needs_metadata_normalization(schema.field(0).data_type()));
2399 let mut props = WriterProperties::builder().build();
2400 add_encoded_arrow_schema_to_metadata(&schema, &mut props);
2401 let encoded = props
2402 .key_value_metadata()
2403 .unwrap()
2404 .iter()
2405 .find(|kv| kv.key == crate::arrow::ARROW_SCHEMA_META_KEY)
2406 .unwrap()
2407 .value
2408 .as_deref()
2409 .unwrap();
2410 let decoded = get_arrow_schema_from_metadata(encoded).unwrap();
2411 let DataType::Struct(children) = decoded.field(0).data_type() else {
2412 panic!("nested dictionary not erased")
2413 };
2414 assert_eq!(children[0].name(), "child");
2415 assert_eq!(children[0].data_type(), &DataType::Int32);
2416 assert!(children[0].is_nullable());
2417 assert!(!decoded.field(0).is_nullable());
2418 assert_eq!(decoded.field(1), schema.field(1));
2419 assert_eq!(decoded.field(2), schema.field(2));
2420 assert_eq!(decoded.field(3), schema.field(3));
2421 assert!(!needs_metadata_normalization(schema.field(1).data_type()));
2422 assert!(!needs_metadata_normalization(schema.field(3).data_type()));
2423 let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
2424 assert_eq!(parquet.column(0).path().parts(), &["nested", "child"]);
2425 assert_eq!(parquet.column(0).max_def_level(), 1);
2426 }
2427
2428 #[test]
2429 fn test_metadata() {
2430 let message_type = "
2431 message test_schema {
2432 OPTIONAL BINARY string (STRING);
2433 }
2434 ";
2435 let parquet_group_type = parse_message_type(message_type).unwrap();
2436
2437 let key_value_metadata = vec![
2438 KeyValue::new("foo".to_owned(), Some("bar".to_owned())),
2439 KeyValue::new("baz".to_owned(), None),
2440 ];
2441
2442 let mut expected_metadata: HashMap<String, String> = HashMap::new();
2443 expected_metadata.insert("foo".to_owned(), "bar".to_owned());
2444
2445 let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
2446 let converted_arrow_schema =
2447 parquet_to_arrow_schema(&parquet_schema, Some(&key_value_metadata)).unwrap();
2448
2449 assert_eq!(converted_arrow_schema.metadata(), &expected_metadata);
2450 }
2451
2452 #[test]
2453 fn test_arrow_schema_roundtrip() -> Result<()> {
2454 let meta = |a: &[(&str, &str)]| -> HashMap<String, String> {
2455 a.iter()
2456 .map(|(a, b)| (a.to_string(), b.to_string()))
2457 .collect()
2458 };
2459
2460 let schema = Schema::new_with_metadata(
2461 vec![
2462 Field::new("c1", DataType::Utf8, false)
2463 .with_metadata(meta(&[("Key", "Foo"), (PARQUET_FIELD_ID_META_KEY, "2")])),
2464 Field::new("c2", DataType::Binary, false),
2465 Field::new("c3", DataType::FixedSizeBinary(3), false),
2466 Field::new("c4", DataType::Boolean, false),
2467 Field::new("c5", DataType::Date32, false),
2468 Field::new("c6", DataType::Date64, false),
2469 Field::new("c7", DataType::Time32(TimeUnit::Second), false),
2470 Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false),
2471 Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false),
2472 Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false),
2473 Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
2474 Field::new(
2475 "c16",
2476 DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
2477 false,
2478 ),
2479 Field::new(
2480 "c17",
2481 DataType::Timestamp(TimeUnit::Microsecond, Some("Africa/Johannesburg".into())),
2482 false,
2483 ),
2484 Field::new(
2485 "c18",
2486 DataType::Timestamp(TimeUnit::Nanosecond, None),
2487 false,
2488 ),
2489 Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
2490 Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
2491 Field::new_list(
2492 "c21",
2493 Field::new_list_field(DataType::Boolean, true)
2494 .with_metadata(meta(&[("Key", "Bar"), (PARQUET_FIELD_ID_META_KEY, "5")])),
2495 false,
2496 )
2497 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "4")])),
2498 Field::new(
2499 "c22",
2500 DataType::FixedSizeList(
2501 Arc::new(Field::new_list_field(DataType::Boolean, true)),
2502 5,
2503 ),
2504 false,
2505 ),
2506 Field::new_list(
2507 "c23",
2508 Field::new_large_list(
2509 "inner",
2510 Field::new_list_field(
2511 DataType::Struct(
2512 vec![
2513 Field::new("a", DataType::Int16, true),
2514 Field::new("b", DataType::Float64, false),
2515 Field::new("c", DataType::Float32, false),
2516 Field::new("d", DataType::Float16, false),
2517 ]
2518 .into(),
2519 ),
2520 false,
2521 ),
2522 true,
2523 ),
2524 false,
2525 ),
2526 Field::new(
2527 "c24",
2528 DataType::Struct(Fields::from(vec![
2529 Field::new("a", DataType::Utf8, false),
2530 Field::new("b", DataType::UInt16, false),
2531 ])),
2532 false,
2533 ),
2534 Field::new("c25", DataType::Interval(IntervalUnit::YearMonth), true),
2535 Field::new("c26", DataType::Interval(IntervalUnit::DayTime), true),
2536 #[expect(deprecated)]
2542 Field::new_dict(
2543 "c31",
2544 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2545 true,
2546 123,
2547 true,
2548 )
2549 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "6")])),
2550 Field::new("c32", DataType::LargeBinary, true),
2551 Field::new("c33", DataType::LargeUtf8, true),
2552 Field::new_large_list(
2553 "c34",
2554 Field::new_list(
2555 "inner",
2556 Field::new_list_field(
2557 DataType::Struct(
2558 vec![
2559 Field::new("a", DataType::Int16, true),
2560 Field::new("b", DataType::Float64, true),
2561 ]
2562 .into(),
2563 ),
2564 true,
2565 ),
2566 true,
2567 ),
2568 true,
2569 ),
2570 Field::new("c35", DataType::Null, true),
2571 Field::new("c36", DataType::Decimal128(2, 1), false),
2572 Field::new("c37", DataType::Decimal256(50, 20), false),
2573 Field::new("c38", DataType::Decimal128(18, 12), true),
2574 Field::new_map(
2575 "c39",
2576 "key_value",
2577 Field::new("key", DataType::Utf8, false),
2578 Field::new_list("value", Field::new("element", DataType::Utf8, true), true),
2579 false, true,
2581 ),
2582 Field::new_map(
2583 "c40",
2584 "my_entries",
2585 Field::new("my_key", DataType::Utf8, false)
2586 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "8")])),
2587 Field::new_list(
2588 "my_value",
2589 Field::new_list_field(DataType::Utf8, true)
2590 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "10")])),
2591 true,
2592 )
2593 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "9")])),
2594 false, true,
2596 )
2597 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "7")])),
2598 Field::new_map(
2599 "c41",
2600 "my_entries",
2601 Field::new("my_key", DataType::Utf8, false),
2602 Field::new_list(
2603 "my_value",
2604 Field::new_list_field(DataType::Utf8, true)
2605 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "11")])),
2606 true,
2607 ),
2608 false, false,
2610 ),
2611 Field::new("c42", DataType::Decimal32(5, 2), false),
2612 Field::new("c43", DataType::Decimal64(18, 12), true),
2613 ],
2614 meta(&[("Key", "Value")]),
2615 );
2616
2617 let file = tempfile::tempfile().unwrap();
2619 let writer =
2620 ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema.clone()), None)?;
2621 writer.close()?;
2622
2623 let arrow_reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
2625
2626 let read_schema = arrow_reader.schema();
2628 assert_eq!(&schema, read_schema.as_ref());
2629
2630 let mut stack = Vec::with_capacity(10);
2632 let mut out = Vec::with_capacity(10);
2633
2634 let root = arrow_reader.parquet_schema().root_schema_ptr();
2635 stack.push((root.name().to_string(), root));
2636
2637 while let Some((p, t)) = stack.pop() {
2638 if t.is_group() {
2639 for f in t.get_fields() {
2640 stack.push((format!("{p}.{}", f.name()), f.clone()))
2641 }
2642 }
2643
2644 let info = t.get_basic_info();
2645 if info.has_id() {
2646 out.push(format!("{p} -> {}", info.id()))
2647 }
2648 }
2649 out.sort_unstable();
2650 let out: Vec<_> = out.iter().map(|x| x.as_str()).collect();
2651
2652 assert_eq!(
2653 &out,
2654 &[
2655 "arrow_schema.c1 -> 2",
2656 "arrow_schema.c21 -> 4",
2657 "arrow_schema.c21.list.item -> 5",
2658 "arrow_schema.c31 -> 6",
2659 "arrow_schema.c40 -> 7",
2660 "arrow_schema.c40.my_entries.my_key -> 8",
2661 "arrow_schema.c40.my_entries.my_value -> 9",
2662 "arrow_schema.c40.my_entries.my_value.list.item -> 10",
2663 "arrow_schema.c41.my_entries.my_value.list.item -> 11",
2664 ]
2665 );
2666
2667 Ok(())
2668 }
2669
2670 #[test]
2671 fn test_read_parquet_field_ids_raw() -> Result<()> {
2672 let meta = |a: &[(&str, &str)]| -> HashMap<String, String> {
2673 a.iter()
2674 .map(|(a, b)| (a.to_string(), b.to_string()))
2675 .collect()
2676 };
2677 let schema = Schema::new_with_metadata(
2678 vec![
2679 Field::new("c1", DataType::Utf8, true)
2680 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "1")])),
2681 Field::new("c2", DataType::Utf8, true)
2682 .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "2")])),
2683 ],
2684 HashMap::new(),
2685 );
2686
2687 let writer = ArrowWriter::try_new(vec![], Arc::new(schema.clone()), None)?;
2688 let parquet_bytes = writer.into_inner()?;
2689
2690 let reader =
2691 crate::file::reader::SerializedFileReader::new(bytes::Bytes::from(parquet_bytes))?;
2692 let schema_descriptor = reader.metadata().file_metadata().schema_descr_ptr();
2693
2694 let arrow_schema = crate::arrow::parquet_to_arrow_schema(&schema_descriptor, None)?;
2696
2697 let parq_schema_descr = ArrowSchemaConverter::new()
2698 .with_coerce_types(true)
2699 .convert(&arrow_schema)?;
2700 let parq_fields = parq_schema_descr.root_schema().get_fields();
2701 assert_eq!(parq_fields.len(), 2);
2702 assert_eq!(parq_fields[0].get_basic_info().id(), 1);
2703 assert_eq!(parq_fields[1].get_basic_info().id(), 2);
2704
2705 Ok(())
2706 }
2707
2708 #[test]
2709 fn test_arrow_schema_roundtrip_lists() -> Result<()> {
2710 let metadata = HashMap::from([("Key".to_string(), "Value".to_string())]);
2711
2712 let schema = Schema::new_with_metadata(
2713 vec![
2714 Field::new_list("c21", Field::new("array", DataType::Boolean, true), false),
2715 Field::new(
2716 "c22",
2717 DataType::FixedSizeList(
2718 Arc::new(Field::new("items", DataType::Boolean, false)),
2719 5,
2720 ),
2721 false,
2722 ),
2723 Field::new_list(
2724 "c23",
2725 Field::new_large_list(
2726 "items",
2727 Field::new_struct(
2728 "items",
2729 vec![
2730 Field::new("a", DataType::Int16, true),
2731 Field::new("b", DataType::Float64, false),
2732 ],
2733 true,
2734 ),
2735 true,
2736 ),
2737 true,
2738 ),
2739 ],
2740 metadata,
2741 );
2742
2743 let file = tempfile::tempfile().unwrap();
2745 let writer =
2746 ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema.clone()), None)?;
2747 writer.close()?;
2748
2749 let arrow_reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
2751 let read_schema = arrow_reader.schema();
2752 assert_eq!(&schema, read_schema.as_ref());
2753 Ok(())
2754 }
2755
2756 #[test]
2757 fn test_get_arrow_schema_from_metadata() {
2758 assert!(get_arrow_schema_from_metadata("").is_err());
2759 }
2760
2761 #[test]
2762 #[cfg(feature = "arrow_canonical_extension_types")]
2763 fn arrow_uuid_to_parquet_uuid() -> Result<()> {
2764 use arrow_schema::extension::Uuid;
2765 let arrow_schema = Schema::new(vec![
2766 Field::new("uuid", DataType::FixedSizeBinary(16), false).with_extension_type(Uuid),
2767 ]);
2768
2769 let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?;
2770
2771 assert_eq!(
2772 parquet_schema.column(0).logical_type_ref(),
2773 Some(&LogicalType::Uuid)
2774 );
2775
2776 let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?;
2777 assert_eq!(arrow_schema.field(0).try_extension_type::<Uuid>()?, Uuid);
2778
2779 Ok(())
2780 }
2781
2782 #[test]
2783 #[cfg(feature = "arrow_canonical_extension_types")]
2784 fn arrow_json_to_parquet_json() -> Result<()> {
2785 use arrow_schema::extension::Json;
2786 let arrow_schema = Schema::new(vec![
2787 Field::new("json", DataType::Utf8, false).with_extension_type(Json::default()),
2788 ]);
2789
2790 let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?;
2791
2792 assert_eq!(
2793 parquet_schema.column(0).logical_type_ref(),
2794 Some(&LogicalType::Json)
2795 );
2796
2797 let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?;
2798 assert_eq!(
2799 arrow_schema.field(0).try_extension_type::<Json>()?,
2800 Json::default()
2801 );
2802
2803 Ok(())
2804 }
2805
2806 #[test]
2807 #[cfg(feature = "variant_experimental")]
2808 fn arrow_variant_to_parquet_variant() -> Result<()> {
2809 use parquet_variant_compute::VariantType;
2810 let variant_fields = Fields::from(vec![
2811 Field::new("metadata", DataType::BinaryView, false),
2812 Field::new("value", DataType::BinaryView, true),
2813 ]);
2814 let arrow_schema = Schema::new(vec![
2815 Field::new("v", DataType::Struct(variant_fields), true)
2816 .with_extension_type(VariantType),
2817 ]);
2818
2819 let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?;
2820
2821 let root = parquet_schema.root_schema();
2824 assert_eq!(
2825 root.get_fields()[0].get_basic_info().logical_type_ref(),
2826 Some(&LogicalType::variant(Some(1)))
2827 );
2828
2829 let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?;
2830 assert!(
2831 arrow_schema
2832 .field(0)
2833 .has_valid_extension_type::<VariantType>()
2834 );
2835
2836 Ok(())
2837 }
2838
2839 #[test]
2840 fn test_parquet_to_arrow_field_levels_with_virtual_rejects_non_virtual() {
2841 let message_type = "
2842 message test_schema {
2843 REQUIRED INT32 id;
2844 }
2845 ";
2846 let parquet_schema = Arc::new(parse_message_type(message_type).unwrap());
2847 let descriptor = SchemaDescriptor::new(parquet_schema);
2848
2849 let regular_field = Arc::new(Field::new("regular_column", DataType::Int64, false));
2851 let result = parquet_to_arrow_field_levels_with_virtual(
2852 &descriptor,
2853 ProjectionMask::all(),
2854 None,
2855 &[regular_field],
2856 );
2857
2858 assert!(result.is_err());
2859 assert!(
2860 result
2861 .unwrap_err()
2862 .to_string()
2863 .contains("is not a virtual column")
2864 );
2865 }
2866}