From c8ba7810f8f6756c682a985588088b44de8292ed Mon Sep 17 00:00:00 2001 From: Dan King Date: Tue, 10 Nov 2026 05:04:39 -0500 Subject: [PATCH 01/13] fix(pco): decode into uninitialized buffers safely (#11363) Pcodec 1.1.0 exposes a MaybeUninit-based API for safe writing into uninitialized memory. Let's use it! Signed-off-by: Daniel King --- Cargo.toml | 2 +- encodings/pco/src/array/mod.rs | 12 ++++++++---- encodings/pco/src/probe.rs | 10 +++++++--- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 7fc483deec5..658785c90f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -214,7 +214,7 @@ parquet = "59.2" parquet-variant = "59.2" parquet-variant-compute = "59.2" paste = "1.0.15" -pco = "1.0.1" +pco = "1.1.0" percent-encoding = "2.3.2" pin-project-lite = "0.2.15" primitive-types = { version = "0.14.0" } diff --git a/encodings/pco/src/array/mod.rs b/encodings/pco/src/array/mod.rs index c30c518cc93..2f11389543a 100644 --- a/encodings/pco/src/array/mod.rs +++ b/encodings/pco/src/array/mod.rs @@ -604,9 +604,6 @@ impl PcoData { let old_len = decompressed_values.len(); let new_len = old_len + page_n_values; decompressed_values.reserve(page_n_values); - unsafe { - decompressed_values.set_len(new_len); - } let page: &[u8] = self .pages .get(page_idx) @@ -626,8 +623,15 @@ impl PcoData { let mut pd = cd .page_decompressor(page, page_n_values) .map_err(vortex_err_from_pco)?; - pd.read(&mut decompressed_values[old_len..new_len]) + let progress = pd + .read_uninit(&mut decompressed_values.spare_capacity_mut()[..page_n_values]) .map_err(vortex_err_from_pco)?; + vortex_ensure_eq!(progress.n_processed, page_n_values); + // SAFETY: `read_uninit` initialized all `page_n_values` spare elements, + // as confirmed by the processed count above. + unsafe { + decompressed_values.set_len(new_len); + } chunk_decompressor = Some(cd); } else { diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 00aa801254a..f7a10c8844f 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -21,6 +21,7 @@ use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use crate::Pco; @@ -156,10 +157,13 @@ fn decode_page( .page_decompressor(buffer, n_values) .map_err(vortex_err_from_pco)?; let mut values = BufferMut::::with_capacity_in(n_values, ctx.allocator().clone()); - // SAFETY: the buffer reserves `n_values` elements, and the page decompressor was built for - // exactly that count, so `read` writes every element before anything observes it. + let progress = decoder + .read_uninit(&mut values.spare_capacity_mut()[..n_values]) + .map_err(vortex_err_from_pco)?; + vortex_ensure_eq!(progress.n_processed, n_values); + // SAFETY: `read_uninit` initialized all `n_values` spare elements, as confirmed by the + // processed count above. unsafe { values.set_len(n_values) }; - decoder.read(&mut values).map_err(vortex_err_from_pco)?; Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) } From ced8ec9fc4eeaadf628477bc7f007929c78c04ca Mon Sep 17 00:00:00 2001 From: Tanish Gudise Date: Tue, 10 Nov 2026 05:31:42 -0500 Subject: [PATCH 02/13] fix(cuda): copy encoded device-resident children to the host in into_host (#10916) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Tracking Issue: #10914 `compress-bench --gpu-decompress --gpu-verify` panics with `expected host buffer` on develop for CMSprovider and Euro2016. A run-end encoded validity mask stays encoded on the device after `execute_cuda`. `CanonicalCudaExt::into_host` then executes it on the host while its buffers are still on the device (`validity_into_host` → `child_into_host` → `RunEnd::execute` → `runend_decode_bools` → `BufferHandle::as_host`). A nullable dictionary's validity, a lazy `take` through device-resident codes, fails the same way. ## Changes - `canonical.rs`: `child_into_host` copies every device buffer of a child that isn't canonical to the host before executing it there. - Children that are already canonical take the existing path. - The copies run concurrently, and an array's buffers are replaced only when one of them moves (Sparse rejects buffer replacement). - The decode path is unchanged. - This doesn't add mixed CPU/GPU execution (#7838): the child moves to the host whole, so nothing decodes from device buffers on the host. - Tests (4), each building the child on the device and comparing `into_host` with the expected result: - a primitive array whose validity is a run-end encoded mask; - one whose validity is a sparse mask (Sparse keeps its fill value in a host buffer it doesn't allow to be replaced, so only its children move); - a string dictionary with nullable values and device-resident codes; - a primitive dictionary with nullable values and device-resident codes. Evidence (A40, CUDA 13.0): | Check | develop | this PR | |---|---|---| | The 4 tests | all fail: `expected host buffer` (run-end mask), `as_slice must be called on host buffer` (sparse mask, both dictionaries) | 4 pass | | `--gpu-verify` on CMSprovider and Euro2016, today's preset | both panic | both pass (234 and 22 fields verified) | | `vortex-cuda` lib suite | — | 711 passed, 0 failed | Everything in the table ran on this PR alone. With all three PRs, `--gpu-verify` passes on all 17 GPU-suite datasets. ## AI assistance Developed with Claude Code (Claude Opus 5.5), which wrote most of the code, tests and this description under my direction. I reviewed the diff and the measurements. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Tanish Gudise Co-authored-by: Claude Opus 5.5 Co-authored-by: Alexander Droste --- vortex-cuda/src/canonical.rs | 186 ++++++++++++++++++++++++++++++++++- 1 file changed, 185 insertions(+), 1 deletion(-) diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 138033a08ea..836f2e7544b 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use async_trait::async_trait; +use futures::future::BoxFuture; use futures::future::try_join_all; use vortex::array::ArrayRef; use vortex::array::Canonical; @@ -29,11 +30,17 @@ use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; -/// Copy a canonical child array to the host. +/// Copy a child array to the host, executing it there if CUDA execution left it encoded. async fn child_into_host(child: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let child = if child.is_canonical() { + child + } else { + buffers_into_host(child).await? + }; Ok(child .execute::(ctx)? .into_host(ctx) @@ -41,6 +48,46 @@ async fn child_into_host(child: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResul .into_array()) } +/// Copies every device buffer of `array` and of its children to the host, keeping the encoding. +/// +/// A child that CUDA execution left encoded, such as a validity mask that is still a lazy `take` +/// through device-resident dictionary codes or a run-end encoded mask, can only be executed on +/// the host once all of its buffers are there. +fn buffers_into_host(array: ArrayRef) -> BoxFuture<'static, VortexResult> { + Box::pin(async move { + if array.is_host() { + return Ok(array); + } + let mut array = array; + if array.children_iter().any(|child| !child.is_host()) { + let mut slots = Vec::with_capacity(array.slots().len()); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(buffers_into_host(child.clone()).await?), + None => None, + }); + } + // SAFETY: the slots are byte-for-byte copies; only their placement changes. + array = unsafe { array.with_slots(slots.into())? }; + } + let buffers = array.buffer_handles(); + // Some encodings, such as Sparse, reject buffer replacement, so only replace when a + // buffer actually moves. + if buffers.iter().all(BufferHandle::is_on_host) { + return Ok(array); + } + let buffers = try_join_all(buffers.into_iter().map(|buffer| async move { + if buffer.is_on_host() { + return Ok::<_, VortexError>(buffer); + } + Ok(BufferHandle::new_host(buffer.try_into_host()?.await?)) + })) + .await?; + // SAFETY: the buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_buffers(buffers) } + }) +} + /// Move all canonical data from to_host from device. #[async_trait] pub trait CanonicalCudaExt { @@ -196,11 +243,19 @@ mod tests { use std::sync::LazyLock; use vortex::array::VortexSessionExecute; + use vortex::array::arrays::DictArray; + use vortex::array::assert_arrays_eq; use vortex::dtype::PType; + use vortex::encodings::runend::RunEnd; + use vortex::encodings::sparse::Sparse; + use vortex::scalar::Scalar; use vortex::session::VortexSession; use super::*; use crate::CudaSession; + use crate::executor::CudaExecute; + use crate::executor::CudaExecutionCtx; + use crate::kernel::DictExecutor; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -225,4 +280,133 @@ mod tests { assert!(Canonical::List(list).into_host(&mut ctx).await.is_err()); Ok(()) } + + /// Moves every buffer of `array` and of its children to the device without decoding it. + fn encoded_on_device(array: ArrayRef, cuda_ctx: &CudaExecutionCtx) -> VortexResult { + let mut slots = Vec::new(); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(encoded_on_device(child.clone(), cuda_ctx)?), + None => None, + }); + } + // SAFETY: the slots are byte-for-byte copies; only their placement changes. + let array = unsafe { array.with_slots(slots.into())? }; + // Sparse keeps its fill value in a host buffer that it doesn't allow to be replaced. + if array.as_opt::().is_some() { + return Ok(array); + } + let mut buffers = Vec::new(); + for buffer in array.buffer_handles() { + buffers.push(cuda_ctx.ensure_on_device_sync(buffer)?); + } + // SAFETY: the buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_buffers(buffers) } + } + + /// A GPU decode can leave a validity mask encoded on the device, here run-end encoded. + /// Copying the values to the host used to execute that mask on the host while its buffers + /// were still on the device, which panicked with "expected host buffer". + #[crate::test] + async fn test_into_host_copies_encoded_device_validity() -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let mask = RunEnd::new( + PrimitiveArray::new(Buffer::from(vec![2u32, 5, 6]), Validity::NonNullable).into_array(), + BoolArray::from_iter([true, false, true]).into_array(), + &mut ctx, + ) + .into_array(); + let mask = encoded_on_device(mask, &cuda_ctx)?; + let values = cuda_ctx + .stream() + .copy_to_device(vec![1i32, 2, 3, 4, 5, 6])? + .await?; + let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask)); + + let host = Canonical::Primitive(array) + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, None, None, Some(6)]); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + /// A sparse validity mask keeps its fill value in a host buffer, which Sparse doesn't allow to + /// be replaced, while its patch indices and values are on the device. Only the children move + /// to the host. + #[crate::test] + async fn test_into_host_copies_sparse_device_validity() -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let mask = Sparse::try_new( + PrimitiveArray::new(Buffer::from(vec![1u32, 4]), Validity::NonNullable).into_array(), + BoolArray::from_iter([false, false]).into_array(), + 6, + Scalar::from(true), + &mut ctx, + )? + .into_array(); + let mask = encoded_on_device(mask, &cuda_ctx)?; + let values = cuda_ctx + .stream() + .copy_to_device(vec![1i32, 2, 3, 4, 5, 6])? + .await?; + let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask)); + + let host = Canonical::Primitive(array) + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + /// A dictionary with nullable values returns its validity as a lazy gather through the codes. + /// With the codes on the device, as a file scan leaves them, copying the decoded dictionary to + /// the host used to panic on a host-buffer read. + async fn assert_device_dict_into_host( + codes: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?; + let device_codes = encoded_on_device(codes.into_array(), &cuda_ctx)?; + let dict = DictArray::try_new(device_codes, values)?; + let host = DictExecutor + .execute(dict.into_array(), &mut cuda_ctx) + .await? + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_into_host_string_dict_with_nullable_values() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([ + Some("a"), + None, + Some("a string longer than twelve bytes"), + ]); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u8, 1, 2, 1, 0, 2]), + Validity::NonNullable, + ); + assert_device_dict_into_host(codes, values.into_array()).await + } + + #[crate::test] + async fn test_into_host_primitive_dict_with_nullable_values() -> VortexResult<()> { + let values = PrimitiveArray::from_option_iter([Some(7i32), None, Some(-3)]); + let codes = + PrimitiveArray::new(Buffer::from(vec![2u32, 0, 1, 1, 0]), Validity::NonNullable); + assert_device_dict_into_host(codes, values.into_array()).await + } } From ec0fd925d2e6bcc359c641101840f05ba0f16fb8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 10 Nov 2026 05:58:12 -0500 Subject: [PATCH 03/13] feat(cuda): Execute Widen arrays on CUDA (#11303) ## Summary - Tracking Issue: #10701 Progress towards #11284 (part 1). CUDA has no kernel for `vortex.widen`, and the CPU fallback only runs when every buffer is on the host, so executing a device-resident `Widen` fails with "CPU fallback with device-resident buffers is not supported". After CD5, every DecimalByteParts array with a narrow MSP reads from a file with an integer `Widen` as its MSP, so CD5 can't merge until GPU reads can execute it. ## Changes `Widen` is not fused into the dynamic dispatch path. Dynamic dispatch zero-extends signed narrow BitPacked and FoR children, so fusing a signed `Widen` there needs care to sign-extend instead. The `TODO` in `plan_builder.rs` tracks that. Until then, `WidenExecutor` executes the child on the device and widens it in a separate kernel. Decimal storage that is wider than the `Widen`'s logical type is an error instead of a narrowing copy on the device. Narrowing needs either an overflow check or trust in the stored precision, and that choice is the open question in part 2 of #11284. The decimal path reuses the kernels in `decimal_cast.cu`, and the Arrow export now shares `widen_decimal_values` with it, so there is one copy of the widening code. The tests compare every integer and decimal widening pair against the CPU `Widen` (with nulls, negative values, a FoR child, and the kernel's block boundaries). They only run in the GPU CI. Signed-off-by: Alexander Droste --- vortex-cuda/kernels/src/decimal_cast.cu | 121 ++++----- vortex-cuda/kernels/src/widen.cu | 36 +++ vortex-cuda/src/arrow/canonical.rs | 71 +---- .../src/dynamic_dispatch/plan_builder.rs | 2 + vortex-cuda/src/kernel/decimal_cast.rs | 124 +++++++++ vortex-cuda/src/kernel/encodings/mod.rs | 6 + vortex-cuda/src/kernel/encodings/widen/mod.rs | 135 ++++++++++ .../src/kernel/encodings/widen/tests.rs | 246 ++++++++++++++++++ vortex-cuda/src/kernel/mod.rs | 4 + vortex-cuda/src/lib.rs | 3 + 10 files changed, 608 insertions(+), 140 deletions(-) create mode 100644 vortex-cuda/kernels/src/widen.cu create mode 100644 vortex-cuda/src/kernel/decimal_cast.rs create mode 100644 vortex-cuda/src/kernel/encodings/widen/mod.rs create mode 100644 vortex-cuda/src/kernel/encodings/widen/tests.rs diff --git a/vortex-cuda/kernels/src/decimal_cast.cu b/vortex-cuda/kernels/src/decimal_cast.cu index 45c8c75e2a0..88618176a94 100644 --- a/vortex-cuda/kernels/src/decimal_cast.cu +++ b/vortex-cuda/kernels/src/decimal_cast.cu @@ -6,64 +6,43 @@ #include #include -// Arrow decimal schemas fix the physical values buffer width: -// - Decimal32: 4 bytes per value. -// - Decimal64: 8 bytes per value. -// - Decimal128: 16 bytes per value. -// - Decimal256: 32 bytes per value. +// Widen decimal values buffers to a wider physical width. Two callers use these kernels: +// - The Arrow export casts to the width that the Arrow schema fixes: 4 bytes for Decimal32, +// 8 bytes for Decimal64, 16 bytes for Decimal128, and 32 bytes for Decimal256. +// - The Widen executor casts a narrow decimal child to the canonical width of its decimal +// dtype, from 2 to 32 bytes. // -// Vortex storage width can differ, so export casts to the schema-implied width. -// Rust-side export rejects narrowing casts because detecting overflow on-device -// would require synchronizing an overflow flag back to the host. +// Only widening pairs are instantiated. A narrowing cast can overflow, and detecting that on the +// device would require synchronizing an overflow flag back to the host, so the Rust side rejects +// narrowing before it loads a kernel. -// Low 64-bit conversion for Decimal32/64 outputs. -template -__device__ __forceinline__ int64_t decimal_to_i64(Input value) { - if constexpr (std::is_same_v) { - return value.lo; - } else if constexpr (std::is_same_v) { - return value.parts[0]; - } else { - return static_cast(value); - } -} - -// 128-bit conversion for Decimal128 outputs. -template -__device__ __forceinline__ int128_t decimal_to_i128(Input value) { - if constexpr (std::is_same_v) { - return value; - } else if constexpr (std::is_same_v) { - return int128_t {value.parts[0], value.parts[1]}; - } else { - const int64_t lo = static_cast(value); - const int64_t hi = value < 0 ? -1 : 0; - return int128_t {lo, hi}; - } -} - -// Convert one value to the Arrow schema's physical width. +// Sign-extend one value to the wider output width in a single step, with no intermediate width. +// Native outputs use one `static_cast`. The 128-bit and 256-bit outputs keep the input's words +// and fill the upper words with its sign. template __device__ __forceinline__ Output decimal_cast_value(Input value) { - if constexpr (std::is_same_v) { - return static_cast(decimal_to_i64(value)); - } else if constexpr (std::is_same_v) { - return decimal_to_i64(value); - } else if constexpr (std::is_same_v) { - return decimal_to_i128(value); + static_assert(sizeof(Input) < sizeof(Output), "decimal casts on the GPU only widen"); + + if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + return int128_t {lo, lo < 0 ? -1 : 0}; + } else if constexpr (std::is_same_v && std::is_same_v) { + const int64_t sign = value.hi < 0 ? -1 : 0; + return int256_t {{value.lo, value.hi, sign, sign}}; + } else if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + const int64_t sign = lo < 0 ? -1 : 0; + return int256_t {{lo, sign, sign, sign}}; } else { - static_assert(std::is_same_v); - if constexpr (std::is_same_v) { - return value; - } else { - const int128_t value128 = decimal_to_i128(value); - const int64_t sign = value128.hi < 0 ? -1 : 0; - return int256_t {{value128.lo, value128.hi, sign, sign}}; - } + static_assert(std::is_integral_v && std::is_signed_v); + static_assert(std::is_integral_v && std::is_signed_v); + return static_cast(value); } } -// Cast a contiguous values buffer to the Arrow schema's physical width. +// Cast a contiguous values buffer to the output physical width. template __device__ void decimal_cast_device(const Input *__restrict input, Output *__restrict output, uint64_t array_len) { @@ -76,28 +55,26 @@ decimal_cast_device(const Input *__restrict input, Output *__restrict output, ui } } -// Generate Decimal32/64/128/256 cast kernels for one input storage type. -#define GENERATE_DECIMAL_CAST_KERNELS(input_suffix, InputType) \ - extern "C" __global__ void decimal_cast_##input_suffix##_i32(const InputType *__restrict input, \ - int32_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i64(const InputType *__restrict input, \ - int64_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i128(const InputType *__restrict input, \ - int128_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i256(const InputType *__restrict input, \ - int256_t *__restrict output, \ - uint64_t array_len) { \ +#define GENERATE_DECIMAL_CAST(input_suffix, InputType, output_suffix, OutputType) \ + extern "C" __global__ void decimal_cast_##input_suffix##_##output_suffix( \ + const InputType *__restrict input, \ + OutputType *__restrict output, \ + uint64_t array_len) { \ decimal_cast_device(input, output, array_len); \ } -FOR_EACH_SIGNED_INT(GENERATE_DECIMAL_CAST_KERNELS) -FOR_EACH_LARGE_DECIMAL(GENERATE_DECIMAL_CAST_KERNELS) +GENERATE_DECIMAL_CAST(i8, int8_t, i16, int16_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i128, int128_t, i256, int256_t) diff --git a/vortex-cuda/kernels/src/widen.cu b/vortex-cuda/kernels/src/widen.cu new file mode 100644 index 00000000000..9654b9db3a0 --- /dev/null +++ b/vortex-cuda/kernels/src/widen.cu @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "scalar_kernel.cuh" +#include "types.cuh" +#include + +template +struct WidenOp { + static_assert(sizeof(InputT) < sizeof(OutputT)); + static_assert(std::is_signed_v == std::is_signed_v); + + __device__ OutputT operator()(InputT value) const { + return static_cast(value); + } +}; + +#define GENERATE_WIDEN(input_suffix, InputT, output_suffix, OutputT) \ + extern "C" __global__ void widen_##input_suffix##_##output_suffix(const InputT *__restrict input, \ + OutputT *__restrict output, \ + uint64_t array_len) { \ + scalar_kernel(input, output, array_len, WidenOp {}); \ + } + +GENERATE_WIDEN(i8, int8_t, i16, int16_t) +GENERATE_WIDEN(i8, int8_t, i32, int32_t) +GENERATE_WIDEN(i8, int8_t, i64, int64_t) +GENERATE_WIDEN(i16, int16_t, i32, int32_t) +GENERATE_WIDEN(i16, int16_t, i64, int64_t) +GENERATE_WIDEN(i32, int32_t, i64, int64_t) +GENERATE_WIDEN(u8, uint8_t, u16, uint16_t) +GENERATE_WIDEN(u8, uint8_t, u32, uint32_t) +GENERATE_WIDEN(u8, uint8_t, u64, uint64_t) +GENERATE_WIDEN(u16, uint16_t, u32, uint32_t) +GENERATE_WIDEN(u16, uint16_t, u64, uint64_t) +GENERATE_WIDEN(u32, uint32_t, u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index b3c87d1abbf..67dfe0cd9ae 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -44,16 +44,13 @@ use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; use vortex::array::builtins::ArrayBuiltins; -use vortex::array::match_each_decimal_value_type; use vortex::array::optimizer::ArrayOptimizer; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::DecimalType; -use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; -use vortex::dtype::i256; use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -62,7 +59,6 @@ use vortex::error::vortex_ensure_eq; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex_cuda_driver::CudaSlice; -use vortex_cuda_driver::DeviceRepr; use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::PushKernelArg; use vortex_onpair::OnPair; @@ -92,6 +88,7 @@ use crate::executor::execute_validity_cuda; use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; use crate::kernel::decode_onpair_varbin; +use crate::kernel::widen_decimal_values; use crate::stream::zero_padding; const BITMAP_THREADS_PER_BLOCK: u32 = 256; @@ -475,73 +472,11 @@ async fn export_decimal_values( "cannot export decimal values from {values_type} storage to Arrow {target_type}: narrowing would require a device-to-host overflow check", ); } - let values = ctx.ensure_on_device(values).await?; if values_type == target_type { - return Ok(values); + return ctx.ensure_on_device(values).await; } - match_each_decimal_value_type!(values_type, |S| { - export_decimal_values_from::(values, target_type, len, ctx).await - }) -} - -/// Dispatch from a concrete storage type `S` to the Arrow-required output width. -async fn export_decimal_values_from( - values: BufferHandle, - target_type: DecimalType, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, -{ - match target_type { - DecimalType::I32 => decimal_cast::(values, len, ctx).await, - DecimalType::I64 => decimal_cast::(values, len, ctx).await, - DecimalType::I128 => decimal_cast::(values, len, ctx).await, - DecimalType::I256 => decimal_cast::(values, len, ctx).await, - target_type => { - vortex_bail!("cannot export DecimalArray as Arrow decimal value type {target_type}") - } - } -} - -/// Launches the CUDA kernel that casts from Vortex storage type `S` to Arrow output type `D`. -/// -/// The caller must ensure this is not a narrowing cast. -async fn decimal_cast( - values: BufferHandle, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, - D: NativeDecimalType + DeviceRepr, -{ - if len == 0 { - return ctx - .ensure_on_device(BufferHandle::new_host( - Buffer::::empty().into_byte_buffer(), - )) - .await; - } - - let mut output_buffer = ctx.device_alloc::(len)?; - - let values_view = values.cuda_view::()?; - let len_u64 = len as u64; - let cuda_function = ctx.load_function_with_suffixes( - "decimal_cast", - &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], - )?; - - ctx.launch_kernel(&cuda_function, len, |args| { - args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); - })?; - - Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( - output_buffer, - )?))) + widen_decimal_values(values, values_type, target_type, len, ctx).await } /// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index d0405f7ee41..9498f1c8419 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -133,6 +133,8 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool { _ => false, }; } + // TODO(connor): Add Widen fusion after fixing sign extension of signed narrow BitPacked and + // FoR children in dynamic_dispatch.cu. Those paths currently zero-extend to the output width. (id == FoR.id() && array.as_::().references().is_global()) || id == ZigZag.id() || id == Primitive.id() diff --git a/vortex-cuda/src/kernel/decimal_cast.rs b/vortex-cuda/src/kernel/decimal_cast.rs new file mode 100644 index 00000000000..8f1fb0d9f87 --- /dev/null +++ b/vortex-cuda/src/kernel/decimal_cast.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Widening casts of decimal values buffers on the device. +//! +//! The Arrow export widens decimal values to the width that the Arrow schema fixes, and the +//! Widen executor widens a narrow decimal child to the canonical width of its decimal dtype. Both +//! launch the `decimal_cast` kernels through this module. +//! +//! These casts only widen. A narrowing cast can overflow, and to detect that on the device, the +//! kernel would need an overflow flag that synchronizes with the host. So this module rejects a +//! narrowing cast, and each caller decides how to handle values that are wider than its target. + +use std::sync::Arc; + +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_decimal_value_type; +use vortex::buffer::Buffer; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::i256; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_ensure_eq; +use vortex::error::vortex_panic; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; + +/// Widen `len` decimal values from `values_type` to the wider `target_type` on the device. +/// +/// The values can be on the host or on the device, and the result is on the device. Each value is +/// sign-extended, so the result holds the same values as the input. +/// +/// # Errors +/// +/// Returns an error if `target_type` is not wider than `values_type`, if the buffer does not +/// contain `len` complete values, or if the copy to the device or the kernel launch fails. +pub(crate) async fn widen_decimal_values( + values: BufferHandle, + values_type: DecimalType, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!( + values_type < target_type, + "a decimal cast on the GPU must widen to a type wider than {values_type}, got \ + {target_type}" + ); + vortex_ensure_eq!(values.len(), len * values_type.byte_width()); + + let values = ctx.ensure_on_device(values).await?; + + match_each_decimal_value_type!(values_type, |S| { + decimal_cast_to::(values, target_type, len, ctx).await + }) +} + +/// Dispatch [`decimal_cast`] from `S` on the runtime `target_type`. +async fn decimal_cast_to( + values: BufferHandle, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, +{ + match target_type { + DecimalType::I8 => vortex_panic!( + "widen_decimal_values ensures a target wider than the values type, and no decimal \ + type is narrower than i8, got {target_type}" + ), + DecimalType::I16 => decimal_cast::(values, len, ctx).await, + DecimalType::I32 => decimal_cast::(values, len, ctx).await, + DecimalType::I64 => decimal_cast::(values, len, ctx).await, + DecimalType::I128 => decimal_cast::(values, len, ctx).await, + DecimalType::I256 => decimal_cast::(values, len, ctx).await, + } +} + +/// Launch the `decimal_cast__` kernel, which sign-extends device values of type `S` to the +/// wider type `D`. +/// +/// `D` **must** be wider than `S`. The `decimal_cast.cu` kernels only instantiate widening pairs, so +/// loading any other pair fails. [`widen_decimal_values`] checks this before it dispatches. +async fn decimal_cast( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, + D: NativeDecimalType + DeviceRepr, +{ + if len == 0 { + return ctx + .ensure_on_device(BufferHandle::new_host( + Buffer::::empty().into_byte_buffer(), + )) + .await; + } + + let mut output_buffer = ctx.device_alloc::(len)?; + + let values_view = values.cuda_view::()?; + let len_u64 = len as u64; + let cuda_function = ctx.load_function_with_suffixes( + "decimal_cast", + &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], + )?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output_buffer, + )?))) +} diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index fb292bc445c..fbc2defb696 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! CUDA executors for encoded arrays. +//! +//! Executors decode children through the CUDA execution context and return canonical arrays. + mod alp; mod bitpacked; mod date_time_parts; @@ -11,6 +15,7 @@ mod fsst; mod onpair; mod runend; mod sequence; +mod widen; mod zigzag; mod zstd; mod zstd_buffers; @@ -29,6 +34,7 @@ pub(crate) use onpair::OnPairExecutor; pub(crate) use onpair::decode_onpair_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; +pub(crate) use widen::WidenExecutor; pub(crate) use zigzag::ZigZagExecutor; pub(crate) use zstd::ZstdExecutor; pub use zstd::ZstdKernelPrep; diff --git a/vortex-cuda/src/kernel/encodings/widen/mod.rs b/vortex-cuda/src/kernel/encodings/widen/mod.rs new file mode 100644 index 00000000000..be324b868b7 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/mod.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Restores the logical width of integer and decimal [`Widen`] arrays on the device. +//! +//! Children execute at their own width before widening, so compressed signed values retain their +//! sign. Execution keeps the child's validity and allocates a wider values buffer, unless decimal +//! storage already has the output width. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::Widen; +use vortex::array::arrays::widen::WidenArraySlotsExt; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::dtype::DType; +use vortex::dtype::NativePType; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::kernel::widen_decimal_values; + +/// CUDA execution for integer and decimal arrays stored at a narrower width. +#[derive(Debug)] +pub(crate) struct WidenExecutor; + +#[async_trait] +impl CudaExecute for WidenExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let array = array.try_downcast::().map_err(|array| { + vortex_err!( + "WidenExecutor requires a Widen array, got {}", + array.encoding_id() + ) + })?; + let child = array.values().clone().execute_cuda(ctx).await?; + + match (array.dtype(), child) { + (DType::Primitive(output_ptype, _), Canonical::Primitive(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let values = ctx.ensure_on_device(parts.buffer).await?; + let output = match_each_integer_ptype!(parts.ptype, |S| { + match_each_integer_ptype!(*output_ptype, |D| { + widen_integers::(values, len, ctx)? + }) + }); + + Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + output, + *output_ptype, + parts.validity, + ))) + } + (DType::Decimal(decimal_dtype, _), Canonical::Decimal(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let target_type = decimal_dtype.canonical_values_type(); + vortex_ensure!( + parts.values_type <= target_type, + "CUDA Widen requires decimal storage no wider than {target_type}, got {}", + parts.values_type + ); + + let values = if parts.values_type == target_type { + ctx.ensure_on_device(parts.values).await? + } else { + widen_decimal_values(parts.values, parts.values_type, target_type, len, ctx) + .await? + }; + + Ok(Canonical::Decimal(DecimalArray::try_new_handle( + values, + target_type, + *decimal_dtype, + parts.validity, + )?)) + } + (dtype, child) => vortex_bail!( + "CUDA Widen requires an integer or decimal child matching {dtype}, got {}", + child.dtype() + ), + } + } +} + +/// Widens `len` device values of type `S` into a new buffer of type `D`. +/// +/// `D` must be strictly wider than `S` with the same signedness, as the [`Widen`] invariant +/// guarantees. The `widen.cu` kernels only instantiate those pairs. The other monomorphized pairs +/// cannot run because a valid `Widen` never dispatches to them. +fn widen_integers( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativePType + DeviceRepr, + D: NativePType + DeviceRepr, +{ + let input = values.cuda_view::()?; + let mut output = ctx.device_alloc::(len)?; + let len_u64 = len as u64; + let function = ctx.load_function("widen", &[S::PTYPE, D::PTYPE])?; + ctx.launch_kernel(&function, len, |args| { + args.arg(&input).arg(&mut output).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output, + )?))) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cuda/src/kernel/encodings/widen/tests.rs b/vortex-cuda/src/kernel/encodings/widen/tests.rs new file mode 100644 index 00000000000..9b0781eeb73 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/tests.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Device execution and CPU agreement for Widen arrays. + +use num_traits::Bounded; +use rstest::rstest; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::WidenArray; +use vortex::array::arrays::decimal::DecimalArrayExt; +use vortex::array::assert_arrays_eq; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::dtype::DType; +use vortex::dtype::DecimalDType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; +use vortex::dtype::PType; +use vortex::encodings::fastlanes::FoR; +use vortex::error::VortexResult; + +use crate::CanonicalCudaExt; +use crate::CudaDispatchMode; +use crate::CudaExecutionCtx; +use crate::CudaSession; +use crate::cuda_session; +use crate::executor::CudaArrayExt; + +/// Copies fixture values to the device while preserving their dtype and validity. +async fn upload(child: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResult { + match child.execute::(ctx.execution_ctx())? { + Canonical::Primitive(child) => { + let parts = child.into_data_parts(); + Ok(PrimitiveArray::from_buffer_handle( + ctx.ensure_on_device(parts.buffer).await?, + parts.ptype, + parts.validity, + ) + .into_array()) + } + Canonical::Decimal(child) => { + let parts = child.into_data_parts(); + Ok(DecimalArray::try_new_handle( + ctx.ensure_on_device(parts.values).await?, + parts.values_type, + parts.decimal_dtype, + parts.validity, + )? + .into_array()) + } + _ => unreachable!("Widen test fixtures are primitive or decimal"), + } +} + +async fn check_widen( + host_child: ArrayRef, + device_child: ArrayRef, + dtype: DType, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<()> { + let host = WidenArray::try_new(host_child, dtype.clone())?.into_array(); + let expected = host.execute::(ctx.execution_ctx())?.into_array(); + let device = WidenArray::try_new(device_child, dtype.clone())?.into_array(); + assert!(!device.is_host()); + + let result = device.execute_cuda(ctx).await?; + assert_eq!(result.dtype(), &dtype); + assert_eq!(result.len(), expected.len()); + match &result { + Canonical::Primitive(array) => assert!(!array.buffer_handle().is_on_host()), + Canonical::Decimal(array) => { + assert!(!array.buffer_handle().is_on_host()); + assert_eq!( + array.values_type(), + dtype + .as_decimal_opt() + .expect("the result is decimal") + .canonical_values_type() + ); + } + _ => unreachable!("Widen returns primitive or decimal arrays"), + } + + let actual = result.into_host(ctx.execution_ctx()).await?.into_array(); + assert_arrays_eq!(actual, expected, ctx.execution_ctx()); + + Ok(()) +} + +fn validity(nullable: bool, len: usize) -> Validity { + if nullable { + Validity::from_iter((0..len).map(|i| i % 3 != 0)) + } else { + Validity::NonNullable + } +} + +#[rstest] +#[case::i8_i16(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I16)] +#[case::i8_i32(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I32)] +#[case::i8_i64(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I64)] +#[case::i16_i32(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I32)] +#[case::i16_i64(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I64)] +#[case::i32_i64(Buffer::from_iter([i32::MIN, i32::MAX, -1, 0, 1]), PType::I64)] +#[case::u8_u16(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U16)] +#[case::u8_u32(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U32)] +#[case::u8_u64(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U64)] +#[case::u16_u32(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U32)] +#[case::u16_u64(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U64)] +#[case::u32_u64(Buffer::from_iter([u32::MIN, u32::MAX, 0, 1, 1 << 31]), PType::U64)] +#[crate::test] +async fn test_integer_pairs( + #[case] values: Buffer, + #[case] output: PType, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let validity = validity(nullable, values.len()); + let child = PrimitiveArray::new(values, validity).into_array(); + let dtype = DType::Primitive(output, child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[case::i8_p4([-99i8, 99], 2, 4)] +#[case::i8_p9([-99i8, 99], 2, 9)] +#[case::i8_p18([-99i8, 99], 2, 18)] +#[case::i8_p38([-99i8, 99], 2, 38)] +#[case::i8_p76([-99i8, 99], 2, 76)] +#[case::i16_p9([-9_999i16, 9_999], 4, 9)] +#[case::i16_p18([-9_999i16, 9_999], 4, 18)] +#[case::i16_p38([-9_999i16, 9_999], 4, 38)] +#[case::i16_p76([-9_999i16, 9_999], 4, 76)] +#[case::i32_p18([-999_999_999i32, 999_999_999], 9, 18)] +#[case::i32_p38([-999_999_999i32, 999_999_999], 9, 38)] +#[case::i32_p76([-999_999_999i32, 999_999_999], 9, 76)] +#[case::i64_p38([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 38)] +#[case::i64_p76([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 76)] +#[case::i128_p76([-(1i128 << 100), (1i128 << 100)], 38, 76)] +#[crate::test] +async fn test_decimal_pairs( + #[case] bounds: [T; 2], + #[case] child_precision: u8, + #[case] precision: u8, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let mut values = [ + T::default(), + bounds[0], + T::default(), + T::default(), + bounds[1], + ]; + if nullable { + values[0] = T::min_value(); + values[3] = T::max_value(); + } + + let child = DecimalArray::new( + Buffer::from_iter(values), + DecimalDType::new(child_precision, 1), + validity(nullable, 5), + ) + .into_array(); + let dtype = DType::Decimal(DecimalDType::new(precision, 1), child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[crate::test] +async fn test_compressed_child( + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let child = PrimitiveArray::from_iter([0i8, 1, 50, 99]).into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + let host = FoR::try_new(child, (-100i8).into())?.into_array(); + let device = FoR::try_new(device_child, (-100i8).into())?.into_array(); + + check_widen(host, device, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::one(1)] +#[case::before_vector(15)] +#[case::vector(16)] +#[case::after_vector(17)] +#[case::before_block(2047)] +#[case::block(2048)] +#[case::after_block(2049)] +#[crate::test] +async fn test_kernel_boundaries(#[case] len: usize) -> VortexResult<()> { + let child = PrimitiveArray::from_iter((0..len).map(|i| if i % 2 == 0 { -99i8 } else { 99 })) + .into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::equal_width(false)] +#[case::narrowing(true)] +#[crate::test] +async fn test_legacy_decimal_storage(#[case] narrowing: bool) -> VortexResult<()> { + let dtype = DecimalDType::new(9, 1); + let child = if narrowing { + DecimalArray::from_iter([-99i128, 0, 99], dtype) + } else { + DecimalArray::from_iter([-99i64, 0, 99], dtype) + } + .into_array(); + let logical = DType::Decimal(DecimalDType::new(18, 1), Nullability::NonNullable); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + if narrowing { + let array = WidenArray::try_new(device_child, logical)?.into_array(); + let error = array + .execute_cuda(&mut ctx) + .await + .expect_err("i128 storage requires narrowing to i64"); + assert!( + error + .to_string() + .contains("CUDA Widen requires decimal storage no wider than i64, got i128"), + "{error}" + ); + } else { + check_widen(child, device_child, logical, &mut ctx).await?; + } + + Ok(()) +} diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index f1eb24b1fad..94d15ffa6d0 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -19,6 +19,10 @@ use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::sys::CUevent_flags; mod arrays; + +mod decimal_cast; +pub(crate) use decimal_cast::widen_decimal_values; + // Fatbins let the CUDA driver select among the configured real and virtual architectures. mod embedded_kernels { include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index bdf9af1421b..0e60e61b2c1 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -54,6 +54,7 @@ use kernel::OnPairExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; +use kernel::WidenExecutor; use kernel::ZigZagExecutor; use kernel::ZstdBuffersExecutor; use kernel::ZstdExecutor; @@ -80,6 +81,7 @@ use vortex::array::arrays::Filter; use vortex::array::arrays::List; use vortex::array::arrays::Masked; use vortex::array::arrays::Slice; +use vortex::array::arrays::Widen; use vortex::encodings::alp::ALP; use vortex::encodings::datetime_parts::DateTimeParts; use vortex::encodings::decimal_byte_parts::DecimalByteParts; @@ -124,6 +126,7 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(OnPair.id(), &OnPairExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); + session.register_kernel(Widen.id(), &WidenExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor); session.register_kernel(Zstd.id(), &ZstdExecutor); session.register_kernel(ZstdBuffers.id(), &ZstdBuffersExecutor); From 500f8b18fbc1ff319efdbf36fd88b61b223cd495 Mon Sep 17 00:00:00 2001 From: ChristopherWirt Date: Tue, 10 Nov 2026 11:17:35 +0000 Subject: [PATCH 04/13] perf(spark): avoid redundant metadata HEADs when opening Hadoop streams (#11368) ## Summary - Tracking Issue: #11365 - Fixes #11365 Open pooled Hadoop streams with the known file length so S3A versions that honor the hint can avoid redundant metadata HEAD requests. Ports the Hadoop-readable changes from [QC #106](https://github.com/quantcast/vortex/pull/106) onto upstream `develop` at `2e1b05b84`. ## Changes - Use `fs.openFile()` with optional known-length and random-read hints; other filesystems may ignore them. - Preserve original asynchronous-open exception types, restore interruption and close any stream produced after an interrupted open. - Add focused tests for supplied/discovered lengths, every pooled stream receiving the options, optional keys, exception identity and late-stream cleanup. Signed-off-by: Christopher Wirt --- .../dev/vortex/spark/io/HadoopReadable.java | 24 +++- .../vortex/spark/io/HadoopReadableTest.java | 132 ++++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java index 26da27450b3..09f4831f69b 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java @@ -5,13 +5,18 @@ import dev.vortex.io.PooledReadable; import java.io.IOException; +import java.io.InterruptedIOException; import java.io.UncheckedIOException; import java.nio.ByteBuffer; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FSDataInputStream; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.fs.StreamCapabilities; +import org.apache.hadoop.io.IOUtils; +import org.apache.hadoop.util.functional.FutureIO; /** A {@link dev.vortex.io.NativeReadable} that serves Vortex's reads from a Hadoop {@link FileSystem}. */ public final class HadoopReadable extends PooledReadable { @@ -46,7 +51,24 @@ private HadoopReadable(String name, Path path, FileSystem fs, long length) { @Override protected PositionalStream openStream() throws IOException { - return new HadoopStream(fs.open(path)); + // Given the length, S3A on Hadoop 3.3.5+ skips the metadata HEAD per stream; older Hadoop + // (Spark 3.5) ignores these keys. String keys because 3.3.4 lacks the OpenFileOptions constants. + CompletableFuture opened = fs.openFile(path) + .opt("fs.option.openfile.length", Long.toString(length())) + .opt("fs.option.openfile.read.policy", "random") + .build(); + try { + return new HadoopStream(opened.get()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + InterruptedIOException failure = new InterruptedIOException("Interrupted while opening " + path); + failure.initCause(e); + // Cancelling a future need not stop the underlying open. Close any stream it produces later. + var unused = opened.thenAccept(IOUtils::closeStream); + throw failure; + } catch (ExecutionException e) { + return FutureIO.raiseInnerCause(e); + } } private static final class HadoopStream implements PositionalStream { diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java index afef7223b2a..588b3cbd83d 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java @@ -5,21 +5,36 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import dev.vortex.io.NativeReadable; import java.io.EOFException; +import java.io.FileNotFoundException; import java.io.IOException; +import java.io.InterruptedIOException; +import java.io.UncheckedIOException; +import java.net.URI; import java.nio.ByteBuffer; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.List; +import java.util.UUID; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.FutureDataInputStreamBuilder; +import org.apache.hadoop.fs.RawLocalFileSystem; +import org.apache.hadoop.fs.impl.FutureDataInputStreamBuilderImpl; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -58,6 +73,74 @@ void suppliedLengthIsTrustedOverAnExtraStat() throws IOException { } } + @Test + void everyStreamUsesTheResolvedLengthAndRandomReadPolicy() throws IOException { + byte[] content = content(64); + Path file = write("builder.vortex", content); + for (long suppliedLength : new long[] {-1, content.length}) { + try (TrackingFileSystem fs = trackingFileSystem(); + HadoopReadable readable = + HadoopReadable.open(fs.getConf(), trackedPath(fs, file), suppliedLength)) { + assertEquals(suppliedLength < 0 ? 1 : 0, fs.stats); + // Open one stream directly and let a pooled read open a second. + try (var first = readable.openStream()) { + ByteBuffer buffer = ByteBuffer.allocateDirect(content.length); + readable.readFully(0, buffer); + assertArrayEquals(content, drain(buffer)); + } + assertEquals(2, fs.openOptions.size()); + for (Configuration options : fs.openOptions) { + assertEquals("64", options.get("fs.option.openfile.length")); + assertEquals("random", options.get("fs.option.openfile.read.policy")); + } + } + } + } + + @Test + void failedAsynchronousOpensPreserveTheirOriginalExceptions() throws IOException { + Path file = write("failed-open.vortex", content(64)); + IOException missing = new FileNotFoundException("missing object"); + for (Throwable failure : new Throwable[] { + missing, + new UncheckedIOException(missing), + new IllegalArgumentException("invalid option"), + new AssertionError("open failed") + }) { + try (TrackingFileSystem fs = trackingFileSystem(); + HadoopReadable readable = HadoopReadable.open(fs.getConf(), trackedPath(fs, file), 64)) { + fs.nextOpen = CompletableFuture.failedFuture(failure); + Throwable expected = failure instanceof UncheckedIOException ? missing : failure; + assertSame( + expected, + assertThrows(expected.getClass(), () -> readable.readFully(0, ByteBuffer.allocateDirect(1)))); + } + } + } + + @Test + void interruptedOpenPreservesInterruptAndClosesALateStream() throws IOException { + Path file = write("interrupted-open.vortex", content(64)); + try (TrackingFileSystem fs = trackingFileSystem(); + HadoopReadable readable = HadoopReadable.open(fs.getConf(), trackedPath(fs, file), 64)) { + CompletableFuture pending = new CompletableFuture<>(); + fs.nextOpen = pending; + Thread.currentThread().interrupt(); + try { + InterruptedIOException failure = assertThrows( + InterruptedIOException.class, () -> readable.readFully(0, ByteBuffer.allocateDirect(1))); + assertTrue(Thread.currentThread().isInterrupted()); + assertTrue(failure.getCause() instanceof InterruptedException); + } finally { + Thread.interrupted(); + } + assertFalse(pending.isCancelled()); + FSDataInputStream lateStream = fs.open(new org.apache.hadoop.fs.Path(trackedPath(fs, file))); + pending.complete(lateStream); + assertThrows(IOException.class, lateStream::read); + } + } + @Test void concurrentReadsEachGetTheirOwnStream() throws Exception { byte[] content = content(SIZE); @@ -148,6 +231,55 @@ void createReplacesAnExistingFile() throws IOException { assertArrayEquals(new byte[] {9}, Files.readAllBytes(file)); } + private static TrackingFileSystem trackingFileSystem() throws IOException { + Configuration trackingConf = new Configuration(); + trackingConf.setClass("fs.tracking.impl", TrackingFileSystem.class, FileSystem.class); + return (TrackingFileSystem) FileSystem.get(URI.create("tracking://" + UUID.randomUUID()), trackingConf); + } + + private static String trackedPath(TrackingFileSystem fs, Path file) { + return new org.apache.hadoop.fs.Path("tracking", fs.getUri().getAuthority(), file.toString()).toString(); + } + + /** Hadoop instantiates this filesystem through the test configuration. */ + public static final class TrackingFileSystem extends RawLocalFileSystem { + private URI uri; + private int stats; + private final List openOptions = new ArrayList<>(); + private CompletableFuture nextOpen; + + @Override + public void initialize(URI name, Configuration configuration) throws IOException { + uri = name; + super.initialize(name, configuration); + } + + @Override + public URI getUri() { + return uri != null ? uri : super.getUri(); + } + + @Override + public FileStatus getFileStatus(org.apache.hadoop.fs.Path path) throws IOException { + stats++; + return super.getFileStatus(path); + } + + @Override + public FutureDataInputStreamBuilder openFile(org.apache.hadoop.fs.Path path) { + return new FutureDataInputStreamBuilderImpl(this, path) { + @Override + public CompletableFuture build() throws IOException { + assertTrue(getMandatoryKeys().isEmpty(), "Open options must remain optional for other filesystems"); + openOptions.add(new Configuration(getOptions())); + return nextOpen != null + ? nextOpen + : CompletableFuture.completedFuture(TrackingFileSystem.this.open(path, getBufferSize())); + } + }; + } + } + private Path write(String name, byte[] content) throws IOException { Path file = tempDir.resolve(name); Files.write(file, content); From 30446ea99eb68655e4ea7cbf14ed00e08ce7df71 Mon Sep 17 00:00:00 2001 From: Tanish Gudise Date: Tue, 10 Nov 2026 06:30:32 -0500 Subject: [PATCH 05/13] fix(cuda): bound dictionary gathers by the dictionary length (#10915) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Tracking Issue: #10914 The standalone dictionary kernels in `dict.cu` gathered `values[codes[i]]` for every row. Nullable codes can hold any value at null positions (`DictArray::try_new` doesn't bound-check them), so these gathers read outside the dictionary. The fused planner already rejects nullable-code dictionaries for this reason and sends them to these kernels. Develop's own fuzzer, extended here, crashes on this within 71k runs. ## Changes - `dict.cu`: every kernel (primitive, decimal, bool and the 16-byte view gather) takes `values_len`. A code at or beyond it writes a zero value. In a valid array only null positions can hold such a code, so the zero is never observed. - `dict.rs` and `benches/dynamic_dispatch_cuda.rs`: pass `values_len` at every launch. - Tests (12), each compared with the expected result through `assert_arrays_eq`: - out-of-range codes under nulls for string, decimal, bool and primitive values, with codes up to `u32::MAX` and `u64::MAX`; - the dictionary the fuzzer crashed on (one u8 value, null code 3,960,089,894); - device-resident codes with nullable values and codes; - all-null codes; - slices across 1,023/1,024/1,025 and 2,047/2,048; - u16 codes with 65,536 values and u32 codes with 70,000; - empty, 12-, 13- and 4,097-byte strings; - non-UTF-8 binary values. - `fuzz/src/gpu/mod.rs`: `compress_gpu` also generates string and binary dictionaries, dictionaries with arbitrary u8/u16/u32 codes under nulls, and strings compressed with the CUDA preset. Evidence (A40, CUDA 13.0): | Check | develop's kernels | this PR | |---|---|---| | 5 bug tests: out-of-range codes under nulls (string, decimal, bool, primitive) and the fuzzer's input | 5 fail with `CUDA_ERROR_ILLEGAL_ADDRESS` | 5 pass | | 7 edge-case tests (all-null codes, block-boundary slices, full u16/u32 ranges, inline boundary, binary, device codes) | pass | 7 pass | | `compute-sanitizer --tool memcheck` on the 6 out-of-range tests | 1-2 errors each, `Invalid __global__ read` of 1, 8 or 16 bytes | 0 errors | | `compress_gpu` fuzz (CI flags, ASAN, empty corpus) | crash after 70,941 runs, `CUDA_ERROR_ILLEGAL_ADDRESS` | 100,000 runs in 283 s, 0 crashes | | `vortex-cuda` lib suite | — | 719 passed, 0 failed | Under nsys, the bounds check costs nothing measurable on numeric dictionaries: 1.574 vs 1.563 ms per iteration on CMSprovider and 0.825 vs 0.832 ms on taxi. On develop's kernels the `u32::MAX`/`u64::MAX` tests leave the CUDA context in a sticky `CUDA_ERROR_ILLEGAL_ADDRESS` state. Run in one process, they also fail the unrelated tests that follow, about 336 of them. A regression of this fix would show up the same way, as a wall of failures led by these tests. Everything in the table ran on this PR alone, except the develop fuzz crash: that ran develop's kernels under this PR's extended fuzz target. ## AI assistance Developed with Claude Code (Claude Opus 5.5), which wrote most of the code, tests and this description under my direction. I reviewed the diff and the measurements. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Tanish Gudise Co-authored-by: Claude Opus 5.5 Co-authored-by: Alexander Droste --- fuzz/src/gpu/mod.rs | 88 ++++++- vortex-cuda/benches/dynamic_dispatch_cuda.rs | 2 + vortex-cuda/kernels/src/dict.cu | 15 +- vortex-cuda/src/kernel/arrays/dict.rs | 237 +++++++++++++++++++ 4 files changed, 334 insertions(+), 8 deletions(-) diff --git a/fuzz/src/gpu/mod.rs b/fuzz/src/gpu/mod.rs index 91e116e38cd..c972581d5f1 100644 --- a/fuzz/src/gpu/mod.rs +++ b/fuzz/src/gpu/mod.rs @@ -13,9 +13,17 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::arbitrary::ArbitraryArray; +use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; +use vortex_array::arrays::arbitrary::ArbitraryWith; use vortex_array::arrays::dict::ArbitraryDictArray; +use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::vortex_err; use crate::SESSION; @@ -28,13 +36,16 @@ use crate::error::VortexFuzzResult; pub enum GpuEncodingKind { /// Dictionary encoding with GPU take support. Dict, + /// Dictionary whose nullable codes hold arbitrary values, possibly outside the dictionary, + /// under null positions. + DictGarbageNullCodes, } impl<'a> Arbitrary<'a> for GpuEncodingKind { fn arbitrary(u: &mut Unstructured<'a>) -> Result { - // Currently only Dict is supported - match u.int_in_range(0..=0)? { + match u.int_in_range(0..=1)? { 0 => Ok(GpuEncodingKind::Dict), + 1 => Ok(GpuEncodingKind::DictGarbageNullCodes), _ => unreachable!(), } } @@ -58,14 +69,83 @@ impl<'a> Arbitrary<'a> for FuzzCompressGpu { .0 .into_array() } + GpuEncodingKind::DictGarbageNullCodes => arbitrary_dict_with_garbage_null_codes(u)?, }; Ok(FuzzCompressGpu { array }) } } +/// Generate a random UTF-8 or binary DType. +fn arbitrary_gpu_string_dtype(u: &mut Unstructured) -> Result { + let nullability: Nullability = u.arbitrary()?; + Ok(if u.arbitrary()? { + DType::Utf8(nullability) + } else { + DType::Binary(nullability) + }) +} + +/// Generate a dictionary whose codes are in range at valid positions and arbitrary, including far +/// outside the dictionary, at null positions. +fn arbitrary_dict_with_garbage_null_codes(u: &mut Unstructured) -> Result { + let dtype = if u.arbitrary()? { + arbitrary_gpu_string_dtype(u)? + } else { + arbitrary_gpu_primitive_dtype(u)? + }; + let values_len = u.int_in_range(1..=20)?; + let values = ArbitraryArray::arbitrary_with_config( + u, + &ArbitraryArrayConfig { + dtype: Some(dtype), + len: values_len..=values_len, + }, + )? + .0; + let len = u.int_in_range(0..=200)?; + let valid: Vec = (0..len).map(|_| u.arbitrary()).collect::>()?; + let values_len = u32::try_from(values_len).map_err(|_| arbitrary::Error::IncorrectFormat)?; + let validity: Validity = valid.iter().copied().collect(); + let codes = match u.int_in_range(0..=2)? { + 0 => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + 1 => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + _ => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + }; + let dict = DictArray::try_new(codes.into_array(), values) + .map_err(|_| arbitrary::Error::IncorrectFormat)?; + Ok(dict.into_array()) +} + +/// Codes in `0..values_len` at valid positions and any value of `T` at null positions. +fn garbage_codes<'a, T>(u: &mut Unstructured<'a>, valid: &[bool], values_len: u32) -> Result> +where + T: Arbitrary<'a> + TryFrom, +{ + valid + .iter() + .map(|&is_valid| { + if is_valid { + T::try_from(u.int_in_range(0..=values_len - 1)?) + .map_err(|_| arbitrary::Error::IncorrectFormat) + } else { + u.arbitrary::() + } + }) + .collect() +} + /// Generate a random primitive DType suitable for GPU operations. -fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result { +fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result { let nullability: Nullability = u.arbitrary()?; let ptype = match u.int_in_range(0..=9)? { 0 => PType::U8, @@ -80,7 +160,7 @@ fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result PType::F64, _ => unreachable!(), }; - Ok(vortex_array::dtype::DType::Primitive(ptype, nullability)) + Ok(DType::Primitive(ptype, nullability)) } /// Run the GPU decompression fuzzer. diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index bee5f2e3b77..9b4df35002c 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -625,9 +625,11 @@ mod standalone { { let mut launch = cuda_ctx.stream().launch_builder(&self.dict_fn); let codes_len = self.codes_len as u64; + let values_len = self.values_len as u64; launch.arg(&self.codes_buf); launch.arg(&codes_len); launch.arg(&self.values_alp_buf); + launch.arg(&values_len); launch.arg(&mut self.output_buf); unsafe { launch diff --git a/vortex-cuda/kernels/src/dict.cu b/vortex-cuda/kernels/src/dict.cu index a456b476821..db8b1caf249 100644 --- a/vortex-cuda/kernels/src/dict.cu +++ b/vortex-cuda/kernels/src/dict.cu @@ -8,10 +8,14 @@ #include "config.cuh" #include "types.cuh" +// Codes under null positions may hold any value, and so may corrupt input, so every gather checks +// the code against the dictionary length. A code outside the dictionary writes a zero value instead +// of reading past the end of `values`. template __device__ void dict_kernel(const IndexT *const __restrict codes, uint64_t codes_len, const ValueT *const __restrict values, + uint64_t values_len, ValueT *const __restrict output) { const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; @@ -20,7 +24,7 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { const IndexT code = codes[idx]; - output[idx] = values[code]; + output[idx] = static_cast(code) < values_len ? values[code] : ValueT {}; } } @@ -39,6 +43,7 @@ template __device__ void dict_bool_kernel(const IndexT *const __restrict codes, uint64_t codes_len, const uint8_t *const __restrict values, + uint64_t values_len, uint64_t values_bit_offset, uint8_t *const __restrict output) { const uint64_t output_len = (codes_len + 7) / 8; @@ -53,7 +58,7 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, #pragma unroll for (uint32_t bit = 0; bit < 8; ++bit) { const uint64_t row = row_start + bit; - if (row < codes_len) { + if (row < codes_len && static_cast(codes[row]) < values_len) { const uint64_t value_idx = values_bit_offset + static_cast(codes[row]); const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; packed |= static_cast(value << bit); @@ -69,8 +74,9 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, const IndexType *const __restrict codes, \ uint64_t codes_len, \ const ValueType *const __restrict values, \ + uint64_t values_len, \ ValueType *const __restrict output) { \ - dict_kernel(codes, codes_len, values, output); \ + dict_kernel(codes, codes_len, values, values_len, output); \ } // Generate dict kernel for all index types (unsigned integers) for a given value type @@ -84,9 +90,10 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, extern "C" __global__ void dict_bool_##index_suffix(const IndexType *const __restrict codes, \ uint64_t codes_len, \ const uint8_t *const __restrict values, \ + uint64_t values_len, \ uint64_t values_bit_offset, \ uint8_t *const __restrict output) { \ - dict_bool_kernel(codes, codes_len, values, values_bit_offset, output); \ + dict_bool_kernel(codes, codes_len, values, values_len, values_bit_offset, output); \ } // Generate fixed-width kernels for all native ptypes and decimal values. diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index aca22975bdb..b5b866ad70e 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -143,6 +143,7 @@ async fn execute_dict_bool_typed( let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; let codes_len_u64 = codes_len as u64; + let values_len_u64 = values_len as u64; let values_offset_u64 = values_meta.offset() as u64; let codes_ptype = I::PTYPE.to_string(); @@ -151,6 +152,7 @@ async fn execute_dict_bool_typed( args.arg(&codes_view) .arg(&codes_len_u64) .arg(&values_view) + .arg(&values_len_u64) .arg(&values_offset_u64) .arg(&mut output_slice); })?; @@ -171,6 +173,7 @@ async fn execute_dict_prim_typed VortexResult { assert!(!codes.is_empty()); let codes_len = codes.len(); + let values_len_u64 = values.len() as u64; let PrimitiveDataParts { ptype: value_ptype, @@ -202,6 +205,7 @@ async fn execute_dict_prim_typed( + codes: Vec, + valid: &[bool], + sane: T, + ) -> (PrimitiveArray, PrimitiveArray) { + let validity: Validity = valid.iter().copied().collect(); + let sane_codes: Vec = codes + .iter() + .zip(valid) + .map(|(&code, &valid)| if valid { code } else { sane }) + .collect(); + ( + PrimitiveArray::new(Buffer::from(codes), validity.clone()), + PrimitiveArray::new(Buffer::from(sane_codes), validity), + ) + } + + async fn gpu_dict_to_host( + dict: DictArray, + cuda_ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let canonical = DictExecutor.execute(dict.into_array(), cuda_ctx).await?; + Ok(canonical + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array()) + } + + async fn assert_gpu_matches( + codes: PrimitiveArray, + sane: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let dict = DictArray::try_new(codes.into_array(), values.clone())?; + let expected = DictArray::try_new(sane.into_array(), values)?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = + VarBinViewArray::from_iter_str(["a", "bb", "a string longer than twelve bytes"]); + // `u32::MAX` points 64 GiB past the views, so an unchecked gather faults. + let (codes, sane) = garbage_under_nulls( + vec![0u32, 200, 2, u32::MAX, 1], + &[true, false, true, false, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_primitive_far_out_of_range_codes_under_nulls() -> VortexResult<()> { + // These codes point gigabytes past the dictionary, so an unchecked gather faults. + let values = PrimitiveArray::new(Buffer::from(vec![10i64, 20, 30]), NonNullable); + let (codes, sane) = garbage_under_nulls( + vec![0u32, u32::MAX, 2, 0x8000_0000, 1], + &[true, false, true, false, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_decimal_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = DecimalArray::from_iter([1i128, 2, 3], DecimalDType::new(38, 0)); + let (codes, sane) = garbage_under_nulls(vec![1u32, u32::MAX, 2], &[true, false, true], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_bool_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = BoolArray::from_iter([true, false, true]); + let (codes, sane) = garbage_under_nulls( + vec![2u64, u64::MAX, 0, 1_000_000, 1, 9, 2, 0, 1], + &[true, false, true, false, true, false, true, true, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + /// The shape the `compress_gpu` fuzzer found: a one-value dictionary whose null code is + /// 3,960,089,894. An unchecked gather reads about 4 GB past the dictionary and fails with + /// `CUDA_ERROR_ILLEGAL_ADDRESS`. + #[crate::test] + async fn test_cuda_dict_fuzz_garbage_null_code() -> VortexResult<()> { + let values = PrimitiveArray::new(Buffer::from(vec![0u8]), NonNullable); + let (codes, sane) = garbage_under_nulls(vec![0u32, 3_960_089_894], &[true, false], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_all_null_codes() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str(["x", "a string longer than twelve bytes"]); + let (codes, sane) = + garbage_under_nulls(vec![5_000u16, 7, u16::MAX], &[false, false, false], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_codes_sliced_across_block_boundaries() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str( + (0..50).map(|i| format!("value-{i}-{}", "x".repeat(i % 20))), + ); + let len = 3_000u16; + let valid: Vec = (0..len).map(|i| i % 11 != 0).collect(); + let raw: Vec = (0..len) + .map(|i| { + if valid[usize::from(i)] { + (i * 7) % 50 + } else { + 60_000 + } + }) + .collect(); + let (codes, sane) = garbage_under_nulls(raw, &valid, 0); + for range in [0..1023, 1023..1025, 1000..2049, 2047..2048, 1024..3000] { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let sliced = codes.clone().into_array().slice(range.clone())?; + let sliced_sane = sane.clone().into_array().slice(range.clone())?; + let dict = DictArray::try_new(sliced, values.clone().into_array())?; + let expected = DictArray::try_new(sliced_sane, values.clone().into_array())?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_u16_codes_full_range() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str((0..65_536u32).map(|i| format!("v{i:05}"))); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u16, u16::MAX, 32_768, 1, u16::MAX - 1]), + NonNullable, + ); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_u32_codes_beyond_u16() -> VortexResult<()> { + let values = + VarBinViewArray::from_iter_str((0..70_000u32).map(|i| format!("dictionary-entry-{i}"))); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u32, 69_999, 65_536, 12_345]), + NonNullable, + ); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_inline_boundary_and_long_values() -> VortexResult<()> { + let long = "z".repeat(4_097); + let values = + VarBinViewArray::from_iter_str(["", "twelve-bytes", "thirteen-byte", long.as_str()]); + let codes = PrimitiveArray::new(Buffer::from(vec![3u8, 0, 1, 2, 3, 2, 1, 0]), NonNullable); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_binary_non_utf8_values() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_bin([ + Some(vec![0xffu8, 0xfe, 0x00]), + Some(vec![0x80]), + Some(vec![]), + None, + Some(vec![0xc3; 20]), + ]); + let (codes, sane) = garbage_under_nulls( + vec![0u8, 1, 2, 3, 4, 77, 4, 0], + &[true, true, true, true, true, false, true, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + /// Moves the codes' buffer to the device, as a file scan leaves it. + async fn codes_on_device( + codes: PrimitiveArray, + cuda_ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let handle = cuda_ctx + .ensure_on_device(codes.buffer_handle().clone()) + .await?; + Ok(PrimitiveArray::from_buffer_handle( + handle, + codes.ptype(), + codes.validity()?, + )) + } + + /// Decodes with device-resident codes, as a file scan leaves them, and copies the result back + /// to the host. + async fn assert_gpu_matches_with_device_codes( + codes: PrimitiveArray, + sane: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let device_codes = codes_on_device(codes, &mut cuda_ctx).await?; + let dict = DictArray::try_new(device_codes.into_array(), values.clone())?; + let expected = DictArray::try_new(sane.into_array(), values)?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_nullable_values_and_codes_device_codes() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([Some("x"), None, Some("zz")]); + let (codes, sane) = garbage_under_nulls( + vec![0u16, 9_999, 2, 1, u16::MAX, 0], + &[true, false, true, true, false, true], + 0, + ); + assert_gpu_matches_with_device_codes(codes, sane, values.into_array()).await + } } From 076796690d6cbdfdaddda22748f40b79f6141f58 Mon Sep 17 00:00:00 2001 From: Tanish Gudise Date: Tue, 10 Nov 2026 07:16:18 -0500 Subject: [PATCH 06/13] feat(btrblocks): allow string dictionaries in the CUDA compression preset (#10917) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Tracking Issue: #10914 - Depends on #10915 and #10916; this branch contains both. The diff of its own is one line plus a test. `CompressionMode::Cuda` excludes `StringDictScheme`, but `execute_dict_varbinview` already decodes string dictionaries on the GPU. Keeping the scheme makes GPU-mode files 16.0% smaller across compress-bench's 17-dataset GPU suite (3,351.3 → 2,813.9 MB). Per-dataset gains: CMSprovider 45.7%, Food 52.9%, HashTags 33.2%, both TPC-H l_comment variants 12.2%, Euro2016 10.7%. The issue has the full table. ## Changes - `vortex-btrblocks/src/builder.rs`: drops `StringDictScheme` from the Cuda exclusions and updates the comment. A new test checks that the preset keeps string dictionaries but not binary ones. This is the only change to CPU-compressor code. Evidence with this branch (A40, CUDA 13.0): - `--gpu-verify` passes on all 17 datasets. develop plus only this line fails 5: taxi, CMSprovider, Euro2016, Food, HashTags. - The CPU-fallback census is identical to develop's: 59 constant-array fallbacks, none from dictionaries. - The dictionary gather takes at most 6.1% of decode time and reaches 62-85% of the measured 563 GB/s. - `vortex-cuda` lib suite on this branch: 723 passed, 0 failed. - `compress_gpu` fuzz (CI flags, ASAN): 100,000 runs in 256 s, 0 crashes. ## Benchmarks Setup: - `compress-bench --gpu-decompress --iterations 5`, a fresh process per run, `taskset -c 2-7`, 3 ABBA rounds. - Speedup = develop time / branch time. - The same-design A/A floor is 0.14 in |log|, so a "faster" claim needs a 95% CI above 1 and a speedup above 1.32×. | Dataset | File MB, develop → branch | Speedup [95% CI] | | |---|---|---:|---| | Food | 84.7 → 39.9 | 1.55× [1.41, 1.71] | faster | | CMSprovider | 771.0 → 418.5 | 1.37× [0.66, 2.86] | within noise | | Euro2016 | 163.3 → 145.9 | 1.46× [0.92, 2.32] | within noise | | TPC-H l_comment chunked | 169.6 → 148.8 | 1.36× [0.55, 3.34] | within noise | | taxi | 52.0 → 51.5 | 1.28× [0.32, 5.11] | within noise | | HashTags | 235.2 → 157.0 | 1.22× [0.89, 1.69] | within noise | | TPC-H l_comment canonical | 169.6 → 148.8 | 1.18× [0.71, 1.96] | within noise | | Arade | 144.2 → 141.7 | 1.05× [0.42, 2.68] | within noise | The nine datasets whose files don't change got one AB/BA round each. Eight came out at 1.01-1.18×, and one (wide table cols=1000 chunks=50) at 0.91×, inside the A/A floor. The measurements are from an A40 (GA102); the repo's GPU benchmarks run on A10G. ## API changes `only_cuda_compatible()`, and therefore `cuda_write_strategy`, can now emit string dictionaries. - An `ArrowDeviceArrayStream` with the default `DictionaryExport::Preserve` rejects a batch whose Arrow schema differs from the first. - With this change, 6 of CMSprovider's 14 string columns mix dictionary and constant chunks. - Today's Cuda files already do this in numeric columns: 7 of taxi's 18 numeric columns and 5 of CMSprovider's 12. - Streams with `DictionaryExport::Decode`, which the FFI scan uses, are unaffected. Overlap: #10425 adds a line next to the removed one, so whichever merges second needs a one-line rebase. #9762 changes dictionary selection, so the sizes are as of a59294e35. ## AI assistance Developed with Claude Code (Claude Opus 5.5), which wrote most of the code, tests and this description under my direction. I reviewed the diff and the measurements. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Tanish Gudise Co-authored-by: Claude Opus 5.5 Co-authored-by: Alexander Droste --- vortex-btrblocks/src/builder.rs | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6d261788780..6d1ce2323f8 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -63,10 +63,12 @@ impl CompressionMode { excluded.push(binary::ZstdBuffersScheme.id()); } Self::Cuda => { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode - // paths. Delta has a CUDA decode kernel, but stays excluded until GPU delta - // decode is benchmarked against the schemes it would displace. + // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export, and + // string dictionaries, which the CUDA dictionary gather decodes into views over + // the decoded values. Other string fragmentation schemes and binary dictionaries + // still require unsupported decode paths. Delta has a CUDA decode kernel, but + // stays excluded until GPU delta decode is benchmarked against the schemes it + // would displace. excluded.extend([ integer::DeltaScheme::default().id(), integer::SparseScheme.id(), @@ -76,7 +78,6 @@ impl CompressionMode { float::NullDominatedSparseScheme.id(), string::NullDominatedSparseScheme.id(), fixed_size_list::FixedSizeListSparseScheme.id(), - string::StringDictScheme.id(), binary::BinaryDictScheme.id(), ]); // Both binary Zstd schemes are kept. Buffer-level compression preserves binary @@ -321,4 +322,14 @@ mod tests { assert!(has_delta(builder.clone().with_compact())); assert!(!has_delta(builder.only_cuda_compatible())); } + + #[test] + fn cuda_keeps_string_dictionaries_but_not_binary_ones() { + let cuda = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .only_cuda_compatible(); + let has = |id: SchemeId| cuda.allowed_schemes().iter().any(|s| s.id() == id); + assert!(has(string::StringDictScheme.id())); + assert!(!has(binary::BinaryDictScheme.id())); + } } From 80186ab2e5e14b22cfdfa5e5a997efacb2213bd3 Mon Sep 17 00:00:00 2001 From: Nemo Yu <83347615+HarukiMoriarty@users.noreply.github.com> Date: Tue, 10 Nov 2026 06:25:42 -0600 Subject: [PATCH 07/13] feat(spatial): prune contains and within filters using AABB statistics (#11372) ## Summary `ST_Contains` filters currently scan even when stored geometry bounds prove that a file or zone cannot match. Register one `SpatialContainsPrune` rule that consumes the existing `GeometryAabb` statistic in both file and zone pruning. Containment is directional: - `ST_Contains(column, constant)` can skip a scope whose aggregate box cannot cover the constant's box. - `ST_Contains(constant, column)` can skip a scope whose aggregate box is disjoint from the constant's box. Partial overlap must remain eligible because individual rows can still match. DuckDB already lowers `ST_Within(a, b)` to `ST_Contains(b, a)`, so it uses the same rule. Missing/null bounds, null/empty constants, touching boxes in the constant-container case, and unsupported predicate shapes remain conservative. There is no containment proof for negated filters. The shared column matcher also accepts direct and nested field paths so file-level proofs can use footer statistics. Unit cases cover bounding-box relations and operand order; write/read tests exercise file pruning without zones and zone pruning without file statistics. --------- Signed-off-by: Nemo Yu --- Cargo.lock | 2 + vortex-spatial/Cargo.toml | 2 + vortex-spatial/README.md | 8 +- vortex-spatial/src/lib.rs | 2 + vortex-spatial/src/prune/contains.rs | 416 +++++++++++++++++++++++++++ vortex-spatial/src/prune/mod.rs | 42 +-- 6 files changed, 453 insertions(+), 19 deletions(-) create mode 100644 vortex-spatial/src/prune/contains.rs diff --git a/Cargo.lock b/Cargo.lock index 88710226f57..49bad204b42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12518,6 +12518,8 @@ dependencies = [ "vortex-edition", "vortex-error", "vortex-expr", + "vortex-file", + "vortex-io", "vortex-layout", "vortex-mask", "vortex-session", diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index b48831cf879..bfb82b4e652 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -39,6 +39,8 @@ divan = { workspace = true } mimalloc = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-file = { workspace = true } +vortex-io = { workspace = true } vortex-layout = { workspace = true } vortex-spatial = { path = ".", features = ["_test-harness"] } diff --git a/vortex-spatial/README.md b/vortex-spatial/README.md index c3af7eec152..8aeae021435 100644 --- a/vortex-spatial/README.md +++ b/vortex-spatial/README.md @@ -63,16 +63,22 @@ Avoid per-row `scalar_at` or validity lookups in either path. Materialize coordi The current rules can prune: - positive `ST_Intersects(geometry_column, constant_geometry)` filters when the boxes are disjoint; +- positive `ST_Contains(geometry_column, constant_geometry)` filters when the file or zone box cannot cover the constant's box; +- positive `ST_Contains(constant_geometry, geometry_column)` filters (equivalently `ST_Within(geometry_column, constant_geometry)`) when the boxes are disjoint; - `<`, `<=`, `>`, and `>=` comparisons on `ST_Distance(geometry_column, constant_geometry)` using box-distance bounds. These proofs are conservative. If the expression has a different shape, the constant is null or empty, or the file has no AABB statistic, the chunk is kept. Box overlap alone never proves that two geometries intersect. +Containment is directional. A file or zone box that extends outside a constant's box may still contain rows within that constant, so that alone cannot prune `ST_Within`. Box coverage does not prove geometry containment either; `NOT ST_Contains` and `NOT ST_Within` have no spatial pruning proof. + +The same rules work with file-level AABB statistics, including nested geometry fields. To record file bounds, use `include_nested_statistics().with_file_statistics(...)` with `GeometryAabb.bind(EmptyOptions)` in the aggregate list. This replaces the default file aggregate list. The [containment pruning tests](src/prune/contains.rs) exercise file pruning without zone maps and zone pruning without file statistics. + ## Interoperability and limits Arrow import and export use the standard GeoArrow extension names and preserve CRS metadata. Native import only supports separated coordinates; interleaved GeoArrow arrays are rejected. -DuckDB uses WKB for general geometry values. Its Vortex expression lowering converts supported WKB constants to native scalars and pushes `ST_Contains`, `ST_Distance`, and `ST_Intersects` into a Vortex scan. +DuckDB uses WKB for general geometry values. Its Vortex expression lowering converts supported WKB constants to native scalars and pushes `ST_Contains`, `ST_Within`, `ST_Distance`, and `ST_Intersects` into a Vortex scan. `ST_Within(a, b)` lowers to `ST_Contains(b, a)` and shares its compute and pruning implementation. Some important limits: diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index ecb3e1e55fd..99fc5f8e556 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -23,6 +23,7 @@ use crate::extension::Point; use crate::extension::Polygon; use crate::extension::Rect; use crate::extension::WellKnownBinary; +use crate::prune::SpatialContainsPrune; use crate::prune::SpatialDistancePrune; use crate::prune::SpatialIntersectsPrune; use crate::scalar_fn::area::SpatialArea; @@ -89,6 +90,7 @@ pub fn initialize(session: &VortexSession) { session.aggregate_fns().register(GeometryAabb); // Register the spatial pruning rules that use that AABB. + session.stats().register_rewrite(SpatialContainsPrune); session.stats().register_rewrite(SpatialDistancePrune); session.stats().register_rewrite(SpatialIntersectsPrune); diff --git a/vortex-spatial/src/prune/contains.rs b/vortex-spatial/src/prune/contains.rs new file mode 100644 index 00000000000..f4ff037c6c8 --- /dev/null +++ b/vortex-spatial/src/prune/contains.rs @@ -0,0 +1,416 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Contains` pruning with a geometry column and a constant in either operand order. + +use vortex_array::expr::BoundExpression; +use vortex_array::expr::bound::get_item; +use vortex_array::expr::bound::gt; +use vortex_array::expr::bound::lit; +use vortex_array::expr::bound::lt; +use vortex_array::expr::bound::or; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_error::VortexResult; +use vortex_expr::stats::rewrite::StatsRewriteRule; +use vortex_session::VortexSession; + +use super::aabb_stat; +use super::geometry_and_constant; +use super::is_column; +use super::query_aabb; +use crate::scalar_fn::contains::SpatialContains; + +/// Prunes files and zones for positive `ST_Contains` filters using their aggregate AABB. +/// +/// For `contains(column, constant)`, the aggregate box must cover the constant's box. For +/// `contains(constant, column)` (equivalently `within(column, constant)`), disjoint boxes prove +/// no row can match. Neither box relation proves containment, so this rule supplies no proof +/// for `NOT ST_Contains`. +#[derive(Debug)] +pub struct SpatialContainsPrune; + +impl StatsRewriteRule for SpatialContainsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + SpatialContains.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + session: &VortexSession, + ) -> VortexResult> { + let Some((geom, constant)) = geometry_and_constant(expr)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, session)? else { + return Ok(None); + }; + let bounds = aabb_stat(geom); + let field = |name| get_item(name, bounds.clone()); + let outside = if is_column(expr.child(0)) { + // Every row's box lies inside the aggregate box. If even that union cannot cover + // the constant, no individual geometry can contain it. + [ + gt(field("xmin"), lit(query.min().x)), + gt(field("ymin"), lit(query.min().y)), + lt(field("xmax"), lit(query.max().x)), + lt(field("ymax"), lit(query.max().y)), + ] + } else { + // The union may extend outside the constant while some rows remain inside it. + // Only strict separation proves every row is outside; touching boxes must scan. + [ + gt(field("xmin"), lit(query.max().x)), + gt(field("ymin"), lit(query.max().y)), + lt(field("xmax"), lit(query.min().x)), + lt(field("ymax"), lit(query.min().y)), + ] + }; + let [xmin, ymin, xmax, ymax] = outside; + Ok(Some(or(or(xmin, ymin), or(xmax, ymax)))) + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions as AggregateOptions; + use vortex_array::arrays::StructArray; + use vortex_array::expr::Expression; + use vortex_array::expr::fill_null; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::not; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_buffer::ByteBuffer; + use vortex_buffer::ByteBufferMut; + use vortex_edition::EDITION_DECLARATIONS; + use vortex_edition::EDITION_FAMILIES; + use vortex_edition::EditionSessionExt; + use vortex_edition::declarations::core::CORE_2026_08_3; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_expr::stats::rewrite::BoundExpressionPruneExt; + use vortex_file::OpenOptionsSessionExt; + use vortex_file::WriteOptionsSessionExt; + use vortex_io::runtime::Handle; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; + use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; + use vortex_layout::layouts::repartition::RepartitionStrategy; + use vortex_layout::layouts::repartition::RepartitionWriterOptions; + use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; + use vortex_layout::layouts::zoned::writer::ZonedStrategy; + use vortex_layout::layouts::zoned::zone_map::ZoneMap; + use vortex_mask::Mask; + use vortex_session::VortexSession; + + use crate::aggregate_fn::GeometryAabb; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; + use crate::scalar_fn::contains::SpatialContains; + use crate::scalar_fn::envelope::SpatialEnvelope; + use crate::test_harness::multipoint_column; + use crate::test_harness::nullable_point_column; + use crate::test_harness::nullable_rect_column; + use crate::test_harness::point_column; + use crate::test_harness::rect_column; + use crate::test_harness::spatial_session; + + fn contains(column: Expression, constant: Scalar, column_contains: bool) -> Expression { + let operands = if column_contains { + [column, lit(constant)] + } else { + [lit(constant), column] + }; + SpatialContains.new_expr(EmptyOptions, operands) + } + + #[rstest] + #[case::covers([-1.0, -1.0, 3.0, 3.0], false, false)] + #[case::equal([0.0, 0.0, 2.0, 2.0], false, false)] + #[case::inside([0.5, 0.5, 1.5, 1.5], true, false)] + #[case::overlaps_right([1.0, 0.0, 3.0, 2.0], true, false)] + #[case::overlaps_left([-1.0, 0.0, 1.0, 2.0], true, false)] + #[case::overlaps_above([0.0, 1.0, 2.0, 3.0], true, false)] + #[case::overlaps_below([0.0, -1.0, 2.0, 1.0], true, false)] + #[case::disjoint_right([3.0, 0.0, 4.0, 2.0], true, true)] + #[case::disjoint_left([-2.0, 0.0, -1.0, 2.0], true, true)] + #[case::disjoint_above([0.0, 3.0, 2.0, 4.0], true, true)] + #[case::disjoint_below([0.0, -2.0, 2.0, -1.0], true, true)] + #[case::touches([2.0, 0.0, 3.0, 2.0], true, false)] + fn box_relations_respect_operand_order( + #[case] bounds: [f64; 4], + #[case] prune_column_contains: bool, + #[case] prune_constant_contains: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let expected = if column_contains { + prune_column_contains + } else { + prune_constant_contains + }; + let zones = aabb_zone_map(&dtype, &[bounds])?; + assert_eq!(zones.prune(&proof, &session)?, Mask::from_iter([expected])); + Ok(()) + } + + #[rstest] + fn equal_points_are_not_pruned( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = point_column(vec![0.0], vec![0.0])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let zones = aabb_zone_map(&dtype, &[[0.0; 4]])?; + assert_eq!(zones.prune(&proof, &session)?, Mask::new_false(1)); + Ok(()) + } + + #[rstest] + fn missing_and_null_bounds_keep_zones( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let aabb = GeometryAabb.bind(AggregateOptions); + let null_bounds = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&[( + aabb.to_string().as_str(), + nullable_rect_column(vec![None, None])?, + )])?, + Arc::new([aabb]), + 1, + 2, + )?; + for zones in [empty_zone_map(&dtype)?, null_bounds] { + assert_eq!(zones.prune(&proof, &session)?, Mask::new_false(2)); + } + Ok(()) + } + + #[rstest] + fn null_and_empty_constants_decline( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let empty = multipoint_column(vec![vec![]])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = empty.dtype().clone(); + for query in [Scalar::null(dtype.as_nullable()), empty] { + assert!( + contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .is_none() + ); + } + Ok(()) + } + + #[rstest] + fn negated_and_computed_predicates_decline( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + for predicate in [ + not(contains(root(), query.clone(), column_contains)), + contains( + SpatialEnvelope.new_expr(EmptyOptions, [root()]), + query, + column_contains, + ), + SpatialContains.new_expr(EmptyOptions, [root(), root()]), + ] { + assert!(predicate.bind(&dtype)?.falsify(&session)?.is_none()); + } + Ok(()) + } + + fn file_session(handle: Handle) -> VortexResult { + let session = VortexSession::empty().with_handle(handle); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + vortex_file::register_default_encodings(&session); + crate::initialize(&session); + Ok(session) + } + + #[rstest] + fn prunes_files_without_zone_maps( + #[values(false, true)] nested: bool, + #[values(false, true)] record_bounds: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + let points = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((1.0, 1.0))])?; + let batch = StructArray::from_fields(&[("geom", points)])?.into_array(); + let (batch, column) = if nested { + ( + StructArray::from_fields(&[("location", batch)])?.into_array(), + get_item("geom", get_item("location", root())), + ) + } else { + (batch, get_item("geom", root())) + }; + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))) + .include_nested_statistics() + .with_file_statistics(if record_bounds { + vec![GeometryAabb.bind(AggregateOptions)] + } else { + vec![] + }) + .write(&mut buffer, batch.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + let mut ctx = session.create_execution_ctx(); + let far = + nullable_point_column(vec![Some((10.0, 10.0))])?.execute_scalar(0, &mut ctx)?; + let near = + nullable_point_column(vec![Some((0.0, 0.0))])?.execute_scalar(0, &mut ctx)?; + let enclosing = + rect_column(vec![(-1.0, -1.0, 2.0, 2.0)])?.execute_scalar(0, &mut ctx)?; + + for (query, cannot_match) in [(far, true), (near, false), (enclosing, column_contains)] + { + let predicate = contains(column.clone(), query, column_contains); + let matches = batch + .clone() + .apply(&fill_null(predicate.clone(), lit(false)))? + .execute::(&mut ctx)?; + let expected = record_bounds && cannot_match; + assert_eq!(file.can_prune(&predicate)?, expected, "{predicate}"); + let mask = file + .layout_reader()? + .pruning_evaluation(&(0..3), &predicate.bind(file.dtype())?, Mask::new_true(3))? + .await?; + assert_eq!(mask, Mask::from_iter([!expected; 3]), "{predicate}"); + assert_eq!( + &mask & &matches, + matches, + "pruning dropped matches for {predicate}" + ); + } + Ok(()) + }) + } + + #[rstest] + fn prunes_zones_without_file_statistics( + #[values(false, true)] record_bounds: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + let array = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((1.0, 1.0)), + Some((10.0, 10.0)), + Some((11.0, 11.0)), + ])?; + let aggregates = if record_bounds { + vec![GeometryAabb.bind(AggregateOptions)] + } else { + vec![] + }; + let strategy = RepartitionStrategy::new( + ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(2).vortex_expect("non-zero zone size"), + aggregate_fns: Some(aggregates.into()), + ..Default::default() + }, + ), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: 2, + block_size_target: None, + canonicalize: false, + }, + ); + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(Arc::new(strategy)) + .with_file_statistics(vec![]) + .write(&mut buffer, array.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + assert!(file.file_stats().is_none()); + let mut ctx = session.create_execution_ctx(); + let query = + nullable_point_column(vec![Some((10.0, 10.0))])?.execute_scalar(0, &mut ctx)?; + let predicate = contains(root(), query, column_contains); + let matches = array + .clone() + .apply(&fill_null(predicate.clone(), lit(false)))? + .execute::(&mut ctx)?; + let mask = file + .layout_reader()? + .pruning_evaluation(&(0..4), &predicate.bind(file.dtype())?, Mask::new_true(4))? + .await?; + assert_eq!( + mask, + Mask::from_iter([!record_bounds, !record_bounds, true, true]) + ); + assert_eq!( + &mask & &matches, + matches, + "pruning dropped matches for {predicate}" + ); + Ok(()) + }) + } +} diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index dd1a6951503..9ea202c8cbf 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -1,20 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned -//! bounding box (AABB). +//! File and zone pruning for spatial filters, using the [`GeometryAabb`] axis-aligned +//! bounding box (AABB) statistic. //! //! One module per predicate. To prune a new one: recover the geometry column and the constant -//! (symmetric predicates can use `geometry_and_constant`), reduce the constant to its box with -//! `query_aabb`, build the proof over `aabb_stat` from the box-relation builders (`min_dist_sq`, -//! `max_dist_sq`), and register the rule in [`crate::initialize`]. No new statistic or -//! file-format change is needed. +//! with `geometry_and_constant` (preserving operand order for asymmetric predicates), reduce the +//! constant with `query_aabb`, build the proof over `aabb_stat`, and register the rule in +//! [`crate::initialize`]. No new statistic or file-format change is needed. +mod contains; mod distance; mod intersects; #[cfg(test)] mod test_harness; +pub use contains::SpatialContainsPrune; pub use distance::SpatialDistancePrune; use geo::BoundingRect; use geo::Rect as SpatialRect; @@ -34,6 +35,7 @@ use vortex_array::expr::bound::lit; use vortex_array::expr::bound::lt; use vortex_array::expr::bound::lt_eq; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -44,22 +46,19 @@ use crate::aggregate_fn::GeometryAabb; use crate::extension::is_native_geometry; use crate::extension::single_geometry; -/// Splits a symmetric two-operand spatial predicate into the scope-rooted geometry column and the -/// constant operand's scalar. +/// Splits a two-operand spatial predicate into the geometry column and the constant operand's scalar. /// /// `None` means the rule must decline: the expression doesn't have the `f(column, constant)` -/// shape (in either operand order), or the column's dtype carries no [`GeometryAabb`] statistic. -/// An asymmetric predicate (e.g. a future contains) must recover which operand is the column -/// itself instead of calling this. +/// shape (in either operand order), or the column's dtype does not support [`GeometryAabb`]. +/// Asymmetric predicates must also check which operand is the column when choosing their proof. fn geometry_and_constant( expr: &BoundExpression, ) -> VortexResult> { - // The predicate is symmetric, so the column (scope root) and the constant may be on either - // side. + // Zone predicates use the column as their root; file predicates refer to it through fields. let (lhs, rhs) = (expr.child(0), expr.child(1)); - let (geom, constant) = if lhs.is_root() { + let (geom, constant) = if is_column(lhs) { (lhs, rhs) - } else if rhs.is_root() { + } else if is_column(rhs) { (rhs, lhs) } else { return Ok(None); @@ -74,6 +73,13 @@ fn geometry_and_constant( Ok(constant.as_opt::().map(|scalar| (geom, scalar))) } +fn is_column(mut expr: &BoundExpression) -> bool { + while expr.is::() { + expr = expr.child(0); + } + expr.is_root() +} + /// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent /// (e.g. an empty geometry). /// @@ -93,10 +99,10 @@ fn query_aabb( Ok(single_geometry(constant, &mut exec)?.bounding_rect()) } -/// The chunk's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. +/// The scope's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. /// -/// A chunk written without the statistic reads as null here; every proof built on top must let -/// that null propagate to its root, where the zone map keeps the chunk. +/// A missing statistic reads as null here; every proof built on top must let that null +/// propagate to its root, keeping the file or zone. fn aabb_stat(geom: &BoundExpression) -> BoundExpression { // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so // proofs can `get_item` the coordinate fields. From 37e3e9fb6530fd0af95446625167595556d16408 Mon Sep 17 00:00:00 2001 From: ChristopherWirt Date: Tue, 10 Nov 2026 15:36:14 +0000 Subject: [PATCH 08/13] fix(layout): normalize rebuilt ListView children before writing (#11371) Fixes #11370. Writing sliced `ListViewArray` chunks through `ListLayoutStrategy` can fail with `Array vortex.slice is not registered for serialization`. Rebuilding the view can leave a lazy elements slice, which is passed directly to the child writer. The ListArray path already normalizes its children with `reset_offsets`. This applies the same normalization after rebuilding a ListView, so its elements can be serialized too. --------- Signed-off-by: Christopher Wirt --- vortex-layout/src/layouts/list/reader.rs | 20 ++++++++++++++++---- vortex-layout/src/layouts/list/writer.rs | 13 +++++++------ 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 79f1066fc3d..334a4f21efb 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -623,6 +623,7 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::Expression; @@ -1283,16 +1284,27 @@ mod tests { } #[rstest] - #[case::one_cut(&[2], false)] - #[case::two_cuts(&[1, 3], false)] - #[case::null_and_empty_at_cuts(&[1, 3], true)] + #[case::one_cut(&[2], false, false)] + #[case::two_cuts(&[1, 3], false, false)] + #[case::null_and_empty_at_cuts(&[1, 3], true, false)] + #[case::list_view_chunks(&[1, 3], true, true)] + #[case::non_nullable_list_view_chunks(&[1, 3], false, true)] #[tokio::test] async fn sliced_list_chunks_round_trip( #[case] cuts: &[usize], #[case] nullable: bool, + #[case] as_list_view: bool, ) -> VortexResult<()> { let list = create_ten_element_list_array(nullable); - let chunks = sliced_chunks(&list, cuts)?; + let source = if as_list_view { + let mut ctx = SESSION.create_execution_ctx(); + list.clone() + .execute::(&mut ctx)? + .into_array() + } else { + list.clone() + }; + let chunks = sliced_chunks(&source, cuts)?; assert_round_trips(&chunked_list_strategy(), chunks, list).await } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index dfd7c856e06..dbcdae74169 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -289,15 +289,16 @@ fn canonicalize_to_list_parts( exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let canonical = array.execute_until::(exec_ctx)?; - if let Some(list) = canonical.as_opt::() { - // A sliced list keeps its parent's whole `elements`, but `global_offsets` assumes each - // chunk's offsets span exactly `[0, elements.len()]`. - Ok(list.reset_offsets(false, exec_ctx)?.into_data_parts()) + let list = if let Some(list) = canonical.as_opt::() { + list.into_owned() } else if let Some(view) = canonical.as_opt::() { - Ok(list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts()) + list_from_list_view(view.into_owned(), exec_ctx)? } else { unreachable!("AnyList matcher guarantees List or ListView") - } + }; + // Sliced lists can retain unused elements, and rebuilding a list view can leave a lazy + // elements slice. Normalize both before passing their children to serializing strategies. + Ok(list.reset_offsets(false, exec_ctx)?.into_data_parts()) } /// Rebase a chunk's local `offsets` into global `u64` positions for the whole-column `offsets` From 56108c8c4823f34802463402b26075f4ed624d47 Mon Sep 17 00:00:00 2001 From: ChristopherWirt Date: Tue, 10 Nov 2026 16:49:13 +0000 Subject: [PATCH 09/13] feat(jni): make Java read coalescing configurable (#11369) - Tracking Issue: #11366 - Fixes #11366 Java-backed readers always use `CoalesceConfig::object_storage()` (a 1 MiB gap, a 16 MiB maximum merged span). Callers can't change it. The right values depend on the storage underneath. Local SSD, spinning disk and object stores each trade request count against extra bytes read differently. This PR lets Java and Spark callers set them, keeping the current 1 MiB / 16 MiB as the defaults. --------- Signed-off-by: Christopher Wirt --- .../java/dev/vortex/io/NativeReadable.java | 25 +++++ .../dev/vortex/io/NativeIOBridgeTest.java | 70 ++++++++++++++ java/vortex-spark/README.md | 25 +++++ .../dev/vortex/spark/io/HadoopReadable.java | 28 +++++- .../java/dev/vortex/spark/io/VortexIo.java | 72 ++++++++------ .../vortex/spark/io/HadoopReadableTest.java | 2 + .../dev/vortex/spark/io/VortexIoTest.java | 56 +++++++++++ vortex-jni/src/data_source.rs | 4 +- vortex-jni/src/file.rs | 4 +- vortex-jni/src/io/mod.rs | 15 ++- vortex-jni/src/io/read_at.rs | 95 ++++++++++++++++++- 11 files changed, 358 insertions(+), 38 deletions(-) diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java index 84f264cc952..b92cd333df6 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java +++ b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java @@ -19,6 +19,31 @@ * source or scan built on top of it has been closed. */ public interface NativeReadable extends Closeable { + /** Default maximum gap between coalesced reads, in bytes (1 MiB), matching the native object-storage policy. */ + int DEFAULT_COALESCE_DISTANCE = 1024 * 1024; + + /** Default maximum span of a coalesced read, in bytes (16 MiB). */ + int DEFAULT_COALESCE_MAX_SIZE = 16 * 1024 * 1024; + + /** + * Maximum gap in bytes between reads that native code may merge. Must be non-negative; zero merges only adjacent or + * overlapping reads. Captured once when the source is registered. + * + *

Larger gaps reduce storage requests but may re-read segments already held in the file's segment cache. + */ + default int coalesceDistance() { + return DEFAULT_COALESCE_DISTANCE; + } + + /** + * Maximum payload span in bytes of a merged read, excluding leading alignment padding. Must be positive. Individual + * reads may exceed this size. Captured once when the source is registered. Native code caps values near + * Integer.MAX_VALUE to reserve alignment headroom within the Java ByteBuffer limit. + */ + default int coalesceMaxSize() { + return DEFAULT_COALESCE_MAX_SIZE; + } + /** * Unique name of this source, typically its original location (URI or file path). Used to key per-session caches * and in native error messages, so it must be stable and unique across sources; it must not contain glob characters diff --git a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java index 9b1d4722499..eef5d4a98d2 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java @@ -48,6 +48,8 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; /** Round-trip tests for the caller-provided I/O bridge ({@link NativeReadable} / {@link NativeWritable}). */ public final class NativeIOBridgeTest { @@ -70,17 +72,39 @@ private static final class FileChannelReadable implements NativeReadable { private final String name; private final FileChannel channel; private final long length; + private final int coalesceDistance; + private final int coalesceMaxSize; + private int distanceCalls; + private int maxSizeCalls; FileChannelReadable(Path path) throws IOException { this(path.toString(), path); } FileChannelReadable(String name, Path path) throws IOException { + this(name, path, DEFAULT_COALESCE_DISTANCE, DEFAULT_COALESCE_MAX_SIZE); + } + + FileChannelReadable(String name, Path path, int coalesceDistance, int coalesceMaxSize) throws IOException { + this.coalesceDistance = coalesceDistance; + this.coalesceMaxSize = coalesceMaxSize; this.name = name; this.channel = FileChannel.open(path, StandardOpenOption.READ); this.length = channel.size(); } + @Override + public int coalesceDistance() { + distanceCalls++; + return coalesceDistance; + } + + @Override + public int coalesceMaxSize() { + maxSizeCalls++; + return coalesceMaxSize; + } + @Override public String name() { return name; @@ -388,6 +412,52 @@ public void testLargeRoundTrip() throws IOException { } } + @ParameterizedTest + @CsvSource({"false, 1048576, 33554432", "true, 1048576, 33554432", "false, 0, 1", "true, 0, 1"}) + public void testCoalescingSettingsCapturedOnce(boolean metadata, int distance, int maxSize) throws IOException { + Path path = tempDir.resolve("coalescing.vortex"); + Session session = Session.create(); + writePeopleFile(session, ArrowAllocation.rootAllocator(), path); + + try (FileChannelReadable readable = new FileChannelReadable(path.toString(), path, distance, maxSize)) { + registerReadable(metadata, session, readable); + assertEquals(1, readable.distanceCalls); + assertEquals(1, readable.maxSizeCalls); + } + } + + @ParameterizedTest + @CsvSource({ + "false, -1, 1024, coalesceDistance", + "true, -1, 1024, coalesceDistance", + "false, 0, 0, coalesceMaxSize", + "true, 0, 0, coalesceMaxSize", + "false, 0, -1, coalesceMaxSize", + "true, 0, -1, coalesceMaxSize" + }) + public void testInvalidCoalescingSettingsRejected(boolean metadata, int distance, int maxSize, String setting) + throws IOException { + Path path = tempDir.resolve("invalid_coalescing.vortex"); + Files.write(path, new byte[] {0}); + Session session = Session.create(); + try (FileChannelReadable readable = new FileChannelReadable(path.toString(), path, distance, maxSize)) { + RuntimeException thrown = + assertThrows(RuntimeException.class, () -> registerReadable(metadata, session, readable)); + assertTrue(thrown.getMessage().contains(setting), thrown.getMessage()); + } + } + + /** Registers {@code readable} through the file-metadata or data-source JNI entry point. */ + private static void registerReadable(boolean metadata, Session session, NativeReadable readable) { + if (metadata) { + assertEquals(2, NativeFiles.readStatistics(session, readable).size()); + } else { + assertEquals( + 3L, + DataSource.open(session, readable).rowCount().asOptional().orElseThrow()); + } + } + @Test public void testReadableExceptionPropagates() { Session session = Session.create(); diff --git a/java/vortex-spark/README.md b/java/vortex-spark/README.md index 110988512b3..418aebc80bc 100644 --- a/java/vortex-spark/README.md +++ b/java/vortex-spark/README.md @@ -65,6 +65,31 @@ Catalog tables read and write through the V1 file format; `spark.sql.sources.useV1SourceList=vortex` puts every path on it, which is what the `_metadata` column and dynamic partition pruning need. +## Read coalescing + +Nearby storage reads can be merged to reduce requests. Set the maximum gap and merged span in bytes: + +```java +Dataset result = spark.read() + .format("vortex") + .option("vortex.read.coalesce.distance", 262144) + .option("vortex.read.coalesce.maxSize", 16777216) + .load("s3a://bucket/path"); +``` + +The defaults remain a 1 MiB gap and a 16 MiB maximum merged span. The gap must be non-negative +(zero merges only adjacent or overlapping reads); the maximum span must be positive. Both accept +integers up to 2147483647 bytes. The native reader caps maximums near that limit to leave room +for alignment padding (up to 64 KiB minus one byte). Individual reads can exceed the maximum +merged span, and merged reads may include leading alignment padding. +A wider gap trades fewer storage requests for more bytes read, including segments already held +in the file's segment cache. + +Options are case-insensitive and work with both V1 and V2 reads. + +Java callers supplying their own `NativeReadable` can override `coalesceDistance()` and +`coalesceMaxSize()`. Native code captures these limits once when registering each source. + ## Benchmarks JMH benchmarks live in `common/src/jmh/java`. They measure the per-file open cost of the diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java index 09f4831f69b..5f071004448 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java @@ -22,6 +22,8 @@ public final class HadoopReadable extends PooledReadable { private final Path path; private final FileSystem fs; + private final int coalesceDistance; + private final int coalesceMaxSize; /** Opens a readable over {@code path}, stating the file for its length. */ public static HadoopReadable open(Configuration conf, String path) { @@ -33,20 +35,42 @@ public static HadoopReadable open(Configuration conf, String path) { * stats every file when it plans the scan, so executors need not stat them again. */ public static HadoopReadable open(Configuration conf, String path, long length) { + return open(conf, path, length, DEFAULT_COALESCE_DISTANCE, DEFAULT_COALESCE_MAX_SIZE); + } + + /** + * Opens a readable with the given coalescing gap and maximum merged size, both in bytes. Native code rejects a + * negative gap or non-positive size when the source is registered. + */ + public static HadoopReadable open( + Configuration conf, String path, long length, int coalesceDistance, int coalesceMaxSize) { Path file = new Path(path); try { FileSystem fs = file.getFileSystem(conf); long resolved = length >= 0 ? length : fs.getFileStatus(file).getLen(); - return new HadoopReadable(path, file, fs, resolved); + return new HadoopReadable(path, file, fs, resolved, coalesceDistance, coalesceMaxSize); } catch (IOException e) { throw new UncheckedIOException("Failed to open " + path, e); } } - private HadoopReadable(String name, Path path, FileSystem fs, long length) { + private HadoopReadable( + String name, Path path, FileSystem fs, long length, int coalesceDistance, int coalesceMaxSize) { super(name, length); this.path = path; this.fs = fs; + this.coalesceDistance = coalesceDistance; + this.coalesceMaxSize = coalesceMaxSize; + } + + @Override + public int coalesceDistance() { + return coalesceDistance; + } + + @Override + public int coalesceMaxSize() { + return coalesceMaxSize; } @Override 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 d59f20392ac..2f7680f7922 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 @@ -24,7 +24,7 @@ *

Built on the driver and shipped to executors, so the configuration travels with it, or is broadcast to them. */ public final class VortexIo implements Serializable { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 2L; /** * Option bounding how many concurrent read upcalls the native reader issues against one file. Each in-flight upcall @@ -34,6 +34,16 @@ public final class VortexIo implements Serializable { private static final int DEFAULT_READ_CONCURRENCY = 0; + /** + * Maximum gap between merged reads, in bytes. Defaults to 1 MiB; zero merges only adjacent or overlapping reads. + */ + public static final String COALESCE_DISTANCE_OPTION = "vortex.read.coalesce.distance"; + + /** + * Maximum span of a merged read, in bytes. Must be positive; defaults to 16 MiB. Individual reads may be larger. + */ + public static final String COALESCE_MAX_SIZE_OPTION = "vortex.read.coalesce.maxSize"; + /** * 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 @@ -51,18 +61,25 @@ public final class VortexIo implements Serializable { /** The configuration broadcast to every executor once, or null when this instance carries {@link #conf}. */ private final Broadcast broadcastConf; - private final int readConcurrency; - private final int splitConcurrency; + private final ReadSettings settings; + + private record ReadSettings(int readConcurrency, int splitConcurrency, int coalesceDistance, int coalesceMaxSize) + implements Serializable { + static ReadSettings parse(VortexOptions options) { + Objects.requireNonNull(options, "options"); + return new ReadSettings( + nonNegative(options, READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY), + nonNegative(options, SPLIT_CONCURRENCY_OPTION, DEFAULT_SPLIT_CONCURRENCY), + nonNegative(options, COALESCE_DISTANCE_OPTION, NativeReadable.DEFAULT_COALESCE_DISTANCE), + positive(options, COALESCE_MAX_SIZE_OPTION, NativeReadable.DEFAULT_COALESCE_MAX_SIZE)); + } + } /** Captures {@code hadoopConf} and the read settings found in the format options. */ public static VortexIo create(VortexOptions options, Configuration hadoopConf) { - Objects.requireNonNull(options, "options"); + ReadSettings settings = ReadSettings.parse(options); Objects.requireNonNull(hadoopConf, "hadoopConf"); - return new VortexIo( - new SerializableHadoopConf(hadoopConf), - null, - nonNegative(options, READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY), - nonNegative(options, SPLIT_CONCURRENCY_OPTION, DEFAULT_SPLIT_CONCURRENCY)); + return new VortexIo(new SerializableHadoopConf(hadoopConf), null, settings); } /** @@ -70,34 +87,22 @@ public static VortexIo create(VortexOptions options, Configuration hadoopConf) { * result then shares the executor's copy, instead of reading the whole configuration back from its own bytes. */ public static VortexIo broadcast(VortexOptions options, Configuration hadoopConf, SparkContext context) { - Objects.requireNonNull(options, "options"); + ReadSettings settings = ReadSettings.parse(options); Objects.requireNonNull(hadoopConf, "hadoopConf"); ClassTag tag = ClassTag$.MODULE$.apply(SerializableHadoopConf.class); - return new VortexIo( - null, - context.broadcast(new SerializableHadoopConf(hadoopConf), tag), - nonNegative(options, READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY), - nonNegative(options, SPLIT_CONCURRENCY_OPTION, DEFAULT_SPLIT_CONCURRENCY)); + return new VortexIo(null, context.broadcast(new SerializableHadoopConf(hadoopConf), tag), settings); } /** 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, - DEFAULT_SPLIT_CONCURRENCY); + return create(VortexOptions.empty(), new Configuration()); } private VortexIo( - SerializableHadoopConf conf, - Broadcast broadcastConf, - int readConcurrency, - int splitConcurrency) { + SerializableHadoopConf conf, Broadcast broadcastConf, ReadSettings settings) { this.conf = conf; this.broadcastConf = broadcastConf; - this.readConcurrency = readConcurrency; - this.splitConcurrency = splitConcurrency; + this.settings = settings; } private static int nonNegative(VortexOptions options, String key, int defaultValue) { @@ -108,18 +113,26 @@ private static int nonNegative(VortexOptions options, String key, int defaultVal return parsed; } + private static int positive(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", key, parsed)); + } + return parsed; + } + public Configuration hadoopConf() { return conf != null ? conf.value() : broadcastConf.value().value(); } /** Bound on concurrent read upcalls per file; zero keeps the native default. */ public int readConcurrency() { - return readConcurrency; + return settings.readConcurrency(); } /** The most splits of a file the native scan reads at once; zero keeps the native default. */ public int splitConcurrency() { - return splitConcurrency; + return settings.splitConcurrency(); } /** @@ -133,6 +146,7 @@ public int splitConcurrency() { */ public NativeReadable openReadable(VortexFile file) { VortexFile.requireVortexExtension(file.path()); - return HadoopReadable.open(hadoopConf(), file.path(), file.length()); + return HadoopReadable.open( + hadoopConf(), file.path(), file.length(), settings.coalesceDistance(), settings.coalesceMaxSize()); } } diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java index 588b3cbd83d..e44ea9506fe 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java @@ -52,6 +52,8 @@ void readsEveryRangeOfTheFile() throws IOException { Path file = write("data.vortex", content); try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) { + assertEquals(NativeReadable.DEFAULT_COALESCE_DISTANCE, readable.coalesceDistance()); + assertEquals(NativeReadable.DEFAULT_COALESCE_MAX_SIZE, readable.coalesceMaxSize()); assertEquals(SIZE, readable.length()); assertEquals(file.toString(), readable.name()); 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 019f811fca0..b7c8c2a15ac 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 @@ -6,17 +6,73 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import dev.vortex.io.NativeReadable; import dev.vortex.spark.VortexOptions; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Map; import org.apache.hadoop.conf.Configuration; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; final class VortexIoTest { + @TempDir + Path tempDir; + + @Test + void coalescingDefaultsArePreserved() throws IOException { + try (NativeReadable readable = VortexIo.defaults().openReadable(file())) { + assertEquals(NativeReadable.DEFAULT_COALESCE_DISTANCE, readable.coalesceDistance()); + assertEquals(NativeReadable.DEFAULT_COALESCE_MAX_SIZE, readable.coalesceMaxSize()); + } + } + + @ParameterizedTest + @CsvSource({ + "vortex.read.coalesce.distance, -1", + "vortex.read.coalesce.distance, many", + "vortex.read.coalesce.distance, 2147483648", + "vortex.read.coalesce.maxSize, 0", + "vortex.read.coalesce.maxSize, -1", + "vortex.read.coalesce.maxSize, many", + "vortex.read.coalesce.maxSize, 2147483648" + }) + void invalidCoalescingOptionsAreRejected(String key, String value) { + assertThrows(IllegalArgumentException.class, () -> VortexIo.create(options(key, value), new Configuration())); + // A null context proves validation happens before a broadcast is attempted. + assertThrows( + IllegalArgumentException.class, + () -> VortexIo.broadcast(options(key, value), new Configuration(), null)); + } + + @ParameterizedTest + @CsvSource({"0, 1", "1048576, 33554432", "2147483647, 2147483647"}) + void coalescingOptionsReachReadableAfterSerialization(int distance, int maxSize) + throws IOException, ClassNotFoundException { + VortexIo io = VortexIo.create( + VortexOptions.of(Map.of( + "VORTEX.READ.COALESCE.DISTANCE", Integer.toString(distance), + "vortex.read.coalesce.maxsize", Integer.toString(maxSize))), + new Configuration()); + try (NativeReadable readable = roundTrip(io).openReadable(file())) { + assertEquals(distance, readable.coalesceDistance()); + assertEquals(maxSize, readable.coalesceMaxSize()); + } + } + + private VortexFile file() throws IOException { + Path path = tempDir.resolve("data.vortex"); + Files.write(path, new byte[] {1}); + return new VortexFile(path.toString(), 1); + } + @Test void concurrenciesDefaultToNativeDefaults() { VortexIo io = VortexIo.create(VortexOptions.empty(), new Configuration()); diff --git a/vortex-jni/src/data_source.rs b/vortex-jni/src/data_source.rs index 7e50b0d934b..bf52f76e72c 100644 --- a/vortex-jni/src/data_source.rs +++ b/vortex-jni/src/data_source.rs @@ -36,6 +36,7 @@ use crate::RUNTIME; use crate::errors::try_or_throw; use crate::file::extract_properties; use crate::io::JavaFileSystem; +use crate::io::read_coalesce_config; use crate::object_store::object_store_fs; use crate::session::session_ref; @@ -155,12 +156,13 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_openFiles( if readable.is_null() { throw_runtime!("null readable for path '{path}'"); } + let coalesce_config = read_coalesce_config(env, &readable)?; let readable = Arc::new(env.new_global_ref(&readable)?); // `MultiFileDataSource::with_glob` strips leading slashes (object-store paths // are bucket-relative), so key the registry by the same normalized form. let key = path.trim_start_matches('/').to_string(); - fs.insert(key.clone(), readable, size)?; + fs.insert(key.clone(), readable, size, coalesce_config)?; ordered_paths.push(key); } diff --git a/vortex-jni/src/file.rs b/vortex-jni/src/file.rs index 9799b7d9139..136c8372d21 100644 --- a/vortex-jni/src/file.rs +++ b/vortex-jni/src/file.rs @@ -31,6 +31,7 @@ use vortex::utils::aliases::hash_map::HashMap; use crate::RUNTIME; use crate::errors::try_or_throw; use crate::io::java_readable; +use crate::io::read_coalesce_config; use crate::object_store::object_store_fs; use crate::session::session_ref; use crate::statistics::ColumnStatistics; @@ -193,8 +194,9 @@ fn open_readable_file( u64::try_from(length).map_err(|_| vortex_err!("negative readable length: {length}"))?; let vm = env.get_java_vm()?; + let coalesce_config = read_coalesce_config(env, readable)?; let readable = Arc::new(env.new_global_ref(readable)?); - let source = java_readable(vm, readable, length, session.handle()); + let source = java_readable(vm, readable, length, coalesce_config, session.handle()); // `MultiFileDataSource` strips leading slashes and the scan path keys Java readables by // that same normalized form, so normalize here too — otherwise a read through this function diff --git a/vortex-jni/src/io/mod.rs b/vortex-jni/src/io/mod.rs index 34ab6ab9623..e6f6c84482b 100644 --- a/vortex-jni/src/io/mod.rs +++ b/vortex-jni/src/io/mod.rs @@ -36,6 +36,7 @@ use jni::Env; use jni::JavaVM; pub(crate) use read_at::JavaFileSystem; pub(crate) use read_at::java_readable; +pub(crate) use read_at::read_coalesce_config; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -125,6 +126,8 @@ pub(crate) mod tests { use jni::refs::Global; use vortex::error::VortexResult; use vortex::error::vortex_bail; + use vortex::error::vortex_err; + use vortex::io::CoalesceConfig; use vortex::io::filesystem::FileSystem; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; @@ -211,10 +214,20 @@ pub(crate) mod tests { let runtime = CurrentThreadRuntime::new(); let mut fs = JavaFileSystem::new(vm.clone(), runtime.handle(), None); - fs.insert("data/file.vortex".to_string(), Arc::new(readable), 4)?; + fs.insert( + "data/file.vortex".to_string(), + Arc::new(readable), + 4, + CoalesceConfig::new(0, 1024), + )?; // `open_read` clones the global ref into a `JavaReadable`; both it and the // file system must let go of the Java object once dropped. let reader = runtime.block_on(fs.open_read("data/file.vortex"))?; + let config = reader + .coalesce_config() + .ok_or_else(|| vortex_err!("Java readable has no coalescing config"))?; + assert_eq!(config.distance, 0); + assert_eq!(config.max_size, 1024); drop(fs); drop(reader); diff --git a/vortex-jni/src/io/read_at.rs b/vortex-jni/src/io/read_at.rs index 2c45fc347cd..de6f9aaf29d 100644 --- a/vortex-jni/src/io/read_at.rs +++ b/vortex-jni/src/io/read_at.rs @@ -12,6 +12,7 @@ use futures::StreamExt; use futures::future::BoxFuture; use futures::stream; use futures::stream::BoxStream; +use jni::Env; use jni::JValue; use jni::JavaVM; use jni::objects::JObject; @@ -30,8 +31,47 @@ use vortex::io::runtime::Handle; use vortex::utils::aliases::hash_map::EntryRef; use vortex::utils::aliases::hash_map::HashMap; +use crate::errors::JNIError; use crate::io::with_jvm; +/// Capture and validate a Java readable's coalescing limits before registering it. +pub(crate) fn read_coalesce_config( + env: &mut Env, + readable: &JObject, +) -> Result { + let distance = env + .call_method( + readable, + jni::jni_str!("coalesceDistance"), + jni::jni_sig!("()I"), + &[], + )? + .i()?; + let max_size = env + .call_method( + readable, + jni::jni_str!("coalesceMaxSize"), + jni::jni_sig!("()I"), + &[], + )? + .i()?; + Ok(checked_coalesce_config(distance, max_size)?) +} + +fn checked_coalesce_config(distance: i32, max_size: i32) -> VortexResult { + if distance < 0 { + vortex_bail!("coalesceDistance must be >= 0, got {distance}"); + } + if max_size <= 0 { + vortex_bail!("coalesceMaxSize must be > 0, got {max_size}"); + } + // FileSegmentSource adds leading alignment padding to the coalescing window. Reserve + // headroom for every alignment accepted from a file so merged reads still fit a ByteBuffer. + let max_padding = Alignment::MAX_UNTRUSTED.as_usize() as u64 - 1; + let max_size = (max_size as u64).min(i32::MAX as u64 - max_padding); + Ok(CoalesceConfig::new(distance as u64, max_size)) +} + /// Default number of concurrent `readFully` upcalls to allow across all files of one /// [`JavaFileSystem`]. Matches the object-store default since the backing storage is /// typically remote. @@ -67,6 +107,7 @@ pub(crate) struct JavaReadable { vm: JavaVM, readable: Arc>>, len: u64, + coalesce_config: CoalesceConfig, handle: Handle, limiter: UpcallLimiter, } @@ -76,6 +117,7 @@ impl JavaReadable { vm: JavaVM, readable: Arc>>, len: u64, + coalesce_config: CoalesceConfig, handle: Handle, limiter: UpcallLimiter, ) -> Self { @@ -83,6 +125,7 @@ impl JavaReadable { vm, readable, len, + coalesce_config, handle, limiter, } @@ -91,9 +134,7 @@ impl JavaReadable { impl VortexReadAt for JavaReadable { fn coalesce_config(&self) -> Option { - // Upcalls have a fixed JNI overhead and the backing storage is usually - // remote, so favor fewer, larger reads. - Some(CoalesceConfig::object_storage()) + Some(self.coalesce_config) } fn concurrency(&self) -> usize { @@ -204,12 +245,14 @@ pub(crate) fn java_readable( vm: JavaVM, readable: Arc>>, len: u64, + coalesce_config: CoalesceConfig, handle: Handle, ) -> Arc { Arc::new(JavaReadable::new( vm, readable, len, + coalesce_config, handle, UpcallLimiter::new(DEFAULT_CONCURRENCY), )) @@ -218,6 +261,7 @@ pub(crate) fn java_readable( struct JavaFileEntry { readable: Arc>>, size: u64, + coalesce_config: CoalesceConfig, } /// A [`FileSystem`] over a fixed set of Java-provided readables, keyed by path. @@ -249,12 +293,17 @@ impl JavaFileSystem { path: String, readable: Arc>>, size: u64, + coalesce_config: CoalesceConfig, ) -> VortexResult<()> { match self.files.entry_ref(&path) { EntryRef::Occupied(_) => { vortex_bail!("multiple Java readables normalize to path '{path}'"); } - EntryRef::Vacant(v) => v.insert(JavaFileEntry { readable, size }), + EntryRef::Vacant(v) => v.insert(JavaFileEntry { + readable, + size, + coalesce_config, + }), }; Ok(()) @@ -302,6 +351,7 @@ impl FileSystem for JavaFileSystem { self.vm.clone(), Arc::clone(&entry.readable), entry.size, + entry.coalesce_config, self.handle.clone(), self.limiter.clone(), ))) @@ -311,3 +361,40 @@ impl FileSystem for JavaFileSystem { vortex_bail!("delete('{path}') is not supported by a Java-readable file system") } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::error::VortexResult; + + use super::checked_coalesce_config; + + #[rstest] + #[case(0, 1, 1)] + #[case(64 * 1024, 16 * 1024 * 1024, 16 * 1024 * 1024)] + #[case(i32::MAX, i32::MAX, 2_147_418_112)] + fn test_coalescing_limits( + #[case] distance: i32, + #[case] max_size: i32, + #[case] expected_max_size: u64, + ) -> VortexResult<()> { + let config = checked_coalesce_config(distance, max_size)?; + assert_eq!(config.distance, distance as u64); + assert_eq!(config.max_size, expected_max_size); + Ok(()) + } + + #[rstest] + #[case(-1, 1024, "coalesceDistance")] + #[case(0, 0, "coalesceMaxSize")] + #[case(0, -1, "coalesceMaxSize")] + fn test_invalid_coalescing_limits( + #[case] distance: i32, + #[case] max_size: i32, + #[case] setting: &str, + ) { + let error = + checked_coalesce_config(distance, max_size).expect_err("invalid coalescing limits"); + assert!(error.to_string().contains(setting)); + } +} From b367941e9d490b4a19ef9bdd6653f63d712c581f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 10 Nov 2026 13:18:52 -0500 Subject: [PATCH 10/13] Fix run-end slice export and generic decoding (#11375) ## Summary Run-end decoding currently fails for logical types without a specialized decoder. Arrow export also narrows absolute ends before applying a slice offset, so a small slice can fail even when its relative ends fit the requested integer width. A retained run ending at the offset produces a zero end that Arrow rejects. ## Changes Use a dictionary as the generic decoding fallback so existing take implementations preserve the logical dtype. Piecewise codes avoid building an expanded index buffer during conversion, and empty arrays retain no source children. Specialized decoders remain in place. Normalize Arrow run ends before narrowing, trim the values child to the same runs, and retain the values field's extension metadata so the exported array agrees with its schema. Unadjusted arrays retain the existing buffer-sharing path. Regression coverage includes clipped and boundary slices, large absolute ends exported as Int16, empty slices, repeated slicing, logical dtype preservation, and UUID values-field metadata. DuckDB's batch exporter remains separate. ## API Changes Adds `RunEndArray::to_dict` and renames the public `runend::trimmed_ends_iter` helper to `runend::slice_relative_ends`.

Validation - Formatting with the pinned `nightly-2026-09-10`. - Scoped Clippy for vortex-arrow and vortex-uuid with all targets, all features, and warnings denied. - Commit hook: workspace Rust and TOML formatting. - Pre-push hook: DCO checks and workspace Clippy with all targets, all features, and warnings denied. - Runtime tests and doctests were not rerun.
Signed-off-by: Connor Tsui --- encodings/uuid/src/arrow.rs | 21 +- vortex-array/src/arrays/runend/array/mod.rs | 9 +- .../src/arrays/runend/array/to_dict.rs | 281 ++++++++++++++++++ vortex-array/src/arrays/runend/compress.rs | 8 +- .../src/arrays/runend/compute/take.rs | 4 +- .../src/arrays/runend/decompress_bool.rs | 6 +- vortex-array/src/arrays/runend/iter.rs | 10 +- vortex-array/src/arrays/runend/mod.rs | 2 +- vortex-arrow/src/executor/run_end.rs | 281 ++++++++++++------ 9 files changed, 497 insertions(+), 125 deletions(-) create mode 100644 vortex-array/src/arrays/runend/array/to_dict.rs diff --git a/encodings/uuid/src/arrow.rs b/encodings/uuid/src/arrow.rs index d69c35a7928..1e3f67bf3fe 100644 --- a/encodings/uuid/src/arrow.rs +++ b/encodings/uuid/src/arrow.rs @@ -233,7 +233,6 @@ mod tests { use vortex_array::iter::ArrayIterator; use vortex_arrow::ArrowArrayStreamAdapter; use vortex_arrow::ArrowSessionExt; - use vortex_error::vortex_err; use vortex_session::VortexSession; use super::*; @@ -595,7 +594,11 @@ mod tests { &values, )?); - let expected = run_array.to_data(); + let expected = run_array + .to_data() + .into_builder() + .data_type(field.data_type().clone()) + .build()?; let vortex_array = session.from_arrow_array(run_array, &field)?; assert_eq!(vortex_array.len(), 5); // The array conversion must agree with the dtype conversion. @@ -603,20 +606,8 @@ mod tests { // And the values must round-trip back out through the export plugin. let exported = vortex_array.execute_arrow(Some(&field), &mut ctx)?; - assert_eq!(exported.len(), 5); - let ree = exported - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - vortex_err!( - "expected an Int32 run-end array, got {}", - exported.data_type() - ) - })?; assert_eq!(exported.to_data(), expected); - let values = ree.values().as_fixed_size_binary(); - assert_eq!(values.value(0), b"0123456789abcdef"); - assert_eq!(values.value(1), b"fedcba9876543210"); + Ok(()) } diff --git a/vortex-array/src/arrays/runend/array/mod.rs b/vortex-array/src/arrays/runend/array/mod.rs index ad3a293c258..000963a85da 100644 --- a/vortex-array/src/arrays/runend/array/mod.rs +++ b/vortex-array/src/arrays/runend/array/mod.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Run-end arrays and their execution into logical values. +//! +//! Each end is an exclusive logical index paired with the value for its preceding run. Decoding +//! **must** apply a slice's logical offset before expanding those values. Specialized decoders +//! handle common storage types. The generic fallback gathers values to retain their logical dtype. + use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -52,6 +58,7 @@ use crate::vtable::VTable; use crate::vtable::ValidityVTable; mod plugin; +mod to_dict; /// A [`RunEnd`]-encoded Vortex array. pub type RunEndArray = Array; @@ -482,7 +489,7 @@ pub(super) fn run_end_canonicalize( .execute_as::("values", ctx)?; runend_decode_listview(pends, values, array.offset(), array.len())?.into_array() } - _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()), + _ => array.to_dict(ctx)?.into_array(), }) } diff --git a/vortex-array/src/arrays/runend/array/to_dict.rs b/vortex-array/src/arrays/runend/array/to_dict.rs new file mode 100644 index 00000000000..511cb479bbd --- /dev/null +++ b/vortex-array/src/arrays/runend/array/to_dict.rs @@ -0,0 +1,281 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Express run-end decoding as a gather from the existing values child. +//! +//! Repeated run indices let the generic take machinery handle logical types without specialized +//! run-end decoders. Piecewise codes keep that mapping compact until a consumer needs flat indices. + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::DictArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::runend::RunEndArray; +use crate::arrays::runend::RunEndArrayExt; +use crate::arrays::runend::RunEndArraySlotsExt; +use crate::arrays::runend::slice_relative_ends; +use crate::dtype::Nullability; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; + +impl RunEndArray { + /// Converts to a dictionary with the same logical dtype and slice. + /// + /// Non-empty arrays share the values child without executing it. Codes use one constant piece + /// per overlapping run, so construction needs no expanded index buffer. Executing the resulting + /// [`DictArray`] may still materialize those indices. Empty arrays retain no source children. + pub fn to_dict(&self, ctx: &mut ExecutionCtx) -> VortexResult { + if self.is_empty() { + return DictArray::try_new( + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + Canonical::empty(self.dtype()).into_array(), + ); + } + + let ends = self.ends().clone().execute::(ctx)?; + let codes = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + run_indices(ends.as_slice::(), self.offset(), self.len())? + }); + + DictArray::try_new(codes.into_array(), self.values().clone()) + } +} + +/// Builds repeated physical value indices for a slice covered by validated run ends. +fn run_indices( + ends: &[E], + offset: usize, + len: usize, +) -> VortexResult { + let mut starts = BufferMut::::with_capacity(ends.len()); + let mut lengths = BufferMut::::with_capacity(ends.len()); + let mut run_start = 0; + + for (index, run_end) in slice_relative_ends(ends, offset, len).enumerate() { + let length = run_end - run_start; + if length != 0 { + starts.push(u64::try_from(index)?); + lengths.push(u64::try_from(length)?); + } + run_start = run_end; + + if run_end == len { + break; + } + } + + let multipliers = ConstantArray::new(0u64, starts.len()).into_array(); + + PiecewiseSequenceArray::try_new( + starts.freeze().into_array(), + lengths.freeze().into_array(), + multipliers, + len, + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::ExecutionCtx; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ExtensionArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::PiecewiseSequence; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; + use crate::arrays::dict::DictArraySlotsExt; + use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; + use crate::arrays::runend::RunEnd; + use crate::arrays::runend::RunEndArray; + use crate::arrays::runend::RunEndArraySlotsExt; + use crate::arrays::runend::tests::SESSION; + use crate::assert_arrays_eq; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + + #[track_caller] + fn check_conversion( + array: &RunEndArray, + expected: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let dict = array.to_dict(ctx)?; + + if array.is_empty() { + assert!(dict.codes().is_empty()); + assert!(dict.values().is_empty()); + } else { + assert!(ArrayRef::ptr_eq(dict.values(), array.values())); + let codes = dict.codes().as_::(); + assert_eq!(codes.multipliers().as_constant(), Some(&Scalar::from(0u64))); + } + + assert_arrays_eq!(dict.into_array(), expected, ctx); + + Ok(()) + } + + fn struct_values(storage: ArrayRef) -> VortexResult { + Ok(StructArray::from_fields(&[("value", storage)])?.into_array()) + } + + fn fixed_size_list_values(storage: ArrayRef) -> VortexResult { + Ok( + FixedSizeListArray::try_new(storage.clone(), 1, storage.validity()?, storage.len())? + .into_array(), + ) + } + + fn timestamp_values(storage: ArrayRef) -> VortexResult { + let ext_dtype = Timestamp::try_new_with_tz( + TimeUnit::Microseconds, + Some("UTC".into()), + storage.dtype().nullability(), + )?; + + Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array()) + } + + #[rstest] + #[case::full(PType::U32, 0, 8)] + #[case::clipped_u8(PType::U8, 2, 5)] + #[case::clipped_u16(PType::U16, 2, 5)] + #[case::clipped_u32(PType::U32, 2, 5)] + #[case::clipped_u64(PType::U64, 2, 5)] + #[case::run_boundary(PType::U32, 3, 2)] + #[case::single_row(PType::U32, 1, 1)] + #[case::empty_retained_runs(PType::U32, 8, 0)] + fn sliced_conversion( + #[case] ends_type: PType, + #[case] offset: usize, + #[case] len: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ends = buffer![3u64, 5, 8] + .into_array() + .cast(DType::Primitive(ends_type, Nullability::NonNullable))?; + let values = PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(); + let array = RunEnd::try_new_offset_length(ends, values, offset, len, &mut ctx)?; + let expanded = [ + Some(10i64), + Some(10), + Some(10), // First run. + None, + None, // Null run. + Some(30), + Some(30), + Some(30), // Last run. + ]; + let expected = + PrimitiveArray::from_option_iter(expanded[offset..offset + len].iter().copied()); + + check_conversion(&array, expected.into_array(), &mut ctx) + } + + #[test] + fn codes_retain_only_overlapping_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + buffer![10i64, 20, 30].into_array(), + 3, + 4, + &mut ctx, + )?; + let dict = array.to_dict(&mut ctx)?; + let codes = dict.codes().as_::(); + + assert_arrays_eq!(codes.starts(), buffer![1u64, 2].into_array(), &mut ctx); + assert_arrays_eq!(codes.lengths(), buffer![2u64, 2].into_array(), &mut ctx); + + assert_arrays_eq!( + dict.into_array(), + buffer![20i64, 20, 30, 30].into_array(), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn empty_conversion() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::empty::(Nullability::Nullable).into_array(); + let array = RunEnd::try_new( + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + values.clone(), + &mut ctx, + )?; + + check_conversion(&array, values, &mut ctx) + } + + #[test] + fn repeated_slicing() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + buffer![3u32, 5, 8].into_array(), + PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(), + &mut ctx, + )?; + let sliced = array.slice(2..8)?.execute_until::(&mut ctx)?; + let sliced = sliced + .slice(1..5)? + .execute_until::(&mut ctx)? + .downcast::(); + let expected = PrimitiveArray::from_option_iter([None, None, Some(30i64), Some(30)]); + + check_conversion(&sliced, expected.into_array(), &mut ctx) + } + + #[rstest] + #[case::structure(struct_values)] + #[case::fixed_size_list(fixed_size_list_values)] + #[case::extension(timestamp_values)] + fn fallback_preserves_dtype( + #[case] wrap: fn(ArrayRef) -> VortexResult, + #[values(false, true)] empty: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let storage = PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(); + let values = wrap(storage)?; + let len = if empty { 0 } else { 5 }; + let expected = wrap( + PrimitiveArray::from_option_iter( + [Some(10i64), None, None, Some(30), Some(30)] + .into_iter() + .take(len), + ) + .into_array(), + )?; + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + values, + 2, + len, + &mut ctx, + )?; + + check_conversion(&array, expected.clone(), &mut ctx)?; + assert_arrays_eq!(array.into_array(), expected, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/runend/compress.rs b/vortex-array/src/arrays/runend/compress.rs index 2c929eccb7f..9aa319c26bf 100644 --- a/vortex-array/src/arrays/runend/compress.rs +++ b/vortex-array/src/arrays/runend/compress.rs @@ -25,7 +25,7 @@ use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::buffer::BufferHandle; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -204,7 +204,7 @@ pub fn runend_decode_primitive( Ok(match_each_native_ptype!(values.ptype(), |P| { match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_primitive( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), values.as_slice::

(), validity_mask, values.dtype().nullability(), @@ -381,7 +381,7 @@ pub fn runend_decode_decimal( Ok(match_each_decimal_value_type!(values.values_type(), |D| { let (decoded, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_slice( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), values.buffer::().as_slice(), validity_mask, values.dtype().nullability(), @@ -414,7 +414,7 @@ pub fn runend_decode_varbinview( let (decoded_views, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_slice( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), views, validity_mask, values.dtype().nullability(), diff --git a/vortex-array/src/arrays/runend/compute/take.rs b/vortex-array/src/arrays/runend/compute/take.rs index c7a568fba6f..cbe73efd869 100644 --- a/vortex-array/src/arrays/runend/compute/take.rs +++ b/vortex-array/src/arrays/runend/compute/take.rs @@ -20,7 +20,7 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::runend::RunEnd; use crate::arrays::runend::array::RunEndArrayExt; use crate::arrays::runend::array::RunEndArraySlotsExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::dtype::UnsignedPType; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; @@ -378,7 +378,7 @@ where { let mut table = BufferMut::with_capacity(len); let mut run_start = 0; - for (run_idx, run_end) in trimmed_ends_iter(ends, offset, len).enumerate() { + for (run_idx, run_end) in slice_relative_ends(ends, offset, len).enumerate() { table.push_n(run_idx.as_(), run_end - run_start); run_start = run_end; } diff --git a/vortex-array/src/arrays/runend/decompress_bool.rs b/vortex-array/src/arrays/runend/decompress_bool.rs index 739fb968276..3136631de2d 100644 --- a/vortex-array/src/arrays/runend/decompress_bool.rs +++ b/vortex-array/src/arrays/runend/decompress_bool.rs @@ -18,7 +18,7 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::dtype::DType; use crate::dtype::Nullability; use crate::match_each_unsigned_integer_ptype; @@ -74,7 +74,7 @@ pub fn runend_decode_bools( Ok(match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_bool( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), &values_buf, validity, nullability, @@ -112,7 +112,7 @@ pub fn runend_decode_typed_bool( } /// Fast path for few runs with no offset. Uses direct slice access to minimize overhead. -/// This avoids the `trimmed_ends_iter` iterator chain which adds significant overhead +/// This avoids the `slice_relative_ends` iterator chain which adds significant overhead /// for small numbers of runs. #[allow(clippy::inline_always)] #[inline(always)] diff --git a/vortex-array/src/arrays/runend/iter.rs b/vortex-array/src/arrays/runend/iter.rs index d41ee2da87c..925725959c7 100644 --- a/vortex-array/src/arrays/runend/iter.rs +++ b/vortex-array/src/arrays/runend/iter.rs @@ -7,7 +7,15 @@ use vortex_error::vortex_panic; use crate::dtype::IntegerPType; -pub fn trimmed_ends_iter( +/// Yields exclusive run ends relative to a logical slice. +/// +/// Subtracts `offset` from each end and clamps it to `length`. One end is yielded per input run, +/// including zero-length runs at the slice boundaries, so indices still match the values child. +/// +/// # Panics +/// +/// Panics if `offset` or `length` cannot be represented by `E`, or an input end is before `offset`. +pub fn slice_relative_ends( run_ends: &[E], offset: usize, length: usize, diff --git a/vortex-array/src/arrays/runend/mod.rs b/vortex-array/src/arrays/runend/mod.rs index 84914068f73..77a8f94335c 100644 --- a/vortex-array/src/arrays/runend/mod.rs +++ b/vortex-array/src/arrays/runend/mod.rs @@ -6,7 +6,7 @@ mod arbitrary; #[cfg(feature = "arbitrary")] pub use arbitrary::ArbitraryRunEndArray; pub use array::*; -pub use iter::trimmed_ends_iter; +pub use iter::slice_relative_ends; mod array; pub mod compress; diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index 04801861ea0..0b8e33a0760 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -1,11 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Export run-end arrays without expanding their logical values. +//! +//! Vortex slices retain absolute run ends. Adjusting before narrowing to Arrow's requested integer +//! width keeps small slices exportable even when their source indices are large. The values field +//! retains extension metadata so imports recover the same logical type. + use std::sync::Arc; +use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::RunArray; use arrow_array::cast::AsArray; +use arrow_array::make_array; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_buffer::ArrowNativeType; @@ -16,10 +24,12 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::RunEnd; use vortex_array::arrays::RunEndArray; use vortex_array::arrays::runend::RunEndArrayExt; use vortex_array::arrays::runend::RunEndArraySlotsExt; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -44,7 +54,35 @@ impl Matcher for ArrowRunEndExportable { pub(super) fn to_arrow_run_end( array: ArrayRef, ends_type: &DataType, - values_type: &Field, + values_field: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let data_type = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", ends_type.clone(), false)), + Arc::new(values_field.clone()), + ); + if array.is_empty() { + return Ok(new_null_array(&data_type, 0)); + } + + // Arrow's RunArray constructors synthesize a values field without extension metadata. + // Retain the caller's field so the array remains consistent with its export schema. + let exported = build_run_end_array(array, ends_type, values_field, options, ctx)?; + + Ok(make_array( + exported + .to_data() + .into_builder() + .data_type(data_type) + .build()?, + )) +} + +fn build_run_end_array( + array: ArrayRef, + ends_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -52,118 +90,108 @@ pub(super) fn to_arrow_run_end( let array = match array.try_downcast::() { Ok(constant) => { - return constant_to_run_end(constant, ends_type, values_type, options, ctx); + return constant_to_run_end(constant, ends_type, values_field, 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, options, ctx), + Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_field, options, ctx), Err(array) => array, }; // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = export_values(array, values_type, options, ctx)?; + let flat = array.execute_arrow_with_options(Some(values_field), options, ctx)?; let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), + Arc::new(values_field.clone()), ); arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) } -/// Export the values of a run-end array through the session, so extension metadata on the target -/// values field dispatches to its export plugin instead of the naive Arrow mapping. -fn export_values( - values: ArrayRef, - values_type: &Field, - options: &ArrowExportOptions, - ctx: &mut ExecutionCtx, -) -> VortexResult { - values.execute_arrow_with_options(Some(values_type), options, ctx) -} - fn run_end_to_arrow( array: RunEndArray, ends_type: &DataType, - values_type: &Field, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let length = array.len(); - let offset = array.offset(); + let ends = array.ends().clone().execute::(ctx)?; + let (runs, ends) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + let run_ends = ends.as_slice::(); + let length = E::try_from(array.len())?; + + if array.offset() == 0 + && run_ends.first().is_some_and(|&end| end > 0) + && run_ends.last() == Some(&length) + { + (0..run_ends.len(), ends.clone().into_array()) + } else { + let offset = E::try_from(array.offset())?; + + // A run ending at the offset contributes no rows. Retaining it would give Arrow a + // zero end, and retaining trailing runs would give duplicate ends after clamping. + let start = run_ends.partition_point(|&end| end <= offset); + let end = start + run_ends[start..].partition_point(|&end| end - offset < length); + let end = (end + 1).min(run_ends.len()); + let adjusted = PrimitiveArray::from_iter( + run_ends[start..end] + .iter() + .map(|&end| (end - offset).min(length)), + ) + .into_array(); - 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)?; + (start..end, adjusted) + } + }); + let values = if runs.start == 0 && runs.end == array.values().len() { + array.values().clone() + } else { + array.values().slice(runs)? + }; + + // Narrow only the relative ends: absolute indices can exceed the target's range even when + // this slice fits. Unadjusted ends and values keep the existing buffer-sharing export path. + let arrow_ends = execute_arrow_naive(ends, Some(ends_type), options, ctx)?; + let arrow_values = values.execute_arrow_with_options(Some(values_field), options, ctx)?; match ends_type { - DataType::Int16 => build_run_array::(&arrow_ends, &arrow_values, offset, length), - DataType::Int32 => build_run_array::(&arrow_ends, &arrow_values, offset, length), - DataType::Int64 => build_run_array::(&arrow_ends, &arrow_values, offset, length), + DataType::Int16 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), + DataType::Int32 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), + DataType::Int64 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), _ => vortex_bail!("Unsupported run-end index type: {:?}", ends_type), } } -fn build_run_array( - ends: &ArrowArrayRef, - values: &ArrowArrayRef, - offset: usize, - length: usize, -) -> VortexResult -where - R::Native: std::ops::Sub + Ord, -{ - let offset_native = R::Native::from_usize(offset) - .ok_or_else(|| vortex_err!("Offset {offset} exceeds run-end index capacity"))?; - let length_native = R::Native::from_usize(length) - .ok_or_else(|| vortex_err!("Length {length} exceeds run-end index capacity"))?; - - let ends_prim = ends.as_primitive::(); - if offset == 0 && ends_prim.values().last() == Some(&length_native) { - // Fast path: no trimming or adjustment needed. - return Ok(Arc::new(RunArray::::try_new(ends_prim, values)?) as ArrowArrayRef); - } - - // Trim to only include runs covering the [offset, offset+length) range. - // Runs beyond this would produce duplicate adjusted ends, violating - // Arrow's strict-ordering requirement for RunArray. - // Run ends are strictly increasing, so we can binary search. - let num_runs = (ends_prim - .values() - .partition_point(|&e| e - offset_native < length_native) - + 1) - .min(ends_prim.len()); - - let trimmed_ends = ends.slice(0, num_runs); - let trimmed_values = values.slice(0, num_runs); - - let adjusted = trimmed_ends - .as_primitive::() - .unary(|end| (end - offset_native).min(length_native)); - - Ok(Arc::new(RunArray::::try_new(&adjusted, &trimmed_values)?) as ArrowArrayRef) -} - /// Convert a constant array to a run-end encoded array with a single run. fn constant_to_run_end( array: ConstantArray, ends_type: &DataType, - values_type: &Field, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); let scalar = array.scalar(); - if scalar.is_null() || len == 0 { + if scalar.is_null() { let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), + Arc::new(values_field.clone()), ); return Ok(new_null_array(&ree_type, len)); } - let values = export_values( - ConstantArray::new(scalar.clone(), 1).into_array(), - values_type, + let values = ConstantArray::new(scalar.clone(), 1).execute_arrow_with_options( + Some(values_field), options, ctx, )?; @@ -201,13 +229,18 @@ mod tests { use arrow_schema::DataType; use arrow_schema::Field; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; use vortex_array::dtype::Nullability::Nullable; use vortex_array::dtype::PType; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::TimeUnit; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -215,10 +248,55 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::ArrowSessionExt; use crate::ExecuteArrowExt; use crate::executor::run_end::ConstantArray; use crate::executor::run_end::RunEnd; + #[track_caller] + fn check_temporal_run_end_round_trip(len: usize, sliced: bool) -> VortexResult<()> { + let session = VortexSession::empty(); + let arrow = session.arrow(); + let mut ctx = session.create_execution_ctx(); + let dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let values = ExtensionArray::try_new( + dtype, + PrimitiveArray::from_iter((0..i32::try_from(len)?).map(|i| 7 + i * 4)).into_array(), + )? + .into_array(); + let source = RunEnd::try_new( + PrimitiveArray::from_iter((1..=u32::try_from(len)?).map(|i| i * 2)).into_array(), + values.clone(), + &mut ctx, + )? + .into_array(); + let source = if sliced && !source.is_empty() { + source.slice(1..source.len())? + } else { + source + }; + + let start = u32::from(sliced && len > 0); + let expected = values.take( + PrimitiveArray::from_iter((start..u32::try_from(len * 2)?).map(|i| i / 2)).into_array(), + )?; + let values_field = arrow.to_arrow_field("dates", source.dtype())?; + let field = Field::new( + "dates", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(values_field), + ), + false, + ); + let exported = source.execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, expected, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { @@ -228,10 +306,7 @@ mod tests { ) } - fn execute( - array: vortex_array::ArrayRef, - dt: &DataType, - ) -> VortexResult { + fn execute(array: ArrayRef, dt: &DataType) -> VortexResult { let field = Field::new("", dt.clone(), array.dtype().is_nullable()); array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx()) } @@ -284,7 +359,7 @@ mod tests { ), )] fn constant_to_ree( - #[case] input: vortex_array::ArrayRef, + #[case] input: ArrayRef, #[case] target_type: DataType, #[case] expected: arrow_array::ArrayRef, ) -> VortexResult<()> { @@ -344,37 +419,47 @@ mod tests { Ok(()) } - /// Regression: build_run_array must trim excess trailing runs and - /// respect the `length` parameter. This happens when a vortex - /// RunEndArray is sliced to fewer rows than the physical run_ends cover. #[rstest] - #[case::offset_zero(0, 5, &[3, 5], &[100, 200])] - #[case::nonzero_offset(2, 3, &[1, 3], &[100, 200])] - #[case::all_runs_needed_but_last_exceeds(0, 8, &[3, 5, 8], &[100, 200, 300])] - fn build_run_array_trims_excess_runs( + #[case::offset_zero( &[3, 5, 10], 0, 5, &[3, 5], &[100, 200])] + #[case::nonzero_offset( &[3, 5, 10], 2, 3, &[1, 3], &[100, 200])] + #[case::clipped_last_run( &[3, 5, 10], 0, 8, &[3, 5, 8], &[100, 200, 300])] + #[case::run_boundary( &[3, 5, 10], 3, 4, &[2, 4], &[200, 300])] + #[case::large_absolute_ends(&[40000, 40002], 39999, 3, &[1, 3], &[100, 200])] + #[case::empty_retained_runs(&[3, 5, 10], 11, 0, &[], &[])] + fn run_end_export_trims_runs( + #[case] ends: &[u32], #[case] offset: usize, #[case] length: usize, - #[case] expected_ends: &[i32], + #[case] expected_ends: &[i16], #[case] expected_values: &[i64], ) -> VortexResult<()> { - // 3 runs covering 10 rows: [0..3), [3..5), [5..10) - let ends: arrow_array::ArrayRef = Arc::new(Int32Array::from(vec![3i32, 5, 10])); - let values: arrow_array::ArrayRef = Arc::new(Int64Array::from(vec![100i64, 200, 300])); + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + PrimitiveArray::from_iter(ends.iter().copied()).into_array(), + PrimitiveArray::from_iter((0..ends.len()).map(|i| (i as i64 + 1) * 100)).into_array(), + offset, + length, + &mut ctx, + )?; - let result = super::build_run_array::(&ends, &values, offset, length)?; + let result = execute( + array.into_array(), + &ree_type(DataType::Int16, DataType::Int64), + )?; assert_eq!(result.len(), length); + let expected = RunArray::::try_new( + &Int16Array::from(expected_ends.to_vec()), + &Int64Array::from(expected_values.to_vec()), + )?; + assert_eq!(result.as_ref(), &expected); - let ree = result - .as_any() - .downcast_ref::>() - .ok_or_else(|| vortex_err!("expected Int32 run-end array"))?; - assert_eq!(ree.run_ends().values(), expected_ends); - let values = ree - .values() - .as_any() - .downcast_ref::() - .ok_or_else(|| vortex_err!("expected Int64 values"))?; - assert_eq!(values.values(), expected_values); Ok(()) } + + #[rstest] + #[case::empty(0, false)] + #[case::sliced(3, true)] + fn temporal_run_end_round_trip(#[case] len: usize, #[case] sliced: bool) -> VortexResult<()> { + check_temporal_run_end_round_trip(len, sliced) + } } From 0ab0321c8319d3d953dd89b837b6091f799c95b6 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 10 Nov 2026 18:12:23 -0500 Subject: [PATCH 11/13] Preserve Arrow round-trip types with typed extension plugins (#11360) ## Summary Builds on #11375. Exporting a versioned Vortex UUID to Arrow and importing it back currently drops the declared version, even when every value is unchanged. Nested extensions can also lose metadata or fail to import when Arrow fields disagree with their storage. The full dtype, including nested metadata and nullability, must survive this round trip. The raw import and export traits leave metadata transport to each plugin. `ArrowExtVTable` puts that handling behind a typed registration adapter, so extension authors implement their native conversions while shared code preserves Vortex metadata. ## Changes The session checks that inferred fields reconstruct the complete source dtype and that successful plugin imports return the dtype inferred from their field. These checks apply to raw and typed plugins. Explicit export targets can still request conversions. UUID needs a separate Vortex payload because canonical `arrow.uuid` metadata must be empty. Keeping that representation lets other Arrow tools recognize UUID columns. A field without the Vortex payload imports with an unspecified version, so ordinary Arrow producers remain supported. Nested conversions need the full Arrow field because a `DataType` alone cannot carry extension identity. GeoArrow's required child names also belong to its layout contract. Metadata restoration uses inferred Vortex storage because it can differ from Arrow's physical storage. Regression coverage retains complete dtype and array comparisons while removing redundant parameter products and assertions. Variant remains excluded pending #11359, and equality derivation remains separate in #11115. ## API Changes Extension plugins implement `ArrowExtVTable`, declare `metadata_policy()`, and register through `ArrowSession::register_extension(vtable)`. The typed import callbacks are associated functions. Raw plugins remain available, but inferred fields and imported arrays must satisfy the same dtype checks.

Validation - Formatting with the pinned `nightly-2026-09-10`. - Commit hook: workspace Rust and TOML formatting. - Pre-push hook: DCO checks and workspace Clippy with all targets, all features, and warnings denied. - Runtime tests, doctests, and fuzz runs were not repeated for these refinements.
Signed-off-by: Connor Tsui --- encodings/uuid/src/arrow.rs | 261 ++++++++------- encodings/uuid/src/lib.rs | 5 +- vortex-arrow/src/executor/dictionary.rs | 87 ++++- vortex-arrow/src/executor/mod.rs | 25 +- vortex-arrow/src/executor/struct_.rs | 28 +- vortex-arrow/src/ext_metadata.rs | 252 ++++++++++++++- vortex-arrow/src/extension/adapter.rs | 206 ++++++++++++ vortex-arrow/src/extension/mod.rs | 142 ++++++++ vortex-arrow/src/extension/tests.rs | 302 ++++++++++++++++++ vortex-arrow/src/lib.rs | 5 + vortex-arrow/src/session.rs | 33 +- vortex-json/src/arrow.rs | 61 +--- vortex-json/src/lib.rs | 5 +- vortex-spatial/src/extension/linestring.rs | 75 ++--- vortex-spatial/src/extension/mod.rs | 83 +++++ .../src/extension/multilinestring.rs | 73 ++--- vortex-spatial/src/extension/multipoint.rs | 74 ++--- vortex-spatial/src/extension/multipolygon.rs | 73 ++--- vortex-spatial/src/extension/point.rs | 64 ++-- vortex-spatial/src/extension/polygon.rs | 75 ++--- vortex-spatial/src/extension/rect.rs | 63 ++-- vortex-spatial/src/extension/wkb.rs | 59 ++-- vortex-spatial/src/lib.rs | 46 ++- vortex-tensor/src/lib.rs | 5 +- .../src/types/fixed_shape_tensor/proto.rs | 15 +- .../src/types/fixed_shape_tensor/vtable.rs | 11 +- vortex-tensor/src/types/vector/arrow.rs | 147 +++------ 27 files changed, 1551 insertions(+), 724 deletions(-) create mode 100644 vortex-arrow/src/extension/adapter.rs create mode 100644 vortex-arrow/src/extension/mod.rs create mode 100644 vortex-arrow/src/extension/tests.rs diff --git a/encodings/uuid/src/arrow.rs b/encodings/uuid/src/arrow.rs index 1e3f67bf3fe..4a4011c5859 100644 --- a/encodings/uuid/src/arrow.rs +++ b/encodings/uuid/src/arrow.rs @@ -30,13 +30,12 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_arrow::nulls; @@ -56,30 +55,30 @@ const UUID_BYTE_LEN: i32 = 16; /// The cached Id of Arrow's canonical `arrow.uuid` extension type. static ARROW_UUID: CachedId = CachedId::new(ArrowUuid::NAME); -impl ArrowExportVTable for Uuid { - fn arrow_ext_id(&self) -> Id { - *ARROW_UUID +impl ArrowExtVTable for Uuid { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::PreserveVortex } - fn vortex_id(&self) -> Id { - Uuid.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_UUID } - // Encode all of these. fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, _session: &ArrowSession, ) -> VortexResult> { let mut field = Field::new( - name.to_string(), + name, DataType::FixedSizeBinary(UUID_BYTE_LEN), - dtype.is_nullable(), + ext_dtype.storage_dtype().is_nullable(), ); field .try_with_extension_type(ArrowUuid) - .vortex_expect("FixedSizeBinary[16] is correct type for ArrowUuid"); + .vortex_expect("UUID field was constructed with FixedSizeBinary(16)"); + Ok(Some(field)) } @@ -90,28 +89,13 @@ impl ArrowExportVTable for Uuid { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_uuid = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_uuid { - return Ok(ArrowExport::Unsupported(array)); - } Ok(ArrowExport::Exported(try_fsl_to_fsb(array, options, ctx)?)) } -} - -impl ArrowImportVTable for Uuid { - fn arrow_ext_id(&self) -> Id { - *ARROW_UUID - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if !field.has_valid_extension_type::() { return Ok(None); } @@ -122,24 +106,20 @@ impl ArrowImportVTable for Uuid { field.is_nullable().into(), ); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Uuid, + UuidMetadata::default(), + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, _field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, _session: &ArrowSession, ) -> VortexResult { - let DType::Extension(dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !matches!(array.data_type(), DataType::FixedSizeBinary(UUID_BYTE_LEN)) - || !dtype.is::() - { + if !matches!(array.data_type(), DataType::FixedSizeBinary(UUID_BYTE_LEN)) { return Ok(ArrowImport::Unsupported(array)); } @@ -150,7 +130,7 @@ impl ArrowImportVTable for Uuid { PType::U8, Validity::NonNullable, ); - let validity = nulls(fsb.nulls(), dtype.is_nullable())?; + let validity = nulls(fsb.nulls(), ext_dtype.storage_dtype().is_nullable())?; let storage = FixedSizeListArray::new( u8_array.into_array(), @@ -160,7 +140,7 @@ impl ArrowImportVTable for Uuid { ) .into_array(); Ok(ArrowImport::Imported( - ExtensionArray::new(dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } @@ -221,7 +201,9 @@ mod tests { use arrow_schema::Fields; use arrow_schema::Schema; use rstest::rstest; + use uuid::Version; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::DictArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; @@ -237,6 +219,40 @@ mod tests { use super::*; + #[track_caller] + fn check_version_round_trip( + version: Option, + nullable: bool, + dictionary: bool, + ) -> VortexResult<()> { + let session = uuid_session(); + let arrow = session.arrow(); + let mut ctx = session.create_execution_ctx(); + let array = versioned_uuid_array(version, nullable)?; + let mut field = arrow.to_arrow_field("id", array.dtype())?; + assert!(field.has_valid_extension_type::()); + assert_eq!(arrow.from_arrow_field(&field)?, *array.dtype()); + + let array = if dictionary { + field = field.clone().with_data_type(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(field.data_type().clone()), + )); + DictArray::try_new(PrimitiveArray::from_iter([2u32, 1, 0]).into_array(), array)? + .into_array() + } else { + array + }; + + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) + } + /// A session with the UUID extension dtype and its Arrow plugins registered. fn uuid_session() -> VortexSession { let session = VortexSession::empty(); @@ -288,46 +304,6 @@ mod tests { Ok(()) } - #[test] - fn to_arrow_field_top_level_uuid_carries_extension_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let field = session.to_arrow_field("id", &uuid_dtype(false))?; - assert!(field.has_valid_extension_type::()); - Ok(()) - } - - #[test] - fn to_arrow_field_struct_with_nested_uuid_preserves_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::Struct( - StructFields::from_iter([(FieldName::from("id"), uuid_dtype(false))]), - Nullability::NonNullable, - ); - let field = session.to_arrow_field("row", &dtype)?; - let DataType::Struct(inner) = field.data_type() else { - panic!("expected Struct, got {:?}", field.data_type()); - }; - assert_eq!(inner.len(), 1); - assert_eq!(inner[0].data_type(), &DataType::FixedSizeBinary(16)); - assert!(inner[0].has_valid_extension_type::()); - Ok(()) - } - - #[test] - fn to_arrow_field_list_of_uuid_preserves_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable); - let field = session.to_arrow_field("ids", &dtype)?; - let DataType::List(elem) = field.data_type() else { - panic!("expected List, got {:?}", field.data_type()); - }; - assert!(elem.has_valid_extension_type::()); - Ok(()) - } - #[test] fn to_arrow_field_fixed_size_list_of_uuid_preserves_metadata() -> VortexResult<()> { let vortex_session = uuid_session(); @@ -378,7 +354,7 @@ mod tests { } #[test] - fn to_arrow_schema_struct_of_struct_uuid() -> VortexResult<()> { + fn schema_roundtrip_preserves_nested_uuid() -> VortexResult<()> { let vortex_session = uuid_session(); let session = vortex_session.arrow(); let inner = DType::Struct( @@ -386,7 +362,14 @@ mod tests { Nullability::NonNullable, ); let outer = DType::Struct( - StructFields::from_iter([(FieldName::from("payload"), inner)]), + StructFields::from_iter([ + (FieldName::from("payload"), inner), + (FieldName::from("id"), uuid_dtype(false)), + ( + FieldName::from("ids"), + DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable), + ), + ]), Nullability::NonNullable, ); let schema = session.to_arrow_schema(&outer)?; @@ -395,6 +378,8 @@ mod tests { panic!("expected Struct, got {:?}", payload.data_type()); }; assert!(inner_fields[0].has_valid_extension_type::()); + assert_eq!(session.from_arrow_schema(&schema)?, outer); + Ok(()) } @@ -414,26 +399,6 @@ mod tests { Ok(()) } - #[test] - fn schema_roundtrip_preserves_nested_uuid() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::Struct( - StructFields::from_iter([ - (FieldName::from("id"), uuid_dtype(false)), - ( - FieldName::from("ids"), - DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable), - ), - ]), - Nullability::NonNullable, - ); - let schema = session.to_arrow_schema(&dtype)?; - let roundtripped = session.from_arrow_schema(&schema)?; - assert_eq!(roundtripped, dtype); - Ok(()) - } - #[test] fn to_arrow_datatype_dispatches_plugins() -> VortexResult<()> { let vortex_session = uuid_session(); @@ -471,13 +436,7 @@ mod tests { let mut ctx = vortex_session.create_execution_ctx(); let session = vortex_session.arrow(); - let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); - field.try_with_extension_type(ArrowUuid)?; - let arrow_array: ArrowArrayRef = Arc::new(FixedSizeBinaryArray::try_from_iter( - [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), - )?); - - let vortex_array = session.from_arrow_array(arrow_array, &field)?; + let vortex_array = uuid_array(&session)?; let vortex_ext = vortex_array.dtype().as_extension(); assert!(vortex_ext.is::()); @@ -656,12 +615,8 @@ mod tests { // extension survives the export as well. let mut ctx = vortex_session.create_execution_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::(); - let uuids = values.values().as_fixed_size_binary(); - assert_eq!(uuids.value(0), b"0123456789abcdef"); - assert_eq!(uuids.value(1), b"fedcba9876543210"); + Ok(()) } @@ -849,4 +804,78 @@ mod tests { Ok(()) } + + /// Builds three UUIDs with valid version bits, with a null middle row when nullable. + fn versioned_uuid_array(version: Option, nullable: bool) -> VortexResult { + let dtype = ExtDType::try_with_vtable( + Uuid, + UuidMetadata { version }, + uuid_dtype(nullable).as_extension().storage_dtype().clone(), + )? + .erased(); + + let bytes = match version { + Some(Version::Max) => [u8::MAX; 16], + Some(Version::Nil) | None => [0; 16], + Some(version) => uuid::Builder::from_bytes([0; 16]) + .with_version(version) + .with_variant(uuid::Variant::RFC4122) + .into_uuid() + .into_bytes(), + }; + + let validity = if nullable { + Validity::from_iter([ + true, // First UUID. + false, // Null middle UUID. + true, // Last UUID. + ]) + } else { + Validity::NonNullable + }; + let storage = FixedSizeListArray::try_new( + PrimitiveArray::from_iter((0..3).flat_map(|_| bytes)).into_array(), + 16, + validity, + 3, + )?; + + Ok(ExtensionArray::try_new(dtype, storage.into_array())?.into_array()) + } + + #[rstest] + #[case::unspecified(None, false, false)] + #[case::nil(Some(Version::Nil), false, false)] + #[case::mac(Some(Version::Mac), false, false)] + #[case::dce(Some(Version::Dce), false, false)] + #[case::md5(Some(Version::Md5), false, false)] + #[case::random(Some(Version::Random), false, false)] + #[case::sha1(Some(Version::Sha1), false, false)] + #[case::sort_mac(Some(Version::SortMac), false, false)] + #[case::sort_rand(Some(Version::SortRand), false, false)] + #[case::custom(Some(Version::Custom), false, false)] + #[case::max(Some(Version::Max), false, false)] + #[case::nullable(Some(Version::Random), true, false)] + #[case::dictionary(Some(Version::Random), false, true)] + #[case::nullable_dictionary(Some(Version::Random), true, true)] + fn version_round_trip( + #[case] version: Option, + #[case] nullable: bool, + #[case] dictionary: bool, + ) -> VortexResult<()> { + check_version_round_trip(version, nullable, dictionary) + } + + #[test] + fn mismatched_metadata_is_rejected() -> VortexResult<()> { + let session = uuid_session(); + let mut ctx = session.create_execution_ctx(); + let array = uuid_array(&session.arrow())?; + let target_array = versioned_uuid_array(Some(Version::Random), false)?; + let field = session.arrow().to_arrow_field("id", target_array.dtype())?; + + assert!(array.execute_arrow(Some(&field), &mut ctx).is_err()); + + Ok(()) + } } diff --git a/encodings/uuid/src/lib.rs b/encodings/uuid/src/lib.rs index b079120d88c..7db124c50a4 100644 --- a/encodings/uuid/src/lib.rs +++ b/encodings/uuid/src/lib.rs @@ -24,8 +24,6 @@ mod arrow; mod metadata; mod vtable; -use std::sync::Arc; - pub use metadata::UuidMetadata; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; @@ -44,8 +42,7 @@ pub struct Uuid; pub fn initialize(session: &VortexSession) { session.dtypes().register(Uuid); session.aggregate_fns().register_storage_ordered(Uuid); - session.arrow().register_exporter(Arc::new(Uuid)); - session.arrow().register_importer(Arc::new(Uuid)); + session.arrow().register_extension(Uuid); } #[cfg(test)] diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index e4fc05bf4bc..30e89093216 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -10,9 +10,9 @@ use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_schema::DataType; +use arrow_schema::Field; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; @@ -24,6 +24,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_dictionary`] requires for export. @@ -40,27 +41,36 @@ impl Matcher for ArrowDictExportable { pub(super) fn to_arrow_dictionary( array: ArrayRef, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let array = array.execute_until::(ctx)?; + // Keep an existing empty dictionary: canonicalization can discard its codes, and Arrow + // cannot repack every extension's storage type into a dictionary. + let array = if array.is::() || array.is::() { + array + } else { + array.execute_until::(ctx)? + }; let array = match array.try_downcast::() { - Ok(dict) => return dict_to_dict(dict, codes_type, values_type, options, ctx), + Ok(dict) => return dict_to_dict(dict, codes_type, values_field, options, ctx), Err(array) => array, }; let array = match array.try_downcast::() { - Ok(constant) => return constant_to_dict(constant, codes_type, values_type, options, ctx), + Ok(constant) => return constant_to_dict(constant, codes_type, values_field, 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), options, ctx)?; + let array = array.execute_arrow_with_options(Some(values_field), options, ctx)?; arrow_cast::cast( &array, - &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), + &DataType::Dictionary( + Box::new(codes_type.clone()), + Box::new(values_field.data_type().clone()), + ), ) .map_err(VortexError::from) } @@ -69,21 +79,22 @@ pub(super) fn to_arrow_dictionary( fn constant_to_dict( array: ConstantArray, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); let scalar = array.scalar(); if scalar.is_null() { - let dict_type = - DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())); + let dict_type = DataType::Dictionary( + Box::new(codes_type.clone()), + Box::new(values_field.data_type().clone()), + ); return Ok(new_null_array(&dict_type, len)); } - let values = execute_arrow_naive( - ConstantArray::new(scalar.clone(), 1).into_array(), - Some(values_type), + let values = ConstantArray::new(scalar.clone(), 1).execute_arrow_with_options( + Some(values_field), options, ctx, )?; @@ -95,12 +106,17 @@ fn constant_to_dict( fn dict_to_dict( array: DictArray, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { 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)?; + let values = + array + .values() + .clone() + .execute_arrow_with_options(Some(values_field), options, ctx)?; + make_dict_array(codes_type, codes, values) } @@ -169,20 +185,54 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::mask::Mask; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::ArrowSessionExt; use crate::ExecuteArrowExt; use crate::executor::dictionary::ConstantArray; use crate::executor::dictionary::DictArray; + #[track_caller] + fn check_fixed_size_list_dictionary_round_trip(len: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0..i32::try_from(len * 2)?).into_array(), + 2, + Validity::NonNullable, + len, + )? + .into_array(); + let source = DictArray::try_new( + PrimitiveArray::from_iter(0..u32::try_from(len)?).into_array(), + values, + )? + .into_array(); + let arrow = SESSION.arrow(); + let values_field = arrow.to_arrow_field("item", source.dtype())?; + let field = Field::new( + "dict", + dict_type(DataType::Int32, values_field.data_type().clone()), + false, + ); + + let exported = source.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, source, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); fn dict_type(codes: DataType, values: DataType) -> DataType { @@ -269,4 +319,11 @@ mod tests { assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array); Ok(()) } + + #[rstest] + #[case::empty(0)] + #[case::nonempty(3)] + fn fixed_size_list_values_round_trip(#[case] len: usize) -> VortexResult<()> { + check_fixed_size_list_dictionary_round_trip(len) + } } diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 803f038962f..386c8723dc9 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -173,6 +173,11 @@ fn execute_arrow_array( } }; + if let DataType::Dictionary(codes, values) = target_field.data_type() { + let values_field = target_field.clone().with_data_type(values.as_ref().clone()); + return to_arrow_dictionary(array, codes, &values_field, options, ctx); + } + if let Some(arrow_ext_name) = target_field.extension_type_name() { return execute_arrow_extension(array, arrow_ext_name, target_field, options, ctx); } @@ -297,9 +302,9 @@ fn execute_arrow_ext_storage( /// 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. +/// An explicit target may omit nested extension metadata or name an unregistered Vortex ID, so +/// export compares its inferred type with the storage layout. Each nested field is checked when +/// it is exported. fn eq_storage_layout(lhs: &DType, rhs: &DType) -> bool { match (lhs, rhs) { (DType::List(lhs, _), DType::List(rhs, _)) => eq_nested_storage_layout(lhs, rhs), @@ -402,9 +407,17 @@ pub(crate) fn execute_arrow_naive( 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, options, ctx) - } + DataType::Dictionary(codes_type, values_type) => to_arrow_dictionary( + array.clone(), + codes_type, + &Field::new( + "", + values_type.as_ref().clone(), + array.dtype().is_nullable(), + ), + options, + ctx, + ), dt @ DataType::Decimal32(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal64(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal128(..) => to_arrow_decimal(array, dt, ctx), diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index 8559a0a327e..c48b119835d 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -170,14 +170,12 @@ fn create_from_fields( arrow_arrays.push(arrow_field); } - Ok(Arc::new(unsafe { - ArrowStructArray::new_unchecked_with_length( - fields.clone(), - arrow_arrays, - null_buffer, - len, - ) - })) + Ok(Arc::new(ArrowStructArray::try_new_with_length( + fields.clone(), + arrow_arrays, + null_buffer, + len, + )?)) } Err(names) => { // No target fields specified - use preferred types for each child. The inferred field @@ -213,14 +211,12 @@ fn create_from_fields( } let arrow_fields = Fields::from(arrow_fields); - Ok(Arc::new(unsafe { - ArrowStructArray::new_unchecked_with_length( - arrow_fields, - arrow_arrays, - null_buffer, - len, - ) - })) + Ok(Arc::new(ArrowStructArray::try_new_with_length( + arrow_fields, + arrow_arrays, + null_buffer, + len, + )?)) } } } diff --git a/vortex-arrow/src/ext_metadata.rs b/vortex-arrow/src/ext_metadata.rs index 85d7b47610d..37ae679684c 100644 --- a/vortex-arrow/src/ext_metadata.rs +++ b/vortex-arrow/src/ext_metadata.rs @@ -1,30 +1,121 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The Arrow field metadata that holds a Vortex extension type without an Arrow plugin. +//! Encoding Vortex extension types in Arrow field metadata. //! -//! Arrow field metadata must be UTF-8, but Vortex metadata is arbitrary bytes, so the metadata is -//! always base64, and [`EXTENSION_METADATA_ENCODING_KEY`] marks it as base64. The marker is needed -//! because other producers write the same Arrow extension keys with plain text metadata. +//! Typed plugins use [`crate::ArrowSession::register_extension`], whose adapter handles Vortex +//! metadata without exposing its keys or encoding to the callbacks. +//! +//! A custom extension with no Arrow [`ExtensionType`] implementation can use +//! [`set_arrow_extension_name`] to name its storage field. An extension without an Arrow plugin +//! uses [`ArrowSession::to_arrow_storage_field`] instead. +//! +//! Field metadata must be UTF-8, so Vortex payloads use marked base64. The Vortex ID path also +//! accepts unmarked UTF-8 metadata from older producers. A separate Vortex payload requires the +//! marker, and an absent payload imposes no constraint on a target's Vortex metadata. +//! +//! [`ArrowSession::to_arrow_storage_field`]: crate::ArrowSession::to_arrow_storage_field +//! [`ExtensionType`]: arrow_schema::extension::ExtensionType use arrow_schema::Field; use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use base64::Engine; use base64::prelude::BASE64_STANDARD; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDTypePlugin; use vortex_array::dtype::extension::ExtDTypeRef; +use vortex_array::dtype::extension::ExtVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; -/// The Arrow field metadata key that names the encoding of `ARROW:extension:metadata`. +/// The Arrow field metadata key that names the encoding of a Vortex metadata payload. /// -/// The only encoding is `base64`. Without this key, the metadata is its own UTF-8 bytes. +/// This applies to `ARROW:extension:metadata` and [`VORTEX_EXTENSION_METADATA_KEY`]. The only +/// encoding is `base64`. Without this key, `ARROW:extension:metadata` is read as raw UTF-8 for +/// backward compatibility. pub const EXTENSION_METADATA_ENCODING_KEY: &str = "vortex:extension:metadata:encoding"; +/// Vortex metadata retained alongside an Arrow extension type. +/// +/// The payload is base64 and **must** carry [`EXTENSION_METADATA_ENCODING_KEY`]. It is separate +/// from the Arrow extension's own metadata, whose format is defined by that extension. +pub const VORTEX_EXTENSION_METADATA_KEY: &str = "vortex:extension:metadata"; + /// The value of [`EXTENSION_METADATA_ENCODING_KEY`] for standard padded base64. const BASE64_ENCODING: &str = "base64"; +/// Sets an Arrow extension name while preserving the field's other metadata. +/// +/// Prefer [`Field::try_with_extension_type`] when an Arrow extension has an +/// [`ExtensionType`](arrow_schema::extension::ExtensionType) implementation, which also validates +/// its storage and native metadata. +pub fn set_arrow_extension_name(field: &mut Field, name: impl Into) { + field + .metadata_mut() + .insert(EXTENSION_TYPE_NAME_KEY.to_owned(), name.into()); +} + +/// Writes Vortex metadata alongside an Arrow field's native metadata. +/// +/// Serialization, keys, and encoding are owned by this module. Callers decide whether the Arrow +/// representation needs a separate Vortex payload. A failed serialization leaves `field` unchanged. +pub(crate) fn write_vortex_metadata( + field: &mut Field, + vtable: &V, + metadata: &V::Metadata, +) -> VortexResult<()> { + let payload = BASE64_STANDARD.encode(vtable.serialize_metadata(metadata)?); + field.metadata_mut().extend([ + (VORTEX_EXTENSION_METADATA_KEY.to_owned(), payload), + ( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + BASE64_ENCODING.to_owned(), + ), + ]); + + Ok(()) +} + +/// Restores a Vortex dtype using metadata and the importer's inferred Vortex storage. +/// +/// Returns `None` when the field has no Vortex payload. A present payload **must** have a +/// supported encoding marker and metadata compatible with `storage_dtype`. Native Arrow metadata +/// is not read. Reconstruction goes through the dtype plugin so its storage checks also apply. +pub(crate) fn read_vortex_metadata( + field: &Field, + vtable: &V, + storage_dtype: &DType, +) -> VortexResult> { + let Some(payload) = field.metadata().get(VORTEX_EXTENSION_METADATA_KEY) else { + return Ok(None); + }; + + vortex_ensure_eq!( + field + .metadata() + .get(EXTENSION_METADATA_ENCODING_KEY) + .map(String::as_str) + .unwrap_or_default(), + BASE64_ENCODING, + "Arrow field {:?} must mark its Vortex metadata as base64", + field.name() + ); + let metadata = decode_base64_metadata(field, payload)?; + + vtable + .deserialize(&metadata, storage_dtype.clone()) + .map(Some) + .map_err(|err| { + err.with_context(format!( + "Invalid Vortex metadata on Arrow field {:?}", + field.name() + )) + }) +} + /// Returns the Arrow field metadata that holds `ext_dtype`. pub fn encode_arrow_ext_metadata(ext_dtype: &ExtDTypeRef) -> VortexResult<[(String, String); 3]> { let metadata = BASE64_STANDARD.encode(ext_dtype.serialize_metadata()?); @@ -57,15 +148,152 @@ pub fn decode_arrow_ext_metadata(field: &Field) -> VortexResult> { .map(String::as_str) { None => Ok(metadata.as_bytes().to_vec()), - Some(BASE64_ENCODING) => BASE64_STANDARD.decode(metadata).map_err(|err| { - vortex_err!( - "Arrow field {:?} must carry base64 extension metadata, got {metadata:?}: {err}", - field.name() - ) - }), + Some(BASE64_ENCODING) => decode_base64_metadata(field, metadata), Some(encoding) => vortex_bail!( "Arrow field {:?} must have the extension metadata encoding {BASE64_ENCODING:?}, got {encoding:?}", field.name() ), } } + +fn decode_base64_metadata(field: &Field, metadata: &str) -> VortexResult> { + BASE64_STANDARD.decode(metadata).map_err(|err| { + vortex_err!( + "Arrow field {:?} must carry base64 extension metadata, got {metadata:?}: {err}", + field.name() + ) + }) +} + +#[cfg(test)] +mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + use rstest::rstest; + use vortex_array::dtype::PType; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::TimeUnit; + use vortex_error::VortexResult; + + use super::*; + + #[track_caller] + fn check_metadata_round_trip(metadata: TimeUnit, native_type: DataType) -> VortexResult<()> { + let mut field = Field::new("date", native_type, true); + let inferred = Date::try_new(metadata, true.into())?.erased(); + field.metadata_mut().extend([ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "example.date".to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "native metadata".to_owned(), + ), + ("application".to_owned(), "retained".to_owned()), // Unrelated application metadata. + ]); + let original = field.clone(); + + write_vortex_metadata(&mut field, &Date, &metadata)?; + + assert_eq!( + read_vortex_metadata(&field, &Date, inferred.storage_dtype())?, + Some(inferred) + ); + assert_eq!(field.data_type(), original.data_type()); + assert_eq!(field.name(), original.name()); + assert_eq!(field.is_nullable(), original.is_nullable()); + for (key, value) in original.metadata() { + assert_eq!(field.metadata().get(key), Some(value)); + } + + Ok(()) + } + + #[rstest] + #[case::days(TimeUnit::Days, DataType::Date32)] + #[case::milliseconds(TimeUnit::Milliseconds, DataType::Date64)] + fn typed_metadata_preserves_native_field( + #[case] metadata: TimeUnit, + #[case] native_type: DataType, + ) -> VortexResult<()> { + check_metadata_round_trip(metadata, native_type) + } + + #[test] + fn absent_vortex_metadata_ignores_native_metadata() -> VortexResult<()> { + let mut field = Field::new("date", DataType::Date32, false); + field.metadata_mut().extend([ + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "native metadata".to_owned(), + ), + ( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + "native encoding".to_owned(), + ), + ]); + + assert_eq!( + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into()))?, + None + ); + + Ok(()) + } + + #[rstest] + #[case::missing_encoding("AA==", None)] + #[case::unsupported_encoding("AA==", Some("hex"))] + #[case::invalid_base64("!", Some("base64"))] + #[case::invalid_typed_metadata("/w==", Some("base64"))] + #[case::empty_typed_metadata("", Some("base64"))] + fn malformed_metadata_is_rejected(#[case] payload: &str, #[case] encoding: Option<&str>) { + let mut field = Field::new("date", DataType::Date32, false); + field + .metadata_mut() + .insert(VORTEX_EXTENSION_METADATA_KEY.to_owned(), payload.to_owned()); + if let Some(encoding) = encoding { + field.metadata_mut().insert( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + encoding.to_owned(), + ); + } + + let error = + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into())) + .expect_err("the metadata is malformed"); + assert!(error.to_string().contains("date")); + } + + #[test] + fn metadata_must_match_inferred_storage() -> VortexResult<()> { + let mut field = Field::new("date", DataType::Date32, false); + write_vortex_metadata(&mut field, &Date, &TimeUnit::Milliseconds)?; + + assert!( + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into())) + .is_err() + ); + + Ok(()) + } + + #[test] + fn extension_name_preserves_other_metadata() { + let mut field = Field::new("value", DataType::Float32, false); + field + .metadata_mut() + .insert("application".to_owned(), "retained".to_owned()); + + set_arrow_extension_name(&mut field, "example.vector"); + + assert_eq!(field.extension_type_name(), Some("example.vector")); + assert_eq!( + field.metadata().get("application").map(String::as_str), + Some("retained") + ); + } +} diff --git a/vortex-arrow/src/extension/adapter.rs b/vortex-arrow/src/extension/adapter.rs new file mode 100644 index 00000000000..28a5310cb36 --- /dev/null +++ b/vortex-arrow/src/extension/adapter.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Connects typed extension conversions to the Arrow session's importer and exporter registries. +//! +//! [`ArrowExtVTable`] defines conversion in both directions. This adapter provides the separate +//! interfaces that [`ArrowSession`] registers and applies the vtable's [`ArrowMetadataPolicy`]. +//! Keeping metadata handling here prevents plugins from forgetting it. The session checks that +//! completed fields and imported arrays preserve their dtypes. +//! +//! [`VORTEX_EXTENSION_METADATA_KEY`] names `vortex:extension:metadata`. This separate field entry +//! leaves `ARROW:extension:metadata` for the Arrow extension's own metadata. The codec in +//! [`crate::ext_metadata`] uses `vortex:extension:metadata:encoding` to mark the payload as base64 +//! and checks that marker before decoding. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::registry::Id; + +use super::ArrowExtVTable; +use super::ArrowMetadataPolicy; +use crate::ArrowExport; +use crate::ArrowExportOptions; +use crate::ArrowExportVTable; +use crate::ArrowImport; +use crate::ArrowImportVTable; +use crate::ArrowSession; +use crate::ArrowSessionExt; +use crate::ext_metadata::VORTEX_EXTENSION_METADATA_KEY; +use crate::ext_metadata::read_vortex_metadata; +use crate::ext_metadata::write_vortex_metadata; + +impl ArrowSession { + /// Registers both directions of a typed extension using the vtable's metadata policy. + /// + /// The adapter handles Vortex metadata encoding, decoding, and target checks. The session + /// verifies complete dtype identity for both typed and raw plugins. + pub fn register_extension(&self, vtable: V) { + let adapter = Arc::new(ExtensionPlugin { vtable }); + let exporter = Arc::clone(&adapter); + + self.register_exporter(exporter); + self.register_importer(adapter); + } +} + +/// Registers one typed extension as both an importer and an exporter with the same metadata policy. +#[derive(Debug)] +struct ExtensionPlugin { + vtable: V, +} + +impl ExtensionPlugin { + /// Rejects export targets whose Vortex metadata disagrees with `ext_dtype`. + /// + /// A target's metadata determines the meaning of exported values on import. Under + /// [`ArrowMetadataPolicy::PreserveVortex`], a target with [`VORTEX_EXTENSION_METADATA_KEY`] + /// **must** carry matching metadata. An absent payload adds no Vortex metadata constraint. + /// Comparison uses decoded metadata so it does not depend on serialization details. + fn validate_export_target( + &self, + ext_dtype: &ExtDType, + target: &Field, + ctx: &ExecutionCtx, + ) -> VortexResult<()> { + if V::metadata_policy() != ArrowMetadataPolicy::PreserveVortex + || !target + .metadata() + .contains_key(VORTEX_EXTENSION_METADATA_KEY) + { + return Ok(()); + } + + let native_dtype = + V::from_arrow_field(target, &ctx.session().arrow())?.ok_or_else(|| { + vortex_err!( + "Arrow field {:?} must have a supported native representation, got {}", + target.name(), + target.data_type() + ) + })?; + if let Some(actual) = + read_vortex_metadata(target, &self.vtable, native_dtype.storage_dtype())? + { + vortex_ensure_eq!( + actual.metadata::(), + ext_dtype.metadata(), + "Arrow field {:?} must carry the expected Vortex metadata", + target.name() + ); + } + + Ok(()) + } +} + +impl ArrowExportVTable for ExtensionPlugin { + fn arrow_ext_id(&self) -> Id { + self.vtable.arrow_ext_id() + } + + fn vortex_id(&self) -> Id { + self.vtable.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(None); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(None); + }; + + let Some(mut field) = self.vtable.to_arrow_field(name, &ext_dtype, session)? else { + return Ok(None); + }; + + // Import routing determines whether the native field needs a separate Vortex payload. + // The session checks the completed field after any payload has been added. + if V::metadata_policy() == ArrowMetadataPolicy::PreserveVortex + && session.from_arrow_field(&field)? != *dtype + { + write_vortex_metadata(&mut field, &self.vtable, ext_dtype.metadata())?; + } + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(ext_dtype) = array.dtype().as_extension_opt() else { + return Ok(ArrowExport::Unsupported(array)); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + + self.validate_export_target(&ext_dtype, target, ctx)?; + + self.vtable.execute_arrow(array, target, options, ctx) + } +} + +impl ArrowImportVTable for ExtensionPlugin { + fn arrow_ext_id(&self) -> Id { + self.vtable.arrow_ext_id() + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + let Some(native_dtype) = V::from_arrow_field(field, session)? else { + return Ok(None); + }; + + // The plugin can map Arrow's physical storage to a different Vortex storage dtype. Restore + // metadata against that inferred storage so the dtype plugin can validate the combination. + let restored_dtype = if V::metadata_policy() == ArrowMetadataPolicy::PreserveVortex { + read_vortex_metadata(field, &self.vtable, native_dtype.storage_dtype())? + } else { + None + }; + + Ok(Some(DType::Extension( + restored_dtype.unwrap_or_else(|| native_dtype.erased()), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(ArrowImport::Unsupported(array)); + }; + + V::from_arrow_array(array, field, &ext_dtype, session) + } +} diff --git a/vortex-arrow/src/extension/mod.rs b/vortex-arrow/src/extension/mod.rs new file mode 100644 index 00000000000..928585dbd6f --- /dev/null +++ b/vortex-arrow/src/extension/mod.rs @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interoperability for Vortex extension types. +//! +//! Implement [`ArrowExtVTable`] and register it with [`ArrowSession::register_extension`] to enable +//! Arrow import and export. Registration wraps the callbacks in an adapter that handles Vortex +//! metadata, while the session validates schema round trips. Plugins only implement conversions. +//! +//! A Vortex extension does not need an equivalent standard Arrow type. It can use a custom Arrow +//! extension name, or decline native export so the session exports storage with the Vortex ID and +//! metadata. When a plugin supplies an Arrow representation, its metadata may not describe the full +//! Vortex dtype. Handling that gap in the adapter prevents plugins from silently losing metadata. +//! +//! Declare an [`ArrowMetadataPolicy`] in [`ArrowExtVTable::metadata_policy`] based on whether the +//! Arrow representation preserves the full Vortex dtype on its own. Callbacks leave transport to +//! the adapter and import arrays with the restored dtype they receive. Plugins remain responsible +//! for preserving values and lengths. + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; +use vortex_error::VortexResult; +use vortex_session::registry::Id; + +use crate::ArrowExport; +use crate::ArrowExportOptions; +use crate::ArrowImport; +use crate::ArrowSession; + +/// Whether an Arrow representation needs separate Vortex metadata or not to restore the full dtype. +/// +/// On schema export, the session imports the generated field and checks that it reproduces the +/// original dtype, including metadata, storage, and nullability. On successful array import, the +/// session checks that the returned array has the dtype inferred from the field. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ArrowMetadataPolicy { + /// Rely on the Arrow representation alone to preserve the complete Vortex dtype. + /// + /// Choose this only when the plugin's Arrow field can reconstruct the dtype without additional + /// Vortex metadata. The adapter neither writes nor reads a separate Vortex metadata payload. + ArrowNative, + + /// Preserve Vortex metadata separately when the Arrow representation alone is insufficient. + /// + /// If the exported field does not reconstruct the original dtype, the adapter adds a Vortex + /// metadata payload before the session checks the completed field. The Arrow extension's own + /// metadata stays unchanged. On import, the payload restores metadata over the plugin's inferred + /// storage. + /// + /// An absent payload uses the native interpretation. A required empty payload is still + /// written, because a codec can distinguish it from absence. + PreserveVortex, +} + +/// Native Arrow conversion callbacks for one Vortex extension type. +/// +/// Register through [`ArrowSession::register_extension`], and use the session's conversion APIs. +/// Calling these callbacks directly bypasses metadata transport and schema validation. Callbacks +/// **must** describe only the native Arrow representation, leaving Vortex metadata transport to the +/// adapter. Register the owning dtype in the Vortex session as well to enable its Vortex ID path. +/// +/// Register the dtype and its Arrow conversions in the same session: +/// +/// ``` +/// use vortex_array::dtype::session::DTypeSessionExt; +/// use vortex_arrow::{ArrowExtVTable, ArrowSessionExt}; +/// use vortex_session::VortexSession; +/// +/// fn register(session: &VortexSession, vtable: V) { +/// session.dtypes().register(vtable.clone()); +/// session.arrow().register_extension(vtable); +/// } +/// ``` +pub trait ArrowExtVTable: ExtVTable { + /// The metadata policy used by the adapter for this extension type. + /// + /// The result **must** remain the same across calls and sessions. + fn metadata_policy() -> ArrowMetadataPolicy; + + /// The native Arrow extension ID used for both import and export. + fn arrow_ext_id(&self) -> Id; + + /// Builds the native field, or declines so that the session uses the Vortex ID path. + /// + /// Use [`Field::try_with_extension_type`] for standard Arrow extension metadata, or + /// [`crate::ext_metadata::set_arrow_extension_name`] for a custom name. The adapter adds any + /// Vortex metadata before the session validates the completed field against the source dtype. + fn to_arrow_field( + &self, + name: &str, + ext_dtype: &ExtDType, + session: &ArrowSession, + ) -> VortexResult>; + + /// Converts an array of this extension type to the native target representation. + /// + /// The adapter checks the source extension type before calling this method. Return + /// [`ArrowExport::Unsupported`] when the target is unsupported. Child exports **must** receive + /// the same `options` through [`crate::ExecuteArrowExt::execute_arrow_with_options`]. + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult; + + /// Interprets the native field, without applying Vortex side metadata. + /// + /// Infer the Vortex storage representation here, which can differ from Arrow's physical type. + /// Return `None` when the field is unsupported. + fn from_arrow_field( + field: &Field, + session: &ArrowSession, + ) -> VortexResult>>; + + /// Imports native values with the extension dtype restored by the adapter. + /// + /// `ext_dtype` includes any Vortex metadata restored by the adapter, so this callback does not + /// need to decode it. Inferring the dtype again from the native Arrow representation could lose + /// that metadata. + /// + /// A successful result **must** have exactly `ext_dtype` and the input length. Return + /// [`ArrowImport::Unsupported`] when the representation is unsupported. Import storage through + /// the session with a nullability rather than the extension field, to avoid dispatching this + /// plugin recursively. + fn from_arrow_array( + array: ArrowArrayRef, + field: &Field, + ext_dtype: &ExtDType, + session: &ArrowSession, + ) -> VortexResult; +} + +mod adapter; + +#[cfg(test)] +mod tests; diff --git a/vortex-arrow/src/extension/tests.rs b/vortex-arrow/src/extension/tests.rs new file mode 100644 index 00000000000..0490c13a685 --- /dev/null +++ b/vortex-arrow/src/extension/tests.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Int64Array; +use arrow_schema::DataType; +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::dtype::session::DTypeSessionExt; +use vortex_array::scalar::ScalarValue; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::*; +use crate::ArrowExportVTable; +use crate::ArrowImportVTable; +use crate::ArrowSessionExt; +use crate::ext_metadata::VORTEX_EXTENSION_METADATA_KEY; +use crate::ext_metadata::set_arrow_extension_name; + +/// Exports a tagged field through its registered policy and checks its complete dtype on import. +#[track_caller] +fn round_trip_field(tag: u8, nullable: bool) -> VortexResult { + let session = tagged_session::(); + let arrow = session.arrow(); + let dtype = DType::Extension(Tagged::::dtype(tag, nullable)?.erased()); + + let field = arrow.to_arrow_field("tag", &dtype)?; + assert_eq!(arrow.from_arrow_field(&field)?, dtype); + + Ok(field) +} + +fn tagged_session() -> VortexSession { + let session = VortexSession::empty(); + session.dtypes().register(Tagged::); + session + .arrow() + .register_extension(Tagged::); + + session +} + +fn raw_tagged_session() -> VortexSession { + let session = VortexSession::empty(); + let plugin = Arc::new(RawTagged); + let exporter = Arc::clone(&plugin); + session.arrow().register_exporter(exporter); + session.arrow().register_importer(plugin); + + session +} + +/// The native field infers tag 255. Tag zero deliberately has an empty serialized payload. +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +struct Tagged; + +impl Tagged { + fn dtype(tag: u8, nullable: bool) -> VortexResult> { + ExtDType::try_with_vtable(Self, tag, DType::Primitive(PType::I64, nullable.into())) + } +} + +impl ExtVTable for Tagged { + type Metadata = u8; + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> Id { + static ID: CachedId = CachedId::new("test.typed_tag"); + *ID + } + + fn serialize_metadata(&self, tag: &u8) -> VortexResult> { + Ok(if *tag == 0 { vec![] } else { vec![*tag] }) + } + + fn deserialize_metadata(&self, bytes: &[u8]) -> VortexResult { + match bytes { + [] => Ok(0), + [tag] => Ok(*tag), + _ => vortex_bail!("tag must contain at most one byte"), + } + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + vortex_ensure_eq!( + *ext_dtype.storage_dtype(), + DType::Primitive(PType::I64, ext_dtype.storage_dtype().nullability()) + ); + + Ok(()) + } + + fn unpack_native<'a>( + _ext_dtype: &ExtDType, + value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(value) + } +} + +impl ArrowExtVTable for Tagged { + fn metadata_policy() -> ArrowMetadataPolicy { + if PRESERVE_VORTEX { + ArrowMetadataPolicy::PreserveVortex + } else { + ArrowMetadataPolicy::ArrowNative + } + } + + fn arrow_ext_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + ext_dtype: &ExtDType, + _session: &ArrowSession, + ) -> VortexResult> { + let mut field = Field::new( + name, + DataType::Int64, + ext_dtype.storage_dtype().is_nullable(), + ); + set_arrow_extension_name(&mut field, self.arrow_ext_id().to_string()); + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + _options: &ArrowExportOptions, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(ArrowExport::Unsupported(array)) + } + + fn from_arrow_field( + field: &Field, + _session: &ArrowSession, + ) -> VortexResult>> { + Ok(Some(Self::dtype(255, field.is_nullable())?)) + } + + fn from_arrow_array( + array: ArrowArrayRef, + _field: &Field, + _ext_dtype: &ExtDType, + _session: &ArrowSession, + ) -> VortexResult { + Ok(ArrowImport::Unsupported(array)) + } +} + +/// Bypasses the typed adapter and deliberately drops the extension when importing arrays. +#[derive(Debug)] +struct RawTagged; + +impl ArrowExportVTable for RawTagged { + fn arrow_ext_id(&self) -> Id { + Tagged::.id() + } + + fn vortex_id(&self) -> Id { + Tagged::.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + as ArrowExtVTable>::to_arrow_field( + &Tagged, + name, + &dtype.as_extension().clone().downcast::>(), + session, + ) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + _options: &ArrowExportOptions, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(ArrowExport::Unsupported(array)) + } +} + +impl ArrowImportVTable for RawTagged { + fn arrow_ext_id(&self) -> Id { + Tagged::.id() + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + Ok( + as ArrowExtVTable>::from_arrow_field(field, session)? + .map(|dtype| DType::Extension(dtype.erased())), + ) + } + + fn from_arrow_array( + &self, + _array: ArrowArrayRef, + _field: &Field, + _dtype: &DType, + _session: &ArrowSession, + ) -> VortexResult { + Ok(ArrowImport::Imported( + PrimitiveArray::from_iter([1i64]).into_array(), + )) + } +} + +#[rstest] +#[case::native(false, 255, false, None)] +#[case::nullable_native(false, 255, true, None)] +#[case::already_exact(true, 255, false, None)] +#[case::empty_metadata(true, 0, true, Some(""))] +#[case::nonempty_metadata(true, 7, false, Some("Bw=="))] +fn policy_preserves_dtype( + #[case] preserve_vortex: bool, + #[case] tag: u8, + #[case] nullable: bool, + #[case] expected_payload: Option<&str>, +) -> VortexResult<()> { + let field = if preserve_vortex { + round_trip_field::(tag, nullable)? + } else { + round_trip_field::(tag, nullable)? + }; + + assert_eq!( + field + .metadata() + .get(VORTEX_EXTENSION_METADATA_KEY) + .map(String::as_str), + expected_payload + ); + + Ok(()) +} + +#[test] +fn lossy_native_policy_is_rejected() -> VortexResult<()> { + let error = + round_trip_field::(7, false).expect_err("the native field cannot express tag 7"); + + assert!(error.to_string().contains("must preserve its Vortex dtype")); + + Ok(()) +} + +#[test] +fn raw_exporter_must_preserve_dtype() -> VortexResult<()> { + let session = raw_tagged_session(); + let dtype = DType::Extension(Tagged::::dtype(7, false)?.erased()); + let error = session + .arrow() + .to_arrow_field("tag", &dtype) + .expect_err("the raw exporter's field loses tag 7"); + + assert!(error.to_string().contains("must preserve its Vortex dtype")); + + Ok(()) +} + +#[test] +fn raw_importer_must_preserve_dtype() -> VortexResult<()> { + let session = raw_tagged_session(); + let dtype = DType::Extension(Tagged::::dtype(255, false)?.erased()); + let field = session.arrow().to_arrow_field("tag", &dtype)?; + let array: ArrowArrayRef = Arc::new(Int64Array::from(vec![1])); + let error = session + .arrow() + .from_arrow_array(array, &field) + .expect_err("the raw importer discards the inferred extension dtype"); + + assert!( + error + .to_string() + .contains("must preserve the dtype inferred from its field") + ); + + Ok(()) +} diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 0ff749bcb08..671f2dfbfcb 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -20,6 +20,11 @@ pub mod convert; mod datum; pub mod dtype; mod executor; + +mod extension; +pub use extension::ArrowExtVTable; +pub use extension::ArrowMetadataPolicy; + pub mod ext_metadata; mod iter; mod null_buffer; diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index ef6aa4029e9..2a2a2d9dc3d 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -158,8 +158,12 @@ impl<'a> IntoArrowField<'a> for Nullability { /// Plugin layer for exporting a Vortex array to an Arrow extension type. /// +/// Prefer [`crate::ArrowExtVTable`] for automatic Vortex metadata transport. Both interfaces must +/// satisfy the dtype round-trip contract checked by [`ArrowSession`]. +/// /// This is purely an implementation trait, its methods should not be called directly. Instead, -/// use [`ExecuteArrowExt`] to export an array, and the methods on [`ArrowSession`] to convert a dtype. +/// 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 { @@ -199,6 +203,9 @@ pub trait ArrowExportVTable: 'static + Send + Sync + Debug { /// Plugin layer for importing an Arrow extension-typed array into a Vortex array. /// +/// Prefer [`crate::ArrowExtVTable`] for automatic Vortex metadata transport. Both interfaces must +/// return arrays with the dtype inferred from their fields, as checked by [`ArrowSession`]. +/// /// Plugins are dispatched by `arrow_ext_id`. The session peels Arrow dictionaries before it /// dispatches, so a plugin never sees a [`DataType::Dictionary`] field or array. It sees the value /// field, with the dictionary field's metadata, and the values array. The session then wraps the @@ -258,6 +265,15 @@ type ArrowImporterRegistry = ArcSwapMap>; /// Session-scoped registry of Arrow extension plugins. /// +/// All plugins **must** preserve the complete dtype across [`Self::to_arrow_field`] and +/// [`Self::from_arrow_field`]. This includes extension IDs, metadata, storage, and nullability at +/// every nesting level. Exporting an array to that field and importing it with the same field must +/// preserve its dtype and logical values. The owning dtype and Arrow plugins must be registered in +/// the importing session. +/// +/// The session checks inferred fields and successful plugin imports, including raw plugins. +/// Explicit export targets may request a different dtype. Unknown Vortex IDs import as storage. +/// /// A plugin converts one Vortex extension type to and from an Arrow extension type. The crate that /// owns the type registers its plugin, as `vortex-uuid` does. Temporal types need no plugin, /// because Arrow has native types for them. @@ -485,6 +501,15 @@ impl ArrowSession { for plugin in self.exporters_by_vortex(&ext_dtype.id()).iter() { if let Some(field) = plugin.to_arrow_field(name, &dtype, self)? { + // Check the completed field through normal routing because another importer can + // take precedence over the exporter that inferred it. + vortex_ensure_eq!( + self.from_arrow_field(&field)?, + dtype, + "Arrow field {:?} must preserve its Vortex dtype", + name + ); + return Ok(Some(field)); } } @@ -804,6 +829,12 @@ impl ArrowSession { arr.len(), len ); + vortex_ensure_eq!( + arr.dtype(), + &dtype, + "Arrow importer must preserve the dtype inferred from its field" + ); + return Ok(arr); } ArrowImport::Unsupported(arr) => { diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 30c311a72d7..0fa7b31992b 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -15,12 +15,11 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexExpect; @@ -39,28 +38,21 @@ fn has_valid_json_extension(field: &Field) -> bool { && ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata()).is_ok() } -impl ArrowExportVTable for Json { - fn arrow_ext_id(&self) -> Id { - *ARROW_JSON +impl ArrowExtVTable for Json { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - Json.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_JSON } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let DType::Extension(ext_dtype) = dtype else { - return Ok(None); - }; - if !ext_dtype.is::() { - return Ok(None); - } - let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field .try_with_extension_type(ArrowJson::default()) @@ -75,15 +67,6 @@ impl ArrowExportVTable for Json { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_json = array - .dtype() - .as_extension_opt() - .map(|ext_dtype| ext_dtype.is::()) - .unwrap_or(false); - if !is_json { - return Ok(ArrowExport::Unsupported(array)); - } - ArrowJson::try_new_from_field_metadata(target.data_type(), target.metadata())?; let executed = array.execute::(ctx)?; @@ -98,52 +81,40 @@ impl ArrowExportVTable for Json { Ok(ArrowExport::Exported(storage)) } -} - -impl ArrowImportVTable for Json { - fn arrow_ext_id(&self) -> Id { - *ARROW_JSON - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if !has_valid_json_extension(field) { return Ok(None); } - Ok(Some(DType::Extension( - ExtDType::::try_new(EmptyMetadata, DType::Utf8(field.is_nullable().into()))? - .erased(), - ))) + Ok(Some(ExtDType::::try_new( + EmptyMetadata, + DType::Utf8(field.is_nullable().into()), + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let DType::Extension(ext_dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() || !has_valid_json_extension(field) { + if !has_valid_json_extension(field) { return Ok(ArrowImport::Unsupported(array)); } let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::new(ext_dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::LazyLock; diff --git a/vortex-json/src/lib.rs b/vortex-json/src/lib.rs index 15c06c82b9f..2759846d1e1 100644 --- a/vortex-json/src/lib.rs +++ b/vortex-json/src/lib.rs @@ -13,8 +13,6 @@ mod arrow; mod dtype; mod json_to_variant; -use std::sync::Arc; - pub mod editions; pub use dtype::Json; pub use json_to_variant::JsonToVariant; @@ -31,8 +29,7 @@ use vortex_session::VortexSession; /// Register JSON extension support with a session. pub fn initialize(session: &VortexSession) { session.dtypes().register(Json); - session.arrow().register_exporter(Arc::new(Json)); - session.arrow().register_importer(Arc::new(Json)); + session.arrow().register_extension(Json); session.scalar_fns().register(JsonToVariant); // JSON is an opt-in durable dtype, so it belongs to an independently enabled edition family. diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 7e7f2de49d9..78f48f9f33a 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -40,9 +40,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -255,27 +255,31 @@ impl LineStringData { } } -impl ArrowExportVTable for LineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_LINESTRING +impl ArrowExtVTable for LineString { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = linestring_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(linestring_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = linestring_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = linestring_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -287,15 +291,6 @@ impl ArrowExportVTable for LineString { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_linestring = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_linestring { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(linestring_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -320,22 +315,11 @@ impl ArrowExportVTable for LineString { Ok(ArrowExport::Exported(Arc::new(linestrings.into_arrow()))) } -} - -impl ArrowImportVTable for LineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_LINESTRING - } - /// Import a `geoarrow.linestring` field as the [`LineString`] dtype. Keyed off the standard - /// GeoArrow name, so any producer resolves here. Accepts the full `LineStringType` extension, or - /// — for a metadata-less geometry literal — the name alone, inferring the dimension from the - /// coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(linestring_meta) = field.try_extension_type::() { vortex_ensure!( @@ -368,23 +352,20 @@ impl ArrowImportVTable for LineString { }; let storage_dtype = linestring_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(LineString, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + LineString, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -392,7 +373,7 @@ impl ArrowImportVTable for LineString { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index b31887bd83b..33e7eb61e31 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Spatial extension types and metadata shared by their Arrow plugins. +//! +//! Geometry types derive their coordinate dimension from storage and carry an optional coordinate +//! reference system in [`SpatialMetadata`]. Shared metadata handling keeps these properties +//! consistent across the GeoArrow fields used to import and export each geometry. + pub(crate) mod coordinate; mod linestring; mod multilinestring; @@ -352,13 +358,23 @@ mod tests { use prost::Message; use rstest::rstest; + use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::dtype::proto::dtype as pb; use vortex_array::dtype::proto::dtype::d_type::DtypeType; + use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -374,6 +390,55 @@ mod tests { use crate::extension::coordinate::coordinate_storage_dtype; use crate::test_harness::spatial_session; + #[track_caller] + fn check_nested_arrow_round_trip( + vtable: impl ExtVTable, + depth: usize, + nullable: bool, + fixed: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let mut ctx = session.create_execution_ctx(); + let mut storage = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + for _ in 0..depth { + storage = DType::List(Arc::new(storage), Nullability::NonNullable); + } + + storage = storage.with_nullability(nullable.into()); + let dtype = ExtDType::try_with_vtable( + vtable, + SpatialMetadata::new(Some("EPSG:4326".into())), + storage, + )? + .erased(); + let mut builder = + builder_with_capacity_in(dtype.storage_dtype(), 2, BufferAllocatorRef::static_ref()); + builder.append_zeros(2); + let array = ExtensionArray::try_new(dtype, builder.finish())?.into_array(); + let array = if fixed { + FixedSizeListArray::try_new(array, 1, Validity::NonNullable, 2)?.into_array() + } else { + ListArray::try_new( + array, + PrimitiveArray::from_iter([0u32, 1, 2]).into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array() + }; + + let arrow = session.arrow(); + let field = arrow.to_arrow_field("geometry", array.dtype())?; + assert_eq!(arrow.from_arrow_field(&field)?, *array.dtype()); + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + exported.to_data().validate_full()?; + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); /// Test shim: decode with an explicitly constructed session. @@ -607,4 +672,22 @@ mod tests { ))); Ok(()) } + + #[rstest] + #[case::line(LineString, 1, false, false)] + #[case::points(MultiPoint, 1, false, false)] + #[case::polygon(Polygon, 2, false, false)] + #[case::multiline(MultiLineString, 2, false, false)] + #[case::multipolygon(MultiPolygon, 3, false, false)] + #[case::nullable(LineString, 1, true, false)] + #[case::fixed_size_list(LineString, 1, false, true)] + #[case::nullable_fixed_list(LineString, 1, true, true)] + fn nested_arrow_round_trip( + #[case] vtable: impl ExtVTable, + #[case] depth: usize, + #[case] nullable: bool, + #[case] fixed: bool, + ) -> VortexResult<()> { + check_nested_arrow_round_trip(vtable, depth, nullable, fixed) + } } diff --git a/vortex-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index a4ea09631ea..6d16cea1e60 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -33,9 +33,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -185,27 +185,31 @@ impl MultiLineStringData { } } -impl ArrowExportVTable for MultiLineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTILINESTRING +impl ArrowExtVTable for MultiLineString { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multilinestring_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multilinestring_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multilinestring_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multilinestring_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -217,15 +221,6 @@ impl ArrowExportVTable for MultiLineString { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multilinestring = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multilinestring { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multilinestring_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -252,20 +247,11 @@ impl ArrowExportVTable for MultiLineString { multilinestrings.into_arrow(), ))) } -} - -impl ArrowImportVTable for MultiLineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTILINESTRING - } - /// Import a `geoarrow.multilinestring` field (matched by GeoArrow name). Accepts the full - /// `MultiLineStringType`, or a metadata-less literal (name only), inferring the dimension. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multilinestring_meta) = field.try_extension_type::() { vortex_ensure!( @@ -301,23 +287,20 @@ impl ArrowImportVTable for MultiLineString { }; let storage_dtype = multilinestring_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiLineString, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiLineString, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -325,7 +308,7 @@ impl ArrowImportVTable for MultiLineString { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index e57cc57c5be..67bc2f27986 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -33,9 +33,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -172,27 +172,31 @@ impl MultiPointData { } } -impl ArrowExportVTable for MultiPoint { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOINT +impl ArrowExtVTable for MultiPoint { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multipoint_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multipoint_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multipoint_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multipoint_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -204,15 +208,6 @@ impl ArrowExportVTable for MultiPoint { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multipoint = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multipoint { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multipoint_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -237,21 +232,11 @@ impl ArrowExportVTable for MultiPoint { Ok(ArrowExport::Exported(Arc::new(multipoints.into_arrow()))) } -} - -impl ArrowImportVTable for MultiPoint { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOINT - } - /// Import a `geoarrow.multipoint` field as the [`MultiPoint`] dtype (matched by GeoArrow name). - /// Accepts the full `MultiPointType`, or a metadata-less literal (name only), inferring the - /// dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multipoint_meta) = field.try_extension_type::() { vortex_ensure!( @@ -282,23 +267,20 @@ impl ArrowImportVTable for MultiPoint { }; let storage_dtype = multipoint_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiPoint, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiPoint, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -306,7 +288,7 @@ impl ArrowImportVTable for MultiPoint { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 711bdb016d3..a1f8552f587 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -32,9 +32,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -182,27 +182,31 @@ impl MultiPolygonData { } } -impl ArrowExportVTable for MultiPolygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOLYGON +impl ArrowExtVTable for MultiPolygon { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multipolygon_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multipolygon_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multipolygon_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multipolygon_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -214,15 +218,6 @@ impl ArrowExportVTable for MultiPolygon { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multipolygon = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multipolygon { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multipolygon_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -247,20 +242,11 @@ impl ArrowExportVTable for MultiPolygon { Ok(ArrowExport::Exported(Arc::new(multipolygons.into_arrow()))) } -} - -impl ArrowImportVTable for MultiPolygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOLYGON - } - /// Import a `geoarrow.multipolygon` field (matched by GeoArrow name). Accepts the full - /// `MultiPolygonType`, or a metadata-less literal (name only), inferring the dimension. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multipolygon_meta) = field.try_extension_type::() { vortex_ensure!( @@ -299,23 +285,20 @@ impl ArrowImportVTable for MultiPolygon { }; let storage_dtype = multipolygon_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiPolygon, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiPolygon, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -323,7 +306,7 @@ impl ArrowImportVTable for MultiPolygon { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index 46b465289ec..5fc9a31985b 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -31,9 +31,9 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -151,26 +151,25 @@ pub(crate) fn point_geometries( .collect() } -impl ArrowExportVTable for Point { - fn arrow_ext_id(&self) -> Id { - *ARROW_POINT +impl ArrowExtVTable for Point { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_POINT } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = coordinate_dimension(ext_type.storage_dtype())?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = coordinate_dimension(ext_dtype.storage_dtype())?; - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(point_type(spatial_metadata, dimension))?; Ok(Some(field)) @@ -183,15 +182,6 @@ impl ArrowExportVTable for Point { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_point = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_point { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(point_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -217,22 +207,11 @@ impl ArrowExportVTable for Point { Ok(ArrowExport::Exported(points.into_arrow())) } -} - -impl ArrowImportVTable for Point { - fn arrow_ext_id(&self) -> Id { - *ARROW_POINT - } - /// Import a `geoarrow.point` field as the [`Point`] dtype. Keyed off the standard GeoArrow name, - /// so any producer (DataFusion, DuckDB, geoarrow-rs, …) resolves here. Accepts the full - /// `PointType` extension, or — for a metadata-less geometry literal — the name alone, inferring - /// the dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(point_meta) = field.try_extension_type::() { vortex_ensure!( @@ -262,23 +241,20 @@ impl ArrowImportVTable for Point { }; let storage_dtype = coordinate_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Point, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Point, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::Struct(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -286,7 +262,7 @@ impl ArrowImportVTable for Point { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index b43101799cd..08a6120adec 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -38,9 +38,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -244,27 +244,31 @@ impl PolygonData { } } -impl ArrowExportVTable for Polygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_POLYGON +impl ArrowExtVTable for Polygon { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_POLYGON } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = polygon_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(polygon_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = polygon_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = polygon_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -276,15 +280,6 @@ impl ArrowExportVTable for Polygon { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_polygon = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_polygon { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(polygon_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -309,22 +304,11 @@ impl ArrowExportVTable for Polygon { Ok(ArrowExport::Exported(Arc::new(polygons.into_arrow()))) } -} - -impl ArrowImportVTable for Polygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_POLYGON - } - /// Import a `geoarrow.polygon` field as the [`Polygon`] dtype. Keyed off the standard GeoArrow - /// name, so any producer (DataFusion, DuckDB, geoarrow-rs, …) resolves here. Accepts the full - /// `PolygonType` extension, or — for a metadata-less geometry literal — the name alone, inferring - /// the dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(polygon_meta) = field.try_extension_type::() { vortex_ensure!( @@ -361,23 +345,20 @@ impl ArrowImportVTable for Polygon { }; let storage_dtype = polygon_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Polygon, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Polygon, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -385,7 +366,7 @@ impl ArrowImportVTable for Polygon { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index 9680595cbe5..bccda5917ec 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -40,9 +40,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; @@ -205,26 +205,25 @@ pub(crate) fn rect_geometries( .collect() } -impl ArrowExportVTable for Rect { - fn arrow_ext_id(&self) -> Id { - *ARROW_BOX +impl ArrowExtVTable for Rect { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_BOX } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = box_dimension(ext_type.storage_dtype())?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = box_dimension(ext_dtype.storage_dtype())?; - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(box_type(spatial_metadata, dimension))?; Ok(Some(field)) @@ -237,15 +236,6 @@ impl ArrowExportVTable for Rect { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_box = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_box { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(box_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -267,21 +257,11 @@ impl ArrowExportVTable for Rect { Ok(ArrowExport::Exported(Arc::new(rects.into_arrow()))) } -} - -impl ArrowImportVTable for Rect { - fn arrow_ext_id(&self) -> Id { - *ARROW_BOX - } - /// Import a `geoarrow.box` field as the [`Rect`] dtype. Keyed off the standard GeoArrow name, - /// so any producer resolves here. Accepts the full `BoxType` extension, or — for a - /// metadata-less literal — the name alone, inferring the dimension from the field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(box_meta) = field.try_extension_type::() { ( box_meta.dimension().into(), @@ -305,23 +285,20 @@ impl ArrowImportVTable for Rect { }; let storage_dtype = box_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Rect, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Rect, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::Struct(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -329,7 +306,7 @@ impl ArrowImportVTable for Rect { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index 47713736c7a..d40b31449a7 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -26,9 +26,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; @@ -168,25 +168,24 @@ impl ExtVTable for WellKnownBinary { static ARROW_WKB: CachedId = CachedId::new(WkbType::NAME); -impl ArrowExportVTable for WellKnownBinary { - fn arrow_ext_id(&self) -> Id { - *ARROW_WKB +impl ArrowExtVTable for WellKnownBinary { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_WKB } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let spatial_metadata = ext_dtype.metadata(); - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(wkb_type(spatial_metadata))?; Ok(Some(field)) @@ -199,15 +198,6 @@ impl ArrowExportVTable for WellKnownBinary { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_wkb = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_wkb { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(wkb_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -247,41 +237,30 @@ impl ArrowExportVTable for WellKnownBinary { Ok(ArrowExport::Exported(arrow_ref)) } -} - -impl ArrowImportVTable for WellKnownBinary { - fn arrow_ext_id(&self) -> Id { - *ARROW_WKB - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let Ok(wkb_meta) = field.try_extension_type::() else { return Ok(None); }; let storage_dtype = DType::Binary(field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(WellKnownBinary, spatial_metadata(&wkb_meta), storage_dtype)? - .erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + WellKnownBinary, + spatial_metadata(&wkb_meta), + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!( array.data_type(), DataType::Binary | DataType::LargeBinary | DataType::BinaryView @@ -292,7 +271,7 @@ impl ArrowImportVTable for WellKnownBinary { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::new(ext_dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 99fc5f8e556..5704e476915 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -3,8 +3,6 @@ #![doc = include_str!(concat!("../", env!("CARGO_PKG_README")))] -use std::sync::Arc; - use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; @@ -48,31 +46,25 @@ mod tests; /// Set up a session with spatial extension types, Arrow interop, functions, and pruning. pub fn initialize(session: &VortexSession) { - // Register the spatial extension types. - session.dtypes().register(WellKnownBinary); - session.arrow().register_exporter(Arc::new(WellKnownBinary)); - session.arrow().register_importer(Arc::new(WellKnownBinary)); - session.dtypes().register(Point); - session.arrow().register_exporter(Arc::new(Point)); - session.arrow().register_importer(Arc::new(Point)); - session.dtypes().register(LineString); - session.arrow().register_exporter(Arc::new(LineString)); - session.arrow().register_importer(Arc::new(LineString)); - session.dtypes().register(MultiPoint); - session.arrow().register_exporter(Arc::new(MultiPoint)); - session.arrow().register_importer(Arc::new(MultiPoint)); - session.dtypes().register(Polygon); - session.arrow().register_exporter(Arc::new(Polygon)); - session.arrow().register_importer(Arc::new(Polygon)); - session.dtypes().register(MultiLineString); - session.arrow().register_exporter(Arc::new(MultiLineString)); - session.arrow().register_importer(Arc::new(MultiLineString)); - session.dtypes().register(MultiPolygon); - session.arrow().register_exporter(Arc::new(MultiPolygon)); - session.arrow().register_importer(Arc::new(MultiPolygon)); - session.dtypes().register(Rect); - session.arrow().register_exporter(Arc::new(Rect)); - session.arrow().register_importer(Arc::new(Rect)); + let dtypes = session.dtypes(); + let arrow = session.arrow(); + + dtypes.register(WellKnownBinary); + arrow.register_extension(WellKnownBinary); + dtypes.register(Point); + arrow.register_extension(Point); + dtypes.register(LineString); + arrow.register_extension(LineString); + dtypes.register(MultiPoint); + arrow.register_extension(MultiPoint); + dtypes.register(Polygon); + arrow.register_extension(Polygon); + dtypes.register(MultiLineString); + arrow.register_extension(MultiLineString); + dtypes.register(MultiPolygon); + arrow.register_extension(MultiPolygon); + dtypes.register(Rect); + arrow.register_extension(Rect); // Register the geometry scalar functions. session.scalar_fns().register(SpatialArea); diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index b2939c21d82..59bb0cd8f37 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -12,8 +12,6 @@ allow(clippy::unwrap_used, clippy::expect_used, clippy::unwrap_in_result) )] -use std::sync::Arc; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; @@ -56,8 +54,7 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(FixedShapeTensor); let arrow_session = session.arrow(); - arrow_session.register_exporter(Arc::new(Vector)); - arrow_session.register_importer(Arc::new(Vector)); + arrow_session.register_extension(Vector); let session_fns = session.scalar_fns(); diff --git a/vortex-tensor/src/types/fixed_shape_tensor/proto.rs b/vortex-tensor/src/types/fixed_shape_tensor/proto.rs index 2bd608d717d..7633a8842c4 100644 --- a/vortex-tensor/src/types/fixed_shape_tensor/proto.rs +++ b/vortex-tensor/src/types/fixed_shape_tensor/proto.rs @@ -4,7 +4,6 @@ //! Protobuf serialization for [`FixedShapeTensorMetadata`]. use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -36,12 +35,12 @@ struct FixedShapeTensorMetadataProto { } /// Serializes [`FixedShapeTensorMetadata`] to protobuf bytes. -pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { +pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> VortexResult> { let logical_shape = metadata .logical_shape() .iter() - .map(|&d| u32::try_from(d).vortex_expect("dimension size exceeds u32")) - .collect(); + .map(|&dimension| u32::try_from(dimension)) + .collect::, _>>()?; let dim_names = metadata.dim_names().map(|n| n.to_vec()).unwrap_or_default(); @@ -49,9 +48,10 @@ pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { .permutation() .map(|p| { p.iter() - .map(|&i| u32::try_from(i).vortex_expect("permutation index exceeds u32")) - .collect() + .map(|&index| u32::try_from(index)) + .collect::, _>>() }) + .transpose()? .unwrap_or_default(); let proto = FixedShapeTensorMetadataProto { @@ -59,7 +59,8 @@ pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { dim_names, permutation, }; - proto.encode_to_vec() + + Ok(proto.encode_to_vec()) } /// Deserializes [`FixedShapeTensorMetadata`] from protobuf bytes. diff --git a/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs b/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs index b98e3dd0acf..ad29b091479 100644 --- a/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs +++ b/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs @@ -28,7 +28,7 @@ impl ExtVTable for FixedShapeTensor { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(proto::serialize(metadata)) + proto::serialize(metadata) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { @@ -119,4 +119,13 @@ mod tests { ) -> VortexResult<()> { assert_roundtrip(&metadata?) } + + #[test] + #[cfg(target_pointer_width = "64")] + fn unrepresentable_dimension_returns_error() -> VortexResult<()> { + let metadata = FixedShapeTensorMetadata::new(vec![0, u32::MAX as usize + 1]); + assert!(FixedShapeTensor.serialize_metadata(&metadata).is_err()); + + Ok(()) + } } diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index fafe66297c3..d7fadcfc471 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -13,7 +13,6 @@ use arrow_array::Array; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::DataType; use arrow_schema::Field; -use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use vortex_array::ArrayRef; use vortex_array::EmptyMetadata; use vortex_array::ExecutionCtx; @@ -22,14 +21,14 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; +use vortex_arrow::ext_metadata::set_arrow_extension_name; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -41,18 +40,6 @@ pub const ARROW_VECTOR_EXTENSION_NAME: &str = "vortex.tensor.vector"; static ARROW_VECTOR: CachedId = CachedId::new(ARROW_VECTOR_EXTENSION_NAME); -#[expect( - clippy::disallowed_types, - reason = "Arrow's Field::set_metadata requires std::collections::HashMap" -)] -fn vector_extension_metadata() -> std::collections::HashMap { - [( - EXTENSION_TYPE_NAME_KEY.to_string(), - ARROW_VECTOR_EXTENSION_NAME.to_string(), - )] - .into() -} - fn is_supported_float(data_type: &DataType) -> bool { matches!( data_type, @@ -60,32 +47,24 @@ fn is_supported_float(data_type: &DataType) -> bool { ) } -impl ArrowExportVTable for Vector { - fn arrow_ext_id(&self) -> Id { - *ARROW_VECTOR +impl ArrowExtVTable for Vector { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - Vector.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_VECTOR } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let DType::Extension(dtype) = dtype else { - return Ok(None); - }; - - if !dtype.is::() { - return Ok(None); - } + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; + set_arrow_extension_name(&mut field, ARROW_VECTOR_EXTENSION_NAME); - // Delegate to Arrow encoding of storage type. - let mut field = session.to_arrow_field(name, dtype.storage_dtype())?; - field.set_metadata(vector_extension_metadata()); Ok(Some(field)) } @@ -96,14 +75,6 @@ impl ArrowExportVTable for Vector { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - if !array - .dtype() - .as_extension_opt() - .is_some_and(|ext| ext.is::()) - { - return Ok(ArrowExport::Unsupported(array)); - } - let executed = array.execute::(ctx)?; let storage = executed.storage_array().clone(); @@ -111,21 +82,15 @@ impl ArrowExportVTable for Vector { Ok(ArrowExport::Exported(arrow_storage)) } -} - -impl ArrowImportVTable for Vector { - fn arrow_ext_id(&self) -> Id { - *ARROW_VECTOR - } fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if field.extension_type_name() != Some(ARROW_VECTOR_EXTENSION_NAME) { return Ok(None); } + let DataType::FixedSizeList(elem, list_size) = field.data_type() else { return Ok(None); }; @@ -138,24 +103,20 @@ impl ArrowImportVTable for Vector { *list_size as u32, field.is_nullable().into(), ); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Vector, EmptyMetadata, storage_dtype)?.erased(), - ))) + + Ok(Some(ExtDType::try_with_vtable( + Vector, + EmptyMetadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, _field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let DType::Extension(dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !dtype.is::() { - return Ok(ArrowImport::Unsupported(array)); - } let DataType::FixedSizeList(elem, _) = array.data_type() else { return Ok(ArrowImport::Unsupported(array)); }; @@ -163,9 +124,9 @@ impl ArrowImportVTable for Vector { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array(array, dtype.is_nullable())?; + let storage = session.from_arrow_array(array, ext_dtype.storage_dtype().is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } @@ -190,7 +151,6 @@ mod tests { use vortex_array::dtype::StructFields; use vortex_array::dtype::extension::ExtDType; use vortex_array::validity::Validity; - use vortex_arrow::ArrowExport; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; @@ -235,8 +195,7 @@ mod tests { /// 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.arrow().register_extension(Vector); session } @@ -264,15 +223,15 @@ mod tests { fn from_arrow_field_recovers_vector_dtype() -> VortexResult<()> { let vortex_session = session_with_vector(); let session = vortex_session.arrow(); - let arrow_field = Field::new( + let mut arrow_field = Field::new( "embedding", DataType::FixedSizeList( Arc::new(Field::new("item", DataType::Float32, false)), DIM as i32, ), true, - ) - .with_metadata(vector_extension_metadata()); + ); + set_arrow_extension_name(&mut arrow_field, ARROW_VECTOR_EXTENSION_NAME); let dtype = session.from_arrow_field(&arrow_field)?; assert_eq!(dtype, vector_dtype(true)); Ok(()) @@ -362,7 +321,7 @@ mod tests { ), false, ); - arrow_field.set_metadata(vector_extension_metadata()); + set_arrow_extension_name(&mut arrow_field, ARROW_VECTOR_EXTENSION_NAME); let arrow = arrow_fsl_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], DIM as i32); let imported = session.from_arrow_array(arrow, &arrow_field)?; assert!( @@ -381,20 +340,16 @@ mod tests { } #[test] - fn execute_arrow_returns_unsupported_for_non_vector() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + fn execute_arrow_rejects_incompatible_non_vector() -> VortexResult<()> { + let session = session_with_vector(); + let mut ctx = session.create_execution_ctx(); let primitive = PrimitiveArray::from_iter([1_i32, 2, 3]).into_array(); - let target = Field::new("ints", DataType::Int32, false); - let result = ::execute_arrow( - &Vector, - primitive.clone(), - &target, - &ArrowExportOptions::default(), - &mut ctx, - )?; - assert!( - matches!(result, ArrowExport::Unsupported(arr) if arr.dtype() == primitive.dtype()) - ); + let target = session + .arrow() + .to_arrow_field("embedding", &vector_dtype(false))?; + + assert!(primitive.execute_arrow(Some(&target), &mut ctx).is_err()); + Ok(()) } @@ -404,36 +359,10 @@ mod tests { let field = Field::new("embedding", DataType::Int32, false); let int_array: ArrowArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let result = ::from_arrow_array( - &Vector, + let result = ::from_arrow_array( int_array, &field, - &dtype, - &VortexSession::empty().arrow(), - )?; - assert!(matches!(result, ArrowImport::Unsupported(_))); - Ok(()) - } - - #[test] - fn from_arrow_array_returns_unsupported_for_non_vector_dtype() -> VortexResult<()> { - use vortex_uuid::Uuid; - use vortex_uuid::UuidMetadata; - let uuid_storage = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), - 16, - Nullability::NonNullable, - ); - let uuid_ext = - ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), uuid_storage)?.erased(); - - let fsl_arrow = arrow_fsl_f32(vec![1.0, 2.0, 3.0], DIM as i32); - let field = Field::new("embedding", fsl_arrow.data_type().clone(), false); - let result = ::from_arrow_array( - &Vector, - fsl_arrow, - &field, - &DType::Extension(uuid_ext), + &dtype.as_extension().clone().downcast::(), &VortexSession::empty().arrow(), )?; assert!(matches!(result, ArrowImport::Unsupported(_))); From 49acec03272162ef7b7f24c7f8dd039ecd50f148 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Tue, 10 Nov 2026 20:29:07 -0500 Subject: [PATCH 12/13] perf(compressor): Improve string dictionary selection (#9762) String Dictionary selection now probes a distributed sample before building an exact candidate. The previous estimator missed profitable dictionaries when columns contained thousands of repeated strings. Samples above 75% distinct values reject Dictionary; an incomplete probe uses the existing estimator. Admitted candidates have explicit limits for dictionary entries, value bytes, code bytes, and input bytes. The compressor compares the candidate's actual size and reuses it if it wins. The probe reads at most 8,192 rows and 1 MiB of values. A row that exceeds the byte budget is excluded from the completed-row count, including the final sampled row. String samples compact their payload before size estimates so unrelated retained bytes do not distort the score. --------- Signed-off-by: Will Manning Signed-off-by: Robert Kruszewski --- vortex-array/src/builders/dict/bytes.rs | 32 +- vortex-array/src/builders/dict/mod.rs | 15 + vortex-btrblocks/Cargo.toml | 5 + .../benches/string_dict_selection.rs | 182 +++++++ vortex-btrblocks/src/schemes/string/sparse.rs | 24 +- .../tests/string_dict_selection.rs | 206 +++++++ vortex-compressor/src/builtins/dict/binary.rs | 66 +-- .../src/builtins/dict/bytes_candidate.rs | 503 ++++++++++++++++++ .../builtins/dict/bytes_candidate/tests.rs | 305 +++++++++++ vortex-compressor/src/builtins/dict/mod.rs | 1 + vortex-compressor/src/builtins/dict/string.rs | 66 +-- vortex-compressor/src/compressor/cascade.rs | 21 +- vortex-compressor/src/compressor/mod.rs | 5 + vortex-compressor/src/compressor/sample.rs | 29 +- vortex-compressor/src/compressor/select.rs | 13 +- vortex-compressor/src/compressor/tests.rs | 178 +++++++ vortex-compressor/src/scheme/estimate.rs | 2 +- vortex-compressor/src/stats/cache.rs | 18 + vortex-layout/src/layouts/dict/writer.rs | 3 +- 19 files changed, 1528 insertions(+), 146 deletions(-) create mode 100644 vortex-btrblocks/benches/string_dict_selection.rs create mode 100644 vortex-btrblocks/tests/string_dict_selection.rs create mode 100644 vortex-compressor/src/builtins/dict/bytes_candidate.rs create mode 100644 vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 1df5dddb5f2..daff5f2a957 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -25,6 +25,7 @@ use vortex_utils::aliases::hash_map::HashTableEntry; use super::DictConstraints; use super::DictEncoder; +use super::codes_ptype_for_max_len; use crate::ArrayRef; use crate::ArrayView; use crate::IntoArray; @@ -59,18 +60,18 @@ pub fn bytes_dict_builder( constraints: &DictConstraints, allocator: BufferAllocatorRef, ) -> Box { - match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + match codes_ptype_for_max_len(constraints.max_len as u64) { + PType::U8 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, )), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + PType::U16 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, )), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + PType::U32 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, @@ -363,6 +364,7 @@ mod test { use std::sync::Arc; use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; @@ -370,6 +372,7 @@ mod test { use vortex_session::VortexSession; use super::BytesDictBuilder; + use super::bytes_dict_builder; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; @@ -379,11 +382,13 @@ mod test { use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; use crate::buffer::BufferHandle; + use crate::builders::dict::DictConstraints; use crate::builders::dict::UNCONSTRAINED; use crate::builders::dict::dict_encode; use crate::builders::dict::dict_encoder_in; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -510,6 +515,25 @@ mod test { Ok(()) } + #[rstest] + #[case(usize::from(u8::MAX) + 1, PType::U8)] + #[case(usize::from(u8::MAX) + 2, PType::U16)] + #[case(usize::from(u16::MAX) + 1, PType::U16)] + #[case(usize::from(u16::MAX) + 2, PType::U32)] + fn selects_narrowest_code_type(#[case] max_len: usize, #[case] expected: PType) { + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len, + }; + let encoder = bytes_dict_builder( + DType::Utf8(Nullability::NonNullable), + &constraints, + BufferAllocatorRef::statically_allocated(), + ); + + assert_eq!(encoder.codes_ptype(), expected); + } + #[test] fn max_dict_bytes_cannot_exceed_the_view_offset_range() { let builder = BytesDictBuilder::::new_in( diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 8e75ea7a903..b2bc39db1e4 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -23,7 +23,9 @@ mod primitive; #[derive(Clone)] pub struct DictConstraints { + /// Limits the encoded storage for dictionary values. pub max_bytes: usize, + /// Limits the number of dictionary values. pub max_len: usize, } @@ -32,6 +34,19 @@ pub const UNCONSTRAINED: DictConstraints = DictConstraints { max_len: usize::MAX, }; +/// Returns the narrowest unsigned code type that can index `max_len` dictionary values. +/// +/// Codes count from zero, so a code type indexes one more value than its maximum: `u8` codes +/// index up to 256 values. +pub const fn codes_ptype_for_max_len(max_len: u64) -> PType { + match max_len { + len if len <= 1 << u8::BITS => PType::U8, + len if len <= 1 << u16::BITS => PType::U16, + len if len <= 1 << u32::BITS => PType::U32, + _ => PType::U64, + } +} + pub trait DictEncoder: Send { /// Assign dictionary codes to the given input array. fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index f1fa752fb92..badd7743984 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -62,3 +62,8 @@ test = false name = "compress_listview" harness = false test = false + +[[bench]] +name = "string_dict_selection" +harness = false +test = false diff --git a/vortex-btrblocks/benches/string_dict_selection.rs b/vortex-btrblocks/benches/string_dict_selection.rs new file mode 100644 index 00000000000..39bcd9aebb4 --- /dev/null +++ b/vortex-btrblocks/benches/string_dict_selection.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(not(codspeed))] +mod benchmarks { + use std::sync::LazyLock; + + use divan::Bencher; + use divan::counter::ItemsCount; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::session::AggregateFnSession; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::memory::MemorySession; + use vortex_array::optimizer::kernels::KernelSession; + use vortex_array::session::ArraySession; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::CompressionSession; + use vortex_error::VortexExpect; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() + .with::() + .with::() + .with::() + }); + + #[derive(Clone, Copy)] + enum Distribution { + Uniform, + Clustered, + VariedPrefix, + } + + #[derive(Clone, Copy)] + struct Case { + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + nullable: bool, + } + + impl std::fmt::Debug for Case { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.name) + } + } + + impl Case { + fn value_index(self, index: usize) -> usize { + match self.distribution { + Distribution::Uniform | Distribution::VariedPrefix => index % self.distinct_values, + Distribution::Clustered => index * self.distinct_values / self.rows, + } + } + + fn value(self, index: usize) -> String { + let value_index = self.value_index(index); + let mut value = match self.distribution { + Distribution::VariedPrefix => format!( + "{:016x}-value-{value_index:08x}", + value_index.wrapping_mul(0x9e3779b9) + ), + _ if self.value_length <= 12 => format!("{value_index:08x}"), + _ => format!("common-prefix-value-{value_index:08x}"), + }; + value.extend(std::iter::repeat_n( + 'x', + self.value_length.saturating_sub(value.len()), + )); + value + } + + fn make_array(self) -> ArrayRef { + let values = (0..self.rows) + .map(|index| self.value(index)) + .collect::>(); + let nullability = if self.nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + + VarBinViewArray::from_iter( + values.iter().enumerate().map(|(index, value)| { + (!self.nullable || index % 10 != 0).then_some(value.as_str()) + }), + DType::Utf8(nullability), + ) + .into_array() + } + } + + const fn case( + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + ) -> Case { + Case { + name, + rows, + distinct_values, + value_length, + distribution, + nullable: false, + } + } + + const CASES: [Case; 9] = [ + case("Inline4096", 65_536, 4096, 8, Distribution::Uniform), + case("Outlined16", 65_536, 16, 28, Distribution::Uniform), + case("Outlined4096", 65_536, 4096, 28, Distribution::Uniform), + case("Outlined8192", 65_536, 8192, 28, Distribution::Uniform), + Case { + nullable: true, + ..case( + "NullableOutlined4096", + 65_536, + 4096, + 28, + Distribution::Uniform, + ) + }, + case("Clustered4096", 65_536, 4096, 28, Distribution::Clustered), + case("Long256", 65_536, 4096, 256, Distribution::Uniform), + case( + "UniqueCommonPrefix", + 65_536, + 65_536, + 28, + Distribution::Uniform, + ), + case( + "UniqueVariedPrefix", + 65_536, + 65_536, + 31, + Distribution::VariedPrefix, + ), + ]; + + #[divan::bench(args = CASES)] + fn compress(bencher: Bencher, case: Case) { + let array = case.make_array(); + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| compressor.compress(array, ctx)); + } + + #[divan::bench(args = CASES)] + fn decompress(bencher: Bencher, case: Case) { + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor + .compress(&case.make_array(), &mut ctx) + .vortex_expect("benchmark input must compress"); + bencher + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + } +} + +fn main() { + divan::main(); +} diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 604018f98ba..976546ddba8 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -9,8 +9,11 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::PType; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -71,22 +74,29 @@ impl Scheme for NullDominatedSparseScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let len = data.array_len() as f64; + let len = data.array_len(); let stats = data.varbinview_stats(exec_ctx); - let value_count = stats.value_count(); + let value_count = stats.value_count() as usize; // All-null arrays should be compressed as constant instead anyways. if value_count == 0 { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // If the majority (90%) of values is null, this will compress well. - if stats.null_count() as f64 / len > 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64)); + // Only null-dominated arrays (90% null) go this route. + if stats.null_count() as f64 / len as f64 <= 0.9 { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // Otherwise we don't go this route. - CompressionEstimate::Verdict(EstimateVerdict::Skip) + // Sparse drops the views of null rows and adds an index per value. It keeps every value's + // bytes, so long values bound the ratio well below the fraction of rows that are null. + let Ok(before) = uncompressed_size_in_bytes(data.array(), exec_ctx) else { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + }; + let null_views = (len - value_count) * size_of::(); + let indices = value_count * PType::min_unsigned_ptype_for_value(len as u64).byte_width(); + let after = before.saturating_sub(null_views) + indices; + CompressionEstimate::Verdict(EstimateVerdict::Ratio(before as f64 / after as f64)) } fn compress( diff --git a/vortex-btrblocks/tests/string_dict_selection.rs b/vortex-btrblocks/tests/string_dict_selection.rs new file mode 100644 index 00000000000..30c748a5b71 --- /dev/null +++ b/vortex-btrblocks/tests/string_dict_selection.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Regression tests for bounded variable-width dictionary selection. + +#![cfg(test)] +#![allow(clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Dict; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::schemes::binary::BinaryDictScheme; +use vortex_btrblocks::schemes::binary::VarBinScheme; +use vortex_btrblocks::schemes::string::FSSTScheme; +use vortex_btrblocks::schemes::string::NullDominatedSparseScheme; +use vortex_btrblocks::schemes::string::StringDictScheme; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +#[derive(Clone, Copy, Debug)] +enum Distribution { + Uniform, + Clustered, +} + +#[rstest] +#[case::outlined_utf8(4096, 28, false, Distribution::Uniform)] +#[case::nullable_utf8(4096, 28, true, Distribution::Uniform)] +#[case::outlined_utf8_8192(8192, 28, false, Distribution::Uniform)] +#[case::clustered_utf8(4096, 28, false, Distribution::Clustered)] +fn test_dict_compressed_with_more_values_than_sample( + #[case] distinct_count: usize, + #[case] value_length: usize, + #[case] nullable: bool, + #[case] distribution: Distribution, +) -> VortexResult<()> { + let distinct_values = (0..distinct_count) + .map(|value| { + let mut string = format!("common-prefix-value-{value:08x}"); + string.extend(std::iter::repeat_n( + 'x', + value_length.saturating_sub(string.len()), + )); + string + }) + .collect::>(); + let values = (0..65_536) + .map(|index| { + let value_index = match distribution { + Distribution::Uniform => index % distinct_values.len(), + Distribution::Clustered => index * distinct_values.len() / 65_536, + }; + (!nullable || index % 10 != 0).then_some(distinct_values[value_index].as_bytes()) + }) + .collect::>(); + let nullability = if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter(values, DType::Utf8(nullability)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_unique_strings_with_common_prefix_not_dict_compressed() -> VortexResult<()> { + let values = (0usize..4096) + .map(|value| Some(format!("common-prefix-value-{value:08x}"))) + .collect::>(); + let array = + VarBinViewArray::from_iter(values, DType::Utf8(Nullability::NonNullable)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + !compressed.is::(), + "expected a non-Dict encoding, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_sample_fallback_can_select_dict() -> VortexResult<()> { + let suffix = "x".repeat((1 << 20) + 1); + let distinct_values = [format!("first-{suffix}"), format!("second-{suffix}")]; + let values = (0..4) + .map(|index| distinct_values[index % distinct_values.len()].as_str()) + .collect::>(); + let array = VarBinViewArray::from_iter_str(values).into_array(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_binary_dict_compressed_with_more_values_than_sample() -> VortexResult<()> { + let distinct_values = (0u32..4096) + .map(|value| { + format!( + "binary-value-{:08x}-{value:08x}", + value.wrapping_mul(0x9e37_79b9) + ) + }) + .collect::>(); + let values = (0..65_536) + .map(|index| Some(distinct_values[index % distinct_values.len()].as_bytes())) + .collect::>(); + let array = + VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&BinaryDictScheme) + .with_new_scheme(&VarBinScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +/// Sparse keeps every value's bytes, so it must not win a null-dominated column of long values +/// on the fraction of null rows alone. +#[test] +fn test_null_dominated_long_values_not_left_uncompressed() -> VortexResult<()> { + let values = (0..65_536) + .map(|index| { + (index % 25 == 0).then(|| { + format!( + "https://example.com/media/{index:08x}/{}", + "photo/".repeat(40) + ) + }) + }) + .collect::>(); + let array = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let value_bytes = values + .iter() + .flatten() + .map(|value| value.len() as u64) + .sum::(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .with_new_scheme(&NullDominatedSparseScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + compressed.nbytes() < value_bytes / 2, + "{} compressed to {} bytes for {value_bytes} bytes of values", + compressed.encoding_id(), + compressed.nbytes() + ); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..83880cff102 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -10,30 +10,21 @@ 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::Dict; -use vortex_array::arrays::DictArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::dict::DictArrayExt; -use vortex_array::arrays::dict::DictArraySlotsExt; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::builders::dict::dict_encode; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use super::bytes_candidate::bytes_dictionary_estimate; +use super::bytes_candidate::compress_bytes_dictionary; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; /// Dictionary encoding for low-cardinality binary values. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -52,12 +43,6 @@ impl Scheme for BinaryDictScheme { vec![Dict.id()] } - fn stats_options(&self) -> GenerateStatsOptions { - GenerateStatsOptions { - count_distinct_values: true, - } - } - /// Children: values=0, codes=1. fn num_children(&self) -> usize { 2 @@ -77,27 +62,11 @@ impl Scheme for BinaryDictScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, + _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let stats = data.varbinview_stats(exec_ctx); - - if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( - "this must be present since `DictScheme` declared that we need distinct values", - ); - - // If > 50% of the values are distinct, skip dictionary scheme. - if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + bytes_dictionary_estimate(&BinaryDictScheme) } fn compress( @@ -107,29 +76,6 @@ impl Scheme for BinaryDictScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let dict = dict_encode(data.array(), exec_ctx)?; - - // Values = child 0. - let compressed_values = - compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?; - - // Codes = child 1. - let narrowed_codes = dict - .codes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_codes = - compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?; - - // SAFETY: compressing codes or values does not alter the invariants. - unsafe { - Ok( - DictArray::new_unchecked(compressed_codes, compressed_values) - .set_all_values_referenced(dict.has_all_values_referenced()) - .into_array(), - ) - } + compress_bytes_dictionary(&BinaryDictScheme, compressor, data, compress_ctx, exec_ctx) } } diff --git a/vortex-compressor/src/builtins/dict/bytes_candidate.rs b/vortex-compressor/src/builtins/dict/bytes_candidate.rs new file mode 100644 index 00000000000..3ad83ec56a4 --- /dev/null +++ b/vortex-compressor/src/builtins/dict/bytes_candidate.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded trial compression for variable-width dictionary selection. +//! +//! Ordinary samples can miss repeated values when a column has thousands of distinct values, so +//! a sample-based estimate under-rates dictionaries for those columns. This module instead +//! builds one bounded dictionary candidate from the complete input. +//! +//! Before paying for that trial, a distributed probe of complete values estimates the number of +//! distinct values in the whole input. A cost model turns that estimate into an optimistic +//! compression ratio, and the trial runs only when that ratio could beat the best estimate so +//! far. A probe that cannot finish within its budget, and a dictionary too large for the trial, +//! fall back to the ordinary sample estimator. +//! +//! The trial's completed output is kept. The compression phase reuses it if the dictionary wins, +//! and the compressor falls back to it if another scheme wins but produces a larger output. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::dict::DictArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::builders::dict::DictConstraints; +use vortex_array::builders::dict::codes_ptype_for_max_len; +use vortex_array::builders::dict::dict_encode; +use vortex_array::builders::dict::dict_encode_with_constraints; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::CascadingCompressor; +use crate::compressor::TrialOutput; +use crate::compressor::estimate_compression_ratio_with_sampling; +use crate::compressor::exact_uncompressed_size; +use crate::compressor::sample_slices; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; + +/// Limits the number of entries so that trial codes remain two bytes wide. +const MAX_DICTIONARY_ENTRIES: usize = 1 << u16::BITS; + +/// Holds the width of a trial dictionary's codes. +const TRIAL_CODE_BYTES: usize = codes_ptype_for_max_len(MAX_DICTIONARY_ENTRIES as u64).byte_width(); + +/// Limits the value and view storage in a trial dictionary. +/// +/// This is larger than the 1 MiB that `DictLayoutConstraints` gives a dictionary shared across a +/// column's chunks: one array's dictionary of 50,000 28-byte values still beats FSST, and a +/// smaller limit would hand it back to sampling, which cannot see that. +const MAX_DICTIONARY_BYTES: usize = 1 << 22; + +/// Limits code storage while the trial scans the complete input. +const MAX_DICTIONARY_CODE_BYTES: usize = 1 << 24; + +/// Limits complete input bytes hashed while constructing a trial dictionary. +const MAX_DICTIONARY_INPUT_BYTES: usize = 1 << 28; + +/// Limits complete value bytes hashed by the preliminary probe. +const MAX_PROBE_BYTES: usize = 1 << 20; + +/// Samples this many adjacent rows from each probe range. +const PROBE_RANGE_SIZE: u32 = 64; + +/// Distributes this many probe ranges across the input. +const PROBE_RANGE_COUNT: u32 = 128; + +/// Samples this many adjacent rows from each range in the head probe that screens out unique +/// inputs before the full probe. +const HEAD_PROBE_RANGE_SIZE: u32 = 16; + +/// Requires a head probe this large before it may reject an incomplete input. +/// +/// A distributed sample of 1024 values without a single repeat means the input holds hundreds of +/// thousands of distinct values, or nearly one per row. +const MIN_HEAD_PROBE_ROWS: usize = 1024; + +/// Bounds how well a string scheme compresses a dictionary's values. +/// +/// Dictionary values are distinct, so a scheme can only shrink them one value at a time. FSST +/// writes a 1-byte code for at most 8 bytes and OnPair a 2-byte code for at most 16, so neither +/// can exceed 8× on them, however well it compresses the repetitive input. +const MAX_VALUES_RATIO: f64 = 8.0; + +/// Returns a deferred estimate that can reuse a completed dictionary result. +pub(super) fn bytes_dictionary_estimate(scheme: &'static dyn Scheme) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + move |compressor, data, best_so_far, compress_ctx, exec_ctx| { + estimate_bytes_dictionary( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ) + }, + ))) +} + +/// Dictionary-encodes the input, reusing the trial result when the estimate cached one. +pub(super) fn compress_bytes_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(trial) = data.get::() + && trial.scheme == scheme.id() + { + return Ok(trial.array.clone()); + } + + let dict = dict_encode(data.array(), exec_ctx)?; + compress_dictionary(scheme, compressor, &dict, compress_ctx, exec_ctx) +} + +/// Compresses the value and code children of a dictionary. +fn compress_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + dict: &DictArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // Values = child 0. + let compressed_values = + compressor.compress_child(dict.values(), &compress_ctx, scheme.id(), 0, exec_ctx)?; + + // Codes = child 1. + let narrowed_codes = dict + .codes() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_codes = + compressor.compress_child(&narrowed_codes, &compress_ctx, scheme.id(), 1, exec_ctx)?; + + // SAFETY: compressing codes or values does not alter the invariants. + unsafe { + Ok( + DictArray::new_unchecked(compressed_codes, compressed_values) + .set_all_values_referenced(dict.has_all_values_referenced()) + .into_array(), + ) + } +} + +/// Resolves the deferred dictionary estimate. +fn estimate_bytes_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let profile = InputProfile::new(data, exec_ctx)?; + let array = data.array_as_varbinview(); + if profile.valid_rows == 0 { + return Ok(EstimateVerdict::Skip); + } + + if ProbeSample::new(&array, &profile.validity, HEAD_PROBE_RANGE_SIZE).is_some_and(|head| { + head.is_unique() && admit_trial(&profile, &head, best_so_far) == Admission::Reject + }) { + return Ok(EstimateVerdict::Skip); + } + + let Some(sample) = ProbeSample::new(&array, &profile.validity, PROBE_RANGE_SIZE) else { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + }; + match admit_trial(&profile, &sample, best_so_far) { + Admission::Trial => {} + Admission::Reject => return Ok(EstimateVerdict::Skip), + Admission::TooLarge => { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + } + } + + let Some(dictionary) = build_candidate(data.array(), exec_ctx)? else { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + }; + let compressed = compress_dictionary(scheme, compressor, &dictionary, compress_ctx, exec_ctx)?; + let score = EstimateScore::from_sample_sizes(profile.uncompressed_nbytes, compressed.nbytes()); + // Keep the trial even when its ratio loses: the compressor prefers it to a winner whose + // actual output turns out larger. + data.get_or_insert_with(|| TrialOutput { + scheme: scheme.id(), + array: compressed, + }); + if !score_is_best(score, best_so_far) { + return Ok(EstimateVerdict::Skip); + } + Ok(score_verdict(score)) +} + +/// Runs the ordinary sample estimate and applies the current threshold. +fn sample_verdict( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let score = estimate_compression_ratio_with_sampling( + compressor, + scheme, + data.array(), + compress_ctx, + exec_ctx, + )?; + if score_is_best(score, best_so_far) { + Ok(score_verdict(score)) + } else { + Ok(EstimateVerdict::Skip) + } +} + +/// Returns whether a score is valid and beats the current best score. +fn score_is_best(score: EstimateScore, best_so_far: Option) -> bool { + score.is_valid() && best_so_far.is_none_or(|best| score.beats(best)) +} + +/// Converts a score into a terminal estimate. +fn score_verdict(score: EstimateScore) -> EstimateVerdict { + match score { + EstimateScore::FiniteCompression(ratio) => EstimateVerdict::Ratio(ratio), + EstimateScore::ZeroBytes => EstimateVerdict::Skip, + } +} + +/// Describes the complete input with one pass over its views. +struct InputProfile { + /// Counts all rows, including nulls. + rows: usize, + /// Counts non-null rows. + valid_rows: usize, + /// Sums the lengths of non-null values. + value_bytes: usize, + /// Measures the bytes the rows reference, as every other estimate measures its input. + uncompressed_nbytes: u64, + /// Marks non-null rows. + validity: Mask, +} + +impl InputProfile { + /// Profiles the complete input. + fn new(data: &ArrayAndStats, exec_ctx: &mut ExecutionCtx) -> VortexResult { + let array = data.array_as_varbinview(); + let validity = array.validity()?.execute_mask(array.len(), exec_ctx)?; + let views = array.views(); + let value_bytes = match validity.bit_buffer() { + AllOr::All => views.iter().map(|view| view.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => bits + .set_indices() + .map(|index| views[index].len() as usize) + .sum(), + }; + + Ok(Self { + rows: array.len(), + valid_rows: validity.true_count(), + value_bytes, + uncompressed_nbytes: exact_uncompressed_size(data.array(), exec_ctx)?, + validity, + }) + } + + /// Returns the mean length of a non-null value. + fn mean_value_bytes(&self) -> f64 { + self.value_bytes as f64 / self.valid_rows as f64 + } +} + +/// Summarizes complete values from distributed ranges of the input. +#[derive(Debug, PartialEq)] +struct ProbeSample { + /// Counts probed non-null rows. + rows: usize, + /// Counts distinct probed non-null values. + distinct: usize, + /// Counts distinct values that occur exactly once in the probe. + singletons: usize, + /// Counts distinct values that occur exactly twice in the probe. + doubletons: usize, + /// Counts adjacent row pairs within a probe range. + adjacent_pairs: usize, + /// Counts adjacent row pairs within a probe range that hold different values. + adjacent_changes: usize, + /// Reports whether the probe covered every row of the input. + complete: bool, +} + +impl ProbeSample { + /// Probes `range_size` adjacent rows from each range, or returns `None` when the probed values + /// exceed the probe budget. + fn new(array: &ArrayView<'_, VarBinView>, validity: &Mask, range_size: u32) -> Option { + let views = array.views(); + let ranges = sample_slices(array.len(), range_size, PROBE_RANGE_COUNT); + let complete = matches!(ranges.as_slice(), [(0, end)] if *end == array.len()); + + let probe_bytes = ranges + .iter() + .flat_map(|&(start, end)| start..end) + .filter(|&index| validity.value(index)) + .map(|index| views[index].len() as usize) + .sum::(); + if probe_bytes > MAX_PROBE_BYTES { + return None; + } + + let mut counts = HashMap::<&[u8], u32>::new(); + let mut adjacent_pairs = 0; + let mut adjacent_changes = 0; + for (start, end) in ranges { + let mut previous = None; + for index in start..end { + let value = validity + .value(index) + .then(|| view_bytes(array, &views[index])); + if let Some(previous) = previous { + adjacent_pairs += 1; + adjacent_changes += usize::from(previous != value); + } + previous = Some(value); + if let Some(value) = value { + *counts.entry(value).or_default() += 1; + } + } + } + + let rows = counts.values().map(|&count| count as usize).sum::(); + if rows == 0 { + return None; + } + Some(Self { + rows, + distinct: counts.len(), + singletons: counts.values().filter(|&&count| count == 1).count(), + doubletons: counts.values().filter(|&&count| count == 2).count(), + adjacent_pairs, + adjacent_changes, + complete, + }) + } + + /// Reports whether every probed value is distinct, in a probe large enough to tell. + fn is_unique(&self) -> bool { + self.distinct == self.rows && (self.complete || self.rows >= MIN_HEAD_PROBE_ROWS) + } + + /// Estimates the number of distinct non-null values in the complete input. + /// + /// Uses the Chao1 estimator, which tends to under-count. An under-count only costs a trial + /// that the complete dictionary then rejects, while an over-count would skip a dictionary + /// that wins. + fn estimated_distinct(&self, valid_rows: usize) -> f64 { + let distinct = self.distinct as f64; + if self.complete { + return distinct; + } + let singletons = self.singletons as f64; + let doubletons = self.doubletons as f64; + let unseen = if self.doubletons > 0 { + singletons * singletons / (2.0 * doubletons) + } else { + singletons * (singletons - 1.0) / 2.0 + }; + (distinct + unseen).min(valid_rows as f64) + } + + /// Estimates the fraction of rows that start a new run of equal values. + fn run_start_fraction(&self) -> f64 { + if self.adjacent_pairs == 0 { + return 1.0; + } + (self.adjacent_changes as f64 / self.adjacent_pairs as f64) + .max(1.0 / f64::from(PROBE_RANGE_SIZE)) + } +} + +/// Decides whether to build a trial dictionary. +#[derive(Debug, PartialEq, Eq)] +enum Admission { + /// Builds the trial dictionary. + Trial, + /// Rejects Dictionary because it cannot beat the best estimate so far. + Reject, + /// Leaves the decision to sampling because the trial would exceed its bounds. + /// + /// A dictionary larger than the trial allows can still win with wider codes. + TooLarge, +} + +/// Applies the trial bounds and the cost model. +fn admit_trial( + profile: &InputProfile, + sample: &ProbeSample, + best_so_far: Option, +) -> Admission { + let best_ratio = match best_so_far { + Some(EstimateScore::FiniteCompression(ratio)) => ratio.max(1.0), + Some(EstimateScore::ZeroBytes) => return Admission::Reject, + None => 1.0, + }; + let distinct = sample.estimated_distinct(profile.valid_rows); + let has_nulls = profile.valid_rows < profile.rows; + let entries = distinct + f64::from(u8::from(has_nulls)); + let values_ratio = best_ratio.min(MAX_VALUES_RATIO); + if optimistic_ratio(profile, sample, entries, values_ratio) <= best_ratio { + return Admission::Reject; + } + + if entries > MAX_DICTIONARY_ENTRIES as f64 + || distinct * profile.mean_value_bytes() > MAX_DICTIONARY_BYTES as f64 + || profile.value_bytes > MAX_DICTIONARY_INPUT_BYTES + || profile.rows.saturating_mul(TRIAL_CODE_BYTES) > MAX_DICTIONARY_CODE_BYTES + { + return Admission::TooLarge; + } + Admission::Trial +} + +/// Estimates the compression ratio of a dictionary with `entries` values. +/// +/// The model assumes dictionary values compress by `values_ratio`, and that codes compress to their +/// bit width over the runs the probe saw. +fn optimistic_ratio( + profile: &InputProfile, + sample: &ProbeSample, + entries: f64, + values_ratio: f64, +) -> f64 { + let values_bytes = entries * (profile.mean_value_bytes() + size_of::() as f64); + let code_bits = entries.max(2.0).log2().ceil(); + let code_bytes = profile.rows as f64 * sample.run_start_fraction() * code_bits / 8.0; + profile.uncompressed_nbytes as f64 / (values_bytes / values_ratio + code_bytes).max(1.0) +} + +/// Builds a dictionary within the trial storage bounds, or returns `None` if it exceeds them. +fn build_candidate( + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult> { + let constraints = DictConstraints { + max_bytes: MAX_DICTIONARY_BYTES, + max_len: MAX_DICTIONARY_ENTRIES, + }; + let dictionary = dict_encode_with_constraints(array, &constraints, exec_ctx)?; + Ok((dictionary.len() == array.len()).then_some(dictionary)) +} + +/// Returns the complete value referenced by a binary view. +fn view_bytes<'a>(array: &'a ArrayView<'a, VarBinView>, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let reference = view.as_view(); + &array.buffer(reference.buffer_index as usize)[reference.as_range()] + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs b/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs new file mode 100644 index 00000000000..07120149ae0 --- /dev/null +++ b/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs @@ -0,0 +1,305 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::Admission; +use super::HEAD_PROBE_RANGE_SIZE; +use super::InputProfile; +use super::MAX_DICTIONARY_ENTRIES; +use super::MAX_PROBE_BYTES; +use super::PROBE_RANGE_SIZE; +use super::ProbeSample; +use super::admit_trial; +use super::build_candidate; +use super::bytes_dictionary_estimate; +use crate::CascadingCompressor; +use crate::builtins::StringDictScheme; +use crate::compressor::TrialOutput; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::stats::ArrayAndStats; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// Builds compression data for candidate tests. +fn string_data(values: &[Option]) -> ArrayAndStats { + let nullability = if values.iter().any(Option::is_none) { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(nullability), + ) + .into_array(); + ArrayAndStats::new(array, Default::default()) +} + +/// Returns `rows` values that cycle through `distinct` outlined strings. +fn cycled(rows: usize, distinct: usize) -> Vec> { + (0..rows) + .map(|index| { + let value = index % distinct; + Some(format!( + "{:016x}-value-{value:08x}", + value.wrapping_mul(0x9e37_79b9_7f4a_7c15) + )) + }) + .collect() +} + +/// Profiles and probes the input. +fn profile_and_probe( + data: &ArrayAndStats, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult<(InputProfile, Option)> { + let profile = InputProfile::new(data, exec_ctx)?; + let sample = ProbeSample::new( + &data.array_as_varbinview(), + &profile.validity, + PROBE_RANGE_SIZE, + ); + Ok((profile, sample)) +} + +/// Resolves the Dictionary estimate against an optional threshold. +fn estimate_dictionary( + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let callback = match bytes_dictionary_estimate(&StringDictScheme) { + CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, + estimate => return Err(vortex_err!("dictionary estimation returned {estimate:?}")), + }; + callback( + compressor, + data, + best_so_far, + CompressorContext::new(), + exec_ctx, + ) +} + +#[test] +fn complete_probe_counts_distinct_exactly() -> VortexResult<()> { + let data = string_data(&cycled(4096, 100)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert!(sample.complete); + assert_eq!(sample.estimated_distinct(profile.valid_rows), 100.0); + Ok(()) +} + +#[rstest] +#[case(20_000)] +#[case(50_000)] +fn partial_probe_estimates_distinct_beyond_probe_size(#[case] distinct: usize) -> VortexResult<()> { + let data = string_data(&cycled(1 << 20, distinct)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + let estimate = sample.estimated_distinct(profile.valid_rows); + + assert!(!sample.complete); + // The probe sees at most 8192 rows, so its raw distinct count cannot reach `distinct`. + assert!(sample.distinct < distinct); + let error = (estimate - distinct as f64).abs() / distinct as f64; + assert!(error < 0.25, "estimated {estimate} for {distinct}"); + Ok(()) +} + +#[test] +fn probe_is_inconclusive_above_byte_budget() -> VortexResult<()> { + let data = string_data(&[Some("a".to_owned()), Some("x".repeat(MAX_PROBE_BYTES))]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (_, sample) = profile_and_probe(&data, &mut exec_ctx)?; + + assert_eq!(sample, None); + Ok(()) +} + +#[test] +fn nulls_count_as_one_dictionary_entry() -> VortexResult<()> { + let values = (0..8192) + .map(|index| (index % 10 == 0).then(|| format!("unique-non-null-value-{index:08x}"))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(profile.valid_rows, 820); + assert_eq!(sample.rows, 820); + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + Ok(()) +} + +#[test] +fn unique_values_are_not_admitted() -> VortexResult<()> { + let data = string_data(&cycled(8192, 8192)); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Reject); + assert!(matches!( + estimate_dictionary(&compressor, &data, None, &mut exec_ctx)?, + EstimateVerdict::Skip + )); + Ok(()) +} + +#[test] +fn threshold_rejects_trial_before_building() -> VortexResult<()> { + let data = string_data(&cycled(65_536, 4096)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + assert_eq!( + admit_trial( + &profile, + &sample, + Some(EstimateScore::FiniteCompression(1000.0)) + ), + Admission::Reject + ); + Ok(()) +} + +#[test] +fn too_many_estimated_entries_defer_to_sampling() -> VortexResult<()> { + // Each value repeats 16 times, so a dictionary wins even though it exceeds the trial bounds. + let distinct = 2 * MAX_DICTIONARY_ENTRIES; + let values = (0..16 * distinct) + .map(|index| Some(format!("{:016x}", (index * 7919) % distinct))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::TooLarge); + Ok(()) +} + +#[test] +fn candidate_stops_when_dictionary_storage_exceeds_limit() -> VortexResult<()> { + let suffix = "x".repeat(70_000); + let values = (0..128) + .map(|index| Some(format!("{:08x}-{suffix}", index % 64))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + + assert!(build_candidate(data.array(), &mut exec_ctx)?.is_none()); + Ok(()) +} + +#[test] +fn candidate_uses_full_u16_code_range() -> VortexResult<()> { + let mut values = (0..=u16::MAX) + .map(|value| Some(format!("{value:08x}"))) + .collect::>(); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert!(build_candidate(string_data(&values).array(), &mut exec_ctx)?.is_some()); + + values.push(Some("past-u16-code-range".to_owned())); + assert!(build_candidate(string_data(&values).array(), &mut exec_ctx)?.is_none()); + Ok(()) +} + +#[test] +fn completed_candidate_is_reused() -> VortexResult<()> { + let data = string_data(&cycled(65_536, 4096)); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let verdict = estimate_dictionary(&compressor, &data, None, &mut exec_ctx)?; + assert!(matches!(verdict, EstimateVerdict::Ratio(_))); + let cached = data + .get::() + .ok_or_else(|| vortex_err!("completed candidate was not cached"))?; + let compressed = + StringDictScheme.compress(&compressor, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(ArrayRef::ptr_eq(&cached.array, &compressed)); + Ok(()) +} + +/// Repeats far apart still reach the head probe: each value below recurs half the input away. +#[rstest] +#[case::unique(65_536, 65_536, true)] +#[case::two_repeats_per_value(65_536, 32_768, false)] +#[case::four_repeats_per_value(65_536, 16_384, false)] +#[case::two_repeats_in_a_shorter_input(16_384, 8_192, false)] +fn head_probe_screens_unique_inputs( + #[case] rows: usize, + #[case] distinct: usize, + #[case] expected: bool, +) -> VortexResult<()> { + let data = string_data(&cycled(rows, distinct)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let profile = InputProfile::new(&data, &mut exec_ctx)?; + let head = ProbeSample::new( + &data.array_as_varbinview(), + &profile.validity, + HEAD_PROBE_RANGE_SIZE, + ) + .ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert!(!head.complete); + assert_eq!(head.is_unique(), expected); + if expected { + assert_eq!(admit_trial(&profile, &head, None), Admission::Reject); + } + Ok(()) +} + +/// A mostly-null column of distinct long values compresses far better as sparse values than as a +/// dictionary, whose distinct values no string scheme compresses more than eight-fold. +#[test] +fn values_ratio_bound_rejects_sparse_unique_values() -> VortexResult<()> { + let suffix = "x".repeat(240); + let values = (0..40_960) + .map(|index| (index % 23 == 0).then(|| format!("{index:016x}{suffix}"))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + assert_eq!( + admit_trial( + &profile, + &sample, + Some(EstimateScore::FiniteCompression(22.0)) + ), + Admission::Reject + ); + Ok(()) +} diff --git a/vortex-compressor/src/builtins/dict/mod.rs b/vortex-compressor/src/builtins/dict/mod.rs index 4862df2b211..a49b9ea1a22 100644 --- a/vortex-compressor/src/builtins/dict/mod.rs +++ b/vortex-compressor/src/builtins/dict/mod.rs @@ -4,6 +4,7 @@ //! Dictionary encoding schemes for binary, integer, float, and string arrays. mod binary; +mod bytes_candidate; mod float; mod integer; mod string; diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..d339b60a9d6 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -10,30 +10,21 @@ 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::Dict; -use vortex_array::arrays::DictArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::dict::DictArrayExt; -use vortex_array::arrays::dict::DictArraySlotsExt; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::builders::dict::dict_encode; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use super::bytes_candidate::bytes_dictionary_estimate; +use super::bytes_candidate::compress_bytes_dictionary; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; /// Dictionary encoding for low-cardinality string values. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -52,12 +43,6 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } - fn stats_options(&self) -> GenerateStatsOptions { - GenerateStatsOptions { - count_distinct_values: true, - } - } - /// Children: values=0, codes=1. fn num_children(&self) -> usize { 2 @@ -77,27 +62,11 @@ impl Scheme for StringDictScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, + _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let stats = data.varbinview_stats(exec_ctx); - - if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( - "this must be present since `DictScheme` declared that we need distinct values", - ); - - // If > 50% of the values are distinct, skip dictionary scheme. - if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + bytes_dictionary_estimate(&StringDictScheme) } fn compress( @@ -107,29 +76,6 @@ impl Scheme for StringDictScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let dict = dict_encode(data.array(), exec_ctx)?; - - // Values = child 0. - let compressed_values = - compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?; - - // Codes = child 1. - let narrowed_codes = dict - .codes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_codes = - compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?; - - // SAFETY: compressing codes or values does not alter the invariants. - unsafe { - Ok( - DictArray::new_unchecked(compressed_codes, compressed_values) - .set_all_values_referenced(dict.has_all_values_referenced()) - .into_array(), - ) - } + compress_bytes_dictionary(&StringDictScheme, compressor, data, compress_ctx, exec_ctx) } } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index de4599a9cc7..8686ad411b4 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -43,6 +43,19 @@ use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; use crate::trace; +/// A complete compressed output that a scheme produced while estimating its ratio. +/// +/// An estimate callback may store one in its [`ArrayAndStats`]. Its size is a measurement, while +/// the winning scheme was picked on an estimate, so the compressor keeps the trial output +/// whenever it is smaller than what the winner actually produced. +#[derive(Debug)] +pub(crate) struct TrialOutput { + /// The scheme that produced the output. + pub(crate) scheme: SchemeId, + /// The compressed array. + pub(crate) array: ArrayRef, +} + /// The outcome of [`CascadingCompressor::choose_and_compress`]. enum Selection { /// A scheme, or the compressor's own empty, null or constant handling, produced a new array. @@ -357,6 +370,12 @@ impl CascadingCompressor { trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); })?; + // The winner was picked on estimates. A trial output is measured, so it wins when smaller. + let compressed = match data.get::() { + Some(trial) if trial.array.nbytes() < compressed.nbytes() => trial.array.clone(), + _ => compressed, + }; + let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); @@ -380,7 +399,7 @@ impl CascadingCompressor { /// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches /// on each chunk and computed otherwise. A dtype that cannot be measured, such as variant, is /// sized by the bytes it holds. -pub(super) fn exact_uncompressed_size( +pub(crate) fn exact_uncompressed_size( array: &ArrayRef, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..8b2117e7b83 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,11 @@ mod sample; mod select; mod structural; +pub(crate) use cascade::TrialOutput; +pub(crate) use cascade::exact_uncompressed_size; +pub(crate) use sample::estimate_compression_ratio_with_sampling; +pub(crate) use sample::sample_slices; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index dcdce0dec1f..714af01e49a 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -43,12 +43,7 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A return input.clone(); } - let slices = stratified_slices( - input.len(), - sample_size, - sample_count, - &mut StdRng::seed_from_u64(SAMPLE_SEED), - ); + let slices = sample_slices(input.len(), sample_size, sample_count); // For every slice, grab the relevant slice and repack into a new PrimitiveArray. let chunks = slices.into_iter().map(|(start, end)| { @@ -60,6 +55,20 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A unsafe { ChunkedArray::new_unchecked(chunks, input.dtype().clone()) }.into_array() } +/// Returns deterministic stratified sample ranges for an array length. +pub(crate) fn sample_slices( + length: usize, + sample_size: u32, + sample_count: u32, +) -> Vec<(usize, usize)> { + stratified_slices( + length, + sample_size, + sample_count, + &mut StdRng::seed_from_u64(SAMPLE_SEED), + ) +} + /// Computes the number of sample chunks to cover approximately 1% of `len` elements, /// with a minimum of `SAMPLE_SIZE * SAMPLE_COUNT` (1024) values. pub(crate) fn sample_count_approx_one_percent(len: usize) -> u32 { @@ -147,7 +156,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> /// # Errors /// /// Returns an error if sample compression fails. -pub(super) fn estimate_compression_ratio_with_sampling( +pub(crate) fn estimate_compression_ratio_with_sampling( compressor: &CascadingCompressor, scheme: &S, array: &ArrayRef, @@ -163,6 +172,12 @@ pub(super) fn estimate_compression_ratio_with_sampling( canonical.into_array() }; + // A sample of only nulls compresses to almost nothing whatever the scheme, so its ratio says + // nothing about how the scheme handles the array's values. + if !compress_ctx.is_sample() && sample_array.all_invalid(exec_ctx)? { + return Ok(EstimateScore::ZeroBytes); + } + let sample_data = ArrayAndStats::new(sample_array, scheme.stats_options()); let error_ctx = trace::enabled_error_context(&compress_ctx); let sample_ctx = compress_ctx.with_sampling(); diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..4afdcf48eea 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -55,11 +55,13 @@ impl CascadingCompressor { /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any /// expensive computations if we don't have to. /// - /// Pass 2 evaluates the deferred work and, for each [`DeferredEstimate::Callback`], passes the - /// current best [`EstimateScore`] as an early-exit hint so the callback can return - /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. + /// Pass 2 evaluates the deferred work. Every [`DeferredEstimate::Sample`] runs before any + /// [`DeferredEstimate::Callback`], and each callback receives the current best + /// [`EstimateScore`] as an early-exit hint so it can return [`EstimateVerdict::Skip`] without + /// doing expensive work when it cannot beat the threshold. /// - /// Ties are broken by registration order within each pass. + /// Ties are broken by evaluation order: immediate verdicts, then samples, then callbacks, each + /// in registration order. /// /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio pub(super) fn choose_best_scheme( @@ -96,7 +98,8 @@ impl CascadingCompressor { } // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can - // short-circuit with `Skip` when they cannot beat it. + // short-circuit with `Skip` when they cannot beat it, so samples run first to tighten it. + deferred.sort_by_key(|(_, estimate)| matches!(estimate, DeferredEstimate::Callback(_))); for (scheme, deferred_estimate) in deferred { let _span = trace::scheme_eval_span(scheme.id()).entered(); let threshold: Option = best.map(|(_, score)| score); diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 29e32b433c7..190b36ee520 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -15,6 +15,7 @@ use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_si use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; @@ -28,6 +29,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; +use vortex_array::builders::dict::dict_encode; use vortex_array::dtype::DType; use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; @@ -47,6 +49,7 @@ use vortex_session::VortexSession; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; +use super::TrialOutput; use super::sample::estimate_compression_ratio_with_sampling; use super::select::WinnerEstimate; use super::structural; @@ -336,6 +339,126 @@ impl Scheme for ZeroBytesSamplingScheme { } } +/// Samples to a finite ratio by writing a single byte. +#[derive(Debug)] +struct FiniteSamplingScheme; + +impl Scheme for FiniteSamplingScheme { + fn scheme_name(&self) -> &'static str { + "test.finite_sampling" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(PrimitiveArray::new(buffer![0u8], Validity::NonNullable).into_array()) + } +} + +/// Promises a ratio it does not deliver: its output is the input. +#[derive(Debug)] +struct OverpromisingScheme; + +impl Scheme for OverpromisingScheme { + fn scheme_name(&self) -> &'static str { + "test.overpromising" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Ratio(100.0)) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +/// Stores a dictionary as its trial output, then reports that it cannot beat the threshold. +#[derive(Debug)] +struct TrialStoringScheme; + +impl Scheme for TrialStoringScheme { + fn scheme_name(&self) -> &'static str { + "test.trial_storing" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, data, _best_so_far, _ctx, exec_ctx| { + let array = dict_encode(data.array(), exec_ctx)?.into_array(); + data.get_or_insert_with(|| TrialOutput { + scheme: TrialStoringScheme.id(), + array, + }); + Ok(EstimateVerdict::Skip) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + #[test] fn test_self_exclusion() { let c = compressor(); @@ -675,6 +798,61 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { Ok(()) } +#[test] +fn callbacks_observe_samples_registered_after_them() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + let compressor = + CascadingCompressor::new(vec![&ThresholdObservingScheme, &FiniteSamplingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&ThresholdObservingScheme, &FiniteSamplingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + let observed = *OBSERVED_THRESHOLD.lock(); + assert!(matches!( + observed, + Some(Some(EstimateScore::FiniteCompression(_))) + )); + Ok(()) +} + +#[test] +fn smaller_trial_output_replaces_an_overpromising_winner() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&OverpromisingScheme, &TrialStoringScheme]); + let array = PrimitiveArray::from_iter((0..1024i64).map(|value| value % 4)).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor.compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn all_null_sample_gives_no_estimate() -> VortexResult<()> { + // Only the first row is valid, and the sampler does not land on it. + let array = + PrimitiveArray::from_option_iter((0..100_000).map(|index| (index == 0).then_some(index))) + .into_array(); + let compressor = CascadingCompressor::new(vec![&FiniteSamplingScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let score = estimate_compression_ratio_with_sampling( + &compressor, + &FiniteSamplingScheme, + &array, + CompressorContext::new(), + &mut exec_ctx, + )?; + + assert_eq!(score, EstimateScore::ZeroBytes); + Ok(()) +} + #[test] fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { // Both schemes produce the same `Ratio(2.0)`, one from pass 1 (immediate) and one from diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/estimate.rs index 20b9dc77e9b..13fe744816d 100644 --- a/vortex-compressor/src/scheme/estimate.rs +++ b/vortex-compressor/src/scheme/estimate.rs @@ -101,7 +101,7 @@ pub enum DeferredEstimate { pub enum EstimateScore { /// A finite compression ratio. Higher means a smaller amount of data, so it is better. FiniteCompression(f64), - /// Trial compression produced a 0-byte output. + /// Trial compression produced a 0-byte output, or sampled only null values. /// /// This has no finite ratio and is not eligible for scheme selection. /// diff --git a/vortex-compressor/src/stats/cache.rs b/vortex-compressor/src/stats/cache.rs index 2652403d3e6..07d0bf1cca1 100644 --- a/vortex-compressor/src/stats/cache.rs +++ b/vortex-compressor/src/stats/cache.rs @@ -67,6 +67,19 @@ impl StatsCache { new_arc } } + + /// Returns a cached value when one exists. + fn get(&self) -> Option> { + let type_id = TypeId::of::(); + let guard = self.entries.lock(); + let position = guard.iter().position(|(id, _)| *id == type_id)?; + Some( + Arc::clone(&guard[position].1) + .downcast::() + .ok() + .vortex_expect("we just checked the TypeID"), + ) + } } /// An array bundled with its lazily-computed statistics cache. @@ -207,4 +220,9 @@ impl ArrayAndStats { pub fn get_or_insert_with(&self, f: impl FnOnce() -> T) -> Arc { self.cache.get_or_insert_with::(f) } + + /// Returns a custom cached value when one exists. + pub(crate) fn get(&self) -> Option> { + self.cache.get::() + } } diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index b82e301a849..f40b790604c 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -73,7 +73,8 @@ pub struct DictLayoutConstraints { /// /// The codes dtype is chosen upfront and is at most [`PType::U16`]. Primitive dictionaries /// use [`PType::U8`] when `max_len <= 256` or the value domain fits in 256 entries including - /// any null entry. String and binary dictionaries use [`PType::U8`] when `max_len <= 255`. + /// any null entry. String and binary dictionaries use [`PType::U8`] when `max_len <= 256`, as + /// [`codes_ptype_for_max_len`](vortex_array::builders::dict::codes_ptype_for_max_len) picks. /// /// Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility. pub max_len: u16, From cb613ac1c0adc56cbe63dfc79e4bbc431ad0bc52 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 10 Nov 2026 23:14:13 -0500 Subject: [PATCH 13/13] Preserve DuckDB dictionaries across run-end chunks (#11378) ## Summary Follow-up to #11375 and [the DuckDB exporter review](https://github.com/vortex-data/vortex/pull/11360#discussion_r4526272214). Chunked run-end inputs currently lose their cached dictionary when an output batch crosses a chunk boundary. Exposing the existing batch codes through `DictionaryParts` lets those ranges reuse the values dictionary, including ranges shared with ordinary dictionary chunks. ## Changes Keep the run-end specialization because routing through `to_dict()` would make the generic exporter materialize the entire codes array. The existing fill loop also remains: using `slice_relative_ends` increased standalone single-row-run export time in all five trial pairs (median paired ratio 1.064). Empty inputs now avoid run lookup and values conversion, and zero-length requests avoid scalar access. Constant batches include requests ending exactly at a run boundary. Compact inline regressions cover empty requests, slice offsets, exact run boundaries, nulls, successive batches, and shared or different dictionaries across chunks. The existing null and chunked-values tests remain. Measurements on `702fd58dea`, before the test and module cleanup, showed about 15% lower batch-export time for single-row runs and 16% for long runs. Rust allocations per export loop fall from 146 to 112 and from 212 to 188, respectively. Standalone export and construction timings show no clear material change. Native DuckDB allocations were not measured.
Validation - After cleanup, `cargo nextest run -p vortex-duckdb -E 'test(exporter::run_end::)'`: 7 passed, 265 outside the filter. - Before cleanup, the full exporter filter passed all 80 tests. - `cargo clippy -p vortex-duckdb --all-targets --all-features -- -D warnings`: passed. - `cargo +nightly-2026-09-10 fmt -p vortex-duckdb`: passed. - Commit hook: workspace Rust and TOML formatting passed. - Pre-push hook: DCO checks and workspace Clippy with all targets, all features, and warnings denied passed. - Broader runtime suites and doctests were not run.
Benchmark methodology and full results Baseline: `b367941e9d490b4a19ef9bdd6653f63d712c581f` with the same benchmark harness. Candidate: `702fd58dea3c0eeb214d590b3cc6d6b5b3c89c9b`. This comparison isolates the follow-up from the Arrow extension/property stack. The [measurement harness](https://github.com/vortex-data/vortex/blob/702fd58dea3c0eeb214d590b3cc6d6b5b3c89c9b/vortex-duckdb/benches/run_end_export.rs) and its Cargo target remain available in that measured commit. The permanent benchmark target was removed during review, along with duplicate test setup and the module split. The production algorithm is unchanged, but these measurements were not rerun after the file move. macOS 27.0.1, aarch64, rustc 1.99.1, LLVM 23.1.1, Cargo bench profile, and DuckDB 2.0.0. Build both revisions with `cargo bench -p vortex-duckdb --bench run_end_export --no-run`. Run their saved executables with `--bench --sample-count 300 --sample-size 20 --color never` in five pairs, alternating baseline/candidate order. The table shows the median of each executable's five run medians, plus the median of the five paired candidate/baseline ratios. These two aggregations need not agree exactly. Inputs are non-null i64 run values with run lengths 1 or 8192. Short-run sources have 16,383 runs, and long-run sources have eight. Chunked inputs repeat the same source four times. Every input is sliced from row 17 to 13 rows before its end, so exports span multiple batches and chunk boundaries fall inside output batches. Fixtures execute the slice to its encoded representation before exporter construction. Construction uses a fresh conversion cache per iteration. Batch export uses a warm cache and excludes exporter construction and data-chunk creation. `AllocProfiler` counts Rust allocations in the measured region, including its timing overhead. It does not count native DuckDB allocations. Neither `xctrace` nor `heaptrack` is available on this machine. Allocation figures below are allocation counts and requested bytes, not peak resident memory. | Phase | Input | Run length | Baseline (us) | Candidate (us) | Paired ratio | All five paired ratios | | --- | --- | ---: | ---: | ---: | ---: | --- | | batches | standalone | 1 | 14.02 | 13.51 | 0.961 | 0.961, 0.954, 0.968, 0.949, 1.028 | | batches | standalone | 8192 | 3.349 | 3.293 | 0.980 | 0.866, 1.007, 1.015, 0.926, 0.980 | | batches | shared chunks | 1 | 48.3 | 41.15 | 0.851 | 0.849, 0.851, 0.858, 0.852, 0.850 | | batches | shared chunks | 8192 | 19.74 | 16.57 | 0.839 | 0.839, 0.839, 0.833, 0.864, 0.832 | | construct | standalone | 1 | 0.8631 | 0.8464 | 0.981 | 0.951, 0.979, 0.984, 0.981, 1.016 | | construct | standalone | 8192 | 0.7922 | 0.7714 | 0.981 | 0.934, 0.982, 0.981, 0.966, 1.004 | | construct | shared chunks | 1 | 51.69 | 51.68 | 1.000 | 1.012, 0.999, 0.997, 1.000, 1.000 | | construct | shared chunks | 8192 | 12.9 | 12.88 | 0.998 | 1.002, 0.997, 0.998, 0.998, 0.998 | | Phase | Input | Run length | Baseline Rust allocations | Candidate Rust allocations | Baseline bytes | Candidate bytes | | --- | --- | ---: | ---: | ---: | ---: | ---: | | Batches | Standalone | 1 | 0 | 0 | 0 | 0 | | Batches | Standalone | 8192 | 0 | 0 | 0 | 0 | | Batches | Shared chunks | 1 | 146 | 112 | 193.1 KB | 130.5 KB | | Batches | Shared chunks | 8192 | 212 | 188 | 200.2 KB | 139.3 KB | | Construct | Standalone | 1 | 5 | 5 | 308 B | 308 B | | Construct | Standalone | 8192 | 5 | 5 | 308 B | 308 B | | Construct | Shared chunks | 1 | 32 | 32 | 265.8 KB | 265.8 KB | | Construct | Shared chunks | 8192 | 22 | 22 | 1.05 MB | 1.05 MB | The shared-chunk batch timing improved in every pair. Standalone controls and construction moved much less, and the standalone long-run ratios varied in direction. No compiler-output analysis was performed, so those small movements are not attributed to a particular code-generation change. These microbenchmarks do not establish an end-to-end query speedup.
--------- Signed-off-by: Connor Tsui --- vortex-duckdb/src/exporter/mod.rs | 8 +- vortex-duckdb/src/exporter/run_end.rs | 271 ++++++++++++++++++++------ 2 files changed, 214 insertions(+), 65 deletions(-) diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index 59a0e7fe728..fe0333ef1cf 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -157,7 +157,7 @@ impl ArrayExporter { /// the offset, len and `WritableVector` as options. Not sure what it should return though? /// This would allow Vortex extension authors to plug into the DuckDB exporter system. pub trait ColumnExporter: 'static { - /// The dictionary parts of this exporter, if it exports DuckDB dictionary vectors. + /// The dictionary parts of this exporter, if it can export DuckDB dictionary vectors. /// /// A chunked exporter uses them to export a range across chunks as one dictionary vector. fn as_dictionary(&self) -> Option<&dyn DictionaryParts> { @@ -175,6 +175,9 @@ pub trait ColumnExporter: 'static { } /// The values and codes of an exporter that emits DuckDB dictionary vectors. +/// +/// Codes can be generated per request rather than stored for the whole array. This capability also +/// applies when the exporter uses a constant vector for a request contained in one run. pub trait DictionaryParts { /// The DuckDB vector of the dictionary values. fn values(&self) -> &ReusableDict; @@ -183,6 +186,9 @@ pub trait DictionaryParts { fn values_key(&self) -> usize; /// Write the codes of rows `offset..offset + codes.len()` into `codes`. + /// + /// The offset is relative to the exported array, including any slice. The range must lie within + /// that array, and every code must index [`Self::values`], including codes for null rows. fn write_codes(&self, offset: usize, codes: &mut [u32]); } diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index fbe0120850a..d455dc3cc1e 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -1,15 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Batch-sized run selections over a cached DuckDB values dictionary. +//! +//! [`DictionaryParts`] lets the chunked exporter join run-end and dictionary chunks that share +//! values without flattening them. Codes are expanded only for the requested batch. Routing through +//! [`RunEndArray::to_dict`] would make the generic dictionary exporter materialize all codes at setup. + use std::marker::PhantomData; use vortex::array::ArrayRef; +use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::match_each_integer_ptype; -use vortex::array::search_sorted::SearchSorted; -use vortex::array::search_sorted::SearchSortedSide; use vortex::dtype::IntegerPType; use vortex::encodings::runend::RunEndArray; use vortex::encodings::runend::RunEndArrayExt; @@ -22,17 +27,18 @@ use crate::duckdb::ReusableDict; use crate::duckdb::SelectionVector; use crate::duckdb::VectorRef; use crate::exporter::ColumnExporter; +use crate::exporter::DictionaryParts; use crate::exporter::cache::ConversionCache; use crate::exporter::cached_values_dict; use crate::exporter::canonical; -/// We export run-end arrays to a DuckDB dictionary vector. Values are exported -/// into a ReusableDict with SelectionVector applied in export(). +/// Keeps the values dictionary cached while generating selections for each requested batch. struct RunEndExporter { ends: PrimitiveArray, ends_type: PhantomData, values: ArrayRef, values_dict: ReusableDict, + /// Added to request offsets because sliced run ends retain their source coordinates. run_end_offset: usize, } @@ -42,8 +48,11 @@ pub(crate) fn new_exporter_with_flatten( ctx: &mut ExecutionCtx, flatten: bool, ) -> VortexResult> { - // Our canonicalization is faster than creating a dictionary vector and - // letting duckdb flatten it for us. + if array.is_empty() { + return canonical::new_exporter(Canonical::empty(array.dtype()).into_array(), cache, ctx); + } + + // Canonicalize directly to avoid building a dictionary that DuckDB would immediately flatten. if flatten { return canonical::new_exporter(array.into_array(), cache, ctx); } @@ -65,7 +74,73 @@ pub(crate) fn new_exporter_with_flatten( }) } +impl RunEndExporter { + /// Finds the run containing an absolute row offset. Run ends are exclusive. + /// + /// The offset must refer to an existing row, rather than the end of an empty request. + fn first_run(&self, offset: usize) -> usize { + let offset = + E::from_usize(offset).vortex_expect("the request offset fits the run-end type"); + self.ends + .as_slice::() + .partition_point(|&end| end <= offset) + } + + /// Fills a non-empty request with indices into the original values child. + /// + /// `offset` is absolute, and `first_run` must be the index returned by [`Self::first_run`]. + fn write_run_codes(&self, offset: usize, first_run: usize, codes: &mut [u32]) { + let mut position = + E::from_usize(offset).vortex_expect("the request offset fits the run-end type"); + let end_offset = position + + E::from_usize(codes.len()).vortex_expect("the request length fits the run-end type"); + let mut remaining = codes; + + for (index, &end) in self.ends.as_slice::()[first_run..].iter().enumerate() { + let next_end = end.min(end_offset); + let run_len = (next_end - position) + .to_usize() + .vortex_expect("run length fits usize"); + let code = u32::try_from(first_run + index).vortex_expect("run index exceeds u32"); + remaining[..run_len].fill(code); + remaining = &mut remaining[run_len..]; + if remaining.is_empty() { + return; + } + position = next_end; + } + + assert!( + remaining.is_empty(), + "run ends must cover the requested range" + ); + } +} + +impl DictionaryParts for RunEndExporter { + fn values(&self) -> &ReusableDict { + &self.values_dict + } + + fn values_key(&self) -> usize { + self.values.addr() + } + + fn write_codes(&self, offset: usize, codes: &mut [u32]) { + if codes.is_empty() { + return; + } + + let offset = self.run_end_offset + offset; + self.write_run_codes(offset, self.first_run(offset), codes); + } +} + impl ColumnExporter for RunEndExporter { + fn as_dictionary(&self) -> Option<&dyn DictionaryParts> { + Some(self) + } + fn export( &self, offset: usize, @@ -73,76 +148,75 @@ impl ColumnExporter for RunEndExporter { vector: &mut VectorRef, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let ends_slice = self.ends.as_slice::(); - - // Adjust offset to account for the run-end offset. - let mut offset = E::from_usize(self.run_end_offset + offset) - .vortex_expect("RunEndExporter::export: offset is not a valid value"); - // Compute the final end offset. - let end_offset = offset + E::from_usize(len).vortex_expect("len is not end type"); - - // Find the run that contains the start offset. - let start_run_idx = ends_slice - .search_sorted(&offset, SearchSortedSide::Right)? - .to_ends_index(ends_slice.len()); - - // Find the final run in case we can short-circuit and return a constant vector. - let end_run_idx = ends_slice - .search_sorted(&end_offset, SearchSortedSide::Right)? - .to_ends_index(ends_slice.len()); - - if start_run_idx == end_run_idx { - // NOTE(ngates): would be great if we could just export and set type == CONSTANT - // self.values_exporter.export(start_run_idx, 1, vector, cache); - let constant = self.values.execute_scalar(start_run_idx, ctx)?; - let value = constant.try_to_duckdb_scalar()?; - vector.reference_value(&value); + if len == 0 { return Ok(()); } - // Build up a selection vector - let mut sel_vec = SelectionVector::with_capacity(len); - let mut sel_vec_slice = unsafe { sel_vec.as_slice_mut(len) }; - - for (run_idx, &next_end) in ends_slice[start_run_idx..=end_run_idx].iter().enumerate() { - let next_end = next_end.min(end_offset); - let run_len = (next_end - offset) - .to_usize() - .vortex_expect("run_len is usize"); - - let global_run_idx = - u32::try_from(start_run_idx + run_idx).vortex_expect("run index exceeds u32"); - sel_vec_slice[..run_len].fill(global_run_idx); - sel_vec_slice = &mut sel_vec_slice[run_len..]; + let offset = self.run_end_offset + offset; + let first_run = self.first_run(offset); + let end = + E::from_usize(offset + len).vortex_expect("the request end fits the run-end type"); - offset = next_end; + // The request end is exclusive, so equality still leaves every row inside this run. + if end <= self.ends.as_slice::()[first_run] { + let constant = self.values.execute_scalar(first_run, ctx)?; + let value = constant.try_to_duckdb_scalar()?; + vector.reference_value(&value); + return Ok(()); } - debug_assert!(sel_vec_slice.is_empty()); - vector.reuse_dictionary(&self.values_dict, &sel_vec, len); + let mut selection = SelectionVector::with_capacity(len); + // SAFETY: the selection has capacity for `len` entries. The validated run ends cover the + // requested range, and write_run_codes fills every entry before DuckDB reads the selection. + let codes = unsafe { selection.as_slice_mut(len) }; + self.write_run_codes(offset, first_run, codes); + vector.reuse_dictionary(&self.values_dict, &selection, len); + Ok(()) } } #[cfg(test)] mod tests { + use rstest::rstest; + use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; + use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; + use vortex::buffer::Buffer; use vortex::buffer::buffer; use vortex::encodings::runend::RunEnd; use vortex::error::VortexResult; use crate::SESSION; - use crate::cpp::duckdb_type::DUCKDB_TYPE_INTEGER; use crate::duckdb::DataChunk; use crate::duckdb::LogicalType; use crate::exporter::ArrayExporter; use crate::exporter::ConversionCache; use crate::exporter::new_array_exporter; + #[track_caller] + fn check_exports(array: ArrayRef, requests: &[(usize, usize, &str)]) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let exporter = new_array_exporter(array, &ConversionCache::default(), &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::int32()]); + + for &(offset, len, expected) in requests { + chunk.reset(); + exporter.export(offset, len, chunk.get_vector_mut(0), &mut ctx)?; + chunk.set_len(len); + assert_eq!( + String::try_from(&*chunk)?, + format!("Chunk - [1 Columns]\n{expected}\n") + ); + } + + Ok(()) + } + #[test] fn test_one_chunk_null() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -150,22 +224,10 @@ mod tests { let array = PrimitiveArray::from_option_iter(source); let array = RunEnd::encode(array.into_array(), &mut ctx)?; - let mut chunk = DataChunk::new([LogicalType::new(DUCKDB_TYPE_INTEGER)]); - new_array_exporter(array.into_array(), &ConversionCache::default(), &mut ctx)?.export( - 0, - 5, - chunk.get_vector_mut(0), - &mut ctx, - )?; - chunk.set_len(5); - let chunk_str = String::try_from(&*chunk)?; - assert_eq!( - chunk_str, - r#"Chunk - [1 Columns] -- DICTIONARY INTEGER: 5 = [ 0, 1, NULL, 3, NULL] -"# - ); - Ok(()) + check_exports( + array.into_array(), + &[(0, 5, "- DICTIONARY INTEGER: 5 = [ 0, 1, NULL, 3, NULL]")], + ) } #[test] @@ -191,4 +253,85 @@ mod tests { assert!(!exporter.export(&mut chunk, None)?); Ok(()) } + + #[test] + fn sliced_batches_keep_run_boundaries() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(), + 2, + 5, + &mut ctx, + )?; + + check_exports( + array.into_array(), + &[ + (0, 1, "- CONSTANT INTEGER: 1 = [ 10]"), // Exact first run end. + (1, 2, "- CONSTANT INTEGER: 2 = [ NULL]"), // Exact null run. + (2, 3, "- DICTIONARY INTEGER: 3 = [ NULL, 30, 30]"), // Clipped batch. + (5, 0, "- FLAT INTEGER: 0 = [ ]"), // Empty request at the end. + ], + ) + } + + #[test] + fn empty_arrays_do_not_cache_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = buffer![10i32].into_array(); + let retained = + RunEnd::try_new_offset_length(buffer![3u32].into_array(), values, 3, 0, &mut ctx)?; + let empty = RunEnd::try_new( + Buffer::::empty().into_array(), + Buffer::::empty().into_array(), + &mut ctx, + )?; + + for array in [empty, retained] { + let cache = ConversionCache::default(); + let exporter = new_array_exporter(array.into_array(), &cache, &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::int32()]); + exporter.export(0, 0, chunk.get_vector_mut(0), &mut ctx)?; + assert!(cache.dict_cache.is_empty()); + } + + Ok(()) + } + + #[rstest] + #[case::runend(false, true)] + #[case::dictionary(true, true)] + #[case::different_values(false, false)] + fn chunk_boundaries(#[case] dictionary: bool, #[case] shared: bool) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(); + let first = RunEnd::try_new_offset_length( + buffer![2u32, 4, 6].into_array(), + values.clone(), + 2, + 4, + &mut ctx, + )? + .into_array(); + let next_values = if shared { + values + } else { + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array() + }; + let second = if dictionary { + DictArray::try_new(buffer![0u32, 1, 1, 2].into_array(), next_values)?.into_array() + } else { + RunEnd::try_new(buffer![1u32, 3, 4].into_array(), next_values, &mut ctx)?.into_array() + }; + let dtype = first.dtype().clone(); + let array = ChunkedArray::try_new(vec![first, second], dtype)?.into_array(); + let crossing = if shared { + "- DICTIONARY INTEGER: 6 = [ NULL, 30, 30, 10, NULL, NULL]" + } else { + "- FLAT INTEGER: 6 = [ NULL, 30, 30, 10, NULL, NULL]" + }; + + check_exports(array, &[(1, 6, crossing)]) + } }