diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 50e2784cbbe..5f8d7a350dc 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -40,7 +40,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; @@ -91,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; @@ -130,17 +136,17 @@ impl ParquetVariant { } else { Validity::NonNullable }); - let metadata = - ArrayRef::from_arrow(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + let metadata = 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| 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| session.from_arrow_array_nullable(tv.as_ref(), typed_value_nullable)) .transpose()?; ParquetVariant::try_new(validity, metadata, value, typed_value).map(IntoArray::into_array) } @@ -508,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; @@ -526,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"))?; @@ -577,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!( @@ -609,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(), @@ -700,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"))?; @@ -737,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"))?; @@ -809,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 c68ffb435e7..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,6 +263,7 @@ impl ArrowImportVTable for ParquetVariant { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { if !dtype.is_variant() || field @@ -276,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 3665934c58d..0fee80bfc68 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -46,7 +46,6 @@ use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -120,9 +119,14 @@ 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 { - ArrayRef::from_arrow(arrow_output.as_ref(), true)? + ctx.session() + .arrow() + .from_arrow_array_nullable(arrow_output.as_ref(), true)? }; vortex_ensure_eq!( @@ -164,9 +168,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()) } } @@ -320,7 +324,7 @@ mod tests { use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_array::validity::Validity; - use vortex_arrow::FromArrowArray; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -343,7 +347,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 { @@ -360,13 +364,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 { @@ -766,15 +770,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( @@ -785,7 +798,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")); @@ -902,7 +915,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/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp index 43b986f4b42..08fdfa510e6 100644 --- a/lang/cpp/include/vortex/array.hpp +++ b/lang/cpp/include/vortex/array.hpp @@ -135,7 +135,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 b4adf69fe5c..d9b1c2d1988 100644 --- a/lang/cpp/src/array.cpp +++ b/lang/cpp/src/array.cpp @@ -177,9 +177,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-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 6fd2fcc73b9..0105b69ef62 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -1,6 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! 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; use arrow_array::AnyDictionaryArray; @@ -153,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) } } }; @@ -177,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) } } @@ -237,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) } } }; @@ -259,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 @@ -290,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) } } @@ -403,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) } } @@ -554,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) } } @@ -597,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/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 a7438ab6f2c..cee29b489b1 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -222,7 +222,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; @@ -239,7 +238,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] @@ -417,7 +416,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..81310e8aedf 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). @@ -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; @@ -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 76b33557289..8ae24dd34ad 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 5942d6aef76..dcfbe142746 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_geo::extension::GeoMetadata; use vortex_geo::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 5a328a4d19b..fa6b50714b8 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 d0ffb472ebe..ab69a3bf653 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -610,7 +610,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; @@ -618,7 +617,6 @@ mod tests { use vortex::metrics::DefaultMetricsRegistry; use vortex::scan::selection::Selection; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; use super::*; use crate::VortexAccessPlan; @@ -736,7 +734,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-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 9d57159c36f..8e666e92b53 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_geo::scalar_fn::contains::GeoContains; use vortex_geo::scalar_fn::distance::GeoDistance; use vortex_geo::scalar_fn::intersects::GeoIntersects; +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 geo_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_geo_column(ctx.fields, col_ref.name.as_ref()) => diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index fd7851a27c9..e82edd4059e 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -701,12 +701,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); uint8_t vx_array_get_u8(const vx_array *array, size_t index); diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index ae2a2e3d93a..b8e045fb50a 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 paste::paste; use vortex::array::ArrayRef; use vortex::array::Canonical; @@ -34,7 +35,7 @@ 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 vortex_arrow::ArrowSessionExt; use crate::box_wrapper; use crate::dtype::vx_dtype; @@ -398,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, @@ -411,12 +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 = ArrayRef::from_arrow(arrow_array.as_ref(), nullable)?; + let vortex_array = session.arrow().from_arrow_array(arrow_array, &field)?; Ok(vx_array::new(vortex_array)) }) } @@ -669,6 +673,7 @@ mod tests { use crate::expression::vx_expression_free; 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; @@ -1036,9 +1041,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, @@ -1068,6 +1075,7 @@ mod tests { vx_array_free(b); vx_array_free(vx); + vx_session_free(session); } } diff --git a/vortex-geo/src/extension/linestring.rs b/vortex-geo/src/extension/linestring.rs index aed69315251..0dc6327d9b5 100644 --- a/vortex-geo/src/extension/linestring.rs +++ b/vortex-geo/src/extension/linestring.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; @@ -157,7 +156,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(), + ) } } @@ -285,6 +287,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)); @@ -296,7 +299,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-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index b05069f8ca5..0b59c8b0449 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -57,7 +57,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; @@ -208,7 +208,10 @@ pub(crate) fn single_geometry( /// 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(&GeoMetadata::default()); let binary = BinaryArray::from(vec![Some(bytes)]); let wkb = GenericWkbArray::::try_from(( @@ -221,7 +224,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) + session.from_arrow_array_nullable(native.to_array_ref().as_ref(), false) }; let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { @@ -316,11 +319,14 @@ pub(crate) fn geoarrow_metadata(geo_metadata: &GeoMetadata) -> Arc { /// Serialize a native geometry array to WKB (a `WkbView` array) via geoarrow's cast. /// Shared by the `to_wkb` methods on the geometry extension types. -pub(crate) fn geoarrow_to_wkb(geo_array: &dyn GeoArrowArray) -> VortexResult { +pub(crate) fn geoarrow_to_wkb( + geo_array: &dyn GeoArrowArray, + session: &ArrowSession, +) -> VortexResult { let wkb_type = GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); let wkb = cast(geo_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 [`GeoMetadata`] from GeoArrow metadata. @@ -341,6 +347,7 @@ pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata { mod tests { use prost::Message; use vortex_array::dtype::DType; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -349,9 +356,16 @@ mod tests { use super::MultiPoint; use super::Point; use super::Polygon; - use super::native_geometry_scalar_from_wkb; use crate::extension::GeoMetadata; + /// 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 = GeoMetadata { diff --git a/vortex-geo/src/extension/multilinestring.rs b/vortex-geo/src/extension/multilinestring.rs index 736ef210e3a..230fdc1e89b 100644 --- a/vortex-geo/src/extension/multilinestring.rs +++ b/vortex-geo/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; @@ -164,7 +163,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(), + ) } } @@ -295,6 +297,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)); @@ -306,7 +309,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-geo/src/extension/multipoint.rs b/vortex-geo/src/extension/multipoint.rs index 6ffb9d6f144..95cb43bba16 100644 --- a/vortex-geo/src/extension/multipoint.rs +++ b/vortex-geo/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-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs index 91b79799640..e2c36e1975d 100644 --- a/vortex-geo/src/extension/multipolygon.rs +++ b/vortex-geo/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-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 1f4324af237..6fc7210deb1 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/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-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index 0482096c905..5a867f08824 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/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; @@ -157,7 +156,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(), + ) } } @@ -289,6 +291,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)); @@ -300,7 +303,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-geo/src/extension/rect.rs b/vortex-geo/src/extension/rect.rs index d6172fb9371..5e7a326d1e6 100644 --- a/vortex-geo/src/extension/rect.rs +++ b/vortex-geo/src/extension/rect.rs @@ -42,7 +42,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; @@ -297,6 +296,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)); @@ -308,7 +308,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-geo/src/extension/wkb.rs b/vortex-geo/src/extension/wkb.rs index 027f3746aa8..8c7dbd30656 100644 --- a/vortex-geo/src/extension/wkb.rs +++ b/vortex-geo/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-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 855a6af1867..2da87f830cd 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -135,6 +135,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; @@ -158,7 +159,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-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 3f3842e4e1c..919485d6f0e 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -134,6 +134,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; @@ -179,7 +180,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-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 b010ae3ca66..f47b7956572 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; @@ -462,7 +461,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-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 803e20e29dc..795dc401ec5 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -338,7 +338,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; @@ -371,7 +370,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(); @@ -387,7 +385,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();