1use std::sync::Arc;
20
21use arrow_schema::DataType;
22use arrow_schema::Field;
23use arrow_schema::FieldRef;
24use arrow_schema::Fields;
25use arrow_schema::Schema;
26use arrow_schema::SchemaBuilder;
27use arrow_schema::SchemaRef;
28use arrow_schema::TimeUnit as ArrowTimeUnit;
29use vortex_array::dtype::DType;
30use vortex_array::dtype::DecimalDType;
31use vortex_array::dtype::FieldName;
32use vortex_array::dtype::Nullability;
33use vortex_array::dtype::PType;
34use vortex_array::dtype::StructFields;
35use vortex_array::extension::datetime::AnyTemporal;
36use vortex_array::extension::datetime::Date;
37use vortex_array::extension::datetime::TemporalMetadata;
38use vortex_array::extension::datetime::Time;
39use vortex_array::extension::datetime::TimeUnit;
40use vortex_array::extension::datetime::Timestamp;
41use vortex_error::VortexExpect;
42use vortex_error::VortexResult;
43use vortex_error::vortex_bail;
44use vortex_error::vortex_ensure;
45use vortex_error::vortex_ensure_eq;
46use vortex_error::vortex_err;
47use vortex_error::vortex_panic;
48
49#[deprecated(
51 note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
52)]
53pub trait FromArrowType<T>: Sized {
54 #[deprecated(
56 note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
57 )]
58 fn from_arrow(value: T) -> Self;
59}
60
61#[deprecated(
63 note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
64)]
65pub trait TryFromArrowType<T>: Sized {
66 #[deprecated(
68 note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
69 )]
70 fn try_from_arrow(value: T) -> VortexResult<Self>;
71}
72
73#[deprecated(note = "Use `ArrowSession::to_arrow_schema` / `to_arrow_datatype` instead")]
78pub trait ToArrowType {
79 #[deprecated(note = "Use `ArrowSession::to_arrow_schema` instead")]
85 fn to_arrow_schema(&self) -> VortexResult<Schema>;
86
87 #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")]
94 fn to_arrow_dtype(&self) -> VortexResult<DataType>;
95}
96
97pub(crate) fn from_arrow_time_unit(value: ArrowTimeUnit) -> TimeUnit {
99 match value {
100 ArrowTimeUnit::Second => TimeUnit::Seconds,
101 ArrowTimeUnit::Millisecond => TimeUnit::Milliseconds,
102 ArrowTimeUnit::Microsecond => TimeUnit::Microseconds,
103 ArrowTimeUnit::Nanosecond => TimeUnit::Nanoseconds,
104 }
105}
106
107pub(crate) fn to_arrow_time_unit(value: TimeUnit) -> VortexResult<ArrowTimeUnit> {
113 Ok(match value {
114 TimeUnit::Seconds => ArrowTimeUnit::Second,
115 TimeUnit::Milliseconds => ArrowTimeUnit::Millisecond,
116 TimeUnit::Microseconds => ArrowTimeUnit::Microsecond,
117 TimeUnit::Nanoseconds => ArrowTimeUnit::Nanosecond,
118 _ => vortex_bail!("Cannot convert {value} to Arrow TimeUnit"),
119 })
120}
121
122pub(crate) fn ptype_from_arrow(value: &DataType) -> VortexResult<PType> {
124 match value {
125 DataType::Int8 => Ok(PType::I8),
126 DataType::Int16 => Ok(PType::I16),
127 DataType::Int32 => Ok(PType::I32),
128 DataType::Int64 => Ok(PType::I64),
129 DataType::UInt8 => Ok(PType::U8),
130 DataType::UInt16 => Ok(PType::U16),
131 DataType::UInt32 => Ok(PType::U32),
132 DataType::UInt64 => Ok(PType::U64),
133 DataType::Float16 => Ok(PType::F16),
134 DataType::Float32 => Ok(PType::F32),
135 DataType::Float64 => Ok(PType::F64),
136 _ => Err(vortex_err!(
137 "Arrow datatype {:?} cannot be converted to ptype",
138 value
139 )),
140 }
141}
142
143pub(crate) fn decimal_dtype_from_arrow(value: &DataType) -> VortexResult<DecimalDType> {
145 match value {
146 DataType::Decimal32(precision, scale)
147 | DataType::Decimal64(precision, scale)
148 | DataType::Decimal128(precision, scale)
149 | DataType::Decimal256(precision, scale) => DecimalDType::try_new(*precision, *scale),
150
151 _ => Err(vortex_err!(
152 "Arrow datatype {:?} cannot be converted to DecimalDType",
153 value
154 )),
155 }
156}
157
158pub(crate) fn from_arrow_data_type(
161 data_type: &DataType,
162 nullability: Nullability,
163) -> VortexResult<DType> {
164 if data_type.is_integer() || data_type.is_floating() {
165 return Ok(DType::Primitive(ptype_from_arrow(data_type)?, nullability));
166 }
167
168 Ok(match data_type {
169 DataType::Null => DType::Null,
170 DataType::Decimal32(precision, scale)
171 | DataType::Decimal64(precision, scale)
172 | DataType::Decimal128(precision, scale)
173 | DataType::Decimal256(precision, scale) => {
174 DType::Decimal(DecimalDType::new(*precision, *scale), nullability)
175 }
176 DataType::Boolean => DType::Bool(nullability),
177 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DType::Utf8(nullability),
178 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
179 DType::Binary(nullability)
180 }
181 DataType::Date32 => DType::Extension(Date::new(TimeUnit::Days, nullability).erased()),
182 DataType::Date64 => {
183 DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased())
184 }
185 DataType::Time32(unit) => {
186 DType::Extension(Time::new(from_arrow_time_unit(*unit), nullability).erased())
187 }
188 DataType::Time64(unit) => {
189 DType::Extension(Time::new(from_arrow_time_unit(*unit), nullability).erased())
190 }
191 DataType::Timestamp(unit, tz) => DType::Extension(
192 Timestamp::new_with_tz(from_arrow_time_unit(*unit), tz.clone(), nullability).erased(),
193 ),
194 DataType::List(e)
195 | DataType::LargeList(e)
196 | DataType::ListView(e)
197 | DataType::LargeListView(e) => {
198 DType::List(Arc::new(from_arrow_field_naive(e.as_ref())?), nullability)
199 }
200 DataType::FixedSizeList(e, size) => DType::FixedSizeList(
201 Arc::new(from_arrow_field_naive(e.as_ref())?),
202 *size as u32,
203 nullability,
204 ),
205 DataType::Struct(f) => DType::Struct(from_arrow_fields_naive(f)?, nullability),
206 DataType::Map(entries, keys_sorted) => {
207 vortex_ensure!(
208 !entries.is_nullable(),
209 "Arrow map entries field must be non-nullable"
210 );
211 let DataType::Struct(fields) = entries.data_type() else {
212 vortex_bail!(
213 "Arrow map entries field must have Struct type, got {:?}",
214 entries.data_type()
215 );
216 };
217 vortex_ensure_eq!(
218 fields.len(),
219 2,
220 InvalidArgument: "Arrow map entries struct must contain exactly two fields"
221 );
222 let key = &fields[0];
223 let value = &fields[1];
224 vortex_ensure!(
225 !key.is_nullable(),
226 "Arrow map key field must be non-nullable"
227 );
228 DType::map(
229 from_arrow_field_naive(key.as_ref())?,
230 from_arrow_field_naive(value.as_ref())?,
231 *keys_sorted,
232 nullability,
233 )?
234 }
235 DataType::Dictionary(_, value_type) => {
236 from_arrow_data_type(value_type.as_ref(), nullability)?
237 }
238 DataType::RunEndEncoded(_, value_type) => {
239 from_arrow_data_type(value_type.data_type(), nullability)?
240 }
241 _ => vortex_bail!("Arrow data type not supported: {data_type:?}"),
242 })
243}
244
245pub(crate) fn from_arrow_field_naive(field: &Field) -> VortexResult<DType> {
248 if field
249 .metadata()
250 .get("ARROW:extension:name")
251 .map(|s| s.as_str())
252 == Some("arrow.parquet.variant")
253 {
254 return Ok(DType::Variant(field.is_nullable().into()));
255 }
256 from_arrow_data_type(field.data_type(), field.is_nullable().into())
257}
258
259pub(crate) fn from_arrow_fields_naive(fields: &Fields) -> VortexResult<StructFields> {
261 fields
262 .into_iter()
263 .map(|f| {
264 Ok((
265 FieldName::from(f.name().as_str()),
266 from_arrow_field_naive(f.as_ref())?,
267 ))
268 })
269 .collect::<VortexResult<StructFields>>()
270}
271
272pub(crate) fn from_arrow_schema_naive(schema: &Schema) -> VortexResult<DType> {
274 Ok(DType::Struct(
275 from_arrow_fields_naive(schema.fields())?,
276 Nullability::NonNullable, ))
278}
279
280fn to_arrow_schema_naive(dtype: &DType) -> VortexResult<Schema> {
283 let DType::Struct(struct_dtype, nullable) = dtype else {
284 vortex_bail!("only DType::Struct can be converted to arrow schema");
285 };
286
287 if *nullable != Nullability::NonNullable {
288 vortex_bail!("top-level struct in Schema must be NonNullable");
289 }
290
291 let mut builder = SchemaBuilder::with_capacity(struct_dtype.names().len());
292 for (field_name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) {
293 let field = if field_dtype.is_variant() {
294 let storage = DataType::Struct(variant_storage_fields_minimal());
295 Field::new(field_name.as_ref(), storage, field_dtype.is_nullable()).with_metadata(
296 [(
297 "ARROW:extension:name".to_owned(),
298 "arrow.parquet.variant".to_owned(),
299 )]
300 .into(),
301 )
302 } else {
303 Field::new(
304 field_name.as_ref(),
305 to_data_type_naive(&field_dtype)?,
306 field_dtype.is_nullable(),
307 )
308 };
309 builder.push(field);
310 }
311
312 Ok(builder.finish())
313}
314
315pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult<DataType> {
317 Ok(match dtype {
318 DType::Null => DataType::Null,
319 DType::Bool(_) => DataType::Boolean,
320 DType::Primitive(ptype, _) => match ptype {
321 PType::U8 => DataType::UInt8,
322 PType::U16 => DataType::UInt16,
323 PType::U32 => DataType::UInt32,
324 PType::U64 => DataType::UInt64,
325 PType::I8 => DataType::Int8,
326 PType::I16 => DataType::Int16,
327 PType::I32 => DataType::Int32,
328 PType::I64 => DataType::Int64,
329 PType::F16 => DataType::Float16,
330 PType::F32 => DataType::Float32,
331 PType::F64 => DataType::Float64,
332 },
333 DType::Decimal(dt, _) => {
334 let precision = dt.precision();
335 let scale = dt.scale();
336
337 match precision {
338 0..=38 => DataType::Decimal128(precision, scale),
345 39.. => DataType::Decimal256(precision, scale),
347 }
348 }
349 DType::Utf8(_) => DataType::Utf8View,
350 DType::Binary(_) => DataType::BinaryView,
351 DType::List(elem_dtype, _) => DataType::List(FieldRef::new(Field::new_list_field(
355 to_data_type_naive(elem_dtype)?,
356 elem_dtype.nullability().into(),
357 ))),
358 DType::FixedSizeList(elem_dtype, size, _) => DataType::FixedSizeList(
359 FieldRef::new(Field::new_list_field(
360 to_data_type_naive(elem_dtype)?,
361 elem_dtype.nullability().into(),
362 )),
363 *size as i32,
364 ),
365 DType::Map(map_dtype, _) => {
366 let key = Field::new("key", to_data_type_naive(&map_dtype.key_dtype())?, false);
367 let value_dtype = map_dtype.value_dtype();
368 let value = Field::new(
369 "value",
370 to_data_type_naive(&value_dtype)?,
371 value_dtype.is_nullable(),
372 );
373 let entries = Field::new_struct("entries", Fields::from(vec![key, value]), false);
374 DataType::Map(FieldRef::new(entries), map_dtype.keys_sorted())
375 }
376 DType::Struct(struct_dtype, _) => {
377 let mut fields = Vec::with_capacity(struct_dtype.names().len());
378 for (field_name, field_dt) in struct_dtype.names().iter().zip(struct_dtype.fields()) {
379 fields.push(FieldRef::from(Field::new(
380 field_name.as_ref(),
381 to_data_type_naive(&field_dt)?,
382 field_dt.is_nullable(),
383 )));
384 }
385
386 DataType::Struct(Fields::from(fields))
387 }
388 DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
389 DType::Variant(_) => vortex_bail!(
390 "DType::Variant requires Arrow Field metadata; use to_arrow_schema or a Field helper"
391 ),
392 DType::Extension(ext_dtype) => {
393 if let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() {
398 return Ok(match temporal {
399 TemporalMetadata::Timestamp(unit, tz) => {
400 DataType::Timestamp(to_arrow_time_unit(*unit)?, tz.clone())
401 }
402 TemporalMetadata::Date(unit) => match unit {
403 TimeUnit::Days => DataType::Date32,
404 TimeUnit::Milliseconds => DataType::Date64,
405 TimeUnit::Nanoseconds | TimeUnit::Microseconds | TimeUnit::Seconds => {
406 vortex_panic!(InvalidArgument: "Invalid TimeUnit {} for {}", unit, ext_dtype.id())
407 }
408 },
409 TemporalMetadata::Time(unit) => match unit {
410 TimeUnit::Seconds => DataType::Time32(ArrowTimeUnit::Second),
411 TimeUnit::Milliseconds => DataType::Time32(ArrowTimeUnit::Millisecond),
412 TimeUnit::Microseconds => DataType::Time64(ArrowTimeUnit::Microsecond),
413 TimeUnit::Nanoseconds => DataType::Time64(ArrowTimeUnit::Nanosecond),
414 TimeUnit::Days => {
415 vortex_panic!(InvalidArgument: "Invalid TimeUnit {} for {}", unit, ext_dtype.id())
416 }
417 },
418 });
419 };
420
421 vortex_bail!("Unsupported extension type \"{}\"", ext_dtype.id())
422 }
423 })
424}
425
426fn variant_storage_fields_minimal() -> Fields {
427 Fields::from(vec![
428 Field::new("metadata", DataType::Binary, false),
429 Field::new("value", DataType::Binary, true),
430 ])
431}
432
433#[allow(deprecated)]
435mod deprecated_impls {
436 use super::*;
437
438 impl TryFromArrowType<&DataType> for PType {
439 fn try_from_arrow(value: &DataType) -> VortexResult<Self> {
440 ptype_from_arrow(value)
441 }
442 }
443
444 impl TryFromArrowType<&DataType> for DecimalDType {
445 fn try_from_arrow(value: &DataType) -> VortexResult<Self> {
446 decimal_dtype_from_arrow(value)
447 }
448 }
449
450 impl FromArrowType<&ArrowTimeUnit> for TimeUnit {
451 fn from_arrow(value: &ArrowTimeUnit) -> Self {
452 from_arrow_time_unit(*value)
453 }
454 }
455
456 impl FromArrowType<ArrowTimeUnit> for TimeUnit {
457 fn from_arrow(value: ArrowTimeUnit) -> Self {
458 from_arrow_time_unit(value)
459 }
460 }
461
462 impl FromArrowType<SchemaRef> for DType {
463 fn from_arrow(value: SchemaRef) -> Self {
464 Self::from_arrow(value.as_ref())
465 }
466 }
467
468 impl TryFromArrowType<SchemaRef> for DType {
469 fn try_from_arrow(value: SchemaRef) -> VortexResult<Self> {
470 from_arrow_schema_naive(value.as_ref())
471 }
472 }
473
474 impl FromArrowType<&Schema> for DType {
475 fn from_arrow(value: &Schema) -> Self {
476 from_arrow_schema_naive(value).vortex_expect("arrow schema to dtype")
477 }
478 }
479
480 impl TryFromArrowType<&Schema> for DType {
481 fn try_from_arrow(value: &Schema) -> VortexResult<Self> {
482 from_arrow_schema_naive(value)
483 }
484 }
485
486 impl FromArrowType<&Fields> for StructFields {
487 fn from_arrow(value: &Fields) -> Self {
488 from_arrow_fields_naive(value).vortex_expect("arrow fields to struct fields")
489 }
490 }
491
492 impl TryFromArrowType<&Fields> for StructFields {
493 fn try_from_arrow(value: &Fields) -> VortexResult<Self> {
494 from_arrow_fields_naive(value)
495 }
496 }
497
498 impl FromArrowType<(&DataType, Nullability)> for DType {
499 fn from_arrow(value: (&DataType, Nullability)) -> Self {
500 from_arrow_data_type(value.0, value.1).vortex_expect("arrow data type to dtype")
501 }
502 }
503
504 impl TryFromArrowType<(&DataType, Nullability)> for DType {
505 fn try_from_arrow(
506 (data_type, nullability): (&DataType, Nullability),
507 ) -> VortexResult<Self> {
508 from_arrow_data_type(data_type, nullability)
509 }
510 }
511
512 impl FromArrowType<&Field> for DType {
513 fn from_arrow(field: &Field) -> Self {
514 from_arrow_field_naive(field).vortex_expect("arrow field to dtype")
515 }
516 }
517
518 impl TryFromArrowType<&Field> for DType {
519 fn try_from_arrow(field: &Field) -> VortexResult<Self> {
520 from_arrow_field_naive(field)
521 }
522 }
523
524 impl ToArrowType for DType {
525 fn to_arrow_schema(&self) -> VortexResult<Schema> {
526 to_arrow_schema_naive(self)
527 }
528
529 fn to_arrow_dtype(&self) -> VortexResult<DataType> {
530 to_data_type_naive(self)
531 }
532 }
533}
534
535#[cfg(test)]
536mod test {
537 #![expect(deprecated, reason = "tests for deprecated traits and methods")]
538 use arrow_schema::DataType;
539 use arrow_schema::Field;
540 use arrow_schema::FieldRef;
541 use arrow_schema::Fields;
542 use arrow_schema::Schema;
543 use rstest::fixture;
544 use rstest::rstest;
545 use vortex_array::dtype::DType;
546 use vortex_array::dtype::FieldName;
547 use vortex_array::dtype::FieldNames;
548 use vortex_array::dtype::Nullability;
549 use vortex_array::dtype::PType;
550 use vortex_array::dtype::StructFields;
551
552 use super::*;
553
554 #[test]
555 fn test_dtype_conversion_success() {
556 assert_eq!(DType::Null.to_arrow_dtype().unwrap(), DataType::Null);
557
558 assert_eq!(
559 DType::Bool(Nullability::NonNullable)
560 .to_arrow_dtype()
561 .unwrap(),
562 DataType::Boolean
563 );
564
565 assert_eq!(
566 DType::Primitive(PType::U64, Nullability::NonNullable)
567 .to_arrow_dtype()
568 .unwrap(),
569 DataType::UInt64
570 );
571
572 assert_eq!(
573 DType::Utf8(Nullability::NonNullable)
574 .to_arrow_dtype()
575 .unwrap(),
576 DataType::Utf8View
577 );
578
579 assert_eq!(
580 DType::Binary(Nullability::NonNullable)
581 .to_arrow_dtype()
582 .unwrap(),
583 DataType::BinaryView
584 );
585
586 assert_eq!(
587 DType::struct_(
588 [
589 ("field_a", DType::Bool(false.into())),
590 ("field_b", DType::Utf8(true.into()))
591 ],
592 Nullability::NonNullable,
593 )
594 .to_arrow_dtype()
595 .unwrap(),
596 DataType::Struct(Fields::from(vec![
597 FieldRef::from(Field::new("field_a", DataType::Boolean, false)),
598 FieldRef::from(Field::new("field_b", DataType::Utf8View, true)),
599 ]))
600 );
601 }
602
603 #[rstest]
604 #[case(1, DataType::Decimal128(1, 0))]
605 #[case(38, DataType::Decimal128(38, 0))]
606 #[case(39, DataType::Decimal256(39, 0))]
607 #[case(76, DataType::Decimal256(76, 0))]
608 fn test_decimal_dtype_to_arrow(#[case] precision: u8, #[case] expected: DataType) {
609 use vortex_array::dtype::DecimalDType;
610
611 let dtype = DType::Decimal(DecimalDType::new(precision, 0), Nullability::NonNullable);
612 assert_eq!(dtype.to_arrow_dtype().unwrap(), expected);
613 }
614
615 #[test]
616 fn test_variant_dtype_to_arrow_dtype_errors() {
617 let err = DType::Variant(Nullability::NonNullable)
618 .to_arrow_dtype()
619 .unwrap_err()
620 .to_string();
621 assert!(err.contains("Variant"));
622 }
623
624 #[test]
625 fn infer_nullable_list_element() {
626 let list_non_nullable = DType::List(
627 Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
628 Nullability::Nullable,
629 );
630
631 let arrow_list_non_nullable = list_non_nullable.to_arrow_dtype().unwrap();
632
633 let list_nullable = DType::List(
634 Arc::new(DType::Primitive(PType::I64, Nullability::Nullable)),
635 Nullability::Nullable,
636 );
637 let arrow_list_nullable = list_nullable.to_arrow_dtype().unwrap();
638
639 assert_ne!(arrow_list_non_nullable, arrow_list_nullable);
640 assert_eq!(
641 arrow_list_nullable,
642 DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
643 );
644 assert_eq!(
645 arrow_list_non_nullable,
646 DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))),
647 );
648 }
649
650 #[fixture]
651 fn the_struct() -> StructFields {
652 StructFields::new(
653 FieldNames::from([
654 FieldName::from("field_a"),
655 FieldName::from("field_b"),
656 FieldName::from("field_c"),
657 ]),
658 vec![
659 DType::Bool(Nullability::NonNullable),
660 DType::Utf8(Nullability::NonNullable),
661 DType::Primitive(PType::I32, Nullability::Nullable),
662 ],
663 )
664 }
665
666 #[rstest]
667 fn test_schema_conversion(the_struct: StructFields) {
668 let schema_nonnull = DType::Struct(the_struct, Nullability::NonNullable);
669
670 assert_eq!(
671 schema_nonnull.to_arrow_schema().unwrap(),
672 Schema::new(Fields::from(vec![
673 Field::new("field_a", DataType::Boolean, false),
674 Field::new("field_b", DataType::Utf8View, false),
675 Field::new("field_c", DataType::Int32, true),
676 ]))
677 );
678 }
679
680 #[test]
681 fn test_schema_variant_field_metadata() {
682 let dtype = DType::struct_(
683 [("v", DType::Variant(Nullability::NonNullable))],
684 Nullability::NonNullable,
685 );
686 let schema = dtype.to_arrow_schema().unwrap();
687 let field = schema.field(0);
688 assert_eq!(
689 field
690 .metadata()
691 .get("ARROW:extension:name")
692 .map(|s| s.as_str()),
693 Some("arrow.parquet.variant")
694 );
695 assert!(matches!(field.data_type(), DataType::Struct(_)));
696 assert!(!field.is_nullable());
697 }
698
699 #[rstest]
700 #[should_panic]
701 fn test_schema_conversion_panics(the_struct: StructFields) {
702 let schema_null = DType::Struct(the_struct, Nullability::Nullable);
703 schema_null.to_arrow_schema().unwrap();
704 }
705
706 #[test]
707 fn test_unicode_field_names_roundtrip() {
708 let unicode_field_name = "\u{5}=A";
713 let original_dtype = DType::struct_(
714 [(
715 unicode_field_name,
716 DType::Primitive(PType::I8, Nullability::Nullable),
717 )],
718 Nullability::NonNullable,
719 );
720
721 let arrow_dtype = original_dtype.to_arrow_dtype().unwrap();
722 let roundtripped_dtype = DType::from_arrow((&arrow_dtype, Nullability::NonNullable));
723
724 assert_eq!(original_dtype, roundtripped_dtype);
725 }
726
727 #[test]
728 fn test_unicode_field_names_nested_roundtrip() {
729 let inner_struct = DType::struct_(
733 [(
734 "\u{6}=inner",
735 DType::Primitive(PType::I32, Nullability::Nullable),
736 )],
737 Nullability::Nullable,
738 );
739 let original_dtype =
740 DType::struct_([("\u{7}=outer", inner_struct)], Nullability::NonNullable);
741
742 let arrow_dtype = original_dtype.to_arrow_dtype().unwrap();
743 let roundtripped_dtype = DType::from_arrow((&arrow_dtype, Nullability::NonNullable));
744
745 assert_eq!(original_dtype, roundtripped_dtype);
746 }
747
748 #[test]
749 fn map_dtype_roundtrip_uses_conventional_export_names_and_positional_import() -> VortexResult<()>
750 {
751 let dtype = DType::map(
752 DType::Primitive(PType::I32, Nullability::NonNullable),
753 DType::Utf8(Nullability::Nullable),
754 true,
755 Nullability::Nullable,
756 )?;
757
758 let arrow = dtype.to_arrow_dtype()?;
759 let DataType::Map(entries, keys_sorted) = &arrow else {
760 panic!("expected Map, got {arrow:?}");
761 };
762 assert!(*keys_sorted);
763 assert_eq!(entries.name(), "entries");
764 assert!(!entries.is_nullable());
765 let DataType::Struct(fields) = entries.data_type() else {
766 panic!("expected map entries to be a struct");
767 };
768 assert_eq!(fields[0].name(), "key");
769 assert!(!fields[0].is_nullable());
770 assert_eq!(fields[1].name(), "value");
771 assert!(fields[1].is_nullable());
772 assert_eq!(
773 DType::try_from_arrow((&arrow, Nullability::Nullable))?,
774 dtype
775 );
776
777 let positional = DataType::Map(
778 Arc::new(Field::new_struct(
779 "anything",
780 Fields::from(vec![
781 Field::new("first", DataType::Int32, false),
782 Field::new("second", DataType::Utf8, true),
783 ]),
784 false,
785 )),
786 false,
787 );
788 assert_eq!(
789 DType::try_from_arrow((&positional, Nullability::NonNullable))?,
790 DType::map(
791 DType::Primitive(PType::I32, Nullability::NonNullable),
792 DType::Utf8(Nullability::Nullable),
793 false,
794 Nullability::NonNullable,
795 )?
796 );
797
798 Ok(())
799 }
800
801 #[test]
802 fn map_dtype_import_rejects_invalid_arrow_shape() {
803 let invalid_entries = [
804 Field::new_struct(
805 "entries",
806 Fields::from(vec![
807 Field::new("key", DataType::Int32, false),
808 Field::new("value", DataType::Utf8, true),
809 ]),
810 true,
811 ),
812 Field::new("entries", DataType::Int32, false),
813 Field::new_struct(
814 "entries",
815 Fields::from(vec![Field::new("key", DataType::Int32, false)]),
816 false,
817 ),
818 Field::new_struct(
819 "entries",
820 Fields::from(vec![
821 Field::new("key", DataType::Int32, true),
822 Field::new("value", DataType::Utf8, true),
823 ]),
824 false,
825 ),
826 ];
827
828 for entries in invalid_entries {
829 let data_type = DataType::Map(Arc::new(entries), false);
830 assert!(DType::try_from_arrow((&data_type, Nullability::NonNullable)).is_err());
831 }
832 }
833
834 #[rstest]
837 #[case::duration(DataType::Duration(ArrowTimeUnit::Microsecond))]
838 #[case::interval(DataType::Interval(arrow_schema::IntervalUnit::DayTime))]
839 #[case::fixed_size_binary(DataType::FixedSizeBinary(3))]
840 fn test_try_from_arrow_unsupported_type_errors(#[case] data_type: DataType) {
841 let err = DType::try_from_arrow((&data_type, Nullability::NonNullable))
842 .expect_err("unsupported Arrow type should not convert")
843 .to_string();
844 assert!(err.contains("not supported"), "unexpected error: {err}");
845
846 let field = Field::new("c0", data_type, true);
848 assert!(DType::try_from_arrow(&field).is_err());
849 let schema = Schema::new(vec![field]);
850 assert!(DType::try_from_arrow(&schema).is_err());
851 }
852
853 #[test]
854 fn test_try_from_arrow_supported_type_succeeds() -> VortexResult<()> {
855 let dtype = DType::try_from_arrow((&DataType::Int32, Nullability::Nullable))?;
856 assert_eq!(dtype, DType::Primitive(PType::I32, Nullability::Nullable));
857 Ok(())
858 }
859}