From a3012f78f3d5d8aa2056b417694c95cbb8f75d4a Mon Sep 17 00:00:00 2001 From: Will Manning Date: Tue, 1 Sep 2026 20:09:29 -0400 Subject: [PATCH 1/6] perf: Add pending list membership and sparse extraction Signed-off-by: Will Manning --- Cargo.lock | 11 +- Cargo.toml | 2 +- encodings/fastlanes/Cargo.toml | 9 + .../fastlanes/benches/bitpacking_filter.rs | 131 ++++++ .../benches/bitpacking_list_contains.rs | 342 ++++++++++++++++ .../fastlanes/benches/bitpacking_take.rs | 108 +++++ .../src/bitpacking/compute/compare_fused.rs | 56 ++- .../src/bitpacking/compute/filter.rs | 66 ++- .../bitpacking/compute/list_contains/mod.rs | 117 ++++++ .../bitpacking/compute/list_contains/tests.rs | 317 ++++++++++++++ .../fastlanes/src/bitpacking/compute/mod.rs | 36 ++ .../fastlanes/src/bitpacking/compute/take.rs | 136 ++++-- .../src/bitpacking/vtable/kernels.rs | 7 + .../sequence/src/compute/list_contains.rs | 136 +++++- vortex-array/Cargo.toml | 4 + vortex-array/benches/list_contains.rs | 247 +++++++++++ .../arrays/primitive/compute/list_contains.rs | 336 +++++++++++++++ .../src/arrays/primitive/compute/mod.rs | 1 + .../src/arrays/primitive/vtable/kernel.rs | 7 + vortex-array/src/expr/exprs.rs | 2 + vortex-array/src/scalar/typed_view/list.rs | 5 + .../fns/list_contains/integer_membership.rs | 66 +++ .../src/scalar_fn/fns/list_contains/kernel.rs | 41 +- .../src/scalar_fn/fns/list_contains/mod.rs | 386 ++++++++++++++---- vortex-datafusion/src/convert/exprs.rs | 119 +++++- vortex-duckdb/src/convert/expr.rs | 37 +- .../src/e2e_test/vortex_scan_test.rs | 14 + 27 files changed, 2567 insertions(+), 172 deletions(-) create mode 100644 encodings/fastlanes/benches/bitpacking_filter.rs create mode 100644 encodings/fastlanes/benches/bitpacking_list_contains.rs create mode 100644 encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs create mode 100644 encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs create mode 100644 vortex-array/benches/list_contains.rs create mode 100644 vortex-array/src/arrays/primitive/compute/list_contains.rs create mode 100644 vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs diff --git a/Cargo.lock b/Cargo.lock index 0a367df1534..343f8413bc3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4042,9 +4042,9 @@ dependencies = [ [[package]] name = "fastlanes" -version = "0.7.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34f6c951d711d8a10f08524071f6171bc95c5f857e8b39d91e36ad6353fc4f4f" +checksum = "9e218082bba8aee4ba5704355ea1ce1f742facbcf748eadbca90de56aa551cd9" dependencies = [ "const_for", "num-traits", @@ -10468,7 +10468,7 @@ dependencies = [ "anyhow", "arrow-array 59.2.0", "codspeed-divan-compat", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "mimalloc", "parquet 59.2.0", "rand 0.10.2", @@ -10871,7 +10871,7 @@ dependencies = [ "bindgen", "codspeed-criterion-compat-walltime", "cudarc", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "futures", "itertools 0.14.0", "kanal", @@ -11045,7 +11045,7 @@ name = "vortex-fastlanes" version = "0.1.0" dependencies = [ "codspeed-divan-compat", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "itertools 0.14.0", "lending-iterator", "num-traits", @@ -11054,6 +11054,7 @@ dependencies = [ "rstest", "vortex-alp", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-fastlanes", diff --git a/Cargo.toml b/Cargo.toml index 294ed20d996..f7fbd95300a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -161,7 +161,7 @@ datafusion-sqllogictest = { version = "55.0.0" } divan = { package = "codspeed-divan-compat", version = "5.0.0" } enum-iterator = "2.0.0" env_logger = "0.11" -fastlanes = "0.7.0" +fastlanes = "0.7.1" flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..732baa7e37f 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -39,6 +39,7 @@ rand = { workspace = true } rstest = { workspace = true } vortex-alp = { path = "../alp" } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-bench-support = { workspace = true } vortex-fastlanes = { path = ".", features = ["_test-harness"] } [features] @@ -48,6 +49,14 @@ _test-harness = ["dep:rand"] name = "bitpacking_take" harness = false +[[bench]] +name = "bitpacking_list_contains" +harness = false + +[[bench]] +name = "bitpacking_filter" +harness = false + [[bench]] name = "canonicalize_bench" harness = false diff --git a/encodings/fastlanes/benches/bitpacking_filter.rs b/encodings/fastlanes/benches/bitpacking_filter.rs new file mode 100644 index 00000000000..51ac0828bf9 --- /dev/null +++ b/encodings/fastlanes/benches/bitpacking_filter.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures selective filtering around the sparse extraction thresholds. + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +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::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPackedData; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +const NUM_ARRAY_CHUNKS: usize = 64; +// Keep the array density below the outer full-decode policy. +const NUM_SELECTED_CHUNKS: usize = 8; +const CHUNK_SIZE: usize = 1_024; +const LEN: usize = NUM_ARRAY_CHUNKS * CHUNK_SIZE; + +trait BenchInt: NativePType { + fn from_counter(value: u64) -> Self; +} + +macro_rules! impl_bench_int { + ($($T:ty),+) => { + $(impl BenchInt for $T { + fn from_counter(value: u64) -> Self { + value as $T + } + })+ + }; +} + +impl_bench_int!(u8, u16, u32, u64); + +fn fixture(bit_width: usize, selected_per_chunk: usize) -> (ArrayRef, Mask) { + let limit = if bit_width == 64 { + u64::MAX + } else { + 1_u64 << bit_width + }; + let values: BufferMut = (0..LEN) + .map(|index| T::from_counter(index as u64 % limit)) + .collect(); + let packed = BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + bit_width as u8, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array(); + let indices = (0..NUM_SELECTED_CHUNKS).flat_map(|chunk| { + (0..selected_per_chunk) + .map(move |index| chunk * CHUNK_SIZE + index * CHUNK_SIZE / selected_per_chunk) + }); + (packed, Mask::from_indices(LEN, indices)) +} + +macro_rules! bench_width { + ($module:ident, $T:ty, $bit_width:expr, [$($selected:expr),+ $(,)?]) => { + mod $module { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = [$($selected),+])] + fn filter(bencher: Bencher, selected_per_chunk: usize) { + let (packed, mask) = fixture::<$T>($bit_width, selected_per_chunk); + bencher + .counter(ItemsCount::new(LEN)) + .with_inputs(|| (mask.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(mask, ctx)| { + packed + .filter(mask.clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); + } + } + }; +} + +macro_rules! bench_type { + ($module:ident, $T:ty, [$(($width_module:ident, $bit_width:expr)),+ $(,)?], $selected:tt) => { + mod $module { + use super::*; + + $(bench_width!($width_module, $T, $bit_width, $selected);)+ + } + }; +} + +bench_type!(u8, u8, [(width1, 1), (width4, 4), (width7, 7)], [8, 16, 24]); +bench_type!( + u16, + u16, + [(width1, 1), (width8, 8), (width15, 15)], + [8, 32, 48] +); +bench_type!( + u32, + u32, + [(width1, 1), (width16, 16), (width31, 31)], + [8, 64, 80, 96] +); +bench_type!( + u64, + u64, + [(width1, 1), (width32, 32), (width63, 63)], + [8, 128, 160, 192] +); diff --git a/encodings/fastlanes/benches/bitpacking_list_contains.rs b/encodings/fastlanes/benches/bitpacking_list_contains.rs new file mode 100644 index 00000000000..27439139eda --- /dev/null +++ b/encodings/fastlanes/benches/bitpacking_list_contains.rs @@ -0,0 +1,342 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures compressed constant-list membership. +//! +//! FastLanes evaluates at most four distinct integer members during unpacking. Larger sets use the +//! frozen generic path. Every path runs on each real CPU feature leg in CodSpeed. +//! +//! Run with `cargo bench -p vortex-fastlanes --bench bitpacking_list_contains`. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hint::black_box; +use std::sync::Arc; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::Alignment; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::BitPackedData; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +trait BenchInt: IntegerPType + Copy + Into { + fn from_counter(value: u64) -> Self; +} + +impl BenchInt for u8 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u16 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u32 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u64 { + fn from_counter(value: u64) -> Self { + value + } +} + +#[derive(Clone, Copy)] +struct PackedCase { + name: &'static str, + ptype: PType, + bit_width: u8, + len: usize, + member_count: usize, + member_stride: u64, + patch_every: Option, +} + +impl Display for PackedCase { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "{}_{}_w{}_n{}", + self.name, self.ptype, self.bit_width, self.len + ) + } +} + +const fn strided_case( + name: &'static str, + ptype: PType, + bit_width: u8, + len: usize, + count: usize, + stride: u64, +) -> PackedCase { + PackedCase { + name, + ptype, + bit_width, + len, + member_count: count, + member_stride: stride, + patch_every: None, + } +} + +const fn patched_case( + name: &'static str, + ptype: PType, + bit_width: u8, + len: usize, + count: usize, + stride: u64, + patch_every: usize, +) -> PackedCase { + PackedCase { + name, + ptype, + bit_width, + len, + member_count: count, + member_stride: stride, + patch_every: Some(patch_every), + } +} + +const PACKED_CASES: &[PackedCase] = &[ + strided_case("direct_u8_m4", PType::U8, 6, 65_536, 4, 2), + strided_case("fallback_u8_m5", PType::U8, 6, 65_536, 5, 2), + strided_case("direct_u16_m4", PType::U16, 8, 65_536, 4, 2), + strided_case("fallback_u16_m5", PType::U16, 8, 65_536, 5, 2), + strided_case("direct_u32_m4", PType::U32, 8, 65_536, 4, 2), + strided_case("fallback_u32_m5", PType::U32, 8, 65_536, 5, 2), + strided_case("direct_u64_m4", PType::U64, 40, 65_536, 4, 2), + strided_case("fallback_u64_m5", PType::U64, 40, 65_536, 5, 2), + strided_case("short_direct_u32_m4", PType::U32, 10, 1_024, 4, 2), + strided_case("short_fallback_u32_m5", PType::U32, 10, 1_024, 5, 2), + strided_case("wide_direct_u32_m4", PType::U32, 31, 65_536, 4, 2), + patched_case("patch_sparse_u32_m4", PType::U32, 8, 65_536, 4, 2, 64), +]; + +fn page_aligned(array: BitPackedArray) -> BitPackedArray { + let ptype = array.dtype().as_ptype(); + let parts = BitPacked::into_parts(array); + BitPacked::try_new( + parts.packed.ensure_aligned(Alignment::new(4_096)).unwrap(), + ptype, + parts.validity, + parts.patches, + parts.bit_width, + parts.len, + parts.offset, + ) + .unwrap() +} + +fn generated_values(case: PackedCase, ordinary_members: &[u64]) -> Vec { + let domain_size = 1u64 << case.bit_width; + let patch_hit = domain_size; + let patch_miss = patch_hit + 1; + let mut state = 0x9E37_79B9_7F4A_7C15u64; + (0..case.len) + .map(|index| { + if let Some(patch_every) = case.patch_every + && index.is_multiple_of(patch_every) + { + return if (index / patch_every).is_multiple_of(2) { + patch_hit + } else { + patch_miss + }; + } + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + let is_hit = (state >> 32).is_multiple_of(2); + if is_hit { + let member_index = + usize::try_from(state % u64::try_from(ordinary_members.len()).unwrap()) + .unwrap(); + ordinary_members[member_index] + } else { + let mut candidate = state.rotate_left(17) % domain_size; + while ordinary_members.contains(&candidate) { + candidate = (candidate + 1) % domain_size; + } + candidate + } + }) + .collect() +} + +fn list_scalar(members: &[u64]) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(T::PTYPE, Nullability::NonNullable)), + members + .iter() + .map(|value| T::from_counter(*value).into()) + .collect(), + Nullability::NonNullable, + ) +} + +fn frozen_generic_membership( + values: ArrayRef, + members: &[u64], +) -> VortexResult { + fn balanced_or(arrays: &[ArrayRef]) -> VortexResult { + if let [array] = arrays { + return Ok(array.clone()); + } + let (left, right) = arrays.split_at(arrays.len() / 2); + balanced_or(left)?.binary(balanced_or(right)?, Operator::Or) + } + + let len = values.len(); + let nullability = values.dtype().nullability(); + let false_scalar = Scalar::bool(false, nullability); + let comparisons = members + .iter() + .map(|member| { + Binary::try_new( + ConstantArray::new(T::from_counter(*member).into(), len).into_array(), + values.clone(), + Operator::Eq, + )? + .into_array() + .fill_null(false_scalar.clone()) + }) + .collect::>>()?; + + if comparisons.is_empty() { + Ok(ConstantArray::new(false_scalar, len).into_array()) + } else { + balanced_or(&comparisons) + } +} + +fn execute_generic_baseline( + values: &BitPackedArray, + members: &[u64], + ctx: &mut vortex_array::ExecutionCtx, +) -> BoolArray { + frozen_generic_membership::(values.clone().into_array(), members) + .unwrap() + .execute::(ctx) + .unwrap() +} + +fn packed_input( + case: PackedCase, +) -> (BitPackedArray, Vec, BoolArray, VortexSession) { + let session = array_session(); + vortex_fastlanes::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let in_domain_member_count = case.member_count - usize::from(case.patch_every.is_some()); + let ordinary_members = (0..in_domain_member_count) + .map(|index| u64::try_from(index).unwrap() * case.member_stride) + .collect::>(); + let mut members = ordinary_members.clone(); + if case.patch_every.is_some() { + members.push(1u64 << case.bit_width); + } + let generated = generated_values(case, &ordinary_members); + let expected = BoolArray::from_iter(generated.iter().map(|value| members.contains(value))); + let values: BufferMut = generated.into_iter().map(T::from_counter).collect(); + let packed = page_aligned( + BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + case.bit_width, + &mut ctx, + ) + .unwrap(), + ); + if case.patch_every.is_some() { + assert!(packed.patches().is_some()); + } + (packed, members, expected, session) +} + +fn bench_packed_current(bencher: Bencher, case: PackedCase) { + let (packed, members, expected, session) = packed_input::(case); + let contains = packed + .into_array() + .apply(&list_contains(lit(list_scalar::(&members)), root())) + .unwrap(); + let mut ctx = session.create_execution_ctx(); + let actual = contains.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!(actual, expected, &mut ctx); + bencher + .counter(ItemsCount::new(case.len)) + .bench_local(|| black_box(contains.clone().execute::(&mut ctx).unwrap())); +} + +fn bench_packed_generic_baseline(bencher: Bencher, case: PackedCase) { + let (packed, members, expected, session) = packed_input::(case); + let mut ctx = session.create_execution_ctx(); + let actual = execute_generic_baseline::(&packed, &members, &mut ctx); + assert_arrays_eq!(actual, expected, &mut ctx); + // The frozen pre-change implementation built this array tree during execution. + bencher + .counter(ItemsCount::new(case.len)) + .bench_local(|| black_box(execute_generic_baseline::(&packed, &members, &mut ctx))); +} + +macro_rules! dispatch_packed { + ($bencher:expr, $case:expr, $function:ident) => { + match $case.ptype { + PType::U8 => $function::($bencher, $case), + PType::U16 => $function::($bencher, $case), + PType::U32 => $function::($bencher, $case), + PType::U64 => $function::($bencher, $case), + _ => unreachable!("benchmark case uses an unsigned integer type"), + } + }; +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = PACKED_CASES)] +fn packed_current(bencher: Bencher, case: PackedCase) { + dispatch_packed!(bencher, case, bench_packed_current); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = PACKED_CASES)] +fn packed_generic_baseline(bencher: Bencher, case: PackedCase) { + dispatch_packed!(bencher, case, bench_packed_generic_baseline); +} diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index eb072017ae3..b4d43ac0411 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -7,18 +7,23 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::ItemsCount; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::BitPackedData; use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; use vortex_session::VortexSession; @@ -32,6 +37,109 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +const NUM_ARRAY_CHUNKS: usize = 64; +// Keep the selected count below the outer full-decode policy. +const NUM_SELECTED_CHUNKS: usize = 8; +const CHUNK_SIZE: usize = 1_024; +const THRESHOLD_FIXTURE_LEN: usize = NUM_ARRAY_CHUNKS * CHUNK_SIZE; + +trait BenchInt: NativePType { + fn from_counter(value: u64) -> Self; +} + +macro_rules! impl_bench_int { + ($($T:ty),+) => { + $(impl BenchInt for $T { + fn from_counter(value: u64) -> Self { + value as $T + } + })+ + }; +} + +impl_bench_int!(u8, u16, u32, u64); + +fn threshold_fixture( + bit_width: usize, + selected_per_chunk: usize, +) -> (ArrayRef, ArrayRef) { + let limit = if bit_width == 64 { + u64::MAX + } else { + 1_u64 << bit_width + }; + let values: BufferMut = (0..THRESHOLD_FIXTURE_LEN) + .map(|index| T::from_counter(index as u64 % limit)) + .collect(); + let packed = BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + bit_width as u8, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array(); + let indices = PrimitiveArray::from_iter((0..NUM_SELECTED_CHUNKS).flat_map(|chunk| { + (0..selected_per_chunk) + .map(move |index| (chunk * CHUNK_SIZE + index * CHUNK_SIZE / selected_per_chunk) as u32) + })) + .into_array(); + (packed, indices) +} + +macro_rules! bench_width { + ($module:ident, $T:ty, $bit_width:expr, [$($selected:expr),+ $(,)?]) => { + mod $module { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = [$($selected),+])] + fn threshold(bencher: Bencher, selected_per_chunk: usize) { + let (packed, indices) = threshold_fixture::<$T>($bit_width, selected_per_chunk); + bencher + .counter(ItemsCount::new(indices.len())) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + packed + .take(indices.clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); + } + } + }; +} + +macro_rules! bench_type { + ($module:ident, $T:ty, [$(($width_module:ident, $bit_width:expr)),+ $(,)?], $selected:tt) => { + mod $module { + use super::*; + + $(bench_width!($width_module, $T, $bit_width, $selected);)+ + } + }; +} + +bench_type!(u8, u8, [(width1, 1), (width4, 4), (width7, 7)], [8, 16, 24]); +bench_type!( + u16, + u16, + [(width1, 1), (width8, 8), (width15, 15)], + [8, 32, 48] +); +bench_type!( + u32, + u32, + [(width1, 1), (width16, 16), (width31, 31)], + [8, 64, 80, 96, 112] +); +bench_type!( + u64, + u64, + [(width1, 1), (width32, 32), (width63, 63)], + [8, 128, 160, 192] +); + #[divan::bench] fn take_10_stratified(bencher: Bencher) { let values = fixture(65_536, 8); diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index 1259ed815fe..97b802611b1 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Fused compare kernel for [`BitPackedArray`] against a constant. +//! Fused predicate kernel for [`BitPackedArray`]. //! //! Where [`super::stream_predicate`] unpacks a full 1024-element FastLanes block into a scratch //! buffer and *then* folds a predicate over it, this path hands the comparison down into the -//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which compares each value against -//! the constant *as it is unpacked*, accumulating the boolean results straight into a 1024-bit +//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which evaluates each value +//! *as it is unpacked*, accumulating the boolean results straight into a 1024-bit //! mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per lane, no //! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates //! that mask into logical row order. @@ -20,7 +20,7 @@ //! slot with no per-block temporary and only one shared scratch `[u64; 16]`. The leading `offset` //! garbage rows are represented as the final [`BitBuffer`] bit offset, which naturally handles //! sub-byte slices without copy-aligning. Inline patches are spliced in afterwards by overwriting -//! the bits at the patched indices with `cmp(patch_value, rhs)`. +//! the bits at the patched indices with the predicate result. //! //! [`BitPackedArray`]: crate::BitPackedArray //! [`BitBuffer`]: vortex_buffer::BitBuffer @@ -70,6 +70,46 @@ pub(super) fn stream_compare_fused( cmp: F, ctx: &mut ExecutionCtx, ) -> VortexResult +where + T: NativePType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, + F: Fn(T, T) -> bool + Copy, +{ + stream_compare_fused_inner(array, rhs, nullability, cmp, ctx) +} + +/// Evaluates `predicate` while FastLanes unpacks each value. +pub(super) fn stream_predicate_fused( + array: ArrayView<'_, BitPacked>, + nullability: Nullability, + predicate: F, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, + F: Fn(T) -> bool + Copy, +{ + stream_compare_fused_inner( + array, + T::default(), + nullability, + move |value, _| predicate(value), + ctx, + ) +} + +fn stream_compare_fused_inner( + array: ArrayView<'_, BitPacked>, + rhs: T, + nullability: Nullability, + cmp: F, + ctx: &mut ExecutionCtx, +) -> VortexResult where T: NativePType + BitPackedIter @@ -84,7 +124,7 @@ where // A degenerate width has no packed payload for the fused kernel to consume; defer to the scalar // streaming predicate, which handles every layout (including the empty array). if len == 0 || bit_width == 0 { - return stream_predicate::(array, nullability, move |v| cmp(v, rhs), ctx); + return stream_predicate::(array, nullability, move |value| cmp(value, rhs), ctx); } // Over-allocate to whole 1024-bit blocks in padded coordinates so every block - including the @@ -121,12 +161,12 @@ where let mut bits = BitBufferMut::from_buffer(words.into_byte_buffer(), offset, len); - // Patched indices hold placeholder packed values, so their fused result is meaningless; - // overwrite each with the comparison against the real patch value. + // Patched indices hold placeholder packed values, so their fused result is meaningless. + // Overwrite each result with the predicate for the real patch value. // TODO(joe): apply patches per `packed_chunked`. if let Some(p) = array.patches() { let p_idx = p.indices().clone().execute::(ctx)?; - // TODO(joe): push down cmp?? + // TODO(joe): push down the predicate. let p_val = p.values().clone().execute::(ctx)?; let p_off = p.offset(); match_each_unsigned_integer_ptype!(p_idx.ptype(), |I| { diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1b9422f86..5544e460a68 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -22,7 +22,8 @@ use vortex_mask::Mask; use vortex_mask::MaskValuesRef; use super::chunked_indices; -use super::take::UNPACK_CHUNK_THRESHOLD; +use super::unpack_chunk_threshold; +use super::unpack_indices_into; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; @@ -150,8 +151,10 @@ fn filter_with_indices( &mut values.as_mut_slice()[values_len..], ); } - } else if indices_within_chunk.len() > UNPACK_CHUNK_THRESHOLD { + } else if indices_within_chunk.len() > unpack_chunk_threshold::() { // Unpack into a temporary chunk and then copy the values. + // SAFETY: The validated bit width fits `T`. The source and destination contain + // one complete FastLanes block. The call initializes every destination value. unsafe { let dst: &mut [MaybeUninit] = &mut unpacked; let dst: &mut [T] = std::mem::transmute(dst); @@ -160,13 +163,11 @@ fn filter_with_indices( values.extend_trusted( indices_within_chunk .iter() + // SAFETY: The preceding unpack initialized the complete temporary block. .map(|&idx| unsafe { unpacked.get_unchecked(idx).assume_init() }), ); } else { - // Otherwise, unpack each element individually. - values.extend_trusted(indices_within_chunk.iter().map(|&idx| unsafe { - BitPacking::unchecked_unpack_single(bit_width, packed, idx) - })); + unpack_indices_into(&mut values, bit_width, packed, indices_within_chunk); } }, ); @@ -186,9 +187,12 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; + use super::filter_with_indices; + use super::unpack_chunk_threshold; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -198,6 +202,56 @@ mod tests { session }); + #[test] + fn sparse_extraction_covers_batch_and_full_chunk_paths() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + macro_rules! check_type { + ($T:ty, $bit_width:expr) => {{ + let values = (0..2_048) + .map(|index| (index % 127) as $T) + .collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + $bit_width, + &mut ctx, + )?; + let threshold = unpack_chunk_threshold::<$T>(); + + for selected in [threshold, threshold + 1] { + let indices = (0..selected) + .map(|index| index * 1_024 / selected) + .collect::>(); + let actual = filter_with_indices::<$T>(&packed, &indices); + let expected = indices + .iter() + .map(|&index| values[index]) + .collect::>(); + assert_eq!(actual.as_slice(), expected); + } + }}; + } + + check_type!(u8, 7); + check_type!(u16, 15); + check_type!(u32, 31); + check_type!(u64, 63); + Ok(()) + } + + #[test] + fn sparse_extraction_supports_zero_width() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter([0u32; 2_048]).into_array(), + 0, + &mut ctx, + )?; + let actual = filter_with_indices::(&packed, &[0, 17, 1_023, 1_024, 2_047]); + assert_eq!(actual.as_slice(), &[0; 5]); + Ok(()) + } + #[test] fn take_indices() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs b/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs new file mode 100644 index 00000000000..a631a1b5891 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use fastlanes::BitPacking; +use fastlanes::BitPackingCompare; +use fastlanes::FastLanesComparable; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PhysicalPType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::list_contains::IntegerMembership; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +use vortex_array::scalar_fn::fns::list_contains::evaluate_constant_list_generic; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; + +use super::compare_fused::stream_predicate_fused; +use crate::BitPacked; +use crate::unpack_iter::BitPacked as BitPackedIter; + +impl ListContainsElementKernel for BitPacked { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + list_contains_compressed(list, element, ctx) + } +} + +fn list_contains_compressed( + list: &ArrayRef, + element: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let nullability = list.dtype().nullability() | element.dtype().nullability(); + + match_each_integer_ptype!(element.dtype().as_ptype(), |T| { + list_contains_typed::(list, element, nullability, ctx) + }) +} + +fn list_contains_typed( + list: &ArrayRef, + element: ArrayView<'_, BitPacked>, + nullability: vortex_array::dtype::Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + T: IntegerPType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, +{ + let Some(membership) = IntegerMembership::::try_from_constant_list(list, element.dtype())? + else { + return evaluate_constant_list_generic(list, element.array(), nullability); + }; + let result = match membership.members() { + [] => BoolArray::new( + BitBuffer::new_unset(element.len()), + element.validity()?.union_nullability(nullability), + ) + .into_array(), + [member] => { + let member = *member; + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(member), + ctx, + )? + } + [first, second] => { + let (first, second) = (*first, *second); + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(first) | value.is_eq(second), + ctx, + )? + } + [first, second, third] => { + let (first, second, third) = (*first, *second, *third); + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(first) | value.is_eq(second) | value.is_eq(third), + ctx, + )? + } + [first, second, third, fourth] => { + let (first, second, third, fourth) = (*first, *second, *third, *fourth); + stream_predicate_fused::( + element, + nullability, + move |value| { + value.is_eq(first) + | value.is_eq(second) + | value.is_eq(third) + | value.is_eq(fourth) + }, + ctx, + )? + } + _ => return Ok(None), + }; + Ok(Some(result)) +} + +#[cfg(test)] +mod tests; diff --git a/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs b/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs new file mode 100644 index 00000000000..3ef71a2cee8 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs @@ -0,0 +1,317 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::slice::SliceKernel; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +#[cfg(not(codspeed))] +use vortex_array::test_harness::trace::trace_op; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::BitPacked; +use crate::BitPackedArray; +use crate::BitPackedArrayExt; +use crate::BitPackedData; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +fn member_list( + values: impl IntoIterator>, + member_nullability: Nullability, +) -> Scalar +where + T: NativePType + Into, +{ + let member_dtype = DType::Primitive(T::PTYPE, member_nullability); + let members = values + .into_iter() + .map(|value| { + value + .map(|value| Scalar::primitive(value, member_nullability)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(); + Scalar::list(Arc::new(member_dtype), members, Nullability::NonNullable) +} + +fn list_array(list: Scalar, len: usize) -> ArrayRef { + ConstantArray::new(list, len).into_array() +} + +fn execute_direct( + list: &ArrayRef, + element: &BitPackedArray, + ctx: &mut vortex_array::ExecutionCtx, +) -> VortexResult { + ::list_contains(list, element.as_view(), ctx)? + .ok_or_else(|| vortex_err!("BitPacked list_contains kernel declined a supported input"))? + .execute::(ctx) +} + +macro_rules! integer_type_test { + ($name:ident, $T:ty, $bit_width:expr) => { + #[test] + fn $name() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048) + .map(|value| (value % 64) as $T) + .collect::>(); + let members = [1 as $T, 3 as $T, 63 as $T]; + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), $bit_width, &mut ctx)?; + let list = list_array( + member_list(members.into_iter().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = + BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + }; +} + +integer_type_test!(test_integer_type_u8, u8, 6); +integer_type_test!(test_integer_type_u16, u16, 6); +integer_type_test!(test_integer_type_u32, u32, 6); +integer_type_test!(test_integer_type_u64, u64, 6); +// BitPacked encoding rejects negative integers. These cases verify signed PType dispatch with the +// representable nonnegative domain. +integer_type_test!(test_integer_type_i8, i8, 6); +integer_type_test!(test_integer_type_i16, i16, 6); +integer_type_test!(test_integer_type_i32, i32, 6); +integer_type_test!(test_integer_type_i64, i64, 6); + +#[rstest] +#[case::one(vec![3])] +#[case::two(vec![3, 7])] +#[case::three(vec![3, 7, 11])] +#[case::four(vec![3, 7, 11, 15])] +#[case::duplicate_source(vec![3, 3, 7, 7])] +fn test_member_cardinalities(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let list = list_array( + member_list(members.iter().copied().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_many_member_public_expression_path() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let members = (0..5).map(|value| value * 2).collect::>(); + let values = (0..4_096).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let expression = list_contains( + lit(member_list( + members.iter().copied().map(Some), + Nullability::NonNullable, + )), + root(), + ); + + let actual = packed + .into_array() + .apply(&expression)? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::present([true; 128], vec![0])] +#[case::absent([false; 128], vec![1])] +fn test_zero_bit_width( + #[case] expected: [bool; 128], + #[case] members: Vec, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([0i32; 128]); + let packed = BitPackedData::encode(&primitive.into_array(), 0, &mut ctx)?; + let list = list_array( + member_list(members.into_iter().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(expected); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::fused(vec![3, 100_388])] +#[case::fallback({ + let mut members = (0..4).collect::>(); + members.push(100_388); + members +})] +fn test_sliced_patched_array(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..5_000) + .map(|index| { + if index % 97 == 0 { + 100_000 + index + } else { + index % 100 + } + }) + .collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + assert!(packed.patches().is_some(), "test setup requires patches"); + let range = 333..4_333; + let sliced = ::slice(packed.as_view(), range.clone(), &mut ctx)? + .ok_or_else(|| vortex_err!("BitPacked slice kernel declined a supported input"))?; + let list = member_list(members.iter().copied().map(Some), Nullability::NonNullable); + + let actual = sliced + .into_array() + .apply(&list_contains(lit(list), root()))? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values[range].iter().map(|value| members.contains(value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::nullable_needles( + vec![Some(1), Some(3)], + Nullability::NonNullable, + vec![Some(1), None, Some(2)], + vec![Some(true), None, Some(false)], +)] +#[case::nullable_members( + vec![Some(1), None, Some(3)], + Nullability::Nullable, + vec![Some(1), Some(2), Some(3)], + vec![Some(true), Some(false), Some(true)], +)] +#[case::all_null_members( + vec![None, None], + Nullability::Nullable, + vec![Some(1), None, Some(2)], + vec![Some(false), None, Some(false)], +)] +fn test_null_semantics( + #[case] members: Vec>, + #[case] member_nullability: Nullability, + #[case] values: Vec>, + #[case] expected: Vec>, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_option_iter(values); + let packed = BitPackedData::encode(&primitive.into_array(), 3, &mut ctx)?; + let list = list_array(member_list(members, member_nullability), packed.len()); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(expected); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_wrong_integer_type_declines_without_panic() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]); + let packed = BitPackedData::encode(&primitive.into_array(), 2, &mut ctx)?; + let list = list_array( + member_list([Some(1i64), Some(3)], Nullability::NonNullable), + packed.len(), + ); + + let result = + ::list_contains(&list, packed.as_view(), &mut ctx)?; + assert!(result.is_none()); + Ok(()) +} + +#[test] +fn test_nonconstant_list_declines() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]); + let packed = BitPackedData::encode(&primitive.into_array(), 2, &mut ctx)?; + let list = ListArray::from_iter_slow::( + vec![vec![1i32], vec![2], vec![3]], + Arc::new(DType::Primitive(i32::PTYPE, Nullability::NonNullable)), + )? + .into_array(); + + let result = + ::list_contains(&list, packed.as_view(), &mut ctx)?; + assert!(result.is_none()); + Ok(()) +} + +#[test] +#[cfg(not(codspeed))] +fn test_registered_kernel_executes_through_expression() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let members = [0, 99]; + let expression = list_contains( + lit(member_list( + members.into_iter().map(Some), + Nullability::NonNullable, + )), + root(), + ); + let contains = packed.into_array().apply(&expression)?; + + let traced = trace_op(|| contains.execute::(&mut ctx))?; + let trace = traced.trace.to_string(); + let applied = trace + .lines() + .filter(|line| { + line.contains("child_execute_parent session[") + && line.contains("slot=1") + && line.contains("parent=vortex.list.contains") + && line.contains("child=fastlanes.bitpacked") + }) + .collect::>(); + // A silent fallback preserves values but loses compressed-domain execution. + assert!(!applied.is_empty(), "{trace}"); + + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(traced.output, expected, &mut ctx); + Ok(()) +} diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..8666c094d73 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -1,16 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::size_of; + +use fastlanes::BitPacking; +use vortex_array::dtype::NativePType; +use vortex_buffer::BufferMut; + mod between; mod cast; mod compare; mod compare_fused; mod filter; pub(crate) mod is_constant; +pub(crate) mod list_contains; mod slice; mod stream_predicate; mod take; +const fn unpack_chunk_threshold() -> usize { + // FastLanes and Vortex benchmarks set conservative crossovers for each physical type. + match size_of::() { + 1 => 16, + 2 => 32, + 4 => 64, + 8 => 160, + _ => unreachable!(), + } +} + +fn unpack_indices_into( + output: &mut BufferMut, + bit_width: usize, + packed: &[T], + indices: &[usize], +) { + let output_len = output.len(); + let destination = &mut output.spare_capacity_mut()[..indices.len()]; + + // SAFETY: `bit_width` comes from validated data and fits `T`. + // `packed` contains one complete block, and each index is block-relative. + // The destination length equals the index length, and `output` reserves enough space. + unsafe { + T::unchecked_unpack_indices(bit_width, packed, indices, destination); + output.set_len(output_len + indices.len()); + } +} + // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. fn chunked_indices( mut indices: impl Iterator, diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..3c9bf79ab1c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::mem; use std::mem::MaybeUninit; use fastlanes::BitPacking; @@ -23,15 +22,11 @@ use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use super::chunked_indices; +use super::unpack_chunk_threshold; +use super::unpack_indices_into; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::bitpack_decompress; - -// TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. -/// assuming the buffer is already allocated (which will happen at most once) then unpacking -/// all 1024 elements takes ~8.8x as long as unpacking a single element on an M2 Macbook Air. -/// see -pub(super) const UNPACK_CHUNK_THRESHOLD: usize = 8; +const FULL_ARRAY_DECODE_RATIO: usize = 8; impl TakeExecute for BitPacked { fn take( @@ -40,7 +35,7 @@ impl TakeExecute for BitPacked { ctx: &mut ExecutionCtx, ) -> VortexResult> { // If the indices are large enough, it's faster to flatten and take the primitive array. - if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() { + if indices.len() * FULL_ARRAY_DECODE_RATIO > array.len() { let prim = array.array().clone().execute::(ctx)?; return prim.into_array().take(indices.clone()).map(Some); } @@ -98,41 +93,22 @@ fn take_primitive( chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; - let mut have_unpacked = false; - let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); - - // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets - for offset_chunk in offset_chunks { - if !have_unpacked { - unsafe { - let dst: &mut [MaybeUninit] = &mut unpacked; - let dst: &mut [T] = mem::transmute(dst); - BitPacking::unchecked_unpack(bit_width, packed, dst); - } - have_unpacked = true; - } - - for &index in offset_chunk { - output.push(unsafe { unpacked[index].assume_init() }); - } - } - - // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it - if !remainder.is_empty() { - if have_unpacked { - // we already bulk unpacked this chunk, so we can just push the remaining elements - for &index in remainder { - output.push(unsafe { unpacked[index].assume_init() }); - } - } else { - // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place, - // so we need to unpack each one individually - for &index in remainder { - output.push(unsafe { - bitpack_decompress::unpack_single_primitive::(packed, bit_width, index) - }); - } + if indices_within_chunk.len() > unpack_chunk_threshold::() { + // SAFETY: The validated bit width fits `T`. The source and destination contain one + // complete FastLanes block. The call initializes every destination value. + unsafe { + let dst: &mut [MaybeUninit] = &mut unpacked; + let dst: &mut [T] = std::mem::transmute(dst); + BitPacking::unchecked_unpack(bit_width, packed, dst); } + output.extend_trusted( + indices_within_chunk + .iter() + // SAFETY: The preceding unpack initialized the complete temporary block. + .map(|&index| unsafe { unpacked.get_unchecked(index).assume_init() }), + ); + } else { + unpack_indices_into(&mut output, bit_width, packed, indices_within_chunk); } }); @@ -157,7 +133,7 @@ fn take_primitive( #[cfg(test)] #[expect(clippy::cast_possible_truncation)] -mod test { +mod tests { use std::sync::LazyLock; use rand::RngExt; @@ -172,12 +148,14 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; use crate::bitpacking::compute::take::take_primitive; + use crate::bitpacking::compute::take::unpack_chunk_threshold; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -185,6 +163,76 @@ mod test { session }); + #[test] + fn sparse_extraction_covers_batch_and_full_chunk_paths() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + macro_rules! check_type { + ($T:ty, $bit_width:expr) => {{ + let values = (0..2_048) + .map(|index| (index % 127) as $T) + .collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + $bit_width, + &mut ctx, + )?; + let threshold = unpack_chunk_threshold::<$T>(); + + for selected in [threshold, threshold + 1] { + let indices = (0..selected) + .map(|index| (index * 1_024 / selected) as u32) + .collect::>(); + let actual = take_primitive::<$T, u32>( + packed.as_view(), + &PrimitiveArray::from_iter(indices.iter().copied()), + Validity::NonNullable, + &mut ctx, + )?; + let expected = indices + .iter() + .map(|&index| values[index as usize]) + .collect::>(); + assert_eq!(actual.as_slice::<$T>(), expected); + } + }}; + } + + check_type!(u8, 7); + check_type!(u16, 15); + check_type!(u32, 31); + check_type!(u64, 63); + Ok(()) + } + + #[test] + fn sparse_take_preserves_nullable_order_and_duplicates() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..8_192).map(|index| index as u32).collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + 13, + &mut ctx, + )?; + let indices = [ + Some(3_073u32), + Some(2), + None, + Some(2), + Some(1_025), + Some(3_073), + Some(0), + ]; + let actual = packed + .take(PrimitiveArray::from_option_iter(indices).into_array())? + .execute::(&mut ctx)?; + let expected = PrimitiveArray::from_option_iter( + indices.map(|index| index.map(|index| values[index as usize])), + ); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + #[test] fn take_indices() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index eb0dd9b7a23..9a0add2130b 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -16,6 +16,8 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; use vortex_session::VortexSession; use crate::BitPacked; @@ -36,4 +38,9 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Filter.id(), BitPacked, FilterExecuteAdaptor(BitPacked)); kernels.register_execute_parent_kernel(Slice.id(), BitPacked, SliceExecuteAdaptor(BitPacked)); kernels.register_execute_parent_kernel(Dict.id(), BitPacked, TakeExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel( + ListContains.id(), + BitPacked, + ListContainsElementExecuteAdaptor(BitPacked), + ); } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index 80ffcad24cd..34cc5b276e4 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -5,10 +5,11 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::array::Sequence; @@ -20,14 +21,22 @@ impl ListContainsElementReduce for Sequence { list: &ArrayRef, element: ArrayView<'_, Self>, ) -> VortexResult> { - let Some(list_scalar) = list.as_constant() else { + let Some(list_array) = list.as_opt::() else { return Ok(None); }; + let DType::List(member_dtype, _) = list.dtype() else { + return Ok(None); + }; + if !member_dtype.eq_ignore_nullability(element.dtype()) { + return Ok(None); + } - let list_elements = list_scalar - .as_list() - .elements() - .vortex_expect("non-null element (checked in entry)"); + let Some(list_elements) = list_array.scalar().as_list().elements() else { + return Ok(None); + }; + if list_elements.is_empty() { + return Ok(None); + } let nullability = list.dtype().nullability() | element.dtype().nullability(); @@ -54,6 +63,12 @@ impl ListContainsElementReduce for Sequence { } } + if set_indices.is_empty() { + return Ok(Some( + ConstantArray::new(Scalar::bool(false, nullability), element.len()).into_array(), + )); + } + Ok(Some( BoolArray::from_indices(element.len(), set_indices, nullability.into()).into_array(), )) @@ -65,16 +80,24 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Constant; + use vortex_array::arrays::ConstantArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType::I32; + use vortex_array::dtype::PType::I64; use vortex_array::expr::list_contains; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Sequence; @@ -139,4 +162,105 @@ mod tests { let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } + + #[test] + fn test_no_intersection_reduces_to_constant() { + let list_scalar = Scalar::list( + Arc::new(I32.into()), + vec![7.into(), 42.into()], + Nullability::NonNullable, + ); + let array = Sequence::try_new_typed(1i32, 1, Nullability::NonNullable, 3) + .unwrap() + .into_array(); + + let result = array + .apply(&list_contains(lit(list_scalar), root())) + .unwrap(); + + assert!(result.is::()); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut SESSION.create_execution_ctx() + ); + } + + #[rstest] + #[case::null_list( + Scalar::null(DType::List(Arc::new(I32.into()), Nullability::Nullable)), + [None, None, None] + )] + #[case::empty_list( + Scalar::list(Arc::new(I32.into()), vec![], Nullability::Nullable), + [Some(false), Some(false), Some(false)] + )] + fn test_constant_list_semantics( + #[case] list_scalar: Scalar, + #[case] expected: [Option; 3], + ) { + let array = Sequence::try_new_typed(1i32, 1, Nullability::NonNullable, 3) + .unwrap() + .into_array(); + let expr = list_contains(lit(list_scalar), root()); + + let result = array.apply(&expr).unwrap(); + + assert!(result.is::()); + assert_arrays_eq!( + result, + BoolArray::from_iter(expected), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn test_nullable_members() -> VortexResult<()> { + let member_dtype = DType::Primitive(I32, Nullability::Nullable); + let list = ConstantArray::new( + Scalar::list( + Arc::new(member_dtype.clone()), + vec![ + Scalar::primitive(1i32, Nullability::Nullable), + Scalar::null(member_dtype), + Scalar::primitive(3i32, Nullability::Nullable), + ], + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + let sequence = Sequence::try_new_typed(1i32, 1, Nullability::NonNullable, 3)?; + + let result = + ::list_contains(&list, sequence.as_view())? + .vortex_expect("matching integer types are supported"); + + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_wrong_integer_type_declines() -> VortexResult<()> { + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I64, Nullability::NonNullable)), + vec![1i64.into(), 3i64.into()], + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + let sequence = Sequence::try_new_typed(1i32, 1, Nullability::NonNullable, 3)?; + + let result = + ::list_contains(&list, sequence.as_view())?; + + assert!(result.is_none()); + Ok(()) + } } diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index f6b06544baf..5a91561c64f 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -291,6 +291,10 @@ harness = false name = "list_length" harness = false +[[bench]] +name = "list_contains" +harness = false + [[bench]] name = "list_sum" harness = false diff --git a/vortex-array/benches/list_contains.rs b/vortex-array/benches/list_contains.rs new file mode 100644 index 00000000000..725c2125682 --- /dev/null +++ b/vortex-array/benches/list_contains.rs @@ -0,0 +1,247 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares the Primitive constant-list membership dispatch paths. +//! +//! Primitive arrays use direct comparisons for at most four distinct integer members. Larger sets +//! use the frozen generic path. Every path runs on each real CPU feature leg in CodSpeed. +//! +//! Run with `cargo bench -p vortex-array --bench list_contains`. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hint::black_box; +use std::sync::Arc; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +trait BenchInt: IntegerPType + Copy + Into { + fn from_counter(value: u64) -> Self; +} + +impl BenchInt for u8 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u16 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u32 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u64 { + fn from_counter(value: u64) -> Self { + value + } +} + +#[derive(Clone, Copy)] +struct PrimitiveCase { + name: &'static str, + len: usize, + member_count: usize, +} + +impl Display for PrimitiveCase { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "{}_m{}_n{}", + self.name, self.member_count, self.len + ) + } +} + +const fn primitive_case(name: &'static str, len: usize, member_count: usize) -> PrimitiveCase { + PrimitiveCase { + name, + len, + member_count, + } +} + +const LONG_M1: PrimitiveCase = primitive_case("long", 65_536, 1); +const LONG_M4: PrimitiveCase = primitive_case("long", 65_536, 4); +const LONG_M5: PrimitiveCase = primitive_case("long", 65_536, 5); +const SHORT_M4: PrimitiveCase = primitive_case("short", 1_024, 4); + +const CURRENT: &[PrimitiveCase] = &[LONG_M1, LONG_M4, LONG_M5, SHORT_M4]; + +fn primitive_input( + case: PrimitiveCase, +) -> (PrimitiveArray, Vec, BoolArray, VortexSession) { + let members = (0..case.member_count) + .map(|index| T::from_counter(u64::try_from(index).unwrap() * 2)) + .collect::>(); + let domain_bits = T::PTYPE.bit_width().min(12); + let domain_size = 1u64 << domain_bits; + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let generated = (0..case.len) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + if (state >> 32).is_multiple_of(2) { + let member_index = + usize::try_from(state % u64::try_from(members.len()).unwrap()).unwrap(); + members[member_index] + } else { + let mut candidate = state.rotate_left(17) % domain_size; + while members.contains(&T::from_counter(candidate)) { + candidate = (candidate + 1) % domain_size; + } + T::from_counter(candidate) + } + }) + .collect::>(); + let expected = BoolArray::from_iter(generated.iter().map(|value| members.contains(value))); + ( + PrimitiveArray::new::(generated, Validity::NonNullable), + members, + expected, + array_session(), + ) +} + +fn list_scalar(members: &[T]) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(T::PTYPE, Nullability::NonNullable)), + members.iter().copied().map(Into::into).collect(), + Nullability::NonNullable, + ) +} + +fn frozen_generic_membership( + values: ArrayRef, + members: &[T], +) -> VortexResult { + fn balanced_or(arrays: &[ArrayRef]) -> VortexResult { + if let [array] = arrays { + return Ok(array.clone()); + } + let (left, right) = arrays.split_at(arrays.len() / 2); + balanced_or(left)?.binary(balanced_or(right)?, Operator::Or) + } + + let len = values.len(); + let nullability = values.dtype().nullability(); + let false_scalar = Scalar::bool(false, nullability); + let comparisons = members + .iter() + .map(|member| { + let member: Scalar = (*member).into(); + Binary::try_new( + ConstantArray::new(member, len).into_array(), + values.clone(), + Operator::Eq, + )? + .into_array() + .fill_null(false_scalar.clone()) + }) + .collect::>>()?; + + if comparisons.is_empty() { + Ok(ConstantArray::new(false_scalar, len).into_array()) + } else { + balanced_or(&comparisons) + } +} + +fn execute_generic_baseline( + values: &PrimitiveArray, + members: &[T], + ctx: &mut vortex_array::ExecutionCtx, +) -> BoolArray { + frozen_generic_membership(values.clone().into_array(), members) + .unwrap() + .execute::(ctx) + .unwrap() +} + +fn bench_current(bencher: Bencher, case: PrimitiveCase) { + let (array, members, expected, session) = primitive_input::(case); + let contains = array + .into_array() + .apply(&list_contains(lit(list_scalar(&members)), root())) + .unwrap(); + let mut ctx = session.create_execution_ctx(); + let actual = contains.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!(actual, expected, &mut ctx); + + bencher + .counter(ItemsCount::new(case.len)) + .bench_local(|| black_box(contains.clone().execute::(&mut ctx).unwrap())); +} + +fn bench_generic_baseline(bencher: Bencher, case: PrimitiveCase) { + let (array, members, expected, session) = primitive_input::(case); + let mut ctx = session.create_execution_ctx(); + let actual = execute_generic_baseline(&array, &members, &mut ctx); + assert_arrays_eq!(actual, expected, &mut ctx); + + // The frozen pre-change implementation built this array tree during execution. + bencher + .counter(ItemsCount::new(case.len)) + .bench_local(|| black_box(execute_generic_baseline(&array, &members, &mut ctx))); +} + +macro_rules! primitive_benchmarks { + ($type_name:ident, $ty:ty, $current:ident) => { + mod $type_name { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = $current)] + fn current(bencher: Bencher, case: PrimitiveCase) { + bench_current::<$ty>(bencher, case); + } + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = $current)] + fn generic_baseline(bencher: Bencher, case: PrimitiveCase) { + bench_generic_baseline::<$ty>(bencher, case); + } + } + }; +} + +primitive_benchmarks!(u8_cases, u8, CURRENT); +primitive_benchmarks!(u16_cases, u16, CURRENT); +primitive_benchmarks!(u32_cases, u32, CURRENT); +primitive_benchmarks!(u64_cases, u64, CURRENT); diff --git a/vortex-array/src/arrays/primitive/compute/list_contains.rs b/vortex-array/src/arrays/primitive/compute/list_contains.rs new file mode 100644 index 00000000000..243c87be23b --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/list_contains.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ArrayView; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Primitive; +use crate::dtype::IntegerPType; +use crate::dtype::NativePType; +use crate::match_each_integer_ptype; +use crate::scalar_fn::fns::list_contains::IntegerMembership; +use crate::scalar_fn::fns::list_contains::ListContainsElementKernel; +use crate::scalar_fn::fns::list_contains::evaluate_constant_list_generic; + +impl ListContainsElementKernel for Primitive { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + evaluate_constant_list_membership(list, element) + } +} + +fn evaluate_constant_list_membership( + list: &ArrayRef, + element: ArrayView<'_, Primitive>, +) -> VortexResult> { + if !element.ptype().is_int() { + return Ok(None); + } + + match_each_integer_ptype!(element.ptype(), |T| { + evaluate_integer_membership::(list, element) + }) +} + +fn evaluate_integer_membership( + list: &ArrayRef, + element: ArrayView<'_, Primitive>, +) -> VortexResult> { + let nullability = list.dtype().nullability() | element.dtype().nullability(); + let Some(membership) = IntegerMembership::::try_from_constant_list(list, element.dtype())? + else { + return evaluate_constant_list_generic(list, element.array(), nullability); + }; + let values = element.as_slice::(); + let bits = match membership.members() { + [] => BitBuffer::new_unset(values.len()), + [member] => collect_direct(values, move |value| value.is_eq(*member)), + [first, second] => collect_direct(values, move |value| { + value.is_eq(*first) | value.is_eq(*second) + }), + [first, second, third] => collect_direct(values, move |value| { + value.is_eq(*first) | value.is_eq(*second) | value.is_eq(*third) + }), + [first, second, third, fourth] => collect_direct(values, move |value| { + value.is_eq(*first) | value.is_eq(*second) | value.is_eq(*third) | value.is_eq(*fourth) + }), + _ => return Ok(None), + }; + Ok(Some( + BoolArray::new(bits, element.validity()?.union_nullability(nullability)).into_array(), + )) +} + +fn collect_direct(values: &[T], mut predicate: impl FnMut(T) -> bool) -> BitBuffer { + BitBuffer::collect_bool(values.len(), |index| { + // SAFETY: collect_bool visits each valid index once. + predicate(unsafe { *values.get_unchecked(index) }) + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_error::VortexExpect; + + use super::*; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + use crate::arrays::Constant; + use crate::arrays::ConstantArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType::F32; + use crate::dtype::PType::I32; + use crate::expr::list_contains; + use crate::expr::lit; + use crate::expr::root; + use crate::scalar::Scalar; + #[cfg(not(codspeed))] + use crate::test_harness::trace::trace_op; + + fn list(values: impl IntoIterator, len: usize) -> ArrayRef { + ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + values + .into_iter() + .map(|value| Scalar::primitive(value, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ), + len, + ) + .into_array() + } + + #[rstest] + #[case::one(vec![3])] + #[case::two(vec![3, 7])] + #[case::three(vec![3, 7, 11])] + #[case::four(vec![3, 7, 11, 15])] + #[case::duplicate_source(vec![3, 3, 7, 7])] + fn test_membership_plans(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [0, 3, 7, 15, 31, 90_000, 310_000]; + let element = PrimitiveArray::from_iter(values); + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + + let actual = ::list_contains( + &list(members, element.len()), + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_five_members_use_generic_fallback() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [0i32, 3, 99]; + let element = PrimitiveArray::from_iter(values); + let members = [0, 3, 6, 9, 12]; + + let actual = ::list_contains( + &list(members, element.len()), + element.as_view(), + &mut ctx, + )? + .vortex_expect("larger constant lists use the generic fallback"); + + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + #[cfg(not(codspeed))] + fn test_registered_kernel_executes_through_expression() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [0i32, 1, 2, 3]; + let element = PrimitiveArray::from_iter(values); + let members = [1, 3]; + let contains = element.into_array().apply(&list_contains( + lit(list(members, values.len()) + .as_constant() + .vortex_expect("constant list")), + root(), + ))?; + + let traced = trace_op(|| contains.execute::(&mut ctx))?; + let trace = traced.trace.to_string(); + let applied = trace + .lines() + .filter(|line| { + line.contains("child_execute_parent session[") + && line.contains("slot=1") + && line.contains("parent=vortex.list.contains") + && line.contains("child=vortex.primitive") + }) + .collect::>(); + // A silent fallback preserves values but loses the membership optimization. + assert!(!applied.is_empty(), "{trace}"); + + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + assert_arrays_eq!(traced.output, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_float_falls_back_through_expression() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [1.5f32, 2.5, 3.5]; + let element = PrimitiveArray::from_iter(values); + let members = [1.5f32, 3.5]; + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(F32, Nullability::NonNullable)), + members.into_iter().map(Scalar::from).collect(), + Nullability::NonNullable, + ), + element.len(), + ) + .into_array(); + let list_scalar = list.as_constant().vortex_expect("list is constant"); + + let actual = element + .into_array() + .apply(&list_contains(lit(list_scalar), root()))? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_null_needles() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let element = PrimitiveArray::from_option_iter([Some(1), None, Some(2)]); + let expected = BoolArray::from_iter([Some(true), None, Some(false)]); + + let actual = ::list_contains( + &list([1, 3], element.len()), + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nullable_list_preserves_output_nullability() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + [1, 3].into_iter().map(Scalar::from).collect(), + Nullability::Nullable, + ), + 3, + ) + .into_array(); + let element = PrimitiveArray::from_iter([1, 2, 3]); + + let actual = ::list_contains( + &list, + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + assert_arrays_eq!( + actual, + BoolArray::from_iter([Some(true), Some(false), Some(true)]), + &mut ctx + ); + Ok(()) + } + + #[rstest] + #[case::null_list(true)] + #[case::empty_list(false)] + fn test_constant_list_adaptor(#[case] null_list: bool) -> VortexResult<()> { + let member_dtype = DType::Primitive(I32, Nullability::NonNullable); + let list = if null_list { + Scalar::null(DType::List(Arc::new(member_dtype), Nullability::Nullable)) + } else { + Scalar::list(Arc::new(member_dtype), vec![], Nullability::NonNullable) + }; + let needles = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); + + let mut ctx = crate::array_session().create_execution_ctx(); + let contains = needles + .apply(&list_contains(lit(list), root()))? + .execute::(&mut ctx)?; + let expected = if null_list { + BoolArray::from_iter([None, None, None]) + } else { + BoolArray::from_iter([Some(false), Some(false), Some(false)]) + }; + + assert!(contains.is::()); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::mixed( + vec![Some(1), None, Some(3)], + [Some(true), None, Some(true)] + )] + #[case::all_null(vec![None, None], [Some(false), None, Some(false)])] + fn test_nullable_members( + #[case] members: Vec>, + #[case] expected: [Option; 3], + ) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let member_dtype = DType::Primitive(I32, Nullability::Nullable); + let list = ConstantArray::new( + Scalar::list( + Arc::new(member_dtype.clone()), + members + .into_iter() + .map(|member| { + member + .map(|value| Scalar::primitive(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(), + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + let element = PrimitiveArray::from_option_iter([Some(1), None, Some(3)]); + + let actual = ::list_contains( + &list, + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 382b42ee6e2..7f1dcdcb4cf 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -5,6 +5,7 @@ mod between; mod cast; mod fill_null; mod fixed_width; +mod list_contains; mod mask; pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..3f13282c334 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -15,6 +15,8 @@ use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; @@ -31,6 +33,11 @@ pub(crate) fn initialize(session: &VortexSession) { Primitive, FillNullExecuteAdaptor(Primitive), ); + kernels.register_execute_parent_kernel( + ListContains.id(), + Primitive, + ListContainsElementExecuteAdaptor(Primitive), + ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index fb8bfe227aa..16dcd3256a7 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -1100,6 +1100,8 @@ pub fn bound_dynamic( /// Creates an expression that checks if a value is contained in a list. /// /// Returns a boolean array indicating whether the value appears in each list. +/// A null list produces null. An empty list produces false, including for a null value. +/// A null value produces null for a nonempty list. Null list members do not match any value. /// /// ```rust /// # use vortex_array::expr::{list_contains, lit, root}; diff --git a/vortex-array/src/scalar/typed_view/list.rs b/vortex-array/src/scalar/typed_view/list.rs index f97857c92ed..ef9948b4218 100644 --- a/vortex-array/src/scalar/typed_view/list.rs +++ b/vortex-array/src/scalar/typed_view/list.rs @@ -142,6 +142,11 @@ impl<'a> ListScalar<'a> { self.elements.is_none() } + #[inline] + pub(crate) fn values(&self) -> Option<&'a [Option]> { + self.elements + } + /// Returns the data type of the list's elements. pub fn element_dtype(&self) -> &DType { self.dtype diff --git a/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs new file mode 100644 index 00000000000..83ea016b6c5 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::arrays::Constant; +use crate::dtype::DType; +use crate::dtype::IntegerPType; + +const MAX_SOURCE_MEMBERS: usize = 4; + +/// A prepared integer set for constant-list membership kernels. +/// +/// The set sorts and deduplicates its members. +#[doc(hidden)] +pub struct IntegerMembership { + members: Box<[T]>, +} + +impl IntegerMembership { + fn new(mut members: Vec) -> Self { + members.sort_unstable(); + members.dedup(); + Self { + members: members.into_boxed_slice(), + } + } + + /// Extracts an integer set from a compatible constant list. + pub fn try_from_constant_list( + list: &ArrayRef, + element_dtype: &DType, + ) -> VortexResult> { + let Some(list_array) = list.as_opt::() else { + return Ok(None); + }; + let DType::List(member_dtype, _) = list.dtype() else { + return Ok(None); + }; + if !member_dtype.eq_ignore_nullability(element_dtype) { + return Ok(None); + } + let Some(elements) = list_array.scalar().as_list().values() else { + return Ok(None); + }; + if elements.len() > MAX_SOURCE_MEMBERS { + return Ok(None); + } + + let members = elements + .iter() + .filter_map(|value| value.as_ref()) + .map(|value| { + // The validated list scalar stores primitive values of `member_dtype`. + value.as_primitive().cast::() + }) + .collect::>>()?; + Ok(Some(Self::new(members))) + } + + /// Returns the prepared members. + pub fn members(&self) -> &[T] { + &self.members + } +} diff --git a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs b/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs index 563600bfeee..ab04ac5ab52 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs @@ -6,24 +6,53 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; use crate::arrays::ScalarFn; use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::ScalarFnArrayView; +use crate::dtype::DType; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::scalar::Scalar; use crate::scalar_fn::fns::list_contains::ListContains as ListContainsExpr; +fn constant_list_result( + list: &ArrayRef, + element_len: usize, + element_nullability: crate::dtype::Nullability, +) -> Option { + let list_array = list.as_opt::()?; + let list_scalar = list_array.scalar().as_list(); + let DType::List(_, list_nullability) = list.dtype() else { + return None; + }; + let nullability = *list_nullability | element_nullability; + + if list_scalar.is_null() { + return Some( + ConstantArray::new(Scalar::null(DType::Bool(nullability)), element_len).into_array(), + ); + } + if list_scalar.is_empty() { + return Some( + ConstantArray::new(Scalar::bool(false, nullability), element_len).into_array(), + ); + } + None +} + /// Check list-contains without reading buffers (metadata-only). /// /// This trait dispatches on the **element** (needle) child at index 1 of the `ListContains` /// expression. `Self::Array` is the concrete element encoding, while the list (haystack) is /// passed as an opaque `&ArrayRef`. /// -/// A future `ListContainsListReduce` could dispatch on the list side (child 0) for encodings -/// with specialized list representations. +/// The parent adaptor resolves null and empty constant lists before delegation. /// /// Return `None` if the operation cannot be resolved from metadata alone. pub trait ListContainsElementReduce: VTable { @@ -38,6 +67,8 @@ pub trait ListContainsElementReduce: VTable { /// Like [`ListContainsElementReduce`], this dispatches on the **element** (needle) child at /// index 1. Unlike the reduce variant, implementations may read and execute on buffers via /// the provided [`ExecutionCtx`]. +/// +/// The parent adaptor resolves null and empty constant lists before delegation. pub trait ListContainsElementKernel: VTable { fn list_contains( list: &ArrayRef, @@ -70,6 +101,9 @@ where .as_opt::() .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray"); let list = scalar_fn_array.get_child(0); + if let Some(result) = constant_list_result(list, array.len(), array.dtype().nullability()) { + return Ok(Some(result)); + } ::list_contains(list, array) } } @@ -99,6 +133,9 @@ where .as_opt::() .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray"); let list = scalar_fn_array.get_child(0); + if let Some(result) = constant_list_result(list, array.len(), array.dtype().nullability()) { + return Ok(Some(result)); + } ::list_contains(list, array, ctx) } } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index d2508014089..6e984c17b79 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod integer_membership; mod kernel; use std::ops::BitOr; use arrow_buffer::bit_iterator::BitIndexIterator; +pub use integer_membership::IntegerMembership; pub use kernel::*; use num_traits::Zero; use vortex_buffer::BitBuffer; @@ -13,6 +15,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::iter::ReduceBalancedIterExt; @@ -56,7 +59,8 @@ impl ListContains { /// /// # Errors /// - /// Returns an error if the children have different lengths or `list` is not a list array. + /// Returns an error if the children have different lengths, `list` is not a list array, or + /// the list member type differs from the needle type. pub fn try_new(list: ArrayRef, needle: ArrayRef) -> VortexResult { ScalarFnArray::try_new(ListContains.bind(EmptyOptions), vec![list, needle]) } @@ -100,16 +104,17 @@ impl ScalarFnVTable for ListContains { let list_dtype = &arg_dtypes[0]; let needle_dtype = &arg_dtypes[1]; - let nullability = match list_dtype { - DType::List(_, list_nullability) => list_nullability, - _ => { - vortex_bail!( - "First argument to ListContains must be a List, got {:?}", - list_dtype - ); - } + let DType::List(member_dtype, list_nullability) = list_dtype else { + vortex_bail!("First argument to ListContains must be a List, got {list_dtype}"); + }; + if !member_dtype.eq_ignore_nullability(needle_dtype) { + vortex_bail!( + "Element type {} of list does not match search value {}", + member_dtype, + needle_dtype + ); } - .bitor(needle_dtype.nullability()); + let nullability = list_nullability.bitor(needle_dtype.nullability()); Ok(DType::Bool(nullability)) } @@ -146,8 +151,7 @@ impl ScalarFnVTable for ListContains { fn compute_contains_scalar(list: &Scalar, needle: &Scalar) -> VortexResult { let nullability = list.dtype().nullability() | needle.dtype().nullability(); - // Handle null list or null needle - if list.is_null() || needle.is_null() { + if list.is_null() { return Ok(Scalar::null(DType::Bool(nullability))); } @@ -155,6 +159,12 @@ fn compute_contains_scalar(list: &Scalar, needle: &Scalar) -> VortexResult(ctx)?; + return list_false_if_empty_else_null(&list_array, nullability, ctx); + } + if let Some(value_scalar) = value.as_constant() { list_contains_scalar(array, &value_scalar, nullability, ctx) } else if let Some(list_scalar) = array.as_constant() { @@ -195,6 +218,30 @@ fn compute_list_contains( } } +/// Evaluates the generic constant-list path for an encoding-specific kernel. +#[doc(hidden)] +pub fn evaluate_constant_list_generic( + list: &ArrayRef, + values: &ArrayRef, + nullability: Nullability, +) -> VortexResult> { + let Some(list_array) = list.as_opt::() else { + return Ok(None); + }; + let DType::List(member_dtype, _) = list.dtype() else { + return Ok(None); + }; + if !member_dtype.eq_ignore_nullability(values.dtype()) { + return Ok(None); + } + let list_scalar = list_array.scalar().as_list(); + if list_scalar.is_null() { + return Ok(None); + } + + constant_list_scalar_contains(&list_scalar, values, nullability).map(Some) +} + /// There is a constant list scalar (haystack) being compared to an array of needles. fn constant_list_scalar_contains( list_scalar: &ListScalar, @@ -206,8 +253,13 @@ fn constant_list_scalar_contains( let len = values.len(); let false_scalar = Scalar::bool(false, nullability); + if elements.is_empty() { + return Ok(ConstantArray::new(false_scalar, len).into_array()); + } + let result = elements .iter() + .filter(|element| !element.is_null()) .map(|element| { Binary::try_new( ConstantArray::new(element.clone(), len).into_array(), @@ -221,7 +273,12 @@ fn constant_list_scalar_contains( .into_iter() .try_reduce_balanced(|acc, res| acc.binary(res, Operator::Or))?; - Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array())) + let result = result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array()); + if values.dtype().is_nullable() { + result.mask(values.validity()?.to_array(len)) + } else { + Ok(result) + } } /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. @@ -244,6 +301,9 @@ fn list_contains_scalar( // Must return false when a list is empty (but valid), or null when the list itself is null. return list_false_or_null(&list_array, nullability); } + if value.is_null() { + return list_false_if_empty_else_null(&list_array, nullability, ctx); + } let rhs = ConstantArray::new(value.clone(), elems.len()); let matching_elements = @@ -266,13 +326,7 @@ fn list_contains_scalar( list_false_or_null(&list_array, nullability) } // No elements match, and all comparisons are valid (result in `false`). - Some(false) => { - // False, but match the nullability to the input list array. - Ok( - ConstantArray::new(Scalar::bool(false, nullability), list_array.len()) - .into_array(), - ) - } + Some(false) => list_false_or_null(&list_array, nullability), // All elements match, and all comparisons are valid (result in `true`). Some(true) => { // True, unless the list itself is empty or NULL. @@ -294,9 +348,9 @@ fn list_contains_scalar( // Process based on the offset and size types. let list_matches = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(sizes.ptype(), |S| { - process_matches::(matches, list_array.len(), offsets, sizes) + process_matches::(&matches, list_array.len(), offsets, sizes, ctx) }) - }); + })?; Ok(BoolArray::new( list_matches, @@ -305,33 +359,58 @@ fn list_contains_scalar( .into_array()) } +/// Returns false for valid empty lists and null for all other lists. +fn list_false_if_empty_else_null( + list_array: &ListViewArray, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let sizes = list_array.sizes().clone().execute::(ctx)?; + let empty = match_each_integer_ptype!(sizes.ptype(), |S| { + Mask::from_iter(sizes.as_slice::().iter().map(|size| size.is_zero())) + }); + let valid = list_array.validity()?.execute_mask(list_array.len(), ctx)? & ∅ + + Ok(BoolArray::new( + BitBuffer::new_unset(list_array.len()), + Validity::from_mask(valid, nullability), + ) + .into_array()) +} + /// Returns a [`BitBuffer`] where each bit represents if a list contains the scalar, derived from a /// [`BoolArray`] of matches on the child elements array. fn process_matches( - matches: BoolArray, + matches: &BoolArray, list_array_len: usize, offsets: PrimitiveArray, sizes: PrimitiveArray, -) -> BitBuffer + ctx: &mut ExecutionCtx, +) -> VortexResult where O: IntegerPType, S: IntegerPType, { let offsets_slice = offsets.as_slice::(); let sizes_slice = sizes.as_slice::(); - let bits = matches.bit_buffer_view(); + let value_bits = matches.to_bit_buffer(); + let valid_matches = match matches.validity()? { + Validity::NonNullable | Validity::AllValid => value_bits, + Validity::AllInvalid => BitBuffer::new_unset(matches.len()), + validity => value_bits & validity.execute_mask(matches.len(), ctx)?.into_bit_buffer(), + }; - (0..list_array_len) + Ok((0..list_array_len) .map(|i| { let offset = offsets_slice[i].as_(); let size = sizes_slice[i].as_(); // BitIndexIterator yields indices of true bits only. If `.next()` returns // `Some(_)`, at least one element in this list's range matches. - let mut set_bits = BitIndexIterator::new(bits.inner(), offset, size); + let mut set_bits = BitIndexIterator::new(valid_matches.inner(), offset, size); set_bits.next().is_some() }) - .collect::() + .collect::()) } /// Returns a `Bool` array with `false` for lists that are valid, @@ -414,6 +493,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::Constant; + use crate::arrays::Dict; + use crate::arrays::DictArray; use crate::arrays::ListArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; @@ -435,6 +517,7 @@ mod tests { use crate::scalar::Scalar; use crate::scalar_fn::fns::list_contains::BoolArray; use crate::scalar_fn::fns::list_contains::ConstantArray; + use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListViewArray; use crate::scalar_fn::fns::list_contains::PrimitiveArray; use crate::stats::StatsSession; @@ -540,8 +623,10 @@ mod tests { ); } - #[test] - pub fn test_nullable() { + #[rstest] + #[case::match_present(2, Some(true))] + #[case::match_absent(4, Some(false))] + pub fn test_nullable(#[case] needle: i32, #[case] expected_first: Option) { let arr = ListArray::try_new( PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2]).into_array(), PrimitiveArray::from_iter(vec![0, 5, 5]).into_array(), @@ -550,18 +635,13 @@ mod tests { .unwrap() .into_array(); - let expr = list_contains(root(), lit(2)); + let expr = list_contains(root(), lit(needle)); let item = arr.apply(&expr).unwrap(); - assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::Nullable) - ); - assert!( - !item - .is_valid(1, &mut array_session().create_execution_ctx()) - .unwrap() + assert_arrays_eq!( + item, + BoolArray::from_iter([expected_first, None]), + &mut array_session().create_execution_ctx() ); } @@ -587,6 +667,54 @@ mod tests { ); } + #[test] + fn test_return_type_rejects_mismatched_member_type() { + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + vec![], + Nullability::NonNullable, + ), + 1, + ) + .into_array(); + let needle = + ConstantArray::new(Scalar::utf8("needle", Nullability::NonNullable), 1).into_array(); + + let error = ListContains::try_new(list, needle).unwrap_err(); + + assert!( + error + .to_string() + .contains("Element type i32 of list does not match search value utf8") + ); + } + + #[test] + fn test_dictionary_needles_preserve_dictionary_pushdown() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let codes = PrimitiveArray::from_iter([0u8, 1, 2, 0]).into_array(); + let needles = DictArray::try_new(codes, values)?.into_array(); + let list = Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + vec![1.into(), 3.into()], + Nullability::NonNullable, + ); + let contains = needles.apply(&list_contains(lit(list), root()))?; + + // Dictionary preservation avoids materializing repeated needle values. + assert!(contains.is::()); + let actual = contains.execute::(&mut ctx)?; + + assert_arrays_eq!( + actual, + BoolArray::from_iter([true, false, true, true]), + &mut ctx + ); + Ok(()) + } + #[test] pub fn list_falsification() -> VortexResult<()> { let expr = list_contains( @@ -639,36 +767,47 @@ mod tests { assert_eq!(expr2.to_string(), "vortex.list.contains($, 42i32)"); } - #[test] - pub fn test_constant_scalars() { - let arr = test_array(); - - // Both list and needle are constants - should use scalar optimization - let list_scalar = Scalar::list( - Arc::new(DType::Primitive(I32, Nullability::NonNullable)), - vec![1.into(), 2.into(), 3.into()], - Nullability::NonNullable, - ); - - // Test contains true - let expr = list_contains(lit(list_scalar.clone()), lit(2i32)); - let result = arr.clone().apply(&expr).unwrap(); - assert_eq!( - result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::NonNullable) - ); + #[rstest] + #[case::present(false, vec![1, 2, 3], Some(2), Some(true))] + #[case::absent(false, vec![1, 2, 3], Some(42), Some(false))] + #[case::null_list(true, vec![], Some(1), None)] + #[case::empty_list_null_needle(false, vec![], None, Some(false))] + #[case::nonempty_list_null_needle(false, vec![1], None, None)] + fn test_constant_scalar_null_semantics( + #[case] null_list: bool, + #[case] members: Vec, + #[case] needle: Option, + #[case] expected: Option, + ) -> VortexResult<()> { + let member_dtype = DType::Primitive(I32, Nullability::NonNullable); + let list_dtype = DType::List(Arc::new(member_dtype.clone()), Nullability::Nullable); + let list = if null_list { + Scalar::null(list_dtype) + } else { + Scalar::list( + Arc::new(member_dtype), + members.into_iter().map(Scalar::from).collect(), + Nullability::Nullable, + ) + }; + let needle = needle + .map(|value| Scalar::primitive(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Primitive(I32, Nullability::Nullable))); + let expected = expected + .map(|value| Scalar::bool(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))); + + let contains = ListContains::try_new( + ConstantArray::new(list, 1).into_array(), + ConstantArray::new(needle, 1).into_array(), + )? + .into_array(); - // Test contains false - let expr = list_contains(lit(list_scalar), lit(42i32)); - let result = arr.apply(&expr).unwrap(); assert_eq!( - result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(false, Nullability::NonNullable) + contains.execute_scalar(0, &mut array_session().create_execution_ctx())?, + expected ); + Ok(()) } // -- Tests migrated from compute/list_contains.rs -- @@ -749,7 +888,7 @@ mod tests { #[case( null_strings(vec![vec![], vec![None, None], vec![None, None, None]]), None, - bool_array(vec![false, true, true], Validity::AllInvalid) + BoolArray::from_iter([Some(false), None, None]) )] #[case( null_strings(vec![vec![], vec![None, None], vec![None, None, None]]), @@ -777,23 +916,80 @@ mod tests { assert_arrays_eq!(result, expected, &mut ctx); } - #[test] - fn test_constant_list() { + #[rstest] + #[case::empty( + Vec::>::new(), + [Some(false), Some(false), Some(false)] + )] + #[case::nonempty( + vec![Some("a"), Some("c")], + [Some(true), None, Some(false)] + )] + #[case::all_null( + vec![None, None], + [Some(false), None, Some(false)] + )] + fn test_constant_string_list_nullable_needles( + #[case] members: Vec>, + #[case] expected: [Option; 3], + ) { let mut ctx = array_session().create_execution_ctx(); - let list_array = ConstantArray::new( - Scalar::list( - Arc::new(DType::Primitive(I32, Nullability::NonNullable)), - vec![1i32.into(), 2i32.into(), 3i32.into()], - Nullability::NonNullable, - ), - 2, + let member_dtype = DType::Utf8(Nullability::Nullable); + let list = Scalar::list( + Arc::new(member_dtype.clone()), + members + .into_iter() + .map(|member| { + member + .map(|value| Scalar::utf8(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(), + Nullability::NonNullable, + ); + let needles = VarBinArray::from_iter( + [Some("a"), None, Some("b")], + DType::Utf8(Nullability::Nullable), ) .into_array(); - let expr = list_contains(root(), lit(2i32)); - let contains = list_array.apply(&expr).unwrap(); - let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(contains, expected, &mut ctx); + let result = needles.apply(&list_contains(lit(list), root())).unwrap(); + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(result, expected, &mut ctx); + } + + #[rstest] + #[case::empty(Vec::<&'static str>::new(), Some(false))] + #[case::nonempty(vec!["a"], None)] + fn test_constant_string_list_all_null_needles_reduces_to_constant( + #[case] members: Vec<&'static str>, + #[case] expected: Option, + ) { + let mut ctx = array_session().create_execution_ctx(); + let list = Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + members.into_iter().map(Scalar::from).collect(), + Nullability::NonNullable, + ); + let needles = VarBinArray::from_iter( + [None::<&str>, None, None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + + let result = needles + .apply(&list_contains(lit(list), root())) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + assert!(result.is::()); + assert_arrays_eq!( + result, + BoolArray::from_iter([expected, expected, expected]), + &mut ctx + ); } #[test] @@ -818,6 +1014,27 @@ mod tests { assert_arrays_eq!(contains, expected, &mut ctx); } + #[test] + fn test_nonconstant_all_null_needles() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lists = ListArray::try_new( + PrimitiveArray::from_iter([1i32]).into_array(), + PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(), + Validity::Array(BoolArray::from(BitBuffer::from(vec![true, true, false])).into_array()), + )? + .into_array(); + let needles = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); + + let contains = ListContains::try_new(lists, needles)?.into_array(); + + assert_arrays_eq!( + contains, + BoolArray::from_iter([Some(false), None, None]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_list_array_element() { let mut ctx = array_session().create_execution_ctx(); @@ -887,8 +1104,9 @@ mod tests { ); assert_arrays_eq!(result, expected, &mut ctx); - // Searching for non-null - let expr2 = list_contains(root(), lit(42i32)); + // Null primitive payloads default to zero. Searching for zero verifies that invalid + // comparison values do not become matches. + let expr2 = list_contains(root(), lit(0i32)); let result2 = list_array.into_array().apply(&expr2).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]); diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 2ab975ecfd7..c227663f10d 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -384,12 +384,16 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } }) .try_collect()?; + let Some(first) = list_elements.first() else { + return Err(exec_datafusion_err!("Cannot push down an empty IN list")); + }; + if list_elements.iter().any(Scalar::is_null) { + return Err(exec_datafusion_err!( + "Cannot push down an IN list that contains null" + )); + } - let list = Scalar::list( - list_elements[0].dtype().clone(), - list_elements, - Nullability::Nullable, - ); + let list = Scalar::list(first.dtype().clone(), list_elements, Nullability::Nullable); let expr = list_contains(lit(list), value); return Ok(if in_list.negated() { not(expr) } else { expr }); @@ -433,6 +437,17 @@ impl ExpressionConvertor for DefaultExpressionConvertor { return Ok(TreeNodeRecursion::Stop); } + if let Some(in_list) = node.downcast_ref::() + && !can_in_list_be_pushed_down(in_list, input_schema) + { + scan_projection.extend(collect_columns(node).into_iter().map(|column| { + (column.name().to_string(), get_item(column.name(), root())) + })); + + leftover_projection.push(projection_expr.clone()); + return Ok(TreeNodeRecursion::Stop); + } + // DataFusion assumes different decimal types can be coerced. // Vortex expects a perfect match so we don't push it down. if let Some(binary_expr) = node.downcast_ref::() @@ -555,11 +570,7 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo } else if let Some(is_not_null) = expr.downcast_ref::() { can_be_pushed_down_impl(is_not_null.arg(), schema) } else if let Some(in_list) = expr.downcast_ref::() { - can_be_pushed_down_impl(in_list.expr(), schema) - && in_list - .list() - .iter() - .all(|e| can_be_pushed_down_impl(e, schema)) + can_in_list_be_pushed_down(in_list, schema) } else if let Some(scalar_fn) = expr.downcast_ref::() { can_scalar_fn_be_pushed_down(scalar_fn, schema) } else if let Some(case_expr) = expr.downcast_ref::() { @@ -570,9 +581,27 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo } } -/// Checks if an expression type is one that convert() can handle. -/// This is less restrictive than can_be_pushed_down since it only checks -/// expression types, not data type support. +fn can_in_list_be_pushed_down(in_list: &df_expr::InListExpr, schema: &Schema) -> bool { + can_be_pushed_down_impl(in_list.expr(), schema) + && is_convertible_in_list(in_list) + && in_list + .list() + .iter() + .all(|expr| can_be_pushed_down_impl(expr, schema)) +} + +fn is_convertible_in_list(in_list: &df_expr::InListExpr) -> bool { + !in_list.list().is_empty() + && in_list.list().iter().all(|expr| { + expr.downcast_ref::() + .is_some_and(|literal| !literal.value().is_null()) + }) +} + +/// Checks if an expression can be converted without schema information. +/// +/// This is less restrictive than `can_be_pushed_down_impl` because it does not check data type +/// support. fn is_convertible_expr(expr: &Arc) -> bool { // Expression types that convert() handles expr.downcast_ref::().is_some() @@ -584,7 +613,9 @@ fn is_convertible_expr(expr: &Arc) -> bool { .is_some_and(|e| is_convertible_expr(e.expr())) || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() + || expr + .downcast_ref::() + .is_some_and(is_convertible_in_list) || expr.downcast_ref::().is_some_and(|sf| { ScalarFunctionExpr::try_downcast_func::(sf).is_some() || ScalarFunctionExpr::try_downcast_func::(sf).is_some() @@ -787,6 +818,19 @@ mod tests { ) } + fn in_list_expr( + values: impl IntoIterator, + negated: bool, + schema: &Schema, + ) -> Arc { + let value = Arc::new(df_expr::Column::new("id", 0)) as Arc; + let list = values + .into_iter() + .map(|value| Arc::new(df_expr::Literal::new(value)) as Arc) + .collect(); + Arc::new(df_expr::InListExpr::try_new(value, list, negated, schema).unwrap()) + } + #[test] fn test_make_vortex_predicate_empty() { let expr_convertor = DefaultExpressionConvertor::default(); @@ -1079,6 +1123,53 @@ mod tests { assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema)); } + #[rstest] + #[case::in_nonempty(vec![ScalarValue::Int32(Some(1))], false, true)] + #[case::not_in_nonempty(vec![ScalarValue::Int32(Some(1))], true, true)] + #[case::in_empty(vec![], false, false)] + #[case::not_in_empty(vec![], true, false)] + #[case::in_null(vec![ScalarValue::Int32(None)], false, false)] + #[case::not_in_null(vec![ScalarValue::Int32(None)], true, false)] + fn test_can_be_pushed_down_in_list( + #[case] values: Vec, + #[case] negated: bool, + #[case] expected: bool, + test_schema: Schema, + ) { + let expression = in_list_expr(values, negated, &test_schema); + let cast_expression = Arc::new(df_expr::CastExpr::new_with_target_field( + Arc::clone(&expression), + Arc::new(Field::new("matches", DataType::Boolean, true)), + None, + )) as Arc; + + assert_eq!(can_be_pushed_down_impl(&expression, &test_schema), expected); + assert_eq!( + can_be_pushed_down_impl(&cast_expression, &test_schema), + expected + ); + } + + #[rstest] + #[case::empty(vec![], false)] + #[case::null(vec![ScalarValue::Int32(None)], true)] + fn test_split_projection_keeps_unsafe_in_list_in_datafusion( + #[case] values: Vec, + #[case] negated: bool, + test_schema: Schema, + ) { + let expression = in_list_expr(values, negated, &test_schema); + let source_projection = + ProjectionExprs::new([ProjectionExpr::new(expression, "matches".to_string())]); + let output_schema = Schema::new(vec![Field::new("matches", DataType::Boolean, true)]); + + let processed = DefaultExpressionConvertor::default() + .split_projection(source_projection.clone(), &test_schema, &output_schema) + .unwrap(); + + assert_eq!(processed.leftover_projection, source_projection); + } + #[rstest] fn test_can_be_pushed_down_like_supported(test_schema: Schema) { let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 51aeebebb6d..cc5d490350f 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -386,6 +386,31 @@ fn can_push_cast(cast: &duckdb::BoundCast<'_>, target: &duckdb::LogicalTypeRef) !cast.is_try && target.is_primitive_integer() && cast.child.return_type().is_primitive_integer() } +fn can_push_in_list(operator: &BoundOperator<'_>) -> bool { + let mut children = operator.children(); + let Some(element) = children.next() else { + return false; + }; + if !can_push_expression(element) { + return false; + } + + let mut member_count = 0; + for child in children { + let Some(BoundConstant(constant)) = child.as_class() else { + return false; + }; + let Ok(member) = Scalar::try_from(constant.value) else { + return false; + }; + if member.is_null() { + return false; + } + member_count += 1; + } + member_count > 0 +} + // Called before pushdown_complex_filter or a table filter expression call. // As we support complex filter pushdown, Duckdb pushes expressions to Vortex. // However, it doesn't know what type of expressions we can handle. Here we list @@ -433,13 +458,18 @@ pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { // columns are native. } ExpressionClass::BoundOperator(op) => { + if matches!( + op.op, + DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_IN + | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_NOT_IN + ) { + return can_push_in_list(&op); + } if !matches!( op.op, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_OPERATOR_NOT | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_OPERATOR_IS_NULL | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_OPERATOR_IS_NOT_NULL - | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_IN - | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_NOT_IN ) { return false; } @@ -740,6 +770,9 @@ fn try_from_compare_in( else { return Ok(None); }; + if list_elements.is_empty() || list_elements.iter().any(Scalar::is_null) { + return Ok(None); + } let list = Scalar::list( Arc::new(list_elements[0].dtype().clone()), list_elements, diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 0876be1ca4c..49b97fa2f80 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -281,6 +281,20 @@ fn test_issue_5927_not_in_does_not_panic() { assert_eq!(sum, -4); } +#[test] +fn test_not_in_with_null_is_not_pushed_down() { + let file = RUNTIME.block_on(async { + let numbers = buffer![1i32, 42, 100, -5, 0]; + write_single_column_vortex_file("number", numbers).await + }); + let count: i64 = scan_vortex_file_single_row::( + file, + "SELECT COUNT(*) FROM ? WHERE number NOT IN (42, NULL)", + 0, + ); + assert_eq!(count, 0); +} + #[test] fn test_vortex_scan_floats() { let file = RUNTIME.block_on(async { From 8db2c3db5b558be5f1e95203090e227158b79db2 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Tue, 1 Sep 2026 20:09:37 -0400 Subject: [PATCH 2/6] feat(velox): Add engine-owned scan adapter Signed-off-by: Will Manning --- Cargo.lock | 20 + Cargo.toml | 1 + vortex-ffi/src/array.rs | 7 + vortex-ffi/src/data_source.rs | 10 +- vortex-ffi/src/dtype.rs | 7 + vortex-ffi/src/expression.rs | 26 +- vortex-ffi/src/lib.rs | 10 + vortex-ffi/src/scan.rs | 16 + vortex-velox/Cargo.toml | 47 + vortex-velox/README.md | 69 ++ vortex-velox/cinclude/vortex_velox.h | 347 ++++++ vortex-velox/src/api.rs | 855 ++++++++++++++ vortex-velox/src/array.rs | 856 ++++++++++++++ vortex-velox/src/lib.rs | 124 ++ vortex-velox/src/projection.rs | 126 ++ vortex-velox/src/read_at.rs | 1012 +++++++++++++++++ vortex-velox/src/schema.rs | 66 ++ vortex-velox/src/source.rs | 278 +++++ vortex-velox/src/visitor.rs | 706 ++++++++++++ vortex-velox/tests/abi_contract.rs | 208 ++++ vortex-velox/tests/velox_include_contract.cpp | 42 + 21 files changed, 4831 insertions(+), 2 deletions(-) create mode 100644 vortex-velox/Cargo.toml create mode 100644 vortex-velox/README.md create mode 100644 vortex-velox/cinclude/vortex_velox.h create mode 100644 vortex-velox/src/api.rs create mode 100644 vortex-velox/src/array.rs create mode 100644 vortex-velox/src/lib.rs create mode 100644 vortex-velox/src/projection.rs create mode 100644 vortex-velox/src/read_at.rs create mode 100644 vortex-velox/src/schema.rs create mode 100644 vortex-velox/src/source.rs create mode 100644 vortex-velox/src/visitor.rs create mode 100644 vortex-velox/tests/abi_contract.rs create mode 100644 vortex-velox/tests/velox_include_contract.cpp diff --git a/Cargo.lock b/Cargo.lock index 343f8413bc3..4660765b51f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11690,6 +11690,26 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "vortex-velox" +version = "0.1.0" +dependencies = [ + "arrow-array 59.2.0", + "arrow-buffer 59.2.0", + "arrow-data 59.2.0", + "arrow-schema 59.2.0", + "bytes", + "futures", + "rstest", + "vortex", + "vortex-array", + "vortex-arrow", + "vortex-buffer", + "vortex-error", + "vortex-ffi", + "vortex-io", +] + [[package]] name = "vortex-web-wasm" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index f7fbd95300a..efc9dc0e0c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ members = [ "vortex-cuda/macros", "vortex-cuda/nvcomp", "vortex-ffi", + "vortex-velox", "fuzz", "vortex-jni", "vortex-python-abi", diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 3d4b9758567..58c69263bfa 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -68,6 +68,13 @@ box_wrapper!( vx_array ); +/// Create an FFI array handle from an owned Vortex array. +/// +/// Layered FFI crates use this function when their host-specific scan path produces an array. +pub fn vx_array_new_with(array: ArrayRef) -> *const vx_array { + vx_array::new(array) +} + /// Borrow the [`ArrayRef`] behind a [`vx_array`] handle, erroring on a null pointer. /// /// A building block for FFI crates layered on top of the base Vortex C API. diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index 50d31b3036d..87c0bcfe36f 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -38,7 +38,15 @@ box_wrapper!( /// /// Copying a vx_data_source via vx_data_source_clone is a cheap operation. MultiLayoutDataSource, - vx_data_source); + vx_data_source +); + +/// Create an FFI data-source handle from a configured multi-layout data source. +/// +/// Layered FFI crates use this function after their host-specific I/O adapter constructs a source. +pub fn vx_data_source_new_with(data_source: MultiLayoutDataSource) -> *const vx_data_source { + vx_data_source::new(data_source) +} /// Options for creating a data source. #[repr(C)] diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index 7156bc5334a..c8121693e69 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -34,6 +34,13 @@ box_wrapper!( vx_dtype ); +/// Create an FFI dtype handle from an owned Vortex dtype. +/// +/// Layered FFI crates use this function after they validate their host-specific type identifiers. +pub fn vx_dtype_new_with(dtype: DType) -> *const vx_dtype { + vx_dtype::new(dtype) +} + /// The variant tag for a Vortex data type. #[non_exhaustive] #[repr(C)] diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index f15db42006f..9b9d6d46248 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -40,7 +40,31 @@ box_wrapper!( /// Operations on expressions don't take ownership of input values, and so /// input values must be freed by the caller. Expression, - vx_expression); + vx_expression +); + +/// Create an FFI expression handle from an owned Vortex expression. +/// +/// Layered FFI crates use this function for host-specific expression constructors. +pub fn vx_expression_new_with(expression: Expression) -> *mut vx_expression { + vx_expression::new(expression) +} + +/// Borrow an expression from a layered FFI crate. +/// +/// # Safety +/// +/// `expression` must point to a live expression handle for the returned reference lifetime. +pub unsafe fn vx_expression_ref<'a>( + expression: *const vx_expression, +) -> vortex::error::VortexResult<&'a Expression> { + let expression = unsafe { + expression + .as_ref() + .ok_or_else(|| vortex::error::vortex_err!("Vortex expression must not be null"))? + }; + Ok(&expression.0) +} /// Create a root expression. A root expression, applied to an array in /// vx_array_apply, takes the array itself as opposed to functions like diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1ef423bb383..9e1c2f0b213 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -27,14 +27,24 @@ use std::sync::Arc; use std::sync::LazyLock; pub use array::vx_array; +pub use array::vx_array_new_with; pub use array::vx_array_ref; +pub use data_source::vx_data_source; +pub use data_source::vx_data_source_new_with; pub use dtype::vx_dtype; +pub use dtype::vx_dtype_new_with; pub use error::try_or; pub use error::vx_error; pub use error::vx_error_free; +pub use expression::vx_expression; +pub use expression::vx_expression_new_with; +pub use expression::vx_expression_ref; pub use log::vx_log_level; +pub use scalar::vx_scalar; +pub use scan::vx_data_source_scan_with; pub use scan::vx_partition; pub use scan::vx_partition_into_array_stream; +pub use scan::vx_scan; pub use session::vx_session; pub use session::vx_session_free; pub use session::vx_session_new_with; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index d5ea24a9ac1..798e2bbbe42 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -225,6 +225,22 @@ fn write_estimate>(estimate: Precision, out: &mut vx_estimate) { } } +/// Start a scan from a request that a layered FFI crate already validated. +/// +/// # Safety +/// +/// `data_source` must point to a live data-source handle created by this crate. +pub unsafe fn vx_data_source_scan_with( + data_source: *const vx_data_source, + request: ScanRequest, +) -> VortexResult<*mut vx_scan> { + vortex_ensure!(!data_source.is_null(), "null vx_data_source"); + RUNTIME.block_on(async { + let scan = vx_data_source::as_ref(data_source).scan(request).await?; + Ok(vx_scan::new(VxScan::Pending(scan))) + }) +} + /// Scan a data source. /// /// A scan may be consumed only once. diff --git a/vortex-velox/Cargo.toml b/vortex-velox/Cargo.toml new file mode 100644 index 00000000000..09ee82cdeed --- /dev/null +++ b/vortex-velox/Cargo.toml @@ -0,0 +1,47 @@ +[package] +name = "vortex-velox" +description = "Native adapter contract between Vortex and Velox" +readme = "README.md" +publish = false +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +keywords = { workspace = true } +include = [ + "cinclude/*.h", + "src/**/*.rs", + "tests/abi_contract.rs", + "tests/velox_include_contract.cpp", + "Cargo.toml", + "README.md", +] +edition = { workspace = true } +rust-version = { workspace = true } +categories = { workspace = true } + +[dependencies] +arrow-array = { workspace = true } +arrow-buffer = { workspace = true } +arrow-data = { workspace = true } +arrow-schema = { workspace = true } +bytes = { workspace = true } +futures = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-arrow = { workspace = true } +vortex-error = { workspace = true } +vortex-ffi = { path = "../vortex-ffi" } +vortex-io = { workspace = true } +vortex = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } + +[lib] +name = "vortex_velox" +crate-type = ["rlib", "staticlib"] + +[lints] +workspace = true diff --git a/vortex-velox/README.md b/vortex-velox/README.md new file mode 100644 index 00000000000..a96950f44e3 --- /dev/null +++ b/vortex-velox/README.md @@ -0,0 +1,69 @@ +# vortex-velox + +`vortex-velox` defines the native adapter contract between Vortex and Velox. +The crate keeps Vortex API use inside the Vortex repository. Velox implements +the engine side of the versioned C ABI. + +The crate is private and ships as part of a pinned Vortex source release. It is +not a general Vortex C API. + +Native primitive visits copy values and validity into compact owned buffers. +The value allocation uses `uint64_t` alignment. The ABI reports the exact +allocation size that the owner retains. This contract avoids an understated +charge for a slice of a larger Vortex allocation. + +Before Arrow fallback conversion, the host reserves a conservative byte count. +The adapter stops before Arrow allocations if the host rejects the reservation. +After conversion, the adapter refunds the difference from retained payload +capacities. If actual capacity exceeds the reservation, the adapter requests +the difference before it returns outputs. Arrow release frees the final charge. + +## Contract boundary + +The adapter exposes a versioned C ABI in `cinclude/vortex_velox.h`. Velox calls +only `vx_velox_*` symbols. The static archive can contain general `vx_*` symbols +from linked Vortex FFI objects. Opaque handle layouts stay inside Vortex. + +The adapter accepts host callbacks for random reads. Velox can implement those +callbacks with `dwio::common::BufferedInput`, so Vortex uses the existing cache +and file-system path. Vortex checks the host cancellation callback before each +read callback. This check does not interrupt an active callback or CPU work. + +Vortex can call read callbacks concurrently. The host callback context must be +thread-safe. Each thread owns its callback error string until its next callback. +Callbacks must catch exceptions and must not unwind across the C ABI. + +Each source reports natural row ranges. Velox maps byte ownership to these row +ranges. The adapter maps projection and filter expressions to Vortex scan +requests. Metadata exclusion returns only splits that Vortex proves cannot +contain a match. + +Source schemas cross the boundary through the Arrow C Data Interface. The +adapter does not maintain a second recursive type protocol. + +Native array visits cover primitive arrays and structural wrappers. The Arrow +C Data fallback covers arrays without a native visit. Both paths transfer +ownership through explicit release callbacks. + +## Build + +Build the static adapter library from the workspace root: + +```bash +cargo build --locked --package vortex-velox +``` + +The Vortex `rust-toolchain.toml` file selects the Rust toolchain. Velox can use +a pinned Vortex archive or a local checkout. + +Run the adapter tests with: + +```bash +cargo test --locked --package vortex-velox +``` + +The package tests compile the C and C++ header contracts. Production builds do +not compile or link contract helper symbols. + +The crate does not construct Velox vectors. Velox owns vector construction, +lazy-load policy, value hooks, mutation semantics, and memory accounting. diff --git a/vortex-velox/cinclude/vortex_velox.h b/vortex-velox/cinclude/vortex_velox.h new file mode 100644 index 00000000000..eeae773bbca --- /dev/null +++ b/vortex-velox/cinclude/vortex_velox.h @@ -0,0 +1,347 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include +#include +#include + +#include "vortex.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Velox must call only vx_velox_* adapter symbols. The static archive can also + * contain general vx_* symbols from linked Vortex FFI objects. + */ + +#define VX_VELOX_ABI_VERSION 1u +#define VX_VELOX_CAPABILITY_BATCH_READ (UINT64_C(1) << 0) +#define VX_VELOX_CAPABILITY_CALLBACK_SOURCE (UINT64_C(1) << 1) +#define VX_VELOX_CAPABILITY_NATURAL_SPLITS (UINT64_C(1) << 2) +#define VX_VELOX_CAPABILITY_PRIMITIVE_VISITOR (UINT64_C(1) << 3) +#define VX_VELOX_CAPABILITY_ARROW_SCHEMA (UINT64_C(1) << 4) +#define VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT (UINT64_C(1) << 5) +#define VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION (UINT64_C(1) << 6) +#define VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING (UINT64_C(1) << 7) +/* Vortex checks cancellation before each host read callback. */ +#define VX_VELOX_CAPABILITY_READ_CANCELLATION (UINT64_C(1) << 8) + +typedef struct vx_velox_read_at vx_velox_read_at; +typedef struct vx_velox_source vx_velox_source; + +typedef uint32_t vx_velox_ptype; +#define VX_VELOX_PTYPE_U8 UINT32_C(0) +#define VX_VELOX_PTYPE_U16 UINT32_C(1) +#define VX_VELOX_PTYPE_U32 UINT32_C(2) +#define VX_VELOX_PTYPE_U64 UINT32_C(3) +#define VX_VELOX_PTYPE_I8 UINT32_C(4) +#define VX_VELOX_PTYPE_I16 UINT32_C(5) +#define VX_VELOX_PTYPE_I32 UINT32_C(6) +#define VX_VELOX_PTYPE_I64 UINT32_C(7) +#define VX_VELOX_PTYPE_F16 UINT32_C(8) +#define VX_VELOX_PTYPE_F32 UINT32_C(9) +#define VX_VELOX_PTYPE_F64 UINT32_C(10) + +typedef uint32_t vx_velox_binary_operator; +#define VX_VELOX_OPERATOR_EQ UINT32_C(0) +#define VX_VELOX_OPERATOR_NOT_EQ UINT32_C(1) +#define VX_VELOX_OPERATOR_GT UINT32_C(2) +#define VX_VELOX_OPERATOR_GTE UINT32_C(3) +#define VX_VELOX_OPERATOR_LT UINT32_C(4) +#define VX_VELOX_OPERATOR_LTE UINT32_C(5) +#define VX_VELOX_OPERATOR_KLEENE_AND UINT32_C(6) +#define VX_VELOX_OPERATOR_KLEENE_OR UINT32_C(7) + +typedef uint32_t vx_velox_scan_selection_include; +#define VX_VELOX_SELECTION_ALL UINT32_C(0) +#define VX_VELOX_SELECTION_INCLUDE UINT32_C(1) +#define VX_VELOX_SELECTION_EXCLUDE UINT32_C(2) + +typedef struct vx_velox_scan_selection { + const uint64_t *indices; + size_t length; + vx_velox_scan_selection_include include; +} vx_velox_scan_selection; + +typedef struct vx_velox_scan_options { + size_t struct_size; + uint32_t abi_version; + const vx_expression *projection; + const vx_expression *filter; + uint64_t row_range_begin; + uint64_t row_range_end; + vx_velox_scan_selection selection; + uint64_t limit; + bool ordered; +} vx_velox_scan_options; + +typedef struct vx_velox_read_request { + size_t struct_size; + uint64_t offset; + size_t length; + size_t alignment; +} vx_velox_read_request; + +typedef struct vx_velox_buffer { + size_t struct_size; + const uint8_t *data; + size_t length; + void *owner; + void (*release)(void *owner); +} vx_velox_buffer; + +/** + * Callbacks for positional reads through the host engine. + * + * Vortex can call size, read_ranges, is_cancelled, and last_error concurrently. + * The context and callbacks must be thread-safe. A non-zero concurrency value + * limits requests in one callback and gives Vortex a scheduling hint. It is not + * a synchronization guarantee. + * + * last_error must return the calling thread's most recent callback error. The + * returned string must remain valid until the next callback on that thread. + * Every callback must catch C++ exceptions. No callback can unwind across this + * C ABI. release_context runs after the final callback and can run on any + * thread. + * + * Vortex checks is_cancelled before each read_ranges call. The check does not + * interrupt an active callback or CPU work. + */ +typedef struct vx_velox_read_at_callbacks { + size_t struct_size; + uint32_t abi_version; + void *context; + int32_t (*size)(void *context, uint64_t *size_out); + int32_t (*read_ranges)(void *context, + const vx_velox_read_request *requests, + size_t request_count, + vx_velox_buffer *outputs); + const char *(*last_error)(void *context); + void (*release_context)(void *context); + int32_t (*is_cancelled)(void *context); + size_t concurrency; +} vx_velox_read_at_callbacks; + +typedef struct vx_velox_natural_split { + size_t struct_size; + uint64_t row_begin; + uint64_t row_end; +} vx_velox_natural_split; + +typedef uint32_t vx_velox_primitive_type; +#define VX_VELOX_PRIMITIVE_U8 UINT32_C(0) +#define VX_VELOX_PRIMITIVE_U16 UINT32_C(1) +#define VX_VELOX_PRIMITIVE_U32 UINT32_C(2) +#define VX_VELOX_PRIMITIVE_U64 UINT32_C(3) +#define VX_VELOX_PRIMITIVE_I8 UINT32_C(4) +#define VX_VELOX_PRIMITIVE_I16 UINT32_C(5) +#define VX_VELOX_PRIMITIVE_I32 UINT32_C(6) +#define VX_VELOX_PRIMITIVE_I64 UINT32_C(7) +#define VX_VELOX_PRIMITIVE_F16 UINT32_C(8) +#define VX_VELOX_PRIMITIVE_F32 UINT32_C(9) +#define VX_VELOX_PRIMITIVE_F64 UINT32_C(10) + +typedef uint32_t vx_velox_validity_kind; +#define VX_VELOX_VALIDITY_NON_NULLABLE UINT32_C(0) +#define VX_VELOX_VALIDITY_ALL_VALID UINT32_C(1) +#define VX_VELOX_VALIDITY_ALL_INVALID UINT32_C(2) +#define VX_VELOX_VALIDITY_BITMAP UINT32_C(3) + +typedef struct vx_velox_buffer_owner { + size_t struct_size; + const void *owner; + void (*retain)(const void *owner); + void (*release)(const void *owner); + size_t retained_bytes; +} vx_velox_buffer_owner; + +/** + * A compact primitive payload and its owner. + * + * Vortex copies values into an allocation with uint64_t alignment. The values + * allocation rounds values_length up to that alignment. Vortex copies a bitmap + * into a compact byte allocation with validity_bit_offset set to zero. + * buffers.retained_bytes is the exact sum of these allocation sizes. + * + * The pointers remain valid through visit_primitive. The host must call retain + * before it stores a pointer beyond that callback. + */ +typedef struct vx_velox_primitive_view { + size_t struct_size; + vx_velox_primitive_type primitive_type; + size_t length; + const uint8_t *values; + size_t values_length; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t values_alignment; + size_t validity_alignment; +} vx_velox_primitive_view; + +typedef struct vx_velox_visit_request { + size_t struct_size; + const uint64_t *rows; + size_t row_count; +} vx_velox_visit_request; + +/** + * Host callbacks for one Vortex array visit. + * + * One array visit calls visit_primitive synchronously. If the host shares this + * table between simultaneous visits, callbacks can occur concurrently. + * last_error returns the calling thread's most recent visitor error. Its string + * remains valid until the next callback on that thread. + * + * Every callback must catch C++ exceptions. No callback can unwind across this + * C ABI. The host owns context and must keep it live until each visit returns. + */ +typedef struct vx_velox_visitor { + size_t struct_size; + uint32_t abi_version; + void *context; + int32_t (*visit_primitive)(void *context, const vx_velox_primitive_view *view); + const char *(*last_error)(void *context); +} vx_velox_visitor; + +/** + * Host memory callbacks for one Arrow C Data export. + * + * Before Arrow conversion, Vortex requests a conservative reservation through + * report_allocation. A rejection stops the export before Arrow allocations. + * Vortex calls report_free after conversion to refund unused reservation bytes. + * The remaining charge equals retained Arrow payload capacities. It excludes + * the schema and small Arrow C Data metadata allocations. + * If the actual charge exceeds the reservation, Vortex requests the difference + * before it returns outputs. A rejection aborts the export and frees the data. + * + * A final Arrow release calls report_free for the remaining charge. It also + * calls release_context. These calls can occur on any thread. All callbacks + * must be thread-safe and must not unwind across the C ABI. + * + * last_error returns the calling thread's most recent allocation error. The + * string remains valid until the next callback on that thread. + */ +typedef struct vx_velox_arrow_memory_callbacks { + size_t struct_size; + uint32_t abi_version; + void *context; + void (*retain_context)(void *context); + void (*release_context)(void *context); + int32_t (*report_allocation)(void *context, size_t retained_bytes); + void (*report_free)(void *context, size_t retained_bytes); + const char *(*last_error)(void *context); +} vx_velox_arrow_memory_callbacks; + +uint32_t vx_velox_abi_version(void); +uint64_t vx_velox_capabilities(void); + +vx_view vx_velox_error_message(const vx_error *error); +void vx_velox_error_free(const vx_error *error); +vx_session *vx_velox_session_new(void); +void vx_velox_session_free(const vx_session *session); + +const vx_dtype *vx_velox_dtype_new_primitive(vx_velox_ptype ptype, + bool nullable, + vx_error **error_out); +void vx_velox_dtype_free(const vx_dtype *dtype); +vx_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); +vx_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); +vx_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); +vx_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); +vx_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); +vx_scalar *vx_velox_scalar_new_f32(float value, bool nullable); +vx_scalar *vx_velox_scalar_new_f64(double value, bool nullable); +vx_scalar *vx_velox_scalar_new_utf8(vx_view value, bool nullable, vx_error **error_out); +vx_scalar * +vx_velox_scalar_new_binary(const uint8_t *data, size_t length, bool nullable, vx_error **error_out); +vx_scalar *vx_velox_scalar_new_list(const vx_dtype *element_dtype, + const vx_scalar *const *elements, + size_t length, + bool nullable, + vx_error **error_out); +void vx_velox_scalar_free(const vx_scalar *scalar); + +vx_expression *vx_velox_expression_root(void); +vx_expression *vx_velox_expression_literal(const vx_scalar *scalar, vx_error **error_out); +vx_expression *vx_velox_expression_get_item(vx_view name, const vx_expression *child); +vx_expression *vx_velox_expression_binary(vx_velox_binary_operator operation, + const vx_expression *left, + const vx_expression *right, + vx_error **error_out); +vx_expression *vx_velox_expression_and(const vx_expression *const *expressions, size_t length); +vx_expression *vx_velox_expression_or(const vx_expression *const *expressions, size_t length); +vx_expression *vx_velox_expression_not(const vx_expression *child); +vx_expression *vx_velox_expression_is_null(const vx_expression *child); +vx_expression *vx_velox_expression_list_contains(const vx_expression *list, const vx_expression *value); +void vx_velox_expression_free(const vx_expression *expression); +vx_expression *vx_velox_expression_select_with_row_index(const vx_view *names, + size_t length, + vx_view row_index_name, + vx_error **error_out); + +/* + * On success, the reader owns context and calls release_context once. On + * failure, the caller still owns context. + */ +vx_velox_read_at *vx_velox_read_at_new(const vx_velox_read_at_callbacks *callbacks, vx_error **error_out); +void vx_velox_read_at_free(vx_velox_read_at *reader); +uint64_t vx_velox_read_at_size(const vx_velox_read_at *reader, vx_error **error_out); + +vx_velox_source * +vx_velox_source_new(const vx_session *session, const vx_velox_read_at *reader, vx_error **error_out); +void vx_velox_source_free(vx_velox_source *source); +uint64_t vx_velox_source_row_count(const vx_velox_source *source); +uint64_t vx_velox_source_file_size(const vx_velox_source *source); +int32_t vx_velox_source_export_schema(const vx_velox_source *source, + FFI_ArrowSchema *schema_out, + vx_error **error_out); +size_t vx_velox_source_natural_split_count(const vx_velox_source *source); +int32_t vx_velox_source_natural_split_at(const vx_velox_source *source, + size_t index, + vx_velox_natural_split *split_out, + vx_error **error_out); +int32_t vx_velox_source_prune_natural_splits(const vx_velox_source *source, + const vx_expression *expression, + size_t first_split, + size_t split_count, + uint8_t *pruned_out, + vx_error **error_out); +const vx_data_source *vx_velox_source_data_source(const vx_velox_source *source, vx_error **error_out); + +void vx_velox_data_source_free(const vx_data_source *data_source); +vx_scan *vx_velox_data_source_scan(const vx_data_source *data_source, + const vx_velox_scan_options *options, + vx_error **error_out); +void vx_velox_scan_free(const vx_scan *scan); +vx_partition *vx_velox_scan_next_partition(vx_scan *scan, vx_error **error_out); +void vx_velox_partition_free(const vx_partition *partition); +const vx_array *vx_velox_partition_next(vx_partition *partition, vx_error **error_out); +void vx_velox_array_free(const vx_array *array); +size_t vx_velox_array_len(const vx_array *array); +const vx_array *vx_velox_array_slice(const vx_array *array, size_t begin, size_t end, vx_error **error_out); +const vx_array *vx_velox_array_get_field(const vx_session *session, + const vx_array *array, + size_t index, + vx_error **error_out); +size_t vx_velox_array_invalid_count(const vx_session *session, const vx_array *array, vx_error **error_out); +int32_t vx_velox_array_visit(const vx_session *session, + const vx_array *array, + const vx_velox_visit_request *request, + const vx_velox_visitor *visitor, + vx_error **error_out); +int32_t vx_velox_array_export_arrow(const vx_session *session, + const vx_array *array, + const vx_velox_arrow_memory_callbacks *memory_callbacks, + FFI_ArrowSchema *schema_out, + FFI_ArrowArray *array_out, + vx_error **error_out); + +#ifdef __cplusplus +} +#endif diff --git a/vortex-velox/src/api.rs b/vortex-velox/src/api.rs new file mode 100644 index 00000000000..485b347dd03 --- /dev/null +++ b/vortex-velox/src/api.rs @@ -0,0 +1,855 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::size_of; +use std::ptr; +use std::slice; + +use vortex::buffer::Buffer; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::dtype::PType; +use vortex::expr::root; +use vortex::scalar_fn::ScalarFnVTableExt; +use vortex::scalar_fn::fns::binary::Binary; +use vortex::scalar_fn::fns::operators::Operator; +use vortex::scan::ScanRequest; +use vortex::scan::selection::Selection; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_error::vortex_bail; +use vortex_ffi::try_or; +use vortex_ffi::vx_array; +use vortex_ffi::vx_data_source; +use vortex_ffi::vx_data_source_scan_with; +use vortex_ffi::vx_dtype; +use vortex_ffi::vx_dtype_new_with; +use vortex_ffi::vx_error; +use vortex_ffi::vx_expression; +use vortex_ffi::vx_expression_new_with; +use vortex_ffi::vx_expression_ref; +use vortex_ffi::vx_partition; +use vortex_ffi::vx_scalar; +use vortex_ffi::vx_scan; +use vortex_ffi::vx_session; +use vortex_ffi::vx_view; + +// The base FFI wrappers are opaque C handles despite their private Rust payloads. +#[allow(improper_ctypes)] +mod ffi { + use super::*; + + unsafe extern "C-unwind" { + pub fn vx_error_message(error: *const vx_error) -> vx_view; + pub fn vx_error_free(error: *const vx_error); + pub fn vx_session_new() -> *mut vx_session; + pub fn vx_session_free(session: *const vx_session); + pub fn vx_dtype_free(dtype: *const vx_dtype); + pub fn vx_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_i8(value: i8, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_i16(value: i16, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_i32(value: i32, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_i64(value: i64, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_f32(value: f32, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_f64(value: f64, nullable: bool) -> *mut vx_scalar; + pub fn vx_scalar_new_utf8( + value: vx_view, + nullable: bool, + error_out: *mut *mut vx_error, + ) -> *mut vx_scalar; + pub fn vx_scalar_new_binary( + data: *const u8, + length: usize, + nullable: bool, + error_out: *mut *mut vx_error, + ) -> *mut vx_scalar; + pub fn vx_scalar_new_list( + element_dtype: *const vx_dtype, + elements: *const *const vx_scalar, + length: usize, + nullable: bool, + error_out: *mut *mut vx_error, + ) -> *mut vx_scalar; + pub fn vx_scalar_free(scalar: *const vx_scalar); + pub fn vx_expression_literal( + scalar: *const vx_scalar, + error_out: *mut *mut vx_error, + ) -> *mut vx_expression; + pub fn vx_expression_free(expression: *const vx_expression); + pub fn vx_data_source_free(data_source: *const vx_data_source); + pub fn vx_scan_free(scan: *const vx_scan); + pub fn vx_scan_next_partition( + scan: *mut vx_scan, + error_out: *mut *mut vx_error, + ) -> *mut vx_partition; + pub fn vx_partition_free(partition: *const vx_partition); + pub fn vx_partition_next( + partition: *mut vx_partition, + error_out: *mut *mut vx_error, + ) -> *const vx_array; + pub fn vx_array_free(array: *const vx_array); + pub fn vx_array_len(array: *const vx_array) -> usize; + pub fn vx_array_slice( + array: *const vx_array, + begin: usize, + end: usize, + error_out: *mut *mut vx_error, + ) -> *const vx_array; + } + + unsafe extern "C" { + pub fn vx_expression_root() -> *mut vx_expression; + pub fn vx_expression_get_item( + name: vx_view, + child: *const vx_expression, + ) -> *mut vx_expression; + pub fn vx_expression_and( + expressions: *const *const vx_expression, + length: usize, + ) -> *mut vx_expression; + pub fn vx_expression_or( + expressions: *const *const vx_expression, + length: usize, + ) -> *mut vx_expression; + pub fn vx_expression_not(child: *const vx_expression) -> *mut vx_expression; + pub fn vx_expression_is_null(child: *const vx_expression) -> *mut vx_expression; + pub fn vx_expression_list_contains( + list: *const vx_expression, + value: *const vx_expression, + ) -> *mut vx_expression; + } +} + +/// A fixed-width primitive type identifier for Velox scalar construction. +pub type vx_velox_ptype = u32; +/// Unsigned 8-bit integer type identifier. +pub const VX_VELOX_PTYPE_U8: vx_velox_ptype = 0; +/// Unsigned 16-bit integer type identifier. +pub const VX_VELOX_PTYPE_U16: vx_velox_ptype = 1; +/// Unsigned 32-bit integer type identifier. +pub const VX_VELOX_PTYPE_U32: vx_velox_ptype = 2; +/// Unsigned 64-bit integer type identifier. +pub const VX_VELOX_PTYPE_U64: vx_velox_ptype = 3; +/// Signed 8-bit integer type identifier. +pub const VX_VELOX_PTYPE_I8: vx_velox_ptype = 4; +/// Signed 16-bit integer type identifier. +pub const VX_VELOX_PTYPE_I16: vx_velox_ptype = 5; +/// Signed 32-bit integer type identifier. +pub const VX_VELOX_PTYPE_I32: vx_velox_ptype = 6; +/// Signed 64-bit integer type identifier. +pub const VX_VELOX_PTYPE_I64: vx_velox_ptype = 7; +/// 16-bit floating-point type identifier. +pub const VX_VELOX_PTYPE_F16: vx_velox_ptype = 8; +/// 32-bit floating-point type identifier. +pub const VX_VELOX_PTYPE_F32: vx_velox_ptype = 9; +/// 64-bit floating-point type identifier. +pub const VX_VELOX_PTYPE_F64: vx_velox_ptype = 10; + +/// A fixed-width binary expression operator identifier. +pub type vx_velox_binary_operator = u32; +/// Equality operator identifier. +pub const VX_VELOX_OPERATOR_EQ: vx_velox_binary_operator = 0; +/// Inequality operator identifier. +pub const VX_VELOX_OPERATOR_NOT_EQ: vx_velox_binary_operator = 1; +/// Greater-than operator identifier. +pub const VX_VELOX_OPERATOR_GT: vx_velox_binary_operator = 2; +/// Greater-than-or-equal operator identifier. +pub const VX_VELOX_OPERATOR_GTE: vx_velox_binary_operator = 3; +/// Less-than operator identifier. +pub const VX_VELOX_OPERATOR_LT: vx_velox_binary_operator = 4; +/// Less-than-or-equal operator identifier. +pub const VX_VELOX_OPERATOR_LTE: vx_velox_binary_operator = 5; +/// Kleene logical AND operator identifier. +pub const VX_VELOX_OPERATOR_KLEENE_AND: vx_velox_binary_operator = 6; +/// Kleene logical OR operator identifier. +pub const VX_VELOX_OPERATOR_KLEENE_OR: vx_velox_binary_operator = 7; + +/// A fixed-width row-selection mode identifier. +pub type vx_velox_scan_selection_include = u32; +/// Include every row. +pub const VX_VELOX_SELECTION_ALL: vx_velox_scan_selection_include = 0; +/// Include the supplied row indexes. +pub const VX_VELOX_SELECTION_INCLUDE: vx_velox_scan_selection_include = 1; +/// Exclude the supplied row indexes. +pub const VX_VELOX_SELECTION_EXCLUDE: vx_velox_scan_selection_include = 2; + +/// A stable row selection for one scan request. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_scan_selection { + /// The selected row indexes. + pub indices: *const u64, + /// The number of selected row indexes. + pub length: usize, + /// The selection mode. + pub include: vx_velox_scan_selection_include, +} + +/// Stable options for one Vortex scan. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_scan_options { + /// Set this field to `sizeof(vx_velox_scan_options)`. + pub struct_size: usize, + /// Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + pub abi_version: u32, + /// The projected expression, or null for every field. + pub projection: *const vx_expression, + /// The exact filter expression, or null for no filter. + pub filter: *const vx_expression, + /// The first row in the scan range. + pub row_range_begin: u64, + /// One past the final row in the scan range. + pub row_range_end: u64, + /// An optional row-index selection. + pub selection: vx_velox_scan_selection, + /// The maximum output row count, or zero for no limit. + pub limit: u64, + /// Return rows in storage order. + pub ordered: bool, +} + +impl Default for vx_velox_scan_selection { + fn default() -> Self { + Self { + indices: ptr::null(), + length: 0, + include: VX_VELOX_SELECTION_ALL, + } + } +} + +/// Return the message stored in an adapter error. +/// +/// # Safety +/// +/// `error` must point to a live error handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_error_message(error: *const vx_error) -> vx_view { + unsafe { ffi::vx_error_message(error) } +} + +/// Free an adapter error. +/// +/// # Safety +/// +/// `error` must be null or an owned error handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_error_free(error: *const vx_error) { + unsafe { ffi::vx_error_free(error) }; +} + +/// Create a default Vortex session for Velox. +#[unsafe(no_mangle)] +pub extern "C-unwind" fn vx_velox_session_new() -> *mut vx_session { + unsafe { ffi::vx_session_new() } +} + +/// Free a Vortex session. +/// +/// # Safety +/// +/// `session` must be null or an owned session handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_session_free(session: *const vx_session) { + unsafe { ffi::vx_session_free(session) }; +} + +fn primitive_type(ptype: vx_velox_ptype) -> vortex_error::VortexResult { + match ptype { + VX_VELOX_PTYPE_U8 => Ok(PType::U8), + VX_VELOX_PTYPE_U16 => Ok(PType::U16), + VX_VELOX_PTYPE_U32 => Ok(PType::U32), + VX_VELOX_PTYPE_U64 => Ok(PType::U64), + VX_VELOX_PTYPE_I8 => Ok(PType::I8), + VX_VELOX_PTYPE_I16 => Ok(PType::I16), + VX_VELOX_PTYPE_I32 => Ok(PType::I32), + VX_VELOX_PTYPE_I64 => Ok(PType::I64), + VX_VELOX_PTYPE_F16 => Ok(PType::F16), + VX_VELOX_PTYPE_F32 => Ok(PType::F32), + VX_VELOX_PTYPE_F64 => Ok(PType::F64), + _ => vortex_bail!("Unknown Vortex Velox primitive type identifier: {ptype}"), + } +} + +/// Create a primitive dtype for a list literal. +/// +/// # Safety +/// +/// `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_dtype_new_primitive( + ptype: vx_velox_ptype, + nullable: bool, + error_out: *mut *mut vx_error, +) -> *const vx_dtype { + try_or(error_out, ptr::null(), || { + Ok(vx_dtype_new_with(DType::Primitive( + primitive_type(ptype)?, + Nullability::from(nullable), + ))) + }) +} + +/// Free a dtype. +/// +/// # Safety +/// +/// `dtype` must be null or an owned dtype handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_dtype_free(dtype: *const vx_dtype) { + unsafe { ffi::vx_dtype_free(dtype) }; +} + +/// Create a Boolean scalar. +#[unsafe(no_mangle)] +pub extern "C-unwind" fn vx_velox_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_scalar { + unsafe { ffi::vx_scalar_new_bool(value, nullable) } +} + +macro_rules! scalar_primitive_wrapper { + ($name:ident, $source:ident, $type:ty, $description:literal) => { + #[doc = $description] + #[unsafe(no_mangle)] + pub extern "C-unwind" fn $name(value: $type, nullable: bool) -> *mut vx_scalar { + unsafe { ffi::$source(value, nullable) } + } + }; +} + +scalar_primitive_wrapper!( + vx_velox_scalar_new_i8, + vx_scalar_new_i8, + i8, + "Create an i8 scalar." +); +scalar_primitive_wrapper!( + vx_velox_scalar_new_i16, + vx_scalar_new_i16, + i16, + "Create an i16 scalar." +); +scalar_primitive_wrapper!( + vx_velox_scalar_new_i32, + vx_scalar_new_i32, + i32, + "Create an i32 scalar." +); +scalar_primitive_wrapper!( + vx_velox_scalar_new_i64, + vx_scalar_new_i64, + i64, + "Create an i64 scalar." +); +scalar_primitive_wrapper!( + vx_velox_scalar_new_f32, + vx_scalar_new_f32, + f32, + "Create an f32 scalar." +); +scalar_primitive_wrapper!( + vx_velox_scalar_new_f64, + vx_scalar_new_f64, + f64, + "Create an f64 scalar." +); + +/// Create a UTF-8 scalar. +/// +/// # Safety +/// +/// `value` and `error_out` must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scalar_new_utf8( + value: vx_view, + nullable: bool, + error_out: *mut *mut vx_error, +) -> *mut vx_scalar { + unsafe { ffi::vx_scalar_new_utf8(value, nullable, error_out) } +} + +/// Create a binary scalar. +/// +/// # Safety +/// +/// `data` must identify `length` bytes. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scalar_new_binary( + data: *const u8, + length: usize, + nullable: bool, + error_out: *mut *mut vx_error, +) -> *mut vx_scalar { + unsafe { ffi::vx_scalar_new_binary(data, length, nullable, error_out) } +} + +/// Create a list scalar. +/// +/// # Safety +/// +/// Every pointer must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scalar_new_list( + element_dtype: *const vx_dtype, + elements: *const *const vx_scalar, + length: usize, + nullable: bool, + error_out: *mut *mut vx_error, +) -> *mut vx_scalar { + unsafe { ffi::vx_scalar_new_list(element_dtype, elements, length, nullable, error_out) } +} + +/// Free a scalar. +/// +/// # Safety +/// +/// `scalar` must be null or an owned scalar handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scalar_free(scalar: *const vx_scalar) { + unsafe { ffi::vx_scalar_free(scalar) }; +} + +/// Create a literal expression. +/// +/// # Safety +/// +/// `scalar` must point to a live scalar. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_expression_literal( + scalar: *const vx_scalar, + error_out: *mut *mut vx_error, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_literal(scalar, error_out) } +} + +/// Create a root expression. +#[unsafe(no_mangle)] +pub extern "C" fn vx_velox_expression_root() -> *mut vx_expression { + unsafe { ffi::vx_expression_root() } +} + +/// Create a field expression. +/// +/// # Safety +/// +/// `child` must point to a live expression. `name` must identify valid UTF-8. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_get_item( + name: vx_view, + child: *const vx_expression, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_get_item(name, child) } +} + +fn binary_operator(operator: vx_velox_binary_operator) -> vortex_error::VortexResult { + match operator { + VX_VELOX_OPERATOR_EQ => Ok(Operator::Eq), + VX_VELOX_OPERATOR_NOT_EQ => Ok(Operator::NotEq), + VX_VELOX_OPERATOR_GT => Ok(Operator::Gt), + VX_VELOX_OPERATOR_GTE => Ok(Operator::Gte), + VX_VELOX_OPERATOR_LT => Ok(Operator::Lt), + VX_VELOX_OPERATOR_LTE => Ok(Operator::Lte), + VX_VELOX_OPERATOR_KLEENE_AND => Ok(Operator::And), + VX_VELOX_OPERATOR_KLEENE_OR => Ok(Operator::Or), + _ => vortex_bail!("Unknown Vortex Velox binary operator identifier: {operator}"), + } +} + +/// Create a binary expression. +/// +/// # Safety +/// +/// Both operands must point to live expressions. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_expression_binary( + operator: vx_velox_binary_operator, + left: *const vx_expression, + right: *const vx_expression, + error_out: *mut *mut vx_error, +) -> *mut vx_expression { + try_or(error_out, ptr::null_mut(), || { + let operator = binary_operator(operator)?; + let left = unsafe { vx_expression_ref(left)? }.clone(); + let right = unsafe { vx_expression_ref(right)? }.clone(); + Ok(vx_expression_new_with( + Binary.new_expr(operator, [left, right]), + )) + }) +} + +/// Create a conjunction from expressions. +/// +/// # Safety +/// +/// `expressions` must identify `length` live expression pointers. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_and( + expressions: *const *const vx_expression, + length: usize, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_and(expressions, length) } +} + +/// Create a disjunction from expressions. +/// +/// # Safety +/// +/// `expressions` must identify `length` live expression pointers. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_or( + expressions: *const *const vx_expression, + length: usize, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_or(expressions, length) } +} + +/// Create a logical negation. +/// +/// # Safety +/// +/// `child` must point to a live expression. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_not( + child: *const vx_expression, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_not(child) } +} + +/// Create a null test. +/// +/// # Safety +/// +/// `child` must point to a live expression. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_is_null( + child: *const vx_expression, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_is_null(child) } +} + +/// Create a list membership test. +/// +/// # Safety +/// +/// Both operands must point to live expressions. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_expression_list_contains( + list: *const vx_expression, + value: *const vx_expression, +) -> *mut vx_expression { + unsafe { ffi::vx_expression_list_contains(list, value) } +} + +/// Free an expression. +/// +/// # Safety +/// +/// `expression` must be null or an owned expression handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_expression_free(expression: *const vx_expression) { + unsafe { ffi::vx_expression_free(expression) }; +} + +/// Free a data source. +/// +/// # Safety +/// +/// `data_source` must be null or an owned data-source handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_data_source_free(data_source: *const vx_data_source) { + unsafe { ffi::vx_data_source_free(data_source) }; +} + +unsafe fn scan_options(options: &vx_velox_scan_options) -> vortex_error::VortexResult { + if options.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox scan options are too small: expected at least {}, got {}", + size_of::(), + options.struct_size + ); + } + if options.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + options.abi_version + ); + } + let projection = if options.projection.is_null() { + root() + } else { + unsafe { vx_expression_ref(options.projection)? }.clone() + }; + let filter = if options.filter.is_null() { + None + } else { + Some(unsafe { vx_expression_ref(options.filter)? }.clone()) + }; + let indices = if options.selection.length == 0 { + &[] + } else { + if options.selection.indices.is_null() { + vortex_bail!("Vortex Velox scan selection indices must not be null"); + } + unsafe { slice::from_raw_parts(options.selection.indices, options.selection.length) } + }; + let selection = match options.selection.include { + VX_VELOX_SELECTION_ALL => { + if !indices.is_empty() { + vortex_bail!("An all-rows selection must not contain row indices"); + } + Selection::All + } + VX_VELOX_SELECTION_INCLUDE => { + Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::copy_from(indices))?) + } + VX_VELOX_SELECTION_EXCLUDE => { + Selection::ExcludeByIndex(StrictSortedBuffer::try_new(Buffer::copy_from(indices))?) + } + include => vortex_bail!("Unknown Vortex Velox scan selection identifier: {include}"), + }; + if options.row_range_end != 0 && options.row_range_begin > options.row_range_end { + vortex_bail!( + "Vortex Velox row range is invalid: {}..{}", + options.row_range_begin, + options.row_range_end + ); + } + Ok(ScanRequest { + projection, + filter, + row_range: (options.row_range_begin != 0 || options.row_range_end != 0) + .then_some(options.row_range_begin..options.row_range_end), + selection, + limit: (options.limit != 0).then_some(options.limit), + ordered: options.ordered, + partition_selection: Selection::All, + partition_range: None, + }) +} + +/// Start a scan through the stable adapter options. +/// +/// # Safety +/// +/// Every pointer must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_data_source_scan( + data_source: *const vx_data_source, + options: *const vx_velox_scan_options, + error_out: *mut *mut vx_error, +) -> *mut vx_scan { + try_or(error_out, ptr::null_mut(), || { + let request = if options.is_null() { + ScanRequest::default() + } else { + unsafe { scan_options(&*options)? } + }; + unsafe { vx_data_source_scan_with(data_source, request) } + }) +} + +/// Free a scan. +/// +/// # Safety +/// +/// `scan` must be null or an owned scan handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scan_free(scan: *const vx_scan) { + unsafe { ffi::vx_scan_free(scan) }; +} + +/// Return the next partition from a scan. +/// +/// # Safety +/// +/// `scan` must point to a live scan. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scan_next_partition( + scan: *mut vx_scan, + error_out: *mut *mut vx_error, +) -> *mut vx_partition { + unsafe { ffi::vx_scan_next_partition(scan, error_out) } +} + +/// Free a partition. +/// +/// # Safety +/// +/// `partition` must be null or an owned partition handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_partition_free(partition: *const vx_partition) { + unsafe { ffi::vx_partition_free(partition) }; +} + +/// Return the next array from a partition. +/// +/// # Safety +/// +/// `partition` must point to a live partition. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_partition_next( + partition: *mut vx_partition, + error_out: *mut *mut vx_error, +) -> *const vx_array { + unsafe { ffi::vx_partition_next(partition, error_out) } +} + +/// Free an array. +/// +/// # Safety +/// +/// `array` must be null or an owned array handle. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_free(array: *const vx_array) { + unsafe { ffi::vx_array_free(array) }; +} + +/// Return an array length. +/// +/// # Safety +/// +/// `array` must point to a live array. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_len(array: *const vx_array) -> usize { + unsafe { ffi::vx_array_len(array) } +} + +/// Slice an array. +/// +/// # Safety +/// +/// `array` must point to a live array. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_slice( + array: *const vx_array, + begin: usize, + end: usize, + error_out: *mut *mut vx_error, +) -> *const vx_array { + unsafe { ffi::vx_array_slice(array, begin, end, error_out) } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + + #[test] + fn translates_scan_options() -> VortexResult<()> { + let options = vx_velox_scan_options { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + projection: ptr::null(), + filter: ptr::null(), + row_range_begin: 10, + row_range_end: 20, + selection: vx_velox_scan_selection::default(), + limit: 5, + ordered: true, + }; + let translated = unsafe { scan_options(&options)? }; + assert_eq!(translated.row_range, Some(10..20)); + assert_eq!(translated.limit, Some(5)); + assert!(translated.ordered); + Ok(()) + } + + #[test] + fn builds_list_membership_expression() { + let mut error = ptr::null_mut(); + let dtype = + unsafe { vx_velox_dtype_new_primitive(VX_VELOX_PTYPE_I64, false, &raw mut error) }; + assert!(error.is_null()); + let values = [ + vx_velox_scalar_new_i64(10, false), + vx_velox_scalar_new_i64(20, false), + ]; + let list = unsafe { + vx_velox_scalar_new_list( + dtype, + values.as_ptr().cast(), + values.len(), + false, + &raw mut error, + ) + }; + assert!(error.is_null()); + let list_literal = unsafe { vx_velox_expression_literal(list, &raw mut error) }; + assert!(error.is_null()); + let root = vx_velox_expression_root(); + let name = vx_view { + ptr: c"value".as_ptr(), + len: 5, + }; + let value = unsafe { vx_velox_expression_get_item(name, root) }; + let membership = unsafe { vx_velox_expression_list_contains(list_literal, value) }; + assert!(!membership.is_null()); + + unsafe { + vx_velox_expression_free(membership); + vx_velox_expression_free(value); + vx_velox_expression_free(root); + vx_velox_expression_free(list_literal); + vx_velox_scalar_free(list); + for value in values { + vx_velox_scalar_free(value); + } + vx_velox_dtype_free(dtype); + } + } + + #[test] + fn rejects_wrong_scan_abi_before_source_access() { + let options = vx_velox_scan_options { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION + 1, + projection: ptr::null(), + filter: ptr::null(), + row_range_begin: 0, + row_range_end: 0, + selection: vx_velox_scan_selection::default(), + limit: 0, + ordered: false, + }; + let mut error = ptr::null_mut(); + let scan = + unsafe { vx_velox_data_source_scan(ptr::null(), &raw const options, &raw mut error) }; + assert!(scan.is_null()); + assert!(!error.is_null()); + unsafe { vx_velox_error_free(error) }; + } + + #[test] + fn rejects_unknown_fixed_width_identifiers() { + let mut error = ptr::null_mut(); + let dtype = unsafe { vx_velox_dtype_new_primitive(u32::MAX, false, &raw mut error) }; + assert!(dtype.is_null()); + assert!(!error.is_null()); + unsafe { vx_velox_error_free(error) }; + + error = ptr::null_mut(); + let expression = unsafe { + vx_velox_expression_binary(u32::MAX, ptr::null(), ptr::null(), &raw mut error) + }; + assert!(expression.is_null()); + assert!(!error.is_null()); + unsafe { vx_velox_error_free(error) }; + + let options = vx_velox_scan_options { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + projection: ptr::null(), + filter: ptr::null(), + row_range_begin: 0, + row_range_end: 0, + selection: vx_velox_scan_selection { + include: u32::MAX, + ..Default::default() + }, + limit: 0, + ordered: false, + }; + assert!(unsafe { scan_options(&options) }.is_err()); + } +} diff --git a/vortex-velox/src/array.rs b/vortex-velox/src/array.rs new file mode 100644 index 00000000000..4f8bf4f3dc0 --- /dev/null +++ b/vortex-velox/src/array.rs @@ -0,0 +1,856 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::c_char; +use std::ffi::c_void; +use std::mem::size_of; +use std::ptr; + +use arrow_array::Array; +use arrow_array::ffi::FFI_ArrowArray; +use arrow_array::ffi::FFI_ArrowSchema; +use arrow_buffer::BooleanBuffer; +use arrow_buffer::Buffer; +use arrow_buffer::NullBuffer; +use arrow_data::ArrayData; +use arrow_data::ArrayDataBuilder; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::struct_::StructArrayExt; +use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_ffi::try_or; +use vortex_ffi::vx_array; +use vortex_ffi::vx_array_new_with; +use vortex_ffi::vx_array_ref; +use vortex_ffi::vx_error; +use vortex_ffi::vx_session; +use vortex_ffi::vx_session_ref; + +/// Host memory callbacks for one Arrow C Data export. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_arrow_memory_callbacks { + /// Set this field to `sizeof(vx_velox_arrow_memory_callbacks)`. + pub struct_size: usize, + /// Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + pub abi_version: u32, + /// An opaque host context. + pub context: *mut c_void, + /// Retain the host context until the Arrow array release callback runs. + pub retain_context: Option, + /// Release one host context reference. + pub release_context: Option, + /// Reserve Arrow payload bytes before conversion. Zero means success. + pub report_allocation: + Option i32>, + /// Free retained Arrow payload bytes. + pub report_free: Option, + /// Return the last callback error as a null-terminated string. + pub last_error: Option *const c_char>, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ArrowMemoryCallbacksPrefix { + struct_size: usize, + abi_version: u32, + context: *mut c_void, + retain_context: Option, + release_context: Option, + report_allocation: + Option i32>, + report_free: Option, +} + +struct ArrowMemoryOwner { + original_private_data: *mut c_void, + original_release: unsafe extern "C" fn(array: *mut FFI_ArrowArray), + callbacks: vx_velox_arrow_memory_callbacks, + retained_bytes: usize, +} + +struct ArrowMemoryReservation { + callbacks: vx_velox_arrow_memory_callbacks, + retained_bytes: usize, + active: bool, +} + +// SAFETY: The callback contract permits the retained context to move between threads. +unsafe impl Send for ArrowMemoryOwner {} +// SAFETY: Release accesses the immutable callback table once through exclusive ownership. +unsafe impl Sync for ArrowMemoryOwner {} + +unsafe extern "C" fn release_accounted_arrow(array: *mut FFI_ArrowArray) { + if array.is_null() { + return; + } + // SAFETY: Arrow calls this function once with the exported top-level array. + let array = unsafe { &mut *array }; + // SAFETY: `private_data` was created from this exact box in `attach_memory_owner`. + let owner = unsafe { Box::from_raw(array.private_data.cast::()) }; + array.private_data = owner.original_private_data; + array.release = Some(owner.original_release); + // SAFETY: The original release function owns the restored original private data. + unsafe { (owner.original_release)(array) }; + if let Some(report_free) = owner.callbacks.report_free { + // SAFETY: The retained callback context stays live until this release function ends. + unsafe { report_free(owner.callbacks.context, owner.retained_bytes) }; + } + if let Some(release_context) = owner.callbacks.release_context { + // SAFETY: This release matches the retain before export. + unsafe { release_context(owner.callbacks.context) }; + } +} + +impl ArrowMemoryReservation { + fn try_new( + callbacks: vx_velox_arrow_memory_callbacks, + retained_bytes: usize, + ) -> VortexResult { + let retain_context = callbacks + .retain_context + .ok_or_else(|| vortex_err!("Missing Arrow context retain callback"))?; + let release_context = callbacks + .release_context + .ok_or_else(|| vortex_err!("Missing Arrow context release callback"))?; + let report_allocation = callbacks + .report_allocation + .ok_or_else(|| vortex_err!("Missing Arrow allocation callback"))?; + // SAFETY: The callback contract accepts one retained context reference. + unsafe { retain_context(callbacks.context) }; + // SAFETY: The retained context stays live for this callback. + let status = unsafe { report_allocation(callbacks.context, retained_bytes) }; + if status != 0 { + let message = callback_error(&callbacks, status); + // SAFETY: The reservation failed, so release the temporary context reference. + unsafe { release_context(callbacks.context) }; + vortex_bail!("{}", message); + } + Ok(Self { + callbacks, + retained_bytes, + active: true, + }) + } + + fn reconcile(&mut self, actual_retained_bytes: usize) -> VortexResult<()> { + match actual_retained_bytes.cmp(&self.retained_bytes) { + std::cmp::Ordering::Less => { + let released = self.retained_bytes - actual_retained_bytes; + let report_free = self + .callbacks + .report_free + .ok_or_else(|| vortex_err!("Missing Arrow free callback"))?; + // SAFETY: The retained callback context stays live for this callback. + unsafe { report_free(self.callbacks.context, released) }; + } + std::cmp::Ordering::Greater => { + let additional = actual_retained_bytes - self.retained_bytes; + let report_allocation = self + .callbacks + .report_allocation + .ok_or_else(|| vortex_err!("Missing Arrow allocation callback"))?; + // SAFETY: The retained callback context stays live for this callback. + let status = unsafe { report_allocation(self.callbacks.context, additional) }; + if status != 0 { + vortex_bail!("{}", callback_error(&self.callbacks, status)); + } + } + std::cmp::Ordering::Equal => {} + } + self.retained_bytes = actual_retained_bytes; + Ok(()) + } +} + +impl Drop for ArrowMemoryReservation { + fn drop(&mut self) { + if !self.active { + return; + } + if let Some(report_free) = self.callbacks.report_free { + // SAFETY: The retained callback context stays live for this callback. + unsafe { report_free(self.callbacks.context, self.retained_bytes) }; + } + if let Some(release_context) = self.callbacks.release_context { + // SAFETY: This release matches the reservation retain. + unsafe { release_context(self.callbacks.context) }; + } + } +} + +unsafe fn parse_memory_callbacks( + callbacks: *const vx_velox_arrow_memory_callbacks, +) -> VortexResult { + if callbacks.is_null() { + vortex_bail!("Arrow memory callbacks must not be null"); + } + // SAFETY: The caller guarantees that the pointer identifies at least `struct_size` bytes. + let struct_size = unsafe { ptr::read(callbacks.cast::()) }; + if struct_size < size_of::() { + vortex_bail!( + "Vortex Velox Arrow memory callback structure is too small: expected at least {}, got {}", + size_of::(), + struct_size + ); + } + // SAFETY: The checked size covers the required callback prefix. + let prefix = unsafe { ptr::read(callbacks.cast::()) }; + if prefix.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + prefix.abi_version + ); + } + if prefix.retain_context.is_none() + || prefix.release_context.is_none() + || prefix.report_allocation.is_none() + || prefix.report_free.is_none() + { + vortex_bail!("Vortex Velox Arrow memory callbacks are incomplete"); + } + let last_error = if struct_size >= size_of::() { + // SAFETY: The checked size covers the optional tail field. + unsafe { ptr::read(ptr::addr_of!((*callbacks).last_error)) } + } else { + None + }; + Ok(vx_velox_arrow_memory_callbacks { + struct_size, + abi_version: prefix.abi_version, + context: prefix.context, + retain_context: prefix.retain_context, + release_context: prefix.release_context, + report_allocation: prefix.report_allocation, + report_free: prefix.report_free, + last_error, + }) +} + +fn callback_error(callbacks: &vx_velox_arrow_memory_callbacks, status: i32) -> String { + let Some(last_error) = callbacks.last_error else { + return format!("Velox Arrow allocation callback failed with status {status}"); + }; + // SAFETY: The callback contract returns null or a valid null-terminated string. + let message = unsafe { last_error(callbacks.context) }; + if message.is_null() { + return format!("Velox Arrow allocation callback failed with status {status}"); + } + // SAFETY: The callback keeps the string valid until the next callback. + unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned() +} + +fn copy_nulls(nulls: &NullBuffer, data_offset: usize) -> VortexResult { + let bit_length = data_offset + .checked_add(nulls.len()) + .ok_or_else(|| vortex_err!("Arrow validity bit count overflow"))?; + let mut bytes = vec![0_u8; bit_length.div_ceil(8)]; + for index in 0..nulls.len() { + if nulls.is_valid(index) { + let bit = data_offset + index; + bytes[bit / 8] |= 1 << (bit % 8); + } + } + Ok(NullBuffer::new(BooleanBuffer::new( + Buffer::from(bytes), + data_offset, + nulls.len(), + ))) +} + +fn copy_arrow_data(data: &ArrayData) -> VortexResult { + let buffers = data + .buffers() + .iter() + .map(|buffer| Buffer::from_slice_ref(buffer.as_slice())) + .collect(); + let children = data + .child_data() + .iter() + .map(copy_arrow_data) + .collect::>>()?; + let nulls = data + .nulls() + .map(|nulls| copy_nulls(nulls, data.offset())) + .transpose()?; + Ok(ArrayDataBuilder::new(data.data_type().clone()) + .len(data.len()) + .offset(data.offset()) + .buffers(buffers) + .child_data(children) + .nulls(nulls) + .build()?) +} + +fn variadic_ffi_buffer_bytes(data: &ArrayData) -> usize { + let own_bytes = if arrow_data::layout(data.data_type()).variadic { + let mut lengths = Vec::::new(); + #[expect( + clippy::same_item_push, + reason = "Match the Arrow FFI vector growth to account its exact retained capacity" + )] + for _ in data.buffers().iter().skip(1) { + lengths.push(0); + } + lengths.capacity() * size_of::() + } else { + 0 + }; + own_bytes + + data + .child_data() + .iter() + .map(variadic_ffi_buffer_bytes) + .sum::() +} + +fn attach_memory_owner( + array: &mut FFI_ArrowArray, + mut reservation: ArrowMemoryReservation, +) -> VortexResult<()> { + let original_release = array + .release + .ok_or_else(|| vortex_err!("Exported Arrow array has no release callback"))?; + let owner = Box::new(ArrowMemoryOwner { + original_private_data: array.private_data, + original_release, + callbacks: reservation.callbacks, + retained_bytes: reservation.retained_bytes, + }); + reservation.active = false; + array.private_data = Box::into_raw(owner).cast(); + array.release = Some(release_accounted_arrow); + Ok(()) +} + +const ARROW_RESERVATION_OVERHEAD: usize = 64 * 1024; + +fn conservative_arrow_reservation( + array: &vortex::array::ArrayRef, + execution: &mut vortex::array::ExecutionCtx, +) -> VortexResult { + uncompressed_size_in_bytes(array, execution)? + .checked_mul(2) + .and_then(|bytes| bytes.checked_add(ARROW_RESERVATION_OVERHEAD)) + .ok_or_else(|| vortex_err!("Arrow reservation size overflow")) +} + +/// Return one struct field with the supplied session. +/// +/// # Safety +/// +/// The session and array pointers must identify live handles. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_get_field( + session: *const vx_session, + array: *const vx_array, + index: usize, + error_out: *mut *mut vx_error, +) -> *const vx_array { + try_or(error_out, ptr::null(), || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + let mut execution = session.create_execution_ctx(); + let struct_array = array.clone().execute::(&mut execution)?; + let field = struct_array + .unmasked_field_opt(index) + .ok_or_else(|| vortex_err!("Field index out of bounds: {index}"))? + .clone(); + Ok(vx_array_new_with(field)) + }) +} + +/// Return the invalid value count with the supplied session. +/// +/// # Safety +/// +/// The session and array pointers must identify live handles. `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_invalid_count( + session: *const vx_session, + array: *const vx_array, + error_out: *mut *mut vx_error, +) -> usize { + try_or(error_out, 0, || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + array.invalid_count(&mut session.create_execution_ctx()) + }) +} + +/// Export one Vortex array through the Arrow C Data Interface. +/// +/// The caller owns both outputs and must call their release callbacks. The memory callbacks reserve +/// a conservative payload charge before Arrow conversion. The adapter refunds the difference after +/// it knows the retained payload capacities. It requests a deficit before it returns the outputs. +/// The charge excludes schema and small FFI metadata. +/// +/// # Safety +/// +/// The session and array pointers must identify live handles. `memory_callbacks` must identify its +/// declared `struct_size` bytes for this call. Its callback context must remain valid through every +/// retained reference. Its callbacks and returned error strings must satisfy the header contract +/// and must not unwind. Both output pointers must identify uninitialized writable structures. +/// `error_out` must be null or identify writable storage for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_export_arrow( + session: *const vx_session, + array: *const vx_array, + memory_callbacks: *const vx_velox_arrow_memory_callbacks, + schema_out: *mut FFI_ArrowSchema, + array_out: *mut FFI_ArrowArray, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + // SAFETY: The caller provides the callback table described by this function's contract. + let memory_callbacks = unsafe { parse_memory_callbacks(memory_callbacks)? }; + if schema_out.is_null() { + return Err(vortex_err!("Arrow schema output must not be null")); + } + if array_out.is_null() { + return Err(vortex_err!("Arrow array output must not be null")); + } + + let mut execution = session.create_execution_ctx(); + let reserved_bytes = conservative_arrow_reservation(array, &mut execution)?; + let mut reservation = ArrowMemoryReservation::try_new(memory_callbacks, reserved_bytes)?; + let mut arrow = session + .arrow() + .execute_arrow(array.clone(), None, &mut execution)?; + if arrow.offset() != 0 { + let length = u64::try_from(array.len()) + .map_err(|_| vortex_err!("Array length does not fit u64: {}", array.len()))?; + let compact = array.take(PrimitiveArray::from_iter(0..length).into_array())?; + arrow = session + .arrow() + .execute_arrow(compact, None, &mut execution)?; + } + let schema = FFI_ArrowSchema::try_from(arrow.data_type())?; + let data = copy_arrow_data(&arrow.to_data())?; + let retained_bytes = data + .get_buffer_memory_size() + .checked_add(variadic_ffi_buffer_bytes(&data)) + .ok_or_else(|| vortex_err!("Arrow retained memory size overflow"))?; + let mut array = FFI_ArrowArray::new(&data); + drop(data); + drop(arrow); + reservation.reconcile(retained_bytes)?; + attach_memory_owner(&mut array, reservation)?; + unsafe { + ptr::write(schema_out, schema); + ptr::write(array_out, array); + } + Ok(0) + }) +} + +#[cfg(test)] +mod tests { + use std::ptr; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use arrow_array::Int32Array; + use arrow_array::StructArray as ArrowStructArray; + use arrow_array::array::make_array; + use arrow_array::ffi::from_ffi; + use vortex::array::IntoArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::StructArray; + use vortex::array::validity::Validity; + use vortex_error::VortexResult; + use vortex_ffi::vx_array_new_with; + use vortex_ffi::vx_error_free; + use vortex_ffi::vx_session_free; + use vortex_ffi::vx_session_new_with; + + use super::*; + use crate::api::vx_velox_array_free; + + #[derive(Default)] + struct MemoryCapture { + allocated: AtomicUsize, + freed: AtomicUsize, + first_allocation: AtomicUsize, + allocation_calls: AtomicUsize, + allocation_attempts: AtomicUsize, + retained_contexts: AtomicUsize, + released_contexts: AtomicUsize, + reject_allocation: AtomicBool, + } + + unsafe extern "C" fn retain_context(context: *mut c_void) { + // SAFETY: The test context is an `Arc` pointer. + unsafe { Arc::increment_strong_count(context.cast::()) }; + // SAFETY: The test keeps one independent `Arc` reference live. + let capture = unsafe { &*context.cast::() }; + capture.retained_contexts.fetch_add(1, Ordering::Relaxed); + } + + unsafe extern "C" fn release_context(context: *mut c_void) { + // SAFETY: The test keeps one independent `Arc` reference live. + let capture = unsafe { &*context.cast::() }; + capture.released_contexts.fetch_add(1, Ordering::Relaxed); + // SAFETY: This release matches one `retain_context` call. + drop(unsafe { Arc::from_raw(context.cast::()) }); + } + + unsafe extern "C" fn report_allocation(context: *mut c_void, bytes: usize) -> i32 { + // SAFETY: The test context is a live `MemoryCapture`. + let capture = unsafe { &*context.cast::() }; + capture.allocation_attempts.fetch_add(1, Ordering::Relaxed); + if capture.reject_allocation.load(Ordering::Relaxed) { + return 7; + } + capture + .first_allocation + .compare_exchange(0, bytes, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + capture.allocation_calls.fetch_add(1, Ordering::Relaxed); + capture.allocated.fetch_add(bytes, Ordering::Relaxed); + 0 + } + + unsafe extern "C" fn report_free(context: *mut c_void, bytes: usize) { + // SAFETY: The test context is a live `MemoryCapture`. + let capture = unsafe { &*context.cast::() }; + capture.freed.fetch_add(bytes, Ordering::Relaxed); + } + + unsafe extern "C" fn last_error(_context: *mut c_void) -> *const c_char { + c"allocation rejected".as_ptr() + } + + fn memory_callbacks(capture: &Arc) -> vx_velox_arrow_memory_callbacks { + vx_velox_arrow_memory_callbacks { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: Arc::as_ptr(capture).cast_mut().cast(), + retain_context: Some(retain_context), + release_context: Some(release_context), + report_allocation: Some(report_allocation), + report_free: Some(report_free), + last_error: Some(last_error), + } + } + + #[test] + fn accepts_required_memory_callback_prefix() -> VortexResult<()> { + let capture = Arc::new(MemoryCapture::default()); + let callbacks = ArrowMemoryCallbacksPrefix { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: Arc::as_ptr(&capture).cast_mut().cast(), + retain_context: Some(retain_context), + release_context: Some(release_context), + report_allocation: Some(report_allocation), + report_free: Some(report_free), + }; + // SAFETY: The test prefix reports its exact initialized size. + let normalized = unsafe { + parse_memory_callbacks( + (&raw const callbacks).cast::(), + )? + }; + assert!(normalized.last_error.is_none()); + Ok(()) + } + + #[test] + fn exports_one_array_to_arrow() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with(PrimitiveArray::from_iter([1_i32, 2, 3]).into_array()); + let mut schema = FFI_ArrowSchema::empty(); + let mut arrow_array = FFI_ArrowArray::empty(); + let mut error = ptr::null_mut(); + let capture = Arc::new(MemoryCapture::default()); + let callbacks = memory_callbacks(&capture); + + let status = unsafe { + vx_velox_array_export_arrow( + session, + array, + &raw const callbacks, + &raw mut schema, + &raw mut arrow_array, + &raw mut error, + ) + }; + if !error.is_null() { + unsafe { vx_error_free(error) }; + } + assert_eq!(status, 0); + assert!(error.is_null()); + + let data = unsafe { from_ffi(arrow_array, &schema)? }; + let arrow = make_array(data); + let values = arrow + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Expected an Arrow Int32 array"))?; + assert_eq!(values.values(), &[1, 2, 3]); + assert!(capture.allocated.load(Ordering::Relaxed) > 0); + assert!(capture.first_allocation.load(Ordering::Relaxed) >= ARROW_RESERVATION_OVERHEAD); + assert_eq!(capture.allocation_calls.load(Ordering::Relaxed), 1); + assert!(capture.freed.load(Ordering::Relaxed) < capture.allocated.load(Ordering::Relaxed)); + drop(arrow); + assert_eq!( + capture.freed.load(Ordering::Relaxed), + capture.allocated.load(Ordering::Relaxed) + ); + assert_eq!(capture.retained_contexts.load(Ordering::Relaxed), 1); + assert_eq!(capture.released_contexts.load(Ordering::Relaxed), 1); + + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn exports_sliced_array_with_zero_arrow_offset() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let source = PrimitiveArray::from_iter([1_i32, 2, 3, 4]).into_array(); + let array = vx_array_new_with(source.slice(1..3)?); + let mut error = ptr::null_mut(); + + let mut schema = FFI_ArrowSchema::empty(); + let mut arrow_array = FFI_ArrowArray::empty(); + let capture = Arc::new(MemoryCapture::default()); + let callbacks = memory_callbacks(&capture); + let status = unsafe { + vx_velox_array_export_arrow( + session, + array, + &raw const callbacks, + &raw mut schema, + &raw mut arrow_array, + &raw mut error, + ) + }; + if !error.is_null() { + unsafe { vx_error_free(error) }; + } + assert_eq!(status, 0); + assert!(error.is_null()); + + let data = unsafe { from_ffi(arrow_array, &schema)? }; + assert_eq!(data.offset(), 0); + let arrow = make_array(data); + let values = arrow + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Expected an Arrow Int32 array"))?; + assert_eq!(values.values(), &[2, 3]); + + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn reserves_and_reconciles_nested_arrow_fallback() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let inner = StructArray::try_new( + ["value"].into(), + vec![PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]).into_array()], + 3, + Validity::NonNullable, + )? + .into_array(); + let outer = StructArray::try_new(["nested"].into(), vec![inner], 3, Validity::NonNullable)? + .into_array(); + let array = vx_array_new_with(outer); + let capture = Arc::new(MemoryCapture::default()); + let callbacks = memory_callbacks(&capture); + let mut schema = FFI_ArrowSchema::empty(); + let mut arrow_array = FFI_ArrowArray::empty(); + let mut error = ptr::null_mut(); + + let status = unsafe { + vx_velox_array_export_arrow( + session, + array, + &raw const callbacks, + &raw mut schema, + &raw mut arrow_array, + &raw mut error, + ) + }; + assert_eq!(status, 0); + assert!(error.is_null()); + assert_eq!(capture.allocation_calls.load(Ordering::Relaxed), 1); + assert!(capture.first_allocation.load(Ordering::Relaxed) >= ARROW_RESERVATION_OVERHEAD); + + let data = unsafe { from_ffi(arrow_array, &schema)? }; + let arrow = make_array(data); + let outer = arrow + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Expected an outer Arrow struct"))?; + let inner = outer + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Expected a nested Arrow struct"))?; + let values = inner + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Expected nested Arrow i32 values"))?; + assert_eq!(values.value(0), 1); + assert!(values.is_null(1)); + assert_eq!(values.value(2), 3); + drop(arrow); + assert_eq!( + capture.freed.load(Ordering::Relaxed), + capture.allocated.load(Ordering::Relaxed) + ); + + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn copies_external_arrow_buffers_before_accounting() -> VortexResult<()> { + struct ExternalBytes { + bytes: Box<[u8]>, + drops: Arc, + } + + impl AsRef<[u8]> for ExternalBytes { + fn as_ref(&self) -> &[u8] { + &self.bytes + } + } + + impl Drop for ExternalBytes { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + let values = [1_i32, 2, 3]; + // SAFETY: `values` is live and readable for its byte size. + let bytes = unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), size_of_val(&values)) + }; + let drops = Arc::new(AtomicUsize::new(0)); + let external = Buffer::from(bytes::Bytes::from_owner(ExternalBytes { + bytes: bytes.into(), + drops: Arc::clone(&drops), + })); + let data_type = Int32Array::from(values.to_vec()).data_type().clone(); + let source = ArrayDataBuilder::new(data_type) + .len(values.len()) + .add_buffer(external) + .build()?; + + let copied = copy_arrow_data(&source)?; + assert!(copied.buffers()[0].capacity() >= size_of_val(&values)); + assert_ne!(copied.buffers()[0].as_ptr(), source.buffers()[0].as_ptr()); + drop(source); + assert_eq!(drops.load(Ordering::Relaxed), 1); + assert_eq!(copied.buffers()[0].as_slice(), bytes); + Ok(()) + } + + #[test] + fn rejects_arrow_allocation_and_releases_context() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with(PrimitiveArray::from_iter([1_i32, 2]).into_array()); + let capture = Arc::new(MemoryCapture::default()); + capture.reject_allocation.store(true, Ordering::Relaxed); + let callbacks = memory_callbacks(&capture); + let mut schema = FFI_ArrowSchema::empty(); + let mut arrow_array = FFI_ArrowArray::empty(); + let mut error = ptr::null_mut(); + + let status = unsafe { + vx_velox_array_export_arrow( + session, + array, + &raw const callbacks, + &raw mut schema, + &raw mut arrow_array, + &raw mut error, + ) + }; + assert_eq!(status, 1); + assert!(!error.is_null()); + assert!(arrow_array.release.is_none()); + assert_eq!(capture.allocated.load(Ordering::Relaxed), 0); + assert_eq!(capture.freed.load(Ordering::Relaxed), 0); + assert_eq!(capture.allocation_calls.load(Ordering::Relaxed), 0); + assert_eq!(capture.allocation_attempts.load(Ordering::Relaxed), 1); + assert!(capture.first_allocation.load(Ordering::Relaxed) == 0); + assert_eq!(capture.retained_contexts.load(Ordering::Relaxed), 1); + assert_eq!(capture.released_contexts.load(Ordering::Relaxed), 1); + + unsafe { + vx_error_free(error); + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn accesses_fields_and_invalid_counts_with_session() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let values = PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]).into_array(); + let structure = + StructArray::try_new(["value"].into(), vec![values], 3, Validity::NonNullable)? + .into_array(); + let array = vx_array_new_with(structure); + let mut error = ptr::null_mut(); + + let field = unsafe { vx_velox_array_get_field(session, array, 0, &raw mut error) }; + assert!(!field.is_null()); + assert!(error.is_null()); + assert_eq!(unsafe { vx_array_ref(field)? }.len(), 3); + assert_eq!( + unsafe { vx_velox_array_invalid_count(session, field, &raw mut error) }, + 1 + ); + assert!(error.is_null()); + + let missing = unsafe { vx_velox_array_get_field(session, array, 1, &raw mut error) }; + assert!(missing.is_null()); + assert!(!error.is_null()); + unsafe { vx_error_free(error) }; + error = ptr::null_mut(); + assert_eq!( + unsafe { vx_velox_array_invalid_count(ptr::null(), field, &raw mut error) }, + 0 + ); + assert!(!error.is_null()); + + unsafe { + vx_error_free(error); + vx_velox_array_free(field); + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } +} diff --git a/vortex-velox/src/lib.rs b/vortex-velox/src/lib.rs new file mode 100644 index 00000000000..bb974262d1b --- /dev/null +++ b/vortex-velox/src/lib.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![deny(missing_docs)] +#![expect(non_camel_case_types)] +#![forbid(clippy::todo)] +#![forbid(clippy::unimplemented)] + +//! Native adapter contract between Vortex and Velox. + +mod api; +mod array; +mod projection; +mod read_at; +mod schema; +mod source; +mod visitor; + +pub use api::*; +pub use array::vx_velox_array_export_arrow; +pub use array::vx_velox_array_get_field; +pub use array::vx_velox_array_invalid_count; +pub use array::vx_velox_arrow_memory_callbacks; +pub use projection::vx_velox_expression_select_with_row_index; +pub use read_at::vx_velox_buffer; +pub use read_at::vx_velox_read_at; +pub use read_at::vx_velox_read_at_callbacks; +pub use read_at::vx_velox_read_request; +pub use schema::vx_velox_source_export_schema; +pub use source::vx_velox_natural_split; +pub use source::vx_velox_source; +pub use source::vx_velox_source_prune_natural_splits; +pub use visitor::vx_velox_buffer_owner; +pub use visitor::vx_velox_primitive_type; +pub use visitor::vx_velox_primitive_view; +pub use visitor::vx_velox_validity_kind; +pub use visitor::vx_velox_visit_request; +pub use visitor::vx_velox_visitor; + +/// The current major version of the Vortex and Velox adapter ABI. +pub const VX_VELOX_ABI_VERSION: u32 = 1; + +/// The adapter supports batched host range reads. +pub const VX_VELOX_CAPABILITY_BATCH_READ: u64 = 1 << 0; + +/// The adapter can open callback-backed Vortex sources. +pub const VX_VELOX_CAPABILITY_CALLBACK_SOURCE: u64 = 1 << 1; + +/// The adapter reports stable natural row splits. +pub const VX_VELOX_CAPABILITY_NATURAL_SPLITS: u64 = 1 << 2; + +/// The adapter can visit canonical primitive values in retained blocks. +pub const VX_VELOX_CAPABILITY_PRIMITIVE_VISITOR: u64 = 1 << 3; + +/// The adapter can export source schemas through the Arrow C Data Interface. +pub const VX_VELOX_CAPABILITY_ARROW_SCHEMA: u64 = 1 << 4; + +/// The adapter can export one Vortex array through the Arrow C Data Interface. +pub const VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT: u64 = 1 << 5; + +/// The adapter can project absolute file-row indexes with scan fields. +pub const VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION: u64 = 1 << 6; + +/// The adapter can prove that natural splits cannot match an expression. +pub const VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING: u64 = 1 << 7; + +/// The callback reader observes host cancellation before each host read callback. +/// +/// This capability does not claim cancellation during cached scans or CPU execution. +pub const VX_VELOX_CAPABILITY_READ_CANCELLATION: u64 = 1 << 8; + +/// Return the adapter ABI version. +#[unsafe(no_mangle)] +pub extern "C" fn vx_velox_abi_version() -> u32 { + VX_VELOX_ABI_VERSION +} + +/// Return the capabilities implemented by this adapter build. +#[unsafe(no_mangle)] +pub extern "C" fn vx_velox_capabilities() -> u64 { + VX_VELOX_CAPABILITY_BATCH_READ + | VX_VELOX_CAPABILITY_CALLBACK_SOURCE + | VX_VELOX_CAPABILITY_NATURAL_SPLITS + | VX_VELOX_CAPABILITY_PRIMITIVE_VISITOR + | VX_VELOX_CAPABILITY_ARROW_SCHEMA + | VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT + | VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION + | VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING + | VX_VELOX_CAPABILITY_READ_CANCELLATION +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reports_contract() { + assert_eq!(vx_velox_abi_version(), VX_VELOX_ABI_VERSION); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_BATCH_READ, + VX_VELOX_CAPABILITY_BATCH_READ + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_ARROW_SCHEMA, + VX_VELOX_CAPABILITY_ARROW_SCHEMA + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT, + VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION, + VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING, + VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_READ_CANCELLATION, + VX_VELOX_CAPABILITY_READ_CANCELLATION + ); + } +} diff --git a/vortex-velox/src/projection.rs b/vortex-velox/src/projection.rs new file mode 100644 index 00000000000..c93314d9b73 --- /dev/null +++ b/vortex-velox/src/projection.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; +use std::slice; + +use vortex::dtype::FieldName; +use vortex::dtype::Nullability; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::expr::get_item; +use vortex::expr::pack; +use vortex::expr::root; +use vortex::layout::layouts::row_idx::row_idx; +use vortex_ffi::try_or; +use vortex_ffi::vx_error; +use vortex_ffi::vx_expression; +use vortex_ffi::vx_expression_new_with; +use vortex_ffi::vx_view; + +unsafe fn projection( + names: *const vx_view, + len: usize, + row_index_name: vx_view, +) -> VortexResult<*mut vx_expression> { + vortex_ensure!(!row_index_name.ptr.is_null() || row_index_name.len == 0); + // SAFETY: The caller keeps this view valid for the duration of this call. + let row_index_name = unsafe { row_index_name.as_str() }?; + vortex_ensure!( + !row_index_name.is_empty(), + "row index field name must not be empty" + ); + + let names = if names.is_null() { + vortex_ensure!(len == 0, "null field names pointer with non-zero length"); + &[] + } else { + // SAFETY: The caller provides `len` initialized views when the pointer is non-null. + unsafe { slice::from_raw_parts(names, len) } + }; + + let mut fields = Vec::with_capacity(len + 1); + fields.push((FieldName::from(row_index_name), row_idx())); + for name in names { + // SAFETY: Each caller-provided view remains valid for this call. + let name = unsafe { name.as_str() }?; + vortex_ensure!( + name != row_index_name, + "row index field name conflicts with projected field: {name}" + ); + fields.push((FieldName::from(name), get_item(name, root()))); + } + Ok(vx_expression_new_with(pack( + fields, + Nullability::NonNullable, + ))) +} + +/// Create a struct projection with an absolute file-row index as its first field. +/// +/// The remaining fields select the supplied names from the scan root. The +/// returned expression stays owned by the caller. +/// +/// # Safety +/// +/// `names` must be null when `len` is zero or point to `len` valid views. +/// Every view and `row_index_name` must remain valid for this call. +/// `error_out` must be null or point to writable storage. No input operation can unwind. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_expression_select_with_row_index( + names: *const vx_view, + len: usize, + row_index_name: vx_view, + error_out: *mut *mut vx_error, +) -> *mut vx_expression { + try_or(error_out, ptr::null_mut(), || unsafe { + projection(names, len, row_index_name) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::vx_velox_expression_free; + + fn view(value: &str) -> vx_view { + vx_view { + ptr: value.as_ptr().cast(), + len: value.len(), + } + } + + #[test] + fn creates_row_index_projection() { + let names = [view("a"), view("b")]; + let mut error = ptr::null_mut(); + let expression = unsafe { + vx_velox_expression_select_with_row_index( + names.as_ptr(), + names.len(), + view("$row_index"), + &raw mut error, + ) + }; + assert!(error.is_null()); + assert!(!expression.is_null()); + unsafe { vx_velox_expression_free(expression) }; + } + + #[test] + fn rejects_name_collision() { + let names = [view("$row_index")]; + let mut error = ptr::null_mut(); + let expression = unsafe { + vx_velox_expression_select_with_row_index( + names.as_ptr(), + names.len(), + view("$row_index"), + &raw mut error, + ) + }; + assert!(expression.is_null()); + assert!(!error.is_null()); + unsafe { vortex_ffi::vx_error_free(error) }; + } +} diff --git a/vortex-velox/src/read_at.rs b/vortex-velox/src/read_at.rs new file mode 100644 index 00000000000..89098dc5ce7 --- /dev/null +++ b/vortex-velox/src/read_at.rs @@ -0,0 +1,1012 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::c_char; +use std::ffi::c_void; +use std::mem::size_of; +use std::slice; +use std::sync::Arc; + +use bytes::Bytes; +use futures::FutureExt; +use futures::StreamExt; +use futures::future::BoxFuture; +use futures::stream; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_ffi::ffi_runtime; +use vortex_ffi::try_or; +use vortex_ffi::vx_error; +use vortex_io::ReadAtRequest; +use vortex_io::ReadAtStream; +use vortex_io::VortexReadAt; +use vortex_io::runtime::BlockingRuntime; + +/// A positional read request passed to the Velox callback. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default)] +pub struct vx_velox_read_request { + /// Set this field to `sizeof(vx_velox_read_request)`. + pub struct_size: usize, + /// The file offset in bytes. + pub offset: u64, + /// The exact requested length in bytes. + pub length: usize, + /// The required buffer alignment in bytes. + pub alignment: usize, +} + +/// A retained host buffer returned by the Velox callback. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_buffer { + /// Set this field to `sizeof(vx_velox_buffer)`. + pub struct_size: usize, + /// The first byte of the returned range. + pub data: *const u8, + /// The number of returned bytes. + pub length: usize, + /// An opaque owner passed to `release`. + pub owner: *mut c_void, + /// Release the owner after Vortex no longer needs the bytes. + pub release: Option, +} + +impl Default for vx_velox_buffer { + fn default() -> Self { + Self { + struct_size: size_of::(), + data: std::ptr::null(), + length: 0, + owner: std::ptr::null_mut(), + release: None, + } + } +} + +/// Velox callbacks that provide a Vortex positional reader. +/// +/// Vortex can call these functions concurrently. The context and every callback must be +/// thread-safe. `concurrency` limits one callback batch and gives Vortex a scheduling hint. It does +/// not provide synchronization. `last_error` must return the calling thread's most recent callback +/// error. Its string must remain valid until the next callback on that thread. Every callback must +/// catch foreign exceptions and must not unwind across this ABI. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_read_at_callbacks { + /// Set this field to `sizeof(vx_velox_read_at_callbacks)`. + pub struct_size: usize, + /// Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + pub abi_version: u32, + /// An opaque callback context. + pub context: *mut c_void, + /// Return the file size through `size_out`. Zero means success. + pub size: Option i32>, + /// Read every request and populate the matching output. Zero means success. + pub read_ranges: Option< + unsafe extern "C" fn( + context: *mut c_void, + requests: *const vx_velox_read_request, + request_count: usize, + outputs: *mut vx_velox_buffer, + ) -> i32, + >, + /// Return the last callback error as a null-terminated string. + pub last_error: Option *const c_char>, + /// Release the callback context. + pub release_context: Option, + /// Return a non-zero value after the host cancels the scan. + pub is_cancelled: Option i32>, + /// Limit one callback batch and give Vortex a preferred concurrency value. + pub concurrency: usize, +} + +struct CallbackState { + callbacks: vx_velox_read_at_callbacks, +} + +// SAFETY: The public ABI requires the callback context and functions to permit concurrent calls. +unsafe impl Send for CallbackState {} +// SAFETY: The public ABI requires the callback context and functions to permit concurrent calls. +unsafe impl Sync for CallbackState {} + +impl Drop for CallbackState { + fn drop(&mut self) { + if let Some(release_context) = self.callbacks.release_context { + // SAFETY: The callback contract keeps `context` valid until this call. + unsafe { release_context(self.callbacks.context) }; + } + } +} + +#[derive(Clone)] +pub(crate) struct CallbackReadAt { + state: Arc, +} + +impl CallbackReadAt { + fn try_new(callbacks: vx_velox_read_at_callbacks) -> VortexResult { + if callbacks.struct_size < size_of::() { + vortex_bail!( + "Velox read callback structure is too small: expected at least {}, got {}", + size_of::(), + callbacks.struct_size + ); + } + if callbacks.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + callbacks.abi_version + ); + } + if callbacks.size.is_none() { + vortex_bail!("Velox read callbacks require a size function"); + } + if callbacks.read_ranges.is_none() { + vortex_bail!("Velox read callbacks require a read_ranges function"); + } + if callbacks.is_cancelled.is_none() { + vortex_bail!("Velox read callbacks require an is_cancelled function"); + } + + Ok(Self { + state: Arc::new(CallbackState { callbacks }), + }) + } + + fn ensure_not_cancelled(&self) -> VortexResult<()> { + let is_cancelled = self + .state + .callbacks + .is_cancelled + .vortex_expect("is_cancelled is validated when the reader is created"); + // SAFETY: The callback context stays live while the reader owns its callback state. + if unsafe { is_cancelled(self.state.callbacks.context) } != 0 { + vortex_bail!("Velox cancelled the Vortex read"); + } + Ok(()) + } + + fn last_error(&self, operation: &str, status: i32) -> String { + let Some(last_error) = self.state.callbacks.last_error else { + return format!("Velox {operation} callback failed with status {status}"); + }; + + // SAFETY: The callback contract returns null or a valid null-terminated string. + let message = unsafe { last_error(self.state.callbacks.context) }; + if message.is_null() { + return format!("Velox {operation} callback failed with status {status}"); + } + + // SAFETY: The callback contract keeps the string valid until the next callback call. + unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned() + } + + fn read_batch(&self, requests: Arc<[ReadAtRequest]>) -> Vec> { + if requests.is_empty() { + return Vec::new(); + } + if let Err(error) = self.ensure_not_cancelled() { + let message = error.to_string(); + return requests + .iter() + .map(|_| Err(vortex_err!("{}", message))) + .collect(); + } + + let raw_requests = requests + .iter() + .map(|request| vx_velox_read_request { + struct_size: size_of::(), + offset: request.offset, + length: request.length, + alignment: usize::from(request.alignment), + }) + .collect::>(); + let mut outputs = vec![vx_velox_buffer::default(); requests.len()]; + let read_ranges = self + .state + .callbacks + .read_ranges + .vortex_expect("read_ranges is validated when the reader is created"); + + // SAFETY: The slices remain valid for the duration of the callback. + let status = unsafe { + read_ranges( + self.state.callbacks.context, + raw_requests.as_ptr(), + raw_requests.len(), + outputs.as_mut_ptr(), + ) + }; + if status != 0 { + let message = self.last_error("read_ranges", status); + release_outputs(&mut outputs); + return requests + .iter() + .map(|_| Err(vortex_err!("{}", message))) + .collect(); + } + + outputs + .into_iter() + .zip(requests.iter()) + .map(|(output, request)| output.into_handle(request)) + .collect() + } +} + +impl VortexReadAt for CallbackReadAt { + fn concurrency(&self) -> usize { + self.state.callbacks.concurrency.max(1) + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let reader = self.clone(); + async move { + reader.ensure_not_cancelled()?; + let mut size = 0; + let size_callback = reader + .state + .callbacks + .size + .vortex_expect("size is validated when the reader is created"); + // SAFETY: `size` remains valid for the duration of the callback. + let status = unsafe { size_callback(reader.state.callbacks.context, &raw mut size) }; + if status != 0 { + vortex_bail!("{}", reader.last_error("size", status)); + } + Ok(size) + } + .boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let reader = self.clone(); + async move { + let request = ReadAtRequest::new(offset, length, alignment); + reader + .read_batch(Arc::from([request])) + .pop() + .ok_or_else(|| vortex_err!("Velox read callback returned no result"))? + } + .boxed() + } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + let pairs = requests + .chunks(self.concurrency()) + .flat_map(|requests| { + let requests: Arc<[ReadAtRequest]> = Arc::from(requests); + let results = self.read_batch(Arc::clone(&requests)); + requests.iter().copied().zip(results).collect::>() + }) + .collect::>(); + stream::iter(pairs).boxed() + } +} + +impl vx_velox_read_at { + pub(crate) fn reader(&self) -> CallbackReadAt { + self.0.clone() + } +} + +struct ForeignBuffer { + data: *const u8, + length: usize, + owner: *mut c_void, + release: unsafe extern "C" fn(owner: *mut c_void), +} + +// SAFETY: The buffer contract keeps immutable bytes valid until `release` runs. +unsafe impl Send for ForeignBuffer {} +// SAFETY: The buffer contract permits immutable byte access from concurrent threads. +unsafe impl Sync for ForeignBuffer {} + +impl AsRef<[u8]> for ForeignBuffer { + fn as_ref(&self) -> &[u8] { + // SAFETY: The buffer contract guarantees a valid immutable range for this lifetime. + unsafe { slice::from_raw_parts(self.data, self.length) } + } +} + +impl Drop for ForeignBuffer { + fn drop(&mut self) { + // SAFETY: The owner is released exactly once when the final `Bytes` reference drops. + unsafe { (self.release)(self.owner) }; + } +} + +impl vx_velox_buffer { + fn into_handle(self, request: &ReadAtRequest) -> VortexResult { + if self.struct_size < size_of::() { + self.release_if_present(); + vortex_bail!( + "Velox read callback returned a buffer structure that is too small: expected at least {}, got {}", + size_of::(), + self.struct_size + ); + } + if self.length != request.length { + self.release_if_present(); + vortex_bail!( + "Velox read callback returned {} bytes for a {} byte request at offset {}", + self.length, + request.length, + request.offset + ); + } + if self.length == 0 { + self.release_if_present(); + return Ok(BufferHandle::new_host(ByteBuffer::empty_aligned( + request.alignment, + ))); + } + if self.data.is_null() { + self.release_if_present(); + vortex_bail!( + "Velox read callback returned a null buffer for {} bytes at offset {}", + request.length, + request.offset + ); + } + let Some(release) = self.release else { + vortex_bail!( + "Velox read callback returned bytes without an owner release callback at offset {}", + request.offset + ); + }; + + let owner = ForeignBuffer { + data: self.data, + length: self.length, + owner: self.owner, + release, + }; + let bytes = Bytes::from_owner(owner); + Ok(BufferHandle::new_host( + ByteBuffer::from(bytes).aligned(request.alignment), + )) + } + + fn release_if_present(self) { + if let Some(release) = self.release { + // SAFETY: Failed validation still transfers the returned owner to this adapter. + unsafe { release(self.owner) }; + } + } +} + +fn release_outputs(outputs: &mut [vx_velox_buffer]) { + for output in outputs { + let owned = std::mem::take(output); + owned.release_if_present(); + } +} + +/// An opaque Vortex positional reader backed by Velox callbacks. +pub struct vx_velox_read_at(CallbackReadAt); + +/// Create a Vortex positional reader from Velox callbacks. +/// +/// # Safety +/// +/// `callbacks` must point to a valid callback structure. Every callback and its context must be +/// thread-safe and must not unwind. `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_read_at_new( + callbacks: *const vx_velox_read_at_callbacks, + error_out: *mut *mut vx_error, +) -> *mut vx_velox_read_at { + try_or(error_out, std::ptr::null_mut(), || { + let callbacks = unsafe { + callbacks + .as_ref() + .ok_or_else(|| vortex_err!("Velox read callbacks must not be null"))? + }; + let reader = CallbackReadAt::try_new(*callbacks)?; + Ok(Box::into_raw(Box::new(vx_velox_read_at(reader)))) + }) +} + +/// Free a Vortex positional reader. +/// +/// # Safety +/// +/// `reader` must be null or a pointer returned by [`vx_velox_read_at_new`]. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_read_at_free(reader: *mut vx_velox_read_at) { + if !reader.is_null() { + // SAFETY: The caller transfers the unique pointer returned by the constructor. + drop(unsafe { Box::from_raw(reader) }); + } +} + +/// Return the size of a callback-backed source. +/// +/// This entry point validates the host callback contract before file-reader code consumes the +/// source. +/// +/// # Safety +/// +/// `reader` must point to a live reader. `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_read_at_size( + reader: *const vx_velox_read_at, + error_out: *mut *mut vx_error, +) -> u64 { + try_or(error_out, 0, || { + let reader = unsafe { + reader + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox reader must not be null"))? + }; + ffi_runtime().block_on(reader.0.size()) + }) +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + use std::ffi::CString; + use std::sync::Barrier; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use arrow_array::ffi::FFI_ArrowSchema; + use arrow_schema::Schema; + use futures::executor::block_on; + use vortex::array::IntoArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::StructArray; + use vortex::array::validity::Validity; + use vortex::expr::and; + use vortex::expr::col; + use vortex::expr::gt; + use vortex::expr::lit; + use vortex::expr::lt; + use vortex::file::WriteOptionsSessionExt; + use vortex_buffer::Alignment; + use vortex_error::vortex_ensure; + use vortex_ffi::vx_array_ref; + use vortex_ffi::vx_expression_new_with; + use vortex_ffi::vx_session_free; + use vortex_ffi::vx_session_new_with; + use vortex_ffi::vx_session_ref; + + use super::*; + use crate::api::vx_velox_array_free; + use crate::api::vx_velox_data_source_free; + use crate::api::vx_velox_data_source_scan; + use crate::api::vx_velox_expression_free; + use crate::api::vx_velox_partition_free; + use crate::api::vx_velox_partition_next; + use crate::api::vx_velox_scan_free; + use crate::api::vx_velox_scan_next_partition; + use crate::api::vx_velox_scan_options; + use crate::api::vx_velox_scan_selection; + use crate::schema::vx_velox_source_export_schema; + use crate::source::vx_velox_natural_split; + use crate::source::vx_velox_source_data_source; + use crate::source::vx_velox_source_file_size; + use crate::source::vx_velox_source_free; + use crate::source::vx_velox_source_natural_split_at; + use crate::source::vx_velox_source_natural_split_count; + use crate::source::vx_velox_source_new; + use crate::source::vx_velox_source_prune_natural_splits; + use crate::source::vx_velox_source_row_count; + + struct TestContext { + bytes: Arc<[u8]>, + error: CString, + calls: AtomicUsize, + cancelled: AtomicBool, + fail_after_first_output: AtomicBool, + releases: Arc, + context_releases: Arc, + } + + struct TestOwner { + bytes: Arc<[u8]>, + releases: Arc, + } + + struct ConcurrentErrorContext { + barrier: Barrier, + context_releases: Arc, + } + + thread_local! { + static CONCURRENT_ERROR: Cell<*const c_char> = const { Cell::new(std::ptr::null()) }; + } + + unsafe extern "C" fn concurrent_size(_context: *mut c_void, size_out: *mut u64) -> i32 { + // SAFETY: The callback contract supplies a valid output pointer. + unsafe { size_out.write(2) }; + 0 + } + + unsafe extern "C" fn concurrent_read_ranges( + context: *mut c_void, + requests: *const vx_velox_read_request, + request_count: usize, + _outputs: *mut vx_velox_buffer, + ) -> i32 { + // SAFETY: The test passes a live context and one readable request. + let context = unsafe { &*context.cast::() }; + // SAFETY: The callback contract supplies `request_count` readable requests. + let requests = unsafe { slice::from_raw_parts(requests, request_count) }; + let message = match requests.first().map(|request| request.offset) { + Some(0) => c"read zero failed".as_ptr(), + Some(1) => c"read one failed".as_ptr(), + _ => c"unexpected read failed".as_ptr(), + }; + CONCURRENT_ERROR.with(|error| error.set(message)); + context.barrier.wait(); + 1 + } + + unsafe extern "C" fn concurrent_last_error(_context: *mut c_void) -> *const c_char { + CONCURRENT_ERROR.with(|error| error.get()) + } + + unsafe extern "C" fn concurrent_release_context(context: *mut c_void) { + // SAFETY: The test created this context with `Box::into_raw`. + let context = unsafe { Box::from_raw(context.cast::()) }; + context.context_releases.fetch_add(1, Ordering::Relaxed); + } + + unsafe extern "C" fn never_cancelled(_context: *mut c_void) -> i32 { + 0 + } + + unsafe extern "C" fn test_size(context: *mut c_void, size_out: *mut u64) -> i32 { + // SAFETY: Tests pass a `TestContext` and a valid output pointer. + let context = unsafe { &*context.cast::() }; + // SAFETY: The callback contract supplies a valid output pointer. + unsafe { size_out.write(context.bytes.len() as u64) }; + 0 + } + + unsafe extern "C" fn test_read_ranges( + context: *mut c_void, + requests: *const vx_velox_read_request, + request_count: usize, + outputs: *mut vx_velox_buffer, + ) -> i32 { + // SAFETY: Tests pass valid callback arguments. + let context = unsafe { &*context.cast::() }; + context.calls.fetch_add(1, Ordering::Relaxed); + // SAFETY: The callback contract supplies arrays with `request_count` entries. + let requests = unsafe { slice::from_raw_parts(requests, request_count) }; + // SAFETY: The callback contract supplies writable outputs with matching length. + let outputs = unsafe { slice::from_raw_parts_mut(outputs, request_count) }; + for (request, output) in requests.iter().zip(outputs) { + let Ok(start) = usize::try_from(request.offset) else { + return 1; + }; + let Some(end) = start.checked_add(request.length) else { + return 1; + }; + if end > context.bytes.len() { + return 1; + } + let owner = Box::new(TestOwner { + bytes: Arc::clone(&context.bytes), + releases: Arc::clone(&context.releases), + }); + output.data = owner.bytes[start..end].as_ptr(); + output.length = request.length; + output.owner = Box::into_raw(owner).cast(); + output.release = Some(test_release_buffer); + if context.fail_after_first_output.load(Ordering::Relaxed) { + return 1; + } + } + 0 + } + + unsafe extern "C" fn test_release_buffer(owner: *mut c_void) { + // SAFETY: The test callback created this owner with `Box::into_raw`. + let owner = unsafe { Box::from_raw(owner.cast::()) }; + owner.releases.fetch_add(1, Ordering::Relaxed); + } + + unsafe extern "C" fn test_last_error(context: *mut c_void) -> *const c_char { + // SAFETY: Tests pass a `TestContext`. + let context = unsafe { &*context.cast::() }; + context.error.as_ptr() + } + + unsafe extern "C" fn test_release_context(context: *mut c_void) { + // SAFETY: The test created this context with `Box::into_raw`. + let context = unsafe { Box::from_raw(context.cast::()) }; + context.context_releases.fetch_add(1, Ordering::Relaxed); + } + + unsafe extern "C" fn test_is_cancelled(context: *mut c_void) -> i32 { + // SAFETY: Tests pass a `TestContext`. + let context = unsafe { &*context.cast::() }; + i32::from(context.cancelled.load(Ordering::Relaxed)) + } + + fn callbacks( + bytes: &[u8], + releases: Arc, + context_releases: Arc, + ) -> vx_velox_read_at_callbacks { + let context = Box::new(TestContext { + bytes: Arc::from(bytes), + error: c"test read failed".to_owned(), + calls: AtomicUsize::new(0), + cancelled: AtomicBool::new(false), + fail_after_first_output: AtomicBool::new(false), + releases, + context_releases, + }); + vx_velox_read_at_callbacks { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: Box::into_raw(context).cast(), + size: Some(test_size), + read_ranges: Some(test_read_ranges), + last_error: Some(test_last_error), + release_context: Some(test_release_context), + is_cancelled: Some(test_is_cancelled), + concurrency: 8, + } + } + + #[test] + fn batches_reads_and_releases_owners() -> VortexResult<()> { + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let read_at = CallbackReadAt::try_new(callbacks( + b"abcdefgh", + Arc::clone(&releases), + Arc::clone(&context_releases), + ))?; + + assert_eq!(block_on(read_at.size())?, 8); + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(1, 3, Alignment::none()), + ReadAtRequest::new(5, 2, Alignment::none()), + ]); + let results = block_on(read_at.read_ranges(requests).collect::>()); + assert_eq!(results.len(), 2); + let mut results = results.into_iter(); + let (_, first) = results + .next() + .vortex_expect("the first read result is present"); + let (_, second) = results + .next() + .vortex_expect("the second read result is present"); + let first = first?; + let second = second?; + assert_eq!(first.to_host_sync().as_ref(), b"bcd"); + assert_eq!(second.to_host_sync().as_ref(), b"fg"); + assert_eq!(read_at.state.callbacks.concurrency, 8); + assert_eq!(releases.load(Ordering::Relaxed), 0); + + drop((first, second)); + assert_eq!(releases.load(Ordering::Relaxed), 2); + drop(read_at); + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn respects_callback_batch_limit() -> VortexResult<()> { + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let mut callbacks = callbacks( + b"abcdefgh", + Arc::clone(&releases), + Arc::clone(&context_releases), + ); + callbacks.concurrency = 1; + // SAFETY: The callback context stays owned by `callbacks` until reader destruction. + let context = unsafe { &*callbacks.context.cast::() }; + let read_at = CallbackReadAt::try_new(callbacks)?; + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(0, 2, Alignment::none()), + ReadAtRequest::new(2, 2, Alignment::none()), + ]); + let results = block_on(read_at.read_ranges(requests).collect::>()); + assert_eq!(results.len(), 2); + assert!(results.iter().all(|(_, result)| result.is_ok())); + assert_eq!(context.calls.load(Ordering::Relaxed), 2); + drop(results); + assert_eq!(releases.load(Ordering::Relaxed), 2); + drop(read_at); + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn rejects_incomplete_callback_table() { + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let mut callbacks = callbacks(b"abc", releases, Arc::clone(&context_releases)); + callbacks.struct_size -= 1; + let result = CallbackReadAt::try_new(callbacks); + assert!(result.is_err()); + + // The constructor did not take ownership after validation failed. + unsafe { test_release_context(callbacks.context) }; + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + } + + #[test] + fn releases_partial_outputs_after_callback_failure() -> VortexResult<()> { + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let callbacks = callbacks( + b"abcdefgh", + Arc::clone(&releases), + Arc::clone(&context_releases), + ); + // SAFETY: The callback context stays owned by `callbacks` until reader destruction. + unsafe { &*callbacks.context.cast::() } + .fail_after_first_output + .store(true, Ordering::Relaxed); + let read_at = CallbackReadAt::try_new(callbacks)?; + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(0, 2, Alignment::none()), + ReadAtRequest::new(2, 2, Alignment::none()), + ]); + let results = block_on(read_at.read_ranges(requests).collect::>()); + assert_eq!(results.len(), 2); + assert!(results.iter().all(|(_, result)| result.is_err())); + assert_eq!(releases.load(Ordering::Relaxed), 1); + drop(read_at); + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn observes_host_cancellation_before_callbacks() -> VortexResult<()> { + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let callbacks = callbacks( + b"abcdefgh", + Arc::clone(&releases), + Arc::clone(&context_releases), + ); + // SAFETY: The callback context stays owned by `callbacks` until reader destruction. + let context = unsafe { &*callbacks.context.cast::() }; + context.cancelled.store(true, Ordering::Relaxed); + let read_at = CallbackReadAt::try_new(callbacks)?; + + let error = block_on(read_at.size()).expect_err("cancelled size must fail"); + assert!(error.to_string().contains("cancelled")); + assert_eq!(context.calls.load(Ordering::Relaxed), 0); + assert_eq!(releases.load(Ordering::Relaxed), 0); + drop(read_at); + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn preserves_per_thread_errors_across_concurrent_callbacks() -> VortexResult<()> { + let context_releases = Arc::new(AtomicUsize::new(0)); + let context = Box::new(ConcurrentErrorContext { + barrier: Barrier::new(2), + context_releases: Arc::clone(&context_releases), + }); + let reader = CallbackReadAt::try_new(vx_velox_read_at_callbacks { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: Box::into_raw(context).cast(), + size: Some(concurrent_size), + read_ranges: Some(concurrent_read_ranges), + last_error: Some(concurrent_last_error), + release_context: Some(concurrent_release_context), + is_cancelled: Some(never_cancelled), + concurrency: 2, + })?; + + let first = reader.clone(); + let first = std::thread::spawn(move || { + block_on(first.read_at(0, 1, Alignment::none())) + .expect_err("the first callback must fail") + .to_string() + }); + let second = reader.clone(); + let second = std::thread::spawn(move || { + block_on(second.read_at(1, 1, Alignment::none())) + .expect_err("the second callback must fail") + .to_string() + }); + + let first = first + .join() + .map_err(|_| vortex_err!("The first callback thread panicked"))?; + let second = second + .join() + .map_err(|_| vortex_err!("The second callback thread panicked"))?; + assert!(first.contains("zero")); + assert!(second.contains("one")); + drop(reader); + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn opens_source_and_reports_natural_splits() -> VortexResult<()> { + let session_handle = vx_session_new_with(|session| session); + // SAFETY: The test owns the live session handle. + let session = unsafe { vx_session_ref(session_handle)? }.clone(); + const ROW_COUNT: u64 = 300_000; + const ROWS_PER_NATURAL_SPLIT: u64 = 100_000; + let values = PrimitiveArray::from_iter(0_i64..i64::try_from(ROW_COUNT)?).into_array(); + let array = StructArray::try_new( + ["value"].into(), + vec![values], + usize::try_from(ROW_COUNT)?, + Validity::NonNullable, + )? + .into_array(); + let mut bytes = Vec::new(); + session + .write_options() + .blocking(ffi_runtime()) + .write(&mut bytes, array.to_array_iterator())?; + + let releases = Arc::new(AtomicUsize::new(0)); + let context_releases = Arc::new(AtomicUsize::new(0)); + let reader = CallbackReadAt::try_new(callbacks( + &bytes, + Arc::clone(&releases), + Arc::clone(&context_releases), + ))?; + let reader_handle = Box::into_raw(Box::new(vx_velox_read_at(reader))); + let mut error = std::ptr::null_mut(); + // SAFETY: The test owns all handles and output pointers. + let source = unsafe { vx_velox_source_new(session_handle, reader_handle, &raw mut error) }; + vortex_ensure!(error.is_null(), "source open returned an error"); + vortex_ensure!(!source.is_null(), "source open returned null"); + + // SAFETY: The source stays live for all calls. + unsafe { + assert_eq!(vx_velox_source_row_count(source), ROW_COUNT); + assert_eq!(vx_velox_source_file_size(source), bytes.len() as u64); + } + let mut schema = FFI_ArrowSchema::empty(); + // SAFETY: The source and outputs stay live for this call. + let status = + unsafe { vx_velox_source_export_schema(source, &raw mut schema, &raw mut error) }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "source schema returned an error"); + let schema = Schema::try_from(&schema)?; + assert_eq!(schema.fields().len(), 1); + assert_eq!(schema.field(0).name(), "value"); + // SAFETY: The source stays live for this call. + let split_count = unsafe { vx_velox_source_natural_split_count(source) }; + assert!(split_count > 0); + let mut previous_end = 0; + for index in 0..split_count { + let mut split = vx_velox_natural_split { + struct_size: size_of::(), + ..Default::default() + }; + // SAFETY: The source and outputs stay live for this call. + let status = unsafe { + vx_velox_source_natural_split_at(source, index, &raw mut split, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "natural split lookup returned an error"); + assert_eq!(split.row_begin, previous_end); + assert!(split.row_end > split.row_begin); + previous_end = split.row_end; + } + assert_eq!(previous_end, ROW_COUNT); + + let expression = vx_expression_new_with(and( + gt( + col("value"), + lit(i64::try_from(ROWS_PER_NATURAL_SPLIT - 1)?), + ), + lt( + col("value"), + lit(i64::try_from(2 * ROWS_PER_NATURAL_SPLIT)?), + ), + )); + let mut pruned = vec![0; split_count]; + // SAFETY: The source, expression, output, and error pointer stay live for this call. + let status = unsafe { + vx_velox_source_prune_natural_splits( + source, + expression, + 0, + split_count, + pruned.as_mut_ptr(), + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "natural split pruning returned an error"); + assert_eq!(pruned.first(), Some(&1)); + assert_eq!(pruned.last(), Some(&1)); + assert!(pruned[1..pruned.len() - 1].contains(&0)); + // SAFETY: The test owns this expression handle. + unsafe { vx_velox_expression_free(expression) }; + + // SAFETY: The source and error output stay live for this call. + let data_source = unsafe { vx_velox_source_data_source(source, &raw mut error) }; + vortex_ensure!(error.is_null(), "data source conversion returned an error"); + vortex_ensure!( + !data_source.is_null(), + "data source conversion returned null" + ); + + let scan_options = vx_velox_scan_options { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + projection: std::ptr::null(), + filter: std::ptr::null(), + row_range_begin: 0, + row_range_end: 0, + selection: vx_velox_scan_selection::default(), + limit: 0, + ordered: false, + }; + // SAFETY: The data source and scan options stay live for this call. + let scan = unsafe { + vx_velox_data_source_scan(data_source, &raw const scan_options, &raw mut error) + }; + vortex_ensure!(error.is_null(), "data source scan returned an error"); + vortex_ensure!(!scan.is_null(), "data source scan returned null"); + let mut scanned_rows = 0; + loop { + // SAFETY: The scan stays live and is consumed from one thread. + let partition = unsafe { vx_velox_scan_next_partition(scan, &raw mut error) }; + vortex_ensure!(error.is_null(), "partition lookup returned an error"); + if partition.is_null() { + break; + } + loop { + // SAFETY: The partition stays live and is consumed from one thread. + let array = unsafe { vx_velox_partition_next(partition, &raw mut error) }; + vortex_ensure!(error.is_null(), "partition scan returned an error"); + if array.is_null() { + break; + } + // SAFETY: The returned array stays live until the matching free call. + scanned_rows += unsafe { vx_array_ref(array)? }.len(); + // SAFETY: The scan returned this owned array handle. + unsafe { vx_velox_array_free(array) }; + } + // SAFETY: The scan returned this owned partition handle. + unsafe { vx_velox_partition_free(partition) }; + } + assert_eq!(scanned_rows, usize::try_from(ROW_COUNT)?); + + // SAFETY: Each owned handle is freed exactly once. + unsafe { + vx_velox_scan_free(scan); + vx_velox_data_source_free(data_source); + vx_velox_source_free(source); + vx_velox_read_at_free(reader_handle); + vx_session_free(session_handle); + } + assert_eq!(context_releases.load(Ordering::Relaxed), 1); + assert!(releases.load(Ordering::Relaxed) > 0); + Ok(()) + } +} diff --git a/vortex-velox/src/schema.rs b/vortex-velox/src/schema.rs new file mode 100644 index 00000000000..aedbb2bdf66 --- /dev/null +++ b/vortex-velox/src/schema.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; + +use arrow_array::ffi::FFI_ArrowSchema; +use vortex_arrow::ArrowSessionExt; +use vortex_error::vortex_err; +use vortex_ffi::try_or; +use vortex_ffi::vx_error; + +use crate::source::vx_velox_source; + +/// Export an opened source schema through the Arrow C Data Interface. +/// +/// The caller owns the output and must invoke its release callback. +/// +/// # Safety +/// +/// `source` must point to a live source. `schema_out` must identify uninitialized writable +/// storage. `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_source_export_schema( + source: *const vx_velox_source, + schema_out: *mut FFI_ArrowSchema, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let source = unsafe { + source + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox source must not be null"))? + }; + if schema_out.is_null() { + return Err(vortex_err!("Arrow schema output must not be null")); + } + let arrow_schema = source + .file() + .session() + .arrow() + .to_arrow_schema(source.file().dtype())?; + let schema = FFI_ArrowSchema::try_from(&arrow_schema)?; + unsafe { ptr::write(schema_out, schema) }; + Ok(0) + }) +} + +#[cfg(test)] +mod tests { + use arrow_array::ffi::FFI_ArrowSchema; + use vortex_error::VortexResult; + + use super::*; + + #[test] + fn rejects_null_source() -> VortexResult<()> { + let mut schema = FFI_ArrowSchema::empty(); + let mut error = ptr::null_mut(); + let status = + unsafe { vx_velox_source_export_schema(ptr::null(), &raw mut schema, &raw mut error) }; + assert_eq!(status, 1); + assert!(!error.is_null()); + unsafe { vortex_ffi::vx_error_free(error) }; + Ok(()) + } +} diff --git a/vortex-velox/src/source.rs b/vortex-velox/src/source.rs new file mode 100644 index 00000000000..7aace08a49c --- /dev/null +++ b/vortex-velox/src/source.rs @@ -0,0 +1,278 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::task::Poll; + +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::layout::scan::multi::MultiLayoutDataSource; +use vortex::mask::Mask; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_ffi::ffi_runtime; +use vortex_ffi::try_or; +use vortex_ffi::vx_data_source; +use vortex_ffi::vx_data_source_new_with; +use vortex_ffi::vx_error; +use vortex_ffi::vx_expression; +use vortex_ffi::vx_expression_ref; +use vortex_ffi::vx_session; +use vortex_ffi::vx_session_ref; +use vortex_io::VortexReadAt; +use vortex_io::runtime::BlockingRuntime; + +use crate::read_at::vx_velox_read_at; + +/// A stable natural row range reported by a Vortex file. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct vx_velox_natural_split { + /// Set this field to `sizeof(vx_velox_natural_split)`. + pub struct_size: usize, + /// The first row in the split. + pub row_begin: u64, + /// One past the final row in the split. + pub row_end: u64, +} + +/// An opened Vortex file that uses Velox callbacks for all reads. +pub struct vx_velox_source { + file: VortexFile, + file_size: u64, + natural_splits: Vec>, +} + +impl vx_velox_source { + pub(crate) fn file(&self) -> &VortexFile { + &self.file + } +} + +/// Open a Vortex source through a callback reader. +/// +/// The source retains the session and reader state. The caller can free both input handles after +/// this function returns. +/// +/// # Safety +/// +/// `session` and `reader` must point to live handles. `error_out` must be null or valid for one +/// error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_source_new( + session: *const vx_session, + reader: *const vx_velox_read_at, + error_out: *mut *mut vx_error, +) -> *mut vx_velox_source { + try_or(error_out, std::ptr::null_mut(), || { + let session = unsafe { vx_session_ref(session)? }.clone(); + let reader = unsafe { + reader + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox reader must not be null"))? + } + .reader(); + let file_size = ffi_runtime().block_on(reader.size())?; + let file = ffi_runtime().block_on( + session + .open_options() + .with_file_size(file_size) + .with_layout_reader_cache() + .open_read(reader), + )?; + let natural_splits = file.splits()?; + Ok(Box::into_raw(Box::new(vx_velox_source { + file, + file_size, + natural_splits, + }))) + }) +} + +/// Free a callback-backed Vortex source and release its callback-owned input buffers. +/// +/// # Safety +/// +/// `source` must be null or a pointer returned by [`vx_velox_source_new`]. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_source_free(source: *mut vx_velox_source) { + if !source.is_null() { + // SAFETY: The caller transfers the unique pointer returned by the constructor. + drop(unsafe { Box::from_raw(source) }); + // The runtime defers some stream drops to its next turn. Run one turn so the C ABI + // releases every callback-owned buffer before this function returns. + let mut yielded = false; + ffi_runtime().block_on(futures::future::poll_fn(|context| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + context.waker().wake_by_ref(); + Poll::Pending + } + })); + } +} + +/// Return the file row count. +/// +/// # Safety +/// +/// `source` must point to a live source. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_source_row_count(source: *const vx_velox_source) -> u64 { + // SAFETY: The caller provides a live source. + unsafe { &*source }.file.row_count() +} + +/// Return the file size in bytes. +/// +/// # Safety +/// +/// `source` must point to a live source. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_source_file_size(source: *const vx_velox_source) -> u64 { + // SAFETY: The caller provides a live source. + unsafe { &*source }.file_size +} + +/// Return the number of natural row splits. +/// +/// # Safety +/// +/// `source` must point to a live source. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_source_natural_split_count( + source: *const vx_velox_source, +) -> usize { + // SAFETY: The caller provides a live source. + unsafe { &*source }.natural_splits.len() +} + +/// Write one natural row split. +/// +/// # Safety +/// +/// `source` must point to a live source. `split_out` must point to a structure with a valid size. +/// `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_source_natural_split_at( + source: *const vx_velox_source, + index: usize, + split_out: *mut vx_velox_natural_split, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let source = unsafe { + source + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox source must not be null"))? + }; + let split_out = unsafe { + split_out + .as_mut() + .ok_or_else(|| vortex_err!("Natural split output must not be null"))? + }; + if split_out.struct_size < size_of::() { + vortex_bail!( + "Natural split structure is too small: expected at least {}, got {}", + size_of::(), + split_out.struct_size + ); + } + let split = source + .natural_splits + .get(index) + .ok_or_else(|| vortex_err!("Natural split index out of bounds: {}", index))?; + split_out.row_begin = split.start; + split_out.row_end = split.end; + Ok(0) + }) +} + +/// Evaluate whether natural splits cannot match an expression. +/// +/// Each output byte is one when the matching split cannot produce a true expression result. Zero +/// means that the split can match or that available statistics cannot prove exclusion. +/// +/// # Safety +/// +/// `source` and `expression` must point to live handles. `pruned_out` must identify `split_count` +/// writable bytes unless `split_count` is zero. `error_out` must be null or valid for one error +/// pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_source_prune_natural_splits( + source: *const vx_velox_source, + expression: *const vx_expression, + first_split: usize, + split_count: usize, + pruned_out: *mut u8, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let source = unsafe { + source + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox source must not be null"))? + }; + let split_end = first_split + .checked_add(split_count) + .ok_or_else(|| vortex_err!("Natural split range overflow"))?; + if split_end > source.natural_splits.len() { + vortex_bail!( + "Natural split range out of bounds: {}..{} for {} splits", + first_split, + split_end, + source.natural_splits.len() + ); + } + if split_count == 0 { + return Ok(0); + } + if pruned_out.is_null() { + vortex_bail!("Natural split pruning output must not be null"); + } + let expression = unsafe { vx_expression_ref(expression)? }; + let bound = expression.bind(source.file.dtype())?; + let reader = source.file.layout_reader()?; + let output = unsafe { std::slice::from_raw_parts_mut(pruned_out, split_count) }; + for (decision, row_range) in output.iter_mut().zip( + source.natural_splits[first_split..split_end] + .iter() + .cloned(), + ) { + let row_count = usize::try_from(row_range.end - row_range.start)?; + let mask = reader.pruning_evaluation(&row_range, &bound, Mask::new_true(row_count))?; + *decision = u8::from(ffi_runtime().block_on(mask)?.all_false()); + } + Ok(0) + }) +} + +/// Create a standard Vortex data source for this file. +/// +/// The caller owns the returned handle and must free it through `vx_data_source_free`. +/// +/// # Safety +/// +/// `source` must point to a live source. `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_source_data_source( + source: *const vx_velox_source, + error_out: *mut *mut vx_error, +) -> *const vx_data_source { + try_or(error_out, std::ptr::null(), || { + let source = unsafe { + source + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox source must not be null"))? + }; + let data_source = MultiLayoutDataSource::new_with_first( + source.file.layout_reader()?, + Vec::new(), + vec![Some(source.file_size)], + source.file.session(), + ); + Ok(vx_data_source_new_with(data_source)) + }) +} diff --git a/vortex-velox/src/visitor.rs b/vortex-velox/src/visitor.rs new file mode 100644 index 00000000000..0dfc4f13455 --- /dev/null +++ b/vortex-velox/src/visitor.rs @@ -0,0 +1,706 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::c_char; +use std::ffi::c_void; +use std::mem::size_of; +use std::slice; +use std::sync::Arc; + +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::primitive::PrimitiveArrayExt; +use vortex::buffer::ByteBuffer; +use vortex::dtype::PType; +use vortex::mask::Mask; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_ffi::try_or; +use vortex_ffi::vx_array; +use vortex_ffi::vx_array_ref; +use vortex_ffi::vx_error; +use vortex_ffi::vx_session; +use vortex_ffi::vx_session_ref; + +/// A fixed-width primitive value identifier in a semantic visitor block. +pub type vx_velox_primitive_type = u32; +/// Unsigned 8-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_U8: vx_velox_primitive_type = 0; +/// Unsigned 16-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_U16: vx_velox_primitive_type = 1; +/// Unsigned 32-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_U32: vx_velox_primitive_type = 2; +/// Unsigned 64-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_U64: vx_velox_primitive_type = 3; +/// Signed 8-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I8: vx_velox_primitive_type = 4; +/// Signed 16-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I16: vx_velox_primitive_type = 5; +/// Signed 32-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I32: vx_velox_primitive_type = 6; +/// Signed 64-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I64: vx_velox_primitive_type = 7; +/// IEEE 754 binary16 primitive identifier. +pub const VX_VELOX_PRIMITIVE_F16: vx_velox_primitive_type = 8; +/// IEEE 754 binary32 primitive identifier. +pub const VX_VELOX_PRIMITIVE_F32: vx_velox_primitive_type = 9; +/// IEEE 754 binary64 primitive identifier. +pub const VX_VELOX_PRIMITIVE_F64: vx_velox_primitive_type = 10; + +fn primitive_type_id(value: PType) -> vx_velox_primitive_type { + match value { + PType::U8 => VX_VELOX_PRIMITIVE_U8, + PType::U16 => VX_VELOX_PRIMITIVE_U16, + PType::U32 => VX_VELOX_PRIMITIVE_U32, + PType::U64 => VX_VELOX_PRIMITIVE_U64, + PType::I8 => VX_VELOX_PRIMITIVE_I8, + PType::I16 => VX_VELOX_PRIMITIVE_I16, + PType::I32 => VX_VELOX_PRIMITIVE_I32, + PType::I64 => VX_VELOX_PRIMITIVE_I64, + PType::F16 => VX_VELOX_PRIMITIVE_F16, + PType::F32 => VX_VELOX_PRIMITIVE_F32, + PType::F64 => VX_VELOX_PRIMITIVE_F64, + } +} + +/// A fixed-width validity representation identifier for one visitor block. +pub type vx_velox_validity_kind = u32; +/// The type is not nullable. +pub const VX_VELOX_VALIDITY_NON_NULLABLE: vx_velox_validity_kind = 0; +/// Every value is valid. +pub const VX_VELOX_VALIDITY_ALL_VALID: vx_velox_validity_kind = 1; +/// Every value is null. +pub const VX_VELOX_VALIDITY_ALL_INVALID: vx_velox_validity_kind = 2; +/// A packed bitmap contains one valid bit per value. +pub const VX_VELOX_VALIDITY_BITMAP: vx_velox_validity_kind = 3; + +/// A retained owner for buffers in a visitor block. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_buffer_owner { + /// Set this field to `sizeof(vx_velox_buffer_owner)`. + pub struct_size: usize, + /// An opaque retained object. + pub owner: *const c_void, + /// Add one owner reference before the callback returns. + pub retain: Option, + /// Release one retained owner reference. + pub release: Option, + /// The allocated number of payload bytes retained by this compact owner. + pub retained_bytes: usize, +} + +/// A canonical primitive block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_primitive_view { + /// Set this field to `sizeof(vx_velox_primitive_view)`. + pub struct_size: usize, + /// The physical type of each value. + pub primitive_type: vx_velox_primitive_type, + /// The number of logical values in the block. + pub length: usize, + /// The first value byte. + pub values: *const u8, + /// The number of value bytes. + pub values_length: usize, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains all pointers in this view. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty values buffer. + pub values_alignment: usize, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + +/// A single-shot subset request for the semantic visitor. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_visit_request { + /// Set this field to `sizeof(vx_velox_visit_request)`. + pub struct_size: usize, + /// Unique, increasing source positions. Null selects every row. + pub rows: *const u64, + /// The number of source positions. + pub row_count: usize, +} + +/// Host callbacks for Vortex array traversal. +/// +/// One array visit calls the primitive callback synchronously. Shared tables can receive concurrent +/// callbacks from simultaneous visits. `last_error` must return the calling thread's most recent +/// error. The string must remain valid until the next callback on that thread. Callbacks must catch +/// foreign exceptions and must not unwind across this ABI. The host owns the context. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_visitor { + /// Set this field to `sizeof(vx_velox_visitor)`. + pub struct_size: usize, + /// Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + pub abi_version: u32, + /// An opaque callback context. + pub context: *mut c_void, + /// Consume one canonical primitive block. Zero means success. + pub visit_primitive: Option< + unsafe extern "C" fn(context: *mut c_void, view: *const vx_velox_primitive_view) -> i32, + >, + /// Return the last callback error as a null-terminated string. + pub last_error: Option *const c_char>, +} + +struct PrimitiveOwner { + values: Box<[u64]>, + values_length: usize, + validity: Option>, + retained_bytes: usize, +} + +impl PrimitiveOwner { + fn try_new( + host_values: &ByteBuffer, + validity: Option<&vortex::buffer::BitBuffer>, + length: usize, + ) -> VortexResult { + let values_length = host_values.len(); + let mut values = vec![0_u64; values_length.div_ceil(size_of::())].into_boxed_slice(); + let values_allocation = values + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + if values_length != 0 { + // SAFETY: The byte view spans the complete initialized `u64` allocation. + let values_bytes = unsafe { + slice::from_raw_parts_mut(values.as_mut_ptr().cast::(), values_allocation) + }; + values_bytes[..values_length].copy_from_slice(host_values.as_slice()); + } + let validity = validity.map(|validity| { + let mut compact = vec![0_u8; length.div_ceil(8)].into_boxed_slice(); + for (index, is_valid) in validity.into_iter().take(length).enumerate() { + if is_valid { + compact[index / 8] |= 1 << (index % 8); + } + } + compact + }); + let retained_bytes = values_allocation + .checked_add(validity.as_ref().map_or(0, |validity| validity.len())) + .ok_or_else(|| vortex_err!("Primitive visitor retained byte count overflow"))?; + Ok(Self { + values, + values_length, + validity, + retained_bytes, + }) + } + + fn values(&self) -> *const u8 { + if self.values_length == 0 { + std::ptr::null() + } else { + self.values.as_ptr().cast() + } + } + + fn validity(&self) -> *const u8 { + self.validity + .as_ref() + .filter(|validity| !validity.is_empty()) + .map_or(std::ptr::null(), |validity| validity.as_ptr()) + } + + fn values_length(&self) -> usize { + self.values_length + } + + fn validity_length(&self) -> usize { + self.validity.as_ref().map_or(0, |validity| validity.len()) + } + + fn retained_bytes(&self) -> usize { + self.retained_bytes + } +} + +fn pointer_alignment(pointer: *const u8) -> usize { + if pointer.is_null() { + return 0; + } + 1usize << pointer.addr().trailing_zeros() +} + +unsafe extern "C" fn retain_primitive_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_primitive_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +fn validate_visitor(visitor: &vx_velox_visitor) -> VortexResult<()> { + if visitor.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox visitor structure is too small: expected at least {}, got {}", + size_of::(), + visitor.struct_size + ); + } + if visitor.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + visitor.abi_version + ); + } + if visitor.visit_primitive.is_none() { + vortex_bail!("Vortex Velox visitor requires a primitive callback"); + } + Ok(()) +} + +fn callback_error(visitor: &vx_velox_visitor, status: i32) -> String { + let Some(last_error) = visitor.last_error else { + return format!("Velox primitive visitor failed with status {status}"); + }; + // SAFETY: The callback contract returns null or a valid null-terminated string. + let message = unsafe { last_error(visitor.context) }; + if message.is_null() { + return format!("Velox primitive visitor failed with status {status}"); + } + // SAFETY: The callback keeps the string valid until the next callback. + unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned() +} + +fn selected_array( + array: &vortex::array::ArrayRef, + request: &vx_velox_visit_request, +) -> VortexResult { + if request.rows.is_null() { + if request.row_count != 0 { + vortex_bail!("A null visitor row pointer requires a zero row count"); + } + return Ok(array.clone()); + } + // SAFETY: The caller supplies `row_count` readable positions. + let rows = unsafe { slice::from_raw_parts(request.rows, request.row_count) }; + let mut previous = None; + for row in rows { + let position = usize::try_from(*row) + .map_err(|_| vortex_err!("Visitor row does not fit usize: {}", row))?; + if position >= array.len() { + vortex_bail!( + "Visitor row is out of bounds: row {}, array length {}", + row, + array.len() + ); + } + if previous.is_some_and(|previous| previous >= *row) { + vortex_bail!("Visitor rows must be unique and increasing"); + } + previous = Some(*row); + } + let dense = rows.len() == array.len() + && rows + .iter() + .enumerate() + .all(|(position, row)| *row == position as u64); + if dense { + return Ok(array.clone()); + } + array.take(PrimitiveArray::from_iter(rows.iter().copied()).into_array()) +} + +fn visit_primitive( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + visitor: &vx_velox_visitor, +) -> VortexResult<()> { + let mut execution = session.create_execution_ctx(); + let Canonical::Primitive(primitive) = array.execute::(&mut execution)? else { + vortex_bail!("Primitive visitor received a non-primitive array"); + }; + let values = primitive.buffer_handle().clone(); + let host_values = values.try_to_host_sync()?; + let mask = primitive + .validity()? + .execute_mask(primitive.len(), &mut execution)?; + let (validity_kind, validity) = if !primitive.dtype().is_nullable() { + (VX_VELOX_VALIDITY_NON_NULLABLE, None) + } else { + match mask { + Mask::AllTrue(_) => (VX_VELOX_VALIDITY_ALL_VALID, None), + Mask::AllFalse(_) => (VX_VELOX_VALIDITY_ALL_INVALID, None), + Mask::Values(values) => (VX_VELOX_VALIDITY_BITMAP, Some(values.bit_buffer().clone())), + } + }; + let owner = Arc::new(PrimitiveOwner::try_new( + &host_values, + validity.as_ref(), + primitive.len(), + )?); + let values_length = owner.values_length(); + let validity_length = owner.validity_length(); + let values = owner.values(); + let validity = owner.validity(); + let view = vx_velox_primitive_view { + struct_size: size_of::(), + primitive_type: primitive_type_id(primitive.ptype()), + length: primitive.len(), + values, + values_length, + validity_kind, + validity, + validity_length, + validity_bit_offset: 0, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&owner).cast(), + retain: Some(retain_primitive_owner), + release: Some(release_primitive_owner), + retained_bytes: owner.retained_bytes(), + }, + values_alignment: pointer_alignment(values), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_primitive + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a primitive callback"))?; + // SAFETY: The view and its local owner stay live until the callback returns. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) +} + +/// Visit one Vortex array through host semantic callbacks. +/// +/// The request selects source positions once. Callback block positions are compact and follow the +/// request order. +/// +/// # Safety +/// +/// Every pointer must be null or valid for the documented access. The array and session handles +/// must remain live until this call returns. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_visit( + session: *const vx_session, + array: *const vx_array, + request: *const vx_velox_visit_request, + visitor: *const vx_velox_visitor, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + let request = unsafe { + request + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visit request must not be null"))? + }; + if request.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox visit request is too small: expected at least {}, got {}", + size_of::(), + request.struct_size + ); + } + let visitor = unsafe { + visitor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? + }; + validate_visitor(visitor)?; + visit_primitive(selected_array(array, request)?, session, visitor)?; + Ok(0) + }) +} + +#[cfg(test)] +mod tests { + use std::ptr; + + use rstest::rstest; + use vortex::array::IntoArray; + use vortex::array::arrays::PrimitiveArray; + use vortex_error::VortexResult; + use vortex_error::vortex_ensure; + use vortex_ffi::vx_array_new_with; + use vortex_ffi::vx_session_free; + use vortex_ffi::vx_session_new_with; + + use super::*; + use crate::api::vx_velox_array_free; + + #[rstest] + #[case(PType::U8, VX_VELOX_PRIMITIVE_U8)] + #[case(PType::U16, VX_VELOX_PRIMITIVE_U16)] + #[case(PType::U32, VX_VELOX_PRIMITIVE_U32)] + #[case(PType::U64, VX_VELOX_PRIMITIVE_U64)] + #[case(PType::I8, VX_VELOX_PRIMITIVE_I8)] + #[case(PType::I16, VX_VELOX_PRIMITIVE_I16)] + #[case(PType::I32, VX_VELOX_PRIMITIVE_I32)] + #[case(PType::I64, VX_VELOX_PRIMITIVE_I64)] + #[case(PType::F16, VX_VELOX_PRIMITIVE_F16)] + #[case(PType::F32, VX_VELOX_PRIMITIVE_F32)] + #[case(PType::F64, VX_VELOX_PRIMITIVE_F64)] + fn maps_primitive_types(#[case] input: PType, #[case] expected: vx_velox_primitive_type) { + assert_eq!(primitive_type_id(input), expected); + } + + #[derive(Default)] + struct Capture { + primitive_type: Option, + length: usize, + values: *const u8, + values_length: usize, + values_alignment: usize, + validity: *const u8, + validity_length: usize, + validity_bit_offset: usize, + validity_alignment: usize, + retained_bytes: usize, + validity_kind: Option, + owner: Option, + } + + unsafe extern "C" fn capture_primitive( + context: *mut c_void, + view: *const vx_velox_primitive_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live `Capture` and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.primitive_type = Some(view.primitive_type); + capture.length = view.length; + capture.values = view.values; + capture.values_length = view.values_length; + capture.values_alignment = view.values_alignment; + capture.validity = view.validity; + capture.validity_length = view.validity_length; + capture.validity_bit_offset = view.validity_bit_offset; + capture.validity_alignment = view.validity_alignment; + capture.retained_bytes = view.buffers.retained_bytes; + capture.validity_kind = Some(view.validity_kind); + capture.owner = Some(view.buffers); + 0 + } + + #[test] + fn visits_sparse_nullable_values_with_retained_buffers() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with( + PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40)]).into_array(), + ); + let rows = [1_u64, 3]; + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: rows.as_ptr(), + row_count: rows.len(), + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + }; + let mut error = ptr::null_mut(); + // SAFETY: Every handle and callback object stays live for this call. + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); + assert_eq!(capture.length, 2); + assert_eq!(capture.values_length, 2 * size_of::()); + assert!(capture.values_alignment.is_power_of_two()); + assert_eq!(capture.values.addr() % capture.values_alignment, 0); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.validity_length, 1); + assert_eq!(capture.validity_bit_offset, 0); + assert!(capture.validity_alignment.is_power_of_two()); + assert_eq!(capture.validity.addr() % capture.validity_alignment, 0); + assert_eq!(capture.retained_bytes, capture.values_length + 1); + // SAFETY: The callback retained the owner before storing these pointers. + let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 2) }; + assert_eq!(values, [0, 40]); + // SAFETY: The retained validity pointer has one readable byte. + let validity = unsafe { *capture.validity }; + assert_eq!(validity & 0b11, 0b10); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_primitive`. + unsafe { release(owner.owner) }; + // SAFETY: Each owned handle is freed exactly once. + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn copies_sliced_values_and_reports_compact_allocation() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let source = PrimitiveArray::from_iter(0_i32..16); + let source_values = source.buffer_handle().try_to_host_sync()?; + // SAFETY: The source contains sixteen i32 values. The fifth value is in bounds. + let source_slice = unsafe { source_values.as_ptr().add(5 * size_of::()) }; + let array = vx_array_new_with(source.into_array().slice(5..8)?); + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: ptr::null(), + row_count: 0, + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + }; + let mut error = ptr::null_mut(); + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_eq!(capture.values_length, 3 * size_of::()); + assert_eq!(capture.retained_bytes, 2 * size_of::()); + assert_ne!(capture.values, source_slice); + // SAFETY: The retained compact values contain three i32 values. + let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 3) }; + assert_eq!(values, [5, 6, 7]); + assert!(capture.values_alignment.is_power_of_two()); + assert_eq!(capture.values.addr() % capture.values_alignment, 0); + assert_eq!(capture.validity_alignment, 0); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + unsafe { release(owner.owner) }; + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn copies_validity_into_compact_storage() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let session_ref = unsafe { vx_session_ref(session)? }; + let primitive = PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]); + let mut execution = session_ref.create_execution_ctx(); + let Mask::Values(mask) = primitive + .validity()? + .execute_mask(primitive.len(), &mut execution)? + else { + vortex_bail!("Expected bitmap validity"); + }; + let expected_validity = mask.bit_buffer().inner().as_ptr(); + let array = vx_array_new_with(primitive.into_array()); + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: ptr::null(), + row_count: 0, + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + }; + let mut error = ptr::null_mut(); + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_ne!(capture.validity, expected_validity); + assert_eq!(capture.validity_bit_offset, 0); + assert_eq!(capture.retained_bytes, 2 * size_of::() + 1); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + unsafe { release(owner.owner) }; + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) + } + + #[test] + fn rejects_unsorted_rows() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(); + let rows = [2_u64, 1]; + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: rows.as_ptr(), + row_count: rows.len(), + }; + match selected_array(&array, &request) { + Ok(_) => vortex_bail!("unsorted rows unexpectedly succeeded"), + Err(error) => assert!(error.to_string().contains("unique and increasing")), + } + Ok(()) + } +} diff --git a/vortex-velox/tests/abi_contract.rs b/vortex-velox/tests/abi_contract.rs new file mode 100644 index 00000000000..40101f17150 --- /dev/null +++ b/vortex-velox/tests/abi_contract.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(test)] +mod tests { + use std::fmt::Write as _; + use std::io::Write as _; + use std::mem::size_of; + use std::process::Command; + use std::process::Stdio; + + use vortex_velox::*; + + fn compile_stdin( + compiler: &str, + arguments: &[&str], + source: &str, + ) -> Result<(), Box> { + let mut child = Command::new(compiler) + .args(arguments) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + child + .stdin + .take() + .ok_or_else(|| std::io::Error::other("compiler stdin is unavailable"))? + .write_all(source.as_bytes())?; + let output = child.wait_with_output()?; + assert!( + output.status.success(), + "{compiler} rejected the adapter header:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + Ok(()) + } + + #[test] + fn c_header_matches_rust_layout() -> Result<(), Box> { + let mut source = + String::from("#include \n#include \n#include \"vortex_velox.h\"\n"); + source.push_str( + "_Static_assert(sizeof(vx_velox_ptype) == sizeof(uint32_t), \"ptype width\");\n", + ); + source.push_str( + "_Static_assert(sizeof(vx_velox_binary_operator) == sizeof(uint32_t), \"operator width\");\n", + ); + source.push_str( + "_Static_assert(sizeof(vx_velox_scan_selection_include) == sizeof(uint32_t), \"selection width\");\n", + ); + source.push_str( + "_Static_assert(sizeof(vx_velox_primitive_type) == sizeof(uint32_t), \"primitive width\");\n", + ); + source.push_str( + "_Static_assert(sizeof(vx_velox_validity_kind) == sizeof(uint32_t), \"validity width\");\n", + ); + source.push_str( + "_Static_assert(VX_VELOX_PTYPE_F64 == 10, \"ptype value\");\n\ + _Static_assert(VX_VELOX_OPERATOR_KLEENE_OR == 7, \"operator value\");\n\ + _Static_assert(VX_VELOX_SELECTION_EXCLUDE == 2, \"selection value\");\n\ + _Static_assert(VX_VELOX_PRIMITIVE_F64 == 10, \"primitive value\");\n\ + _Static_assert(VX_VELOX_VALIDITY_BITMAP == 3, \"validity value\");\n", + ); + + macro_rules! check_layout { + ($type:ty, [$($field:ident),+ $(,)?]) => {{ + writeln!( + source, + "_Static_assert(sizeof({0}) == {1}, \"{0} size\");", + stringify!($type), + size_of::<$type>() + )?; + $( + writeln!( + source, + "_Static_assert(offsetof({0}, {1}) == {2}, \"{0}.{1} offset\");", + stringify!($type), + stringify!($field), + std::mem::offset_of!($type, $field) + )?; + )+ + }}; + } + + check_layout!(vx_velox_scan_selection, [indices, length, include]); + check_layout!( + vx_velox_scan_options, + [ + struct_size, + abi_version, + projection, + filter, + row_range_begin, + row_range_end, + selection, + limit, + ordered, + ] + ); + check_layout!( + vx_velox_read_request, + [struct_size, offset, length, alignment] + ); + check_layout!(vx_velox_buffer, [struct_size, data, length, owner, release]); + check_layout!( + vx_velox_read_at_callbacks, + [ + struct_size, + abi_version, + context, + size, + read_ranges, + last_error, + release_context, + is_cancelled, + concurrency, + ] + ); + check_layout!(vx_velox_natural_split, [struct_size, row_begin, row_end]); + check_layout!( + vx_velox_buffer_owner, + [struct_size, owner, retain, release, retained_bytes] + ); + check_layout!( + vx_velox_primitive_view, + [ + struct_size, + primitive_type, + length, + values, + values_length, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + values_alignment, + validity_alignment, + ] + ); + check_layout!(vx_velox_visit_request, [struct_size, rows, row_count]); + check_layout!( + vx_velox_visitor, + [ + struct_size, + abi_version, + context, + visit_primitive, + last_error, + ] + ); + check_layout!( + vx_velox_arrow_memory_callbacks, + [ + struct_size, + abi_version, + context, + retain_context, + release_context, + report_allocation, + report_free, + last_error, + ] + ); + + let manifest = env!("CARGO_MANIFEST_DIR"); + let include = format!("-I{manifest}/cinclude"); + let base_include = format!("-I{manifest}/../vortex-ffi/cinclude"); + let compiler = std::env::var("CC").unwrap_or_else(|_| "cc".to_owned()); + compile_stdin( + &compiler, + &[ + "-std=c11", + "-fsyntax-only", + "-x", + "c", + &include, + &base_include, + "-", + ], + &source, + ) + } + + #[test] + fn header_compiles_with_host_arrow_declarations() -> Result<(), Box> { + let manifest = env!("CARGO_MANIFEST_DIR"); + let include = format!("-I{manifest}/cinclude"); + let base_include = format!("-I{manifest}/../vortex-ffi/cinclude"); + let compiler = std::env::var("CXX").unwrap_or_else(|_| "c++".to_owned()); + let source = + std::fs::read_to_string(format!("{manifest}/tests/velox_include_contract.cpp"))?; + compile_stdin( + &compiler, + &[ + "-std=c++20", + "-fsyntax-only", + "-x", + "c++", + &include, + &base_include, + "-", + ], + &source, + ) + } +} diff --git a/vortex-velox/tests/velox_include_contract.cpp b/vortex-velox/tests/velox_include_contract.cpp new file mode 100644 index 00000000000..332248ce665 --- /dev/null +++ b/vortex-velox/tests/velox_include_contract.cpp @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +struct ArrowSchema; +struct ArrowArray; +struct ArrowArrayStream; + +#define USE_OWN_ARROW +typedef struct ArrowSchema FFI_ArrowSchema; +typedef struct ArrowArray FFI_ArrowArray; +typedef struct ArrowArrayStream FFI_ArrowArrayStream; +#include "vortex_velox.h" +#undef USE_OWN_ARROW + +static_assert(VX_VELOX_ABI_VERSION == 1u); +static_assert(VX_VELOX_SELECTION_ALL == 0); +static_assert(VX_VELOX_OPERATOR_EQ == 0); + +void vx_velox_compile_velox_include_contract() { + vx_velox_read_at_callbacks callbacks {}; + callbacks.struct_size = sizeof(callbacks); + callbacks.abi_version = VX_VELOX_ABI_VERSION; + + vx_velox_scan_options options {}; + options.struct_size = sizeof(options); + options.abi_version = VX_VELOX_ABI_VERSION; + options.selection.include = VX_VELOX_SELECTION_ALL; + + const vx_dtype *(*new_primitive)(vx_velox_ptype, bool, vx_error **) = + vx_velox_dtype_new_primitive; + vx_expression *(*new_binary)(vx_velox_binary_operator, + const vx_expression *, + const vx_expression *, + vx_error **) = vx_velox_expression_binary; + + (void)callbacks; + (void)options; + (void)new_primitive; + (void)new_binary; + (void)vx_velox_source_export_schema; + (void)vx_velox_data_source_scan; +} From 79d3fc7c58864fdafccba3436fa76403ab2da727 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 07:28:41 -0400 Subject: [PATCH 3/6] feat(vortex-velox): Checkpoint native export paths --- Cargo.lock | 1 + encodings/fsst/src/compute/mod.rs | 1 + .../src/compute/uncompressed_size_in_bytes.rs | 94 + encodings/fsst/src/lib.rs | 9 + .../fns/list_contains/integer_membership.rs | 6 +- .../src/scalar_fn/fns/list_contains/mod.rs | 2 + .../schemes/string/scheme_selection_tests.rs | 43 +- vortex-buffer/src/buffer_mut.rs | 12 + vortex-compressor/src/builtins/dict/binary.rs | 11 +- vortex-compressor/src/builtins/dict/mod.rs | 25 + vortex-compressor/src/builtins/dict/string.rs | 11 +- vortex-compressor/src/compressor/constant.rs | 4 +- vortex-compressor/src/compressor/sample.rs | 1 + vortex-compressor/src/stats/varbinview.rs | 91 +- vortex-layout/src/layouts/dict/writer.rs | 84 +- vortex-velox/Cargo.toml | 1 + vortex-velox/cinclude/vortex_velox.h | 230 +- vortex-velox/src/api.rs | 68 + vortex-velox/src/array.rs | 43 +- vortex-velox/src/lib.rs | 98 +- vortex-velox/src/projection.rs | 80 +- vortex-velox/src/schema.rs | 4 + vortex-velox/src/temporal.rs | 150 + vortex-velox/src/visitor.rs | 3604 ++++++++++++++++- vortex-velox/tests/abi_contract.rs | 125 +- vortex-velox/tests/velox_include_contract.cpp | 2 +- 26 files changed, 4500 insertions(+), 300 deletions(-) create mode 100644 encodings/fsst/src/compute/uncompressed_size_in_bytes.rs create mode 100644 vortex-velox/src/temporal.rs diff --git a/Cargo.lock b/Cargo.lock index 4660765b51f..ebc429a4746 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11706,6 +11706,7 @@ dependencies = [ "vortex-arrow", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-ffi", "vortex-io", ] diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index e23d4272866..8343fc64e00 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -6,6 +6,7 @@ mod cast; mod compare; mod filter; mod like; +pub(crate) mod uncompressed_size_in_bytes; use vortex_array::ArrayRef; use vortex_array::ArrayView; diff --git a/encodings/fsst/src/compute/uncompressed_size_in_bytes.rs b/encodings/fsst/src/compute/uncompressed_size_in_bytes.rs new file mode 100644 index 00000000000..b3235a8bc08 --- /dev/null +++ b/encodings/fsst/src/compute/uncompressed_size_in_bytes.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::size_of; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::FSST; +use crate::canonical::FsstDecodePlan; + +#[derive(Debug)] +pub(crate) struct FsstUncompressedSizeKernel; + +impl DynAggregateKernel for FsstUncompressedSizeKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() { + return Ok(None); + } + let Some(fsst) = batch.as_opt::() else { + return Ok(None); + }; + Ok(Some(Scalar::from(uncompressed_size(fsst, ctx)?))) + } +} + +fn uncompressed_size(fsst: ArrayView<'_, FSST>, ctx: &mut ExecutionCtx) -> VortexResult { + let plan = FsstDecodePlan::new(fsst, ctx)?; + let views_size = fsst + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("FSST view size overflow"))?; + let validity_size = match fsst.validity()? { + Validity::NonNullable | Validity::AllValid | Validity::AllInvalid => 0, + Validity::Array(validity) => validity.len().div_ceil(u8::BITS as usize), + }; + views_size + .checked_add(plan.total_size) + .and_then(|size| size.checked_add(validity_size)) + .and_then(|size| u64::try_from(size).ok()) + .ok_or_else(|| vortex_err!("FSST uncompressed size overflow")) +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; + use vortex_array::array_session; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; + + use crate::fsst_compress; + use crate::fsst_train_compressor; + + #[test] + fn matches_canonical_size_for_nullable_strings() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let input = VarBinArray::from_iter( + [ + Some("short"), + None, + Some("a string that uses an outlined view"), + Some("another outlined string"), + ], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&input, &mut ctx)?; + let encoded = fsst_compress(&input, &compressor, &mut ctx)?.into_array(); + let expected = uncompressed_size_in_bytes(&input, &mut ctx)?; + let actual = uncompressed_size_in_bytes(&encoded, &mut ctx)?; + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/encodings/fsst/src/lib.rs b/encodings/fsst/src/lib.rs index 70dcc705249..cc6a62db13a 100644 --- a/encodings/fsst/src/lib.rs +++ b/encodings/fsst/src/lib.rs @@ -27,6 +27,10 @@ mod tests; pub use array::*; pub use compress::*; +use vortex_array::ArrayVTable; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; @@ -34,4 +38,9 @@ use vortex_session::VortexSession; pub fn initialize(session: &VortexSession) { session.arrays().register(FSST); kernel::initialize(session); + session.aggregate_fns().register_aggregate_kernel( + FSST.id(), + Some(UncompressedSizeInBytes.id()), + &compute::uncompressed_size_in_bytes::FsstUncompressedSizeKernel, + ); } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs index 83ea016b6c5..7f27e417356 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs @@ -8,7 +8,9 @@ use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::IntegerPType; -const MAX_SOURCE_MEMBERS: usize = 4; +/// Maximum source members for the direct integer membership kernel. +#[doc(hidden)] +pub const MAX_DIRECT_INTEGER_MEMBERS: usize = 4; /// A prepared integer set for constant-list membership kernels. /// @@ -44,7 +46,7 @@ impl IntegerMembership { let Some(elements) = list_array.scalar().as_list().values() else { return Ok(None); }; - if elements.len() > MAX_SOURCE_MEMBERS { + if elements.len() > MAX_DIRECT_INTEGER_MEMBERS { return Ok(None); } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 6e984c17b79..85339ad4ab9 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -8,6 +8,8 @@ use std::ops::BitOr; use arrow_buffer::bit_iterator::BitIndexIterator; pub use integer_membership::IntegerMembership; +#[doc(hidden)] +pub use integer_membership::MAX_DIRECT_INTEGER_MEMBERS; pub use kernel::*; use num_traits::Zero; use vortex_buffer::BitBuffer; diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index ec51c5104bb..5a168b6675f 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -42,7 +42,48 @@ fn test_dict_compressed() -> VortexResult<()> { let array_ref = array.into_array(); let compressed = BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; - assert!(compressed.is::()); + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + Ok(()) +} + +#[test] +fn test_dict_compressed_with_more_values_than_sample() -> VortexResult<()> { + let distinct_values = (0..4096) + .map(|value| format!("repeated string value {value:04}")) + .collect::>(); + let strings = (0..65_536) + .map(|index| Some(distinct_values[index % distinct_values.len()].as_str())) + .collect::>(); + let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); + let compressed = BtrBlocksCompressor::default() + .compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + Ok(()) +} + +#[test] +fn test_unique_strings_not_dict_compressed() -> VortexResult<()> { + let strings = (0..4096) + .map(|value| Some(format!("unique string value {value:04}"))) + .collect::>(); + let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); + let compressed = BtrBlocksCompressor::default() + .compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!( + !compressed.is::(), + "expected a non-Dict encoding, got {}", + compressed.encoding_id() + ); Ok(()) } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 27c3dad3532..bb5c7971a4f 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -408,6 +408,11 @@ impl BufferMut { self.capacity } + /// Returns the complete size of the retained backing allocation in bytes. + pub fn allocation_size(&self) -> usize { + self.allocation.size() + } + /// Returns a raw pointer to the buffer's data. #[allow(clippy::inline_always)] #[inline(always)] @@ -1002,6 +1007,13 @@ mod test { assert_eq!(buf.alignment(), Alignment::new(1024)); } + #[test] + fn allocation_size_includes_alignment_storage() { + let buffer = BufferMut::::with_capacity_aligned(1, Alignment::new(256)); + + assert!(buffer.allocation_size() > buffer.capacity()); + } + #[test] fn growth_preserves_alignment_and_values() { let alignment = Alignment::new(4096); diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..e5bfb2fbbc4 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -24,10 +24,10 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::builtins::dict::varbinview_dict_compression_ratio; 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; @@ -87,17 +87,18 @@ impl Scheme for BinaryDictScheme { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( + let distinct_values_count = stats.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 { + if distinct_values_count > stats.value_count() / 2 { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + CompressionEstimate::Verdict(EstimateVerdict::Ratio(varbinview_dict_compression_ratio( + &stats, + ))) } fn compress( diff --git a/vortex-compressor/src/builtins/dict/mod.rs b/vortex-compressor/src/builtins/dict/mod.rs index 4862df2b211..06acae1c8dc 100644 --- a/vortex-compressor/src/builtins/dict/mod.rs +++ b/vortex-compressor/src/builtins/dict/mod.rs @@ -14,3 +14,28 @@ pub use float::dictionary_encode as float_dictionary_encode; pub use integer::IntDictScheme; pub use integer::dictionary_encode as integer_dictionary_encode; pub use string::StringDictScheme; + +use vortex_array::arrays::varbinview::BinaryView; +use vortex_error::VortexExpect; + +use crate::stats::StringStats; + +/// Estimates the ratio for canonical variable-width values encoded as a dictionary. +fn varbinview_dict_compression_ratio(stats: &StringStats) -> f64 { + let distinct_count = stats + .distinct_count() + .vortex_expect("distinct value count must be available"); + let distinct_value_bytes = stats + .distinct_value_bytes() + .vortex_expect("distinct value bytes must be available"); + let row_count = u64::from(stats.value_count()) + u64::from(stats.null_count()); + + let view_size = size_of::() as u64; + let canonical_bytes = row_count * view_size + stats.value_bytes(); + let dictionary_values_bytes = u64::from(distinct_count) * view_size + distinct_value_bytes; + let code_bits = u64::from(u32::BITS - distinct_count.leading_zeros()); + let code_bytes = (row_count * code_bits).div_ceil(8); + + // Child compression can only improve this conservative estimate. + canonical_bytes as f64 / (dictionary_values_bytes + code_bytes) as f64 +} diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..854bab5bf3b 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -24,10 +24,10 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::builtins::dict::varbinview_dict_compression_ratio; 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; @@ -87,17 +87,18 @@ impl Scheme for StringDictScheme { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( + let distinct_values_count = stats.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 { + if distinct_values_count > stats.value_count() / 2 { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + CompressionEstimate::Verdict(EstimateVerdict::Ratio(varbinview_dict_compression_ratio( + &stats, + ))) } fn compress( diff --git a/vortex-compressor/src/compressor/constant.rs b/vortex-compressor/src/compressor/constant.rs index 7e8edc73f2f..ddce5ed23d3 100644 --- a/vortex-compressor/src/compressor/constant.rs +++ b/vortex-compressor/src/compressor/constant.rs @@ -76,9 +76,7 @@ pub(crate) fn is_constant_for_compression( if dtype.is_utf8() || dtype.is_binary() { let stats = data.varbinview_stats(exec_ctx); - // The estimated distinct count is a lower bound on the actual distinct count, so a value - // above 1 proves the array is not constant without scanning it. - if stats.estimated_distinct_count().is_some_and(|c| c > 1) { + if stats.distinct_count().is_some_and(|count| count > 1) { return Ok(false); } diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..b7eb1a93d38 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -162,6 +162,7 @@ pub(super) fn estimate_compression_ratio_with_sampling( let sample_count = sample_count_approx_one_percent(array.len()); // `ArrayAndStats` expects a canonical array (so that it can easily compute lazy stats). let canonical: Canonical = sample(array, SAMPLE_SIZE, sample_count).execute(exec_ctx)?; + let canonical = canonical.compact(exec_ctx)?; canonical.into_array() }; diff --git a/vortex-compressor/src/stats/varbinview.rs b/vortex-compressor/src/stats/varbinview.rs index 486d431e684..2ee84ea4b13 100644 --- a/vortex-compressor/src/stats/varbinview.rs +++ b/vortex-compressor/src/stats/varbinview.rs @@ -5,6 +5,8 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -15,32 +17,57 @@ use super::GenerateStatsOptions; /// Array of variable-length byte/string values, and relevant stats for compression. #[derive(Clone, Debug)] pub struct StringStats { - /// The estimated number of distinct values, or `None` if not computed. + /// The number of distinct values, or `None` if not computed. /// This _must_ be non-zero. - estimated_distinct_count: Option, + distinct_count: Option, + /// The visible bytes across the distinct values. + distinct_value_bytes: Option, + /// The visible bytes across all values. + value_bytes: u64, /// The number of non-null values. value_count: u32, /// The number of null values. null_count: u32, } -/// Estimate the number of distinct values in the var bin view array. -fn estimate_distinct_count(varbinview: &VarBinViewArray) -> VortexResult { +/// Returns the bytes referenced by a variable-width view. +fn view_bytes<'a>(buffers: &[&'a ByteBuffer], view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let reference = view.as_view(); + &buffers[reference.buffer_index as usize][reference.as_range()] + } +} + +/// Counts distinct values and their visible bytes. +fn count_distinct_values( + varbinview: &VarBinViewArray, + ctx: &mut ExecutionCtx, +) -> VortexResult<(u32, u64)> { let views = varbinview.views(); - // Iterate the views. Two values which are equal must have the same first 8-bytes. - // NOTE: there are cases where this performs pessimally, e.g. when we have strings that all - // share a 4-byte prefix and have the same length. + let buffers = varbinview + .data_buffers() + .iter() + .map(|buffer| buffer.as_host()) + .collect::>(); + let validity = varbinview + .as_ref() + .validity()? + .execute_mask(varbinview.len(), ctx)?; let mut distinct = HashSet::with_capacity(views.len() / 2); - views.iter().for_each(|&view| { - #[expect( - clippy::cast_possible_truncation, - reason = "approximate uniqueness with view prefix" - )] - let len_and_prefix = view.as_u128() as u64; - distinct.insert(len_and_prefix); - }); - - Ok(u32::try_from(distinct.len())?) + let mut distinct_value_bytes = 0u64; + + for (index, view) in views.iter().enumerate() { + if validity.value(index) { + let bytes = view_bytes(&buffers, view); + if distinct.insert(bytes) { + distinct_value_bytes += bytes.len() as u64; + } + } + } + + Ok((u32::try_from(distinct.len())?, distinct_value_bytes)) } impl StringStats { @@ -55,15 +82,21 @@ impl StringStats { .compute_null_count(ctx) .ok_or_else(|| vortex_err!("Failed to compute null_count"))?; let value_count = input.len() - null_count; - let estimated_distinct_count = opts + let distinct_values = opts .count_distinct_values - .then(|| estimate_distinct_count(input)) + .then(|| count_distinct_values(input, ctx)) .transpose()?; + let (distinct_count, distinct_value_bytes) = distinct_values + .map(|(count, bytes)| (Some(count), Some(bytes))) + .unwrap_or((None, None)); + let value_bytes = input.views().iter().map(|view| u64::from(view.len())).sum(); Ok(Self { value_count: u32::try_from(value_count)?, null_count: u32::try_from(null_count)?, - estimated_distinct_count, + distinct_count, + distinct_value_bytes, + value_bytes, }) } } @@ -84,11 +117,19 @@ impl StringStats { .vortex_expect("StringStats::generate_opts should not fail") } - /// Returns the estimated number of distinct values, or `None` if not computed. - /// - /// This estimation is always going to be less than or equal to the actual distinct count. - pub fn estimated_distinct_count(&self) -> Option { - self.estimated_distinct_count + /// Returns the number of distinct values, or `None` if not computed. + pub fn distinct_count(&self) -> Option { + self.distinct_count + } + + /// Returns the visible bytes across the distinct values. + pub fn distinct_value_bytes(&self) -> Option { + self.distinct_value_bytes + } + + /// Returns the visible bytes across all values. + pub fn value_bytes(&self) -> u64 { + self.value_bytes } /// Returns the number of non-null values. diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index bd68864725e..21bc19eaf59 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -17,12 +17,14 @@ use futures::channel::oneshot; use futures::future::BoxFuture; use futures::pin_mut; use futures::stream::BoxStream; +use futures::stream::iter; use futures::stream::once; use futures::try_join; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Dict; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; @@ -53,6 +55,8 @@ use crate::sequence::SequentialStream; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; +const DICT_LAYOUT_PROBE_ROWS: usize = 65_536; + /// Constraints for dictionary layout encoding. /// /// Note that [`max_len`](Self::max_len) is limited to `u16` (65,535 entries) by design. Since @@ -151,10 +155,11 @@ impl LayoutStrategy for DictStrategy { let dtype = stream.dtype().clone(); // 0. decide if chunks are eligible for dict encoding - let (stream, first_chunk) = peek_first_chunk(stream).await?; + let (stream, probe_chunk) = + peek_probe_chunk(stream, &dtype, DICT_LAYOUT_PROBE_ROWS).await?; let stream = SequentialStreamAdapter::new(dtype.clone(), stream).sendable(); - let should_fallback = match first_chunk { + let should_fallback = match probe_chunk { None => true, // empty stream Some(chunk) => { let mut exec_ctx = session.create_execution_ctx(); @@ -512,19 +517,43 @@ impl Stream for DictionaryTransformer { } } -async fn peek_first_chunk( +async fn peek_probe_chunk( mut stream: BoxStream<'static, SequencedChunk>, + dtype: &DType, + probe_rows: usize, ) -> VortexResult<(BoxStream<'static, SequencedChunk>, Option)> { - match stream.next().await { - None => Ok((stream.boxed(), None)), - Some(Err(e)) => Err(e), - Some(Ok((sequence_id, chunk))) => { - let chunk_clone = chunk.clone(); - let reconstructed_stream = - once(async move { Ok((sequence_id, chunk_clone)) }).chain(stream); - Ok((reconstructed_stream.boxed(), Some(chunk))) + let mut buffered = Vec::new(); + let mut buffered_rows = 0usize; + + while buffered_rows < probe_rows { + match stream.next().await { + None => break, + Some(Err(error)) => return Err(error), + Some(Ok((sequence_id, chunk))) => { + buffered_rows += chunk.len(); + buffered.push((sequence_id, chunk)); + } } } + + if buffered.is_empty() { + return Ok((stream.boxed(), None)); + } + + let probe_chunks = buffered + .iter() + .map(|(_, chunk)| chunk.clone()) + .collect::>(); + let probe_chunk = if probe_chunks.len() == 1 { + probe_chunks + .into_iter() + .next() + .vortex_expect("one probe chunk") + } else { + ChunkedArray::try_new(probe_chunks, dtype.clone())?.into_array() + }; + let reconstructed_stream = iter(buffered.into_iter().map(Ok)).chain(stream); + Ok((reconstructed_stream.boxed(), Some(probe_chunk))) } pub fn dict_layout_supported(dtype: &DType) -> bool { @@ -597,10 +626,14 @@ mod tests { use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; use vortex_array::session::ArraySession; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::DICT_LAYOUT_PROBE_ROWS; use super::DictionaryTransformer; use super::dict_encode_stream; + use super::peek_probe_chunk; use crate::sequence::SequenceId; use crate::sequence::SequentialStream; use crate::sequence::SequentialStreamAdapter; @@ -609,6 +642,35 @@ mod tests { static SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); + #[tokio::test] + async fn test_probe_uses_bounded_prefix_and_replays_stream() -> VortexResult<()> { + let dtype = DType::Utf8(NonNullable); + let chunks = (0..10) + .map(|chunk_index| { + let value = format!("value_{chunk_index}"); + VarBinArray::from(vec![value.as_str(); 8192]).into_array() + }) + .collect::>(); + let mut pointer = SequenceId::root(); + let stream = futures::stream::iter( + chunks + .into_iter() + .map(move |chunk| Ok((pointer.advance(), chunk))), + ) + .boxed(); + + let (mut replayed, probe) = + peek_probe_chunk(stream, &dtype, DICT_LAYOUT_PROBE_ROWS).await?; + assert_eq!(probe.vortex_expect("probe chunk").len(), 65_536); + + let mut replayed_rows = 0; + while let Some(item) = replayed.next().await { + replayed_rows += item?.1.len(); + } + assert_eq!(replayed_rows, 81_920); + Ok(()) + } + /// Regression test for a bug where the codes stream dtype was hardcoded to U16 instead of /// using the actual codes dtype from the array. When `max_len <= 255`, the dict encoder /// produces U8 codes, but the stream was incorrectly typed as U16, causing a dtype mismatch diff --git a/vortex-velox/Cargo.toml b/vortex-velox/Cargo.toml index 09ee82cdeed..b0f3bcf3315 100644 --- a/vortex-velox/Cargo.toml +++ b/vortex-velox/Cargo.toml @@ -32,6 +32,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-arrow = { workspace = true } vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } vortex-ffi = { path = "../vortex-ffi" } vortex-io = { workspace = true } vortex = { workspace = true } diff --git a/vortex-velox/cinclude/vortex_velox.h b/vortex-velox/cinclude/vortex_velox.h index eeae773bbca..98a0948163f 100644 --- a/vortex-velox/cinclude/vortex_velox.h +++ b/vortex-velox/cinclude/vortex_velox.h @@ -17,7 +17,7 @@ extern "C" { * contain general vx_* symbols from linked Vortex FFI objects. */ -#define VX_VELOX_ABI_VERSION 1u +#define VX_VELOX_ABI_VERSION 5u #define VX_VELOX_CAPABILITY_BATCH_READ (UINT64_C(1) << 0) #define VX_VELOX_CAPABILITY_CALLBACK_SOURCE (UINT64_C(1) << 1) #define VX_VELOX_CAPABILITY_NATURAL_SPLITS (UINT64_C(1) << 2) @@ -28,9 +28,21 @@ extern "C" { #define VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING (UINT64_C(1) << 7) /* Vortex checks cancellation before each host read callback. */ #define VX_VELOX_CAPABILITY_READ_CANCELLATION (UINT64_C(1) << 8) +#define VX_VELOX_CAPABILITY_EXPORT_CURSOR (UINT64_C(1) << 9) +#define VX_VELOX_CAPABILITY_PLAIN_PROJECTION (UINT64_C(1) << 10) +#define VX_VELOX_CAPABILITY_VARBIN_VISITOR (UINT64_C(1) << 11) +#define VX_VELOX_CAPABILITY_DICTIONARY_VISITOR (UINT64_C(1) << 12) +#define VX_VELOX_CAPABILITY_CONSTANT_VISITOR (UINT64_C(1) << 13) +#define VX_VELOX_CAPABILITY_BOOL_VISITOR (UINT64_C(1) << 14) +#define VX_VELOX_CAPABILITY_DATE_VISITOR (UINT64_C(1) << 15) +#define VX_VELOX_CAPABILITY_DECIMAL_VISITOR (UINT64_C(1) << 16) +#define VX_VELOX_CAPABILITY_STRUCT_VISITOR (UINT64_C(1) << 17) +#define VX_VELOX_CAPABILITY_LIST_VISITOR (UINT64_C(1) << 18) +#define VX_VELOX_CAPABILITY_MAP_VISITOR (UINT64_C(1) << 19) typedef struct vx_velox_read_at vx_velox_read_at; typedef struct vx_velox_source vx_velox_source; +typedef struct vx_velox_export_cursor vx_velox_export_cursor; typedef uint32_t vx_velox_ptype; #define VX_VELOX_PTYPE_U8 UINT32_C(0) @@ -143,6 +155,7 @@ typedef uint32_t vx_velox_primitive_type; #define VX_VELOX_PRIMITIVE_F16 UINT32_C(8) #define VX_VELOX_PRIMITIVE_F32 UINT32_C(9) #define VX_VELOX_PRIMITIVE_F64 UINT32_C(10) +#define VX_VELOX_PRIMITIVE_I128 UINT32_C(11) typedef uint32_t vx_velox_validity_kind; #define VX_VELOX_VALIDITY_NON_NULLABLE UINT32_C(0) @@ -163,7 +176,8 @@ typedef struct vx_velox_buffer_owner { * * Vortex copies values into an allocation with uint64_t alignment. The values * allocation rounds values_length up to that alignment. Vortex copies a bitmap - * into a compact byte allocation with validity_bit_offset set to zero. + * into a uint64_t-aligned, word-padded allocation. A window rebases the pointer + * and reports its remaining bit offset. * buffers.retained_bytes is the exact sum of these allocation sizes. * * The pointers remain valid through visit_primitive. The host must call retain @@ -172,6 +186,8 @@ typedef struct vx_velox_buffer_owner { typedef struct vx_velox_primitive_view { size_t struct_size; vx_velox_primitive_type primitive_type; + uint32_t decimal_precision; + int32_t decimal_scale; size_t length; const uint8_t *values; size_t values_length; @@ -184,6 +200,173 @@ typedef struct vx_velox_primitive_view { size_t validity_alignment; } vx_velox_primitive_view; +typedef uint32_t vx_velox_varbin_kind; +#define VX_VELOX_VARBIN_UTF8 UINT32_C(0) +#define VX_VELOX_VARBIN_BINARY UINT32_C(1) + +typedef struct vx_velox_byte_buffer_view { + const uint8_t *data; + size_t length; +} vx_velox_byte_buffer_view; + +/** + * A stable 16-byte variable-width binary view. + * + * Values of 12 bytes or fewer occupy data directly. Longer values store four + * prefix bytes, a uint32_t buffer index, and a uint32_t byte offset in data. + * The host must reject lengths above INT32_MAX. Each outlined range fits its + * payload buffer. + */ +typedef struct vx_velox_binary_view { + uint32_t length; + uint8_t data[12]; +} vx_velox_binary_view; + +/** + * A canonical UTF-8 or binary payload and its owner. + * + * The view buffer and each payload buffer stay valid through visit_varbin. + * The host must retain buffers before it stores pointers after the callback. + * buffers.retained_bytes includes all retained allocation capacities. + */ +typedef struct vx_velox_varbin_view { + size_t struct_size; + vx_velox_varbin_kind kind; + size_t length; + const vx_velox_binary_view *views; + size_t views_length; + const vx_velox_byte_buffer_view *data_buffers; + size_t data_buffer_count; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t views_alignment; + size_t validity_alignment; +} vx_velox_varbin_view; + +/** + * A canonical packed Boolean payload and its owner. + * + * The value and validity buffers use least-significant-bit-first order. + * The host must retain buffers before it stores pointers after the callback. + */ +typedef struct vx_velox_bool_view { + size_t struct_size; + size_t length; + const uint8_t *values; + size_t values_length; + size_t values_bit_offset; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t values_alignment; + size_t validity_alignment; +} vx_velox_bool_view; + +/** + * A dictionary payload for one output window. + * + * codes owns the integer code buffers. values remains valid only during the + * callback. The host can visit the prepared cursor during that callback. + */ +typedef struct vx_velox_dictionary_view { + size_t struct_size; + size_t length; + vx_velox_primitive_view codes; + const vx_velox_export_cursor *values; + size_t values_length; +} vx_velox_dictionary_view; + +/** + * A constant payload for one output window. + * + * value contains one canonical value. It remains valid only during the + * callback. The host can visit the prepared cursor during that callback. + */ +typedef struct vx_velox_constant_view { + size_t struct_size; + size_t length; + const vx_velox_export_cursor *value; +} vx_velox_constant_view; + +/** + * A canonical struct payload for one output window. + * + * fields contains borrowed prepared cursors in declaration order. The host + * visits each field cursor at offset for length rows during this callback. + * buffers owns only the parent validity buffer. + */ +typedef struct vx_velox_struct_view { + size_t struct_size; + size_t length; + size_t offset; + const vx_velox_export_cursor *const *fields; + size_t field_count; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t validity_alignment; +} vx_velox_struct_view; + +/** + * A canonical list window. + * + * offsets and sizes start at the requested parent window. Offset values remain + * absolute against the complete elements cursor. buffers retains the complete + * prepared metadata allocation. The host must retain buffers before it stores + * a metadata pointer. elements is borrowed during the callback. Vectors + * imported from elements retain their own owners. + */ +typedef struct vx_velox_list_view { + size_t struct_size; + size_t length; + const int32_t *offsets; + const int32_t *sizes; + const vx_velox_export_cursor *elements; + size_t elements_length; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t offsets_alignment; + size_t sizes_alignment; + size_t validity_alignment; +} vx_velox_list_view; + +/** + * A canonical map window. + * + * offsets and sizes start at the requested parent window. Offset values remain + * absolute against the complete key and value cursors. buffers retains the + * complete prepared metadata allocation. The child cursors are borrowed during + * the callback. Vectors imported from them retain their own owners. + */ +typedef struct vx_velox_map_view { + size_t struct_size; + size_t length; + const int32_t *offsets; + const int32_t *sizes; + const vx_velox_export_cursor *keys; + const vx_velox_export_cursor *values; + size_t entries_length; + bool keys_sorted; + vx_velox_validity_kind validity_kind; + const uint8_t *validity; + size_t validity_length; + size_t validity_bit_offset; + vx_velox_buffer_owner buffers; + size_t offsets_alignment; + size_t sizes_alignment; + size_t validity_alignment; +} vx_velox_map_view; + typedef struct vx_velox_visit_request { size_t struct_size; const uint64_t *rows; @@ -193,7 +376,7 @@ typedef struct vx_velox_visit_request { /** * Host callbacks for one Vortex array visit. * - * One array visit calls visit_primitive synchronously. If the host shares this + * One array visit calls the matching callback synchronously. If the host shares this * table between simultaneous visits, callbacks can occur concurrently. * last_error returns the calling thread's most recent visitor error. Its string * remains valid until the next callback on that thread. @@ -207,6 +390,13 @@ typedef struct vx_velox_visitor { void *context; int32_t (*visit_primitive)(void *context, const vx_velox_primitive_view *view); const char *(*last_error)(void *context); + int32_t (*visit_varbin)(void *context, const vx_velox_varbin_view *view); + int32_t (*visit_dictionary)(void *context, const vx_velox_dictionary_view *view); + int32_t (*visit_constant)(void *context, const vx_velox_constant_view *view); + int32_t (*visit_bool)(void *context, const vx_velox_bool_view *view); + int32_t (*visit_struct)(void *context, const vx_velox_struct_view *view); + int32_t (*visit_list)(void *context, const vx_velox_list_view *view); + int32_t (*visit_map)(void *context, const vx_velox_map_view *view); } vx_velox_visitor; /** @@ -244,6 +434,7 @@ uint64_t vx_velox_capabilities(void); vx_view vx_velox_error_message(const vx_error *error); void vx_velox_error_free(const vx_error *error); vx_session *vx_velox_session_new(void); +vx_session *vx_velox_session_clone(const vx_session *session); void vx_velox_session_free(const vx_session *session); const vx_dtype *vx_velox_dtype_new_primitive(vx_velox_ptype ptype, @@ -254,6 +445,9 @@ vx_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); vx_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); vx_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); vx_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); +vx_scalar *vx_velox_scalar_new_date_days(int32_t value, + bool nullable, + vx_error **error_out); vx_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); vx_scalar *vx_velox_scalar_new_f32(float value, bool nullable); vx_scalar *vx_velox_scalar_new_f64(double value, bool nullable); @@ -279,7 +473,11 @@ vx_expression *vx_velox_expression_or(const vx_expression *const *expressions, s vx_expression *vx_velox_expression_not(const vx_expression *child); vx_expression *vx_velox_expression_is_null(const vx_expression *child); vx_expression *vx_velox_expression_list_contains(const vx_expression *list, const vx_expression *value); +bool vx_velox_can_push_down_integer_values(size_t value_count); void vx_velox_expression_free(const vx_expression *expression); +vx_expression *vx_velox_expression_select(const vx_view *names, + size_t length, + vx_error **error_out); vx_expression *vx_velox_expression_select_with_row_index(const vx_view *names, size_t length, vx_view row_index_name, @@ -335,6 +533,32 @@ int32_t vx_velox_array_visit(const vx_session *session, const vx_velox_visit_request *request, const vx_velox_visitor *visitor, vx_error **error_out); + +/** + * Create one prepared exporter for several engine-sized output windows. + * + * memory_callbacks must identify a complete, thread-safe callback table. + * The exporter retains its callback context until the last buffer owner releases it. + */ +vx_velox_export_cursor *vx_velox_export_cursor_new(const vx_session *session, + const vx_array *array, + const vx_velox_arrow_memory_callbacks *memory_callbacks, + vx_error **error_out); + +/** Free one prepared exporter. */ +void vx_velox_export_cursor_free(vx_velox_export_cursor *cursor); + +/** + * Visit one contiguous output window from a prepared exporter. + * + * Concurrent visit calls are valid. Do not free the cursor before all visits return. + */ +int32_t vx_velox_export_cursor_visit(const vx_velox_export_cursor *cursor, + size_t offset, + size_t length, + const vx_velox_visitor *visitor, + vx_error **error_out); + int32_t vx_velox_array_export_arrow(const vx_session *session, const vx_array *array, const vx_velox_arrow_memory_callbacks *memory_callbacks, diff --git a/vortex-velox/src/api.rs b/vortex-velox/src/api.rs index 485b347dd03..944809c53cc 100644 --- a/vortex-velox/src/api.rs +++ b/vortex-velox/src/api.rs @@ -10,8 +10,11 @@ use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::expr::root; +use vortex::extension::datetime::Date; +use vortex::extension::datetime::TimeUnit; use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::fns::binary::Binary; +use vortex::scalar_fn::fns::list_contains::MAX_DIRECT_INTEGER_MEMBERS; use vortex::scalar_fn::fns::operators::Operator; use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; @@ -42,6 +45,7 @@ mod ffi { pub fn vx_error_message(error: *const vx_error) -> vx_view; pub fn vx_error_free(error: *const vx_error); pub fn vx_session_new() -> *mut vx_session; + pub fn vx_session_clone(session: *const vx_session) -> *mut vx_session; pub fn vx_session_free(session: *const vx_session); pub fn vx_dtype_free(dtype: *const vx_dtype); pub fn vx_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_scalar; @@ -69,6 +73,11 @@ mod ffi { nullable: bool, error_out: *mut *mut vx_error, ) -> *mut vx_scalar; + pub fn vx_scalar_new_extension( + dtype: *const vx_dtype, + storage: *const vx_scalar, + error_out: *mut *mut vx_error, + ) -> *mut vx_scalar; pub fn vx_scalar_free(scalar: *const vx_scalar); pub fn vx_expression_literal( scalar: *const vx_scalar, @@ -244,6 +253,18 @@ pub extern "C-unwind" fn vx_velox_session_new() -> *mut vx_session { unsafe { ffi::vx_session_new() } } +/// Clone a Vortex session. +/// +/// # Safety +/// +/// `session` must point to a live Vortex session. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_session_clone( + session: *const vx_session, +) -> *mut vx_session { + unsafe { ffi::vx_session_clone(session) } +} + /// Free a Vortex session. /// /// # Safety @@ -334,6 +355,29 @@ scalar_primitive_wrapper!( i32, "Create an i32 scalar." ); + +/// Create a date scalar that stores days since the Unix epoch. +/// +/// # Safety +/// +/// `error_out` must be null or valid for one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_scalar_new_date_days( + value: i32, + nullable: bool, + error_out: *mut *mut vx_error, +) -> *mut vx_scalar { + let dtype = vx_dtype_new_with(DType::Extension( + Date::new(TimeUnit::Days, Nullability::from(nullable)).erased(), + )); + let storage = unsafe { ffi::vx_scalar_new_i32(value, nullable) }; + let scalar = unsafe { ffi::vx_scalar_new_extension(dtype, storage, error_out) }; + unsafe { + ffi::vx_scalar_free(storage); + ffi::vx_dtype_free(dtype); + } + scalar +} scalar_primitive_wrapper!( vx_velox_scalar_new_i64, vx_scalar_new_i64, @@ -539,6 +583,12 @@ pub unsafe extern "C" fn vx_velox_expression_list_contains( unsafe { ffi::vx_expression_list_contains(list, value) } } +/// Return whether Vortex can push down an integer value set without generic expansion. +#[unsafe(no_mangle)] +pub extern "C" fn vx_velox_can_push_down_integer_values(value_count: usize) -> bool { + value_count <= MAX_DIRECT_INTEGER_MEMBERS +} + /// Free an expression. /// /// # Safety @@ -735,6 +785,18 @@ mod tests { use super::*; + #[test] + fn clones_adapter_session() { + let session = vx_velox_session_new(); + assert!(!session.is_null()); + let cloned = unsafe { vx_velox_session_clone(session) }; + assert!(!cloned.is_null()); + unsafe { + vx_velox_session_free(cloned); + vx_velox_session_free(session); + } + } + #[test] fn translates_scan_options() -> VortexResult<()> { let options = vx_velox_scan_options { @@ -799,6 +861,12 @@ mod tests { } } + #[test] + fn reports_direct_integer_membership_limit() { + assert!(vx_velox_can_push_down_integer_values(4)); + assert!(!vx_velox_can_push_down_integer_values(5)); + } + #[test] fn rejects_wrong_scan_abi_before_source_access() { let options = vx_velox_scan_options { diff --git a/vortex-velox/src/array.rs b/vortex-velox/src/array.rs index 4f8bf4f3dc0..541ebe58452 100644 --- a/vortex-velox/src/array.rs +++ b/vortex-velox/src/array.rs @@ -32,6 +32,8 @@ use vortex_ffi::vx_error; use vortex_ffi::vx_session; use vortex_ffi::vx_session_ref; +use crate::temporal::validate_velox_arrow_data; + /// Host memory callbacks for one Arrow C Data export. #[repr(C)] #[derive(Clone, Copy)] @@ -75,12 +77,17 @@ struct ArrowMemoryOwner { retained_bytes: usize, } -struct ArrowMemoryReservation { +pub(crate) struct ArrowMemoryReservation { callbacks: vx_velox_arrow_memory_callbacks, retained_bytes: usize, active: bool, } +// SAFETY: The callback contract requires a thread-safe retained context. +unsafe impl Send for ArrowMemoryReservation {} +// SAFETY: The reservation only reads its callback table until exclusive drop. +unsafe impl Sync for ArrowMemoryReservation {} + // SAFETY: The callback contract permits the retained context to move between threads. unsafe impl Send for ArrowMemoryOwner {} // SAFETY: Release accesses the immutable callback table once through exclusive ownership. @@ -109,7 +116,7 @@ unsafe extern "C" fn release_accounted_arrow(array: *mut FFI_ArrowArray) { } impl ArrowMemoryReservation { - fn try_new( + pub(crate) fn try_new( callbacks: vx_velox_arrow_memory_callbacks, retained_bytes: usize, ) -> VortexResult { @@ -139,7 +146,7 @@ impl ArrowMemoryReservation { }) } - fn reconcile(&mut self, actual_retained_bytes: usize) -> VortexResult<()> { + pub(crate) fn reconcile(&mut self, actual_retained_bytes: usize) -> VortexResult<()> { match actual_retained_bytes.cmp(&self.retained_bytes) { std::cmp::Ordering::Less => { let released = self.retained_bytes - actual_retained_bytes; @@ -185,7 +192,7 @@ impl Drop for ArrowMemoryReservation { } } -unsafe fn parse_memory_callbacks( +pub(crate) unsafe fn parse_memory_callbacks( callbacks: *const vx_velox_arrow_memory_callbacks, ) -> VortexResult { if callbacks.is_null() { @@ -253,15 +260,15 @@ fn copy_nulls(nulls: &NullBuffer, data_offset: usize) -> VortexResult VortexResult { @@ -423,7 +430,7 @@ pub unsafe extern "C-unwind" fn vx_velox_array_export_arrow( } let mut execution = session.create_execution_ctx(); - let reserved_bytes = conservative_arrow_reservation(array, &mut execution)?; + let reserved_bytes = conservative_export_reservation(array, &mut execution)?; let mut reservation = ArrowMemoryReservation::try_new(memory_callbacks, reserved_bytes)?; let mut arrow = session .arrow() @@ -436,6 +443,7 @@ pub unsafe extern "C-unwind" fn vx_velox_array_export_arrow( .arrow() .execute_arrow(compact, None, &mut execution)?; } + validate_velox_arrow_data(&arrow.to_data())?; let schema = FFI_ArrowSchema::try_from(arrow.data_type())?; let data = copy_arrow_data(&arrow.to_data())?; let retained_bytes = data @@ -457,6 +465,7 @@ pub unsafe extern "C-unwind" fn vx_velox_array_export_arrow( #[cfg(test)] mod tests { + use std::mem::align_of; use std::ptr; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -774,6 +783,22 @@ mod tests { Ok(()) } + #[test] + fn copies_arrow_validity_into_word_padded_storage() -> VortexResult<()> { + let source = Int32Array::from(vec![Some(1), None, Some(3)]).to_data(); + let copied = copy_arrow_data(&source)?; + let nulls = copied + .nulls() + .ok_or_else(|| vortex_err!("Copied Arrow data omitted validity"))?; + + assert_eq!(nulls.buffer().len(), size_of::()); + assert_eq!(nulls.buffer().as_ptr().addr() % align_of::(), 0); + assert!(nulls.is_valid(0)); + assert!(!nulls.is_valid(1)); + assert!(nulls.is_valid(2)); + Ok(()) + } + #[test] fn rejects_arrow_allocation_and_releases_context() -> VortexResult<()> { let session = vx_session_new_with(|session| session); diff --git a/vortex-velox/src/lib.rs b/vortex-velox/src/lib.rs index bb974262d1b..a745c27b026 100644 --- a/vortex-velox/src/lib.rs +++ b/vortex-velox/src/lib.rs @@ -14,6 +14,7 @@ mod projection; mod read_at; mod schema; mod source; +mod temporal; mod visitor; pub use api::*; @@ -21,6 +22,7 @@ pub use array::vx_velox_array_export_arrow; pub use array::vx_velox_array_get_field; pub use array::vx_velox_array_invalid_count; pub use array::vx_velox_arrow_memory_callbacks; +pub use projection::vx_velox_expression_select; pub use projection::vx_velox_expression_select_with_row_index; pub use read_at::vx_velox_buffer; pub use read_at::vx_velox_read_at; @@ -30,15 +32,26 @@ pub use schema::vx_velox_source_export_schema; pub use source::vx_velox_natural_split; pub use source::vx_velox_source; pub use source::vx_velox_source_prune_natural_splits; +pub use visitor::vx_velox_binary_view; +pub use visitor::vx_velox_bool_view; pub use visitor::vx_velox_buffer_owner; +pub use visitor::vx_velox_byte_buffer_view; +pub use visitor::vx_velox_constant_view; +pub use visitor::vx_velox_dictionary_view; +pub use visitor::vx_velox_export_cursor; +pub use visitor::vx_velox_list_view; +pub use visitor::vx_velox_map_view; pub use visitor::vx_velox_primitive_type; pub use visitor::vx_velox_primitive_view; +pub use visitor::vx_velox_struct_view; pub use visitor::vx_velox_validity_kind; +pub use visitor::vx_velox_varbin_kind; +pub use visitor::vx_velox_varbin_view; pub use visitor::vx_velox_visit_request; pub use visitor::vx_velox_visitor; /// The current major version of the Vortex and Velox adapter ABI. -pub const VX_VELOX_ABI_VERSION: u32 = 1; +pub const VX_VELOX_ABI_VERSION: u32 = 5; /// The adapter supports batched host range reads. pub const VX_VELOX_CAPABILITY_BATCH_READ: u64 = 1 << 0; @@ -69,6 +82,34 @@ pub const VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING: u64 = 1 << 7; /// This capability does not claim cancellation during cached scans or CPU execution. pub const VX_VELOX_CAPABILITY_READ_CANCELLATION: u64 = 1 << 8; +/// The adapter retains one prepared array across several Velox output windows. +pub const VX_VELOX_CAPABILITY_EXPORT_CURSOR: u64 = 1 << 9; + +/// The adapter can omit row-index projection from scans with contiguous rows. +pub const VX_VELOX_CAPABILITY_PLAIN_PROJECTION: u64 = 1 << 10; + +/// The adapter can visit canonical UTF-8 and binary values in retained blocks. +pub const VX_VELOX_CAPABILITY_VARBIN_VISITOR: u64 = 1 << 11; + +/// The adapter can preserve dictionary arrays during native export. +pub const VX_VELOX_CAPABILITY_DICTIONARY_VISITOR: u64 = 1 << 12; + +/// The adapter can preserve constant arrays during native export. +pub const VX_VELOX_CAPABILITY_CONSTANT_VISITOR: u64 = 1 << 13; + +/// The adapter can visit canonical packed Boolean values in retained blocks. +pub const VX_VELOX_CAPABILITY_BOOL_VISITOR: u64 = 1 << 14; +/// The adapter can export Vortex day-based dates through the primitive visitor. +pub const VX_VELOX_CAPABILITY_DATE_VISITOR: u64 = 1 << 15; +/// The adapter can normalize Vortex decimals for the primitive visitor. +pub const VX_VELOX_CAPABILITY_DECIMAL_VISITOR: u64 = 1 << 16; +/// The adapter can preserve canonical struct children during native export. +pub const VX_VELOX_CAPABILITY_STRUCT_VISITOR: u64 = 1 << 17; +/// The adapter can preserve canonical list children during native export. +pub const VX_VELOX_CAPABILITY_LIST_VISITOR: u64 = 1 << 18; +/// The adapter can preserve canonical map children during native export. +pub const VX_VELOX_CAPABILITY_MAP_VISITOR: u64 = 1 << 19; + /// Return the adapter ABI version. #[unsafe(no_mangle)] pub extern "C" fn vx_velox_abi_version() -> u32 { @@ -87,6 +128,17 @@ pub extern "C" fn vx_velox_capabilities() -> u64 { | VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION | VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING | VX_VELOX_CAPABILITY_READ_CANCELLATION + | VX_VELOX_CAPABILITY_EXPORT_CURSOR + | VX_VELOX_CAPABILITY_PLAIN_PROJECTION + | VX_VELOX_CAPABILITY_VARBIN_VISITOR + | VX_VELOX_CAPABILITY_DICTIONARY_VISITOR + | VX_VELOX_CAPABILITY_CONSTANT_VISITOR + | VX_VELOX_CAPABILITY_BOOL_VISITOR + | VX_VELOX_CAPABILITY_DATE_VISITOR + | VX_VELOX_CAPABILITY_DECIMAL_VISITOR + | VX_VELOX_CAPABILITY_STRUCT_VISITOR + | VX_VELOX_CAPABILITY_LIST_VISITOR + | VX_VELOX_CAPABILITY_MAP_VISITOR } #[cfg(test)] @@ -120,5 +172,49 @@ mod tests { vx_velox_capabilities() & VX_VELOX_CAPABILITY_READ_CANCELLATION, VX_VELOX_CAPABILITY_READ_CANCELLATION ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_EXPORT_CURSOR, + VX_VELOX_CAPABILITY_EXPORT_CURSOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_PLAIN_PROJECTION, + VX_VELOX_CAPABILITY_PLAIN_PROJECTION + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_VARBIN_VISITOR, + VX_VELOX_CAPABILITY_VARBIN_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_DICTIONARY_VISITOR, + VX_VELOX_CAPABILITY_DICTIONARY_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_CONSTANT_VISITOR, + VX_VELOX_CAPABILITY_CONSTANT_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_BOOL_VISITOR, + VX_VELOX_CAPABILITY_BOOL_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_DATE_VISITOR, + VX_VELOX_CAPABILITY_DATE_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_DECIMAL_VISITOR, + VX_VELOX_CAPABILITY_DECIMAL_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_STRUCT_VISITOR, + VX_VELOX_CAPABILITY_STRUCT_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_LIST_VISITOR, + VX_VELOX_CAPABILITY_LIST_VISITOR + ); + assert_eq!( + vx_velox_capabilities() & VX_VELOX_CAPABILITY_MAP_VISITOR, + VX_VELOX_CAPABILITY_MAP_VISITOR + ); } } diff --git a/vortex-velox/src/projection.rs b/vortex-velox/src/projection.rs index c93314d9b73..8ed84369a0c 100644 --- a/vortex-velox/src/projection.rs +++ b/vortex-velox/src/projection.rs @@ -18,7 +18,37 @@ use vortex_ffi::vx_expression; use vortex_ffi::vx_expression_new_with; use vortex_ffi::vx_view; -unsafe fn projection( +unsafe fn field_names(names: *const vx_view, len: usize) -> VortexResult> { + let names = if names.is_null() { + vortex_ensure!(len == 0, "null field names pointer with non-zero length"); + &[] + } else { + // SAFETY: The caller provides `len` initialized views when the pointer is non-null. + unsafe { slice::from_raw_parts(names, len) } + }; + + names + .iter() + .map(|name| { + // SAFETY: Each caller-provided view remains valid for this call. + unsafe { name.as_str() }.map(FieldName::from) + }) + .collect() +} + +unsafe fn projection(names: *const vx_view, len: usize) -> VortexResult<*mut vx_expression> { + let names = unsafe { field_names(names, len) }?; + let fields = names + .into_iter() + .map(|name| (name.clone(), get_item(name, root()))) + .collect::>(); + Ok(vx_expression_new_with(pack( + fields, + Nullability::NonNullable, + ))) +} + +unsafe fn projection_with_row_index( names: *const vx_view, len: usize, row_index_name: vx_view, @@ -31,24 +61,15 @@ unsafe fn projection( "row index field name must not be empty" ); - let names = if names.is_null() { - vortex_ensure!(len == 0, "null field names pointer with non-zero length"); - &[] - } else { - // SAFETY: The caller provides `len` initialized views when the pointer is non-null. - unsafe { slice::from_raw_parts(names, len) } - }; - + let names = unsafe { field_names(names, len) }?; let mut fields = Vec::with_capacity(len + 1); fields.push((FieldName::from(row_index_name), row_idx())); for name in names { - // SAFETY: Each caller-provided view remains valid for this call. - let name = unsafe { name.as_str() }?; vortex_ensure!( - name != row_index_name, + name.as_ref() != row_index_name, "row index field name conflicts with projected field: {name}" ); - fields.push((FieldName::from(name), get_item(name, root()))); + fields.push((name.clone(), get_item(name, root()))); } Ok(vx_expression_new_with(pack( fields, @@ -56,6 +77,26 @@ unsafe fn projection( ))) } +/// Create a struct projection from the supplied field names. +/// +/// The returned expression stays owned by the caller. +/// +/// # Safety +/// +/// `names` must be null when `len` is zero or point to `len` valid views. +/// Every view must remain valid for this call. +/// `error_out` must be null or point to writable storage. No input operation can unwind. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_expression_select( + names: *const vx_view, + len: usize, + error_out: *mut *mut vx_error, +) -> *mut vx_expression { + try_or(error_out, ptr::null_mut(), || unsafe { + projection(names, len) + }) +} + /// Create a struct projection with an absolute file-row index as its first field. /// /// The remaining fields select the supplied names from the scan root. The @@ -74,7 +115,7 @@ pub unsafe extern "C-unwind" fn vx_velox_expression_select_with_row_index( error_out: *mut *mut vx_error, ) -> *mut vx_expression { try_or(error_out, ptr::null_mut(), || unsafe { - projection(names, len, row_index_name) + projection_with_row_index(names, len, row_index_name) }) } @@ -90,6 +131,17 @@ mod tests { } } + #[test] + fn creates_projection_without_row_index() { + let names = [view("a"), view("b")]; + let mut error = ptr::null_mut(); + let expression = + unsafe { vx_velox_expression_select(names.as_ptr(), names.len(), &raw mut error) }; + assert!(error.is_null()); + assert!(!expression.is_null()); + unsafe { vx_velox_expression_free(expression) }; + } + #[test] fn creates_row_index_projection() { let names = [view("a"), view("b")]; diff --git a/vortex-velox/src/schema.rs b/vortex-velox/src/schema.rs index aedbb2bdf66..d6277a4144c 100644 --- a/vortex-velox/src/schema.rs +++ b/vortex-velox/src/schema.rs @@ -10,6 +10,7 @@ use vortex_ffi::try_or; use vortex_ffi::vx_error; use crate::source::vx_velox_source; +use crate::temporal::validate_velox_arrow_type; /// Export an opened source schema through the Arrow C Data Interface. /// @@ -39,6 +40,9 @@ pub unsafe extern "C-unwind" fn vx_velox_source_export_schema( .session() .arrow() .to_arrow_schema(source.file().dtype())?; + for field in arrow_schema.fields() { + validate_velox_arrow_type(field.data_type())?; + } let schema = FFI_ArrowSchema::try_from(&arrow_schema)?; unsafe { ptr::write(schema_out, schema) }; Ok(0) diff --git a/vortex-velox/src/temporal.rs b/vortex-velox/src/temporal.rs new file mode 100644 index 00000000000..071869d1126 --- /dev/null +++ b/vortex-velox/src/temporal.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use arrow_data::ArrayData; +use arrow_schema::DataType; +use arrow_schema::TimeUnit; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +const VELOX_TIMESTAMP_MAX_SECONDS: i64 = i64::MAX / 1_000; +const VELOX_TIMESTAMP_MIN_SECONDS: i64 = i64::MIN / 1_000 - 1; + +pub(crate) fn validate_velox_arrow_type(data_type: &DataType) -> VortexResult<()> { + match data_type { + DataType::Timestamp(_, Some(timezone)) => { + vortex_bail!( + "Velox Vortex scans do not support timestamp timezone metadata: {timezone}" + ) + } + DataType::List(field) + | DataType::ListView(field) + | DataType::FixedSizeList(field, _) + | DataType::LargeList(field) + | DataType::LargeListView(field) + | DataType::Map(field, _) => validate_velox_arrow_type(field.data_type()), + DataType::Struct(fields) => { + for field in fields { + validate_velox_arrow_type(field.data_type())?; + } + Ok(()) + } + DataType::Union(fields, _) => { + for (_, field) in fields.iter() { + validate_velox_arrow_type(field.data_type())?; + } + Ok(()) + } + DataType::Dictionary(_, values) => validate_velox_arrow_type(values), + DataType::RunEndEncoded(_, values) => validate_velox_arrow_type(values.data_type()), + _ => Ok(()), + } +} + +pub(crate) fn validate_velox_arrow_data(data: &ArrayData) -> VortexResult<()> { + validate_velox_arrow_type(data.data_type())?; + if matches!( + data.data_type(), + DataType::Timestamp(TimeUnit::Second, None) + ) { + let values = data + .buffers() + .first() + .ok_or_else(|| vortex_err!("Arrow timestamp array lacks a values buffer"))? + .typed_data::(); + let end = data + .offset() + .checked_add(data.len()) + .ok_or_else(|| vortex_err!("Arrow timestamp array range overflows"))?; + let values = values.get(data.offset()..end).ok_or_else(|| { + vortex_err!( + "Arrow timestamp values buffer is too small: expected {}, got {}", + end, + values.len() + ) + })?; + for (index, value) in values.iter().enumerate() { + if data.nulls().is_some_and(|nulls| !nulls.is_valid(index)) { + continue; + } + if !(VELOX_TIMESTAMP_MIN_SECONDS..=VELOX_TIMESTAMP_MAX_SECONDS).contains(value) { + vortex_bail!("Timestamp seconds exceed the Velox range: {value}"); + } + } + } + for child in data.child_data() { + validate_velox_arrow_data(child)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::Array; + use arrow_array::ArrayRef; + use arrow_array::StructArray; + use arrow_array::TimestampMillisecondArray; + use arrow_array::TimestampSecondArray; + use arrow_buffer::NullBuffer; + use arrow_buffer::ScalarBuffer; + use arrow_schema::Field; + use vortex_error::VortexResult; + + use super::*; + + #[test] + fn rejects_nested_timestamp_timezone_metadata() -> VortexResult<()> { + let data_type = DataType::Struct( + vec![Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + )] + .into(), + ); + let error = match validate_velox_arrow_type(&data_type) { + Ok(()) => vortex_bail!("Timestamp timezone metadata unexpectedly passed validation"), + Err(error) => error, + }; + assert!(error.to_string().contains("timezone metadata: UTC")); + Ok(()) + } + + #[test] + fn validates_nested_second_timestamp_range_and_nulls() -> VortexResult<()> { + let valid = Arc::new(TimestampSecondArray::from(vec![ + Some(VELOX_TIMESTAMP_MIN_SECONDS), + None, + Some(VELOX_TIMESTAMP_MAX_SECONDS), + ])) as ArrayRef; + let valid = StructArray::from(vec![( + Arc::new(Field::new("timestamp", valid.data_type().clone(), true)), + valid, + )]); + validate_velox_arrow_data(&valid.to_data())?; + + let invalid = TimestampSecondArray::from(vec![Some(VELOX_TIMESTAMP_MAX_SECONDS + 1), None]); + let error = match validate_velox_arrow_data(&invalid.to_data()) { + Ok(()) => vortex_bail!("Out-of-range timestamp unexpectedly passed validation"), + Err(error) => error, + }; + assert!(error.to_string().contains("exceed the Velox range")); + + let null_out_of_range = TimestampSecondArray::new( + ScalarBuffer::from(vec![VELOX_TIMESTAMP_MAX_SECONDS + 1]), + Some(NullBuffer::new_null(1)), + ) + .to_data(); + validate_velox_arrow_data(&null_out_of_range)?; + Ok(()) + } + + #[test] + fn accepts_full_millisecond_storage_range() -> VortexResult<()> { + let timestamps = TimestampMillisecondArray::from(vec![i64::MIN, i64::MAX]); + validate_velox_arrow_data(×tamps.to_data()) + } +} diff --git a/vortex-velox/src/visitor.rs b/vortex-velox/src/visitor.rs index 0dfc4f13455..33a1398f83a 100644 --- a/vortex-velox/src/visitor.rs +++ b/vortex-velox/src/visitor.rs @@ -3,21 +3,57 @@ use std::ffi::c_char; use std::ffi::c_void; +use std::mem::MaybeUninit; +use std::mem::align_of; use std::mem::size_of; +use std::mem::size_of_val; +use std::ptr; use std::slice; use std::sync::Arc; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; +use vortex::array::arrays::Constant; +use vortex::array::arrays::ConstantArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::Dict; +use vortex::array::arrays::Extension; +use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::ListView; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::MapArray; use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::decimal::DecimalArrayExt; +use vortex::array::arrays::extension::ExtensionArrayExt; +use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; +use vortex::array::arrays::map::MapArrayExt; +use vortex::array::arrays::map::MapArraySlotsExt; use vortex::array::arrays::primitive::PrimitiveArrayExt; +use vortex::array::arrays::struct_::StructArrayExt; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_unsigned_integer_ptype; +use vortex::buffer::Buffer; +use vortex::buffer::BufferMut; use vortex::buffer::ByteBuffer; +use vortex::dtype::DType; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; use vortex::dtype::PType; +use vortex::extension::datetime::Date; +use vortex::extension::datetime::TimeUnit; use vortex::mask::Mask; +use vortex_array::ArrayView; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; use vortex_ffi::try_or; use vortex_ffi::vx_array; use vortex_ffi::vx_array_ref; @@ -25,6 +61,11 @@ use vortex_ffi::vx_error; use vortex_ffi::vx_session; use vortex_ffi::vx_session_ref; +use crate::array::ArrowMemoryReservation; +use crate::array::conservative_export_reservation; +use crate::array::parse_memory_callbacks; +use crate::array::vx_velox_arrow_memory_callbacks; + /// A fixed-width primitive value identifier in a semantic visitor block. pub type vx_velox_primitive_type = u32; /// Unsigned 8-bit primitive identifier. @@ -49,6 +90,8 @@ pub const VX_VELOX_PRIMITIVE_F16: vx_velox_primitive_type = 8; pub const VX_VELOX_PRIMITIVE_F32: vx_velox_primitive_type = 9; /// IEEE 754 binary64 primitive identifier. pub const VX_VELOX_PRIMITIVE_F64: vx_velox_primitive_type = 10; +/// Signed 128-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I128: vx_velox_primitive_type = 11; fn primitive_type_id(value: PType) -> vx_velox_primitive_type { match value { @@ -89,7 +132,7 @@ pub struct vx_velox_buffer_owner { pub retain: Option, /// Release one retained owner reference. pub release: Option, - /// The allocated number of payload bytes retained by this compact owner. + /// The exact sum of the value and validity allocation sizes retained by this owner. pub retained_bytes: usize, } @@ -101,6 +144,10 @@ pub struct vx_velox_primitive_view { pub struct_size: usize, /// The physical type of each value. pub primitive_type: vx_velox_primitive_type, + /// The logical decimal precision, or zero for a non-decimal block. + pub decimal_precision: u32, + /// The logical decimal scale, or zero for a non-decimal block. + pub decimal_scale: i32, /// The number of logical values in the block. pub length: usize, /// The first value byte. @@ -123,6 +170,225 @@ pub struct vx_velox_primitive_view { pub validity_alignment: usize, } +/// Identifies the logical type of a variable-width binary block. +pub type vx_velox_varbin_kind = u32; +/// Identifies UTF-8 values. +pub const VX_VELOX_VARBIN_UTF8: vx_velox_varbin_kind = 0; +/// Identifies arbitrary binary values. +pub const VX_VELOX_VARBIN_BINARY: vx_velox_varbin_kind = 1; + +/// Describes one retained payload buffer. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_byte_buffer_view { + /// The first payload byte. + pub data: *const u8, + /// The number of visible payload bytes. + pub length: usize, +} + +/// Defines the stable 16-byte variable-width binary view contract. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_binary_view { + /// Stores the logical byte length. + pub length: u32, + /// Stores inline bytes, or prefix, buffer index, and offset for outlined values. + pub data: [u8; 12], +} + +/// A canonical variable-width binary block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_varbin_view { + /// Set this field to `sizeof(vx_velox_varbin_view)`. + pub struct_size: usize, + /// Identifies UTF-8 or binary values. + pub kind: vx_velox_varbin_kind, + /// The number of logical values in the block. + pub length: usize, + /// The first 16-byte binary view. + pub views: *const vx_velox_binary_view, + /// The number of readable bytes in `views`. + pub views_length: usize, + /// The retained payload buffer descriptors. + pub data_buffers: *const vx_velox_byte_buffer_view, + /// The number of payload buffer descriptors. + pub data_buffer_count: usize, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of readable validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains all pointers in this view. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty view buffer. + pub views_alignment: usize, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + +/// A canonical packed Boolean block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_bool_view { + /// Set this field to `sizeof(vx_velox_bool_view)`. + pub struct_size: usize, + /// The number of logical Boolean values. + pub length: usize, + /// The first packed value byte. + pub values: *const u8, + /// The number of readable value bytes. + pub values_length: usize, + /// The first logical value bit within `values`. + pub values_bit_offset: usize, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of readable validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains all pointers in this view. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty value buffer. + pub values_alignment: usize, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + +/// A dictionary block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_dictionary_view { + /// Set this field to `sizeof(vx_velox_dictionary_view)`. + pub struct_size: usize, + /// The number of logical dictionary codes. + pub length: usize, + /// The canonical integer codes for this output window. + pub codes: vx_velox_primitive_view, + /// A borrowed prepared cursor for the dictionary values. + pub values: *const vx_velox_export_cursor, + /// The number of dictionary values. + pub values_length: usize, +} + +/// A constant block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_constant_view { + /// Set this field to `sizeof(vx_velox_constant_view)`. + pub struct_size: usize, + /// The number of repeated logical values. + pub length: usize, + /// A borrowed prepared cursor with one canonical value. + pub value: *const vx_velox_export_cursor, +} + +/// A canonical struct block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_struct_view { + /// Set this field to `sizeof(vx_velox_struct_view)`. + pub struct_size: usize, + /// The number of logical struct values in this window. + pub length: usize, + /// The first logical row in each field cursor. + pub offset: usize, + /// Borrowed prepared cursors in field order. + pub fields: *const *const vx_velox_export_cursor, + /// The number of field cursors. + pub field_count: usize, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of readable validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains the parent validity buffer. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + +/// A canonical list block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_list_view { + /// Set this field to `sizeof(vx_velox_list_view)`. + pub struct_size: usize, + /// The number of logical lists in this window. + pub length: usize, + /// One non-negative element offset per list. Values remain absolute against `elements`. + pub offsets: *const i32, + /// One non-negative element count per list. + pub sizes: *const i32, + /// A borrowed prepared cursor for all referenced elements. + pub elements: *const vx_velox_export_cursor, + /// The number of values in the element cursor. + pub elements_length: usize, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of readable validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains the complete offsets, sizes, and parent validity allocations. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty offsets buffer. + pub offsets_alignment: usize, + /// The guaranteed byte alignment of a non-empty sizes buffer. + pub sizes_alignment: usize, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + +/// A canonical map block delivered to Velox. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct vx_velox_map_view { + /// Set this field to `sizeof(vx_velox_map_view)`. + pub struct_size: usize, + /// The number of logical maps in this window. + pub length: usize, + /// One non-negative entry offset per map. Values remain absolute against the child cursors. + pub offsets: *const i32, + /// One non-negative entry count per map. + pub sizes: *const i32, + /// A borrowed prepared cursor for all map keys. + pub keys: *const vx_velox_export_cursor, + /// A borrowed prepared cursor for all map values. + pub values: *const vx_velox_export_cursor, + /// The number of entries in each child cursor. + pub entries_length: usize, + /// True when each map asserts sorted keys. + pub keys_sorted: bool, + /// The validity representation. + pub validity_kind: vx_velox_validity_kind, + /// The first validity byte when `validity_kind` is `Bitmap`. + pub validity: *const u8, + /// The number of readable validity bytes. + pub validity_length: usize, + /// The first logical validity bit within `validity`. + pub validity_bit_offset: usize, + /// Retains the complete offsets, sizes, and parent validity allocations. + pub buffers: vx_velox_buffer_owner, + /// The guaranteed byte alignment of a non-empty offsets buffer. + pub offsets_alignment: usize, + /// The guaranteed byte alignment of a non-empty sizes buffer. + pub sizes_alignment: usize, + /// The guaranteed byte alignment of a non-empty validity buffer. + pub validity_alignment: usize, +} + /// A single-shot subset request for the semantic visitor. #[repr(C)] #[derive(Clone, Copy, Debug)] @@ -137,7 +403,7 @@ pub struct vx_velox_visit_request { /// Host callbacks for Vortex array traversal. /// -/// One array visit calls the primitive callback synchronously. Shared tables can receive concurrent +/// One array visit calls the matching callback synchronously. Shared tables can receive concurrent /// callbacks from simultaneous visits. `last_error` must return the calling thread's most recent /// error. The string must remain valid until the next callback on that thread. Callbacks must catch /// foreign exceptions and must not unwind across this ABI. The host owns the context. @@ -156,241 +422,1878 @@ pub struct vx_velox_visitor { >, /// Return the last callback error as a null-terminated string. pub last_error: Option *const c_char>, + /// Consume one canonical variable-width binary block. Zero means success. + pub visit_varbin: Option< + unsafe extern "C" fn(context: *mut c_void, view: *const vx_velox_varbin_view) -> i32, + >, + /// Consume one dictionary block. Zero means success. + pub visit_dictionary: Option< + unsafe extern "C" fn(context: *mut c_void, view: *const vx_velox_dictionary_view) -> i32, + >, + /// Consume one constant block. Zero means success. + pub visit_constant: Option< + unsafe extern "C" fn(context: *mut c_void, view: *const vx_velox_constant_view) -> i32, + >, + /// Consume one canonical packed Boolean block. Zero means success. + pub visit_bool: + Option i32>, + /// Consume one canonical struct block. Zero means success. + pub visit_struct: Option< + unsafe extern "C" fn(context: *mut c_void, view: *const vx_velox_struct_view) -> i32, + >, + /// Consume one canonical list block. Zero means success. + pub visit_list: + Option i32>, + /// Consume one canonical map block. Zero means success. + pub visit_map: + Option i32>, } -struct PrimitiveOwner { - values: Box<[u64]>, - values_length: usize, - validity: Option>, +/// Retains one prepared Vortex array across several Velox output windows. +#[repr(C)] +pub struct vx_velox_export_cursor { + export: CursorExport, +} + +enum CursorExport { + Primitive(PrimitiveExport), + Bool(BoolExport), + VarBin(VarBinExport), + Dictionary(DictionaryExport), + Constant(ConstantExport), + Struct(StructExport), + List(ListExport), + Map(MapExport), +} + +struct PackedBits(Box<[u64]>); + +impl PackedBits { + fn try_new(bits: vortex::buffer::BitBuffer) -> VortexResult<(Self, usize)> { + let compact = bits + .chunks() + .iter_padded() + .collect::>() + .into_boxed_slice(); + let allocation = size_of_val(compact.as_ref()); + Ok((Self(compact), allocation)) + } + + fn as_ptr(&self) -> *const u8 { + self.0.as_ptr().cast() + } + + fn len(&self) -> usize { + size_of_val(self.0.as_ref()) + } +} + +struct BoolOwner { + values: PackedBits, + validity: Option, retained_bytes: usize, + memory_reservation: Option, } -impl PrimitiveOwner { +impl BoolOwner { fn try_new( - host_values: &ByteBuffer, - validity: Option<&vortex::buffer::BitBuffer>, - length: usize, + values: vortex::buffer::BitBuffer, + validity: Option, ) -> VortexResult { - let values_length = host_values.len(); - let mut values = vec![0_u64; values_length.div_ceil(size_of::())].into_boxed_slice(); - let values_allocation = values - .len() - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; - if values_length != 0 { - // SAFETY: The byte view spans the complete initialized `u64` allocation. - let values_bytes = unsafe { - slice::from_raw_parts_mut(values.as_mut_ptr().cast::(), values_allocation) - }; - values_bytes[..values_length].copy_from_slice(host_values.as_slice()); - } - let validity = validity.map(|validity| { - let mut compact = vec![0_u8; length.div_ceil(8)].into_boxed_slice(); - for (index, is_valid) in validity.into_iter().take(length).enumerate() { - if is_valid { - compact[index / 8] |= 1 << (index % 8); - } + let (values, values_allocation) = PackedBits::try_new(values)?; + let (validity, validity_allocation) = match validity { + Some(validity) => { + let (validity, allocation) = PackedBits::try_new(validity)?; + (Some(validity), allocation) } - compact - }); + None => (None, 0), + }; let retained_bytes = values_allocation - .checked_add(validity.as_ref().map_or(0, |validity| validity.len())) - .ok_or_else(|| vortex_err!("Primitive visitor retained byte count overflow"))?; + .checked_add(validity_allocation) + .ok_or_else(|| vortex_err!("Boolean visitor retained byte count overflow"))?; Ok(Self { values, - values_length, validity, retained_bytes, + memory_reservation: None, }) } - fn values(&self) -> *const u8 { - if self.values_length == 0 { - std::ptr::null() - } else { - self.values.as_ptr().cast() - } - } - - fn validity(&self) -> *const u8 { - self.validity - .as_ref() - .filter(|validity| !validity.is_empty()) - .map_or(std::ptr::null(), |validity| validity.as_ptr()) - } - - fn values_length(&self) -> usize { - self.values_length - } - - fn validity_length(&self) -> usize { - self.validity.as_ref().map_or(0, |validity| validity.len()) + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); } +} - fn retained_bytes(&self) -> usize { - self.retained_bytes - } +enum PrimitiveValues { + Compact64(Box<[MaybeUninit]>), + Compact128(Box<[MaybeUninit]>), + Retained(ByteBuffer), } -fn pointer_alignment(pointer: *const u8) -> usize { - if pointer.is_null() { - return 0; +impl PrimitiveValues { + fn as_ptr(&self) -> *const u8 { + match self { + Self::Compact64(values) => values.as_ptr().cast(), + Self::Compact128(values) => values.as_ptr().cast(), + Self::Retained(values) => values.as_ptr(), + } } - 1usize << pointer.addr().trailing_zeros() } -unsafe extern "C" fn retain_primitive_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; +struct PrimitiveOwner { + values: PrimitiveValues, + values_length: usize, + validity: Option, + retained_bytes: usize, + memory_reservation: Option, } -unsafe extern "C" fn release_primitive_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); +enum RetainedBytes { + Retained(ByteBuffer), + Compact(Box<[u8]>), } -fn validate_visitor(visitor: &vx_velox_visitor) -> VortexResult<()> { - if visitor.struct_size < size_of::() { - vortex_bail!( - "Vortex Velox visitor structure is too small: expected at least {}, got {}", - size_of::(), - visitor.struct_size - ); +impl RetainedBytes { + fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { + let buffer = handle.try_into_host_sync()?; + let length = buffer.len(); + match buffer.try_into_mut() { + Ok(buffer) => { + let allocation_size = buffer.allocation_size(); + Ok((Self::Retained(buffer.freeze()), allocation_size)) + } + Err(buffer) => { + let compact = buffer.as_slice().to_vec().into_boxed_slice(); + Ok((Self::Compact(compact), length)) + } + } } - if visitor.abi_version != crate::VX_VELOX_ABI_VERSION { - vortex_bail!( - "Unsupported Vortex Velox ABI version: expected {}, got {}", - crate::VX_VELOX_ABI_VERSION, - visitor.abi_version - ); + + fn as_ptr(&self) -> *const u8 { + match self { + Self::Retained(buffer) => buffer.as_ptr(), + Self::Compact(buffer) => buffer.as_ptr(), + } } - if visitor.visit_primitive.is_none() { - vortex_bail!("Vortex Velox visitor requires a primitive callback"); + + fn len(&self) -> usize { + match self { + Self::Retained(buffer) => buffer.len(), + Self::Compact(buffer) => buffer.len(), + } } - Ok(()) } -fn callback_error(visitor: &vx_velox_visitor, status: i32) -> String { - let Some(last_error) = visitor.last_error else { - return format!("Velox primitive visitor failed with status {status}"); - }; - // SAFETY: The callback contract returns null or a valid null-terminated string. - let message = unsafe { last_error(visitor.context) }; - if message.is_null() { - return format!("Velox primitive visitor failed with status {status}"); - } - // SAFETY: The callback keeps the string valid until the next callback. - unsafe { std::ffi::CStr::from_ptr(message) } - .to_string_lossy() - .into_owned() +enum RetainedViews { + Retained(ByteBuffer), + Compact(Box<[vx_velox_binary_view]>), } -fn selected_array( - array: &vortex::array::ArrayRef, - request: &vx_velox_visit_request, -) -> VortexResult { - if request.rows.is_null() { - if request.row_count != 0 { - vortex_bail!("A null visitor row pointer requires a zero row count"); - } - return Ok(array.clone()); - } - // SAFETY: The caller supplies `row_count` readable positions. - let rows = unsafe { slice::from_raw_parts(request.rows, request.row_count) }; - let mut previous = None; - for row in rows { - let position = usize::try_from(*row) - .map_err(|_| vortex_err!("Visitor row does not fit usize: {}", row))?; - if position >= array.len() { +impl RetainedViews { + fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { + let buffer = handle.try_into_host_sync()?; + if !buffer + .len() + .is_multiple_of(size_of::()) + { vortex_bail!( - "Visitor row is out of bounds: row {}, array length {}", - row, - array.len() + "Vortex variable-width view buffer has an invalid byte length: {}", + buffer.len() ); } - if previous.is_some_and(|previous| previous >= *row) { - vortex_bail!("Visitor rows must be unique and increasing"); + match buffer.try_into_mut() { + Ok(buffer) => { + let allocation_size = buffer.allocation_size(); + Ok((Self::Retained(buffer.freeze()), allocation_size)) + } + Err(buffer) => { + let length = buffer.len() / size_of::(); + let mut compact = vec![ + vx_velox_binary_view { + length: 0, + data: [0; 12], + }; + length + ] + .into_boxed_slice(); + if !buffer.is_empty() { + // SAFETY: Both byte ranges have the checked identical size. + unsafe { + ptr::copy_nonoverlapping( + buffer.as_ptr(), + compact.as_mut_ptr().cast::(), + buffer.len(), + ) + }; + } + let allocation = size_of_val(compact.as_ref()); + Ok((Self::Compact(compact), allocation)) + } } - previous = Some(*row); - } - let dense = rows.len() == array.len() - && rows - .iter() - .enumerate() - .all(|(position, row)| *row == position as u64); - if dense { - return Ok(array.clone()); } - array.take(PrimitiveArray::from_iter(rows.iter().copied()).into_array()) -} -fn visit_primitive( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - visitor: &vx_velox_visitor, -) -> VortexResult<()> { - let mut execution = session.create_execution_ctx(); - let Canonical::Primitive(primitive) = array.execute::(&mut execution)? else { - vortex_bail!("Primitive visitor received a non-primitive array"); - }; - let values = primitive.buffer_handle().clone(); - let host_values = values.try_to_host_sync()?; - let mask = primitive - .validity()? - .execute_mask(primitive.len(), &mut execution)?; - let (validity_kind, validity) = if !primitive.dtype().is_nullable() { - (VX_VELOX_VALIDITY_NON_NULLABLE, None) - } else { - match mask { - Mask::AllTrue(_) => (VX_VELOX_VALIDITY_ALL_VALID, None), - Mask::AllFalse(_) => (VX_VELOX_VALIDITY_ALL_INVALID, None), - Mask::Values(values) => (VX_VELOX_VALIDITY_BITMAP, Some(values.bit_buffer().clone())), + fn as_ptr(&self) -> *const vx_velox_binary_view { + match self { + Self::Retained(buffer) => buffer.as_ptr().cast(), + Self::Compact(buffer) => buffer.as_ptr(), } - }; - let owner = Arc::new(PrimitiveOwner::try_new( - &host_values, - validity.as_ref(), - primitive.len(), - )?); - let values_length = owner.values_length(); - let validity_length = owner.validity_length(); - let values = owner.values(); - let validity = owner.validity(); - let view = vx_velox_primitive_view { - struct_size: size_of::(), - primitive_type: primitive_type_id(primitive.ptype()), - length: primitive.len(), - values, - values_length, - validity_kind, - validity, - validity_length, - validity_bit_offset: 0, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&owner).cast(), - retain: Some(retain_primitive_owner), - release: Some(release_primitive_owner), - retained_bytes: owner.retained_bytes(), - }, - values_alignment: pointer_alignment(values), - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor - .visit_primitive - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a primitive callback"))?; - // SAFETY: The view and its local owner stay live until the callback returns. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); } - Ok(()) } -/// Visit one Vortex array through host semantic callbacks. -/// -/// The request selects source positions once. Callback block positions are compact and follow the -/// request order. +struct VarBinOwner { + views: RetainedViews, + _data: Box<[RetainedBytes]>, + descriptors: Box<[vx_velox_byte_buffer_view]>, + validity: Option, + retained_bytes: usize, + memory_reservation: Option, +} + +// SAFETY: The owner never mutates its buffers or pointer descriptors after construction. +// Every descriptor points into an immutable allocation that the same owner retains. +unsafe impl Send for VarBinOwner {} +// SAFETY: Shared access only reads immutable buffers and descriptors retained by this owner. +unsafe impl Sync for VarBinOwner {} + +impl VarBinOwner { + fn try_new( + views: BufferHandle, + mut buffers: Arc<[BufferHandle]>, + validity: Option, + length: usize, + ) -> VortexResult { + let (views, views_allocation) = RetainedViews::try_new(views)?; + let handles = if let Some(handles) = Arc::get_mut(&mut buffers) { + handles + .iter_mut() + .map(|handle| { + std::mem::replace(handle, BufferHandle::new_host(ByteBuffer::empty())) + }) + .collect::>() + } else { + buffers.iter().cloned().collect::>() + }; + let mut data_allocation = 0usize; + let data = handles + .into_iter() + .map(|handle| { + let (buffer, allocation) = RetainedBytes::try_new(handle)?; + data_allocation = data_allocation + .checked_add(allocation) + .ok_or_else(|| vortex_err!("Vortex string payload allocation overflow"))?; + Ok(buffer) + }) + .collect::>>()? + .into_boxed_slice(); + let descriptors = data + .iter() + .map(|buffer| vx_velox_byte_buffer_view { + data: buffer.as_ptr(), + length: buffer.len(), + }) + .collect::>() + .into_boxed_slice(); + let descriptor_allocation = size_of_val(descriptors.as_ref()); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = views_allocation + .checked_add(data_allocation) + .and_then(|bytes| bytes.checked_add(descriptor_allocation)) + .and_then(|bytes| bytes.checked_add(validity_allocation)) + .ok_or_else(|| vortex_err!("Vortex string retained byte count overflow"))?; + Ok(Self { + views, + _data: data, + descriptors, + validity, + retained_bytes, + memory_reservation: None, + }) + } + + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); + } +} + +fn retain_validity( + validity: Option, + length: usize, +) -> VortexResult<(Option, usize)> { + let Some(validity) = validity else { + return Ok((None, 0)); + }; + if validity.len() < length { + vortex_bail!( + "Vortex validity length is too small: {} for {length} values", + validity.len() + ); + } + let validity = if validity.len() == length { + validity + } else { + validity.slice(..length) + }; + let (validity, allocation) = PackedBits::try_new(validity)?; + Ok((Some(validity), allocation)) +} + +impl PrimitiveOwner { + fn try_allocate( + values_length: usize, + values_alignment: usize, + validity: Option, + length: usize, + ) -> VortexResult { + let (values, values_allocation) = if values_alignment > align_of::() { + if values_alignment > align_of::() { + vortex_bail!( + "Primitive visitor does not support value alignment {values_alignment}" + ); + } + let values = + vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] + .into_boxed_slice(); + let allocation = values + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + (PrimitiveValues::Compact128(values), allocation) + } else { + let values = + vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] + .into_boxed_slice(); + let allocation = values + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + (PrimitiveValues::Compact64(values), allocation) + }; + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = values_allocation + .checked_add(validity_allocation) + .ok_or_else(|| vortex_err!("Primitive visitor retained byte count overflow"))?; + Ok(Self { + values, + values_length, + validity, + retained_bytes, + memory_reservation: None, + }) + } + + fn try_new( + host_values: ByteBuffer, + values_alignment: usize, + validity: Option, + length: usize, + retain_values: bool, + ) -> VortexResult { + let values_length = host_values.len(); + let host_values = if retain_values { + match host_values.try_into_mut() { + Ok(values) => { + let values_allocation = values.allocation_size(); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = values_allocation + .checked_add(validity_allocation) + .ok_or_else(|| { + vortex_err!("Primitive visitor retained byte count overflow") + })?; + return Ok(Self { + values: PrimitiveValues::Retained(values.freeze()), + values_length, + validity, + retained_bytes, + memory_reservation: None, + }); + } + Err(values) => values, + } + } else { + host_values + }; + let mut owner = Self::try_allocate(values_length, values_alignment, validity, length)?; + if !host_values.is_empty() { + let (values_pointer, values_capacity) = match &mut owner.values { + PrimitiveValues::Compact64(values) => ( + values.as_mut_ptr().cast::(), + values.len() * size_of::(), + ), + PrimitiveValues::Compact128(values) => ( + values.as_mut_ptr().cast::(), + values.len() * size_of::(), + ), + PrimitiveValues::Retained(_) => { + unreachable!("a newly allocated primitive owner must be compact") + } + }; + // SAFETY: The byte view spans the complete compact allocation. + let values_bytes = + unsafe { slice::from_raw_parts_mut(values_pointer, values_capacity) }; + values_bytes[..values_length].copy_from_slice(host_values.as_slice()); + } + Ok(owner) + } + + fn try_new_bitpacked_i64( + array: ArrayView<'_, BitPacked>, + validity: Option, + ) -> VortexResult { + let values_length = array + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + let mut owner = + Self::try_allocate(values_length, align_of::(), validity, array.len())?; + // SAFETY: The allocation uses `u64` alignment and contains at least `values_length` bytes. + // The output slice covers exactly `array.len()` values and remains uniquely borrowed. + let output = unsafe { + slice::from_raw_parts_mut( + match &mut owner.values { + PrimitiveValues::Compact64(values) => { + values.as_mut_ptr().cast::>() + } + PrimitiveValues::Compact128(_) | PrimitiveValues::Retained(_) => { + unreachable!("a newly allocated primitive owner must be compact") + } + }, + array.len(), + ) + }; + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + array.unpacked_chunks(&mut scratch)?.decode_into(output); + Ok(owner) + } + + fn values(&self) -> *const u8 { + if self.values_length == 0 { + ptr::null() + } else { + self.values.as_ptr() + } + } + + fn retained_bytes(&self) -> usize { + self.retained_bytes + } + + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); + } +} + +fn pointer_alignment(pointer: *const u8) -> usize { + if pointer.is_null() { + return 0; + } + 1usize << pointer.addr().trailing_zeros() +} + +fn primitive_width(primitive_type: vx_velox_primitive_type) -> VortexResult { + Ok(match primitive_type { + VX_VELOX_PRIMITIVE_U8 | VX_VELOX_PRIMITIVE_I8 => 1, + VX_VELOX_PRIMITIVE_U16 | VX_VELOX_PRIMITIVE_I16 | VX_VELOX_PRIMITIVE_F16 => 2, + VX_VELOX_PRIMITIVE_U32 | VX_VELOX_PRIMITIVE_I32 | VX_VELOX_PRIMITIVE_F32 => 4, + VX_VELOX_PRIMITIVE_U64 | VX_VELOX_PRIMITIVE_I64 | VX_VELOX_PRIMITIVE_F64 => 8, + VX_VELOX_PRIMITIVE_I128 => 16, + _ => vortex_bail!("Unknown Vortex Velox primitive type: {primitive_type}"), + }) +} + +fn cast_decimal_values(values: Buffer, validity: &Mask) -> VortexResult +where + T: NativeDecimalType, + S: NativeDecimalType, +{ + let mut output = BufferMut::::with_capacity(values.len()); + for (index, value) in values.into_iter().enumerate() { + if !validity.value(index) { + output.push(T::default()); + continue; + } + output.push(::from(value).ok_or_else(|| { + vortex_err!( + "Decimal value cannot be represented as {}", + std::any::type_name::() + ) + })?); + } + Ok(output.freeze().into_byte_buffer()) +} + +fn normalized_decimal_values(array: &DecimalArray, validity: &Mask) -> VortexResult +where + T: NativeDecimalType, +{ + if array.values_type() == T::DECIMAL_TYPE { + return array.buffer_handle().clone().try_into_host_sync(); + } + match array.values_type() { + DecimalType::I8 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I16 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I32 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I64 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I128 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I256 => cast_decimal_values::( + array.buffer::(), + validity, + ), + } +} + +struct PrimitiveExport { + primitive_type: vx_velox_primitive_type, + decimal_precision: u32, + decimal_scale: i32, + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl PrimitiveExport { + fn try_new_decimal( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let retain_values = memory_callbacks.is_some(); + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let decimal = array.execute::(&mut execution)?; + let decimal_precision = u32::from(decimal.precision()); + let decimal_scale = i32::from(decimal.scale()); + let length = decimal.len(); + let mask = decimal + .as_ref() + .validity()? + .execute_mask(length, &mut execution)?; + let (primitive_type, host_values) = match decimal.precision() { + 1..=18 => ( + VX_VELOX_PRIMITIVE_I64, + normalized_decimal_values::(&decimal, &mask)?, + ), + 19..=38 => ( + VX_VELOX_PRIMITIVE_I128, + normalized_decimal_values::(&decimal, &mask)?, + ), + precision => { + vortex_bail!("Vortex Velox visitor does not support decimal precision {precision}") + } + }; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = PrimitiveOwner::try_new( + host_values, + primitive_width(primitive_type)?, + validity, + length, + retain_values, + )?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes())?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + primitive_type, + decimal_precision, + decimal_scale, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let retain_values = memory_callbacks.is_some(); + let direct_bitpacked = array.as_opt::().filter(|bitpacked| { + array.dtype().as_ptype() == PType::I64 && bitpacked.patches().is_none() + }); + let values_length = + array + .len() + .checked_mul(array.dtype().element_size().ok_or_else(|| { + vortex_err!("Primitive visitor received a variable-width array") + })?) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + let values_allocation = values_length + .checked_add(size_of::() - 1) + .ok_or_else(|| vortex_err!("Primitive visitor value allocation overflow"))? + / size_of::() + * size_of::(); + let validity_allocation = if array.dtype().is_nullable() { + array + .len() + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor validity allocation overflow"))? + } else { + 0 + }; + let peak_reservation = + if direct_bitpacked.is_some() { + values_allocation.checked_add(validity_allocation.checked_mul(2).ok_or_else( + || vortex_err!("Primitive visitor validity reservation overflow"), + )?) + } else { + values_allocation + .checked_add(validity_allocation) + .and_then(|bytes| bytes.checked_mul(2)) + } + .ok_or_else(|| vortex_err!("Primitive visitor memory reservation overflow"))?; + let mut memory_reservation = match (memory_callbacks, peak_reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + + let mut execution = session.create_execution_ctx(); + let (primitive_type, length, validity_kind, mut owner) = if let Some(bitpacked) = + direct_bitpacked + { + let primitive_type = primitive_type_id(array.dtype().as_ptype()); + let length = array.len(); + let mask = bitpacked.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(array.dtype().is_nullable(), mask); + let owner = PrimitiveOwner::try_new_bitpacked_i64(bitpacked, validity)?; + (primitive_type, length, validity_kind, owner) + } else { + let Canonical::Primitive(primitive) = array.execute::(&mut execution)? + else { + vortex_bail!("Primitive visitor received a non-primitive array"); + }; + let primitive_type = primitive_type_id(primitive.ptype()); + let length = primitive.len(); + let mask = primitive.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = + exported_validity(primitive.dtype().is_nullable(), mask); + let host_values = primitive.into_data_parts().buffer.try_into_host_sync()?; + let owner = PrimitiveOwner::try_new( + host_values, + primitive_width(primitive_type)?, + validity, + length, + retain_values, + )?; + (primitive_type, length, validity_kind, owner) + }; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes())?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + primitive_type, + decimal_precision: 0, + decimal_scale: 0, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn view(&self, offset: usize, length: usize) -> VortexResult { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let width = primitive_width(self.primitive_type)?; + let byte_offset = offset + .checked_mul(width) + .ok_or_else(|| vortex_err!("Vortex Velox value offset overflow"))?; + let values_length = length + .checked_mul(width) + .ok_or_else(|| vortex_err!("Vortex Velox value length overflow"))?; + let values = if values_length == 0 { + ptr::null() + } else { + // SAFETY: The checked export range lies within the retained primitive buffer. + unsafe { self.owner.values().add(byte_offset) } + }; + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Primitive validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + Ok(vx_velox_primitive_view { + struct_size: size_of::(), + primitive_type: self.primitive_type, + decimal_precision: self.decimal_precision, + decimal_scale: self.decimal_scale, + length, + values, + values_length, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_primitive_owner), + release: Some(release_primitive_owner), + retained_bytes: self.owner.retained_bytes(), + }, + values_alignment: pointer_alignment(values), + validity_alignment: pointer_alignment(validity), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let view = self.view(offset, length)?; + let callback = visitor + .visit_primitive + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a primitive callback"))?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct BoolExport { + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl BoolExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let Canonical::Bool(boolean) = array.execute::(&mut execution)? else { + vortex_bail!("Boolean visitor received a non-Boolean array"); + }; + let length = boolean.len(); + let mask = boolean.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = BoolOwner::try_new(boolean.into_bit_buffer(), validity)?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes)?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (values, values_length, values_bit_offset) = + packed_bits_window(&self.owner.values, offset, length)?; + let (validity, validity_length, validity_bit_offset) = match &self.owner.validity { + Some(validity) => packed_bits_window(validity, offset, length)?, + None => (ptr::null(), 0, 0), + }; + let view = vx_velox_bool_view { + struct_size: size_of::(), + length, + values, + values_length, + values_bit_offset, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_bool_owner), + release: Some(release_bool_owner), + retained_bytes: self.owner.retained_bytes, + }, + values_alignment: pointer_alignment(values), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_bool + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a Boolean callback"))?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +fn packed_bits_window( + bits: &PackedBits, + offset: usize, + length: usize, +) -> VortexResult<(*const u8, usize, usize)> { + if length == 0 { + return Ok((ptr::null(), 0, 0)); + } + let word_bits = u64::BITS as usize; + let byte_offset = offset / word_bits * size_of::(); + let bit_offset = offset % word_bits; + let required_length = bit_offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Packed Boolean window overflow"))? + .div_ceil(u8::BITS as usize); + let byte_length = bits + .len() + .checked_sub(byte_offset) + .ok_or_else(|| vortex_err!("Packed Boolean window exceeds its owner"))?; + if byte_length < required_length { + vortex_bail!("Packed Boolean window exceeds its readable bytes"); + } + // SAFETY: The caller validated the logical window against the owner length. + let values = unsafe { bits.as_ptr().add(byte_offset) }; + Ok((values, byte_length, bit_offset)) +} + +struct VarBinExport { + kind: vx_velox_varbin_kind, + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl VarBinExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let varbin = array.execute::(&mut execution)?; + let length = varbin.len(); + let parts = varbin.into_data_parts(); + let kind = match parts.dtype { + DType::Utf8(_) => VX_VELOX_VARBIN_UTF8, + DType::Binary(_) => VX_VELOX_VARBIN_BINARY, + dtype => vortex_bail!("Variable-width visitor received an invalid type: {dtype}"), + }; + let mask = parts.validity.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes)?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + kind, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let view_byte_offset = offset + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Vortex string view offset overflow"))?; + let views_length = length + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Vortex string view length overflow"))?; + let views = if views_length == 0 { + ptr::null() + } else { + // SAFETY: The checked export range lies within the retained view buffer. + unsafe { + self.owner + .views + .as_ptr() + .cast::() + .add(view_byte_offset) + .cast() + } + }; + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("String validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let data_buffers = if self.owner.descriptors.is_empty() { + ptr::null() + } else { + self.owner.descriptors.as_ptr() + }; + let view = vx_velox_varbin_view { + struct_size: size_of::(), + kind: self.kind, + length, + views, + views_length, + data_buffers, + data_buffer_count: self.owner.descriptors.len(), + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_varbin_owner), + release: Some(release_varbin_owner), + retained_bytes: self.owner.retained_bytes, + }, + views_alignment: pointer_alignment(views.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor.visit_varbin.ok_or_else(|| { + vortex_err!("Vortex Velox visitor requires a variable-width callback") + })?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct DictionaryExport { + codes: PrimitiveExport, + values_length: usize, + values: Box, +} + +impl DictionaryExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let dictionary = array.as_::(); + let values = dictionary.values().clone(); + Ok(Self { + codes: PrimitiveExport::try_new(dictionary.codes().clone(), session, memory_callbacks)?, + values_length: values.len(), + values: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new_canonical(values, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let codes = self.codes.view(offset, length)?; + let view = vx_velox_dictionary_view { + struct_size: size_of::(), + length, + codes, + values: &raw const *self.values, + values_length: self.values_length, + }; + let callback = visitor + .visit_dictionary + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a dictionary callback"))?; + // SAFETY: The borrowed child cursor and every code buffer remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct ConstantExport { + length: usize, + value: Box, +} + +impl ConstantExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let length = array.len(); + let scalar = array.as_::().scalar().clone(); + let value = ConstantArray::new(scalar, 1).into_array(); + Ok(Self { + length, + value: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new_canonical(value, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let view = vx_velox_constant_view { + struct_size: size_of::(), + length, + value: &raw const *self.value, + }; + let callback = visitor + .visit_constant + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a constant callback"))?; + // SAFETY: The borrowed child cursor remains live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct StructOwner { + validity: Option, + retained_bytes: usize, + _memory_reservation: Option, +} + +struct StructExport { + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, + fields: Box<[vx_velox_export_cursor]>, + field_pointers: Box<[*const vx_velox_export_cursor]>, +} + +impl StructExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let is_nullable = array.dtype().is_nullable(); + let mut execution = session.create_execution_ctx(); + let struct_array = array.execute::(&mut execution)?; + let length = struct_array.len(); + let mask = struct_array + .struct_validity() + .execute_mask(length, &mut execution)?; + let validity_reservation = if matches!(mask, Mask::Values(_)) { + length + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Struct validity reservation overflow"))? + } else { + 0 + }; + let mut memory_reservation = match (memory_callbacks, validity_reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let (validity, retained_bytes) = retain_validity(validity, length)?; + if let Some(reservation) = memory_reservation.as_mut() { + reservation.reconcile(retained_bytes)?; + } + let owner = Arc::new(StructOwner { + validity, + retained_bytes, + _memory_reservation: memory_reservation, + }); + let fields = struct_array + .iter_unmasked_fields() + .map(|field| { + Ok(vx_velox_export_cursor { + export: CursorExport::try_new(field.clone(), session, memory_callbacks)?, + }) + }) + .collect::>>()? + .into_boxed_slice(); + let field_pointers = fields + .iter() + .map(|field| field as *const vx_velox_export_cursor) + .collect::>() + .into_boxed_slice(); + Ok(Self { + length, + validity_kind, + owner, + fields, + field_pointers, + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Struct validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let view = vx_velox_struct_view { + struct_size: size_of::(), + length, + offset, + fields: if self.field_pointers.is_empty() { + ptr::null() + } else { + self.field_pointers.as_ptr() + }, + field_count: self.fields.len(), + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_struct_owner), + release: Some(release_struct_owner), + retained_bytes: self.owner.retained_bytes, + }, + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_struct + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a struct callback"))?; + // SAFETY: The borrowed field cursors and parent validity remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct ListOwner { + offsets: Box<[i32]>, + sizes: Box<[i32]>, + validity: Option, + retained_bytes: usize, + _memory_reservation: Option, +} + +struct ListMetadata { + length: usize, + elements_length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +struct ListExport { + length: usize, + elements_length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, + elements: Box, +} + +fn list_metadata_value(value: T, name: &str) -> VortexResult +where + T: Copy + std::fmt::Display, + i32: TryFrom, +{ + i32::try_from(value) + .map_err(|_| vortex_err!("Vortex list {name} exceeds the Velox vector limit: {value}")) +} + +fn list_metadata_values(values: PrimitiveArray, name: &str) -> VortexResult> { + let values = values.reinterpret_cast(values.ptype().to_unsigned()); + match_each_unsigned_integer_ptype!(values.ptype(), |P| { + values + .as_slice::

() + .iter() + .map(|&value| list_metadata_value(value, name)) + .collect::>>() + .map(Vec::into_boxed_slice) + }) +} + +fn prepare_list_metadata( + list: &ListViewArray, + session: &vortex::session::VortexSession, + memory_callbacks: Option, +) -> VortexResult { + let is_nullable = list.dtype().is_nullable(); + let mut execution = session.create_execution_ctx(); + let length = list.len(); + let elements_length = list.elements().len(); + if elements_length > i32::MAX as usize { + vortex_bail!("Vortex list elements exceed the Velox vector limit: {elements_length}"); + } + let mask = list + .listview_validity() + .execute_mask(length, &mut execution)?; + let validity_reservation = if matches!(mask, Mask::Values(_)) { + length + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("List validity reservation overflow"))? + } else { + 0 + }; + let metadata_reservation = length + .checked_mul(2 * size_of::()) + .ok_or_else(|| vortex_err!("List metadata reservation overflow"))?; + let reservation = metadata_reservation + .checked_add(validity_reservation) + .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; + let mut memory_reservation = match (memory_callbacks, reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + let offsets = list_metadata_values( + list.offsets() + .clone() + .execute::(&mut execution)?, + "offset", + )?; + let sizes = list_metadata_values( + list.sizes() + .clone() + .execute::(&mut execution)?, + "size", + )?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = size_of_val(offsets.as_ref()) + .checked_add(size_of_val(sizes.as_ref())) + .and_then(|bytes| bytes.checked_add(validity_allocation)) + .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; + if let Some(reservation) = memory_reservation.as_mut() { + reservation.reconcile(retained_bytes)?; + } + Ok(ListMetadata { + length, + elements_length, + validity_kind, + owner: Arc::new(ListOwner { + offsets, + sizes, + validity, + retained_bytes, + _memory_reservation: memory_reservation, + }), + }) +} + +impl ListExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let list = array.execute::(&mut execution)?; + let elements = list.elements().clone(); + let metadata = prepare_list_metadata(&list, session, memory_callbacks)?; + Ok(Self { + length: metadata.length, + elements_length: metadata.elements_length, + validity_kind: metadata.validity_kind, + owner: metadata.owner, + elements: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(elements, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("List validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let offsets = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.offsets.as_ptr().add(offset) } + }; + let sizes = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.sizes.as_ptr().add(offset) } + }; + let view = vx_velox_list_view { + struct_size: size_of::(), + length, + offsets, + sizes, + elements: &raw const *self.elements, + elements_length: self.elements_length, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_list_owner), + release: Some(release_list_owner), + retained_bytes: self.owner.retained_bytes, + }, + offsets_alignment: pointer_alignment(offsets.cast()), + sizes_alignment: pointer_alignment(sizes.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_list + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a list callback"))?; + // SAFETY: The borrowed element cursor and parent buffers remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct MapExport { + length: usize, + entries_length: usize, + keys_sorted: bool, + validity_kind: vx_velox_validity_kind, + owner: Arc, + keys: Box, + values: Box, +} + +impl MapExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let map = array.execute::(&mut execution)?; + let keys_sorted = map.keys_sorted(); + let entries = map.entries().clone().downcast::(); + let entry_values = entries.elements().clone(); + let entry_struct = entry_values.execute::(&mut execution)?; + let fields = entry_struct.iter_unmasked_fields().collect::>(); + if fields.len() != 2 { + vortex_bail!( + "Vortex map entries require two fields, got {}", + fields.len() + ); + } + let metadata = prepare_list_metadata(&entries, session, memory_callbacks)?; + Ok(Self { + length: metadata.length, + entries_length: metadata.elements_length, + keys_sorted, + validity_kind: metadata.validity_kind, + owner: metadata.owner, + keys: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(fields[0].clone(), session, memory_callbacks)?, + }), + values: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(fields[1].clone(), session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Map validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let offsets = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.offsets.as_ptr().add(offset) } + }; + let sizes = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.sizes.as_ptr().add(offset) } + }; + let view = vx_velox_map_view { + struct_size: size_of::(), + length, + offsets, + sizes, + keys: &raw const *self.keys, + values: &raw const *self.values, + entries_length: self.entries_length, + keys_sorted: self.keys_sorted, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_list_owner), + release: Some(release_list_owner), + retained_bytes: self.owner.retained_bytes, + }, + offsets_alignment: pointer_alignment(offsets.cast()), + sizes_alignment: pointer_alignment(sizes.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_map + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a map callback"))?; + // SAFETY: The borrowed child cursors and parent buffers remain live through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +impl CursorExport { + fn date_storage( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + ) -> VortexResult> { + let DType::Extension(ext_dtype) = array.dtype() else { + return Ok(None); + }; + let Some(time_unit) = ext_dtype.metadata_opt::() else { + return Ok(None); + }; + if *time_unit != TimeUnit::Days { + vortex_bail!( + "Vortex Velox visitor does not support date unit {time_unit}; Velox DATE uses days" + ); + } + + if let Some(extension) = array.as_opt::() { + return Ok(Some(extension.storage_array().clone())); + } + let mut execution = session.create_execution_ctx(); + let extension = array.execute::(&mut execution)?; + Ok(Some(extension.storage_array().clone())) + } + + fn try_new_canonical( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + if matches!(array.dtype(), DType::Map(..)) { + Ok(Self::Map(MapExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::List(..)) { + Ok(Self::List(ListExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Struct(..)) { + Ok(Self::Struct(StructExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Decimal(..)) { + Ok(Self::Primitive(PrimitiveExport::try_new_decimal( + array, + session, + memory_callbacks, + )?)) + } else if let Some(storage) = Self::date_storage(array.clone(), session)? { + Ok(Self::Primitive(PrimitiveExport::try_new( + storage, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Bool(_)) { + Ok(Self::Bool(BoolExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { + Ok(Self::VarBin(VarBinExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else { + Ok(Self::Primitive(PrimitiveExport::try_new( + array, + session, + memory_callbacks, + )?)) + } + } + + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + if array.is::() { + Ok(Self::Dictionary(DictionaryExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if array.is::() { + Ok(Self::Constant(ConstantExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else { + Self::try_new_canonical(array, session, memory_callbacks) + } + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + match self { + Self::Primitive(export) => export.visit(offset, length, visitor), + Self::Bool(export) => export.visit(offset, length, visitor), + Self::VarBin(export) => export.visit(offset, length, visitor), + Self::Dictionary(export) => export.visit(offset, length, visitor), + Self::Constant(export) => export.visit(offset, length, visitor), + Self::Struct(export) => export.visit(offset, length, visitor), + Self::List(export) => export.visit(offset, length, visitor), + Self::Map(export) => export.visit(offset, length, visitor), + } + } +} + +fn exported_validity( + is_nullable: bool, + mask: Mask, +) -> (vx_velox_validity_kind, Option) { + if !is_nullable { + return (VX_VELOX_VALIDITY_NON_NULLABLE, None); + } + match mask { + Mask::AllTrue(_) => (VX_VELOX_VALIDITY_ALL_VALID, None), + Mask::AllFalse(_) => (VX_VELOX_VALIDITY_ALL_INVALID, None), + Mask::Values(values) => (VX_VELOX_VALIDITY_BITMAP, Some(values.bit_buffer().clone())), + } +} + +unsafe extern "C" fn retain_primitive_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_primitive_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_bool_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_bool_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_varbin_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_varbin_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_struct_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_struct_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_list_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_list_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +fn validate_visitor(visitor: &vx_velox_visitor) -> VortexResult<()> { + if visitor.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox visitor structure is too small: expected at least {}, got {}", + size_of::(), + visitor.struct_size + ); + } + if visitor.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + visitor.abi_version + ); + } + Ok(()) +} + +fn callback_error(visitor: &vx_velox_visitor, status: i32) -> String { + let Some(last_error) = visitor.last_error else { + return format!("Velox visitor failed with status {status}"); + }; + // SAFETY: The callback contract returns null or a valid null-terminated string. + let message = unsafe { last_error(visitor.context) }; + if message.is_null() { + return format!("Velox visitor failed with status {status}"); + } + // SAFETY: The callback keeps the string valid until the next callback. + unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned() +} + +fn selected_array( + array: &vortex::array::ArrayRef, + request: &vx_velox_visit_request, +) -> VortexResult { + if request.rows.is_null() { + if request.row_count != 0 { + vortex_bail!("A null visitor row pointer requires a zero row count"); + } + return Ok(array.clone()); + } + // SAFETY: The caller supplies `row_count` readable positions. + let rows = unsafe { slice::from_raw_parts(request.rows, request.row_count) }; + let mut previous = None; + for row in rows { + let position = usize::try_from(*row) + .map_err(|_| vortex_err!("Visitor row does not fit usize: {}", row))?; + if position >= array.len() { + vortex_bail!( + "Visitor row is out of bounds: row {}, array length {}", + row, + array.len() + ); + } + if previous.is_some_and(|previous| previous >= *row) { + vortex_bail!("Visitor rows must be unique and increasing"); + } + previous = Some(*row); + } + let dense = rows.len() == array.len() + && rows + .iter() + .enumerate() + .all(|(position, row)| *row == position as u64); + if dense { + return Ok(array.clone()); + } + array.take(PrimitiveArray::from_iter(rows.iter().copied()).into_array()) +} + +fn visit_array( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + visitor: &vx_velox_visitor, +) -> VortexResult<()> { + let length = array.len(); + CursorExport::try_new_canonical(array, session, None)?.visit(0, length, visitor) +} + +/// Create one export cursor for several Velox output windows. +/// +/// # Safety +/// +/// The session and array pointers must identify live handles. +/// The memory callbacks must identify a complete, thread-safe callback table. +/// `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_export_cursor_new( + session: *const vx_session, + array: *const vx_array, + memory_callbacks: *const vx_velox_arrow_memory_callbacks, + error_out: *mut *mut vx_error, +) -> *mut vx_velox_export_cursor { + try_or(error_out, ptr::null_mut(), || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + let memory_callbacks = unsafe { parse_memory_callbacks(memory_callbacks)? }; + Ok(Box::into_raw(Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(array.clone(), session, Some(memory_callbacks))?, + }))) + }) +} + +/// Free one export cursor. +/// +/// # Safety +/// +/// The pointer must be null or come from [`vx_velox_export_cursor_new`]. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_export_cursor_free(cursor: *mut vx_velox_export_cursor) { + if !cursor.is_null() { + // SAFETY: The pointer came from `Box::into_raw` and is freed once. + drop(unsafe { Box::from_raw(cursor) }); + } +} + +/// Visit one contiguous range from a retained export cursor. +/// +/// # Safety +/// +/// The cursor and visitor pointers must remain live until this call returns. +/// Concurrent calls are valid. The caller must not free the cursor before all calls return. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_export_cursor_visit( + cursor: *const vx_velox_export_cursor, + offset: usize, + length: usize, + visitor: *const vx_velox_visitor, + error_out: *mut *mut vx_error, +) -> i32 { + try_or(error_out, 1, || { + let cursor = unsafe { + cursor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox export cursor must not be null"))? + }; + let visitor = unsafe { + visitor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? + }; + validate_visitor(visitor)?; + cursor.export.visit(offset, length, visitor)?; + Ok(0) + }) +} + +/// Visit one Vortex array through host semantic callbacks. +/// +/// The request selects source positions once. Callback block positions are compact and follow the +/// request order. /// /// # Safety /// @@ -425,20 +2328,40 @@ pub unsafe extern "C-unwind" fn vx_velox_array_visit( .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? }; validate_visitor(visitor)?; - visit_primitive(selected_array(array, request)?, session, visitor)?; + visit_array(selected_array(array, request)?, session, visitor)?; Ok(0) }) } #[cfg(test)] mod tests { + use std::mem::align_of; use std::ptr; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use rstest::rstest; + use vortex::array::ArrayRef; use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::DecimalArray; + use vortex::array::arrays::DictArray; + use vortex::array::arrays::ListViewArray; + use vortex::array::arrays::MapArray; use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::StructArray; + use vortex::array::arrays::TemporalArray; + use vortex::array::arrays::VarBinViewArray; + use vortex::array::validity::Validity; + use vortex::buffer::buffer; + use vortex::dtype::DecimalDType; + use vortex::dtype::FieldNames; + use vortex::dtype::MapDType; + use vortex::dtype::Nullability; + use vortex::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_ensure; + use vortex_fastlanes::BitPackedData; use vortex_ffi::vx_array_new_with; use vortex_ffi::vx_session_free; use vortex_ffi::vx_session_new_with; @@ -446,6 +2369,41 @@ mod tests { use super::*; use crate::api::vx_velox_array_free; + #[derive(Default)] + struct TestMemory { + retained_bytes: AtomicUsize, + } + + unsafe extern "C" fn retain_test_memory(_context: *mut c_void) {} + + unsafe extern "C" fn release_test_memory(_context: *mut c_void) {} + + unsafe extern "C" fn reserve_test_memory(context: *mut c_void, bytes: usize) -> i32 { + // SAFETY: The test context stays live through every callback. + let memory = unsafe { &*context.cast::() }; + memory.retained_bytes.fetch_add(bytes, Ordering::Relaxed); + 0 + } + + unsafe extern "C" fn free_test_memory(context: *mut c_void, bytes: usize) { + // SAFETY: The test context stays live through every callback. + let memory = unsafe { &*context.cast::() }; + memory.retained_bytes.fetch_sub(bytes, Ordering::Relaxed); + } + + fn test_memory_callbacks(memory: &mut TestMemory) -> vx_velox_arrow_memory_callbacks { + vx_velox_arrow_memory_callbacks { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (memory as *mut TestMemory).cast(), + retain_context: Some(retain_test_memory), + release_context: Some(release_test_memory), + report_allocation: Some(reserve_test_memory), + report_free: Some(free_test_memory), + last_error: None, + } + } + #[rstest] #[case(PType::U8, VX_VELOX_PRIMITIVE_U8)] #[case(PType::U16, VX_VELOX_PRIMITIVE_U16)] @@ -462,51 +2420,1228 @@ mod tests { assert_eq!(primitive_type_id(input), expected); } + #[test] + fn date_days_use_i32_storage_and_millisecond_dates_are_rejected() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let days = TemporalArray::new_date( + PrimitiveArray::from_option_iter([Some(-1_i32), None, Some(19_000)]).into_array(), + TimeUnit::Days, + ) + .into_array(); + let CursorExport::Primitive(export) = + CursorExport::try_new_canonical(days, &session, None)? + else { + vortex_bail!("date visitor did not produce primitive storage"); + }; + assert_eq!(export.primitive_type, VX_VELOX_PRIMITIVE_I32); + let view = export.view(0, 3)?; + // SAFETY: The export owns three readable i32 values. + let values = unsafe { slice::from_raw_parts(view.values.cast::(), 3) }; + assert_eq!(values, [-1, 0, 19_000]); + assert_eq!(view.validity_kind, VX_VELOX_VALIDITY_BITMAP); + + let milliseconds = TemporalArray::new_date( + PrimitiveArray::from_iter([86_400_000_i64]).into_array(), + TimeUnit::Milliseconds, + ) + .into_array(); + let error = match CursorExport::try_new_canonical(milliseconds, &session, None) { + Ok(_) => vortex_bail!("millisecond date visitor unexpectedly succeeded"), + Err(error) => error, + }; + assert!(error.to_string().contains("Velox DATE uses days")); + Ok(()) + } + + #[test] + fn decimals_normalize_to_velox_storage_widths() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let short = DecimalArray::new( + buffer![1_i8, -2, 3], + DecimalDType::new(18, 2), + Validity::NonNullable, + ) + .into_array(); + let short = PrimitiveExport::try_new_decimal(short, &session, None)?; + assert_eq!(short.primitive_type, VX_VELOX_PRIMITIVE_I64); + let short_view = short.view(0, 3)?; + assert_eq!(short_view.decimal_precision, 18); + assert_eq!(short_view.decimal_scale, 2); + // SAFETY: The export owns three readable i64 values. + let short_values = unsafe { slice::from_raw_parts(short_view.values.cast::(), 3) }; + assert_eq!(short_values, [1, -2, 3]); + + let nullable_short = DecimalArray::new( + buffer![1_i128, i128::MAX], + DecimalDType::new(18, 2), + Validity::from_iter([true, false]), + ) + .into_array(); + let nullable_short = PrimitiveExport::try_new_decimal(nullable_short, &session, None)?; + let nullable_short_view = nullable_short.view(0, 2)?; + // SAFETY: The export owns two readable i64 values. + let nullable_short_values = + unsafe { slice::from_raw_parts(nullable_short_view.values.cast::(), 2) }; + assert_eq!(nullable_short_values, [1, 0]); + assert_eq!(nullable_short_view.validity_kind, VX_VELOX_VALIDITY_BITMAP); + + let long = DecimalArray::new( + buffer![1_i64, -2, 3], + DecimalDType::new(30, 4), + Validity::NonNullable, + ) + .into_array(); + let long = PrimitiveExport::try_new_decimal(long, &session, None)?; + assert_eq!(long.primitive_type, VX_VELOX_PRIMITIVE_I128); + let long_view = long.view(0, 3)?; + assert_eq!(long_view.decimal_precision, 30); + assert_eq!(long_view.decimal_scale, 4); + // SAFETY: The export owns three readable i128 values. + let long_values = unsafe { slice::from_raw_parts(long_view.values.cast::(), 3) }; + assert_eq!(long_values, [1, -2, 3]); + + let unsupported = DecimalArray::new( + buffer![1_i8], + DecimalDType::new(39, 0), + Validity::NonNullable, + ) + .into_array(); + let error = match PrimitiveExport::try_new_decimal(unsupported, &session, None) { + Ok(_) => vortex_bail!("precision 39 decimal visitor unexpectedly succeeded"), + Err(error) => error, + }; + assert!(error.to_string().contains("decimal precision 39")); + Ok(()) + } + + #[test] + fn dictionary_export_preserves_code_width_and_nullable_children() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let code_cases: [(ArrayRef, vx_velox_primitive_type); 4] = [ + (buffer![0_u8, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U8), + (buffer![0_u16, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U16), + (buffer![0_u32, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U32), + (buffer![0_u64, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U64), + ]; + for (codes, expected_type) in code_cases { + let dictionary = DictArray::try_new(codes, buffer![10_i64, 20].into_array())?; + let CursorExport::Dictionary(export) = + CursorExport::try_new(dictionary.into_array(), &session, None)? + else { + vortex_bail!("dictionary export lost its outer encoding"); + }; + assert_eq!(export.codes.primitive_type, expected_type); + assert_eq!(export.values_length, 2); + assert!(matches!(export.values.export, CursorExport::Primitive(_))); + } + + let codes = PrimitiveArray::from_option_iter([Some(0_u8), None, Some(1)]).into_array(); + let values = PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(); + let dictionary = DictArray::try_new(codes, values)?; + let CursorExport::Dictionary(export) = + CursorExport::try_new(dictionary.into_array(), &session, None)? + else { + vortex_bail!("nullable dictionary export lost its outer encoding"); + }; + assert_eq!(export.codes.validity_kind, VX_VELOX_VALIDITY_BITMAP); + let CursorExport::Primitive(values) = &export.values.export else { + vortex_bail!("nullable dictionary values lost their primitive representation"); + }; + assert_eq!(values.validity_kind, VX_VELOX_VALIDITY_BITMAP); + Ok(()) + } + + #[test] + fn constant_export_preserves_null_value() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let constant = ConstantArray::new(Scalar::null_native::(), 10).into_array(); + let CursorExport::Constant(export) = CursorExport::try_new(constant, &session, None)? + else { + vortex_bail!("constant export lost its outer encoding"); + }; + assert_eq!(export.length, 10); + let CursorExport::Primitive(value) = &export.value.export else { + vortex_bail!("null constant lost its primitive representation"); + }; + assert_eq!(value.length, 1); + assert_eq!(value.validity_kind, VX_VELOX_VALIDITY_ALL_INVALID); + Ok(()) + } + + #[test] + fn struct_export_preserves_children_and_nonzero_window() -> VortexResult<()> { + #[derive(Default)] + struct StructCapture { + length: usize, + offset: usize, + fields: *const *const vx_velox_export_cursor, + field_count: usize, + validity: *const u8, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_struct( + context: *mut c_void, + view: *const vx_velox_struct_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offset = view.offset; + capture.fields = view.fields; + capture.field_count = view.field_count; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let length: usize = 130; + let dictionary = DictArray::try_new( + PrimitiveArray::from_iter((0..length).map(|index| [0_u8, 1][index % 2])).into_array(), + buffer![10_i64, 20].into_array(), + )? + .into_array(); + let constant = ConstantArray::new(Scalar::from(7_i64), length).into_array(); + let parent_validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); + let struct_array = StructArray::new( + FieldNames::from(["dictionary", "constant"]), + [dictionary, constant], + length, + parent_validity, + ) + .into_array(); + let CursorExport::Struct(export) = CursorExport::try_new(struct_array, &session, None)? + else { + vortex_bail!("struct export lost its outer encoding"); + }; + assert!(matches!( + export.fields[0].export, + CursorExport::Dictionary(_) + )); + assert!(matches!(export.fields[1].export, CursorExport::Constant(_))); + + let mut capture = StructCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: Some(capture_struct), + visit_list: None, + visit_map: None, + }; + export.visit(65, 63, &visitor)?; + assert_eq!(capture.length, 63); + assert_eq!(capture.offset, 65); + assert_eq!(capture.field_count, 2); + assert_eq!(capture.validity_bit_offset, 1); + // SAFETY: The export retains both field cursors until it is dropped below. + assert_eq!(unsafe { *capture.fields }, &raw const export.fields[0]); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("struct callback returned no validity owner"))?; + drop(export); + // SAFETY: The callback retained the parent owner before the cursor was dropped. + assert!( + unsafe { + *capture + .validity + .add(capture.validity_bit_offset / u8::BITS as usize) + } != 0 + ); + let release = owner + .release + .ok_or_else(|| vortex_err!("struct owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + Ok(()) + } + + #[test] + fn list_export_preserves_elements_window_and_accounting() -> VortexResult<()> { + #[derive(Default)] + struct ListCapture { + length: usize, + offsets: *const i32, + sizes: *const i32, + elements_length: usize, + validity: *const u8, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_list( + context: *mut c_void, + view: *const vx_velox_list_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offsets = view.offsets; + capture.sizes = view.sizes; + capture.elements_length = view.elements_length; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let length = 130; + let elements = DictArray::try_new( + buffer![0_u8, 1, 0, 1, 0, 1].into_array(), + PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(), + )? + .into_array(); + let offsets = PrimitiveArray::from_iter((0..length).map(|index| [0_u32, 2, 4][index % 3])); + let sizes = + PrimitiveArray::from_iter((0..length).map(|index| if index % 10 == 0 { 0 } else { 2 })); + let validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); + let list = ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity) + .into_array(); + let mut memory = TestMemory::default(); + let CursorExport::List(export) = + CursorExport::try_new(list, &session, Some(test_memory_callbacks(&mut memory)))? + else { + vortex_bail!("list export lost its outer encoding"); + }; + assert!(matches!( + export.elements.export, + CursorExport::Dictionary(_) + )); + let expected_parent_bytes = + length * 2 * size_of::() + length.div_ceil(u64::BITS as usize) * size_of::(); + assert_eq!(export.owner.retained_bytes, expected_parent_bytes); + + let mut capture = ListCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: Some(capture_list), + visit_map: None, + }; + export.visit(65, 63, &visitor)?; + assert_eq!(capture.length, 63); + assert_eq!(capture.elements_length, 6); + assert_eq!(capture.validity_bit_offset, 1); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.offsets }, 4); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.sizes }, 2); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("list callback returned no owner"))?; + drop(export); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_parent_bytes + ); + // SAFETY: The callback retained the owner before the export was dropped. + assert_eq!(unsafe { *capture.offsets.add(1) }, 0); + let release = owner + .release + .ok_or_else(|| vortex_err!("list owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn map_export_preserves_children_window_and_accounting() -> VortexResult<()> { + #[derive(Default)] + struct MapCapture { + length: usize, + offsets: *const i32, + sizes: *const i32, + keys: *const vx_velox_export_cursor, + values: *const vx_velox_export_cursor, + entries_length: usize, + keys_sorted: bool, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_map( + context: *mut c_void, + view: *const vx_velox_map_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offsets = view.offsets; + capture.sizes = view.sizes; + capture.keys = view.keys; + capture.values = view.values; + capture.entries_length = view.entries_length; + capture.keys_sorted = view.keys_sorted; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let keys = DictArray::try_new( + buffer![0_u8, 1, 0, 1, 0, 1].into_array(), + buffer![10_i64, 20].into_array(), + )? + .into_array(); + let values = ConstantArray::new(Scalar::from(7_i64), 6).into_array(); + let entries = StructArray::new( + FieldNames::from(["key", "value"]), + [keys, values], + 6, + Validity::NonNullable, + ) + .into_array(); + let entry_lists = ListViewArray::new( + entries, + buffer![0_u32, 2, 4].into_array(), + buffer![2_u32, 2, 2].into_array(), + Validity::from_iter([true, false, true]), + ); + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I64, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::NonNullable), + true, + )?; + let map = MapArray::try_new(map_dtype, entry_lists)?.into_array(); + let mut memory = TestMemory::default(); + let CursorExport::Map(export) = + CursorExport::try_new(map, &session, Some(test_memory_callbacks(&mut memory)))? + else { + vortex_bail!("map export lost its outer encoding"); + }; + assert!(matches!(export.keys.export, CursorExport::Dictionary(_))); + assert!(matches!(export.values.export, CursorExport::Constant(_))); + let expected_parent_bytes = 3 * 2 * size_of::() + size_of::(); + assert_eq!(export.owner.retained_bytes, expected_parent_bytes); + + let mut capture = MapCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: Some(capture_map), + }; + export.visit(1, 2, &visitor)?; + assert_eq!(capture.length, 2); + assert_eq!(capture.entries_length, 6); + assert!(capture.keys_sorted); + assert_eq!(capture.validity_bit_offset, 1); + assert_eq!(capture.keys, &raw const *export.keys); + assert_eq!(capture.values, &raw const *export.values); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.offsets }, 2); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.sizes }, 2); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("map callback returned no owner"))?; + drop(export); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_parent_bytes + ); + // SAFETY: The callback retained the owner before the export was dropped. + assert_eq!(unsafe { *capture.offsets.add(1) }, 4); + let release = owner + .release + .ok_or_else(|| vortex_err!("map owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[derive(Default)] + struct Capture { + primitive_type: Option, + length: usize, + values: *const u8, + values_length: usize, + values_alignment: usize, + validity: *const u8, + validity_length: usize, + validity_bit_offset: usize, + validity_alignment: usize, + retained_bytes: usize, + validity_kind: Option, + owner: Option, + } + + unsafe extern "C" fn capture_primitive( + context: *mut c_void, + view: *const vx_velox_primitive_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live `Capture` and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.primitive_type = Some(view.primitive_type); + capture.length = view.length; + capture.values = view.values; + capture.values_length = view.values_length; + capture.values_alignment = view.values_alignment; + capture.validity = view.validity; + capture.validity_length = view.validity_length; + capture.validity_bit_offset = view.validity_bit_offset; + capture.validity_alignment = view.validity_alignment; + capture.retained_bytes = view.buffers.retained_bytes; + capture.validity_kind = Some(view.validity_kind); + capture.owner = Some(view.buffers); + 0 + } + + fn release_capture(capture: &Capture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_primitive`. + unsafe { release(owner.owner) }; + Ok(()) + } + #[derive(Default)] - struct Capture { - primitive_type: Option, + struct VarBinCapture { + struct_size: usize, + kind: Option, length: usize, - values: *const u8, - values_length: usize, - values_alignment: usize, + views: *const vx_velox_binary_view, + views_length: usize, + views_alignment: usize, + data_buffers: *const vx_velox_byte_buffer_view, + data_buffer_count: usize, validity: *const u8, validity_length: usize, validity_bit_offset: usize, validity_alignment: usize, - retained_bytes: usize, validity_kind: Option, + retained_bytes: usize, owner: Option, } - unsafe extern "C" fn capture_primitive( + unsafe extern "C" fn capture_varbin( context: *mut c_void, - view: *const vx_velox_primitive_view, + view: *const vx_velox_varbin_view, ) -> i32 { if context.is_null() || view.is_null() { return 1; } - // SAFETY: The test passes pointers to live `Capture` and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; let Some(retain) = view.buffers.retain else { return 2; }; // SAFETY: The visitor owner is live for the callback. unsafe { retain(view.buffers.owner) }; - capture.primitive_type = Some(view.primitive_type); + capture.struct_size = view.struct_size; + capture.kind = Some(view.kind); capture.length = view.length; - capture.values = view.values; - capture.values_length = view.values_length; - capture.values_alignment = view.values_alignment; + capture.views = view.views; + capture.views_length = view.views_length; + capture.views_alignment = view.views_alignment; + capture.data_buffers = view.data_buffers; + capture.data_buffer_count = view.data_buffer_count; capture.validity = view.validity; capture.validity_length = view.validity_length; capture.validity_bit_offset = view.validity_bit_offset; capture.validity_alignment = view.validity_alignment; + capture.validity_kind = Some(view.validity_kind); capture.retained_bytes = view.buffers.retained_bytes; + capture.owner = Some(view.buffers); + 0 + } + + fn release_varbin_capture(capture: &VarBinCapture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained string owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("string owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_varbin`. + unsafe { release(owner.owner) }; + Ok(()) + } + + #[derive(Default)] + struct BoolCapture { + length: usize, + values: *const u8, + values_bit_offset: usize, + validity: *const u8, + validity_bit_offset: usize, + validity_kind: Option, + retained_bytes: usize, + owner: Option, + } + + unsafe extern "C" fn capture_bool( + context: *mut c_void, + view: *const vx_velox_bool_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.values = view.values; + capture.values_bit_offset = view.values_bit_offset; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; capture.validity_kind = Some(view.validity_kind); + capture.retained_bytes = view.buffers.retained_bytes; capture.owner = Some(view.buffers); 0 } + fn release_bool_capture(capture: &BoolCapture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained Boolean owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("Boolean owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_bool`. + unsafe { release(owner.owner) }; + Ok(()) + } + + #[expect( + clippy::host_endian_bytes, + reason = "The Vortex binary-view fields use the host C ABI layout" + )] + unsafe fn captured_varbin_value(capture: &VarBinCapture, index: usize) -> Option<&[u8]> { + if capture.validity_kind == Some(VX_VELOX_VALIDITY_BITMAP) { + let bit_index = capture.validity_bit_offset + index; + // SAFETY: The callback contract retains the bitmap for every captured row. + let byte = unsafe { *capture.validity.add(bit_index / 8) }; + if byte & (1 << (bit_index % 8)) == 0 { + return None; + } + } + // SAFETY: The callback contract retains `length` readable views. + let view = unsafe { &*capture.views.add(index) }; + let length = view.length as usize; + const INLINE_LENGTH: usize = size_of::() - size_of::(); + if length <= INLINE_LENGTH { + return Some(&view.data[..length]); + } + let buffer_index = + u32::from_ne_bytes([view.data[4], view.data[5], view.data[6], view.data[7]]) as usize; + let offset = + u32::from_ne_bytes([view.data[8], view.data[9], view.data[10], view.data[11]]) as usize; + // SAFETY: The callback contract retains all payload descriptors. + let buffer = unsafe { &*capture.data_buffers.add(buffer_index) }; + // SAFETY: Canonical Vortex views contain validated payload ranges. + Some(unsafe { slice::from_raw_parts(buffer.data.add(offset), length) }) + } + + #[rstest] + #[case(DType::Utf8(Nullability::Nullable), VX_VELOX_VARBIN_UTF8)] + #[case(DType::Binary(Nullability::Nullable), VX_VELOX_VARBIN_BINARY)] + fn varbin_cursor_retains_mixed_views_across_nonzero_window( + #[case] dtype: DType, + #[case] expected_kind: vx_velox_varbin_kind, + ) -> VortexResult<()> { + let utf8_expected: [Option<&[u8]>; 7] = [ + Some(b""), + Some(b"a"), + None, + Some(b"abcdefghijkl"), + Some(b"abcdefghijklm"), + Some("vortex 🌀 outlined".as_bytes()), + Some(b"tail"), + ]; + let binary_expected: [Option<&[u8]>; 7] = [ + Some(b""), + Some(b"\xff"), + None, + Some(b"abcdefghijkl"), + Some(b"\x00abcdefghijklm"), + Some(b"\xff\x00 binary outlined value"), + Some(b"tail"), + ]; + let expected = if matches!(dtype, DType::Utf8(_)) { + utf8_expected + } else { + binary_expected + }; + let session = vx_session_new_with(|session| session); + let varbin = VarBinViewArray::from_iter(expected, dtype); + let array = vx_array_new_with(varbin.into_array()); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "string cursor creation failed"); + vortex_ensure!(error.is_null(), "string cursor returned an error"); + + let mut capture = VarBinCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: Some(capture_varbin), + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 1, 5, &raw const visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "string export window returned an error"); + assert_eq!(capture.struct_size, size_of::()); + assert_eq!(capture.kind, Some(expected_kind)); + assert_eq!(capture.length, 5); + assert_eq!(capture.views_length, 5 * size_of::()); + assert!(capture.views_alignment >= align_of::()); + assert_eq!(capture.views.addr() % align_of::(), 0); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.validity_bit_offset, 1); + assert!(capture.validity_length >= 1); + assert!(capture.validity_alignment >= align_of::()); + assert_eq!(capture.validity.addr() % align_of::(), 0); + assert!(capture.data_buffer_count >= 1); + assert!(!capture.data_buffers.is_null()); + assert_eq!( + capture.retained_bytes, + memory.retained_bytes.load(Ordering::Relaxed) + ); + + // SAFETY: Each owned handle is freed once. The callback retained the string owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + capture.retained_bytes + ); + for (index, expected) in expected[1..6].iter().enumerate() { + // SAFETY: The retained owner keeps every captured pointer live. + let actual = unsafe { captured_varbin_value(&capture, index) }; + assert_eq!(actual, *expected); + } + release_varbin_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn varbin_shared_buffers_compact_into_exact_owned_storage() -> VortexResult<()> { + let length = 130_usize; + let strings = VarBinViewArray::from_iter( + (0..length).map(|index| { + (index % 11 != 0).then(|| format!("outlined string value {index:03}")) + }), + DType::Utf8(Nullability::Nullable), + ); + let parts = strings.into_data_parts(); + let views_length = parts.views.try_to_host_sync()?.len(); + let data_length = parts + .buffers + .iter() + .map(|buffer| Ok(buffer.try_to_host_sync()?.len())) + .sum::>()?; + let descriptor_length = parts.buffers.len() * size_of::(); + let validity_length = length.div_ceil(u64::BITS as usize) * size_of::(); + let expected_retained = views_length + data_length + descriptor_length + validity_length; + + let retained_views = parts.views.clone(); + let retained_buffers = Arc::<[BufferHandle]>::clone(&parts.buffers); + let mut execution = vortex::session::VortexSession::empty().create_execution_ctx(); + let mask = parts.validity.execute_mask(length, &mut execution)?; + let (_, validity) = exported_validity(true, mask); + let owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; + + assert!(matches!(owner.views, RetainedViews::Compact(_))); + assert!( + owner + ._data + .iter() + .all(|buffer| matches!(buffer, RetainedBytes::Compact(_))) + ); + assert_eq!(owner.retained_bytes, expected_retained); + drop(retained_views); + drop(retained_buffers); + Ok(()) + } + + #[test] + fn retained_varbin_buffers_report_complete_unique_allocations() -> VortexResult<()> { + let alignment = vortex::buffer::Alignment::new(256); + let mut payload = BufferMut::::with_capacity_aligned(17, alignment); + payload.extend(0..17); + let expected_payload_allocation = payload.allocation_size(); + let (retained_payload, payload_allocation) = + RetainedBytes::try_new(BufferHandle::new_host(payload.freeze()))?; + assert!(matches!(retained_payload, RetainedBytes::Retained(_))); + assert_eq!(payload_allocation, expected_payload_allocation); + assert!(payload_allocation > 17); + + let mut views = BufferMut::::with_capacity_aligned( + 2 * size_of::(), + alignment, + ); + views.extend(std::iter::repeat_n( + 0, + 2 * size_of::(), + )); + let expected_views_allocation = views.allocation_size(); + let (retained_views, views_allocation) = + RetainedViews::try_new(BufferHandle::new_host(views.freeze()))?; + assert!(matches!(retained_views, RetainedViews::Retained(_))); + assert_eq!(views_allocation, expected_views_allocation); + assert!(views_allocation > 2 * size_of::()); + Ok(()) + } + + #[test] + fn word_aligned_windows_rebase_validity_buffers() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let primitive = PrimitiveArray::from_option_iter( + (0..130).map(|index| (index % 7 != 0).then_some(index as i64)), + ) + .into_array(); + let primitive = PrimitiveExport::try_new(primitive, &session, None)?; + let primitive_first = primitive.view(0, 64)?; + let primitive_second = primitive.view(64, 64)?; + assert_eq!(primitive_first.validity_bit_offset, 0); + assert_eq!(primitive_second.validity_bit_offset, 0); + // SAFETY: Both pointers lie in the retained validity allocation. + assert_eq!(primitive_second.validity, unsafe { + primitive_first.validity.add(size_of::()) + }); + + let strings = VarBinViewArray::from_iter( + (0..130).map(|index| (index % 11 != 0).then(|| format!("value-{index}"))), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let strings = VarBinExport::try_new(strings, &session, None)?; + let mut first = VarBinCapture::default(); + let first_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut first).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: Some(capture_varbin), + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + strings.visit(0, 64, &first_visitor)?; + + let mut second = VarBinCapture::default(); + let second_visitor = vx_velox_visitor { + context: (&raw mut second).cast(), + ..first_visitor + }; + strings.visit(64, 64, &second_visitor)?; + assert_eq!(first.validity_bit_offset, 0); + assert_eq!(second.validity_bit_offset, 0); + // SAFETY: Both pointers lie in the retained validity allocation. + assert_eq!(second.validity, unsafe { + first.validity.add(size_of::()) + }); + release_varbin_capture(&first)?; + release_varbin_capture(&second)?; + Ok(()) + } + + #[test] + fn bool_cursor_retains_nonzero_window_and_exact_accounting() -> VortexResult<()> { + let expected = (0..130) + .map(|index| (index % 11 != 0).then_some(index % 3 == 0)) + .collect::>(); + let session = vx_session_new_with(|session| session); + let boolean = BoolArray::from_iter(expected.iter().copied()); + let array = vx_array_new_with(boolean.into_array()); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "Boolean cursor creation failed"); + vortex_ensure!(error.is_null(), "Boolean cursor returned an error"); + + let mut capture = BoolCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: Some(capture_bool), + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 65, 63, &raw const visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "Boolean export window returned an error"); + assert_eq!(capture.length, 63); + assert_eq!(capture.values_bit_offset, 1); + assert_eq!(capture.validity_bit_offset, 1); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.retained_bytes, 6 * size_of::()); + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 48); + + // SAFETY: Each owned handle is freed once. The callback retained the Boolean owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + capture.retained_bytes + ); + for (relative_index, expected) in expected[65..128].iter().enumerate() { + let value_bit = capture.values_bit_offset + relative_index; + let validity_bit = capture.validity_bit_offset + relative_index; + // SAFETY: The retained buffers cover every captured value and validity bit. + let (actual, is_valid) = unsafe { + ( + *capture.values.add(value_bit / 8) & (1 << (value_bit % 8)) != 0, + *capture.validity.add(validity_bit / 8) & (1 << (validity_bit % 8)) != 0, + ) + }; + assert_eq!(is_valid, expected.is_some()); + if let Some(expected) = expected { + assert_eq!(actual, *expected); + } + } + release_bool_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn export_cursor_reuses_one_prepared_array_across_windows() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with( + PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40), Some(50)]) + .into_array(), + ); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); + vortex_ensure!(error.is_null(), "export cursor returned an error"); + assert!(memory.retained_bytes.load(Ordering::Relaxed) >= 48); + + let mut first = Capture::default(); + let first_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut first).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 1, 2, &raw const first_visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "first export window returned an error"); + assert_eq!(first.length, 2); + assert_eq!(first.validity_bit_offset, 1); + // SAFETY: The callback retained two readable i64 values. + let first_values = unsafe { slice::from_raw_parts(first.values.cast::(), 2) }; + assert_eq!(first_values, [0, 30]); + assert_eq!( + first.retained_bytes, + memory.retained_bytes.load(Ordering::Relaxed) + ); + let owner = first + .owner + .ok_or_else(|| vortex_err!("first export window returned no owner"))? + .owner; + release_capture(&first)?; + + let mut second = Capture::default(); + let second_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut second).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 3, 2, &raw const second_visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "second export window returned an error"); + assert_eq!(second.length, 2); + assert_eq!(second.validity_bit_offset, 3); + assert_eq!( + second + .owner + .ok_or_else(|| vortex_err!("second export window returned no owner"))? + .owner, + owner + ); + + // SAFETY: Each owned handle is freed exactly once. The second callback retained the owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + // SAFETY: The retained cursor owner keeps these two i64 values live. + let second_values = unsafe { slice::from_raw_parts(second.values.cast::(), 2) }; + assert_eq!(second_values, [40, 50]); + release_capture(&second)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn export_cursor_decodes_sliced_bitpacked_into_exact_owner() -> VortexResult<()> { + let session = vx_session_new_with(|session| { + vortex_fastlanes::initialize(&session); + session + }); + let session_ref = unsafe { vx_session_ref(session)? }; + let values = (0..2_050).map(|index| (index % 7 != 0).then_some(i64::from(index % 100))); + let primitive = PrimitiveArray::from_option_iter(values).into_array(); + let mut execution = session_ref.create_execution_ctx(); + let bitpacked = BitPackedData::encode(&primitive, 7, &mut execution)?; + vortex_ensure!( + bitpacked.patches().is_none(), + "test bit-packed array unexpectedly contains patches" + ); + let slice_begin = 113; + let slice_end = 1_941; + let sliced = bitpacked.into_array().slice(slice_begin..slice_end)?; + let array = vx_array_new_with(sliced); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); + vortex_ensure!(error.is_null(), "export cursor returned an error"); + let sliced_length = slice_end - slice_begin; + let expected_retained = sliced_length * size_of::() + + sliced_length.div_ceil(u64::BITS as usize) * size_of::(); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_retained + ); + + let window_offset = 997; + let window_length = 6; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit( + cursor, + window_offset, + window_length, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "export window returned an error"); + assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!( + capture.validity_bit_offset, + window_offset % u64::BITS as usize + ); + assert_eq!(capture.retained_bytes, expected_retained); + // SAFETY: The callback retained `window_length` readable i64 values. + let actual = unsafe { slice::from_raw_parts(capture.values.cast::(), window_length) }; + for (relative_index, value) in actual.iter().enumerate() { + let sliced_index = window_offset + relative_index; + let source_index = slice_begin + sliced_index; + // SAFETY: The retained bitmap covers every row in the sliced array. + let validity_index = capture.validity_bit_offset + relative_index; + let validity_byte = unsafe { *capture.validity.add(validity_index / 8) }; + let is_valid = validity_byte & (1 << (validity_index % 8)) != 0; + assert_eq!(is_valid, source_index % 7 != 0); + if is_valid { + assert_eq!(*value, i64::try_from(source_index % 100)?); + } + } + + // SAFETY: Each owned handle is freed exactly once. The callback retained the owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_retained + ); + release_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn patched_bitpacked_uses_retained_canonical_fallback() -> VortexResult<()> { + let session = vx_session_new_with(|session| { + vortex_fastlanes::initialize(&session); + session + }); + let session_ref = unsafe { vx_session_ref(session)? }; + let expected = [1_u64, 2, 3, u64::MAX]; + let primitive = PrimitiveArray::from_iter(expected).into_array(); + let mut execution = session_ref.create_execution_ctx(); + let bitpacked = BitPackedData::encode(&primitive, 2, &mut execution)?; + vortex_ensure!( + bitpacked.patches().is_some(), + "test bit-packed array unexpectedly omitted patches" + ); + let mut memory = TestMemory::default(); + let export = PrimitiveExport::try_new( + bitpacked.into_array(), + session_ref, + Some(test_memory_callbacks(&mut memory)), + )?; + assert!(matches!(export.owner.values, PrimitiveValues::Retained(_))); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + export.owner.retained_bytes() + ); + // SAFETY: The export owner contains `expected.len()` initialized u64 values. + let actual = + unsafe { slice::from_raw_parts(export.owner.values().cast::(), expected.len()) }; + assert_eq!(actual, expected); + drop(export); + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + unsafe { vx_session_free(session) }; + Ok(()) + } + #[test] fn visits_sparse_nullable_values_with_retained_buffers() -> VortexResult<()> { let session = vx_session_new_with(|session| session); @@ -526,6 +3661,13 @@ mod tests { context: (&raw mut capture).cast(), visit_primitive: Some(capture_primitive), last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, }; let mut error = ptr::null_mut(); // SAFETY: Every handle and callback object stays live for this call. @@ -546,15 +3688,18 @@ mod tests { assert!(capture.values_alignment.is_power_of_two()); assert_eq!(capture.values.addr() % capture.values_alignment, 0); assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); - assert_eq!(capture.validity_length, 1); + assert_eq!(capture.validity_length, size_of::()); assert_eq!(capture.validity_bit_offset, 0); assert!(capture.validity_alignment.is_power_of_two()); assert_eq!(capture.validity.addr() % capture.validity_alignment, 0); - assert_eq!(capture.retained_bytes, capture.values_length + 1); + assert_eq!( + capture.retained_bytes, + capture.values_length + size_of::() + ); // SAFETY: The callback retained the owner before storing these pointers. let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 2) }; assert_eq!(values, [0, 40]); - // SAFETY: The retained validity pointer has one readable byte. + // SAFETY: The retained validity pointer has one readable word. let validity = unsafe { *capture.validity }; assert_eq!(validity & 0b11, 0b10); @@ -575,12 +3720,13 @@ mod tests { } #[test] - fn copies_sliced_values_and_reports_compact_allocation() -> VortexResult<()> { + fn copies_sliced_values_into_exact_owned_storage() -> VortexResult<()> { let session = vx_session_new_with(|session| session); let source = PrimitiveArray::from_iter(0_i32..16); let source_values = source.buffer_handle().try_to_host_sync()?; // SAFETY: The source contains sixteen i32 values. The fifth value is in bounds. let source_slice = unsafe { source_values.as_ptr().add(5 * size_of::()) }; + drop(source_values); let array = vx_array_new_with(source.into_array().slice(5..8)?); let request = vx_velox_visit_request { struct_size: size_of::(), @@ -594,6 +3740,13 @@ mod tests { context: (&raw mut capture).cast(), visit_primitive: Some(capture_primitive), last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, }; let mut error = ptr::null_mut(); let status = unsafe { @@ -610,7 +3763,12 @@ mod tests { assert_eq!(capture.values_length, 3 * size_of::()); assert_eq!(capture.retained_bytes, 2 * size_of::()); assert_ne!(capture.values, source_slice); - // SAFETY: The retained compact values contain three i32 values. + // SAFETY: Each owned handle is freed exactly once. The callback retained the value owner. + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + // SAFETY: The retained compact buffer contains three i32 values. let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 3) }; assert_eq!(values, [5, 6, 7]); assert!(capture.values_alignment.is_power_of_two()); @@ -624,15 +3782,11 @@ mod tests { .release .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; unsafe { release(owner.owner) }; - unsafe { - vx_velox_array_free(array); - vx_session_free(session); - } Ok(()) } #[test] - fn copies_validity_into_compact_storage() -> VortexResult<()> { + fn copies_validity_into_word_padded_storage() -> VortexResult<()> { let session = vx_session_new_with(|session| session); let session_ref = unsafe { vx_session_ref(session)? }; let primitive = PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]); @@ -657,6 +3811,13 @@ mod tests { context: (&raw mut capture).cast(), visit_primitive: Some(capture_primitive), last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, }; let mut error = ptr::null_mut(); let status = unsafe { @@ -672,7 +3833,12 @@ mod tests { vortex_ensure!(error.is_null(), "visitor returned an error"); assert_ne!(capture.validity, expected_validity); assert_eq!(capture.validity_bit_offset, 0); - assert_eq!(capture.retained_bytes, 2 * size_of::() + 1); + assert_eq!(capture.validity_length, size_of::()); + assert!(capture.validity_alignment >= align_of::()); + assert_eq!( + capture.retained_bytes, + capture.values_length.div_ceil(size_of::()) * size_of::() + size_of::() + ); let owner = capture .owner diff --git a/vortex-velox/tests/abi_contract.rs b/vortex-velox/tests/abi_contract.rs index 40101f17150..f3d526e1214 100644 --- a/vortex-velox/tests/abi_contract.rs +++ b/vortex-velox/tests/abi_contract.rs @@ -55,12 +55,27 @@ mod tests { source.push_str( "_Static_assert(sizeof(vx_velox_validity_kind) == sizeof(uint32_t), \"validity width\");\n", ); + source.push_str( + "_Static_assert(sizeof(vx_velox_varbin_kind) == sizeof(uint32_t), \"varbin width\");\n", + ); source.push_str( "_Static_assert(VX_VELOX_PTYPE_F64 == 10, \"ptype value\");\n\ _Static_assert(VX_VELOX_OPERATOR_KLEENE_OR == 7, \"operator value\");\n\ _Static_assert(VX_VELOX_SELECTION_EXCLUDE == 2, \"selection value\");\n\ _Static_assert(VX_VELOX_PRIMITIVE_F64 == 10, \"primitive value\");\n\ - _Static_assert(VX_VELOX_VALIDITY_BITMAP == 3, \"validity value\");\n", + _Static_assert(VX_VELOX_PRIMITIVE_I128 == 11, \"i128 primitive value\");\n\ + _Static_assert(VX_VELOX_VALIDITY_BITMAP == 3, \"validity value\");\n\ + _Static_assert(VX_VELOX_VARBIN_BINARY == 1, \"varbin value\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_VARBIN_VISITOR == (UINT64_C(1) << 11), \"varbin capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_DICTIONARY_VISITOR == (UINT64_C(1) << 12), \"dictionary capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_CONSTANT_VISITOR == (UINT64_C(1) << 13), \"constant capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_BOOL_VISITOR == (UINT64_C(1) << 14), \"Boolean capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_DATE_VISITOR == (UINT64_C(1) << 15), \"date capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_DECIMAL_VISITOR == (UINT64_C(1) << 16), \"decimal capability\");\n\ + _Static_assert(VX_VELOX_CAPABILITY_STRUCT_VISITOR == (UINT64_C(1) << 17), \"struct capability\");\n", + ); + source.push_str( + "_Static_assert(VX_VELOX_CAPABILITY_LIST_VISITOR == (UINT64_C(1) << 18), \"list capability\");\n", ); macro_rules! check_layout { @@ -127,9 +142,110 @@ mod tests { [ struct_size, primitive_type, + decimal_precision, + decimal_scale, + length, + values, + values_length, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + values_alignment, + validity_alignment, + ] + ); + check_layout!(vx_velox_byte_buffer_view, [data, length]); + check_layout!(vx_velox_binary_view, [length, data]); + check_layout!( + vx_velox_varbin_view, + [ + struct_size, + kind, + length, + views, + views_length, + data_buffers, + data_buffer_count, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + views_alignment, + validity_alignment, + ] + ); + check_layout!( + vx_velox_dictionary_view, + [struct_size, length, codes, values, values_length] + ); + check_layout!(vx_velox_constant_view, [struct_size, length, value]); + check_layout!( + vx_velox_struct_view, + [ + struct_size, + length, + offset, + fields, + field_count, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + validity_alignment, + ] + ); + check_layout!( + vx_velox_list_view, + [ + struct_size, + length, + offsets, + sizes, + elements, + elements_length, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + offsets_alignment, + sizes_alignment, + validity_alignment, + ] + ); + check_layout!( + vx_velox_map_view, + [ + struct_size, + length, + offsets, + sizes, + keys, + values, + entries_length, + keys_sorted, + validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers, + offsets_alignment, + sizes_alignment, + validity_alignment, + ] + ); + check_layout!( + vx_velox_bool_view, + [ + struct_size, length, values, values_length, + values_bit_offset, validity_kind, validity, validity_length, @@ -148,6 +264,13 @@ mod tests { context, visit_primitive, last_error, + visit_varbin, + visit_dictionary, + visit_constant, + visit_bool, + visit_struct, + visit_list, + visit_map, ] ); check_layout!( diff --git a/vortex-velox/tests/velox_include_contract.cpp b/vortex-velox/tests/velox_include_contract.cpp index 332248ce665..d4b927d8d43 100644 --- a/vortex-velox/tests/velox_include_contract.cpp +++ b/vortex-velox/tests/velox_include_contract.cpp @@ -12,7 +12,7 @@ typedef struct ArrowArrayStream FFI_ArrowArrayStream; #include "vortex_velox.h" #undef USE_OWN_ARROW -static_assert(VX_VELOX_ABI_VERSION == 1u); +static_assert(VX_VELOX_ABI_VERSION == 5u); static_assert(VX_VELOX_SELECTION_ALL == 0); static_assert(VX_VELOX_OPERATOR_EQ == 0); From aceea3d221e7e69c62266c37deaabbcb27402c91 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 13:36:25 -0400 Subject: [PATCH 4/6] refactor(vortex-velox): Own adapter ABI support Signed-off-by: Will Manning --- Cargo.lock | 2 +- vortex-ffi/src/array.rs | 7 - vortex-ffi/src/data_source.rs | 10 +- vortex-ffi/src/dtype.rs | 7 - vortex-ffi/src/expression.rs | 26 +- vortex-ffi/src/lib.rs | 10 - vortex-ffi/src/scan.rs | 16 - vortex-velox/Cargo.toml | 2 +- vortex-velox/README.md | 6 +- vortex-velox/cinclude/vortex_velox.h | 269 +++++---- vortex-velox/cinclude/vortex_velox_test.h | 30 + vortex-velox/src/api.rs | 265 +++------ vortex-velox/src/array.rs | 42 +- vortex-velox/src/ffi.rs | 550 ++++++++++++++++++ vortex-velox/src/lib.rs | 2 + vortex-velox/src/projection.rs | 44 +- vortex-velox/src/read_at.rs | 21 +- vortex-velox/src/schema.rs | 8 +- vortex-velox/src/source.rs | 32 +- vortex-velox/src/test_support.rs | 239 ++++++++ vortex-velox/src/visitor.rs | 32 +- vortex-velox/tests/abi_contract.rs | 22 +- vortex-velox/tests/velox_include_contract.cpp | 17 +- 23 files changed, 1160 insertions(+), 499 deletions(-) create mode 100644 vortex-velox/cinclude/vortex_velox_test.h create mode 100644 vortex-velox/src/ffi.rs create mode 100644 vortex-velox/src/test_support.rs diff --git a/Cargo.lock b/Cargo.lock index ebc429a4746..619e90ee5f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11698,6 +11698,7 @@ dependencies = [ "arrow-buffer 59.2.0", "arrow-data 59.2.0", "arrow-schema 59.2.0", + "async-fs", "bytes", "futures", "rstest", @@ -11707,7 +11708,6 @@ dependencies = [ "vortex-buffer", "vortex-error", "vortex-fastlanes", - "vortex-ffi", "vortex-io", ] diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 58c69263bfa..3d4b9758567 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -68,13 +68,6 @@ box_wrapper!( vx_array ); -/// Create an FFI array handle from an owned Vortex array. -/// -/// Layered FFI crates use this function when their host-specific scan path produces an array. -pub fn vx_array_new_with(array: ArrayRef) -> *const vx_array { - vx_array::new(array) -} - /// Borrow the [`ArrayRef`] behind a [`vx_array`] handle, erroring on a null pointer. /// /// A building block for FFI crates layered on top of the base Vortex C API. diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index 87c0bcfe36f..50d31b3036d 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -38,15 +38,7 @@ box_wrapper!( /// /// Copying a vx_data_source via vx_data_source_clone is a cheap operation. MultiLayoutDataSource, - vx_data_source -); - -/// Create an FFI data-source handle from a configured multi-layout data source. -/// -/// Layered FFI crates use this function after their host-specific I/O adapter constructs a source. -pub fn vx_data_source_new_with(data_source: MultiLayoutDataSource) -> *const vx_data_source { - vx_data_source::new(data_source) -} + vx_data_source); /// Options for creating a data source. #[repr(C)] diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index c8121693e69..7156bc5334a 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -34,13 +34,6 @@ box_wrapper!( vx_dtype ); -/// Create an FFI dtype handle from an owned Vortex dtype. -/// -/// Layered FFI crates use this function after they validate their host-specific type identifiers. -pub fn vx_dtype_new_with(dtype: DType) -> *const vx_dtype { - vx_dtype::new(dtype) -} - /// The variant tag for a Vortex data type. #[non_exhaustive] #[repr(C)] diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 9b9d6d46248..f15db42006f 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -40,31 +40,7 @@ box_wrapper!( /// Operations on expressions don't take ownership of input values, and so /// input values must be freed by the caller. Expression, - vx_expression -); - -/// Create an FFI expression handle from an owned Vortex expression. -/// -/// Layered FFI crates use this function for host-specific expression constructors. -pub fn vx_expression_new_with(expression: Expression) -> *mut vx_expression { - vx_expression::new(expression) -} - -/// Borrow an expression from a layered FFI crate. -/// -/// # Safety -/// -/// `expression` must point to a live expression handle for the returned reference lifetime. -pub unsafe fn vx_expression_ref<'a>( - expression: *const vx_expression, -) -> vortex::error::VortexResult<&'a Expression> { - let expression = unsafe { - expression - .as_ref() - .ok_or_else(|| vortex::error::vortex_err!("Vortex expression must not be null"))? - }; - Ok(&expression.0) -} + vx_expression); /// Create a root expression. A root expression, applied to an array in /// vx_array_apply, takes the array itself as opposed to functions like diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 9e1c2f0b213..1ef423bb383 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -27,24 +27,14 @@ use std::sync::Arc; use std::sync::LazyLock; pub use array::vx_array; -pub use array::vx_array_new_with; pub use array::vx_array_ref; -pub use data_source::vx_data_source; -pub use data_source::vx_data_source_new_with; pub use dtype::vx_dtype; -pub use dtype::vx_dtype_new_with; pub use error::try_or; pub use error::vx_error; pub use error::vx_error_free; -pub use expression::vx_expression; -pub use expression::vx_expression_new_with; -pub use expression::vx_expression_ref; pub use log::vx_log_level; -pub use scalar::vx_scalar; -pub use scan::vx_data_source_scan_with; pub use scan::vx_partition; pub use scan::vx_partition_into_array_stream; -pub use scan::vx_scan; pub use session::vx_session; pub use session::vx_session_free; pub use session::vx_session_new_with; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 798e2bbbe42..d5ea24a9ac1 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -225,22 +225,6 @@ fn write_estimate>(estimate: Precision, out: &mut vx_estimate) { } } -/// Start a scan from a request that a layered FFI crate already validated. -/// -/// # Safety -/// -/// `data_source` must point to a live data-source handle created by this crate. -pub unsafe fn vx_data_source_scan_with( - data_source: *const vx_data_source, - request: ScanRequest, -) -> VortexResult<*mut vx_scan> { - vortex_ensure!(!data_source.is_null(), "null vx_data_source"); - RUNTIME.block_on(async { - let scan = vx_data_source::as_ref(data_source).scan(request).await?; - Ok(vx_scan::new(VxScan::Pending(scan))) - }) -} - /// Scan a data source. /// /// A scan may be consumed only once. diff --git a/vortex-velox/Cargo.toml b/vortex-velox/Cargo.toml index b0f3bcf3315..8a1de5108a7 100644 --- a/vortex-velox/Cargo.toml +++ b/vortex-velox/Cargo.toml @@ -26,6 +26,7 @@ arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } +async-fs = { workspace = true } bytes = { workspace = true } futures = { workspace = true } vortex-array = { workspace = true } @@ -33,7 +34,6 @@ vortex-buffer = { workspace = true } vortex-arrow = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } -vortex-ffi = { path = "../vortex-ffi" } vortex-io = { workspace = true } vortex = { workspace = true } diff --git a/vortex-velox/README.md b/vortex-velox/README.md index a96950f44e3..d55ef1f0bff 100644 --- a/vortex-velox/README.md +++ b/vortex-velox/README.md @@ -20,9 +20,9 @@ the difference before it returns outputs. Arrow release frees the final charge. ## Contract boundary -The adapter exposes a versioned C ABI in `cinclude/vortex_velox.h`. Velox calls -only `vx_velox_*` symbols. The static archive can contain general `vx_*` symbols -from linked Vortex FFI objects. Opaque handle layouts stay inside Vortex. +The adapter exposes a versioned C ABI in `cinclude/vortex_velox.h`. The header +and static library are standalone. Velox calls only `vx_velox_*` symbols. +Opaque handle layouts stay inside Vortex. The adapter accepts host callbacks for random reads. Velox can implement those callbacks with `dwio::common::BufferedInput`, so Vortex uses the existing cache diff --git a/vortex-velox/cinclude/vortex_velox.h b/vortex-velox/cinclude/vortex_velox.h index 98a0948163f..62e6ff9c76f 100644 --- a/vortex-velox/cinclude/vortex_velox.h +++ b/vortex-velox/cinclude/vortex_velox.h @@ -6,16 +6,14 @@ #include #include -#include "vortex.h" +struct ArrowSchema; +struct ArrowArray; #ifdef __cplusplus extern "C" { #endif -/* - * Velox must call only vx_velox_* adapter symbols. The static archive can also - * contain general vx_* symbols from linked Vortex FFI objects. - */ +/* Velox calls only vx_velox_* adapter symbols. */ #define VX_VELOX_ABI_VERSION 5u #define VX_VELOX_CAPABILITY_BATCH_READ (UINT64_C(1) << 0) @@ -27,22 +25,36 @@ extern "C" { #define VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION (UINT64_C(1) << 6) #define VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING (UINT64_C(1) << 7) /* Vortex checks cancellation before each host read callback. */ -#define VX_VELOX_CAPABILITY_READ_CANCELLATION (UINT64_C(1) << 8) -#define VX_VELOX_CAPABILITY_EXPORT_CURSOR (UINT64_C(1) << 9) -#define VX_VELOX_CAPABILITY_PLAIN_PROJECTION (UINT64_C(1) << 10) -#define VX_VELOX_CAPABILITY_VARBIN_VISITOR (UINT64_C(1) << 11) -#define VX_VELOX_CAPABILITY_DICTIONARY_VISITOR (UINT64_C(1) << 12) -#define VX_VELOX_CAPABILITY_CONSTANT_VISITOR (UINT64_C(1) << 13) -#define VX_VELOX_CAPABILITY_BOOL_VISITOR (UINT64_C(1) << 14) -#define VX_VELOX_CAPABILITY_DATE_VISITOR (UINT64_C(1) << 15) -#define VX_VELOX_CAPABILITY_DECIMAL_VISITOR (UINT64_C(1) << 16) -#define VX_VELOX_CAPABILITY_STRUCT_VISITOR (UINT64_C(1) << 17) -#define VX_VELOX_CAPABILITY_LIST_VISITOR (UINT64_C(1) << 18) -#define VX_VELOX_CAPABILITY_MAP_VISITOR (UINT64_C(1) << 19) +#define VX_VELOX_CAPABILITY_READ_CANCELLATION (UINT64_C(1) << 8) +#define VX_VELOX_CAPABILITY_EXPORT_CURSOR (UINT64_C(1) << 9) +#define VX_VELOX_CAPABILITY_PLAIN_PROJECTION (UINT64_C(1) << 10) +#define VX_VELOX_CAPABILITY_VARBIN_VISITOR (UINT64_C(1) << 11) +#define VX_VELOX_CAPABILITY_DICTIONARY_VISITOR (UINT64_C(1) << 12) +#define VX_VELOX_CAPABILITY_CONSTANT_VISITOR (UINT64_C(1) << 13) +#define VX_VELOX_CAPABILITY_BOOL_VISITOR (UINT64_C(1) << 14) +#define VX_VELOX_CAPABILITY_DATE_VISITOR (UINT64_C(1) << 15) +#define VX_VELOX_CAPABILITY_DECIMAL_VISITOR (UINT64_C(1) << 16) +#define VX_VELOX_CAPABILITY_STRUCT_VISITOR (UINT64_C(1) << 17) +#define VX_VELOX_CAPABILITY_LIST_VISITOR (UINT64_C(1) << 18) +#define VX_VELOX_CAPABILITY_MAP_VISITOR (UINT64_C(1) << 19) typedef struct vx_velox_read_at vx_velox_read_at; typedef struct vx_velox_source vx_velox_source; typedef struct vx_velox_export_cursor vx_velox_export_cursor; +typedef struct vx_velox_error vx_velox_error; +typedef struct vx_velox_session vx_velox_session; +typedef struct vx_velox_dtype vx_velox_dtype; +typedef struct vx_velox_scalar vx_velox_scalar; +typedef struct vx_velox_expression vx_velox_expression; +typedef struct vx_velox_data_source vx_velox_data_source; +typedef struct vx_velox_scan vx_velox_scan; +typedef struct vx_velox_partition vx_velox_partition; +typedef struct vx_velox_array vx_velox_array; + +typedef struct vx_velox_view { + const char *ptr; + size_t len; +} vx_velox_view; typedef uint32_t vx_velox_ptype; #define VX_VELOX_PTYPE_U8 UINT32_C(0) @@ -58,14 +70,14 @@ typedef uint32_t vx_velox_ptype; #define VX_VELOX_PTYPE_F64 UINT32_C(10) typedef uint32_t vx_velox_binary_operator; -#define VX_VELOX_OPERATOR_EQ UINT32_C(0) -#define VX_VELOX_OPERATOR_NOT_EQ UINT32_C(1) -#define VX_VELOX_OPERATOR_GT UINT32_C(2) -#define VX_VELOX_OPERATOR_GTE UINT32_C(3) -#define VX_VELOX_OPERATOR_LT UINT32_C(4) -#define VX_VELOX_OPERATOR_LTE UINT32_C(5) -#define VX_VELOX_OPERATOR_KLEENE_AND UINT32_C(6) -#define VX_VELOX_OPERATOR_KLEENE_OR UINT32_C(7) +#define VX_VELOX_OPERATOR_EQ UINT32_C(0) +#define VX_VELOX_OPERATOR_NOT_EQ UINT32_C(1) +#define VX_VELOX_OPERATOR_GT UINT32_C(2) +#define VX_VELOX_OPERATOR_GTE UINT32_C(3) +#define VX_VELOX_OPERATOR_LT UINT32_C(4) +#define VX_VELOX_OPERATOR_LTE UINT32_C(5) +#define VX_VELOX_OPERATOR_KLEENE_AND UINT32_C(6) +#define VX_VELOX_OPERATOR_KLEENE_OR UINT32_C(7) typedef uint32_t vx_velox_scan_selection_include; #define VX_VELOX_SELECTION_ALL UINT32_C(0) @@ -81,8 +93,8 @@ typedef struct vx_velox_scan_selection { typedef struct vx_velox_scan_options { size_t struct_size; uint32_t abi_version; - const vx_expression *projection; - const vx_expression *filter; + const vx_velox_expression *projection; + const vx_velox_expression *filter; uint64_t row_range_begin; uint64_t row_range_end; vx_velox_scan_selection selection; @@ -144,17 +156,17 @@ typedef struct vx_velox_natural_split { } vx_velox_natural_split; typedef uint32_t vx_velox_primitive_type; -#define VX_VELOX_PRIMITIVE_U8 UINT32_C(0) -#define VX_VELOX_PRIMITIVE_U16 UINT32_C(1) -#define VX_VELOX_PRIMITIVE_U32 UINT32_C(2) -#define VX_VELOX_PRIMITIVE_U64 UINT32_C(3) -#define VX_VELOX_PRIMITIVE_I8 UINT32_C(4) -#define VX_VELOX_PRIMITIVE_I16 UINT32_C(5) -#define VX_VELOX_PRIMITIVE_I32 UINT32_C(6) -#define VX_VELOX_PRIMITIVE_I64 UINT32_C(7) -#define VX_VELOX_PRIMITIVE_F16 UINT32_C(8) -#define VX_VELOX_PRIMITIVE_F32 UINT32_C(9) -#define VX_VELOX_PRIMITIVE_F64 UINT32_C(10) +#define VX_VELOX_PRIMITIVE_U8 UINT32_C(0) +#define VX_VELOX_PRIMITIVE_U16 UINT32_C(1) +#define VX_VELOX_PRIMITIVE_U32 UINT32_C(2) +#define VX_VELOX_PRIMITIVE_U64 UINT32_C(3) +#define VX_VELOX_PRIMITIVE_I8 UINT32_C(4) +#define VX_VELOX_PRIMITIVE_I16 UINT32_C(5) +#define VX_VELOX_PRIMITIVE_I32 UINT32_C(6) +#define VX_VELOX_PRIMITIVE_I64 UINT32_C(7) +#define VX_VELOX_PRIMITIVE_F16 UINT32_C(8) +#define VX_VELOX_PRIMITIVE_F32 UINT32_C(9) +#define VX_VELOX_PRIMITIVE_F64 UINT32_C(10) #define VX_VELOX_PRIMITIVE_I128 UINT32_C(11) typedef uint32_t vx_velox_validity_kind; @@ -431,108 +443,111 @@ typedef struct vx_velox_arrow_memory_callbacks { uint32_t vx_velox_abi_version(void); uint64_t vx_velox_capabilities(void); -vx_view vx_velox_error_message(const vx_error *error); -void vx_velox_error_free(const vx_error *error); -vx_session *vx_velox_session_new(void); -vx_session *vx_velox_session_clone(const vx_session *session); -void vx_velox_session_free(const vx_session *session); - -const vx_dtype *vx_velox_dtype_new_primitive(vx_velox_ptype ptype, - bool nullable, - vx_error **error_out); -void vx_velox_dtype_free(const vx_dtype *dtype); -vx_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); -vx_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); -vx_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); -vx_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); -vx_scalar *vx_velox_scalar_new_date_days(int32_t value, - bool nullable, - vx_error **error_out); -vx_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); -vx_scalar *vx_velox_scalar_new_f32(float value, bool nullable); -vx_scalar *vx_velox_scalar_new_f64(double value, bool nullable); -vx_scalar *vx_velox_scalar_new_utf8(vx_view value, bool nullable, vx_error **error_out); -vx_scalar * -vx_velox_scalar_new_binary(const uint8_t *data, size_t length, bool nullable, vx_error **error_out); -vx_scalar *vx_velox_scalar_new_list(const vx_dtype *element_dtype, - const vx_scalar *const *elements, - size_t length, - bool nullable, - vx_error **error_out); -void vx_velox_scalar_free(const vx_scalar *scalar); - -vx_expression *vx_velox_expression_root(void); -vx_expression *vx_velox_expression_literal(const vx_scalar *scalar, vx_error **error_out); -vx_expression *vx_velox_expression_get_item(vx_view name, const vx_expression *child); -vx_expression *vx_velox_expression_binary(vx_velox_binary_operator operation, - const vx_expression *left, - const vx_expression *right, - vx_error **error_out); -vx_expression *vx_velox_expression_and(const vx_expression *const *expressions, size_t length); -vx_expression *vx_velox_expression_or(const vx_expression *const *expressions, size_t length); -vx_expression *vx_velox_expression_not(const vx_expression *child); -vx_expression *vx_velox_expression_is_null(const vx_expression *child); -vx_expression *vx_velox_expression_list_contains(const vx_expression *list, const vx_expression *value); -bool vx_velox_can_push_down_integer_values(size_t value_count); -void vx_velox_expression_free(const vx_expression *expression); -vx_expression *vx_velox_expression_select(const vx_view *names, +vx_velox_view vx_velox_error_message(const vx_velox_error *error); +void vx_velox_error_free(const vx_velox_error *error); +vx_velox_session *vx_velox_session_new(void); +vx_velox_session *vx_velox_session_clone(const vx_velox_session *session); +void vx_velox_session_free(const vx_velox_session *session); + +const vx_velox_dtype * +vx_velox_dtype_new_primitive(vx_velox_ptype ptype, bool nullable, vx_velox_error **error_out); +void vx_velox_dtype_free(const vx_velox_dtype *dtype); +vx_velox_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_date_days(int32_t value, bool nullable, vx_velox_error **error_out); +vx_velox_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_f32(float value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_f64(double value, bool nullable); +vx_velox_scalar *vx_velox_scalar_new_utf8(vx_velox_view value, bool nullable, vx_velox_error **error_out); +vx_velox_scalar * +vx_velox_scalar_new_binary(const uint8_t *data, size_t length, bool nullable, vx_velox_error **error_out); +vx_velox_scalar *vx_velox_scalar_new_list(const vx_velox_dtype *element_dtype, + const vx_velox_scalar *const *elements, size_t length, - vx_error **error_out); -vx_expression *vx_velox_expression_select_with_row_index(const vx_view *names, - size_t length, - vx_view row_index_name, - vx_error **error_out); + bool nullable, + vx_velox_error **error_out); +void vx_velox_scalar_free(const vx_velox_scalar *scalar); + +vx_velox_expression *vx_velox_expression_root(void); +vx_velox_expression *vx_velox_expression_literal(const vx_velox_scalar *scalar, vx_velox_error **error_out); +vx_velox_expression *vx_velox_expression_get_item(vx_velox_view name, const vx_velox_expression *child); +vx_velox_expression *vx_velox_expression_binary(vx_velox_binary_operator operation, + const vx_velox_expression *left, + const vx_velox_expression *right, + vx_velox_error **error_out); +vx_velox_expression *vx_velox_expression_and(const vx_velox_expression *const *expressions, size_t length); +vx_velox_expression *vx_velox_expression_or(const vx_velox_expression *const *expressions, size_t length); +vx_velox_expression *vx_velox_expression_not(const vx_velox_expression *child); +vx_velox_expression *vx_velox_expression_is_null(const vx_velox_expression *child); +vx_velox_expression *vx_velox_expression_list_contains(const vx_velox_expression *list, + const vx_velox_expression *value); +bool vx_velox_can_push_down_integer_values(size_t value_count); +void vx_velox_expression_free(const vx_velox_expression *expression); +vx_velox_expression * +vx_velox_expression_select(const vx_velox_view *names, size_t length, vx_velox_error **error_out); +vx_velox_expression *vx_velox_expression_select_with_row_index(const vx_velox_view *names, + size_t length, + vx_velox_view row_index_name, + vx_velox_error **error_out); /* * On success, the reader owns context and calls release_context once. On * failure, the caller still owns context. */ -vx_velox_read_at *vx_velox_read_at_new(const vx_velox_read_at_callbacks *callbacks, vx_error **error_out); +vx_velox_read_at *vx_velox_read_at_new(const vx_velox_read_at_callbacks *callbacks, + vx_velox_error **error_out); void vx_velox_read_at_free(vx_velox_read_at *reader); -uint64_t vx_velox_read_at_size(const vx_velox_read_at *reader, vx_error **error_out); +uint64_t vx_velox_read_at_size(const vx_velox_read_at *reader, vx_velox_error **error_out); -vx_velox_source * -vx_velox_source_new(const vx_session *session, const vx_velox_read_at *reader, vx_error **error_out); +vx_velox_source *vx_velox_source_new(const vx_velox_session *session, + const vx_velox_read_at *reader, + vx_velox_error **error_out); void vx_velox_source_free(vx_velox_source *source); uint64_t vx_velox_source_row_count(const vx_velox_source *source); uint64_t vx_velox_source_file_size(const vx_velox_source *source); int32_t vx_velox_source_export_schema(const vx_velox_source *source, - FFI_ArrowSchema *schema_out, - vx_error **error_out); + struct ArrowSchema *schema_out, + vx_velox_error **error_out); size_t vx_velox_source_natural_split_count(const vx_velox_source *source); int32_t vx_velox_source_natural_split_at(const vx_velox_source *source, size_t index, vx_velox_natural_split *split_out, - vx_error **error_out); + vx_velox_error **error_out); int32_t vx_velox_source_prune_natural_splits(const vx_velox_source *source, - const vx_expression *expression, + const vx_velox_expression *expression, size_t first_split, size_t split_count, uint8_t *pruned_out, - vx_error **error_out); -const vx_data_source *vx_velox_source_data_source(const vx_velox_source *source, vx_error **error_out); - -void vx_velox_data_source_free(const vx_data_source *data_source); -vx_scan *vx_velox_data_source_scan(const vx_data_source *data_source, - const vx_velox_scan_options *options, - vx_error **error_out); -void vx_velox_scan_free(const vx_scan *scan); -vx_partition *vx_velox_scan_next_partition(vx_scan *scan, vx_error **error_out); -void vx_velox_partition_free(const vx_partition *partition); -const vx_array *vx_velox_partition_next(vx_partition *partition, vx_error **error_out); -void vx_velox_array_free(const vx_array *array); -size_t vx_velox_array_len(const vx_array *array); -const vx_array *vx_velox_array_slice(const vx_array *array, size_t begin, size_t end, vx_error **error_out); -const vx_array *vx_velox_array_get_field(const vx_session *session, - const vx_array *array, - size_t index, - vx_error **error_out); -size_t vx_velox_array_invalid_count(const vx_session *session, const vx_array *array, vx_error **error_out); -int32_t vx_velox_array_visit(const vx_session *session, - const vx_array *array, + vx_velox_error **error_out); +const vx_velox_data_source *vx_velox_source_data_source(const vx_velox_source *source, + vx_velox_error **error_out); + +void vx_velox_data_source_free(const vx_velox_data_source *data_source); +vx_velox_scan *vx_velox_data_source_scan(const vx_velox_data_source *data_source, + const vx_velox_scan_options *options, + vx_velox_error **error_out); +void vx_velox_scan_free(const vx_velox_scan *scan); +vx_velox_partition *vx_velox_scan_next_partition(vx_velox_scan *scan, vx_velox_error **error_out); +void vx_velox_partition_free(const vx_velox_partition *partition); +const vx_velox_array *vx_velox_partition_next(vx_velox_partition *partition, vx_velox_error **error_out); +void vx_velox_array_free(const vx_velox_array *array); +size_t vx_velox_array_len(const vx_velox_array *array); +const vx_velox_array * +vx_velox_array_slice(const vx_velox_array *array, size_t begin, size_t end, vx_velox_error **error_out); +const vx_velox_array *vx_velox_array_get_field(const vx_velox_session *session, + const vx_velox_array *array, + size_t index, + vx_velox_error **error_out); +size_t vx_velox_array_invalid_count(const vx_velox_session *session, + const vx_velox_array *array, + vx_velox_error **error_out); +int32_t vx_velox_array_visit(const vx_velox_session *session, + const vx_velox_array *array, const vx_velox_visit_request *request, const vx_velox_visitor *visitor, - vx_error **error_out); + vx_velox_error **error_out); /** * Create one prepared exporter for several engine-sized output windows. @@ -540,10 +555,10 @@ int32_t vx_velox_array_visit(const vx_session *session, * memory_callbacks must identify a complete, thread-safe callback table. * The exporter retains its callback context until the last buffer owner releases it. */ -vx_velox_export_cursor *vx_velox_export_cursor_new(const vx_session *session, - const vx_array *array, +vx_velox_export_cursor *vx_velox_export_cursor_new(const vx_velox_session *session, + const vx_velox_array *array, const vx_velox_arrow_memory_callbacks *memory_callbacks, - vx_error **error_out); + vx_velox_error **error_out); /** Free one prepared exporter. */ void vx_velox_export_cursor_free(vx_velox_export_cursor *cursor); @@ -557,14 +572,14 @@ int32_t vx_velox_export_cursor_visit(const vx_velox_export_cursor *cursor, size_t offset, size_t length, const vx_velox_visitor *visitor, - vx_error **error_out); + vx_velox_error **error_out); -int32_t vx_velox_array_export_arrow(const vx_session *session, - const vx_array *array, +int32_t vx_velox_array_export_arrow(const vx_velox_session *session, + const vx_velox_array *array, const vx_velox_arrow_memory_callbacks *memory_callbacks, - FFI_ArrowSchema *schema_out, - FFI_ArrowArray *array_out, - vx_error **error_out); + struct ArrowSchema *schema_out, + struct ArrowArray *array_out, + vx_velox_error **error_out); #ifdef __cplusplus } diff --git a/vortex-velox/cinclude/vortex_velox_test.h b/vortex-velox/cinclude/vortex_velox_test.h new file mode 100644 index 00000000000..48e1f43d895 --- /dev/null +++ b/vortex-velox/cinclude/vortex_velox_test.h @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex_velox.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* This interface exists only for Velox tests and benchmark fixture generation. */ +typedef struct vx_velox_test_writer vx_velox_test_writer; + +const vx_velox_array *vx_velox_test_array_from_arrow_apply(const vx_velox_session *session, + struct ArrowArray *array, + struct ArrowSchema *schema, + const vx_velox_expression *expression, + vx_velox_error **error_out); + +vx_velox_test_writer *vx_velox_test_writer_new(vx_velox_view path, vx_velox_error **error_out); +int32_t vx_velox_test_writer_push(vx_velox_test_writer *writer, + struct ArrowArray *array, + struct ArrowSchema *schema, + vx_velox_error **error_out); +int32_t vx_velox_test_writer_close(vx_velox_test_writer *writer, vx_velox_error **error_out); +void vx_velox_test_writer_abort(vx_velox_test_writer *writer); + +#ifdef __cplusplus +} +#endif diff --git a/vortex-velox/src/api.rs b/vortex-velox/src/api.rs index 944809c53cc..e80d510b14f 100644 --- a/vortex-velox/src/api.rs +++ b/vortex-velox/src/api.rs @@ -20,113 +20,23 @@ use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_error::vortex_bail; -use vortex_ffi::try_or; -use vortex_ffi::vx_array; -use vortex_ffi::vx_data_source; -use vortex_ffi::vx_data_source_scan_with; -use vortex_ffi::vx_dtype; -use vortex_ffi::vx_dtype_new_with; -use vortex_ffi::vx_error; -use vortex_ffi::vx_expression; -use vortex_ffi::vx_expression_new_with; -use vortex_ffi::vx_expression_ref; -use vortex_ffi::vx_partition; -use vortex_ffi::vx_scalar; -use vortex_ffi::vx_scan; -use vortex_ffi::vx_session; -use vortex_ffi::vx_view; - -// The base FFI wrappers are opaque C handles despite their private Rust payloads. -#[allow(improper_ctypes)] -mod ffi { - use super::*; - - unsafe extern "C-unwind" { - pub fn vx_error_message(error: *const vx_error) -> vx_view; - pub fn vx_error_free(error: *const vx_error); - pub fn vx_session_new() -> *mut vx_session; - pub fn vx_session_clone(session: *const vx_session) -> *mut vx_session; - pub fn vx_session_free(session: *const vx_session); - pub fn vx_dtype_free(dtype: *const vx_dtype); - pub fn vx_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_i8(value: i8, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_i16(value: i16, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_i32(value: i32, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_i64(value: i64, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_f32(value: f32, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_f64(value: f64, nullable: bool) -> *mut vx_scalar; - pub fn vx_scalar_new_utf8( - value: vx_view, - nullable: bool, - error_out: *mut *mut vx_error, - ) -> *mut vx_scalar; - pub fn vx_scalar_new_binary( - data: *const u8, - length: usize, - nullable: bool, - error_out: *mut *mut vx_error, - ) -> *mut vx_scalar; - pub fn vx_scalar_new_list( - element_dtype: *const vx_dtype, - elements: *const *const vx_scalar, - length: usize, - nullable: bool, - error_out: *mut *mut vx_error, - ) -> *mut vx_scalar; - pub fn vx_scalar_new_extension( - dtype: *const vx_dtype, - storage: *const vx_scalar, - error_out: *mut *mut vx_error, - ) -> *mut vx_scalar; - pub fn vx_scalar_free(scalar: *const vx_scalar); - pub fn vx_expression_literal( - scalar: *const vx_scalar, - error_out: *mut *mut vx_error, - ) -> *mut vx_expression; - pub fn vx_expression_free(expression: *const vx_expression); - pub fn vx_data_source_free(data_source: *const vx_data_source); - pub fn vx_scan_free(scan: *const vx_scan); - pub fn vx_scan_next_partition( - scan: *mut vx_scan, - error_out: *mut *mut vx_error, - ) -> *mut vx_partition; - pub fn vx_partition_free(partition: *const vx_partition); - pub fn vx_partition_next( - partition: *mut vx_partition, - error_out: *mut *mut vx_error, - ) -> *const vx_array; - pub fn vx_array_free(array: *const vx_array); - pub fn vx_array_len(array: *const vx_array) -> usize; - pub fn vx_array_slice( - array: *const vx_array, - begin: usize, - end: usize, - error_out: *mut *mut vx_error, - ) -> *const vx_array; - } - unsafe extern "C" { - pub fn vx_expression_root() -> *mut vx_expression; - pub fn vx_expression_get_item( - name: vx_view, - child: *const vx_expression, - ) -> *mut vx_expression; - pub fn vx_expression_and( - expressions: *const *const vx_expression, - length: usize, - ) -> *mut vx_expression; - pub fn vx_expression_or( - expressions: *const *const vx_expression, - length: usize, - ) -> *mut vx_expression; - pub fn vx_expression_not(child: *const vx_expression) -> *mut vx_expression; - pub fn vx_expression_is_null(child: *const vx_expression) -> *mut vx_expression; - pub fn vx_expression_list_contains( - list: *const vx_expression, - value: *const vx_expression, - ) -> *mut vx_expression; - } -} +use crate::ffi; +use crate::ffi::try_or; +use crate::ffi::vx_data_source_scan_with; +use crate::ffi::vx_dtype_new_with; +use crate::ffi::vx_expression_new_with; +use crate::ffi::vx_expression_ref; +use crate::ffi::vx_velox_array; +use crate::ffi::vx_velox_data_source; +use crate::ffi::vx_velox_dtype; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_expression; +use crate::ffi::vx_velox_partition; +use crate::ffi::vx_velox_scalar; +use crate::ffi::vx_velox_scan; +use crate::ffi::vx_velox_session; +use crate::ffi::vx_velox_view; /// A fixed-width primitive type identifier for Velox scalar construction. pub type vx_velox_ptype = u32; @@ -202,9 +112,9 @@ pub struct vx_velox_scan_options { /// Set this field to [`crate::VX_VELOX_ABI_VERSION`]. pub abi_version: u32, /// The projected expression, or null for every field. - pub projection: *const vx_expression, + pub projection: *const vx_velox_expression, /// The exact filter expression, or null for no filter. - pub filter: *const vx_expression, + pub filter: *const vx_velox_expression, /// The first row in the scan range. pub row_range_begin: u64, /// One past the final row in the scan range. @@ -233,7 +143,9 @@ impl Default for vx_velox_scan_selection { /// /// `error` must point to a live error handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_error_message(error: *const vx_error) -> vx_view { +pub unsafe extern "C-unwind" fn vx_velox_error_message( + error: *const vx_velox_error, +) -> vx_velox_view { unsafe { ffi::vx_error_message(error) } } @@ -243,14 +155,14 @@ pub unsafe extern "C-unwind" fn vx_velox_error_message(error: *const vx_error) - /// /// `error` must be null or an owned error handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_error_free(error: *const vx_error) { +pub unsafe extern "C-unwind" fn vx_velox_error_free(error: *const vx_velox_error) { unsafe { ffi::vx_error_free(error) }; } /// Create a default Vortex session for Velox. #[unsafe(no_mangle)] -pub extern "C-unwind" fn vx_velox_session_new() -> *mut vx_session { - unsafe { ffi::vx_session_new() } +pub extern "C-unwind" fn vx_velox_session_new() -> *mut vx_velox_session { + ffi::vx_session_new() } /// Clone a Vortex session. @@ -260,8 +172,8 @@ pub extern "C-unwind" fn vx_velox_session_new() -> *mut vx_session { /// `session` must point to a live Vortex session. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_session_clone( - session: *const vx_session, -) -> *mut vx_session { + session: *const vx_velox_session, +) -> *mut vx_velox_session { unsafe { ffi::vx_session_clone(session) } } @@ -271,7 +183,7 @@ pub unsafe extern "C-unwind" fn vx_velox_session_clone( /// /// `session` must be null or an owned session handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_session_free(session: *const vx_session) { +pub unsafe extern "C-unwind" fn vx_velox_session_free(session: *const vx_velox_session) { unsafe { ffi::vx_session_free(session) }; } @@ -301,8 +213,8 @@ fn primitive_type(ptype: vx_velox_ptype) -> vortex_error::VortexResult { pub unsafe extern "C-unwind" fn vx_velox_dtype_new_primitive( ptype: vx_velox_ptype, nullable: bool, - error_out: *mut *mut vx_error, -) -> *const vx_dtype { + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_dtype { try_or(error_out, ptr::null(), || { Ok(vx_dtype_new_with(DType::Primitive( primitive_type(ptype)?, @@ -317,13 +229,16 @@ pub unsafe extern "C-unwind" fn vx_velox_dtype_new_primitive( /// /// `dtype` must be null or an owned dtype handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_dtype_free(dtype: *const vx_dtype) { +pub unsafe extern "C-unwind" fn vx_velox_dtype_free(dtype: *const vx_velox_dtype) { unsafe { ffi::vx_dtype_free(dtype) }; } /// Create a Boolean scalar. #[unsafe(no_mangle)] -pub extern "C-unwind" fn vx_velox_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_scalar { +pub extern "C-unwind" fn vx_velox_scalar_new_bool( + value: bool, + nullable: bool, +) -> *mut vx_velox_scalar { unsafe { ffi::vx_scalar_new_bool(value, nullable) } } @@ -331,7 +246,7 @@ macro_rules! scalar_primitive_wrapper { ($name:ident, $source:ident, $type:ty, $description:literal) => { #[doc = $description] #[unsafe(no_mangle)] - pub extern "C-unwind" fn $name(value: $type, nullable: bool) -> *mut vx_scalar { + pub extern "C-unwind" fn $name(value: $type, nullable: bool) -> *mut vx_velox_scalar { unsafe { ffi::$source(value, nullable) } } }; @@ -365,8 +280,8 @@ scalar_primitive_wrapper!( pub unsafe extern "C-unwind" fn vx_velox_scalar_new_date_days( value: i32, nullable: bool, - error_out: *mut *mut vx_error, -) -> *mut vx_scalar { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { let dtype = vx_dtype_new_with(DType::Extension( Date::new(TimeUnit::Days, Nullability::from(nullable)).erased(), )); @@ -404,10 +319,10 @@ scalar_primitive_wrapper!( /// `value` and `error_out` must satisfy the adapter header contract. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_scalar_new_utf8( - value: vx_view, + value: vx_velox_view, nullable: bool, - error_out: *mut *mut vx_error, -) -> *mut vx_scalar { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { unsafe { ffi::vx_scalar_new_utf8(value, nullable, error_out) } } @@ -421,8 +336,8 @@ pub unsafe extern "C-unwind" fn vx_velox_scalar_new_binary( data: *const u8, length: usize, nullable: bool, - error_out: *mut *mut vx_error, -) -> *mut vx_scalar { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { unsafe { ffi::vx_scalar_new_binary(data, length, nullable, error_out) } } @@ -433,12 +348,12 @@ pub unsafe extern "C-unwind" fn vx_velox_scalar_new_binary( /// Every pointer must satisfy the adapter header contract. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_scalar_new_list( - element_dtype: *const vx_dtype, - elements: *const *const vx_scalar, + element_dtype: *const vx_velox_dtype, + elements: *const *const vx_velox_scalar, length: usize, nullable: bool, - error_out: *mut *mut vx_error, -) -> *mut vx_scalar { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { unsafe { ffi::vx_scalar_new_list(element_dtype, elements, length, nullable, error_out) } } @@ -448,7 +363,7 @@ pub unsafe extern "C-unwind" fn vx_velox_scalar_new_list( /// /// `scalar` must be null or an owned scalar handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_scalar_free(scalar: *const vx_scalar) { +pub unsafe extern "C-unwind" fn vx_velox_scalar_free(scalar: *const vx_velox_scalar) { unsafe { ffi::vx_scalar_free(scalar) }; } @@ -459,16 +374,16 @@ pub unsafe extern "C-unwind" fn vx_velox_scalar_free(scalar: *const vx_scalar) { /// `scalar` must point to a live scalar. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_expression_literal( - scalar: *const vx_scalar, - error_out: *mut *mut vx_error, -) -> *mut vx_expression { + scalar: *const vx_velox_scalar, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_literal(scalar, error_out) } } /// Create a root expression. #[unsafe(no_mangle)] -pub extern "C" fn vx_velox_expression_root() -> *mut vx_expression { - unsafe { ffi::vx_expression_root() } +pub extern "C" fn vx_velox_expression_root() -> *mut vx_velox_expression { + ffi::vx_expression_root() } /// Create a field expression. @@ -478,9 +393,9 @@ pub extern "C" fn vx_velox_expression_root() -> *mut vx_expression { /// `child` must point to a live expression. `name` must identify valid UTF-8. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_get_item( - name: vx_view, - child: *const vx_expression, -) -> *mut vx_expression { + name: vx_velox_view, + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_get_item(name, child) } } @@ -506,10 +421,10 @@ fn binary_operator(operator: vx_velox_binary_operator) -> vortex_error::VortexRe #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_expression_binary( operator: vx_velox_binary_operator, - left: *const vx_expression, - right: *const vx_expression, - error_out: *mut *mut vx_error, -) -> *mut vx_expression { + left: *const vx_velox_expression, + right: *const vx_velox_expression, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_expression { try_or(error_out, ptr::null_mut(), || { let operator = binary_operator(operator)?; let left = unsafe { vx_expression_ref(left)? }.clone(); @@ -527,9 +442,9 @@ pub unsafe extern "C-unwind" fn vx_velox_expression_binary( /// `expressions` must identify `length` live expression pointers. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_and( - expressions: *const *const vx_expression, + expressions: *const *const vx_velox_expression, length: usize, -) -> *mut vx_expression { +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_and(expressions, length) } } @@ -540,9 +455,9 @@ pub unsafe extern "C" fn vx_velox_expression_and( /// `expressions` must identify `length` live expression pointers. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_or( - expressions: *const *const vx_expression, + expressions: *const *const vx_velox_expression, length: usize, -) -> *mut vx_expression { +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_or(expressions, length) } } @@ -553,8 +468,8 @@ pub unsafe extern "C" fn vx_velox_expression_or( /// `child` must point to a live expression. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_not( - child: *const vx_expression, -) -> *mut vx_expression { + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_not(child) } } @@ -565,8 +480,8 @@ pub unsafe extern "C" fn vx_velox_expression_not( /// `child` must point to a live expression. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_is_null( - child: *const vx_expression, -) -> *mut vx_expression { + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_is_null(child) } } @@ -577,9 +492,9 @@ pub unsafe extern "C" fn vx_velox_expression_is_null( /// Both operands must point to live expressions. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_velox_expression_list_contains( - list: *const vx_expression, - value: *const vx_expression, -) -> *mut vx_expression { + list: *const vx_velox_expression, + value: *const vx_velox_expression, +) -> *mut vx_velox_expression { unsafe { ffi::vx_expression_list_contains(list, value) } } @@ -595,7 +510,7 @@ pub extern "C" fn vx_velox_can_push_down_integer_values(value_count: usize) -> b /// /// `expression` must be null or an owned expression handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_expression_free(expression: *const vx_expression) { +pub unsafe extern "C-unwind" fn vx_velox_expression_free(expression: *const vx_velox_expression) { unsafe { ffi::vx_expression_free(expression) }; } @@ -605,7 +520,9 @@ pub unsafe extern "C-unwind" fn vx_velox_expression_free(expression: *const vx_e /// /// `data_source` must be null or an owned data-source handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_data_source_free(data_source: *const vx_data_source) { +pub unsafe extern "C-unwind" fn vx_velox_data_source_free( + data_source: *const vx_velox_data_source, +) { unsafe { ffi::vx_data_source_free(data_source) }; } @@ -684,10 +601,10 @@ unsafe fn scan_options(options: &vx_velox_scan_options) -> vortex_error::VortexR /// Every pointer must satisfy the adapter header contract. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_data_source_scan( - data_source: *const vx_data_source, + data_source: *const vx_velox_data_source, options: *const vx_velox_scan_options, - error_out: *mut *mut vx_error, -) -> *mut vx_scan { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scan { try_or(error_out, ptr::null_mut(), || { let request = if options.is_null() { ScanRequest::default() @@ -704,7 +621,7 @@ pub unsafe extern "C-unwind" fn vx_velox_data_source_scan( /// /// `scan` must be null or an owned scan handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_scan_free(scan: *const vx_scan) { +pub unsafe extern "C-unwind" fn vx_velox_scan_free(scan: *const vx_velox_scan) { unsafe { ffi::vx_scan_free(scan) }; } @@ -715,9 +632,9 @@ pub unsafe extern "C-unwind" fn vx_velox_scan_free(scan: *const vx_scan) { /// `scan` must point to a live scan. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_scan_next_partition( - scan: *mut vx_scan, - error_out: *mut *mut vx_error, -) -> *mut vx_partition { + scan: *mut vx_velox_scan, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_partition { unsafe { ffi::vx_scan_next_partition(scan, error_out) } } @@ -727,7 +644,7 @@ pub unsafe extern "C-unwind" fn vx_velox_scan_next_partition( /// /// `partition` must be null or an owned partition handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_partition_free(partition: *const vx_partition) { +pub unsafe extern "C-unwind" fn vx_velox_partition_free(partition: *const vx_velox_partition) { unsafe { ffi::vx_partition_free(partition) }; } @@ -738,9 +655,9 @@ pub unsafe extern "C-unwind" fn vx_velox_partition_free(partition: *const vx_par /// `partition` must point to a live partition. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_partition_next( - partition: *mut vx_partition, - error_out: *mut *mut vx_error, -) -> *const vx_array { + partition: *mut vx_velox_partition, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { unsafe { ffi::vx_partition_next(partition, error_out) } } @@ -750,7 +667,7 @@ pub unsafe extern "C-unwind" fn vx_velox_partition_next( /// /// `array` must be null or an owned array handle. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_array_free(array: *const vx_array) { +pub unsafe extern "C-unwind" fn vx_velox_array_free(array: *const vx_velox_array) { unsafe { ffi::vx_array_free(array) }; } @@ -760,7 +677,7 @@ pub unsafe extern "C-unwind" fn vx_velox_array_free(array: *const vx_array) { /// /// `array` must point to a live array. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_array_len(array: *const vx_array) -> usize { +pub unsafe extern "C-unwind" fn vx_velox_array_len(array: *const vx_velox_array) -> usize { unsafe { ffi::vx_array_len(array) } } @@ -771,11 +688,11 @@ pub unsafe extern "C-unwind" fn vx_velox_array_len(array: *const vx_array) -> us /// `array` must point to a live array. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_array_slice( - array: *const vx_array, + array: *const vx_velox_array, begin: usize, end: usize, - error_out: *mut *mut vx_error, -) -> *const vx_array { + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { unsafe { ffi::vx_array_slice(array, begin, end, error_out) } } @@ -840,7 +757,7 @@ mod tests { let list_literal = unsafe { vx_velox_expression_literal(list, &raw mut error) }; assert!(error.is_null()); let root = vx_velox_expression_root(); - let name = vx_view { + let name = vx_velox_view { ptr: c"value".as_ptr(), len: 5, }; diff --git a/vortex-velox/src/array.rs b/vortex-velox/src/array.rs index 541ebe58452..b141c49fc3a 100644 --- a/vortex-velox/src/array.rs +++ b/vortex-velox/src/array.rs @@ -24,14 +24,14 @@ use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_ffi::try_or; -use vortex_ffi::vx_array; -use vortex_ffi::vx_array_new_with; -use vortex_ffi::vx_array_ref; -use vortex_ffi::vx_error; -use vortex_ffi::vx_session; -use vortex_ffi::vx_session_ref; +use crate::ffi::try_or; +use crate::ffi::vx_array_new_with; +use crate::ffi::vx_array_ref; +use crate::ffi::vx_session_ref; +use crate::ffi::vx_velox_array; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_session; use crate::temporal::validate_velox_arrow_data; /// Host memory callbacks for one Arrow C Data export. @@ -358,11 +358,11 @@ pub(crate) fn conservative_export_reservation( /// The session and array pointers must identify live handles. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_array_get_field( - session: *const vx_session, - array: *const vx_array, + session: *const vx_velox_session, + array: *const vx_velox_array, index: usize, - error_out: *mut *mut vx_error, -) -> *const vx_array { + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { try_or(error_out, ptr::null(), || { let session = unsafe { vx_session_ref(session)? }; let array = unsafe { vx_array_ref(array)? }; @@ -383,9 +383,9 @@ pub unsafe extern "C-unwind" fn vx_velox_array_get_field( /// The session and array pointers must identify live handles. `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_array_invalid_count( - session: *const vx_session, - array: *const vx_array, - error_out: *mut *mut vx_error, + session: *const vx_velox_session, + array: *const vx_velox_array, + error_out: *mut *mut vx_velox_error, ) -> usize { try_or(error_out, 0, || { let session = unsafe { vx_session_ref(session)? }; @@ -410,12 +410,12 @@ pub unsafe extern "C-unwind" fn vx_velox_array_invalid_count( /// `error_out` must be null or identify writable storage for one error pointer. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_array_export_arrow( - session: *const vx_session, - array: *const vx_array, + session: *const vx_velox_session, + array: *const vx_velox_array, memory_callbacks: *const vx_velox_arrow_memory_callbacks, schema_out: *mut FFI_ArrowSchema, array_out: *mut FFI_ArrowArray, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let session = unsafe { vx_session_ref(session)? }; @@ -481,13 +481,13 @@ mod tests { use vortex::array::arrays::StructArray; use vortex::array::validity::Validity; use vortex_error::VortexResult; - use vortex_ffi::vx_array_new_with; - use vortex_ffi::vx_error_free; - use vortex_ffi::vx_session_free; - use vortex_ffi::vx_session_new_with; use super::*; use crate::api::vx_velox_array_free; + use crate::ffi::vx_array_new_with; + use crate::ffi::vx_error_free; + use crate::ffi::vx_session_free; + use crate::ffi::vx_session_new_with; #[derive(Default)] struct MemoryCapture { diff --git a/vortex-velox/src/ffi.rs b/vortex-velox/src/ffi.rs new file mode 100644 index 00000000000..ee159e4aafd --- /dev/null +++ b/vortex-velox/src/ffi.rs @@ -0,0 +1,550 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Private C-handle support for the Velox adapter. + +use std::any::Any; +use std::ffi::c_char; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::ptr; +use std::slice; +use std::sync::Arc; +use std::sync::LazyLock; + +use futures::StreamExt; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::stream::SendableArrayStream; +use vortex::dtype::DType; +use vortex::dtype::FieldName; +use vortex::dtype::Nullability; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::expr::Expression; +use vortex::expr::and_collect; +use vortex::expr::get_item; +use vortex::expr::is_null; +use vortex::expr::list_contains; +use vortex::expr::lit; +use vortex::expr::not; +use vortex::expr::or_collect; +use vortex::expr::root; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::current::CurrentThreadRuntime; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::scan::multi::MultiLayoutDataSource; +use vortex::scalar::Scalar; +use vortex::scalar::ScalarValue; +use vortex::scan::DataSource; +use vortex::scan::DataSourceScanRef; +use vortex::scan::PartitionRef; +use vortex::scan::PartitionStream; +use vortex::scan::ScanRequest; +use vortex::session::VortexSession; + +static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); + +pub(crate) fn ffi_runtime() -> &'static CurrentThreadRuntime { + &RUNTIME +} + +macro_rules! ffi_handle { + ($name:ident, $inner:ty, $free:ident) => { + #[repr(transparent)] + pub struct $name($inner); + + #[allow(dead_code)] + impl $name { + pub(crate) fn new(value: $inner) -> *mut Self { + Box::into_raw(Box::new(Self(value))) + } + + pub(crate) unsafe fn as_ref<'a>(pointer: *const Self) -> &'a $inner { + // SAFETY: Callers validate pointer ownership and lifetime at the C boundary. + &unsafe { &*pointer }.0 + } + + pub(crate) unsafe fn as_mut<'a>(pointer: *mut Self) -> &'a mut $inner { + // SAFETY: Callers validate unique pointer ownership at the C boundary. + &mut unsafe { &mut *pointer }.0 + } + } + + pub(crate) unsafe fn $free(pointer: *const $name) { + if !pointer.is_null() { + // SAFETY: The caller transfers one handle created by `new`. + drop(unsafe { Box::from_raw(pointer.cast_mut()) }); + } + } + }; +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_velox_view { + pub ptr: *const c_char, + pub len: usize, +} + +impl vx_velox_view { + pub(crate) fn from_str(value: &str) -> Self { + Self { + ptr: value.as_ptr().cast(), + len: value.len(), + } + } + + pub(crate) unsafe fn as_bytes<'a>(&self) -> VortexResult<&'a [u8]> { + if self.ptr.is_null() { + vortex_ensure!(self.len == 0, "null view pointer with non-zero length"); + return Ok(&[]); + } + // SAFETY: The caller provides `len` readable bytes. + Ok(unsafe { slice::from_raw_parts(self.ptr.cast(), self.len) }) + } + + pub(crate) unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> { + std::str::from_utf8(unsafe { self.as_bytes() }?) + .map_err(|error| vortex_err!("invalid UTF-8: {error}")) + } +} + +pub(crate) struct AdapterError { + message: Arc, +} + +ffi_handle!(vx_velox_error, AdapterError, vx_error_free); +ffi_handle!(vx_velox_session, VortexSession, vx_session_free); +ffi_handle!(vx_velox_dtype, DType, vx_dtype_free); +ffi_handle!(vx_velox_scalar, Scalar, vx_scalar_free); +ffi_handle!(vx_velox_expression, Expression, vx_expression_free); +ffi_handle!( + vx_velox_data_source, + MultiLayoutDataSource, + vx_data_source_free +); +ffi_handle!(vx_velox_array, ArrayRef, vx_array_free); + +pub(crate) enum ScanState { + Pending(DataSourceScanRef), + Started(PartitionStream), + Finished, +} + +ffi_handle!(vx_velox_scan, ScanState, vx_scan_free); + +pub(crate) enum PartitionState { + Pending(PartitionRef), + Started(SendableArrayStream), + Finished, +} + +ffi_handle!(vx_velox_partition, PartitionState, vx_partition_free); + +fn clear_error(error_out: *mut *mut vx_velox_error) { + if !error_out.is_null() { + // SAFETY: The caller provides writable storage for one pointer. + unsafe { error_out.write(ptr::null_mut()) }; + } +} + +fn write_error(error_out: *mut *mut vx_velox_error, message: impl Into>) { + if !error_out.is_null() { + // SAFETY: The caller provides writable storage for one pointer. + unsafe { + error_out.write(vx_velox_error::new(AdapterError { + message: message.into(), + })) + }; + } +} + +fn panic_message(payload: &(dyn Any + Send)) -> String { + if let Some(message) = payload.downcast_ref::<&str>() { + format!("panic in Vortex Velox adapter: {message}") + } else if let Some(message) = payload.downcast_ref::() { + format!("panic in Vortex Velox adapter: {message}") + } else { + "panic in Vortex Velox adapter".to_string() + } +} + +pub(crate) fn try_or( + error_out: *mut *mut vx_velox_error, + error_value: T, + function: impl FnOnce() -> VortexResult, +) -> T { + match catch_unwind(AssertUnwindSafe(function)) { + Ok(Ok(value)) => { + clear_error(error_out); + value + } + Ok(Err(error)) => { + write_error(error_out, error.to_string()); + error_value + } + Err(payload) => { + write_error(error_out, panic_message(payload.as_ref())); + error_value + } + } +} + +pub(crate) unsafe fn vx_error_message(error: *const vx_velox_error) -> vx_velox_view { + vx_velox_view::from_str(&unsafe { vx_velox_error::as_ref(error) }.message) +} + +pub(crate) fn vx_session_new() -> *mut vx_velox_session { + vx_velox_session::new(VortexSession::default().with_handle(RUNTIME.handle())) +} + +pub(crate) unsafe fn vx_session_clone(session: *const vx_velox_session) -> *mut vx_velox_session { + vx_velox_session::new(unsafe { vx_velox_session::as_ref(session) }.clone()) +} + +pub(crate) unsafe fn vx_session_ref<'a>( + session: *const vx_velox_session, +) -> VortexResult<&'a VortexSession> { + vortex_ensure!(!session.is_null(), "Vortex Velox session must not be null"); + Ok(unsafe { vx_velox_session::as_ref(session) }) +} + +#[cfg(test)] +pub(crate) fn vx_session_new_with( + configure: impl FnOnce(VortexSession) -> VortexSession, +) -> *mut vx_velox_session { + vx_velox_session::new(configure( + VortexSession::default().with_handle(RUNTIME.handle()), + )) +} + +pub(crate) fn vx_dtype_new_with(dtype: DType) -> *const vx_velox_dtype { + vx_velox_dtype::new(dtype) +} + +pub(crate) unsafe fn vx_scalar_new_bool(value: bool, nullable: bool) -> *mut vx_velox_scalar { + vx_velox_scalar::new(Scalar::bool(value, Nullability::from(nullable))) +} + +macro_rules! scalar_primitive { + ($name:ident, $type:ty) => { + pub(crate) unsafe fn $name(value: $type, nullable: bool) -> *mut vx_velox_scalar { + vx_velox_scalar::new(Scalar::primitive(value, Nullability::from(nullable))) + } + }; +} + +scalar_primitive!(vx_scalar_new_i8, i8); +scalar_primitive!(vx_scalar_new_i16, i16); +scalar_primitive!(vx_scalar_new_i32, i32); +scalar_primitive!(vx_scalar_new_i64, i64); +scalar_primitive!(vx_scalar_new_f32, f32); +scalar_primitive!(vx_scalar_new_f64, f64); + +pub(crate) unsafe fn vx_scalar_new_utf8( + value: vx_velox_view, + nullable: bool, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { + try_or(error_out, ptr::null_mut(), || { + Ok(vx_velox_scalar::new(Scalar::utf8( + unsafe { value.as_str() }?.to_owned(), + Nullability::from(nullable), + ))) + }) +} + +pub(crate) unsafe fn vx_scalar_new_binary( + data: *const u8, + length: usize, + nullable: bool, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { + try_or(error_out, ptr::null_mut(), || { + if length != 0 { + vortex_ensure!(!data.is_null(), "binary data pointer must not be null"); + } + let bytes = if length == 0 { + &[] + } else { + // SAFETY: The caller provides `length` readable bytes. + unsafe { slice::from_raw_parts(data, length) } + }; + Ok(vx_velox_scalar::new(Scalar::binary( + bytes.to_vec(), + Nullability::from(nullable), + ))) + }) +} + +pub(crate) unsafe fn vx_scalar_new_list( + element_dtype: *const vx_velox_dtype, + elements: *const *const vx_velox_scalar, + length: usize, + nullable: bool, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { + try_or(error_out, ptr::null_mut(), || { + vortex_ensure!( + !element_dtype.is_null(), + "list element dtype must not be null" + ); + if length != 0 { + vortex_ensure!(!elements.is_null(), "list elements must not be null"); + } + let values = if length == 0 { + Vec::new() + } else { + unsafe { slice::from_raw_parts(elements, length) } + .iter() + .enumerate() + .map(|(index, scalar)| { + vortex_ensure!(!scalar.is_null(), "list scalar {index} must not be null"); + Ok(unsafe { vx_velox_scalar::as_ref(*scalar) } + .clone() + .into_value()) + }) + .collect::>>()? + }; + Ok(vx_velox_scalar::new(Scalar::try_new( + DType::List( + Arc::new(unsafe { vx_velox_dtype::as_ref(element_dtype) }.clone()), + Nullability::from(nullable), + ), + Some(ScalarValue::Tuple(values)), + )?)) + }) +} + +pub(crate) unsafe fn vx_scalar_new_extension( + dtype: *const vx_velox_dtype, + storage: *const vx_velox_scalar, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_scalar { + try_or(error_out, ptr::null_mut(), || { + vortex_ensure!(!dtype.is_null(), "extension dtype must not be null"); + vortex_ensure!(!storage.is_null(), "extension storage must not be null"); + let dtype = unsafe { vx_velox_dtype::as_ref(dtype) }; + let storage = unsafe { vx_velox_scalar::as_ref(storage) }; + let DType::Extension(extension) = dtype else { + vortex_bail!("dtype is not an extension type: {dtype}"); + }; + vortex_ensure!( + storage + .dtype() + .eq_ignore_nullability(extension.storage_dtype()), + "storage dtype {} does not match extension storage dtype {}", + storage.dtype(), + extension.storage_dtype() + ); + Ok(vx_velox_scalar::new(Scalar::try_new( + dtype.clone(), + storage.value().cloned(), + )?)) + }) +} + +pub(crate) unsafe fn vx_expression_literal( + scalar: *const vx_velox_scalar, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_expression { + try_or(error_out, ptr::null_mut(), || { + vortex_ensure!(!scalar.is_null(), "literal scalar must not be null"); + Ok(vx_velox_expression::new(lit(unsafe { + vx_velox_scalar::as_ref(scalar) + } + .clone()))) + }) +} + +pub(crate) fn vx_expression_new_with(expression: Expression) -> *mut vx_velox_expression { + vx_velox_expression::new(expression) +} + +pub(crate) unsafe fn vx_expression_ref<'a>( + expression: *const vx_velox_expression, +) -> VortexResult<&'a Expression> { + vortex_ensure!( + !expression.is_null(), + "Vortex Velox expression must not be null" + ); + Ok(unsafe { vx_velox_expression::as_ref(expression) }) +} + +pub(crate) fn vx_expression_root() -> *mut vx_velox_expression { + vx_velox_expression::new(root()) +} + +pub(crate) unsafe fn vx_expression_get_item( + name: vx_velox_view, + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { + let Ok(name) = (unsafe { name.as_str() }) else { + return ptr::null_mut(); + }; + vx_velox_expression::new(get_item( + FieldName::from(name), + unsafe { vx_velox_expression::as_ref(child) }.clone(), + )) +} + +unsafe fn expression_slice<'a>( + expressions: *const *const vx_velox_expression, + length: usize, +) -> &'a [*const vx_velox_expression] { + if length == 0 { + &[] + } else { + // SAFETY: The caller provides `length` expression pointers. + unsafe { slice::from_raw_parts(expressions, length) } + } +} + +pub(crate) unsafe fn vx_expression_and( + expressions: *const *const vx_velox_expression, + length: usize, +) -> *mut vx_velox_expression { + and_collect( + unsafe { expression_slice(expressions, length) } + .iter() + .map(|expression| unsafe { vx_velox_expression::as_ref(*expression) }.clone()), + ) + .map_or(ptr::null_mut(), vx_velox_expression::new) +} + +pub(crate) unsafe fn vx_expression_or( + expressions: *const *const vx_velox_expression, + length: usize, +) -> *mut vx_velox_expression { + or_collect( + unsafe { expression_slice(expressions, length) } + .iter() + .map(|expression| unsafe { vx_velox_expression::as_ref(*expression) }.clone()), + ) + .map_or(ptr::null_mut(), vx_velox_expression::new) +} + +pub(crate) unsafe fn vx_expression_not( + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { + vx_velox_expression::new(not(unsafe { vx_velox_expression::as_ref(child) }.clone())) +} + +pub(crate) unsafe fn vx_expression_is_null( + child: *const vx_velox_expression, +) -> *mut vx_velox_expression { + vx_velox_expression::new(is_null( + unsafe { vx_velox_expression::as_ref(child) }.clone(), + )) +} + +pub(crate) unsafe fn vx_expression_list_contains( + list: *const vx_velox_expression, + value: *const vx_velox_expression, +) -> *mut vx_velox_expression { + vx_velox_expression::new(list_contains( + unsafe { vx_velox_expression::as_ref(list) }.clone(), + unsafe { vx_velox_expression::as_ref(value) }.clone(), + )) +} + +pub(crate) fn vx_data_source_new_with( + data_source: MultiLayoutDataSource, +) -> *const vx_velox_data_source { + vx_velox_data_source::new(data_source) +} + +pub(crate) unsafe fn vx_data_source_scan_with( + data_source: *const vx_velox_data_source, + request: ScanRequest, +) -> VortexResult<*mut vx_velox_scan> { + vortex_ensure!( + !data_source.is_null(), + "Vortex Velox data source must not be null" + ); + RUNTIME.block_on(async { + let scan = unsafe { vx_velox_data_source::as_ref(data_source) } + .scan(request) + .await?; + Ok(vx_velox_scan::new(ScanState::Pending(scan))) + }) +} + +pub(crate) unsafe fn vx_scan_next_partition( + scan: *mut vx_velox_scan, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_partition { + try_or(error_out, ptr::null_mut(), || { + vortex_ensure!(!scan.is_null(), "Vortex Velox scan must not be null"); + let scan = unsafe { vx_velox_scan::as_mut(scan) }; + let state = std::mem::replace(scan, ScanState::Finished); + let mut stream = match state { + ScanState::Pending(scan) => scan.partitions(), + ScanState::Started(stream) => stream, + ScanState::Finished => return Ok(ptr::null_mut()), + }; + match RUNTIME.block_on(stream.next()) { + Some(partition) => { + *scan = ScanState::Started(stream); + Ok(vx_velox_partition::new(PartitionState::Pending(partition?))) + } + None => Ok(ptr::null_mut()), + } + }) +} + +pub(crate) unsafe fn vx_partition_next( + partition: *mut vx_velox_partition, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { + try_or(error_out, ptr::null(), || { + vortex_ensure!( + !partition.is_null(), + "Vortex Velox partition must not be null" + ); + let partition = unsafe { vx_velox_partition::as_mut(partition) }; + let state = std::mem::replace(partition, PartitionState::Finished); + let mut stream = match state { + PartitionState::Pending(partition) => partition.execute()?, + PartitionState::Started(stream) => stream, + PartitionState::Finished => return Ok(ptr::null()), + }; + match RUNTIME.block_on(stream.next()) { + Some(array) => { + *partition = PartitionState::Started(stream); + Ok(vx_velox_array::new(array?)) + } + None => Ok(ptr::null()), + } + }) +} + +pub(crate) fn vx_array_new_with(array: ArrayRef) -> *const vx_velox_array { + vx_velox_array::new(array) +} + +pub(crate) unsafe fn vx_array_ref<'a>(array: *const vx_velox_array) -> VortexResult<&'a ArrayRef> { + vortex_ensure!(!array.is_null(), "Vortex Velox array must not be null"); + Ok(unsafe { vx_velox_array::as_ref(array) }) +} + +pub(crate) unsafe fn vx_array_len(array: *const vx_velox_array) -> usize { + unsafe { vx_velox_array::as_ref(array) }.len() +} + +pub(crate) unsafe fn vx_array_slice( + array: *const vx_velox_array, + begin: usize, + end: usize, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { + try_or(error_out, ptr::null(), || { + let array = unsafe { vx_array_ref(array) }?; + vortex_ensure!(begin <= end, "array slice begin exceeds end"); + vortex_ensure!(end <= array.len(), "array slice end exceeds array length"); + Ok(vx_velox_array::new(array.slice(begin..end)?)) + }) +} diff --git a/vortex-velox/src/lib.rs b/vortex-velox/src/lib.rs index a745c27b026..3ead79e66f2 100644 --- a/vortex-velox/src/lib.rs +++ b/vortex-velox/src/lib.rs @@ -10,11 +10,13 @@ mod api; mod array; +mod ffi; mod projection; mod read_at; mod schema; mod source; mod temporal; +mod test_support; mod visitor; pub use api::*; diff --git a/vortex-velox/src/projection.rs b/vortex-velox/src/projection.rs index 8ed84369a0c..3a524a2bedc 100644 --- a/vortex-velox/src/projection.rs +++ b/vortex-velox/src/projection.rs @@ -12,13 +12,14 @@ use vortex::expr::get_item; use vortex::expr::pack; use vortex::expr::root; use vortex::layout::layouts::row_idx::row_idx; -use vortex_ffi::try_or; -use vortex_ffi::vx_error; -use vortex_ffi::vx_expression; -use vortex_ffi::vx_expression_new_with; -use vortex_ffi::vx_view; -unsafe fn field_names(names: *const vx_view, len: usize) -> VortexResult> { +use crate::ffi::try_or; +use crate::ffi::vx_expression_new_with; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_expression; +use crate::ffi::vx_velox_view; + +unsafe fn field_names(names: *const vx_velox_view, len: usize) -> VortexResult> { let names = if names.is_null() { vortex_ensure!(len == 0, "null field names pointer with non-zero length"); &[] @@ -36,7 +37,10 @@ unsafe fn field_names(names: *const vx_view, len: usize) -> VortexResult VortexResult<*mut vx_expression> { +unsafe fn projection( + names: *const vx_velox_view, + len: usize, +) -> VortexResult<*mut vx_velox_expression> { let names = unsafe { field_names(names, len) }?; let fields = names .into_iter() @@ -49,10 +53,10 @@ unsafe fn projection(names: *const vx_view, len: usize) -> VortexResult<*mut vx_ } unsafe fn projection_with_row_index( - names: *const vx_view, + names: *const vx_velox_view, len: usize, - row_index_name: vx_view, -) -> VortexResult<*mut vx_expression> { + row_index_name: vx_velox_view, +) -> VortexResult<*mut vx_velox_expression> { vortex_ensure!(!row_index_name.ptr.is_null() || row_index_name.len == 0); // SAFETY: The caller keeps this view valid for the duration of this call. let row_index_name = unsafe { row_index_name.as_str() }?; @@ -88,10 +92,10 @@ unsafe fn projection_with_row_index( /// `error_out` must be null or point to writable storage. No input operation can unwind. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_expression_select( - names: *const vx_view, + names: *const vx_velox_view, len: usize, - error_out: *mut *mut vx_error, -) -> *mut vx_expression { + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_expression { try_or(error_out, ptr::null_mut(), || unsafe { projection(names, len) }) @@ -109,11 +113,11 @@ pub unsafe extern "C-unwind" fn vx_velox_expression_select( /// `error_out` must be null or point to writable storage. No input operation can unwind. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_expression_select_with_row_index( - names: *const vx_view, + names: *const vx_velox_view, len: usize, - row_index_name: vx_view, - error_out: *mut *mut vx_error, -) -> *mut vx_expression { + row_index_name: vx_velox_view, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_expression { try_or(error_out, ptr::null_mut(), || unsafe { projection_with_row_index(names, len, row_index_name) }) @@ -124,8 +128,8 @@ mod tests { use super::*; use crate::api::vx_velox_expression_free; - fn view(value: &str) -> vx_view { - vx_view { + fn view(value: &str) -> vx_velox_view { + vx_velox_view { ptr: value.as_ptr().cast(), len: value.len(), } @@ -173,6 +177,6 @@ mod tests { }; assert!(expression.is_null()); assert!(!error.is_null()); - unsafe { vortex_ffi::vx_error_free(error) }; + unsafe { crate::ffi::vx_error_free(error) }; } } diff --git a/vortex-velox/src/read_at.rs b/vortex-velox/src/read_at.rs index 89098dc5ce7..19995da0fb4 100644 --- a/vortex-velox/src/read_at.rs +++ b/vortex-velox/src/read_at.rs @@ -19,14 +19,15 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_ffi::ffi_runtime; -use vortex_ffi::try_or; -use vortex_ffi::vx_error; use vortex_io::ReadAtRequest; use vortex_io::ReadAtStream; use vortex_io::VortexReadAt; use vortex_io::runtime::BlockingRuntime; +use crate::ffi::ffi_runtime; +use crate::ffi::try_or; +use crate::ffi::vx_velox_error; + /// A positional read request passed to the Velox callback. #[repr(C)] #[derive(Clone, Copy, Debug, Default)] @@ -410,7 +411,7 @@ pub struct vx_velox_read_at(CallbackReadAt); #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_read_at_new( callbacks: *const vx_velox_read_at_callbacks, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> *mut vx_velox_read_at { try_or(error_out, std::ptr::null_mut(), || { let callbacks = unsafe { @@ -447,7 +448,7 @@ pub unsafe extern "C" fn vx_velox_read_at_free(reader: *mut vx_velox_read_at) { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_read_at_size( reader: *const vx_velox_read_at, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> u64 { try_or(error_out, 0, || { let reader = unsafe { @@ -483,11 +484,6 @@ mod tests { use vortex::file::WriteOptionsSessionExt; use vortex_buffer::Alignment; use vortex_error::vortex_ensure; - use vortex_ffi::vx_array_ref; - use vortex_ffi::vx_expression_new_with; - use vortex_ffi::vx_session_free; - use vortex_ffi::vx_session_new_with; - use vortex_ffi::vx_session_ref; use super::*; use crate::api::vx_velox_array_free; @@ -500,6 +496,11 @@ mod tests { use crate::api::vx_velox_scan_next_partition; use crate::api::vx_velox_scan_options; use crate::api::vx_velox_scan_selection; + use crate::ffi::vx_array_ref; + use crate::ffi::vx_expression_new_with; + use crate::ffi::vx_session_free; + use crate::ffi::vx_session_new_with; + use crate::ffi::vx_session_ref; use crate::schema::vx_velox_source_export_schema; use crate::source::vx_velox_natural_split; use crate::source::vx_velox_source_data_source; diff --git a/vortex-velox/src/schema.rs b/vortex-velox/src/schema.rs index d6277a4144c..d7c0b113e0c 100644 --- a/vortex-velox/src/schema.rs +++ b/vortex-velox/src/schema.rs @@ -6,9 +6,9 @@ use std::ptr; use arrow_array::ffi::FFI_ArrowSchema; use vortex_arrow::ArrowSessionExt; use vortex_error::vortex_err; -use vortex_ffi::try_or; -use vortex_ffi::vx_error; +use crate::ffi::try_or; +use crate::ffi::vx_velox_error; use crate::source::vx_velox_source; use crate::temporal::validate_velox_arrow_type; @@ -24,7 +24,7 @@ use crate::temporal::validate_velox_arrow_type; pub unsafe extern "C-unwind" fn vx_velox_source_export_schema( source: *const vx_velox_source, schema_out: *mut FFI_ArrowSchema, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let source = unsafe { @@ -64,7 +64,7 @@ mod tests { unsafe { vx_velox_source_export_schema(ptr::null(), &raw mut schema, &raw mut error) }; assert_eq!(status, 1); assert!(!error.is_null()); - unsafe { vortex_ffi::vx_error_free(error) }; + unsafe { crate::ffi::vx_error_free(error) }; Ok(()) } } diff --git a/vortex-velox/src/source.rs b/vortex-velox/src/source.rs index 7aace08a49c..15cfa32d9aa 100644 --- a/vortex-velox/src/source.rs +++ b/vortex-velox/src/source.rs @@ -10,18 +10,18 @@ use vortex::layout::scan::multi::MultiLayoutDataSource; use vortex::mask::Mask; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_ffi::ffi_runtime; -use vortex_ffi::try_or; -use vortex_ffi::vx_data_source; -use vortex_ffi::vx_data_source_new_with; -use vortex_ffi::vx_error; -use vortex_ffi::vx_expression; -use vortex_ffi::vx_expression_ref; -use vortex_ffi::vx_session; -use vortex_ffi::vx_session_ref; use vortex_io::VortexReadAt; use vortex_io::runtime::BlockingRuntime; +use crate::ffi::ffi_runtime; +use crate::ffi::try_or; +use crate::ffi::vx_data_source_new_with; +use crate::ffi::vx_expression_ref; +use crate::ffi::vx_session_ref; +use crate::ffi::vx_velox_data_source; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_expression; +use crate::ffi::vx_velox_session; use crate::read_at::vx_velox_read_at; /// A stable natural row range reported by a Vortex file. @@ -60,9 +60,9 @@ impl vx_velox_source { /// error pointer. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_source_new( - session: *const vx_session, + session: *const vx_velox_session, reader: *const vx_velox_read_at, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> *mut vx_velox_source { try_or(error_out, std::ptr::null_mut(), || { let session = unsafe { vx_session_ref(session)? }.clone(); @@ -160,7 +160,7 @@ pub unsafe extern "C-unwind" fn vx_velox_source_natural_split_at( source: *const vx_velox_source, index: usize, split_out: *mut vx_velox_natural_split, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let source = unsafe { @@ -203,11 +203,11 @@ pub unsafe extern "C-unwind" fn vx_velox_source_natural_split_at( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_source_prune_natural_splits( source: *const vx_velox_source, - expression: *const vx_expression, + expression: *const vx_velox_expression, first_split: usize, split_count: usize, pruned_out: *mut u8, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let source = unsafe { @@ -259,8 +259,8 @@ pub unsafe extern "C-unwind" fn vx_velox_source_prune_natural_splits( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_source_data_source( source: *const vx_velox_source, - error_out: *mut *mut vx_error, -) -> *const vx_data_source { + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_data_source { try_or(error_out, std::ptr::null(), || { let source = unsafe { source diff --git a/vortex-velox/src/test_support.rs b/vortex-velox/src/test_support.rs new file mode 100644 index 00000000000..e68ba0a3b46 --- /dev/null +++ b/vortex-velox/src/test_support.rs @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Private support for the Velox tests and benchmarks. + +use std::ptr; + +use arrow_array::array::make_array; +use arrow_array::ffi::FFI_ArrowArray; +use arrow_array::ffi::FFI_ArrowSchema; +use arrow_array::ffi::from_ffi; +use arrow_schema::Field; +use futures::SinkExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::channel::mpsc::Sender; +use vortex::array::ArrayRef; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::dtype::DType; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::file::WriteSummary; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::Task; +use vortex::io::session::RuntimeSessionExt; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +use crate::ffi::ffi_runtime; +use crate::ffi::try_or; +use crate::ffi::vx_array_new_with; +use crate::ffi::vx_session_new; +use crate::ffi::vx_session_ref; +use crate::ffi::vx_velox_array; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_expression; +use crate::ffi::vx_velox_view; + +struct TestSink { + input: Sender>, + writer: Task>, + dtype: DType, +} + +impl TestSink { + fn try_new(session: &VortexSession, path: String, dtype: DType) -> VortexResult { + let (input, output) = mpsc::channel(32); + let stream = ArrayStreamAdapter::new(dtype.clone(), output.into_stream()); + let writer_session = session.clone(); + let writer = session.handle().spawn(async move { + let mut file = async_fs::File::create(path).await?; + writer_session + .write_options() + .with_strategy(WriteStrategyBuilder::default().build()) + .write(&mut file, stream) + .await + }); + Ok(Self { + input, + writer, + dtype, + }) + } + + fn push(&mut self, array: ArrayRef) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == &self.dtype, + "array dtype {} does not match writer dtype {}", + array.dtype(), + self.dtype + ); + ffi_runtime() + .block_on(self.input.send(Ok(array))) + .map_err(|error| vortex_err!("Vortex test writer send failed: {error}")) + } + + fn close(self) -> VortexResult<()> { + drop(self.input); + ffi_runtime().block_on(async { + self.writer.await?; + Ok(()) + }) + } +} + +/// An opaque file writer for Velox tests and benchmarks. +pub struct vx_velox_test_writer { + session: *mut crate::ffi::vx_velox_session, + path: String, + sink: Option, +} + +impl Drop for vx_velox_test_writer { + fn drop(&mut self) { + // SAFETY: The writer owns this session handle. + unsafe { crate::ffi::vx_session_free(self.session) }; + } +} + +unsafe fn import_arrow( + session: &VortexSession, + array: *mut FFI_ArrowArray, + schema: *mut FFI_ArrowSchema, +) -> VortexResult { + vortex_ensure!(!array.is_null(), "Arrow array must not be null"); + vortex_ensure!(!schema.is_null(), "Arrow schema must not be null"); + // SAFETY: The caller transfers both initialized Arrow C Data structures. + let array = unsafe { ptr::replace(array, FFI_ArrowArray::empty()) }; + // SAFETY: The caller transfers both initialized Arrow C Data structures. + let schema = unsafe { ptr::replace(schema, FFI_ArrowSchema::empty()) }; + let array_data = unsafe { from_ffi(array, &schema) }?; + let field = Field::try_from(&schema)?.with_nullable(false); + let arrow_array = make_array(array_data); + session.arrow().from_arrow_array(arrow_array, &field) +} + +/// Import one Arrow C Data batch and apply an expression for adapter tests. +/// +/// # Safety +/// +/// Every pointer must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_test_array_from_arrow_apply( + session: *const crate::ffi::vx_velox_session, + array: *mut FFI_ArrowArray, + schema: *mut FFI_ArrowSchema, + expression: *const vx_velox_expression, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { + try_or(error_out, ptr::null(), || { + let session = unsafe { vx_session_ref(session) }?; + vortex_ensure!( + !expression.is_null(), + "Vortex test expression must not be null" + ); + let array = unsafe { import_arrow(session, array, schema) }?; + let expression = unsafe { vx_velox_expression::as_ref(expression) }; + Ok(vx_array_new_with(array.apply(expression)?)) + }) +} + +/// Create a private Vortex file writer for Velox tests. +/// +/// # Safety +/// +/// `path` and `error_out` must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_test_writer_new( + path: vx_velox_view, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_test_writer { + try_or(error_out, ptr::null_mut(), || { + let path = unsafe { path.as_str() }?.to_owned(); + vortex_ensure!( + !path.is_empty(), + "Vortex test writer path must not be empty" + ); + Ok(Box::into_raw(Box::new(vx_velox_test_writer { + session: vx_session_new(), + path, + sink: None, + }))) + }) +} + +/// Push one Arrow C Data batch into a private Vortex test writer. +/// +/// # Safety +/// +/// Every pointer must satisfy the adapter header contract. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_test_writer_push( + writer: *mut vx_velox_test_writer, + array: *mut FFI_ArrowArray, + schema: *mut FFI_ArrowSchema, + error_out: *mut *mut vx_velox_error, +) -> i32 { + try_or(error_out, 1, || { + let writer = unsafe { + writer + .as_mut() + .ok_or_else(|| vortex_err!("Vortex test writer must not be null"))? + }; + let session = unsafe { vx_session_ref(writer.session) }?; + let array = unsafe { import_arrow(session, array, schema) }?; + if writer.sink.is_none() { + writer.sink = Some(TestSink::try_new( + session, + writer.path.clone(), + array.dtype().clone(), + )?); + } + writer + .sink + .as_mut() + .ok_or_else(|| vortex_err!("Vortex test writer did not create a sink"))? + .push(array)?; + Ok(0) + }) +} + +/// Close a private Vortex test writer. +/// +/// # Safety +/// +/// `writer` must transfer one live writer. `error_out` must be null or writable. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_test_writer_close( + writer: *mut vx_velox_test_writer, + error_out: *mut *mut vx_velox_error, +) -> i32 { + try_or(error_out, 1, || { + vortex_ensure!(!writer.is_null(), "Vortex test writer must not be null"); + // SAFETY: The caller transfers one writer from `vx_velox_test_writer_new`. + let mut writer = unsafe { Box::from_raw(writer) }; + writer + .sink + .take() + .ok_or_else(|| vortex_err!("Vortex test writer received no batches"))? + .close()?; + Ok(0) + }) +} + +/// Abort and free a private Vortex test writer. +/// +/// # Safety +/// +/// `writer` must be null or transfer one live writer. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_test_writer_abort(writer: *mut vx_velox_test_writer) { + if !writer.is_null() { + // SAFETY: The caller transfers one writer from `vx_velox_test_writer_new`. + drop(unsafe { Box::from_raw(writer) }); + } +} diff --git a/vortex-velox/src/visitor.rs b/vortex-velox/src/visitor.rs index 33a1398f83a..be7dce2346e 100644 --- a/vortex-velox/src/visitor.rs +++ b/vortex-velox/src/visitor.rs @@ -54,17 +54,17 @@ use vortex_error::vortex_err; use vortex_fastlanes::BitPacked; use vortex_fastlanes::BitPackedArrayExt; use vortex_fastlanes::FL_CHUNK_SIZE; -use vortex_ffi::try_or; -use vortex_ffi::vx_array; -use vortex_ffi::vx_array_ref; -use vortex_ffi::vx_error; -use vortex_ffi::vx_session; -use vortex_ffi::vx_session_ref; use crate::array::ArrowMemoryReservation; use crate::array::conservative_export_reservation; use crate::array::parse_memory_callbacks; use crate::array::vx_velox_arrow_memory_callbacks; +use crate::ffi::try_or; +use crate::ffi::vx_array_ref; +use crate::ffi::vx_session_ref; +use crate::ffi::vx_velox_array; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_session; /// A fixed-width primitive value identifier in a semantic visitor block. pub type vx_velox_primitive_type = u32; @@ -2231,10 +2231,10 @@ fn visit_array( /// `error_out` must be null or valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_export_cursor_new( - session: *const vx_session, - array: *const vx_array, + session: *const vx_velox_session, + array: *const vx_velox_array, memory_callbacks: *const vx_velox_arrow_memory_callbacks, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> *mut vx_velox_export_cursor { try_or(error_out, ptr::null_mut(), || { let session = unsafe { vx_session_ref(session)? }; @@ -2271,7 +2271,7 @@ pub unsafe extern "C-unwind" fn vx_velox_export_cursor_visit( offset: usize, length: usize, visitor: *const vx_velox_visitor, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let cursor = unsafe { @@ -2301,11 +2301,11 @@ pub unsafe extern "C-unwind" fn vx_velox_export_cursor_visit( /// must remain live until this call returns. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_velox_array_visit( - session: *const vx_session, - array: *const vx_array, + session: *const vx_velox_session, + array: *const vx_velox_array, request: *const vx_velox_visit_request, visitor: *const vx_velox_visitor, - error_out: *mut *mut vx_error, + error_out: *mut *mut vx_velox_error, ) -> i32 { try_or(error_out, 1, || { let session = unsafe { vx_session_ref(session)? }; @@ -2362,12 +2362,12 @@ mod tests { use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_fastlanes::BitPackedData; - use vortex_ffi::vx_array_new_with; - use vortex_ffi::vx_session_free; - use vortex_ffi::vx_session_new_with; use super::*; use crate::api::vx_velox_array_free; + use crate::ffi::vx_array_new_with; + use crate::ffi::vx_session_free; + use crate::ffi::vx_session_new_with; #[derive(Default)] struct TestMemory { diff --git a/vortex-velox/tests/abi_contract.rs b/vortex-velox/tests/abi_contract.rs index f3d526e1214..250451fdd22 100644 --- a/vortex-velox/tests/abi_contract.rs +++ b/vortex-velox/tests/abi_contract.rs @@ -289,19 +289,10 @@ mod tests { let manifest = env!("CARGO_MANIFEST_DIR"); let include = format!("-I{manifest}/cinclude"); - let base_include = format!("-I{manifest}/../vortex-ffi/cinclude"); let compiler = std::env::var("CC").unwrap_or_else(|_| "cc".to_owned()); compile_stdin( &compiler, - &[ - "-std=c11", - "-fsyntax-only", - "-x", - "c", - &include, - &base_include, - "-", - ], + &["-std=c11", "-fsyntax-only", "-x", "c", &include, "-"], &source, ) } @@ -310,21 +301,12 @@ mod tests { fn header_compiles_with_host_arrow_declarations() -> Result<(), Box> { let manifest = env!("CARGO_MANIFEST_DIR"); let include = format!("-I{manifest}/cinclude"); - let base_include = format!("-I{manifest}/../vortex-ffi/cinclude"); let compiler = std::env::var("CXX").unwrap_or_else(|_| "c++".to_owned()); let source = std::fs::read_to_string(format!("{manifest}/tests/velox_include_contract.cpp"))?; compile_stdin( &compiler, - &[ - "-std=c++20", - "-fsyntax-only", - "-x", - "c++", - &include, - &base_include, - "-", - ], + &["-std=c++20", "-fsyntax-only", "-x", "c++", &include, "-"], &source, ) } diff --git a/vortex-velox/tests/velox_include_contract.cpp b/vortex-velox/tests/velox_include_contract.cpp index d4b927d8d43..f32a5dfff8c 100644 --- a/vortex-velox/tests/velox_include_contract.cpp +++ b/vortex-velox/tests/velox_include_contract.cpp @@ -3,14 +3,7 @@ struct ArrowSchema; struct ArrowArray; -struct ArrowArrayStream; - -#define USE_OWN_ARROW -typedef struct ArrowSchema FFI_ArrowSchema; -typedef struct ArrowArray FFI_ArrowArray; -typedef struct ArrowArrayStream FFI_ArrowArrayStream; #include "vortex_velox.h" -#undef USE_OWN_ARROW static_assert(VX_VELOX_ABI_VERSION == 5u); static_assert(VX_VELOX_SELECTION_ALL == 0); @@ -26,12 +19,12 @@ void vx_velox_compile_velox_include_contract() { options.abi_version = VX_VELOX_ABI_VERSION; options.selection.include = VX_VELOX_SELECTION_ALL; - const vx_dtype *(*new_primitive)(vx_velox_ptype, bool, vx_error **) = + const vx_velox_dtype *(*new_primitive)(vx_velox_ptype, bool, vx_velox_error **) = vx_velox_dtype_new_primitive; - vx_expression *(*new_binary)(vx_velox_binary_operator, - const vx_expression *, - const vx_expression *, - vx_error **) = vx_velox_expression_binary; + vx_velox_expression *(*new_binary)(vx_velox_binary_operator, + const vx_velox_expression *, + const vx_velox_expression *, + vx_velox_error **) = vx_velox_expression_binary; (void)callbacks; (void)options; From cd0029e0cd48bc6b7e2a69fee137d58f838b7e8c Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 15:34:43 -0400 Subject: [PATCH 5/6] feat(vortex-velox): Generate the adapter ABI Signed-off-by: Will Manning --- .github/workflows/ci.yml | 4 +- Cargo.lock | 1 + vortex-velox/Cargo.toml | 3 + vortex-velox/build.rs | 68 + vortex-velox/cbindgen.toml | 76 + vortex-velox/cinclude/vortex_velox.h | 2106 +++++++++++++---- vortex-velox/src/lib.rs | 5 +- vortex-velox/src/read_at.rs | 7 + vortex-velox/src/source.rs | 19 + vortex-velox/tests/abi_contract.rs | 5 +- vortex-velox/tests/velox_include_contract.cpp | 2 +- 11 files changed, 1830 insertions(+), 466 deletions(-) create mode 100644 vortex-velox/build.rs create mode 100644 vortex-velox/cbindgen.toml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c647a95d42e..614650bde47 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -747,9 +747,9 @@ jobs: run: | BASE="${{ github.event.pull_request.base.sha || 'HEAD^' }}" cargo run --profile ci -p xtask -- check-editions --base "$BASE" - - name: "regenerate FFI header file" + - name: "regenerate C API header files" run: | - cargo +$NIGHTLY_TOOLCHAIN build --profile ci -p vortex-ffi + cargo +$NIGHTLY_TOOLCHAIN build --profile ci -p vortex-ffi -p vortex-velox - name: "Make sure no files changed after regenerating" run: | git status --porcelain diff --git a/Cargo.lock b/Cargo.lock index 619e90ee5f9..a967a2dd743 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11700,6 +11700,7 @@ dependencies = [ "arrow-schema 59.2.0", "async-fs", "bytes", + "cbindgen", "futures", "rstest", "vortex", diff --git a/vortex-velox/Cargo.toml b/vortex-velox/Cargo.toml index 8a1de5108a7..22b88c855f3 100644 --- a/vortex-velox/Cargo.toml +++ b/vortex-velox/Cargo.toml @@ -40,6 +40,9 @@ vortex = { workspace = true } [dev-dependencies] rstest = { workspace = true } +[build-dependencies] +cbindgen = { workspace = true } + [lib] name = "vortex_velox" crate-type = ["rlib", "staticlib"] diff --git a/vortex-velox/build.rs b/vortex-velox/build.rs new file mode 100644 index 00000000000..bf165604802 --- /dev/null +++ b/vortex-velox/build.rs @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::env; +use std::path::PathBuf; +use std::process::Command; +use std::process::exit; + +fn main() { + println!("cargo:rerun-if-changed=src"); + println!("cargo:rerun-if-changed=cbindgen.toml"); + println!("cargo:rerun-if-changed=Cargo.toml"); + println!("cargo:rerun-if-changed=build.rs"); + for variable in ["MIRI", "MIRIFLAGS", "CARGO_ENCODED_RUSTFLAGS"] { + println!("cargo:rerun-if-env-changed={variable}"); + } + + let rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap_or_default(); + if rustflags.contains("sanitizer") || rustflags.contains("address") { + println!("cargo:info=Skipping header generation under a sanitizer"); + return; + } + if env::var("MIRI").is_ok() || env::var("MIRIFLAGS").is_ok() { + println!("cargo:info=Skipping header generation under Miri"); + return; + } + + let rustc = Command::new("rustc").arg("-V").output(); + let is_nightly = rustc + .as_ref() + .map(|output| String::from_utf8_lossy(&output.stdout).contains("nightly")) + .unwrap_or(false); + if !is_nightly { + println!("cargo:info=Skipping header generation outside nightly Rust"); + return; + } + + let crate_dir = env::var("CARGO_MANIFEST_DIR").unwrap(); + let output_file = PathBuf::from(&crate_dir) + .join("cinclude") + .join("vortex_velox.h"); + let config = cbindgen::Config::from_file("cbindgen.toml").unwrap(); + let bindings = cbindgen::Builder::new() + .with_crate(&crate_dir) + .with_config(config) + .generate(); + + match bindings { + Ok(bindings) => { + bindings.write_to_file(&output_file); + if let Ok(status) = Command::new("clang-format") + .arg("-i") + .arg("--style=file") + .arg(&output_file) + .status() + && !status.success() + { + println!("cargo:warning=clang-format exited with status {status}"); + } + } + Err(error) => { + println!("cargo:error=Failed to generate vortex_velox.h: {error}"); + exit(1); + } + } +} diff --git a/vortex-velox/cbindgen.toml b/vortex-velox/cbindgen.toml new file mode 100644 index 00000000000..9cc2f35bbd0 --- /dev/null +++ b/vortex-velox/cbindgen.toml @@ -0,0 +1,76 @@ +language = "C" +braces = "SameLine" +pragma_once = false +cpp_compat = true +usize_is_size_t = true +style = "type" + +header = """ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include +#include +#include + +// THIS FILE IS AUTO-GENERATED. DO NOT EDIT IT DIRECTLY. + +typedef struct ArrowSchema ArrowSchema; +typedef struct ArrowArray ArrowArray; + +typedef struct vx_velox_error vx_velox_error; +typedef struct vx_velox_session vx_velox_session; +typedef struct vx_velox_dtype vx_velox_dtype; +typedef struct vx_velox_scalar vx_velox_scalar; +typedef struct vx_velox_expression vx_velox_expression; +typedef struct vx_velox_data_source vx_velox_data_source; +typedef struct vx_velox_scan vx_velox_scan; +typedef struct vx_velox_partition vx_velox_partition; +typedef struct vx_velox_array vx_velox_array; +typedef struct vx_velox_export_cursor vx_velox_export_cursor; +""" + +[export] +exclude = [ + "AdapterError", + "ArrayRef", + "CursorExport", + "DType", + "Expression", + "MultiLayoutDataSource", + "Nullability", + "PType", + "PartitionState", + "Primitive", + "Scalar", + "ScanState", + "VortexSession", + "vx_velox_array", + "vx_velox_data_source", + "vx_velox_dtype", + "vx_velox_error", + "vx_velox_expression", + "vx_velox_export_cursor", + "vx_velox_partition", + "vx_velox_scalar", + "vx_velox_scan", + "vx_velox_session", + "vx_velox_test_array_from_arrow_apply", + "vx_velox_test_writer", + "vx_velox_test_writer_abort", + "vx_velox_test_writer_close", + "vx_velox_test_writer_new", + "vx_velox_test_writer_push", +] + +[export.rename] +"FFI_ArrowSchema" = "ArrowSchema" +"FFI_ArrowArray" = "ArrowArray" +"f16" = "uint16_t" + +[parse] +parse_deps = false + +[parse.expand] +crates = ["vortex-velox"] diff --git a/vortex-velox/cinclude/vortex_velox.h b/vortex-velox/cinclude/vortex_velox.h index 62e6ff9c76f..68ed4bdd35a 100644 --- a/vortex-velox/cinclude/vortex_velox.h +++ b/vortex-velox/cinclude/vortex_velox.h @@ -6,41 +6,11 @@ #include #include -struct ArrowSchema; -struct ArrowArray; +// THIS FILE IS AUTO-GENERATED. DO NOT EDIT IT DIRECTLY. -#ifdef __cplusplus -extern "C" { -#endif - -/* Velox calls only vx_velox_* adapter symbols. */ - -#define VX_VELOX_ABI_VERSION 5u -#define VX_VELOX_CAPABILITY_BATCH_READ (UINT64_C(1) << 0) -#define VX_VELOX_CAPABILITY_CALLBACK_SOURCE (UINT64_C(1) << 1) -#define VX_VELOX_CAPABILITY_NATURAL_SPLITS (UINT64_C(1) << 2) -#define VX_VELOX_CAPABILITY_PRIMITIVE_VISITOR (UINT64_C(1) << 3) -#define VX_VELOX_CAPABILITY_ARROW_SCHEMA (UINT64_C(1) << 4) -#define VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT (UINT64_C(1) << 5) -#define VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION (UINT64_C(1) << 6) -#define VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING (UINT64_C(1) << 7) -/* Vortex checks cancellation before each host read callback. */ -#define VX_VELOX_CAPABILITY_READ_CANCELLATION (UINT64_C(1) << 8) -#define VX_VELOX_CAPABILITY_EXPORT_CURSOR (UINT64_C(1) << 9) -#define VX_VELOX_CAPABILITY_PLAIN_PROJECTION (UINT64_C(1) << 10) -#define VX_VELOX_CAPABILITY_VARBIN_VISITOR (UINT64_C(1) << 11) -#define VX_VELOX_CAPABILITY_DICTIONARY_VISITOR (UINT64_C(1) << 12) -#define VX_VELOX_CAPABILITY_CONSTANT_VISITOR (UINT64_C(1) << 13) -#define VX_VELOX_CAPABILITY_BOOL_VISITOR (UINT64_C(1) << 14) -#define VX_VELOX_CAPABILITY_DATE_VISITOR (UINT64_C(1) << 15) -#define VX_VELOX_CAPABILITY_DECIMAL_VISITOR (UINT64_C(1) << 16) -#define VX_VELOX_CAPABILITY_STRUCT_VISITOR (UINT64_C(1) << 17) -#define VX_VELOX_CAPABILITY_LIST_VISITOR (UINT64_C(1) << 18) -#define VX_VELOX_CAPABILITY_MAP_VISITOR (UINT64_C(1) << 19) +typedef struct ArrowSchema ArrowSchema; +typedef struct ArrowArray ArrowArray; -typedef struct vx_velox_read_at vx_velox_read_at; -typedef struct vx_velox_source vx_velox_source; -typedef struct vx_velox_export_cursor vx_velox_export_cursor; typedef struct vx_velox_error vx_velox_error; typedef struct vx_velox_session vx_velox_session; typedef struct vx_velox_dtype vx_velox_dtype; @@ -50,523 +20,1727 @@ typedef struct vx_velox_data_source vx_velox_data_source; typedef struct vx_velox_scan vx_velox_scan; typedef struct vx_velox_partition vx_velox_partition; typedef struct vx_velox_array vx_velox_array; +typedef struct vx_velox_export_cursor vx_velox_export_cursor; + + +#include +#include +#include +#include +#include + +/** + * The current major version of the Vortex and Velox adapter ABI. + */ +#define VX_VELOX_ABI_VERSION 6 + +/** + * The adapter supports batched host range reads. + */ +#define VX_VELOX_CAPABILITY_BATCH_READ (1 << 0) + +/** + * The adapter can open callback-backed Vortex sources. + */ +#define VX_VELOX_CAPABILITY_CALLBACK_SOURCE (1 << 1) + +/** + * The adapter reports stable natural row splits. + */ +#define VX_VELOX_CAPABILITY_NATURAL_SPLITS (1 << 2) + +/** + * The adapter can visit canonical primitive values in retained blocks. + */ +#define VX_VELOX_CAPABILITY_PRIMITIVE_VISITOR (1 << 3) + +/** + * The adapter can export source schemas through the Arrow C Data Interface. + */ +#define VX_VELOX_CAPABILITY_ARROW_SCHEMA (1 << 4) + +/** + * The adapter can export one Vortex array through the Arrow C Data Interface. + */ +#define VX_VELOX_CAPABILITY_ARRAY_ARROW_EXPORT (1 << 5) + +/** + * The adapter can project absolute file-row indexes with scan fields. + */ +#define VX_VELOX_CAPABILITY_ROW_INDEX_PROJECTION (1 << 6) + +/** + * The adapter can prove that natural splits cannot match an expression. + */ +#define VX_VELOX_CAPABILITY_NATURAL_SPLIT_PRUNING (1 << 7) + +/** + * The callback reader observes host cancellation before each host read callback. + * + * This capability does not claim cancellation during cached scans or CPU execution. + */ +#define VX_VELOX_CAPABILITY_READ_CANCELLATION (1 << 8) + +/** + * The adapter retains one prepared array across several Velox output windows. + */ +#define VX_VELOX_CAPABILITY_EXPORT_CURSOR (1 << 9) + +/** + * The adapter can omit row-index projection from scans with contiguous rows. + */ +#define VX_VELOX_CAPABILITY_PLAIN_PROJECTION (1 << 10) + +/** + * The adapter can visit canonical UTF-8 and binary values in retained blocks. + */ +#define VX_VELOX_CAPABILITY_VARBIN_VISITOR (1 << 11) + +/** + * The adapter can preserve dictionary arrays during native export. + */ +#define VX_VELOX_CAPABILITY_DICTIONARY_VISITOR (1 << 12) -typedef struct vx_velox_view { - const char *ptr; - size_t len; +/** + * The adapter can preserve constant arrays during native export. + */ +#define VX_VELOX_CAPABILITY_CONSTANT_VISITOR (1 << 13) + +/** + * The adapter can visit canonical packed Boolean values in retained blocks. + */ +#define VX_VELOX_CAPABILITY_BOOL_VISITOR (1 << 14) + +/** + * The adapter can export Vortex day-based dates through the primitive visitor. + */ +#define VX_VELOX_CAPABILITY_DATE_VISITOR (1 << 15) + +/** + * The adapter can normalize Vortex decimals for the primitive visitor. + */ +#define VX_VELOX_CAPABILITY_DECIMAL_VISITOR (1 << 16) + +/** + * The adapter can preserve canonical struct children during native export. + */ +#define VX_VELOX_CAPABILITY_STRUCT_VISITOR (1 << 17) + +/** + * The adapter can preserve canonical list children during native export. + */ +#define VX_VELOX_CAPABILITY_LIST_VISITOR (1 << 18) + +/** + * The adapter can preserve canonical map children during native export. + */ +#define VX_VELOX_CAPABILITY_MAP_VISITOR (1 << 19) + +/** + * Natural splits include a stable byte-range assignment token. + */ +#define VX_VELOX_CAPABILITY_SPLIT_ASSIGNMENT (1 << 20) + +/** + * An opaque Vortex positional reader backed by Velox callbacks. + */ +typedef struct vx_velox_read_at vx_velox_read_at; + +/** + * An opened Vortex file that uses Velox callbacks for all reads. + */ +typedef struct vx_velox_source vx_velox_source; + +typedef struct { + const char *ptr; + size_t len; } vx_velox_view; +/** + * A fixed-width primitive type identifier for Velox scalar construction. + */ typedef uint32_t vx_velox_ptype; -#define VX_VELOX_PTYPE_U8 UINT32_C(0) -#define VX_VELOX_PTYPE_U16 UINT32_C(1) -#define VX_VELOX_PTYPE_U32 UINT32_C(2) -#define VX_VELOX_PTYPE_U64 UINT32_C(3) -#define VX_VELOX_PTYPE_I8 UINT32_C(4) -#define VX_VELOX_PTYPE_I16 UINT32_C(5) -#define VX_VELOX_PTYPE_I32 UINT32_C(6) -#define VX_VELOX_PTYPE_I64 UINT32_C(7) -#define VX_VELOX_PTYPE_F16 UINT32_C(8) -#define VX_VELOX_PTYPE_F32 UINT32_C(9) -#define VX_VELOX_PTYPE_F64 UINT32_C(10) +/** + * A fixed-width binary expression operator identifier. + */ typedef uint32_t vx_velox_binary_operator; -#define VX_VELOX_OPERATOR_EQ UINT32_C(0) -#define VX_VELOX_OPERATOR_NOT_EQ UINT32_C(1) -#define VX_VELOX_OPERATOR_GT UINT32_C(2) -#define VX_VELOX_OPERATOR_GTE UINT32_C(3) -#define VX_VELOX_OPERATOR_LT UINT32_C(4) -#define VX_VELOX_OPERATOR_LTE UINT32_C(5) -#define VX_VELOX_OPERATOR_KLEENE_AND UINT32_C(6) -#define VX_VELOX_OPERATOR_KLEENE_OR UINT32_C(7) +/** + * A fixed-width row-selection mode identifier. + */ typedef uint32_t vx_velox_scan_selection_include; -#define VX_VELOX_SELECTION_ALL UINT32_C(0) -#define VX_VELOX_SELECTION_INCLUDE UINT32_C(1) -#define VX_VELOX_SELECTION_EXCLUDE UINT32_C(2) - -typedef struct vx_velox_scan_selection { - const uint64_t *indices; - size_t length; - vx_velox_scan_selection_include include; + +/** + * A stable row selection for one scan request. + */ +typedef struct { + /** + * The selected row indexes. + */ + const uint64_t *indices; + /** + * The number of selected row indexes. + */ + size_t length; + /** + * The selection mode. + */ + vx_velox_scan_selection_include include; } vx_velox_scan_selection; -typedef struct vx_velox_scan_options { - size_t struct_size; - uint32_t abi_version; - const vx_velox_expression *projection; - const vx_velox_expression *filter; - uint64_t row_range_begin; - uint64_t row_range_end; - vx_velox_scan_selection selection; - uint64_t limit; - bool ordered; +/** + * Stable options for one Vortex scan. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_scan_options)`. + */ + size_t struct_size; + /** + * Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + */ + uint32_t abi_version; + /** + * The projected expression, or null for every field. + */ + const vx_velox_expression *projection; + /** + * The exact filter expression, or null for no filter. + */ + const vx_velox_expression *filter; + /** + * The first row in the scan range. + */ + uint64_t row_range_begin; + /** + * One past the final row in the scan range. + */ + uint64_t row_range_end; + /** + * An optional row-index selection. + */ + vx_velox_scan_selection selection; + /** + * The maximum output row count, or zero for no limit. + */ + uint64_t limit; + /** + * Return rows in storage order. + */ + bool ordered; } vx_velox_scan_options; -typedef struct vx_velox_read_request { - size_t struct_size; - uint64_t offset; - size_t length; - size_t alignment; +/** + * Host memory callbacks for one Arrow C Data export. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_arrow_memory_callbacks)`. + */ + size_t struct_size; + /** + * Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + */ + uint32_t abi_version; + /** + * An opaque host context. + */ + void *context; + /** + * Retain the host context until the Arrow array release callback runs. + */ + void (*retain_context)(void *context); + /** + * Release one host context reference. + */ + void (*release_context)(void *context); + /** + * Reserve Arrow payload bytes before conversion. Zero means success. + */ + int32_t (*report_allocation)(void *context, size_t retained_bytes); + /** + * Free retained Arrow payload bytes. + */ + void (*report_free)(void *context, size_t retained_bytes); + /** + * Return the last callback error as a null-terminated string. + */ + const char *(*last_error)(void *context); +} vx_velox_arrow_memory_callbacks; + +/** + * A positional read request passed to the Velox callback. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_read_request)`. + */ + size_t struct_size; + /** + * The file offset in bytes. + */ + uint64_t offset; + /** + * The exact requested length in bytes. + */ + size_t length; + /** + * The required buffer alignment in bytes. + */ + size_t alignment; } vx_velox_read_request; -typedef struct vx_velox_buffer { - size_t struct_size; - const uint8_t *data; - size_t length; - void *owner; - void (*release)(void *owner); +/** + * A retained host buffer returned by the Velox callback. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_buffer)`. + */ + size_t struct_size; + /** + * The first byte of the returned range. + */ + const uint8_t *data; + /** + * The number of returned bytes. + */ + size_t length; + /** + * An opaque owner passed to `release`. + */ + void *owner; + /** + * Release the owner after Vortex no longer needs the bytes. + */ + void (*release)(void *owner); } vx_velox_buffer; /** - * Callbacks for positional reads through the host engine. - * - * Vortex can call size, read_ranges, is_cancelled, and last_error concurrently. - * The context and callbacks must be thread-safe. A non-zero concurrency value - * limits requests in one callback and gives Vortex a scheduling hint. It is not - * a synchronization guarantee. - * - * last_error must return the calling thread's most recent callback error. The - * returned string must remain valid until the next callback on that thread. - * Every callback must catch C++ exceptions. No callback can unwind across this - * C ABI. release_context runs after the final callback and can run on any - * thread. - * - * Vortex checks is_cancelled before each read_ranges call. The check does not - * interrupt an active callback or CPU work. - */ -typedef struct vx_velox_read_at_callbacks { - size_t struct_size; - uint32_t abi_version; - void *context; - int32_t (*size)(void *context, uint64_t *size_out); - int32_t (*read_ranges)(void *context, - const vx_velox_read_request *requests, - size_t request_count, - vx_velox_buffer *outputs); - const char *(*last_error)(void *context); - void (*release_context)(void *context); - int32_t (*is_cancelled)(void *context); - size_t concurrency; -} vx_velox_read_at_callbacks; - -typedef struct vx_velox_natural_split { - size_t struct_size; - uint64_t row_begin; - uint64_t row_end; -} vx_velox_natural_split; + * Velox callbacks that provide a Vortex positional reader. + * + * Vortex can call these functions concurrently. The context and every callback must be + * thread-safe. `concurrency` limits one callback batch and gives Vortex a scheduling hint. It does + * not provide synchronization. `last_error` must return the calling thread's most recent callback + * error. Its string must remain valid until the next callback on that thread. Every callback must + * catch foreign exceptions and must not unwind across this ABI. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_read_at_callbacks)`. + */ + size_t struct_size; + /** + * Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + */ + uint32_t abi_version; + /** + * An opaque callback context. + */ + void *context; + /** + * Return the file size through `size_out`. Zero means success. + */ + int32_t (*size)(void *context, uint64_t *size_out); + /** + * Read every request and populate the matching output. Zero means success. + */ + int32_t (*read_ranges)(void *context, + const vx_velox_read_request *requests, + size_t request_count, + vx_velox_buffer *outputs); + /** + * Return the last callback error as a null-terminated string. + */ + const char *(*last_error)(void *context); + /** + * Release the callback context. + */ + void (*release_context)(void *context); + /** + * Return a non-zero value after the host cancels the scan. + */ + int32_t (*is_cancelled)(void *context); + /** + * Limit one callback batch and give Vortex a preferred concurrency value. + */ + size_t concurrency; +} vx_velox_read_at_callbacks; + +/** + * A stable natural row range reported by a Vortex file. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_natural_split)`. + */ + size_t struct_size; + /** + * The first row in the split. + */ + uint64_t row_begin; + /** + * One past the final row in the split. + */ + uint64_t row_end; + /** + * The file byte that assigns this split to one external byte range. + */ + uint64_t assignment_byte; +} vx_velox_natural_split; + +/** + * A fixed-width primitive value identifier in a semantic visitor block. + */ +typedef uint32_t vx_velox_primitive_type; + +/** + * A fixed-width validity representation identifier for one visitor block. + */ +typedef uint32_t vx_velox_validity_kind; + +/** + * A retained owner for buffers in a visitor block. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_buffer_owner)`. + */ + size_t struct_size; + /** + * An opaque retained object. + */ + const void *owner; + /** + * Add one owner reference before the callback returns. + */ + void (*retain)(const void *owner); + /** + * Release one retained owner reference. + */ + void (*release)(const void *owner); + /** + * The exact sum of the value and validity allocation sizes retained by this owner. + */ + size_t retained_bytes; +} vx_velox_buffer_owner; + +/** + * A canonical primitive block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_primitive_view)`. + */ + size_t struct_size; + /** + * The physical type of each value. + */ + vx_velox_primitive_type primitive_type; + /** + * The logical decimal precision, or zero for a non-decimal block. + */ + uint32_t decimal_precision; + /** + * The logical decimal scale, or zero for a non-decimal block. + */ + int32_t decimal_scale; + /** + * The number of logical values in the block. + */ + size_t length; + /** + * The first value byte. + */ + const uint8_t *values; + /** + * The number of value bytes. + */ + size_t values_length; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains all pointers in this view. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty values buffer. + */ + size_t values_alignment; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_primitive_view; + +/** + * Identifies the logical type of a variable-width binary block. + */ +typedef uint32_t vx_velox_varbin_kind; + +/** + * Defines the stable 16-byte variable-width binary view contract. + */ +typedef struct { + /** + * Stores the logical byte length. + */ + uint32_t length; + /** + * Stores inline bytes, or prefix, buffer index, and offset for outlined values. + */ + uint8_t data[12]; +} vx_velox_binary_view; + +/** + * Describes one retained payload buffer. + */ +typedef struct { + /** + * The first payload byte. + */ + const uint8_t *data; + /** + * The number of visible payload bytes. + */ + size_t length; +} vx_velox_byte_buffer_view; + +/** + * A canonical variable-width binary block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_varbin_view)`. + */ + size_t struct_size; + /** + * Identifies UTF-8 or binary values. + */ + vx_velox_varbin_kind kind; + /** + * The number of logical values in the block. + */ + size_t length; + /** + * The first 16-byte binary view. + */ + const vx_velox_binary_view *views; + /** + * The number of readable bytes in `views`. + */ + size_t views_length; + /** + * The retained payload buffer descriptors. + */ + const vx_velox_byte_buffer_view *data_buffers; + /** + * The number of payload buffer descriptors. + */ + size_t data_buffer_count; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of readable validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains all pointers in this view. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty view buffer. + */ + size_t views_alignment; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_varbin_view; + +/** + * A dictionary block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_dictionary_view)`. + */ + size_t struct_size; + /** + * The number of logical dictionary codes. + */ + size_t length; + /** + * The canonical integer codes for this output window. + */ + vx_velox_primitive_view codes; + /** + * A borrowed prepared cursor for the dictionary values. + */ + const vx_velox_export_cursor *values; + /** + * The number of dictionary values. + */ + size_t values_length; +} vx_velox_dictionary_view; + +/** + * A constant block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_constant_view)`. + */ + size_t struct_size; + /** + * The number of repeated logical values. + */ + size_t length; + /** + * A borrowed prepared cursor with one canonical value. + */ + const vx_velox_export_cursor *value; +} vx_velox_constant_view; + +/** + * A canonical packed Boolean block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_bool_view)`. + */ + size_t struct_size; + /** + * The number of logical Boolean values. + */ + size_t length; + /** + * The first packed value byte. + */ + const uint8_t *values; + /** + * The number of readable value bytes. + */ + size_t values_length; + /** + * The first logical value bit within `values`. + */ + size_t values_bit_offset; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of readable validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains all pointers in this view. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty value buffer. + */ + size_t values_alignment; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_bool_view; + +/** + * A canonical struct block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_struct_view)`. + */ + size_t struct_size; + /** + * The number of logical struct values in this window. + */ + size_t length; + /** + * The first logical row in each field cursor. + */ + size_t offset; + /** + * Borrowed prepared cursors in field order. + */ + const vx_velox_export_cursor *const *fields; + /** + * The number of field cursors. + */ + size_t field_count; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of readable validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains the parent validity buffer. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_struct_view; + +/** + * A canonical list block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_list_view)`. + */ + size_t struct_size; + /** + * The number of logical lists in this window. + */ + size_t length; + /** + * One non-negative element offset per list. Values remain absolute against `elements`. + */ + const int32_t *offsets; + /** + * One non-negative element count per list. + */ + const int32_t *sizes; + /** + * A borrowed prepared cursor for all referenced elements. + */ + const vx_velox_export_cursor *elements; + /** + * The number of values in the element cursor. + */ + size_t elements_length; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of readable validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains the complete offsets, sizes, and parent validity allocations. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty offsets buffer. + */ + size_t offsets_alignment; + /** + * The guaranteed byte alignment of a non-empty sizes buffer. + */ + size_t sizes_alignment; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_list_view; + +/** + * A canonical map block delivered to Velox. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_map_view)`. + */ + size_t struct_size; + /** + * The number of logical maps in this window. + */ + size_t length; + /** + * One non-negative entry offset per map. Values remain absolute against the child cursors. + */ + const int32_t *offsets; + /** + * One non-negative entry count per map. + */ + const int32_t *sizes; + /** + * A borrowed prepared cursor for all map keys. + */ + const vx_velox_export_cursor *keys; + /** + * A borrowed prepared cursor for all map values. + */ + const vx_velox_export_cursor *values; + /** + * The number of entries in each child cursor. + */ + size_t entries_length; + /** + * True when each map asserts sorted keys. + */ + bool keys_sorted; + /** + * The validity representation. + */ + vx_velox_validity_kind validity_kind; + /** + * The first validity byte when `validity_kind` is `Bitmap`. + */ + const uint8_t *validity; + /** + * The number of readable validity bytes. + */ + size_t validity_length; + /** + * The first logical validity bit within `validity`. + */ + size_t validity_bit_offset; + /** + * Retains the complete offsets, sizes, and parent validity allocations. + */ + vx_velox_buffer_owner buffers; + /** + * The guaranteed byte alignment of a non-empty offsets buffer. + */ + size_t offsets_alignment; + /** + * The guaranteed byte alignment of a non-empty sizes buffer. + */ + size_t sizes_alignment; + /** + * The guaranteed byte alignment of a non-empty validity buffer. + */ + size_t validity_alignment; +} vx_velox_map_view; + +/** + * Host callbacks for Vortex array traversal. + * + * One array visit calls the matching callback synchronously. Shared tables can receive concurrent + * callbacks from simultaneous visits. `last_error` must return the calling thread's most recent + * error. The string must remain valid until the next callback on that thread. Callbacks must catch + * foreign exceptions and must not unwind across this ABI. The host owns the context. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_visitor)`. + */ + size_t struct_size; + /** + * Set this field to [`crate::VX_VELOX_ABI_VERSION`]. + */ + uint32_t abi_version; + /** + * An opaque callback context. + */ + void *context; + /** + * Consume one canonical primitive block. Zero means success. + */ + int32_t (*visit_primitive)(void *context, const vx_velox_primitive_view *view); + /** + * Return the last callback error as a null-terminated string. + */ + const char *(*last_error)(void *context); + /** + * Consume one canonical variable-width binary block. Zero means success. + */ + int32_t (*visit_varbin)(void *context, const vx_velox_varbin_view *view); + /** + * Consume one dictionary block. Zero means success. + */ + int32_t (*visit_dictionary)(void *context, const vx_velox_dictionary_view *view); + /** + * Consume one constant block. Zero means success. + */ + int32_t (*visit_constant)(void *context, const vx_velox_constant_view *view); + /** + * Consume one canonical packed Boolean block. Zero means success. + */ + int32_t (*visit_bool)(void *context, const vx_velox_bool_view *view); + /** + * Consume one canonical struct block. Zero means success. + */ + int32_t (*visit_struct)(void *context, const vx_velox_struct_view *view); + /** + * Consume one canonical list block. Zero means success. + */ + int32_t (*visit_list)(void *context, const vx_velox_list_view *view); + /** + * Consume one canonical map block. Zero means success. + */ + int32_t (*visit_map)(void *context, const vx_velox_map_view *view); +} vx_velox_visitor; + +/** + * A single-shot subset request for the semantic visitor. + */ +typedef struct { + /** + * Set this field to `sizeof(vx_velox_visit_request)`. + */ + size_t struct_size; + /** + * Unique, increasing source positions. Null selects every row. + */ + const uint64_t *rows; + /** + * The number of source positions. + */ + size_t row_count; +} vx_velox_visit_request; + +/** + * Unsigned 8-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_U8 0 + +/** + * Unsigned 16-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_U16 1 + +/** + * Unsigned 32-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_U32 2 + +/** + * Unsigned 64-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_U64 3 + +/** + * Signed 8-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_I8 4 + +/** + * Signed 16-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_I16 5 + +/** + * Signed 32-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_I32 6 + +/** + * Signed 64-bit integer type identifier. + */ +#define VX_VELOX_PTYPE_I64 7 + +/** + * 16-bit floating-point type identifier. + */ +#define VX_VELOX_PTYPE_F16 8 + +/** + * 32-bit floating-point type identifier. + */ +#define VX_VELOX_PTYPE_F32 9 + +/** + * 64-bit floating-point type identifier. + */ +#define VX_VELOX_PTYPE_F64 10 + +/** + * Equality operator identifier. + */ +#define VX_VELOX_OPERATOR_EQ 0 + +/** + * Inequality operator identifier. + */ +#define VX_VELOX_OPERATOR_NOT_EQ 1 + +/** + * Greater-than operator identifier. + */ +#define VX_VELOX_OPERATOR_GT 2 + +/** + * Greater-than-or-equal operator identifier. + */ +#define VX_VELOX_OPERATOR_GTE 3 + +/** + * Less-than operator identifier. + */ +#define VX_VELOX_OPERATOR_LT 4 + +/** + * Less-than-or-equal operator identifier. + */ +#define VX_VELOX_OPERATOR_LTE 5 + +/** + * Kleene logical AND operator identifier. + */ +#define VX_VELOX_OPERATOR_KLEENE_AND 6 + +/** + * Kleene logical OR operator identifier. + */ +#define VX_VELOX_OPERATOR_KLEENE_OR 7 + +/** + * Include every row. + */ +#define VX_VELOX_SELECTION_ALL 0 + +/** + * Include the supplied row indexes. + */ +#define VX_VELOX_SELECTION_INCLUDE 1 + +/** + * Exclude the supplied row indexes. + */ +#define VX_VELOX_SELECTION_EXCLUDE 2 + +/** + * Unsigned 8-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_U8 0 + +/** + * Unsigned 16-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_U16 1 + +/** + * Unsigned 32-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_U32 2 + +/** + * Unsigned 64-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_U64 3 + +/** + * Signed 8-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_I8 4 + +/** + * Signed 16-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_I16 5 + +/** + * Signed 32-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_I32 6 + +/** + * Signed 64-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_I64 7 + +/** + * IEEE 754 binary16 primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_F16 8 + +/** + * IEEE 754 binary32 primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_F32 9 + +/** + * IEEE 754 binary64 primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_F64 10 + +/** + * Signed 128-bit primitive identifier. + */ +#define VX_VELOX_PRIMITIVE_I128 11 + +/** + * The type is not nullable. + */ +#define VX_VELOX_VALIDITY_NON_NULLABLE 0 + +/** + * Every value is valid. + */ +#define VX_VELOX_VALIDITY_ALL_VALID 1 + +/** + * Every value is null. + */ +#define VX_VELOX_VALIDITY_ALL_INVALID 2 + +/** + * A packed bitmap contains one valid bit per value. + */ +#define VX_VELOX_VALIDITY_BITMAP 3 + +/** + * Identifies UTF-8 values. + */ +#define VX_VELOX_VARBIN_UTF8 0 + +/** + * Identifies arbitrary binary values. + */ +#define VX_VELOX_VARBIN_BINARY 1 + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +/** + * Return the adapter ABI version. + */ +uint32_t vx_velox_abi_version(void); + +/** + * Return the capabilities implemented by this adapter build. + */ +uint64_t vx_velox_capabilities(void); + +/** + * Return the message stored in an adapter error. + * + * # Safety + * + * `error` must point to a live error handle. + */ +vx_velox_view vx_velox_error_message(const vx_velox_error *error); + +/** + * Free an adapter error. + * + * # Safety + * + * `error` must be null or an owned error handle. + */ +void vx_velox_error_free(const vx_velox_error *error); + +/** + * Create a default Vortex session for Velox. + */ +vx_velox_session *vx_velox_session_new(void); + +/** + * Clone a Vortex session. + * + * # Safety + * + * `session` must point to a live Vortex session. + */ +vx_velox_session *vx_velox_session_clone(const vx_velox_session *session); + +/** + * Free a Vortex session. + * + * # Safety + * + * `session` must be null or an owned session handle. + */ +void vx_velox_session_free(const vx_velox_session *session); + +/** + * Create a primitive dtype for a list literal. + * + * # Safety + * + * `error_out` must be null or valid for one error pointer. + */ +const vx_velox_dtype *vx_velox_dtype_new_primitive(vx_velox_ptype ptype, + bool nullable, + vx_velox_error **error_out); + +/** + * Free a dtype. + * + * # Safety + * + * `dtype` must be null or an owned dtype handle. + */ +void vx_velox_dtype_free(const vx_velox_dtype *dtype); + +/** + * Create a Boolean scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); + +/** + *Create an i8 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); + +/** + *Create an i16 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); + +/** + *Create an i32 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); + +/** + * Create a date scalar that stores days since the Unix epoch. + * + * # Safety + * + * `error_out` must be null or valid for one error pointer. + */ +vx_velox_scalar *vx_velox_scalar_new_date_days(int32_t value, + bool nullable, + vx_velox_error **error_out); + +/** + *Create an i64 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); + +/** + *Create an f32 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_f32(float value, bool nullable); + +/** + *Create an f64 scalar. + */ +vx_velox_scalar *vx_velox_scalar_new_f64(double value, bool nullable); + +/** + * Create a UTF-8 scalar. + * + * # Safety + * + * `value` and `error_out` must satisfy the adapter header contract. + */ +vx_velox_scalar *vx_velox_scalar_new_utf8(vx_velox_view value, + bool nullable, + vx_velox_error **error_out); + +/** + * Create a binary scalar. + * + * # Safety + * + * `data` must identify `length` bytes. `error_out` must be null or valid. + */ +vx_velox_scalar *vx_velox_scalar_new_binary(const uint8_t *data, + size_t length, + bool nullable, + vx_velox_error **error_out); + +/** + * Create a list scalar. + * + * # Safety + * + * Every pointer must satisfy the adapter header contract. + */ +vx_velox_scalar *vx_velox_scalar_new_list(const vx_velox_dtype *element_dtype, + const vx_velox_scalar *const *elements, + size_t length, + bool nullable, + vx_velox_error **error_out); + +/** + * Free a scalar. + * + * # Safety + * + * `scalar` must be null or an owned scalar handle. + */ +void vx_velox_scalar_free(const vx_velox_scalar *scalar); + +/** + * Create a literal expression. + * + * # Safety + * + * `scalar` must point to a live scalar. `error_out` must be null or valid. + */ +vx_velox_expression *vx_velox_expression_literal(const vx_velox_scalar *scalar, + vx_velox_error **error_out); + +/** + * Create a root expression. + */ +vx_velox_expression *vx_velox_expression_root(void); + +/** + * Create a field expression. + * + * # Safety + * + * `child` must point to a live expression. `name` must identify valid UTF-8. + */ +vx_velox_expression *vx_velox_expression_get_item(vx_velox_view name, + const vx_velox_expression *child); + +/** + * Create a binary expression. + * + * # Safety + * + * Both operands must point to live expressions. `error_out` must be null or valid. + */ +vx_velox_expression *vx_velox_expression_binary(vx_velox_binary_operator operator_, + const vx_velox_expression *left, + const vx_velox_expression *right, + vx_velox_error **error_out); + +/** + * Create a conjunction from expressions. + * + * # Safety + * + * `expressions` must identify `length` live expression pointers. + */ +vx_velox_expression *vx_velox_expression_and(const vx_velox_expression *const *expressions, + size_t length); -typedef uint32_t vx_velox_primitive_type; -#define VX_VELOX_PRIMITIVE_U8 UINT32_C(0) -#define VX_VELOX_PRIMITIVE_U16 UINT32_C(1) -#define VX_VELOX_PRIMITIVE_U32 UINT32_C(2) -#define VX_VELOX_PRIMITIVE_U64 UINT32_C(3) -#define VX_VELOX_PRIMITIVE_I8 UINT32_C(4) -#define VX_VELOX_PRIMITIVE_I16 UINT32_C(5) -#define VX_VELOX_PRIMITIVE_I32 UINT32_C(6) -#define VX_VELOX_PRIMITIVE_I64 UINT32_C(7) -#define VX_VELOX_PRIMITIVE_F16 UINT32_C(8) -#define VX_VELOX_PRIMITIVE_F32 UINT32_C(9) -#define VX_VELOX_PRIMITIVE_F64 UINT32_C(10) -#define VX_VELOX_PRIMITIVE_I128 UINT32_C(11) +/** + * Create a disjunction from expressions. + * + * # Safety + * + * `expressions` must identify `length` live expression pointers. + */ +vx_velox_expression *vx_velox_expression_or(const vx_velox_expression *const *expressions, + size_t length); -typedef uint32_t vx_velox_validity_kind; -#define VX_VELOX_VALIDITY_NON_NULLABLE UINT32_C(0) -#define VX_VELOX_VALIDITY_ALL_VALID UINT32_C(1) -#define VX_VELOX_VALIDITY_ALL_INVALID UINT32_C(2) -#define VX_VELOX_VALIDITY_BITMAP UINT32_C(3) - -typedef struct vx_velox_buffer_owner { - size_t struct_size; - const void *owner; - void (*retain)(const void *owner); - void (*release)(const void *owner); - size_t retained_bytes; -} vx_velox_buffer_owner; +/** + * Create a logical negation. + * + * # Safety + * + * `child` must point to a live expression. + */ +vx_velox_expression *vx_velox_expression_not(const vx_velox_expression *child); /** - * A compact primitive payload and its owner. - * - * Vortex copies values into an allocation with uint64_t alignment. The values - * allocation rounds values_length up to that alignment. Vortex copies a bitmap - * into a uint64_t-aligned, word-padded allocation. A window rebases the pointer - * and reports its remaining bit offset. - * buffers.retained_bytes is the exact sum of these allocation sizes. - * - * The pointers remain valid through visit_primitive. The host must call retain - * before it stores a pointer beyond that callback. - */ -typedef struct vx_velox_primitive_view { - size_t struct_size; - vx_velox_primitive_type primitive_type; - uint32_t decimal_precision; - int32_t decimal_scale; - size_t length; - const uint8_t *values; - size_t values_length; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t values_alignment; - size_t validity_alignment; -} vx_velox_primitive_view; + * Create a null test. + * + * # Safety + * + * `child` must point to a live expression. + */ +vx_velox_expression *vx_velox_expression_is_null(const vx_velox_expression *child); -typedef uint32_t vx_velox_varbin_kind; -#define VX_VELOX_VARBIN_UTF8 UINT32_C(0) -#define VX_VELOX_VARBIN_BINARY UINT32_C(1) +/** + * Create a list membership test. + * + * # Safety + * + * Both operands must point to live expressions. + */ +vx_velox_expression *vx_velox_expression_list_contains(const vx_velox_expression *list, + const vx_velox_expression *value); -typedef struct vx_velox_byte_buffer_view { - const uint8_t *data; - size_t length; -} vx_velox_byte_buffer_view; +/** + * Return whether Vortex can push down an integer value set without generic expansion. + */ +bool vx_velox_can_push_down_integer_values(size_t value_count); /** - * A stable 16-byte variable-width binary view. + * Free an expression. + * + * # Safety * - * Values of 12 bytes or fewer occupy data directly. Longer values store four - * prefix bytes, a uint32_t buffer index, and a uint32_t byte offset in data. - * The host must reject lengths above INT32_MAX. Each outlined range fits its - * payload buffer. + * `expression` must be null or an owned expression handle. */ -typedef struct vx_velox_binary_view { - uint32_t length; - uint8_t data[12]; -} vx_velox_binary_view; +void vx_velox_expression_free(const vx_velox_expression *expression); /** - * A canonical UTF-8 or binary payload and its owner. - * - * The view buffer and each payload buffer stay valid through visit_varbin. - * The host must retain buffers before it stores pointers after the callback. - * buffers.retained_bytes includes all retained allocation capacities. - */ -typedef struct vx_velox_varbin_view { - size_t struct_size; - vx_velox_varbin_kind kind; - size_t length; - const vx_velox_binary_view *views; - size_t views_length; - const vx_velox_byte_buffer_view *data_buffers; - size_t data_buffer_count; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t views_alignment; - size_t validity_alignment; -} vx_velox_varbin_view; + * Free a data source. + * + * # Safety + * + * `data_source` must be null or an owned data-source handle. + */ +void vx_velox_data_source_free(const vx_velox_data_source *data_source); /** - * A canonical packed Boolean payload and its owner. - * - * The value and validity buffers use least-significant-bit-first order. - * The host must retain buffers before it stores pointers after the callback. - */ -typedef struct vx_velox_bool_view { - size_t struct_size; - size_t length; - const uint8_t *values; - size_t values_length; - size_t values_bit_offset; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t values_alignment; - size_t validity_alignment; -} vx_velox_bool_view; + * Start a scan through the stable adapter options. + * + * # Safety + * + * Every pointer must satisfy the adapter header contract. + */ +vx_velox_scan *vx_velox_data_source_scan(const vx_velox_data_source *data_source, + const vx_velox_scan_options *options, + vx_velox_error **error_out); /** - * A dictionary payload for one output window. + * Free a scan. * - * codes owns the integer code buffers. values remains valid only during the - * callback. The host can visit the prepared cursor during that callback. + * # Safety + * + * `scan` must be null or an owned scan handle. */ -typedef struct vx_velox_dictionary_view { - size_t struct_size; - size_t length; - vx_velox_primitive_view codes; - const vx_velox_export_cursor *values; - size_t values_length; -} vx_velox_dictionary_view; +void vx_velox_scan_free(const vx_velox_scan *scan); /** - * A constant payload for one output window. + * Return the next partition from a scan. + * + * # Safety * - * value contains one canonical value. It remains valid only during the - * callback. The host can visit the prepared cursor during that callback. + * `scan` must point to a live scan. `error_out` must be null or valid. */ -typedef struct vx_velox_constant_view { - size_t struct_size; - size_t length; - const vx_velox_export_cursor *value; -} vx_velox_constant_view; +vx_velox_partition *vx_velox_scan_next_partition(vx_velox_scan *scan, vx_velox_error **error_out); /** - * A canonical struct payload for one output window. - * - * fields contains borrowed prepared cursors in declaration order. The host - * visits each field cursor at offset for length rows during this callback. - * buffers owns only the parent validity buffer. - */ -typedef struct vx_velox_struct_view { - size_t struct_size; - size_t length; - size_t offset; - const vx_velox_export_cursor *const *fields; - size_t field_count; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t validity_alignment; -} vx_velox_struct_view; + * Free a partition. + * + * # Safety + * + * `partition` must be null or an owned partition handle. + */ +void vx_velox_partition_free(const vx_velox_partition *partition); /** - * A canonical list window. - * - * offsets and sizes start at the requested parent window. Offset values remain - * absolute against the complete elements cursor. buffers retains the complete - * prepared metadata allocation. The host must retain buffers before it stores - * a metadata pointer. elements is borrowed during the callback. Vectors - * imported from elements retain their own owners. - */ -typedef struct vx_velox_list_view { - size_t struct_size; - size_t length; - const int32_t *offsets; - const int32_t *sizes; - const vx_velox_export_cursor *elements; - size_t elements_length; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t offsets_alignment; - size_t sizes_alignment; - size_t validity_alignment; -} vx_velox_list_view; + * Return the next array from a partition. + * + * # Safety + * + * `partition` must point to a live partition. `error_out` must be null or valid. + */ +const vx_velox_array *vx_velox_partition_next(vx_velox_partition *partition, + vx_velox_error **error_out); /** - * A canonical map window. - * - * offsets and sizes start at the requested parent window. Offset values remain - * absolute against the complete key and value cursors. buffers retains the - * complete prepared metadata allocation. The child cursors are borrowed during - * the callback. Vectors imported from them retain their own owners. - */ -typedef struct vx_velox_map_view { - size_t struct_size; - size_t length; - const int32_t *offsets; - const int32_t *sizes; - const vx_velox_export_cursor *keys; - const vx_velox_export_cursor *values; - size_t entries_length; - bool keys_sorted; - vx_velox_validity_kind validity_kind; - const uint8_t *validity; - size_t validity_length; - size_t validity_bit_offset; - vx_velox_buffer_owner buffers; - size_t offsets_alignment; - size_t sizes_alignment; - size_t validity_alignment; -} vx_velox_map_view; - -typedef struct vx_velox_visit_request { - size_t struct_size; - const uint64_t *rows; - size_t row_count; -} vx_velox_visit_request; + * Free an array. + * + * # Safety + * + * `array` must be null or an owned array handle. + */ +void vx_velox_array_free(const vx_velox_array *array); /** - * Host callbacks for one Vortex array visit. - * - * One array visit calls the matching callback synchronously. If the host shares this - * table between simultaneous visits, callbacks can occur concurrently. - * last_error returns the calling thread's most recent visitor error. Its string - * remains valid until the next callback on that thread. - * - * Every callback must catch C++ exceptions. No callback can unwind across this - * C ABI. The host owns context and must keep it live until each visit returns. - */ -typedef struct vx_velox_visitor { - size_t struct_size; - uint32_t abi_version; - void *context; - int32_t (*visit_primitive)(void *context, const vx_velox_primitive_view *view); - const char *(*last_error)(void *context); - int32_t (*visit_varbin)(void *context, const vx_velox_varbin_view *view); - int32_t (*visit_dictionary)(void *context, const vx_velox_dictionary_view *view); - int32_t (*visit_constant)(void *context, const vx_velox_constant_view *view); - int32_t (*visit_bool)(void *context, const vx_velox_bool_view *view); - int32_t (*visit_struct)(void *context, const vx_velox_struct_view *view); - int32_t (*visit_list)(void *context, const vx_velox_list_view *view); - int32_t (*visit_map)(void *context, const vx_velox_map_view *view); -} vx_velox_visitor; + * Return an array length. + * + * # Safety + * + * `array` must point to a live array. + */ +size_t vx_velox_array_len(const vx_velox_array *array); /** - * Host memory callbacks for one Arrow C Data export. + * Slice an array. * - * Before Arrow conversion, Vortex requests a conservative reservation through - * report_allocation. A rejection stops the export before Arrow allocations. - * Vortex calls report_free after conversion to refund unused reservation bytes. - * The remaining charge equals retained Arrow payload capacities. It excludes - * the schema and small Arrow C Data metadata allocations. - * If the actual charge exceeds the reservation, Vortex requests the difference - * before it returns outputs. A rejection aborts the export and frees the data. - * - * A final Arrow release calls report_free for the remaining charge. It also - * calls release_context. These calls can occur on any thread. All callbacks - * must be thread-safe and must not unwind across the C ABI. - * - * last_error returns the calling thread's most recent allocation error. The - * string remains valid until the next callback on that thread. - */ -typedef struct vx_velox_arrow_memory_callbacks { - size_t struct_size; - uint32_t abi_version; - void *context; - void (*retain_context)(void *context); - void (*release_context)(void *context); - int32_t (*report_allocation)(void *context, size_t retained_bytes); - void (*report_free)(void *context, size_t retained_bytes); - const char *(*last_error)(void *context); -} vx_velox_arrow_memory_callbacks; + * # Safety + * + * `array` must point to a live array. `error_out` must be null or valid. + */ +const vx_velox_array *vx_velox_array_slice(const vx_velox_array *array, + size_t begin, + size_t end, + vx_velox_error **error_out); -uint32_t vx_velox_abi_version(void); -uint64_t vx_velox_capabilities(void); +/** + * Return one struct field with the supplied session. + * + * # Safety + * + * The session and array pointers must identify live handles. `error_out` must be null or valid. + */ +const vx_velox_array *vx_velox_array_get_field(const vx_velox_session *session, + const vx_velox_array *array, + size_t index, + vx_velox_error **error_out); -vx_velox_view vx_velox_error_message(const vx_velox_error *error); -void vx_velox_error_free(const vx_velox_error *error); -vx_velox_session *vx_velox_session_new(void); -vx_velox_session *vx_velox_session_clone(const vx_velox_session *session); -void vx_velox_session_free(const vx_velox_session *session); +/** + * Return the invalid value count with the supplied session. + * + * # Safety + * + * The session and array pointers must identify live handles. `error_out` must be null or valid. + */ +size_t vx_velox_array_invalid_count(const vx_velox_session *session, + const vx_velox_array *array, + vx_velox_error **error_out); -const vx_velox_dtype * -vx_velox_dtype_new_primitive(vx_velox_ptype ptype, bool nullable, vx_velox_error **error_out); -void vx_velox_dtype_free(const vx_velox_dtype *dtype); -vx_velox_scalar *vx_velox_scalar_new_bool(bool value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_i8(int8_t value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_i16(int16_t value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_i32(int32_t value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_date_days(int32_t value, bool nullable, vx_velox_error **error_out); -vx_velox_scalar *vx_velox_scalar_new_i64(int64_t value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_f32(float value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_f64(double value, bool nullable); -vx_velox_scalar *vx_velox_scalar_new_utf8(vx_velox_view value, bool nullable, vx_velox_error **error_out); -vx_velox_scalar * -vx_velox_scalar_new_binary(const uint8_t *data, size_t length, bool nullable, vx_velox_error **error_out); -vx_velox_scalar *vx_velox_scalar_new_list(const vx_velox_dtype *element_dtype, - const vx_velox_scalar *const *elements, - size_t length, - bool nullable, - vx_velox_error **error_out); -void vx_velox_scalar_free(const vx_velox_scalar *scalar); +/** + * Export one Vortex array through the Arrow C Data Interface. + * + * The caller owns both outputs and must call their release callbacks. The memory callbacks reserve + * a conservative payload charge before Arrow conversion. The adapter refunds the difference after + * it knows the retained payload capacities. It requests a deficit before it returns the outputs. + * The charge excludes schema and small FFI metadata. + * + * # Safety + * + * The session and array pointers must identify live handles. `memory_callbacks` must identify its + * declared `struct_size` bytes for this call. Its callback context must remain valid through every + * retained reference. Its callbacks and returned error strings must satisfy the header contract + * and must not unwind. Both output pointers must identify uninitialized writable structures. + * `error_out` must be null or identify writable storage for one error pointer. + */ +int32_t vx_velox_array_export_arrow(const vx_velox_session *session, + const vx_velox_array *array, + const vx_velox_arrow_memory_callbacks *memory_callbacks, + ArrowSchema *schema_out, + ArrowArray *array_out, + vx_velox_error **error_out); -vx_velox_expression *vx_velox_expression_root(void); -vx_velox_expression *vx_velox_expression_literal(const vx_velox_scalar *scalar, vx_velox_error **error_out); -vx_velox_expression *vx_velox_expression_get_item(vx_velox_view name, const vx_velox_expression *child); -vx_velox_expression *vx_velox_expression_binary(vx_velox_binary_operator operation, - const vx_velox_expression *left, - const vx_velox_expression *right, +/** + * Create a struct projection from the supplied field names. + * + * The returned expression stays owned by the caller. + * + * # Safety + * + * `names` must be null when `len` is zero or point to `len` valid views. + * Every view must remain valid for this call. + * `error_out` must be null or point to writable storage. No input operation can unwind. + */ +vx_velox_expression *vx_velox_expression_select(const vx_velox_view *names, + size_t len, vx_velox_error **error_out); -vx_velox_expression *vx_velox_expression_and(const vx_velox_expression *const *expressions, size_t length); -vx_velox_expression *vx_velox_expression_or(const vx_velox_expression *const *expressions, size_t length); -vx_velox_expression *vx_velox_expression_not(const vx_velox_expression *child); -vx_velox_expression *vx_velox_expression_is_null(const vx_velox_expression *child); -vx_velox_expression *vx_velox_expression_list_contains(const vx_velox_expression *list, - const vx_velox_expression *value); -bool vx_velox_can_push_down_integer_values(size_t value_count); -void vx_velox_expression_free(const vx_velox_expression *expression); -vx_velox_expression * -vx_velox_expression_select(const vx_velox_view *names, size_t length, vx_velox_error **error_out); + +/** + * Create a struct projection with an absolute file-row index as its first field. + * + * The remaining fields select the supplied names from the scan root. The + * returned expression stays owned by the caller. + * + * # Safety + * + * `names` must be null when `len` is zero or point to `len` valid views. + * Every view and `row_index_name` must remain valid for this call. + * `error_out` must be null or point to writable storage. No input operation can unwind. + */ vx_velox_expression *vx_velox_expression_select_with_row_index(const vx_velox_view *names, - size_t length, + size_t len, vx_velox_view row_index_name, vx_velox_error **error_out); -/* - * On success, the reader owns context and calls release_context once. On - * failure, the caller still owns context. +/** + * Create a Vortex positional reader from Velox callbacks. + * + * # Safety + * + * `callbacks` must point to a valid callback structure. Every callback and its context must be + * thread-safe and must not unwind. `error_out` must be null or valid for one error pointer. */ vx_velox_read_at *vx_velox_read_at_new(const vx_velox_read_at_callbacks *callbacks, vx_velox_error **error_out); + +/** + * Free a Vortex positional reader. + * + * # Safety + * + * `reader` must be null or a pointer returned by [`vx_velox_read_at_new`]. + */ void vx_velox_read_at_free(vx_velox_read_at *reader); + +/** + * Return the size of a callback-backed source. + * + * This entry point validates the host callback contract before file-reader code consumes the + * source. + * + * # Safety + * + * `reader` must point to a live reader. `error_out` must be null or valid for one error pointer. + */ uint64_t vx_velox_read_at_size(const vx_velox_read_at *reader, vx_velox_error **error_out); +/** + * Export an opened source schema through the Arrow C Data Interface. + * + * The caller owns the output and must invoke its release callback. + * + * # Safety + * + * `source` must point to a live source. `schema_out` must identify uninitialized writable + * storage. `error_out` must be null or valid for one error pointer. + */ +int32_t vx_velox_source_export_schema(const vx_velox_source *source, + ArrowSchema *schema_out, + vx_velox_error **error_out); + +/** + * Open a Vortex source through a callback reader. + * + * The source retains the session and reader state. The caller can free both input handles after + * this function returns. + * + * # Safety + * + * `session` and `reader` must point to live handles. `error_out` must be null or valid for one + * error pointer. + */ vx_velox_source *vx_velox_source_new(const vx_velox_session *session, const vx_velox_read_at *reader, vx_velox_error **error_out); + +/** + * Free a callback-backed Vortex source and release its callback-owned input buffers. + * + * # Safety + * + * `source` must be null or a pointer returned by [`vx_velox_source_new`]. + */ void vx_velox_source_free(vx_velox_source *source); + +/** + * Return the file row count. + * + * # Safety + * + * `source` must point to a live source. + */ uint64_t vx_velox_source_row_count(const vx_velox_source *source); + +/** + * Return the file size in bytes. + * + * # Safety + * + * `source` must point to a live source. + */ uint64_t vx_velox_source_file_size(const vx_velox_source *source); -int32_t vx_velox_source_export_schema(const vx_velox_source *source, - struct ArrowSchema *schema_out, - vx_velox_error **error_out); + +/** + * Return the number of natural row splits. + * + * # Safety + * + * `source` must point to a live source. + */ size_t vx_velox_source_natural_split_count(const vx_velox_source *source); + +/** + * Write one natural row split. + * + * # Safety + * + * `source` must point to a live source. `split_out` must point to a structure with a valid size. + * `error_out` must be null or valid for one error pointer. + */ int32_t vx_velox_source_natural_split_at(const vx_velox_source *source, size_t index, vx_velox_natural_split *split_out, vx_velox_error **error_out); + +/** + * Evaluate whether natural splits cannot match an expression. + * + * Each output byte is one when the matching split cannot produce a true expression result. Zero + * means that the split can match or that available statistics cannot prove exclusion. + * + * # Safety + * + * `source` and `expression` must point to live handles. `pruned_out` must identify `split_count` + * writable bytes unless `split_count` is zero. `error_out` must be null or valid for one error + * pointer. + */ int32_t vx_velox_source_prune_natural_splits(const vx_velox_source *source, const vx_velox_expression *expression, size_t first_split, size_t split_count, uint8_t *pruned_out, vx_velox_error **error_out); + +/** + * Create a standard Vortex data source for this file. + * + * The caller owns the returned handle and must free it through `vx_data_source_free`. + * + * # Safety + * + * `source` must point to a live source. `error_out` must be null or valid for one error pointer. + */ const vx_velox_data_source *vx_velox_source_data_source(const vx_velox_source *source, vx_velox_error **error_out); -void vx_velox_data_source_free(const vx_velox_data_source *data_source); -vx_velox_scan *vx_velox_data_source_scan(const vx_velox_data_source *data_source, - const vx_velox_scan_options *options, - vx_velox_error **error_out); -void vx_velox_scan_free(const vx_velox_scan *scan); -vx_velox_partition *vx_velox_scan_next_partition(vx_velox_scan *scan, vx_velox_error **error_out); -void vx_velox_partition_free(const vx_velox_partition *partition); -const vx_velox_array *vx_velox_partition_next(vx_velox_partition *partition, vx_velox_error **error_out); -void vx_velox_array_free(const vx_velox_array *array); -size_t vx_velox_array_len(const vx_velox_array *array); -const vx_velox_array * -vx_velox_array_slice(const vx_velox_array *array, size_t begin, size_t end, vx_velox_error **error_out); -const vx_velox_array *vx_velox_array_get_field(const vx_velox_session *session, - const vx_velox_array *array, - size_t index, - vx_velox_error **error_out); -size_t vx_velox_array_invalid_count(const vx_velox_session *session, - const vx_velox_array *array, - vx_velox_error **error_out); -int32_t vx_velox_array_visit(const vx_velox_session *session, - const vx_velox_array *array, - const vx_velox_visit_request *request, - const vx_velox_visitor *visitor, - vx_velox_error **error_out); - /** - * Create one prepared exporter for several engine-sized output windows. + * Create one export cursor for several Velox output windows. * - * memory_callbacks must identify a complete, thread-safe callback table. - * The exporter retains its callback context until the last buffer owner releases it. + * # Safety + * + * The session and array pointers must identify live handles. + * The memory callbacks must identify a complete, thread-safe callback table. + * `error_out` must be null or valid. */ vx_velox_export_cursor *vx_velox_export_cursor_new(const vx_velox_session *session, const vx_velox_array *array, const vx_velox_arrow_memory_callbacks *memory_callbacks, vx_velox_error **error_out); -/** Free one prepared exporter. */ +/** + * Free one export cursor. + * + * # Safety + * + * The pointer must be null or come from [`vx_velox_export_cursor_new`]. + */ void vx_velox_export_cursor_free(vx_velox_export_cursor *cursor); /** - * Visit one contiguous output window from a prepared exporter. + * Visit one contiguous range from a retained export cursor. + * + * # Safety * - * Concurrent visit calls are valid. Do not free the cursor before all visits return. + * The cursor and visitor pointers must remain live until this call returns. + * Concurrent calls are valid. The caller must not free the cursor before all calls return. */ int32_t vx_velox_export_cursor_visit(const vx_velox_export_cursor *cursor, size_t offset, @@ -574,13 +1748,23 @@ int32_t vx_velox_export_cursor_visit(const vx_velox_export_cursor *cursor, const vx_velox_visitor *visitor, vx_velox_error **error_out); -int32_t vx_velox_array_export_arrow(const vx_velox_session *session, - const vx_velox_array *array, - const vx_velox_arrow_memory_callbacks *memory_callbacks, - struct ArrowSchema *schema_out, - struct ArrowArray *array_out, - vx_velox_error **error_out); +/** + * Visit one Vortex array through host semantic callbacks. + * + * The request selects source positions once. Callback block positions are compact and follow the + * request order. + * + * # Safety + * + * Every pointer must be null or valid for the documented access. The array and session handles + * must remain live until this call returns. + */ +int32_t vx_velox_array_visit(const vx_velox_session *session, + const vx_velox_array *array, + const vx_velox_visit_request *request, + const vx_velox_visitor *visitor, + vx_velox_error **error_out); #ifdef __cplusplus -} -#endif +} // extern "C" +#endif // __cplusplus diff --git a/vortex-velox/src/lib.rs b/vortex-velox/src/lib.rs index 3ead79e66f2..403ca0a179e 100644 --- a/vortex-velox/src/lib.rs +++ b/vortex-velox/src/lib.rs @@ -53,7 +53,7 @@ pub use visitor::vx_velox_visit_request; pub use visitor::vx_velox_visitor; /// The current major version of the Vortex and Velox adapter ABI. -pub const VX_VELOX_ABI_VERSION: u32 = 5; +pub const VX_VELOX_ABI_VERSION: u32 = 6; /// The adapter supports batched host range reads. pub const VX_VELOX_CAPABILITY_BATCH_READ: u64 = 1 << 0; @@ -111,6 +111,8 @@ pub const VX_VELOX_CAPABILITY_STRUCT_VISITOR: u64 = 1 << 17; pub const VX_VELOX_CAPABILITY_LIST_VISITOR: u64 = 1 << 18; /// The adapter can preserve canonical map children during native export. pub const VX_VELOX_CAPABILITY_MAP_VISITOR: u64 = 1 << 19; +/// Natural splits include a stable byte-range assignment token. +pub const VX_VELOX_CAPABILITY_SPLIT_ASSIGNMENT: u64 = 1 << 20; /// Return the adapter ABI version. #[unsafe(no_mangle)] @@ -141,6 +143,7 @@ pub extern "C" fn vx_velox_capabilities() -> u64 { | VX_VELOX_CAPABILITY_STRUCT_VISITOR | VX_VELOX_CAPABILITY_LIST_VISITOR | VX_VELOX_CAPABILITY_MAP_VISITOR + | VX_VELOX_CAPABILITY_SPLIT_ASSIGNMENT } #[cfg(test)] diff --git a/vortex-velox/src/read_at.rs b/vortex-velox/src/read_at.rs index 19995da0fb4..58b18a1955b 100644 --- a/vortex-velox/src/read_at.rs +++ b/vortex-velox/src/read_at.rs @@ -901,6 +901,7 @@ mod tests { let split_count = unsafe { vx_velox_source_natural_split_count(source) }; assert!(split_count > 0); let mut previous_end = 0; + let mut previous_assignment_byte = 0; for index in 0..split_count { let mut split = vx_velox_natural_split { struct_size: size_of::(), @@ -914,7 +915,13 @@ mod tests { vortex_ensure!(error.is_null(), "natural split lookup returned an error"); assert_eq!(split.row_begin, previous_end); assert!(split.row_end > split.row_begin); + assert!(split.assignment_byte >= previous_assignment_byte); + assert!(split.assignment_byte < bytes.len() as u64); + if index == 0 { + assert_eq!(split.assignment_byte, 0); + } previous_end = split.row_end; + previous_assignment_byte = split.assignment_byte; } assert_eq!(previous_end, ROW_COUNT); diff --git a/vortex-velox/src/source.rs b/vortex-velox/src/source.rs index 15cfa32d9aa..fabb60d9165 100644 --- a/vortex-velox/src/source.rs +++ b/vortex-velox/src/source.rs @@ -8,6 +8,7 @@ use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::layout::scan::multi::MultiLayoutDataSource; use vortex::mask::Mask; +use vortex_error::VortexExpect; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_io::VortexReadAt; @@ -34,6 +35,8 @@ pub struct vx_velox_natural_split { pub row_begin: u64, /// One past the final row in the split. pub row_end: u64, + /// The file byte that assigns this split to one external byte range. + pub assignment_byte: u64, } /// An opened Vortex file that uses Velox callbacks for all reads. @@ -186,10 +189,26 @@ pub unsafe extern "C-unwind" fn vx_velox_source_natural_split_at( .ok_or_else(|| vortex_err!("Natural split index out of bounds: {}", index))?; split_out.row_begin = split.start; split_out.row_end = split.end; + split_out.assignment_byte = + split_assignment_byte(index, split, source.file.row_count(), source.file_size); Ok(0) }) } +fn split_assignment_byte(index: usize, split: &Range, row_count: u64, file_size: u64) -> u64 { + if index == 0 && split.start == 0 { + return 0; + } + if row_count == 0 { + return 0; + } + + let midpoint_row = split.start + (split.end - split.start) / 2; + let assignment_byte = + (u128::from(midpoint_row) * u128::from(file_size)) / u128::from(row_count); + u64::try_from(assignment_byte).vortex_expect("The split assignment byte must fit in u64") +} + /// Evaluate whether natural splits cannot match an expression. /// /// Each output byte is one when the matching split cannot produce a true expression result. Zero diff --git a/vortex-velox/tests/abi_contract.rs b/vortex-velox/tests/abi_contract.rs index 250451fdd22..6fb966615f8 100644 --- a/vortex-velox/tests/abi_contract.rs +++ b/vortex-velox/tests/abi_contract.rs @@ -132,7 +132,10 @@ mod tests { concurrency, ] ); - check_layout!(vx_velox_natural_split, [struct_size, row_begin, row_end]); + check_layout!( + vx_velox_natural_split, + [struct_size, row_begin, row_end, assignment_byte] + ); check_layout!( vx_velox_buffer_owner, [struct_size, owner, retain, release, retained_bytes] diff --git a/vortex-velox/tests/velox_include_contract.cpp b/vortex-velox/tests/velox_include_contract.cpp index f32a5dfff8c..ffd341d7069 100644 --- a/vortex-velox/tests/velox_include_contract.cpp +++ b/vortex-velox/tests/velox_include_contract.cpp @@ -5,7 +5,7 @@ struct ArrowSchema; struct ArrowArray; #include "vortex_velox.h" -static_assert(VX_VELOX_ABI_VERSION == 5u); +static_assert(VX_VELOX_ABI_VERSION == 6u); static_assert(VX_VELOX_SELECTION_ALL == 0); static_assert(VX_VELOX_OPERATOR_EQ == 0); From 1a8ab432c1e0b09ed135e7a1b957fc459b3076af Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 15:55:36 -0400 Subject: [PATCH 6/6] refactor(vortex-velox): Separate visitor implementation Signed-off-by: Will Manning --- vortex-velox/src/visitor.rs | 3500 +--------------------- vortex-velox/src/visitor/export.rs | 1970 ++++++++++++ vortex-velox/src/visitor/export/tests.rs | 1523 ++++++++++ 3 files changed, 3495 insertions(+), 3498 deletions(-) create mode 100644 vortex-velox/src/visitor/export.rs create mode 100644 vortex-velox/src/visitor/export/tests.rs diff --git a/vortex-velox/src/visitor.rs b/vortex-velox/src/visitor.rs index be7dce2346e..1e6ef97a814 100644 --- a/vortex-velox/src/visitor.rs +++ b/vortex-velox/src/visitor.rs @@ -3,68 +3,10 @@ use std::ffi::c_char; use std::ffi::c_void; -use std::mem::MaybeUninit; -use std::mem::align_of; -use std::mem::size_of; -use std::mem::size_of_val; -use std::ptr; -use std::slice; -use std::sync::Arc; -use vortex::array::Canonical; -use vortex::array::IntoArray; -use vortex::array::VortexSessionExecute; -use vortex::array::arrays::Constant; -use vortex::array::arrays::ConstantArray; -use vortex::array::arrays::DecimalArray; -use vortex::array::arrays::Dict; -use vortex::array::arrays::Extension; -use vortex::array::arrays::ExtensionArray; -use vortex::array::arrays::ListView; -use vortex::array::arrays::ListViewArray; -use vortex::array::arrays::MapArray; -use vortex::array::arrays::PrimitiveArray; -use vortex::array::arrays::StructArray; -use vortex::array::arrays::VarBinViewArray; -use vortex::array::arrays::decimal::DecimalArrayExt; -use vortex::array::arrays::extension::ExtensionArrayExt; -use vortex::array::arrays::listview::ListViewArrayExt; -use vortex::array::arrays::listview::ListViewArraySlotsExt; -use vortex::array::arrays::map::MapArrayExt; -use vortex::array::arrays::map::MapArraySlotsExt; -use vortex::array::arrays::primitive::PrimitiveArrayExt; -use vortex::array::arrays::struct_::StructArrayExt; -use vortex::array::buffer::BufferHandle; -use vortex::array::match_each_unsigned_integer_ptype; -use vortex::buffer::Buffer; -use vortex::buffer::BufferMut; -use vortex::buffer::ByteBuffer; -use vortex::dtype::DType; -use vortex::dtype::DecimalType; -use vortex::dtype::NativeDecimalType; -use vortex::dtype::PType; -use vortex::extension::datetime::Date; -use vortex::extension::datetime::TimeUnit; -use vortex::mask::Mask; -use vortex_array::ArrayView; -use vortex_array::arrays::dict::DictArraySlotsExt; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; -use vortex_fastlanes::BitPacked; -use vortex_fastlanes::BitPackedArrayExt; -use vortex_fastlanes::FL_CHUNK_SIZE; +mod export; -use crate::array::ArrowMemoryReservation; -use crate::array::conservative_export_reservation; -use crate::array::parse_memory_callbacks; -use crate::array::vx_velox_arrow_memory_callbacks; -use crate::ffi::try_or; -use crate::ffi::vx_array_ref; -use crate::ffi::vx_session_ref; -use crate::ffi::vx_velox_array; -use crate::ffi::vx_velox_error; -use crate::ffi::vx_velox_session; +pub use export::vx_velox_export_cursor; /// A fixed-width primitive value identifier in a semantic visitor block. pub type vx_velox_primitive_type = u32; @@ -93,22 +35,6 @@ pub const VX_VELOX_PRIMITIVE_F64: vx_velox_primitive_type = 10; /// Signed 128-bit primitive identifier. pub const VX_VELOX_PRIMITIVE_I128: vx_velox_primitive_type = 11; -fn primitive_type_id(value: PType) -> vx_velox_primitive_type { - match value { - PType::U8 => VX_VELOX_PRIMITIVE_U8, - PType::U16 => VX_VELOX_PRIMITIVE_U16, - PType::U32 => VX_VELOX_PRIMITIVE_U32, - PType::U64 => VX_VELOX_PRIMITIVE_U64, - PType::I8 => VX_VELOX_PRIMITIVE_I8, - PType::I16 => VX_VELOX_PRIMITIVE_I16, - PType::I32 => VX_VELOX_PRIMITIVE_I32, - PType::I64 => VX_VELOX_PRIMITIVE_I64, - PType::F16 => VX_VELOX_PRIMITIVE_F16, - PType::F32 => VX_VELOX_PRIMITIVE_F32, - PType::F64 => VX_VELOX_PRIMITIVE_F64, - } -} - /// A fixed-width validity representation identifier for one visitor block. pub type vx_velox_validity_kind = u32; /// The type is not nullable. @@ -448,3425 +374,3 @@ pub struct vx_velox_visitor { pub visit_map: Option i32>, } - -/// Retains one prepared Vortex array across several Velox output windows. -#[repr(C)] -pub struct vx_velox_export_cursor { - export: CursorExport, -} - -enum CursorExport { - Primitive(PrimitiveExport), - Bool(BoolExport), - VarBin(VarBinExport), - Dictionary(DictionaryExport), - Constant(ConstantExport), - Struct(StructExport), - List(ListExport), - Map(MapExport), -} - -struct PackedBits(Box<[u64]>); - -impl PackedBits { - fn try_new(bits: vortex::buffer::BitBuffer) -> VortexResult<(Self, usize)> { - let compact = bits - .chunks() - .iter_padded() - .collect::>() - .into_boxed_slice(); - let allocation = size_of_val(compact.as_ref()); - Ok((Self(compact), allocation)) - } - - fn as_ptr(&self) -> *const u8 { - self.0.as_ptr().cast() - } - - fn len(&self) -> usize { - size_of_val(self.0.as_ref()) - } -} - -struct BoolOwner { - values: PackedBits, - validity: Option, - retained_bytes: usize, - memory_reservation: Option, -} - -impl BoolOwner { - fn try_new( - values: vortex::buffer::BitBuffer, - validity: Option, - ) -> VortexResult { - let (values, values_allocation) = PackedBits::try_new(values)?; - let (validity, validity_allocation) = match validity { - Some(validity) => { - let (validity, allocation) = PackedBits::try_new(validity)?; - (Some(validity), allocation) - } - None => (None, 0), - }; - let retained_bytes = values_allocation - .checked_add(validity_allocation) - .ok_or_else(|| vortex_err!("Boolean visitor retained byte count overflow"))?; - Ok(Self { - values, - validity, - retained_bytes, - memory_reservation: None, - }) - } - - fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { - self.memory_reservation = Some(reservation); - } -} - -enum PrimitiveValues { - Compact64(Box<[MaybeUninit]>), - Compact128(Box<[MaybeUninit]>), - Retained(ByteBuffer), -} - -impl PrimitiveValues { - fn as_ptr(&self) -> *const u8 { - match self { - Self::Compact64(values) => values.as_ptr().cast(), - Self::Compact128(values) => values.as_ptr().cast(), - Self::Retained(values) => values.as_ptr(), - } - } -} - -struct PrimitiveOwner { - values: PrimitiveValues, - values_length: usize, - validity: Option, - retained_bytes: usize, - memory_reservation: Option, -} - -enum RetainedBytes { - Retained(ByteBuffer), - Compact(Box<[u8]>), -} - -impl RetainedBytes { - fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { - let buffer = handle.try_into_host_sync()?; - let length = buffer.len(); - match buffer.try_into_mut() { - Ok(buffer) => { - let allocation_size = buffer.allocation_size(); - Ok((Self::Retained(buffer.freeze()), allocation_size)) - } - Err(buffer) => { - let compact = buffer.as_slice().to_vec().into_boxed_slice(); - Ok((Self::Compact(compact), length)) - } - } - } - - fn as_ptr(&self) -> *const u8 { - match self { - Self::Retained(buffer) => buffer.as_ptr(), - Self::Compact(buffer) => buffer.as_ptr(), - } - } - - fn len(&self) -> usize { - match self { - Self::Retained(buffer) => buffer.len(), - Self::Compact(buffer) => buffer.len(), - } - } -} - -enum RetainedViews { - Retained(ByteBuffer), - Compact(Box<[vx_velox_binary_view]>), -} - -impl RetainedViews { - fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { - let buffer = handle.try_into_host_sync()?; - if !buffer - .len() - .is_multiple_of(size_of::()) - { - vortex_bail!( - "Vortex variable-width view buffer has an invalid byte length: {}", - buffer.len() - ); - } - match buffer.try_into_mut() { - Ok(buffer) => { - let allocation_size = buffer.allocation_size(); - Ok((Self::Retained(buffer.freeze()), allocation_size)) - } - Err(buffer) => { - let length = buffer.len() / size_of::(); - let mut compact = vec![ - vx_velox_binary_view { - length: 0, - data: [0; 12], - }; - length - ] - .into_boxed_slice(); - if !buffer.is_empty() { - // SAFETY: Both byte ranges have the checked identical size. - unsafe { - ptr::copy_nonoverlapping( - buffer.as_ptr(), - compact.as_mut_ptr().cast::(), - buffer.len(), - ) - }; - } - let allocation = size_of_val(compact.as_ref()); - Ok((Self::Compact(compact), allocation)) - } - } - } - - fn as_ptr(&self) -> *const vx_velox_binary_view { - match self { - Self::Retained(buffer) => buffer.as_ptr().cast(), - Self::Compact(buffer) => buffer.as_ptr(), - } - } -} - -struct VarBinOwner { - views: RetainedViews, - _data: Box<[RetainedBytes]>, - descriptors: Box<[vx_velox_byte_buffer_view]>, - validity: Option, - retained_bytes: usize, - memory_reservation: Option, -} - -// SAFETY: The owner never mutates its buffers or pointer descriptors after construction. -// Every descriptor points into an immutable allocation that the same owner retains. -unsafe impl Send for VarBinOwner {} -// SAFETY: Shared access only reads immutable buffers and descriptors retained by this owner. -unsafe impl Sync for VarBinOwner {} - -impl VarBinOwner { - fn try_new( - views: BufferHandle, - mut buffers: Arc<[BufferHandle]>, - validity: Option, - length: usize, - ) -> VortexResult { - let (views, views_allocation) = RetainedViews::try_new(views)?; - let handles = if let Some(handles) = Arc::get_mut(&mut buffers) { - handles - .iter_mut() - .map(|handle| { - std::mem::replace(handle, BufferHandle::new_host(ByteBuffer::empty())) - }) - .collect::>() - } else { - buffers.iter().cloned().collect::>() - }; - let mut data_allocation = 0usize; - let data = handles - .into_iter() - .map(|handle| { - let (buffer, allocation) = RetainedBytes::try_new(handle)?; - data_allocation = data_allocation - .checked_add(allocation) - .ok_or_else(|| vortex_err!("Vortex string payload allocation overflow"))?; - Ok(buffer) - }) - .collect::>>()? - .into_boxed_slice(); - let descriptors = data - .iter() - .map(|buffer| vx_velox_byte_buffer_view { - data: buffer.as_ptr(), - length: buffer.len(), - }) - .collect::>() - .into_boxed_slice(); - let descriptor_allocation = size_of_val(descriptors.as_ref()); - let (validity, validity_allocation) = retain_validity(validity, length)?; - let retained_bytes = views_allocation - .checked_add(data_allocation) - .and_then(|bytes| bytes.checked_add(descriptor_allocation)) - .and_then(|bytes| bytes.checked_add(validity_allocation)) - .ok_or_else(|| vortex_err!("Vortex string retained byte count overflow"))?; - Ok(Self { - views, - _data: data, - descriptors, - validity, - retained_bytes, - memory_reservation: None, - }) - } - - fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { - self.memory_reservation = Some(reservation); - } -} - -fn retain_validity( - validity: Option, - length: usize, -) -> VortexResult<(Option, usize)> { - let Some(validity) = validity else { - return Ok((None, 0)); - }; - if validity.len() < length { - vortex_bail!( - "Vortex validity length is too small: {} for {length} values", - validity.len() - ); - } - let validity = if validity.len() == length { - validity - } else { - validity.slice(..length) - }; - let (validity, allocation) = PackedBits::try_new(validity)?; - Ok((Some(validity), allocation)) -} - -impl PrimitiveOwner { - fn try_allocate( - values_length: usize, - values_alignment: usize, - validity: Option, - length: usize, - ) -> VortexResult { - let (values, values_allocation) = if values_alignment > align_of::() { - if values_alignment > align_of::() { - vortex_bail!( - "Primitive visitor does not support value alignment {values_alignment}" - ); - } - let values = - vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] - .into_boxed_slice(); - let allocation = values - .len() - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; - (PrimitiveValues::Compact128(values), allocation) - } else { - let values = - vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] - .into_boxed_slice(); - let allocation = values - .len() - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; - (PrimitiveValues::Compact64(values), allocation) - }; - let (validity, validity_allocation) = retain_validity(validity, length)?; - let retained_bytes = values_allocation - .checked_add(validity_allocation) - .ok_or_else(|| vortex_err!("Primitive visitor retained byte count overflow"))?; - Ok(Self { - values, - values_length, - validity, - retained_bytes, - memory_reservation: None, - }) - } - - fn try_new( - host_values: ByteBuffer, - values_alignment: usize, - validity: Option, - length: usize, - retain_values: bool, - ) -> VortexResult { - let values_length = host_values.len(); - let host_values = if retain_values { - match host_values.try_into_mut() { - Ok(values) => { - let values_allocation = values.allocation_size(); - let (validity, validity_allocation) = retain_validity(validity, length)?; - let retained_bytes = values_allocation - .checked_add(validity_allocation) - .ok_or_else(|| { - vortex_err!("Primitive visitor retained byte count overflow") - })?; - return Ok(Self { - values: PrimitiveValues::Retained(values.freeze()), - values_length, - validity, - retained_bytes, - memory_reservation: None, - }); - } - Err(values) => values, - } - } else { - host_values - }; - let mut owner = Self::try_allocate(values_length, values_alignment, validity, length)?; - if !host_values.is_empty() { - let (values_pointer, values_capacity) = match &mut owner.values { - PrimitiveValues::Compact64(values) => ( - values.as_mut_ptr().cast::(), - values.len() * size_of::(), - ), - PrimitiveValues::Compact128(values) => ( - values.as_mut_ptr().cast::(), - values.len() * size_of::(), - ), - PrimitiveValues::Retained(_) => { - unreachable!("a newly allocated primitive owner must be compact") - } - }; - // SAFETY: The byte view spans the complete compact allocation. - let values_bytes = - unsafe { slice::from_raw_parts_mut(values_pointer, values_capacity) }; - values_bytes[..values_length].copy_from_slice(host_values.as_slice()); - } - Ok(owner) - } - - fn try_new_bitpacked_i64( - array: ArrayView<'_, BitPacked>, - validity: Option, - ) -> VortexResult { - let values_length = array - .len() - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; - let mut owner = - Self::try_allocate(values_length, align_of::(), validity, array.len())?; - // SAFETY: The allocation uses `u64` alignment and contains at least `values_length` bytes. - // The output slice covers exactly `array.len()` values and remains uniquely borrowed. - let output = unsafe { - slice::from_raw_parts_mut( - match &mut owner.values { - PrimitiveValues::Compact64(values) => { - values.as_mut_ptr().cast::>() - } - PrimitiveValues::Compact128(_) | PrimitiveValues::Retained(_) => { - unreachable!("a newly allocated primitive owner must be compact") - } - }, - array.len(), - ) - }; - let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; - array.unpacked_chunks(&mut scratch)?.decode_into(output); - Ok(owner) - } - - fn values(&self) -> *const u8 { - if self.values_length == 0 { - ptr::null() - } else { - self.values.as_ptr() - } - } - - fn retained_bytes(&self) -> usize { - self.retained_bytes - } - - fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { - self.memory_reservation = Some(reservation); - } -} - -fn pointer_alignment(pointer: *const u8) -> usize { - if pointer.is_null() { - return 0; - } - 1usize << pointer.addr().trailing_zeros() -} - -fn primitive_width(primitive_type: vx_velox_primitive_type) -> VortexResult { - Ok(match primitive_type { - VX_VELOX_PRIMITIVE_U8 | VX_VELOX_PRIMITIVE_I8 => 1, - VX_VELOX_PRIMITIVE_U16 | VX_VELOX_PRIMITIVE_I16 | VX_VELOX_PRIMITIVE_F16 => 2, - VX_VELOX_PRIMITIVE_U32 | VX_VELOX_PRIMITIVE_I32 | VX_VELOX_PRIMITIVE_F32 => 4, - VX_VELOX_PRIMITIVE_U64 | VX_VELOX_PRIMITIVE_I64 | VX_VELOX_PRIMITIVE_F64 => 8, - VX_VELOX_PRIMITIVE_I128 => 16, - _ => vortex_bail!("Unknown Vortex Velox primitive type: {primitive_type}"), - }) -} - -fn cast_decimal_values(values: Buffer, validity: &Mask) -> VortexResult -where - T: NativeDecimalType, - S: NativeDecimalType, -{ - let mut output = BufferMut::::with_capacity(values.len()); - for (index, value) in values.into_iter().enumerate() { - if !validity.value(index) { - output.push(T::default()); - continue; - } - output.push(::from(value).ok_or_else(|| { - vortex_err!( - "Decimal value cannot be represented as {}", - std::any::type_name::() - ) - })?); - } - Ok(output.freeze().into_byte_buffer()) -} - -fn normalized_decimal_values(array: &DecimalArray, validity: &Mask) -> VortexResult -where - T: NativeDecimalType, -{ - if array.values_type() == T::DECIMAL_TYPE { - return array.buffer_handle().clone().try_into_host_sync(); - } - match array.values_type() { - DecimalType::I8 => cast_decimal_values::(array.buffer::(), validity), - DecimalType::I16 => cast_decimal_values::(array.buffer::(), validity), - DecimalType::I32 => cast_decimal_values::(array.buffer::(), validity), - DecimalType::I64 => cast_decimal_values::(array.buffer::(), validity), - DecimalType::I128 => cast_decimal_values::(array.buffer::(), validity), - DecimalType::I256 => cast_decimal_values::( - array.buffer::(), - validity, - ), - } -} - -struct PrimitiveExport { - primitive_type: vx_velox_primitive_type, - decimal_precision: u32, - decimal_scale: i32, - length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, -} - -impl PrimitiveExport { - fn try_new_decimal( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let retain_values = memory_callbacks.is_some(); - let mut execution = session.create_execution_ctx(); - let mut memory_reservation = match memory_callbacks { - Some(callbacks) => Some(ArrowMemoryReservation::try_new( - callbacks, - conservative_export_reservation(&array, &mut execution)?, - )?), - None => None, - }; - let is_nullable = array.dtype().is_nullable(); - let decimal = array.execute::(&mut execution)?; - let decimal_precision = u32::from(decimal.precision()); - let decimal_scale = i32::from(decimal.scale()); - let length = decimal.len(); - let mask = decimal - .as_ref() - .validity()? - .execute_mask(length, &mut execution)?; - let (primitive_type, host_values) = match decimal.precision() { - 1..=18 => ( - VX_VELOX_PRIMITIVE_I64, - normalized_decimal_values::(&decimal, &mask)?, - ), - 19..=38 => ( - VX_VELOX_PRIMITIVE_I128, - normalized_decimal_values::(&decimal, &mask)?, - ), - precision => { - vortex_bail!("Vortex Velox visitor does not support decimal precision {precision}") - } - }; - let (validity_kind, validity) = exported_validity(is_nullable, mask); - let mut owner = PrimitiveOwner::try_new( - host_values, - primitive_width(primitive_type)?, - validity, - length, - retain_values, - )?; - if let Some(mut reservation) = memory_reservation.take() { - reservation.reconcile(owner.retained_bytes())?; - owner.set_memory_reservation(reservation); - } - Ok(Self { - primitive_type, - decimal_precision, - decimal_scale, - length, - validity_kind, - owner: Arc::new(owner), - }) - } - - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let retain_values = memory_callbacks.is_some(); - let direct_bitpacked = array.as_opt::().filter(|bitpacked| { - array.dtype().as_ptype() == PType::I64 && bitpacked.patches().is_none() - }); - let values_length = - array - .len() - .checked_mul(array.dtype().element_size().ok_or_else(|| { - vortex_err!("Primitive visitor received a variable-width array") - })?) - .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; - let values_allocation = values_length - .checked_add(size_of::() - 1) - .ok_or_else(|| vortex_err!("Primitive visitor value allocation overflow"))? - / size_of::() - * size_of::(); - let validity_allocation = if array.dtype().is_nullable() { - array - .len() - .div_ceil(u64::BITS as usize) - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Primitive visitor validity allocation overflow"))? - } else { - 0 - }; - let peak_reservation = - if direct_bitpacked.is_some() { - values_allocation.checked_add(validity_allocation.checked_mul(2).ok_or_else( - || vortex_err!("Primitive visitor validity reservation overflow"), - )?) - } else { - values_allocation - .checked_add(validity_allocation) - .and_then(|bytes| bytes.checked_mul(2)) - } - .ok_or_else(|| vortex_err!("Primitive visitor memory reservation overflow"))?; - let mut memory_reservation = match (memory_callbacks, peak_reservation) { - (Some(callbacks), bytes) if bytes != 0 => { - Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) - } - _ => None, - }; - - let mut execution = session.create_execution_ctx(); - let (primitive_type, length, validity_kind, mut owner) = if let Some(bitpacked) = - direct_bitpacked - { - let primitive_type = primitive_type_id(array.dtype().as_ptype()); - let length = array.len(); - let mask = bitpacked.validity()?.execute_mask(length, &mut execution)?; - let (validity_kind, validity) = exported_validity(array.dtype().is_nullable(), mask); - let owner = PrimitiveOwner::try_new_bitpacked_i64(bitpacked, validity)?; - (primitive_type, length, validity_kind, owner) - } else { - let Canonical::Primitive(primitive) = array.execute::(&mut execution)? - else { - vortex_bail!("Primitive visitor received a non-primitive array"); - }; - let primitive_type = primitive_type_id(primitive.ptype()); - let length = primitive.len(); - let mask = primitive.validity()?.execute_mask(length, &mut execution)?; - let (validity_kind, validity) = - exported_validity(primitive.dtype().is_nullable(), mask); - let host_values = primitive.into_data_parts().buffer.try_into_host_sync()?; - let owner = PrimitiveOwner::try_new( - host_values, - primitive_width(primitive_type)?, - validity, - length, - retain_values, - )?; - (primitive_type, length, validity_kind, owner) - }; - if let Some(mut reservation) = memory_reservation.take() { - reservation.reconcile(owner.retained_bytes())?; - owner.set_memory_reservation(reservation); - } - Ok(Self { - primitive_type, - decimal_precision: 0, - decimal_scale: 0, - length, - validity_kind, - owner: Arc::new(owner), - }) - } - - fn view(&self, offset: usize, length: usize) -> VortexResult { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let width = primitive_width(self.primitive_type)?; - let byte_offset = offset - .checked_mul(width) - .ok_or_else(|| vortex_err!("Vortex Velox value offset overflow"))?; - let values_length = length - .checked_mul(width) - .ok_or_else(|| vortex_err!("Vortex Velox value length overflow"))?; - let values = if values_length == 0 { - ptr::null() - } else { - // SAFETY: The checked export range lies within the retained primitive buffer. - unsafe { self.owner.values().add(byte_offset) } - }; - let (validity, validity_length, validity_bit_offset) = - if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { - packed_bits_window( - self.owner - .validity - .as_ref() - .ok_or_else(|| vortex_err!("Primitive validity bitmap is missing"))?, - offset, - length, - )? - } else { - (ptr::null(), 0, 0) - }; - Ok(vx_velox_primitive_view { - struct_size: size_of::(), - primitive_type: self.primitive_type, - decimal_precision: self.decimal_precision, - decimal_scale: self.decimal_scale, - length, - values, - values_length, - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_primitive_owner), - release: Some(release_primitive_owner), - retained_bytes: self.owner.retained_bytes(), - }, - values_alignment: pointer_alignment(values), - validity_alignment: pointer_alignment(validity), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let view = self.view(offset, length)?; - let callback = visitor - .visit_primitive - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a primitive callback"))?; - // SAFETY: The cursor retains every buffer in the view through this callback. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct BoolExport { - length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, -} - -impl BoolExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let mut execution = session.create_execution_ctx(); - let mut memory_reservation = match memory_callbacks { - Some(callbacks) => Some(ArrowMemoryReservation::try_new( - callbacks, - conservative_export_reservation(&array, &mut execution)?, - )?), - None => None, - }; - let is_nullable = array.dtype().is_nullable(); - let Canonical::Bool(boolean) = array.execute::(&mut execution)? else { - vortex_bail!("Boolean visitor received a non-Boolean array"); - }; - let length = boolean.len(); - let mask = boolean.validity()?.execute_mask(length, &mut execution)?; - let (validity_kind, validity) = exported_validity(is_nullable, mask); - let mut owner = BoolOwner::try_new(boolean.into_bit_buffer(), validity)?; - if let Some(mut reservation) = memory_reservation.take() { - reservation.reconcile(owner.retained_bytes)?; - owner.set_memory_reservation(reservation); - } - Ok(Self { - length, - validity_kind, - owner: Arc::new(owner), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let (values, values_length, values_bit_offset) = - packed_bits_window(&self.owner.values, offset, length)?; - let (validity, validity_length, validity_bit_offset) = match &self.owner.validity { - Some(validity) => packed_bits_window(validity, offset, length)?, - None => (ptr::null(), 0, 0), - }; - let view = vx_velox_bool_view { - struct_size: size_of::(), - length, - values, - values_length, - values_bit_offset, - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_bool_owner), - release: Some(release_bool_owner), - retained_bytes: self.owner.retained_bytes, - }, - values_alignment: pointer_alignment(values), - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor - .visit_bool - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a Boolean callback"))?; - // SAFETY: The cursor retains every buffer in the view through this callback. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -fn packed_bits_window( - bits: &PackedBits, - offset: usize, - length: usize, -) -> VortexResult<(*const u8, usize, usize)> { - if length == 0 { - return Ok((ptr::null(), 0, 0)); - } - let word_bits = u64::BITS as usize; - let byte_offset = offset / word_bits * size_of::(); - let bit_offset = offset % word_bits; - let required_length = bit_offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Packed Boolean window overflow"))? - .div_ceil(u8::BITS as usize); - let byte_length = bits - .len() - .checked_sub(byte_offset) - .ok_or_else(|| vortex_err!("Packed Boolean window exceeds its owner"))?; - if byte_length < required_length { - vortex_bail!("Packed Boolean window exceeds its readable bytes"); - } - // SAFETY: The caller validated the logical window against the owner length. - let values = unsafe { bits.as_ptr().add(byte_offset) }; - Ok((values, byte_length, bit_offset)) -} - -struct VarBinExport { - kind: vx_velox_varbin_kind, - length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, -} - -impl VarBinExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let mut execution = session.create_execution_ctx(); - let mut memory_reservation = match memory_callbacks { - Some(callbacks) => Some(ArrowMemoryReservation::try_new( - callbacks, - conservative_export_reservation(&array, &mut execution)?, - )?), - None => None, - }; - let is_nullable = array.dtype().is_nullable(); - let varbin = array.execute::(&mut execution)?; - let length = varbin.len(); - let parts = varbin.into_data_parts(); - let kind = match parts.dtype { - DType::Utf8(_) => VX_VELOX_VARBIN_UTF8, - DType::Binary(_) => VX_VELOX_VARBIN_BINARY, - dtype => vortex_bail!("Variable-width visitor received an invalid type: {dtype}"), - }; - let mask = parts.validity.execute_mask(length, &mut execution)?; - let (validity_kind, validity) = exported_validity(is_nullable, mask); - let mut owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; - if let Some(mut reservation) = memory_reservation.take() { - reservation.reconcile(owner.retained_bytes)?; - owner.set_memory_reservation(reservation); - } - Ok(Self { - kind, - length, - validity_kind, - owner: Arc::new(owner), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let view_byte_offset = offset - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Vortex string view offset overflow"))?; - let views_length = length - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Vortex string view length overflow"))?; - let views = if views_length == 0 { - ptr::null() - } else { - // SAFETY: The checked export range lies within the retained view buffer. - unsafe { - self.owner - .views - .as_ptr() - .cast::() - .add(view_byte_offset) - .cast() - } - }; - let (validity, validity_length, validity_bit_offset) = - if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { - packed_bits_window( - self.owner - .validity - .as_ref() - .ok_or_else(|| vortex_err!("String validity bitmap is missing"))?, - offset, - length, - )? - } else { - (ptr::null(), 0, 0) - }; - let data_buffers = if self.owner.descriptors.is_empty() { - ptr::null() - } else { - self.owner.descriptors.as_ptr() - }; - let view = vx_velox_varbin_view { - struct_size: size_of::(), - kind: self.kind, - length, - views, - views_length, - data_buffers, - data_buffer_count: self.owner.descriptors.len(), - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_varbin_owner), - release: Some(release_varbin_owner), - retained_bytes: self.owner.retained_bytes, - }, - views_alignment: pointer_alignment(views.cast()), - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor.visit_varbin.ok_or_else(|| { - vortex_err!("Vortex Velox visitor requires a variable-width callback") - })?; - // SAFETY: The cursor retains every buffer in the view through this callback. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct DictionaryExport { - codes: PrimitiveExport, - values_length: usize, - values: Box, -} - -impl DictionaryExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let dictionary = array.as_::(); - let values = dictionary.values().clone(); - Ok(Self { - codes: PrimitiveExport::try_new(dictionary.codes().clone(), session, memory_callbacks)?, - values_length: values.len(), - values: Box::new(vx_velox_export_cursor { - export: CursorExport::try_new_canonical(values, session, memory_callbacks)?, - }), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let codes = self.codes.view(offset, length)?; - let view = vx_velox_dictionary_view { - struct_size: size_of::(), - length, - codes, - values: &raw const *self.values, - values_length: self.values_length, - }; - let callback = visitor - .visit_dictionary - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a dictionary callback"))?; - // SAFETY: The borrowed child cursor and every code buffer remain live through this call. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct ConstantExport { - length: usize, - value: Box, -} - -impl ConstantExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let length = array.len(); - let scalar = array.as_::().scalar().clone(); - let value = ConstantArray::new(scalar, 1).into_array(); - Ok(Self { - length, - value: Box::new(vx_velox_export_cursor { - export: CursorExport::try_new_canonical(value, session, memory_callbacks)?, - }), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let view = vx_velox_constant_view { - struct_size: size_of::(), - length, - value: &raw const *self.value, - }; - let callback = visitor - .visit_constant - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a constant callback"))?; - // SAFETY: The borrowed child cursor remains live through this call. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct StructOwner { - validity: Option, - retained_bytes: usize, - _memory_reservation: Option, -} - -struct StructExport { - length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, - fields: Box<[vx_velox_export_cursor]>, - field_pointers: Box<[*const vx_velox_export_cursor]>, -} - -impl StructExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let is_nullable = array.dtype().is_nullable(); - let mut execution = session.create_execution_ctx(); - let struct_array = array.execute::(&mut execution)?; - let length = struct_array.len(); - let mask = struct_array - .struct_validity() - .execute_mask(length, &mut execution)?; - let validity_reservation = if matches!(mask, Mask::Values(_)) { - length - .div_ceil(u64::BITS as usize) - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("Struct validity reservation overflow"))? - } else { - 0 - }; - let mut memory_reservation = match (memory_callbacks, validity_reservation) { - (Some(callbacks), bytes) if bytes != 0 => { - Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) - } - _ => None, - }; - let (validity_kind, validity) = exported_validity(is_nullable, mask); - let (validity, retained_bytes) = retain_validity(validity, length)?; - if let Some(reservation) = memory_reservation.as_mut() { - reservation.reconcile(retained_bytes)?; - } - let owner = Arc::new(StructOwner { - validity, - retained_bytes, - _memory_reservation: memory_reservation, - }); - let fields = struct_array - .iter_unmasked_fields() - .map(|field| { - Ok(vx_velox_export_cursor { - export: CursorExport::try_new(field.clone(), session, memory_callbacks)?, - }) - }) - .collect::>>()? - .into_boxed_slice(); - let field_pointers = fields - .iter() - .map(|field| field as *const vx_velox_export_cursor) - .collect::>() - .into_boxed_slice(); - Ok(Self { - length, - validity_kind, - owner, - fields, - field_pointers, - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let (validity, validity_length, validity_bit_offset) = - if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { - packed_bits_window( - self.owner - .validity - .as_ref() - .ok_or_else(|| vortex_err!("Struct validity bitmap is missing"))?, - offset, - length, - )? - } else { - (ptr::null(), 0, 0) - }; - let view = vx_velox_struct_view { - struct_size: size_of::(), - length, - offset, - fields: if self.field_pointers.is_empty() { - ptr::null() - } else { - self.field_pointers.as_ptr() - }, - field_count: self.fields.len(), - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_struct_owner), - release: Some(release_struct_owner), - retained_bytes: self.owner.retained_bytes, - }, - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor - .visit_struct - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a struct callback"))?; - // SAFETY: The borrowed field cursors and parent validity remain live through this call. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct ListOwner { - offsets: Box<[i32]>, - sizes: Box<[i32]>, - validity: Option, - retained_bytes: usize, - _memory_reservation: Option, -} - -struct ListMetadata { - length: usize, - elements_length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, -} - -struct ListExport { - length: usize, - elements_length: usize, - validity_kind: vx_velox_validity_kind, - owner: Arc, - elements: Box, -} - -fn list_metadata_value(value: T, name: &str) -> VortexResult -where - T: Copy + std::fmt::Display, - i32: TryFrom, -{ - i32::try_from(value) - .map_err(|_| vortex_err!("Vortex list {name} exceeds the Velox vector limit: {value}")) -} - -fn list_metadata_values(values: PrimitiveArray, name: &str) -> VortexResult> { - let values = values.reinterpret_cast(values.ptype().to_unsigned()); - match_each_unsigned_integer_ptype!(values.ptype(), |P| { - values - .as_slice::

() - .iter() - .map(|&value| list_metadata_value(value, name)) - .collect::>>() - .map(Vec::into_boxed_slice) - }) -} - -fn prepare_list_metadata( - list: &ListViewArray, - session: &vortex::session::VortexSession, - memory_callbacks: Option, -) -> VortexResult { - let is_nullable = list.dtype().is_nullable(); - let mut execution = session.create_execution_ctx(); - let length = list.len(); - let elements_length = list.elements().len(); - if elements_length > i32::MAX as usize { - vortex_bail!("Vortex list elements exceed the Velox vector limit: {elements_length}"); - } - let mask = list - .listview_validity() - .execute_mask(length, &mut execution)?; - let validity_reservation = if matches!(mask, Mask::Values(_)) { - length - .div_ceil(u64::BITS as usize) - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("List validity reservation overflow"))? - } else { - 0 - }; - let metadata_reservation = length - .checked_mul(2 * size_of::()) - .ok_or_else(|| vortex_err!("List metadata reservation overflow"))?; - let reservation = metadata_reservation - .checked_add(validity_reservation) - .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; - let mut memory_reservation = match (memory_callbacks, reservation) { - (Some(callbacks), bytes) if bytes != 0 => { - Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) - } - _ => None, - }; - let offsets = list_metadata_values( - list.offsets() - .clone() - .execute::(&mut execution)?, - "offset", - )?; - let sizes = list_metadata_values( - list.sizes() - .clone() - .execute::(&mut execution)?, - "size", - )?; - let (validity_kind, validity) = exported_validity(is_nullable, mask); - let (validity, validity_allocation) = retain_validity(validity, length)?; - let retained_bytes = size_of_val(offsets.as_ref()) - .checked_add(size_of_val(sizes.as_ref())) - .and_then(|bytes| bytes.checked_add(validity_allocation)) - .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; - if let Some(reservation) = memory_reservation.as_mut() { - reservation.reconcile(retained_bytes)?; - } - Ok(ListMetadata { - length, - elements_length, - validity_kind, - owner: Arc::new(ListOwner { - offsets, - sizes, - validity, - retained_bytes, - _memory_reservation: memory_reservation, - }), - }) -} - -impl ListExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let mut execution = session.create_execution_ctx(); - let list = array.execute::(&mut execution)?; - let elements = list.elements().clone(); - let metadata = prepare_list_metadata(&list, session, memory_callbacks)?; - Ok(Self { - length: metadata.length, - elements_length: metadata.elements_length, - validity_kind: metadata.validity_kind, - owner: metadata.owner, - elements: Box::new(vx_velox_export_cursor { - export: CursorExport::try_new(elements, session, memory_callbacks)?, - }), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let (validity, validity_length, validity_bit_offset) = - if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { - packed_bits_window( - self.owner - .validity - .as_ref() - .ok_or_else(|| vortex_err!("List validity bitmap is missing"))?, - offset, - length, - )? - } else { - (ptr::null(), 0, 0) - }; - let offsets = if length == 0 { - ptr::null() - } else { - // SAFETY: The checked range lies within the metadata arrays. - unsafe { self.owner.offsets.as_ptr().add(offset) } - }; - let sizes = if length == 0 { - ptr::null() - } else { - // SAFETY: The checked range lies within the metadata arrays. - unsafe { self.owner.sizes.as_ptr().add(offset) } - }; - let view = vx_velox_list_view { - struct_size: size_of::(), - length, - offsets, - sizes, - elements: &raw const *self.elements, - elements_length: self.elements_length, - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_list_owner), - release: Some(release_list_owner), - retained_bytes: self.owner.retained_bytes, - }, - offsets_alignment: pointer_alignment(offsets.cast()), - sizes_alignment: pointer_alignment(sizes.cast()), - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor - .visit_list - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a list callback"))?; - // SAFETY: The borrowed element cursor and parent buffers remain live through this call. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -struct MapExport { - length: usize, - entries_length: usize, - keys_sorted: bool, - validity_kind: vx_velox_validity_kind, - owner: Arc, - keys: Box, - values: Box, -} - -impl MapExport { - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - let mut execution = session.create_execution_ctx(); - let map = array.execute::(&mut execution)?; - let keys_sorted = map.keys_sorted(); - let entries = map.entries().clone().downcast::(); - let entry_values = entries.elements().clone(); - let entry_struct = entry_values.execute::(&mut execution)?; - let fields = entry_struct.iter_unmasked_fields().collect::>(); - if fields.len() != 2 { - vortex_bail!( - "Vortex map entries require two fields, got {}", - fields.len() - ); - } - let metadata = prepare_list_metadata(&entries, session, memory_callbacks)?; - Ok(Self { - length: metadata.length, - entries_length: metadata.elements_length, - keys_sorted, - validity_kind: metadata.validity_kind, - owner: metadata.owner, - keys: Box::new(vx_velox_export_cursor { - export: CursorExport::try_new(fields[0].clone(), session, memory_callbacks)?, - }), - values: Box::new(vx_velox_export_cursor { - export: CursorExport::try_new(fields[1].clone(), session, memory_callbacks)?, - }), - }) - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - let end = offset - .checked_add(length) - .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; - if end > self.length { - vortex_bail!( - "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", - self.length - ); - } - let (validity, validity_length, validity_bit_offset) = - if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { - packed_bits_window( - self.owner - .validity - .as_ref() - .ok_or_else(|| vortex_err!("Map validity bitmap is missing"))?, - offset, - length, - )? - } else { - (ptr::null(), 0, 0) - }; - let offsets = if length == 0 { - ptr::null() - } else { - // SAFETY: The checked range lies within the metadata arrays. - unsafe { self.owner.offsets.as_ptr().add(offset) } - }; - let sizes = if length == 0 { - ptr::null() - } else { - // SAFETY: The checked range lies within the metadata arrays. - unsafe { self.owner.sizes.as_ptr().add(offset) } - }; - let view = vx_velox_map_view { - struct_size: size_of::(), - length, - offsets, - sizes, - keys: &raw const *self.keys, - values: &raw const *self.values, - entries_length: self.entries_length, - keys_sorted: self.keys_sorted, - validity_kind: self.validity_kind, - validity, - validity_length, - validity_bit_offset, - buffers: vx_velox_buffer_owner { - struct_size: size_of::(), - owner: Arc::as_ptr(&self.owner).cast(), - retain: Some(retain_list_owner), - release: Some(release_list_owner), - retained_bytes: self.owner.retained_bytes, - }, - offsets_alignment: pointer_alignment(offsets.cast()), - sizes_alignment: pointer_alignment(sizes.cast()), - validity_alignment: pointer_alignment(validity), - }; - let callback = visitor - .visit_map - .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a map callback"))?; - // SAFETY: The borrowed child cursors and parent buffers remain live through this callback. - let status = unsafe { callback(visitor.context, &raw const view) }; - if status != 0 { - vortex_bail!("{}", callback_error(visitor, status)); - } - Ok(()) - } -} - -impl CursorExport { - fn date_storage( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - ) -> VortexResult> { - let DType::Extension(ext_dtype) = array.dtype() else { - return Ok(None); - }; - let Some(time_unit) = ext_dtype.metadata_opt::() else { - return Ok(None); - }; - if *time_unit != TimeUnit::Days { - vortex_bail!( - "Vortex Velox visitor does not support date unit {time_unit}; Velox DATE uses days" - ); - } - - if let Some(extension) = array.as_opt::() { - return Ok(Some(extension.storage_array().clone())); - } - let mut execution = session.create_execution_ctx(); - let extension = array.execute::(&mut execution)?; - Ok(Some(extension.storage_array().clone())) - } - - fn try_new_canonical( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - if matches!(array.dtype(), DType::Map(..)) { - Ok(Self::Map(MapExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else if matches!(array.dtype(), DType::List(..)) { - Ok(Self::List(ListExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else if matches!(array.dtype(), DType::Struct(..)) { - Ok(Self::Struct(StructExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else if matches!(array.dtype(), DType::Decimal(..)) { - Ok(Self::Primitive(PrimitiveExport::try_new_decimal( - array, - session, - memory_callbacks, - )?)) - } else if let Some(storage) = Self::date_storage(array.clone(), session)? { - Ok(Self::Primitive(PrimitiveExport::try_new( - storage, - session, - memory_callbacks, - )?)) - } else if matches!(array.dtype(), DType::Bool(_)) { - Ok(Self::Bool(BoolExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { - Ok(Self::VarBin(VarBinExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else { - Ok(Self::Primitive(PrimitiveExport::try_new( - array, - session, - memory_callbacks, - )?)) - } - } - - fn try_new( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - memory_callbacks: Option, - ) -> VortexResult { - if array.is::() { - Ok(Self::Dictionary(DictionaryExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else if array.is::() { - Ok(Self::Constant(ConstantExport::try_new( - array, - session, - memory_callbacks, - )?)) - } else { - Self::try_new_canonical(array, session, memory_callbacks) - } - } - - fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { - match self { - Self::Primitive(export) => export.visit(offset, length, visitor), - Self::Bool(export) => export.visit(offset, length, visitor), - Self::VarBin(export) => export.visit(offset, length, visitor), - Self::Dictionary(export) => export.visit(offset, length, visitor), - Self::Constant(export) => export.visit(offset, length, visitor), - Self::Struct(export) => export.visit(offset, length, visitor), - Self::List(export) => export.visit(offset, length, visitor), - Self::Map(export) => export.visit(offset, length, visitor), - } - } -} - -fn exported_validity( - is_nullable: bool, - mask: Mask, -) -> (vx_velox_validity_kind, Option) { - if !is_nullable { - return (VX_VELOX_VALIDITY_NON_NULLABLE, None); - } - match mask { - Mask::AllTrue(_) => (VX_VELOX_VALIDITY_ALL_VALID, None), - Mask::AllFalse(_) => (VX_VELOX_VALIDITY_ALL_INVALID, None), - Mask::Values(values) => (VX_VELOX_VALIDITY_BITMAP, Some(values.bit_buffer().clone())), - } -} - -unsafe extern "C" fn retain_primitive_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; -} - -unsafe extern "C" fn release_primitive_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); -} - -unsafe extern "C" fn retain_bool_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; -} - -unsafe extern "C" fn release_bool_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); -} - -unsafe extern "C" fn retain_varbin_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; -} - -unsafe extern "C" fn release_varbin_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); -} - -unsafe extern "C" fn retain_struct_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; -} - -unsafe extern "C" fn release_struct_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); -} - -unsafe extern "C" fn retain_list_owner(owner: *const c_void) { - // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. - unsafe { Arc::increment_strong_count(owner.cast::()) }; -} - -unsafe extern "C" fn release_list_owner(owner: *const c_void) { - // SAFETY: Each release matches a prior retain of this `Arc` pointer. - drop(unsafe { Arc::from_raw(owner.cast::()) }); -} - -fn validate_visitor(visitor: &vx_velox_visitor) -> VortexResult<()> { - if visitor.struct_size < size_of::() { - vortex_bail!( - "Vortex Velox visitor structure is too small: expected at least {}, got {}", - size_of::(), - visitor.struct_size - ); - } - if visitor.abi_version != crate::VX_VELOX_ABI_VERSION { - vortex_bail!( - "Unsupported Vortex Velox ABI version: expected {}, got {}", - crate::VX_VELOX_ABI_VERSION, - visitor.abi_version - ); - } - Ok(()) -} - -fn callback_error(visitor: &vx_velox_visitor, status: i32) -> String { - let Some(last_error) = visitor.last_error else { - return format!("Velox visitor failed with status {status}"); - }; - // SAFETY: The callback contract returns null or a valid null-terminated string. - let message = unsafe { last_error(visitor.context) }; - if message.is_null() { - return format!("Velox visitor failed with status {status}"); - } - // SAFETY: The callback keeps the string valid until the next callback. - unsafe { std::ffi::CStr::from_ptr(message) } - .to_string_lossy() - .into_owned() -} - -fn selected_array( - array: &vortex::array::ArrayRef, - request: &vx_velox_visit_request, -) -> VortexResult { - if request.rows.is_null() { - if request.row_count != 0 { - vortex_bail!("A null visitor row pointer requires a zero row count"); - } - return Ok(array.clone()); - } - // SAFETY: The caller supplies `row_count` readable positions. - let rows = unsafe { slice::from_raw_parts(request.rows, request.row_count) }; - let mut previous = None; - for row in rows { - let position = usize::try_from(*row) - .map_err(|_| vortex_err!("Visitor row does not fit usize: {}", row))?; - if position >= array.len() { - vortex_bail!( - "Visitor row is out of bounds: row {}, array length {}", - row, - array.len() - ); - } - if previous.is_some_and(|previous| previous >= *row) { - vortex_bail!("Visitor rows must be unique and increasing"); - } - previous = Some(*row); - } - let dense = rows.len() == array.len() - && rows - .iter() - .enumerate() - .all(|(position, row)| *row == position as u64); - if dense { - return Ok(array.clone()); - } - array.take(PrimitiveArray::from_iter(rows.iter().copied()).into_array()) -} - -fn visit_array( - array: vortex::array::ArrayRef, - session: &vortex::session::VortexSession, - visitor: &vx_velox_visitor, -) -> VortexResult<()> { - let length = array.len(); - CursorExport::try_new_canonical(array, session, None)?.visit(0, length, visitor) -} - -/// Create one export cursor for several Velox output windows. -/// -/// # Safety -/// -/// The session and array pointers must identify live handles. -/// The memory callbacks must identify a complete, thread-safe callback table. -/// `error_out` must be null or valid. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_export_cursor_new( - session: *const vx_velox_session, - array: *const vx_velox_array, - memory_callbacks: *const vx_velox_arrow_memory_callbacks, - error_out: *mut *mut vx_velox_error, -) -> *mut vx_velox_export_cursor { - try_or(error_out, ptr::null_mut(), || { - let session = unsafe { vx_session_ref(session)? }; - let array = unsafe { vx_array_ref(array)? }; - let memory_callbacks = unsafe { parse_memory_callbacks(memory_callbacks)? }; - Ok(Box::into_raw(Box::new(vx_velox_export_cursor { - export: CursorExport::try_new(array.clone(), session, Some(memory_callbacks))?, - }))) - }) -} - -/// Free one export cursor. -/// -/// # Safety -/// -/// The pointer must be null or come from [`vx_velox_export_cursor_new`]. -#[unsafe(no_mangle)] -pub unsafe extern "C" fn vx_velox_export_cursor_free(cursor: *mut vx_velox_export_cursor) { - if !cursor.is_null() { - // SAFETY: The pointer came from `Box::into_raw` and is freed once. - drop(unsafe { Box::from_raw(cursor) }); - } -} - -/// Visit one contiguous range from a retained export cursor. -/// -/// # Safety -/// -/// The cursor and visitor pointers must remain live until this call returns. -/// Concurrent calls are valid. The caller must not free the cursor before all calls return. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_export_cursor_visit( - cursor: *const vx_velox_export_cursor, - offset: usize, - length: usize, - visitor: *const vx_velox_visitor, - error_out: *mut *mut vx_velox_error, -) -> i32 { - try_or(error_out, 1, || { - let cursor = unsafe { - cursor - .as_ref() - .ok_or_else(|| vortex_err!("Vortex Velox export cursor must not be null"))? - }; - let visitor = unsafe { - visitor - .as_ref() - .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? - }; - validate_visitor(visitor)?; - cursor.export.visit(offset, length, visitor)?; - Ok(0) - }) -} - -/// Visit one Vortex array through host semantic callbacks. -/// -/// The request selects source positions once. Callback block positions are compact and follow the -/// request order. -/// -/// # Safety -/// -/// Every pointer must be null or valid for the documented access. The array and session handles -/// must remain live until this call returns. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_velox_array_visit( - session: *const vx_velox_session, - array: *const vx_velox_array, - request: *const vx_velox_visit_request, - visitor: *const vx_velox_visitor, - error_out: *mut *mut vx_velox_error, -) -> i32 { - try_or(error_out, 1, || { - let session = unsafe { vx_session_ref(session)? }; - let array = unsafe { vx_array_ref(array)? }; - let request = unsafe { - request - .as_ref() - .ok_or_else(|| vortex_err!("Vortex Velox visit request must not be null"))? - }; - if request.struct_size < size_of::() { - vortex_bail!( - "Vortex Velox visit request is too small: expected at least {}, got {}", - size_of::(), - request.struct_size - ); - } - let visitor = unsafe { - visitor - .as_ref() - .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? - }; - validate_visitor(visitor)?; - visit_array(selected_array(array, request)?, session, visitor)?; - Ok(0) - }) -} - -#[cfg(test)] -mod tests { - use std::mem::align_of; - use std::ptr; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - use rstest::rstest; - use vortex::array::ArrayRef; - use vortex::array::IntoArray; - use vortex::array::arrays::BoolArray; - use vortex::array::arrays::DecimalArray; - use vortex::array::arrays::DictArray; - use vortex::array::arrays::ListViewArray; - use vortex::array::arrays::MapArray; - use vortex::array::arrays::PrimitiveArray; - use vortex::array::arrays::StructArray; - use vortex::array::arrays::TemporalArray; - use vortex::array::arrays::VarBinViewArray; - use vortex::array::validity::Validity; - use vortex::buffer::buffer; - use vortex::dtype::DecimalDType; - use vortex::dtype::FieldNames; - use vortex::dtype::MapDType; - use vortex::dtype::Nullability; - use vortex::scalar::Scalar; - use vortex_error::VortexResult; - use vortex_error::vortex_ensure; - use vortex_fastlanes::BitPackedData; - - use super::*; - use crate::api::vx_velox_array_free; - use crate::ffi::vx_array_new_with; - use crate::ffi::vx_session_free; - use crate::ffi::vx_session_new_with; - - #[derive(Default)] - struct TestMemory { - retained_bytes: AtomicUsize, - } - - unsafe extern "C" fn retain_test_memory(_context: *mut c_void) {} - - unsafe extern "C" fn release_test_memory(_context: *mut c_void) {} - - unsafe extern "C" fn reserve_test_memory(context: *mut c_void, bytes: usize) -> i32 { - // SAFETY: The test context stays live through every callback. - let memory = unsafe { &*context.cast::() }; - memory.retained_bytes.fetch_add(bytes, Ordering::Relaxed); - 0 - } - - unsafe extern "C" fn free_test_memory(context: *mut c_void, bytes: usize) { - // SAFETY: The test context stays live through every callback. - let memory = unsafe { &*context.cast::() }; - memory.retained_bytes.fetch_sub(bytes, Ordering::Relaxed); - } - - fn test_memory_callbacks(memory: &mut TestMemory) -> vx_velox_arrow_memory_callbacks { - vx_velox_arrow_memory_callbacks { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (memory as *mut TestMemory).cast(), - retain_context: Some(retain_test_memory), - release_context: Some(release_test_memory), - report_allocation: Some(reserve_test_memory), - report_free: Some(free_test_memory), - last_error: None, - } - } - - #[rstest] - #[case(PType::U8, VX_VELOX_PRIMITIVE_U8)] - #[case(PType::U16, VX_VELOX_PRIMITIVE_U16)] - #[case(PType::U32, VX_VELOX_PRIMITIVE_U32)] - #[case(PType::U64, VX_VELOX_PRIMITIVE_U64)] - #[case(PType::I8, VX_VELOX_PRIMITIVE_I8)] - #[case(PType::I16, VX_VELOX_PRIMITIVE_I16)] - #[case(PType::I32, VX_VELOX_PRIMITIVE_I32)] - #[case(PType::I64, VX_VELOX_PRIMITIVE_I64)] - #[case(PType::F16, VX_VELOX_PRIMITIVE_F16)] - #[case(PType::F32, VX_VELOX_PRIMITIVE_F32)] - #[case(PType::F64, VX_VELOX_PRIMITIVE_F64)] - fn maps_primitive_types(#[case] input: PType, #[case] expected: vx_velox_primitive_type) { - assert_eq!(primitive_type_id(input), expected); - } - - #[test] - fn date_days_use_i32_storage_and_millisecond_dates_are_rejected() -> VortexResult<()> { - let session = vortex::session::VortexSession::empty(); - let days = TemporalArray::new_date( - PrimitiveArray::from_option_iter([Some(-1_i32), None, Some(19_000)]).into_array(), - TimeUnit::Days, - ) - .into_array(); - let CursorExport::Primitive(export) = - CursorExport::try_new_canonical(days, &session, None)? - else { - vortex_bail!("date visitor did not produce primitive storage"); - }; - assert_eq!(export.primitive_type, VX_VELOX_PRIMITIVE_I32); - let view = export.view(0, 3)?; - // SAFETY: The export owns three readable i32 values. - let values = unsafe { slice::from_raw_parts(view.values.cast::(), 3) }; - assert_eq!(values, [-1, 0, 19_000]); - assert_eq!(view.validity_kind, VX_VELOX_VALIDITY_BITMAP); - - let milliseconds = TemporalArray::new_date( - PrimitiveArray::from_iter([86_400_000_i64]).into_array(), - TimeUnit::Milliseconds, - ) - .into_array(); - let error = match CursorExport::try_new_canonical(milliseconds, &session, None) { - Ok(_) => vortex_bail!("millisecond date visitor unexpectedly succeeded"), - Err(error) => error, - }; - assert!(error.to_string().contains("Velox DATE uses days")); - Ok(()) - } - - #[test] - fn decimals_normalize_to_velox_storage_widths() -> VortexResult<()> { - let session = vortex::session::VortexSession::empty(); - let short = DecimalArray::new( - buffer![1_i8, -2, 3], - DecimalDType::new(18, 2), - Validity::NonNullable, - ) - .into_array(); - let short = PrimitiveExport::try_new_decimal(short, &session, None)?; - assert_eq!(short.primitive_type, VX_VELOX_PRIMITIVE_I64); - let short_view = short.view(0, 3)?; - assert_eq!(short_view.decimal_precision, 18); - assert_eq!(short_view.decimal_scale, 2); - // SAFETY: The export owns three readable i64 values. - let short_values = unsafe { slice::from_raw_parts(short_view.values.cast::(), 3) }; - assert_eq!(short_values, [1, -2, 3]); - - let nullable_short = DecimalArray::new( - buffer![1_i128, i128::MAX], - DecimalDType::new(18, 2), - Validity::from_iter([true, false]), - ) - .into_array(); - let nullable_short = PrimitiveExport::try_new_decimal(nullable_short, &session, None)?; - let nullable_short_view = nullable_short.view(0, 2)?; - // SAFETY: The export owns two readable i64 values. - let nullable_short_values = - unsafe { slice::from_raw_parts(nullable_short_view.values.cast::(), 2) }; - assert_eq!(nullable_short_values, [1, 0]); - assert_eq!(nullable_short_view.validity_kind, VX_VELOX_VALIDITY_BITMAP); - - let long = DecimalArray::new( - buffer![1_i64, -2, 3], - DecimalDType::new(30, 4), - Validity::NonNullable, - ) - .into_array(); - let long = PrimitiveExport::try_new_decimal(long, &session, None)?; - assert_eq!(long.primitive_type, VX_VELOX_PRIMITIVE_I128); - let long_view = long.view(0, 3)?; - assert_eq!(long_view.decimal_precision, 30); - assert_eq!(long_view.decimal_scale, 4); - // SAFETY: The export owns three readable i128 values. - let long_values = unsafe { slice::from_raw_parts(long_view.values.cast::(), 3) }; - assert_eq!(long_values, [1, -2, 3]); - - let unsupported = DecimalArray::new( - buffer![1_i8], - DecimalDType::new(39, 0), - Validity::NonNullable, - ) - .into_array(); - let error = match PrimitiveExport::try_new_decimal(unsupported, &session, None) { - Ok(_) => vortex_bail!("precision 39 decimal visitor unexpectedly succeeded"), - Err(error) => error, - }; - assert!(error.to_string().contains("decimal precision 39")); - Ok(()) - } - - #[test] - fn dictionary_export_preserves_code_width_and_nullable_children() -> VortexResult<()> { - let session = vortex::session::VortexSession::empty(); - let code_cases: [(ArrayRef, vx_velox_primitive_type); 4] = [ - (buffer![0_u8, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U8), - (buffer![0_u16, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U16), - (buffer![0_u32, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U32), - (buffer![0_u64, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U64), - ]; - for (codes, expected_type) in code_cases { - let dictionary = DictArray::try_new(codes, buffer![10_i64, 20].into_array())?; - let CursorExport::Dictionary(export) = - CursorExport::try_new(dictionary.into_array(), &session, None)? - else { - vortex_bail!("dictionary export lost its outer encoding"); - }; - assert_eq!(export.codes.primitive_type, expected_type); - assert_eq!(export.values_length, 2); - assert!(matches!(export.values.export, CursorExport::Primitive(_))); - } - - let codes = PrimitiveArray::from_option_iter([Some(0_u8), None, Some(1)]).into_array(); - let values = PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(); - let dictionary = DictArray::try_new(codes, values)?; - let CursorExport::Dictionary(export) = - CursorExport::try_new(dictionary.into_array(), &session, None)? - else { - vortex_bail!("nullable dictionary export lost its outer encoding"); - }; - assert_eq!(export.codes.validity_kind, VX_VELOX_VALIDITY_BITMAP); - let CursorExport::Primitive(values) = &export.values.export else { - vortex_bail!("nullable dictionary values lost their primitive representation"); - }; - assert_eq!(values.validity_kind, VX_VELOX_VALIDITY_BITMAP); - Ok(()) - } - - #[test] - fn constant_export_preserves_null_value() -> VortexResult<()> { - let session = vortex::session::VortexSession::empty(); - let constant = ConstantArray::new(Scalar::null_native::(), 10).into_array(); - let CursorExport::Constant(export) = CursorExport::try_new(constant, &session, None)? - else { - vortex_bail!("constant export lost its outer encoding"); - }; - assert_eq!(export.length, 10); - let CursorExport::Primitive(value) = &export.value.export else { - vortex_bail!("null constant lost its primitive representation"); - }; - assert_eq!(value.length, 1); - assert_eq!(value.validity_kind, VX_VELOX_VALIDITY_ALL_INVALID); - Ok(()) - } - - #[test] - fn struct_export_preserves_children_and_nonzero_window() -> VortexResult<()> { - #[derive(Default)] - struct StructCapture { - length: usize, - offset: usize, - fields: *const *const vx_velox_export_cursor, - field_count: usize, - validity: *const u8, - validity_bit_offset: usize, - owner: Option, - } - - unsafe extern "C" fn capture_struct( - context: *mut c_void, - view: *const vx_velox_struct_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live capture and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.length = view.length; - capture.offset = view.offset; - capture.fields = view.fields; - capture.field_count = view.field_count; - capture.validity = view.validity; - capture.validity_bit_offset = view.validity_bit_offset; - capture.owner = Some(view.buffers); - 0 - } - - let session = vortex::session::VortexSession::empty(); - let length: usize = 130; - let dictionary = DictArray::try_new( - PrimitiveArray::from_iter((0..length).map(|index| [0_u8, 1][index % 2])).into_array(), - buffer![10_i64, 20].into_array(), - )? - .into_array(); - let constant = ConstantArray::new(Scalar::from(7_i64), length).into_array(); - let parent_validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); - let struct_array = StructArray::new( - FieldNames::from(["dictionary", "constant"]), - [dictionary, constant], - length, - parent_validity, - ) - .into_array(); - let CursorExport::Struct(export) = CursorExport::try_new(struct_array, &session, None)? - else { - vortex_bail!("struct export lost its outer encoding"); - }; - assert!(matches!( - export.fields[0].export, - CursorExport::Dictionary(_) - )); - assert!(matches!(export.fields[1].export, CursorExport::Constant(_))); - - let mut capture = StructCapture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: Some(capture_struct), - visit_list: None, - visit_map: None, - }; - export.visit(65, 63, &visitor)?; - assert_eq!(capture.length, 63); - assert_eq!(capture.offset, 65); - assert_eq!(capture.field_count, 2); - assert_eq!(capture.validity_bit_offset, 1); - // SAFETY: The export retains both field cursors until it is dropped below. - assert_eq!(unsafe { *capture.fields }, &raw const export.fields[0]); - let owner = capture - .owner - .ok_or_else(|| vortex_err!("struct callback returned no validity owner"))?; - drop(export); - // SAFETY: The callback retained the parent owner before the cursor was dropped. - assert!( - unsafe { - *capture - .validity - .add(capture.validity_bit_offset / u8::BITS as usize) - } != 0 - ); - let release = owner - .release - .ok_or_else(|| vortex_err!("struct owner returned no release callback"))?; - // SAFETY: This release matches the callback retain above. - unsafe { release(owner.owner) }; - Ok(()) - } - - #[test] - fn list_export_preserves_elements_window_and_accounting() -> VortexResult<()> { - #[derive(Default)] - struct ListCapture { - length: usize, - offsets: *const i32, - sizes: *const i32, - elements_length: usize, - validity: *const u8, - validity_bit_offset: usize, - owner: Option, - } - - unsafe extern "C" fn capture_list( - context: *mut c_void, - view: *const vx_velox_list_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live capture and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.length = view.length; - capture.offsets = view.offsets; - capture.sizes = view.sizes; - capture.elements_length = view.elements_length; - capture.validity = view.validity; - capture.validity_bit_offset = view.validity_bit_offset; - capture.owner = Some(view.buffers); - 0 - } - - let session = vortex::session::VortexSession::empty(); - let length = 130; - let elements = DictArray::try_new( - buffer![0_u8, 1, 0, 1, 0, 1].into_array(), - PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(), - )? - .into_array(); - let offsets = PrimitiveArray::from_iter((0..length).map(|index| [0_u32, 2, 4][index % 3])); - let sizes = - PrimitiveArray::from_iter((0..length).map(|index| if index % 10 == 0 { 0 } else { 2 })); - let validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); - let list = ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity) - .into_array(); - let mut memory = TestMemory::default(); - let CursorExport::List(export) = - CursorExport::try_new(list, &session, Some(test_memory_callbacks(&mut memory)))? - else { - vortex_bail!("list export lost its outer encoding"); - }; - assert!(matches!( - export.elements.export, - CursorExport::Dictionary(_) - )); - let expected_parent_bytes = - length * 2 * size_of::() + length.div_ceil(u64::BITS as usize) * size_of::(); - assert_eq!(export.owner.retained_bytes, expected_parent_bytes); - - let mut capture = ListCapture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: Some(capture_list), - visit_map: None, - }; - export.visit(65, 63, &visitor)?; - assert_eq!(capture.length, 63); - assert_eq!(capture.elements_length, 6); - assert_eq!(capture.validity_bit_offset, 1); - // SAFETY: The retained owner keeps both metadata arrays live. - assert_eq!(unsafe { *capture.offsets }, 4); - // SAFETY: The retained owner keeps both metadata arrays live. - assert_eq!(unsafe { *capture.sizes }, 2); - let owner = capture - .owner - .ok_or_else(|| vortex_err!("list callback returned no owner"))?; - drop(export); - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - expected_parent_bytes - ); - // SAFETY: The callback retained the owner before the export was dropped. - assert_eq!(unsafe { *capture.offsets.add(1) }, 0); - let release = owner - .release - .ok_or_else(|| vortex_err!("list owner returned no release callback"))?; - // SAFETY: This release matches the callback retain above. - unsafe { release(owner.owner) }; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[test] - fn map_export_preserves_children_window_and_accounting() -> VortexResult<()> { - #[derive(Default)] - struct MapCapture { - length: usize, - offsets: *const i32, - sizes: *const i32, - keys: *const vx_velox_export_cursor, - values: *const vx_velox_export_cursor, - entries_length: usize, - keys_sorted: bool, - validity_bit_offset: usize, - owner: Option, - } - - unsafe extern "C" fn capture_map( - context: *mut c_void, - view: *const vx_velox_map_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live capture and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.length = view.length; - capture.offsets = view.offsets; - capture.sizes = view.sizes; - capture.keys = view.keys; - capture.values = view.values; - capture.entries_length = view.entries_length; - capture.keys_sorted = view.keys_sorted; - capture.validity_bit_offset = view.validity_bit_offset; - capture.owner = Some(view.buffers); - 0 - } - - let session = vortex::session::VortexSession::empty(); - let keys = DictArray::try_new( - buffer![0_u8, 1, 0, 1, 0, 1].into_array(), - buffer![10_i64, 20].into_array(), - )? - .into_array(); - let values = ConstantArray::new(Scalar::from(7_i64), 6).into_array(); - let entries = StructArray::new( - FieldNames::from(["key", "value"]), - [keys, values], - 6, - Validity::NonNullable, - ) - .into_array(); - let entry_lists = ListViewArray::new( - entries, - buffer![0_u32, 2, 4].into_array(), - buffer![2_u32, 2, 2].into_array(), - Validity::from_iter([true, false, true]), - ); - let map_dtype = MapDType::try_new( - DType::Primitive(PType::I64, Nullability::NonNullable), - DType::Primitive(PType::I64, Nullability::NonNullable), - true, - )?; - let map = MapArray::try_new(map_dtype, entry_lists)?.into_array(); - let mut memory = TestMemory::default(); - let CursorExport::Map(export) = - CursorExport::try_new(map, &session, Some(test_memory_callbacks(&mut memory)))? - else { - vortex_bail!("map export lost its outer encoding"); - }; - assert!(matches!(export.keys.export, CursorExport::Dictionary(_))); - assert!(matches!(export.values.export, CursorExport::Constant(_))); - let expected_parent_bytes = 3 * 2 * size_of::() + size_of::(); - assert_eq!(export.owner.retained_bytes, expected_parent_bytes); - - let mut capture = MapCapture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: Some(capture_map), - }; - export.visit(1, 2, &visitor)?; - assert_eq!(capture.length, 2); - assert_eq!(capture.entries_length, 6); - assert!(capture.keys_sorted); - assert_eq!(capture.validity_bit_offset, 1); - assert_eq!(capture.keys, &raw const *export.keys); - assert_eq!(capture.values, &raw const *export.values); - // SAFETY: The retained owner keeps both metadata arrays live. - assert_eq!(unsafe { *capture.offsets }, 2); - // SAFETY: The retained owner keeps both metadata arrays live. - assert_eq!(unsafe { *capture.sizes }, 2); - let owner = capture - .owner - .ok_or_else(|| vortex_err!("map callback returned no owner"))?; - drop(export); - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - expected_parent_bytes - ); - // SAFETY: The callback retained the owner before the export was dropped. - assert_eq!(unsafe { *capture.offsets.add(1) }, 4); - let release = owner - .release - .ok_or_else(|| vortex_err!("map owner returned no release callback"))?; - // SAFETY: This release matches the callback retain above. - unsafe { release(owner.owner) }; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[derive(Default)] - struct Capture { - primitive_type: Option, - length: usize, - values: *const u8, - values_length: usize, - values_alignment: usize, - validity: *const u8, - validity_length: usize, - validity_bit_offset: usize, - validity_alignment: usize, - retained_bytes: usize, - validity_kind: Option, - owner: Option, - } - - unsafe extern "C" fn capture_primitive( - context: *mut c_void, - view: *const vx_velox_primitive_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live `Capture` and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.primitive_type = Some(view.primitive_type); - capture.length = view.length; - capture.values = view.values; - capture.values_length = view.values_length; - capture.values_alignment = view.values_alignment; - capture.validity = view.validity; - capture.validity_length = view.validity_length; - capture.validity_bit_offset = view.validity_bit_offset; - capture.validity_alignment = view.validity_alignment; - capture.retained_bytes = view.buffers.retained_bytes; - capture.validity_kind = Some(view.validity_kind); - capture.owner = Some(view.buffers); - 0 - } - - fn release_capture(capture: &Capture) -> VortexResult<()> { - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; - // SAFETY: This release matches the retain in `capture_primitive`. - unsafe { release(owner.owner) }; - Ok(()) - } - - #[derive(Default)] - struct VarBinCapture { - struct_size: usize, - kind: Option, - length: usize, - views: *const vx_velox_binary_view, - views_length: usize, - views_alignment: usize, - data_buffers: *const vx_velox_byte_buffer_view, - data_buffer_count: usize, - validity: *const u8, - validity_length: usize, - validity_bit_offset: usize, - validity_alignment: usize, - validity_kind: Option, - retained_bytes: usize, - owner: Option, - } - - unsafe extern "C" fn capture_varbin( - context: *mut c_void, - view: *const vx_velox_varbin_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live capture and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.struct_size = view.struct_size; - capture.kind = Some(view.kind); - capture.length = view.length; - capture.views = view.views; - capture.views_length = view.views_length; - capture.views_alignment = view.views_alignment; - capture.data_buffers = view.data_buffers; - capture.data_buffer_count = view.data_buffer_count; - capture.validity = view.validity; - capture.validity_length = view.validity_length; - capture.validity_bit_offset = view.validity_bit_offset; - capture.validity_alignment = view.validity_alignment; - capture.validity_kind = Some(view.validity_kind); - capture.retained_bytes = view.buffers.retained_bytes; - capture.owner = Some(view.buffers); - 0 - } - - fn release_varbin_capture(capture: &VarBinCapture) -> VortexResult<()> { - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained string owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("string owner did not return a release callback"))?; - // SAFETY: This release matches the retain in `capture_varbin`. - unsafe { release(owner.owner) }; - Ok(()) - } - - #[derive(Default)] - struct BoolCapture { - length: usize, - values: *const u8, - values_bit_offset: usize, - validity: *const u8, - validity_bit_offset: usize, - validity_kind: Option, - retained_bytes: usize, - owner: Option, - } - - unsafe extern "C" fn capture_bool( - context: *mut c_void, - view: *const vx_velox_bool_view, - ) -> i32 { - if context.is_null() || view.is_null() { - return 1; - } - // SAFETY: The test passes pointers to live capture and view objects. - let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; - let Some(retain) = view.buffers.retain else { - return 2; - }; - // SAFETY: The visitor owner is live for the callback. - unsafe { retain(view.buffers.owner) }; - capture.length = view.length; - capture.values = view.values; - capture.values_bit_offset = view.values_bit_offset; - capture.validity = view.validity; - capture.validity_bit_offset = view.validity_bit_offset; - capture.validity_kind = Some(view.validity_kind); - capture.retained_bytes = view.buffers.retained_bytes; - capture.owner = Some(view.buffers); - 0 - } - - fn release_bool_capture(capture: &BoolCapture) -> VortexResult<()> { - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained Boolean owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("Boolean owner did not return a release callback"))?; - // SAFETY: This release matches the retain in `capture_bool`. - unsafe { release(owner.owner) }; - Ok(()) - } - - #[expect( - clippy::host_endian_bytes, - reason = "The Vortex binary-view fields use the host C ABI layout" - )] - unsafe fn captured_varbin_value(capture: &VarBinCapture, index: usize) -> Option<&[u8]> { - if capture.validity_kind == Some(VX_VELOX_VALIDITY_BITMAP) { - let bit_index = capture.validity_bit_offset + index; - // SAFETY: The callback contract retains the bitmap for every captured row. - let byte = unsafe { *capture.validity.add(bit_index / 8) }; - if byte & (1 << (bit_index % 8)) == 0 { - return None; - } - } - // SAFETY: The callback contract retains `length` readable views. - let view = unsafe { &*capture.views.add(index) }; - let length = view.length as usize; - const INLINE_LENGTH: usize = size_of::() - size_of::(); - if length <= INLINE_LENGTH { - return Some(&view.data[..length]); - } - let buffer_index = - u32::from_ne_bytes([view.data[4], view.data[5], view.data[6], view.data[7]]) as usize; - let offset = - u32::from_ne_bytes([view.data[8], view.data[9], view.data[10], view.data[11]]) as usize; - // SAFETY: The callback contract retains all payload descriptors. - let buffer = unsafe { &*capture.data_buffers.add(buffer_index) }; - // SAFETY: Canonical Vortex views contain validated payload ranges. - Some(unsafe { slice::from_raw_parts(buffer.data.add(offset), length) }) - } - - #[rstest] - #[case(DType::Utf8(Nullability::Nullable), VX_VELOX_VARBIN_UTF8)] - #[case(DType::Binary(Nullability::Nullable), VX_VELOX_VARBIN_BINARY)] - fn varbin_cursor_retains_mixed_views_across_nonzero_window( - #[case] dtype: DType, - #[case] expected_kind: vx_velox_varbin_kind, - ) -> VortexResult<()> { - let utf8_expected: [Option<&[u8]>; 7] = [ - Some(b""), - Some(b"a"), - None, - Some(b"abcdefghijkl"), - Some(b"abcdefghijklm"), - Some("vortex 🌀 outlined".as_bytes()), - Some(b"tail"), - ]; - let binary_expected: [Option<&[u8]>; 7] = [ - Some(b""), - Some(b"\xff"), - None, - Some(b"abcdefghijkl"), - Some(b"\x00abcdefghijklm"), - Some(b"\xff\x00 binary outlined value"), - Some(b"tail"), - ]; - let expected = if matches!(dtype, DType::Utf8(_)) { - utf8_expected - } else { - binary_expected - }; - let session = vx_session_new_with(|session| session); - let varbin = VarBinViewArray::from_iter(expected, dtype); - let array = vx_array_new_with(varbin.into_array()); - let mut error = ptr::null_mut(); - let mut memory = TestMemory::default(); - let memory_callbacks = test_memory_callbacks(&mut memory); - // SAFETY: The session and array handles remain live until cursor creation finishes. - let cursor = unsafe { - vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) - }; - vortex_ensure!(!cursor.is_null(), "string cursor creation failed"); - vortex_ensure!(error.is_null(), "string cursor returned an error"); - - let mut capture = VarBinCapture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: Some(capture_varbin), - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - // SAFETY: The cursor and callback state remain live through the call. - let status = unsafe { - vx_velox_export_cursor_visit(cursor, 1, 5, &raw const visitor, &raw mut error) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "string export window returned an error"); - assert_eq!(capture.struct_size, size_of::()); - assert_eq!(capture.kind, Some(expected_kind)); - assert_eq!(capture.length, 5); - assert_eq!(capture.views_length, 5 * size_of::()); - assert!(capture.views_alignment >= align_of::()); - assert_eq!(capture.views.addr() % align_of::(), 0); - assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); - assert_eq!(capture.validity_bit_offset, 1); - assert!(capture.validity_length >= 1); - assert!(capture.validity_alignment >= align_of::()); - assert_eq!(capture.validity.addr() % align_of::(), 0); - assert!(capture.data_buffer_count >= 1); - assert!(!capture.data_buffers.is_null()); - assert_eq!( - capture.retained_bytes, - memory.retained_bytes.load(Ordering::Relaxed) - ); - - // SAFETY: Each owned handle is freed once. The callback retained the string owner. - unsafe { - vx_velox_export_cursor_free(cursor); - vx_velox_array_free(array); - vx_session_free(session); - } - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - capture.retained_bytes - ); - for (index, expected) in expected[1..6].iter().enumerate() { - // SAFETY: The retained owner keeps every captured pointer live. - let actual = unsafe { captured_varbin_value(&capture, index) }; - assert_eq!(actual, *expected); - } - release_varbin_capture(&capture)?; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[test] - fn varbin_shared_buffers_compact_into_exact_owned_storage() -> VortexResult<()> { - let length = 130_usize; - let strings = VarBinViewArray::from_iter( - (0..length).map(|index| { - (index % 11 != 0).then(|| format!("outlined string value {index:03}")) - }), - DType::Utf8(Nullability::Nullable), - ); - let parts = strings.into_data_parts(); - let views_length = parts.views.try_to_host_sync()?.len(); - let data_length = parts - .buffers - .iter() - .map(|buffer| Ok(buffer.try_to_host_sync()?.len())) - .sum::>()?; - let descriptor_length = parts.buffers.len() * size_of::(); - let validity_length = length.div_ceil(u64::BITS as usize) * size_of::(); - let expected_retained = views_length + data_length + descriptor_length + validity_length; - - let retained_views = parts.views.clone(); - let retained_buffers = Arc::<[BufferHandle]>::clone(&parts.buffers); - let mut execution = vortex::session::VortexSession::empty().create_execution_ctx(); - let mask = parts.validity.execute_mask(length, &mut execution)?; - let (_, validity) = exported_validity(true, mask); - let owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; - - assert!(matches!(owner.views, RetainedViews::Compact(_))); - assert!( - owner - ._data - .iter() - .all(|buffer| matches!(buffer, RetainedBytes::Compact(_))) - ); - assert_eq!(owner.retained_bytes, expected_retained); - drop(retained_views); - drop(retained_buffers); - Ok(()) - } - - #[test] - fn retained_varbin_buffers_report_complete_unique_allocations() -> VortexResult<()> { - let alignment = vortex::buffer::Alignment::new(256); - let mut payload = BufferMut::::with_capacity_aligned(17, alignment); - payload.extend(0..17); - let expected_payload_allocation = payload.allocation_size(); - let (retained_payload, payload_allocation) = - RetainedBytes::try_new(BufferHandle::new_host(payload.freeze()))?; - assert!(matches!(retained_payload, RetainedBytes::Retained(_))); - assert_eq!(payload_allocation, expected_payload_allocation); - assert!(payload_allocation > 17); - - let mut views = BufferMut::::with_capacity_aligned( - 2 * size_of::(), - alignment, - ); - views.extend(std::iter::repeat_n( - 0, - 2 * size_of::(), - )); - let expected_views_allocation = views.allocation_size(); - let (retained_views, views_allocation) = - RetainedViews::try_new(BufferHandle::new_host(views.freeze()))?; - assert!(matches!(retained_views, RetainedViews::Retained(_))); - assert_eq!(views_allocation, expected_views_allocation); - assert!(views_allocation > 2 * size_of::()); - Ok(()) - } - - #[test] - fn word_aligned_windows_rebase_validity_buffers() -> VortexResult<()> { - let session = vortex::session::VortexSession::empty(); - let primitive = PrimitiveArray::from_option_iter( - (0..130).map(|index| (index % 7 != 0).then_some(index as i64)), - ) - .into_array(); - let primitive = PrimitiveExport::try_new(primitive, &session, None)?; - let primitive_first = primitive.view(0, 64)?; - let primitive_second = primitive.view(64, 64)?; - assert_eq!(primitive_first.validity_bit_offset, 0); - assert_eq!(primitive_second.validity_bit_offset, 0); - // SAFETY: Both pointers lie in the retained validity allocation. - assert_eq!(primitive_second.validity, unsafe { - primitive_first.validity.add(size_of::()) - }); - - let strings = VarBinViewArray::from_iter( - (0..130).map(|index| (index % 11 != 0).then(|| format!("value-{index}"))), - DType::Utf8(Nullability::Nullable), - ) - .into_array(); - let strings = VarBinExport::try_new(strings, &session, None)?; - let mut first = VarBinCapture::default(); - let first_visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut first).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: Some(capture_varbin), - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - strings.visit(0, 64, &first_visitor)?; - - let mut second = VarBinCapture::default(); - let second_visitor = vx_velox_visitor { - context: (&raw mut second).cast(), - ..first_visitor - }; - strings.visit(64, 64, &second_visitor)?; - assert_eq!(first.validity_bit_offset, 0); - assert_eq!(second.validity_bit_offset, 0); - // SAFETY: Both pointers lie in the retained validity allocation. - assert_eq!(second.validity, unsafe { - first.validity.add(size_of::()) - }); - release_varbin_capture(&first)?; - release_varbin_capture(&second)?; - Ok(()) - } - - #[test] - fn bool_cursor_retains_nonzero_window_and_exact_accounting() -> VortexResult<()> { - let expected = (0..130) - .map(|index| (index % 11 != 0).then_some(index % 3 == 0)) - .collect::>(); - let session = vx_session_new_with(|session| session); - let boolean = BoolArray::from_iter(expected.iter().copied()); - let array = vx_array_new_with(boolean.into_array()); - let mut error = ptr::null_mut(); - let mut memory = TestMemory::default(); - let memory_callbacks = test_memory_callbacks(&mut memory); - // SAFETY: The session and array handles remain live until cursor creation finishes. - let cursor = unsafe { - vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) - }; - vortex_ensure!(!cursor.is_null(), "Boolean cursor creation failed"); - vortex_ensure!(error.is_null(), "Boolean cursor returned an error"); - - let mut capture = BoolCapture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: None, - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: Some(capture_bool), - visit_struct: None, - visit_list: None, - visit_map: None, - }; - // SAFETY: The cursor and callback state remain live through the call. - let status = unsafe { - vx_velox_export_cursor_visit(cursor, 65, 63, &raw const visitor, &raw mut error) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "Boolean export window returned an error"); - assert_eq!(capture.length, 63); - assert_eq!(capture.values_bit_offset, 1); - assert_eq!(capture.validity_bit_offset, 1); - assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); - assert_eq!(capture.retained_bytes, 6 * size_of::()); - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 48); - - // SAFETY: Each owned handle is freed once. The callback retained the Boolean owner. - unsafe { - vx_velox_export_cursor_free(cursor); - vx_velox_array_free(array); - vx_session_free(session); - } - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - capture.retained_bytes - ); - for (relative_index, expected) in expected[65..128].iter().enumerate() { - let value_bit = capture.values_bit_offset + relative_index; - let validity_bit = capture.validity_bit_offset + relative_index; - // SAFETY: The retained buffers cover every captured value and validity bit. - let (actual, is_valid) = unsafe { - ( - *capture.values.add(value_bit / 8) & (1 << (value_bit % 8)) != 0, - *capture.validity.add(validity_bit / 8) & (1 << (validity_bit % 8)) != 0, - ) - }; - assert_eq!(is_valid, expected.is_some()); - if let Some(expected) = expected { - assert_eq!(actual, *expected); - } - } - release_bool_capture(&capture)?; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[test] - fn export_cursor_reuses_one_prepared_array_across_windows() -> VortexResult<()> { - let session = vx_session_new_with(|session| session); - let array = vx_array_new_with( - PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40), Some(50)]) - .into_array(), - ); - let mut error = ptr::null_mut(); - let mut memory = TestMemory::default(); - let memory_callbacks = test_memory_callbacks(&mut memory); - // SAFETY: The session and array handles remain live until cursor creation finishes. - let cursor = unsafe { - vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) - }; - vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); - vortex_ensure!(error.is_null(), "export cursor returned an error"); - assert!(memory.retained_bytes.load(Ordering::Relaxed) >= 48); - - let mut first = Capture::default(); - let first_visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut first).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - // SAFETY: The cursor and callback state remain live through the call. - let status = unsafe { - vx_velox_export_cursor_visit(cursor, 1, 2, &raw const first_visitor, &raw mut error) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "first export window returned an error"); - assert_eq!(first.length, 2); - assert_eq!(first.validity_bit_offset, 1); - // SAFETY: The callback retained two readable i64 values. - let first_values = unsafe { slice::from_raw_parts(first.values.cast::(), 2) }; - assert_eq!(first_values, [0, 30]); - assert_eq!( - first.retained_bytes, - memory.retained_bytes.load(Ordering::Relaxed) - ); - let owner = first - .owner - .ok_or_else(|| vortex_err!("first export window returned no owner"))? - .owner; - release_capture(&first)?; - - let mut second = Capture::default(); - let second_visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut second).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - // SAFETY: The cursor and callback state remain live through the call. - let status = unsafe { - vx_velox_export_cursor_visit(cursor, 3, 2, &raw const second_visitor, &raw mut error) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "second export window returned an error"); - assert_eq!(second.length, 2); - assert_eq!(second.validity_bit_offset, 3); - assert_eq!( - second - .owner - .ok_or_else(|| vortex_err!("second export window returned no owner"))? - .owner, - owner - ); - - // SAFETY: Each owned handle is freed exactly once. The second callback retained the owner. - unsafe { - vx_velox_export_cursor_free(cursor); - vx_velox_array_free(array); - vx_session_free(session); - } - // SAFETY: The retained cursor owner keeps these two i64 values live. - let second_values = unsafe { slice::from_raw_parts(second.values.cast::(), 2) }; - assert_eq!(second_values, [40, 50]); - release_capture(&second)?; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[test] - fn export_cursor_decodes_sliced_bitpacked_into_exact_owner() -> VortexResult<()> { - let session = vx_session_new_with(|session| { - vortex_fastlanes::initialize(&session); - session - }); - let session_ref = unsafe { vx_session_ref(session)? }; - let values = (0..2_050).map(|index| (index % 7 != 0).then_some(i64::from(index % 100))); - let primitive = PrimitiveArray::from_option_iter(values).into_array(); - let mut execution = session_ref.create_execution_ctx(); - let bitpacked = BitPackedData::encode(&primitive, 7, &mut execution)?; - vortex_ensure!( - bitpacked.patches().is_none(), - "test bit-packed array unexpectedly contains patches" - ); - let slice_begin = 113; - let slice_end = 1_941; - let sliced = bitpacked.into_array().slice(slice_begin..slice_end)?; - let array = vx_array_new_with(sliced); - let mut error = ptr::null_mut(); - let mut memory = TestMemory::default(); - let memory_callbacks = test_memory_callbacks(&mut memory); - // SAFETY: The session and array handles remain live until cursor creation finishes. - let cursor = unsafe { - vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) - }; - vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); - vortex_ensure!(error.is_null(), "export cursor returned an error"); - let sliced_length = slice_end - slice_begin; - let expected_retained = sliced_length * size_of::() - + sliced_length.div_ceil(u64::BITS as usize) * size_of::(); - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - expected_retained - ); - - let window_offset = 997; - let window_length = 6; - let mut capture = Capture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - // SAFETY: The cursor and callback state remain live through the call. - let status = unsafe { - vx_velox_export_cursor_visit( - cursor, - window_offset, - window_length, - &raw const visitor, - &raw mut error, - ) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "export window returned an error"); - assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); - assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); - assert_eq!( - capture.validity_bit_offset, - window_offset % u64::BITS as usize - ); - assert_eq!(capture.retained_bytes, expected_retained); - // SAFETY: The callback retained `window_length` readable i64 values. - let actual = unsafe { slice::from_raw_parts(capture.values.cast::(), window_length) }; - for (relative_index, value) in actual.iter().enumerate() { - let sliced_index = window_offset + relative_index; - let source_index = slice_begin + sliced_index; - // SAFETY: The retained bitmap covers every row in the sliced array. - let validity_index = capture.validity_bit_offset + relative_index; - let validity_byte = unsafe { *capture.validity.add(validity_index / 8) }; - let is_valid = validity_byte & (1 << (validity_index % 8)) != 0; - assert_eq!(is_valid, source_index % 7 != 0); - if is_valid { - assert_eq!(*value, i64::try_from(source_index % 100)?); - } - } - - // SAFETY: Each owned handle is freed exactly once. The callback retained the owner. - unsafe { - vx_velox_export_cursor_free(cursor); - vx_velox_array_free(array); - vx_session_free(session); - } - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - expected_retained - ); - release_capture(&capture)?; - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - Ok(()) - } - - #[test] - fn patched_bitpacked_uses_retained_canonical_fallback() -> VortexResult<()> { - let session = vx_session_new_with(|session| { - vortex_fastlanes::initialize(&session); - session - }); - let session_ref = unsafe { vx_session_ref(session)? }; - let expected = [1_u64, 2, 3, u64::MAX]; - let primitive = PrimitiveArray::from_iter(expected).into_array(); - let mut execution = session_ref.create_execution_ctx(); - let bitpacked = BitPackedData::encode(&primitive, 2, &mut execution)?; - vortex_ensure!( - bitpacked.patches().is_some(), - "test bit-packed array unexpectedly omitted patches" - ); - let mut memory = TestMemory::default(); - let export = PrimitiveExport::try_new( - bitpacked.into_array(), - session_ref, - Some(test_memory_callbacks(&mut memory)), - )?; - assert!(matches!(export.owner.values, PrimitiveValues::Retained(_))); - assert_eq!( - memory.retained_bytes.load(Ordering::Relaxed), - export.owner.retained_bytes() - ); - // SAFETY: The export owner contains `expected.len()` initialized u64 values. - let actual = - unsafe { slice::from_raw_parts(export.owner.values().cast::(), expected.len()) }; - assert_eq!(actual, expected); - drop(export); - assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); - unsafe { vx_session_free(session) }; - Ok(()) - } - - #[test] - fn visits_sparse_nullable_values_with_retained_buffers() -> VortexResult<()> { - let session = vx_session_new_with(|session| session); - let array = vx_array_new_with( - PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40)]).into_array(), - ); - let rows = [1_u64, 3]; - let request = vx_velox_visit_request { - struct_size: size_of::(), - rows: rows.as_ptr(), - row_count: rows.len(), - }; - let mut capture = Capture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - let mut error = ptr::null_mut(); - // SAFETY: Every handle and callback object stays live for this call. - let status = unsafe { - vx_velox_array_visit( - session, - array, - &raw const request, - &raw const visitor, - &raw mut error, - ) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "visitor returned an error"); - assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); - assert_eq!(capture.length, 2); - assert_eq!(capture.values_length, 2 * size_of::()); - assert!(capture.values_alignment.is_power_of_two()); - assert_eq!(capture.values.addr() % capture.values_alignment, 0); - assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); - assert_eq!(capture.validity_length, size_of::()); - assert_eq!(capture.validity_bit_offset, 0); - assert!(capture.validity_alignment.is_power_of_two()); - assert_eq!(capture.validity.addr() % capture.validity_alignment, 0); - assert_eq!( - capture.retained_bytes, - capture.values_length + size_of::() - ); - // SAFETY: The callback retained the owner before storing these pointers. - let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 2) }; - assert_eq!(values, [0, 40]); - // SAFETY: The retained validity pointer has one readable word. - let validity = unsafe { *capture.validity }; - assert_eq!(validity & 0b11, 0b10); - - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; - // SAFETY: This release matches the retain in `capture_primitive`. - unsafe { release(owner.owner) }; - // SAFETY: Each owned handle is freed exactly once. - unsafe { - vx_velox_array_free(array); - vx_session_free(session); - } - Ok(()) - } - - #[test] - fn copies_sliced_values_into_exact_owned_storage() -> VortexResult<()> { - let session = vx_session_new_with(|session| session); - let source = PrimitiveArray::from_iter(0_i32..16); - let source_values = source.buffer_handle().try_to_host_sync()?; - // SAFETY: The source contains sixteen i32 values. The fifth value is in bounds. - let source_slice = unsafe { source_values.as_ptr().add(5 * size_of::()) }; - drop(source_values); - let array = vx_array_new_with(source.into_array().slice(5..8)?); - let request = vx_velox_visit_request { - struct_size: size_of::(), - rows: ptr::null(), - row_count: 0, - }; - let mut capture = Capture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - let mut error = ptr::null_mut(); - let status = unsafe { - vx_velox_array_visit( - session, - array, - &raw const request, - &raw const visitor, - &raw mut error, - ) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "visitor returned an error"); - assert_eq!(capture.values_length, 3 * size_of::()); - assert_eq!(capture.retained_bytes, 2 * size_of::()); - assert_ne!(capture.values, source_slice); - // SAFETY: Each owned handle is freed exactly once. The callback retained the value owner. - unsafe { - vx_velox_array_free(array); - vx_session_free(session); - } - // SAFETY: The retained compact buffer contains three i32 values. - let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 3) }; - assert_eq!(values, [5, 6, 7]); - assert!(capture.values_alignment.is_power_of_two()); - assert_eq!(capture.values.addr() % capture.values_alignment, 0); - assert_eq!(capture.validity_alignment, 0); - - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; - unsafe { release(owner.owner) }; - Ok(()) - } - - #[test] - fn copies_validity_into_word_padded_storage() -> VortexResult<()> { - let session = vx_session_new_with(|session| session); - let session_ref = unsafe { vx_session_ref(session)? }; - let primitive = PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]); - let mut execution = session_ref.create_execution_ctx(); - let Mask::Values(mask) = primitive - .validity()? - .execute_mask(primitive.len(), &mut execution)? - else { - vortex_bail!("Expected bitmap validity"); - }; - let expected_validity = mask.bit_buffer().inner().as_ptr(); - let array = vx_array_new_with(primitive.into_array()); - let request = vx_velox_visit_request { - struct_size: size_of::(), - rows: ptr::null(), - row_count: 0, - }; - let mut capture = Capture::default(); - let visitor = vx_velox_visitor { - struct_size: size_of::(), - abi_version: crate::VX_VELOX_ABI_VERSION, - context: (&raw mut capture).cast(), - visit_primitive: Some(capture_primitive), - last_error: None, - visit_varbin: None, - visit_dictionary: None, - visit_constant: None, - visit_bool: None, - visit_struct: None, - visit_list: None, - visit_map: None, - }; - let mut error = ptr::null_mut(); - let status = unsafe { - vx_velox_array_visit( - session, - array, - &raw const request, - &raw const visitor, - &raw mut error, - ) - }; - assert_eq!(status, 0); - vortex_ensure!(error.is_null(), "visitor returned an error"); - assert_ne!(capture.validity, expected_validity); - assert_eq!(capture.validity_bit_offset, 0); - assert_eq!(capture.validity_length, size_of::()); - assert!(capture.validity_alignment >= align_of::()); - assert_eq!( - capture.retained_bytes, - capture.values_length.div_ceil(size_of::()) * size_of::() + size_of::() - ); - - let owner = capture - .owner - .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; - let release = owner - .release - .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; - unsafe { release(owner.owner) }; - unsafe { - vx_velox_array_free(array); - vx_session_free(session); - } - Ok(()) - } - - #[test] - fn rejects_unsorted_rows() -> VortexResult<()> { - let array = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(); - let rows = [2_u64, 1]; - let request = vx_velox_visit_request { - struct_size: size_of::(), - rows: rows.as_ptr(), - row_count: rows.len(), - }; - match selected_array(&array, &request) { - Ok(_) => vortex_bail!("unsorted rows unexpectedly succeeded"), - Err(error) => assert!(error.to_string().contains("unique and increasing")), - } - Ok(()) - } -} diff --git a/vortex-velox/src/visitor/export.rs b/vortex-velox/src/visitor/export.rs new file mode 100644 index 00000000000..c121f2c9e0a --- /dev/null +++ b/vortex-velox/src/visitor/export.rs @@ -0,0 +1,1970 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::c_void; +use std::mem::MaybeUninit; +use std::mem::align_of; +use std::mem::size_of; +use std::mem::size_of_val; +use std::ptr; +use std::slice; +use std::sync::Arc; + +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::Constant; +use vortex::array::arrays::ConstantArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::Dict; +use vortex::array::arrays::Extension; +use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::ListView; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::MapArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::decimal::DecimalArrayExt; +use vortex::array::arrays::extension::ExtensionArrayExt; +use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; +use vortex::array::arrays::map::MapArrayExt; +use vortex::array::arrays::map::MapArraySlotsExt; +use vortex::array::arrays::primitive::PrimitiveArrayExt; +use vortex::array::arrays::struct_::StructArrayExt; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_unsigned_integer_ptype; +use vortex::buffer::Buffer; +use vortex::buffer::BufferMut; +use vortex::buffer::ByteBuffer; +use vortex::dtype::DType; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::PType; +use vortex::extension::datetime::Date; +use vortex::extension::datetime::TimeUnit; +use vortex::mask::Mask; +use vortex_array::ArrayView; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; + +use super::*; +use crate::array::ArrowMemoryReservation; +use crate::array::conservative_export_reservation; +use crate::array::parse_memory_callbacks; +use crate::array::vx_velox_arrow_memory_callbacks; +use crate::ffi::try_or; +use crate::ffi::vx_array_ref; +use crate::ffi::vx_session_ref; +use crate::ffi::vx_velox_array; +use crate::ffi::vx_velox_error; +use crate::ffi::vx_velox_session; + +fn primitive_type_id(value: PType) -> vx_velox_primitive_type { + match value { + PType::U8 => VX_VELOX_PRIMITIVE_U8, + PType::U16 => VX_VELOX_PRIMITIVE_U16, + PType::U32 => VX_VELOX_PRIMITIVE_U32, + PType::U64 => VX_VELOX_PRIMITIVE_U64, + PType::I8 => VX_VELOX_PRIMITIVE_I8, + PType::I16 => VX_VELOX_PRIMITIVE_I16, + PType::I32 => VX_VELOX_PRIMITIVE_I32, + PType::I64 => VX_VELOX_PRIMITIVE_I64, + PType::F16 => VX_VELOX_PRIMITIVE_F16, + PType::F32 => VX_VELOX_PRIMITIVE_F32, + PType::F64 => VX_VELOX_PRIMITIVE_F64, + } +} + +/// Retains one prepared Vortex array across several Velox output windows. +#[repr(C)] +pub struct vx_velox_export_cursor { + export: CursorExport, +} + +enum CursorExport { + Primitive(PrimitiveExport), + Bool(BoolExport), + VarBin(VarBinExport), + Dictionary(DictionaryExport), + Constant(ConstantExport), + Struct(StructExport), + List(ListExport), + Map(MapExport), +} + +struct PackedBits(Box<[u64]>); + +impl PackedBits { + fn try_new(bits: vortex::buffer::BitBuffer) -> VortexResult<(Self, usize)> { + let compact = bits + .chunks() + .iter_padded() + .collect::>() + .into_boxed_slice(); + let allocation = size_of_val(compact.as_ref()); + Ok((Self(compact), allocation)) + } + + fn as_ptr(&self) -> *const u8 { + self.0.as_ptr().cast() + } + + fn len(&self) -> usize { + size_of_val(self.0.as_ref()) + } +} + +struct BoolOwner { + values: PackedBits, + validity: Option, + retained_bytes: usize, + memory_reservation: Option, +} + +impl BoolOwner { + fn try_new( + values: vortex::buffer::BitBuffer, + validity: Option, + ) -> VortexResult { + let (values, values_allocation) = PackedBits::try_new(values)?; + let (validity, validity_allocation) = match validity { + Some(validity) => { + let (validity, allocation) = PackedBits::try_new(validity)?; + (Some(validity), allocation) + } + None => (None, 0), + }; + let retained_bytes = values_allocation + .checked_add(validity_allocation) + .ok_or_else(|| vortex_err!("Boolean visitor retained byte count overflow"))?; + Ok(Self { + values, + validity, + retained_bytes, + memory_reservation: None, + }) + } + + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); + } +} + +enum PrimitiveValues { + Compact64(Box<[MaybeUninit]>), + Compact128(Box<[MaybeUninit]>), + Retained(ByteBuffer), +} + +impl PrimitiveValues { + fn as_ptr(&self) -> *const u8 { + match self { + Self::Compact64(values) => values.as_ptr().cast(), + Self::Compact128(values) => values.as_ptr().cast(), + Self::Retained(values) => values.as_ptr(), + } + } +} + +struct PrimitiveOwner { + values: PrimitiveValues, + values_length: usize, + validity: Option, + retained_bytes: usize, + memory_reservation: Option, +} + +enum RetainedBytes { + Retained(ByteBuffer), + Compact(Box<[u8]>), +} + +impl RetainedBytes { + fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { + let buffer = handle.try_into_host_sync()?; + let length = buffer.len(); + match buffer.try_into_mut() { + Ok(buffer) => { + let allocation_size = buffer.allocation_size(); + Ok((Self::Retained(buffer.freeze()), allocation_size)) + } + Err(buffer) => { + let compact = buffer.as_slice().to_vec().into_boxed_slice(); + Ok((Self::Compact(compact), length)) + } + } + } + + fn as_ptr(&self) -> *const u8 { + match self { + Self::Retained(buffer) => buffer.as_ptr(), + Self::Compact(buffer) => buffer.as_ptr(), + } + } + + fn len(&self) -> usize { + match self { + Self::Retained(buffer) => buffer.len(), + Self::Compact(buffer) => buffer.len(), + } + } +} + +enum RetainedViews { + Retained(ByteBuffer), + Compact(Box<[vx_velox_binary_view]>), +} + +impl RetainedViews { + fn try_new(handle: BufferHandle) -> VortexResult<(Self, usize)> { + let buffer = handle.try_into_host_sync()?; + if !buffer + .len() + .is_multiple_of(size_of::()) + { + vortex_bail!( + "Vortex variable-width view buffer has an invalid byte length: {}", + buffer.len() + ); + } + match buffer.try_into_mut() { + Ok(buffer) => { + let allocation_size = buffer.allocation_size(); + Ok((Self::Retained(buffer.freeze()), allocation_size)) + } + Err(buffer) => { + let length = buffer.len() / size_of::(); + let mut compact = vec![ + vx_velox_binary_view { + length: 0, + data: [0; 12], + }; + length + ] + .into_boxed_slice(); + if !buffer.is_empty() { + // SAFETY: Both byte ranges have the checked identical size. + unsafe { + ptr::copy_nonoverlapping( + buffer.as_ptr(), + compact.as_mut_ptr().cast::(), + buffer.len(), + ) + }; + } + let allocation = size_of_val(compact.as_ref()); + Ok((Self::Compact(compact), allocation)) + } + } + } + + fn as_ptr(&self) -> *const vx_velox_binary_view { + match self { + Self::Retained(buffer) => buffer.as_ptr().cast(), + Self::Compact(buffer) => buffer.as_ptr(), + } + } +} + +struct VarBinOwner { + views: RetainedViews, + _data: Box<[RetainedBytes]>, + descriptors: Box<[vx_velox_byte_buffer_view]>, + validity: Option, + retained_bytes: usize, + memory_reservation: Option, +} + +// SAFETY: The owner never mutates its buffers or pointer descriptors after construction. +// Every descriptor points into an immutable allocation that the same owner retains. +unsafe impl Send for VarBinOwner {} +// SAFETY: Shared access only reads immutable buffers and descriptors retained by this owner. +unsafe impl Sync for VarBinOwner {} + +impl VarBinOwner { + fn try_new( + views: BufferHandle, + mut buffers: Arc<[BufferHandle]>, + validity: Option, + length: usize, + ) -> VortexResult { + let (views, views_allocation) = RetainedViews::try_new(views)?; + let handles = if let Some(handles) = Arc::get_mut(&mut buffers) { + handles + .iter_mut() + .map(|handle| { + std::mem::replace(handle, BufferHandle::new_host(ByteBuffer::empty())) + }) + .collect::>() + } else { + buffers.iter().cloned().collect::>() + }; + let mut data_allocation = 0usize; + let data = handles + .into_iter() + .map(|handle| { + let (buffer, allocation) = RetainedBytes::try_new(handle)?; + data_allocation = data_allocation + .checked_add(allocation) + .ok_or_else(|| vortex_err!("Vortex string payload allocation overflow"))?; + Ok(buffer) + }) + .collect::>>()? + .into_boxed_slice(); + let descriptors = data + .iter() + .map(|buffer| vx_velox_byte_buffer_view { + data: buffer.as_ptr(), + length: buffer.len(), + }) + .collect::>() + .into_boxed_slice(); + let descriptor_allocation = size_of_val(descriptors.as_ref()); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = views_allocation + .checked_add(data_allocation) + .and_then(|bytes| bytes.checked_add(descriptor_allocation)) + .and_then(|bytes| bytes.checked_add(validity_allocation)) + .ok_or_else(|| vortex_err!("Vortex string retained byte count overflow"))?; + Ok(Self { + views, + _data: data, + descriptors, + validity, + retained_bytes, + memory_reservation: None, + }) + } + + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); + } +} + +fn retain_validity( + validity: Option, + length: usize, +) -> VortexResult<(Option, usize)> { + let Some(validity) = validity else { + return Ok((None, 0)); + }; + if validity.len() < length { + vortex_bail!( + "Vortex validity length is too small: {} for {length} values", + validity.len() + ); + } + let validity = if validity.len() == length { + validity + } else { + validity.slice(..length) + }; + let (validity, allocation) = PackedBits::try_new(validity)?; + Ok((Some(validity), allocation)) +} + +impl PrimitiveOwner { + fn try_allocate( + values_length: usize, + values_alignment: usize, + validity: Option, + length: usize, + ) -> VortexResult { + let (values, values_allocation) = if values_alignment > align_of::() { + if values_alignment > align_of::() { + vortex_bail!( + "Primitive visitor does not support value alignment {values_alignment}" + ); + } + let values = + vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] + .into_boxed_slice(); + let allocation = values + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + (PrimitiveValues::Compact128(values), allocation) + } else { + let values = + vec![MaybeUninit::::uninit(); values_length.div_ceil(size_of::())] + .into_boxed_slice(); + let allocation = values + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + (PrimitiveValues::Compact64(values), allocation) + }; + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = values_allocation + .checked_add(validity_allocation) + .ok_or_else(|| vortex_err!("Primitive visitor retained byte count overflow"))?; + Ok(Self { + values, + values_length, + validity, + retained_bytes, + memory_reservation: None, + }) + } + + fn try_new( + host_values: ByteBuffer, + values_alignment: usize, + validity: Option, + length: usize, + retain_values: bool, + ) -> VortexResult { + let values_length = host_values.len(); + let host_values = if retain_values { + match host_values.try_into_mut() { + Ok(values) => { + let values_allocation = values.allocation_size(); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = values_allocation + .checked_add(validity_allocation) + .ok_or_else(|| { + vortex_err!("Primitive visitor retained byte count overflow") + })?; + return Ok(Self { + values: PrimitiveValues::Retained(values.freeze()), + values_length, + validity, + retained_bytes, + memory_reservation: None, + }); + } + Err(values) => values, + } + } else { + host_values + }; + let mut owner = Self::try_allocate(values_length, values_alignment, validity, length)?; + if !host_values.is_empty() { + let (values_pointer, values_capacity) = match &mut owner.values { + PrimitiveValues::Compact64(values) => ( + values.as_mut_ptr().cast::(), + values.len() * size_of::(), + ), + PrimitiveValues::Compact128(values) => ( + values.as_mut_ptr().cast::(), + values.len() * size_of::(), + ), + PrimitiveValues::Retained(_) => { + unreachable!("a newly allocated primitive owner must be compact") + } + }; + // SAFETY: The byte view spans the complete compact allocation. + let values_bytes = + unsafe { slice::from_raw_parts_mut(values_pointer, values_capacity) }; + values_bytes[..values_length].copy_from_slice(host_values.as_slice()); + } + Ok(owner) + } + + fn try_new_bitpacked_i64( + array: ArrayView<'_, BitPacked>, + validity: Option, + ) -> VortexResult { + let values_length = array + .len() + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + let mut owner = + Self::try_allocate(values_length, align_of::(), validity, array.len())?; + // SAFETY: The allocation uses `u64` alignment and contains at least `values_length` bytes. + // The output slice covers exactly `array.len()` values and remains uniquely borrowed. + let output = unsafe { + slice::from_raw_parts_mut( + match &mut owner.values { + PrimitiveValues::Compact64(values) => { + values.as_mut_ptr().cast::>() + } + PrimitiveValues::Compact128(_) | PrimitiveValues::Retained(_) => { + unreachable!("a newly allocated primitive owner must be compact") + } + }, + array.len(), + ) + }; + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + array.unpacked_chunks(&mut scratch)?.decode_into(output); + Ok(owner) + } + + fn values(&self) -> *const u8 { + if self.values_length == 0 { + ptr::null() + } else { + self.values.as_ptr() + } + } + + fn retained_bytes(&self) -> usize { + self.retained_bytes + } + + fn set_memory_reservation(&mut self, reservation: ArrowMemoryReservation) { + self.memory_reservation = Some(reservation); + } +} + +fn pointer_alignment(pointer: *const u8) -> usize { + if pointer.is_null() { + return 0; + } + 1usize << pointer.addr().trailing_zeros() +} + +fn primitive_width(primitive_type: vx_velox_primitive_type) -> VortexResult { + Ok(match primitive_type { + VX_VELOX_PRIMITIVE_U8 | VX_VELOX_PRIMITIVE_I8 => 1, + VX_VELOX_PRIMITIVE_U16 | VX_VELOX_PRIMITIVE_I16 | VX_VELOX_PRIMITIVE_F16 => 2, + VX_VELOX_PRIMITIVE_U32 | VX_VELOX_PRIMITIVE_I32 | VX_VELOX_PRIMITIVE_F32 => 4, + VX_VELOX_PRIMITIVE_U64 | VX_VELOX_PRIMITIVE_I64 | VX_VELOX_PRIMITIVE_F64 => 8, + VX_VELOX_PRIMITIVE_I128 => 16, + _ => vortex_bail!("Unknown Vortex Velox primitive type: {primitive_type}"), + }) +} + +fn cast_decimal_values(values: Buffer, validity: &Mask) -> VortexResult +where + T: NativeDecimalType, + S: NativeDecimalType, +{ + let mut output = BufferMut::::with_capacity(values.len()); + for (index, value) in values.into_iter().enumerate() { + if !validity.value(index) { + output.push(T::default()); + continue; + } + output.push(::from(value).ok_or_else(|| { + vortex_err!( + "Decimal value cannot be represented as {}", + std::any::type_name::() + ) + })?); + } + Ok(output.freeze().into_byte_buffer()) +} + +fn normalized_decimal_values(array: &DecimalArray, validity: &Mask) -> VortexResult +where + T: NativeDecimalType, +{ + if array.values_type() == T::DECIMAL_TYPE { + return array.buffer_handle().clone().try_into_host_sync(); + } + match array.values_type() { + DecimalType::I8 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I16 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I32 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I64 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I128 => cast_decimal_values::(array.buffer::(), validity), + DecimalType::I256 => cast_decimal_values::( + array.buffer::(), + validity, + ), + } +} + +struct PrimitiveExport { + primitive_type: vx_velox_primitive_type, + decimal_precision: u32, + decimal_scale: i32, + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl PrimitiveExport { + fn try_new_decimal( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let retain_values = memory_callbacks.is_some(); + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let decimal = array.execute::(&mut execution)?; + let decimal_precision = u32::from(decimal.precision()); + let decimal_scale = i32::from(decimal.scale()); + let length = decimal.len(); + let mask = decimal + .as_ref() + .validity()? + .execute_mask(length, &mut execution)?; + let (primitive_type, host_values) = match decimal.precision() { + 1..=18 => ( + VX_VELOX_PRIMITIVE_I64, + normalized_decimal_values::(&decimal, &mask)?, + ), + 19..=38 => ( + VX_VELOX_PRIMITIVE_I128, + normalized_decimal_values::(&decimal, &mask)?, + ), + precision => { + vortex_bail!("Vortex Velox visitor does not support decimal precision {precision}") + } + }; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = PrimitiveOwner::try_new( + host_values, + primitive_width(primitive_type)?, + validity, + length, + retain_values, + )?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes())?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + primitive_type, + decimal_precision, + decimal_scale, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let retain_values = memory_callbacks.is_some(); + let direct_bitpacked = array.as_opt::().filter(|bitpacked| { + array.dtype().as_ptype() == PType::I64 && bitpacked.patches().is_none() + }); + let values_length = + array + .len() + .checked_mul(array.dtype().element_size().ok_or_else(|| { + vortex_err!("Primitive visitor received a variable-width array") + })?) + .ok_or_else(|| vortex_err!("Primitive visitor value byte count overflow"))?; + let values_allocation = values_length + .checked_add(size_of::() - 1) + .ok_or_else(|| vortex_err!("Primitive visitor value allocation overflow"))? + / size_of::() + * size_of::(); + let validity_allocation = if array.dtype().is_nullable() { + array + .len() + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Primitive visitor validity allocation overflow"))? + } else { + 0 + }; + let peak_reservation = + if direct_bitpacked.is_some() { + values_allocation.checked_add(validity_allocation.checked_mul(2).ok_or_else( + || vortex_err!("Primitive visitor validity reservation overflow"), + )?) + } else { + values_allocation + .checked_add(validity_allocation) + .and_then(|bytes| bytes.checked_mul(2)) + } + .ok_or_else(|| vortex_err!("Primitive visitor memory reservation overflow"))?; + let mut memory_reservation = match (memory_callbacks, peak_reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + + let mut execution = session.create_execution_ctx(); + let (primitive_type, length, validity_kind, mut owner) = if let Some(bitpacked) = + direct_bitpacked + { + let primitive_type = primitive_type_id(array.dtype().as_ptype()); + let length = array.len(); + let mask = bitpacked.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(array.dtype().is_nullable(), mask); + let owner = PrimitiveOwner::try_new_bitpacked_i64(bitpacked, validity)?; + (primitive_type, length, validity_kind, owner) + } else { + let Canonical::Primitive(primitive) = array.execute::(&mut execution)? + else { + vortex_bail!("Primitive visitor received a non-primitive array"); + }; + let primitive_type = primitive_type_id(primitive.ptype()); + let length = primitive.len(); + let mask = primitive.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = + exported_validity(primitive.dtype().is_nullable(), mask); + let host_values = primitive.into_data_parts().buffer.try_into_host_sync()?; + let owner = PrimitiveOwner::try_new( + host_values, + primitive_width(primitive_type)?, + validity, + length, + retain_values, + )?; + (primitive_type, length, validity_kind, owner) + }; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes())?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + primitive_type, + decimal_precision: 0, + decimal_scale: 0, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn view(&self, offset: usize, length: usize) -> VortexResult { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let width = primitive_width(self.primitive_type)?; + let byte_offset = offset + .checked_mul(width) + .ok_or_else(|| vortex_err!("Vortex Velox value offset overflow"))?; + let values_length = length + .checked_mul(width) + .ok_or_else(|| vortex_err!("Vortex Velox value length overflow"))?; + let values = if values_length == 0 { + ptr::null() + } else { + // SAFETY: The checked export range lies within the retained primitive buffer. + unsafe { self.owner.values().add(byte_offset) } + }; + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Primitive validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + Ok(vx_velox_primitive_view { + struct_size: size_of::(), + primitive_type: self.primitive_type, + decimal_precision: self.decimal_precision, + decimal_scale: self.decimal_scale, + length, + values, + values_length, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_primitive_owner), + release: Some(release_primitive_owner), + retained_bytes: self.owner.retained_bytes(), + }, + values_alignment: pointer_alignment(values), + validity_alignment: pointer_alignment(validity), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let view = self.view(offset, length)?; + let callback = visitor + .visit_primitive + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a primitive callback"))?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct BoolExport { + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl BoolExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let Canonical::Bool(boolean) = array.execute::(&mut execution)? else { + vortex_bail!("Boolean visitor received a non-Boolean array"); + }; + let length = boolean.len(); + let mask = boolean.validity()?.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = BoolOwner::try_new(boolean.into_bit_buffer(), validity)?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes)?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (values, values_length, values_bit_offset) = + packed_bits_window(&self.owner.values, offset, length)?; + let (validity, validity_length, validity_bit_offset) = match &self.owner.validity { + Some(validity) => packed_bits_window(validity, offset, length)?, + None => (ptr::null(), 0, 0), + }; + let view = vx_velox_bool_view { + struct_size: size_of::(), + length, + values, + values_length, + values_bit_offset, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_bool_owner), + release: Some(release_bool_owner), + retained_bytes: self.owner.retained_bytes, + }, + values_alignment: pointer_alignment(values), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_bool + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a Boolean callback"))?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +fn packed_bits_window( + bits: &PackedBits, + offset: usize, + length: usize, +) -> VortexResult<(*const u8, usize, usize)> { + if length == 0 { + return Ok((ptr::null(), 0, 0)); + } + let word_bits = u64::BITS as usize; + let byte_offset = offset / word_bits * size_of::(); + let bit_offset = offset % word_bits; + let required_length = bit_offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Packed Boolean window overflow"))? + .div_ceil(u8::BITS as usize); + let byte_length = bits + .len() + .checked_sub(byte_offset) + .ok_or_else(|| vortex_err!("Packed Boolean window exceeds its owner"))?; + if byte_length < required_length { + vortex_bail!("Packed Boolean window exceeds its readable bytes"); + } + // SAFETY: The caller validated the logical window against the owner length. + let values = unsafe { bits.as_ptr().add(byte_offset) }; + Ok((values, byte_length, bit_offset)) +} + +struct VarBinExport { + kind: vx_velox_varbin_kind, + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +impl VarBinExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let mut memory_reservation = match memory_callbacks { + Some(callbacks) => Some(ArrowMemoryReservation::try_new( + callbacks, + conservative_export_reservation(&array, &mut execution)?, + )?), + None => None, + }; + let is_nullable = array.dtype().is_nullable(); + let varbin = array.execute::(&mut execution)?; + let length = varbin.len(); + let parts = varbin.into_data_parts(); + let kind = match parts.dtype { + DType::Utf8(_) => VX_VELOX_VARBIN_UTF8, + DType::Binary(_) => VX_VELOX_VARBIN_BINARY, + dtype => vortex_bail!("Variable-width visitor received an invalid type: {dtype}"), + }; + let mask = parts.validity.execute_mask(length, &mut execution)?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let mut owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; + if let Some(mut reservation) = memory_reservation.take() { + reservation.reconcile(owner.retained_bytes)?; + owner.set_memory_reservation(reservation); + } + Ok(Self { + kind, + length, + validity_kind, + owner: Arc::new(owner), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let view_byte_offset = offset + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Vortex string view offset overflow"))?; + let views_length = length + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Vortex string view length overflow"))?; + let views = if views_length == 0 { + ptr::null() + } else { + // SAFETY: The checked export range lies within the retained view buffer. + unsafe { + self.owner + .views + .as_ptr() + .cast::() + .add(view_byte_offset) + .cast() + } + }; + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("String validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let data_buffers = if self.owner.descriptors.is_empty() { + ptr::null() + } else { + self.owner.descriptors.as_ptr() + }; + let view = vx_velox_varbin_view { + struct_size: size_of::(), + kind: self.kind, + length, + views, + views_length, + data_buffers, + data_buffer_count: self.owner.descriptors.len(), + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_varbin_owner), + release: Some(release_varbin_owner), + retained_bytes: self.owner.retained_bytes, + }, + views_alignment: pointer_alignment(views.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor.visit_varbin.ok_or_else(|| { + vortex_err!("Vortex Velox visitor requires a variable-width callback") + })?; + // SAFETY: The cursor retains every buffer in the view through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct DictionaryExport { + codes: PrimitiveExport, + values_length: usize, + values: Box, +} + +impl DictionaryExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let dictionary = array.as_::(); + let values = dictionary.values().clone(); + Ok(Self { + codes: PrimitiveExport::try_new(dictionary.codes().clone(), session, memory_callbacks)?, + values_length: values.len(), + values: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new_canonical(values, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let codes = self.codes.view(offset, length)?; + let view = vx_velox_dictionary_view { + struct_size: size_of::(), + length, + codes, + values: &raw const *self.values, + values_length: self.values_length, + }; + let callback = visitor + .visit_dictionary + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a dictionary callback"))?; + // SAFETY: The borrowed child cursor and every code buffer remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct ConstantExport { + length: usize, + value: Box, +} + +impl ConstantExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let length = array.len(); + let scalar = array.as_::().scalar().clone(); + let value = ConstantArray::new(scalar, 1).into_array(); + Ok(Self { + length, + value: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new_canonical(value, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let view = vx_velox_constant_view { + struct_size: size_of::(), + length, + value: &raw const *self.value, + }; + let callback = visitor + .visit_constant + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a constant callback"))?; + // SAFETY: The borrowed child cursor remains live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct StructOwner { + validity: Option, + retained_bytes: usize, + _memory_reservation: Option, +} + +struct StructExport { + length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, + fields: Box<[vx_velox_export_cursor]>, + field_pointers: Box<[*const vx_velox_export_cursor]>, +} + +impl StructExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let is_nullable = array.dtype().is_nullable(); + let mut execution = session.create_execution_ctx(); + let struct_array = array.execute::(&mut execution)?; + let length = struct_array.len(); + let mask = struct_array + .struct_validity() + .execute_mask(length, &mut execution)?; + let validity_reservation = if matches!(mask, Mask::Values(_)) { + length + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("Struct validity reservation overflow"))? + } else { + 0 + }; + let mut memory_reservation = match (memory_callbacks, validity_reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let (validity, retained_bytes) = retain_validity(validity, length)?; + if let Some(reservation) = memory_reservation.as_mut() { + reservation.reconcile(retained_bytes)?; + } + let owner = Arc::new(StructOwner { + validity, + retained_bytes, + _memory_reservation: memory_reservation, + }); + let fields = struct_array + .iter_unmasked_fields() + .map(|field| { + Ok(vx_velox_export_cursor { + export: CursorExport::try_new(field.clone(), session, memory_callbacks)?, + }) + }) + .collect::>>()? + .into_boxed_slice(); + let field_pointers = fields + .iter() + .map(|field| field as *const vx_velox_export_cursor) + .collect::>() + .into_boxed_slice(); + Ok(Self { + length, + validity_kind, + owner, + fields, + field_pointers, + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Struct validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let view = vx_velox_struct_view { + struct_size: size_of::(), + length, + offset, + fields: if self.field_pointers.is_empty() { + ptr::null() + } else { + self.field_pointers.as_ptr() + }, + field_count: self.fields.len(), + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_struct_owner), + release: Some(release_struct_owner), + retained_bytes: self.owner.retained_bytes, + }, + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_struct + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a struct callback"))?; + // SAFETY: The borrowed field cursors and parent validity remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct ListOwner { + offsets: Box<[i32]>, + sizes: Box<[i32]>, + validity: Option, + retained_bytes: usize, + _memory_reservation: Option, +} + +struct ListMetadata { + length: usize, + elements_length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, +} + +struct ListExport { + length: usize, + elements_length: usize, + validity_kind: vx_velox_validity_kind, + owner: Arc, + elements: Box, +} + +fn list_metadata_value(value: T, name: &str) -> VortexResult +where + T: Copy + std::fmt::Display, + i32: TryFrom, +{ + i32::try_from(value) + .map_err(|_| vortex_err!("Vortex list {name} exceeds the Velox vector limit: {value}")) +} + +fn list_metadata_values(values: PrimitiveArray, name: &str) -> VortexResult> { + let values = values.reinterpret_cast(values.ptype().to_unsigned()); + match_each_unsigned_integer_ptype!(values.ptype(), |P| { + values + .as_slice::

() + .iter() + .map(|&value| list_metadata_value(value, name)) + .collect::>>() + .map(Vec::into_boxed_slice) + }) +} + +fn prepare_list_metadata( + list: &ListViewArray, + session: &vortex::session::VortexSession, + memory_callbacks: Option, +) -> VortexResult { + let is_nullable = list.dtype().is_nullable(); + let mut execution = session.create_execution_ctx(); + let length = list.len(); + let elements_length = list.elements().len(); + if elements_length > i32::MAX as usize { + vortex_bail!("Vortex list elements exceed the Velox vector limit: {elements_length}"); + } + let mask = list + .listview_validity() + .execute_mask(length, &mut execution)?; + let validity_reservation = if matches!(mask, Mask::Values(_)) { + length + .div_ceil(u64::BITS as usize) + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("List validity reservation overflow"))? + } else { + 0 + }; + let metadata_reservation = length + .checked_mul(2 * size_of::()) + .ok_or_else(|| vortex_err!("List metadata reservation overflow"))?; + let reservation = metadata_reservation + .checked_add(validity_reservation) + .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; + let mut memory_reservation = match (memory_callbacks, reservation) { + (Some(callbacks), bytes) if bytes != 0 => { + Some(ArrowMemoryReservation::try_new(callbacks, bytes)?) + } + _ => None, + }; + let offsets = list_metadata_values( + list.offsets() + .clone() + .execute::(&mut execution)?, + "offset", + )?; + let sizes = list_metadata_values( + list.sizes() + .clone() + .execute::(&mut execution)?, + "size", + )?; + let (validity_kind, validity) = exported_validity(is_nullable, mask); + let (validity, validity_allocation) = retain_validity(validity, length)?; + let retained_bytes = size_of_val(offsets.as_ref()) + .checked_add(size_of_val(sizes.as_ref())) + .and_then(|bytes| bytes.checked_add(validity_allocation)) + .ok_or_else(|| vortex_err!("List retained byte count overflow"))?; + if let Some(reservation) = memory_reservation.as_mut() { + reservation.reconcile(retained_bytes)?; + } + Ok(ListMetadata { + length, + elements_length, + validity_kind, + owner: Arc::new(ListOwner { + offsets, + sizes, + validity, + retained_bytes, + _memory_reservation: memory_reservation, + }), + }) +} + +impl ListExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let list = array.execute::(&mut execution)?; + let elements = list.elements().clone(); + let metadata = prepare_list_metadata(&list, session, memory_callbacks)?; + Ok(Self { + length: metadata.length, + elements_length: metadata.elements_length, + validity_kind: metadata.validity_kind, + owner: metadata.owner, + elements: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(elements, session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("List validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let offsets = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.offsets.as_ptr().add(offset) } + }; + let sizes = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.sizes.as_ptr().add(offset) } + }; + let view = vx_velox_list_view { + struct_size: size_of::(), + length, + offsets, + sizes, + elements: &raw const *self.elements, + elements_length: self.elements_length, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_list_owner), + release: Some(release_list_owner), + retained_bytes: self.owner.retained_bytes, + }, + offsets_alignment: pointer_alignment(offsets.cast()), + sizes_alignment: pointer_alignment(sizes.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_list + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a list callback"))?; + // SAFETY: The borrowed element cursor and parent buffers remain live through this call. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +struct MapExport { + length: usize, + entries_length: usize, + keys_sorted: bool, + validity_kind: vx_velox_validity_kind, + owner: Arc, + keys: Box, + values: Box, +} + +impl MapExport { + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + let mut execution = session.create_execution_ctx(); + let map = array.execute::(&mut execution)?; + let keys_sorted = map.keys_sorted(); + let entries = map.entries().clone().downcast::(); + let entry_values = entries.elements().clone(); + let entry_struct = entry_values.execute::(&mut execution)?; + let fields = entry_struct.iter_unmasked_fields().collect::>(); + if fields.len() != 2 { + vortex_bail!( + "Vortex map entries require two fields, got {}", + fields.len() + ); + } + let metadata = prepare_list_metadata(&entries, session, memory_callbacks)?; + Ok(Self { + length: metadata.length, + entries_length: metadata.elements_length, + keys_sorted, + validity_kind: metadata.validity_kind, + owner: metadata.owner, + keys: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(fields[0].clone(), session, memory_callbacks)?, + }), + values: Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(fields[1].clone(), session, memory_callbacks)?, + }), + }) + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + let end = offset + .checked_add(length) + .ok_or_else(|| vortex_err!("Vortex Velox export range overflow"))?; + if end > self.length { + vortex_bail!( + "Vortex Velox export range is out of bounds: {offset}..{end}, array length {}", + self.length + ); + } + let (validity, validity_length, validity_bit_offset) = + if self.validity_kind == VX_VELOX_VALIDITY_BITMAP { + packed_bits_window( + self.owner + .validity + .as_ref() + .ok_or_else(|| vortex_err!("Map validity bitmap is missing"))?, + offset, + length, + )? + } else { + (ptr::null(), 0, 0) + }; + let offsets = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.offsets.as_ptr().add(offset) } + }; + let sizes = if length == 0 { + ptr::null() + } else { + // SAFETY: The checked range lies within the metadata arrays. + unsafe { self.owner.sizes.as_ptr().add(offset) } + }; + let view = vx_velox_map_view { + struct_size: size_of::(), + length, + offsets, + sizes, + keys: &raw const *self.keys, + values: &raw const *self.values, + entries_length: self.entries_length, + keys_sorted: self.keys_sorted, + validity_kind: self.validity_kind, + validity, + validity_length, + validity_bit_offset, + buffers: vx_velox_buffer_owner { + struct_size: size_of::(), + owner: Arc::as_ptr(&self.owner).cast(), + retain: Some(retain_list_owner), + release: Some(release_list_owner), + retained_bytes: self.owner.retained_bytes, + }, + offsets_alignment: pointer_alignment(offsets.cast()), + sizes_alignment: pointer_alignment(sizes.cast()), + validity_alignment: pointer_alignment(validity), + }; + let callback = visitor + .visit_map + .ok_or_else(|| vortex_err!("Vortex Velox visitor requires a map callback"))?; + // SAFETY: The borrowed child cursors and parent buffers remain live through this callback. + let status = unsafe { callback(visitor.context, &raw const view) }; + if status != 0 { + vortex_bail!("{}", callback_error(visitor, status)); + } + Ok(()) + } +} + +impl CursorExport { + fn date_storage( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + ) -> VortexResult> { + let DType::Extension(ext_dtype) = array.dtype() else { + return Ok(None); + }; + let Some(time_unit) = ext_dtype.metadata_opt::() else { + return Ok(None); + }; + if *time_unit != TimeUnit::Days { + vortex_bail!( + "Vortex Velox visitor does not support date unit {time_unit}; Velox DATE uses days" + ); + } + + if let Some(extension) = array.as_opt::() { + return Ok(Some(extension.storage_array().clone())); + } + let mut execution = session.create_execution_ctx(); + let extension = array.execute::(&mut execution)?; + Ok(Some(extension.storage_array().clone())) + } + + fn try_new_canonical( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + if matches!(array.dtype(), DType::Map(..)) { + Ok(Self::Map(MapExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::List(..)) { + Ok(Self::List(ListExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Struct(..)) { + Ok(Self::Struct(StructExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Decimal(..)) { + Ok(Self::Primitive(PrimitiveExport::try_new_decimal( + array, + session, + memory_callbacks, + )?)) + } else if let Some(storage) = Self::date_storage(array.clone(), session)? { + Ok(Self::Primitive(PrimitiveExport::try_new( + storage, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Bool(_)) { + Ok(Self::Bool(BoolExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { + Ok(Self::VarBin(VarBinExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else { + Ok(Self::Primitive(PrimitiveExport::try_new( + array, + session, + memory_callbacks, + )?)) + } + } + + fn try_new( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + memory_callbacks: Option, + ) -> VortexResult { + if array.is::() { + Ok(Self::Dictionary(DictionaryExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else if array.is::() { + Ok(Self::Constant(ConstantExport::try_new( + array, + session, + memory_callbacks, + )?)) + } else { + Self::try_new_canonical(array, session, memory_callbacks) + } + } + + fn visit(&self, offset: usize, length: usize, visitor: &vx_velox_visitor) -> VortexResult<()> { + match self { + Self::Primitive(export) => export.visit(offset, length, visitor), + Self::Bool(export) => export.visit(offset, length, visitor), + Self::VarBin(export) => export.visit(offset, length, visitor), + Self::Dictionary(export) => export.visit(offset, length, visitor), + Self::Constant(export) => export.visit(offset, length, visitor), + Self::Struct(export) => export.visit(offset, length, visitor), + Self::List(export) => export.visit(offset, length, visitor), + Self::Map(export) => export.visit(offset, length, visitor), + } + } +} + +fn exported_validity( + is_nullable: bool, + mask: Mask, +) -> (vx_velox_validity_kind, Option) { + if !is_nullable { + return (VX_VELOX_VALIDITY_NON_NULLABLE, None); + } + match mask { + Mask::AllTrue(_) => (VX_VELOX_VALIDITY_ALL_VALID, None), + Mask::AllFalse(_) => (VX_VELOX_VALIDITY_ALL_INVALID, None), + Mask::Values(values) => (VX_VELOX_VALIDITY_BITMAP, Some(values.bit_buffer().clone())), + } +} + +unsafe extern "C" fn retain_primitive_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_primitive_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_bool_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_bool_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_varbin_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_varbin_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_struct_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_struct_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +unsafe extern "C" fn retain_list_owner(owner: *const c_void) { + // SAFETY: The visitor receives a pointer from `Arc::as_ptr` while one strong reference lives. + unsafe { Arc::increment_strong_count(owner.cast::()) }; +} + +unsafe extern "C" fn release_list_owner(owner: *const c_void) { + // SAFETY: Each release matches a prior retain of this `Arc` pointer. + drop(unsafe { Arc::from_raw(owner.cast::()) }); +} + +fn validate_visitor(visitor: &vx_velox_visitor) -> VortexResult<()> { + if visitor.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox visitor structure is too small: expected at least {}, got {}", + size_of::(), + visitor.struct_size + ); + } + if visitor.abi_version != crate::VX_VELOX_ABI_VERSION { + vortex_bail!( + "Unsupported Vortex Velox ABI version: expected {}, got {}", + crate::VX_VELOX_ABI_VERSION, + visitor.abi_version + ); + } + Ok(()) +} + +fn callback_error(visitor: &vx_velox_visitor, status: i32) -> String { + let Some(last_error) = visitor.last_error else { + return format!("Velox visitor failed with status {status}"); + }; + // SAFETY: The callback contract returns null or a valid null-terminated string. + let message = unsafe { last_error(visitor.context) }; + if message.is_null() { + return format!("Velox visitor failed with status {status}"); + } + // SAFETY: The callback keeps the string valid until the next callback. + unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned() +} + +fn selected_array( + array: &vortex::array::ArrayRef, + request: &vx_velox_visit_request, +) -> VortexResult { + if request.rows.is_null() { + if request.row_count != 0 { + vortex_bail!("A null visitor row pointer requires a zero row count"); + } + return Ok(array.clone()); + } + // SAFETY: The caller supplies `row_count` readable positions. + let rows = unsafe { slice::from_raw_parts(request.rows, request.row_count) }; + let mut previous = None; + for row in rows { + let position = usize::try_from(*row) + .map_err(|_| vortex_err!("Visitor row does not fit usize: {}", row))?; + if position >= array.len() { + vortex_bail!( + "Visitor row is out of bounds: row {}, array length {}", + row, + array.len() + ); + } + if previous.is_some_and(|previous| previous >= *row) { + vortex_bail!("Visitor rows must be unique and increasing"); + } + previous = Some(*row); + } + let dense = rows.len() == array.len() + && rows + .iter() + .enumerate() + .all(|(position, row)| *row == position as u64); + if dense { + return Ok(array.clone()); + } + array.take(PrimitiveArray::from_iter(rows.iter().copied()).into_array()) +} + +fn visit_array( + array: vortex::array::ArrayRef, + session: &vortex::session::VortexSession, + visitor: &vx_velox_visitor, +) -> VortexResult<()> { + let length = array.len(); + CursorExport::try_new_canonical(array, session, None)?.visit(0, length, visitor) +} + +/// Create one export cursor for several Velox output windows. +/// +/// # Safety +/// +/// The session and array pointers must identify live handles. +/// The memory callbacks must identify a complete, thread-safe callback table. +/// `error_out` must be null or valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_export_cursor_new( + session: *const vx_velox_session, + array: *const vx_velox_array, + memory_callbacks: *const vx_velox_arrow_memory_callbacks, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_export_cursor { + try_or(error_out, ptr::null_mut(), || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + let memory_callbacks = unsafe { parse_memory_callbacks(memory_callbacks)? }; + Ok(Box::into_raw(Box::new(vx_velox_export_cursor { + export: CursorExport::try_new(array.clone(), session, Some(memory_callbacks))?, + }))) + }) +} + +/// Free one export cursor. +/// +/// # Safety +/// +/// The pointer must be null or come from [`vx_velox_export_cursor_new`]. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_velox_export_cursor_free(cursor: *mut vx_velox_export_cursor) { + if !cursor.is_null() { + // SAFETY: The pointer came from `Box::into_raw` and is freed once. + drop(unsafe { Box::from_raw(cursor) }); + } +} + +/// Visit one contiguous range from a retained export cursor. +/// +/// # Safety +/// +/// The cursor and visitor pointers must remain live until this call returns. +/// Concurrent calls are valid. The caller must not free the cursor before all calls return. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_export_cursor_visit( + cursor: *const vx_velox_export_cursor, + offset: usize, + length: usize, + visitor: *const vx_velox_visitor, + error_out: *mut *mut vx_velox_error, +) -> i32 { + try_or(error_out, 1, || { + let cursor = unsafe { + cursor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox export cursor must not be null"))? + }; + let visitor = unsafe { + visitor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? + }; + validate_visitor(visitor)?; + cursor.export.visit(offset, length, visitor)?; + Ok(0) + }) +} + +/// Visit one Vortex array through host semantic callbacks. +/// +/// The request selects source positions once. Callback block positions are compact and follow the +/// request order. +/// +/// # Safety +/// +/// Every pointer must be null or valid for the documented access. The array and session handles +/// must remain live until this call returns. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_velox_array_visit( + session: *const vx_velox_session, + array: *const vx_velox_array, + request: *const vx_velox_visit_request, + visitor: *const vx_velox_visitor, + error_out: *mut *mut vx_velox_error, +) -> i32 { + try_or(error_out, 1, || { + let session = unsafe { vx_session_ref(session)? }; + let array = unsafe { vx_array_ref(array)? }; + let request = unsafe { + request + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visit request must not be null"))? + }; + if request.struct_size < size_of::() { + vortex_bail!( + "Vortex Velox visit request is too small: expected at least {}, got {}", + size_of::(), + request.struct_size + ); + } + let visitor = unsafe { + visitor + .as_ref() + .ok_or_else(|| vortex_err!("Vortex Velox visitor must not be null"))? + }; + validate_visitor(visitor)?; + visit_array(selected_array(array, request)?, session, visitor)?; + Ok(0) + }) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-velox/src/visitor/export/tests.rs b/vortex-velox/src/visitor/export/tests.rs new file mode 100644 index 00000000000..b8976c95436 --- /dev/null +++ b/vortex-velox/src/visitor/export/tests.rs @@ -0,0 +1,1523 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::align_of; +use std::ptr; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex::array::ArrayRef; +use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::DictArray; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::MapArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::TemporalArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::validity::Validity; +use vortex::buffer::buffer; +use vortex::dtype::DecimalDType; +use vortex::dtype::FieldNames; +use vortex::dtype::MapDType; +use vortex::dtype::Nullability; +use vortex::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_fastlanes::BitPackedData; + +use super::*; +use crate::api::vx_velox_array_free; +use crate::ffi::vx_array_new_with; +use crate::ffi::vx_session_free; +use crate::ffi::vx_session_new_with; + +#[derive(Default)] +struct TestMemory { + retained_bytes: AtomicUsize, +} + +unsafe extern "C" fn retain_test_memory(_context: *mut c_void) {} + +unsafe extern "C" fn release_test_memory(_context: *mut c_void) {} + +unsafe extern "C" fn reserve_test_memory(context: *mut c_void, bytes: usize) -> i32 { + // SAFETY: The test context stays live through every callback. + let memory = unsafe { &*context.cast::() }; + memory.retained_bytes.fetch_add(bytes, Ordering::Relaxed); + 0 +} + +unsafe extern "C" fn free_test_memory(context: *mut c_void, bytes: usize) { + // SAFETY: The test context stays live through every callback. + let memory = unsafe { &*context.cast::() }; + memory.retained_bytes.fetch_sub(bytes, Ordering::Relaxed); +} + +fn test_memory_callbacks(memory: &mut TestMemory) -> vx_velox_arrow_memory_callbacks { + vx_velox_arrow_memory_callbacks { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (memory as *mut TestMemory).cast(), + retain_context: Some(retain_test_memory), + release_context: Some(release_test_memory), + report_allocation: Some(reserve_test_memory), + report_free: Some(free_test_memory), + last_error: None, + } +} + +#[rstest] +#[case(PType::U8, VX_VELOX_PRIMITIVE_U8)] +#[case(PType::U16, VX_VELOX_PRIMITIVE_U16)] +#[case(PType::U32, VX_VELOX_PRIMITIVE_U32)] +#[case(PType::U64, VX_VELOX_PRIMITIVE_U64)] +#[case(PType::I8, VX_VELOX_PRIMITIVE_I8)] +#[case(PType::I16, VX_VELOX_PRIMITIVE_I16)] +#[case(PType::I32, VX_VELOX_PRIMITIVE_I32)] +#[case(PType::I64, VX_VELOX_PRIMITIVE_I64)] +#[case(PType::F16, VX_VELOX_PRIMITIVE_F16)] +#[case(PType::F32, VX_VELOX_PRIMITIVE_F32)] +#[case(PType::F64, VX_VELOX_PRIMITIVE_F64)] +fn maps_primitive_types(#[case] input: PType, #[case] expected: vx_velox_primitive_type) { + assert_eq!(primitive_type_id(input), expected); +} + +#[test] +fn date_days_use_i32_storage_and_millisecond_dates_are_rejected() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let days = TemporalArray::new_date( + PrimitiveArray::from_option_iter([Some(-1_i32), None, Some(19_000)]).into_array(), + TimeUnit::Days, + ) + .into_array(); + let CursorExport::Primitive(export) = CursorExport::try_new_canonical(days, &session, None)? + else { + vortex_bail!("date visitor did not produce primitive storage"); + }; + assert_eq!(export.primitive_type, VX_VELOX_PRIMITIVE_I32); + let view = export.view(0, 3)?; + // SAFETY: The export owns three readable i32 values. + let values = unsafe { slice::from_raw_parts(view.values.cast::(), 3) }; + assert_eq!(values, [-1, 0, 19_000]); + assert_eq!(view.validity_kind, VX_VELOX_VALIDITY_BITMAP); + + let milliseconds = TemporalArray::new_date( + PrimitiveArray::from_iter([86_400_000_i64]).into_array(), + TimeUnit::Milliseconds, + ) + .into_array(); + let error = match CursorExport::try_new_canonical(milliseconds, &session, None) { + Ok(_) => vortex_bail!("millisecond date visitor unexpectedly succeeded"), + Err(error) => error, + }; + assert!(error.to_string().contains("Velox DATE uses days")); + Ok(()) +} + +#[test] +fn decimals_normalize_to_velox_storage_widths() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let short = DecimalArray::new( + buffer![1_i8, -2, 3], + DecimalDType::new(18, 2), + Validity::NonNullable, + ) + .into_array(); + let short = PrimitiveExport::try_new_decimal(short, &session, None)?; + assert_eq!(short.primitive_type, VX_VELOX_PRIMITIVE_I64); + let short_view = short.view(0, 3)?; + assert_eq!(short_view.decimal_precision, 18); + assert_eq!(short_view.decimal_scale, 2); + // SAFETY: The export owns three readable i64 values. + let short_values = unsafe { slice::from_raw_parts(short_view.values.cast::(), 3) }; + assert_eq!(short_values, [1, -2, 3]); + + let nullable_short = DecimalArray::new( + buffer![1_i128, i128::MAX], + DecimalDType::new(18, 2), + Validity::from_iter([true, false]), + ) + .into_array(); + let nullable_short = PrimitiveExport::try_new_decimal(nullable_short, &session, None)?; + let nullable_short_view = nullable_short.view(0, 2)?; + // SAFETY: The export owns two readable i64 values. + let nullable_short_values = + unsafe { slice::from_raw_parts(nullable_short_view.values.cast::(), 2) }; + assert_eq!(nullable_short_values, [1, 0]); + assert_eq!(nullable_short_view.validity_kind, VX_VELOX_VALIDITY_BITMAP); + + let long = DecimalArray::new( + buffer![1_i64, -2, 3], + DecimalDType::new(30, 4), + Validity::NonNullable, + ) + .into_array(); + let long = PrimitiveExport::try_new_decimal(long, &session, None)?; + assert_eq!(long.primitive_type, VX_VELOX_PRIMITIVE_I128); + let long_view = long.view(0, 3)?; + assert_eq!(long_view.decimal_precision, 30); + assert_eq!(long_view.decimal_scale, 4); + // SAFETY: The export owns three readable i128 values. + let long_values = unsafe { slice::from_raw_parts(long_view.values.cast::(), 3) }; + assert_eq!(long_values, [1, -2, 3]); + + let unsupported = DecimalArray::new( + buffer![1_i8], + DecimalDType::new(39, 0), + Validity::NonNullable, + ) + .into_array(); + let error = match PrimitiveExport::try_new_decimal(unsupported, &session, None) { + Ok(_) => vortex_bail!("precision 39 decimal visitor unexpectedly succeeded"), + Err(error) => error, + }; + assert!(error.to_string().contains("decimal precision 39")); + Ok(()) +} + +#[test] +fn dictionary_export_preserves_code_width_and_nullable_children() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let code_cases: [(ArrayRef, vx_velox_primitive_type); 4] = [ + (buffer![0_u8, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U8), + (buffer![0_u16, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U16), + (buffer![0_u32, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U32), + (buffer![0_u64, 1, 0].into_array(), VX_VELOX_PRIMITIVE_U64), + ]; + for (codes, expected_type) in code_cases { + let dictionary = DictArray::try_new(codes, buffer![10_i64, 20].into_array())?; + let CursorExport::Dictionary(export) = + CursorExport::try_new(dictionary.into_array(), &session, None)? + else { + vortex_bail!("dictionary export lost its outer encoding"); + }; + assert_eq!(export.codes.primitive_type, expected_type); + assert_eq!(export.values_length, 2); + assert!(matches!(export.values.export, CursorExport::Primitive(_))); + } + + let codes = PrimitiveArray::from_option_iter([Some(0_u8), None, Some(1)]).into_array(); + let values = PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(); + let dictionary = DictArray::try_new(codes, values)?; + let CursorExport::Dictionary(export) = + CursorExport::try_new(dictionary.into_array(), &session, None)? + else { + vortex_bail!("nullable dictionary export lost its outer encoding"); + }; + assert_eq!(export.codes.validity_kind, VX_VELOX_VALIDITY_BITMAP); + let CursorExport::Primitive(values) = &export.values.export else { + vortex_bail!("nullable dictionary values lost their primitive representation"); + }; + assert_eq!(values.validity_kind, VX_VELOX_VALIDITY_BITMAP); + Ok(()) +} + +#[test] +fn constant_export_preserves_null_value() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let constant = ConstantArray::new(Scalar::null_native::(), 10).into_array(); + let CursorExport::Constant(export) = CursorExport::try_new(constant, &session, None)? else { + vortex_bail!("constant export lost its outer encoding"); + }; + assert_eq!(export.length, 10); + let CursorExport::Primitive(value) = &export.value.export else { + vortex_bail!("null constant lost its primitive representation"); + }; + assert_eq!(value.length, 1); + assert_eq!(value.validity_kind, VX_VELOX_VALIDITY_ALL_INVALID); + Ok(()) +} + +#[test] +fn struct_export_preserves_children_and_nonzero_window() -> VortexResult<()> { + #[derive(Default)] + struct StructCapture { + length: usize, + offset: usize, + fields: *const *const vx_velox_export_cursor, + field_count: usize, + validity: *const u8, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_struct( + context: *mut c_void, + view: *const vx_velox_struct_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offset = view.offset; + capture.fields = view.fields; + capture.field_count = view.field_count; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let length: usize = 130; + let dictionary = DictArray::try_new( + PrimitiveArray::from_iter((0..length).map(|index| [0_u8, 1][index % 2])).into_array(), + buffer![10_i64, 20].into_array(), + )? + .into_array(); + let constant = ConstantArray::new(Scalar::from(7_i64), length).into_array(); + let parent_validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); + let struct_array = StructArray::new( + FieldNames::from(["dictionary", "constant"]), + [dictionary, constant], + length, + parent_validity, + ) + .into_array(); + let CursorExport::Struct(export) = CursorExport::try_new(struct_array, &session, None)? else { + vortex_bail!("struct export lost its outer encoding"); + }; + assert!(matches!( + export.fields[0].export, + CursorExport::Dictionary(_) + )); + assert!(matches!(export.fields[1].export, CursorExport::Constant(_))); + + let mut capture = StructCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: Some(capture_struct), + visit_list: None, + visit_map: None, + }; + export.visit(65, 63, &visitor)?; + assert_eq!(capture.length, 63); + assert_eq!(capture.offset, 65); + assert_eq!(capture.field_count, 2); + assert_eq!(capture.validity_bit_offset, 1); + // SAFETY: The export retains both field cursors until it is dropped below. + assert_eq!(unsafe { *capture.fields }, &raw const export.fields[0]); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("struct callback returned no validity owner"))?; + drop(export); + // SAFETY: The callback retained the parent owner before the cursor was dropped. + assert!( + unsafe { + *capture + .validity + .add(capture.validity_bit_offset / u8::BITS as usize) + } != 0 + ); + let release = owner + .release + .ok_or_else(|| vortex_err!("struct owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + Ok(()) +} + +#[test] +fn list_export_preserves_elements_window_and_accounting() -> VortexResult<()> { + #[derive(Default)] + struct ListCapture { + length: usize, + offsets: *const i32, + sizes: *const i32, + elements_length: usize, + validity: *const u8, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_list( + context: *mut c_void, + view: *const vx_velox_list_view, + ) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offsets = view.offsets; + capture.sizes = view.sizes; + capture.elements_length = view.elements_length; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let length = 130; + let elements = DictArray::try_new( + buffer![0_u8, 1, 0, 1, 0, 1].into_array(), + PrimitiveArray::from_option_iter([Some(10_i64), None]).into_array(), + )? + .into_array(); + let offsets = PrimitiveArray::from_iter((0..length).map(|index| [0_u32, 2, 4][index % 3])); + let sizes = + PrimitiveArray::from_iter((0..length).map(|index| if index % 10 == 0 { 0 } else { 2 })); + let validity = Validity::from_iter((0..length).map(|index| index % 9 != 0)); + let list = ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity) + .into_array(); + let mut memory = TestMemory::default(); + let CursorExport::List(export) = + CursorExport::try_new(list, &session, Some(test_memory_callbacks(&mut memory)))? + else { + vortex_bail!("list export lost its outer encoding"); + }; + assert!(matches!( + export.elements.export, + CursorExport::Dictionary(_) + )); + let expected_parent_bytes = + length * 2 * size_of::() + length.div_ceil(u64::BITS as usize) * size_of::(); + assert_eq!(export.owner.retained_bytes, expected_parent_bytes); + + let mut capture = ListCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: Some(capture_list), + visit_map: None, + }; + export.visit(65, 63, &visitor)?; + assert_eq!(capture.length, 63); + assert_eq!(capture.elements_length, 6); + assert_eq!(capture.validity_bit_offset, 1); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.offsets }, 4); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.sizes }, 2); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("list callback returned no owner"))?; + drop(export); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_parent_bytes + ); + // SAFETY: The callback retained the owner before the export was dropped. + assert_eq!(unsafe { *capture.offsets.add(1) }, 0); + let release = owner + .release + .ok_or_else(|| vortex_err!("list owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn map_export_preserves_children_window_and_accounting() -> VortexResult<()> { + #[derive(Default)] + struct MapCapture { + length: usize, + offsets: *const i32, + sizes: *const i32, + keys: *const vx_velox_export_cursor, + values: *const vx_velox_export_cursor, + entries_length: usize, + keys_sorted: bool, + validity_bit_offset: usize, + owner: Option, + } + + unsafe extern "C" fn capture_map(context: *mut c_void, view: *const vx_velox_map_view) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.offsets = view.offsets; + capture.sizes = view.sizes; + capture.keys = view.keys; + capture.values = view.values; + capture.entries_length = view.entries_length; + capture.keys_sorted = view.keys_sorted; + capture.validity_bit_offset = view.validity_bit_offset; + capture.owner = Some(view.buffers); + 0 + } + + let session = vortex::session::VortexSession::empty(); + let keys = DictArray::try_new( + buffer![0_u8, 1, 0, 1, 0, 1].into_array(), + buffer![10_i64, 20].into_array(), + )? + .into_array(); + let values = ConstantArray::new(Scalar::from(7_i64), 6).into_array(); + let entries = StructArray::new( + FieldNames::from(["key", "value"]), + [keys, values], + 6, + Validity::NonNullable, + ) + .into_array(); + let entry_lists = ListViewArray::new( + entries, + buffer![0_u32, 2, 4].into_array(), + buffer![2_u32, 2, 2].into_array(), + Validity::from_iter([true, false, true]), + ); + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I64, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::NonNullable), + true, + )?; + let map = MapArray::try_new(map_dtype, entry_lists)?.into_array(); + let mut memory = TestMemory::default(); + let CursorExport::Map(export) = + CursorExport::try_new(map, &session, Some(test_memory_callbacks(&mut memory)))? + else { + vortex_bail!("map export lost its outer encoding"); + }; + assert!(matches!(export.keys.export, CursorExport::Dictionary(_))); + assert!(matches!(export.values.export, CursorExport::Constant(_))); + let expected_parent_bytes = 3 * 2 * size_of::() + size_of::(); + assert_eq!(export.owner.retained_bytes, expected_parent_bytes); + + let mut capture = MapCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: Some(capture_map), + }; + export.visit(1, 2, &visitor)?; + assert_eq!(capture.length, 2); + assert_eq!(capture.entries_length, 6); + assert!(capture.keys_sorted); + assert_eq!(capture.validity_bit_offset, 1); + assert_eq!(capture.keys, &raw const *export.keys); + assert_eq!(capture.values, &raw const *export.values); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.offsets }, 2); + // SAFETY: The retained owner keeps both metadata arrays live. + assert_eq!(unsafe { *capture.sizes }, 2); + let owner = capture + .owner + .ok_or_else(|| vortex_err!("map callback returned no owner"))?; + drop(export); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_parent_bytes + ); + // SAFETY: The callback retained the owner before the export was dropped. + assert_eq!(unsafe { *capture.offsets.add(1) }, 4); + let release = owner + .release + .ok_or_else(|| vortex_err!("map owner returned no release callback"))?; + // SAFETY: This release matches the callback retain above. + unsafe { release(owner.owner) }; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[derive(Default)] +struct Capture { + primitive_type: Option, + length: usize, + values: *const u8, + values_length: usize, + values_alignment: usize, + validity: *const u8, + validity_length: usize, + validity_bit_offset: usize, + validity_alignment: usize, + retained_bytes: usize, + validity_kind: Option, + owner: Option, +} + +unsafe extern "C" fn capture_primitive( + context: *mut c_void, + view: *const vx_velox_primitive_view, +) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live `Capture` and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.primitive_type = Some(view.primitive_type); + capture.length = view.length; + capture.values = view.values; + capture.values_length = view.values_length; + capture.values_alignment = view.values_alignment; + capture.validity = view.validity; + capture.validity_length = view.validity_length; + capture.validity_bit_offset = view.validity_bit_offset; + capture.validity_alignment = view.validity_alignment; + capture.retained_bytes = view.buffers.retained_bytes; + capture.validity_kind = Some(view.validity_kind); + capture.owner = Some(view.buffers); + 0 +} + +fn release_capture(capture: &Capture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_primitive`. + unsafe { release(owner.owner) }; + Ok(()) +} + +#[derive(Default)] +struct VarBinCapture { + struct_size: usize, + kind: Option, + length: usize, + views: *const vx_velox_binary_view, + views_length: usize, + views_alignment: usize, + data_buffers: *const vx_velox_byte_buffer_view, + data_buffer_count: usize, + validity: *const u8, + validity_length: usize, + validity_bit_offset: usize, + validity_alignment: usize, + validity_kind: Option, + retained_bytes: usize, + owner: Option, +} + +unsafe extern "C" fn capture_varbin( + context: *mut c_void, + view: *const vx_velox_varbin_view, +) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.struct_size = view.struct_size; + capture.kind = Some(view.kind); + capture.length = view.length; + capture.views = view.views; + capture.views_length = view.views_length; + capture.views_alignment = view.views_alignment; + capture.data_buffers = view.data_buffers; + capture.data_buffer_count = view.data_buffer_count; + capture.validity = view.validity; + capture.validity_length = view.validity_length; + capture.validity_bit_offset = view.validity_bit_offset; + capture.validity_alignment = view.validity_alignment; + capture.validity_kind = Some(view.validity_kind); + capture.retained_bytes = view.buffers.retained_bytes; + capture.owner = Some(view.buffers); + 0 +} + +fn release_varbin_capture(capture: &VarBinCapture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained string owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("string owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_varbin`. + unsafe { release(owner.owner) }; + Ok(()) +} + +#[derive(Default)] +struct BoolCapture { + length: usize, + values: *const u8, + values_bit_offset: usize, + validity: *const u8, + validity_bit_offset: usize, + validity_kind: Option, + retained_bytes: usize, + owner: Option, +} + +unsafe extern "C" fn capture_bool(context: *mut c_void, view: *const vx_velox_bool_view) -> i32 { + if context.is_null() || view.is_null() { + return 1; + } + // SAFETY: The test passes pointers to live capture and view objects. + let (capture, view) = unsafe { (&mut *context.cast::(), &*view) }; + let Some(retain) = view.buffers.retain else { + return 2; + }; + // SAFETY: The visitor owner is live for the callback. + unsafe { retain(view.buffers.owner) }; + capture.length = view.length; + capture.values = view.values; + capture.values_bit_offset = view.values_bit_offset; + capture.validity = view.validity; + capture.validity_bit_offset = view.validity_bit_offset; + capture.validity_kind = Some(view.validity_kind); + capture.retained_bytes = view.buffers.retained_bytes; + capture.owner = Some(view.buffers); + 0 +} + +fn release_bool_capture(capture: &BoolCapture) -> VortexResult<()> { + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained Boolean owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("Boolean owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_bool`. + unsafe { release(owner.owner) }; + Ok(()) +} + +#[expect( + clippy::host_endian_bytes, + reason = "The Vortex binary-view fields use the host C ABI layout" +)] +unsafe fn captured_varbin_value(capture: &VarBinCapture, index: usize) -> Option<&[u8]> { + if capture.validity_kind == Some(VX_VELOX_VALIDITY_BITMAP) { + let bit_index = capture.validity_bit_offset + index; + // SAFETY: The callback contract retains the bitmap for every captured row. + let byte = unsafe { *capture.validity.add(bit_index / 8) }; + if byte & (1 << (bit_index % 8)) == 0 { + return None; + } + } + // SAFETY: The callback contract retains `length` readable views. + let view = unsafe { &*capture.views.add(index) }; + let length = view.length as usize; + const INLINE_LENGTH: usize = size_of::() - size_of::(); + if length <= INLINE_LENGTH { + return Some(&view.data[..length]); + } + let buffer_index = + u32::from_ne_bytes([view.data[4], view.data[5], view.data[6], view.data[7]]) as usize; + let offset = + u32::from_ne_bytes([view.data[8], view.data[9], view.data[10], view.data[11]]) as usize; + // SAFETY: The callback contract retains all payload descriptors. + let buffer = unsafe { &*capture.data_buffers.add(buffer_index) }; + // SAFETY: Canonical Vortex views contain validated payload ranges. + Some(unsafe { slice::from_raw_parts(buffer.data.add(offset), length) }) +} + +#[rstest] +#[case(DType::Utf8(Nullability::Nullable), VX_VELOX_VARBIN_UTF8)] +#[case(DType::Binary(Nullability::Nullable), VX_VELOX_VARBIN_BINARY)] +fn varbin_cursor_retains_mixed_views_across_nonzero_window( + #[case] dtype: DType, + #[case] expected_kind: vx_velox_varbin_kind, +) -> VortexResult<()> { + let utf8_expected: [Option<&[u8]>; 7] = [ + Some(b""), + Some(b"a"), + None, + Some(b"abcdefghijkl"), + Some(b"abcdefghijklm"), + Some("vortex 🌀 outlined".as_bytes()), + Some(b"tail"), + ]; + let binary_expected: [Option<&[u8]>; 7] = [ + Some(b""), + Some(b"\xff"), + None, + Some(b"abcdefghijkl"), + Some(b"\x00abcdefghijklm"), + Some(b"\xff\x00 binary outlined value"), + Some(b"tail"), + ]; + let expected = if matches!(dtype, DType::Utf8(_)) { + utf8_expected + } else { + binary_expected + }; + let session = vx_session_new_with(|session| session); + let varbin = VarBinViewArray::from_iter(expected, dtype); + let array = vx_array_new_with(varbin.into_array()); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "string cursor creation failed"); + vortex_ensure!(error.is_null(), "string cursor returned an error"); + + let mut capture = VarBinCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: Some(capture_varbin), + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = + unsafe { vx_velox_export_cursor_visit(cursor, 1, 5, &raw const visitor, &raw mut error) }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "string export window returned an error"); + assert_eq!(capture.struct_size, size_of::()); + assert_eq!(capture.kind, Some(expected_kind)); + assert_eq!(capture.length, 5); + assert_eq!(capture.views_length, 5 * size_of::()); + assert!(capture.views_alignment >= align_of::()); + assert_eq!(capture.views.addr() % align_of::(), 0); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.validity_bit_offset, 1); + assert!(capture.validity_length >= 1); + assert!(capture.validity_alignment >= align_of::()); + assert_eq!(capture.validity.addr() % align_of::(), 0); + assert!(capture.data_buffer_count >= 1); + assert!(!capture.data_buffers.is_null()); + assert_eq!( + capture.retained_bytes, + memory.retained_bytes.load(Ordering::Relaxed) + ); + + // SAFETY: Each owned handle is freed once. The callback retained the string owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + capture.retained_bytes + ); + for (index, expected) in expected[1..6].iter().enumerate() { + // SAFETY: The retained owner keeps every captured pointer live. + let actual = unsafe { captured_varbin_value(&capture, index) }; + assert_eq!(actual, *expected); + } + release_varbin_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn varbin_shared_buffers_compact_into_exact_owned_storage() -> VortexResult<()> { + let length = 130_usize; + let strings = VarBinViewArray::from_iter( + (0..length) + .map(|index| (index % 11 != 0).then(|| format!("outlined string value {index:03}"))), + DType::Utf8(Nullability::Nullable), + ); + let parts = strings.into_data_parts(); + let views_length = parts.views.try_to_host_sync()?.len(); + let data_length = parts + .buffers + .iter() + .map(|buffer| Ok(buffer.try_to_host_sync()?.len())) + .sum::>()?; + let descriptor_length = parts.buffers.len() * size_of::(); + let validity_length = length.div_ceil(u64::BITS as usize) * size_of::(); + let expected_retained = views_length + data_length + descriptor_length + validity_length; + + let retained_views = parts.views.clone(); + let retained_buffers = Arc::<[BufferHandle]>::clone(&parts.buffers); + let mut execution = vortex::session::VortexSession::empty().create_execution_ctx(); + let mask = parts.validity.execute_mask(length, &mut execution)?; + let (_, validity) = exported_validity(true, mask); + let owner = VarBinOwner::try_new(parts.views, parts.buffers, validity, length)?; + + assert!(matches!(owner.views, RetainedViews::Compact(_))); + assert!( + owner + ._data + .iter() + .all(|buffer| matches!(buffer, RetainedBytes::Compact(_))) + ); + assert_eq!(owner.retained_bytes, expected_retained); + drop(retained_views); + drop(retained_buffers); + Ok(()) +} + +#[test] +fn retained_varbin_buffers_report_complete_unique_allocations() -> VortexResult<()> { + let alignment = vortex::buffer::Alignment::new(256); + let mut payload = BufferMut::::with_capacity_aligned(17, alignment); + payload.extend(0..17); + let expected_payload_allocation = payload.allocation_size(); + let (retained_payload, payload_allocation) = + RetainedBytes::try_new(BufferHandle::new_host(payload.freeze()))?; + assert!(matches!(retained_payload, RetainedBytes::Retained(_))); + assert_eq!(payload_allocation, expected_payload_allocation); + assert!(payload_allocation > 17); + + let mut views = + BufferMut::::with_capacity_aligned(2 * size_of::(), alignment); + views.extend(std::iter::repeat_n( + 0, + 2 * size_of::(), + )); + let expected_views_allocation = views.allocation_size(); + let (retained_views, views_allocation) = + RetainedViews::try_new(BufferHandle::new_host(views.freeze()))?; + assert!(matches!(retained_views, RetainedViews::Retained(_))); + assert_eq!(views_allocation, expected_views_allocation); + assert!(views_allocation > 2 * size_of::()); + Ok(()) +} + +#[test] +fn word_aligned_windows_rebase_validity_buffers() -> VortexResult<()> { + let session = vortex::session::VortexSession::empty(); + let primitive = PrimitiveArray::from_option_iter( + (0..130).map(|index| (index % 7 != 0).then_some(index as i64)), + ) + .into_array(); + let primitive = PrimitiveExport::try_new(primitive, &session, None)?; + let primitive_first = primitive.view(0, 64)?; + let primitive_second = primitive.view(64, 64)?; + assert_eq!(primitive_first.validity_bit_offset, 0); + assert_eq!(primitive_second.validity_bit_offset, 0); + // SAFETY: Both pointers lie in the retained validity allocation. + assert_eq!(primitive_second.validity, unsafe { + primitive_first.validity.add(size_of::()) + }); + + let strings = VarBinViewArray::from_iter( + (0..130).map(|index| (index % 11 != 0).then(|| format!("value-{index}"))), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let strings = VarBinExport::try_new(strings, &session, None)?; + let mut first = VarBinCapture::default(); + let first_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut first).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: Some(capture_varbin), + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + strings.visit(0, 64, &first_visitor)?; + + let mut second = VarBinCapture::default(); + let second_visitor = vx_velox_visitor { + context: (&raw mut second).cast(), + ..first_visitor + }; + strings.visit(64, 64, &second_visitor)?; + assert_eq!(first.validity_bit_offset, 0); + assert_eq!(second.validity_bit_offset, 0); + // SAFETY: Both pointers lie in the retained validity allocation. + assert_eq!(second.validity, unsafe { + first.validity.add(size_of::()) + }); + release_varbin_capture(&first)?; + release_varbin_capture(&second)?; + Ok(()) +} + +#[test] +fn bool_cursor_retains_nonzero_window_and_exact_accounting() -> VortexResult<()> { + let expected = (0..130) + .map(|index| (index % 11 != 0).then_some(index % 3 == 0)) + .collect::>(); + let session = vx_session_new_with(|session| session); + let boolean = BoolArray::from_iter(expected.iter().copied()); + let array = vx_array_new_with(boolean.into_array()); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "Boolean cursor creation failed"); + vortex_ensure!(error.is_null(), "Boolean cursor returned an error"); + + let mut capture = BoolCapture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: None, + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: Some(capture_bool), + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = + unsafe { vx_velox_export_cursor_visit(cursor, 65, 63, &raw const visitor, &raw mut error) }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "Boolean export window returned an error"); + assert_eq!(capture.length, 63); + assert_eq!(capture.values_bit_offset, 1); + assert_eq!(capture.validity_bit_offset, 1); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.retained_bytes, 6 * size_of::()); + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 48); + + // SAFETY: Each owned handle is freed once. The callback retained the Boolean owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + capture.retained_bytes + ); + for (relative_index, expected) in expected[65..128].iter().enumerate() { + let value_bit = capture.values_bit_offset + relative_index; + let validity_bit = capture.validity_bit_offset + relative_index; + // SAFETY: The retained buffers cover every captured value and validity bit. + let (actual, is_valid) = unsafe { + ( + *capture.values.add(value_bit / 8) & (1 << (value_bit % 8)) != 0, + *capture.validity.add(validity_bit / 8) & (1 << (validity_bit % 8)) != 0, + ) + }; + assert_eq!(is_valid, expected.is_some()); + if let Some(expected) = expected { + assert_eq!(actual, *expected); + } + } + release_bool_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn export_cursor_reuses_one_prepared_array_across_windows() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with( + PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40), Some(50)]) + .into_array(), + ); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); + vortex_ensure!(error.is_null(), "export cursor returned an error"); + assert!(memory.retained_bytes.load(Ordering::Relaxed) >= 48); + + let mut first = Capture::default(); + let first_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut first).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 1, 2, &raw const first_visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "first export window returned an error"); + assert_eq!(first.length, 2); + assert_eq!(first.validity_bit_offset, 1); + // SAFETY: The callback retained two readable i64 values. + let first_values = unsafe { slice::from_raw_parts(first.values.cast::(), 2) }; + assert_eq!(first_values, [0, 30]); + assert_eq!( + first.retained_bytes, + memory.retained_bytes.load(Ordering::Relaxed) + ); + let owner = first + .owner + .ok_or_else(|| vortex_err!("first export window returned no owner"))? + .owner; + release_capture(&first)?; + + let mut second = Capture::default(); + let second_visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut second).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit(cursor, 3, 2, &raw const second_visitor, &raw mut error) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "second export window returned an error"); + assert_eq!(second.length, 2); + assert_eq!(second.validity_bit_offset, 3); + assert_eq!( + second + .owner + .ok_or_else(|| vortex_err!("second export window returned no owner"))? + .owner, + owner + ); + + // SAFETY: Each owned handle is freed exactly once. The second callback retained the owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + // SAFETY: The retained cursor owner keeps these two i64 values live. + let second_values = unsafe { slice::from_raw_parts(second.values.cast::(), 2) }; + assert_eq!(second_values, [40, 50]); + release_capture(&second)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn export_cursor_decodes_sliced_bitpacked_into_exact_owner() -> VortexResult<()> { + let session = vx_session_new_with(|session| { + vortex_fastlanes::initialize(&session); + session + }); + let session_ref = unsafe { vx_session_ref(session)? }; + let values = (0..2_050).map(|index| (index % 7 != 0).then_some(i64::from(index % 100))); + let primitive = PrimitiveArray::from_option_iter(values).into_array(); + let mut execution = session_ref.create_execution_ctx(); + let bitpacked = BitPackedData::encode(&primitive, 7, &mut execution)?; + vortex_ensure!( + bitpacked.patches().is_none(), + "test bit-packed array unexpectedly contains patches" + ); + let slice_begin = 113; + let slice_end = 1_941; + let sliced = bitpacked.into_array().slice(slice_begin..slice_end)?; + let array = vx_array_new_with(sliced); + let mut error = ptr::null_mut(); + let mut memory = TestMemory::default(); + let memory_callbacks = test_memory_callbacks(&mut memory); + // SAFETY: The session and array handles remain live until cursor creation finishes. + let cursor = unsafe { + vx_velox_export_cursor_new(session, array, &raw const memory_callbacks, &raw mut error) + }; + vortex_ensure!(!cursor.is_null(), "export cursor creation failed"); + vortex_ensure!(error.is_null(), "export cursor returned an error"); + let sliced_length = slice_end - slice_begin; + let expected_retained = sliced_length * size_of::() + + sliced_length.div_ceil(u64::BITS as usize) * size_of::(); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_retained + ); + + let window_offset = 997; + let window_length = 6; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + // SAFETY: The cursor and callback state remain live through the call. + let status = unsafe { + vx_velox_export_cursor_visit( + cursor, + window_offset, + window_length, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "export window returned an error"); + assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!( + capture.validity_bit_offset, + window_offset % u64::BITS as usize + ); + assert_eq!(capture.retained_bytes, expected_retained); + // SAFETY: The callback retained `window_length` readable i64 values. + let actual = unsafe { slice::from_raw_parts(capture.values.cast::(), window_length) }; + for (relative_index, value) in actual.iter().enumerate() { + let sliced_index = window_offset + relative_index; + let source_index = slice_begin + sliced_index; + // SAFETY: The retained bitmap covers every row in the sliced array. + let validity_index = capture.validity_bit_offset + relative_index; + let validity_byte = unsafe { *capture.validity.add(validity_index / 8) }; + let is_valid = validity_byte & (1 << (validity_index % 8)) != 0; + assert_eq!(is_valid, source_index % 7 != 0); + if is_valid { + assert_eq!(*value, i64::try_from(source_index % 100)?); + } + } + + // SAFETY: Each owned handle is freed exactly once. The callback retained the owner. + unsafe { + vx_velox_export_cursor_free(cursor); + vx_velox_array_free(array); + vx_session_free(session); + } + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + expected_retained + ); + release_capture(&capture)?; + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn patched_bitpacked_uses_retained_canonical_fallback() -> VortexResult<()> { + let session = vx_session_new_with(|session| { + vortex_fastlanes::initialize(&session); + session + }); + let session_ref = unsafe { vx_session_ref(session)? }; + let expected = [1_u64, 2, 3, u64::MAX]; + let primitive = PrimitiveArray::from_iter(expected).into_array(); + let mut execution = session_ref.create_execution_ctx(); + let bitpacked = BitPackedData::encode(&primitive, 2, &mut execution)?; + vortex_ensure!( + bitpacked.patches().is_some(), + "test bit-packed array unexpectedly omitted patches" + ); + let mut memory = TestMemory::default(); + let export = PrimitiveExport::try_new( + bitpacked.into_array(), + session_ref, + Some(test_memory_callbacks(&mut memory)), + )?; + assert!(matches!(export.owner.values, PrimitiveValues::Retained(_))); + assert_eq!( + memory.retained_bytes.load(Ordering::Relaxed), + export.owner.retained_bytes() + ); + // SAFETY: The export owner contains `expected.len()` initialized u64 values. + let actual = + unsafe { slice::from_raw_parts(export.owner.values().cast::(), expected.len()) }; + assert_eq!(actual, expected); + drop(export); + assert_eq!(memory.retained_bytes.load(Ordering::Relaxed), 0); + unsafe { vx_session_free(session) }; + Ok(()) +} + +#[test] +fn visits_sparse_nullable_values_with_retained_buffers() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let array = vx_array_new_with( + PrimitiveArray::from_option_iter([Some(10_i64), None, Some(30), Some(40)]).into_array(), + ); + let rows = [1_u64, 3]; + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: rows.as_ptr(), + row_count: rows.len(), + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + let mut error = ptr::null_mut(); + // SAFETY: Every handle and callback object stays live for this call. + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_eq!(capture.primitive_type, Some(VX_VELOX_PRIMITIVE_I64)); + assert_eq!(capture.length, 2); + assert_eq!(capture.values_length, 2 * size_of::()); + assert!(capture.values_alignment.is_power_of_two()); + assert_eq!(capture.values.addr() % capture.values_alignment, 0); + assert_eq!(capture.validity_kind, Some(VX_VELOX_VALIDITY_BITMAP)); + assert_eq!(capture.validity_length, size_of::()); + assert_eq!(capture.validity_bit_offset, 0); + assert!(capture.validity_alignment.is_power_of_two()); + assert_eq!(capture.validity.addr() % capture.validity_alignment, 0); + assert_eq!( + capture.retained_bytes, + capture.values_length + size_of::() + ); + // SAFETY: The callback retained the owner before storing these pointers. + let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 2) }; + assert_eq!(values, [0, 40]); + // SAFETY: The retained validity pointer has one readable word. + let validity = unsafe { *capture.validity }; + assert_eq!(validity & 0b11, 0b10); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + // SAFETY: This release matches the retain in `capture_primitive`. + unsafe { release(owner.owner) }; + // SAFETY: Each owned handle is freed exactly once. + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) +} + +#[test] +fn copies_sliced_values_into_exact_owned_storage() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let source = PrimitiveArray::from_iter(0_i32..16); + let source_values = source.buffer_handle().try_to_host_sync()?; + // SAFETY: The source contains sixteen i32 values. The fifth value is in bounds. + let source_slice = unsafe { source_values.as_ptr().add(5 * size_of::()) }; + drop(source_values); + let array = vx_array_new_with(source.into_array().slice(5..8)?); + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: ptr::null(), + row_count: 0, + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + let mut error = ptr::null_mut(); + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_eq!(capture.values_length, 3 * size_of::()); + assert_eq!(capture.retained_bytes, 2 * size_of::()); + assert_ne!(capture.values, source_slice); + // SAFETY: Each owned handle is freed exactly once. The callback retained the value owner. + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + // SAFETY: The retained compact buffer contains three i32 values. + let values = unsafe { slice::from_raw_parts(capture.values.cast::(), 3) }; + assert_eq!(values, [5, 6, 7]); + assert!(capture.values_alignment.is_power_of_two()); + assert_eq!(capture.values.addr() % capture.values_alignment, 0); + assert_eq!(capture.validity_alignment, 0); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + unsafe { release(owner.owner) }; + Ok(()) +} + +#[test] +fn copies_validity_into_word_padded_storage() -> VortexResult<()> { + let session = vx_session_new_with(|session| session); + let session_ref = unsafe { vx_session_ref(session)? }; + let primitive = PrimitiveArray::from_option_iter([Some(1_i32), None, Some(3)]); + let mut execution = session_ref.create_execution_ctx(); + let Mask::Values(mask) = primitive + .validity()? + .execute_mask(primitive.len(), &mut execution)? + else { + vortex_bail!("Expected bitmap validity"); + }; + let expected_validity = mask.bit_buffer().inner().as_ptr(); + let array = vx_array_new_with(primitive.into_array()); + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: ptr::null(), + row_count: 0, + }; + let mut capture = Capture::default(); + let visitor = vx_velox_visitor { + struct_size: size_of::(), + abi_version: crate::VX_VELOX_ABI_VERSION, + context: (&raw mut capture).cast(), + visit_primitive: Some(capture_primitive), + last_error: None, + visit_varbin: None, + visit_dictionary: None, + visit_constant: None, + visit_bool: None, + visit_struct: None, + visit_list: None, + visit_map: None, + }; + let mut error = ptr::null_mut(); + let status = unsafe { + vx_velox_array_visit( + session, + array, + &raw const request, + &raw const visitor, + &raw mut error, + ) + }; + assert_eq!(status, 0); + vortex_ensure!(error.is_null(), "visitor returned an error"); + assert_ne!(capture.validity, expected_validity); + assert_eq!(capture.validity_bit_offset, 0); + assert_eq!(capture.validity_length, size_of::()); + assert!(capture.validity_alignment >= align_of::()); + assert_eq!( + capture.retained_bytes, + capture.values_length.div_ceil(size_of::()) * size_of::() + size_of::() + ); + + let owner = capture + .owner + .ok_or_else(|| vortex_err!("visitor did not return a retained owner"))?; + let release = owner + .release + .ok_or_else(|| vortex_err!("visitor owner did not return a release callback"))?; + unsafe { release(owner.owner) }; + unsafe { + vx_velox_array_free(array); + vx_session_free(session); + } + Ok(()) +} + +#[test] +fn rejects_unsorted_rows() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(); + let rows = [2_u64, 1]; + let request = vx_velox_visit_request { + struct_size: size_of::(), + rows: rows.as_ptr(), + row_count: rows.len(), + }; + match selected_array(&array, &request) { + Ok(_) => vortex_bail!("unsorted rows unexpectedly succeeded"), + Err(error) => assert!(error.to_string().contains("unique and increasing")), + } + Ok(()) +}