1use std::collections::HashMap;
21use std::sync::Arc;
22
23use arrow_array::types::{Decimal128Type, validate_decimal_precision_and_scale};
24use arrow_array::{
25 BinaryArray, BooleanArray, Date32Array, Datum as ArrowDatum, Decimal128Array,
26 FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array, Int64Array, Scalar, StringArray,
27 TimestampMicrosecondArray, TimestampNanosecondArray,
28};
29use arrow_schema::extension::ExtensionType;
30use arrow_schema::{
31 ArrowError, DataType, Field, FieldRef, Fields, Schema as ArrowSchema, TimeUnit,
32};
33use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
34use parquet::file::statistics::Statistics;
35use uuid::Uuid;
36
37use crate::error::Result;
38use crate::spec::decimal_utils::i128_from_be_bytes;
39use crate::spec::{
40 Datum, FIRST_FIELD_ID, ListType, MapType, NestedField, NestedFieldRef, PrimitiveLiteral,
41 PrimitiveType, Schema, SchemaVisitor, StructType, Type, VariantType,
42};
43use crate::{Error, ErrorKind};
44
45pub const DEFAULT_MAP_FIELD_NAME: &str = "key_value";
47pub const UTC_TIME_ZONE: &str = "+00:00";
49
50#[derive(Debug, Clone, Copy, Default)]
59pub(crate) struct VariantExtensionType;
60
61impl ExtensionType for VariantExtensionType {
62 const NAME: &'static str = "arrow.parquet.variant";
63
64 type Metadata = ();
65
66 fn metadata(&self) -> &Self::Metadata {
67 &()
68 }
69
70 fn serialize_metadata(&self) -> Option<String> {
71 None
72 }
73
74 fn deserialize_metadata(
75 metadata: Option<&str>,
76 ) -> std::result::Result<Self::Metadata, ArrowError> {
77 match metadata {
78 None | Some("") => Ok(()),
79 Some(other) => Err(ArrowError::InvalidArgumentError(format!(
80 "arrow.parquet.variant extension type takes no metadata, got {other:?}"
81 ))),
82 }
83 }
84
85 fn supports_data_type(&self, data_type: &DataType) -> std::result::Result<(), ArrowError> {
86 match data_type {
87 DataType::Struct(_) => Ok(()),
88 other => Err(ArrowError::InvalidArgumentError(format!(
89 "arrow.parquet.variant extension type requires a Struct storage type, got {other}"
90 ))),
91 }
92 }
93
94 fn try_new(
95 data_type: &DataType,
96 _metadata: Self::Metadata,
97 ) -> std::result::Result<Self, ArrowError> {
98 Self.supports_data_type(data_type)?;
99 Ok(Self)
100 }
101}
102
103pub trait ArrowSchemaVisitor {
107 type T;
109
110 type U;
112
113 fn before_field(&mut self, _field: &FieldRef) -> Result<()> {
115 Ok(())
116 }
117
118 fn after_field(&mut self, _field: &FieldRef) -> Result<()> {
120 Ok(())
121 }
122
123 fn before_list_element(&mut self, _field: &FieldRef) -> Result<()> {
125 Ok(())
126 }
127
128 fn after_list_element(&mut self, _field: &FieldRef) -> Result<()> {
130 Ok(())
131 }
132
133 fn before_map_key(&mut self, _field: &FieldRef) -> Result<()> {
135 Ok(())
136 }
137
138 fn after_map_key(&mut self, _field: &FieldRef) -> Result<()> {
140 Ok(())
141 }
142
143 fn before_map_value(&mut self, _field: &FieldRef) -> Result<()> {
145 Ok(())
146 }
147
148 fn after_map_value(&mut self, _field: &FieldRef) -> Result<()> {
150 Ok(())
151 }
152
153 fn schema(&mut self, schema: &ArrowSchema, values: Vec<Self::T>) -> Result<Self::U>;
155
156 fn r#struct(&mut self, fields: &Fields, results: Vec<Self::T>) -> Result<Self::T>;
158
159 fn list(&mut self, list: &DataType, value: Self::T) -> Result<Self::T>;
161
162 fn map(&mut self, map: &DataType, key_value: Self::T, value: Self::T) -> Result<Self::T>;
164
165 fn primitive(&mut self, p: &DataType) -> Result<Self::T>;
167
168 fn variant(&mut self, field: &FieldRef) -> Result<Self::T>
179 where Self: Sized {
180 visit_type(field.data_type(), self)
181 }
182}
183
184fn visit_type<V: ArrowSchemaVisitor>(r#type: &DataType, visitor: &mut V) -> Result<V::T> {
186 match r#type {
187 p if p.is_primitive()
188 || matches!(
189 p,
190 DataType::Boolean
191 | DataType::Utf8
192 | DataType::LargeUtf8
193 | DataType::Utf8View
194 | DataType::Binary
195 | DataType::LargeBinary
196 | DataType::BinaryView
197 | DataType::FixedSizeBinary(_)
198 ) =>
199 {
200 visitor.primitive(p)
201 }
202 DataType::List(element_field) => visit_list(r#type, element_field, visitor),
203 DataType::LargeList(element_field) => visit_list(r#type, element_field, visitor),
204 DataType::FixedSizeList(element_field, _) => visit_list(r#type, element_field, visitor),
205 DataType::Map(field, _) => match field.data_type() {
206 DataType::Struct(fields) => {
207 if fields.len() != 2 {
208 return Err(Error::new(
209 ErrorKind::DataInvalid,
210 "Map field must have exactly 2 fields",
211 ));
212 }
213
214 let key_field = &fields[0];
215 let value_field = &fields[1];
216
217 let key_result = {
218 visitor.before_map_key(key_field)?;
219 let ret = visit_field(key_field, visitor)?;
220 visitor.after_map_key(key_field)?;
221 ret
222 };
223
224 let value_result = {
225 visitor.before_map_value(value_field)?;
226 let ret = visit_field(value_field, visitor)?;
227 visitor.after_map_value(value_field)?;
228 ret
229 };
230
231 visitor.map(r#type, key_result, value_result)
232 }
233 _ => Err(Error::new(
234 ErrorKind::DataInvalid,
235 "Map field must have struct type",
236 )),
237 },
238 DataType::Struct(fields) => visit_struct(fields, visitor),
239 DataType::Dictionary(_key_type, value_type) => visit_type(value_type, visitor),
240 other => Err(Error::new(
241 ErrorKind::DataInvalid,
242 format!("Cannot visit Arrow data type: {other}"),
243 )),
244 }
245}
246
247fn visit_field<V: ArrowSchemaVisitor>(field: &FieldRef, visitor: &mut V) -> Result<V::T> {
250 if field.extension_type_name() == Some(VariantExtensionType::NAME) {
251 visitor.variant(field)
252 } else {
253 visit_type(field.data_type(), visitor)
254 }
255}
256
257fn visit_list<V: ArrowSchemaVisitor>(
259 data_type: &DataType,
260 element_field: &FieldRef,
261 visitor: &mut V,
262) -> Result<V::T> {
263 visitor.before_list_element(element_field)?;
264 let value = visit_field(element_field, visitor)?;
265 visitor.after_list_element(element_field)?;
266 visitor.list(data_type, value)
267}
268
269fn visit_struct<V: ArrowSchemaVisitor>(fields: &Fields, visitor: &mut V) -> Result<V::T> {
271 let mut results = Vec::with_capacity(fields.len());
272 for field in fields {
273 visitor.before_field(field)?;
274 let result = visit_field(field, visitor)?;
275 visitor.after_field(field)?;
276 results.push(result);
277 }
278
279 visitor.r#struct(fields, results)
280}
281
282pub(crate) fn visit_schema<V: ArrowSchemaVisitor>(
284 schema: &ArrowSchema,
285 visitor: &mut V,
286) -> Result<V::U> {
287 let mut results = Vec::with_capacity(schema.fields().len());
288 for field in schema.fields() {
289 visitor.before_field(field)?;
290 let result = visit_field(field, visitor)?;
291 visitor.after_field(field)?;
292 results.push(result);
293 }
294 visitor.schema(schema, results)
295}
296
297pub fn arrow_schema_to_schema(schema: &ArrowSchema) -> Result<Schema> {
303 let mut visitor = ArrowSchemaConverter::new();
304 visit_schema(schema, &mut visitor)
305}
306
307pub fn arrow_schema_to_schema_auto_assign_ids(schema: &ArrowSchema) -> Result<Schema> {
316 let mut visitor = ArrowSchemaConverter::new_with_field_ids_from(FIRST_FIELD_ID);
317 visit_schema(schema, &mut visitor)
318}
319
320pub fn arrow_type_to_type(ty: &DataType) -> Result<Type> {
322 let mut visitor = ArrowSchemaConverter::new();
323 visit_type(ty, &mut visitor)
324}
325
326const ARROW_FIELD_DOC_KEY: &str = "doc";
327
328pub(super) fn get_field_id_from_metadata(field: &FieldRef) -> Result<i32> {
329 if let Some(value) = field.metadata().get(PARQUET_FIELD_ID_META_KEY) {
330 return value.parse::<i32>().map_err(|e| {
331 Error::new(
332 ErrorKind::DataInvalid,
333 "Failed to parse field id".to_string(),
334 )
335 .with_context("value", value)
336 .with_source(e)
337 });
338 }
339 Err(Error::new(
340 ErrorKind::DataInvalid,
341 "Field id not found in metadata",
342 ))
343}
344
345fn get_field_doc(field: &FieldRef) -> Option<String> {
346 if let Some(value) = field.metadata().get(ARROW_FIELD_DOC_KEY) {
347 return Some(value.clone());
348 }
349 None
350}
351
352struct ArrowSchemaConverter {
353 reassign_field_ids_from: Option<i32>,
356 next_field_id: i32,
360}
361
362impl ArrowSchemaConverter {
363 fn new() -> Self {
364 Self {
365 reassign_field_ids_from: None,
366 next_field_id: 0,
367 }
368 }
369
370 fn new_with_field_ids_from(start_from: i32) -> Self {
371 Self {
372 reassign_field_ids_from: Some(start_from),
373 next_field_id: 0,
374 }
375 }
376
377 fn get_field_id(&mut self, field: &FieldRef) -> Result<i32> {
378 if self.reassign_field_ids_from.is_some() {
379 let temp_id = self.next_field_id;
383 self.next_field_id += 1;
384 Ok(temp_id)
385 } else {
386 get_field_id_from_metadata(field)
388 }
389 }
390
391 fn convert_fields(
392 &mut self,
393 fields: &Fields,
394 field_results: &[Type],
395 ) -> Result<Vec<NestedFieldRef>> {
396 let mut results = Vec::with_capacity(fields.len());
397 for i in 0..fields.len() {
398 let field = &fields[i];
399 let field_type = &field_results[i];
400 let id = self.get_field_id(field)?;
401 let doc = get_field_doc(field);
402 let nested_field = NestedField {
403 id,
404 doc,
405 name: field.name().clone(),
406 required: !field.is_nullable(),
407 field_type: Box::new(field_type.clone()),
408 initial_default: None,
409 write_default: None,
410 };
411 results.push(Arc::new(nested_field));
412 }
413 Ok(results)
414 }
415}
416
417impl ArrowSchemaVisitor for ArrowSchemaConverter {
418 type T = Type;
419 type U = Schema;
420
421 fn schema(&mut self, schema: &ArrowSchema, values: Vec<Self::T>) -> Result<Self::U> {
422 let fields = self.convert_fields(schema.fields(), &values)?;
423 let mut builder = Schema::builder().with_fields(fields);
424 if let Some(start_from) = self.reassign_field_ids_from {
425 builder = builder.with_reassigned_field_ids(start_from)
426 }
427 builder.build()
428 }
429
430 fn r#struct(&mut self, fields: &Fields, results: Vec<Self::T>) -> Result<Self::T> {
431 let fields = self.convert_fields(fields, &results)?;
432 Ok(Type::Struct(StructType::new(fields)))
433 }
434
435 fn list(&mut self, list: &DataType, value: Self::T) -> Result<Self::T> {
436 let element_field = match list {
437 DataType::List(element_field) => element_field,
438 DataType::LargeList(element_field) => element_field,
439 DataType::FixedSizeList(element_field, _) => element_field,
440 _ => {
441 return Err(Error::new(
442 ErrorKind::DataInvalid,
443 "List type must have list data type",
444 ));
445 }
446 };
447
448 let id = self.get_field_id(element_field)?;
449 let doc = get_field_doc(element_field);
450 let mut element_field =
451 NestedField::list_element(id, value.clone(), !element_field.is_nullable());
452 if let Some(doc) = doc {
453 element_field = element_field.with_doc(doc);
454 }
455 let element_field = Arc::new(element_field);
456 Ok(Type::List(ListType { element_field }))
457 }
458
459 fn map(&mut self, map: &DataType, key_value: Self::T, value: Self::T) -> Result<Self::T> {
460 match map {
461 DataType::Map(field, _) => match field.data_type() {
462 DataType::Struct(fields) => {
463 if fields.len() != 2 {
464 return Err(Error::new(
465 ErrorKind::DataInvalid,
466 "Map field must have exactly 2 fields",
467 ));
468 }
469
470 let key_field = &fields[0];
471 let value_field = &fields[1];
472
473 let key_id = self.get_field_id(key_field)?;
474 let key_doc = get_field_doc(key_field);
475 let mut key_field = NestedField::map_key_element(key_id, key_value.clone());
476 if let Some(doc) = key_doc {
477 key_field = key_field.with_doc(doc);
478 }
479 let key_field = Arc::new(key_field);
480
481 let value_id = self.get_field_id(value_field)?;
482 let value_doc = get_field_doc(value_field);
483 let mut value_field = NestedField::map_value_element(
484 value_id,
485 value.clone(),
486 !value_field.is_nullable(),
487 );
488 if let Some(doc) = value_doc {
489 value_field = value_field.with_doc(doc);
490 }
491 let value_field = Arc::new(value_field);
492
493 Ok(Type::Map(MapType {
494 key_field,
495 value_field,
496 }))
497 }
498 _ => Err(Error::new(
499 ErrorKind::DataInvalid,
500 "Map field must have struct type",
501 )),
502 },
503 _ => Err(Error::new(
504 ErrorKind::DataInvalid,
505 "Map type must have map data type",
506 )),
507 }
508 }
509
510 fn primitive(&mut self, p: &DataType) -> Result<Self::T> {
511 match p {
512 DataType::Boolean => Ok(Type::Primitive(PrimitiveType::Boolean)),
513 DataType::Int8 | DataType::Int16 | DataType::Int32 => {
514 Ok(Type::Primitive(PrimitiveType::Int))
515 }
516 DataType::UInt8 | DataType::UInt16 => Ok(Type::Primitive(PrimitiveType::Int)),
517 DataType::UInt32 => Ok(Type::Primitive(PrimitiveType::Long)),
518 DataType::Int64 => Ok(Type::Primitive(PrimitiveType::Long)),
519 DataType::UInt64 => {
520 Err(Error::new(
522 ErrorKind::DataInvalid,
523 "UInt64 is not supported. Use Int64 for values ≤ 9,223,372,036,854,775,807 or Decimal(20,0) for full uint64 range.",
524 ))
525 }
526 DataType::Float32 => Ok(Type::Primitive(PrimitiveType::Float)),
527 DataType::Float64 => Ok(Type::Primitive(PrimitiveType::Double)),
528 DataType::Decimal128(p, s) => Type::decimal(*p as u32, *s as u32).map_err(|e| {
529 Error::new(
530 ErrorKind::DataInvalid,
531 "Failed to create decimal type".to_string(),
532 )
533 .with_source(e)
534 }),
535 DataType::Date32 => Ok(Type::Primitive(PrimitiveType::Date)),
536 DataType::Time64(unit) if unit == &TimeUnit::Microsecond => {
537 Ok(Type::Primitive(PrimitiveType::Time))
538 }
539 DataType::Timestamp(unit, None) if unit == &TimeUnit::Microsecond => {
540 Ok(Type::Primitive(PrimitiveType::Timestamp))
541 }
542 DataType::Timestamp(unit, None) if unit == &TimeUnit::Nanosecond => {
543 Ok(Type::Primitive(PrimitiveType::TimestampNs))
544 }
545 DataType::Timestamp(unit, Some(zone))
546 if unit == &TimeUnit::Microsecond
547 && (zone.as_ref() == "UTC" || zone.as_ref() == "+00:00") =>
548 {
549 Ok(Type::Primitive(PrimitiveType::Timestamptz))
550 }
551 DataType::Timestamp(unit, Some(zone))
552 if unit == &TimeUnit::Nanosecond
553 && (zone.as_ref() == "UTC" || zone.as_ref() == "+00:00") =>
554 {
555 Ok(Type::Primitive(PrimitiveType::TimestamptzNs))
556 }
557 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
558 Ok(Type::Primitive(PrimitiveType::Binary))
559 }
560 DataType::FixedSizeBinary(width) => {
561 Ok(Type::Primitive(PrimitiveType::Fixed(*width as u64)))
562 }
563 DataType::Utf8View | DataType::Utf8 | DataType::LargeUtf8 => {
564 Ok(Type::Primitive(PrimitiveType::String))
565 }
566 _ => Err(Error::new(
567 ErrorKind::DataInvalid,
568 format!("Unsupported Arrow data type: {p}"),
569 )),
570 }
571 }
572
573 fn variant(&mut self, field: &FieldRef) -> Result<Self::T> {
574 if !matches!(field.data_type(), DataType::Struct(_)) {
577 return Err(Error::new(
578 ErrorKind::DataInvalid,
579 "arrow.parquet.variant extension requires Struct storage",
580 ));
581 }
582 Ok(Type::Variant(VariantType))
586 }
587}
588
589struct ToArrowSchemaConverter;
590
591enum ArrowSchemaOrFieldOrType {
592 Schema(ArrowSchema),
593 Field(Field),
594 Type(DataType),
595}
596
597impl SchemaVisitor for ToArrowSchemaConverter {
598 type T = ArrowSchemaOrFieldOrType;
599
600 fn schema(
601 &mut self,
602 _schema: &Schema,
603 value: ArrowSchemaOrFieldOrType,
604 ) -> Result<ArrowSchemaOrFieldOrType> {
605 let struct_type = match value {
606 ArrowSchemaOrFieldOrType::Type(DataType::Struct(fields)) => fields,
607 _ => unreachable!(),
608 };
609 Ok(ArrowSchemaOrFieldOrType::Schema(ArrowSchema::new(
610 struct_type,
611 )))
612 }
613
614 fn field(
615 &mut self,
616 field: &NestedFieldRef,
617 value: ArrowSchemaOrFieldOrType,
618 ) -> Result<ArrowSchemaOrFieldOrType> {
619 let ty = match value {
620 ArrowSchemaOrFieldOrType::Type(ty) => ty,
621 _ => unreachable!(),
622 };
623 let metadata = if let Some(doc) = &field.doc {
624 HashMap::from([
625 (PARQUET_FIELD_ID_META_KEY.to_string(), field.id.to_string()),
626 (ARROW_FIELD_DOC_KEY.to_string(), doc.clone()),
627 ])
628 } else {
629 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), field.id.to_string())])
630 };
631 let arrow_field =
632 Field::new(field.name.clone(), ty, !field.required).with_metadata(metadata);
633 let arrow_field = if field.field_type.is_variant() {
636 arrow_field.with_extension_type(VariantExtensionType)
637 } else {
638 arrow_field
639 };
640 Ok(ArrowSchemaOrFieldOrType::Field(arrow_field))
641 }
642
643 fn r#struct(
644 &mut self,
645 _: &StructType,
646 results: Vec<ArrowSchemaOrFieldOrType>,
647 ) -> Result<ArrowSchemaOrFieldOrType> {
648 let fields = results
649 .into_iter()
650 .map(|result| match result {
651 ArrowSchemaOrFieldOrType::Field(field) => field,
652 _ => unreachable!(),
653 })
654 .collect();
655 Ok(ArrowSchemaOrFieldOrType::Type(DataType::Struct(fields)))
656 }
657
658 fn list(&mut self, list: &ListType, value: ArrowSchemaOrFieldOrType) -> Result<Self::T> {
659 let field = match self.field(&list.element_field, value)? {
663 ArrowSchemaOrFieldOrType::Field(field) => field,
664 _ => unreachable!(),
665 };
666 Ok(ArrowSchemaOrFieldOrType::Type(DataType::List(Arc::new(
667 field,
668 ))))
669 }
670
671 fn map(
672 &mut self,
673 map: &MapType,
674 key_value: ArrowSchemaOrFieldOrType,
675 value: ArrowSchemaOrFieldOrType,
676 ) -> Result<ArrowSchemaOrFieldOrType> {
677 let key_field = match self.field(&map.key_field, key_value)? {
678 ArrowSchemaOrFieldOrType::Field(field) => field,
679 _ => unreachable!(),
680 };
681 let value_field = match self.field(&map.value_field, value)? {
682 ArrowSchemaOrFieldOrType::Field(field) => field,
683 _ => unreachable!(),
684 };
685 let field = Field::new(
686 DEFAULT_MAP_FIELD_NAME,
687 DataType::Struct(vec![key_field, value_field].into()),
688 false,
690 );
691
692 Ok(ArrowSchemaOrFieldOrType::Type(DataType::Map(
693 field.into(),
694 false,
695 )))
696 }
697
698 fn primitive(&mut self, p: &PrimitiveType) -> Result<ArrowSchemaOrFieldOrType> {
699 match p {
700 PrimitiveType::Boolean => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Boolean)),
701 PrimitiveType::Int => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Int32)),
702 PrimitiveType::Long => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Int64)),
703 PrimitiveType::Float => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Float32)),
704 PrimitiveType::Double => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Float64)),
705 PrimitiveType::Decimal { precision, scale } => {
706 let (precision, scale) = {
707 let precision: u8 = precision.to_owned().try_into().map_err(|err| {
708 Error::new(
709 ErrorKind::DataInvalid,
710 "incompatible precision for decimal type convert",
711 )
712 .with_source(err)
713 })?;
714 let scale = scale.to_owned().try_into().map_err(|err| {
715 Error::new(
716 ErrorKind::DataInvalid,
717 "incompatible scale for decimal type convert",
718 )
719 .with_source(err)
720 })?;
721 (precision, scale)
722 };
723 validate_decimal_precision_and_scale::<Decimal128Type>(precision, scale).map_err(
724 |err| {
725 Error::new(
726 ErrorKind::DataInvalid,
727 "incompatible precision and scale for decimal type convert",
728 )
729 .with_source(err)
730 },
731 )?;
732 Ok(ArrowSchemaOrFieldOrType::Type(DataType::Decimal128(
733 precision, scale,
734 )))
735 }
736 PrimitiveType::Date => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Date32)),
737 PrimitiveType::Time => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Time64(
738 TimeUnit::Microsecond,
739 ))),
740 PrimitiveType::Timestamp => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Timestamp(
741 TimeUnit::Microsecond,
742 None,
743 ))),
744 PrimitiveType::Timestamptz => Ok(ArrowSchemaOrFieldOrType::Type(
745 DataType::Timestamp(TimeUnit::Microsecond, Some(UTC_TIME_ZONE.into())),
747 )),
748 PrimitiveType::TimestampNs => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Timestamp(
749 TimeUnit::Nanosecond,
750 None,
751 ))),
752 PrimitiveType::TimestamptzNs => Ok(ArrowSchemaOrFieldOrType::Type(
753 DataType::Timestamp(TimeUnit::Nanosecond, Some(UTC_TIME_ZONE.into())),
755 )),
756 PrimitiveType::String => Ok(ArrowSchemaOrFieldOrType::Type(DataType::Utf8)),
757 PrimitiveType::Uuid => Ok(ArrowSchemaOrFieldOrType::Type(DataType::FixedSizeBinary(
758 16,
759 ))),
760 PrimitiveType::Fixed(len) => Ok(ArrowSchemaOrFieldOrType::Type(
761 i32::try_from(*len)
762 .ok()
763 .map(DataType::FixedSizeBinary)
764 .unwrap_or(DataType::LargeBinary),
765 )),
766 PrimitiveType::Binary => Ok(ArrowSchemaOrFieldOrType::Type(DataType::LargeBinary)),
767 }
768 }
769
770 fn variant(&mut self, _v: &VariantType) -> Result<ArrowSchemaOrFieldOrType> {
771 let metadata_field = Field::new("metadata", DataType::Binary, false);
777 let value_field = Field::new("value", DataType::Binary, true);
778 Ok(ArrowSchemaOrFieldOrType::Type(DataType::Struct(
779 vec![metadata_field, value_field].into(),
780 )))
781 }
782}
783
784pub fn schema_to_arrow_schema(schema: &Schema) -> Result<ArrowSchema> {
786 let mut converter = ToArrowSchemaConverter;
787 match crate::spec::visit_schema(schema, &mut converter)? {
788 ArrowSchemaOrFieldOrType::Schema(schema) => Ok(schema),
789 _ => unreachable!(),
790 }
791}
792
793pub fn type_to_arrow_type(ty: &Type) -> Result<DataType> {
795 let mut converter = ToArrowSchemaConverter;
796 match crate::spec::visit_type(ty, &mut converter)? {
797 ArrowSchemaOrFieldOrType::Type(ty) => Ok(ty),
798 _ => unreachable!(),
799 }
800}
801
802pub(crate) fn get_arrow_datum(datum: &Datum) -> Result<Arc<dyn ArrowDatum + Send + Sync>> {
804 match (datum.data_type(), datum.literal()) {
805 (PrimitiveType::Boolean, PrimitiveLiteral::Boolean(value)) => {
806 Ok(Arc::new(BooleanArray::new_scalar(*value)))
807 }
808 (PrimitiveType::Int, PrimitiveLiteral::Int(value)) => {
809 Ok(Arc::new(Int32Array::new_scalar(*value)))
810 }
811 (PrimitiveType::Long, PrimitiveLiteral::Long(value)) => {
812 Ok(Arc::new(Int64Array::new_scalar(*value)))
813 }
814 (PrimitiveType::Float, PrimitiveLiteral::Float(value)) => {
815 Ok(Arc::new(Float32Array::new_scalar(value.into_inner())))
816 }
817 (PrimitiveType::Double, PrimitiveLiteral::Double(value)) => {
818 Ok(Arc::new(Float64Array::new_scalar(value.into_inner())))
819 }
820 (PrimitiveType::String, PrimitiveLiteral::String(value)) => {
821 Ok(Arc::new(StringArray::new_scalar(value.as_str())))
822 }
823 (PrimitiveType::Binary, PrimitiveLiteral::Binary(value)) => {
824 Ok(Arc::new(BinaryArray::new_scalar(value.as_slice())))
825 }
826 (PrimitiveType::Date, PrimitiveLiteral::Int(value)) => {
827 Ok(Arc::new(Date32Array::new_scalar(*value)))
828 }
829 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(value)) => {
830 Ok(Arc::new(TimestampMicrosecondArray::new_scalar(*value)))
831 }
832 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(value)) => Ok(Arc::new(Scalar::new(
833 TimestampMicrosecondArray::new(vec![*value; 1].into(), None).with_timezone_utc(),
834 ))),
835 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(value)) => {
836 Ok(Arc::new(TimestampNanosecondArray::new_scalar(*value)))
837 }
838 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(value)) => Ok(Arc::new(Scalar::new(
839 TimestampNanosecondArray::new(vec![*value; 1].into(), None).with_timezone_utc(),
840 ))),
841 (PrimitiveType::Decimal { precision, scale }, PrimitiveLiteral::Int128(value)) => {
842 let array = Decimal128Array::from_value(*value, 1)
843 .with_precision_and_scale(*precision as _, *scale as _)
844 .unwrap();
845 Ok(Arc::new(Scalar::new(array)))
846 }
847 (PrimitiveType::Uuid, PrimitiveLiteral::UInt128(value)) => {
848 let bytes = Uuid::from_u128(*value).into_bytes();
849 let array = FixedSizeBinaryArray::try_from_iter(vec![bytes].into_iter()).unwrap();
850 Ok(Arc::new(Scalar::new(array)))
851 }
852 (PrimitiveType::Fixed(_), PrimitiveLiteral::Binary(value)) => {
853 let array = FixedSizeBinaryArray::try_from_iter(std::iter::once(value.as_slice()))
854 .map_err(|e| Error::new(ErrorKind::DataInvalid, e.to_string()))?;
855 Ok(Arc::new(Scalar::new(array)))
856 }
857
858 (primitive_type, _) => Err(Error::new(
859 ErrorKind::FeatureUnsupported,
860 format!("Converting datum from type {primitive_type:?} to arrow not supported yet."),
861 )),
862 }
863}
864
865pub(crate) fn get_parquet_stat_min_as_datum(
866 primitive_type: &PrimitiveType,
867 stats: &Statistics,
868) -> Result<Option<Datum>> {
869 Ok(match (primitive_type, stats) {
870 (PrimitiveType::Boolean, Statistics::Boolean(stats)) => {
871 stats.min_opt().map(|val| Datum::bool(*val))
872 }
873 (PrimitiveType::Int, Statistics::Int32(stats)) => {
874 stats.min_opt().map(|val| Datum::int(*val))
875 }
876 (PrimitiveType::Date, Statistics::Int32(stats)) => {
877 stats.min_opt().map(|val| Datum::date(*val))
878 }
879 (PrimitiveType::Long, Statistics::Int64(stats)) => {
880 stats.min_opt().map(|val| Datum::long(*val))
881 }
882 (PrimitiveType::Time, Statistics::Int64(stats)) => {
883 let Some(val) = stats.min_opt() else {
884 return Ok(None);
885 };
886
887 Some(Datum::time_micros(*val)?)
888 }
889 (PrimitiveType::Timestamp, Statistics::Int64(stats)) => {
890 stats.min_opt().map(|val| Datum::timestamp_micros(*val))
891 }
892 (PrimitiveType::Timestamptz, Statistics::Int64(stats)) => {
893 stats.min_opt().map(|val| Datum::timestamptz_micros(*val))
894 }
895 (PrimitiveType::TimestampNs, Statistics::Int64(stats)) => {
896 stats.min_opt().map(|val| Datum::timestamp_nanos(*val))
897 }
898 (PrimitiveType::TimestamptzNs, Statistics::Int64(stats)) => {
899 stats.min_opt().map(|val| Datum::timestamptz_nanos(*val))
900 }
901 (PrimitiveType::Float, Statistics::Float(stats)) => {
902 stats.min_opt().map(|val| Datum::float(*val))
903 }
904 (PrimitiveType::Double, Statistics::Double(stats)) => {
905 stats.min_opt().map(|val| Datum::double(*val))
906 }
907 (PrimitiveType::String, Statistics::ByteArray(stats)) => {
908 let Some(val) = stats.min_opt() else {
909 return Ok(None);
910 };
911
912 Some(Datum::string(val.as_utf8()?))
913 }
914 (
915 PrimitiveType::Decimal {
916 precision: _,
917 scale: _,
918 },
919 Statistics::ByteArray(stats),
920 ) => {
921 let Some(bytes) = stats.min_bytes_opt() else {
922 return Ok(None);
923 };
924 Some(Datum::new(
925 primitive_type.clone(),
926 PrimitiveLiteral::Int128(i128::from_be_bytes(bytes.try_into()?)),
927 ))
928 }
929 (
930 PrimitiveType::Decimal {
931 precision: _,
932 scale: _,
933 },
934 Statistics::FixedLenByteArray(stats),
935 ) => {
936 let Some(bytes) = stats.min_bytes_opt() else {
937 return Ok(None);
938 };
939 Some(Datum::new(
940 primitive_type.clone(),
941 PrimitiveLiteral::Int128(i128_from_be_bytes(bytes).ok_or_else(|| {
942 Error::new(
943 ErrorKind::DataInvalid,
944 format!("Can't convert bytes to i128: {bytes:?}"),
945 )
946 })?),
947 ))
948 }
949 (
950 PrimitiveType::Decimal {
951 precision: _,
952 scale: _,
953 },
954 Statistics::Int32(stats),
955 ) => stats.min_opt().map(|val| {
956 Datum::new(
957 primitive_type.clone(),
958 PrimitiveLiteral::Int128(i128::from(*val)),
959 )
960 }),
961
962 (
963 PrimitiveType::Decimal {
964 precision: _,
965 scale: _,
966 },
967 Statistics::Int64(stats),
968 ) => stats.min_opt().map(|val| {
969 Datum::new(
970 primitive_type.clone(),
971 PrimitiveLiteral::Int128(i128::from(*val)),
972 )
973 }),
974 (PrimitiveType::Uuid, Statistics::FixedLenByteArray(stats)) => {
975 let Some(bytes) = stats.min_bytes_opt() else {
976 return Ok(None);
977 };
978 if bytes.len() != 16 {
979 return Err(Error::new(
980 ErrorKind::Unexpected,
981 "Invalid length of uuid bytes.",
982 ));
983 }
984 Some(Datum::uuid(Uuid::from_bytes(
985 bytes[..16].try_into().unwrap(),
986 )))
987 }
988 (PrimitiveType::Fixed(len), Statistics::FixedLenByteArray(stat)) => {
989 let Some(bytes) = stat.min_bytes_opt() else {
990 return Ok(None);
991 };
992 if bytes.len() != *len as usize {
993 return Err(Error::new(
994 ErrorKind::Unexpected,
995 "Invalid length of fixed bytes.",
996 ));
997 }
998 Some(Datum::fixed(bytes.to_vec()))
999 }
1000 (PrimitiveType::Binary, Statistics::ByteArray(stat)) => {
1001 return Ok(stat
1002 .min_bytes_opt()
1003 .map(|bytes| Datum::binary(bytes.to_vec())));
1004 }
1005 _ => {
1006 return Ok(None);
1007 }
1008 })
1009}
1010
1011pub(crate) fn get_parquet_stat_max_as_datum(
1012 primitive_type: &PrimitiveType,
1013 stats: &Statistics,
1014) -> Result<Option<Datum>> {
1015 Ok(match (primitive_type, stats) {
1016 (PrimitiveType::Boolean, Statistics::Boolean(stats)) => {
1017 stats.max_opt().map(|val| Datum::bool(*val))
1018 }
1019 (PrimitiveType::Int, Statistics::Int32(stats)) => {
1020 stats.max_opt().map(|val| Datum::int(*val))
1021 }
1022 (PrimitiveType::Date, Statistics::Int32(stats)) => {
1023 stats.max_opt().map(|val| Datum::date(*val))
1024 }
1025 (PrimitiveType::Long, Statistics::Int64(stats)) => {
1026 stats.max_opt().map(|val| Datum::long(*val))
1027 }
1028 (PrimitiveType::Time, Statistics::Int64(stats)) => {
1029 let Some(val) = stats.max_opt() else {
1030 return Ok(None);
1031 };
1032
1033 Some(Datum::time_micros(*val)?)
1034 }
1035 (PrimitiveType::Timestamp, Statistics::Int64(stats)) => {
1036 stats.max_opt().map(|val| Datum::timestamp_micros(*val))
1037 }
1038 (PrimitiveType::Timestamptz, Statistics::Int64(stats)) => {
1039 stats.max_opt().map(|val| Datum::timestamptz_micros(*val))
1040 }
1041 (PrimitiveType::TimestampNs, Statistics::Int64(stats)) => {
1042 stats.max_opt().map(|val| Datum::timestamp_nanos(*val))
1043 }
1044 (PrimitiveType::TimestamptzNs, Statistics::Int64(stats)) => {
1045 stats.max_opt().map(|val| Datum::timestamptz_nanos(*val))
1046 }
1047 (PrimitiveType::Float, Statistics::Float(stats)) => {
1048 stats.max_opt().map(|val| Datum::float(*val))
1049 }
1050 (PrimitiveType::Double, Statistics::Double(stats)) => {
1051 stats.max_opt().map(|val| Datum::double(*val))
1052 }
1053 (PrimitiveType::String, Statistics::ByteArray(stats)) => {
1054 let Some(val) = stats.max_opt() else {
1055 return Ok(None);
1056 };
1057
1058 Some(Datum::string(val.as_utf8()?))
1059 }
1060 (
1061 PrimitiveType::Decimal {
1062 precision: _,
1063 scale: _,
1064 },
1065 Statistics::ByteArray(stats),
1066 ) => {
1067 let Some(bytes) = stats.max_bytes_opt() else {
1068 return Ok(None);
1069 };
1070 Some(Datum::new(
1071 primitive_type.clone(),
1072 PrimitiveLiteral::Int128(i128::from_be_bytes(bytes.try_into()?)),
1073 ))
1074 }
1075 (
1076 PrimitiveType::Decimal {
1077 precision: _,
1078 scale: _,
1079 },
1080 Statistics::FixedLenByteArray(stats),
1081 ) => {
1082 let Some(bytes) = stats.max_bytes_opt() else {
1083 return Ok(None);
1084 };
1085 Some(Datum::new(
1086 primitive_type.clone(),
1087 PrimitiveLiteral::Int128(i128_from_be_bytes(bytes).ok_or_else(|| {
1088 Error::new(
1089 ErrorKind::DataInvalid,
1090 format!("Can't convert bytes to i128: {bytes:?}"),
1091 )
1092 })?),
1093 ))
1094 }
1095 (
1096 PrimitiveType::Decimal {
1097 precision: _,
1098 scale: _,
1099 },
1100 Statistics::Int32(stats),
1101 ) => stats.max_opt().map(|val| {
1102 Datum::new(
1103 primitive_type.clone(),
1104 PrimitiveLiteral::Int128(i128::from(*val)),
1105 )
1106 }),
1107
1108 (
1109 PrimitiveType::Decimal {
1110 precision: _,
1111 scale: _,
1112 },
1113 Statistics::Int64(stats),
1114 ) => stats.max_opt().map(|val| {
1115 Datum::new(
1116 primitive_type.clone(),
1117 PrimitiveLiteral::Int128(i128::from(*val)),
1118 )
1119 }),
1120 (PrimitiveType::Uuid, Statistics::FixedLenByteArray(stats)) => {
1121 let Some(bytes) = stats.max_bytes_opt() else {
1122 return Ok(None);
1123 };
1124 if bytes.len() != 16 {
1125 return Err(Error::new(
1126 ErrorKind::Unexpected,
1127 "Invalid length of uuid bytes.",
1128 ));
1129 }
1130 Some(Datum::uuid(Uuid::from_bytes(
1131 bytes[..16].try_into().unwrap(),
1132 )))
1133 }
1134 (PrimitiveType::Fixed(len), Statistics::FixedLenByteArray(stat)) => {
1135 let Some(bytes) = stat.max_bytes_opt() else {
1136 return Ok(None);
1137 };
1138 if bytes.len() != *len as usize {
1139 return Err(Error::new(
1140 ErrorKind::Unexpected,
1141 "Invalid length of fixed bytes.",
1142 ));
1143 }
1144 Some(Datum::fixed(bytes.to_vec()))
1145 }
1146 (PrimitiveType::Binary, Statistics::ByteArray(stat)) => {
1147 return Ok(stat
1148 .max_bytes_opt()
1149 .map(|bytes| Datum::binary(bytes.to_vec())));
1150 }
1151 _ => {
1152 return Ok(None);
1153 }
1154 })
1155}
1156
1157impl TryFrom<&ArrowSchema> for Schema {
1158 type Error = Error;
1159
1160 fn try_from(schema: &ArrowSchema) -> Result<Self> {
1161 arrow_schema_to_schema(schema)
1162 }
1163}
1164
1165impl TryFrom<&Schema> for ArrowSchema {
1166 type Error = Error;
1167
1168 fn try_from(schema: &Schema) -> Result<Self> {
1169 schema_to_arrow_schema(schema)
1170 }
1171}
1172
1173pub fn datum_to_arrow_type_with_ree(datum: &Datum) -> DataType {
1195 let make_ree = |values_type: DataType| -> DataType {
1199 let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int32, false));
1200 let values_field = Arc::new(Field::new("values", values_type, true));
1201 DataType::RunEndEncoded(run_ends_field, values_field)
1202 };
1203
1204 match datum.data_type() {
1206 PrimitiveType::Boolean => make_ree(DataType::Boolean),
1207 PrimitiveType::Int => make_ree(DataType::Int32),
1208 PrimitiveType::Long => make_ree(DataType::Int64),
1209 PrimitiveType::Float => make_ree(DataType::Float32),
1210 PrimitiveType::Double => make_ree(DataType::Float64),
1211 PrimitiveType::Date => make_ree(DataType::Date32),
1212 PrimitiveType::Time => make_ree(DataType::Int64),
1213 PrimitiveType::Timestamp => make_ree(DataType::Int64),
1214 PrimitiveType::Timestamptz => make_ree(DataType::Int64),
1215 PrimitiveType::TimestampNs => make_ree(DataType::Int64),
1216 PrimitiveType::TimestamptzNs => make_ree(DataType::Int64),
1217 PrimitiveType::String => make_ree(DataType::Utf8),
1218 PrimitiveType::Uuid => make_ree(DataType::Binary),
1219 PrimitiveType::Fixed(_) => make_ree(DataType::Binary),
1220 PrimitiveType::Binary => make_ree(DataType::Binary),
1221 PrimitiveType::Decimal { precision, scale } => {
1222 make_ree(DataType::Decimal128(*precision as u8, *scale as i8))
1223 }
1224 }
1225}
1226
1227struct MetadataStripVisitor {
1233 field_stack: Vec<Field>,
1235}
1236
1237impl MetadataStripVisitor {
1238 fn new() -> Self {
1239 Self {
1240 field_stack: Vec::new(),
1241 }
1242 }
1243}
1244
1245impl ArrowSchemaVisitor for MetadataStripVisitor {
1246 type T = Field;
1247 type U = ArrowSchema;
1248
1249 fn before_field(&mut self, field: &FieldRef) -> Result<()> {
1250 self.field_stack.push(Field::new(
1252 field.name(),
1253 DataType::Null, field.is_nullable(),
1255 ));
1256 Ok(())
1257 }
1258
1259 fn after_field(&mut self, _field: &FieldRef) -> Result<()> {
1260 Ok(())
1261 }
1262
1263 fn schema(&mut self, _schema: &ArrowSchema, values: Vec<Self::T>) -> Result<Self::U> {
1264 Ok(ArrowSchema::new(values))
1265 }
1266
1267 fn r#struct(&mut self, _fields: &Fields, results: Vec<Self::T>) -> Result<Self::T> {
1268 let field_info = self
1270 .field_stack
1271 .pop()
1272 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in struct"))?;
1273
1274 Ok(Field::new(
1276 field_info.name(),
1277 DataType::Struct(Fields::from(results)),
1278 field_info.is_nullable(),
1279 ))
1280 }
1281
1282 fn list(&mut self, list: &DataType, value: Self::T) -> Result<Self::T> {
1283 let field_info = self
1285 .field_stack
1286 .pop()
1287 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in list"))?;
1288
1289 let list_type = match list {
1291 DataType::List(_) => DataType::List(Arc::new(value)),
1292 DataType::LargeList(_) => DataType::LargeList(Arc::new(value)),
1293 DataType::FixedSizeList(_, size) => DataType::FixedSizeList(Arc::new(value), *size),
1294 _ => {
1295 return Err(Error::new(
1296 ErrorKind::Unexpected,
1297 format!("Expected list type, got {list}"),
1298 ));
1299 }
1300 };
1301
1302 Ok(Field::new(
1303 field_info.name(),
1304 list_type,
1305 field_info.is_nullable(),
1306 ))
1307 }
1308
1309 fn map(&mut self, map: &DataType, key_value: Self::T, value: Self::T) -> Result<Self::T> {
1310 let field_info = self
1312 .field_stack
1313 .pop()
1314 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in map"))?;
1315
1316 let struct_field = Field::new(
1318 DEFAULT_MAP_FIELD_NAME,
1319 DataType::Struct(Fields::from(vec![key_value, value])),
1320 false,
1321 );
1322
1323 let sorted = match map {
1325 DataType::Map(_, sorted) => *sorted,
1326 _ => {
1327 return Err(Error::new(
1328 ErrorKind::Unexpected,
1329 format!("Expected map type, got {map}"),
1330 ));
1331 }
1332 };
1333
1334 Ok(Field::new(
1336 field_info.name(),
1337 DataType::Map(Arc::new(struct_field), sorted),
1338 field_info.is_nullable(),
1339 ))
1340 }
1341
1342 fn primitive(&mut self, p: &DataType) -> Result<Self::T> {
1343 let field_info = self.field_stack.pop().ok_or_else(|| {
1345 Error::new(ErrorKind::Unexpected, "Field stack underflow in primitive")
1346 })?;
1347
1348 Ok(Field::new(
1350 field_info.name(),
1351 p.clone(),
1352 field_info.is_nullable(),
1353 ))
1354 }
1355}
1356
1357pub fn strip_metadata_from_schema(schema: &ArrowSchema) -> Result<ArrowSchema> {
1387 let mut visitor = MetadataStripVisitor::new();
1388 visit_schema(schema, &mut visitor)
1389}
1390
1391#[cfg(test)]
1392mod tests {
1393 use std::collections::HashMap;
1394 use std::sync::Arc;
1395
1396 use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
1397
1398 use super::*;
1399 use crate::spec::decimal_utils::decimal_new;
1400 use crate::spec::{Literal, Schema};
1401
1402 fn simple_field(name: &str, ty: DataType, nullable: bool, value: &str) -> Field {
1404 Field::new(name, ty, nullable).with_metadata(HashMap::from([(
1405 PARQUET_FIELD_ID_META_KEY.to_string(),
1406 value.to_string(),
1407 )]))
1408 }
1409
1410 fn arrow_schema_for_arrow_schema_to_schema_test() -> ArrowSchema {
1411 let fields = Fields::from(vec![
1412 simple_field("key", DataType::Int32, false, "28"),
1413 simple_field("value", DataType::Utf8, true, "29"),
1414 ]);
1415
1416 let r#struct = DataType::Struct(fields);
1417 let map = DataType::Map(
1418 Arc::new(simple_field(DEFAULT_MAP_FIELD_NAME, r#struct, false, "17")),
1419 false,
1420 );
1421 let dictionary = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
1422
1423 let fields = Fields::from(vec![
1424 simple_field("aa", DataType::Int32, false, "18"),
1425 simple_field("bb", DataType::Utf8, true, "19"),
1426 simple_field(
1427 "cc",
1428 DataType::Timestamp(TimeUnit::Microsecond, None),
1429 false,
1430 "20",
1431 ),
1432 ]);
1433
1434 let r#struct = DataType::Struct(fields);
1435
1436 ArrowSchema::new(vec![
1437 simple_field("a", DataType::Int32, false, "2"),
1438 simple_field("b", DataType::Int64, false, "1"),
1439 simple_field("c", DataType::Utf8, false, "3"),
1440 simple_field("n", DataType::Utf8, false, "21"),
1441 simple_field(
1442 "d",
1443 DataType::Timestamp(TimeUnit::Microsecond, None),
1444 true,
1445 "4",
1446 ),
1447 simple_field("e", DataType::Boolean, true, "6"),
1448 simple_field("f", DataType::Float32, false, "5"),
1449 simple_field("g", DataType::Float64, false, "7"),
1450 simple_field("p", DataType::Decimal128(10, 2), false, "27"),
1451 simple_field("h", DataType::Date32, false, "8"),
1452 simple_field("i", DataType::Time64(TimeUnit::Microsecond), false, "9"),
1453 simple_field(
1454 "j",
1455 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
1456 false,
1457 "10",
1458 ),
1459 simple_field(
1460 "k",
1461 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1462 false,
1463 "12",
1464 ),
1465 simple_field("l", DataType::Binary, false, "13"),
1466 simple_field("o", DataType::LargeBinary, false, "22"),
1467 simple_field("m", DataType::FixedSizeBinary(10), false, "11"),
1468 simple_field(
1469 "list",
1470 DataType::List(Arc::new(simple_field(
1471 "element",
1472 DataType::Int32,
1473 false,
1474 "15",
1475 ))),
1476 true,
1477 "14",
1478 ),
1479 simple_field(
1480 "large_list",
1481 DataType::LargeList(Arc::new(simple_field(
1482 "element",
1483 DataType::Utf8,
1484 false,
1485 "23",
1486 ))),
1487 true,
1488 "24",
1489 ),
1490 simple_field(
1491 "fixed_list",
1492 DataType::FixedSizeList(
1493 Arc::new(simple_field("element", DataType::Binary, false, "26")),
1494 10,
1495 ),
1496 true,
1497 "25",
1498 ),
1499 simple_field("map", map, false, "16"),
1500 simple_field("struct", r#struct, false, "17"),
1501 simple_field("dictionary", dictionary, false, "30"),
1502 ])
1503 }
1504
1505 fn iceberg_schema_for_arrow_schema_to_schema_test() -> Schema {
1506 let schema_json = r#"{
1507 "type":"struct",
1508 "schema-id":0,
1509 "fields":[
1510 {
1511 "id":2,
1512 "name":"a",
1513 "required":true,
1514 "type":"int"
1515 },
1516 {
1517 "id":1,
1518 "name":"b",
1519 "required":true,
1520 "type":"long"
1521 },
1522 {
1523 "id":3,
1524 "name":"c",
1525 "required":true,
1526 "type":"string"
1527 },
1528 {
1529 "id":21,
1530 "name":"n",
1531 "required":true,
1532 "type":"string"
1533 },
1534 {
1535 "id":4,
1536 "name":"d",
1537 "required":false,
1538 "type":"timestamp"
1539 },
1540 {
1541 "id":6,
1542 "name":"e",
1543 "required":false,
1544 "type":"boolean"
1545 },
1546 {
1547 "id":5,
1548 "name":"f",
1549 "required":true,
1550 "type":"float"
1551 },
1552 {
1553 "id":7,
1554 "name":"g",
1555 "required":true,
1556 "type":"double"
1557 },
1558 {
1559 "id":27,
1560 "name":"p",
1561 "required":true,
1562 "type":"decimal(10,2)"
1563 },
1564 {
1565 "id":8,
1566 "name":"h",
1567 "required":true,
1568 "type":"date"
1569 },
1570 {
1571 "id":9,
1572 "name":"i",
1573 "required":true,
1574 "type":"time"
1575 },
1576 {
1577 "id":10,
1578 "name":"j",
1579 "required":true,
1580 "type":"timestamptz"
1581 },
1582 {
1583 "id":12,
1584 "name":"k",
1585 "required":true,
1586 "type":"timestamptz"
1587 },
1588 {
1589 "id":13,
1590 "name":"l",
1591 "required":true,
1592 "type":"binary"
1593 },
1594 {
1595 "id":22,
1596 "name":"o",
1597 "required":true,
1598 "type":"binary"
1599 },
1600 {
1601 "id":11,
1602 "name":"m",
1603 "required":true,
1604 "type":"fixed[10]"
1605 },
1606 {
1607 "id":14,
1608 "name":"list",
1609 "required": false,
1610 "type": {
1611 "type": "list",
1612 "element-id": 15,
1613 "element-required": true,
1614 "element": "int"
1615 }
1616 },
1617 {
1618 "id":24,
1619 "name":"large_list",
1620 "required": false,
1621 "type": {
1622 "type": "list",
1623 "element-id": 23,
1624 "element-required": true,
1625 "element": "string"
1626 }
1627 },
1628 {
1629 "id":25,
1630 "name":"fixed_list",
1631 "required": false,
1632 "type": {
1633 "type": "list",
1634 "element-id": 26,
1635 "element-required": true,
1636 "element": "binary"
1637 }
1638 },
1639 {
1640 "id":16,
1641 "name":"map",
1642 "required": true,
1643 "type": {
1644 "type": "map",
1645 "key-id": 28,
1646 "key": "int",
1647 "value-id": 29,
1648 "value-required": false,
1649 "value": "string"
1650 }
1651 },
1652 {
1653 "id":17,
1654 "name":"struct",
1655 "required": true,
1656 "type": {
1657 "type": "struct",
1658 "fields": [
1659 {
1660 "id":18,
1661 "name":"aa",
1662 "required":true,
1663 "type":"int"
1664 },
1665 {
1666 "id":19,
1667 "name":"bb",
1668 "required":false,
1669 "type":"string"
1670 },
1671 {
1672 "id":20,
1673 "name":"cc",
1674 "required":true,
1675 "type":"timestamp"
1676 }
1677 ]
1678 }
1679 },
1680 {
1681 "id":30,
1682 "name":"dictionary",
1683 "required":true,
1684 "type":"string"
1685 }
1686 ],
1687 "identifier-field-ids":[]
1688 }"#;
1689
1690 let schema: Schema = serde_json::from_str(schema_json).unwrap();
1691 schema
1692 }
1693
1694 #[test]
1695 fn test_arrow_schema_to_schema() {
1696 let arrow_schema = arrow_schema_for_arrow_schema_to_schema_test();
1697 let schema = iceberg_schema_for_arrow_schema_to_schema_test();
1698 let converted_schema = arrow_schema_to_schema(&arrow_schema).unwrap();
1699 pretty_assertions::assert_eq!(converted_schema, schema);
1700 }
1701
1702 fn arrow_schema_for_schema_to_arrow_schema_test() -> ArrowSchema {
1703 let fields = Fields::from(vec![
1704 simple_field("key", DataType::Int32, false, "28"),
1705 simple_field("value", DataType::Utf8, true, "29"),
1706 ]);
1707
1708 let r#struct = DataType::Struct(fields);
1709 let map = DataType::Map(
1710 Arc::new(Field::new(DEFAULT_MAP_FIELD_NAME, r#struct, false)),
1711 false,
1712 );
1713
1714 let fields = Fields::from(vec![
1715 simple_field("aa", DataType::Int32, false, "18"),
1716 simple_field("bb", DataType::Utf8, true, "19"),
1717 simple_field(
1718 "cc",
1719 DataType::Timestamp(TimeUnit::Microsecond, None),
1720 false,
1721 "20",
1722 ),
1723 ]);
1724
1725 let r#struct = DataType::Struct(fields);
1726
1727 ArrowSchema::new(vec![
1728 simple_field("a", DataType::Int32, false, "2"),
1729 simple_field("b", DataType::Int64, false, "1"),
1730 simple_field("c", DataType::Utf8, false, "3"),
1731 simple_field("n", DataType::Utf8, false, "21"),
1732 simple_field(
1733 "d",
1734 DataType::Timestamp(TimeUnit::Microsecond, None),
1735 true,
1736 "4",
1737 ),
1738 simple_field("e", DataType::Boolean, true, "6"),
1739 simple_field("f", DataType::Float32, false, "5"),
1740 simple_field("g", DataType::Float64, false, "7"),
1741 simple_field("p", DataType::Decimal128(10, 2), false, "27"),
1742 simple_field("h", DataType::Date32, false, "8"),
1743 simple_field("i", DataType::Time64(TimeUnit::Microsecond), false, "9"),
1744 simple_field(
1745 "j",
1746 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1747 false,
1748 "10",
1749 ),
1750 simple_field(
1751 "k",
1752 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1753 false,
1754 "12",
1755 ),
1756 simple_field("l", DataType::LargeBinary, false, "13"),
1757 simple_field("o", DataType::LargeBinary, false, "22"),
1758 simple_field("m", DataType::FixedSizeBinary(10), false, "11"),
1759 simple_field(
1760 "list",
1761 DataType::List(Arc::new(simple_field(
1762 "element",
1763 DataType::Int32,
1764 false,
1765 "15",
1766 ))),
1767 true,
1768 "14",
1769 ),
1770 simple_field(
1771 "large_list",
1772 DataType::List(Arc::new(simple_field(
1773 "element",
1774 DataType::Utf8,
1775 false,
1776 "23",
1777 ))),
1778 true,
1779 "24",
1780 ),
1781 simple_field(
1782 "fixed_list",
1783 DataType::List(Arc::new(simple_field(
1784 "element",
1785 DataType::LargeBinary,
1786 false,
1787 "26",
1788 ))),
1789 true,
1790 "25",
1791 ),
1792 simple_field("map", map, false, "16"),
1793 simple_field("struct", r#struct, false, "17"),
1794 simple_field("uuid", DataType::FixedSizeBinary(16), false, "30"),
1795 Field::new(
1796 "v",
1797 DataType::Struct(Fields::from(vec![
1798 Field::new("metadata", DataType::Binary, false),
1799 Field::new("value", DataType::Binary, true),
1800 ])),
1801 true,
1802 )
1803 .with_metadata(HashMap::from([
1804 (PARQUET_FIELD_ID_META_KEY.to_string(), "31".to_string()),
1805 (
1806 arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
1807 "arrow.parquet.variant".to_string(),
1808 ),
1809 ])),
1810 ])
1811 }
1812
1813 fn iceberg_schema_for_schema_to_arrow_schema() -> Schema {
1814 let schema_json = r#"{
1815 "type":"struct",
1816 "schema-id":0,
1817 "fields":[
1818 {
1819 "id":2,
1820 "name":"a",
1821 "required":true,
1822 "type":"int"
1823 },
1824 {
1825 "id":1,
1826 "name":"b",
1827 "required":true,
1828 "type":"long"
1829 },
1830 {
1831 "id":3,
1832 "name":"c",
1833 "required":true,
1834 "type":"string"
1835 },
1836 {
1837 "id":21,
1838 "name":"n",
1839 "required":true,
1840 "type":"string"
1841 },
1842 {
1843 "id":4,
1844 "name":"d",
1845 "required":false,
1846 "type":"timestamp"
1847 },
1848 {
1849 "id":6,
1850 "name":"e",
1851 "required":false,
1852 "type":"boolean"
1853 },
1854 {
1855 "id":5,
1856 "name":"f",
1857 "required":true,
1858 "type":"float"
1859 },
1860 {
1861 "id":7,
1862 "name":"g",
1863 "required":true,
1864 "type":"double"
1865 },
1866 {
1867 "id":27,
1868 "name":"p",
1869 "required":true,
1870 "type":"decimal(10,2)"
1871 },
1872 {
1873 "id":8,
1874 "name":"h",
1875 "required":true,
1876 "type":"date"
1877 },
1878 {
1879 "id":9,
1880 "name":"i",
1881 "required":true,
1882 "type":"time"
1883 },
1884 {
1885 "id":10,
1886 "name":"j",
1887 "required":true,
1888 "type":"timestamptz"
1889 },
1890 {
1891 "id":12,
1892 "name":"k",
1893 "required":true,
1894 "type":"timestamptz"
1895 },
1896 {
1897 "id":13,
1898 "name":"l",
1899 "required":true,
1900 "type":"binary"
1901 },
1902 {
1903 "id":22,
1904 "name":"o",
1905 "required":true,
1906 "type":"binary"
1907 },
1908 {
1909 "id":11,
1910 "name":"m",
1911 "required":true,
1912 "type":"fixed[10]"
1913 },
1914 {
1915 "id":14,
1916 "name":"list",
1917 "required": false,
1918 "type": {
1919 "type": "list",
1920 "element-id": 15,
1921 "element-required": true,
1922 "element": "int"
1923 }
1924 },
1925 {
1926 "id":24,
1927 "name":"large_list",
1928 "required": false,
1929 "type": {
1930 "type": "list",
1931 "element-id": 23,
1932 "element-required": true,
1933 "element": "string"
1934 }
1935 },
1936 {
1937 "id":25,
1938 "name":"fixed_list",
1939 "required": false,
1940 "type": {
1941 "type": "list",
1942 "element-id": 26,
1943 "element-required": true,
1944 "element": "binary"
1945 }
1946 },
1947 {
1948 "id":16,
1949 "name":"map",
1950 "required": true,
1951 "type": {
1952 "type": "map",
1953 "key-id": 28,
1954 "key": "int",
1955 "value-id": 29,
1956 "value-required": false,
1957 "value": "string"
1958 }
1959 },
1960 {
1961 "id":17,
1962 "name":"struct",
1963 "required": true,
1964 "type": {
1965 "type": "struct",
1966 "fields": [
1967 {
1968 "id":18,
1969 "name":"aa",
1970 "required":true,
1971 "type":"int"
1972 },
1973 {
1974 "id":19,
1975 "name":"bb",
1976 "required":false,
1977 "type":"string"
1978 },
1979 {
1980 "id":20,
1981 "name":"cc",
1982 "required":true,
1983 "type":"timestamp"
1984 }
1985 ]
1986 }
1987 },
1988 {
1989 "id":30,
1990 "name":"uuid",
1991 "required":true,
1992 "type":"uuid"
1993 },
1994 {
1995 "id":31,
1996 "name":"v",
1997 "required":false,
1998 "type":"variant"
1999 }
2000 ],
2001 "identifier-field-ids":[]
2002 }"#;
2003
2004 let schema: Schema = serde_json::from_str(schema_json).unwrap();
2005 schema
2006 }
2007
2008 #[test]
2009 fn test_schema_to_arrow_schema() {
2010 let arrow_schema = arrow_schema_for_schema_to_arrow_schema_test();
2011 let schema = iceberg_schema_for_schema_to_arrow_schema();
2012 let converted_arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2013 assert_eq!(converted_arrow_schema, arrow_schema);
2014 }
2015
2016 #[test]
2017 fn test_variant_type_to_arrow_type() {
2018 let arrow_type = type_to_arrow_type(&Type::Variant(VariantType)).unwrap();
2021 assert_eq!(
2022 arrow_type,
2023 DataType::Struct(Fields::from(vec![
2024 Field::new("metadata", DataType::Binary, false),
2025 Field::new("value", DataType::Binary, true),
2026 ]))
2027 );
2028 }
2029
2030 #[test]
2031 fn test_variant_field_carries_arrow_extension_type() {
2032 let schema = Schema::builder()
2035 .with_fields(vec![
2036 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2037 ])
2038 .build()
2039 .unwrap();
2040
2041 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2042 let field = arrow_schema.field_with_name("v").unwrap();
2043
2044 assert_eq!(field.extension_type_name(), Some("arrow.parquet.variant"));
2045 assert_eq!(
2047 field.metadata().get(PARQUET_FIELD_ID_META_KEY),
2048 Some(&"1".to_string())
2049 );
2050 assert_eq!(
2051 field.data_type(),
2052 &DataType::Struct(Fields::from(vec![
2053 Field::new("metadata", DataType::Binary, false),
2054 Field::new("value", DataType::Binary, true),
2055 ]))
2056 );
2057 }
2058
2059 #[test]
2060 fn test_variant_nested_in_list_and_map_carries_arrow_extension_type() {
2061 let schema = Schema::builder()
2065 .with_fields(vec![
2066 NestedField::optional(
2067 1,
2068 "l",
2069 Type::List(ListType::new(
2070 NestedField::optional(2, "element", Type::Variant(VariantType)).into(),
2071 )),
2072 )
2073 .into(),
2074 NestedField::optional(
2075 3,
2076 "m",
2077 Type::Map(MapType::new(
2078 NestedField::map_key_element(4, Type::Primitive(PrimitiveType::String))
2079 .into(),
2080 NestedField::map_value_element(5, Type::Variant(VariantType), false).into(),
2081 )),
2082 )
2083 .into(),
2084 ])
2085 .build()
2086 .unwrap();
2087
2088 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2089
2090 let DataType::List(element) = arrow_schema.field_with_name("l").unwrap().data_type() else {
2091 panic!("expected a list");
2092 };
2093 assert_eq!(element.extension_type_name(), Some("arrow.parquet.variant"));
2094
2095 let DataType::Map(entries, _) = arrow_schema.field_with_name("m").unwrap().data_type()
2096 else {
2097 panic!("expected a map");
2098 };
2099 let DataType::Struct(kv) = entries.data_type() else {
2100 panic!("expected a key_value struct");
2101 };
2102 let value = kv.iter().find(|f| f.name() == "value").unwrap();
2103 assert_eq!(value.extension_type_name(), Some("arrow.parquet.variant"));
2104 }
2105
2106 fn variant_storage() -> DataType {
2109 DataType::Struct(Fields::from(vec![
2110 Field::new("metadata", DataType::Binary, false),
2111 Field::new("value", DataType::Binary, true),
2112 ]))
2113 }
2114
2115 #[test]
2116 fn test_variant_arrow_field_folds_to_iceberg_variant() {
2117 let field = simple_field("v", variant_storage(), true, "1")
2121 .with_extension_type(VariantExtensionType);
2122 let arrow_schema = ArrowSchema::new(vec![field]);
2123
2124 let converted = arrow_schema_to_schema(&arrow_schema).unwrap();
2125 let expected = Schema::builder()
2126 .with_fields(vec![
2127 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2128 ])
2129 .build()
2130 .unwrap();
2131 pretty_assertions::assert_eq!(converted, expected);
2132 }
2133
2134 #[test]
2135 fn test_variant_schema_round_trips() {
2136 let schema = Schema::builder()
2139 .with_fields(vec![
2140 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2141 NestedField::optional(
2142 2,
2143 "s",
2144 Type::Struct(StructType::new(vec![
2145 NestedField::optional(3, "sv", Type::Variant(VariantType)).into(),
2146 ])),
2147 )
2148 .into(),
2149 NestedField::optional(
2150 4,
2151 "l",
2152 Type::List(ListType::new(
2153 NestedField::optional(5, "element", Type::Variant(VariantType)).into(),
2154 )),
2155 )
2156 .into(),
2157 NestedField::optional(
2158 6,
2159 "m",
2160 Type::Map(MapType::new(
2161 NestedField::map_key_element(7, Type::Primitive(PrimitiveType::String))
2162 .into(),
2163 NestedField::map_value_element(8, Type::Variant(VariantType), false).into(),
2164 )),
2165 )
2166 .into(),
2167 ])
2168 .build()
2169 .unwrap();
2170
2171 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2172 let round_tripped = arrow_schema_to_schema(&arrow_schema).unwrap();
2173 pretty_assertions::assert_eq!(round_tripped, schema);
2174 }
2175
2176 #[test]
2177 fn test_variant_recognized_with_auto_assigned_ids() {
2178 let field =
2181 Field::new("v", variant_storage(), true).with_extension_type(VariantExtensionType);
2182 let arrow_schema = ArrowSchema::new(vec![field]);
2183
2184 let converted = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap();
2185 let expected = Schema::builder()
2186 .with_fields(vec![
2187 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2188 ])
2189 .build()
2190 .unwrap();
2191 pretty_assertions::assert_eq!(converted, expected);
2192 }
2193
2194 #[test]
2195 fn test_variant_extension_on_non_struct_storage_is_rejected() {
2196 let field = Field::new("v", DataType::Int32, true).with_metadata(HashMap::from([
2199 (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()),
2200 (
2201 arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
2202 VariantExtensionType::NAME.to_string(),
2203 ),
2204 ]));
2205 let arrow_schema = ArrowSchema::new(vec![field]);
2206
2207 let err = arrow_schema_to_schema(&arrow_schema).unwrap_err();
2208 assert!(
2209 err.to_string().contains("requires Struct storage"),
2210 "unexpected error: {err}"
2211 );
2212 }
2213
2214 #[test]
2215 fn test_type_conversion() {
2216 {
2218 let arrow_type = DataType::Int32;
2219 let iceberg_type = Type::Primitive(PrimitiveType::Int);
2220 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2221 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2222 }
2223
2224 {
2226 let arrow_type = DataType::Struct(Fields::from(vec![
2228 Field::new("a", DataType::Int64, false),
2229 Field::new("b", DataType::Utf8, true),
2230 ]));
2231 assert_eq!(
2232 &arrow_type_to_type(&arrow_type).unwrap_err().to_string(),
2233 "DataInvalid => Field id not found in metadata"
2234 );
2235
2236 let arrow_type = DataType::Struct(Fields::from(vec![
2237 Field::new("a", DataType::Int64, false).with_metadata(HashMap::from_iter([(
2238 PARQUET_FIELD_ID_META_KEY.to_string(),
2239 1.to_string(),
2240 )])),
2241 Field::new("b", DataType::Utf8, true).with_metadata(HashMap::from_iter([(
2242 PARQUET_FIELD_ID_META_KEY.to_string(),
2243 2.to_string(),
2244 )])),
2245 ]));
2246 let iceberg_type = Type::Struct(StructType::new(vec![
2247 NestedField {
2248 id: 1,
2249 doc: None,
2250 name: "a".to_string(),
2251 required: true,
2252 field_type: Box::new(Type::Primitive(PrimitiveType::Long)),
2253 initial_default: None,
2254 write_default: None,
2255 }
2256 .into(),
2257 NestedField {
2258 id: 2,
2259 doc: None,
2260 name: "b".to_string(),
2261 required: false,
2262 field_type: Box::new(Type::Primitive(PrimitiveType::String)),
2263 initial_default: None,
2264 write_default: None,
2265 }
2266 .into(),
2267 ]));
2268 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2269 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2270
2271 let iceberg_type = Type::Struct(StructType::new(vec![
2273 NestedField {
2274 id: 1,
2275 doc: None,
2276 name: "a".to_string(),
2277 required: true,
2278 field_type: Box::new(Type::Primitive(PrimitiveType::Long)),
2279 initial_default: Some(Literal::Primitive(PrimitiveLiteral::Int(114514))),
2280 write_default: None,
2281 }
2282 .into(),
2283 NestedField {
2284 id: 2,
2285 doc: None,
2286 name: "b".to_string(),
2287 required: false,
2288 field_type: Box::new(Type::Primitive(PrimitiveType::String)),
2289 initial_default: None,
2290 write_default: Some(Literal::Primitive(PrimitiveLiteral::String(
2291 "514".to_string(),
2292 ))),
2293 }
2294 .into(),
2295 ]));
2296 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2297 }
2298
2299 {
2301 let arrow_type =
2302 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int8));
2303 let iceberg_type = Type::Primitive(PrimitiveType::Int);
2304 assert_eq!(
2305 iceberg_type,
2306 arrow_type_to_type(&arrow_type).unwrap(),
2307 "Expected dictionary conversion to use the contained value"
2308 );
2309
2310 let arrow_type =
2311 DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Boolean));
2312 let iceberg_type = Type::Primitive(PrimitiveType::Boolean);
2313 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2314 }
2315 }
2316
2317 #[test]
2318 fn test_unsigned_integer_type_conversion() {
2319 let test_cases = vec![
2320 (DataType::UInt8, PrimitiveType::Int),
2321 (DataType::UInt16, PrimitiveType::Int),
2322 (DataType::UInt32, PrimitiveType::Long),
2323 ];
2324
2325 for (arrow_type, expected_iceberg_type) in test_cases {
2326 let arrow_field = Field::new("test", arrow_type.clone(), false).with_metadata(
2327 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2328 );
2329 let arrow_schema = ArrowSchema::new(vec![arrow_field]);
2330
2331 let iceberg_schema = arrow_schema_to_schema(&arrow_schema).unwrap();
2332 let iceberg_field = iceberg_schema.as_struct().fields().first().unwrap();
2333
2334 assert!(
2335 matches!(iceberg_field.field_type.as_ref(), Type::Primitive(t) if *t == expected_iceberg_type),
2336 "Expected {arrow_type:?} to map to {expected_iceberg_type:?}"
2337 );
2338 }
2339
2340 {
2342 let arrow_field = Field::new("test", DataType::UInt64, false).with_metadata(
2343 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2344 );
2345 let arrow_schema = ArrowSchema::new(vec![arrow_field]);
2346
2347 let result = arrow_schema_to_schema(&arrow_schema);
2348 assert!(result.is_err());
2349 assert!(
2350 result
2351 .unwrap_err()
2352 .to_string()
2353 .contains("UInt64 is not supported")
2354 );
2355 }
2356 }
2357
2358 #[test]
2359 fn test_datum_conversion() {
2360 {
2361 let datum = Datum::bool(true);
2362 let arrow_datum = get_arrow_datum(&datum).unwrap();
2363 let (array, is_scalar) = arrow_datum.get();
2364 let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
2365 assert!(is_scalar);
2366 assert!(array.value(0));
2367 }
2368 {
2369 let datum = Datum::int(42);
2370 let arrow_datum = get_arrow_datum(&datum).unwrap();
2371 let (array, is_scalar) = arrow_datum.get();
2372 let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
2373 assert!(is_scalar);
2374 assert_eq!(array.value(0), 42);
2375 }
2376 {
2377 let datum = Datum::long(42);
2378 let arrow_datum = get_arrow_datum(&datum).unwrap();
2379 let (array, is_scalar) = arrow_datum.get();
2380 let array = array.as_any().downcast_ref::<Int64Array>().unwrap();
2381 assert!(is_scalar);
2382 assert_eq!(array.value(0), 42);
2383 }
2384 {
2385 let datum = Datum::float(42.42);
2386 let arrow_datum = get_arrow_datum(&datum).unwrap();
2387 let (array, is_scalar) = arrow_datum.get();
2388 let array = array.as_any().downcast_ref::<Float32Array>().unwrap();
2389 assert!(is_scalar);
2390 assert_eq!(array.value(0), 42.42);
2391 }
2392 {
2393 let datum = Datum::double(42.42);
2394 let arrow_datum = get_arrow_datum(&datum).unwrap();
2395 let (array, is_scalar) = arrow_datum.get();
2396 let array = array.as_any().downcast_ref::<Float64Array>().unwrap();
2397 assert!(is_scalar);
2398 assert_eq!(array.value(0), 42.42);
2399 }
2400 {
2401 let datum = Datum::string("abc");
2402 let arrow_datum = get_arrow_datum(&datum).unwrap();
2403 let (array, is_scalar) = arrow_datum.get();
2404 let array = array.as_any().downcast_ref::<StringArray>().unwrap();
2405 assert!(is_scalar);
2406 assert_eq!(array.value(0), "abc");
2407 }
2408 {
2409 let datum = Datum::binary(vec![1, 2, 3, 4]);
2410 let arrow_datum = get_arrow_datum(&datum).unwrap();
2411 let (array, is_scalar) = arrow_datum.get();
2412 let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
2413 assert!(is_scalar);
2414 assert_eq!(array.value(0), &[1, 2, 3, 4]);
2415 }
2416 {
2417 let datum = Datum::date(42);
2418 let arrow_datum = get_arrow_datum(&datum).unwrap();
2419 let (array, is_scalar) = arrow_datum.get();
2420 let array = array.as_any().downcast_ref::<Date32Array>().unwrap();
2421 assert!(is_scalar);
2422 assert_eq!(array.value(0), 42);
2423 }
2424 {
2425 let datum = Datum::timestamp_micros(42);
2426 let arrow_datum = get_arrow_datum(&datum).unwrap();
2427 let (array, is_scalar) = arrow_datum.get();
2428 let array = array
2429 .as_any()
2430 .downcast_ref::<TimestampMicrosecondArray>()
2431 .unwrap();
2432 assert!(is_scalar);
2433 assert_eq!(array.value(0), 42);
2434 }
2435 {
2436 let datum = Datum::timestamptz_micros(42);
2437 let arrow_datum = get_arrow_datum(&datum).unwrap();
2438 let (array, is_scalar) = arrow_datum.get();
2439 let array = array
2440 .as_any()
2441 .downcast_ref::<TimestampMicrosecondArray>()
2442 .unwrap();
2443 assert!(is_scalar);
2444 assert_eq!(array.timezone(), Some("+00:00"));
2445 assert_eq!(array.value(0), 42);
2446 }
2447 {
2448 let datum = Datum::decimal_with_precision(decimal_new(123, 2), 30).unwrap();
2449 let arrow_datum = get_arrow_datum(&datum).unwrap();
2450 let (array, is_scalar) = arrow_datum.get();
2451 let array = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
2452 assert!(is_scalar);
2453 assert_eq!(array.precision(), 30);
2454 assert_eq!(array.scale(), 2);
2455 assert_eq!(array.value(0), 123);
2456 }
2457 {
2458 let datum = Datum::uuid_from_str("42424242-4242-4242-4242-424242424242").unwrap();
2459 let arrow_datum = get_arrow_datum(&datum).unwrap();
2460 let (array, is_scalar) = arrow_datum.get();
2461 let array = array
2462 .as_any()
2463 .downcast_ref::<FixedSizeBinaryArray>()
2464 .unwrap();
2465 assert!(is_scalar);
2466 assert_eq!(array.value(0), [66u8; 16]);
2467 }
2468 {
2469 let datum = Datum::fixed(vec![1u8, 2, 3, 4, 5, 6, 7, 8]);
2470 let arrow_datum = get_arrow_datum(&datum).unwrap();
2471 let (array, is_scalar) = arrow_datum.get();
2472 let array = array
2473 .as_any()
2474 .downcast_ref::<FixedSizeBinaryArray>()
2475 .unwrap();
2476 assert!(is_scalar);
2477 assert_eq!(array.value_length(), 8);
2478 assert_eq!(array.value(0), &[1u8, 2, 3, 4, 5, 6, 7, 8]);
2479 }
2480 }
2481
2482 #[test]
2483 fn test_arrow_schema_to_schema_with_field_id() {
2484 let arrow_schema = ArrowSchema::new(vec![
2487 Field::new("id", DataType::Int64, false),
2488 Field::new("name", DataType::Utf8, true),
2489 Field::new("price", DataType::Decimal128(10, 2), false),
2490 Field::new(
2491 "created_at",
2492 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
2493 true,
2494 ),
2495 Field::new(
2496 "tags",
2497 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
2498 true,
2499 ),
2500 Field::new(
2501 "address",
2502 DataType::Struct(Fields::from(vec![
2503 Field::new("street", DataType::Utf8, true),
2504 Field::new("city", DataType::Utf8, false),
2505 Field::new("zip", DataType::Int32, true),
2506 ])),
2507 true,
2508 ),
2509 Field::new(
2510 "attributes",
2511 DataType::Map(
2512 Arc::new(Field::new(
2513 DEFAULT_MAP_FIELD_NAME,
2514 DataType::Struct(Fields::from(vec![
2515 Field::new("key", DataType::Utf8, false),
2516 Field::new("value", DataType::Utf8, true),
2517 ])),
2518 false,
2519 )),
2520 false,
2521 ),
2522 true,
2523 ),
2524 Field::new(
2525 "orders",
2526 DataType::List(Arc::new(Field::new(
2527 "element",
2528 DataType::Struct(Fields::from(vec![
2529 Field::new("order_id", DataType::Int64, false),
2530 Field::new("amount", DataType::Float64, false),
2531 ])),
2532 true,
2533 ))),
2534 true,
2535 ),
2536 ]);
2537
2538 let schema = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap();
2539
2540 let expected = Schema::builder()
2545 .with_fields(vec![
2546 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(),
2547 NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(),
2548 NestedField::required(
2549 3,
2550 "price",
2551 Type::Primitive(PrimitiveType::Decimal {
2552 precision: 10,
2553 scale: 2,
2554 }),
2555 )
2556 .into(),
2557 NestedField::optional(4, "created_at", Type::Primitive(PrimitiveType::Timestamptz))
2558 .into(),
2559 NestedField::optional(
2560 5,
2561 "tags",
2562 Type::List(ListType {
2563 element_field: NestedField::list_element(
2564 9,
2565 Type::Primitive(PrimitiveType::String),
2566 false,
2567 )
2568 .into(),
2569 }),
2570 )
2571 .into(),
2572 NestedField::optional(
2573 6,
2574 "address",
2575 Type::Struct(StructType::new(vec![
2576 NestedField::optional(10, "street", Type::Primitive(PrimitiveType::String))
2577 .into(),
2578 NestedField::required(11, "city", Type::Primitive(PrimitiveType::String))
2579 .into(),
2580 NestedField::optional(12, "zip", Type::Primitive(PrimitiveType::Int))
2581 .into(),
2582 ])),
2583 )
2584 .into(),
2585 NestedField::optional(
2586 7,
2587 "attributes",
2588 Type::Map(MapType {
2589 key_field: NestedField::map_key_element(
2590 13,
2591 Type::Primitive(PrimitiveType::String),
2592 )
2593 .into(),
2594 value_field: NestedField::map_value_element(
2595 14,
2596 Type::Primitive(PrimitiveType::String),
2597 false,
2598 )
2599 .into(),
2600 }),
2601 )
2602 .into(),
2603 NestedField::optional(
2604 8,
2605 "orders",
2606 Type::List(ListType {
2607 element_field: NestedField::list_element(
2608 15,
2609 Type::Struct(StructType::new(vec![
2610 NestedField::required(
2611 16,
2612 "order_id",
2613 Type::Primitive(PrimitiveType::Long),
2614 )
2615 .into(),
2616 NestedField::required(
2617 17,
2618 "amount",
2619 Type::Primitive(PrimitiveType::Double),
2620 )
2621 .into(),
2622 ])),
2623 false,
2624 )
2625 .into(),
2626 }),
2627 )
2628 .into(),
2629 ])
2630 .build()
2631 .unwrap();
2632
2633 pretty_assertions::assert_eq!(schema, expected);
2634 assert_eq!(schema.highest_field_id(), 17);
2635 }
2636}