1use std::sync::Arc;
17
18use arrow_schema::DataType;
19use arrow_schema::Field;
20use arrow_schema::FieldRef;
21use arrow_schema::Fields;
22use arrow_schema::Schema;
23use arrow_schema::SchemaBuilder;
24use arrow_schema::SchemaRef;
25use arrow_schema::TimeUnit as ArrowTimeUnit;
26use vortex_array::dtype::DType;
27use vortex_array::dtype::DecimalDType;
28use vortex_array::dtype::FieldName;
29use vortex_array::dtype::Nullability;
30use vortex_array::dtype::PType;
31use vortex_array::dtype::StructFields;
32use vortex_array::extension::datetime::AnyTemporal;
33use vortex_array::extension::datetime::Date;
34use vortex_array::extension::datetime::TemporalMetadata;
35use vortex_array::extension::datetime::Time;
36use vortex_array::extension::datetime::TimeUnit;
37use vortex_array::extension::datetime::Timestamp;
38use vortex_error::VortexExpect;
39use vortex_error::VortexResult;
40use vortex_error::vortex_bail;
41use vortex_error::vortex_err;
42use vortex_error::vortex_panic;
43
44pub trait FromArrowType<T>: Sized {
46 fn from_arrow(value: T) -> Self;
48}
49
50pub trait TryFromArrowType<T>: Sized {
52 fn try_from_arrow(value: T) -> VortexResult<Self>;
54}
55
56impl TryFromArrowType<&DataType> for PType {
57 fn try_from_arrow(value: &DataType) -> VortexResult<Self> {
58 match value {
59 DataType::Int8 => Ok(Self::I8),
60 DataType::Int16 => Ok(Self::I16),
61 DataType::Int32 => Ok(Self::I32),
62 DataType::Int64 => Ok(Self::I64),
63 DataType::UInt8 => Ok(Self::U8),
64 DataType::UInt16 => Ok(Self::U16),
65 DataType::UInt32 => Ok(Self::U32),
66 DataType::UInt64 => Ok(Self::U64),
67 DataType::Float16 => Ok(Self::F16),
68 DataType::Float32 => Ok(Self::F32),
69 DataType::Float64 => Ok(Self::F64),
70 _ => Err(vortex_err!(
71 "Arrow datatype {:?} cannot be converted to ptype",
72 value
73 )),
74 }
75 }
76}
77
78impl TryFromArrowType<&DataType> for DecimalDType {
79 fn try_from_arrow(value: &DataType) -> VortexResult<Self> {
80 match value {
81 DataType::Decimal32(precision, scale)
82 | DataType::Decimal64(precision, scale)
83 | DataType::Decimal128(precision, scale)
84 | DataType::Decimal256(precision, scale) => Self::try_new(*precision, *scale),
85
86 _ => Err(vortex_err!(
87 "Arrow datatype {:?} cannot be converted to DecimalDType",
88 value
89 )),
90 }
91 }
92}
93
94impl FromArrowType<&ArrowTimeUnit> for TimeUnit {
95 fn from_arrow(value: &ArrowTimeUnit) -> Self {
96 Self::from_arrow(*value)
97 }
98}
99
100impl FromArrowType<ArrowTimeUnit> for TimeUnit {
101 fn from_arrow(value: ArrowTimeUnit) -> Self {
102 match value {
103 ArrowTimeUnit::Second => Self::Seconds,
104 ArrowTimeUnit::Millisecond => Self::Milliseconds,
105 ArrowTimeUnit::Microsecond => Self::Microseconds,
106 ArrowTimeUnit::Nanosecond => Self::Nanoseconds,
107 }
108 }
109}
110
111pub fn to_arrow_time_unit(value: TimeUnit) -> VortexResult<ArrowTimeUnit> {
117 Ok(match value {
118 TimeUnit::Seconds => ArrowTimeUnit::Second,
119 TimeUnit::Milliseconds => ArrowTimeUnit::Millisecond,
120 TimeUnit::Microseconds => ArrowTimeUnit::Microsecond,
121 TimeUnit::Nanoseconds => ArrowTimeUnit::Nanosecond,
122 _ => vortex_bail!("Cannot convert {value} to Arrow TimeUnit"),
123 })
124}
125
126impl FromArrowType<SchemaRef> for DType {
127 fn from_arrow(value: SchemaRef) -> Self {
128 Self::from_arrow(value.as_ref())
129 }
130}
131
132impl TryFromArrowType<SchemaRef> for DType {
133 fn try_from_arrow(value: SchemaRef) -> VortexResult<Self> {
134 Self::try_from_arrow(value.as_ref())
135 }
136}
137
138impl FromArrowType<&Schema> for DType {
139 fn from_arrow(value: &Schema) -> Self {
140 Self::try_from_arrow(value).vortex_expect("arrow schema to dtype")
141 }
142}
143
144impl TryFromArrowType<&Schema> for DType {
145 fn try_from_arrow(value: &Schema) -> VortexResult<Self> {
146 Ok(Self::Struct(
147 StructFields::try_from_arrow(value.fields())?,
148 Nullability::NonNullable, ))
150 }
151}
152
153impl FromArrowType<&Fields> for StructFields {
154 fn from_arrow(value: &Fields) -> Self {
155 Self::try_from_arrow(value).vortex_expect("arrow fields to struct fields")
156 }
157}
158
159impl TryFromArrowType<&Fields> for StructFields {
160 fn try_from_arrow(value: &Fields) -> VortexResult<Self> {
161 value
162 .into_iter()
163 .map(|f| {
164 Ok((
165 FieldName::from(f.name().as_str()),
166 DType::try_from_arrow(f.as_ref())?,
167 ))
168 })
169 .collect::<VortexResult<StructFields>>()
170 }
171}
172
173impl FromArrowType<(&DataType, Nullability)> for DType {
174 fn from_arrow(value: (&DataType, Nullability)) -> Self {
175 Self::try_from_arrow(value).vortex_expect("arrow data type to dtype")
176 }
177}
178
179impl TryFromArrowType<(&DataType, Nullability)> for DType {
180 fn try_from_arrow((data_type, nullability): (&DataType, Nullability)) -> VortexResult<Self> {
181 if data_type.is_integer() || data_type.is_floating() {
182 return Ok(DType::Primitive(
183 PType::try_from_arrow(data_type)?,
184 nullability,
185 ));
186 }
187
188 Ok(match data_type {
189 DataType::Null => DType::Null,
190 DataType::Decimal32(precision, scale)
191 | DataType::Decimal64(precision, scale)
192 | DataType::Decimal128(precision, scale)
193 | DataType::Decimal256(precision, scale) => {
194 DType::Decimal(DecimalDType::new(*precision, *scale), nullability)
195 }
196 DataType::Boolean => DType::Bool(nullability),
197 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DType::Utf8(nullability),
198 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
199 DType::Binary(nullability)
200 }
201 DataType::Date32 => DType::Extension(Date::new(TimeUnit::Days, nullability).erased()),
202 DataType::Date64 => {
203 DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased())
204 }
205 DataType::Time32(unit) => {
206 DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased())
207 }
208 DataType::Time64(unit) => {
209 DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased())
210 }
211 DataType::Timestamp(unit, tz) => DType::Extension(
212 Timestamp::new_with_tz(TimeUnit::from_arrow(unit), tz.clone(), nullability)
213 .erased(),
214 ),
215 DataType::List(e)
216 | DataType::LargeList(e)
217 | DataType::ListView(e)
218 | DataType::LargeListView(e) => {
219 DType::List(Arc::new(Self::try_from_arrow(e.as_ref())?), nullability)
220 }
221 DataType::FixedSizeList(e, size) => DType::FixedSizeList(
222 Arc::new(Self::try_from_arrow(e.as_ref())?),
223 *size as u32,
224 nullability,
225 ),
226 DataType::Struct(f) => DType::Struct(StructFields::try_from_arrow(f)?, nullability),
227 DataType::Dictionary(_, value_type) => {
228 Self::try_from_arrow((value_type.as_ref(), nullability))?
229 }
230 DataType::RunEndEncoded(_, value_type) => {
231 Self::try_from_arrow((value_type.data_type(), nullability))?
232 }
233 _ => vortex_bail!("Arrow data type not supported: {data_type:?}"),
234 })
235 }
236}
237
238impl FromArrowType<&Field> for DType {
239 fn from_arrow(field: &Field) -> Self {
240 Self::try_from_arrow(field).vortex_expect("arrow field to dtype")
241 }
242}
243
244impl TryFromArrowType<&Field> for DType {
245 fn try_from_arrow(field: &Field) -> VortexResult<Self> {
246 if field
247 .metadata()
248 .get("ARROW:extension:name")
249 .map(|s| s.as_str())
250 == Some("arrow.parquet.variant")
251 {
252 return Ok(DType::Variant(field.is_nullable().into()));
253 }
254 Self::try_from_arrow((field.data_type(), field.is_nullable().into()))
255 }
256}
257
258pub trait ToArrowType {
263 fn to_arrow_schema(&self) -> VortexResult<Schema>;
270
271 #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")]
279 fn to_arrow_dtype(&self) -> VortexResult<DataType>;
280}
281
282impl ToArrowType for DType {
283 fn to_arrow_schema(&self) -> VortexResult<Schema> {
284 let DType::Struct(struct_dtype, nullable) = self else {
285 vortex_bail!("only DType::Struct can be converted to arrow schema");
286 };
287
288 if *nullable != Nullability::NonNullable {
289 vortex_bail!("top-level struct in Schema must be NonNullable");
290 }
291
292 let mut builder = SchemaBuilder::with_capacity(struct_dtype.names().len());
293 for (field_name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) {
294 let field = if field_dtype.is_variant() {
295 let storage = DataType::Struct(variant_storage_fields_minimal());
296 Field::new(field_name.as_ref(), storage, field_dtype.is_nullable()).with_metadata(
297 [(
298 "ARROW:extension:name".to_owned(),
299 "arrow.parquet.variant".to_owned(),
300 )]
301 .into(),
302 )
303 } else {
304 Field::new(
305 field_name.as_ref(),
306 to_data_type_naive(&field_dtype)?,
307 field_dtype.is_nullable(),
308 )
309 };
310 builder.push(field);
311 }
312
313 Ok(builder.finish())
314 }
315
316 fn to_arrow_dtype(&self) -> VortexResult<DataType> {
317 to_data_type_naive(self)
318 }
319}
320
321pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult<DataType> {
323 Ok(match dtype {
324 DType::Null => DataType::Null,
325 DType::Bool(_) => DataType::Boolean,
326 DType::Primitive(ptype, _) => match ptype {
327 PType::U8 => DataType::UInt8,
328 PType::U16 => DataType::UInt16,
329 PType::U32 => DataType::UInt32,
330 PType::U64 => DataType::UInt64,
331 PType::I8 => DataType::Int8,
332 PType::I16 => DataType::Int16,
333 PType::I32 => DataType::Int32,
334 PType::I64 => DataType::Int64,
335 PType::F16 => DataType::Float16,
336 PType::F32 => DataType::Float32,
337 PType::F64 => DataType::Float64,
338 },
339 DType::Decimal(dt, _) => {
340 let precision = dt.precision();
341 let scale = dt.scale();
342
343 match precision {
344 0..=38 => DataType::Decimal128(precision, scale),
351 39.. => DataType::Decimal256(precision, scale),
353 }
354 }
355 DType::Utf8(_) => DataType::Utf8View,
356 DType::Binary(_) => DataType::BinaryView,
357 DType::List(elem_dtype, _) => DataType::List(FieldRef::new(Field::new_list_field(
361 to_data_type_naive(elem_dtype)?,
362 elem_dtype.nullability().into(),
363 ))),
364 DType::FixedSizeList(elem_dtype, size, _) => DataType::FixedSizeList(
365 FieldRef::new(Field::new_list_field(
366 to_data_type_naive(elem_dtype)?,
367 elem_dtype.nullability().into(),
368 )),
369 *size as i32,
370 ),
371 DType::Struct(struct_dtype, _) => {
372 let mut fields = Vec::with_capacity(struct_dtype.names().len());
373 for (field_name, field_dt) in struct_dtype.names().iter().zip(struct_dtype.fields()) {
374 fields.push(FieldRef::from(Field::new(
375 field_name.as_ref(),
376 to_data_type_naive(&field_dt)?,
377 field_dt.is_nullable(),
378 )));
379 }
380
381 DataType::Struct(Fields::from(fields))
382 }
383 DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
384 DType::Variant(_) => vortex_bail!(
385 "DType::Variant requires Arrow Field metadata; use to_arrow_schema or a Field helper"
386 ),
387 DType::Extension(ext_dtype) => {
388 if let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() {
393 return Ok(match temporal {
394 TemporalMetadata::Timestamp(unit, tz) => {
395 DataType::Timestamp(to_arrow_time_unit(*unit)?, tz.clone())
396 }
397 TemporalMetadata::Date(unit) => match unit {
398 TimeUnit::Days => DataType::Date32,
399 TimeUnit::Milliseconds => DataType::Date64,
400 TimeUnit::Nanoseconds | TimeUnit::Microseconds | TimeUnit::Seconds => {
401 vortex_panic!(InvalidArgument: "Invalid TimeUnit {} for {}", unit, ext_dtype.id())
402 }
403 },
404 TemporalMetadata::Time(unit) => match unit {
405 TimeUnit::Seconds => DataType::Time32(ArrowTimeUnit::Second),
406 TimeUnit::Milliseconds => DataType::Time32(ArrowTimeUnit::Millisecond),
407 TimeUnit::Microseconds => DataType::Time64(ArrowTimeUnit::Microsecond),
408 TimeUnit::Nanoseconds => DataType::Time64(ArrowTimeUnit::Nanosecond),
409 TimeUnit::Days => {
410 vortex_panic!(InvalidArgument: "Invalid TimeUnit {} for {}", unit, ext_dtype.id())
411 }
412 },
413 });
414 };
415
416 vortex_bail!("Unsupported extension type \"{}\"", ext_dtype.id())
417 }
418 })
419}
420
421fn variant_storage_fields_minimal() -> Fields {
422 Fields::from(vec![
423 Field::new("metadata", DataType::Binary, false),
424 Field::new("value", DataType::Binary, true),
425 ])
426}
427
428#[cfg(test)]
429mod test {
430 #![expect(deprecated, reason = "tests for deprecated methods")]
431 use arrow_schema::DataType;
432 use arrow_schema::Field;
433 use arrow_schema::FieldRef;
434 use arrow_schema::Fields;
435 use arrow_schema::Schema;
436 use rstest::fixture;
437 use rstest::rstest;
438 use vortex_array::dtype::DType;
439 use vortex_array::dtype::FieldName;
440 use vortex_array::dtype::FieldNames;
441 use vortex_array::dtype::Nullability;
442 use vortex_array::dtype::PType;
443 use vortex_array::dtype::StructFields;
444
445 use super::*;
446
447 #[test]
448 fn test_dtype_conversion_success() {
449 assert_eq!(DType::Null.to_arrow_dtype().unwrap(), DataType::Null);
450
451 assert_eq!(
452 DType::Bool(Nullability::NonNullable)
453 .to_arrow_dtype()
454 .unwrap(),
455 DataType::Boolean
456 );
457
458 assert_eq!(
459 DType::Primitive(PType::U64, Nullability::NonNullable)
460 .to_arrow_dtype()
461 .unwrap(),
462 DataType::UInt64
463 );
464
465 assert_eq!(
466 DType::Utf8(Nullability::NonNullable)
467 .to_arrow_dtype()
468 .unwrap(),
469 DataType::Utf8View
470 );
471
472 assert_eq!(
473 DType::Binary(Nullability::NonNullable)
474 .to_arrow_dtype()
475 .unwrap(),
476 DataType::BinaryView
477 );
478
479 assert_eq!(
480 DType::struct_(
481 [
482 ("field_a", DType::Bool(false.into())),
483 ("field_b", DType::Utf8(true.into()))
484 ],
485 Nullability::NonNullable,
486 )
487 .to_arrow_dtype()
488 .unwrap(),
489 DataType::Struct(Fields::from(vec![
490 FieldRef::from(Field::new("field_a", DataType::Boolean, false)),
491 FieldRef::from(Field::new("field_b", DataType::Utf8View, true)),
492 ]))
493 );
494 }
495
496 #[rstest]
497 #[case(1, DataType::Decimal128(1, 0))]
498 #[case(38, DataType::Decimal128(38, 0))]
499 #[case(39, DataType::Decimal256(39, 0))]
500 #[case(76, DataType::Decimal256(76, 0))]
501 fn test_decimal_dtype_to_arrow(#[case] precision: u8, #[case] expected: DataType) {
502 use vortex_array::dtype::DecimalDType;
503
504 let dtype = DType::Decimal(DecimalDType::new(precision, 0), Nullability::NonNullable);
505 assert_eq!(dtype.to_arrow_dtype().unwrap(), expected);
506 }
507
508 #[test]
509 fn test_variant_dtype_to_arrow_dtype_errors() {
510 let err = DType::Variant(Nullability::NonNullable)
511 .to_arrow_dtype()
512 .unwrap_err()
513 .to_string();
514 assert!(err.contains("Variant"));
515 }
516
517 #[test]
518 fn infer_nullable_list_element() {
519 let list_non_nullable = DType::List(
520 Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
521 Nullability::Nullable,
522 );
523
524 let arrow_list_non_nullable = list_non_nullable.to_arrow_dtype().unwrap();
525
526 let list_nullable = DType::List(
527 Arc::new(DType::Primitive(PType::I64, Nullability::Nullable)),
528 Nullability::Nullable,
529 );
530 let arrow_list_nullable = list_nullable.to_arrow_dtype().unwrap();
531
532 assert_ne!(arrow_list_non_nullable, arrow_list_nullable);
533 assert_eq!(
534 arrow_list_nullable,
535 DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
536 );
537 assert_eq!(
538 arrow_list_non_nullable,
539 DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))),
540 );
541 }
542
543 #[fixture]
544 fn the_struct() -> StructFields {
545 StructFields::new(
546 FieldNames::from([
547 FieldName::from("field_a"),
548 FieldName::from("field_b"),
549 FieldName::from("field_c"),
550 ]),
551 vec![
552 DType::Bool(Nullability::NonNullable),
553 DType::Utf8(Nullability::NonNullable),
554 DType::Primitive(PType::I32, Nullability::Nullable),
555 ],
556 )
557 }
558
559 #[rstest]
560 fn test_schema_conversion(the_struct: StructFields) {
561 let schema_nonnull = DType::Struct(the_struct, Nullability::NonNullable);
562
563 assert_eq!(
564 schema_nonnull.to_arrow_schema().unwrap(),
565 Schema::new(Fields::from(vec![
566 Field::new("field_a", DataType::Boolean, false),
567 Field::new("field_b", DataType::Utf8View, false),
568 Field::new("field_c", DataType::Int32, true),
569 ]))
570 );
571 }
572
573 #[test]
574 fn test_schema_variant_field_metadata() {
575 let dtype = DType::struct_(
576 [("v", DType::Variant(Nullability::NonNullable))],
577 Nullability::NonNullable,
578 );
579 let schema = dtype.to_arrow_schema().unwrap();
580 let field = schema.field(0);
581 assert_eq!(
582 field
583 .metadata()
584 .get("ARROW:extension:name")
585 .map(|s| s.as_str()),
586 Some("arrow.parquet.variant")
587 );
588 assert!(matches!(field.data_type(), DataType::Struct(_)));
589 assert!(!field.is_nullable());
590 }
591
592 #[rstest]
593 #[should_panic]
594 fn test_schema_conversion_panics(the_struct: StructFields) {
595 let schema_null = DType::Struct(the_struct, Nullability::Nullable);
596 schema_null.to_arrow_schema().unwrap();
597 }
598
599 #[test]
600 fn test_unicode_field_names_roundtrip() {
601 let unicode_field_name = "\u{5}=A";
606 let original_dtype = DType::struct_(
607 [(
608 unicode_field_name,
609 DType::Primitive(PType::I8, Nullability::Nullable),
610 )],
611 Nullability::NonNullable,
612 );
613
614 let arrow_dtype = original_dtype.to_arrow_dtype().unwrap();
615 let roundtripped_dtype = DType::from_arrow((&arrow_dtype, Nullability::NonNullable));
616
617 assert_eq!(original_dtype, roundtripped_dtype);
618 }
619
620 #[test]
621 fn test_unicode_field_names_nested_roundtrip() {
622 let inner_struct = DType::struct_(
626 [(
627 "\u{6}=inner",
628 DType::Primitive(PType::I32, Nullability::Nullable),
629 )],
630 Nullability::Nullable,
631 );
632 let original_dtype =
633 DType::struct_([("\u{7}=outer", inner_struct)], Nullability::NonNullable);
634
635 let arrow_dtype = original_dtype.to_arrow_dtype().unwrap();
636 let roundtripped_dtype = DType::from_arrow((&arrow_dtype, Nullability::NonNullable));
637
638 assert_eq!(original_dtype, roundtripped_dtype);
639 }
640
641 #[rstest]
644 #[case::duration(DataType::Duration(ArrowTimeUnit::Microsecond))]
645 #[case::interval(DataType::Interval(arrow_schema::IntervalUnit::DayTime))]
646 #[case::fixed_size_binary(DataType::FixedSizeBinary(3))]
647 fn test_try_from_arrow_unsupported_type_errors(#[case] data_type: DataType) {
648 let err = DType::try_from_arrow((&data_type, Nullability::NonNullable))
649 .expect_err("unsupported Arrow type should not convert")
650 .to_string();
651 assert!(err.contains("not supported"), "unexpected error: {err}");
652
653 let field = Field::new("c0", data_type, true);
655 assert!(DType::try_from_arrow(&field).is_err());
656 let schema = Schema::new(vec![field]);
657 assert!(DType::try_from_arrow(&schema).is_err());
658 }
659
660 #[test]
661 fn test_try_from_arrow_supported_type_succeeds() -> VortexResult<()> {
662 let dtype = DType::try_from_arrow((&DataType::Int32, Nullability::Nullable))?;
663 assert_eq!(dtype, DType::Primitive(PType::I32, Nullability::Nullable));
664 Ok(())
665 }
666}