Skip to main content

vortex_arrow/
dtype.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Convert between Vortex [`vortex_array::dtype::DType`] and Apache Arrow [`arrow_schema::DataType`].
5//!
6//! Apache Arrow's type system includes physical information, which could lead to ambiguities as
7//! Vortex treats encodings as separate from logical types.
8//!
9//! The conversions in this module are "naive": every logical type is encoded in its simplest
10//! corresponding Arrow type, and Arrow extension types (other than the builtin temporal types and
11//! Parquet Variant) are not understood. The authoritative, plugin-aware conversion entry point is
12//! [`ArrowSession`](crate::ArrowSession) — prefer its `to_arrow_schema` / `to_arrow_datatype` /
13//! `from_arrow_schema` / `from_arrow_field` / `from_arrow_datatype` methods over the deprecated
14//! traits in this module.
15//!
16//! For this reason, it's recommended to do as much computation as possible within Vortex, and then
17//! materialize an Arrow ArrayRef at the very end of the processing chain.
18
19use 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/// Trait for converting Arrow types to Vortex types.
50#[deprecated(
51    note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
52)]
53pub trait FromArrowType<T>: Sized {
54    /// Convert the Arrow type to a Vortex type.
55    #[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/// Trait for converting Arrow types to Vortex types.
62#[deprecated(
63    note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead"
64)]
65pub trait TryFromArrowType<T>: Sized {
66    /// Convert the Arrow type to a Vortex type.
67    #[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/// Extension trait converting Vortex [`DType`]s into Arrow schemas and data types.
74///
75/// This mirrors inherent methods that lived on [`DType`] before Arrow interoperability moved
76/// into this crate.
77#[deprecated(note = "Use `ArrowSession::to_arrow_schema` / `to_arrow_datatype` instead")]
78pub trait ToArrowType {
79    /// Convert a Vortex [`DType`] into an Arrow [`Schema`].
80    ///
81    /// This method is not plugin-aware and strips any `ARROW:extension:name` metadata for
82    /// non-builtin extensions (only `arrow.parquet.variant` is special-cased here). Use the
83    /// session method when you need round-trippable extension metadata.
84    #[deprecated(note = "Use `ArrowSession::to_arrow_schema` instead")]
85    fn to_arrow_schema(&self) -> VortexResult<Schema>;
86
87    /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`].
88    ///
89    /// This method has no awareness of registered Arrow extension plugins, so any
90    /// [`DType::Extension`] outside the builtin temporal set will fail or silently lose its
91    /// `ARROW:extension:name` metadata. The session methods recurse through containers
92    /// and dispatch plugins at every extension node.
93    #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")]
94    fn to_arrow_dtype(&self) -> VortexResult<DataType>;
95}
96
97/// Convert an Arrow [`ArrowTimeUnit`] to a Vortex [`TimeUnit`].
98pub(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
107/// Convert a Vortex [`TimeUnit`] to an Arrow [`ArrowTimeUnit`].
108///
109/// # Errors
110///
111/// Returns an error for units with no Arrow equivalent (e.g. [`TimeUnit::Days`]).
112pub(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
122/// Naive conversion of an Arrow [`DataType`] to the Vortex [`PType`] it stores.
123pub(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
143/// Naive conversion of an Arrow decimal [`DataType`] to a [`DecimalDType`].
144pub(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
158/// Naive conversion of an Arrow [`DataType`] to a Vortex [`DType`], with no awareness of Arrow
159/// extension plugins.
160pub(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
245/// Naive conversion of an Arrow [`Field`] to a Vortex [`DType`]. Only the builtin
246/// `arrow.parquet.variant` extension metadata is recognized.
247pub(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
259/// Naive conversion of Arrow [`Fields`] to Vortex [`StructFields`].
260pub(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
272/// Naive conversion of an Arrow [`Schema`] to a Vortex top-level non-nullable struct [`DType`].
273pub(crate) fn from_arrow_schema_naive(schema: &Schema) -> VortexResult<DType> {
274    Ok(DType::Struct(
275        from_arrow_fields_naive(schema.fields())?,
276        Nullability::NonNullable, // Must match From<RecordBatch> for Array
277    ))
278}
279
280/// Naive conversion of a Vortex top-level struct [`DType`] to an Arrow [`Schema`], with the
281/// builtin `arrow.parquet.variant` special-case.
282fn 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
315/// Naive conversion from a Vortex `DType` to the nearest Arrow physical data type.
316pub(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                // This code is commented out until DataFusion improves its support for smaller decimals.
339                // // DECIMAL32_MAX_PRECISION
340                // 0..=9 => DataType::Decimal32(precision, scale),
341                // // DECIMAL64_MAX_PRECISION
342                // 10..=18 => DataType::Decimal64(precision, scale),
343                // DECIMAL128_MAX_PRECISION
344                0..=38 => DataType::Decimal128(precision, scale),
345                // DECIMAL256_MAX_PRECISION
346                39.. => DataType::Decimal256(precision, scale),
347            }
348        }
349        DType::Utf8(_) => DataType::Utf8View,
350        DType::Binary(_) => DataType::BinaryView,
351        // There are four kinds of lists: List (32-bit offsets), Large List (64-bit), List View
352        // (32-bit), Large List View (64-bit). We cannot both guarantee zero-copy and commit to an
353        // Arrow dtype because we do not how large our offsets are.
354        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            // NOTE: Temporal are the only builtin and default-loaded extension types, and they map
394            // directly onto non-extension Arrow physical encodings. For this reason, we
395            // choose to special-case them as part of this function rather than implementing them
396            // as an import/export VTable.
397            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/// Impls of the deprecated conversion traits, delegating to the naive free functions above.
434#[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        // Regression test for https://github.com/vortex-data/vortex/issues/5979.
709
710        // Unicode characters in field names should survive an Arrow roundtrip without
711        // double-escaping.
712        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        // Regression test for https://github.com/vortex-data/vortex/issues/5979.
730
731        // Nested structs with unicode field names should also survive an Arrow roundtrip.
732        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    // Regression test for https://github.com/vortex-data/vortex/issues/8346: unsupported Arrow
835    // types must return an error instead of panicking with `unimplemented!`.
836    #[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        // The same unsupported type nested in a field or schema must also error cleanly.
847        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}