From 19ae8535484287862eb5f7569a9176d0987ac84c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 9 Nov 2026 10:54:34 -0500 Subject: [PATCH 1/6] Compare nested extension types by storage in Vortex ID exports (#11351) ## Summary - Tracking Issue: #11114 Follow-up to #11301, which checks that the Arrow type of a field that names a Vortex ID imports as the storage of the array. That import can lose information about nested extension types: the UUID plugin imports every `arrow.uuid` field without a version, and a session that does not know a nested Vortex ID imports it as plain storage. So since #11301, some valid exports fail. For example, an extension type whose storage is a struct with a `Uuid(v4)` field can no longer export to its own `to_arrow_storage_field`, which worked before. ## Changes The check now compares nested extension types by their storage types, and keeps the exact comparison at the top level. Nested fields keep their checks, because the exporter checks each nested field when it reaches it, and the temporal executor still rejects a time unit that does not match. The alternative was to require every import plugin to give back exactly the type that it exported. `ArrowImportVTable` does not require that today, and a plugin for a canonical Arrow type such as `arrow.uuid` has no standard place for Vortex-only metadata such as the UUID version. Exact round trips for the plugins that Vortex owns are still worth having, but as a separate goal with its own property tests, not as an assumption of this check. This PR also does the comment pass from the review of #11301 on the export dispatch, the map import, and the map test. The map import still runs import plugins for the extension name of the entries field, which `from_arrow_datatype` ignores. The comment now says so instead of claiming that the two match, and a test relies on those plugins running, so the behavior is unchanged. The metadata check that the import and the export both copied is now one helper. New tests cover the nested case (it fails without the fix), an alias of a Vortex ID, and a plain array that exports to a field that names a Vortex ID. Signed-off-by: Connor Tsui --- vortex-arrow/src/session.rs | 107 ++++++++++++++++++++++++-------- vortex-arrow/tests/vortex_id.rs | 74 +++++++++++++++++++++- 2 files changed, 151 insertions(+), 30 deletions(-) diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 651c72452bd..8320feb76f6 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -267,11 +267,13 @@ 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 @@ -529,11 +531,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); } @@ -969,10 +967,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( @@ -1130,8 +1128,8 @@ impl ArrowExporter<'_> { } } - // The registry can hold an alias, so compare the array with the ID of the dtype plugin - // that the extension name resolves to. + // 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 = self .session .ext_dtypes @@ -1149,9 +1147,9 @@ impl ArrowExporter<'_> { 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. + // 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) { @@ -1182,25 +1180,21 @@ impl ArrowExporter<'_> { 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() - ); + ensure_ext_metadata(target, &decode_arrow_ext_metadata(target)?, ext_dtype)?; 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 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()), self, ctx); } vortex_ensure!( - target_dtype.eq_ignore_nullability(ext_dtype.storage_dtype()), + 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() @@ -1215,6 +1209,65 @@ impl ArrowExporter<'_> { } } +/// Fails unless `metadata`, the extension metadata of `field`, is the metadata of `ext_dtype`. +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(()) +} + +/// 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 +} + /// Extension trait for accessing the [`ArrowSession`] on a Vortex session. pub trait ArrowSessionExt: SessionExt { /// Get the Arrow session. diff --git a/vortex-arrow/tests/vortex_id.rs b/vortex-arrow/tests/vortex_id.rs index 584feee6cbd..92d66d0148a 100644 --- a/vortex-arrow/tests/vortex_id.rs +++ b/vortex-arrow/tests/vortex_id.rs @@ -58,16 +58,24 @@ 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); + session + .dtypes() + .registry() + .insert(*TAGGED_ALIAS_ID, Arc::new(Tagged)); vortex_arrow::initialize(&session); 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)] @@ -334,6 +342,66 @@ 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 = SESSION + .arrow() + .execute_arrow(array.clone(), 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 = SESSION + .arrow() + .execute_arrow(storage.clone(), 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(); + vortex_arrow::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let exported = session + .arrow() + .execute_arrow(array.clone(), 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<()> { @@ -426,8 +494,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)); From 31283ee1cb15459b2fd68aa87479b0335ba744b7 Mon Sep 17 00:00:00 2001 From: ChristopherWirt Date: Mon, 9 Nov 2026 16:22:08 +0000 Subject: [PATCH 2/6] feat(scan): let a scan request bound the splits it decodes at once (#11127) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Tracking Issue: #11124 A scan decodes `4 × available_parallelism()` splits at once. In Spark, every task in an executor runs its own scan sized to the whole machine, and each split holds the decoded chunks of every column it overlaps. On wide, nested files that is GiBs of native memory per executor (#11123). This adds a per-scan limit on how many splits are decoded at once. | Layer | API | |---|---| | Rust | `ScanRequest::split_concurrency` → `ScanBuilder::with_some_max_splits_in_flight` | | JNI | `ScanOptions.splitConcurrency()` | | Spark | `vortex.splitConcurrency` read option | **Result:** a full scan of a 180 MB list-layout file, peak live native bytes: | Splits at once | Peak live | Time | |---|---|---| | 48 (default, 4 × 12 cores) | 3.65 GiB | 1.0 s | | 12 | 1.18 GiB | 1.0 s | --------- Signed-off-by: Christopher Wirt --- docs/user-guide/spark.md | 1 + .../main/java/dev/vortex/api/DataSource.java | 4 +- .../main/java/dev/vortex/api/ScanOptions.java | 20 ++++ .../main/java/dev/vortex/jni/NativeScan.java | 4 +- .../java/dev/vortex/api/ScanOptionsTest.java | 16 +++ .../test/java/dev/vortex/api/TestMinimal.java | 12 ++ .../java/dev/vortex/spark/io/VortexIo.java | 47 ++++++-- .../spark/read/VortexPartitionReader.java | 3 + .../vortex/spark/VortexV1FallbackTest.java | 17 +++ .../dev/vortex/spark/io/VortexIoTest.java | 27 ++++- vortex-ffi/src/scan.rs | 1 + vortex-jni/src/scan.rs | 103 +++++++++++++++--- vortex-layout/src/scan/layout.rs | 8 +- vortex-layout/src/scan/multi.rs | 66 ++++++++++- vortex-layout/src/scan/repeated_scan.rs | 37 ++++++- vortex-layout/src/scan/scan_builder.rs | 30 ++++- vortex-layout/src/scan/test.rs | 31 ++++++ vortex-scan/src/lib.rs | 9 ++ 18 files changed, 398 insertions(+), 38 deletions(-) diff --git a/docs/user-guide/spark.md b/docs/user-guide/spark.md index 8bd797cba0a..f7d25e803df 100644 --- a/docs/user-guide/spark.md +++ b/docs/user-guide/spark.md @@ -196,6 +196,7 @@ own storage clients are not used by the Spark connector. | Option | Default | Description | |----------------------------|---------|-------------| | `vortex.readConcurrency` | `0` | Maximum Hadoop read upcalls per file; `0` uses the native default. | +| `vortex.splitConcurrency` | `0` | Most row splits of a file decoded at once by one task; `0` uses the native default of four per core. Lowering this reduces memory held by concurrent split decoding and Arrow conversion. A scan with both a filter and a limit may also retain one trimmed array per split while preparing output. This does not cap total native memory or I/O read-ahead. | | `vortex.workerThreads` | `4` | Background threads driving Vortex futures. JVM-wide, and set once per JVM by the first read or write. | | `vortex.session.provider` | — | Fully-qualified name of a `VortexSessionProvider` supplying a custom session on the driver and on every executor. | | `vortex.aggregatePushdown` | `true` | Answer `COUNT(*)` from file footers. | diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java index 8658cf5198b..22c090c6419 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java @@ -238,6 +238,7 @@ public Scan scan(ScanOptions options) { byte[] selectionRoaringBitmap = selectionRoaringBitmap(options); long limit = options.limit().orElse(0L); boolean ordered = options.ordered(); + int splitConcurrency = options.splitConcurrency().orElse(0); long scanPtr = dev.vortex.jni.NativeScan.create( pointer, @@ -249,7 +250,8 @@ public Scan scan(ScanOptions options) { selectionRoaringBitmap, selectionMode.code(), limit, - ordered); + ordered, + splitConcurrency); return Scan.fromPointer(session, scanPtr); } diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java b/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java index 25f03e76f7f..43e713581ad 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java @@ -9,6 +9,7 @@ import java.io.UncheckedIOException; import java.util.Objects; import java.util.Optional; +import java.util.OptionalInt; import java.util.OptionalLong; import org.immutables.value.Value; import org.roaringbitmap.longlong.Roaring64NavigableMap; @@ -59,6 +60,25 @@ default boolean ordered() { return false; } + /** + * The most splits of each partition the native scan decodes at once, in all, and the most of its batches converted + * to Arrow ahead of the caller. Absent keeps the native default, which scales with the machine's cores. + * + *

Each split being decoded holds the decoded chunks of every column it overlaps, so a scan running alongside + * others, such as one Spark task of several in an executor, can lower this to reduce memory held by concurrent + * decoding and Arrow conversion. A scan with both a filter and a limit may also retain one trimmed array per split + * while preparing output. This does not cap total native memory or I/O read-ahead. + */ + OptionalInt splitConcurrency(); + + @Value.Check + default void validateSplitConcurrency() { + if (splitConcurrency().isPresent() && splitConcurrency().getAsInt() <= 0) { + throw new IllegalArgumentException("split concurrency must be positive, got " + + splitConcurrency().getAsInt()); + } + } + @Value.Check default void validateSelectionPayload() { boolean hasIndices = selectionIndices().isPresent(); diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java index ff80b30caaa..ff4ad73a317 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java @@ -27,6 +27,7 @@ private NativeScan() {} * {@code selectionRoaringBitmap}) * @param limit max rows to return, or {@code 0} for "no limit" * @param ordered true to preserve row order across partitions + * @param splitConcurrency the most splits of each partition read at once, or {@code 0} for the native default */ public static native long create( long dataSourcePointer, @@ -38,7 +39,8 @@ public static native long create( byte[] selectionRoaringBitmap, byte selectionInclude, long limit, - boolean ordered); + boolean ordered, + int splitConcurrency); /** Free a scan pointer. */ public static native void free(long pointer); diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java index c72e67f4051..cb0c6ff5b8a 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java @@ -12,6 +12,7 @@ import dev.vortex.api.ScanOptions.SelectionMode; import java.io.ByteArrayInputStream; import java.io.DataInputStream; +import java.util.OptionalInt; import java.util.OptionalLong; import org.junit.jupiter.api.Test; import org.roaringbitmap.longlong.Roaring64NavigableMap; @@ -39,6 +40,21 @@ public void defaultsReadEveryRowAndColumn() { assertTrue(options.selectionRoaringBitmap().isEmpty()); assertEquals(SelectionMode.INCLUDE_ALL, options.selectionMode()); assertFalse(options.ordered()); + assertEquals(OptionalInt.empty(), options.splitConcurrency()); + } + + @Test + public void splitConcurrencyMustBePositive() { + for (int value : new int[] {0, -1}) { + IllegalArgumentException exception = assertThrows( + IllegalArgumentException.class, + () -> ScanOptions.builder().splitConcurrency(value).build(), + () -> "no exception for " + value); + assertEquals("split concurrency must be positive, got " + value, exception.getMessage()); + } + assertEquals( + OptionalInt.of(2), + ScanOptions.builder().splitConcurrency(2).build().splitConcurrency()); } @Test diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java index 2b270c3e165..f4adefed093 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java @@ -184,6 +184,18 @@ public void testProjectedScanWithFilter() throws Exception { assertEquals(List.of(new Person("John", BigDecimal.valueOf(10_000L, 2), "VA")), people); } + @Test + public void testScanWithSplitConcurrencyReadsEveryRow() throws Exception { + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + DataSource ds = DataSource.open(session, writePath); + + ScanOptions options = + ScanOptions.builder().splitConcurrency(1).ordered(true).build(); + List people = readAll(ds, options, allocator, TestMinimal::readFullBatch); + assertEquals(MINIMAL_DATA, people); + } + @Test public void testSelectionIncludesRows() throws Exception { BufferAllocator allocator = ArrowAllocation.rootAllocator(); diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java index 7655c9e5f88..d59f20392ac 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java @@ -34,6 +34,17 @@ public final class VortexIo implements Serializable { private static final int DEFAULT_READ_CONCURRENCY = 0; + /** + * Option bounding how many splits of a file the native scan reads at once, in all. Each split in flight holds the + * decoded chunks of every column it overlaps, and every task in an executor runs its own scan, which by default + * keeps several splits in flight per core. Lowering this reduces memory held by concurrent split decoding and Arrow + * conversion. A scan with both a filter and a limit may also retain one trimmed array per split while preparing + * output. This does not cap total native memory or I/O read-ahead. Zero keeps the native default. + */ + public static final String SPLIT_CONCURRENCY_OPTION = "vortex.splitConcurrency"; + + private static final int DEFAULT_SPLIT_CONCURRENCY = 0; + /** The configuration this instance carries itself, or null when it reads {@link #broadcastConf}. */ private final SerializableHadoopConf conf; @@ -41,12 +52,17 @@ public final class VortexIo implements Serializable { private final Broadcast broadcastConf; private final int readConcurrency; + private final int splitConcurrency; /** Captures {@code hadoopConf} and the read settings found in the format options. */ public static VortexIo create(VortexOptions options, Configuration hadoopConf) { Objects.requireNonNull(options, "options"); Objects.requireNonNull(hadoopConf, "hadoopConf"); - return new VortexIo(new SerializableHadoopConf(hadoopConf), null, parseReadConcurrency(options)); + return new VortexIo( + new SerializableHadoopConf(hadoopConf), + null, + nonNegative(options, READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY), + nonNegative(options, SPLIT_CONCURRENCY_OPTION, DEFAULT_SPLIT_CONCURRENCY)); } /** @@ -58,26 +74,36 @@ public static VortexIo broadcast(VortexOptions options, Configuration hadoopConf Objects.requireNonNull(hadoopConf, "hadoopConf"); ClassTag tag = ClassTag$.MODULE$.apply(SerializableHadoopConf.class); return new VortexIo( - null, context.broadcast(new SerializableHadoopConf(hadoopConf), tag), parseReadConcurrency(options)); + null, + context.broadcast(new SerializableHadoopConf(hadoopConf), tag), + nonNegative(options, READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY), + nonNegative(options, SPLIT_CONCURRENCY_OPTION, DEFAULT_SPLIT_CONCURRENCY)); } /** Hadoop I/O over a default configuration, for callers that have no Spark session to draw one from. */ public static VortexIo defaults() { - return new VortexIo(new SerializableHadoopConf(new Configuration()), null, DEFAULT_READ_CONCURRENCY); + return new VortexIo( + new SerializableHadoopConf(new Configuration()), + null, + DEFAULT_READ_CONCURRENCY, + DEFAULT_SPLIT_CONCURRENCY); } private VortexIo( - SerializableHadoopConf conf, Broadcast broadcastConf, int readConcurrency) { + SerializableHadoopConf conf, + Broadcast broadcastConf, + int readConcurrency, + int splitConcurrency) { this.conf = conf; this.broadcastConf = broadcastConf; this.readConcurrency = readConcurrency; + this.splitConcurrency = splitConcurrency; } - private static int parseReadConcurrency(VortexOptions options) { - int parsed = options.getInt(READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY); + private static int nonNegative(VortexOptions options, String key, int defaultValue) { + int parsed = options.getInt(key, defaultValue); if (parsed < 0) { - throw new IllegalArgumentException( - String.format(Locale.ROOT, "%s must be >= 0, got %d", READ_CONCURRENCY_OPTION, parsed)); + throw new IllegalArgumentException(String.format(Locale.ROOT, "%s must be >= 0, got %d", key, parsed)); } return parsed; } @@ -91,6 +117,11 @@ public int readConcurrency() { return readConcurrency; } + /** The most splits of a file the native scan reads at once; zero keeps the native default. */ + public int splitConcurrency() { + return splitConcurrency; + } + /** * Opens a byte source for {@code file}, stating it only if the listing that produced it did not report a size. The * caller owns the result and must close it once the scan built on it is done. diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java index 94189713274..b651a026dd4 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java @@ -81,6 +81,9 @@ public VortexPartitionReader( List fileColumns = dataSource.arrowSchema(allocator).getFields(); var options = ScanOptions.builder(); + if (io.splitConcurrency() > 0) { + options.splitConcurrency(io.splitConcurrency()); + } // Always project, an empty read schema included: a query that needs only partition columns or only a // row count must not pull every data column off storage. Struct columns that Spark pruned read only the // nested fields the query uses. 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/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java index a7800f25786..019f811fca0 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java @@ -18,10 +18,30 @@ final class VortexIoTest { @Test - void readConcurrencyDefaultsToNativeDefault() { + void concurrenciesDefaultToNativeDefaults() { VortexIo io = VortexIo.create(VortexOptions.empty(), new Configuration()); assertEquals(0, io.readConcurrency()); + assertEquals(0, io.splitConcurrency()); + } + + @Test + void splitConcurrencyIsParsedAndValidated() { + Configuration conf = new Configuration(); + + assertEquals( + 2, + VortexIo.create(options(VortexIo.SPLIT_CONCURRENCY_OPTION, "2"), conf) + .splitConcurrency()); + assertEquals( + 2, + VortexIo.create(options("VORTEX.SPLITCONCURRENCY", "2"), conf).splitConcurrency()); + assertThrows( + IllegalArgumentException.class, + () -> VortexIo.create(options(VortexIo.SPLIT_CONCURRENCY_OPTION, "-1"), conf)); + assertThrows( + IllegalArgumentException.class, + () -> VortexIo.create(options(VortexIo.SPLIT_CONCURRENCY_OPTION, "many"), conf)); } @Test @@ -54,11 +74,14 @@ void readConcurrencyIsFoundUnderAnyKeyCase() { void settingsAndConfigurationSurviveSerialization() throws IOException, ClassNotFoundException { Configuration conf = new Configuration(); conf.set("fs.s3a.endpoint", "https://storage.example"); - VortexIo io = VortexIo.create(options(VortexIo.READ_CONCURRENCY_OPTION, "4"), conf); + VortexIo io = VortexIo.create( + VortexOptions.of(Map.of(VortexIo.READ_CONCURRENCY_OPTION, "4", VortexIo.SPLIT_CONCURRENCY_OPTION, "3")), + conf); VortexIo shipped = roundTrip(io); assertEquals(4, shipped.readConcurrency()); + assertEquals(3, shipped.splitConcurrency()); assertEquals("https://storage.example", shipped.hadoopConf().get("fs.s3a.endpoint")); } diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 1b842303043..9e1d37b278f 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -208,6 +208,7 @@ fn scan_request(opts: *const vx_scan_options) -> VortexResult { limit, partition_selection: Selection::All, partition_range: None, + split_concurrency: None, }) } diff --git a/vortex-jni/src/scan.rs b/vortex-jni/src/scan.rs index 263a43a2906..1f67bb12dd2 100644 --- a/vortex-jni/src/scan.rs +++ b/vortex-jni/src/scan.rs @@ -11,6 +11,7 @@ //! `FFI_ArrowArrayStream` that Java imports via Arrow's C Data Interface. use std::io::Cursor; +use std::num::NonZeroUsize; use std::ops::Range; use std::ptr; use std::sync::Arc; @@ -27,12 +28,14 @@ use jni::objects::JByteArray; use jni::objects::JClass; use jni::objects::JLongArray; use jni::sys::jboolean; +use jni::sys::jint; use jni::sys::jlong; use vortex::array::VortexSessionExecute; use vortex::array::stream::SendableArrayStream; use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_err; use vortex::expr::Expression; use vortex::expr::root; use vortex::expr::stats::Precision; @@ -55,17 +58,21 @@ use crate::session::session_ref; /// Opaque scan handle. Holds a three-state machine: either the scan is pending (not yet /// realized as a stream), actively streaming partitions, or finished. +/// +/// The scan's `split_concurrency`, if set, also bounds how many of each partition's arrays are +/// converted to Arrow ahead of Java, so the conversion stays within the scan's bound. #[allow(dead_code)] pub(crate) enum NativeScan { - Pending(Box), - Started(PartitionStream), + Pending(Box, Option), + Started(PartitionStream, Option), Finished, } -/// Opaque partition handle with the same three-state machine shape. +/// Opaque partition handle with the same three-state machine shape, carrying its scan's +/// `split_concurrency`. #[allow(dead_code)] pub(crate) enum NativePartition { - Pending(PartitionRef), + Pending(PartitionRef, Option), Started(SendableArrayStream), Finished, } @@ -81,6 +88,7 @@ fn build_scan_request( selection_include: u8, limit: jlong, ordered: jboolean, + split_concurrency: jint, ) -> VortexResult { let projection = if projection_ptr == 0 { root() @@ -113,6 +121,10 @@ fn build_scan_request( }); let limit = (limit > 0).then_some(limit as u64); + let split_concurrency = + NonZeroUsize::new(usize::try_from(split_concurrency).map_err(|_| { + vortex_err!("split concurrency must not be negative, got {split_concurrency}") + })?); Ok(ScanRequest { projection, @@ -123,6 +135,7 @@ fn build_scan_request( limit, partition_selection: Selection::All, partition_range: None, + split_concurrency, }) } @@ -149,6 +162,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_create( selection_include: jni::sys::jbyte, limit: jlong, ordered: jboolean, + split_concurrency: jint, ) -> jlong { try_or_throw(&mut env, |env| { let ds = unsafe { NativeDataSource::from_ptr(data_source_ptr) }; @@ -185,10 +199,12 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_create( selection_include as u8, limit, ordered, + split_concurrency, )?; + let split_concurrency = request.split_concurrency; let scan = RUNTIME.block_on(async { ds.inner().scan(request).await })?; - Ok(Box::into_raw(Box::new(NativeScan::Pending(scan))) as jlong) + Ok(Box::into_raw(Box::new(NativeScan::Pending(scan, split_concurrency))) as jlong) }) } @@ -219,7 +235,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_arrowSchema( } let session = unsafe { session_ref(session_ptr) }; let scan = unsafe { &*(pointer as *const NativeScan) }; - let NativeScan::Pending(scan) = scan else { + let NativeScan::Pending(scan, _) = scan else { throw_runtime!("schema unavailable: scan already started"); }; let arrow_schema = session.arrow().to_arrow_schema(scan.dtype())?; @@ -241,7 +257,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_partitionCount( ) { try_or_throw(&mut env, |env| { let scan = unsafe { &*(pointer as *const NativeScan) }; - let NativeScan::Pending(scan) = scan else { + let NativeScan::Pending(scan, _) = scan else { throw_runtime!("partition count unavailable: scan already started"); }; let (rows, cardinality) = match scan.partition_count() { @@ -271,17 +287,20 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_nextPartition( // only rewrite it on the happy path where we want to keep streaming. ptr::write(slot, NativeScan::Finished); - let mut stream = match owned { - NativeScan::Pending(scan) => scan.partitions(), - NativeScan::Started(stream) => stream, + let (mut stream, split_concurrency) = match owned { + NativeScan::Pending(scan, split_concurrency) => (scan.partitions(), split_concurrency), + NativeScan::Started(stream, split_concurrency) => (stream, split_concurrency), NativeScan::Finished => return Ok(0), }; match RUNTIME.block_on(stream.next()) { Some(partition) => { let partition = partition?; - let handle = Box::into_raw(Box::new(NativePartition::Pending(partition))) as jlong; - ptr::write(slot, NativeScan::Started(stream)); + let handle = Box::into_raw(Box::new(NativePartition::Pending( + partition, + split_concurrency, + ))) as jlong; + ptr::write(slot, NativeScan::Started(stream, split_concurrency)); Ok(handle) } None => Ok(0), @@ -311,7 +330,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_rowCount( ) { try_or_throw(&mut env, |env| { let partition = unsafe { &*(pointer as *const NativePartition) }; - let NativePartition::Pending(partition) = partition else { + let NativePartition::Pending(partition, _) = partition else { throw_runtime!("row count unavailable: partition already started"); }; let (rows, cardinality) = match partition.row_count() { @@ -346,10 +365,13 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( throw_runtime!("null arrow stream address"); } - let partition = match partition { - NativePartition::Pending(p) => p, + let (partition, split_concurrency) = match partition { + NativePartition::Pending(p, split_concurrency) => (p, split_concurrency), _ => throw_runtime!("partition already consumed"), }; + let lookahead = split_concurrency + .map_or(usize::MAX, NonZeroUsize::get) + .min(POOL.worker_count().max(1)); let array_stream = partition.execute()?; let dtype = array_stream.dtype().clone(); @@ -373,7 +395,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( rebase_offsets(RecordBatch::from(arrow.as_struct().clone())) }) }) - .buffered(POOL.worker_count().max(1)); + .buffered(lookahead); // Each JNI stream has a single consumer, so keep only one completed batch ahead of Java // rather than the per-core queue of `block_on_stream_thread_safe`. @@ -392,3 +414,52 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( Ok(()) }); } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::StreamExt; + use futures::stream; + use rstest::rstest; + use vortex::io::runtime::BlockingRuntime; + use vortex::io::runtime::current::CurrentThreadRuntime; + + #[rstest] + #[case(1)] + #[case(2)] + fn split_concurrency_bounds_completed_batches_ahead_of_the_consumer(#[case] max: usize) { + let runtime = CurrentThreadRuntime::new(); + let pool = runtime.new_pool(); + pool.set_workers(2); + let completed = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&completed); + let handle = runtime.handle(); + let batches = stream::iter(0..16) + .map(move |value| { + let counter = Arc::clone(&counter); + handle.spawn(async move { + counter.fetch_add(1, Ordering::SeqCst); + value + }) + }) + .buffered(max); + let mut iter = runtime.block_on_stream_prefetch(batches); + assert_eq!(iter.next(), Some(0)); + + // Let background work finish while the consumer keeps its first batch open. + runtime.block_on(async { + for _ in 0..64 { + runtime.handle().spawn(async {}).await; + } + }); + assert!( + completed.load(Ordering::SeqCst) <= max + 1, + "completed batches escaped the conversion lookahead bound" + ); + + assert_eq!(iter.collect::>(), (1..16).collect::>()); + } +} diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index ae670958ff4..62b1805aacf 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::num::NonZeroUsize; use std::ops::Range; use std::pin::Pin; use std::sync::Arc; @@ -182,6 +183,7 @@ impl DataSource for LayoutReaderDataSource { row_limit, selection: scan_request.selection, ordered: scan_request.ordered, + split_concurrency: scan_request.split_concurrency, metrics_registry: self.metrics_registry.clone(), next_row: row_range.start, end_row: row_range.end, @@ -207,6 +209,7 @@ struct LayoutReaderScan { /// A row budget shared by every split of an unordered limited scan; see [`RowLimit`]. row_limit: Option, ordered: bool, + split_concurrency: Option, selection: Selection, metrics_registry: Option>, next_row: u64, @@ -276,6 +279,7 @@ impl Stream for LayoutReaderScan { limit: split_limit, row_limit: this.row_limit.clone(), ordered: this.ordered, + split_concurrency: this.split_concurrency, row_range, selection: this.selection.clone(), metrics_registry: this.metrics_registry.clone(), @@ -309,6 +313,7 @@ struct LayoutReaderSplit { /// A row budget shared with the scan's other splits (unordered limited scans only). row_limit: Option, ordered: bool, + split_concurrency: Option, row_range: Range, selection: Selection, metrics_registry: Option>, @@ -347,7 +352,8 @@ impl Partition for LayoutReaderSplit { .with_projection(self.projection) .with_some_filter(self.filter) .with_some_metrics_registry(self.metrics_registry) - .with_ordered(self.ordered); + .with_ordered(self.ordered) + .with_some_max_splits_in_flight(self.split_concurrency); let builder = match self.row_limit { Some(row_limit) => builder.with_shared_limit(row_limit), None => builder.with_some_limit(self.limit), diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index c4a0d56fe3b..1f1ee275688 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -26,6 +26,7 @@ use std::any::Any; use std::collections::VecDeque; +use std::num::NonZeroUsize; use std::ops::Range; use std::sync::Arc; @@ -341,6 +342,7 @@ struct BoundScanRequest { partition_range: Option>, ordered: bool, limit: Option, + split_concurrency: Option, } impl BoundScanRequest { @@ -354,6 +356,7 @@ impl BoundScanRequest { partition_range, ordered, limit, + split_concurrency, } = request; Ok(Self { @@ -368,6 +371,7 @@ impl BoundScanRequest { partition_range, ordered, limit, + split_concurrency, }) } } @@ -599,7 +603,8 @@ impl Partition for MultiLayoutPartition { .with_selection(request.selection) .with_projection(request.projection) .with_some_filter(filter) - .with_ordered(request.ordered); + .with_ordered(request.ordered) + .with_some_max_splits_in_flight(request.split_concurrency); builder = match self.row_limit { Some(row_limit) => builder.with_shared_limit(row_limit), None => builder.with_some_limit(request.limit), @@ -637,6 +642,7 @@ mod tests { use vortex_scan::ScanRequest; use super::*; + use crate::scan::test::InFlight; use crate::scan::test::TestLayoutReader; use crate::scan::test::collect_scan_values; use crate::scan::test::new_session; @@ -734,4 +740,62 @@ mod tests { assert_eq!(values.len(), 3); Ok(()) } + + /// Scans 64 one-row splits, returning the rows and the most splits read at once. + fn scan_with_split_concurrency( + split_concurrency: Option, + ) -> VortexResult<(Vec, usize)> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let in_flight = Arc::new(InFlight::default()); + let reader: LayoutReaderRef = Arc::new( + TestLayoutReader::new(64) + .with_split_size(1) + .with_in_flight(Arc::clone(&in_flight)), + ); + let source = MultiLayoutDataSource::new_with_first(reader, vec![], vec![], &session); + + let scan = runtime.block_on(source.scan(ScanRequest { + ordered: true, + split_concurrency, + ..Default::default() + }))?; + let mut values = Vec::new(); + for partition in runtime.block_on(scan.partitions().try_collect::>())? { + values.extend(collect_scan_values( + runtime.block_on_stream(partition.execute()?), + )?); + } + Ok((values, in_flight.peak())) + } + + /// A per-worker concurrency set after the bound leaves the bound in place. + #[test] + fn a_later_per_worker_concurrency_keeps_the_bound() { + let reader: LayoutReaderRef = Arc::new(TestLayoutReader::new(1)); + let builder = ScanBuilder::new(new_session(), reader) + .with_some_max_splits_in_flight(NonZeroUsize::new(2)) + .with_concurrency(8); + + assert_eq!(builder.concurrency(), 8); + assert_eq!(builder.max_splits_in_flight(), NonZeroUsize::new(2)); + } + + /// The scan reads at most `split_concurrency` splits at once, however many cores there are. + #[rstest] + #[case(1)] + #[case(2)] + fn split_concurrency_bounds_the_splits_read_at_once(#[case] max: usize) -> VortexResult<()> { + let (values, peak) = scan_with_split_concurrency(NonZeroUsize::new(max))?; + let (default_values, default_peak) = scan_with_split_concurrency(None)?; + + assert_eq!(values, (0..64).collect::>()); + assert_eq!(values, default_values); + assert_eq!(peak, max, "splits read at once"); + assert!( + default_peak > max, + "the default read {default_peak} at once" + ); + Ok(()) + } } diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 170a1febc46..55a6df4bac2 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -3,6 +3,7 @@ use std::cmp; use std::iter; +use std::num::NonZeroUsize; use std::ops::Range; use std::sync::Arc; use std::task::Poll; @@ -62,8 +63,8 @@ pub struct RepeatedScan { selection: Selection, /// The natural splits of the file. splits: Splits, - /// The number of splits to make progress on concurrently **per-thread**. - concurrency: usize, + /// How many splits to make progress on at once. + concurrency: SplitConcurrency, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: MapFn, /// Maximal number of rows to read (after filtering). @@ -137,7 +138,7 @@ impl RepeatedScan { row_range: Option>, selection: Selection, splits: Splits, - concurrency: usize, + concurrency: SplitConcurrency, map_fn: MapFn, limit: Option, dtype: DType, @@ -242,7 +243,7 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult>> { - let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); + let concurrency = self.concurrency.splits_in_flight(); let handle = self.session.handle(); let ordered = self.ordered; let ctx = self.task_context(); @@ -387,6 +388,23 @@ fn stop_on_error( .boxed() } +/// How many splits a scan makes progress on at once: `per_worker` for each worker thread, as +/// [`get_available_parallelism`] counts them, and at most `max` in all when that is set. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct SplitConcurrency { + pub(crate) per_worker: usize, + pub(crate) max: Option, +} + +impl SplitConcurrency { + fn splits_in_flight(self) -> usize { + let scaled = self + .per_worker + .saturating_mul(get_available_parallelism().unwrap_or(1)); + self.max.map_or(scaled, |max| scaled.min(max.get())) + } +} + /// Run up to `concurrency` futures of `tasks` at once, in order when `ordered` is set. fn buffer( tasks: impl Stream + Send + 'static> + Send + 'static, @@ -411,6 +429,7 @@ fn intersect_ranges(left: Option<&Range>, right: Option>) -> Opt #[cfg(test)] mod tests { + use std::num::NonZeroUsize; use std::sync::Arc; use std::task::Poll; @@ -423,10 +442,20 @@ mod tests { use vortex_error::VortexResult; use vortex_error::vortex_err; + use super::SplitConcurrency; use super::limit_rows; use super::stop_on_error; use crate::scan::limit::RowLimit; + #[test] + fn split_cap_handles_oversized_per_worker_concurrency() { + let concurrency = SplitConcurrency { + per_worker: usize::MAX, + max: NonZeroUsize::new(1), + }; + assert_eq!(concurrency.splits_in_flight(), 1); + } + #[rstest] #[case::whole_array(4)] #[case::trimmed_array(2)] diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 21dc22c04ca..d8c11e5541d 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::num::NonZeroUsize; use std::ops::Range; use std::pin::Pin; use std::sync::Arc; @@ -39,6 +40,7 @@ use crate::layouts::row_idx::RowIdxLayoutReader; use crate::scan::limit::RowLimit; use crate::scan::limit::ScanLimit; use crate::scan::repeated_scan::RepeatedScan; +use crate::scan::repeated_scan::SplitConcurrency; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; @@ -71,7 +73,7 @@ pub struct ScanBuilder { /// uses them instead of walking the layout. natural_splits: Option>, /// The number of splits to make progress on concurrently **per-thread**. - concurrency: usize, + concurrency: SplitConcurrency, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, @@ -98,7 +100,10 @@ impl ScanBuilder { natural_splits: None, // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. - concurrency: 4, + concurrency: SplitConcurrency { + per_worker: 4, + max: None, + }, map_fn: Arc::new(Ok), metrics_registry: None, limit: None, @@ -222,14 +227,31 @@ impl ScanBuilder { /// Returns the per-worker row-split concurrency. pub fn concurrency(&self) -> usize { - self.concurrency + self.concurrency.per_worker + } + + /// Returns the most row splits in progress at once, in all, if a bound is set. + pub fn max_splits_in_flight(&self) -> Option { + self.concurrency.max } /// The number of row splits to make progress on concurrently per-thread, must /// be greater than 0. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); - self.concurrency = concurrency; + self.concurrency.per_worker = concurrency; + self + } + + /// Bound the row splits in progress at once to `max` in all, whatever the per-worker + /// concurrency and the machine's core count. `None` removes the bound. A filtered, limited + /// scan may hold one more trimmed array per split while it is mapped. + /// + /// Each split in progress holds the decoded chunks of every column it overlaps, so a scan + /// that runs alongside others, such as one task of several on a machine, can set this to + /// reduce memory held by concurrent decoding. This does not cap total memory or I/O read-ahead. + pub fn with_some_max_splits_in_flight(mut self, max: Option) -> Self { + self.concurrency.max = max; self } diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 7d030f2c987..53bca5c5b5b 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -4,6 +4,8 @@ use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Poll; use parking_lot::Mutex; @@ -66,6 +68,20 @@ pub struct TestLayoutReader { fail_projection: bool, projection_masks: Option>>>, filter_exprs: Option>>>, + in_flight: Option>, +} + +/// Counts projections running at once, and the most that ever were. +#[derive(Debug, Default)] +pub struct InFlight { + now: AtomicUsize, + peak: AtomicUsize, +} + +impl InFlight { + pub fn peak(&self) -> usize { + self.peak.load(Ordering::SeqCst) + } } impl TestLayoutReader { @@ -82,6 +98,7 @@ impl TestLayoutReader { fail_projection: false, projection_masks: None, filter_exprs: None, + in_flight: None, } } @@ -111,6 +128,13 @@ impl TestLayoutReader { self } + /// Records into `in_flight` how many projections are running at once. Each projection yields + /// once before finishing, so others that the scan has started get to run alongside it. + pub fn with_in_flight(mut self, in_flight: Arc) -> Self { + self.in_flight = Some(in_flight); + self + } + pub fn with_filter_exprs(mut self, filter_exprs: Arc>>) -> Self { self.filter_exprs = Some(filter_exprs); self @@ -213,8 +237,15 @@ impl LayoutReader for TestLayoutReader { let base = self.base; let projection_masks = self.projection_masks.clone(); let fail = self.fail_projection || (self.fail_first_projection && row_range.start == 0); + let in_flight = self.in_flight.clone(); Ok(Box::pin(async move { + if let Some(in_flight) = &in_flight { + let now = in_flight.now.fetch_add(1, Ordering::SeqCst) + 1; + in_flight.peak.fetch_max(now, Ordering::SeqCst); + tokio::task::yield_now().await; + in_flight.now.fetch_sub(1, Ordering::SeqCst); + } let mask = mask.await?; if let Some(projection_masks) = projection_masks { projection_masks.lock().push(mask.true_count()); diff --git a/vortex-scan/src/lib.rs b/vortex-scan/src/lib.rs index 755ec472c6c..a0454c6c1a6 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -28,6 +28,7 @@ pub mod selection; pub mod strict_sorted_buffer; use std::any::Any; +use std::num::NonZeroUsize; use std::ops::Range; use std::sync::Arc; @@ -170,6 +171,13 @@ pub struct ScanRequest { /// order is completion order, so each partition applies the limit locally and the caller must /// trim the concatenated result. pub limit: Option, + /// The most splits of each partition to decode at once, in all. + /// + /// `None` keeps the reader's default, which scales with the machine's cores. Each split being + /// decoded holds the decoded chunks of every column it overlaps, so a scan that runs alongside + /// others, such as one task of several on a machine, can set this to bound its memory. It does + /// not bound I/O: segment reads for later splits may still be coalesced ahead of them. + pub split_concurrency: Option, } impl Default for ScanRequest { @@ -183,6 +191,7 @@ impl Default for ScanRequest { ordered: false, limit: None, partition_range: None, + split_concurrency: None, } } } From ce2b6d027338a2ececdec913e116c16b86c43d8c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 9 Nov 2026 17:40:17 +0000 Subject: [PATCH 3/6] Keep IntDictScheme for chunks DictStrategy rejects (#11331) The default write strategy removed `IntDictScheme` from the one compressor shared by `DictStrategy`'s codes path and its fallback path. The exclusion is only needed for dictionary codes, so they aren't dictionary-encoded twice. But it also applied to every column `DictStrategy` rejects, and those could never use integer dictionaries, even nested inside another encoding. This hits decimals hardest: they always take the fallback, so a `DecimalByteParts` most significant part could never be dict-encoded. --------- Signed-off-by: Robert Kruszewski --- vortex-file/src/strategy.rs | 125 ++++++++++++++-------- vortex-file/src/tests.rs | 44 ++++++++ vortex-layout/src/layouts/zoned/writer.rs | 2 +- 3 files changed, 128 insertions(+), 43 deletions(-) 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 b7ba7ed8587..61f3bbd58cb 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -377,6 +377,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-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() { From a1fbac6799aa6ada3bf1f307f1543e6f4ad988f2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 9 Nov 2026 13:02:04 -0500 Subject: [PATCH 4/6] Export to Arrow with the session of the execution context (#11304) ## Summary - Tracking Issue: #11114 - Progress towards #11286, because the export entry point now takes a target `Field` instead of a `DataType`. An Arrow export could use two different `ArrowSession`s at once, and the API made that easy to do by accident. `ArrowSession::execute_arrow` took the session as its receiver and also an `ExecutionCtx`, which carries its own `VortexSession` and so its own `ArrowSession`. The top-level dispatch and the target checks used the receiver, but plugins such as `Vector` and `Json` export their storage through the session of the ctx. So `ArrowSession::default().execute_arrow(array, None, &mut session.create_execution_ctx())` compiled, and which plugins ran depended on which part of the export was looking. A session that `arrow()` created on first use also could not import what it exported. It was a bare default without the dtype registry, so it wrote an extension type without a plugin under its Vortex ID but read the same field back as plain storage. Only `vortex_arrow::initialize` fixed that, and it replaced the existing session, so calling it after `vortex_uuid::initialize` silently dropped the UUID plugins. ## Changes Exports now go through `ExecuteArrowExt` on the array, which takes its session from the ctx, so there is no second session to pass. Keeping `session.arrow().execute_arrow(..)` and using the session of the ctx inside was rejected, because the receiver would then be silently ignored. A runtime check that both are the same session was rejected too, because it only fails on the paths that tests happen to cover. The old methods are removed instead of deprecated, since a deprecated method that takes a session keeps the same trap open. `arrow()` now builds the session from the dtype registry on first use, through the new `SessionExt::get_or_insert_with`, so every session imports what it exports. That leaves nothing for `initialize` to do, so it is removed instead of fixed. Note that `arrow()` now also creates the `DTypeSession`, so `.with::()` after the first `arrow()` call panics, as it already did after `initialize`. New tests cover a UUID session that imports by Vortex ID without `initialize`, and a dtype that is registered after the first `arrow()` call. ## API Changes - `ArrowSession::execute_arrow`, `ArrowSession::exporter`, and `ArrowExporter` are removed. Use `execute_arrow` or `execute_arrow_with_options` from `vortex_arrow::ExecuteArrowExt` on the array instead. - The deprecated `ArrowArrayExecutor` is removed. - `vortex_arrow::initialize` is removed. - `ArrowSession` has no `Default` or public constructor, so every `ArrowSession` comes from a `VortexSession`. - `vortex-session` gains `SessionExt::get_or_insert_with`. Signed-off-by: Connor Tsui --- benchmarks/compress-bench/src/gpu/vortex.rs | 6 +- encodings/parquet-variant/src/array.rs | 17 +- encodings/parquet-variant/src/arrow.rs | 44 +-- encodings/parquet-variant/src/kernel.rs | 3 +- encodings/pco/src/tests.rs | 17 +- encodings/sparse/src/canonical.rs | 51 +-- encodings/uuid/src/arrow.rs | 63 +++- vortex-arrow/benches/to_arrow.rs | 46 +-- vortex-arrow/src/datum.rs | 22 +- vortex-arrow/src/dtype.rs | 13 +- vortex-arrow/src/executor/byte.rs | 60 +-- vortex-arrow/src/executor/byte_view.rs | 6 +- vortex-arrow/src/executor/decimal.rs | 28 +- vortex-arrow/src/executor/dictionary.rs | 26 +- vortex-arrow/src/executor/fixed_size_list.rs | 24 +- vortex-arrow/src/executor/list.rs | 51 +-- vortex-arrow/src/executor/list_view.rs | 20 +- vortex-arrow/src/executor/map.rs | 18 +- vortex-arrow/src/executor/mod.rs | 357 ++++++++++++++---- vortex-arrow/src/executor/run_end.rs | 30 +- vortex-arrow/src/executor/struct_.rs | 34 +- vortex-arrow/src/lib.rs | 35 +- vortex-arrow/src/options/tests.rs | 14 +- vortex-arrow/src/run_end_import.rs | 5 +- vortex-arrow/src/session.rs | 318 +++------------- vortex-arrow/src/tests.rs | 6 +- vortex-arrow/tests/canonical.rs | 17 +- vortex-arrow/tests/vortex_id.rs | 51 +-- vortex-bench/src/datasets/sparse_wide.rs | 3 +- .../src/datasets/struct_list_of_ints.rs | 5 +- vortex-datafusion/src/convert/schema.rs | 38 +- vortex-datafusion/src/persistent/opener.rs | 11 +- vortex-datafusion/src/v2/source.rs | 8 +- vortex-ffi/src/scan.rs | 5 +- vortex-file/src/v2/scan.rs | 5 +- vortex-file/src/v2/scan/tests.rs | 5 +- vortex-jni/src/scan.rs | 6 +- vortex-json/src/arrow.rs | 30 +- .../examples/reverse_index/pruning_tests.rs | 2 - vortex-layout/src/scan/arrow.rs | 5 +- vortex-python/src/arrays/mod.rs | 16 +- vortex-python/src/iter/mod.rs | 7 +- vortex-session/src/lib.rs | 18 +- vortex-session/src/session.rs | 62 +-- vortex-spatial/benches/area.rs | 2 - vortex-spatial/benches/binary_predicates.rs | 2 - vortex-spatial/benches/collect.rs | 2 - vortex-spatial/benches/convex_hull.rs | 2 - vortex-spatial/benches/distance.rs | 2 - vortex-spatial/benches/envelope.rs | 2 - vortex-spatial/benches/length.rs | 2 - vortex-spatial/benches/make_line.rs | 2 - vortex-spatial/src/extension/linestring.rs | 10 +- .../src/extension/multilinestring.rs | 10 +- vortex-spatial/src/extension/multipoint.rs | 10 +- vortex-spatial/src/extension/multipolygon.rs | 10 +- vortex-spatial/src/extension/point.rs | 10 +- vortex-spatial/src/extension/polygon.rs | 10 +- vortex-spatial/src/extension/rect.rs | 11 +- vortex-spatial/src/extension/wkb.rs | 8 +- vortex-spatial/src/tests/point.rs | 9 +- vortex-spatial/src/tests/rect.rs | 5 +- vortex-spatial/src/tests/wkb.rs | 9 +- vortex-tensor/src/types/vector/arrow.rs | 46 +-- vortex-test/e2e-cuda/src/lib.rs | 26 +- vortex-web/crate/src/wasm.rs | 5 +- vortex/src/lib.rs | 4 +- 67 files changed, 850 insertions(+), 957 deletions(-) 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/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 8320feb76f6..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 @@ -277,27 +283,34 @@ type ArrowImporterRegistry = ArcSwapMap>; /// /// 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() } } @@ -317,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([])) } @@ -507,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 @@ -695,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 @@ -1064,153 +1030,8 @@ 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 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 = 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); - } - - // 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_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 { - ensure_ext_metadata(target, &decode_arrow_ext_metadata(target)?, ext_dtype)?; - - let target_dtype = self - .session - .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()), self, 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()), self, ctx) - } -} - /// Fails unless `metadata`, the extension metadata of `field`, is the metadata of `ext_dtype`. -fn ensure_ext_metadata( +pub(crate) fn ensure_ext_metadata( field: &Field, metadata: &[u8], ext_dtype: &ExtDTypeRef, @@ -1224,59 +1045,18 @@ fn ensure_ext_metadata( Ok(()) } -/// 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 -} - /// 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())) } } @@ -1353,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()); @@ -1374,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), @@ -1395,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), @@ -1469,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]), @@ -1595,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 92d66d0148a..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; @@ -66,7 +66,6 @@ static SESSION: LazyLock = LazyLock::new(|| { .dtypes() .registry() .insert(*TAGGED_ALIAS_ID, Arc::new(Tagged)); - vortex_arrow::initialize(&session); session }); @@ -155,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() @@ -228,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}"); @@ -264,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( @@ -288,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()); @@ -307,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); @@ -327,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!( @@ -352,9 +350,7 @@ fn alias_of_vortex_id_round_trips() -> VortexResult<()> { .insert(EXTENSION_TYPE_NAME_KEY.to_owned(), TAGGED_ALIAS.to_owned()); let mut ctx = SESSION.create_execution_ctx(); - let exported = SESSION - .arrow() - .execute_arrow(array.clone(), Some(&field), &mut 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); @@ -369,9 +365,7 @@ fn plain_array_exports_to_a_vortex_id_field() -> VortexResult<()> { let field = storage_field_as(tagged_i64s(0xff), DataType::Int64)?; let mut ctx = SESSION.create_execution_ctx(); - let exported = SESSION - .arrow() - .execute_arrow(storage.clone(), Some(&field), &mut 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); @@ -390,12 +384,9 @@ fn export_compares_nested_extension_types_by_storage() -> VortexResult<()> { .arrow() .to_arrow_storage_field("t", array.dtype().as_extension())?; let session = VortexSession::empty(); - vortex_arrow::initialize(&session); let mut ctx = session.create_execution_ctx(); - let exported = session - .arrow() - .execute_arrow(array.clone(), Some(&field), &mut 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); @@ -409,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()); @@ -422,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] 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-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/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 1f67bb12dd2..9f7287391f4 100644 --- a/vortex-jni/src/scan.rs +++ b/vortex-jni/src/scan.rs @@ -48,6 +48,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; @@ -388,10 +389,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-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/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 0103a95eb52..0cc6c6c3b35 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 @@ -176,23 +176,13 @@ impl VortexSession { Self(SharedSessionVars::default()) } - /// 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)); } @@ -250,11 +240,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> { @@ -367,6 +370,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"); From 2e1b05b84e7b7a2df5171b1c6c18bfbc93d0c3d7 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 9 Nov 2026 19:49:22 +0000 Subject: [PATCH 5/6] Compress fixed-size lists without their null rows (#10425) Add SparseScheme for FixedSizedList. Instead of compressing just the elements we attempt compressing whole lists using SparseArray Readers need no changes, because `Sparse` already canonicalizes fixed-size lists. On the data from #10417: | | Before | After | |---|---|---| | Mostly-null file (1 row in 20 valid) | 25,296 B | 9,056 B | | Rewrite of both files | 292,808 B | 190,768 B | --------- Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + vortex-btrblocks/Cargo.toml | 1 + vortex-btrblocks/src/builder.rs | 2 + .../src/schemes/fixed_size_list.rs | 160 ++++++++++++++++++ vortex-btrblocks/src/schemes/mod.rs | 1 + vortex-btrblocks/src/session.rs | 3 + vortex-btrblocks/src/tests.rs | 45 +++++ vortex-compressor/src/compressor/cascade.rs | 26 +++ 8 files changed, 239 insertions(+) create mode 100644 vortex-btrblocks/src/schemes/fixed_size_list.rs diff --git a/Cargo.lock b/Cargo.lock index f9b8e501845..88710226f57 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/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( From 7b36f7fe4fdcd41929e03eaf7cb9da5683df69ab Mon Sep 17 00:00:00 2001 From: Dan King Date: Mon, 9 Nov 2026 17:12:04 -0500 Subject: [PATCH 6/6] fix(array): preserve constants when zip selects one branch (#11373) When zip's mask is uniformly true or false (including null masks, which select the false branch), executing the selected ConstantArray materializes it unnecessarily. Scalar access then logs `Scalar function vortex.zip() returned non-constant array from execution over all scalar inputs`. Return the selected branch cast to the output dtype without executing it again. Removing the zip node already makes execution progress; the caller can preserve constants with `Columnar` or continue materializing with `Canonical`. Added a parameterized regression test for true, false, and null masks, with matching branch types and with nullable/non-nullable branches. It checks constant encoding, values, and output nullability. Validation: reviewed the patch and execution paths. Tests, formatting, and lints were not run locally, following AGENTS.md's instruction to leave verification execution to the user. Implemented and opened by Codex at Daniel King's request. Daniel King (@danking) is the responsible human. Signed-off-by: Daniel King --- vortex-array/src/scalar_fn/fns/zip/mod.rs | 30 +++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) 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);