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 primitive_type_to_arrow_type_with_ree(datum.data_type())
1196}
1197
1198pub(crate) fn primitive_type_to_arrow_type_with_ree(primitive_type: &PrimitiveType) -> DataType {
1202 let make_ree = |values_type: DataType| -> DataType {
1206 let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int32, false));
1207 let values_field = Arc::new(Field::new("values", values_type, true));
1208 DataType::RunEndEncoded(run_ends_field, values_field)
1209 };
1210
1211 match primitive_type {
1212 PrimitiveType::Boolean => make_ree(DataType::Boolean),
1213 PrimitiveType::Int => make_ree(DataType::Int32),
1214 PrimitiveType::Long => make_ree(DataType::Int64),
1215 PrimitiveType::Float => make_ree(DataType::Float32),
1216 PrimitiveType::Double => make_ree(DataType::Float64),
1217 PrimitiveType::Date => make_ree(DataType::Date32),
1218 PrimitiveType::Time => make_ree(DataType::Int64),
1219 PrimitiveType::Timestamp => make_ree(DataType::Int64),
1220 PrimitiveType::Timestamptz => make_ree(DataType::Int64),
1221 PrimitiveType::TimestampNs => make_ree(DataType::Int64),
1222 PrimitiveType::TimestamptzNs => make_ree(DataType::Int64),
1223 PrimitiveType::String => make_ree(DataType::Utf8),
1224 PrimitiveType::Uuid => make_ree(DataType::Binary),
1225 PrimitiveType::Fixed(_) => make_ree(DataType::Binary),
1226 PrimitiveType::Binary => make_ree(DataType::Binary),
1227 PrimitiveType::Decimal { precision, scale } => {
1228 make_ree(DataType::Decimal128(*precision as u8, *scale as i8))
1229 }
1230 }
1231}
1232
1233struct MetadataStripVisitor {
1239 field_stack: Vec<Field>,
1241}
1242
1243impl MetadataStripVisitor {
1244 fn new() -> Self {
1245 Self {
1246 field_stack: Vec::new(),
1247 }
1248 }
1249}
1250
1251impl ArrowSchemaVisitor for MetadataStripVisitor {
1252 type T = Field;
1253 type U = ArrowSchema;
1254
1255 fn before_field(&mut self, field: &FieldRef) -> Result<()> {
1256 self.field_stack.push(Field::new(
1258 field.name(),
1259 DataType::Null, field.is_nullable(),
1261 ));
1262 Ok(())
1263 }
1264
1265 fn after_field(&mut self, _field: &FieldRef) -> Result<()> {
1266 Ok(())
1267 }
1268
1269 fn schema(&mut self, _schema: &ArrowSchema, values: Vec<Self::T>) -> Result<Self::U> {
1270 Ok(ArrowSchema::new(values))
1271 }
1272
1273 fn r#struct(&mut self, _fields: &Fields, results: Vec<Self::T>) -> Result<Self::T> {
1274 let field_info = self
1276 .field_stack
1277 .pop()
1278 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in struct"))?;
1279
1280 Ok(Field::new(
1282 field_info.name(),
1283 DataType::Struct(Fields::from(results)),
1284 field_info.is_nullable(),
1285 ))
1286 }
1287
1288 fn list(&mut self, list: &DataType, value: Self::T) -> Result<Self::T> {
1289 let field_info = self
1291 .field_stack
1292 .pop()
1293 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in list"))?;
1294
1295 let list_type = match list {
1297 DataType::List(_) => DataType::List(Arc::new(value)),
1298 DataType::LargeList(_) => DataType::LargeList(Arc::new(value)),
1299 DataType::FixedSizeList(_, size) => DataType::FixedSizeList(Arc::new(value), *size),
1300 _ => {
1301 return Err(Error::new(
1302 ErrorKind::Unexpected,
1303 format!("Expected list type, got {list}"),
1304 ));
1305 }
1306 };
1307
1308 Ok(Field::new(
1309 field_info.name(),
1310 list_type,
1311 field_info.is_nullable(),
1312 ))
1313 }
1314
1315 fn map(&mut self, map: &DataType, key_value: Self::T, value: Self::T) -> Result<Self::T> {
1316 let field_info = self
1318 .field_stack
1319 .pop()
1320 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in map"))?;
1321
1322 let struct_field = Field::new(
1324 DEFAULT_MAP_FIELD_NAME,
1325 DataType::Struct(Fields::from(vec![key_value, value])),
1326 false,
1327 );
1328
1329 let sorted = match map {
1331 DataType::Map(_, sorted) => *sorted,
1332 _ => {
1333 return Err(Error::new(
1334 ErrorKind::Unexpected,
1335 format!("Expected map type, got {map}"),
1336 ));
1337 }
1338 };
1339
1340 Ok(Field::new(
1342 field_info.name(),
1343 DataType::Map(Arc::new(struct_field), sorted),
1344 field_info.is_nullable(),
1345 ))
1346 }
1347
1348 fn primitive(&mut self, p: &DataType) -> Result<Self::T> {
1349 let field_info = self.field_stack.pop().ok_or_else(|| {
1351 Error::new(ErrorKind::Unexpected, "Field stack underflow in primitive")
1352 })?;
1353
1354 Ok(Field::new(
1356 field_info.name(),
1357 p.clone(),
1358 field_info.is_nullable(),
1359 ))
1360 }
1361}
1362
1363pub fn strip_metadata_from_schema(schema: &ArrowSchema) -> Result<ArrowSchema> {
1393 let mut visitor = MetadataStripVisitor::new();
1394 visit_schema(schema, &mut visitor)
1395}
1396
1397#[cfg(test)]
1398mod tests {
1399 use std::collections::HashMap;
1400 use std::sync::Arc;
1401
1402 use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
1403
1404 use super::*;
1405 use crate::spec::decimal_utils::decimal_new;
1406 use crate::spec::{Literal, Schema};
1407
1408 fn simple_field(name: &str, ty: DataType, nullable: bool, value: &str) -> Field {
1410 Field::new(name, ty, nullable).with_metadata(HashMap::from([(
1411 PARQUET_FIELD_ID_META_KEY.to_string(),
1412 value.to_string(),
1413 )]))
1414 }
1415
1416 fn arrow_schema_for_arrow_schema_to_schema_test() -> ArrowSchema {
1417 let fields = Fields::from(vec![
1418 simple_field("key", DataType::Int32, false, "28"),
1419 simple_field("value", DataType::Utf8, true, "29"),
1420 ]);
1421
1422 let r#struct = DataType::Struct(fields);
1423 let map = DataType::Map(
1424 Arc::new(simple_field(DEFAULT_MAP_FIELD_NAME, r#struct, false, "17")),
1425 false,
1426 );
1427 let dictionary = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
1428
1429 let fields = Fields::from(vec![
1430 simple_field("aa", DataType::Int32, false, "18"),
1431 simple_field("bb", DataType::Utf8, true, "19"),
1432 simple_field(
1433 "cc",
1434 DataType::Timestamp(TimeUnit::Microsecond, None),
1435 false,
1436 "20",
1437 ),
1438 ]);
1439
1440 let r#struct = DataType::Struct(fields);
1441
1442 ArrowSchema::new(vec![
1443 simple_field("a", DataType::Int32, false, "2"),
1444 simple_field("b", DataType::Int64, false, "1"),
1445 simple_field("c", DataType::Utf8, false, "3"),
1446 simple_field("n", DataType::Utf8, false, "21"),
1447 simple_field(
1448 "d",
1449 DataType::Timestamp(TimeUnit::Microsecond, None),
1450 true,
1451 "4",
1452 ),
1453 simple_field("e", DataType::Boolean, true, "6"),
1454 simple_field("f", DataType::Float32, false, "5"),
1455 simple_field("g", DataType::Float64, false, "7"),
1456 simple_field("p", DataType::Decimal128(10, 2), false, "27"),
1457 simple_field("h", DataType::Date32, false, "8"),
1458 simple_field("i", DataType::Time64(TimeUnit::Microsecond), false, "9"),
1459 simple_field(
1460 "j",
1461 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
1462 false,
1463 "10",
1464 ),
1465 simple_field(
1466 "k",
1467 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1468 false,
1469 "12",
1470 ),
1471 simple_field("l", DataType::Binary, false, "13"),
1472 simple_field("o", DataType::LargeBinary, false, "22"),
1473 simple_field("m", DataType::FixedSizeBinary(10), false, "11"),
1474 simple_field(
1475 "list",
1476 DataType::List(Arc::new(simple_field(
1477 "element",
1478 DataType::Int32,
1479 false,
1480 "15",
1481 ))),
1482 true,
1483 "14",
1484 ),
1485 simple_field(
1486 "large_list",
1487 DataType::LargeList(Arc::new(simple_field(
1488 "element",
1489 DataType::Utf8,
1490 false,
1491 "23",
1492 ))),
1493 true,
1494 "24",
1495 ),
1496 simple_field(
1497 "fixed_list",
1498 DataType::FixedSizeList(
1499 Arc::new(simple_field("element", DataType::Binary, false, "26")),
1500 10,
1501 ),
1502 true,
1503 "25",
1504 ),
1505 simple_field("map", map, false, "16"),
1506 simple_field("struct", r#struct, false, "17"),
1507 simple_field("dictionary", dictionary, false, "30"),
1508 ])
1509 }
1510
1511 fn iceberg_schema_for_arrow_schema_to_schema_test() -> Schema {
1512 let schema_json = r#"{
1513 "type":"struct",
1514 "schema-id":0,
1515 "fields":[
1516 {
1517 "id":2,
1518 "name":"a",
1519 "required":true,
1520 "type":"int"
1521 },
1522 {
1523 "id":1,
1524 "name":"b",
1525 "required":true,
1526 "type":"long"
1527 },
1528 {
1529 "id":3,
1530 "name":"c",
1531 "required":true,
1532 "type":"string"
1533 },
1534 {
1535 "id":21,
1536 "name":"n",
1537 "required":true,
1538 "type":"string"
1539 },
1540 {
1541 "id":4,
1542 "name":"d",
1543 "required":false,
1544 "type":"timestamp"
1545 },
1546 {
1547 "id":6,
1548 "name":"e",
1549 "required":false,
1550 "type":"boolean"
1551 },
1552 {
1553 "id":5,
1554 "name":"f",
1555 "required":true,
1556 "type":"float"
1557 },
1558 {
1559 "id":7,
1560 "name":"g",
1561 "required":true,
1562 "type":"double"
1563 },
1564 {
1565 "id":27,
1566 "name":"p",
1567 "required":true,
1568 "type":"decimal(10,2)"
1569 },
1570 {
1571 "id":8,
1572 "name":"h",
1573 "required":true,
1574 "type":"date"
1575 },
1576 {
1577 "id":9,
1578 "name":"i",
1579 "required":true,
1580 "type":"time"
1581 },
1582 {
1583 "id":10,
1584 "name":"j",
1585 "required":true,
1586 "type":"timestamptz"
1587 },
1588 {
1589 "id":12,
1590 "name":"k",
1591 "required":true,
1592 "type":"timestamptz"
1593 },
1594 {
1595 "id":13,
1596 "name":"l",
1597 "required":true,
1598 "type":"binary"
1599 },
1600 {
1601 "id":22,
1602 "name":"o",
1603 "required":true,
1604 "type":"binary"
1605 },
1606 {
1607 "id":11,
1608 "name":"m",
1609 "required":true,
1610 "type":"fixed[10]"
1611 },
1612 {
1613 "id":14,
1614 "name":"list",
1615 "required": false,
1616 "type": {
1617 "type": "list",
1618 "element-id": 15,
1619 "element-required": true,
1620 "element": "int"
1621 }
1622 },
1623 {
1624 "id":24,
1625 "name":"large_list",
1626 "required": false,
1627 "type": {
1628 "type": "list",
1629 "element-id": 23,
1630 "element-required": true,
1631 "element": "string"
1632 }
1633 },
1634 {
1635 "id":25,
1636 "name":"fixed_list",
1637 "required": false,
1638 "type": {
1639 "type": "list",
1640 "element-id": 26,
1641 "element-required": true,
1642 "element": "binary"
1643 }
1644 },
1645 {
1646 "id":16,
1647 "name":"map",
1648 "required": true,
1649 "type": {
1650 "type": "map",
1651 "key-id": 28,
1652 "key": "int",
1653 "value-id": 29,
1654 "value-required": false,
1655 "value": "string"
1656 }
1657 },
1658 {
1659 "id":17,
1660 "name":"struct",
1661 "required": true,
1662 "type": {
1663 "type": "struct",
1664 "fields": [
1665 {
1666 "id":18,
1667 "name":"aa",
1668 "required":true,
1669 "type":"int"
1670 },
1671 {
1672 "id":19,
1673 "name":"bb",
1674 "required":false,
1675 "type":"string"
1676 },
1677 {
1678 "id":20,
1679 "name":"cc",
1680 "required":true,
1681 "type":"timestamp"
1682 }
1683 ]
1684 }
1685 },
1686 {
1687 "id":30,
1688 "name":"dictionary",
1689 "required":true,
1690 "type":"string"
1691 }
1692 ],
1693 "identifier-field-ids":[]
1694 }"#;
1695
1696 let schema: Schema = serde_json::from_str(schema_json).unwrap();
1697 schema
1698 }
1699
1700 #[test]
1701 fn test_arrow_schema_to_schema() {
1702 let arrow_schema = arrow_schema_for_arrow_schema_to_schema_test();
1703 let schema = iceberg_schema_for_arrow_schema_to_schema_test();
1704 let converted_schema = arrow_schema_to_schema(&arrow_schema).unwrap();
1705 pretty_assertions::assert_eq!(converted_schema, schema);
1706 }
1707
1708 fn arrow_schema_for_schema_to_arrow_schema_test() -> ArrowSchema {
1709 let fields = Fields::from(vec![
1710 simple_field("key", DataType::Int32, false, "28"),
1711 simple_field("value", DataType::Utf8, true, "29"),
1712 ]);
1713
1714 let r#struct = DataType::Struct(fields);
1715 let map = DataType::Map(
1716 Arc::new(Field::new(DEFAULT_MAP_FIELD_NAME, r#struct, false)),
1717 false,
1718 );
1719
1720 let fields = Fields::from(vec![
1721 simple_field("aa", DataType::Int32, false, "18"),
1722 simple_field("bb", DataType::Utf8, true, "19"),
1723 simple_field(
1724 "cc",
1725 DataType::Timestamp(TimeUnit::Microsecond, None),
1726 false,
1727 "20",
1728 ),
1729 ]);
1730
1731 let r#struct = DataType::Struct(fields);
1732
1733 ArrowSchema::new(vec![
1734 simple_field("a", DataType::Int32, false, "2"),
1735 simple_field("b", DataType::Int64, false, "1"),
1736 simple_field("c", DataType::Utf8, false, "3"),
1737 simple_field("n", DataType::Utf8, false, "21"),
1738 simple_field(
1739 "d",
1740 DataType::Timestamp(TimeUnit::Microsecond, None),
1741 true,
1742 "4",
1743 ),
1744 simple_field("e", DataType::Boolean, true, "6"),
1745 simple_field("f", DataType::Float32, false, "5"),
1746 simple_field("g", DataType::Float64, false, "7"),
1747 simple_field("p", DataType::Decimal128(10, 2), false, "27"),
1748 simple_field("h", DataType::Date32, false, "8"),
1749 simple_field("i", DataType::Time64(TimeUnit::Microsecond), false, "9"),
1750 simple_field(
1751 "j",
1752 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1753 false,
1754 "10",
1755 ),
1756 simple_field(
1757 "k",
1758 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
1759 false,
1760 "12",
1761 ),
1762 simple_field("l", DataType::LargeBinary, false, "13"),
1763 simple_field("o", DataType::LargeBinary, false, "22"),
1764 simple_field("m", DataType::FixedSizeBinary(10), false, "11"),
1765 simple_field(
1766 "list",
1767 DataType::List(Arc::new(simple_field(
1768 "element",
1769 DataType::Int32,
1770 false,
1771 "15",
1772 ))),
1773 true,
1774 "14",
1775 ),
1776 simple_field(
1777 "large_list",
1778 DataType::List(Arc::new(simple_field(
1779 "element",
1780 DataType::Utf8,
1781 false,
1782 "23",
1783 ))),
1784 true,
1785 "24",
1786 ),
1787 simple_field(
1788 "fixed_list",
1789 DataType::List(Arc::new(simple_field(
1790 "element",
1791 DataType::LargeBinary,
1792 false,
1793 "26",
1794 ))),
1795 true,
1796 "25",
1797 ),
1798 simple_field("map", map, false, "16"),
1799 simple_field("struct", r#struct, false, "17"),
1800 simple_field("uuid", DataType::FixedSizeBinary(16), false, "30"),
1801 Field::new(
1802 "v",
1803 DataType::Struct(Fields::from(vec![
1804 Field::new("metadata", DataType::Binary, false),
1805 Field::new("value", DataType::Binary, true),
1806 ])),
1807 true,
1808 )
1809 .with_metadata(HashMap::from([
1810 (PARQUET_FIELD_ID_META_KEY.to_string(), "31".to_string()),
1811 (
1812 arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
1813 "arrow.parquet.variant".to_string(),
1814 ),
1815 ])),
1816 ])
1817 }
1818
1819 fn iceberg_schema_for_schema_to_arrow_schema() -> Schema {
1820 let schema_json = r#"{
1821 "type":"struct",
1822 "schema-id":0,
1823 "fields":[
1824 {
1825 "id":2,
1826 "name":"a",
1827 "required":true,
1828 "type":"int"
1829 },
1830 {
1831 "id":1,
1832 "name":"b",
1833 "required":true,
1834 "type":"long"
1835 },
1836 {
1837 "id":3,
1838 "name":"c",
1839 "required":true,
1840 "type":"string"
1841 },
1842 {
1843 "id":21,
1844 "name":"n",
1845 "required":true,
1846 "type":"string"
1847 },
1848 {
1849 "id":4,
1850 "name":"d",
1851 "required":false,
1852 "type":"timestamp"
1853 },
1854 {
1855 "id":6,
1856 "name":"e",
1857 "required":false,
1858 "type":"boolean"
1859 },
1860 {
1861 "id":5,
1862 "name":"f",
1863 "required":true,
1864 "type":"float"
1865 },
1866 {
1867 "id":7,
1868 "name":"g",
1869 "required":true,
1870 "type":"double"
1871 },
1872 {
1873 "id":27,
1874 "name":"p",
1875 "required":true,
1876 "type":"decimal(10,2)"
1877 },
1878 {
1879 "id":8,
1880 "name":"h",
1881 "required":true,
1882 "type":"date"
1883 },
1884 {
1885 "id":9,
1886 "name":"i",
1887 "required":true,
1888 "type":"time"
1889 },
1890 {
1891 "id":10,
1892 "name":"j",
1893 "required":true,
1894 "type":"timestamptz"
1895 },
1896 {
1897 "id":12,
1898 "name":"k",
1899 "required":true,
1900 "type":"timestamptz"
1901 },
1902 {
1903 "id":13,
1904 "name":"l",
1905 "required":true,
1906 "type":"binary"
1907 },
1908 {
1909 "id":22,
1910 "name":"o",
1911 "required":true,
1912 "type":"binary"
1913 },
1914 {
1915 "id":11,
1916 "name":"m",
1917 "required":true,
1918 "type":"fixed[10]"
1919 },
1920 {
1921 "id":14,
1922 "name":"list",
1923 "required": false,
1924 "type": {
1925 "type": "list",
1926 "element-id": 15,
1927 "element-required": true,
1928 "element": "int"
1929 }
1930 },
1931 {
1932 "id":24,
1933 "name":"large_list",
1934 "required": false,
1935 "type": {
1936 "type": "list",
1937 "element-id": 23,
1938 "element-required": true,
1939 "element": "string"
1940 }
1941 },
1942 {
1943 "id":25,
1944 "name":"fixed_list",
1945 "required": false,
1946 "type": {
1947 "type": "list",
1948 "element-id": 26,
1949 "element-required": true,
1950 "element": "binary"
1951 }
1952 },
1953 {
1954 "id":16,
1955 "name":"map",
1956 "required": true,
1957 "type": {
1958 "type": "map",
1959 "key-id": 28,
1960 "key": "int",
1961 "value-id": 29,
1962 "value-required": false,
1963 "value": "string"
1964 }
1965 },
1966 {
1967 "id":17,
1968 "name":"struct",
1969 "required": true,
1970 "type": {
1971 "type": "struct",
1972 "fields": [
1973 {
1974 "id":18,
1975 "name":"aa",
1976 "required":true,
1977 "type":"int"
1978 },
1979 {
1980 "id":19,
1981 "name":"bb",
1982 "required":false,
1983 "type":"string"
1984 },
1985 {
1986 "id":20,
1987 "name":"cc",
1988 "required":true,
1989 "type":"timestamp"
1990 }
1991 ]
1992 }
1993 },
1994 {
1995 "id":30,
1996 "name":"uuid",
1997 "required":true,
1998 "type":"uuid"
1999 },
2000 {
2001 "id":31,
2002 "name":"v",
2003 "required":false,
2004 "type":"variant"
2005 }
2006 ],
2007 "identifier-field-ids":[]
2008 }"#;
2009
2010 let schema: Schema = serde_json::from_str(schema_json).unwrap();
2011 schema
2012 }
2013
2014 #[test]
2015 fn test_schema_to_arrow_schema() {
2016 let arrow_schema = arrow_schema_for_schema_to_arrow_schema_test();
2017 let schema = iceberg_schema_for_schema_to_arrow_schema();
2018 let converted_arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2019 assert_eq!(converted_arrow_schema, arrow_schema);
2020 }
2021
2022 #[test]
2023 fn test_variant_type_to_arrow_type() {
2024 let arrow_type = type_to_arrow_type(&Type::Variant(VariantType)).unwrap();
2027 assert_eq!(
2028 arrow_type,
2029 DataType::Struct(Fields::from(vec![
2030 Field::new("metadata", DataType::Binary, false),
2031 Field::new("value", DataType::Binary, true),
2032 ]))
2033 );
2034 }
2035
2036 #[test]
2037 fn test_variant_field_carries_arrow_extension_type() {
2038 let schema = Schema::builder()
2041 .with_fields(vec![
2042 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2043 ])
2044 .build()
2045 .unwrap();
2046
2047 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2048 let field = arrow_schema.field_with_name("v").unwrap();
2049
2050 assert_eq!(field.extension_type_name(), Some("arrow.parquet.variant"));
2051 assert_eq!(
2053 field.metadata().get(PARQUET_FIELD_ID_META_KEY),
2054 Some(&"1".to_string())
2055 );
2056 assert_eq!(
2057 field.data_type(),
2058 &DataType::Struct(Fields::from(vec![
2059 Field::new("metadata", DataType::Binary, false),
2060 Field::new("value", DataType::Binary, true),
2061 ]))
2062 );
2063 }
2064
2065 #[test]
2066 fn test_variant_nested_in_list_and_map_carries_arrow_extension_type() {
2067 let schema = Schema::builder()
2071 .with_fields(vec![
2072 NestedField::optional(
2073 1,
2074 "l",
2075 Type::List(ListType::new(
2076 NestedField::optional(2, "element", Type::Variant(VariantType)).into(),
2077 )),
2078 )
2079 .into(),
2080 NestedField::optional(
2081 3,
2082 "m",
2083 Type::Map(MapType::new(
2084 NestedField::map_key_element(4, Type::Primitive(PrimitiveType::String))
2085 .into(),
2086 NestedField::map_value_element(5, Type::Variant(VariantType), false).into(),
2087 )),
2088 )
2089 .into(),
2090 ])
2091 .build()
2092 .unwrap();
2093
2094 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2095
2096 let DataType::List(element) = arrow_schema.field_with_name("l").unwrap().data_type() else {
2097 panic!("expected a list");
2098 };
2099 assert_eq!(element.extension_type_name(), Some("arrow.parquet.variant"));
2100
2101 let DataType::Map(entries, _) = arrow_schema.field_with_name("m").unwrap().data_type()
2102 else {
2103 panic!("expected a map");
2104 };
2105 let DataType::Struct(kv) = entries.data_type() else {
2106 panic!("expected a key_value struct");
2107 };
2108 let value = kv.iter().find(|f| f.name() == "value").unwrap();
2109 assert_eq!(value.extension_type_name(), Some("arrow.parquet.variant"));
2110 }
2111
2112 fn variant_storage() -> DataType {
2115 DataType::Struct(Fields::from(vec![
2116 Field::new("metadata", DataType::Binary, false),
2117 Field::new("value", DataType::Binary, true),
2118 ]))
2119 }
2120
2121 #[test]
2122 fn test_variant_arrow_field_folds_to_iceberg_variant() {
2123 let field = simple_field("v", variant_storage(), true, "1")
2127 .with_extension_type(VariantExtensionType);
2128 let arrow_schema = ArrowSchema::new(vec![field]);
2129
2130 let converted = arrow_schema_to_schema(&arrow_schema).unwrap();
2131 let expected = Schema::builder()
2132 .with_fields(vec![
2133 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2134 ])
2135 .build()
2136 .unwrap();
2137 pretty_assertions::assert_eq!(converted, expected);
2138 }
2139
2140 #[test]
2141 fn test_variant_schema_round_trips() {
2142 let schema = Schema::builder()
2145 .with_fields(vec![
2146 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2147 NestedField::optional(
2148 2,
2149 "s",
2150 Type::Struct(StructType::new(vec![
2151 NestedField::optional(3, "sv", Type::Variant(VariantType)).into(),
2152 ])),
2153 )
2154 .into(),
2155 NestedField::optional(
2156 4,
2157 "l",
2158 Type::List(ListType::new(
2159 NestedField::optional(5, "element", Type::Variant(VariantType)).into(),
2160 )),
2161 )
2162 .into(),
2163 NestedField::optional(
2164 6,
2165 "m",
2166 Type::Map(MapType::new(
2167 NestedField::map_key_element(7, Type::Primitive(PrimitiveType::String))
2168 .into(),
2169 NestedField::map_value_element(8, Type::Variant(VariantType), false).into(),
2170 )),
2171 )
2172 .into(),
2173 ])
2174 .build()
2175 .unwrap();
2176
2177 let arrow_schema = schema_to_arrow_schema(&schema).unwrap();
2178 let round_tripped = arrow_schema_to_schema(&arrow_schema).unwrap();
2179 pretty_assertions::assert_eq!(round_tripped, schema);
2180 }
2181
2182 #[test]
2183 fn test_variant_recognized_with_auto_assigned_ids() {
2184 let field =
2187 Field::new("v", variant_storage(), true).with_extension_type(VariantExtensionType);
2188 let arrow_schema = ArrowSchema::new(vec![field]);
2189
2190 let converted = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap();
2191 let expected = Schema::builder()
2192 .with_fields(vec![
2193 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
2194 ])
2195 .build()
2196 .unwrap();
2197 pretty_assertions::assert_eq!(converted, expected);
2198 }
2199
2200 #[test]
2201 fn test_variant_extension_on_non_struct_storage_is_rejected() {
2202 let field = Field::new("v", DataType::Int32, true).with_metadata(HashMap::from([
2205 (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()),
2206 (
2207 arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
2208 VariantExtensionType::NAME.to_string(),
2209 ),
2210 ]));
2211 let arrow_schema = ArrowSchema::new(vec![field]);
2212
2213 let err = arrow_schema_to_schema(&arrow_schema).unwrap_err();
2214 assert!(
2215 err.to_string().contains("requires Struct storage"),
2216 "unexpected error: {err}"
2217 );
2218 }
2219
2220 #[test]
2221 fn test_type_conversion() {
2222 {
2224 let arrow_type = DataType::Int32;
2225 let iceberg_type = Type::Primitive(PrimitiveType::Int);
2226 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2227 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2228 }
2229
2230 {
2232 let arrow_type = DataType::Struct(Fields::from(vec![
2234 Field::new("a", DataType::Int64, false),
2235 Field::new("b", DataType::Utf8, true),
2236 ]));
2237 assert_eq!(
2238 &arrow_type_to_type(&arrow_type).unwrap_err().to_string(),
2239 "DataInvalid => Field id not found in metadata"
2240 );
2241
2242 let arrow_type = DataType::Struct(Fields::from(vec![
2243 Field::new("a", DataType::Int64, false).with_metadata(HashMap::from_iter([(
2244 PARQUET_FIELD_ID_META_KEY.to_string(),
2245 1.to_string(),
2246 )])),
2247 Field::new("b", DataType::Utf8, true).with_metadata(HashMap::from_iter([(
2248 PARQUET_FIELD_ID_META_KEY.to_string(),
2249 2.to_string(),
2250 )])),
2251 ]));
2252 let iceberg_type = Type::Struct(StructType::new(vec![
2253 NestedField {
2254 id: 1,
2255 doc: None,
2256 name: "a".to_string(),
2257 required: true,
2258 field_type: Box::new(Type::Primitive(PrimitiveType::Long)),
2259 initial_default: None,
2260 write_default: None,
2261 }
2262 .into(),
2263 NestedField {
2264 id: 2,
2265 doc: None,
2266 name: "b".to_string(),
2267 required: false,
2268 field_type: Box::new(Type::Primitive(PrimitiveType::String)),
2269 initial_default: None,
2270 write_default: None,
2271 }
2272 .into(),
2273 ]));
2274 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2275 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2276
2277 let iceberg_type = Type::Struct(StructType::new(vec![
2279 NestedField {
2280 id: 1,
2281 doc: None,
2282 name: "a".to_string(),
2283 required: true,
2284 field_type: Box::new(Type::Primitive(PrimitiveType::Long)),
2285 initial_default: Some(Literal::Primitive(PrimitiveLiteral::Int(114514))),
2286 write_default: None,
2287 }
2288 .into(),
2289 NestedField {
2290 id: 2,
2291 doc: None,
2292 name: "b".to_string(),
2293 required: false,
2294 field_type: Box::new(Type::Primitive(PrimitiveType::String)),
2295 initial_default: None,
2296 write_default: Some(Literal::Primitive(PrimitiveLiteral::String(
2297 "514".to_string(),
2298 ))),
2299 }
2300 .into(),
2301 ]));
2302 assert_eq!(arrow_type, type_to_arrow_type(&iceberg_type).unwrap());
2303 }
2304
2305 {
2307 let arrow_type =
2308 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int8));
2309 let iceberg_type = Type::Primitive(PrimitiveType::Int);
2310 assert_eq!(
2311 iceberg_type,
2312 arrow_type_to_type(&arrow_type).unwrap(),
2313 "Expected dictionary conversion to use the contained value"
2314 );
2315
2316 let arrow_type =
2317 DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Boolean));
2318 let iceberg_type = Type::Primitive(PrimitiveType::Boolean);
2319 assert_eq!(iceberg_type, arrow_type_to_type(&arrow_type).unwrap());
2320 }
2321 }
2322
2323 #[test]
2324 fn test_unsigned_integer_type_conversion() {
2325 let test_cases = vec![
2326 (DataType::UInt8, PrimitiveType::Int),
2327 (DataType::UInt16, PrimitiveType::Int),
2328 (DataType::UInt32, PrimitiveType::Long),
2329 ];
2330
2331 for (arrow_type, expected_iceberg_type) in test_cases {
2332 let arrow_field = Field::new("test", arrow_type.clone(), false).with_metadata(
2333 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2334 );
2335 let arrow_schema = ArrowSchema::new(vec![arrow_field]);
2336
2337 let iceberg_schema = arrow_schema_to_schema(&arrow_schema).unwrap();
2338 let iceberg_field = iceberg_schema.as_struct().fields().first().unwrap();
2339
2340 assert!(
2341 matches!(iceberg_field.field_type.as_ref(), Type::Primitive(t) if *t == expected_iceberg_type),
2342 "Expected {arrow_type:?} to map to {expected_iceberg_type:?}"
2343 );
2344 }
2345
2346 {
2348 let arrow_field = Field::new("test", DataType::UInt64, false).with_metadata(
2349 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2350 );
2351 let arrow_schema = ArrowSchema::new(vec![arrow_field]);
2352
2353 let result = arrow_schema_to_schema(&arrow_schema);
2354 assert!(result.is_err());
2355 assert!(
2356 result
2357 .unwrap_err()
2358 .to_string()
2359 .contains("UInt64 is not supported")
2360 );
2361 }
2362 }
2363
2364 #[test]
2365 fn test_datum_conversion() {
2366 {
2367 let datum = Datum::bool(true);
2368 let arrow_datum = get_arrow_datum(&datum).unwrap();
2369 let (array, is_scalar) = arrow_datum.get();
2370 let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
2371 assert!(is_scalar);
2372 assert!(array.value(0));
2373 }
2374 {
2375 let datum = Datum::int(42);
2376 let arrow_datum = get_arrow_datum(&datum).unwrap();
2377 let (array, is_scalar) = arrow_datum.get();
2378 let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
2379 assert!(is_scalar);
2380 assert_eq!(array.value(0), 42);
2381 }
2382 {
2383 let datum = Datum::long(42);
2384 let arrow_datum = get_arrow_datum(&datum).unwrap();
2385 let (array, is_scalar) = arrow_datum.get();
2386 let array = array.as_any().downcast_ref::<Int64Array>().unwrap();
2387 assert!(is_scalar);
2388 assert_eq!(array.value(0), 42);
2389 }
2390 {
2391 let datum = Datum::float(42.42);
2392 let arrow_datum = get_arrow_datum(&datum).unwrap();
2393 let (array, is_scalar) = arrow_datum.get();
2394 let array = array.as_any().downcast_ref::<Float32Array>().unwrap();
2395 assert!(is_scalar);
2396 assert_eq!(array.value(0), 42.42);
2397 }
2398 {
2399 let datum = Datum::double(42.42);
2400 let arrow_datum = get_arrow_datum(&datum).unwrap();
2401 let (array, is_scalar) = arrow_datum.get();
2402 let array = array.as_any().downcast_ref::<Float64Array>().unwrap();
2403 assert!(is_scalar);
2404 assert_eq!(array.value(0), 42.42);
2405 }
2406 {
2407 let datum = Datum::string("abc");
2408 let arrow_datum = get_arrow_datum(&datum).unwrap();
2409 let (array, is_scalar) = arrow_datum.get();
2410 let array = array.as_any().downcast_ref::<StringArray>().unwrap();
2411 assert!(is_scalar);
2412 assert_eq!(array.value(0), "abc");
2413 }
2414 {
2415 let datum = Datum::binary(vec![1, 2, 3, 4]);
2416 let arrow_datum = get_arrow_datum(&datum).unwrap();
2417 let (array, is_scalar) = arrow_datum.get();
2418 let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
2419 assert!(is_scalar);
2420 assert_eq!(array.value(0), &[1, 2, 3, 4]);
2421 }
2422 {
2423 let datum = Datum::date(42);
2424 let arrow_datum = get_arrow_datum(&datum).unwrap();
2425 let (array, is_scalar) = arrow_datum.get();
2426 let array = array.as_any().downcast_ref::<Date32Array>().unwrap();
2427 assert!(is_scalar);
2428 assert_eq!(array.value(0), 42);
2429 }
2430 {
2431 let datum = Datum::timestamp_micros(42);
2432 let arrow_datum = get_arrow_datum(&datum).unwrap();
2433 let (array, is_scalar) = arrow_datum.get();
2434 let array = array
2435 .as_any()
2436 .downcast_ref::<TimestampMicrosecondArray>()
2437 .unwrap();
2438 assert!(is_scalar);
2439 assert_eq!(array.value(0), 42);
2440 }
2441 {
2442 let datum = Datum::timestamptz_micros(42);
2443 let arrow_datum = get_arrow_datum(&datum).unwrap();
2444 let (array, is_scalar) = arrow_datum.get();
2445 let array = array
2446 .as_any()
2447 .downcast_ref::<TimestampMicrosecondArray>()
2448 .unwrap();
2449 assert!(is_scalar);
2450 assert_eq!(array.timezone(), Some("+00:00"));
2451 assert_eq!(array.value(0), 42);
2452 }
2453 {
2454 let datum = Datum::decimal_with_precision(decimal_new(123, 2), 30).unwrap();
2455 let arrow_datum = get_arrow_datum(&datum).unwrap();
2456 let (array, is_scalar) = arrow_datum.get();
2457 let array = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
2458 assert!(is_scalar);
2459 assert_eq!(array.precision(), 30);
2460 assert_eq!(array.scale(), 2);
2461 assert_eq!(array.value(0), 123);
2462 }
2463 {
2464 let datum = Datum::uuid_from_str("42424242-4242-4242-4242-424242424242").unwrap();
2465 let arrow_datum = get_arrow_datum(&datum).unwrap();
2466 let (array, is_scalar) = arrow_datum.get();
2467 let array = array
2468 .as_any()
2469 .downcast_ref::<FixedSizeBinaryArray>()
2470 .unwrap();
2471 assert!(is_scalar);
2472 assert_eq!(array.value(0), [66u8; 16]);
2473 }
2474 {
2475 let datum = Datum::fixed(vec![1u8, 2, 3, 4, 5, 6, 7, 8]);
2476 let arrow_datum = get_arrow_datum(&datum).unwrap();
2477 let (array, is_scalar) = arrow_datum.get();
2478 let array = array
2479 .as_any()
2480 .downcast_ref::<FixedSizeBinaryArray>()
2481 .unwrap();
2482 assert!(is_scalar);
2483 assert_eq!(array.value_length(), 8);
2484 assert_eq!(array.value(0), &[1u8, 2, 3, 4, 5, 6, 7, 8]);
2485 }
2486 }
2487
2488 #[test]
2489 fn test_arrow_schema_to_schema_with_field_id() {
2490 let arrow_schema = ArrowSchema::new(vec![
2493 Field::new("id", DataType::Int64, false),
2494 Field::new("name", DataType::Utf8, true),
2495 Field::new("price", DataType::Decimal128(10, 2), false),
2496 Field::new(
2497 "created_at",
2498 DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
2499 true,
2500 ),
2501 Field::new(
2502 "tags",
2503 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
2504 true,
2505 ),
2506 Field::new(
2507 "address",
2508 DataType::Struct(Fields::from(vec![
2509 Field::new("street", DataType::Utf8, true),
2510 Field::new("city", DataType::Utf8, false),
2511 Field::new("zip", DataType::Int32, true),
2512 ])),
2513 true,
2514 ),
2515 Field::new(
2516 "attributes",
2517 DataType::Map(
2518 Arc::new(Field::new(
2519 DEFAULT_MAP_FIELD_NAME,
2520 DataType::Struct(Fields::from(vec![
2521 Field::new("key", DataType::Utf8, false),
2522 Field::new("value", DataType::Utf8, true),
2523 ])),
2524 false,
2525 )),
2526 false,
2527 ),
2528 true,
2529 ),
2530 Field::new(
2531 "orders",
2532 DataType::List(Arc::new(Field::new(
2533 "element",
2534 DataType::Struct(Fields::from(vec![
2535 Field::new("order_id", DataType::Int64, false),
2536 Field::new("amount", DataType::Float64, false),
2537 ])),
2538 true,
2539 ))),
2540 true,
2541 ),
2542 ]);
2543
2544 let schema = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap();
2545
2546 let expected = Schema::builder()
2551 .with_fields(vec![
2552 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(),
2553 NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(),
2554 NestedField::required(
2555 3,
2556 "price",
2557 Type::Primitive(PrimitiveType::Decimal {
2558 precision: 10,
2559 scale: 2,
2560 }),
2561 )
2562 .into(),
2563 NestedField::optional(4, "created_at", Type::Primitive(PrimitiveType::Timestamptz))
2564 .into(),
2565 NestedField::optional(
2566 5,
2567 "tags",
2568 Type::List(ListType {
2569 element_field: NestedField::list_element(
2570 9,
2571 Type::Primitive(PrimitiveType::String),
2572 false,
2573 )
2574 .into(),
2575 }),
2576 )
2577 .into(),
2578 NestedField::optional(
2579 6,
2580 "address",
2581 Type::Struct(StructType::new(vec![
2582 NestedField::optional(10, "street", Type::Primitive(PrimitiveType::String))
2583 .into(),
2584 NestedField::required(11, "city", Type::Primitive(PrimitiveType::String))
2585 .into(),
2586 NestedField::optional(12, "zip", Type::Primitive(PrimitiveType::Int))
2587 .into(),
2588 ])),
2589 )
2590 .into(),
2591 NestedField::optional(
2592 7,
2593 "attributes",
2594 Type::Map(MapType {
2595 key_field: NestedField::map_key_element(
2596 13,
2597 Type::Primitive(PrimitiveType::String),
2598 )
2599 .into(),
2600 value_field: NestedField::map_value_element(
2601 14,
2602 Type::Primitive(PrimitiveType::String),
2603 false,
2604 )
2605 .into(),
2606 }),
2607 )
2608 .into(),
2609 NestedField::optional(
2610 8,
2611 "orders",
2612 Type::List(ListType {
2613 element_field: NestedField::list_element(
2614 15,
2615 Type::Struct(StructType::new(vec![
2616 NestedField::required(
2617 16,
2618 "order_id",
2619 Type::Primitive(PrimitiveType::Long),
2620 )
2621 .into(),
2622 NestedField::required(
2623 17,
2624 "amount",
2625 Type::Primitive(PrimitiveType::Double),
2626 )
2627 .into(),
2628 ])),
2629 false,
2630 )
2631 .into(),
2632 }),
2633 )
2634 .into(),
2635 ])
2636 .build()
2637 .unwrap();
2638
2639 pretty_assertions::assert_eq!(schema, expected);
2640 assert_eq!(schema.highest_field_id(), 17);
2641 }
2642}