From 0cf4681db38d1e82fed0cf139acca7671bc77824 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 20 Oct 2026 14:01:14 -0400 Subject: [PATCH 1/2] Preserve unknown extension dtypes across Arrow conversion Export extensions without an Arrow plugin using their storage type and extension identity annotations. Encode serialized metadata using standard padded base64, tagged with vortex:extension:metadata:encoding=base64. Import unhandled Arrow extensions as ForeignExtDType, decoding tagged metadata and preserving untagged metadata verbatim. Explicit Arrow plugins control canonical formats. Temporal mappings and the native Vortex file representation are unchanged. Arrow import retains unknown identity independently of the file deserialization allow_unknown policy, following Arrow's storage fallback model. Expose the existing foreign dtype and opaque metadata for cross-crate use. Cover top-level and nested round-trips, structural storage, nulls, dictionaries, arbitrary bytes, invalid encodings and plugin precedence. ExtVTable still cannot establish whether storage sorting or equality preserves an extension's logical semantics. Validation: pinned nightly formatter; full tests and doctests for both touched crates; clippy with all targets/features and warnings denied. The final Arrow suite passes all 361 tests. Regression cases fail against the original implementations and pass with these fixes. Patch whitespace checks pass. Signed-off-by: Nicholas Gates --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-array/src/dtype/extension/foreign.rs | 10 +- vortex-array/src/dtype/extension/mod.rs | 2 +- vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/mod.rs | 9 + vortex-arrow/src/session.rs | 98 +++++- vortex-arrow/tests/foreign_extension.rs | 365 ++++++++++++++++++++ 8 files changed, 475 insertions(+), 12 deletions(-) create mode 100644 vortex-arrow/tests/foreign_extension.rs diff --git a/Cargo.lock b/Cargo.lock index f87fc92a8cc..ccea2dbc13b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11465,6 +11465,7 @@ dependencies = [ "arrow-data 59.3.1", "arrow-schema 59.3.1", "arrow-select 59.3.1", + "base64 0.23.1", "codspeed-divan-compat", "itertools 0.14.0", "mimalloc", diff --git a/Cargo.toml b/Cargo.toml index 4d2d2fbb573..09252e81123 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -116,6 +116,7 @@ async-lock = "3.4" async-stream = "0.3.6" async-trait = "0.1.89" base16ct = "1.0.0" +base64 = "0.23.1" bigdecimal = "0.4.8" bindgen = "0.72.0" bit-vec = "0.10.1" diff --git a/vortex-array/src/dtype/extension/foreign.rs b/vortex-array/src/dtype/extension/foreign.rs index b7636225ff9..656f05b838b 100644 --- a/vortex-array/src/dtype/extension/foreign.rs +++ b/vortex-array/src/dtype/extension/foreign.rs @@ -14,8 +14,12 @@ use crate::dtype::extension::ExtId; use crate::dtype::extension::ExtVTable; use crate::scalar::ScalarValue; +/// Opaque serialized metadata for an extension without a registered dtype plugin. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct ForeignExtMetadata(pub Vec); +pub struct ForeignExtMetadata( + /// The extension's original serialized metadata bytes. + pub Vec, +); impl Display for ForeignExtMetadata { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { @@ -23,17 +27,19 @@ impl Display for ForeignExtMetadata { } } -/// Placeholder extension dtype used when deserializing an unknown extension ID. +/// Placeholder extension dtype used when deserializing or importing an unknown extension ID. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct ForeignExtDType { id: ExtId, } impl ForeignExtDType { + /// Creates a vtable for the given unknown extension ID. pub fn new(id: ExtId) -> Self { Self { id } } + /// Preserves an unknown extension's ID, serialized metadata and storage dtype. pub fn from_parts( id: ExtId, metadata: Vec, diff --git a/vortex-array/src/dtype/extension/mod.rs b/vortex-array/src/dtype/extension/mod.rs index 44a6e7a01a7..72fc257ea70 100644 --- a/vortex-array/src/dtype/extension/mod.rs +++ b/vortex-array/src/dtype/extension/mod.rs @@ -29,7 +29,7 @@ mod plugin; pub use plugin::*; mod foreign; -pub(crate) use foreign::*; +pub use foreign::*; mod typed; pub use typed::*; diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 8605b37a3f0..2b7160e5d8f 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -25,6 +25,7 @@ arrow-cast = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } +base64 = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } simdutf8 = { workspace = true } diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index ab7c562bf7c..fcde7d12ad7 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -33,12 +33,15 @@ use arrow_schema::Schema; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::List; use vortex_array::arrays::VarBin; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::extension::datetime::AnyTemporal; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; @@ -130,6 +133,12 @@ pub(crate) fn execute_arrow_naive( exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { + if let DType::Extension(ext) = array.dtype() + && !ext.is::() + { + let extension = array.execute::(ctx)?; + return execute_arrow_naive(extension.storage_array().clone(), data_type, exporter, ctx); + } let len = array.len(); let resolved_type: DataType = match data_type { diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 18af691ef1a..f8299b84e32 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -19,6 +19,11 @@ //! Multiple plugins may register against the same key. They are tried in registration order; //! each may return [`ArrowExport::Unsupported`] / [`ArrowImport::Unsupported`] to defer to //! the next. +//! +//! The default extension export uses standard padded base64 for serialized Vortex metadata, +//! marked by `vortex:extension:metadata:encoding = base64`. Import decodes marked metadata; +//! unmarked Arrow metadata is preserved verbatim. Explicit Arrow plugins control their own +//! metadata format, including canonical Arrow formats. use std::any::Any; use std::borrow::Cow; @@ -39,13 +44,17 @@ use arrow_schema::Field; use arrow_schema::FieldRef; use arrow_schema::Fields; use arrow_schema::Schema; +use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; +use base64::Engine as _; +use base64::prelude::BASE64_STANDARD; use tracing::debug; use tracing::trace; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DictArray; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; @@ -56,7 +65,9 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; +use vortex_array::dtype::extension::ExtDTypeRef; use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ForeignExtDType; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -82,6 +93,8 @@ use crate::dtype::to_data_type_naive; use crate::executor::execute_arrow_naive; use crate::run_end_import::run_end_from_arrow; +const VORTEX_METADATA_ENCODING_KEY: &str = "vortex:extension:metadata:encoding"; + /// Outcome of a successful call to [`ArrowExportVTable::execute_arrow`]. /// /// Plugins that don't handle the supplied array return [`Unsupported`][Self::Unsupported] @@ -285,6 +298,9 @@ impl ArrowSession { /// /// For [`DType::Extension`]s, plugins registered against the extension's `Id` /// are tried in registration order; the first plugin to return `Some(field)` wins. + /// Otherwise, the storage field carries the Vortex extension ID and serialized metadata + /// encoded using standard padded base64. The field metadata includes + /// `vortex:extension:metadata:encoding = base64` so import can recover arbitrary bytes. pub fn to_arrow_field(&self, name: &str, dtype: &DType) -> VortexResult { // Handle the structural encodings, which may have recursive types match dtype { @@ -333,7 +349,20 @@ impl ArrowSession { return Ok(field); } } - vortex_bail!("extension type cannot be converted to Arrow without a plugin: {ext}"); + let metadata = BASE64_STANDARD.encode(ext.serialize_metadata()?); + let mut field = self.to_arrow_field(name, ext.storage_dtype())?; + field.set_metadata( + [ + (EXTENSION_TYPE_NAME_KEY.to_string(), ext.id().to_string()), + (EXTENSION_TYPE_METADATA_KEY.to_string(), metadata), + ( + VORTEX_METADATA_ENCODING_KEY.to_string(), + "base64".to_string(), + ), + ] + .into(), + ); + Ok(field) } DType::Variant(_) => { // TODO(Adam): This currently encodes information about parquet-variant @@ -394,9 +423,13 @@ impl ArrowSession { /// Plugins registered against the field's Arrow extension name are tried in /// registration order; the first plugin to return `Some(dtype)` wins. If none /// match (or all return `None`), the builtin `arrow.parquet.variant` extension maps - /// to [`DType::Variant`], and any other field converts through - /// [`Self::from_arrow_datatype`] so extension metadata on nested element/struct - /// fields is preserved. + /// to [`DType::Variant`]. Other unknown extensions become [`ForeignExtDType`], preserving + /// their name, metadata and storage, including nested extension fields. This follows Arrow's + /// storage fallback rule independently of [`vortex_session::VortexSession::allow_unknown`], + /// which controls Vortex deserialization rather than Arrow interoperability. + /// Metadata marked with `vortex:extension:metadata:encoding = base64` is decoded to its + /// original bytes; unmarked metadata is preserved as UTF-8 bytes. Malformed base64 or an + /// unsupported encoding marker returns an error. Registered plugins run before this fallback. #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] pub fn from_arrow_field(&self, field: &Field) -> VortexResult { if let Some(name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) { @@ -409,6 +442,7 @@ impl ArrowSession { if name == "arrow.parquet.variant" { return Ok(DType::Variant(field.is_nullable().into())); } + return self.foreign_ext_dtype(field, name).map(DType::Extension); } if let DataType::Dictionary(_, value_type) = field.data_type() { return self @@ -417,6 +451,32 @@ impl ArrowSession { self.from_arrow_datatype(field.data_type(), field.is_nullable().into()) } + fn foreign_ext_dtype(&self, field: &Field, name: &str) -> VortexResult { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] + let id = ExtId::new(name); + let metadata = field + .metadata() + .get(EXTENSION_TYPE_METADATA_KEY) + .map(String::as_str) + .unwrap_or_default(); + let metadata = match field.metadata().get(VORTEX_METADATA_ENCODING_KEY) { + None => metadata.as_bytes().to_vec(), + Some(encoding) if encoding == "base64" => { + BASE64_STANDARD.decode(metadata).map_err(|err| { + vortex_err!("Invalid base64 metadata for extension {name}: {err}") + })? + } + Some(encoding) => { + vortex_bail!("Unsupported metadata encoding for extension {name}: {encoding}") + } + }; + ForeignExtDType::from_parts( + id, + metadata, + self.from_arrow_datatype(field.data_type(), field.is_nullable().into())?, + ) + } + /// Build the Vortex [`DType`] for an Arrow [`DataType`]. /// /// Recurses into container types ([`DataType::List`] family, [`DataType::FixedSizeList`], @@ -536,8 +596,9 @@ impl ArrowSession { /// If `target` carries an `ARROW:extension:name`, the plugin registry is probed for one that /// can support executing to the target extension type. /// - /// With `target = None` the fallback path picks the array's preferred Arrow physical type - /// and executes directly into that, ignoring extension types. + /// With `target = None`, infer a field through [`Self::to_arrow_field`], including extension + /// metadata. Retain that field separately when exporting a top-level extension: the returned + /// Arrow array only carries its storage data type. /// /// Uses the default [`ArrowExportOptions`]; see [`exporter`](Self::exporter) to /// change them. @@ -594,7 +655,8 @@ impl ArrowSession { /// arrays ([`arrow_array::StructArray`], [`arrow_array::GenericListArray`], /// [`arrow_array::FixedSizeListArray`], [`arrow_array::GenericListViewArray`]) so /// extension fields nested inside containers reach their importers; leaf types fall - /// through to the canonical Arrow → Vortex array conversion. + /// through to the canonical Arrow → Vortex array conversion. Unhandled extension names + /// wrap the converted storage in a [`ForeignExtDType`]. /// /// # Errors /// @@ -654,10 +716,28 @@ impl ArrowSession { } } } - return self.from_arrow_array_canonical(current.as_ref(), field); + return self.from_arrow_array_fallback(current.as_ref(), field); } } - self.from_arrow_array_canonical(array.as_ref(), field) + self.from_arrow_array_fallback(array.as_ref(), field) + } + + #[allow(clippy::wrong_self_convention)] + fn from_arrow_array_fallback( + &self, + array: &dyn ArrowArray, + field: &Field, + ) -> VortexResult { + // Dictionary field metadata describes its values; canonical import routes it there. + if !matches!(field.data_type(), DataType::Dictionary(..)) + && let Some(name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) + && name != "arrow.parquet.variant" + { + let dtype = self.foreign_ext_dtype(field, name)?; + let storage = self.from_arrow_array_canonical(array, field)?; + return Ok(ExtensionArray::try_new(dtype, storage)?.into_array()); + } + self.from_arrow_array_canonical(array, field) } /// Recurse into Arrow container arrays so nested fields with extension metadata reach diff --git a/vortex-arrow/tests/foreign_extension.rs b/vortex-arrow/tests/foreign_extension.rs new file mode 100644 index 00000000000..440f2a9077a --- /dev/null +++ b/vortex-arrow/tests/foreign_extension.rs @@ -0,0 +1,365 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::tests_outside_test_module)] + +//! Unknown Arrow extensions retain their identity while exposing their storage. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::DictionaryArray; +use arrow_array::Int32Array; +use arrow_array::types::Int32Type; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; +use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; +use arrow_schema::extension::ExtensionType; +use arrow_schema::extension::FixedShapeTensor as ArrowTensor; +use base64::Engine as _; +use base64::prelude::BASE64_STANDARD; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ForeignExtDType; +use vortex_array::validity::Validity; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +static OPAQUE_ID: CachedId = CachedId::new("test.opaque"); + +const METADATA_ENCODING_KEY: &str = "vortex:extension:metadata:encoding"; + +fn foreign(storage: ArrayRef, metadata: &[u8]) -> VortexResult { + Ok(ExtensionArray::try_new( + ForeignExtDType::from_parts(*OPAQUE_ID, metadata.to_vec(), storage.dtype().clone())?, + storage, + )? + .into_array()) +} + +fn list(elements: ArrayRef) -> VortexResult { + Ok(ListArray::try_new( + elements, + buffer![0i32, 2, 3].into_array(), + Validity::NonNullable, + &mut array_session().create_execution_ctx(), + )? + .into_array()) +} + +#[rstest] +#[case::top_level("top")] +#[case::struct_child("struct")] +#[case::list_child("list")] +#[case::struct_storage("struct_storage")] +#[case::list_storage("list_storage")] +fn test_foreign_extension_roundtrip( + #[case] placement: &str, + #[values(false, true)] nullable: bool, + #[values(b"".as_slice(), "{\"label\":\"日本語\"}\0".as_bytes(), b"\0\xff\x80\xfe".as_slice())] + metadata: &[u8], +) -> VortexResult<()> { + let storage = if nullable { + VarBinViewArray::from_iter_nullable_bin([Some(b"abc".as_slice()), None, Some(b"xyz")]) + } else { + VarBinViewArray::from_iter_bin([b"abc", b"def", b"xyz"]) + } + .into_array(); + let original = match placement { + "struct_storage" => foreign( + StructArray::from_fields(&[("opaque", storage)])?.into_array(), + metadata, + )?, + "list_storage" => foreign(list(storage)?, metadata)?, + _ => { + let extension = foreign(storage, metadata)?; + match placement { + "struct" => StructArray::from_fields(&[("opaque", extension)])?.into_array(), + "list" => list(extension)?, + _ => extension, + } + } + }; + let session = array_session(); + assert!(!session.allows_unknown()); + let arrow_session = ArrowSession::default(); + let field = arrow_session.to_arrow_field("column", original.dtype())?; + let extension_field = match placement { + "struct" => match field.data_type() { + DataType::Struct(fields) => fields[0].as_ref(), + other => panic!("expected struct, got {other}"), + }, + "list" => match field.data_type() { + DataType::List(element) => element.as_ref(), + other => panic!("expected list, got {other}"), + }, + _ => &field, + }; + assert_eq!( + extension_field.metadata()[EXTENSION_TYPE_NAME_KEY], + "test.opaque" + ); + assert_eq!( + extension_field.metadata()[EXTENSION_TYPE_METADATA_KEY], + BASE64_STANDARD.encode(metadata) + ); + assert_eq!(extension_field.metadata()[METADATA_ENCODING_KEY], "base64"); + assert_eq!(arrow_session.from_arrow_field(&field)?, *original.dtype()); + + let mut ctx = session.create_execution_ctx(); + for target in [None, Some(&field)] { + let arrow = arrow_session.execute_arrow(original.clone(), target, &mut ctx)?; + assert_eq!(arrow.data_type(), field.data_type()); + let imported = arrow_session.from_arrow_array(arrow, &field)?; + assert_eq!(imported.dtype(), original.dtype()); + assert_arrays_eq!(imported, original, &mut ctx); + } + Ok(()) +} + +#[rstest] +fn test_unknown_dictionary_extension_retains_identity( + #[values(false, true)] encoded: bool, +) -> VortexResult<()> { + let arrow_session = ArrowSession::default(); + let values = Arc::new(Int32Array::from(vec![42, 99])); + let dictionary = Arc::new(DictionaryArray::::try_new( + Int32Array::from(vec![Some(0), None, Some(1), Some(0)]), + values, + )?); + let mut field = Field::new( + "opaque", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)), + true, + ) + .with_metadata( + [( + EXTENSION_TYPE_NAME_KEY.to_string(), + "test.opaque".to_string(), + )] + .into(), + ); + let metadata: &[u8] = if encoded { &[0xff] } else { b"" }; + if encoded { + field.metadata_mut().extend([ + (METADATA_ENCODING_KEY.to_string(), "base64".to_string()), + (EXTENSION_TYPE_METADATA_KEY.to_string(), "/w==".to_string()), + ]); + } + let imported = arrow_session.from_arrow_array(dictionary, &field)?; + assert_eq!(imported.dtype(), &arrow_session.from_arrow_field(&field)?); + let expected = foreign( + PrimitiveArray::from_option_iter([Some(42i32), None, Some(99), Some(42)]).into_array(), + metadata, + )?; + assert_arrays_eq!( + imported, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn test_default_metadata_uses_standard_padded_base64() -> VortexResult<()> { + let metadata = (0..=255).collect::>(); + let original = foreign(buffer![1i32].into_array(), &metadata)?; + let session = ArrowSession::default(); + let field = session.to_arrow_field("opaque", original.dtype())?; + assert_eq!(field.metadata()[METADATA_ENCODING_KEY], "base64"); + let imported = session.from_arrow_field(&field)?; + assert_eq!(&imported, original.dtype()); + + let single_byte = foreign(buffer![1i32].into_array(), &[0xff])?; + let field = session.to_arrow_field("opaque", single_byte.dtype())?; + assert_eq!(field.metadata()[EXTENSION_TYPE_METADATA_KEY], "/w=="); + Ok(()) +} + +#[rstest] +#[case::json("{\"shape\":[1]}")] +#[case::base64_looking_text("/w==")] +fn test_unmarked_arrow_metadata_is_preserved(#[case] metadata: &str) -> VortexResult<()> { + let session = ArrowSession::default(); + let field = Field::new("opaque", DataType::Int32, false).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "test.opaque".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + metadata.to_string(), + ), + ] + .into(), + ); + let original = foreign(buffer![42i32].into_array(), metadata.as_bytes())?; + assert_eq!(&session.from_arrow_field(&field)?, original.dtype()); + let imported = session.from_arrow_array(Arc::new(Int32Array::from(vec![42])), &field)?; + assert_arrays_eq!( + imported, + original, + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[rstest] +#[case::invalid_base64("base64", "{\"shape\":[1]}", "Invalid base64 metadata")] +#[case::missing_padding("base64", "/w", "Invalid base64 metadata")] +#[case::unsupported_encoding("hex", "ff", "Unsupported metadata encoding")] +fn test_invalid_metadata_encoding_is_rejected( + #[case] encoding: &str, + #[case] metadata: &str, + #[case] expected_error: &str, +) { + let session = ArrowSession::default(); + let field = Field::new("opaque", DataType::Int32, false).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "test.opaque".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + metadata.to_string(), + ), + (METADATA_ENCODING_KEY.to_string(), encoding.to_string()), + ] + .into(), + ); + let error = session + .from_arrow_field(&field) + .expect_err("invalid metadata encoding"); + assert!(error.to_string().contains(expected_error)); + let error = session + .from_arrow_array(Arc::new(Int32Array::from(vec![42])), &field) + .expect_err("invalid metadata encoding"); + assert!(error.to_string().contains(expected_error)); +} + +/// A test plugin mapping opaque binary metadata to Arrow's canonical tensor JSON. +#[derive(Debug)] +struct TensorPlugin; + +static ARROW_TENSOR_ID: CachedId = CachedId::new(ArrowTensor::NAME); + +impl ArrowExportVTable for TensorPlugin { + fn arrow_ext_id(&self) -> Id { + *ARROW_TENSOR_ID + } + + fn vortex_id(&self) -> Id { + *OPAQUE_ID + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let DType::Extension(dtype) = dtype else { + return Ok(None); + }; + let mut field = session.to_arrow_field(name, dtype.storage_dtype())?; + field.try_with_extension_type(ArrowTensor::try_new( + DataType::Int32, + [1], + Some(vec!["x".to_string()]), + Some(vec![0]), + )?)?; + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = array.execute::(ctx)?; + ArrowSession::default() + .exporter(options) + .execute_arrow(array.storage_array().clone(), Some(target), ctx) + .map(ArrowExport::Exported) + } +} + +impl ArrowImportVTable for TensorPlugin { + fn arrow_ext_id(&self) -> Id { + *ARROW_TENSOR_ID + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + field.try_extension_type::()?; + Ok(Some(DType::Extension(ForeignExtDType::from_parts( + *OPAQUE_ID, + vec![0xff], + session.from_arrow_datatype(field.data_type(), field.is_nullable().into())?, + )?))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + _dtype: &DType, + session: &ArrowSession, + ) -> VortexResult { + foreign( + session.from_arrow_array(array, field.is_nullable())?, + &[0xff], + ) + .map(ArrowImport::Imported) + } +} + +#[test] +fn test_arrow_plugin_overrides_default_metadata_encoding() -> VortexResult<()> { + let session = ArrowSession::default(); + session.register_exporter(Arc::new(TensorPlugin)); + session.register_importer(Arc::new(TensorPlugin)); + let storage = + FixedSizeListArray::try_new(buffer![42i32].into_array(), 1, Validity::NonNullable, 1)? + .into_array(); + let original = foreign(storage, &[0xff])?; + let field = session.to_arrow_field("tensor", original.dtype())?; + assert_eq!(field.metadata()[EXTENSION_TYPE_NAME_KEY], ArrowTensor::NAME); + assert!(field.metadata()[EXTENSION_TYPE_METADATA_KEY].starts_with('{')); + assert!(!field.metadata().contains_key(METADATA_ENCODING_KEY)); + assert_eq!(&session.from_arrow_field(&field)?, original.dtype()); + let mut ctx = array_session().create_execution_ctx(); + let arrow = session.execute_arrow(original.clone(), Some(&field), &mut ctx)?; + let imported = session.from_arrow_array(arrow, &field)?; + assert_arrays_eq!(imported, original, &mut ctx); + Ok(()) +} From 91c46ff08941545642d61692bf88fd956a4655f7 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 20 Oct 2026 14:56:40 -0400 Subject: [PATCH 2/2] Validate Arrow targets before unwrapping extension storage Remove unconditional extension unwrapping from the physical Arrow executor. Perform it in the session fallback after exporter plugins have had a chance to handle the target, and only when the requested Arrow type maps to the extension's storage dtype. Keep the original logical dtype for other conversions so canonical target validation can reject opaque extensions requested as dates, times or timestamps. Preserve explicit struct/list storage exports. Add ten temporal target regression cases covering annotated and unannotated fields. All ten fail before the fix and pass afterward. Also cover valid exports to explicit structural storage targets. Validation: pinned nightly formatter; 373 Arrow tests; two doctests; clippy with all targets/features and warnings denied; git diff --check. Signed-off-by: Nicholas Gates --- vortex-arrow/src/executor/mod.rs | 9 ---- vortex-arrow/src/session.rs | 27 +++++++++++- vortex-arrow/tests/foreign_extension.rs | 56 +++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 11 deletions(-) diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index fcde7d12ad7..ab7c562bf7c 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -33,15 +33,12 @@ use arrow_schema::Schema; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::List; use vortex_array::arrays::VarBin; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::PType; -use vortex_array::extension::datetime::AnyTemporal; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; @@ -133,12 +130,6 @@ pub(crate) fn execute_arrow_naive( exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { - if let DType::Extension(ext) = array.dtype() - && !ext.is::() - { - let extension = array.execute::(ctx)?; - return execute_arrow_naive(extension.storage_array().clone(), data_type, exporter, ctx); - } let len = array.len(); let resolved_type: DataType = match data_type { diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index f8299b84e32..ef14cb15d1f 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -59,6 +59,7 @@ use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; @@ -996,10 +997,32 @@ impl ArrowExporter<'_> { "unsupported Arrow extension type encountered, falling back to naive execution" ); - return execute_arrow_naive(current, Some(target_field.data_type()), self, ctx); + return self.execute_arrow_fallback(current, Some(target_field.data_type()), ctx); } - execute_arrow_naive(array, target.map(|field| field.data_type()), self, ctx) + self.execute_arrow_fallback(array, target.map(|field| field.data_type()), ctx) + } + + fn execute_arrow_fallback( + &self, + array: ArrayRef, + data_type: Option<&DataType>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Only unwrap for a storage target. Other conversions must see the original logical + // dtype so their validators can reject incompatible extension semantics. + if let Some(target) = data_type + && let DType::Extension(ext) = array.dtype() + && self + .session + .from_arrow_datatype(target, ext.nullability())? + == *ext.storage_dtype() + { + let extension = array.execute::(ctx)?; + return execute_arrow_naive(extension.storage_array().clone(), data_type, self, ctx); + } + + execute_arrow_naive(array, data_type, self, ctx) } } diff --git a/vortex-arrow/tests/foreign_extension.rs b/vortex-arrow/tests/foreign_extension.rs index 440f2a9077a..498eae264f9 100644 --- a/vortex-arrow/tests/foreign_extension.rs +++ b/vortex-arrow/tests/foreign_extension.rs @@ -13,6 +13,7 @@ use arrow_array::Int32Array; use arrow_array::types::Int32Type; use arrow_schema::DataType; use arrow_schema::Field; +use arrow_schema::TimeUnit as ArrowTimeUnit; use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use arrow_schema::extension::ExtensionType; @@ -198,6 +199,61 @@ fn test_default_metadata_uses_standard_padded_base64() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::date32(DataType::Date32)] +#[case::date64(DataType::Date64)] +#[case::time32(DataType::Time32(ArrowTimeUnit::Second))] +#[case::time64(DataType::Time64(ArrowTimeUnit::Microsecond))] +#[case::timestamp(DataType::Timestamp(ArrowTimeUnit::Millisecond, None))] +fn test_foreign_extension_rejects_temporal_target( + #[case] target: DataType, + #[values(false, true)] annotated: bool, +) -> VortexResult<()> { + let session = ArrowSession::default(); + let source = foreign(buffer![42i64].into_array(), b"")?; + let field = if annotated { + session + .to_arrow_field("opaque", source.dtype())? + .with_data_type(target) + } else { + Field::new("opaque", target, false) + }; + let error = session + .execute_arrow( + source, + Some(&field), + &mut array_session().create_execution_ctx(), + ) + .expect_err("an opaque extension cannot be reinterpreted as a temporal type"); + assert!( + error.to_string().contains("Cannot convert extension"), + "{error}" + ); + Ok(()) +} + +#[rstest] +#[case::struct_storage(true)] +#[case::list_storage(false)] +fn test_foreign_extension_exports_to_explicit_storage_target( + #[case] struct_storage: bool, +) -> VortexResult<()> { + let values = buffer![1i32, 2, 3].into_array(); + let storage = if struct_storage { + StructArray::from_fields(&[("values", values)])?.into_array() + } else { + list(values)? + }; + let source = foreign(storage.clone(), &[0xff])?; + let session = ArrowSession::default(); + let field = session.to_arrow_field("storage", storage.dtype())?; + let mut ctx = array_session().create_execution_ctx(); + let arrow = session.execute_arrow(source, Some(&field), &mut ctx)?; + let imported = session.from_arrow_array(arrow, &field)?; + assert_arrays_eq!(imported, storage, &mut ctx); + Ok(()) +} + #[rstest] #[case::json("{\"shape\":[1]}")] #[case::base64_looking_text("/w==")]