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 0a367df1534..a967a2dd743 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", @@ -11689,6 +11690,28 @@ 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", + "async-fs", + "bytes", + "cbindgen", + "futures", + "rstest", + "vortex", + "vortex-array", + "vortex-arrow", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-io", +] + [[package]] name = "vortex-web-wasm" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 294ed20d996..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", @@ -161,7 +162,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/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/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..7f27e417356 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs @@ -0,0 +1,68 @@ +// 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; + +/// 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. +/// +/// 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_DIRECT_INTEGER_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..85339ad4ab9 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,15 @@ // 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; +#[doc(hidden)] +pub use integer_membership::MAX_DIRECT_INTEGER_MEMBERS; pub use kernel::*; use num_traits::Zero; use vortex_buffer::BitBuffer; @@ -13,6 +17,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 +61,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 +106,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 +153,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 +161,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 +220,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 +255,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 +275,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 +303,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 +328,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 +350,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 +361,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 +495,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 +519,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 +625,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 +637,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 +669,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 +769,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 +890,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 +918,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 +1016,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 +1106,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-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-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 { 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 new file mode 100644 index 00000000000..22b88c855f3 --- /dev/null +++ b/vortex-velox/Cargo.toml @@ -0,0 +1,51 @@ +[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 } +async-fs = { 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-fastlanes = { workspace = true } +vortex-io = { workspace = true } +vortex = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } + +[build-dependencies] +cbindgen = { 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..d55ef1f0bff --- /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`. 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 +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/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 new file mode 100644 index 00000000000..68ed4bdd35a --- /dev/null +++ b/vortex-velox/cinclude/vortex_velox.h @@ -0,0 +1,1770 @@ +// 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; + + +#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) + +/** + * 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; + +/** + * A fixed-width binary expression operator identifier. + */ +typedef uint32_t vx_velox_binary_operator; + +/** + * A fixed-width row-selection mode identifier. + */ +typedef uint32_t vx_velox_scan_selection_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; + +/** + * 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; + +/** + * 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; + +/** + * 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; + +/** + * 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); + +/** + * 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); + +/** + * Create a logical negation. + * + * # Safety + * + * `child` must point to a live expression. + */ +vx_velox_expression *vx_velox_expression_not(const vx_velox_expression *child); + +/** + * 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); + +/** + * 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); + +/** + * 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); + +/** + * Free an expression. + * + * # Safety + * + * `expression` must be null or an owned expression handle. + */ +void vx_velox_expression_free(const vx_velox_expression *expression); + +/** + * 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); + +/** + * 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); + +/** + * Free a scan. + * + * # Safety + * + * `scan` must be null or an owned scan handle. + */ +void vx_velox_scan_free(const vx_velox_scan *scan); + +/** + * Return the next partition from a scan. + * + * # Safety + * + * `scan` must point to a live scan. `error_out` must be null or valid. + */ +vx_velox_partition *vx_velox_scan_next_partition(vx_velox_scan *scan, vx_velox_error **error_out); + +/** + * Free a partition. + * + * # Safety + * + * `partition` must be null or an owned partition handle. + */ +void vx_velox_partition_free(const vx_velox_partition *partition); + +/** + * 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); + +/** + * Free an array. + * + * # Safety + * + * `array` must be null or an owned array handle. + */ +void vx_velox_array_free(const vx_velox_array *array); + +/** + * Return an array length. + * + * # Safety + * + * `array` must point to a live array. + */ +size_t vx_velox_array_len(const vx_velox_array *array); + +/** + * Slice an array. + * + * # 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); + +/** + * 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); + +/** + * 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); + +/** + * 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); + +/** + * 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); + +/** + * 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 len, + vx_velox_view row_index_name, + vx_velox_error **error_out); + +/** + * 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); + +/** + * 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); + +/** + * 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. + */ +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 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 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. + */ +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_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 +} // extern "C" +#endif // __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 new file mode 100644 index 00000000000..e80d510b14f --- /dev/null +++ b/vortex-velox/src/api.rs @@ -0,0 +1,840 @@ +// 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::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; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_error::vortex_bail; + +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; +/// 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_velox_expression, + /// The exact filter expression, or null for no filter. + 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. + 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_velox_error, +) -> vx_velox_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_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_velox_session { + 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_velox_session, +) -> *mut vx_velox_session { + unsafe { ffi::vx_session_clone(session) } +} + +/// 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_velox_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_velox_error, +) -> *const vx_velox_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_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_velox_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_velox_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." +); + +/// 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_velox_error, +) -> *mut vx_velox_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, + 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_velox_view, + nullable: bool, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_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_velox_error, +) -> *mut vx_velox_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_velox_dtype, + elements: *const *const vx_velox_scalar, + length: usize, + nullable: bool, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_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_velox_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_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_velox_expression { + 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_velox_view, + child: *const vx_velox_expression, +) -> *mut vx_velox_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_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(); + 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_velox_expression, + length: usize, +) -> *mut vx_velox_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_velox_expression, + length: usize, +) -> *mut vx_velox_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_velox_expression, +) -> *mut vx_velox_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_velox_expression, +) -> *mut vx_velox_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_velox_expression, + value: *const vx_velox_expression, +) -> *mut vx_velox_expression { + 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 +/// +/// `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_velox_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_velox_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_velox_data_source, + options: *const vx_velox_scan_options, + 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() + } 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_velox_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_velox_scan, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_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_velox_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_velox_partition, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_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_velox_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_velox_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_velox_array, + begin: usize, + end: usize, + error_out: *mut *mut vx_velox_error, +) -> *const vx_velox_array { + unsafe { ffi::vx_array_slice(array, begin, end, error_out) } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + 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 { + 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_velox_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 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 { + 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..b141c49fc3a --- /dev/null +++ b/vortex-velox/src/array.rs @@ -0,0 +1,881 @@ +// 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 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. +#[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, +} + +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. +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 { + pub(crate) 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, + }) + } + + 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; + 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) }; + } + } +} + +pub(crate) 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 words = vec![0_u64; bit_length.div_ceil(u64::BITS as usize)]; + for index in 0..nulls.len() { + if nulls.is_valid(index) { + let bit = data_offset + index; + words[bit / u64::BITS as usize] |= 1 << (bit % u64::BITS as usize); + } + } + Ok(NullBuffer::new(BooleanBuffer::new( + Buffer::from_vec(words), + 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; + +pub(crate) fn conservative_export_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_velox_session, + array: *const vx_velox_array, + index: usize, + 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)? }; + 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_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)? }; + 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_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_velox_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_export_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)?; + } + 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 + .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::mem::align_of; + 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 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 { + 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 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); + 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/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 new file mode 100644 index 00000000000..403ca0a179e --- /dev/null +++ b/vortex-velox/src/lib.rs @@ -0,0 +1,225 @@ +// 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 ffi; +mod projection; +mod read_at; +mod schema; +mod source; +mod temporal; +mod test_support; +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; +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_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 = 6; + +/// 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; + +/// 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; +/// 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)] +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 + | 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 + | VX_VELOX_CAPABILITY_SPLIT_ASSIGNMENT +} + +#[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 + ); + 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 new file mode 100644 index 00000000000..3a524a2bedc --- /dev/null +++ b/vortex-velox/src/projection.rs @@ -0,0 +1,182 @@ +// 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 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"); + &[] + } 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_velox_view, + len: usize, +) -> VortexResult<*mut vx_velox_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_velox_view, + len: usize, + 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() }?; + vortex_ensure!( + !row_index_name.is_empty(), + "row index field name must not be empty" + ); + + 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 { + vortex_ensure!( + name.as_ref() != row_index_name, + "row index field name conflicts with projected field: {name}" + ); + fields.push((name.clone(), get_item(name, root()))); + } + Ok(vx_expression_new_with(pack( + fields, + Nullability::NonNullable, + ))) +} + +/// 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_velox_view, + len: usize, + error_out: *mut *mut vx_velox_error, +) -> *mut vx_velox_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 +/// 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_velox_view, + len: usize, + 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) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::vx_velox_expression_free; + + fn view(value: &str) -> vx_velox_view { + vx_velox_view { + ptr: value.as_ptr().cast(), + len: value.len(), + } + } + + #[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")]; + 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 { crate::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..58b18a1955b --- /dev/null +++ b/vortex-velox/src/read_at.rs @@ -0,0 +1,1020 @@ +// 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_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)] +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_velox_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_velox_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 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::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; + 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; + let mut previous_assignment_byte = 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); + 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); + + 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..d7c0b113e0c --- /dev/null +++ b/vortex-velox/src/schema.rs @@ -0,0 +1,70 @@ +// 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 crate::ffi::try_or; +use crate::ffi::vx_velox_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. +/// +/// 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_velox_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())?; + 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) + }) +} + +#[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 { crate::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..fabb60d9165 --- /dev/null +++ b/vortex-velox/src/source.rs @@ -0,0 +1,297 @@ +// 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::VortexExpect; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +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. +#[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, + /// 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. +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_velox_session, + reader: *const vx_velox_read_at, + 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(); + 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_velox_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; + 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 +/// 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_velox_expression, + first_split: usize, + split_count: usize, + pruned_out: *mut u8, + error_out: *mut *mut vx_velox_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_velox_error, +) -> *const vx_velox_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/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/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 new file mode 100644 index 00000000000..1e6ef97a814 --- /dev/null +++ b/vortex-velox/src/visitor.rs @@ -0,0 +1,376 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::c_char; +use std::ffi::c_void; + +mod export; + +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; +/// 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; +/// Signed 128-bit primitive identifier. +pub const VX_VELOX_PRIMITIVE_I128: vx_velox_primitive_type = 11; + +/// 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 exact sum of the value and validity allocation sizes retained by this 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 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. + 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, +} + +/// 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)] +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 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. +#[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>, + /// 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>, +} 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(()) +} diff --git a/vortex-velox/tests/abi_contract.rs b/vortex-velox/tests/abi_contract.rs new file mode 100644 index 00000000000..6fb966615f8 --- /dev/null +++ b/vortex-velox/tests/abi_contract.rs @@ -0,0 +1,316 @@ +// 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(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_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 { + ($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, assignment_byte] + ); + check_layout!( + vx_velox_buffer_owner, + [struct_size, owner, retain, release, retained_bytes] + ); + check_layout!( + vx_velox_primitive_view, + [ + 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, + 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, + visit_varbin, + visit_dictionary, + visit_constant, + visit_bool, + visit_struct, + visit_list, + visit_map, + ] + ); + 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 compiler = std::env::var("CC").unwrap_or_else(|_| "cc".to_owned()); + compile_stdin( + &compiler, + &["-std=c11", "-fsyntax-only", "-x", "c", &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 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, "-"], + &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..ffd341d7069 --- /dev/null +++ b/vortex-velox/tests/velox_include_contract.cpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +struct ArrowSchema; +struct ArrowArray; +#include "vortex_velox.h" + +static_assert(VX_VELOX_ABI_VERSION == 6u); +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_velox_dtype *(*new_primitive)(vx_velox_ptype, bool, vx_velox_error **) = + vx_velox_dtype_new_primitive; + 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; + (void)new_primitive; + (void)new_binary; + (void)vx_velox_source_export_schema; + (void)vx_velox_data_source_scan; +}