diff --git a/Cargo.lock b/Cargo.lock index db3052e5132..0a35c84a451 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7624,7 +7624,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009a43e88f0f7b3b3c2c9de5384a0ac2cef8f33996b16f9e6738e24661618775" dependencies = [ "async-trait", - "asyncband 0.7.2", + "asyncband 0.7.3", "bytes", "chrono", "futures", @@ -7726,7 +7726,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ddff5c648dcb3922c2bec71e0d26e881ddf9c367cb7c9a59ed7f2d39a4f38218" dependencies = [ "anyhow", - "asyncband 0.7.2", + "asyncband 0.7.3", "base64 0.23.1", "bytes", "futures", @@ -7765,7 +7765,7 @@ version = "0.59.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b454e68e5f6b331bd079a2ee02ebe05078fc0ddaae6822767ef54ab3c2ac0aa8" dependencies = [ - "asyncband 0.7.2", + "asyncband 0.7.3", "futures", "http", "opendal-core 0.59.3", diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 5d364cae0ef..4155573536b 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -35,12 +35,9 @@ use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; -#[expect( - deprecated, - reason = "TODO(aduffy): figure out what to do with Parquet Variant" -)] -use vortex_arrow::ArrowArrayExecutor; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; use vortex_arrow::to_arrow_null_buffer; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; @@ -436,11 +433,14 @@ pub trait ParquetVariantArrayExt: } /// Converts this storage array to Arrow's canonical Parquet Variant extension storage. - #[expect( - deprecated, - reason = "TODO(aduffy): figure out what to do with Parquet Variant" - )] - fn to_arrow(&self, ctx: &mut ExecutionCtx) -> VortexResult { + fn to_arrow( + &self, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let session = ctx.session().clone(); + let arrow = session.arrow(); + let exporter = arrow.exporter(options); let metadata = self.metadata(); let len = metadata.len(); let nulls = to_arrow_null_buffer(self.parquet_variant_validity(), len, ctx)?; @@ -448,7 +448,7 @@ pub trait ParquetVariantArrayExt: let mut fields = Vec::with_capacity(3); let mut arrays: Vec = Vec::with_capacity(3); - let metadata_arrow = metadata.clone().execute_arrow(None, ctx)?; + let metadata_arrow = exporter.execute_arrow(metadata.clone(), None, ctx)?; fields.push(Arc::new(Field::new( "metadata", metadata_arrow.data_type().clone(), @@ -457,7 +457,7 @@ pub trait ParquetVariantArrayExt: arrays.push(metadata_arrow); if let Some(value) = self.value() { - let value_arrow = value.clone().execute_arrow(None, ctx)?; + let value_arrow = exporter.execute_arrow(value.clone(), None, ctx)?; fields.push(Arc::new(Field::new( "value", value_arrow.data_type().clone(), @@ -467,7 +467,7 @@ pub trait ParquetVariantArrayExt: } if let Some(typed_value) = self.typed_value() { - let tv_arrow = typed_value.clone().execute_arrow(None, ctx)?; + let tv_arrow = exporter.execute_arrow(typed_value.clone(), None, ctx)?; fields.push(Arc::new(Field::new( "typed_value", tv_arrow.data_type().clone(), @@ -514,6 +514,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; + use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -539,7 +540,7 @@ mod tests { .ok_or_else(|| vortex_err!("expected parquet variant child"))?; let mut ctx = SESSION.create_execution_ctx(); - let roundtripped = inner.to_arrow(&mut ctx)?; + let roundtripped = inner.to_arrow(&ArrowExportOptions::default(), &mut ctx)?; let roundtripped = roundtripped.inner(); assert_eq!(struct_array.len(), roundtripped.len()); @@ -647,7 +648,7 @@ mod tests { let pv_array = ParquetVariant::try_new(Validity::NonNullable, metadata, Some(value), None)?; let mut ctx = SESSION.create_execution_ctx(); - let variant_arr = pv_array.to_arrow(&mut ctx)?; + let variant_arr = pv_array.to_arrow(&ArrowExportOptions::default(), &mut ctx)?; let struct_arr = variant_arr.inner(); assert_eq!(struct_arr.num_columns(), 2); @@ -669,7 +670,7 @@ mod tests { )?; let mut ctx = SESSION.create_execution_ctx(); - let variant_arr = pv_array.to_arrow(&mut ctx)?; + let variant_arr = pv_array.to_arrow(&ArrowExportOptions::default(), &mut ctx)?; let struct_arr = variant_arr.inner(); assert_eq!(struct_arr.num_columns(), 3); @@ -756,7 +757,7 @@ mod tests { assert!(parquet_array.typed_value().is_some()); let mut ctx = SESSION.create_execution_ctx(); - let roundtripped = parquet_array.to_arrow(&mut ctx)?; + let roundtripped = parquet_array.to_arrow(&ArrowExportOptions::default(), &mut ctx)?; let roundtripped = roundtripped.inner(); assert_eq!( roundtripped.column_names(), diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index e1c03d704dc..24835d6357a 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::Variant; use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::dtype::DType; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -63,8 +64,12 @@ fn parquet_variant_storage_request(fields: &Fields) -> Option<(bool, bool)> { pub(crate) fn export_storage_to_target( parquet_array: &T, target_fields: &Fields, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { + let session = ctx.session().clone(); + let arrow = session.arrow(); + let exporter = arrow.exporter(options); let mut arrays = Vec::with_capacity(target_fields.len()); for field in target_fields { @@ -82,11 +87,7 @@ pub(crate) fn export_storage_to_target( ); }; - arrays.push(ctx.session().clone().arrow().execute_arrow( - child, - Some(field.as_ref()), - ctx, - )?); + arrays.push(exporter.execute_arrow(child, Some(field.as_ref()), ctx)?); } let nulls = to_arrow_null_buffer( @@ -104,9 +105,10 @@ pub(crate) fn export_storage_to_target( pub(crate) fn export_unshredded_storage_to_target( parquet_array: &T, target_fields: &Fields, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let arrow_variant = parquet_array.to_arrow(ctx)?; + let arrow_variant = parquet_array.to_arrow(options, ctx)?; let unshredded = unshred_variant(&arrow_variant)?; let unshredded_array = if parquet_array.as_ref().dtype().is_nullable() { ParquetVariant::from_arrow_variant_nullable(&unshredded, &ctx.session().arrow())? @@ -114,7 +116,7 @@ pub(crate) fn export_unshredded_storage_to_target( ParquetVariant::from_arrow_variant(&unshredded, &ctx.session().arrow())? }; let unshredded_parquet = unshredded_array.as_::(); - export_storage_to_target(&unshredded_parquet, target_fields, ctx) + export_storage_to_target(&unshredded_parquet, target_fields, options, ctx) } pub(crate) fn parquet_variant_for_export( @@ -178,6 +180,7 @@ impl ArrowExportVTable for ParquetVariant { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { if target @@ -203,6 +206,7 @@ impl ArrowExportVTable for ParquetVariant { return Ok(ArrowExport::Exported(export_unshredded_storage_to_target( &parquet_array, fields, + options, ctx, )?)); } @@ -220,11 +224,13 @@ impl ArrowExportVTable for ParquetVariant { return Ok(ArrowExport::Exported(export_storage_to_target( &parquet_array, fields, + options, ctx, )?)); } - let arrow_variant = Arc::new(parquet_array.to_arrow(ctx)?.into_inner()) as ArrowArrayRef; + let arrow_variant = + Arc::new(parquet_array.to_arrow(options, ctx)?.into_inner()) as ArrowArrayRef; if arrow_variant.data_type() == target.data_type() { Ok(ArrowExport::Exported(arrow_variant)) diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 7b12a1988ef..f7ccad5c505 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -43,6 +43,7 @@ use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::variant_get::VariantGet; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; @@ -106,7 +107,7 @@ impl ExecuteParentKernel for VariantGetKernel { return Ok(None); } - let arrow_variant = array.to_arrow(ctx)?; + let arrow_variant = array.to_arrow(&ArrowExportOptions::default(), ctx)?; let arrow_input: ArrowArrayRef = Arc::new(arrow_variant.into_inner()); let session = ctx.session().clone(); let as_type = to_arrow_as_type(parent.options.dtype(), &session.arrow())?; diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index eab80bf63e1..434eded87c2 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -396,6 +396,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; + use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -548,7 +549,7 @@ mod tests { let inner_pv = vortex_arr.as_opt::().unwrap(); let mut ctx = array_session().create_execution_ctx(); - let roundtripped = inner_pv.to_arrow(&mut ctx)?; + let roundtripped = inner_pv.to_arrow(&ArrowExportOptions::default(), &mut ctx)?; assert_eq!(roundtripped.inner().null_count(), 2); Ok(()) diff --git a/encodings/uuid/src/arrow.rs b/encodings/uuid/src/arrow.rs index ba405438dbf..fba82141b8a 100644 --- a/encodings/uuid/src/arrow.rs +++ b/encodings/uuid/src/arrow.rs @@ -33,6 +33,7 @@ use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -86,6 +87,7 @@ impl ArrowExportVTable for Uuid { &self, array: ArrayRef, _target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_uuid = array @@ -96,7 +98,7 @@ impl ArrowExportVTable for Uuid { if !is_uuid { return Ok(ArrowExport::Unsupported(array)); } - Ok(ArrowExport::Exported(try_fsl_to_fsb(array, ctx)?)) + Ok(ArrowExport::Exported(try_fsl_to_fsb(array, options, ctx)?)) } } @@ -165,7 +167,11 @@ impl ArrowImportVTable for Uuid { /// Reinterpret a Vortex UUID extension array's `FixedSizeList` storage as an Arrow /// `FixedSizeBinary[16]` array, sharing the underlying byte buffer. -fn try_fsl_to_fsb(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { +fn try_fsl_to_fsb( + array: ArrayRef, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, +) -> VortexResult { let executed = array.execute::(ctx)?; let storage = executed.storage_array().clone(); let storage_arrow_type = DataType::FixedSizeList( @@ -180,9 +186,11 @@ fn try_fsl_to_fsb(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult( pub(super) fn to_arrow_byte_view( array: ArrayRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { // First we cast the array into the desired ByteView type. @@ -62,7 +65,11 @@ pub(super) fn to_arrow_byte_view( let array = array.execute::(ctx)?; let varbinview = array.execute::(ctx)?; - execute_varbinview_to_arrow::(&varbinview, ctx) + if exporter.options().get_or_default::().0 { + execute_varbinview_to_arrow::(&varbinview, ctx) + } else { + canonical_varbinview_to_arrow::(&varbinview, ctx) + } } #[cfg(test)] diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 1a17d9de1dd..8f16ec28528 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -23,7 +23,8 @@ use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use crate::ArrowArrayExecutor; +use crate::ArrowExporter; +use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_dictionary`] requires for export. struct ArrowDictExportable; @@ -40,22 +41,23 @@ pub(super) fn to_arrow_dictionary( array: ArrayRef, codes_type: &DataType, values_type: &DataType, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let array = array.execute_until::(ctx)?; let array = match array.try_downcast::() { - Ok(dict) => return dict_to_dict(dict, codes_type, values_type, ctx), + Ok(dict) => return dict_to_dict(dict, codes_type, values_type, exporter, ctx), Err(array) => array, }; let array = match array.try_downcast::() { - Ok(constant) => return constant_to_dict(constant, codes_type, values_type, ctx), + Ok(constant) => return constant_to_dict(constant, codes_type, values_type, exporter, ctx), Err(array) => array, }; // Otherwise, we should try and build a dictionary. // Arrow hides this functionality inside the cast module! - let array = array.execute_arrow(Some(values_type), ctx)?; + let array = execute_arrow_naive(array, Some(values_type), exporter, ctx)?; arrow_cast::cast( &array, &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), @@ -68,6 +70,7 @@ fn constant_to_dict( array: ConstantArray, codes_type: &DataType, values_type: &DataType, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -78,9 +81,12 @@ fn constant_to_dict( return Ok(new_null_array(&dict_type, len)); } - let values = ConstantArray::new(scalar.clone(), 1) - .into_array() - .execute_arrow(Some(values_type), ctx)?; + let values = execute_arrow_naive( + ConstantArray::new(scalar.clone(), 1).into_array(), + Some(values_type), + exporter, + ctx, + )?; let codes = zeroed_codes_array(codes_type, len)?; make_dict_array(codes_type, codes, values) } @@ -90,13 +96,11 @@ fn dict_to_dict( array: DictArray, codes_type: &DataType, values_type: &DataType, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let codes = array.codes().clone().execute_arrow(Some(codes_type), ctx)?; - let values = array - .values() - .clone() - .execute_arrow(Some(values_type), ctx)?; + let codes = execute_arrow_naive(array.codes().clone(), Some(codes_type), exporter, ctx)?; + let values = execute_arrow_naive(array.values().clone(), Some(values_type), exporter, ctx)?; make_dict_array(codes_type, codes, values) } diff --git a/vortex-arrow/src/executor/fixed_size_list.rs b/vortex-arrow/src/executor/fixed_size_list.rs index 7220bbe13f7..b830af710dd 100644 --- a/vortex-arrow/src/executor/fixed_size_list.rs +++ b/vortex-arrow/src/executor/fixed_size_list.rs @@ -14,13 +14,14 @@ use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use crate::ArrowExporter; use crate::executor::validity::to_arrow_null_buffer; -use crate::session::ArrowSessionExt; pub(super) fn to_arrow_fixed_list( array: ArrayRef, list_size: i32, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -31,18 +32,25 @@ pub(super) fn to_arrow_fixed_list( // Check for Vortex FixedSizeListArray and convert directly. if let Some(array) = array.as_opt::() { - return list_to_list(&array.into_owned(), elements_field, list_size, ctx); + return list_to_list( + &array.into_owned(), + elements_field, + list_size, + exporter, + ctx, + ); } // Otherwise, we execute the array to become a FixedSizeListArray. let fixed_size_list = array.execute::(ctx)?; - list_to_list(&fixed_size_list, elements_field, list_size, ctx) + list_to_list(&fixed_size_list, elements_field, list_size, exporter, ctx) } fn list_to_list( array: &FixedSizeListArray, elements_field: &FieldRef, list_size: i32, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -52,11 +60,8 @@ fn list_to_list( list_size ); - let elements = ctx.session().clone().arrow().execute_arrow( - array.elements().clone(), - Some(elements_field.as_ref()), - ctx, - )?; + let elements = + exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Cannot convert FixedSizeListArray to non-nullable Arrow array when elements are nullable" diff --git a/vortex-arrow/src/executor/list.rs b/vortex-arrow/src/executor/list.rs index f995661a0c1..a50f255f7bb 100644 --- a/vortex-arrow/src/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -31,8 +31,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use crate::ArrowExporter; use crate::executor::validity::to_arrow_null_buffer; -use crate::session::ArrowSessionExt; /// Matches the encodings [`to_arrow_list`] requires for export. struct ArrowListExportable; @@ -50,6 +50,7 @@ impl Matcher for ArrowListExportable { pub(super) fn to_arrow_list( array: ArrayRef, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -62,14 +63,19 @@ pub(super) fn to_arrow_list( // If the Vortex array is already in List format, we can directly convert it. if let Some(list) = array.as_opt::() { - return list_to_list::(&list.into_owned(), elements_field, ctx); + return list_to_list::(&list.into_owned(), elements_field, exporter, ctx); } // Converting each chunk individually, then using the fast concat logic from arrow if let Some(chunked) = array.as_opt::() { let mut arrow_chunks: Vec = Vec::with_capacity(chunked.nchunks()); for chunk in chunked.chunks() { - arrow_chunks.push(to_arrow_list::(chunk.clone(), elements_field, ctx)?); + arrow_chunks.push(to_arrow_list::( + chunk.clone(), + elements_field, + exporter, + ctx, + )?); } let refs = arrow_chunks.iter().map(|a| a.as_ref()).collect::>(); @@ -90,7 +96,7 @@ pub(super) fn to_arrow_list( } else { list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? }; - list_view_zctl::(zctl, elements_field, ctx) + list_view_zctl::(zctl, elements_field, exporter, ctx) } #[allow(rustdoc::broken_intra_doc_links)] @@ -98,6 +104,7 @@ pub(super) fn to_arrow_list( fn list_to_list( array: &ListArray, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { // We must cast the offsets to the required offset type. @@ -109,11 +116,8 @@ fn list_to_list( .to_buffer::() .into_arrow_offset_buffer(); - let elements = ctx.session().clone().arrow().execute_arrow( - array.elements().clone(), - Some(elements_field.as_ref()), - ctx, - )?; + let elements = + exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Cannot convert to non-nullable Arrow array with null elements" @@ -133,16 +137,14 @@ fn list_to_list( fn list_view_zctl( array: ListViewArray, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { assert!(array.is_zero_copy_to_list()); if array.is_empty() { - let elements = ctx.session().clone().arrow().execute_arrow( - array.elements().clone(), - Some(elements_field.as_ref()), - ctx, - )?; + let elements = + exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; return Ok(Arc::new(GenericListArray::::new( Arc::clone(elements_field), OffsetBuffer::new_empty(), @@ -191,11 +193,7 @@ fn list_view_zctl( }); // Extract the elements array. - let elements = ctx.session().clone().arrow().execute_arrow( - elements, - Some(elements_field.as_ref()), - ctx, - )?; + let elements = exporter.execute_arrow(elements, Some(elements_field.as_ref()), ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Cannot convert to non-nullable Arrow array with null elements" diff --git a/vortex-arrow/src/executor/list_view.rs b/vortex-arrow/src/executor/list_view.rs index 4699e33c54b..a19de07aa47 100644 --- a/vortex-arrow/src/executor/list_view.rs +++ b/vortex-arrow/src/executor/list_view.rs @@ -24,12 +24,13 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use crate::ArrowExporter; use crate::executor::validity::to_arrow_null_buffer; -use crate::session::ArrowSessionExt; pub(super) fn to_arrow_list_view( array: ArrayRef, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -67,12 +68,13 @@ pub(super) fn to_arrow_list_view( array }; - list_view_to_list_view::(array, elements_field, ctx) + list_view_to_list_view::(array, elements_field, exporter, ctx) } fn list_view_to_list_view( array: ListViewArray, elements_field: &FieldRef, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let ListViewDataParts { @@ -84,11 +86,7 @@ fn list_view_to_list_view( } = array.into_data_parts(); let n_elements = elements.len(); - let elements = ctx.session().clone().arrow().execute_arrow( - elements, - Some(elements_field.as_ref()), - ctx, - )?; + let elements = exporter.execute_arrow(elements, Some(elements_field.as_ref()), ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Elements field is non-nullable but elements array contains nulls" diff --git a/vortex-arrow/src/executor/map.rs b/vortex-arrow/src/executor/map.rs index c6da7206695..ae65af526bc 100644 --- a/vortex-arrow/src/executor/map.rs +++ b/vortex-arrow/src/executor/map.rs @@ -19,6 +19,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use crate::ArrowExporter; use crate::executor::list::to_arrow_list; /// Converts a Vortex Map array into an Arrow [`MapArray`](ArrowMapArray). @@ -30,6 +31,7 @@ pub(super) fn to_arrow_map( array: ArrayRef, entries_field: &FieldRef, keys_sorted: bool, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let array = match array.try_downcast::() { @@ -44,7 +46,7 @@ pub(super) fn to_arrow_map( let entries = array.entries().clone(); let entries_list_type = DataType::List(Arc::clone(entries_field)); - let entries_list = to_arrow_list::(entries, entries_field, ctx)?; + let entries_list = to_arrow_list::(entries, entries_field, exporter, ctx)?; vortex_ensure_eq!( entries_list.data_type(), &entries_list_type, diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 37556b09ee1..983a60d842b 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -43,6 +43,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use crate::ArrowExporter; use crate::executor::bool::to_arrow_bool; use crate::executor::byte::to_arrow_byte_array; use crate::executor::byte_view::to_arrow_byte_view; @@ -126,6 +127,7 @@ impl ArrowArrayExecutor for ArrayRef { pub(crate) fn execute_arrow_naive( array: ArrayRef, data_type: Option<&DataType>, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -153,21 +155,27 @@ pub(crate) fn execute_arrow_naive( DataType::LargeBinary => to_arrow_byte_array::(array, ctx), DataType::Utf8 => to_arrow_byte_array::(array, ctx), DataType::LargeUtf8 => to_arrow_byte_array::(array, ctx), - DataType::BinaryView => to_arrow_byte_view::(array, ctx), - DataType::Utf8View => to_arrow_byte_view::(array, ctx), + DataType::BinaryView => to_arrow_byte_view::(array, exporter, ctx), + DataType::Utf8View => to_arrow_byte_view::(array, exporter, ctx), // TODO(joe): pass down preferred - DataType::List(elements_field) => to_arrow_list::(array, elements_field, ctx), + DataType::List(elements_field) => { + to_arrow_list::(array, elements_field, exporter, ctx) + } // TODO(joe): pass down preferred - DataType::LargeList(elements_field) => to_arrow_list::(array, elements_field, ctx), + DataType::LargeList(elements_field) => { + to_arrow_list::(array, elements_field, exporter, ctx) + } // TODO(joe): pass down preferred DataType::FixedSizeList(elements_field, list_size) => { - to_arrow_fixed_list(array, *list_size, elements_field, ctx) + to_arrow_fixed_list(array, *list_size, elements_field, exporter, ctx) } // TODO(joe): pass down preferred - DataType::ListView(elements_field) => to_arrow_list_view::(array, elements_field, ctx), + DataType::ListView(elements_field) => { + to_arrow_list_view::(array, elements_field, exporter, ctx) + } // TODO(joe): pass down preferred DataType::LargeListView(elements_field) => { - to_arrow_list_view::(array, elements_field, ctx) + to_arrow_list_view::(array, elements_field, exporter, ctx) } DataType::Struct(fields) => { let fields = if data_type.is_none() { @@ -175,11 +183,11 @@ pub(crate) fn execute_arrow_naive( } else { Some(fields) }; - to_arrow_struct(array, fields, ctx) + to_arrow_struct(array, fields, exporter, ctx) } // TODO(joe): pass down preferred DataType::Dictionary(codes_type, values_type) => { - to_arrow_dictionary(array, codes_type, values_type, ctx) + to_arrow_dictionary(array, codes_type, values_type, exporter, ctx) } dt @ DataType::Decimal32(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal64(..) => to_arrow_decimal(array, dt, ctx), @@ -187,13 +195,13 @@ pub(crate) fn execute_arrow_naive( dt @ DataType::Decimal256(..) => to_arrow_decimal(array, dt, ctx), // TODO(joe): pass down preferred DataType::RunEndEncoded(ends_type, values_type) => { - to_arrow_run_end(array, ends_type.data_type(), values_type, ctx) + to_arrow_run_end(array, ends_type.data_type(), values_type, exporter, ctx) } dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), DataType::Map(entries_field, keys_sorted) => { - to_arrow_map(array, entries_field, *keys_sorted, ctx) + to_arrow_map(array, entries_field, *keys_sorted, exporter, ctx) } DataType::FixedSizeBinary(_) | DataType::Duration(_) diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index 7dc16b303cc..e1f6928c09f 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -26,8 +26,8 @@ use vortex_runend::RunEndArray; use vortex_runend::RunEndArrayExt; use vortex_runend::RunEndArraySlotsExt; -use crate::ArrowArrayExecutor; -use crate::session::ArrowSessionExt; +use crate::ArrowExporter; +use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_run_end`] requires for export. struct ArrowRunEndExportable; @@ -44,21 +44,24 @@ pub(super) fn to_arrow_run_end( array: ArrayRef, ends_type: &DataType, values_type: &Field, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let array = array.execute_until::(ctx)?; let array = match array.try_downcast::() { - Ok(constant) => return constant_to_run_end(constant, ends_type, values_type, ctx), + Ok(constant) => { + return constant_to_run_end(constant, ends_type, values_type, exporter, ctx); + } Err(array) => array, }; let array = match array.try_downcast::() { - Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, ctx), + Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, exporter, ctx), Err(array) => array, }; // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = export_values(array, values_type, ctx)?; + let flat = export_values(array, values_type, exporter, ctx)?; let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), Arc::new(values_type.clone()), @@ -71,25 +74,24 @@ pub(super) fn to_arrow_run_end( fn export_values( values: ArrayRef, values_type: &Field, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { - ctx.session() - .clone() - .arrow() - .execute_arrow(values, Some(values_type), ctx) + exporter.execute_arrow(values, Some(values_type), ctx) } fn run_end_to_arrow( array: RunEndArray, ends_type: &DataType, values_type: &Field, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let length = array.len(); let offset = array.offset(); - let arrow_ends = array.ends().clone().execute_arrow(Some(ends_type), ctx)?; - let arrow_values = export_values(array.values().clone(), values_type, ctx)?; + let arrow_ends = execute_arrow_naive(array.ends().clone(), Some(ends_type), exporter, ctx)?; + let arrow_values = export_values(array.values().clone(), values_type, exporter, ctx)?; match ends_type { DataType::Int16 => build_run_array::(&arrow_ends, &arrow_values, offset, length), @@ -144,6 +146,7 @@ fn constant_to_run_end( array: ConstantArray, ends_type: &DataType, values_type: &Field, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -160,6 +163,7 @@ fn constant_to_run_end( let values = export_values( ConstantArray::new(scalar.clone(), 1).into_array(), values_type, + exporter, ctx, )?; diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index 11fcb24e56b..7e25345aac0 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -27,7 +27,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use crate::ArrowArrayExecutor; +use crate::ArrowExporter; use crate::executor::infer_nearest_arrow_field; use crate::executor::validity::to_arrow_null_buffer; use crate::session::ArrowSessionExt; @@ -51,6 +51,7 @@ impl Matcher for ArrowStructExportable { pub(super) fn to_arrow_struct( array: ArrayRef, target_fields: Option<&Fields>, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -82,6 +83,7 @@ pub(super) fn to_arrow_struct( &fields, validity, len, + exporter, ctx, ); } @@ -100,6 +102,7 @@ pub(super) fn to_arrow_struct( &array.children(), None, // Pack is never null, len, + exporter, ctx, ); } @@ -130,6 +133,7 @@ pub(super) fn to_arrow_struct( &fields, validity, len, + exporter, ctx, ) } @@ -139,6 +143,7 @@ fn create_from_fields( vortex_fields: &[ArrayRef], null_buffer: Option, len: usize, + exporter: &ArrowExporter<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { match fields { @@ -153,11 +158,8 @@ fn create_from_fields( for (field, vx_field) in fields.iter().zip_eq(vortex_fields.iter()) { // Route through the session with the full Field (not just data_type) so any // ARROW:extension:name metadata reaches the export-plugin dispatcher. - let arrow_field = ctx.session().clone().arrow().execute_arrow( - vx_field.clone(), - Some(field.as_ref()), - ctx, - )?; + let arrow_field = + exporter.execute_arrow(vx_field.clone(), Some(field.as_ref()), ctx)?; vortex_ensure!( field.is_nullable() || arrow_field.null_count() == 0, "Cannot convert field '{}' to non-nullable Arrow field because it contains nulls", @@ -183,7 +185,7 @@ fn create_from_fields( let mut arrow_fields = Vec::with_capacity(vortex_fields.len()); for (name, vx_field) in names.iter().zip_eq(vortex_fields.iter()) { let inferred = infer_nearest_arrow_field(vx_field, name.as_ref(), ctx)?; - let arrow_array = vx_field.clone().execute_arrow(None, ctx)?; + let arrow_array = exporter.execute_arrow(vx_field.clone(), None, ctx)?; // The executed array is authoritative for the physical type; only the metadata is // taken from the inferred field. arrow_fields.push(Arc::new( diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 9bab3273ecf..7b08d5e7c6f 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -30,6 +30,7 @@ pub mod dtype; mod executor; mod iter; mod null_buffer; +mod options; mod run_end_import; mod scalar; mod session; @@ -47,6 +48,7 @@ pub use executor::*; pub use iter::*; pub use null_buffer::to_arrow_null_buffer; pub use null_buffer::to_null_buffer; +pub use options::*; pub use scalar::ToArrowDatum; pub use session::*; diff --git a/vortex-arrow/src/options/mod.rs b/vortex-arrow/src/options/mod.rs new file mode 100644 index 00000000000..5cef3b5f45b --- /dev/null +++ b/vortex-arrow/src/options/mod.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Per-export options, including options defined by external Arrow export plugins. + +use std::any::Any; +use std::any::TypeId; +use std::fmt; +use std::sync::Arc; + +/// Typed options propagated to every nested array in an Arrow export. +/// +/// Export plugins can define their own option types without modifying this crate. Each type has +/// at most one value; absent options leave the exporter's default behavior unchanged. +/// +/// ``` +/// use vortex_arrow::{ArrowExportOptions, CompactBuffers}; +/// +/// let options = ArrowExportOptions::default().with(CompactBuffers(false)); +/// assert_eq!(options.get::(), Some(&CompactBuffers(false))); +/// ``` +#[derive(Clone, Default)] +pub struct ArrowExportOptions { + values: Vec<(TypeId, Arc)>, +} + +impl ArrowExportOptions { + /// Set an option, replacing any previous value of the same type. + pub fn with(mut self, option: T) -> Self { + let id = TypeId::of::(); + if let Some((_, value)) = self.values.iter_mut().find(|(key, _)| *key == id) { + *value = Arc::new(option); + } else { + self.values.push((id, Arc::new(option))); + } + self + } + + /// Look up an option defined by this crate or an external exporter. + pub fn get(&self) -> Option<&T> { + self.values + .iter() + .find(|(key, _)| *key == TypeId::of::()) + .and_then(|(_, value)| value.downcast_ref()) + } + + /// Look up an option, falling back to its default when unset. + pub fn get_or_default(&self) -> T { + self.get::().cloned().unwrap_or_default() + } +} + +impl fmt::Debug for ArrowExportOptions { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ArrowExportOptions") + .field("option_count", &self.values.len()) + .finish_non_exhaustive() + } +} + +/// Whether to compact string and binary view buffers during Arrow export. +/// +/// Enabled by default. Disabling compaction avoids scans and copies, but the exported array may +/// retain backing data outside the selected rows. Values and the Arrow type are unchanged. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CompactBuffers(pub bool); + +impl Default for CompactBuffers { + fn default() -> Self { + Self(true) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-arrow/src/options/tests.rs b/vortex-arrow/src/options/tests.rs new file mode 100644 index 00000000000..79eb64b4e75 --- /dev/null +++ b/vortex-arrow/src/options/tests.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicU32; +use std::sync::atomic::Ordering; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::RunArray; +use arrow_array::cast::AsArray; +use arrow_array::types::BinaryViewType; +use arrow_array::types::Int32Type; +use arrow_array::types::StringViewType; +use arrow_array::types::UInt8Type; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::FieldRef; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_runend::RunEnd; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::*; +use crate::ArrowExport; +use crate::ArrowExportVTable; +use crate::ArrowSession; +use crate::ArrowSessionExt; +use crate::executor::byte_view::canonical_varbinview_to_arrow; + +#[test] +fn typed_options_are_independent() { + struct CustomOption(u32); + let original = ArrowExportOptions::default().with(CompactBuffers(false)); + let changed = original + .clone() + .with(CompactBuffers(true)) + .with(CustomOption(42)); + assert_eq!( + original.get::(), + Some(&CompactBuffers(false)) + ); + assert!(original.get::().is_none()); + assert_eq!(changed.get::(), Some(&CompactBuffers(true))); + assert_eq!(changed.get::().map(|value| value.0), Some(42)); + assert!( + ArrowExportOptions::default() + .get::() + .is_none() + ); +} + +#[test] +fn get_or_default_falls_back_to_default() { + assert_eq!( + ArrowExportOptions::default().get_or_default::(), + CompactBuffers(true) + ); + assert_eq!( + ArrowExportOptions::default() + .with(CompactBuffers(false)) + .get_or_default::(), + CompactBuffers(false) + ); +} + +#[derive(Debug)] +struct Plugin; + +struct PluginOption(Arc); + +static TEST_OPTION_ID: CachedId = CachedId::new("test.options"); + +impl ArrowExportVTable for Plugin { + fn arrow_ext_id(&self) -> Id { + *TEST_OPTION_ID + } + + fn vortex_id(&self) -> Id { + *TEST_OPTION_ID + } + + fn to_arrow_field( + &self, + _name: &str, + _dtype: &DType, + _session: &ArrowSession, + ) -> VortexResult> { + Ok(None) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + options: &ArrowExportOptions, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + options + .get::() + .expect("external option") + .0 + .store(42, Ordering::Relaxed); + Ok(ArrowExport::Unsupported(array)) + } +} + +#[test] +fn external_options_reach_nested_exporters() -> VortexResult<()> { + let session = array_session(); + session.arrow().register_exporter(Arc::new(Plugin)); + let mut ctx = session.create_execution_ctx(); + let child = Field::new("item", DataType::Int32, false).with_metadata( + [("ARROW:extension:name".to_owned(), "test.options".to_owned())] + .into_iter() + .collect(), + ); + let field = Field::new("list", DataType::List(Arc::new(child)), false); + let array = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0i32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let called = Arc::new(AtomicU32::new(0)); + let options = ArrowExportOptions::default().with(PluginOption(Arc::clone(&called))); + let arrow = session + .arrow() + .exporter(&options) + .execute_arrow(array, Some(&field), &mut ctx)?; + assert_eq!(arrow.len(), 1); + assert_eq!(called.load(Ordering::Relaxed), 42); + Ok(()) +} + +/// Three rows taken from distant positions of one large view buffer, so compaction has to copy +/// rather than just trim it. +fn sparse_views(binary: bool, ctx: &mut ExecutionCtx) -> VortexResult { + let value = "x".repeat(256); + let values = (0..128).map(|i| (i != 1).then(|| format!("{i}{value}"))); + let array = if binary { + VarBinViewArray::from_iter_nullable_bin(values) + } else { + VarBinViewArray::from_iter_nullable_str(values) + }; + array + .into_array() + .take(buffer![0u32, 1, 12].into_array())? + .execute::(ctx) +} + +fn view_buffer_ptrs(array: &ArrowArrayRef) -> Vec<*const u8> { + let buffers = match array.data_type() { + DataType::BinaryView => array.as_binary_view().data_buffers(), + DataType::Utf8View => array.as_string_view().data_buffers(), + dt => unreachable!("not a view type: {dt}"), + }; + buffers.iter().map(|buffer| buffer.as_ptr()).collect() +} + +#[derive(Clone, Copy, Debug)] +enum Shape { + Flat, + List, + ListView, + FixedSizeList, + RunEnd, + Struct, + ListOfStruct, + Dictionary, +} + +impl Shape { + /// Wrap a three-row `leaf` into this shape, returning the array and its Arrow target type. + fn wrap( + self, + leaf: ArrayRef, + leaf_type: DataType, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(ArrayRef, DataType)> { + let field: FieldRef = Arc::new(Field::new("item", leaf_type.clone(), true)); + Ok(match self { + Shape::Flat => (leaf, leaf_type), + Shape::List => ( + ListArray::try_new(leaf, buffer![0i32, 3].into_array(), Validity::NonNullable)? + .into_array(), + DataType::List(field), + ), + Shape::ListView => ( + ListViewArray::try_new( + leaf, + buffer![0i32].into_array(), + buffer![3i32].into_array(), + Validity::NonNullable, + )? + .into_array(), + DataType::ListView(field), + ), + Shape::FixedSizeList => ( + FixedSizeListArray::try_new(leaf, 3, Validity::NonNullable, 1)?.into_array(), + DataType::FixedSizeList(field, 3), + ), + Shape::RunEnd => ( + RunEnd::try_new(buffer![1u32, 2, 3].into_array(), leaf, ctx)?.into_array(), + DataType::RunEndEncoded( + Arc::new(Field::new("ends", DataType::Int32, false)), + field, + ), + ), + Shape::Struct => ( + StructArray::try_new( + FieldNames::from(["item"]), + vec![leaf], + 3, + Validity::NonNullable, + )? + .into_array(), + DataType::Struct(vec![field].into()), + ), + Shape::ListOfStruct => { + let (array, struct_type) = Shape::Struct.wrap(leaf, leaf_type, ctx)?; + Shape::List.wrap(array, struct_type, ctx)? + } + Shape::Dictionary => ( + DictArray::try_new(buffer![0u8, 1, 2].into_array(), leaf)?.into_array(), + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(leaf_type)), + ), + }) + } + + /// Extract the leaf column from an exported array of this shape. + fn leaf(self, arrow: ArrowArrayRef) -> ArrowArrayRef { + match self { + Shape::Flat => arrow, + Shape::List => Arc::clone(arrow.as_list::().values()), + Shape::ListView => Arc::clone(arrow.as_list_view::().values()), + Shape::FixedSizeList => Arc::clone(arrow.as_fixed_size_list().values()), + Shape::RunEnd => Arc::clone( + arrow + .as_any() + .downcast_ref::>() + .expect("run array") + .values(), + ), + Shape::Struct => Arc::clone(arrow.as_struct().column(0)), + Shape::ListOfStruct => Shape::Struct.leaf(Shape::List.leaf(arrow)), + Shape::Dictionary => Arc::clone(arrow.as_dictionary::().values()), + } + } +} + +#[rstest] +fn compaction_option_reaches_nested_views( + #[values( + Shape::Flat, + Shape::List, + Shape::ListView, + Shape::FixedSizeList, + Shape::RunEnd, + Shape::Struct, + Shape::ListOfStruct, + Shape::Dictionary + )] + shape: Shape, + #[values(false, true)] binary: bool, + #[values(None, Some(false), Some(true))] compact: Option, +) -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let options = compact.map_or_else(ArrowExportOptions::default, |compact| { + ArrowExportOptions::default().with(CompactBuffers(compact)) + }); + + let view = sparse_views(binary, &mut ctx)?; + let backing = view.data_buffers()[0].as_host().as_ptr(); + let expected = if binary { + canonical_varbinview_to_arrow::(&view, &mut ctx)? + } else { + canonical_varbinview_to_arrow::(&view, &mut ctx)? + }; + + let (array, target) = shape.wrap(view.into_array(), expected.data_type().clone(), &mut ctx)?; + let arrow = session.arrow().exporter(&options).execute_arrow( + array, + Some(&Field::new("", target, true)), + &mut ctx, + )?; + + let leaf = shape.leaf(arrow); + assert_eq!(leaf.to_data(), expected.to_data()); + assert_eq!( + view_buffer_ptrs(&leaf).contains(&backing), + compact == Some(false) + ); + Ok(()) +} + +#[rstest] +fn execute_arrow_compacts_by_default(#[values(false, true)] binary: bool) -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let view = sparse_views(binary, &mut ctx)?; + let backing = view.data_buffers()[0].as_host().as_ptr(); + + let arrow = session + .arrow() + .execute_arrow(view.into_array(), None, &mut ctx)?; + + assert!(!view_buffer_ptrs(&arrow).contains(&backing)); + Ok(()) +} diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 9915f63a105..641a8dfbeb7 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -70,6 +70,7 @@ use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; +use crate::ArrowExportOptions; use crate::IntoVortexArray; use crate::convert::from_arrow_dyn; use crate::convert::map_from_arrow_parts; @@ -170,11 +171,13 @@ pub trait ArrowExportVTable: 'static + Send + Sync + Debug { /// Convert a Vortex array into an Arrow array shaped to `target`. /// /// Returns ownership of `array` via [`ArrowExport::Unsupported`] when the plugin cannot - /// handle the input. + /// handle the input. Plugins that export child or storage arrays must pass `options` on to + /// those exports. fn execute_arrow( &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult; } @@ -531,63 +534,47 @@ impl ArrowSession { /// /// With `target = None` the fallback path picks the array's preferred Arrow physical type /// and executes directly into that, ignoring extension types. - #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] + /// + /// Uses the default [`ArrowExportOptions`]; see [`exporter`](Self::exporter) to + /// change them. pub fn execute_arrow( &self, array: ArrayRef, target: Option<&Field>, ctx: &mut ExecutionCtx, ) -> VortexResult { - // NOTE(aduffy): this looks strange, but we do this to keep target_field as &Field so - // we can avoid cloning target when it is provided. It contains a HashMap internally that - // can be expensive to copy. - let arrow_field; - let target_field = match target { - Some(field) => field, - None => { - let session = ctx.session().clone(); - arrow_field = session.arrow().to_arrow_field("", array.dtype())?; - &arrow_field - } - }; - - if let Some(arrow_ext_name) = target_field.metadata().get(EXTENSION_TYPE_NAME_KEY) { - // There can be multiple plugins that report support for a particular extension type. - // We try them in order until one of them reports a successful conversion. - let len = array.len(); - let mut current = array; - - for plugin in self.exporters(&Id::new(arrow_ext_name)).iter() { - trace!( - plugin = ?plugin, - extension_name = arrow_ext_name, - "probing plugin for converting Arrow array" - ); - - match plugin.execute_arrow(current, target_field, ctx)? { - ArrowExport::Exported(arrow) => { - vortex_ensure_eq!( - arrow.len(), - len, - "Arrow array length does not match Vortex array length after conversion to {:?}", - arrow - ); - return Ok(arrow); - } - ArrowExport::Unsupported(array) => current = array, - } - } - - debug!( - extension_id = arrow_ext_name, - data_type = ?target_field.data_type(), - "unsupported Arrow extension type encountered, falling back to naive execution" - ); + self.exporter(&ArrowExportOptions::default()) + .execute_arrow(array, target, ctx) + } - return execute_arrow_naive(current, Some(target_field.data_type()), ctx); + /// Return an exporter that applies `options` to its exports. + /// + /// ``` + /// use arrow_array::Array; + /// use vortex_array::IntoArray; + /// use vortex_array::VortexSessionExecute; + /// use vortex_array::array_session; + /// use vortex_array::arrays::VarBinViewArray; + /// use vortex_arrow::ArrowExportOptions; + /// use vortex_arrow::ArrowSessionExt; + /// use vortex_arrow::CompactBuffers; + /// + /// let session = array_session(); + /// let array = VarBinViewArray::from_iter_str(["a", "b"]).into_array(); + /// let options = ArrowExportOptions::default().with(CompactBuffers(false)); + /// let arrow = session.arrow().exporter(&options).execute_arrow( + /// array, + /// None, + /// &mut session.create_execution_ctx(), + /// )?; + /// assert_eq!(arrow.len(), 2); + /// # Ok::<(), vortex_error::VortexError>(()) + /// ``` + pub fn exporter<'a>(&'a self, options: &'a ArrowExportOptions) -> ArrowExporter<'a> { + ArrowExporter { + session: self, + options, } - - execute_arrow_naive(array, target.map(|field| field.data_type()), ctx) } /// Decode an Arrow array into a Vortex array. @@ -810,6 +797,81 @@ impl SessionVar for ArrowSession { } } +/// An [`ArrowSession`] bound to [`ArrowExportOptions`], created by [`ArrowSession::exporter`]. +#[derive(Clone, Copy, Debug)] +pub struct ArrowExporter<'a> { + session: &'a ArrowSession, + options: &'a ArrowExportOptions, +} + +impl ArrowExporter<'_> { + pub(crate) fn options(&self) -> &ArrowExportOptions { + self.options + } + + /// Execute a Vortex array into an Arrow array with the bound options. + /// + /// Behaves like [`ArrowSession::execute_arrow`]. The options are propagated to nested arrays + /// and to registered export plugins. + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] + pub fn execute_arrow( + &self, + array: ArrayRef, + target: Option<&Field>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NOTE(aduffy): this looks strange, but we do this to keep target_field as &Field so + // we can avoid cloning target when it is provided. It contains a HashMap internally that + // can be expensive to copy. + let arrow_field; + let target_field = match target { + Some(field) => field, + None => { + arrow_field = self.session.to_arrow_field("", array.dtype())?; + &arrow_field + } + }; + + if let Some(arrow_ext_name) = target_field.metadata().get(EXTENSION_TYPE_NAME_KEY) { + // There can be multiple plugins that report support for a particular extension type. + // We try them in order until one of them reports a successful conversion. + let len = array.len(); + let mut current = array; + + for plugin in self.session.exporters(&Id::new(arrow_ext_name)).iter() { + trace!( + plugin = ?plugin, + extension_name = arrow_ext_name, + "probing plugin for converting Arrow array" + ); + + match plugin.execute_arrow(current, target_field, self.options, ctx)? { + ArrowExport::Exported(arrow) => { + vortex_ensure_eq!( + arrow.len(), + len, + "Arrow array length does not match Vortex array length after conversion to {:?}", + arrow + ); + return Ok(arrow); + } + ArrowExport::Unsupported(array) => current = array, + } + } + + debug!( + extension_id = arrow_ext_name, + data_type = ?target_field.data_type(), + "unsupported Arrow extension type encountered, falling back to naive execution" + ); + + return execute_arrow_naive(current, Some(target_field.data_type()), self, ctx); + } + + execute_arrow_naive(array, target.map(|field| field.data_type()), self, ctx) + } +} + /// Extension trait for accessing the [`ArrowSession`] on a Vortex session. pub trait ArrowSessionExt: SessionExt { /// Get the Arrow session. diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 474000795e2..9509816f4ad 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -17,6 +17,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -71,6 +72,7 @@ impl ArrowExportVTable for Json { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_json = array @@ -93,9 +95,11 @@ impl ArrowExportVTable for Json { ); let session = ctx.session().clone(); - let storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; Ok(ArrowExport::Exported(storage)) } @@ -159,10 +163,14 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; + use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSessionExt; + use vortex_arrow::CompactBuffers; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -234,4 +242,42 @@ mod tests { assert_eq!(strings.value(1), "{\"id\":2}"); Ok(()) } + + /// The plugin must forward export options to its storage export. + #[test] + fn export_options_reach_json_storage() -> VortexResult<()> { + let session = vortex_array::array_session(); + initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let views = VarBinViewArray::from_iter_str((0..64).map(|i| format!("{{\"id\":{i:>32}}}"))); + let backing = views.data_buffers()[0].as_host().as_ptr(); + // Rows far apart in one buffer, so compaction has to copy. + let storage = views + .into_array() + .take(PrimitiveArray::from_iter([0u32, 48]).into_array())?; + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage.dtype().clone())?.erased(); + let array = ExtensionArray::new(ext_dtype, storage).into_array(); + let field = session.arrow().to_arrow_field("data", array.dtype())?; + + let retains_backing = |options: &ArrowExportOptions, + ctx: &mut vortex_array::ExecutionCtx| + -> VortexResult { + let exported = session.arrow().exporter(options).execute_arrow( + array.clone(), + Some(&field), + ctx, + )?; + Ok(exported + .as_string_view() + .data_buffers() + .iter() + .any(|buffer| buffer.as_ptr() == backing)) + }; + + let uncompacted = ArrowExportOptions::default().with(CompactBuffers(false)); + assert!(retains_backing(&uncompacted, &mut ctx)?); + assert!(!retains_backing(&ArrowExportOptions::default(), &mut ctx)?); + Ok(()) + } } diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index c1d6b374586..0cd15667b37 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -39,6 +39,7 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -272,6 +273,7 @@ impl ArrowExportVTable for LineString { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_linestring = array @@ -299,9 +301,11 @@ impl ArrowExportVTable for LineString { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through GeoArrow's line-string array; `into_arrow` is concrete, so wrap in `Arc`. let linestrings = LineStringArray::try_from((arrow_storage.as_ref(), linestring_meta)) diff --git a/vortex-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index fb53863ecc7..93ff0e31c8e 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -32,6 +32,7 @@ use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -202,6 +203,7 @@ impl ArrowExportVTable for MultiLineString { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_multilinestring = array @@ -229,9 +231,11 @@ impl ArrowExportVTable for MultiLineString { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; let multilinestrings = MultiLineStringArray::try_from((arrow_storage.as_ref(), multilinestring_meta)) diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index 4afc9052582..c0a2f9d95a5 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -32,6 +32,7 @@ use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -190,6 +191,7 @@ impl ArrowExportVTable for MultiPoint { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_multipoint = array @@ -217,9 +219,11 @@ impl ArrowExportVTable for MultiPoint { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through GeoArrow's multipoint array; `into_arrow` is concrete, so wrap in `Arc`. let multipoints = MultiPointArray::try_from((arrow_storage.as_ref(), multipoint_meta)) diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index a3291ba68ea..f6891e7708e 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -31,6 +31,7 @@ use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -200,6 +201,7 @@ impl ArrowExportVTable for MultiPolygon { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_multipolygon = array @@ -227,9 +229,11 @@ impl ArrowExportVTable for MultiPolygon { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; let multipolygons = MultiPolygonArray::try_from((arrow_storage.as_ref(), multipolygon_meta)) diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e6a00fe8fea..3f3ab8c02ff 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -30,6 +30,7 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -179,6 +180,7 @@ impl ArrowExportVTable for Point { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_point = array @@ -206,9 +208,11 @@ impl ArrowExportVTable for Point { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through the GeoArrow point array type: this validates that the storage is // the separated-coordinate struct layout expected for a `PointType` extension field. diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index 0a0763a7140..e733c01031b 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -37,6 +37,7 @@ use vortex_array::memory::BufferAllocatorRef; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -255,6 +256,7 @@ impl ArrowExportVTable for Polygon { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_polygon = array @@ -282,9 +284,11 @@ impl ArrowExportVTable for Polygon { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through GeoArrow's polygon array; `into_arrow` is concrete, so wrap in `Arc`. let polygons = PolygonArray::try_from((arrow_storage.as_ref(), polygon_meta)) diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index 21e33808cc6..2c53020ee97 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -39,6 +39,7 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -234,6 +235,7 @@ impl ArrowExportVTable for Rect { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_box = array @@ -258,9 +260,11 @@ impl ArrowExportVTable for Rect { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through GeoArrow's rect array; `into_arrow` is concrete, so wrap in `Arc`. let rects = RectArray::try_from((arrow_storage.as_ref(), box_meta)) diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index a6dddc70497..8b0b42fb97b 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -195,6 +196,7 @@ impl ArrowExportVTable for WellKnownBinary { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let is_wkb = array @@ -219,9 +221,11 @@ impl ArrowExportVTable for WellKnownBinary { target.is_nullable(), ); let session = ctx.session().clone(); - let arrow_storage = session - .arrow() - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(&storage_field), ctx)?; // Round-trip through the GeoArrow WKB array types: this validates that the storage // is a binary-family Arrow array and produces the canonical physical representation diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index 804ae26b6ad..ef071bbdf25 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -24,6 +24,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; @@ -92,6 +93,7 @@ impl ArrowExportVTable for Vector { &self, array: ArrayRef, target: &Field, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { if !array @@ -106,7 +108,11 @@ impl ArrowExportVTable for Vector { let storage = executed.storage_array().clone(); let session = ctx.session().clone(); - let arrow_storage = session.arrow().execute_arrow(storage, Some(target), ctx)?; + let arrow_storage = + session + .arrow() + .exporter(options) + .execute_arrow(storage, Some(target), ctx)?; Ok(ArrowExport::Exported(arrow_storage)) } @@ -360,6 +366,7 @@ mod tests { &Vector, primitive.clone(), &target, + &ArrowExportOptions::default(), &mut ctx, )?; assert!(