diff --git a/Cargo.lock b/Cargo.lock index fa0f248cbf0..8451ac1a3e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11578,6 +11578,7 @@ dependencies = [ "vortex-error", "vortex-fastlanes", "vortex-fsst", + "vortex-mask", "vortex-onpair", "vortex-pco", "vortex-sequence", diff --git a/benchmarks/compress-bench/src/gpu/vortex.rs b/benchmarks/compress-bench/src/gpu/vortex.rs index 414126c15e3..2b3b9787b01 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -29,7 +29,7 @@ use vortex::file::WriteOptionsSessionExt; use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex::layout::layouts::compressed::CompressingStrategy; use vortex::layout::scan::split_by::SplitBy; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressed; @@ -303,11 +303,11 @@ fn verify_field( batch_index: usize, field_index: usize, ) -> Result<()> { - let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?; + let expected = host.clone().execute_arrow(None, ctx)?; // Pin the Arrow target type so the two sides cannot land on different but equivalent // encodings of the same logical values. let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); - let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?; + let actual = gpu.execute_arrow(Some(&target), ctx)?; if expected.to_data() == actual.to_data() { return Ok(()); diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 96b36237059..aa0369dc98f 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -38,7 +38,7 @@ use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSession; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_arrow::to_arrow_null_buffer; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; @@ -449,9 +449,6 @@ pub trait ParquetVariantArrayExt: 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)?; @@ -459,7 +456,9 @@ pub trait ParquetVariantArrayExt: let mut fields = Vec::with_capacity(3); let mut arrays: Vec = Vec::with_capacity(3); - let metadata_arrow = exporter.execute_arrow(metadata.clone(), None, ctx)?; + let metadata_arrow = metadata + .clone() + .execute_arrow_with_options(None, options, ctx)?; fields.push(Arc::new(Field::new( "metadata", metadata_arrow.data_type().clone(), @@ -468,7 +467,9 @@ pub trait ParquetVariantArrayExt: arrays.push(metadata_arrow); if let Some(value) = self.value() { - let value_arrow = exporter.execute_arrow(value.clone(), None, ctx)?; + let value_arrow = value + .clone() + .execute_arrow_with_options(None, options, ctx)?; fields.push(Arc::new(Field::new( "value", value_arrow.data_type().clone(), @@ -478,7 +479,9 @@ pub trait ParquetVariantArrayExt: } if let Some(typed_value) = self.typed_value() { - let tv_arrow = exporter.execute_arrow(typed_value.clone(), None, ctx)?; + let tv_arrow = typed_value + .clone() + .execute_arrow_with_options(None, options, ctx)?; fields.push(Arc::new(Field::new( "typed_value", tv_arrow.data_type().clone(), diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 12c7c7e386a..430df4f8284 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -45,6 +45,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_arrow::to_arrow_null_buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -84,9 +85,6 @@ pub(crate) fn export_storage_to_target( 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 { @@ -104,7 +102,7 @@ pub(crate) fn export_storage_to_target( ); }; - arrays.push(exporter.execute_arrow(child, Some(field.as_ref()), ctx)?); + arrays.push(child.execute_arrow_with_options(Some(field.as_ref()), options, ctx)?); } let nulls = to_arrow_null_buffer( @@ -471,6 +469,7 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -490,9 +489,7 @@ mod tests { fn export_variant(session: &VortexSession, array: ArrayRef) -> VortexResult { let mut ctx = session.create_execution_ctx(); let field = session.arrow().to_arrow_field("variant", array.dtype())?; - let result = session - .arrow() - .execute_arrow(array, Some(&field), &mut ctx)?; + let result = array.execute_arrow(Some(&field), &mut ctx)?; Ok(ArrowVariantArray::try_new(result.as_ref())?) } @@ -784,9 +781,7 @@ mod tests { assert!(imported.as_opt::().is_some()); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(imported, Some(&field), &mut ctx)?; + let exported = imported.execute_arrow(Some(&field), &mut ctx)?; let exported = exported.as_struct(); assert_struct_arrays_eq(exported, &storage); @@ -805,9 +800,7 @@ mod tests { let canonical = VariantArray::try_new(core_storage, None)?.into_array(); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(canonical, Some(&field), &mut ctx)?; + let exported = canonical.execute_arrow(Some(&field), &mut ctx)?; let exported = exported.as_struct(); assert_struct_arrays_eq(exported, &storage); @@ -855,9 +848,7 @@ mod tests { .into(), ); - let exported = session - .arrow() - .execute_arrow(canonical, Some(&field), &mut ctx)?; + let exported = canonical.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), field.data_type()); let actual = session.arrow().from_arrow_array(exported, &field)?; @@ -883,9 +874,7 @@ mod tests { let expected = array.clone(); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(array.clone(), None, &mut ctx)?; + let exported = array.clone().execute_arrow(None, &mut ctx)?; let field = Field::new( "", @@ -926,9 +915,7 @@ mod tests { .into_array(); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(expected.clone(), None, &mut ctx)?; + let exported = expected.clone().execute_arrow(None, &mut ctx)?; assert_eq!( exported.data_type(), @@ -999,9 +986,7 @@ mod tests { ); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(expected.clone(), Some(&field), &mut ctx)?; + let exported = expected.clone().execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), field.data_type()); let actual = session.arrow().from_arrow_array(exported, &field)?; @@ -1044,9 +1029,8 @@ mod tests { ); let mut ctx = session.create_execution_ctx(); - let err = session - .arrow() - .execute_arrow(array, Some(&field), &mut ctx) + let err = array + .execute_arrow(Some(&field), &mut ctx) .unwrap_err() .to_string(); assert!( @@ -1094,9 +1078,7 @@ mod tests { .into(), ); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(array, Some(&field), &mut ctx)?; + let exported = array.execute_arrow(Some(&field), &mut ctx)?; assert_eq!( exported.data_type(), field.data_type(), diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 2d4d387c52d..d28bd43a45d 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -46,6 +46,7 @@ use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -154,7 +155,7 @@ fn json_strings_to_variant( ) -> VortexResult { let nullable = parent.dtype().is_nullable(); let session = ctx.session().clone(); - let arrow_strings = session.arrow().execute_arrow(strings, None, ctx)?; + let arrow_strings = strings.execute_arrow(None, ctx)?; // Any row that fails to parse as JSON fails the whole conversion. let arrow_variant = parquet_variant_compute::json_to_variant(&arrow_strings)?; diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index b73104cf6c5..9a71ad8d923 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -35,6 +35,8 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +use vortex_arrow::ExecuteArrowExt; + use crate::Pco; #[test] fn test_compress_decompress() { @@ -207,16 +209,11 @@ fn test_serde() -> VortexResult<()> { &SESSION, )?; let data_type = SESSION.arrow().to_arrow_field("", data.dtype())?; - let pco_arrow = SESSION - .arrow() - .execute_arrow(pco, Some(&data_type), &mut ctx)?; - let decoded_arrow = SESSION - .arrow() - .execute_arrow(decoded, Some(&data_type), &mut ctx)?; - let expected_arrow = - SESSION - .arrow() - .execute_arrow(data.into_array(), Some(&data_type), &mut ctx)?; + let pco_arrow = pco.execute_arrow(Some(&data_type), &mut ctx)?; + let decoded_arrow = decoded.execute_arrow(Some(&data_type), &mut ctx)?; + let expected_arrow = data + .into_array() + .execute_arrow(Some(&data_type), &mut ctx)?; assert_eq!(pco_arrow.as_ref(), expected_arrow.as_ref()); assert_eq!(decoded_arrow.as_ref(), expected_arrow.as_ref()); Ok(()) diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 418abea5a51..4fa1dc6dfb0 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -910,6 +910,7 @@ mod test { use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; use vortex_buffer::buffer_mut; @@ -1136,33 +1137,23 @@ mod test { let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar, &mut ctx).unwrap(); - let expected = SESSION - .arrow() - .execute_arrow( - DecimalArray::new( - buffer![100i128, 200, 123, 123, 123, 123, 123, 300, 4000, 123], - decimal_dtype, - // NB: patch indices: [0, 1, 7, 8]; patch validity: [Valid, Valid, Valid, Invalid]; ergo 0, 1, 7 are valid. - Validity::from_mask(Mask::from_excluded_indices(10, vec![8]), Nullable), - ) - .into_array(), - None, - &mut ctx, - ) - .unwrap(); + let expected = DecimalArray::new( + buffer![100i128, 200, 123, 123, 123, 123, 123, 300, 4000, 123], + decimal_dtype, + // NB: patch indices: [0, 1, 7, 8]; patch validity: [Valid, Valid, Valid, Invalid]; ergo 0, 1, 7 are valid. + Validity::from_mask(Mask::from_excluded_indices(10, vec![8]), Nullable), + ) + .into_array() + .execute_arrow(None, &mut ctx) + .unwrap(); - let actual = SESSION - .arrow() - .execute_arrow( - sparse_struct - .as_array() - .clone() - .execute::(&mut ctx) - .unwrap() - .into_array(), - None, - &mut ctx, - ) + let actual = sparse_struct + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap() + .into_array() + .execute_arrow(None, &mut ctx) .unwrap(); assert_eq!(expected.data_type(), actual.data_type()); @@ -1931,12 +1922,8 @@ mod test { // Note that the preferred arrow list representation is `List` (not `ListView`). let arrow_dtype = SESSION.arrow().to_arrow_field("", expected.dtype())?; - let actual = SESSION - .arrow() - .execute_arrow(actual, Some(&arrow_dtype), &mut ctx)?; - let expected = SESSION - .arrow() - .execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?; + let actual = actual.execute_arrow(Some(&arrow_dtype), &mut ctx)?; + let expected = expected.execute_arrow(Some(&arrow_dtype), &mut ctx)?; assert_eq!(actual.data_type(), expected.data_type()); Ok(()) diff --git a/encodings/uuid/src/arrow.rs b/encodings/uuid/src/arrow.rs index 55202147633..d69c35a7928 100644 --- a/encodings/uuid/src/arrow.rs +++ b/encodings/uuid/src/arrow.rs @@ -38,7 +38,7 @@ use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_arrow::nulls; use vortex_buffer::Alignment; use vortex_buffer::Buffer; @@ -185,12 +185,7 @@ fn try_fsl_to_fsb( storage.dtype().is_nullable(), ); - let session = ctx.session().clone(); - let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + let arrow_storage = storage.execute_arrow_with_options(Some(&storage_field), options, ctx)?; let fsl = arrow_storage.as_fixed_size_list(); let bytes = fsl @@ -231,11 +226,13 @@ mod tests { use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::assert_arrays_eq; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::StructFields; use vortex_array::iter::ArrayIterator; use vortex_arrow::ArrowArrayStreamAdapter; + use vortex_arrow::ArrowSessionExt; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -261,6 +258,37 @@ mod tests { ) } + /// A session that only [`crate::initialize`] sets up imports an `arrow.uuid` field through the + /// UUID plugin, and a field that names the Vortex ID of the UUID type through the dtype + /// registry. + #[test] + fn uuid_session_imports_by_plugin_and_by_vortex_id() -> VortexResult<()> { + let vortex_session = uuid_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let dtype = uuid_dtype(false); + + let mut plugin_field = Field::new("id", DataType::FixedSizeBinary(16), false); + plugin_field.try_with_extension_type(ArrowUuid)?; + assert_eq!(session.from_arrow_field(&plugin_field)?, dtype); + + let vortex_id_field = session.to_arrow_storage_field("id", dtype.as_extension())?; + assert_eq!(session.from_arrow_field(&vortex_id_field)?, dtype); + + let array = uuid_array(&session)?; + let exported = array + .clone() + .execute_arrow(Some(&vortex_id_field), &mut ctx)?; + assert_eq!(exported.data_type(), vortex_id_field.data_type()); + assert_arrays_eq!( + session.from_arrow_array(exported, &vortex_id_field)?, + array, + &mut ctx + ); + + Ok(()) + } + #[test] fn to_arrow_field_top_level_uuid_carries_extension_metadata() -> VortexResult<()> { let vortex_session = uuid_session(); @@ -455,7 +483,7 @@ mod tests { let vortex_ext = vortex_array.dtype().as_extension(); assert!(vortex_ext.is::()); - let exported = session.execute_arrow(vortex_array, None, &mut ctx)?; + let exported = vortex_array.execute_arrow(None, &mut ctx)?; assert_eq!(exported.data_type(), &DataType::FixedSizeBinary(16)); let fsb = exported.as_fixed_size_binary(); assert_eq!(fsb.len(), 2); @@ -492,7 +520,7 @@ mod tests { )? .into_array(); - let exported = session.execute_arrow(struct_array, None, &mut ctx)?; + let exported = struct_array.execute_arrow(None, &mut ctx)?; let DataType::Struct(fields) = exported.data_type() else { panic!("expected Struct, got {:?}", exported.data_type()); }; @@ -521,7 +549,7 @@ mod tests { )? .into_array(); - let exported = session.execute_arrow(list, None, &mut ctx)?; + let exported = list.execute_arrow(None, &mut ctx)?; let DataType::List(elem) = exported.data_type() else { panic!("expected List, got {:?}", exported.data_type()); }; @@ -574,7 +602,7 @@ mod tests { assert_eq!(vortex_array.dtype(), &dtype); // And the values must round-trip back out through the export plugin. - let exported = session.execute_arrow(vortex_array, Some(&field), &mut ctx)?; + let exported = vortex_array.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.len(), 5); let ree = exported .as_any() @@ -636,7 +664,7 @@ mod tests { // Arrow's dictionary values data type does carry the nested element field, so the // extension survives the export as well. let mut ctx = vortex_session.create_execution_ctx(); - let exported = session.execute_arrow(array, Some(&field), &mut ctx)?; + let exported = array.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), field.data_type()); assert_eq!(exported.to_data(), expected); let values = exported.as_any_dictionary().values().as_list::(); @@ -665,8 +693,7 @@ mod tests { assert_eq!(array.dtype(), &session.from_arrow_schema(&schema)?); let ids = array.as_::().unmasked_field_by_name("id")?.clone(); assert_eq!(ids.dtype(), &uuid_dtype(false)); - let exported = session.execute_arrow( - ids, + let exported = ids.execute_arrow( Some(&uuid_field), &mut vortex_session.create_execution_ctx(), )?; @@ -718,7 +745,7 @@ mod tests { assert_eq!(by_nullability.dtype(), array.dtype()); let mut ctx = vortex_session.create_execution_ctx(); for imported in [array, by_nullability] { - let exported = session.execute_arrow(imported, None, &mut ctx)?; + let exported = imported.execute_arrow(None, &mut ctx)?; assert_eq!(exported.to_data(), arrow_struct.to_data()); } Ok(()) @@ -767,9 +794,7 @@ mod tests { ); assert_eq!(array.len(), 2); let mut ctx = vortex_session.create_execution_ctx(); - let exported = vortex_session - .arrow() - .execute_arrow(array, None, &mut ctx)?; + let exported = array.execute_arrow(None, &mut ctx)?; let ids = exported.as_struct().column(0).as_fixed_size_binary(); assert_eq!(ids.len(), 2); assert_eq!(ids.null_count(), 0); @@ -826,7 +851,7 @@ mod tests { assert!(map_dtype.key_dtype().is_extension()); assert!(map_dtype.value_dtype().is_extension()); - let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let exported = vortex.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), field.data_type()); assert_eq!(exported.as_map().value_offsets(), &[0, 2]); assert_eq!(exported.to_data(), expected); diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java index 13776acf7d0..047c610d170 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java @@ -70,6 +70,23 @@ void readsAndWritesThroughTheV1FileFormat() { .toList()); } + @Test + void aSplitConcurrencyReadsEveryRowThroughV1AndV2() { + Path output = write("split_concurrency", false); + try { + for (String useV1 : new String[] {"", "vortex"}) { + spark.conf().set("spark.sql.sources.useV1SourceList", useV1); + Dataset read = spark.read() + .format("vortex") + .option("vortex.splitConcurrency", "1") + .load(output.toString()); + assertEquals(780L, read.selectExpr("sum(id)").first().getLong(0), "useV1SourceList=" + useV1); + } + } finally { + spark.conf().set("spark.sql.sources.useV1SourceList", "vortex"); + } + } + @Test void readsMissingColumnsWithCodegen() { Path output = tempDir.resolve("v1_merged"); diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 8de3a6ec03c..332d13c8d42 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -141,11 +141,11 @@ impl ScalarFnVTable for Zip { let return_dtype = zip_return_dtype(if_true.dtype(), if_false.dtype())?; if mask.all_true() { - return if_true.cast(return_dtype)?.execute(ctx); + return if_true.cast(return_dtype); } if mask.all_false() { - return if_false.cast(return_dtype)?.execute(ctx); + return if_false.cast(return_dtype); } if !if_true.is_canonical() || !if_false.is_canonical() { @@ -345,6 +345,7 @@ pub(crate) fn zip_validity( mod tests { use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -353,6 +354,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; @@ -437,6 +439,30 @@ mod tests { assert_eq!(result.dtype(), if_true.dtype()); } + #[rstest] + #[case::all_true(Some(true), 10)] + #[case::all_false(Some(false), 20)] + #[case::null_mask(None, 20)] + fn test_zip_constant_mask_preserves_constant( + #[case] mask: Option, + #[case] expected: i32, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let if_true = ConstantArray::new(10i32, 4).into_array(); + let if_false = ConstantArray::new(Scalar::primitive(20i32, nullability), 4).into_array(); + let result = ConstantArray::new(mask, 4) + .into_array() + .zip(if_true, if_false)? + .execute::(&mut ctx)? + .into_array(); + + assert!(result.is::()); + let expected = ConstantArray::new(Scalar::primitive(expected, nullability), 4); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + #[test] fn test_invalid_lengths() { let mask = Mask::new_false(4); diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index c948380a602..2341f477a02 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -39,13 +39,8 @@ use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -#[expect( - deprecated, - reason = "benchmark comparing deprecated method with new one" -)] -use vortex_arrow::ArrowArrayExecutor; -use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_buffer::Buffer; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_fsst::fsst_compress; @@ -68,12 +63,16 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() - .with::(); + .with::(); vortex_decimal_byte_parts::initialize(&session); vortex_fsst::initialize(&session); vortex_onpair::initialize(&session); session.arrays().register(Zstd); + + // The Arrow session is otherwise created on its first use, which would be inside the first + // benchmark that exports to Arrow, and so part of its measurement. + session.arrow(); + session }); @@ -426,12 +425,7 @@ fn string_export(bencher: Bencher, case: StringExportCase) { bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) - .bench_values(|(array, mut ctx)| { - SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut ctx) - .unwrap() - }); + .bench_values(|(array, mut ctx)| array.execute_arrow(Some(&field), &mut ctx).unwrap()); } /// Measures a direct append to an offset builder. @@ -629,22 +623,7 @@ fn decimal_export(bencher: Bencher, case: DecimalExportCase) { bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) - .bench_values(|(array, mut ctx)| { - SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut ctx) - .unwrap() - }); -} - -#[divan::bench] -fn to_arrow_array(bencher: Bencher) { - bencher - .with_inputs(|| (array(), SESSION.create_execution_ctx())) - .bench_values(|(array, mut ctx)| { - #[expect(deprecated, reason = "benchmarking deprecated code path")] - array.execute_arrow(None, &mut ctx).unwrap() - }); + .bench_values(|(array, mut ctx)| array.execute_arrow(Some(&field), &mut ctx).unwrap()); } #[allow(non_snake_case)] @@ -652,10 +631,5 @@ fn to_arrow_array(bencher: Bencher) { fn ArrowExportVTable_execute_arrow(bencher: Bencher) { bencher .with_inputs(|| (array(), SESSION.create_execution_ctx())) - .bench_values(|(array, mut ctx)| { - SESSION - .arrow() - .execute_arrow(array, None, &mut ctx) - .unwrap() - }) + .bench_values(|(array, mut ctx)| array.execute_arrow(None, &mut ctx).unwrap()) } diff --git a/vortex-arrow/src/datum.rs b/vortex-arrow/src/datum.rs index 9f20563e461..145a6b5176e 100644 --- a/vortex-arrow/src/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::ArrowSessionExt; +use crate::ExecuteArrowExt; /// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute. /// @@ -33,17 +34,14 @@ pub struct Datum { impl Datum { /// Create a new [`Datum`] from an [`ArrayRef`], which can then be passed to Arrow compute. pub fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let session = ctx.session().clone(); if array.is::() { Ok(Self { - array: session - .arrow() - .execute_arrow(array.slice(0..1)?, None, ctx)?, + array: array.slice(0..1)?.execute_arrow(None, ctx)?, is_scalar: true, }) } else { Ok(Self { - array: session.arrow().execute_arrow(array.clone(), None, ctx)?, + array: array.clone().execute_arrow(None, ctx)?, is_scalar: false, }) } @@ -52,9 +50,8 @@ impl Datum { /// Create a new [`Datum`] from an `DynArrayData`, which can then be passed to Arrow compute. /// This not try and convert the array to a scalar if it is constant. pub fn try_new_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let session = ctx.session().clone(); Ok(Self { - array: session.arrow().execute_arrow(array.clone(), None, ctx)?, + array: array.clone().execute_arrow(None, ctx)?, is_scalar: false, }) } @@ -64,7 +61,6 @@ impl Datum { target_datatype: &DataType, ctx: &mut ExecutionCtx, ) -> VortexResult { - let session = ctx.session().clone(); let target_field = Field::new( String::new(), target_datatype.clone(), @@ -73,18 +69,12 @@ impl Datum { if array.is::() { Ok(Self { - array: session.arrow().execute_arrow( - array.slice(0..1)?, - Some(&target_field), - ctx, - )?, + array: array.slice(0..1)?.execute_arrow(Some(&target_field), ctx)?, is_scalar: true, }) } else { Ok(Self { - array: session - .arrow() - .execute_arrow(array.clone(), Some(&target_field), ctx)?, + array: array.clone().execute_arrow(Some(&target_field), ctx)?, is_scalar: false, }) } diff --git a/vortex-arrow/src/dtype.rs b/vortex-arrow/src/dtype.rs index 660babe06f0..57e765272e7 100644 --- a/vortex-arrow/src/dtype.rs +++ b/vortex-arrow/src/dtype.rs @@ -328,9 +328,10 @@ mod test { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; + use vortex_session::VortexSession; use super::*; - use crate::ArrowSession; + use crate::ArrowSessionExt; #[test] fn test_dtype_conversion_success() { @@ -456,7 +457,8 @@ mod test { let schema_nonnull = DType::Struct(the_struct, Nullability::NonNullable); assert_eq!( - ArrowSession::default() + VortexSession::empty() + .arrow() .to_arrow_schema(&schema_nonnull) .unwrap(), Schema::new(Fields::from(vec![ @@ -630,7 +632,12 @@ mod test { let field = Field::new("c0", data_type, true); assert!(from_arrow_field_naive(&field).is_err()); let schema = Schema::new(vec![field]); - assert!(ArrowSession::default().from_arrow_schema(&schema).is_err()); + assert!( + VortexSession::empty() + .arrow() + .from_arrow_schema(&schema) + .is_err() + ); } #[test] diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index cbc8b75d33c..e3bdf595cd2 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -205,7 +205,7 @@ mod tests { use vortex_onpair::onpair_compress; use vortex_session::VortexSession; - use crate::ArrowSessionExt; + use crate::ExecuteArrowExt; static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty(); @@ -270,16 +270,9 @@ mod tests { TraceOptions { resolution: TraceResolution::Attempts, }, - || { - SESSION - .arrow() - .execute_arrow(source, Some(&field), &mut ctx) - }, + || source.execute_arrow(Some(&field), &mut ctx), )?; - let expected = - SESSION - .arrow() - .execute_arrow(input.slice(range)?, Some(&field), &mut ctx)?; + let expected = input.slice(range)?.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(trace.output.to_data(), expected.to_data()); // CodSpeed disables tracing; normal tests check that export avoids executing // the root encoding to an intermediate string view. @@ -308,9 +301,7 @@ mod tests { let masked = MaskFn::try_new(varbin.into_array(), mask.into_array())?.into_array(); let field = Field::new("s", DataType::Utf8, true); - let arrow = SESSION - .arrow() - .execute_arrow(masked, Some(&field), &mut ctx)?; + let arrow = masked.execute_arrow(Some(&field), &mut ctx)?; let strings = arrow.as_string::(); assert_eq!(strings.len(), 3); @@ -357,9 +348,9 @@ mod tests { ) { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("test_field", target_dtype.clone(), true); - let arrow = SESSION - .arrow() - .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx) + let arrow = vortex_array + .into_array() + .execute_arrow(Some(&field), &mut ctx) .unwrap(); assert_eq!(arrow.data_type(), &target_dtype); @@ -405,9 +396,9 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("test_field", target_dtype.clone(), true); - let arrow = SESSION - .arrow() - .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx) + let arrow = vortex_array + .into_array() + .execute_arrow(Some(&field), &mut ctx) .unwrap(); assert_eq!(arrow.data_type(), &target_dtype); @@ -441,10 +432,9 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("test_field", target_dtype, true); - let result = - SESSION - .arrow() - .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx); + let result = vortex_array + .into_array() + .execute_arrow(Some(&field), &mut ctx); assert!(result.is_err()); } @@ -465,9 +455,7 @@ mod tests { &mut ctx, ); let field = Field::new("test_field", target_dtype, false); - let result = SESSION - .arrow() - .execute_arrow(array.into_array(), Some(&field), &mut ctx); + let result = array.into_array().execute_arrow(Some(&field), &mut ctx); assert!(result.is_err()); } @@ -488,9 +476,7 @@ mod tests { &mut ctx, ); let field = Field::new("test_field", target_dtype.clone(), false); - let arrow = SESSION - .arrow() - .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + let arrow = array.into_array().execute_arrow(Some(&field), &mut ctx)?; assert_eq!(arrow.len(), 2); assert_eq!(arrow.null_count(), 0); @@ -520,9 +506,7 @@ mod tests { &mut ctx, ); let field = Field::new("test_field", target_dtype, true); - let arrow = SESSION - .arrow() - .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + let arrow = array.into_array().execute_arrow(Some(&field), &mut ctx)?; assert_eq!(arrow.len(), 3); let values = match arrow.data_type() { @@ -542,11 +526,9 @@ mod tests { fn incompatible_vortex_dtype_returns_error(#[case] target_dtype: DataType) { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("test_field", target_dtype, true); - let result = SESSION.arrow().execute_arrow( - PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), - Some(&field), - &mut ctx, - ); + let result = PrimitiveArray::from_iter([1i32, 2, 3]) + .into_array() + .execute_arrow(Some(&field), &mut ctx); assert!(result.is_err()); } @@ -561,9 +543,7 @@ mod tests { .filter(Mask::from_iter([true, false, false]))?; let mut ctx = SESSION.create_execution_ctx(); - let arrow = SESSION - .arrow() - .execute_arrow(filtered.into_array(), None, &mut ctx)?; + let arrow = filtered.into_array().execute_arrow(None, &mut ctx)?; assert_eq!(arrow.as_string_view().value(0), "selected"); assert!( diff --git a/vortex-arrow/src/executor/byte_view.rs b/vortex-arrow/src/executor/byte_view.rs index e99980867fc..139512e4622 100644 --- a/vortex-arrow/src/executor/byte_view.rs +++ b/vortex-arrow/src/executor/byte_view.rs @@ -13,7 +13,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_buffer::Buffer; use vortex_error::VortexResult; -use crate::ArrowExporter; +use crate::ArrowExportOptions; use crate::CompactBuffers; use crate::dtype::from_arrow_data_type; use crate::null_buffer::to_null_buffer; @@ -53,7 +53,7 @@ pub fn execute_varbinview_to_arrow( pub(super) fn to_arrow_byte_view( array: ArrayRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { // First we cast the array into the desired ByteView type. @@ -65,7 +65,7 @@ pub(super) fn to_arrow_byte_view( let array = array.execute::(ctx)?; let varbinview = array.execute::(ctx)?; - if exporter.options().get_or_default::().0 { + if options.get_or_default::().0 { execute_varbinview_to_arrow::(&varbinview, ctx) } else { canonical_varbinview_to_arrow::(&varbinview, ctx) diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs index 912bdb8f7da..7fa55ddef4f 100644 --- a/vortex-arrow/src/executor/decimal.rs +++ b/vortex-arrow/src/executor/decimal.rs @@ -273,6 +273,7 @@ mod tests { use arrow_array::StructArray as ArrowStructArray; use arrow_buffer::i256; use arrow_schema::DataType; + use arrow_schema::Field; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -293,7 +294,7 @@ mod tests { use vortex_error::vortex_panic; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -347,9 +348,10 @@ mod tests { DecimalDType::new(19, 2), Validity::NonNullable, ); - let arrow = decimal_vortex - .into_array() - .execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?; + let arrow = decimal_vortex.into_array().execute_arrow( + Some(&Field::new("", DataType::Decimal128(19, 2), false)), + &mut ctx, + )?; assert_arrow_eq( &arrow, Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)?, @@ -386,7 +388,8 @@ mod tests { decimal.append_value(12); let decimal = decimal.finish(); - let actual = decimal.execute_arrow(Some(&data_type), &mut ctx)?; + let field = Field::new("", data_type.clone(), false); + let actual = decimal.execute_arrow(Some(&field), &mut ctx)?; let expected: arrow_array::ArrayRef = match data_type { DataType::Decimal32(precision, scale) => Arc::new( Decimal32Array::from(vec![10, 11, 12]) @@ -439,8 +442,9 @@ mod tests { .execute::(&mut ctx)? .into_array(); - let actual = array.execute_arrow(Some(&data_type), &mut ctx)?; - let expected = canonical.execute_arrow(Some(&data_type), &mut ctx)?; + let field = Field::new("", data_type.clone(), nullable); + let actual = array.execute_arrow(Some(&field), &mut ctx)?; + let expected = canonical.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(actual.data_type(), &data_type); assert_eq!(&actual, &expected); @@ -475,7 +479,8 @@ mod tests { ); let sliced = widen(values, 18)?.slice(1..4)?; - let arrow = sliced.execute_arrow(Some(&DataType::Decimal128(18, 2)), &mut ctx)?; + let field = Field::new("", DataType::Decimal128(18, 2), true); + let arrow = sliced.execute_arrow(Some(&field), &mut ctx)?; assert_arrow_eq( &arrow, Decimal128Array::from(vec![None, Some(30), Some(40)]) @@ -522,8 +527,8 @@ mod tests { let values = DecimalArray::new(buffer![1i128, i128::MAX], DecimalDType::new(9, 2), validity); - let arrow = - widen(values, 18)?.execute_arrow(Some(&DataType::Decimal64(18, 2)), &mut ctx)?; + let field = Field::new("", DataType::Decimal64(18, 2), true); + let arrow = widen(values, 18)?.execute_arrow(Some(&field), &mut ctx)?; assert_arrow_eq( &arrow, Decimal64Array::from(expected).with_precision_and_scale(18, 2)?, @@ -540,10 +545,11 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let values = DecimalArray::new(buffer![i128::MAX, 0], DecimalDType::new(9, 2), validity); + let field = Field::new("", DataType::Decimal64(9, 2), values.dtype().is_nullable()); let err = values .into_array() - .execute_arrow(Some(&DataType::Decimal64(9, 2)), &mut ctx) + .execute_arrow(Some(&field), &mut ctx) .expect_err("a valid value outside the target width must fail"); assert!( err.to_string() diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 343c08357af..e4fc05bf4bc 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -23,7 +23,7 @@ use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use crate::ArrowExporter; +use crate::ArrowExportOptions; use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_dictionary`] requires for export. @@ -41,23 +41,23 @@ pub(super) fn to_arrow_dictionary( array: ArrayRef, codes_type: &DataType, values_type: &DataType, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, 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, exporter, ctx), + Ok(dict) => return dict_to_dict(dict, codes_type, values_type, options, ctx), Err(array) => array, }; let array = match array.try_downcast::() { - Ok(constant) => return constant_to_dict(constant, codes_type, values_type, exporter, ctx), + Ok(constant) => return constant_to_dict(constant, codes_type, values_type, options, ctx), Err(array) => array, }; // Otherwise, we should try and build a dictionary. // Arrow hides this functionality inside the cast module! - let array = execute_arrow_naive(array, Some(values_type), exporter, ctx)?; + let array = execute_arrow_naive(array, Some(values_type), options, ctx)?; arrow_cast::cast( &array, &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), @@ -70,7 +70,7 @@ fn constant_to_dict( array: ConstantArray, codes_type: &DataType, values_type: &DataType, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -84,7 +84,7 @@ fn constant_to_dict( let values = execute_arrow_naive( ConstantArray::new(scalar.clone(), 1).into_array(), Some(values_type), - exporter, + options, ctx, )?; let codes = zeroed_codes_array(codes_type, len)?; @@ -96,11 +96,11 @@ fn dict_to_dict( array: DictArray, codes_type: &DataType, values_type: &DataType, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - 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)?; + let codes = execute_arrow_naive(array.codes().clone(), Some(codes_type), options, ctx)?; + let values = execute_arrow_naive(array.values().clone(), Some(values_type), options, ctx)?; make_dict_array(codes_type, codes, values) } @@ -164,6 +164,7 @@ mod tests { use arrow_array::types::UInt8Type; use arrow_array::types::UInt32Type; use arrow_schema::DataType; + use arrow_schema::Field; use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -178,7 +179,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; use crate::executor::dictionary::ConstantArray; use crate::executor::dictionary::DictArray; @@ -192,7 +193,8 @@ mod tests { array: vortex_array::ArrayRef, dt: &DataType, ) -> VortexResult { - array.execute_arrow(Some(dt), &mut SESSION.create_execution_ctx()) + let field = Field::new("", dt.clone(), array.dtype().is_nullable()); + array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx()) } fn dict_basic_input() -> vortex_array::ArrayRef { diff --git a/vortex-arrow/src/executor/fixed_size_list.rs b/vortex-arrow/src/executor/fixed_size_list.rs index b830af710dd..5424afa2b35 100644 --- a/vortex-arrow/src/executor/fixed_size_list.rs +++ b/vortex-arrow/src/executor/fixed_size_list.rs @@ -14,14 +14,15 @@ use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrowExporter; +use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::validity::to_arrow_null_buffer; pub(super) fn to_arrow_fixed_list( array: ArrayRef, list_size: i32, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -32,25 +33,19 @@ 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, - exporter, - ctx, - ); + return list_to_list(&array.into_owned(), elements_field, list_size, options, 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, exporter, ctx) + list_to_list(&fixed_size_list, elements_field, list_size, options, ctx) } fn list_to_list( array: &FixedSizeListArray, elements_field: &FieldRef, list_size: i32, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -60,8 +55,11 @@ fn list_to_list( list_size ); - let elements = - exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; + let elements = array.elements().clone().execute_arrow_with_options( + Some(elements_field.as_ref()), + options, + 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 965b5e6a09a..5f654dd4e61 100644 --- a/vortex-arrow/src/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -29,7 +29,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrowExporter; +use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::validity::to_arrow_null_buffer; /// Matches the encodings [`to_arrow_list`] requires for export. @@ -48,7 +49,7 @@ impl Matcher for ArrowListExportable { pub(super) fn to_arrow_list( array: ArrayRef, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -61,7 +62,7 @@ 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, exporter, ctx); + return list_to_list::(&list.into_owned(), elements_field, options, ctx); } // Otherwise the array is canonical: a ListViewArray, which we rebuild to ZCTL if needed. @@ -80,7 +81,7 @@ pub(super) fn to_arrow_list( } else { list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? }; - list_view_zctl::(zctl, elements_field, exporter, ctx) + list_view_zctl::(zctl, elements_field, options, ctx) } #[allow(rustdoc::broken_intra_doc_links)] @@ -88,7 +89,7 @@ pub(super) fn to_arrow_list( fn list_to_list( array: &ListArray, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { // We must cast the offsets to the required offset type. @@ -103,8 +104,11 @@ fn list_to_list( // checked cast keeps every value, so the order stays the same. let offsets = unsafe { offsets.into_arrow_offset_buffer_unchecked() }; - let elements = - exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; + let elements = array.elements().clone().execute_arrow_with_options( + Some(elements_field.as_ref()), + options, + ctx, + )?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Cannot convert to non-nullable Arrow array with null elements" @@ -124,14 +128,17 @@ fn list_to_list( fn list_view_zctl( array: ListViewArray, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { assert!(array.is_zero_copy_to_list()); if array.is_empty() { - let elements = - exporter.execute_arrow(array.elements().clone(), Some(elements_field.as_ref()), ctx)?; + let elements = array.elements().clone().execute_arrow_with_options( + Some(elements_field.as_ref()), + options, + ctx, + )?; return Ok(Arc::new(GenericListArray::::new( Arc::clone(elements_field), OffsetBuffer::new_empty(), @@ -180,7 +187,8 @@ fn list_view_zctl( }); // Extract the elements array. - let elements = exporter.execute_arrow(elements, Some(elements_field.as_ref()), ctx)?; + let elements = + elements.execute_arrow_with_options(Some(elements_field.as_ref()), options, ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Cannot convert to non-nullable Arrow array with null elements" @@ -228,7 +236,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; use crate::executor::list::ListArray; use crate::executor::list::ListViewArray; @@ -254,9 +262,10 @@ mod tests { let chunked = ChunkedArray::try_new(vec![first.into_array(), second.into_array()], dtype)?; let field = Field::new("item", DataType::Int32, false); - let actual = chunked - .into_array() - .execute_arrow(Some(&DataType::List(field.clone().into())), &mut ctx)?; + let actual = chunked.into_array().execute_arrow( + Some(&Field::new("", DataType::List(field.clone().into()), true)), + &mut ctx, + )?; let expected = GenericListArray::::try_new( field.into(), @@ -291,7 +300,7 @@ mod tests { let arrow_dt = DataType::List(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt.clone(), true)), &mut ctx)?; // Verify the type is correct. assert_eq!(arrow_array.data_type(), &arrow_dt); @@ -357,7 +366,7 @@ mod tests { let arrow_dt = DataType::LargeList(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt.clone(), true)), &mut ctx)?; // Verify the type is correct. assert_eq!(arrow_array.data_type(), &arrow_dt); @@ -407,7 +416,7 @@ mod tests { let arrow_dt = DataType::List(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx)?; let list = arrow_array .as_any() @@ -456,7 +465,7 @@ mod tests { let field = Field::new("item", DataType::Int32, false); let arrow_dt = DataType::List(field.into()); - let arrow_array = sliced.execute_arrow(Some(&arrow_dt), &mut ctx)?; + let arrow_array = sliced.execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx)?; let arrow_list = arrow_array .as_any() @@ -511,7 +520,7 @@ mod tests { let field = Field::new("item", DataType::Int32, false); let arrow_dt = DataType::List(field.into()); - let arrow_array = masked.execute_arrow(Some(&arrow_dt), &mut ctx)?; + let arrow_array = masked.execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx)?; let arrow_list = arrow_array .as_any() @@ -557,7 +566,7 @@ mod tests { let arrow_dt = DataType::List(Field::new("item", DataType::Int32, false).into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt.clone(), true)), &mut ctx)?; assert_eq!(arrow_array.len(), 0); assert_eq!(arrow_array.data_type(), &arrow_dt); Ok(()) diff --git a/vortex-arrow/src/executor/list_view.rs b/vortex-arrow/src/executor/list_view.rs index 2bae37f5835..72b8587af9e 100644 --- a/vortex-arrow/src/executor/list_view.rs +++ b/vortex-arrow/src/executor/list_view.rs @@ -24,13 +24,14 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use crate::ArrowExporter; +use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::validity::to_arrow_null_buffer; pub(super) fn to_arrow_list_view( array: ArrayRef, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -68,13 +69,13 @@ pub(super) fn to_arrow_list_view( array }; - list_view_to_list_view::(array, elements_field, exporter, ctx) + list_view_to_list_view::(array, elements_field, options, ctx) } fn list_view_to_list_view( array: ListViewArray, elements_field: &FieldRef, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let ListViewDataParts { @@ -86,7 +87,8 @@ fn list_view_to_list_view( } = array.into_data_parts(); let n_elements = elements.len(); - let elements = exporter.execute_arrow(elements, Some(elements_field.as_ref()), ctx)?; + let elements = + elements.execute_arrow_with_options(Some(elements_field.as_ref()), options, ctx)?; vortex_ensure!( elements_field.is_nullable() || elements.null_count() == 0, "Elements field is non-nullable but elements array contains nulls" @@ -142,7 +144,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; use crate::executor::list_view::ListViewArray; use crate::executor::list_view::PrimitiveArray; @@ -174,7 +176,7 @@ mod tests { let arrow_dt = DataType::ListView(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx)?; let listview = arrow_array .as_any() @@ -216,7 +218,7 @@ mod tests { let arrow_dt = DataType::ListView(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt.clone(), true)), &mut ctx)?; // Verify the type is correct. assert_eq!(arrow_array.data_type(), &arrow_dt); @@ -280,7 +282,7 @@ mod tests { let arrow_dt = DataType::LargeListView(field.into()); let arrow_array = list_array .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt.clone(), true)), &mut ctx)?; // Verify the type is correct. assert_eq!(arrow_array.data_type(), &arrow_dt); diff --git a/vortex-arrow/src/executor/map.rs b/vortex-arrow/src/executor/map.rs index 4e5646cd7d9..b41c6a11328 100644 --- a/vortex-arrow/src/executor/map.rs +++ b/vortex-arrow/src/executor/map.rs @@ -20,7 +20,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use crate::ArrowExporter; +use crate::ArrowExportOptions; use crate::executor::list::to_arrow_list; /// Converts a Vortex Map array into an Arrow [`MapArray`](ArrowMapArray). @@ -32,7 +32,7 @@ pub(super) fn to_arrow_map( array: ArrayRef, entries_field: &FieldRef, keys_sorted: bool, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure!( @@ -53,7 +53,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, exporter, ctx)?; + let entries_list = to_arrow_list::(entries, entries_field, options, ctx)?; vortex_ensure_eq!( entries_list.data_type(), &entries_list_type, @@ -113,6 +113,7 @@ mod tests { use vortex_mask::Mask; use vortex_session::VortexSession; + use crate::ExecuteArrowExt; use crate::convert::from_arrow_map; use crate::session::ArrowSessionExt as _; @@ -188,7 +189,7 @@ mod tests { )? ); - let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let exported = vortex.execute_arrow(Some(&field), &mut ctx)?; let map = exported.as_map(); assert_eq!(map.value_offsets(), &[0, 2, 2, 2, 4]); assert!(map.is_valid(0)); @@ -393,7 +394,7 @@ mod tests { assert_eq!(entries.size_at(1, &mut ctx), 1); assert_eq!(entries.elements().len(), 3); - let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let exported = vortex.execute_arrow(Some(&field), &mut ctx)?; let map = exported.as_map(); assert_eq!(map.value_offsets(), &[0, 2, 3]); @@ -421,7 +422,6 @@ mod tests { fn filtered_map_with_entry_gap_exports_correct_rows() -> VortexResult<()> { let vortex_session = VortexSession::empty(); let mut ctx = vortex_session.create_execution_ctx(); - let session = vortex_session.arrow(); let map_dtype = MapDType::try_new( DType::Primitive(PType::I32, NonNullable), DType::Utf8(Nullable), @@ -458,7 +458,7 @@ mod tests { let filtered = source.filter(Mask::from_iter([true, false, true]))?; let field = i32_utf8_map_field(false, true); - let exported = session.execute_arrow(filtered, Some(&field), &mut ctx)?; + let exported = filtered.execute_arrow(Some(&field), &mut ctx)?; let map = exported.as_map(); assert_eq!(map.value_offsets(), &[0, 2, 3]); @@ -500,8 +500,8 @@ mod tests { )?; let vortex = session.from_arrow_array(Arc::new(arrow), &source_field)?; - let error = session - .execute_arrow(vortex, Some(&target_field), &mut ctx) + let error = vortex + .execute_arrow(Some(&target_field), &mut ctx) .unwrap_err(); assert!(error.to_string().contains("keys_sorted=true")); diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 7e6918e021f..803f038962f 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -1,10 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#![expect( - deprecated, - reason = "This module defines and implements a deprecated trait `ArrowArrayExecutor`" -)] +//! Execution of Vortex arrays into Arrow arrays. +//! +//! [`ExecuteArrowExt`] is the entry point for every export of a Vortex array to Arrow. Most targets +//! go to the canonical export, which maps the array to the Arrow type of the target. Only a target +//! with an Arrow extension name goes to the export plugins of the [`ArrowSession`] first. Each +//! submodule exports one family of Arrow types. +//! +//! [`ArrowSession`]: crate::ArrowSession pub mod bool; mod byte; @@ -23,27 +27,32 @@ mod temporal; mod validity; use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::RecordBatch; -use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::FieldRef; -use arrow_schema::Schema; -use itertools::Itertools; +use tracing::debug; +use tracing::trace; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::List; use vortex_array::arrays::VarBin; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDTypeRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use vortex_session::registry::Id; -use crate::ArrowExporter; +use crate::ArrowExport; +use crate::ArrowExportOptions; use crate::executor::bool::to_arrow_bool; use crate::executor::byte::to_arrow_byte_array; use crate::executor::byte_view::to_arrow_byte_view; @@ -60,74 +69,283 @@ use crate::executor::struct_::to_arrow_struct; use crate::executor::temporal::to_arrow_date; use crate::executor::temporal::to_arrow_time; use crate::executor::temporal::to_arrow_timestamp; +use crate::ext_metadata::decode_arrow_ext_metadata; use crate::session::ArrowSessionExt; +use crate::session::ensure_ext_metadata; -/// Trait for executing a Vortex array to produce an Arrow array. -#[deprecated(note = "Use an `ArrowSession` to perform conversions to/from Arrow arrays")] -pub trait ArrowArrayExecutor: Sized { - /// Execute the array to produce an Arrow array. +/// Execution of a Vortex array into an Arrow array, for every type that implements [`IntoArray`]. +/// +/// An export uses the [`ArrowSession`] of its [`ExecutionCtx`] for every decision: the export +/// plugins, the checks of the target, and the Arrow types of nested arrays. It takes no other +/// session, so one export can never mix the plugins of two sessions. +/// +/// ``` +/// # use arrow_array::Array; +/// # use vortex_array::IntoArray; +/// # use vortex_array::VortexSessionExecute; +/// # use vortex_array::arrays::VarBinViewArray; +/// use vortex_arrow::ExecuteArrowExt; +/// # use vortex_session::VortexSession; +/// +/// let session = VortexSession::empty(); +/// let array = VarBinViewArray::from_iter_str(["a", "b"]).into_array(); +/// let arrow = array.execute_arrow(None, &mut session.create_execution_ctx())?; +/// assert_eq!(arrow.len(), 2); +/// # Ok::<(), vortex_error::VortexError>(()) +/// ``` +/// +/// [`ArrowSession`]: crate::ArrowSession +pub trait ExecuteArrowExt: IntoArray + Sized { + /// Execute this array into an Arrow array of the type of `target`, with the default + /// [`ArrowExportOptions`]. + /// + /// A `target` with an Arrow extension name goes to the plugins registered for that name, in + /// registration order. If none of them accepts the array, but the name is the Vortex ID of the + /// array, the array exports as its storage (see [`ArrowSession`]). /// - /// If a [`DataType`] is given, the array will be converted to the desired Arrow type. - /// If `None`, the array's preferred (cheapest) Arrow type will be used. - #[deprecated(note = "Use an `ArrowSession` to perform conversions to/from Arrow arrays")] + /// Without a `target`, the array exports to the field from + /// [`ArrowSession::default_arrow_export_field`], so it fails for the same types. + /// + /// [`ArrowSession`]: crate::ArrowSession + /// [`ArrowSession::default_arrow_export_field`]: crate::ArrowSession::default_arrow_export_field fn execute_arrow( self, - data_type: Option<&DataType>, + target: Option<&Field>, ctx: &mut ExecutionCtx, - ) -> VortexResult; + ) -> VortexResult { + self.execute_arrow_with_options(target, &ArrowExportOptions::default(), ctx) + } - /// Execute the array to produce an Arrow `RecordBatch` with the given schema. - #[deprecated(note = "Use an `ArrowSession` to perform conversions to/from Arrow arrays")] - fn execute_record_batch( + /// Execute this array into an Arrow array with `options`, which reach every nested array and + /// every export plugin. + /// + /// Otherwise, this behaves like [`execute_arrow`](Self::execute_arrow). + /// + /// ``` + /// # use arrow_array::Array; + /// # use vortex_array::IntoArray; + /// # use vortex_array::VortexSessionExecute; + /// # use vortex_array::arrays::VarBinViewArray; + /// # use vortex_arrow::ArrowExportOptions; + /// # use vortex_arrow::CompactBuffers; + /// use vortex_arrow::ExecuteArrowExt; + /// # use vortex_session::VortexSession; + /// + /// let session = VortexSession::empty(); + /// let array = VarBinViewArray::from_iter_str(["a", "b"]).into_array(); + /// let options = ArrowExportOptions::default().with(CompactBuffers(false)); + /// let arrow = + /// array.execute_arrow_with_options(None, &options, &mut session.create_execution_ctx())?; + /// assert_eq!(arrow.len(), 2); + /// # Ok::<(), vortex_error::VortexError>(()) + /// ``` + fn execute_arrow_with_options( self, - schema: &Schema, + target: Option<&Field>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, - ) -> VortexResult { - let array = self.execute_arrow(Some(&DataType::Struct(schema.fields.clone())), ctx)?; - Ok(RecordBatch::from(array.as_struct())) + ) -> VortexResult { + execute_arrow_array(self.into_array(), target, options, ctx) } +} - /// Execute the array to produce Arrow `RecordBatch`'s with the given schema. - #[deprecated(note = "Use an `ArrowSession` to perform conversions to/from Arrow arrays")] - fn execute_record_batches( - self, - schema: &Schema, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; +impl ExecuteArrowExt for T {} + +/// Export `array` to `target`, or to the default field of its dtype if there is no target. +fn execute_arrow_array( + array: ArrayRef, + target: Option<&Field>, + options: &ArrowExportOptions, + 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 = ctx + .session() + .arrow() + .default_arrow_export_field(array.dtype())?; + &arrow_field + } + }; + + if let Some(arrow_ext_name) = target_field.extension_type_name() { + return execute_arrow_extension(array, arrow_ext_name, target_field, options, ctx); + } + + execute_arrow_naive(array, target.map(Field::data_type), options, ctx) } -#[expect(deprecated, reason = "backward compatibility")] -impl ArrowArrayExecutor for ArrayRef { - fn execute_arrow( - self, - data_type: Option<&DataType>, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let target = data_type.map(|dt| Field::new("", dt.clone(), self.dtype().is_nullable())); - let session = ctx.session().clone(); - session.arrow().execute_arrow(self, target.as_ref(), ctx) +/// Export `array` to `target`, whose Arrow extension name is `arrow_ext_name`, through the first +/// export plugin that accepts it, or as its storage if `arrow_ext_name` is its Vortex ID. +#[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] +fn execute_arrow_extension( + array: ArrayRef, + arrow_ext_name: &str, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let arrow_ext_id = Id::new(arrow_ext_name); + let exporters = ctx.session().arrow().exporters(&arrow_ext_id); + + // 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 exporters.iter() { + trace!( + plugin = ?plugin, + extension_name = arrow_ext_name, + "probing plugin for converting Arrow array" + ); + + match plugin.execute_arrow(current, target, 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, + } } - fn execute_record_batches( - self, - schema: &Schema, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - self.to_array_iterator() - .map(|a| a?.execute_record_batch(schema, ctx)) - .try_collect() + // The extension name can be an alias, which the dtype registry resolves to the Vortex ID of its + // type, so the array can match either of them. + let named_vortex_id = ctx.session().arrow().known_vortex_id(&arrow_ext_id); + + if let Some(ext_dtype) = current + .dtype() + .as_extension_opt() + .filter(|ext_dtype| { + ext_dtype.id() == arrow_ext_id || Some(ext_dtype.id()) == named_vortex_id + }) + .cloned() + { + return execute_arrow_ext_storage(current, &ext_dtype, target, options, ctx); + } + + // An extension array of another type must not export to a field that names this Vortex ID. A + // plain array can, because a plugin such as `Vector` exports its storage to the same field that + // it received. + if let (Some(ext_dtype), Some(named_vortex_id)) = + (current.dtype().as_extension_opt(), named_vortex_id) + { + vortex_bail!( + "Arrow field {:?} names the extension type {named_vortex_id}, so the array must have that type, got {ext_dtype}", + target.name() + ); + } + + debug!( + extension_id = arrow_ext_name, + data_type = ?target.data_type(), + "unsupported Arrow extension type encountered, falling back to naive execution" + ); + + execute_arrow_naive(current, Some(target.data_type()), options, ctx) +} + +/// Export `array`, whose type is `ext_dtype`, to `target`, whose extension name is the Vortex ID of +/// `ext_dtype`. [`ArrowSession`](crate::ArrowSession) describes the checks on `target`. +fn execute_arrow_ext_storage( + array: ArrayRef, + ext_dtype: &ExtDTypeRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, +) -> VortexResult { + ensure_ext_metadata(target, &decode_arrow_ext_metadata(target)?, ext_dtype)?; + + let target_dtype = ctx + .session() + .arrow() + .from_arrow_datatype(target.data_type(), ext_dtype.nullability())?; + + // If the Arrow type of the target imports as the extension type itself, such as an Arrow + // timestamp, export the array to it. The executor of that Arrow type then checks the parameters + // of the type, such as the time unit of a timestamp. + if target_dtype.eq_ignore_nullability(array.dtype()) { + return execute_arrow_naive(array, Some(target.data_type()), options, ctx); + } + + vortex_ensure!( + eq_storage_layout(&target_dtype, ext_dtype.storage_dtype()), + "Arrow field {:?} must have an Arrow type that imports as the storage of {ext_dtype}, got {}", + target.name(), + target.data_type() + ); + + let storage = array + .execute::(ctx)? + .storage_array() + .clone(); + + execute_arrow_naive(storage, Some(target.data_type()), options, ctx) +} + +/// Returns whether `lhs` and `rhs` have the same storage layout, which means that they are equal, +/// ignoring nullability, after every nested extension type is replaced by its storage type. The +/// types at the top level are compared as they are. +/// +/// An export compares the type that its target imports as with the storage type this way, because +/// an import can lose the metadata of a nested extension type, such as the version of a UUID, or +/// not know its Vortex ID. The export of each nested field checks that field when it reaches it. +fn eq_storage_layout(lhs: &DType, rhs: &DType) -> bool { + match (lhs, rhs) { + (DType::List(lhs, _), DType::List(rhs, _)) => eq_nested_storage_layout(lhs, rhs), + (DType::FixedSizeList(lhs, lhs_size, _), DType::FixedSizeList(rhs, rhs_size, _)) => { + lhs_size == rhs_size && eq_nested_storage_layout(lhs, rhs) + } + (DType::Map(lhs, _), DType::Map(rhs, _)) => { + lhs.keys_sorted() == rhs.keys_sorted() + && eq_nested_storage_layout(&lhs.key_dtype(), &rhs.key_dtype()) + && eq_nested_storage_layout(&lhs.value_dtype(), &rhs.value_dtype()) + } + (DType::Struct(lhs, _), DType::Struct(rhs, _)) => { + lhs.names() == rhs.names() + && lhs + .fields() + .zip(rhs.fields()) + .all(|(lhs, rhs)| eq_nested_storage_layout(&lhs, &rhs)) + } + _ => lhs.eq_ignore_nullability(rhs), } } +/// [`eq_storage_layout`] for two nested types, which replaces their own extension types by their +/// storage types too. +fn eq_nested_storage_layout(lhs: &DType, rhs: &DType) -> bool { + eq_storage_layout(innermost_storage_dtype(lhs), innermost_storage_dtype(rhs)) +} + +/// Returns `dtype` without its extension types, which can wrap each other. +fn innermost_storage_dtype(mut dtype: &DType) -> &DType { + while let Some(ext_dtype) = dtype.as_extension_opt() { + dtype = ext_dtype.storage_dtype(); + } + + dtype +} + /// Execute an arbitrary Vortex array into an Arrow array, dispatched by [`DataType`]. This pathway /// is naive to any extension type information, and only seeks to map canonical Vortex arrays to /// some target Arrow physical encoding. /// -/// Public callers should go through the `ArrowSession` instead. +/// Public callers should go through [`ExecuteArrowExt`] instead. pub(crate) fn execute_arrow_naive( array: ArrayRef, data_type: Option<&DataType>, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -155,27 +373,25 @@ 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, exporter, ctx), - DataType::Utf8View => to_arrow_byte_view::(array, exporter, ctx), + DataType::BinaryView => to_arrow_byte_view::(array, options, ctx), + DataType::Utf8View => to_arrow_byte_view::(array, options, ctx), // TODO(joe): pass down preferred - DataType::List(elements_field) => { - to_arrow_list::(array, elements_field, exporter, ctx) - } + DataType::List(elements_field) => to_arrow_list::(array, elements_field, options, ctx), // TODO(joe): pass down preferred DataType::LargeList(elements_field) => { - to_arrow_list::(array, elements_field, exporter, ctx) + to_arrow_list::(array, elements_field, options, ctx) } // TODO(joe): pass down preferred DataType::FixedSizeList(elements_field, list_size) => { - to_arrow_fixed_list(array, *list_size, elements_field, exporter, ctx) + to_arrow_fixed_list(array, *list_size, elements_field, options, ctx) } // TODO(joe): pass down preferred DataType::ListView(elements_field) => { - to_arrow_list_view::(array, elements_field, exporter, ctx) + to_arrow_list_view::(array, elements_field, options, ctx) } // TODO(joe): pass down preferred DataType::LargeListView(elements_field) => { - to_arrow_list_view::(array, elements_field, exporter, ctx) + to_arrow_list_view::(array, elements_field, options, ctx) } DataType::Struct(fields) => { let fields = if data_type.is_none() { @@ -183,11 +399,11 @@ pub(crate) fn execute_arrow_naive( } else { Some(fields) }; - to_arrow_struct(array, fields, exporter, ctx) + to_arrow_struct(array, fields, options, ctx) } // TODO(joe): pass down preferred DataType::Dictionary(codes_type, values_type) => { - to_arrow_dictionary(array, codes_type, values_type, exporter, ctx) + to_arrow_dictionary(array, codes_type, values_type, options, ctx) } dt @ DataType::Decimal32(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal64(..) => to_arrow_decimal(array, dt, ctx), @@ -195,13 +411,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, exporter, ctx) + to_arrow_run_end(array, ends_type.data_type(), values_type, options, 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, exporter, ctx) + to_arrow_map(array, entries_field, *keys_sorted, options, ctx) } DataType::FixedSizeBinary(_) | DataType::Duration(_) @@ -283,10 +499,7 @@ pub(crate) fn infer_nearest_arrow_field( let data_type = infer_nearest_arrow_type(array, ctx)?; return Ok(Field::new(name, data_type, array.dtype().is_nullable())); } - ctx.session() - .clone() - .arrow() - .to_arrow_field(name, array.dtype()) + ctx.session().arrow().to_arrow_field(name, array.dtype()) } #[cfg(test)] @@ -315,7 +528,7 @@ mod tests { use vortex_array::arrays::VarBinViewArray; use vortex_session::VortexSession; - use crate::ArrowSessionExt; + use crate::ExecuteArrowExt; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -372,7 +585,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("f", target.clone(), array.dtype().is_nullable()); - let result = SESSION.arrow().execute_arrow(array, Some(&field), &mut ctx); + let result = array.execute_arrow(Some(&field), &mut ctx); assert!( result.is_err(), @@ -397,7 +610,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let field = Field::new("f", target, array.dtype().is_nullable()); - let result = SESSION.arrow().execute_arrow(array, Some(&field), &mut ctx); + let result = array.execute_arrow(Some(&field), &mut ctx); assert_eq!(result.unwrap().to_data(), expected.to_data()); } diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index 7a735ce8bc7..04801861ea0 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -26,7 +26,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use crate::ArrowExporter; +use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_run_end`] requires for export. @@ -44,24 +45,24 @@ pub(super) fn to_arrow_run_end( array: ArrayRef, ends_type: &DataType, values_type: &Field, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, 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, exporter, ctx); + return constant_to_run_end(constant, ends_type, values_type, options, ctx); } Err(array) => array, }; let array = match array.try_downcast::() { - Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, exporter, ctx), + Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, options, ctx), Err(array) => array, }; // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = export_values(array, values_type, exporter, ctx)?; + let flat = export_values(array, values_type, options, ctx)?; let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), Arc::new(values_type.clone()), @@ -74,24 +75,24 @@ pub(super) fn to_arrow_run_end( fn export_values( values: ArrayRef, values_type: &Field, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - exporter.execute_arrow(values, Some(values_type), ctx) + values.execute_arrow_with_options(Some(values_type), options, ctx) } fn run_end_to_arrow( array: RunEndArray, ends_type: &DataType, values_type: &Field, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let length = array.len(); let offset = array.offset(); - 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)?; + let arrow_ends = execute_arrow_naive(array.ends().clone(), Some(ends_type), options, ctx)?; + let arrow_values = export_values(array.values().clone(), values_type, options, ctx)?; match ends_type { DataType::Int16 => build_run_array::(&arrow_ends, &arrow_values, offset, length), @@ -146,7 +147,7 @@ fn constant_to_run_end( array: ConstantArray, ends_type: &DataType, values_type: &Field, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -163,7 +164,7 @@ fn constant_to_run_end( let values = export_values( ConstantArray::new(scalar.clone(), 1).into_array(), values_type, - exporter, + options, ctx, )?; @@ -214,7 +215,7 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; use crate::executor::run_end::ConstantArray; use crate::executor::run_end::RunEnd; @@ -231,7 +232,8 @@ mod tests { array: vortex_array::ArrayRef, dt: &DataType, ) -> VortexResult { - array.execute_arrow(Some(dt), &mut SESSION.create_execution_ctx()) + let field = Field::new("", dt.clone(), array.dtype().is_nullable()); + array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx()) } fn constant_i32_with_i16_ends() -> arrow_array::ArrayRef { diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index 213770e5771..8559a0a327e 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -27,7 +27,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use crate::ArrowExporter; +use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::infer_nearest_arrow_field; use crate::executor::validity::to_arrow_null_buffer; use crate::session::ArrowSessionExt; @@ -51,7 +52,7 @@ impl Matcher for ArrowStructExportable { pub(super) fn to_arrow_struct( array: ArrayRef, target_fields: Option<&Fields>, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); @@ -83,7 +84,7 @@ pub(super) fn to_arrow_struct( &fields, validity, len, - exporter, + options, ctx, ); } @@ -102,7 +103,7 @@ pub(super) fn to_arrow_struct( &array.children(), None, // Pack is never null, len, - exporter, + options, ctx, ); } @@ -131,7 +132,7 @@ pub(super) fn to_arrow_struct( &fields, validity, len, - exporter, + options, ctx, ) } @@ -141,7 +142,7 @@ fn create_from_fields( vortex_fields: &[ArrayRef], null_buffer: Option, len: usize, - exporter: &ArrowExporter<'_>, + options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { match fields { @@ -156,8 +157,11 @@ 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 = - exporter.execute_arrow(vx_field.clone(), Some(field.as_ref()), ctx)?; + let arrow_field = vx_field.clone().execute_arrow_with_options( + Some(field.as_ref()), + options, + ctx, + )?; vortex_ensure!( field.is_nullable() || arrow_field.null_count() == 0, "Cannot convert field '{}' to non-nullable Arrow field because it contains nulls", @@ -192,7 +196,9 @@ fn create_from_fields( .extension_type_name() .is_some() .then_some(&inferred); - let arrow_array = exporter.execute_arrow(vx_field.clone(), target, ctx)?; + let arrow_array = vx_field + .clone() + .execute_arrow_with_options(target, options, ctx)?; // The executed array is authoritative for the physical type; only the metadata is // taken from the inferred field. arrow_fields.push(Arc::new( @@ -249,7 +255,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrowArrayExecutor; + use crate::ExecuteArrowExt; use crate::convert::from_arrow_dyn; use crate::dtype::to_data_type_naive; @@ -272,7 +278,7 @@ mod tests { let actual = struct_a .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx)?; + .execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx)?; let expected = ArrowStructArray::new( vec![Field::new("xs", DataType::Int64, false)].into(), vec![Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4]))], @@ -300,7 +306,7 @@ mod tests { let error = struct_a .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx) + .execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx) .unwrap_err(); assert!( error @@ -330,7 +336,7 @@ mod tests { let err = struct_a .into_array() - .execute_arrow(Some(&arrow_dt), &mut ctx) + .execute_arrow(Some(&Field::new("", arrow_dt, true)), &mut ctx) .err() .unwrap(); assert!( @@ -377,7 +383,7 @@ mod tests { assert_eq!( &array .into_array() - .execute_arrow(Some(&arrow_dtype), &mut ctx)?, + .execute_arrow(Some(&Field::new("", arrow_dtype, true)), &mut ctx)?, &arrow_array ); Ok(()) diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index b297dc31876..0ff749bcb08 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -3,23 +3,18 @@ //! Apache Arrow interoperability for Vortex. //! -//! This crate owns every conversion between Vortex and Arrow. The authoritative entry point is -//! the [`ArrowSession`]: importing Arrow schemas, fields, and data types into Vortex -//! ([`ArrowSession::from_arrow_schema`], [`ArrowSession::from_arrow_field`], +//! This crate owns every conversion between Vortex and Arrow. Exporting a Vortex array to Arrow goes +//! through [`ExecuteArrowExt`], which uses the [`ArrowSession`] of its execution context. Every other +//! conversion goes through the [`ArrowSession`] itself: 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`]), -//! 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). +//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_record_batch`]), and exporting +//! Vortex dtypes to Arrow ([`ArrowSession::to_arrow_schema`], [`ArrowSession::to_arrow_field`], +//! [`ArrowSession::to_arrow_datatype`]). //! -//! The [`ArrowSession`] session variable is created lazily on first use of -//! [`ArrowSessionExt::arrow`], so no explicit registration is required for plain conversions. -//! To import extension types that have no Arrow plugin, register the session with [`initialize`] -//! (see [`ArrowSession`]). - -use vortex_array::dtype::session::DTypeSessionExt; -use vortex_session::VortexSession; +//! The [`ArrowSession`] session variable is created on first use of [`ArrowSessionExt::arrow`], so +//! no explicit registration is required. It can look up the extension types of its Vortex session +//! by their Vortex IDs (see [`ArrowSession`]). pub mod convert; mod datum; @@ -45,13 +40,3 @@ pub use null_buffer::to_null_buffer; pub use options::*; pub use scalar::ToArrowDatum; pub use session::*; - -/// Register vortex-arrow's session extensions into `session`. -/// -/// This eagerly registers the [`ArrowSession`], which is otherwise created lazily on first use. -/// Only a session registered here can look up the extension types of `session` by their Vortex IDs. -/// A session that is created lazily on first use cannot, so it imports an extension type without an -/// Arrow plugin as plain storage. -pub fn initialize(session: &VortexSession) { - session.register(ArrowSession::new(&session.dtypes())); -} diff --git a/vortex-arrow/src/options/tests.rs b/vortex-arrow/src/options/tests.rs index 78bab8581d6..c24efed54b4 100644 --- a/vortex-arrow/src/options/tests.rs +++ b/vortex-arrow/src/options/tests.rs @@ -44,6 +44,7 @@ use crate::ArrowExport; use crate::ArrowExportVTable; use crate::ArrowSession; use crate::ArrowSessionExt; +use crate::ExecuteArrowExt; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -144,10 +145,7 @@ fn external_options_reach_nested_exporters() -> VortexResult<()> { .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)?; + let arrow = array.execute_arrow_with_options(Some(&field), &options, &mut ctx)?; assert_eq!(arrow.len(), 1); assert_eq!(arrow.null_count(), 0); assert_eq!( @@ -322,9 +320,9 @@ fn compaction_option_reaches_nested_views( let expected = expected_sparse_views(binary); let (array, target) = shape.wrap(view.into_array(), expected.data_type().clone(), &mut ctx)?; - let arrow = SESSION.arrow().exporter(&options).execute_arrow( - array, + let arrow = array.execute_arrow_with_options( Some(&Field::new("", target, true)), + &options, &mut ctx, )?; @@ -343,9 +341,7 @@ fn execute_arrow_compacts_by_default(#[values(false, true)] binary: bool) -> Vor 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)?; + let arrow = view.into_array().execute_arrow(None, &mut ctx)?; assert_eq!(arrow.to_data(), expected_sparse_views(binary).to_data()); assert!(!view_buffer_ptrs(&arrow).contains(&backing)); diff --git a/vortex-arrow/src/run_end_import.rs b/vortex-arrow/src/run_end_import.rs index 27b39233bef..692aeb3f46e 100644 --- a/vortex-arrow/src/run_end_import.rs +++ b/vortex-arrow/src/run_end_import.rs @@ -90,6 +90,7 @@ mod tests { use vortex_session::VortexSession; use crate::ArrowSessionExt; + use crate::ExecuteArrowExt; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -284,9 +285,7 @@ mod tests { fn execute(array: ArrayRef, dt: &DataType) -> VortexResult { let field = Field::new("", dt.clone(), true); - SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut SESSION.create_execution_ctx()) + array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx()) } #[test] diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 651c72452bd..ef6aa4029e9 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -44,8 +44,6 @@ use arrow_schema::FieldRef; use arrow_schema::Fields; use arrow_schema::Schema; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; -use tracing::debug; -use tracing::trace; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -55,7 +53,6 @@ use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::StructArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; @@ -65,6 +62,7 @@ use vortex_array::dtype::StructFields; use vortex_array::dtype::extension::ExtDTypeRef; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::session::DTypeSession; +use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::dtype::session::ExtDTypeRegistry; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::validity::Validity; @@ -89,7 +87,6 @@ use crate::convert::trim_map_entries; use crate::convert::trim_offsets; use crate::dtype::from_arrow_data_type; use crate::dtype::to_data_type_naive; -use crate::executor::execute_arrow_naive; use crate::ext_metadata::decode_arrow_ext_metadata; use crate::ext_metadata::encode_arrow_ext_metadata; use crate::run_end_import::run_end_from_arrow; @@ -162,7 +159,9 @@ impl<'a> IntoArrowField<'a> for Nullability { /// Plugin layer for exporting a Vortex array to an Arrow extension type. /// /// This is purely an implementation trait, its methods should not be called directly. Instead, -/// use the methods on [`ArrowSession`]. +/// use [`ExecuteArrowExt`] to export an array, and the methods on [`ArrowSession`] to convert a dtype. +/// +/// [`ExecuteArrowExt`]: crate::ExecuteArrowExt pub trait ArrowExportVTable: 'static + Send + Sync + Debug { /// The Arrow extension ID this plugin produces. fn arrow_ext_id(&self) -> Id; @@ -185,7 +184,10 @@ pub trait ArrowExportVTable: 'static + Send + Sync + Debug { /// /// Returns ownership of `array` via [`ArrowExport::Unsupported`] when the plugin cannot /// handle the input. Plugins that export child or storage arrays must pass `options` on to - /// those exports. + /// those exports, through [`ExecuteArrowExt::execute_arrow_with_options`]. The [`ArrowSession`] + /// of `ctx` is the session of the export, so a plugin that needs one must use it. + /// + /// [`ExecuteArrowExt::execute_arrow_with_options`]: crate::ExecuteArrowExt::execute_arrow_with_options fn execute_arrow( &self, array: ArrayRef, @@ -260,6 +262,10 @@ type ArrowImporterRegistry = ArcSwapMap>; /// owns the type registers its plugin, as `vortex-uuid` does. Temporal types need no plugin, /// because Arrow has native types for them. /// +/// Every [`ArrowSession`] belongs to a [`VortexSession`], which creates it on first use of +/// [`ArrowSessionExt::arrow`]. An export takes no session, because it uses the session of its +/// [`ExecutionCtx`] (see [`ExecuteArrowExt`]). So one export cannot mix the plugins of two sessions. +/// /// An extension type without a plugin exports as its storage, and the Arrow field records the type /// so that an import can restore it. `ARROW:extension:name` holds the Vortex ID, and /// `ARROW:extension:metadata` holds the metadata in base64, since Arrow metadata must be UTF-8 @@ -267,35 +273,44 @@ type ArrowImporterRegistry = ArcSwapMap>; /// schemas, so [`ExtDTypeRef::serialize_metadata`] must give equal bytes for equal types. /// /// A caller can also export an array to a field that names its Vortex ID, such as a field from -/// [`Self::to_arrow_storage_field`]. The export then checks that an import of the result gives back -/// the same type. The field must have the same metadata, and its Arrow type must convert to the -/// storage type (or to the type itself, such as an Arrow timestamp). Without these checks, values -/// could silently change meaning, such as milliseconds that read back as seconds. For the same -/// reason, an extension array cannot export to a field that names a different known Vortex ID. +/// [`Self::to_arrow_storage_field`]. The export then checks that the field describes the type of the +/// array. The field must have the same metadata, and its Arrow type must import as the storage type +/// (or as the type itself, such as an Arrow timestamp). Nested extension types count as their +/// storage types in this check, and the export checks each nested field when it reaches it. Without +/// these checks, values could silently change meaning, such as milliseconds that read back as +/// seconds. For the same reason, an extension array cannot export to a field that names a different +/// known Vortex ID. /// /// On import, the plugins run first. If none applies, a field that names a known Vortex ID imports /// as that type over its storage, and a field with any other name imports as plain storage. A -/// Vortex ID is known if its type is registered in the session (see [`initialize`]). +/// Vortex ID is known if its type is registered in the [`DTypeSession`] of the same +/// [`VortexSession`], even after the first use of the Arrow session. The Arrow session keeps the +/// registry of the [`DTypeSession`] that existed at its first use, so it does not see a new +/// [`DTypeSession`] that later replaces it through [`VortexSession::register`]. /// +/// [`ExecuteArrowExt`]: crate::ExecuteArrowExt /// [`ext_metadata`]: crate::ext_metadata -/// [`initialize`]: crate::initialize -#[derive(Clone, Debug, Default)] +/// [`VortexSession`]: vortex_session::VortexSession +/// [`VortexSession::register`]: vortex_session::VortexSession::register +#[derive(Clone, Debug)] pub struct ArrowSession { exporters: ArrowExporterRegistry, exporters_by_vortex: VortexExporterRegistry, importers: ArrowImporterRegistry, - /// The extension dtypes that the session can look up by Vortex ID. [`Self::new`] shares the - /// registry of a [`DTypeSession`], so dtypes that are registered there later are also found. - /// The default session has none. + /// The extension dtypes that the session can look up by Vortex ID. This is the registry of a + /// [`DTypeSession`], so dtypes that are registered there later are also found. ext_dtypes: ExtDTypeRegistry, } impl ArrowSession { - /// Create a session that can look up the extension dtypes of `dtype_session` by Vortex ID. - pub(crate) fn new(dtype_session: &DTypeSession) -> Self { + /// Create a session without plugins that can look up the extension dtypes of `dtype_session` + /// by Vortex ID. + fn new(dtype_session: &DTypeSession) -> Self { Self { + exporters: ArrowExporterRegistry::default(), + exporters_by_vortex: VortexExporterRegistry::default(), + importers: ArrowImporterRegistry::default(), ext_dtypes: dtype_session.registry().clone(), - ..Self::default() } } @@ -315,7 +330,7 @@ impl ArrowSession { self.importers.push(importer.arrow_ext_id(), importer); } - fn exporters(&self, id: &Id) -> Arc<[ArrowExportVTableRef]> { + pub(crate) fn exporters(&self, id: &Id) -> Arc<[ArrowExportVTableRef]> { self.exporters.get(id).unwrap_or_else(|| Arc::from([])) } @@ -505,6 +520,14 @@ impl ArrowSession { Ok(field) } + /// Returns the Vortex ID that the extension name `id` resolves to, or `None` if the session does + /// not know it. The registry can hold an alias, so the result can differ from `id`. + pub(crate) fn known_vortex_id(&self, id: &Id) -> Option { + self.ext_dtypes + .get(id) + .map(|ext_dtype_plugin| ext_dtype_plugin.id()) + } + /// Returns the extension type that `field` names by Vortex ID, to wrap around `storage_dtype`. /// /// Returns `None` if the session does not know that Vortex ID, or if `storage_dtype` already is @@ -529,11 +552,7 @@ impl ArrowSession { .as_extension_opt() .filter(|ext_dtype| ext_dtype.id() == ext_dtype_plugin.id()); if let Some(existing_ext_dtype) = existing_ext_dtype { - vortex_ensure!( - metadata == existing_ext_dtype.serialize_metadata()?, - "Arrow field {:?} must carry the metadata of {existing_ext_dtype}, got {metadata:?}", - field.name() - ); + ensure_ext_metadata(field, &metadata, existing_ext_dtype)?; return Ok(None); } @@ -697,61 +716,6 @@ impl ArrowSession { Ok(StructArray::try_new(names, columns, length, Validity::NonNullable)?.into_array()) } - /// Execute a Vortex array into an Arrow array. - /// - /// If `target` carries an `ARROW:extension:name`, the plugin registry is probed for one that - /// can support executing to the target extension type. If no plugin accepts it, but the - /// extension name is the Vortex ID of the array, the array exports its storage (see - /// [`ArrowSession`]). - /// - /// With `target = None`, the export uses the field from [`Self::default_arrow_export_field`], - /// so it fails for the same types. If that field has an Arrow extension name, the plugins - /// export the array to it, as above. Otherwise, the array exports to its preferred Arrow type, - /// and each nested extension type exports through its plugin, or as its storage plus Arrow - /// extension metadata. - /// - /// 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 { - self.exporter(&ArrowExportOptions::default()) - .execute_arrow(array, target, ctx) - } - - /// Return an exporter that applies `options` to its exports. - /// - /// ``` - /// use arrow_array::Array; - /// use vortex_array::IntoArray; - /// use vortex_array::VortexSessionExecute; - /// use vortex_session::VortexSession; - /// use vortex_array::arrays::VarBinViewArray; - /// use vortex_arrow::ArrowExportOptions; - /// use vortex_arrow::ArrowSessionExt; - /// use vortex_arrow::CompactBuffers; - /// - /// let session = VortexSession::empty(); - /// 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, - } - } - /// Decode an Arrow array into a Vortex array. /// /// `field` describes the Arrow type the array is imported from: pass an Arrow [`Field`] when @@ -969,10 +933,10 @@ impl ArrowSession { DataType::Map(entries_field, keys_sorted) => { let map = array.as_map(); let (entries, offsets) = trim_map_entries(map); + // `map_from_arrow_parts` needs the entries as a plain struct, so a Vortex ID on the + // entries field is ignored. Import plugins for its extension name still run, while + // `from_arrow_datatype` ignores the entries field completely. let entries_array: ArrowArrayRef = Arc::new(entries); - // Map entries must import as a plain struct, so a Vortex ID on their field is - // ignored, as `from_arrow_datatype` also does. Wrapping the entries in that - // extension type would only make `map_from_arrow_parts` fail. let entries = self.from_arrow_array_with_importers(entries_array, entries_field.as_ref())?; map_from_arrow_parts( @@ -1066,164 +1030,33 @@ 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.default_arrow_export_field(array.dtype())?; - &arrow_field - } - }; - - if let Some(arrow_ext_name) = target_field.metadata().get(EXTENSION_TYPE_NAME_KEY) { - let arrow_ext_id = Id::new(arrow_ext_name); - - // 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(&arrow_ext_id).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, - } - } - - // The registry can hold an alias, so compare the array with the ID of the dtype plugin - // that the extension name resolves to. - let named_vortex_id = self - .session - .ext_dtypes - .get(&arrow_ext_id) - .map(|dtype_plugin| dtype_plugin.id()); - - if let Some(ext_dtype) = current - .dtype() - .as_extension_opt() - .filter(|ext_dtype| { - ext_dtype.id() == arrow_ext_id || Some(ext_dtype.id()) == named_vortex_id - }) - .cloned() - { - return self.execute_arrow_ext_storage(current, &ext_dtype, target_field, ctx); - } - - // Only an extension array fails on a target that names another Vortex ID. A plain - // array still exports to it, because plugins export their storage to the target that - // they receive. - if let (Some(ext_dtype), Some(named_vortex_id)) = - (current.dtype().as_extension_opt(), named_vortex_id) - { - vortex_bail!( - "Arrow field {:?} names the extension type {named_vortex_id}, so the array must have that type, got {ext_dtype}", - target_field.name() - ); - } - - 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) - } - - /// Export `array`, whose type is `ext_dtype`, to `target`, whose extension name is the Vortex - /// ID of `ext_dtype`. [`ArrowSession`] describes the checks on `target`. - fn execute_arrow_ext_storage( - &self, - array: ArrayRef, - ext_dtype: &ExtDTypeRef, - target: &Field, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let target_metadata = decode_arrow_ext_metadata(target)?; - vortex_ensure!( - target_metadata == ext_dtype.serialize_metadata()?, - "Arrow field {:?} must carry the metadata of {ext_dtype}, got {target_metadata:?}", - target.name() - ); - - let target_dtype = self - .session - .from_arrow_datatype(target.data_type(), ext_dtype.nullability())?; - - // An Arrow type that already holds the extension type, such as an Arrow timestamp, exports - // the array itself, so that the executor of that type checks its own parameters. - if target_dtype.eq_ignore_nullability(array.dtype()) { - return execute_arrow_naive(array, Some(target.data_type()), self, ctx); - } - - vortex_ensure!( - target_dtype.eq_ignore_nullability(ext_dtype.storage_dtype()), - "Arrow field {:?} must have an Arrow type that imports as the storage of {ext_dtype}, got {}", - target.name(), - target.data_type() - ); - - let storage = array - .execute::(ctx)? - .storage_array() - .clone(); - - execute_arrow_naive(storage, Some(target.data_type()), self, ctx) - } +/// Fails unless `metadata`, the extension metadata of `field`, is the metadata of `ext_dtype`. +pub(crate) fn ensure_ext_metadata( + field: &Field, + metadata: &[u8], + ext_dtype: &ExtDTypeRef, +) -> VortexResult<()> { + vortex_ensure!( + metadata == ext_dtype.serialize_metadata()?, + "Arrow field {:?} must carry the metadata of {ext_dtype}, got {metadata:?}", + field.name() + ); + + Ok(()) } /// Extension trait for accessing the [`ArrowSession`] on a Vortex session. pub trait ArrowSessionExt: SessionExt { - /// Get the Arrow session. + /// Get the Arrow session, creating it on first use. + /// + /// The new session has no plugins, and it can look up the extension types of + /// [`DTypeSessionExt::dtypes`] by Vortex ID. fn arrow(&self) -> SessionGuard<'_, ArrowSession>; } impl ArrowSessionExt for S { fn arrow(&self) -> SessionGuard<'_, ArrowSession> { - self.get::() + self.get_or_insert_with(|| ArrowSession::new(&self.dtypes())) } } @@ -1300,7 +1133,7 @@ mod tests { )) } }; - let session = ArrowSession::default(); + let session = SESSION.arrow(); let sliced = session.from_arrow_array(lists(0..1000).slice(start, len), false)?; let fresh = session.from_arrow_array(lists(start as i64..(start + len) as i64), false)?; assert_eq!(sliced.nbytes(), fresh.nbytes()); @@ -1321,7 +1154,7 @@ mod tests { #[test] fn from_arrow_fields_matches_schema_conversion() -> VortexResult<()> { - let session = ArrowSession::default(); + let session = SESSION.arrow(); let fields = Fields::from(vec![ Field::new("a", DataType::Int32, false), Field::new("b", DataType::Utf8View, true), @@ -1342,7 +1175,7 @@ mod tests { #[test] fn from_arrow_field_maps_variant_without_importer() -> VortexResult<()> { - let session = ArrowSession::default(); + let session = SESSION.arrow(); let storage = DataType::Struct( vec![ Field::new("metadata", DataType::BinaryView, false), @@ -1416,7 +1249,8 @@ mod tests { #[case::plain(None)] #[case::extension_name(Some(*REJECT_DICTIONARY_ID))] fn dictionary_imports_as_dict_encoding(#[case] extension_name: Option) -> VortexResult<()> { - let session = ArrowSession::default(); + let vortex_session = VortexSession::empty(); + let session = vortex_session.arrow(); session.register_importer(Arc::new(DictionaryRejectingImporter)); let dict: ArrowArrayRef = Arc::new(DictionaryArray::::try_new( Int32Array::from(vec![0, 1, 0, 1]), @@ -1542,7 +1376,8 @@ mod tests { #[case] array: ArrowArrayRef, #[values(false, true)] decline: bool, ) { - let session = ArrowSession::default(); + let vortex_session = VortexSession::empty(); + let session = vortex_session.arrow(); session.register_importer(Arc::new(TruncatingImporter { decline })); let err = session diff --git a/vortex-arrow/src/tests.rs b/vortex-arrow/src/tests.rs index d68dcffc87e..17ccde552d7 100644 --- a/vortex-arrow/src/tests.rs +++ b/vortex-arrow/src/tests.rs @@ -20,12 +20,14 @@ use vortex_array::assert_arrays_eq; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrowSession; +use crate::ArrowSessionExt; use crate::convert::from_arrow_dyn; fn import(array: ArrowArrayRef, nullable: bool, use_session: bool) -> VortexResult { if use_session { - ArrowSession::default().from_arrow_array(array, nullable) + VortexSession::empty() + .arrow() + .from_arrow_array(array, nullable) } else { from_arrow_dyn(array.as_ref(), nullable) } diff --git a/vortex-arrow/tests/canonical.rs b/vortex-arrow/tests/canonical.rs index d18a35b1880..bb506df58b0 100644 --- a/vortex-arrow/tests/canonical.rs +++ b/vortex-arrow/tests/canonical.rs @@ -28,6 +28,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -56,9 +57,9 @@ fn test_canonicalize_nested_struct() { ]) .unwrap(); - let arrow_struct = SESSION - .arrow() - .execute_arrow(nested_struct_array.into_array(), None, &mut ctx) + let arrow_struct = nested_struct_array + .into_array() + .execute_arrow(None, &mut ctx) .unwrap() .as_any() .downcast_ref::() @@ -147,10 +148,7 @@ fn roundtrip_struct() { .arrow() .from_arrow_array(Arc::clone(&arrow_struct), true) .unwrap(); - let vortex_struct = SESSION - .arrow() - .execute_arrow(vortex_struct, None, &mut ctx) - .unwrap(); + let vortex_struct = vortex_struct.execute_arrow(None, &mut ctx).unwrap(); assert_eq!(&arrow_struct, &vortex_struct); } @@ -177,9 +175,8 @@ fn roundtrip_list() { .from_arrow_array(Arc::clone(&arrow_list), true) .unwrap(); - let rt_arrow_list = SESSION - .arrow() - .execute_arrow(vortex_list, Some(&list_field), &mut ctx) + let rt_arrow_list = vortex_list + .execute_arrow(Some(&list_field), &mut ctx) .unwrap(); assert_eq!(&arrow_list, &rt_arrow_list); diff --git a/vortex-arrow/tests/vortex_id.rs b/vortex-arrow/tests/vortex_id.rs index 584feee6cbd..6267ab88baa 100644 --- a/vortex-arrow/tests/vortex_id.rs +++ b/vortex-arrow/tests/vortex_id.rs @@ -48,9 +48,9 @@ use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::Datum; +use vortex_arrow::ExecuteArrowExt; use vortex_arrow::ext_metadata::EXTENSION_METADATA_ENCODING_KEY; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -58,16 +58,23 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -/// A session that can look up [`Tagged`] by its Vortex ID. +/// A session that can look up [`Tagged`] by its Vortex ID and by [`TAGGED_ALIAS`]. static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty(); session.dtypes().register(Tagged); - vortex_arrow::initialize(&session); + session + .dtypes() + .registry() + .insert(*TAGGED_ALIAS_ID, Arc::new(Tagged)); session }); static TAGGED_ID: CachedId = CachedId::new("test.tagged"); +/// Another name that the dtype registry of [`SESSION`] resolves to [`Tagged`]. +const TAGGED_ALIAS: &str = "test.tagged.alias"; +static TAGGED_ALIAS_ID: CachedId = CachedId::new(TAGGED_ALIAS); + /// An extension type without an Arrow plugin. It accepts any storage, and its metadata is one tag /// byte. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -147,6 +154,8 @@ fn tagged_timestamps() -> VortexResult { /// The field from [`ArrowSession::to_arrow_storage_field`] for the type of `array`, with the Arrow /// type `data_type`. It takes a [`VortexResult`] so that an rstest case can pass a fixture call. +/// +/// [`ArrowSession::to_arrow_storage_field`]: vortex_arrow::ArrowSession::to_arrow_storage_field fn storage_field_as(array: VortexResult, data_type: DataType) -> VortexResult { let field = SESSION .arrow() @@ -220,9 +229,8 @@ fn export_rejects_target_that_changes_the_type( #[case] target: VortexResult, #[case] message: &str, ) -> VortexResult<()> { - let err = SESSION - .arrow() - .execute_arrow(array?, Some(&target?), &mut SESSION.create_execution_ctx()) + let err = array? + .execute_arrow(Some(&target?), &mut SESSION.create_execution_ctx()) .expect_err("a target that changes the type must be rejected"); assert!(err.to_string().contains(message), "unexpected error: {err}"); @@ -256,7 +264,7 @@ fn export_rejects_storage_with_extension_name() -> VortexResult<()> { SESSION.arrow().to_arrow_datatype(array.dtype()).map(drop) })] #[case::execute_arrow(|array: &ArrayRef, ctx: &mut ExecutionCtx| { - SESSION.arrow().execute_arrow(array.clone(), None, ctx).map(drop) + array.clone().execute_arrow(None, ctx).map(drop) })] #[case::datum(|array: &ArrayRef, ctx: &mut ExecutionCtx| Datum::try_new(array, ctx).map(drop))] fn exports_without_a_field_reject_types_without_a_plugin( @@ -280,7 +288,7 @@ fn nested_export_without_a_target_keeps_ext_metadata() -> VortexResult<()> { let array = struct_of_tagged()?; let field = SESSION.arrow().to_arrow_field("", array.dtype())?; let mut ctx = SESSION.create_execution_ctx(); - let exported = SESSION.arrow().execute_arrow(array, None, &mut ctx)?; + let exported = array.execute_arrow(None, &mut ctx)?; assert_eq!(exported.data_type(), field.data_type()); @@ -299,7 +307,7 @@ fn round_trip_keeps_the_extension_type(#[case] array: VortexResult) -> let session = SESSION.arrow(); let mut ctx = SESSION.create_execution_ctx(); let field = session.to_arrow_field("t", array.dtype())?; - let exported = session.execute_arrow(array.clone(), Some(&field), &mut ctx)?; + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; assert_eq!(session.from_arrow_field(&field)?, *array.dtype()); assert_arrays_eq!(session.from_arrow_array(exported, &field)?, array, &mut ctx); @@ -319,9 +327,7 @@ fn timestamp_round_trips_through_a_vortex_id_field( let field = storage_field_as(Ok(timestamps.clone()), data_type.clone())?; let mut ctx = SESSION.create_execution_ctx(); - let exported = SESSION - .arrow() - .execute_arrow(timestamps.clone(), Some(&field), &mut ctx)?; + let exported = timestamps.clone().execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), &data_type); assert_eq!( @@ -334,6 +340,59 @@ fn timestamp_round_trips_through_a_vortex_id_field( Ok(()) } +/// A field that names an alias of a Vortex ID exports and imports like a field that names the ID. +#[test] +fn alias_of_vortex_id_round_trips() -> VortexResult<()> { + let array = tagged_i64s(0xff)?; + let mut field = storage_field_as(Ok(array.clone()), DataType::Int64)?; + field + .metadata_mut() + .insert(EXTENSION_TYPE_NAME_KEY.to_owned(), TAGGED_ALIAS.to_owned()); + let mut ctx = SESSION.create_execution_ctx(); + + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + let imported = SESSION.arrow().from_arrow_array(exported, &field)?; + assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) +} + +/// A plain array exports to a field that names a Vortex ID, because a plugin such as `Vector` +/// exports its storage to the same field that it received. +#[test] +fn plain_array_exports_to_a_vortex_id_field() -> VortexResult<()> { + let storage = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); + let field = storage_field_as(tagged_i64s(0xff), DataType::Int64)?; + let mut ctx = SESSION.create_execution_ctx(); + + let exported = storage.clone().execute_arrow(Some(&field), &mut ctx)?; + let plain_field = Field::new("t", DataType::Int64, true); + let imported = SESSION.arrow().from_arrow_array(exported, &plain_field)?; + assert_arrays_eq!(imported, storage, &mut ctx); + + Ok(()) +} + +/// The export to a field that names a Vortex ID compares nested extension types by their storage, +/// because the import of the field may not restore them. Here the session does not know the nested +/// Vortex ID, and an import plugin that loses metadata, such as the version of a UUID, has the same +/// effect. +#[test] +fn export_compares_nested_extension_types_by_storage() -> VortexResult<()> { + let array = tagged_array(0x01, struct_of_tagged()?)?; + let field = SESSION + .arrow() + .to_arrow_storage_field("t", array.dtype().as_extension())?; + let session = VortexSession::empty(); + let mut ctx = session.create_execution_ctx(); + + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + let imported = SESSION.arrow().from_arrow_array(exported, &field)?; + assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) +} + /// A session that does not know the Vortex ID imports the field as plain storage. #[test] fn unknown_vortex_id_imports_as_storage() -> VortexResult<()> { @@ -341,7 +400,8 @@ fn unknown_vortex_id_imports_as_storage() -> VortexResult<()> { .arrow() .to_arrow_field("t", tagged_i64s(0xff)?.dtype())?; let arrow: ArrowArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); - let session = ArrowSession::default(); + let unaware_session = VortexSession::empty(); + let session = unaware_session.arrow(); let storage = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]); assert_eq!(session.from_arrow_field(&field)?, *storage.dtype()); @@ -354,6 +414,21 @@ fn unknown_vortex_id_imports_as_storage() -> VortexResult<()> { Ok(()) } +/// The Arrow session of a Vortex session knows the extension types of its dtype session, including +/// the ones that are registered after the Arrow session is created. +#[test] +fn vortex_id_imports_types_registered_after_the_arrow_session() -> VortexResult<()> { + let array = tagged_i64s(0xff)?; + let session = VortexSession::empty(); + let field = session.arrow().to_arrow_field("t", array.dtype())?; + + session.dtypes().register(Tagged); + + assert_eq!(session.arrow().from_arrow_field(&field)?, *array.dtype()); + + Ok(()) +} + /// Metadata without the encoding key is its own UTF-8 bytes, as other producers of these fields /// write it. #[test] @@ -426,8 +501,8 @@ fn run_end_field_that_names_its_type_twice_wraps_it_once() -> VortexResult<()> { Ok(()) } -/// The entries field of a map is part of the map type, so a Vortex ID on it is ignored by both the -/// dtype import and the array import. +/// The entries field of a map is part of the map type, so the dtype import and the array import +/// both ignore a Vortex ID on it. #[test] fn map_entries_ignore_a_vortex_id() -> VortexResult<()> { let key_field = Arc::new(Field::new("key", DataType::Int64, false)); diff --git a/vortex-bench/src/datasets/sparse_wide.rs b/vortex-bench/src/datasets/sparse_wide.rs index a275a3a8eb5..a0b7229518c 100644 --- a/vortex-bench/src/datasets/sparse_wide.rs +++ b/vortex-bench/src/datasets/sparse_wide.rs @@ -33,6 +33,7 @@ use vortex::buffer::Buffer; use vortex::dtype::Nullability; use vortex::mask::Mask; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use crate::IdempotentPath; use crate::SESSION; @@ -184,7 +185,7 @@ impl Dataset for SparseWide { fn write_parquet(path: &Path, array: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { let schema = SESSION.arrow().to_arrow_schema(array.dtype())?; let root = Field::new_struct("", schema.fields, false); - let arrow = SESSION.arrow().execute_arrow(array, Some(&root), ctx)?; + let arrow = array.execute_arrow(Some(&root), ctx)?; let batch = RecordBatch::from(arrow.as_struct()); let properties = WriterProperties::builder() .set_compression(Compression::ZSTD(ZstdLevel::try_new(3)?)) diff --git a/vortex-bench/src/datasets/struct_list_of_ints.rs b/vortex-bench/src/datasets/struct_list_of_ints.rs index f854b59c707..50abc12f8dd 100644 --- a/vortex-bench/src/datasets/struct_list_of_ints.rs +++ b/vortex-bench/src/datasets/struct_list_of_ints.rs @@ -27,6 +27,7 @@ use vortex::array::arrays::listview::recursive_list_from_list_view; use vortex::array::validity::Validity; use vortex::dtype::FieldNames; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use crate::IdempotentPath; use crate::SESSION; @@ -149,9 +150,7 @@ impl Dataset for StructListOfInts { let converted = recursive_list_from_list_view(chunk.clone(), &mut ctx)?; let schema = SESSION.arrow().to_arrow_schema(converted.dtype())?; let schema = Field::new_struct("", schema.fields, false); - let batch = SESSION - .arrow() - .execute_arrow(converted, Some(&schema), &mut ctx)?; + let batch = converted.execute_arrow(Some(&schema), &mut ctx)?; let batch = RecordBatch::from(batch.as_struct()); if writer.is_none() { diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 09e97324fc3..f1fa752fb92 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -28,6 +28,7 @@ vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } +vortex-mask = { workspace = true } vortex-onpair = { workspace = true } vortex-pco = { workspace = true, optional = true } vortex-sequence = { workspace = true } diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 9b6e9aaf7c4..6d261788780 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -18,6 +18,7 @@ use crate::SchemeExt; use crate::SchemeId; use crate::allowed_ids::AllowedIds; use crate::schemes::binary; +use crate::schemes::fixed_size_list; use crate::schemes::float; use crate::schemes::integer; use crate::schemes::string; @@ -74,6 +75,7 @@ impl CompressionMode { float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), string::NullDominatedSparseScheme.id(), + fixed_size_list::FixedSizeListSparseScheme.id(), string::StringDictScheme.id(), binary::BinaryDictScheme.id(), ]); diff --git a/vortex-btrblocks/src/schemes/fixed_size_list.rs b/vortex-btrblocks/src/schemes/fixed_size_list.rs new file mode 100644 index 00000000000..318e761b373 --- /dev/null +++ b/vortex-btrblocks/src/schemes/fixed_size_list.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Sparse encoding for fixed-size lists with many null rows. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_sparse::Sparse; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::SchemeExt; +use crate::schemes::integer::IntDictScheme; +use crate::schemes::integer::IntRLEScheme; +use crate::schemes::integer::RunEndScheme; +use crate::schemes::integer::SparseScheme as IntSparseScheme; + +/// The smallest estimated ratio of the dense list's size to the sparse one's worth encoding for. +const MIN_ESTIMATED_RATIO: f64 = 1.1; + +/// Sparse encoding of a fixed-size list that stores only its valid rows. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct FixedSizeListSparseScheme; + +impl Scheme for FixedSizeListSparseScheme { + fn scheme_name(&self) -> &'static str { + "vortex.fixed_size_list.sparse" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches!(canonical, Canonical::FixedSizeList(_)) && canonical.dtype().is_nullable() + } + + fn produced_encodings(&self) -> Vec { + vec![Sparse.id()] + } + + /// Children: values=0, indices=1. + fn num_children(&self) -> usize { + 2 + } + + /// Sparse indices (child 1) are monotonically increasing positions with all unique values. + fn descendant_exclusions(&self) -> Vec { + vec![ + DescendantExclusion { + excluded: IntDictScheme.id(), + children: ChildSelection::One(1), + }, + DescendantExclusion { + excluded: RunEndScheme.id(), + children: ChildSelection::One(1), + }, + DescendantExclusion { + excluded: IntRLEScheme.id(), + children: ChildSelection::One(1), + }, + DescendantExclusion { + excluded: IntSparseScheme.id(), + children: ChildSelection::One(1), + }, + ] + } + + /// Compares the list's size with the valid rows plus their bit-packed positions. + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + let array = data.array(); + let len = array.len(); + let Ok(mask) = validity_mask(array, exec_ctx) else { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + }; + let valid = mask.true_count(); + if valid == 0 || valid == len { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + + let dense_bytes = array.nbytes() as f64; + let row_bytes = dense_bytes / len as f64; + let index_bytes = f64::from(usize::BITS - (len - 1).leading_zeros()) / 8.0; + let sparse_bytes = valid as f64 * (row_bytes + index_bytes); + + let ratio = dense_bytes / sparse_bytes; + if ratio < MIN_ESTIMATED_RATIO { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = data.array(); + let mask = validity_mask(array, exec_ctx)?; + let AllOr::Some(valid_indices) = mask.indices() else { + // The estimate skips lists without both valid and null rows. + return Ok(array.clone()); + }; + + let values = array.filter(mask.clone())?; + let compressed_values = + compressor.compress_child(&values, &compress_ctx, self.id(), 0, exec_ctx)?; + + let indices = PrimitiveArray::new( + valid_indices + .iter() + .map(|&idx| idx as u64) + .collect::>(), + Validity::NonNullable, + ) + .narrow(exec_ctx)?; + let compressed_indices = compressor.compress_child( + &indices.into_array(), + &compress_ctx, + self.id(), + 1, + exec_ctx, + )?; + + Ok(Sparse::try_new( + compressed_indices, + compressed_values, + array.len(), + Scalar::null(array.dtype().clone()), + exec_ctx, + )? + .into_array()) + } +} + +/// Returns which rows of `array` are valid. +fn validity_mask(array: &ArrayRef, exec_ctx: &mut ExecutionCtx) -> VortexResult { + array.validity()?.execute_mask(array.len(), exec_ctx) +} diff --git a/vortex-btrblocks/src/schemes/mod.rs b/vortex-btrblocks/src/schemes/mod.rs index a0e9b042a66..1cab9b72b56 100644 --- a/vortex-btrblocks/src/schemes/mod.rs +++ b/vortex-btrblocks/src/schemes/mod.rs @@ -9,6 +9,7 @@ pub mod integer; pub mod string; pub mod decimal; +pub mod fixed_size_list; pub mod temporal; pub(crate) mod patches; diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs index eea3c9f8e0f..1c8849e011e 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -14,6 +14,7 @@ use crate::Scheme; use crate::SchemeExt; use crate::schemes::binary; use crate::schemes::decimal; +use crate::schemes::fixed_size_list; use crate::schemes::float; use crate::schemes::integer; use crate::schemes::string; @@ -95,6 +96,8 @@ impl Default for CompressionSession { &decimal::DECIMAL_V1, // Temporal schemes. &temporal::TemporalScheme, + // Fixed-size list schemes. + &fixed_size_list::FixedSizeListSparseScheme, ], } } diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index fea1433dbe7..a77e6c365c4 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -21,6 +21,7 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -32,6 +33,7 @@ use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_datetime_parts::DateTimeParts; #[cfg(feature = "zstd")] use vortex_edition::EDITION_DECLARATIONS; @@ -46,6 +48,7 @@ use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexResult; use vortex_fastlanes::Delta; use vortex_session::VortexSession; +use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; use crate::BtrBlocksCompressorBuilder; @@ -163,6 +166,48 @@ fn test_delta_unaligned_roundtrip(#[case] input: PrimitiveArray) -> VortexResult Ok(()) } +/// Fixed-size lists of random bytes, like UUIDs, with the rows `is_valid` picks valid. The null +/// rows hold zeros, as Arrow's fixed-size binary arrays do. +fn random_byte_lists(rows: usize, is_valid: impl Fn(usize) -> bool) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let elements = (0..rows * 16) + .map(|i| { + if is_valid(i / 16) { + rng.next_u32().to_le_bytes()[0] + } else { + 0 + } + }) + .collect::>(); + let validity = Validity::from_iter((0..rows).map(&is_valid)); + FixedSizeListArray::new(elements.into_array(), 16, validity, rows).into_array() +} + +#[rstest] +#[case::mostly_null(|row| row % 20 == 0, true)] +#[case::half_null(|row| row % 2 == 0, true)] +// Dropping one row in twenty saves less than storing the valid rows' positions costs. +#[case::few_nulls(|row| row % 20 != 0, false)] +#[case::all_valid(|_| true, false)] +fn test_fixed_size_list_drops_null_rows( + #[case] is_valid: fn(usize) -> bool, + #[case] sparse: bool, +) -> VortexResult<()> { + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + let input = random_byte_lists(8_000, is_valid); + let compressed = assert_roundtrip(&compressor, &input)?; + assert_eq!( + compressed.is::(), + sparse, + "{}", + compressed.display_tree() + ); + + Ok(()) +} + #[cfg(feature = "zstd")] #[rstest] #[case::array_level(CORE_2026_08_3, vortex_zstd::Zstd.id())] diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index f86ab6b136b..de4599a9cc7 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -175,6 +175,32 @@ impl CascadingCompressor { } Canonical::Map(map_array) => self.compress_map_array(map_array, compress_ctx, exec_ctx), Canonical::FixedSizeList(fsl_array) => { + // Canonicalize the elements first, so that schemes see the list's real size and + // their output is accepted only if it beats it. + let elements = fsl_array + .elements() + .clone() + .execute::(exec_ctx)? + .0 + .compact(exec_ctx)? + .into_array(); + let fsl_array = FixedSizeListArray::try_new( + elements, + fsl_array.list_size(), + fsl_array.validity()?, + fsl_array.len(), + )?; + + // Schemes over the whole list can drop the elements of null rows, which + // compressing the elements alone has to keep. + if let Selection::Compressed(compressed) = self.choose_and_compress( + Canonical::FixedSizeList(fsl_array.clone()), + compress_ctx, + exec_ctx, + )? { + return Ok(compressed); + } + let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; Ok(FixedSizeListArray::try_new( diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 399f096604c..22c143ad2a1 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -291,7 +291,9 @@ mod tests { use vortex::dtype::extension::ExtVTable; use vortex::error::VortexResult; use vortex::scalar::ScalarValue; + use vortex::session::VortexSession; use vortex::session::registry::CachedId; + use vortex_arrow::ArrowSessionExt; use super::*; @@ -312,7 +314,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -333,7 +336,7 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow())?; assert_eq!(physical_schema, logical_schema); Ok(()) } @@ -360,7 +363,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -376,7 +380,8 @@ mod tests { Nullability::NonNullable, ); - let result = calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()); + let result = + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()); assert!( result .unwrap_err() @@ -396,7 +401,8 @@ mod tests { Nullability::NonNullable, ); - let result = calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()); + let result = + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()); assert!( result .unwrap_err() @@ -443,7 +449,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -473,7 +480,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -505,7 +513,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -516,7 +525,8 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()); + let result = + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()); assert!(result.is_err()); assert!( result @@ -552,7 +562,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -594,7 +605,8 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow()) + .unwrap(); assert_eq!(physical_schema, logical_schema); } @@ -627,7 +639,7 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow())?; assert_eq!(physical_schema, logical_schema); Ok(()) } @@ -719,7 +731,7 @@ mod tests { ); let physical_schema = - calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + calculate_physical_schema(&dtype, &logical_schema, &VortexSession::empty().arrow())?; assert_eq!(physical_schema, logical_schema); Ok(()) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index f56634b73d8..a23c8e02661 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -46,6 +46,7 @@ use futures::TryStreamExt; use futures::stream; use object_store::path::Path; use tracing::Instrument; +use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; use vortex::error::VortexError; use vortex::error::VortexExpect; @@ -69,6 +70,7 @@ use vortex::metrics::MetricsRegistry; use vortex::scan::planning::planner::Planner; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; use vortex_utils::parallelism::get_available_parallelism; @@ -483,14 +485,9 @@ impl FileOpener for VortexOpener { let native_limit = limit .filter(|_| scan_v2 && residual_filter.is_none()) .map(|limit| (limit, Arc::clone(&output_schema))); - let map_batch = move |chunk| { + let map_batch = move |chunk: ArrayRef| { let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; + let arrow = chunk.execute_arrow(Some(&stream_target_field), &mut ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) }; let stream = if scan_v2 { diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index c62f88449db..22a74463c6f 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -138,6 +138,7 @@ use vortex::scan::PartitionStream; use vortex::scan::ScanRequest; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_utils::parallelism::get_available_parallelism; use super::dynamic_filters::convert_dynamic_filter; @@ -589,11 +590,8 @@ impl VortexDataSource { let _timer = metrics.elapsed_compute().timer(); let mut ctx = session.create_execution_ctx(); result.and_then(|chunk| { - let arrow = session.arrow().execute_arrow( - chunk, - Some(target_field.as_ref()), - &mut ctx, - )?; + let arrow = + chunk.execute_arrow(Some(target_field.as_ref()), &mut ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) }) }) diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 9e1d37b278f..16f6803e23f 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -33,6 +33,7 @@ use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use crate::RUNTIME; use crate::array::vx_array; @@ -378,9 +379,7 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( let on_chunk = move |chunk: VortexResult| -> VortexResult { let chunk: ArrayRef = chunk?; let mut ctx: ExecutionCtx = session.create_execution_ctx(); - let arrow = session - .arrow() - .execute_arrow(chunk, Some(&target), &mut ctx)?; + let arrow = chunk.execute_arrow(Some(&target), &mut ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) }; diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 8cf7bc5bc78..1c15a1a0ba7 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -38,13 +38,74 @@ const ONE_MEG: u64 = 1 << 20; /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. - /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive - /// dictionary encoding. + /// `IntDictScheme` is automatically excluded when compressing dictionary codes to prevent + /// recursive dictionary encoding. BtrBlocks(BtrBlocksCompressorBuilder), /// An opaque compressor used as-is for both data and stats compression. Opaque(Arc), } +/// The compressors that [`WriteStrategyBuilder::build`] derives from a [`CompressorConfig`]. +struct Compressors { + /// Compresses dictionary codes. + codes: Arc, + /// Compresses everything else: chunks that dictionary encoding rejects, dictionary values and + /// stats tables. It also probes columns for dictionary encoding by default. + full: Arc, +} + +impl CompressorConfig { + /// Splits the configured compressor into the codes and full compressors. + /// + /// The codes compressor excludes `IntDictScheme` because `DictStrategy` already + /// dictionary-encoded the column, and dictionary-encoding its integer codes again would be + /// redundant. The full compressor keeps every scheme, so chunks that `DictStrategy` rejects can + /// still use integer dictionaries, e.g. for a decimal's most significant part. An opaque + /// compressor is used as-is for both. + fn into_compressors(self) -> Compressors { + match self { + CompressorConfig::BtrBlocks(builder) => Compressors { + codes: Arc::new( + builder + .clone() + .exclude_schemes([IntDictScheme.id()]) + .build(), + ), + full: Arc::new(builder.build()), + }, + CompressorConfig::Opaque(compressor) => Compressors { + codes: Arc::clone(&compressor), + full: compressor, + }, + } + } +} + +/// Coalesces chunks up to the data block target size, then compresses them with `compressor` +/// and passes them to `buffered`. +fn coalesce_then_compress( + buffered: BufferedStrategy, + compressor: Arc, + row_block_size: usize, + data_block_target_bytes: Option, +) -> RepartitionStrategy { + RepartitionStrategy::new( + CompressingStrategy::new(buffered, compressor), + RepartitionWriterOptions { + // Write stream partitions roughly become segments. Because Vortex never reads less + // than one segment, the size of segments and, therefore, partitions, must be small + // enough to both (1) allow fine-grained random access reads and (2) allow + // sufficient read concurrency for the desired throughput. One megabyte is small + // enough to achieve this for S3 (Durner et al., "Exploiting Cloud Object Storage for + // High-Performance Analytics", VLDB Vol 16, Iss 11). + block_size_minimum: data_block_target_bytes.unwrap_or(0), + block_len_multiple: row_block_size, + block_size_target: data_block_target_bytes, + canonicalize: true, + }, + ) +} + /// Build a new [writer strategy](LayoutStrategy) to compress and reorganize chunks of a Vortex /// file. /// @@ -161,62 +222,42 @@ impl WriteStrategyBuilder { Arc::new(FlatLayoutStrategy::default()) }; - let compressor = self.compressor; - // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); // 6. buffer chunks so they end up with closer segment ids physically let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB - // 5. compress each chunk. - // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already - // dictionary-encodes columns. Allowing IntDictScheme here would redundantly - // dictionary-encode the integer codes produced by that earlier step. - let data_compressor: Arc = match &compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new( - builder - .clone() - .exclude_schemes([IntDictScheme.id()]) - .build(), - ), - CompressorConfig::Opaque(compressor) => Arc::clone(compressor), - }; - let compressing = CompressingStrategy::new(buffered, data_compressor); - - // 4. prior to compression, coalesce up to a minimum size - let coalescing = RepartitionStrategy::new( - compressing, - RepartitionWriterOptions { - // Write stream partitions roughly become segments. Because Vortex never reads less - // than one segment, the size of segments and, therefore, partitions, must be small - // enough to both (1) allow fine-grained random access reads and (2) allow - // sufficient read concurrency for the desired throughput. One megabyte is small - // enough to achieve this for S3 (Durner et al., "Exploiting Cloud Object Storage for - // High-Performance Analytics", VLDB Vol 16, Iss 11). - block_size_minimum: self.data_block_target_bytes.unwrap_or(0), - block_len_multiple: self.row_block_size, - block_size_target: self.data_block_target_bytes, - canonicalize: true, - }, + // 4. coalesce chunks up to a minimum size, then 5. compress each chunk. + // Dict codes and the chunks DictStrategy rejects get separate pipelines because they use + // different compressors; see `CompressorConfig::into_compressors`. + let compressors = self.compressor.into_compressors(); + let codes = coalesce_then_compress( + buffered.clone(), + compressors.codes, + self.row_block_size, + self.data_block_target_bytes, ); + let fallback = coalesce_then_compress( + buffered, + Arc::clone(&compressors.full), + self.row_block_size, + self.data_block_target_bytes, + ); + let full_compressor = compressors.full; // 2.1. | 3.1. compress stats tables and dict values. - let stats_compressor: Arc = match compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), - CompressorConfig::Opaque(compressor) => compressor, - }; - let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); + let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&full_compressor)); // 3. apply dict encoding or fallback let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { probe_compressor } else { - Arc::clone(&stats_compressor) + full_compressor }; let dict = DictStrategy::new( - coalescing.clone(), + codes, compress_then_flat.clone(), - coalescing, + fallback, DictLayoutOptions { probe_len: self.row_block_size, ..Default::default() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 6fa012f4bdd..3a20894e3e7 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -383,6 +383,50 @@ async fn decimal_writer_uses_default_or_session_restrictions( Ok(()) } +/// A wide decimal with unique values is not dictionary encoded by the dict layout, but its byte +/// parts are still compressed with every scheme, so the low-cardinality most significant part can +/// use an integer dictionary. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn rejected_decimal_keeps_nested_int_dict() -> VortexResult<()> { + let distinct = [3i128, -5, 7, 11]; + let values = (0..16_384usize) + .map(|i| { + (distinct[(i * 7 + i / 3) % distinct.len()] << 64) + i128::try_from(i * 31).unwrap() + }) + .collect::>(); + let decimal = DecimalArray::new(values, DecimalDType::new(38, 2), Validity::NonNullable); + let array = StructArray::from_fields(&[("decimal", decimal.into_array())])?.into_array(); + + let mut buffer = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .disable_editions() + .write(&mut buffer, array.clone().to_array_stream()) + .await?; + assert!(!layout_has_dict(summary.footer().layout().as_ref())); + + let actual = SESSION + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let has_dict_msp = actual.depth_first_traversal().any(|array| { + array + .as_opt::() + .is_some_and(|parts| parts.msp().is::()) + }); + assert!( + has_dict_msp, + "expected a dict-encoded msp in {}", + actual.display_tree() + ); + assert_arrays_eq!(actual, array, &mut SESSION.create_execution_ctx()); + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_round_trip_many_types() { diff --git a/vortex-file/src/v2/scan.rs b/vortex-file/src/v2/scan.rs index b5a3f104d9c..8f70d921176 100644 --- a/vortex-file/src/v2/scan.rs +++ b/vortex-file/src/v2/scan.rs @@ -28,6 +28,7 @@ use futures::stream::BoxStream; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_io::request::IoSource; @@ -93,9 +94,7 @@ impl ScanDriver { builder, Arc::new(move |batch| { let mut ctx = session.create_execution_ctx(); - let arrow = session - .arrow() - .execute_arrow(batch, Some(&field), &mut ctx)?; + let arrow = batch.execute_arrow(Some(&field), &mut ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) }), ) diff --git a/vortex-file/src/v2/scan/tests.rs b/vortex-file/src/v2/scan/tests.rs index d02467b63b0..878308f4e42 100644 --- a/vortex-file/src/v2/scan/tests.rs +++ b/vortex-file/src/v2/scan/tests.rs @@ -17,6 +17,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -75,9 +76,7 @@ async fn record_batches_match_arrays() -> VortexResult<()> { let expected = arrays .into_iter() .map(|array| { - let arrow = SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut ctx)?; + let arrow = array.execute_arrow(Some(&field), &mut ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) }) .collect::>>()?; diff --git a/vortex-jni/src/scan.rs b/vortex-jni/src/scan.rs index 9bab53b813c..22d0e9438bf 100644 --- a/vortex-jni/src/scan.rs +++ b/vortex-jni/src/scan.rs @@ -50,6 +50,7 @@ use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use crate::POOL; use crate::RUNTIME; @@ -390,10 +391,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( handle.spawn(async move { let chunk = chunk?; let mut ctx = session.create_execution_ctx(); - let arrow = - session - .arrow() - .execute_arrow(chunk, Some(target.as_ref()), &mut ctx)?; + let arrow = chunk.execute_arrow(Some(target.as_ref()), &mut ctx)?; rebase_offsets(RecordBatch::from(arrow.as_struct().clone())) }) }) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 765cd745836..51020ce6b55 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -561,7 +561,7 @@ mod tests { use vortex::session::SessionExt; use vortex::session::SessionVar; use vortex::session::registry::Id; - use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; use super::LIST_LAYOUT_EDITION; use super::write_session; @@ -756,7 +756,7 @@ mod tests { let mut chunks = partition?.execute()?; while let Some(chunk) = chunks.next().await { let mut ctx = session.create_execution_ctx(); - session.arrow().execute_arrow(chunk?, None, &mut ctx)?; + chunk?.execute_arrow(None, &mut ctx)?; } } } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index cf9cef8c586..30c311a72d7 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -22,7 +22,7 @@ use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -93,13 +93,8 @@ impl ArrowExportVTable for Json { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); - let storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + let storage = storage.execute_arrow_with_options(Some(&storage_field), options, ctx)?; Ok(ArrowExport::Exported(storage)) } @@ -172,6 +167,7 @@ mod tests { use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSessionExt; use vortex_arrow::CompactBuffers; + use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -203,11 +199,7 @@ mod tests { assert!(!field.is_nullable()); ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata())?; - let exported = SESSION.arrow().execute_arrow( - array, - Some(&field), - &mut SESSION.create_execution_ctx(), - )?; + let exported = array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx())?; assert!(exported.data_type().is_string()); @@ -237,11 +229,7 @@ mod tests { &vortex_array::dtype::DType::Utf8(Nullability::NonNullable) ); - let exported = SESSION.arrow().execute_arrow( - imported, - Some(&field), - &mut SESSION.create_execution_ctx(), - )?; + let exported = imported.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx())?; let strings = exported.as_string::(); assert_eq!( strings.iter().collect::>(), @@ -268,11 +256,9 @@ mod tests { let retains_backing = |options: &ArrowExportOptions, ctx: &mut vortex_array::ExecutionCtx| -> VortexResult { - let exported = SESSION.arrow().exporter(options).execute_arrow( - array.clone(), - Some(&field), - ctx, - )?; + let exported = array + .clone() + .execute_arrow_with_options(Some(&field), options, ctx)?; assert_eq!( exported .as_string_view() diff --git a/vortex-layout/examples/reverse_index/pruning_tests.rs b/vortex-layout/examples/reverse_index/pruning_tests.rs index 994313e4b13..26a95e3d2c5 100644 --- a/vortex-layout/examples/reverse_index/pruning_tests.rs +++ b/vortex-layout/examples/reverse_index/pruning_tests.rs @@ -113,7 +113,6 @@ fn session_with_indexed_layout() -> anyhow::Result { .with::() .with::() .with::(); - vortex_arrow::initialize(&session); vortex_sequence::initialize(&session); session.register_edition(&INDEXED_TEST_DECLARATION)?; session.enable_edition(INDEXED_TEST_EDITION)?; @@ -129,7 +128,6 @@ fn session_without_reverse_index() -> VortexSession { .with::() .with::() .with::(); - vortex_arrow::initialize(&session); vortex_sequence::initialize(&session); session } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 09044047e73..8331567b399 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -394,7 +394,7 @@ mod tests { Nullability::NonNullable, ))] fn zone_maps_hold_no_variant_state(#[case] dtype: DType) { - let session = SESSION.clone().with::(); + let session = VortexSession::empty().with::(); let accumulator = AggregateStatsAccumulator::new(&dtype, &default_zoned_aggregate_fns(&dtype, &session)); for aggregate_fn in accumulator.aggregate_fns().iter() { diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index ff0e8622f4b..551b60d3e45 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -14,7 +14,7 @@ use futures::TryStreamExt; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; @@ -69,8 +69,7 @@ pub(super) fn to_record_batch( data_type: &Field, ctx: &mut ExecutionCtx, ) -> VortexResult { - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(chunk, Some(data_type), ctx)?; + let arrow = chunk.execute_arrow(Some(data_type), ctx)?; Ok(RecordBatch::from(arrow.as_struct().clone())) } diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 81de637f501..37d3e64f1ac 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -55,6 +55,7 @@ use vortex::ipc::messages::EncoderMessage; use vortex::ipc::messages::MessageEncoder; use vortex::scalar_fn::fns::operators::Operator; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_python_abi::BUFFER_EXPORT_CAPSULE_NAME; use vortex_python_abi::VORTEX_BUFFER_EXPORT_VERSION; use vortex_python_abi::VORTEX_BUFFER_HOST; @@ -234,22 +235,17 @@ impl ArrowExport { let chunks = chunked_array .iter_chunks() .map(|chunk| { - session().arrow().execute_arrow( - chunk.clone(), - Some(&arrow_field), - &mut session().create_execution_ctx(), - ) + chunk + .clone() + .execute_arrow(Some(&arrow_field), &mut session().create_execution_ctx()) }) .collect::>>()?; return Ok(Self::Chunked(chunks, arrow_field)); } - let arrow_array = session().arrow().execute_arrow( - array, - target_field.as_ref(), - &mut session().create_execution_ctx(), - )?; + let arrow_array = + array.execute_arrow(target_field.as_ref(), &mut session().create_execution_ctx())?; Ok(Self::Single(arrow_array)) } diff --git a/vortex-python/src/iter/mod.rs b/vortex-python/src/iter/mod.rs index 3b8e30cd91e..ad67f4ec53e 100644 --- a/vortex-python/src/iter/mod.rs +++ b/vortex-python/src/iter/mod.rs @@ -28,6 +28,7 @@ use vortex::array::iter::ArrayIteratorExt; use vortex::dtype::DType; use vortex::error::VortexResult; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use crate::arrays::ArrowExport; use crate::arrays::PyArrayRef; @@ -170,11 +171,7 @@ impl PyArrayIterator { Box::new(RecordBatchIterator::new( iter.map(move |chunk| { let target = target.clone(); - session().arrow().execute_arrow( - chunk?, - Some(&target), - &mut session().create_execution_ctx(), - ) + chunk?.execute_arrow(Some(&target), &mut session().create_execution_ctx()) }) .map(|chunk| chunk.map_err(|e| ArrowError::ExternalError(Box::new(e)))) .map(|array| array.map(|a| RecordBatch::from(a.as_struct().clone()))), diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index 2cfbab5860d..ac83bf94e38 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -37,11 +37,23 @@ pub trait SessionExt: Sized + private::Sealed { /// Returns the session variable of type `V`, inserting a default one if it does not exist. /// - /// The default is constructed and inserted copy-on-write: `V::default()` runs without any lock - /// held, so it may freely re-enter the session, and a concurrent insert of the same type is - /// resolved by keeping the first value published. + /// This is [`get_or_insert_with`](Self::get_or_insert_with) with `V::default` as the + /// initializer. fn get(&self) -> SessionGuard<'_, V>; + /// Returns the session variable of type `V`, inserting the result of `init` if it does not + /// exist. + /// + /// Use this instead of [`get`](Self::get) when the first value of `V` depends on the session, + /// for example on another session variable. The value is constructed and inserted + /// copy-on-write: `init` runs without any lock held, so it may freely re-enter the session. A + /// concurrent insert of the same type is resolved by keeping the first value published, so + /// `init` may run even though its result is then dropped. + fn get_or_insert_with( + &self, + init: impl FnOnce() -> V, + ) -> SessionGuard<'_, V>; + /// Returns the session variable of type `V` if it exists. fn get_opt(&self) -> Option>; diff --git a/vortex-session/src/session.rs b/vortex-session/src/session.rs index b0e52d3eba3..184e816bca5 100644 --- a/vortex-session/src/session.rs +++ b/vortex-session/src/session.rs @@ -11,13 +11,13 @@ //! read never takes a lock, it can never deadlock, and because it never holds a lock across the //! returned reference there is no reader/writer contention. //! -//! * **Writes** ([`VortexSession::with_some`], [`SessionExt::get`] on a missing default) are -//! copy-on-write: the map is cloned, the change applied to the private copy, and the new map -//! atomically published. The value is constructed *before* the map is updated, so no user code -//! (in particular, no `Default::default` implementation) ever runs while a lock is held. This is -//! the key difference from the previous `DashMap`-backed session, where -//! `entry().or_insert_with(f)` ran `f` while holding the shard's write lock and could deadlock if -//! `f` re-entered the session. +//! * **Writes** ([`VortexSession::with_some`], and [`SessionExt::get`] or +//! [`SessionExt::get_or_insert_with`] on a missing variable) are copy-on-write: the map is cloned, +//! the change applied to the private copy, and the new map atomically published. The value is +//! constructed *before* the map is updated, so no user code (in particular, no `Default::default` +//! implementation or initializer closure) ever runs while a lock is held. This is the key +//! difference from the previous `DashMap`-backed session, where `entry().or_insert_with(f)` ran +//! `f` while holding the shard's write lock and could deadlock if `f` re-entered the session. //! //! A modified session is produced by mutating it **in place**: [`VortexSession::with_some`] — and //! the configuration `with_*` helpers built on it — apply their change copy-on-write to the shared @@ -185,23 +185,13 @@ impl VortexSession { Self(self.0.detach()) } - /// Inserts `V::default()` if no variable of type `V` is present yet, copy-on-write. - /// - /// The default is constructed *before* the map is updated, so `V::default()` never runs under a - /// lock and is never run more than once — it may therefore freely re-enter the session (read or - /// even register other variables) without risk of deadlock. If a concurrent writer published a - /// value first, that value is kept and the prebuilt default is simply dropped. - fn insert_default(&self) { - let default: Arc = Arc::new(V::default()); - self.0.insert_if_absent(TypeId::of::(), default); - } - /// Inserts a session variable of type `V`, replacing any existing variable of that type. /// /// This is the copy-on-write insert primitive behind [`with_some`](Self::with_some) and /// [`get_mut`](SessionExt::get_mut); it is not public, so a variable can only enter the type-map - /// through those (or through a default inserted by [`get`](SessionExt::get)). The mutation is - /// applied in place to the shared backing store, so it is visible through every clone. + /// through those (or through a value inserted by [`get`](SessionExt::get) or + /// [`get_or_insert_with`](SessionExt::get_or_insert_with)). The mutation is applied in place to + /// the shared backing store, so it is visible through every clone. pub fn register(&self, var: V) { self.0.insert(TypeId::of::(), Arc::new(var)); } @@ -259,11 +249,24 @@ impl SessionExt for VortexSession { } fn get(&self) -> SessionGuard<'_, V> { - if self.get_opt::().is_none() { - self.insert_default::(); + self.get_or_insert_with(V::default) + } + + fn get_or_insert_with( + &self, + init: impl FnOnce() -> V, + ) -> SessionGuard<'_, V> { + if let Some(var) = self.get_opt::() { + return var; } + + // NB: `init` runs before the map is updated, never inside it, because it may re-enter the + // session. If a concurrent writer published a value first, that value is kept. + let var: Arc = Arc::new(init()); + self.0.insert_if_absent(TypeId::of::(), var); + self.get_opt::() - .vortex_expect("variable was just inserted") + .vortex_expect("a variable of type `V` is present after `insert_if_absent`") } fn get_opt(&self) -> Option> { @@ -376,6 +379,19 @@ mod tests { assert!(session.get_opt::().is_some()); } + #[test] + fn get_or_insert_with_inserts_only_when_absent() { + let session = VortexSession::empty(); + + assert_eq!(session.get_or_insert_with(|| Counter { count: 3 }).count, 3); + assert_eq!( + session + .get_or_insert_with::(|| panic!("the variable is already present")) + .count, + 3 + ); + } + #[test] fn register_is_visible_through_clones() { let session = VortexSession::empty(); diff --git a/vortex-spatial/benches/area.rs b/vortex-spatial/benches/area.rs index e2bc0d6a05d..8fe3e61c0cf 100644 --- a/vortex-spatial/benches/area.rs +++ b/vortex-spatial/benches/area.rs @@ -28,7 +28,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -49,7 +48,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/binary_predicates.rs b/vortex-spatial/benches/binary_predicates.rs index d566ba61339..45fa2021387 100644 --- a/vortex-spatial/benches/binary_predicates.rs +++ b/vortex-spatial/benches/binary_predicates.rs @@ -41,7 +41,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_error::VortexResult; @@ -69,7 +68,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs index d08be00a66f..a130f6b8273 100644 --- a/vortex-spatial/benches/collect.rs +++ b/vortex-spatial/benches/collect.rs @@ -30,7 +30,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -54,7 +53,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/convex_hull.rs b/vortex-spatial/benches/convex_hull.rs index 8d18c653c00..00b5aab0038 100644 --- a/vortex-spatial/benches/convex_hull.rs +++ b/vortex-spatial/benches/convex_hull.rs @@ -29,7 +29,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -49,7 +48,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/distance.rs b/vortex-spatial/benches/distance.rs index 0f5f4e38b67..02a498cc617 100644 --- a/vortex-spatial/benches/distance.rs +++ b/vortex-spatial/benches/distance.rs @@ -29,7 +29,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_error::VortexResult; @@ -51,7 +50,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/envelope.rs b/vortex-spatial/benches/envelope.rs index f55b9a60a68..7c27496d25f 100644 --- a/vortex-spatial/benches/envelope.rs +++ b/vortex-spatial/benches/envelope.rs @@ -32,7 +32,6 @@ use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -62,7 +61,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/length.rs b/vortex-spatial/benches/length.rs index d7ba227d1d2..aad8e646d01 100644 --- a/vortex-spatial/benches/length.rs +++ b/vortex-spatial/benches/length.rs @@ -30,7 +30,6 @@ use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; use vortex_array::validity::Validity; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -50,7 +49,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/benches/make_line.rs b/vortex-spatial/benches/make_line.rs index 1d16161adc0..b585b3e6742 100644 --- a/vortex-spatial/benches/make_line.rs +++ b/vortex-spatial/benches/make_line.rs @@ -29,7 +29,6 @@ use vortex_array::dtype::session::DTypeSession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::test_harness::mark_optimized; -use vortex_arrow::ArrowSession; use vortex_edition::EditionSession; use vortex_edition::EnabledEditions; use vortex_expr::stats::session::StatsSession; @@ -49,7 +48,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() .with::() .with::() - .with::() .with::() .with::() .with::(); diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 9bb122eaf24..7e7f2de49d9 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -45,6 +45,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_buffer::Buffer; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -224,8 +225,7 @@ fn linestring_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< &SpatialMetadata::default(), linestring_dimension(storage.dtype())?, ); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + let arrow = storage.clone().execute_arrow(None, ctx)?; LineStringArray::try_from((arrow.as_ref(), linestring_type)) .map_err(|e| vortex_err!("failed to construct LineStringArray: {e}")) } @@ -311,12 +311,8 @@ impl ArrowExportVTable for LineString { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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 404f3656f5c..a4ea09631ea 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -38,6 +38,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -154,8 +155,7 @@ fn multilinestring_array( &SpatialMetadata::default(), multilinestring_dimension(storage.dtype())?, ); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + let arrow = storage.clone().execute_arrow(None, ctx)?; MultiLineStringArray::try_from((arrow.as_ref(), multilinestring_type)) .map_err(|e| vortex_err!("failed to construct MultiLineStringArray: {e}")) } @@ -241,12 +241,8 @@ impl ArrowExportVTable for MultiLineString { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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 bd219dcf4ca..e57cc57c5be 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -38,6 +38,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -141,8 +142,7 @@ fn multipoint_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< &SpatialMetadata::default(), multipoint_dimension(storage.dtype())?, ); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + let arrow = storage.clone().execute_arrow(None, ctx)?; MultiPointArray::try_from((arrow.as_ref(), multipoint_type)) .map_err(|e| vortex_err!("failed to construct MultiPointArray: {e}")) } @@ -228,12 +228,8 @@ impl ArrowExportVTable for MultiPoint { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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 9f67ce94747..711bdb016d3 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -37,6 +37,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -151,8 +152,7 @@ fn multipolygon_array( &SpatialMetadata::default(), multipolygon_dimension(storage.dtype())?, ); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + let arrow = storage.clone().execute_arrow(None, ctx)?; MultiPolygonArray::try_from((arrow.as_ref(), multipolygon_type)) .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}")) } @@ -238,12 +238,8 @@ impl ArrowExportVTable for MultiPolygon { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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 9de00619d8a..46b465289ec 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -36,6 +36,7 @@ use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -129,8 +130,7 @@ fn point_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult BoxType /// Build a geoarrow `RectArray` from a `Rect`'s box `Struct` storage. fn rect_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let box_type = box_type(&SpatialMetadata::default(), box_dimension(storage.dtype())?); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + let arrow = storage.clone().execute_arrow(None, ctx)?; RectArray::try_from((arrow.as_ref(), box_type)) .map_err(|e| vortex_err!("failed to construct RectArray: {e}")) } @@ -259,12 +258,8 @@ impl ArrowExportVTable for Rect { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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 8b0b42fb97b..47713736c7a 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -30,7 +30,7 @@ use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -220,12 +220,8 @@ impl ArrowExportVTable for WellKnownBinary { target.data_type().clone(), target.is_nullable(), ); - let session = ctx.session().clone(); let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(&storage_field), ctx)?; + storage.execute_arrow_with_options(Some(&storage_field), options, 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-spatial/src/tests/point.rs b/vortex-spatial/src/tests/point.rs index aec2e91bcbf..f70882963c8 100644 --- a/vortex-spatial/src/tests/point.rs +++ b/vortex-spatial/src/tests/point.rs @@ -23,6 +23,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -84,9 +85,7 @@ fn exports_to_struct() -> VortexResult<()> { let array = point_column(vec![1.0, 3.0], vec![2.0, 4.0])?; let target = point_field("loc", false, Some("EPSG:4326")); - let exported = SESSION - .arrow() - .execute_arrow(array, Some(&target), &mut ctx)?; + let exported = array.execute_arrow(Some(&target), &mut ctx)?; let points = exported.as_struct(); let ordinates = |name: &str| -> VortexResult> { @@ -168,9 +167,7 @@ fn roundtrips_through_arrow() -> VortexResult<()> { let original = point_column(vec![0.0, 3.0], vec![4.0, 0.0])?; let target = point_field("loc", false, Some("EPSG:4326")); - let exported = SESSION - .arrow() - .execute_arrow(original, Some(&target), &mut ctx)?; + let exported = original.execute_arrow(Some(&target), &mut ctx)?; let reimported = SESSION.arrow().from_arrow_array(exported, &target)?; let ext = reimported diff --git a/vortex-spatial/src/tests/rect.rs b/vortex-spatial/src/tests/rect.rs index 9639e5ba400..e63eb523154 100644 --- a/vortex-spatial/src/tests/rect.rs +++ b/vortex-spatial/src/tests/rect.rs @@ -16,6 +16,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -78,9 +79,7 @@ fn roundtrips_through_arrow() -> VortexResult<()> { let original = rect_column(vec![(0.0, 1.0, 2.0, 3.0), (-5.0, -5.0, 5.0, 5.0)])?; let target = box_field("bbox", GeoArrowDimension::XY, false, Some("EPSG:4326")); - let exported = SESSION - .arrow() - .execute_arrow(original, Some(&target), &mut ctx)?; + let exported = original.execute_arrow(Some(&target), &mut ctx)?; let reimported = SESSION.arrow().from_arrow_array(exported, &target)?; let ext = reimported diff --git a/vortex-spatial/src/tests/wkb.rs b/vortex-spatial/src/tests/wkb.rs index 05b14ad1942..d8ae5455072 100644 --- a/vortex-spatial/src/tests/wkb.rs +++ b/vortex-spatial/src/tests/wkb.rs @@ -34,6 +34,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use wkb::writer::WriteOptions; @@ -146,9 +147,7 @@ fn exports_to_binary_family(#[case] data_type: DataType) -> VortexResult<()> { let field = Field::new("geom", data_type.clone(), false) .with_extension_type(WkbType::new(Default::default())); - let exported = SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut ctx)?; + let exported = array.execute_arrow(Some(&field), &mut ctx)?; assert_eq!(exported.data_type(), &data_type); let values: Vec<&[u8]> = match &data_type { @@ -266,9 +265,7 @@ fn roundtrips_through_arrow() -> VortexResult<()> { let field = Field::new("geom", DataType::BinaryView, false) .with_extension_type(WkbType::new(Default::default())); - let exported = SESSION - .arrow() - .execute_arrow(original, Some(&field), &mut ctx)?; + let exported = original.execute_arrow(Some(&field), &mut ctx)?; let arrow_field = wkb_field("geom", DataType::BinaryView, false, Some("EPSG:4326")); let reimported = SESSION.arrow().from_arrow_array(exported, &arrow_field)?; diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index 6e9a5cd1077..fafe66297c3 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -29,7 +29,7 @@ use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowImportVTable; use vortex_arrow::ArrowSession; -use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -107,12 +107,7 @@ impl ArrowExportVTable for Vector { let executed = array.execute::(ctx)?; let storage = executed.storage_array().clone(); - let session = ctx.session().clone(); - let arrow_storage = - session - .arrow() - .exporter(options) - .execute_arrow(storage, Some(target), ctx)?; + let arrow_storage = storage.execute_arrow_with_options(Some(target), options, ctx)?; Ok(ArrowExport::Exported(arrow_storage)) } @@ -197,9 +192,9 @@ mod tests { use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowImport; - use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::tests::SESSION; @@ -237,16 +232,18 @@ mod tests { )) } - fn session_with_vector() -> ArrowSession { - let session = ArrowSession::default(); - session.register_exporter(Arc::new(Vector)); - session.register_importer(Arc::new(Vector)); + /// A session whose only Arrow plugins are the [`Vector`] exporter and importer. + fn session_with_vector() -> VortexSession { + let session = VortexSession::empty(); + session.arrow().register_exporter(Arc::new(Vector)); + session.arrow().register_importer(Arc::new(Vector)); session } #[test] fn to_arrow_field_attaches_extension_metadata() -> VortexResult<()> { - let session = session_with_vector(); + let vortex_session = session_with_vector(); + let session = vortex_session.arrow(); let field = session.to_arrow_field("embedding", &vector_dtype(false))?; assert_eq!( field.extension_type_name(), @@ -265,7 +262,8 @@ mod tests { #[test] fn from_arrow_field_recovers_vector_dtype() -> VortexResult<()> { - let session = session_with_vector(); + let vortex_session = session_with_vector(); + let session = vortex_session.arrow(); let arrow_field = Field::new( "embedding", DataType::FixedSizeList( @@ -282,7 +280,8 @@ mod tests { #[test] fn schema_roundtrip_preserves_top_level_vector() -> VortexResult<()> { - let session = session_with_vector(); + let vortex_session = session_with_vector(); + let session = vortex_session.arrow(); let dtype = DType::Struct( StructFields::from_iter([(FieldName::from("embedding"), vector_dtype(false))]), Nullability::NonNullable, @@ -295,7 +294,8 @@ mod tests { #[test] fn schema_roundtrip_preserves_nested_struct_vector() -> VortexResult<()> { - let session = session_with_vector(); + let vortex_session = session_with_vector(); + let session = vortex_session.arrow(); let inner = DType::Struct( StructFields::from_iter([(FieldName::from("embedding"), vector_dtype(true))]), Nullability::NonNullable, @@ -312,7 +312,8 @@ mod tests { #[test] fn schema_roundtrip_preserves_list_of_vector() -> VortexResult<()> { - let session = session_with_vector(); + let vortex_session = session_with_vector(); + let session = vortex_session.arrow(); let dtype = DType::Struct( StructFields::from_iter([( FieldName::from("embeddings"), @@ -332,7 +333,7 @@ mod tests { let session = SESSION.arrow(); let original = sample_vector_array(); let field = session.to_arrow_field("embedding", original.dtype())?; - let arrow = session.execute_arrow(original.clone(), Some(&field), &mut ctx)?; + let arrow = original.clone().execute_arrow(Some(&field), &mut ctx)?; assert!(matches!(arrow.data_type(), DataType::FixedSizeList(_, n) if *n == DIM as i32)); @@ -351,7 +352,8 @@ mod tests { // Session with no Vector plugin must not error on a vector.tensor.vector-tagged field; // it should fall through to the canonical Arrow → Vortex path and drop the extension // identity, producing the raw FSL storage instead. - let session = ArrowSession::default(); + let unregistered = VortexSession::empty(); + let session = unregistered.arrow(); let mut arrow_field = Field::new( "embedding", DataType::FixedSizeList( @@ -407,7 +409,7 @@ mod tests { int_array, &field, &dtype, - &ArrowSession::default(), + &VortexSession::empty().arrow(), )?; assert!(matches!(result, ArrowImport::Unsupported(_))); Ok(()) @@ -432,7 +434,7 @@ mod tests { fsl_arrow, &field, &DType::Extension(uuid_ext), - &ArrowSession::default(), + &VortexSession::empty().arrow(), )?; assert!(matches!(result, ArrowImport::Unsupported(_))); Ok(()) @@ -443,7 +445,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let session = SESSION.arrow(); let original = sample_vector_array(); - let arrow = session.execute_arrow(original.clone(), None, &mut ctx)?; + let arrow = original.clone().execute_arrow(None, &mut ctx)?; let field = session.to_arrow_field("v", original.dtype())?; assert_eq!( diff --git a/vortex-test/e2e-cuda/src/lib.rs b/vortex-test/e2e-cuda/src/lib.rs index c573a70cdda..f432a1fc7ab 100644 --- a/vortex-test/e2e-cuda/src/lib.rs +++ b/vortex-test/e2e-cuda/src/lib.rs @@ -47,7 +47,7 @@ use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::varbinview::BinaryView; use vortex::array::stream::ArrayStreamExt; use vortex::array::validity::Validity; -use vortex::arrow::ArrowSessionExt; +use vortex::arrow::ExecuteArrowExt; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; @@ -417,13 +417,9 @@ fn validate_array_inner(ffi_schema: &FFI_ArrowSchema, ffi_array: &mut FFI_ArrowA let array = make_array(array_data); let struct_array = array.as_struct(); - let primitive = SESSION - .arrow() - .execute_arrow( - primitive_array().expect("expected primitive array"), - None, - &mut SESSION.create_execution_ctx(), - ) + let primitive = primitive_array() + .expect("expected primitive array") + .execute_arrow(None, &mut SESSION.create_execution_ctx()) .expect("expected primitive Arrow array"); let bools = BooleanArray::from(vec![true, false, false, true, true]); let sliced_i32 = Int32Array::from(vec![Some(10), None, Some(30), Some(40), None]); @@ -473,17 +469,11 @@ fn validate_array_inner(ffi_schema: &FFI_ArrowSchema, ffi_array: &mut FFI_ArrowA .into_iter() .collect::>(), ); - let list = SESSION - .arrow() - .execute_arrow(list_array(), None, &mut SESSION.create_execution_ctx()) + let list = list_array() + .execute_arrow(None, &mut SESSION.create_execution_ctx()) .expect("expected list Arrow array"); - let fixed_size_list = SESSION - .arrow() - .execute_arrow( - fixed_size_list_as_list_array(), - None, - &mut SESSION.create_execution_ctx(), - ) + let fixed_size_list = fixed_size_list_as_list_array() + .execute_arrow(None, &mut SESSION.create_execution_ctx()) .expect("expected fixed-size-list-as-list Arrow array"); let expected_fields = Fields::from_iter([ diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index f5f91daeb1f..e9f357de202 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -40,6 +40,7 @@ use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::session::VortexSession; use vortex::session::registry::ReadContext; use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ExecuteArrowExt; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::JsFuture; @@ -781,9 +782,7 @@ fn array_to_record_batch( }; let target = Field::new("", data_type, array.dtype().is_nullable()); let mut ctx = session.create_execution_ctx(); - let arrow = session - .arrow() - .execute_arrow(array, Some(&target), &mut ctx)?; + let arrow = array.execute_arrow(Some(&target), &mut ctx)?; match dtype { DType::Struct(..) => Ok(RecordBatch::from(arrow.as_struct().clone())), _ => Ok(RecordBatch::try_new(schema.clone(), vec![arrow])?), diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 30e148db3c7..dc47180cfa1 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -346,7 +346,6 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); - vortex_arrow::initialize(&session); #[cfg(feature = "parquet-variant")] vortex_parquet_variant::initialize(&session); vortex_uuid::initialize(&session); @@ -616,6 +615,7 @@ mod test { use vortex_array::dtype::extension::ExtVTable; use vortex_array::dtype::session::DTypeSessionExt; use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; use vortex_uuid::Uuid; let session = VortexSession::default(); @@ -634,7 +634,7 @@ mod test { assert!(array.dtype().as_extension().is::()); let mut ctx = session.create_execution_ctx(); - let exported = session.arrow().execute_arrow(array, None, &mut ctx)?; + let exported = array.execute_arrow(None, &mut ctx)?; let fsb = exported.as_fixed_size_binary(); assert_eq!(fsb.len(), 2); assert_eq!(fsb.value(0), b"0123456789abcdef");