diff --git a/vortex-array/src/arrow/session.rs b/vortex-array/src/arrow/session.rs index 8279d7a1ca7..0c4cc7b2a6d 100644 --- a/vortex-array/src/arrow/session.rs +++ b/vortex-array/src/arrow/session.rs @@ -57,7 +57,7 @@ use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; -use crate::dtype::arrow::FromArrowType; +use crate::dtype::arrow::TryFromArrowType; use crate::dtype::arrow::to_data_type_naive; use crate::dtype::extension::ExtId; use crate::extension::datetime::AnyTemporal; @@ -313,7 +313,7 @@ impl ArrowSession { /// match (or all return `None`), recurses into container types ([`DataType::List`] /// family, [`DataType::FixedSizeList`], [`DataType::Struct`]) so extension metadata /// on nested element/struct fields is preserved. Leaf types use the canonical - /// Arrow → Vortex mapping via [`DType::from_arrow`]. + /// Arrow → Vortex mapping via [`DType::try_from_arrow`]. pub fn from_arrow_field(&self, field: &Field) -> VortexResult { if let Some(name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) { for plugin in self.importers(&Id::new(name)).iter() { @@ -345,7 +345,7 @@ impl ArrowSession { .collect::>>()?; DType::Struct(StructFields::from_iter(entries), nullability) } - _ => DType::from_arrow(field), + _ => DType::try_from_arrow(field)?, }) } diff --git a/vortex-array/src/dtype/arrow.rs b/vortex-array/src/dtype/arrow.rs index 103c52daf93..9d381b6eab2 100644 --- a/vortex-array/src/dtype/arrow.rs +++ b/vortex-array/src/dtype/arrow.rs @@ -130,36 +130,63 @@ impl FromArrowType for DType { } } +impl TryFromArrowType for DType { + fn try_from_arrow(value: SchemaRef) -> VortexResult { + Self::try_from_arrow(value.as_ref()) + } +} + impl FromArrowType<&Schema> for DType { fn from_arrow(value: &Schema) -> Self { - Self::Struct( - StructFields::from_arrow(value.fields()), + Self::try_from_arrow(value).vortex_expect("arrow schema to dtype") + } +} + +impl TryFromArrowType<&Schema> for DType { + fn try_from_arrow(value: &Schema) -> VortexResult { + Ok(Self::Struct( + StructFields::try_from_arrow(value.fields())?, Nullability::NonNullable, // Must match From for Array - ) + )) } } impl FromArrowType<&Fields> for StructFields { fn from_arrow(value: &Fields) -> Self { - StructFields::from_iter(value.into_iter().map(|f| { - ( - FieldName::from(f.name().as_str()), - DType::from_arrow(f.as_ref()), - ) - })) + Self::try_from_arrow(value).vortex_expect("arrow fields to struct fields") + } +} + +impl TryFromArrowType<&Fields> for StructFields { + fn try_from_arrow(value: &Fields) -> VortexResult { + value + .into_iter() + .map(|f| { + Ok(( + FieldName::from(f.name().as_str()), + DType::try_from_arrow(f.as_ref())?, + )) + }) + .collect::>() } } impl FromArrowType<(&DataType, Nullability)> for DType { - fn from_arrow((data_type, nullability): (&DataType, Nullability)) -> Self { + fn from_arrow(value: (&DataType, Nullability)) -> Self { + Self::try_from_arrow(value).vortex_expect("arrow data type to dtype") + } +} + +impl TryFromArrowType<(&DataType, Nullability)> for DType { + fn try_from_arrow((data_type, nullability): (&DataType, Nullability)) -> VortexResult { if data_type.is_integer() || data_type.is_floating() { - return DType::Primitive( - PType::try_from_arrow(data_type).vortex_expect("arrow float/integer to ptype"), + return Ok(DType::Primitive( + PType::try_from_arrow(data_type)?, nullability, - ); + )); } - match data_type { + Ok(match data_type { DataType::Null => DType::Null, DataType::Decimal32(precision, scale) | DataType::Decimal64(precision, scale) @@ -189,36 +216,42 @@ impl FromArrowType<(&DataType, Nullability)> for DType { | DataType::LargeList(e) | DataType::ListView(e) | DataType::LargeListView(e) => { - DType::List(Arc::new(Self::from_arrow(e.as_ref())), nullability) + DType::List(Arc::new(Self::try_from_arrow(e.as_ref())?), nullability) } DataType::FixedSizeList(e, size) => DType::FixedSizeList( - Arc::new(Self::from_arrow(e.as_ref())), + Arc::new(Self::try_from_arrow(e.as_ref())?), *size as u32, nullability, ), - DataType::Struct(f) => DType::Struct(StructFields::from_arrow(f), nullability), + DataType::Struct(f) => DType::Struct(StructFields::try_from_arrow(f)?, nullability), DataType::Dictionary(_, value_type) => { - Self::from_arrow((value_type.as_ref(), nullability)) + Self::try_from_arrow((value_type.as_ref(), nullability))? } DataType::RunEndEncoded(_, value_type) => { - Self::from_arrow((value_type.data_type(), nullability)) + Self::try_from_arrow((value_type.data_type(), nullability))? } - _ => unimplemented!("Arrow data type not yet supported: {:?}", data_type), - } + _ => vortex_bail!("Arrow data type not supported: {data_type:?}"), + }) } } impl FromArrowType<&Field> for DType { fn from_arrow(field: &Field) -> Self { + Self::try_from_arrow(field).vortex_expect("arrow field to dtype") + } +} + +impl TryFromArrowType<&Field> for DType { + fn try_from_arrow(field: &Field) -> VortexResult { if field .metadata() .get("ARROW:extension:name") .map(|s| s.as_str()) == Some("arrow.parquet.variant") { - return DType::Variant(field.is_nullable().into()); + return Ok(DType::Variant(field.is_nullable().into())); } - Self::from_arrow((field.data_type(), field.is_nullable().into())) + Self::try_from_arrow((field.data_type(), field.is_nullable().into())) } } @@ -594,4 +627,30 @@ mod test { assert_eq!(original_dtype, roundtripped_dtype); } + + // Regression test for https://github.com/vortex-data/vortex/issues/8346: unsupported Arrow + // types must return an error instead of panicking with `unimplemented!`. + #[rstest] + #[case::duration(DataType::Duration(ArrowTimeUnit::Microsecond))] + #[case::interval(DataType::Interval(arrow_schema::IntervalUnit::DayTime))] + #[case::fixed_size_binary(DataType::FixedSizeBinary(3))] + fn test_try_from_arrow_unsupported_type_errors(#[case] data_type: DataType) { + let err = DType::try_from_arrow((&data_type, Nullability::NonNullable)) + .expect_err("unsupported Arrow type should not convert") + .to_string(); + assert!(err.contains("not supported"), "unexpected error: {err}"); + + // The same unsupported type nested in a field or schema must also error cleanly. + let field = Field::new("c0", data_type, true); + assert!(DType::try_from_arrow(&field).is_err()); + let schema = Schema::new(vec![field]); + assert!(DType::try_from_arrow(&schema).is_err()); + } + + #[test] + fn test_try_from_arrow_supported_type_succeeds() -> VortexResult<()> { + let dtype = DType::try_from_arrow((&DataType::Int32, Nullability::Nullable))?; + assert_eq!(dtype, DType::Primitive(PType::I32, Nullability::Nullable)); + Ok(()) + } } diff --git a/vortex-python/src/arrays/from_arrow.rs b/vortex-python/src/arrays/from_arrow.rs index beb142a1625..e8c9ca04d60 100644 --- a/vortex-python/src/arrays/from_arrow.rs +++ b/vortex-python/src/arrays/from_arrow.rs @@ -15,7 +15,7 @@ use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; use vortex::array::arrow::FromArrowArray; use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; +use vortex::dtype::arrow::TryFromArrowType; use vortex::error::VortexError; use vortex::error::VortexResult; @@ -48,16 +48,18 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult>>()?; - let dtype: DType = obj + let arrow_dtype = obj .getattr(intern!(py, "type")) - .and_then(|v| DataType::from_pyarrow(&v.as_borrowed())) - .map(|dt| DType::from_arrow(&Field::new("_", dt, false)))?; + .and_then(|v| DataType::from_pyarrow(&v.as_borrowed()))?; + let dtype = DType::try_from_arrow(&Field::new("_", arrow_dtype, false)) + .map_err(|e| PyValueError::new_err(e.to_string()))?; Ok(PyArrayRef::from( ChunkedArray::try_new(encoded_chunks, dtype)?.into_array(), )) } else if obj.is_instance(table)? { let array_stream = ArrowArrayStreamReader::from_pyarrow(&obj.as_borrowed())?; - let dtype = DType::from_arrow(array_stream.schema()); + let dtype = DType::try_from_arrow(array_stream.schema()) + .map_err(|e| PyValueError::new_err(e.to_string()))?; let chunks = array_stream .into_iter() .map(|b| { diff --git a/vortex-python/src/dtype/mod.rs b/vortex-python/src/dtype/mod.rs index bea25c133c1..59cee1fc857 100644 --- a/vortex-python/src/dtype/mod.rs +++ b/vortex-python/src/dtype/mod.rs @@ -35,7 +35,7 @@ use pyo3::types::PyType; use pyo3::wrap_pyfunction; use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; +use vortex::dtype::arrow::TryFromArrowType; use crate::arrow::FromPyArrow; use crate::arrow::ToPyArrow; @@ -198,10 +198,9 @@ impl PyDType { #[pyo3(from_py_with = import_arrow_dtype)] arrow_dtype: DataType, non_nullable: bool, ) -> PyResult> { - Self::init( - cls.py(), - DType::from_arrow(&Field::new("_", arrow_dtype, !non_nullable)), - ) + let dtype = DType::try_from_arrow(&Field::new("_", arrow_dtype, !non_nullable)) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Self::init(cls.py(), dtype) } } diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 473c95cc68f..a8d61edfc18 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -5,6 +5,7 @@ use arrow_array::RecordBatchReader; use arrow_array::ffi_stream::ArrowArrayStreamReader; use async_fs::File; use pyo3::exceptions::PyTypeError; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::pyfunction; use pyo3_object_store::PyObjectStore; @@ -17,7 +18,7 @@ use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; +use vortex::dtype::arrow::TryFromArrowType; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::file::WriteOptionsSessionExt; @@ -407,7 +408,8 @@ fn try_arrow_stream_to_iterator( if ob.is_instance(pa_table)? || ob.is_instance(pa_record_batch_reader)? { // Convert to Arrow stream using FFI let arrow_stream = ArrowArrayStreamReader::from_pyarrow(ob)?; - let dtype = DType::from_arrow(arrow_stream.schema()); + let dtype = DType::try_from_arrow(arrow_stream.schema()) + .map_err(|e| PyValueError::new_err(e.to_string()))?; // Convert Arrow RecordBatch stream to Vortex ArrayIterator let vortex_iter = arrow_stream diff --git a/vortex-python/test/test_array.py b/vortex-python/test/test_array.py index cd7c9f4f4a8..7458f738e05 100644 --- a/vortex-python/test/test_array.py +++ b/vortex-python/test/test_array.py @@ -43,3 +43,19 @@ def test_empty_array(): def test_scalar_at_out_of_bounds(): a = vortex.array([10, 42, 999, 1992]) _s = a.scalar_at(10) + + +@pytest.mark.parametrize( + "arrow_type", + [ + pa.duration("us"), + pa.month_day_nano_interval(), + pa.binary(3), + ], +) +def test_unsupported_arrow_type_raises_value_error(arrow_type: pa.DataType): + # Regression test for https://github.com/vortex-data/vortex/issues/8346: + # unsupported Arrow types must surface as a clean ValueError, not a PanicException. + table = pa.table({"c0": pa.array([], type=arrow_type)}) + with pytest.raises(ValueError): + _ = vortex.array(table)