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