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/session.rs b/vortex-arrow/src/session.rs index 18af691ef1a..ef14cb15d1f 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,24 +44,31 @@ 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; 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; 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 +94,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 +299,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 +350,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 +424,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 +443,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 +452,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 +597,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 +656,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 +717,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 @@ -916,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); + } + + 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, target.map(|field| field.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 new file mode 100644 index 00000000000..498eae264f9 --- /dev/null +++ b/vortex-arrow/tests/foreign_extension.rs @@ -0,0 +1,421 @@ +// 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::TimeUnit as ArrowTimeUnit; +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::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==")] +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(()) +}