From 20173ac4f3ad6b898bbd6cee1568b91e74274a61 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 16:37:48 +0100 Subject: [PATCH 1/8] Deprecate FromArrowArray in favour of ArrowSession array imports Follow-up to #8918: `FromArrowArray` is now #[deprecated] and every usage is replaced with `ArrowSession` methods, making the session the authoritative Arrow array import path. - `ArrowSession` gains `from_arrow_array_nullable(&dyn Array, bool)` for imports with no Arrow `Field` in hand; nested fields still dispatch extension import plugins. - `ArrowImportVTable::from_arrow_array` now receives the `ArrowSession` so plugins convert storage through the session (uuid, geo, json, tensor, parquet-variant updated). - All consumers (datafusion, python, ffi, tui, bench, compat-gen, layout, facade doc snippet) migrated to `from_arrow_array` / `from_arrow_array_nullable` / `from_arrow_record_batch`. - `vx_array_from_arrow` now imports through the session using the FFI schema's field, so Arrow extension types round-trip. - `IntoArrowArray` was already deprecated and has no remaining users. Co-Authored-By: Claude Fable 5 Signed-off-by: Robert Kruszewski --- encodings/parquet-variant/src/array.rs | 15 ++++-- encodings/parquet-variant/src/arrow.rs | 1 + encodings/parquet-variant/src/kernel.rs | 18 +++++-- vortex-arrow/src/convert.rs | 16 ++++++ vortex-arrow/src/datum.rs | 3 ++ vortex-arrow/src/executor/struct_.rs | 5 +- vortex-arrow/src/iter.rs | 4 +- vortex-arrow/src/lib.rs | 10 +++- vortex-arrow/src/session.rs | 50 ++++++++++++++----- vortex-arrow/src/uuid.rs | 1 + vortex-arrow/tests/canonical.rs | 12 +++-- vortex-bench/src/conversions.rs | 4 +- vortex-bench/src/tpch/tpchgen.rs | 6 ++- vortex-datafusion/src/convert/exprs.rs | 8 +-- vortex-datafusion/src/lib.rs | 8 +-- vortex-datafusion/src/persistent/opener.rs | 5 +- vortex-ffi/src/array.rs | 6 ++- vortex-json/src/arrow.rs | 4 +- vortex-layout/src/scan/arrow.rs | 6 ++- vortex-python/src/arrays/from_arrow.rs | 17 ++++--- vortex-python/src/io.rs | 4 +- vortex-spatial/src/extension/linestring.rs | 9 ++-- vortex-spatial/src/extension/mod.rs | 16 ++++-- .../src/extension/multilinestring.rs | 9 ++-- vortex-spatial/src/extension/multipoint.rs | 9 ++-- vortex-spatial/src/extension/multipolygon.rs | 9 ++-- vortex-spatial/src/extension/point.rs | 9 ++-- vortex-spatial/src/extension/polygon.rs | 9 ++-- vortex-spatial/src/extension/rect.rs | 4 +- vortex-spatial/src/extension/wkb.rs | 4 +- vortex-tensor/src/types/vector/arrow.rs | 15 ++++-- .../fixtures/arrays/datasets/clickbench.rs | 8 ++- .../src/fixtures/arrays/datasets/tpch.rs | 8 ++- vortex-tui/src/convert.rs | 11 ++-- vortex/src/lib.rs | 5 +- 35 files changed, 226 insertions(+), 102 deletions(-) diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 50e2784cbbe..e54a3c9616a 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; @@ -40,7 +41,7 @@ use vortex_array::vtable::validity_to_child; reason = "TODO(aduffy): figure out what to do with Parquet Variant" )] use vortex_arrow::ArrowArrayExecutor; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSession; use vortex_arrow::to_arrow_null_buffer; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; @@ -66,6 +67,10 @@ pub struct ParquetVariantSlots { pub typed_value: Option, } +/// Variant storage children (metadata/value/typed_value) carry no Arrow extension metadata, +/// so [`ParquetVariant::from_arrow_variant`] converts them through a default [`ArrowSession`]. +pub(crate) static ARROW_SESSION: LazyLock = LazyLock::new(ArrowSession::default); + impl ParquetVariant { /// Creates a Parquet Variant array from canonical extension storage slots. /// @@ -130,17 +135,17 @@ impl ParquetVariant { } else { Validity::NonNullable }); - let metadata = - ArrayRef::from_arrow(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + let metadata = ARROW_SESSION + .from_arrow_array_nullable(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; let value = arrow_variant .value_field() - .map(|v| ArrayRef::from_arrow(v as &dyn ArrowArray, value_nullable)) + .map(|v| ARROW_SESSION.from_arrow_array_nullable(v as &dyn ArrowArray, value_nullable)) .transpose()?; let typed_value = arrow_variant .typed_value_field() - .map(|tv| ArrayRef::from_arrow(tv.as_ref(), typed_value_nullable)) + .map(|tv| ARROW_SESSION.from_arrow_array_nullable(tv.as_ref(), typed_value_nullable)) .transpose()?; ParquetVariant::try_new(validity, metadata, value, typed_value).map(IntoArray::into_array) } diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index c68ffb435e7..01673a5d364 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -263,6 +263,7 @@ impl ArrowImportVTable for ParquetVariant { array: ArrowArrayRef, field: &Field, dtype: &DType, + _session: &ArrowSession, ) -> VortexResult { if !dtype.is_variant() || field diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 14305287b63..aea1ce4f52d 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -327,7 +327,6 @@ mod tests { use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; - use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -837,15 +836,24 @@ mod tests { .map(|field| field.is_nullable()) .unwrap_or(false); - let metadata = - ArrayRef::from_arrow(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + let metadata = SESSION + .arrow() + .from_arrow_array_nullable(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; let value = arrow_variant .value_field() - .map(|value| ArrayRef::from_arrow(value as &dyn ArrowArray, value_nullable)) + .map(|value| { + SESSION + .arrow() + .from_arrow_array_nullable(value as &dyn ArrowArray, value_nullable) + }) .transpose()?; let typed_value = arrow_variant .typed_value_field() - .map(|typed_value| ArrayRef::from_arrow(typed_value.as_ref(), typed_value_nullable)) + .map(|typed_value| { + SESSION + .arrow() + .from_arrow_array_nullable(typed_value.as_ref(), typed_value_nullable) + }) .transpose()?; Ok( diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 6fd2fcc73b9..7b0b22ca136 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +// This module hosts the canonical Arrow → Vortex conversion machinery, implemented as the +// deprecated `FromArrowArray` trait. Internal callers go through `from_arrow_dyn` / +// `from_arrow_batch`; external callers use the `ArrowSession` methods. +#![allow(deprecated)] + use std::sync::Arc; use arrow_array::AnyDictionaryArray; @@ -104,6 +109,17 @@ use crate::dtype::from_arrow_time_unit; /// /// This mirrors [`IntoArray`] for Arrow buffer types; a separate trait is required because both /// [`IntoArray`] and the Arrow buffer types are foreign to this crate. +/// Canonical (non-plugin) conversion of an Arrow array into a Vortex array. Internal +/// entry point for callers that don't dispatch Arrow extension plugins. +pub(crate) fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult { + ArrayRef::from_arrow(array, nullable) +} + +/// Canonical (non-plugin) conversion of an Arrow [`RecordBatch`] into a Vortex struct array. +pub(crate) fn from_arrow_batch(batch: &RecordBatch, nullable: bool) -> VortexResult { + ArrayRef::from_arrow(batch, nullable) +} + pub trait IntoVortexArray { /// Convert this Arrow buffer into a non-nullable Vortex array without copying. fn into_array(self) -> ArrayRef; diff --git a/vortex-arrow/src/datum.rs b/vortex-arrow/src/datum.rs index 5019642c452..ef100fe8fc3 100644 --- a/vortex-arrow/src/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -18,6 +18,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::ArrowSessionExt; +#[allow(deprecated)] use crate::FromArrowArray; /// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute. @@ -109,6 +110,7 @@ impl ArrowDatum for Datum { note = "Relies on the hidden global `legacy_session()`; use `from_arrow_columnar` with an explicit `ExecutionCtx` instead" )] #[allow(clippy::disallowed_methods)] +#[allow(deprecated)] pub fn from_arrow_array_with_len(array: A, len: usize, nullable: bool) -> VortexResult where ArrayRef: FromArrowArray, @@ -145,6 +147,7 @@ where /// # Error /// /// The provided array must have length `len` or `1`. +#[allow(deprecated)] pub fn from_arrow_columnar( array: A, len: usize, diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index d5de4317925..78abd8e1a0b 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -224,7 +224,6 @@ mod tests { use arrow_buffer::NullBuffer; use arrow_schema::DataType; use arrow_schema::Field; - use vortex_array as array; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -241,7 +240,7 @@ mod tests { use vortex_error::VortexResult; use crate::ArrowArrayExecutor; - use crate::FromArrowArray; + use crate::convert::from_arrow_dyn; use crate::dtype::to_data_type_naive; #[test] @@ -419,7 +418,7 @@ mod tests { )?; let orig_dtype = array.dtype().clone(); let arrow_array = array.into_array().execute_arrow(None, &mut ctx)?; - let from_arrow = array::ArrayRef::from_arrow(arrow_array.as_ref(), false)?; + let from_arrow = from_arrow_dyn(arrow_array.as_ref(), false)?; assert_eq!(&orig_dtype, from_arrow.dtype()); Ok(()) } diff --git a/vortex-arrow/src/iter.rs b/vortex-arrow/src/iter.rs index 7462c909a30..30fa7119390 100644 --- a/vortex-arrow/src/iter.rs +++ b/vortex-arrow/src/iter.rs @@ -9,7 +9,7 @@ use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::FromArrowArray; +use crate::convert::from_arrow_batch; use crate::dtype::from_arrow_schema_naive; /// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`. @@ -42,7 +42,7 @@ impl Iterator for ArrowArrayStreamAdapter { &from_arrow_schema_naive(b.schema().as_ref()) .vortex_expect("arrow schema to dtype") ); - ArrayRef::from_arrow(b, false) + from_arrow_batch(&b, false) })) } } diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 1a9c526f6b4..865ac9b190b 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -7,8 +7,8 @@ //! the [`ArrowSession`]: importing Arrow schemas, fields, and data types into Vortex //! ([`ArrowSession::from_arrow_schema`], [`ArrowSession::from_arrow_field`], //! [`ArrowSession::from_arrow_datatype`]), importing Arrow arrays and record batches -//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_record_batch`], and the -//! low-level [`FromArrowArray`]), exporting Vortex dtypes to Arrow +//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_array_nullable`], +//! [`ArrowSession::from_arrow_record_batch`]), exporting Vortex dtypes to Arrow //! ([`ArrowSession::to_arrow_schema`], [`ArrowSession::to_arrow_field`], //! [`ArrowSession::to_arrow_datatype`]), and executing Vortex arrays into Arrow //! ([`ArrowSession::execute_arrow`] and the [`ArrowArrayExecutor`] convenience trait). @@ -62,6 +62,9 @@ pub fn initialize(session: &VortexSession) { /// /// Implementations reuse the underlying Arrow buffers without copying wherever the Arrow and /// Vortex memory layouts allow it. +#[deprecated( + note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_array_nullable`, `from_arrow_record_batch`) instead" +)] pub trait FromArrowArray { /// Convert `array` into a Vortex array whose [`DType`](vortex_array::dtype::DType) has the requested /// `nullable` [`Nullability`](vortex_array::dtype::Nullability). @@ -80,6 +83,9 @@ pub trait FromArrowArray { /// Returns an error if `nullable` is `false` but `array` physically contains one or more nulls /// (including an Arrow `NullArray`, which is entirely null), or if the Arrow data type is not /// supported. + #[deprecated( + note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_array_nullable`, `from_arrow_record_batch`) instead" + )] fn from_arrow(array: A, nullable: bool) -> VortexResult where Self: Sized; diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index c2be94531c9..a8f282e44a3 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -24,7 +24,7 @@ use std::any::Any; use std::fmt::Debug; use std::sync::Arc; -use arrow_array::Array as _; +use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::RecordBatch; use arrow_array::RunArray; @@ -70,8 +70,8 @@ use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; -use crate::FromArrowArray; use crate::IntoVortexArray; +use crate::convert::from_arrow_dyn; use crate::convert::map_from_arrow_parts; use crate::convert::nulls; use crate::convert::remove_nulls; @@ -165,12 +165,17 @@ pub trait ArrowImportVTable: 'static + Send + Sync + Debug { /// /// Returns ownership of `array` via [`ArrowImport::Unsupported`] when the plugin cannot /// handle the input. + /// + /// `session` is provided so plugins can convert storage or nested arrays through the + /// session (e.g. [`ArrowSession::from_arrow_array_nullable`]) instead of the deprecated + /// `FromArrowArray` trait. #[allow(clippy::wrong_self_convention)] fn from_arrow_array( &self, array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult; } @@ -578,23 +583,42 @@ impl ArrowSession { let dtype = self.from_arrow_field(field)?; let mut current = array; for plugin in importers.iter() { - match plugin.from_arrow_array(current, field, &dtype)? { + match plugin.from_arrow_array(current, field, &dtype, self)? { ArrowImport::Imported(arr) => return Ok(arr), ArrowImport::Unsupported(arr) => current = arr, } } - return ArrayRef::from_arrow(current.as_ref(), field.is_nullable()); + return self.from_arrow_array_canonical(current.as_ref(), field); } } - self.from_arrow_array_canonical(array, field) + self.from_arrow_array_canonical(array.as_ref(), field) + } + + /// Decode an Arrow array into a Vortex array whose dtype has the requested `nullable`ness. + /// + /// An Arrow array can carry a validity (null) buffer regardless of whether its schema + /// declares the field nullable, so the desired nullability is supplied by the caller. + /// Returns an error if `nullable` is `false` but the array physically contains nulls. + /// + /// No top-level Arrow [`Field`] is available in this form, so no extension plugin is + /// dispatched for the array itself; fields nested inside container data types still carry + /// their metadata and are routed through [`Self::from_arrow_array`]. Prefer the + /// field-aware [`Self::from_arrow_array`] when a [`Field`] is in hand. + pub fn from_arrow_array_nullable( + &self, + array: &dyn ArrowArray, + nullable: bool, + ) -> VortexResult { + let field = Field::new("", array.data_type().clone(), nullable); + self.from_arrow_array_canonical(array, &field) } /// Recurse into Arrow container arrays so nested fields with extension metadata reach - /// their importers, falling through to [`ArrayRef::from_arrow`] for leaf types. + /// their importers, falling through to the canonical conversion for leaf types. #[allow(clippy::wrong_self_convention)] fn from_arrow_array_canonical( &self, - array: ArrowArrayRef, + array: &dyn ArrowArray, field: &Field, ) -> VortexResult { use arrow_array::cast::AsArray; @@ -685,9 +709,9 @@ impl ArrowSession { DataType::RunEndEncoded(ends_field, values_field) => { let values_field = run_end_values_field(values_field, field.is_nullable().into()); match ends_field.data_type() { - DataType::Int16 => self.run_end_from_arrow::(&array, &values_field), - DataType::Int32 => self.run_end_from_arrow::(&array, &values_field), - DataType::Int64 => self.run_end_from_arrow::(&array, &values_field), + DataType::Int16 => self.run_end_from_arrow::(array, &values_field), + DataType::Int32 => self.run_end_from_arrow::(array, &values_field), + DataType::Int64 => self.run_end_from_arrow::(array, &values_field), ends_dt => vortex_bail!( "Arrow run-end array run ends must be Int16, Int32 or Int64, got {ends_dt}" ), @@ -699,12 +723,12 @@ impl ArrowSession { let values = self.from_arrow_array(ArrowArrayRef::clone(dict.values()), &values_field)?; let codes = dict.keys(); - let codes = ArrayRef::from_arrow(codes, codes.is_nullable())?; + let codes = from_arrow_dyn(codes, codes.is_nullable())?; // SAFETY: arrow-rs enforces the dictionary invariants on construction, so the // codes are in-bounds for the values. Ok(unsafe { DictArray::new_unchecked(codes, values) }.into_array()) } - _ => ArrayRef::from_arrow(array.as_ref(), field.is_nullable()), + _ => from_arrow_dyn(array, field.is_nullable()), } } @@ -713,7 +737,7 @@ impl ArrowSession { #[allow(clippy::wrong_self_convention)] fn run_end_from_arrow( &self, - array: &ArrowArrayRef, + array: &dyn ArrowArray, values_field: &Field, ) -> VortexResult where diff --git a/vortex-arrow/src/uuid.rs b/vortex-arrow/src/uuid.rs index b7a6cc77b29..33228816283 100644 --- a/vortex-arrow/src/uuid.rs +++ b/vortex-arrow/src/uuid.rs @@ -129,6 +129,7 @@ impl ArrowImportVTable for Uuid { array: ArrowArrayRef, _field: &Field, dtype: &DType, + _session: &ArrowSession, ) -> VortexResult { let DType::Extension(dtype) = dtype else { return Ok(ArrowImport::Unsupported(array)); diff --git a/vortex-arrow/tests/canonical.rs b/vortex-arrow/tests/canonical.rs index f67e0b73115..18355fb98c8 100644 --- a/vortex-arrow/tests/canonical.rs +++ b/vortex-arrow/tests/canonical.rs @@ -24,13 +24,11 @@ use arrow_buffer::NullBufferBuilder; use arrow_buffer::OffsetBuffer; use arrow_schema::DataType; use arrow_schema::Field; -use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -128,7 +126,10 @@ fn roundtrip_struct() { nulls.finish(), ); - let vortex_struct = ArrayRef::from_arrow(&arrow_struct, true).unwrap(); + let vortex_struct = SESSION + .arrow() + .from_arrow_array_nullable(&arrow_struct, true) + .unwrap(); let vortex_struct = SESSION .arrow() .execute_arrow(vortex_struct, None, &mut ctx) @@ -154,7 +155,10 @@ fn roundtrip_list() { let list_data_type = arrow_list.data_type(); let list_field = Field::new(String::new(), list_data_type.clone(), true); - let vortex_list = ArrayRef::from_arrow(&arrow_list, true).unwrap(); + let vortex_list = SESSION + .arrow() + .from_arrow_array_nullable(&arrow_list, true) + .unwrap(); let rt_arrow_list = SESSION .arrow() diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 0086d321e89..69cba42c6b0 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -54,7 +54,6 @@ use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; use vortex::utils::parallelism::get_available_parallelism; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_spatial::extension::SpatialMetadata; use vortex_spatial::extension::WellKnownBinary; use wkb::Endianness; @@ -118,7 +117,8 @@ pub fn parquet_to_vortex_stream( ) -> impl futures::Stream> { reader.map(move |result| { result.map_err(|e| vortex_err!(External: e)).and_then(|rb| { - let chunk = ArrayRef::from_arrow(rb, false)?; + let schema = rb.schema(); + let chunk = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); // Canonicalize the chunk. diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index cd2f8c24341..2c4a26fd527 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -36,7 +36,6 @@ use vortex::array::stream::ArrayStreamAdapter; use vortex::error::VortexExpect; use vortex::file::WriteOptionsSessionExt; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use crate::CompactionStrategy; use crate::Format; @@ -362,7 +361,10 @@ impl VortexWriter { #[async_trait::async_trait] impl FileWriter for VortexWriter { async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let array = ArrayRef::from_arrow(batch, false)?; + let schema = batch.schema(); + let array = SESSION + .arrow() + .from_arrow_record_batch(batch.clone(), &schema)?; self.sender .as_ref() .vortex_expect("sender closed early") diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 1b4b1126ede..01bde086e06 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -1236,7 +1236,6 @@ mod tests { use vortex::array::Canonical; use vortex::array::VortexSessionExecute as _; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; // Create test data let values = Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20])); @@ -1283,10 +1282,13 @@ mod tests { let vortex_expr = expr_convertor.try_convert_case_expr(&case_expr).unwrap(); // Convert batch to Vortex array - let vortex_array: ArrayRef = ArrayRef::from_arrow(&batch, false).unwrap(); + let session = VortexSession::default(); + let vortex_array: ArrayRef = session + .arrow() + .from_arrow_record_batch(batch.clone(), &batch.schema()) + .unwrap(); // Apply Vortex expression - let session = VortexSession::default(); let mut ctx = session.create_execution_ctx(); let vortex_result = vortex_array .apply(&vortex_expr) diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 5de551f59ad..f036aa20b1f 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -140,12 +140,11 @@ mod common_tests { use object_store::memory::InMemory; use url::Url; use vortex::VortexSessionDefault; - use vortex::array::ArrayRef; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; + use vortex_arrow::ArrowSessionExt; use crate::VortexFormatFactory; use crate::VortexTableOptions; @@ -204,7 +203,10 @@ mod common_tests { where P: Into, { - let array = ArrayRef::from_arrow(batch, false)?; + let schema = batch.schema(); + let array = VX_SESSION + .arrow() + .from_arrow_record_batch(batch.clone(), &schema)?; let mut write = ObjectStoreWrite::new(Arc::clone(&self.store), &path.into()).await?; VX_SESSION .write_options() diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index f5fa2147100..89abfe68d18 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -674,7 +674,6 @@ mod tests { use object_store::memory::InMemory; use rstest::rstest; use vortex::VortexSessionDefault; - use vortex::array::ArrayRef; use vortex::buffer::Buffer; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; @@ -683,7 +682,6 @@ mod tests { use vortex::scan::selection::Selection; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; use super::*; use crate::VortexAccessPlan; @@ -834,7 +832,8 @@ mod tests { path: &str, rb: RecordBatch, ) -> anyhow::Result { - let array = ArrayRef::from_arrow(rb, false)?; + let schema = rb.schema(); + let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; let path = Path::parse(path)?; let mut write = ObjectStoreWrite::new(object_store, &path).await?; diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 4383061b734..6d347a6484b 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -10,6 +10,7 @@ use arrow_array::array::make_array; use arrow_array::ffi::FFI_ArrowArray; use arrow_array::ffi::FFI_ArrowSchema; use arrow_array::ffi::from_ffi; +use arrow_schema::Field; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; @@ -33,9 +34,9 @@ use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::error::vortex_panic; -use vortex_arrow::FromArrowArray; use crate::box_wrapper; +use crate::dtype::ARROW_SESSION; use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_variant; use crate::error::try_or; @@ -414,9 +415,10 @@ pub unsafe extern "C-unwind" fn vx_array_from_arrow( let ffi_array = unsafe { ptr::replace(array, FFI_ArrowArray::empty()) }; let ffi_schema = unsafe { ptr::replace(schema, FFI_ArrowSchema::empty()) }; let array_data = unsafe { from_ffi(ffi_array, &ffi_schema) }?; + let field = Field::try_from(&ffi_schema)?.with_nullable(nullable); drop(ffi_schema); let arrow_array = make_array(array_data); - let vortex_array = ArrayRef::from_arrow(arrow_array.as_ref(), nullable)?; + let vortex_array = ARROW_SESSION.from_arrow_array(arrow_array, &field)?; Ok(vx_array::new(vortex_array)) }) } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 806edf926dd..3c4a61f2ba1 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -22,7 +22,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -127,6 +126,7 @@ impl ArrowImportVTable for Json { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let DType::Extension(ext_dtype) = dtype else { return Ok(ArrowImport::Unsupported(array)); @@ -135,7 +135,7 @@ impl ArrowImportVTable for Json { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::new(ext_dtype.clone(), storage).into_array(), )) diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 663b29d9501..05e5c890558 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -128,7 +128,7 @@ mod tests { use arrow_schema::Schema; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; - use vortex_arrow::FromArrowArray; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use super::*; @@ -155,7 +155,9 @@ mod tests { ); // Convert to Vortex - ArrayRef::from_arrow(&struct_array, true) + SCAN_SESSION + .arrow() + .from_arrow_array_nullable(&struct_array, true) } fn create_arrow_schema() -> Arc { diff --git a/vortex-python/src/arrays/from_arrow.rs b/vortex-python/src/arrays/from_arrow.rs index 2853f398d64..20a6b5c433b 100644 --- a/vortex-python/src/arrays/from_arrow.rs +++ b/vortex-python/src/arrays/from_arrow.rs @@ -10,13 +10,11 @@ use arrow_schema::Field; use pyo3::exceptions::PyValueError; use pyo3::intern; use pyo3::prelude::*; -use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use crate::arrays::PyArrayRef; use crate::arrow::FromPyArrow; @@ -37,7 +35,9 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult> = obj.getattr(intern!(py, "chunks"))?.extract()?; @@ -45,7 +45,10 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult>>()?; let arrow_dtype = obj @@ -67,8 +70,10 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult>>()?; Ok(PyArrayRef::from( diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index bf4773fe82c..901ab1287c8 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -26,7 +26,6 @@ use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use crate::PyVortex; use crate::RUNTIME; @@ -473,7 +472,8 @@ fn try_arrow_stream_to_iterator( .into_iter() .map(|batch_result| -> VortexResult { let batch = batch_result.map_err(VortexError::from)?; - ArrayRef::from_arrow(batch, false) + let schema = batch.schema(); + session().arrow().from_arrow_record_batch(batch, &schema) }); Ok(Box::new(ArrayIteratorAdapter::new(dtype, vortex_iter))) diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 2c6a46ab06a..ffef079122d 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -44,7 +44,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_buffer::Buffer; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -237,7 +236,10 @@ impl TryFrom for LineStringData { impl LineStringData { /// Serialize line strings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&linestring_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &linestring_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -365,6 +367,7 @@ impl ArrowImportVTable for LineString { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -376,7 +379,7 @@ impl ArrowImportVTable for LineString { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index d1e2c37ebf4..89304f961f9 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -13,6 +13,7 @@ mod wkb; use std::fmt::Display; use std::sync::Arc; +use std::sync::LazyLock; use ::wkb::reader::GeometryType; use arrow_array::BinaryArray; @@ -57,7 +58,7 @@ use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSession; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -190,6 +191,10 @@ pub(crate) fn single_geometry( .ok_or_else(|| vortex_err!("spatial: constant operand decoded to no geometry")) } +/// Plan-time geometry literal decoding has no session in scope, so its storage arrays +/// (which carry no Arrow extension metadata) convert through a default [`ArrowSession`]. +static ARROW_SESSION: LazyLock = LazyLock::new(ArrowSession::default); + /// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native /// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only. pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { @@ -205,7 +210,7 @@ pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult VortexResult { let native = cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; - ArrayRef::from_arrow(native.to_array_ref().as_ref(), false) + ARROW_SESSION.from_arrow_array_nullable(native.to_array_ref().as_ref(), false) }; let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { @@ -301,12 +306,15 @@ pub(crate) fn geoarrow_metadata(spatial_metadata: &SpatialMetadata) -> Arc VortexResult { +pub(crate) fn geoarrow_to_wkb( + geoarrow_array: &dyn GeoArrowArray, + session: &ArrowSession, +) -> VortexResult { let wkb_type = GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&SpatialMetadata::default()))); let wkb = cast(geoarrow_array, &wkb_type) .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; - ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + session.from_arrow_array_nullable(wkb.to_array_ref().as_ref(), false) } /// Recover [`SpatialMetadata`] from GeoArrow metadata. diff --git a/vortex-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index a7e7104927a..84be29e58b3 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -37,7 +37,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -167,7 +166,10 @@ impl TryFrom for MultiLineStringData { impl MultiLineStringData { /// Serialize multilinestrings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&multilinestring_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &multilinestring_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -298,6 +300,7 @@ impl ArrowImportVTable for MultiLineString { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -309,7 +312,7 @@ impl ArrowImportVTable for MultiLineString { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index 46597ccdd5c..1caca024301 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -37,7 +37,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -155,7 +154,10 @@ impl TryFrom for MultiPointData { impl MultiPointData { /// Serialize multipoints to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&multipoint_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &multipoint_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -280,6 +282,7 @@ impl ArrowImportVTable for MultiPoint { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -291,7 +294,7 @@ impl ArrowImportVTable for MultiPoint { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 80078a4b07e..23c8cb67b2f 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -36,7 +36,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -165,7 +164,10 @@ impl TryFrom for MultiPolygonData { impl MultiPolygonData { /// Serialize multipolygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&multipolygon_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &multipolygon_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -297,6 +299,7 @@ impl ArrowImportVTable for MultiPolygon { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -308,7 +311,7 @@ impl ArrowImportVTable for MultiPolygon { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index 5dfbbb99949..505896d748e 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -35,7 +35,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -115,7 +114,10 @@ impl TryFrom for PointData { impl PointData { /// Serialize points to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&point_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &point_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -270,6 +272,7 @@ impl ArrowImportVTable for Point { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -281,7 +284,7 @@ impl ArrowImportVTable for Point { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index e727a7ab3cb..6ff784d11b6 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -36,7 +36,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -160,7 +159,10 @@ impl TryFrom for PolygonData { impl PolygonData { /// Serialize polygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&polygon_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb( + &polygon_array(self.0.storage_array(), ctx)?, + &ctx.session().arrow(), + ) } } @@ -292,6 +294,7 @@ impl ArrowImportVTable for Polygon { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -303,7 +306,7 @@ impl ArrowImportVTable for Polygon { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index ff25b4ba6a3..90025bc2b32 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -44,7 +44,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -317,6 +316,7 @@ impl ArrowImportVTable for Rect { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -328,7 +328,7 @@ impl ArrowImportVTable for Rect { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index 977e2855d51..14a2bc41fdc 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -30,7 +30,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -276,6 +275,7 @@ impl ArrowImportVTable for WellKnownBinary { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let Some(ext_dtype) = dtype.as_extension_opt() else { return Ok(ArrowImport::Unsupported(array)); @@ -290,7 +290,7 @@ impl ArrowImportVTable for WellKnownBinary { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::new(ext_dtype.clone(), storage).into_array(), )) diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index 7cec2cd9345..b92d0ee6918 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -29,7 +29,6 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -148,6 +147,7 @@ impl ArrowImportVTable for Vector { array: ArrowArrayRef, _field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { let DType::Extension(dtype) = dtype else { return Ok(ArrowImport::Unsupported(array)); @@ -162,7 +162,8 @@ impl ArrowImportVTable for Vector { return Ok(ArrowImport::Unsupported(array)); } - let storage = ArrayRef::from_arrow(array.as_ref() as &dyn Array, dtype.is_nullable())?; + let storage = + session.from_arrow_array_nullable(array.as_ref() as &dyn Array, dtype.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(dtype.clone(), storage)?.into_array(), )) @@ -374,8 +375,13 @@ mod tests { let field = Field::new("embedding", DataType::Int32, false); let int_array: ArrowArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let result = - ::from_arrow_array(&Vector, int_array, &field, &dtype)?; + let result = ::from_arrow_array( + &Vector, + int_array, + &field, + &dtype, + &ArrowSession::default(), + )?; assert!(matches!(result, ArrowImport::Unsupported(_))); Ok(()) } @@ -399,6 +405,7 @@ mod tests { fsl_arrow, &field, &DType::Extension(uuid_ext), + &ArrowSession::default(), )?; assert!(matches!(result, ArrowImport::Unsupported(_))); Ok(()) diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs index 12d86c41241..eb619909ff4 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs @@ -11,7 +11,7 @@ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSession; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -195,10 +195,14 @@ impl DatasetFixture for ClickBenchHits5kFixture { .collect::, _>>() .map_err(|e| vortex_err!("failed to read parquet batches: {e}"))?; + let arrow = ArrowSession::default(); Ok(ChunkedArray::from_iter( batches .into_iter() - .map(|batch| ArrayRef::from_arrow(batch, false)) + .map(|batch| { + let schema = batch.schema(); + arrow.from_arrow_record_batch(batch, &schema) + }) .collect::>>()?, ) .into_array()) diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs index 59e97a7e587..d67f39674b6 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs @@ -8,7 +8,7 @@ use tpchgen_arrow::RecordBatchIterator; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSession; use vortex_error::VortexResult; use crate::fixtures::DatasetFixture; @@ -17,10 +17,14 @@ const SCALE_FACTOR: f64 = 0.01; fn collect_batches_as_vortex(iter: impl RecordBatchIterator) -> VortexResult { let batches: Vec = iter.collect(); + let arrow = ArrowSession::default(); Ok(ChunkedArray::from_iter( batches .into_iter() - .map(|batch| ArrayRef::from_arrow(batch, false)) + .map(|batch| { + let schema = batch.schema(); + arrow.from_arrow_record_batch(batch, &schema) + }) .collect::>>()?, ) .into_array()) diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index ce85a7b4f69..bf56a38763f 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -12,7 +12,6 @@ use indicatif::ProgressBar; use parquet::arrow::ParquetRecordBatchStreamBuilder; use tokio::fs::File; use tokio::io::AsyncWriteExt; -use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::error::VortexExpect; @@ -20,8 +19,8 @@ use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; use vortex::session::VortexSession; +use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; /// Compression strategy to use when converting Parquet files to Vortex format. #[derive(Clone, Copy, Debug, Default, ValueEnum)] @@ -73,12 +72,16 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow let dtype = session .arrow() .from_arrow_schema(parquet.schema().as_ref())?; + let arrow_session = ArrowSession::clone(&session.arrow()); let mut vortex_stream = parquet .build()? - .map(|record_batch| { + .map(move |record_batch| { record_batch .map_err(|e| vortex_err!(External: e)) - .and_then(|rb| ArrayRef::from_arrow(rb, false)) + .and_then(|rb| { + let schema = rb.schema(); + arrow_session.from_arrow_record_batch(rb, &schema) + }) }) .boxed(); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index ecbe4a9cd2a..289500e2543 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -342,7 +342,6 @@ impl VortexSessionDefault for VortexSession { mod test { use std::path::PathBuf; - use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -375,7 +374,6 @@ mod test { use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex::array::arrays::ChunkedArray; use vortex::arrow::ArrowSessionExt; - use vortex::arrow::FromArrowArray; use vortex::session::VortexSession; let session = VortexSession::default(); @@ -391,7 +389,8 @@ mod test { let chunks: Vec<_> = reader .map(|record_batch| { let batch = record_batch?; - ArrayRef::from_arrow(batch, false) + let schema = batch.schema(); + session.arrow().from_arrow_record_batch(batch, &schema) }) .collect::>()?; let vortex_array = ChunkedArray::try_new(chunks, dtype)?.into_array(); From 685aa6570abbda986ce7ee289d9417cc91d93dc6 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 17:23:42 +0100 Subject: [PATCH 2/8] Thread ArrowSession explicitly and extract canonical conversion functions - `ParquetVariant::from_arrow_variant{,_nullable}` take an `&ArrowSession` threaded from vtables/kernels, so shredded storage children resolve through the registered import plugins. - `native_geometry_scalar_from_wkb` takes an `&ArrowSession`; vortex-duckdb passes its crate session. - `vx_array_from_arrow` takes a `const vx_session*` (header regenerated). - Extract each `FromArrowArray` impl body into a named, invocable `vortex_arrow::convert::from_arrow_*` function (primitive, decimal, temporal, bytes, byte-view, boolean, struct, list, list-view, fixed-size-list, null, dictionary, run-ends, dyn dispatch, record batch); the deprecated trait impls are now thin shims so the trait can eventually be deleted. Co-Authored-By: Claude Fable 5 Signed-off-by: Robert Kruszewski --- encodings/parquet-variant/src/array.rs | 38 +- encodings/parquet-variant/src/arrow.rs | 10 +- encodings/parquet-variant/src/kernel.rs | 19 +- encodings/parquet-variant/src/operations.rs | 23 +- encodings/parquet-variant/src/vtable.rs | 6 +- vortex-arrow/src/convert.rs | 661 ++++++++++++-------- vortex-arrow/src/lib.rs | 2 +- vortex-duckdb/src/convert/expr.rs | 4 +- vortex-ffi/src/array.rs | 11 +- vortex-spatial/src/extension/mod.rs | 22 +- vortex-spatial/src/scalar_fn/contains.rs | 4 +- vortex-spatial/src/scalar_fn/intersects.rs | 4 +- 12 files changed, 478 insertions(+), 326 deletions(-) diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index e54a3c9616a..5f8d7a350dc 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; -use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; @@ -67,10 +66,6 @@ pub struct ParquetVariantSlots { pub typed_value: Option, } -/// Variant storage children (metadata/value/typed_value) carry no Arrow extension metadata, -/// so [`ParquetVariant::from_arrow_variant`] converts them through a default [`ArrowSession`]. -pub(crate) static ARROW_SESSION: LazyLock = LazyLock::new(ArrowSession::default); - impl ParquetVariant { /// Creates a Parquet Variant array from canonical extension storage slots. /// @@ -96,20 +91,26 @@ impl ParquetVariant { ) } - /// Converts an Arrow `parquet_variant_compute::VariantArray` into Parquet Variant storage. - pub fn from_arrow_variant(arrow_variant: &ArrowVariantArray) -> VortexResult { - Self::from_arrow_variant_impl(arrow_variant, false) + /// Converts an Arrow `parquet_variant_compute::VariantArray` into Parquet Variant storage, + /// converting the storage children through `session`. + pub fn from_arrow_variant( + arrow_variant: &ArrowVariantArray, + session: &ArrowSession, + ) -> VortexResult { + Self::from_arrow_variant_impl(arrow_variant, false, session) } pub(crate) fn from_arrow_variant_nullable( arrow_variant: &ArrowVariantArray, + session: &ArrowSession, ) -> VortexResult { - Self::from_arrow_variant_impl(arrow_variant, true) + Self::from_arrow_variant_impl(arrow_variant, true, session) } fn from_arrow_variant_impl( arrow_variant: &ArrowVariantArray, force_nullable: bool, + session: &ArrowSession, ) -> VortexResult { let storage = arrow_variant.inner(); let mut value_nullable = false; @@ -135,17 +136,17 @@ impl ParquetVariant { } else { Validity::NonNullable }); - let metadata = ARROW_SESSION + let metadata = session .from_arrow_array_nullable(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; let value = arrow_variant .value_field() - .map(|v| ARROW_SESSION.from_arrow_array_nullable(v as &dyn ArrowArray, value_nullable)) + .map(|v| session.from_arrow_array_nullable(v as &dyn ArrowArray, value_nullable)) .transpose()?; let typed_value = arrow_variant .typed_value_field() - .map(|tv| ARROW_SESSION.from_arrow_array_nullable(tv.as_ref(), typed_value_nullable)) + .map(|tv| session.from_arrow_array_nullable(tv.as_ref(), typed_value_nullable)) .transpose()?; ParquetVariant::try_new(validity, metadata, value, typed_value).map(IntoArray::into_array) } @@ -513,6 +514,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -531,7 +533,7 @@ mod tests { fn assert_arrow_variant_storage_roundtrip(struct_array: StructArray) -> VortexResult<()> { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let inner = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant child"))?; @@ -582,7 +584,7 @@ mod tests { builder.append_variant(PqVariant::from(true)); let arrow_variant = builder.build(); - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.len(), 3); assert_eq!( @@ -614,7 +616,7 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.len(), 3); assert_eq!( vortex_arr.dtype(), @@ -705,7 +707,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; @@ -742,7 +744,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; @@ -814,7 +816,7 @@ mod tests { .with_path("a", &DataType::Int32)? .build(); let shredded = shred_variant(&json_to_variant(&json)?, &shredding)?; - let original = ParquetVariant::from_arrow_variant(&shredded)?; + let original = ParquetVariant::from_arrow_variant(&shredded, &SESSION.arrow())?; assert!( original .as_opt::() diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 01673a5d364..23590cbb59d 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -109,9 +109,9 @@ pub(crate) fn export_unshredded_storage_to_target( let arrow_variant = parquet_array.to_arrow(ctx)?; let unshredded = unshred_variant(&arrow_variant)?; let unshredded_array = if parquet_array.as_ref().dtype().is_nullable() { - ParquetVariant::from_arrow_variant_nullable(&unshredded)? + ParquetVariant::from_arrow_variant_nullable(&unshredded, &ctx.session().arrow())? } else { - ParquetVariant::from_arrow_variant(&unshredded)? + ParquetVariant::from_arrow_variant(&unshredded, &ctx.session().arrow())? }; let unshredded_parquet = unshredded_array.as_::(); export_storage_to_target(&unshredded_parquet, target_fields, ctx) @@ -263,7 +263,7 @@ impl ArrowImportVTable for ParquetVariant { array: ArrowArrayRef, field: &Field, dtype: &DType, - _session: &ArrowSession, + session: &ArrowSession, ) -> VortexResult { if !dtype.is_variant() || field @@ -277,9 +277,9 @@ impl ArrowImportVTable for ParquetVariant { let arrow_variant = ArrowVariantArray::try_new(array.as_struct())?; let imported = if dtype.is_nullable() { - ParquetVariant::from_arrow_variant_nullable(&arrow_variant)? + ParquetVariant::from_arrow_variant_nullable(&arrow_variant, session)? } else { - ParquetVariant::from_arrow_variant(&arrow_variant)? + ParquetVariant::from_arrow_variant(&arrow_variant, session)? }; Ok(ArrowImport::Imported(imported.into_array())) } diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index aea1ce4f52d..be6893a719b 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -117,7 +117,10 @@ impl ExecuteParentKernel for VariantGetKernel { let arrow_output = arrow_variant_get(&arrow_input, get_options)?; let output = if parent.options.dtype().is_none_or(DType::is_variant) { let arrow_variant_output = ArrowVariantArray::try_new(arrow_output.as_ref())?; - ParquetVariant::from_arrow_variant_nullable(&arrow_variant_output)? + ParquetVariant::from_arrow_variant_nullable( + &arrow_variant_output, + &ctx.session().arrow(), + )? } else { // Import through the same `as_type` field the cast targeted, so an extension target // dtype comes back as that extension rather than as its bare storage type. @@ -166,9 +169,9 @@ fn json_strings_to_variant( }; if nullable { - ParquetVariant::from_arrow_variant_nullable(&arrow_variant) + ParquetVariant::from_arrow_variant_nullable(&arrow_variant, &session.arrow()) } else { - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &session.arrow()) } } @@ -413,7 +416,7 @@ mod tests { builder.append_variant(PqVariant::from("hello")); builder.append_variant(PqVariant::from(true)); builder.append_variant(PqVariant::from(99i64)); - ParquetVariant::from_arrow_variant(&builder.build()) + ParquetVariant::from_arrow_variant(&builder.build(), &SESSION.arrow()) } fn make_nullable_array() -> VortexResult { @@ -430,13 +433,13 @@ mod tests { Some(NullBuffer::from(vec![true, false, true, false])), )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn make_unshredded_json_array(values: Vec>) -> VortexResult { let json: ArrowArrayRef = Arc::new(StringArray::from(values)); let arrow_variant = json_to_variant(&json)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn parse_path(path: &str) -> VortexResult { @@ -864,7 +867,7 @@ mod tests { fn make_partially_shredded_object_array() -> VortexResult { let arrow_variant = make_partially_shredded_arrow_variant()?; - let parquet_array = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let parquet_array = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = parquet_array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); @@ -981,7 +984,7 @@ mod tests { None, )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn assert_typed_value_i32( diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 16563dc7ea1..8517df27f22 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -372,6 +372,7 @@ fn parquet_variant_to_scalar(variant: PqVariant<'_, '_>) -> VortexResult #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; @@ -393,12 +394,20 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; use crate::operations::parquet_variant_to_scalar; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + fn binary_view_array(values: &[&[u8]]) -> ArrowArrayRef { let mut builder = BinaryViewBuilder::new(); for value in values { @@ -411,7 +420,7 @@ mod tests { arrow_variant: &ArrowVariantArray, rows: impl IntoIterator, ) -> VortexResult<()> { - let vortex_arr = ParquetVariant::from_arrow_variant(arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(arrow_variant, &SESSION.arrow())?; for index in rows { let expected_inner = parquet_variant_to_scalar(arrow_variant.try_value(index)?)?; @@ -443,7 +452,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); @@ -484,7 +493,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let present_variant_null = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; @@ -521,7 +530,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); assert!( @@ -550,7 +559,7 @@ mod tests { builder.append_variant(PqVariant::from(2i32)); let arrow_variant = builder.build(); - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!( vortex_arr.dtype(), @@ -671,7 +680,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let row0 = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let row0 = row0.as_variant().value().unwrap().as_list(); @@ -763,7 +772,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let object = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let object = object.as_variant().value().unwrap().as_struct(); diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index e83fe0ca83a..9e1c03fbeaa 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -322,6 +322,7 @@ mod tests { use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::BitBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; @@ -387,7 +388,10 @@ mod tests { None, )?; - ParquetVariant::from_arrow_variant(&ArrowVariantArray::try_new(&arrow_storage)?) + ParquetVariant::from_arrow_variant( + &ArrowVariantArray::try_new(&arrow_storage)?, + &SESSION.arrow(), + ) } #[fixture] diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 7b0b22ca136..0105b69ef62 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -1,9 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -// This module hosts the canonical Arrow → Vortex conversion machinery, implemented as the -// deprecated `FromArrowArray` trait. Internal callers go through `from_arrow_dyn` / -// `from_arrow_batch`; external callers use the `ArrowSession` methods. +//! Canonical (non-plugin) conversions of Arrow arrays into Vortex arrays. +//! +//! Each `from_arrow_*` function converts one Arrow array shape; [`from_arrow_dyn`] dispatches on +//! the Arrow [`DataType`]. These functions have no awareness of Arrow extension types — the +//! plugin-aware entry points are the [`ArrowSession`](crate::ArrowSession) methods, which fall +//! back to these conversions for non-extension data. The deprecated [`FromArrowArray`] impls are +//! thin shims over these functions and will eventually be removed. #![allow(deprecated)] use std::sync::Arc; @@ -109,17 +113,6 @@ use crate::dtype::from_arrow_time_unit; /// /// This mirrors [`IntoArray`] for Arrow buffer types; a separate trait is required because both /// [`IntoArray`] and the Arrow buffer types are foreign to this crate. -/// Canonical (non-plugin) conversion of an Arrow array into a Vortex array. Internal -/// entry point for callers that don't dispatch Arrow extension plugins. -pub(crate) fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult { - ArrayRef::from_arrow(array, nullable) -} - -/// Canonical (non-plugin) conversion of an Arrow [`RecordBatch`] into a Vortex struct array. -pub(crate) fn from_arrow_batch(batch: &RecordBatch, nullable: bool) -> VortexResult { - ArrayRef::from_arrow(batch, nullable) -} - pub trait IntoVortexArray { /// Convert this Arrow buffer into a non-nullable Vortex array without copying. fn into_array(self) -> ArrayRef; @@ -169,13 +162,27 @@ where } } +/// Zero-copy conversion of an Arrow numeric primitive array into a Vortex array. +/// +/// Use [`from_arrow_temporal`] for Arrow temporal arrays, which carry a logical dtype beyond +/// their physical primitive storage. +pub fn from_arrow_primitive( + value: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult +where + T::Native: NativePType, +{ + let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone()); + let validity = nulls(value.nulls(), nullable)?; + Ok(PrimitiveArray::new(buffer, validity).into_array()) +} + macro_rules! impl_from_arrow_primitive { ($T:path) => { impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef { fn from_arrow(value: &ArrowPrimitiveArray<$T>, nullable: bool) -> VortexResult { - let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone()); - let validity = nulls(value.nulls(), nullable)?; - Ok(PrimitiveArray::new(buffer, validity).into_array()) + from_arrow_primitive(value, nullable) } } }; @@ -193,56 +200,87 @@ impl_from_arrow_primitive!(Float16Type); impl_from_arrow_primitive!(Float32Type); impl_from_arrow_primitive!(Float64Type); +/// Zero-copy conversion of an Arrow `Decimal32` array into a Vortex decimal array. +pub fn from_arrow_decimal32( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal32(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal64` array into a Vortex decimal array. +pub fn from_arrow_decimal64( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal64(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal128` array into a Vortex decimal array. +pub fn from_arrow_decimal128( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal128(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal256` array into a Vortex decimal array. +pub fn from_arrow_decimal256( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + // SAFETY: Our i256 implementation has the same bit-pattern representation of the + // arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values + // of either type. + let buffer = unsafe { std::mem::transmute::, Buffer>(buffer) }; + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - // SAFETY: Our i256 implementation has the same bit-pattern representation of the - // arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values - // of either type. - let buffer = - unsafe { std::mem::transmute::, Buffer>(buffer) }; - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal256(array, nullable) } } @@ -253,7 +291,7 @@ macro_rules! impl_from_arrow_temporal { value: &ArrowPrimitiveArray<$T>, nullable: bool, ) -> vortex_error::VortexResult { - temporal_array(value, nullable) + from_arrow_temporal(value, nullable) } } }; @@ -275,7 +313,9 @@ impl_from_arrow_temporal!(Time64NanosecondType); impl_from_arrow_temporal!(Date32Type); impl_from_arrow_temporal!(Date64Type); -fn temporal_array( +/// Conversion of an Arrow temporal array (timestamp/date/time) into a Vortex temporal +/// extension array. +pub fn from_arrow_temporal( value: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult @@ -306,68 +346,92 @@ where }) } +/// Zero-copy conversion of an Arrow (large) string/binary array into a Vortex `VarBin` array. +pub fn from_arrow_bytes( + value: &GenericByteArray, + nullable: bool, +) -> VortexResult +where + ::Offset: IntegerPType, +{ + let dtype = match T::DATA_TYPE { + DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()), + DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()), + dt => vortex_panic!("Invalid data type for ByteArray: {dt}"), + }; + // SAFETY: Arrow arrays are already validated (valid UTF-8, valid offsets, correct validity). + Ok(unsafe { + VarBinArray::new_unchecked( + value.offsets().clone().into_array(), + ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::()), + dtype, + nulls(value.nulls(), nullable)?, + ) + } + .into_array()) +} + impl FromArrowArray<&GenericByteArray> for ArrayRef where ::Offset: IntegerPType, { fn from_arrow(value: &GenericByteArray, nullable: bool) -> VortexResult { - let dtype = match T::DATA_TYPE { - DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()), - DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()), - dt => vortex_panic!("Invalid data type for ByteArray: {dt}"), - }; - // SAFETY: Arrow arrays are already validated (valid UTF-8, valid offsets, correct validity). - Ok(unsafe { - VarBinArray::new_unchecked( - value.offsets().clone().into_array(), - ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::()), - dtype, - nulls(value.nulls(), nullable)?, - ) - } - .into_array()) + from_arrow_bytes(value, nullable) } } -impl FromArrowArray<&GenericByteViewArray> for ArrayRef { - fn from_arrow(value: &GenericByteViewArray, nullable: bool) -> VortexResult { - let dtype = match T::DATA_TYPE { - DataType::BinaryView => DType::Binary(nullable.into()), - DataType::Utf8View => DType::Utf8(nullable.into()), - dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"), - }; +/// Zero-copy conversion of an Arrow string/binary view array into a Vortex `VarBinView` array. +pub fn from_arrow_byte_view( + value: &GenericByteViewArray, + nullable: bool, +) -> VortexResult { + let dtype = match T::DATA_TYPE { + DataType::BinaryView => DType::Binary(nullable.into()), + DataType::Utf8View => DType::Utf8(nullable.into()), + dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"), + }; - let views_buffer = Buffer::from_byte_buffer( - Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(), - ); + let views_buffer = Buffer::from_byte_buffer( + Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(), + ); - // SAFETY: arrow-rs ByteViewArray already checks the same invariants, we inherit those - // guarantees by zero-copy constructing from one. - Ok(unsafe { - VarBinViewArray::new_unchecked( - views_buffer, - Arc::from( - value - .data_buffers() - .iter() - .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::())) - .collect::>(), - ), - dtype, - nulls(value.nulls(), nullable)?, - ) - .into_array() - }) + // SAFETY: arrow-rs ByteViewArray already checks the same invariants, we inherit those + // guarantees by zero-copy constructing from one. + Ok(unsafe { + VarBinViewArray::new_unchecked( + views_buffer, + Arc::from( + value + .data_buffers() + .iter() + .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::())) + .collect::>(), + ), + dtype, + nulls(value.nulls(), nullable)?, + ) + .into_array() + }) +} + +impl FromArrowArray<&GenericByteViewArray> for ArrayRef { + fn from_arrow(value: &GenericByteViewArray, nullable: bool) -> VortexResult { + from_arrow_byte_view(value, nullable) } } +/// Zero-copy conversion of an Arrow boolean array into a Vortex `Bool` array. +pub fn from_arrow_boolean(value: &ArrowBooleanArray, nullable: bool) -> VortexResult { + Ok(BoolArray::new( + value.values().clone().into(), + nulls(value.nulls(), nullable)?, + ) + .into_array()) +} + impl FromArrowArray<&ArrowBooleanArray> for ArrayRef { fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> VortexResult { - Ok(BoolArray::new( - value.values().clone().into(), - nulls(value.nulls(), nullable)?, - ) - .into_array()) + from_arrow_boolean(value, nullable) } } @@ -419,84 +483,113 @@ pub(crate) fn remove_nulls(data: arrow_data::ArrayData) -> VortexResult VortexResult { + Ok(StructArray::try_new( + value.column_names().iter().copied().collect(), + value + .columns() + .iter() + .zip(value.fields()) + .map(|(c, field)| { + // Arrow pushes down nulls, even into non-nullable fields. So we strip them + // out here because Vortex is a little more strict. + if c.null_count() > 0 && !field.is_nullable() { + let stripped = make_array(remove_nulls(c.into_data())?); + from_arrow_dyn(stripped.as_ref(), false) + } else { + from_arrow_dyn(c.as_ref(), field.is_nullable()) + } + }) + .collect::>>()?, + value.len(), + nulls(value.nulls(), nullable)?, + )? + .into_array()) +} + impl FromArrowArray<&ArrowStructArray> for ArrayRef { fn from_arrow(value: &ArrowStructArray, nullable: bool) -> VortexResult { - Ok(StructArray::try_new( - value.column_names().iter().copied().collect(), - value - .columns() - .iter() - .zip(value.fields()) - .map(|(c, field)| { - // Arrow pushes down nulls, even into non-nullable fields. So we strip them - // out here because Vortex is a little more strict. - if c.null_count() > 0 && !field.is_nullable() { - let stripped = make_array(remove_nulls(c.into_data())?); - Self::from_arrow(stripped.as_ref(), false) - } else { - Self::from_arrow(c.as_ref(), field.is_nullable()) - } - }) - .collect::>>()?, - value.len(), - nulls(value.nulls(), nullable)?, - )? - .into_array()) + from_arrow_struct(value, nullable) } } +/// Conversion of an Arrow (large) list array into a Vortex `List` array. +pub fn from_arrow_list( + value: &GenericListArray, + nullable: bool, +) -> VortexResult { + // Extract the validity of the underlying element array. + let elements_are_nullable = match value.data_type() { + DataType::List(field) => field.is_nullable(), + DataType::LargeList(field) => field.is_nullable(), + dt => vortex_panic!("Invalid data type for ListArray: {dt}"), + }; + + let elements = from_arrow_dyn(value.values().as_ref(), elements_are_nullable)?; + + // `offsets` are always non-nullable. + let offsets = value.offsets().clone().into_array(); + let nulls = nulls(value.nulls(), nullable)?; + + Ok(ListArray::try_new(elements, offsets, nulls)?.into_array()) +} + impl FromArrowArray<&GenericListArray> for ArrayRef { fn from_arrow(value: &GenericListArray, nullable: bool) -> VortexResult { - // Extract the validity of the underlying element array. - let elements_are_nullable = match value.data_type() { - DataType::List(field) => field.is_nullable(), - DataType::LargeList(field) => field.is_nullable(), - dt => vortex_panic!("Invalid data type for ListArray: {dt}"), - }; + from_arrow_list(value, nullable) + } +} - let elements = Self::from_arrow(value.values().as_ref(), elements_are_nullable)?; +/// Conversion of an Arrow (large) list-view array into a Vortex `ListView` array. +pub fn from_arrow_list_view( + array: &GenericListViewArray, + nullable: bool, +) -> VortexResult { + // Extract the validity of the underlying element array. + let elements_are_nullable = match array.data_type() { + DataType::ListView(field) => field.is_nullable(), + DataType::LargeListView(field) => field.is_nullable(), + dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"), + }; - // `offsets` are always non-nullable. - let offsets = value.offsets().clone().into_array(); - let nulls = nulls(value.nulls(), nullable)?; + let elements = from_arrow_dyn(array.values().as_ref(), elements_are_nullable)?; - Ok(ListArray::try_new(elements, offsets, nulls)?.into_array()) - } + // `offsets` and `sizes` are always non-nullable. + let offsets = array.offsets().clone().into_array(); + let sizes = array.sizes().clone().into_array(); + let nulls = nulls(array.nulls(), nullable)?; + + Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array()) } impl FromArrowArray<&GenericListViewArray> for ArrayRef { fn from_arrow(array: &GenericListViewArray, nullable: bool) -> VortexResult { - // Extract the validity of the underlying element array. - let elements_are_nullable = match array.data_type() { - DataType::ListView(field) => field.is_nullable(), - DataType::LargeListView(field) => field.is_nullable(), - dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"), - }; - - let elements = Self::from_arrow(array.values().as_ref(), elements_are_nullable)?; + from_arrow_list_view(array, nullable) + } +} - // `offsets` and `sizes` are always non-nullable. - let offsets = array.offsets().clone().into_array(); - let sizes = array.sizes().clone().into_array(); - let nulls = nulls(array.nulls(), nullable)?; +/// Conversion of an Arrow fixed-size list array into a Vortex `FixedSizeList` array. +pub fn from_arrow_fixed_size_list( + array: &ArrowFixedSizeListArray, + nullable: bool, +) -> VortexResult { + let DataType::FixedSizeList(field, list_size) = array.data_type() else { + vortex_panic!("Invalid data type for ListArray: {}", array.data_type()); + }; - Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array()) - } + Ok(FixedSizeListArray::try_new( + from_arrow_dyn(array.values().as_ref(), field.is_nullable())?, + *list_size as u32, + nulls(array.nulls(), nullable)?, + array.len(), + )? + .into_array()) } impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { fn from_arrow(array: &ArrowFixedSizeListArray, nullable: bool) -> VortexResult { - let DataType::FixedSizeList(field, list_size) = array.data_type() else { - vortex_panic!("Invalid data type for ListArray: {}", array.data_type()); - }; - - Ok(FixedSizeListArray::try_new( - Self::from_arrow(array.values().as_ref(), field.is_nullable())?, - *list_size as u32, - nulls(array.nulls(), nullable)?, - array.len(), - )? - .into_array()) + from_arrow_fixed_size_list(array, nullable) } } @@ -570,24 +663,36 @@ impl FromArrowArray<&ArrowMapArray> for ArrayRef { ) } } +/// Conversion of an Arrow null array into a Vortex `Null` array. +pub fn from_arrow_null(value: &ArrowNullArray, nullable: bool) -> VortexResult { + vortex_ensure!( + nullable, + "Cannot convert an Arrow NullArray into a non-nullable Vortex array" + ); + Ok(NullArray::new(value.len()).into_array()) +} impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { - vortex_ensure!( - nullable, - "Cannot convert an Arrow NullArray into a non-nullable Vortex array" - ); - Ok(NullArray::new(value.len()).into_array()) + from_arrow_null(value, nullable) } } +/// Conversion of an Arrow dictionary array into a Vortex `Dict` array. +pub fn from_arrow_dictionary( + array: &DictionaryArray, + nullable: bool, +) -> VortexResult { + let keys = AnyDictionaryArray::keys(array); + let keys = from_arrow_dyn(keys, keys.is_nullable())?; + let values = from_arrow_dyn(array.values().as_ref(), nullable)?; + // SAFETY: we assume that Arrow has checked the invariants on construction. + Ok(unsafe { DictArray::new_unchecked(keys, values) }) +} + impl FromArrowArray<&DictionaryArray> for DictArray { fn from_arrow(array: &DictionaryArray, nullable: bool) -> VortexResult { - let keys = AnyDictionaryArray::keys(array); - let keys = ArrayRef::from_arrow(keys, keys.is_nullable())?; - let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?; - // SAFETY: we assume that Arrow has checked the invariants on construction. - Ok(unsafe { DictArray::new_unchecked(keys, values) }) + from_arrow_dictionary(array, nullable) } } @@ -613,138 +718,150 @@ pub(crate) fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> VortexResult< } } -impl FromArrowArray<&dyn ArrowArray> for ArrayRef { - fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult { - match array.data_type() { - DataType::Boolean => Self::from_arrow(array.as_boolean(), nullable), - DataType::UInt8 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int8 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Utf8 => Self::from_arrow(array.as_string::(), nullable), - DataType::LargeUtf8 => Self::from_arrow(array.as_string::(), nullable), - DataType::Binary => Self::from_arrow(array.as_binary::(), nullable), - DataType::LargeBinary => Self::from_arrow(array.as_binary::(), nullable), - DataType::BinaryView => Self::from_arrow(array.as_binary_view(), nullable), - DataType::Utf8View => Self::from_arrow(array.as_string_view(), nullable), - DataType::Struct(_) => Self::from_arrow(array.as_struct(), nullable), - DataType::List(_) => Self::from_arrow(array.as_list::(), nullable), - DataType::LargeList(_) => Self::from_arrow(array.as_list::(), nullable), - DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), - DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), - DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), - DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), - DataType::Null => Self::from_arrow(as_null_array(array), nullable), - DataType::Timestamp(u, _) => match u { - ArrowTimeUnit::Second => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Millisecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Microsecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Nanosecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - }, - DataType::Date32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Date64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Time32(u) => match u { - ArrowTimeUnit::Second => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Millisecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(), - }, - DataType::Time64(u) => match u { - ArrowTimeUnit::Microsecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Nanosecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(), - }, - DataType::Decimal32(..) => { - Self::from_arrow(array.as_primitive::(), nullable) +/// Canonical conversion of any supported Arrow array into a Vortex array, dispatching on the +/// Arrow [`DataType`]. +pub fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult { + match array.data_type() { + DataType::Boolean => from_arrow_boolean(array.as_boolean(), nullable), + DataType::UInt8 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int8 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Utf8 => from_arrow_bytes(array.as_string::(), nullable), + DataType::LargeUtf8 => from_arrow_bytes(array.as_string::(), nullable), + DataType::Binary => from_arrow_bytes(array.as_binary::(), nullable), + DataType::LargeBinary => from_arrow_bytes(array.as_binary::(), nullable), + DataType::BinaryView => from_arrow_byte_view(array.as_binary_view(), nullable), + DataType::Utf8View => from_arrow_byte_view(array.as_string_view(), nullable), + DataType::Struct(_) => from_arrow_struct(array.as_struct(), nullable), + DataType::List(_) => from_arrow_list(array.as_list::(), nullable), + DataType::LargeList(_) => from_arrow_list(array.as_list::(), nullable), + DataType::ListView(_) => from_arrow_list_view(array.as_list_view::(), nullable), + DataType::LargeListView(_) => from_arrow_list_view(array.as_list_view::(), nullable), + DataType::FixedSizeList(..) => { + from_arrow_fixed_size_list(array.as_fixed_size_list(), nullable) + } + DataType::Null => from_arrow_null(as_null_array(array), nullable), + DataType::Timestamp(u, _) => match u { + ArrowTimeUnit::Second => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Millisecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Microsecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal64(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + ArrowTimeUnit::Nanosecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal128(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + }, + DataType::Date32 => from_arrow_temporal(array.as_primitive::(), nullable), + DataType::Date64 => from_arrow_temporal(array.as_primitive::(), nullable), + DataType::Time32(u) => match u { + ArrowTimeUnit::Second => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal256(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + ArrowTimeUnit::Millisecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Dictionary(key_type, _) => match key_type.as_ref() { - DataType::Int8 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int16 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int32 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int64 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt8 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt16 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt32 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt64 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"), - }, - dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"), + ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(), + }, + DataType::Time64(u) => match u { + ArrowTimeUnit::Microsecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Nanosecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(), + }, + DataType::Decimal32(..) => { + from_arrow_decimal32(array.as_primitive::(), nullable) + } + DataType::Decimal64(..) => { + from_arrow_decimal64(array.as_primitive::(), nullable) + } + DataType::Decimal128(..) => { + from_arrow_decimal128(array.as_primitive::(), nullable) } + DataType::Decimal256(..) => { + from_arrow_decimal256(array.as_primitive::(), nullable) + } + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::Int8 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int16 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int32 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int64 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt8 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt16 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt32 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt64 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"), + }, + dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"), + } +} + +impl FromArrowArray<&dyn ArrowArray> for ArrayRef { + fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult { + from_arrow_dyn(array, nullable) } } +/// Canonical conversion of an Arrow [`RecordBatch`] into a Vortex struct array. +pub fn from_arrow_batch(batch: &RecordBatch, nullable: bool) -> VortexResult { + from_arrow_struct(&arrow_array::StructArray::from(batch.clone()), nullable) +} + impl FromArrowArray for ArrayRef { fn from_arrow(array: RecordBatch, nullable: bool) -> VortexResult { - ArrayRef::from_arrow(&arrow_array::StructArray::from(array), nullable) + from_arrow_batch(&array, nullable) } } impl FromArrowArray<&RecordBatch> for ArrayRef { fn from_arrow(array: &RecordBatch, nullable: bool) -> VortexResult { - Self::from_arrow(array.clone(), nullable) + from_arrow_batch(array, nullable) } } diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 865ac9b190b..81310e8aedf 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -24,7 +24,7 @@ use vortex_array::legacy_session; use vortex_error::VortexResult; use vortex_session::VortexSession; -mod convert; +pub mod convert; mod datum; pub mod dtype; mod executor; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 5ee407994de..e7b4a31c6a5 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -18,6 +18,7 @@ use vortex::aggregate_fn::fns::max::Max; use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; use vortex::aggregate_fn::fns::sum::Sum; +use vortex::arrow::ArrowSessionExt; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -64,6 +65,7 @@ use vortex_spatial::scalar_fn::contains::SpatialContains; use vortex_spatial::scalar_fn::distance::SpatialDistance; use vortex_spatial::scalar_fn::intersects::SpatialIntersects; +use crate::SESSION; use crate::convert::dtype::FromLogicalType; use crate::cpp::DUCKDB_TYPE; use crate::cpp::DUCKDB_VX_EXPR_TYPE; @@ -158,7 +160,7 @@ fn spatial_operand( let Some(buf) = storage.as_binary_opt().and_then(|b| b.value()) else { return Ok(None); }; - Ok(native_geometry_scalar_from_wkb(buf.as_slice())?.map(lit)) + Ok(native_geometry_scalar_from_wkb(buf.as_slice(), &SESSION.arrow())?.map(lit)) } Some(BoundColumnRef(col_ref)) if is_native_spatial_column(ctx.fields, col_ref.name.as_ref()) => diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 6d347a6484b..d009a9e4c47 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -34,9 +34,9 @@ use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::error::vortex_panic; +use vortex_arrow::ArrowSessionExt; use crate::box_wrapper; -use crate::dtype::ARROW_SESSION; use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_variant; use crate::error::try_or; @@ -399,11 +399,12 @@ pub extern "C-unwind" fn vx_array_new_primitive( /// /// // export an Arrow record batch into (array, schema), then: /// vx_error* error = NULL; -/// const vx_array* vx = vx_array_from_arrow(&array, &schema, false, &error); +/// const vx_array* vx = vx_array_from_arrow(session, &array, &schema, false, &error); /// // ... push it to a sink or write it ... /// vx_array_free(vx); #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_array_from_arrow( + session: *const vx_session, array: *mut FFI_ArrowArray, schema: *mut FFI_ArrowSchema, nullable: bool, @@ -412,13 +413,14 @@ pub unsafe extern "C-unwind" fn vx_array_from_arrow( try_or_default(error_out, || { vortex_ensure!(!array.is_null(), "null arrow array"); vortex_ensure!(!schema.is_null(), "null arrow schema"); + let session = vx_session::as_ref(session); let ffi_array = unsafe { ptr::replace(array, FFI_ArrowArray::empty()) }; let ffi_schema = unsafe { ptr::replace(schema, FFI_ArrowSchema::empty()) }; let array_data = unsafe { from_ffi(ffi_array, &ffi_schema) }?; let field = Field::try_from(&ffi_schema)?.with_nullable(nullable); drop(ffi_schema); let arrow_array = make_array(array_data); - let vortex_array = ARROW_SESSION.from_arrow_array(arrow_array, &field)?; + let vortex_array = session.arrow().from_arrow_array(arrow_array, &field)?; Ok(vx_array::new(vortex_array)) }) } @@ -646,6 +648,7 @@ mod tests { use crate::scalar::*; use crate::session::vx_session_free; use crate::session::vx_session_new; + use crate::session::vx_session_new_with; use crate::tests::assert_error; use crate::tests::assert_no_error; @@ -979,9 +982,11 @@ mod tests { let data = ArrowArrayTrait::into_data(arrow_array::StructArray::from(batch)); let (mut ffi_array, mut ffi_schema) = to_ffi(&data).unwrap(); + let session = vx_session_new_with(|s| s); let mut error = ptr::null_mut(); let vx = unsafe { vx_array_from_arrow( + session, &raw mut ffi_array, &raw mut ffi_schema, false, diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index 89304f961f9..3b8f14c17cc 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -13,7 +13,6 @@ mod wkb; use std::fmt::Display; use std::sync::Arc; -use std::sync::LazyLock; use ::wkb::reader::GeometryType; use arrow_array::BinaryArray; @@ -191,13 +190,12 @@ pub(crate) fn single_geometry( .ok_or_else(|| vortex_err!("spatial: constant operand decoded to no geometry")) } -/// Plan-time geometry literal decoding has no session in scope, so its storage arrays -/// (which carry no Arrow extension metadata) convert through a default [`ArrowSession`]. -static ARROW_SESSION: LazyLock = LazyLock::new(ArrowSession::default); - /// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native /// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only. -pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { +pub fn native_geometry_scalar_from_wkb( + bytes: &[u8], + session: &ArrowSession, +) -> VortexResult> { let metadata = geoarrow_metadata(&SpatialMetadata::default()); let binary = BinaryArray::from(vec![Some(bytes)]); let wkb = GenericWkbArray::::try_from(( @@ -210,7 +208,7 @@ pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult VortexResult { let native = cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; - ARROW_SESSION.from_arrow_array_nullable(native.to_array_ref().as_ref(), false) + session.from_arrow_array_nullable(native.to_array_ref().as_ref(), false) }; let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { @@ -335,6 +333,7 @@ pub(crate) fn spatial_metadata_from_arrow(metadata: &Metadata) -> SpatialMetadat mod tests { use prost::Message; use vortex_array::dtype::DType; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -343,9 +342,16 @@ mod tests { use super::MultiPoint; use super::Point; use super::Polygon; - use super::native_geometry_scalar_from_wkb; use crate::extension::SpatialMetadata; + /// Test shim: decode with an explicitly constructed session. + fn native_geometry_scalar_from_wkb( + bytes: &[u8], + ) -> VortexResult> { + let session = vortex_array::array_session(); + super::native_geometry_scalar_from_wkb(bytes, &session.arrow()) + } + #[test] fn test_metadata() { let meta = SpatialMetadata { diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 599c0eee2be..8850e59f751 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -152,6 +152,7 @@ mod tests { use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -175,7 +176,8 @@ mod tests { let mut buf = Vec::new(); wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; - let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + let session = vortex_array::array_session(); + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf, &session.arrow())? .ok_or_else(|| vortex_err!("unsupported geometry type"))?; Ok(ConstantArray::new(scalar, len).into_array()) } diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index bdabd2b9967..77d33886ff3 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -151,6 +151,7 @@ mod tests { use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -196,7 +197,8 @@ mod tests { let mut buf = Vec::new(); wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; - let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + let session = vortex_array::array_session(); + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf, &session.arrow())? .ok_or_else(|| vortex_err!("unsupported geometry type"))?; Ok(ConstantArray::new(scalar, len).into_array()) } From 86ce596a12a3b1bb80d69a6919854bf4a17835cb Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 22:49:00 +0100 Subject: [PATCH 3/8] Update C++ API for session-taking vx_array_from_arrow Array::from_arrow takes a Session; C header regenerated. Co-Authored-By: Claude Fable 5 Signed-off-by: Robert Kruszewski --- lang/cpp/include/vortex/array.hpp | 2 +- lang/cpp/src/array.cpp | 4 ++-- lang/cpp/tests/arrow.cpp | 2 +- lang/cpp/tests/string_binary.cpp | 4 ++-- vortex-ffi/cinclude/vortex.h | 9 ++++++--- 5 files changed, 12 insertions(+), 9 deletions(-) diff --git a/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp index 47652dd14fa..3c0df8865c7 100644 --- a/lang/cpp/include/vortex/array.hpp +++ b/lang/cpp/include/vortex/array.hpp @@ -136,7 +136,7 @@ class Array { * Import an Arrow array. Consumes both "array" and "schema", do not use * or release them afterwards. For a record batch pass nullable = false. */ - static Array from_arrow(ArrowArray *array, ArrowSchema *schema, bool nullable); + static Array from_arrow(const Session &session, ArrowArray *array, ArrowSchema *schema, bool nullable); size_t size() const; bool nullable() const; diff --git a/lang/cpp/src/array.cpp b/lang/cpp/src/array.cpp index 5e46c953bba..d68d8729380 100644 --- a/lang/cpp/src/array.cpp +++ b/lang/cpp/src/array.cpp @@ -179,9 +179,9 @@ Array Array::primitive_raw(vx_ptype ptype, const void *data, size_t len, const V return Access::adopt(out); } -Array Array::from_arrow(ArrowArray *array, ArrowSchema *schema, bool nullable) { +Array Array::from_arrow(const Session &session, ArrowArray *array, ArrowSchema *schema, bool nullable) { vx_error *error = nullptr; - const vx_array *out = vx_array_from_arrow(array, schema, nullable, &error); + const vx_array *out = vx_array_from_arrow(Access::c_ptr(session), array, schema, nullable, &error); throw_on_error(error); return Access::adopt(out); } diff --git a/lang/cpp/tests/arrow.cpp b/lang/cpp/tests/arrow.cpp index 0270c2e23c1..bfd287fcecb 100644 --- a/lang/cpp/tests/arrow.cpp +++ b/lang/cpp/tests/arrow.cpp @@ -73,7 +73,7 @@ TEST_CASE("Import Arrow array as Vortex array", "[arrow]") { ArrowArrayMove(arr.get(), &raw_arr); ArrowSchemaMove(schema.get(), &raw_schema); - Array vx = Array::from_arrow(&raw_arr, &raw_schema, false); + Array vx = Array::from_arrow(session, &raw_arr, &raw_schema, false); REQUIRE(vx.size() == 3); REQUIRE(vx.has_dtype(DataTypeVariant::Struct)); diff --git a/lang/cpp/tests/string_binary.cpp b/lang/cpp/tests/string_binary.cpp index 16ddb7eaa3c..c2df5a656ed 100644 --- a/lang/cpp/tests/string_binary.cpp +++ b/lang/cpp/tests/string_binary.cpp @@ -42,7 +42,7 @@ Array strings_from_arrow(std::span values, bool with_nul ArrowSchema raw_schema = {}; ArrowArrayMove(arr.get(), &raw_arr); ArrowSchemaMove(schema.get(), &raw_schema); - return Array::from_arrow(&raw_arr, &raw_schema, true); + return Array::from_arrow(Session(), &raw_arr, &raw_schema, true); } Array bytes_from_arrow(std::span values) { @@ -62,7 +62,7 @@ Array bytes_from_arrow(std::span values) { ArrowSchema raw_schema = {}; ArrowArrayMove(arr.get(), &raw_arr); ArrowSchemaMove(schema.get(), &raw_schema); - return Array::from_arrow(&raw_arr, &raw_schema, true); + return Array::from_arrow(Session(), &raw_arr, &raw_schema, true); } TEST_CASE("String view over utf8 array", "[strings]") { diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 5fe99daf1e1..c7bc179e7ee 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -707,12 +707,15 @@ const vx_array *vx_array_new_primitive(vx_ptype ptype, * * // export an Arrow record batch into (array, schema), then: * vx_error* error = NULL; - * const vx_array* vx = vx_array_from_arrow(&array, &schema, false, &error); + * const vx_array* vx = vx_array_from_arrow(session, &array, &schema, false, &error); * // ... push it to a sink or write it ... * vx_array_free(vx); */ -const vx_array * -vx_array_from_arrow(FFI_ArrowArray *array, FFI_ArrowSchema *schema, bool nullable, vx_error **error_out); +const vx_array *vx_array_from_arrow(const vx_session *session, + FFI_ArrowArray *array, + FFI_ArrowSchema *schema, + bool nullable, + vx_error **error_out); /** * Return UTF-8 string at "index" in a canonical Utf8 array. From 1accaaaa1675bf6a1a5d04cee959e270a4ff23f3 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 10 Aug 2026 13:37:05 +0100 Subject: [PATCH 4/8] Convert the remaining FromArrowArray call sites Two `ArrayRef::from_arrow` call sites landed on develop while this branch was open, so the rebase reintroduced deprecation warnings the PR exists to remove. Both read a whole Arrow `RecordBatch`, so they convert to `ArrowSession::from_arrow_record_batch`, matching the other migrated readers. Signed-off-by: "Robert Kruszewski" --- benchmarks/string-bench/src/lib.rs | 12 ++++++++---- vortex-btrblocks/src/trace_tests.rs | 7 +++++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/benchmarks/string-bench/src/lib.rs b/benchmarks/string-bench/src/lib.rs index 52fd05f7943..957375e93d0 100644 --- a/benchmarks/string-bench/src/lib.rs +++ b/benchmarks/string-bench/src/lib.rs @@ -45,7 +45,7 @@ use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::IdempotentPath; use vortex_bench::datasets::Dataset; @@ -299,9 +299,13 @@ async fn read_parquet_projected(path: PathBuf, column: &str) -> Result let chunks: Vec = reader .map(|batch| { - batch - .map_err(anyhow::Error::from) - .and_then(|rb| ArrayRef::from_arrow(rb, false).map_err(anyhow::Error::from)) + batch.map_err(anyhow::Error::from).and_then(|rb| { + let schema = rb.schema(); + SESSION + .arrow() + .from_arrow_record_batch(rb, &schema) + .map_err(anyhow::Error::from) + }) }) .try_collect() .await?; diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index ee91ffab298..d779757bc77 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -49,7 +49,7 @@ use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_array::test_harness::trace::Traced; use vortex_array::test_harness::trace::trace_op; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -122,7 +122,10 @@ fn lineitem() -> VortexResult { .with_batch_size(1 << 12) .next() .expect("at least one batch"); - ArrayRef::from_arrow(&batch, false) + let schema = batch.schema(); + trace_session() + .arrow() + .from_arrow_record_batch(batch, &schema) } fn compressed_lineitem() -> VortexResult { From 8693977dcde20b3f5a5c62de5c742123443ee62b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 10 Aug 2026 17:11:01 +0100 Subject: [PATCH 5/8] fixes Signed-off-by: Robert Kruszewski --- vortex-arrow/src/convert.rs | 68 ++++++++++++++++++---- vortex-arrow/src/datum.rs | 16 +++--- vortex-arrow/src/iter.rs | 111 ++++++++++++++++++++++++++++++++---- 3 files changed, 164 insertions(+), 31 deletions(-) diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 0105b69ef62..90fcdbbc3f1 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -648,21 +648,27 @@ pub(crate) fn map_from_arrow_parts( Ok(MapArray::try_new(map_dtype, entries)?.into_array()) } +/// Conversion of an Arrow map array into a Vortex `Map` array. +pub fn from_arrow_map(array: &ArrowMapArray, nullable: bool) -> VortexResult { + let DataType::Map(_, keys_sorted) = array.data_type() else { + vortex_panic!("Invalid data type for MapArray: {}", array.data_type()); + }; + let entries = from_arrow_struct(array.entries(), false)?; + map_from_arrow_parts( + entries, + array.offsets(), + array.nulls(), + *keys_sorted, + nullable, + ) +} + impl FromArrowArray<&ArrowMapArray> for ArrayRef { fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult { - let DataType::Map(_, keys_sorted) = array.data_type() else { - vortex_panic!("Invalid data type for MapArray: {}", array.data_type()); - }; - let entries = Self::from_arrow(array.entries(), false)?; - map_from_arrow_parts( - entries, - array.offsets(), - array.nulls(), - *keys_sorted, - nullable, - ) + from_arrow_map(array, nullable) } } + /// Conversion of an Arrow null array into a Vortex `Null` array. pub fn from_arrow_null(value: &ArrowNullArray, nullable: bool) -> VortexResult { vortex_ensure!( @@ -748,6 +754,7 @@ pub fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult { from_arrow_fixed_size_list(array.as_fixed_size_list(), nullable) } + DataType::Map(..) => from_arrow_map(array.as_map(), nullable), DataType::Null => from_arrow_null(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { @@ -906,6 +913,8 @@ mod tests { use arrow_array::builder::Int32Builder; use arrow_array::builder::LargeListBuilder; use arrow_array::builder::ListBuilder; + use arrow_array::builder::MapBuilder as ArrowMapBuilder; + use arrow_array::builder::StringBuilder; use arrow_array::builder::StringViewBuilder; use arrow_array::new_null_array; use arrow_array::types::ArrowPrimitiveType; @@ -937,9 +946,12 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; + use vortex_error::VortexResult; use crate::FromArrowArray as _; use crate::IntoVortexArray as _; + use crate::convert::from_arrow_batch; + use crate::convert::from_arrow_dyn; #[rstest] #[case::i8( @@ -1726,6 +1738,40 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap(); } + #[test] + fn dyn_dispatch_imports_map() -> VortexResult<()> { + let mut builder = ArrowMapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("joe"); + builder.values().append_value(1); + builder.append(true)?; + builder.append(false)?; + let arrow = builder.finish(); + + let expected = DType::map( + DType::Utf8(Nullability::NonNullable), + DType::Primitive(PType::I32, Nullability::Nullable), + false, + Nullability::Nullable, + )?; + + // Dispatching on the Arrow data type must cover maps... + assert_eq!(from_arrow_dyn(&arrow, true)?.dtype(), &expected); + + // ...including maps reached by recursing into a container. + let batch = RecordBatch::try_from_iter_with_nullable([( + "maps", + Arc::new(arrow) as Arc, + true, + )])?; + let imported = from_arrow_batch(&batch, false)?; + assert_eq!( + imported.as_::().unmasked_field(0).dtype(), + &expected + ); + + Ok(()) + } + #[test] fn non_nullable_request_rejects_nulls() { // Requesting `nullable = false` on an Arrow array that physically contains nulls is a diff --git a/vortex-arrow/src/datum.rs b/vortex-arrow/src/datum.rs index ef100fe8fc3..ba9d3f27259 100644 --- a/vortex-arrow/src/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -144,20 +144,20 @@ where /// This is useful for compute functions that delegate to Arrow using [Datum], /// which will return a scalar (length 1 Arrow array) if the input array is constant. /// +/// The array is imported through the [`ArrowSession`](crate::ArrowSession) of `ctx`, so nested +/// extension fields reach their registered import plugins. +/// /// # Error /// /// The provided array must have length `len` or `1`. -#[allow(deprecated)] -pub fn from_arrow_columnar( - array: A, +pub fn from_arrow_columnar( + array: &dyn ArrowArray, len: usize, nullable: bool, ctx: &mut ExecutionCtx, -) -> VortexResult -where - ArrayRef: FromArrowArray, -{ - let array = ArrayRef::from_arrow(array, nullable)?; +) -> VortexResult { + let session = ctx.session().clone(); + let array = session.arrow().from_arrow_array_nullable(array, nullable)?; if array.len() == len { return Ok(array); } diff --git a/vortex-arrow/src/iter.rs b/vortex-arrow/src/iter.rs index 30fa7119390..730585ea0e5 100644 --- a/vortex-arrow/src/iter.rs +++ b/vortex-arrow/src/iter.rs @@ -1,26 +1,50 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use arrow_array::RecordBatchReader; use arrow_array::ffi_stream; +use arrow_schema::SchemaRef; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::iter::ArrayIterator; use vortex_error::VortexError; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::convert::from_arrow_batch; -use crate::dtype::from_arrow_schema_naive; +use crate::ArrowSession; -/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`. +/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayIterator`. +/// +/// The stream's schema and batches are imported through the [`ArrowSession`], so Arrow extension +/// types are routed to their registered import plugins instead of being flattened into their +/// storage types. pub struct ArrowArrayStreamAdapter { stream: ffi_stream::ArrowArrayStreamReader, + session: ArrowSession, + schema: SchemaRef, dtype: DType, } impl ArrowArrayStreamAdapter { - pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self { - Self { stream, dtype } + /// Adapt `stream`, importing its schema and each of its batches through `session`. + /// + /// The adapter holds a clone of `session`, which shares the plugin registries with it, so + /// plugins registered after construction are still observed. + /// + /// The schema declared by the stream is the authoritative schema for the import: every batch + /// is converted against it, so extension metadata survives even if an individual batch has + /// lost it. + pub fn try_new( + session: &ArrowSession, + stream: ffi_stream::ArrowArrayStreamReader, + ) -> VortexResult { + let schema = stream.schema(); + let dtype = session.from_arrow_schema(schema.as_ref())?; + Ok(Self { + stream, + session: session.clone(), + schema, + dtype, + }) } } @@ -37,12 +61,75 @@ impl Iterator for ArrowArrayStreamAdapter { let batch = self.stream.next()?; Some(batch.map_err(VortexError::from).and_then(|b| { - debug_assert_eq!( - &self.dtype, - &from_arrow_schema_naive(b.schema().as_ref()) - .vortex_expect("arrow schema to dtype") - ); - from_arrow_batch(&b, false) + self.session + .from_arrow_record_batch(b, self.schema.as_ref()) })) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::FixedSizeBinaryArray; + use arrow_array::RecordBatch; + use arrow_array::RecordBatchIterator; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use arrow_schema::extension::Uuid as ArrowUuid; + use vortex_array::array_session; + use vortex_array::arrays::Struct; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::extension::uuid::Uuid; + use vortex_error::VortexExpect; + + use super::*; + use crate::ArrowSessionExt; + + /// The adapter imports through the [`ArrowSession`], so a UUID column arrives as a Vortex + /// extension array rather than its `FixedSizeList` storage. + #[test] + fn stream_preserves_extension_types() -> VortexResult<()> { + let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); + field.try_with_extension_type(ArrowUuid)?; + let schema = Arc::new(Schema::new(vec![field])); + let ids = FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids)])?; + + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let stream = ffi_stream::ArrowArrayStreamReader::try_new( + ffi_stream::FFI_ArrowArrayStream::new(Box::new(reader)), + )?; + + let vortex_session = array_session(); + let mut adapter = ArrowArrayStreamAdapter::try_new(&vortex_session.arrow(), stream)?; + + let DType::Struct(fields, _) = adapter.dtype().clone() else { + panic!("expected a struct dtype, got {}", adapter.dtype()); + }; + assert!( + fields + .field_by_index(0) + .vortex_expect("one field") + .as_extension() + .is::() + ); + + let array = adapter.next().vortex_expect("one batch")?; + assert_eq!(array.dtype(), adapter.dtype()); + assert!( + array + .as_::() + .unmasked_field(0) + .dtype() + .as_extension() + .is::() + ); + assert!(adapter.next().is_none()); + + Ok(()) + } +} From 123465e3def6710456085dd2a0b6d535f1cc0260 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 10 Aug 2026 18:06:43 +0100 Subject: [PATCH 6/8] fixes Signed-off-by: Robert Kruszewski --- .../src/fixtures/arrays/datasets/clickbench.rs | 3 +-- .../compat-gen/src/fixtures/arrays/datasets/mod.rs | 5 ++++- .../src/fixtures/arrays/datasets/tpch.rs | 14 ++++++++------ vortex-test/compat-gen/src/fixtures/mod.rs | 8 +++++--- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs index eb619909ff4..ba24c20dc82 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs @@ -179,7 +179,7 @@ impl DatasetFixture for ClickBenchHits5kFixture { "5000 rows (5x1000 from random offsets) of ClickBench hits dataset with wide schema of primitives and strings" } - fn build(&self) -> VortexResult { + fn build(&self, arrow: &ArrowSession) -> VortexResult { let path = cached_clickbench_parquet()?; let file_bytes = fs::read(&path) .map_err(|e| vortex_err!("failed to read cached parquet at {}: {e}", path.display()))?; @@ -195,7 +195,6 @@ impl DatasetFixture for ClickBenchHits5kFixture { .collect::, _>>() .map_err(|e| vortex_err!("failed to read parquet batches: {e}"))?; - let arrow = ArrowSession::default(); Ok(ChunkedArray::from_iter( batches .into_iter() diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs index 929e416fe14..1f84d0edeea 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs @@ -18,6 +18,8 @@ pub fn fixtures() -> Vec> { mod tests { use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::WriteStrategyBuilder; + use vortex_array::array_session; + use vortex_arrow::ArrowSessionExt; use super::fixtures; use crate::adapter; @@ -28,11 +30,12 @@ mod tests { #[test] fn roundtrip_non_clickbench_fixtures_to_bytes() { + let session = array_session(); for dataset in fixtures() .into_iter() .filter(|fixture| !is_clickbench_fixture(fixture.name())) { - let array = dataset.build().unwrap(); + let array = dataset.build(&session.arrow()).unwrap(); let regular_bytes = adapter::write_compressed_to_bytes( array.clone(), WriteStrategyBuilder::default().build(), diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs index d67f39674b6..f21fc4ce49d 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs @@ -15,9 +15,11 @@ use crate::fixtures::DatasetFixture; const SCALE_FACTOR: f64 = 0.01; -fn collect_batches_as_vortex(iter: impl RecordBatchIterator) -> VortexResult { +fn collect_batches_as_vortex( + iter: impl RecordBatchIterator, + arrow: &ArrowSession, +) -> VortexResult { let batches: Vec = iter.collect(); - let arrow = ArrowSession::default(); Ok(ChunkedArray::from_iter( batches .into_iter() @@ -41,10 +43,10 @@ impl DatasetFixture for TpchLineitemFixture { "TPC-H lineitem table at scale factor 0.01 with decimals, dates, and strings" } - fn build(&self) -> VortexResult { + fn build(&self, arrow: &ArrowSession) -> VortexResult { let generator = LineItemGenerator::new(SCALE_FACTOR, 1, 1); let arrow_iter = tpchgen_arrow::LineItemArrow::new(generator).with_batch_size(65_536); - collect_batches_as_vortex(arrow_iter) + collect_batches_as_vortex(arrow_iter, arrow) } } @@ -59,10 +61,10 @@ impl DatasetFixture for TpchOrdersFixture { "TPC-H orders table at scale factor 0.01 with decimals, dates, and strings" } - fn build(&self) -> VortexResult { + fn build(&self, arrow: &ArrowSession) -> VortexResult { let generator = OrderGenerator::new(SCALE_FACTOR, 1, 1); let arrow_iter = tpchgen_arrow::OrderArrow::new(generator).with_batch_size(65_536); - collect_batches_as_vortex(arrow_iter) + collect_batches_as_vortex(arrow_iter, arrow) } } diff --git a/vortex-test/compat-gen/src/fixtures/mod.rs b/vortex-test/compat-gen/src/fixtures/mod.rs index 829acad7dd4..291d9a5f5ff 100644 --- a/vortex-test/compat-gen/src/fixtures/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/mod.rs @@ -11,6 +11,8 @@ use vortex::array::ArrayRef; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::WriteStrategyBuilder; use vortex_array::ExecutionCtx; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -64,7 +66,7 @@ pub trait DatasetFixture { fn description(&self) -> &str; /// Build the dataset as a chunked array. Must be deterministic. - fn build(&self) -> VortexResult; + fn build(&self, arrow: &ArrowSession) -> VortexResult; } // --------------------------------------------------------------------------- @@ -132,8 +134,8 @@ impl Fixture for DatasetFixtureAdapter { self.inner.description() } - fn write(&self, dir: &Path, _ctx: &mut ExecutionCtx) -> VortexResult> { - let array = self.inner.build()?; + fn write(&self, dir: &Path, ctx: &mut ExecutionCtx) -> VortexResult> { + let array = self.inner.build(&ctx.session().arrow())?; let path = dir.join(self.name()); if self.compact { let strategy = WriteStrategyBuilder::default() From 662532e3298442b756b630e18283afdf6a1fe2a7 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 11 Aug 2026 15:50:23 +0100 Subject: [PATCH 7/8] less Signed-off-by: Robert Kruszewski --- encodings/parquet-variant/src/array.rs | 6 +- encodings/parquet-variant/src/kernel.rs | 6 +- vortex-arrow/src/datum.rs | 4 +- vortex-arrow/src/lib.rs | 10 +- vortex-arrow/src/run_end_import.rs | 3 +- vortex-arrow/src/session.rs | 174 +++++++++++++----- vortex-arrow/tests/canonical.rs | 20 +- vortex-json/src/arrow.rs | 2 +- vortex-layout/src/scan/arrow.rs | 2 +- vortex-python/src/arrays/from_arrow.rs | 4 +- vortex-spatial/src/extension/linestring.rs | 2 +- vortex-spatial/src/extension/mod.rs | 4 +- .../src/extension/multilinestring.rs | 2 +- vortex-spatial/src/extension/multipoint.rs | 2 +- vortex-spatial/src/extension/multipolygon.rs | 2 +- vortex-spatial/src/extension/point.rs | 2 +- vortex-spatial/src/extension/polygon.rs | 2 +- vortex-spatial/src/extension/rect.rs | 2 +- vortex-spatial/src/extension/wkb.rs | 2 +- vortex-tensor/src/types/vector/arrow.rs | 3 +- 20 files changed, 168 insertions(+), 86 deletions(-) diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 5f8d7a350dc..533819ec852 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -137,16 +137,16 @@ impl ParquetVariant { Validity::NonNullable }); let metadata = session - .from_arrow_array_nullable(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + .from_arrow_array(ArrowArrayRef::clone(arrow_variant.metadata_field()), false)?; let value = arrow_variant .value_field() - .map(|v| session.from_arrow_array_nullable(v as &dyn ArrowArray, value_nullable)) + .map(|v| session.from_arrow_array(ArrowArrayRef::clone(v), value_nullable)) .transpose()?; let typed_value = arrow_variant .typed_value_field() - .map(|tv| session.from_arrow_array_nullable(tv.as_ref(), typed_value_nullable)) + .map(|tv| session.from_arrow_array(ArrowArrayRef::clone(tv), typed_value_nullable)) .transpose()?; ParquetVariant::try_new(validity, metadata, value, typed_value).map(IntoArray::into_array) } diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index be6893a719b..461d400333c 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -841,13 +841,13 @@ mod tests { let metadata = SESSION .arrow() - .from_arrow_array_nullable(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + .from_arrow_array(ArrowArrayRef::clone(arrow_variant.metadata_field()), false)?; let value = arrow_variant .value_field() .map(|value| { SESSION .arrow() - .from_arrow_array_nullable(value as &dyn ArrowArray, value_nullable) + .from_arrow_array(ArrowArrayRef::clone(value), value_nullable) }) .transpose()?; let typed_value = arrow_variant @@ -855,7 +855,7 @@ mod tests { .map(|typed_value| { SESSION .arrow() - .from_arrow_array_nullable(typed_value.as_ref(), typed_value_nullable) + .from_arrow_array(ArrowArrayRef::clone(typed_value), typed_value_nullable) }) .transpose()?; diff --git a/vortex-arrow/src/datum.rs b/vortex-arrow/src/datum.rs index ba9d3f27259..953a9d8d714 100644 --- a/vortex-arrow/src/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -151,13 +151,13 @@ where /// /// The provided array must have length `len` or `1`. pub fn from_arrow_columnar( - array: &dyn ArrowArray, + array: ArrowArrayRef, len: usize, nullable: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { let session = ctx.session().clone(); - let array = session.arrow().from_arrow_array_nullable(array, nullable)?; + let array = session.arrow().from_arrow_array(array, nullable)?; if array.len() == len { return Ok(array); } diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 81310e8aedf..c8536e10edf 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -7,8 +7,8 @@ //! the [`ArrowSession`]: importing Arrow schemas, fields, and data types into Vortex //! ([`ArrowSession::from_arrow_schema`], [`ArrowSession::from_arrow_field`], //! [`ArrowSession::from_arrow_datatype`]), importing Arrow arrays and record batches -//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_array_nullable`], -//! [`ArrowSession::from_arrow_record_batch`]), exporting Vortex dtypes to Arrow +//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_record_batch`]), +//! exporting Vortex dtypes to Arrow //! ([`ArrowSession::to_arrow_schema`], [`ArrowSession::to_arrow_field`], //! [`ArrowSession::to_arrow_datatype`]), and executing Vortex arrays into Arrow //! ([`ArrowSession::execute_arrow`] and the [`ArrowArrayExecutor`] convenience trait). @@ -62,9 +62,7 @@ pub fn initialize(session: &VortexSession) { /// /// Implementations reuse the underlying Arrow buffers without copying wherever the Arrow and /// Vortex memory layouts allow it. -#[deprecated( - note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_array_nullable`, `from_arrow_record_batch`) instead" -)] +#[deprecated(note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_record_batch`) instead")] pub trait FromArrowArray { /// Convert `array` into a Vortex array whose [`DType`](vortex_array::dtype::DType) has the requested /// `nullable` [`Nullability`](vortex_array::dtype::Nullability). @@ -84,7 +82,7 @@ pub trait FromArrowArray { /// (including an Arrow `NullArray`, which is entirely null), or if the Arrow data type is not /// supported. #[deprecated( - note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_array_nullable`, `from_arrow_record_batch`) instead" + note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_record_batch`) instead" )] fn from_arrow(array: A, nullable: bool) -> VortexResult where diff --git a/vortex-arrow/src/run_end_import.rs b/vortex-arrow/src/run_end_import.rs index e29fc0a5e0f..6a301298907 100644 --- a/vortex-arrow/src/run_end_import.rs +++ b/vortex-arrow/src/run_end_import.rs @@ -112,10 +112,9 @@ mod tests { where R::Native: NativePType, { - let field = Field::new("", array.data_type().clone(), nullable); SESSION .arrow() - .from_arrow_array(Arc::new(array.clone()), &field) + .from_arrow_array(Arc::new(array.clone()), nullable) } #[test] diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index a8f282e44a3..d3a6a532233 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -21,6 +21,7 @@ //! the next. use std::any::Any; +use std::borrow::Cow; use std::fmt::Debug; use std::sync::Arc; @@ -104,6 +105,47 @@ pub enum ArrowImport { Imported(ArrayRef), } +/// The Arrow type description of an array being imported by [`ArrowSession::from_arrow_array`]. +/// +/// Callers holding an Arrow [`Field`] (or [`FieldRef`]) should pass it: its +/// `ARROW:extension:name` metadata is what dispatches the array to a registered +/// [`ArrowImportVTable`]. +/// +/// An Arrow array can carry a validity (null) buffer regardless of whether its schema declares +/// the field nullable, so when no [`Field`] is in hand the caller passes the desired nullability +/// instead, as a [`bool`] or a [`Nullability`]. An anonymous field is then synthesized from the +/// array's own data type, which means no extension plugin is dispatched for the array itself; +/// fields nested inside container data types still carry their metadata and are routed through +/// their importers. +pub trait IntoArrowField<'a> { + /// Resolve to the Arrow [`Field`] describing an array of `data_type`. + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field>; +} + +impl<'a> IntoArrowField<'a> for &'a Field { + fn into_arrow_field(self, _data_type: &DataType) -> Cow<'a, Field> { + Cow::Borrowed(self) + } +} + +impl<'a> IntoArrowField<'a> for &'a FieldRef { + fn into_arrow_field(self, _data_type: &DataType) -> Cow<'a, Field> { + Cow::Borrowed(self.as_ref()) + } +} + +impl<'a> IntoArrowField<'a> for bool { + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field> { + Cow::Owned(Field::new("", data_type.clone(), self)) + } +} + +impl<'a> IntoArrowField<'a> for Nullability { + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field> { + self.is_nullable().into_arrow_field(data_type) + } +} + /// Plugin layer for exporting a Vortex array to an Arrow extension type. /// /// This is purely an implementation trait, its methods should not be called directly. Instead, @@ -167,7 +209,8 @@ pub trait ArrowImportVTable: 'static + Send + Sync + Debug { /// handle the input. /// /// `session` is provided so plugins can convert storage or nested arrays through the - /// session (e.g. [`ArrowSession::from_arrow_array_nullable`]) instead of the deprecated + /// session (e.g. [`ArrowSession::from_arrow_array`], passing a nullability rather than a + /// [`Field`] so the plugin is not dispatched again) instead of the deprecated /// `FromArrowArray` trait. #[allow(clippy::wrong_self_convention)] fn from_arrow_array( @@ -496,7 +539,7 @@ impl ArrowSession { ); let mut columns = Vec::with_capacity(schema.fields().len()); for (col, field) in batch.columns().iter().zip(schema.fields().iter()) { - columns.push(self.from_arrow_array(ArrowArrayRef::clone(col), field)?); + columns.push(self.from_arrow_array_inner(ArrowArrayRef::clone(col), field)?); } Ok(StructArray::try_new(names, columns, length, Validity::NonNullable)?.into_array()) } @@ -568,6 +611,11 @@ impl ArrowSession { /// Decode an Arrow array into a Vortex array. /// + /// `field` describes the Arrow type the array is imported from: pass an Arrow [`Field`] when + /// one is available, or just the desired nullability (`true` / `false` / + /// [`Nullability`]) to synthesize an anonymous field from the array's own data type. See + /// [`IntoArrowField`] for the trade-off between the two. + /// /// Routes through the registered import plugin if `field` carries an Arrow extension /// name we recognize, probing each plugin in registration order until one handles the /// input or all return [`ArrowImport::Unsupported`]. Otherwise recurses into container @@ -575,7 +623,29 @@ impl ArrowSession { /// [`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. - pub fn from_arrow_array(&self, array: ArrowArrayRef, field: &Field) -> VortexResult { + /// + /// # Errors + /// + /// Returns an error if the field (or requested nullability) is non-nullable but the array + /// physically contains nulls, or if the Arrow data type is unsupported. + pub fn from_arrow_array<'a>( + &self, + array: ArrowArrayRef, + field: impl IntoArrowField<'a>, + ) -> VortexResult { + let field = field.into_arrow_field(array.data_type()); + self.from_arrow_array_inner(array, field.as_ref()) + } + + /// [`Self::from_arrow_array`] with the Arrow [`Field`] already resolved: probe the import + /// plugins registered for the field's extension name, then fall back to the canonical + /// conversion. Also the recursion point for nested fields, which already have a [`Field`]. + #[allow(clippy::wrong_self_convention)] + fn from_arrow_array_inner( + &self, + array: ArrowArrayRef, + field: &Field, + ) -> VortexResult { if let Some(extension_name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let importers = self.importers(&Id::new(extension_name)); @@ -594,25 +664,6 @@ impl ArrowSession { self.from_arrow_array_canonical(array.as_ref(), field) } - /// Decode an Arrow array into a Vortex array whose dtype has the requested `nullable`ness. - /// - /// An Arrow array can carry a validity (null) buffer regardless of whether its schema - /// declares the field nullable, so the desired nullability is supplied by the caller. - /// Returns an error if `nullable` is `false` but the array physically contains nulls. - /// - /// No top-level Arrow [`Field`] is available in this form, so no extension plugin is - /// dispatched for the array itself; fields nested inside container data types still carry - /// their metadata and are routed through [`Self::from_arrow_array`]. Prefer the - /// field-aware [`Self::from_arrow_array`] when a [`Field`] is in hand. - pub fn from_arrow_array_nullable( - &self, - array: &dyn ArrowArray, - nullable: bool, - ) -> VortexResult { - let field = Field::new("", array.data_type().clone(), nullable); - self.from_arrow_array_canonical(array, &field) - } - /// Recurse into Arrow container arrays so nested fields with extension metadata reach /// their importers, falling through to the canonical conversion for leaf types. #[allow(clippy::wrong_self_convention)] @@ -641,7 +692,7 @@ impl ArrowSession { } else { ArrowArrayRef::clone(col) }; - self.from_arrow_array(inner, child_field.as_ref()) + self.from_arrow_array_inner(inner, child_field.as_ref()) }) .collect::>>()?; let validity = nulls(arrow_struct.nulls(), field.is_nullable())?; @@ -668,8 +719,10 @@ impl ArrowSession { } DataType::FixedSizeList(elem_field, list_size) => { let fsl = array.as_fixed_size_list(); - let elements = - self.from_arrow_array(ArrowArrayRef::clone(fsl.values()), elem_field.as_ref())?; + let elements = self.from_arrow_array_inner( + ArrowArrayRef::clone(fsl.values()), + elem_field.as_ref(), + )?; let validity = nulls(fsl.nulls(), field.is_nullable())?; Ok( FixedSizeListArray::try_new(elements, *list_size as u32, validity, fsl.len())? @@ -697,7 +750,7 @@ impl ArrowSession { DataType::Map(entries_field, keys_sorted) => { let map = array.as_map(); let entries_array: ArrowArrayRef = Arc::new(map.entries().clone()); - let entries = self.from_arrow_array(entries_array, entries_field.as_ref())?; + let entries = self.from_arrow_array_inner(entries_array, entries_field.as_ref())?; map_from_arrow_parts( entries, map.offsets(), @@ -717,11 +770,15 @@ impl ArrowSession { ), } } - DataType::Dictionary(_, values_type) => { + DataType::Dictionary(..) => { let dict = array.as_any_dictionary(); - let values_field = dictionary_values_field(values_type, field.is_nullable()); - let values = - self.from_arrow_array(ArrowArrayRef::clone(dict.values()), &values_field)?; + // Arrow models dictionary values as a bare `DataType`, so there is no field + // metadata to carry an extension name for the values themselves. Fields *nested + // inside* that data type (list elements, struct fields, map entries) do keep + // their metadata, so importing the values by nullability alone still routes them + // back through the plugin-aware conversion. + let values = self + .from_arrow_array(ArrowArrayRef::clone(dict.values()), field.is_nullable())?; let codes = dict.keys(); let codes = from_arrow_dyn(codes, codes.is_nullable())?; // SAFETY: arrow-rs enforces the dictionary invariants on construction, so the @@ -748,7 +805,7 @@ impl ArrowSession { .downcast_ref::>() .ok_or_else(|| vortex_err!("expected an Arrow RunArray, got {}", array.data_type()))?; let values = - self.from_arrow_array(ArrowArrayRef::clone(run_array.values()), values_field)?; + self.from_arrow_array_inner(ArrowArrayRef::clone(run_array.values()), values_field)?; run_end_from_arrow(run_array, values) } } @@ -762,17 +819,6 @@ fn run_end_values_field(values_field: &FieldRef, nullability: Nullability) -> Fi .with_nullable(nullability.into()) } -/// A synthetic field for the values of an Arrow [`DataType::Dictionary`], carrying the dictionary -/// array's own nullability. -/// -/// Arrow models dictionary values as a bare [`DataType`], so there is no field metadata to carry -/// an extension name for the values themselves. Fields *nested inside* that data type (list -/// elements, struct fields, map entries) do keep their metadata, so wrapping the values in a field -/// is enough to route them back through the plugin-aware conversion. -fn dictionary_values_field(values_type: &DataType, nullable: bool) -> Field { - Field::new("", values_type.clone(), nullable) -} - // NOTE(aduffy): We should remove this once we bump Arrow to 0.59.0. This is replicating the // `Field::has_valid_extension_type` method on Arrow added in 58.2.0, we polyfill it here so that // this crate can build with minimal-versions declared. @@ -815,11 +861,13 @@ mod tests { use arrow_array::Int32Array; use arrow_array::ListArray as ArrowListArray; use arrow_array::StringArray; + use arrow_array::StructArray as ArrowStructArray; use arrow_array::cast::AsArray; use arrow_buffer::OffsetBuffer; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::extension::Uuid as ArrowUuid; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -835,6 +883,7 @@ mod tests { use vortex_array::dtype::extension::ExtVTable; use vortex_array::extension::uuid::Uuid; use vortex_array::extension::uuid::UuidMetadata; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use super::*; @@ -1267,6 +1316,47 @@ mod tests { Ok(()) } + /// Importing by nullability instead of by [`Field`] synthesizes an anonymous field from the + /// array's own data type, so extension metadata on *nested* fields still reaches its importer. + /// A [`bool`] and the equivalent [`Nullability`] must agree. + #[rstest] + #[case(true)] + #[case(false)] + fn from_arrow_array_by_nullability(#[case] nullable: bool) -> VortexResult<()> { + let session = ArrowSession::default(); + + let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), false); + uuid_field.try_with_extension_type(ArrowUuid)?; + let uuids: ArrowArrayRef = Arc::new(FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?); + let arrow_struct: ArrowArrayRef = Arc::new(ArrowStructArray::try_new( + Fields::from(vec![uuid_field]), + vec![uuids], + None, + )?); + + let array = session.from_arrow_array(ArrowArrayRef::clone(&arrow_struct), nullable)?; + assert_eq!(array.dtype().nullability(), nullable.into()); + let DType::Struct(fields, _) = array.dtype() else { + panic!("expected a Struct dtype, got {}", array.dtype()); + }; + assert!( + fields + .field_by_index(0) + .vortex_expect("struct dtype has one field") + .as_extension() + .is::(), + "expected the nested Uuid extension to survive, got {}", + array.dtype() + ); + + // The `Nullability` form is equivalent to the `bool` form. + let by_nullability = session.from_arrow_array(arrow_struct, Nullability::from(nullable))?; + assert_eq!(by_nullability.dtype(), array.dtype()); + Ok(()) + } + /// A plain Arrow dictionary imports as a Vortex `Dict` array over the dictionary values, /// matching the dtype the schema conversion reports. #[test] diff --git a/vortex-arrow/tests/canonical.rs b/vortex-arrow/tests/canonical.rs index 18355fb98c8..5b1b70e29ac 100644 --- a/vortex-arrow/tests/canonical.rs +++ b/vortex-arrow/tests/canonical.rs @@ -16,7 +16,6 @@ use arrow_array::PrimitiveArray as ArrowPrimitiveArray; use arrow_array::StringArray; use arrow_array::StringViewArray; use arrow_array::StructArray as ArrowStructArray; -use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use arrow_array::types::Int64Type; use arrow_array::types::UInt64Type; @@ -116,7 +115,7 @@ fn roundtrip_struct() { None, ])); - let arrow_struct = ArrowStructArray::new( + let arrow_struct = Arc::new(ArrowStructArray::new( vec![ Arc::new(Field::new("name", DataType::Utf8View, true)), Arc::new(Field::new("age", DataType::Int32, true)), @@ -124,17 +123,17 @@ fn roundtrip_struct() { .into(), vec![names, ages], nulls.finish(), - ); + )) as ArrowArrayRef; let vortex_struct = SESSION .arrow() - .from_arrow_array_nullable(&arrow_struct, true) + .from_arrow_array(Arc::clone(&arrow_struct), true) .unwrap(); let vortex_struct = SESSION .arrow() .execute_arrow(vortex_struct, None, &mut ctx) .unwrap(); - assert_eq!(&arrow_struct, vortex_struct.as_struct()); + assert_eq!(&arrow_struct, &vortex_struct); } #[test] @@ -146,18 +145,18 @@ fn roundtrip_list() { Some("Mikhail"), ])); - let arrow_list = ArrowListArray::new( + let arrow_list = Arc::new(ArrowListArray::new( Arc::new(Field::new_list_field(DataType::Utf8, true)), OffsetBuffer::from_lengths(vec![0, 2, 1]), names, None, - ); + )) as ArrowArrayRef; let list_data_type = arrow_list.data_type(); let list_field = Field::new(String::new(), list_data_type.clone(), true); let vortex_list = SESSION .arrow() - .from_arrow_array_nullable(&arrow_list, true) + .from_arrow_array(Arc::clone(&arrow_list), true) .unwrap(); let rt_arrow_list = SESSION @@ -165,8 +164,5 @@ fn roundtrip_list() { .execute_arrow(vortex_list, Some(&list_field), &mut ctx) .unwrap(); - assert_eq!( - (Arc::new(arrow_list.clone()) as ArrowArrayRef).as_ref(), - rt_arrow_list.as_ref() - ); + assert_eq!(&arrow_list, &rt_arrow_list); } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 3c4a61f2ba1..474000795e2 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -135,7 +135,7 @@ impl ArrowImportVTable for Json { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::new(ext_dtype.clone(), storage).into_array(), )) diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 05e5c890558..224035a45fe 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -157,7 +157,7 @@ mod tests { // Convert to Vortex SCAN_SESSION .arrow() - .from_arrow_array_nullable(&struct_array, true) + .from_arrow_array(Arc::new(struct_array), true) } fn create_arrow_schema() -> Arc { diff --git a/vortex-python/src/arrays/from_arrow.rs b/vortex-python/src/arrays/from_arrow.rs index 20a6b5c433b..1077d6d3aaa 100644 --- a/vortex-python/src/arrays/from_arrow.rs +++ b/vortex-python/src/arrays/from_arrow.rs @@ -37,7 +37,7 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult> = obj.getattr(intern!(py, "chunks"))?.extract()?; @@ -47,7 +47,7 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult>>()?; diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index ffef079122d..c1d6b374586 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -379,7 +379,7 @@ impl ArrowImportVTable for LineString { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index 3b8f14c17cc..f24f31f02aa 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -208,7 +208,7 @@ pub fn native_geometry_scalar_from_wkb( let to_storage = |target: &GeoArrowType| -> VortexResult { let native = cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; - session.from_arrow_array_nullable(native.to_array_ref().as_ref(), false) + session.from_arrow_array(native.to_array_ref(), false) }; let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { @@ -312,7 +312,7 @@ pub(crate) fn geoarrow_to_wkb( GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&SpatialMetadata::default()))); let wkb = cast(geoarrow_array, &wkb_type) .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; - session.from_arrow_array_nullable(wkb.to_array_ref().as_ref(), false) + session.from_arrow_array(wkb.to_array_ref(), false) } /// Recover [`SpatialMetadata`] from GeoArrow metadata. diff --git a/vortex-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index 84be29e58b3..fb53863ecc7 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -312,7 +312,7 @@ impl ArrowImportVTable for MultiLineString { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index 1caca024301..4afc9052582 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -294,7 +294,7 @@ impl ArrowImportVTable for MultiPoint { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 23c8cb67b2f..a3291ba68ea 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -311,7 +311,7 @@ impl ArrowImportVTable for MultiPolygon { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index 505896d748e..e6a00fe8fea 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -284,7 +284,7 @@ impl ArrowImportVTable for Point { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index 6ff784d11b6..224d1c17912 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -306,7 +306,7 @@ impl ArrowImportVTable for Polygon { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index 90025bc2b32..21e33808cc6 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -328,7 +328,7 @@ impl ArrowImportVTable for Rect { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), )) diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index 14a2bc41fdc..a6dddc70497 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -290,7 +290,7 @@ impl ArrowImportVTable for WellKnownBinary { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array_nullable(array.as_ref(), field.is_nullable())?; + let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::new(ext_dtype.clone(), storage).into_array(), )) diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index b92d0ee6918..8833b38fe03 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -162,8 +162,7 @@ impl ArrowImportVTable for Vector { return Ok(ArrowImport::Unsupported(array)); } - let storage = - session.from_arrow_array_nullable(array.as_ref() as &dyn Array, dtype.is_nullable())?; + let storage = session.from_arrow_array(array, dtype.is_nullable())?; Ok(ArrowImport::Imported( ExtensionArray::try_new(dtype.clone(), storage)?.into_array(), )) From 1368ff3494fab5843713e17fced5b73ca6b18f8c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 10:16:23 +0100 Subject: [PATCH 8/8] less Signed-off-by: Robert Kruszewski --- vortex-ffi/src/array.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index d009a9e4c47..3d4b9758567 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -997,8 +997,6 @@ mod tests { assert!(!vx.is_null()); unsafe { - let session = vx_session_new(); - assert!(vx_array_has_dtype(vx, vx_dtype_variant::DTYPE_STRUCT)); assert_eq!(vx_array_len(vx), 3); assert!(!vx_array_is_nullable(vx));