From 97452ecf027756a8b21eccdb39108e1e28308fd0 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 12:56:57 +0000 Subject: [PATCH 01/13] feat(cuda): export FSST directly to Arrow varbin Decode FSST directly into device-resident Arrow offsets and values while retaining opt-in VarBinView export. Add session-level layout selection, coverage for both layouts, and CUDA benchmarks. Signed-off-by: Alexander Droste --- vortex-cuda/benches/fsst_cuda.rs | 126 ++++++-- vortex-cuda/kernels/src/fsst.cu | 125 ++++++-- vortex-cuda/src/arrow/canonical.rs | 143 +++++++-- vortex-cuda/src/arrow/mod.rs | 14 +- vortex-cuda/src/kernel/encodings/fsst.rs | 357 +++++++++++++++++++++-- vortex-cuda/src/kernel/encodings/mod.rs | 2 + vortex-cuda/src/lib.rs | 1 + vortex-cuda/src/session.rs | 23 ++ 8 files changed, 672 insertions(+), 119 deletions(-) diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 21db8744919..49930c9a30a 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -4,7 +4,6 @@ //! CUDA benchmarks for FSST decompression. #![expect(clippy::unwrap_used)] -#![expect(clippy::cast_possible_truncation)] #[allow(dead_code)] mod bench_config; @@ -18,12 +17,20 @@ use criterion::BenchmarkId; use criterion::Criterion; use criterion::Throughput; use futures::executor::block_on; +use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::match_each_integer_ptype; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArrayExt; use vortex::error::VortexExpect; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; +use vortex_cuda::VarBinExportLayout; +use vortex_cuda::arrow::DeviceArrayExt; +use vortex_cuda::arrow::release_device_array; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -36,30 +43,63 @@ use crate::timed_launch_strategy::TimedLaunchStrategy; // other kernels benchmark. const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")]; +fn session_with_varbin_layout(layout: VarBinExportLayout) -> vortex::session::VortexSession { + vortex::array::array_session().with_some( + CudaSession::try_default() + .vortex_expect("failed to create CUDA session") + .with_varbin_export_layout(layout), + ) +} + +struct FSSTBenchFixture { + utf8: ArrayRef, + binary: ArrayRef, + uncompressed_size: u64, +} + +fn make_fixture(n: usize) -> FSSTBenchFixture { + let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context"); + let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx()); + + let lens = fsst + .uncompressed_lengths() + .clone() + .execute::(setup_ctx.execution_ctx()) + .vortex_expect("canonicalize uncompressed_lengths"); + #[allow(clippy::unnecessary_cast)] + let uncompressed_size = match_each_integer_ptype!(lens.ptype(), |P| { + lens.as_slice::

().iter().map(|x| *x as u64).sum() + }); + + let binary = FSST::try_new( + DType::Binary(Nullability::NonNullable), + fsst.symbols().clone(), + fsst.symbol_lengths().clone(), + fsst.codes(), + fsst.uncompressed_lengths().clone(), + setup_ctx.execution_ctx(), + ) + .vortex_expect("rebuild FSST fixture with Binary dtype") + .into_array(); + + FSSTBenchFixture { + utf8: fsst.into_array(), + binary, + uncompressed_size, + } +} + fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { let mut group = c.benchmark_group("cuda"); for &(n, len_str) in BENCH_SIZES { - let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) - .vortex_expect("failed to create execution context"); - let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx()); - - let lens = fsst - .uncompressed_lengths() - .clone() - .execute::(setup_ctx.execution_ctx()) - .vortex_expect("canonicalize uncompressed_lengths"); - let total_size: usize = match_each_integer_ptype!(lens.ptype(), |P| { - lens.as_slice::

().iter().map(|x| *x as usize).sum() - }); - let uncompressed_size = total_size as u64; - - let fsst_array = fsst.into_array(); - - group.throughput(Throughput::Bytes(uncompressed_size)); + let fixture = make_fixture(n); + + group.throughput(Throughput::Bytes(fixture.uncompressed_size)); group.bench_with_input( - BenchmarkId::new("cuda/fsst/decompress", len_str), - &fsst_array, + BenchmarkId::new("cuda/fsst/decompress_to_varbinview", len_str), + &fixture.utf8, |b, fsst_array| { b.iter_custom(|iters| { let timed = TimedLaunchStrategy::default(); @@ -68,6 +108,7 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { @@ -77,6 +118,51 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { }); }, ); + + for (name, array, layout) in [ + ( + "cuda/fsst/decompress_to_varbin", + &fixture.utf8, + VarBinExportLayout::VarBin, + ), + ( + "cuda/fsst/export_binary", + &fixture.binary, + VarBinExportLayout::VarBin, + ), + ( + "cuda/fsst/export_utf8_view", + &fixture.utf8, + VarBinExportLayout::VarBinView, + ), + ( + "cuda/fsst/export_binary_view", + &fixture.binary, + VarBinExportLayout::VarBinView, + ), + ] { + group.bench_with_input(BenchmarkId::new(name, len_str), array, |b, array| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + + let session = session_with_varbin_layout(layout); + let mut cuda_ctx = CudaSession::create_execution_ctx(&session) + .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) + .with_launch_strategy(Arc::new(timed)); + + for _ in 0..iters { + let mut exported = + block_on((*array).clone().export_device_array(&mut cuda_ctx)) + .vortex_expect("export FSST device array"); + release_device_array(&mut exported); + } + + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }); + } } group.finish(); diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 938f262c3da..fb726eb26df 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -123,13 +123,18 @@ struct Scratch { } }; -template +// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12 +// bytes are stored inline after the u32 length. Longer values store their +// first four bytes, backing-buffer index, and byte offset. +constexpr uint32_t MAX_INLINED_SIZE = 12; + +template struct FSSTArgs { // Compressed FSST code stream, contiguous across all strings. String // `sid`'s codes live in `[codes_offsets[sid], codes_offsets[sid + 1])`. const uint8_t *__restrict codes_bytes; // Per-string offsets into `codes_bytes`, length `num_strings + 1`. - const OffT *__restrict codes_offsets; + const CodeOffsetT *__restrict codes_offsets; // FSST symbol table. const uint64_t *__restrict symbols; // Length in bytes (1..=8) of each entry in `symbols`. The remaining bits @@ -138,19 +143,55 @@ struct FSSTArgs { // Buffer to write decoded data into. uint8_t *__restrict output_bytes; // Per-string offsets into `output_bytes`, length `num_strings + 1`. - const uint64_t *__restrict output_offsets; + const OutputOffsetT *__restrict output_offsets; // Validity of each string. const uint8_t *__restrict validity_bits; + // Optional output views, one 16-byte uint4 per string. A null pointer + // requests bytes-only decoding for heaps that need the host rollover path. + uint4 *__restrict output_views; }; -template -__device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { +// Build one BinaryView from the bytes this thread just decoded. The Rust +// caller only provides output_views when every offset fits in the view's u32 +// fields and the decoded heap is exposed as backing buffer zero. +template +__device__ inline void fsst_write_view(const FSSTArgs &args, uint64_t sid) { + if (args.output_views == nullptr) { + return; + } + + const uint64_t start = args.output_offsets[sid]; + const uint32_t len = (uint32_t)(args.output_offsets[sid + 1] - start); + if (len <= MAX_INLINED_SIZE) { + uint32_t words[3] = {0, 0, 0}; +#pragma unroll + for (uint32_t i = 0; i < MAX_INLINED_SIZE; i++) { + if (i < len) { + words[i >> 2] |= (uint32_t)args.output_bytes[start + i] << (8u * (i & 3u)); + } + } + args.output_views[sid] = make_uint4(len, words[0], words[1], words[2]); + return; + } + + const uint32_t prefix = (uint32_t)args.output_bytes[start] | + ((uint32_t)args.output_bytes[start + 1] << 8u) | + ((uint32_t)args.output_bytes[start + 2] << 16u) | + ((uint32_t)args.output_bytes[start + 3] << 24u); + args.output_views[sid] = make_uint4(len, prefix, 0, (uint32_t)start); +} + +template +__device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) { + if (args.output_views != nullptr) { + args.output_views[sid] = make_uint4(0, 0, 0, 0); + } return; } - OffT in_pos = args.codes_offsets[sid]; - const OffT in_end = args.codes_offsets[sid + 1]; + CodeOffsetT in_pos = args.codes_offsets[sid]; + const CodeOffsetT in_end = args.codes_offsets[sid + 1]; uint64_t out_pos = args.output_offsets[sid]; const uint64_t out_end = args.output_offsets[sid + 1]; @@ -181,46 +222,66 @@ __device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t s sym &= mask; scratch.push(sym, len); - in_pos += (OffT)consumed; + in_pos += (CodeOffsetT)consumed; } // Epilogue: drain everything that's left. while (scratch.cursor > 0) { scratch.drain(args.output_bytes, out_pos, out_end); } + + fsst_write_view(args, sid); } -#define GENERATE_FSST_KERNEL(suffix, OffT) \ +#define FSST_GRID_STRIDE_LOOP(CodeOffsetT, OutputOffsetT, args) \ + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \ + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \ + const uint64_t block_end = (block_start + elements_per_block < num_strings) \ + ? (block_start + elements_per_block) \ + : num_strings; \ + for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \ + fsst_decode_string(args, sid); \ + } + +#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \ extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \ - const OffT *__restrict codes_offsets, \ + const CodeOffsetT *__restrict codes_offsets, \ const uint64_t *__restrict symbols, \ const uint8_t *__restrict symbol_lengths, \ const uint64_t *__restrict output_offsets, \ const uint8_t *__restrict validity_bits, \ uint8_t *__restrict output_bytes, \ + uint4 *__restrict output_views, \ uint64_t num_strings) { \ - const FSSTArgs args = { \ - codes_bytes, \ - codes_offsets, \ - symbols, \ - symbol_lengths, \ - output_bytes, \ - output_offsets, \ - validity_bits, \ + const FSSTArgs args = { \ + codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \ + validity_bits, output_views, \ }; \ - \ - const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \ - const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \ - const uint64_t block_end = (block_start + elements_per_block < num_strings) \ - ? (block_start + elements_per_block) \ - : num_strings; \ - \ - for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \ - fsst_decode_string(args, sid); \ - } \ + FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \ } -GENERATE_FSST_KERNEL(u8, uint8_t) -GENERATE_FSST_KERNEL(u16, uint16_t) -GENERATE_FSST_KERNEL(u32, uint32_t) -GENERATE_FSST_KERNEL(u64, uint64_t) +#define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \ + extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \ + const CodeOffsetT *__restrict codes_offsets, \ + const uint64_t *__restrict symbols, \ + const uint8_t *__restrict symbol_lengths, \ + const int32_t *__restrict output_offsets, \ + const uint8_t *__restrict validity_bits, \ + uint8_t *__restrict output_bytes, \ + uint64_t num_strings) { \ + const FSSTArgs args = { \ + codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \ + validity_bits, nullptr, \ + }; \ + FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \ + } + +GENERATE_FSST_VIEW_KERNEL(u8, uint8_t) +GENERATE_FSST_VIEW_KERNEL(u16, uint16_t) +GENERATE_FSST_VIEW_KERNEL(u32, uint32_t) +GENERATE_FSST_VIEW_KERNEL(u64, uint64_t) + +GENERATE_FSST_VARBIN_KERNEL(u8, uint8_t) +GENERATE_FSST_VARBIN_KERNEL(u16, uint16_t) +GENERATE_FSST_VARBIN_KERNEL(u32, uint32_t) +GENERATE_FSST_VARBIN_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index af141bfe548..4537b4dda31 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -54,6 +54,8 @@ use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::i256; +use vortex::encodings::fsst::FSST; +use vortex::encodings::fsst::FSSTArray; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; @@ -63,6 +65,7 @@ use vortex::extension::datetime::AnyTemporal; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; +use crate::VarBinExportLayout; use crate::arrow::ARROW_DEVICE_CUDA; use crate::arrow::ArrowArray; use crate::arrow::ArrowDeviceArray; @@ -78,6 +81,8 @@ use crate::cub::exclusive_sum_i32; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; +use crate::kernel::FSSTVarBin; +use crate::kernel::decode_fsst_varbin; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -220,6 +225,13 @@ fn export_array( Ok(list_view) => return export_list_view(list_view, ctx).await, Err(array) => array, }; + let array = match array.try_downcast::() { + Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { + return export_fsst_varbin(fsst, ctx).await; + } + Ok(fsst) => fsst.into_array(), + Err(array) => array, + }; let cuda_array = array.execute_cuda(ctx).await?; export_canonical(cuda_array, ctx).await @@ -303,8 +315,8 @@ fn export_canonical( export_fixed_size_list(fixed_size_list, ctx).await } Canonical::VarBinView(varbinview) => { - if matches!(varbinview.dtype(), DType::Binary(_)) { - return export_binary(varbinview, ctx).await; + if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin { + return export_varbin(varbinview, ctx).await; } let len = varbinview.len(); @@ -541,8 +553,8 @@ where Ok(BufferHandle::new_device(Arc::new(output_device))) } -/// Export Vortex binary views as an Arrow Device array with standard `Binary` layout. -async fn export_binary( +/// Export Vortex binary views as an Arrow Device array with standard `Utf8`/`Binary` layout. +async fn export_varbin( varbinview: VarBinViewArray, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { @@ -558,16 +570,44 @@ async fn export_binary( let views = ctx.ensure_on_device(views).await?; let (offsets, values) = export_binary_buffers(&views, &data_buffers, validity_buffer.as_ref(), len, ctx).await?; + export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) +} - let buffers = vec![validity_buffer, Some(offsets), Some(values)]; +async fn export_fsst_varbin( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let FSSTVarBin { + dtype, + len, + offsets, + values, + validity, + } = decode_fsst_varbin(fsst, ctx).await?; + vortex_ensure!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "FSST produced invalid variable-length dtype {dtype}" + ); + let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; + export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) +} +fn export_varbin_buffers( + len: usize, + validity_buffer: Option, + null_count: i64, + offsets: BufferHandle, + values: BufferHandle, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let buffers = vec![validity_buffer, Some(offsets), Some(values)]; let mut private_data = PrivateData::new(buffers, vec![], ctx)?; let sync_event = private_data.sync_event(); let arrow_array = ArrowArray { length: len as i64, null_count, offset: 0, - // Arrow Binary layout: optional null bitmap, i32 offsets, contiguous bytes. + // Arrow Utf8/Binary layout: optional null bitmap, i32 offsets, contiguous bytes. n_buffers: 3, buffers: private_data.buffer_ptrs.as_mut_ptr(), n_children: 0, @@ -1436,6 +1476,7 @@ mod tests { use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::device_buffer::cuda_backing_allocation; use crate::session::CudaSession; + use crate::session::VarBinExportLayout; unsafe fn release_exported_array(array: *mut ArrowArray) { unsafe { @@ -1445,6 +1486,12 @@ mod tests { } } + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + // Assert Arrow Device metadata that consumers use before reading buffers. fn assert_device_metadata( device_array: &ArrowDeviceArray, @@ -1812,7 +1859,7 @@ mod tests { field, Field::new( "", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8View)), + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), true, ) ); @@ -1833,7 +1880,13 @@ mod tests { ); let dictionary = unsafe { &*exported.array.array.dictionary }; - assert_varbinview_layout(dictionary, 2, 0, &[out_of_line.len()])?; + assert_binary_layout( + dictionary, + 2, + 0, + &[0, 5, i32::try_from(5 + out_of_line.len())?], + ["alpha", out_of_line].concat().as_bytes(), + )?; unsafe { release_exported_array(&raw mut exported.array.array) }; Ok(()) @@ -1894,7 +1947,7 @@ mod tests { "", DataType::Struct(Fields::from(vec![Field::new( "dict", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8View)), + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), true, )])), false, @@ -2182,9 +2235,8 @@ mod tests { } #[crate::test] - async fn test_export_varbinview() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + async fn test_export_varbinview_opt_in() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBinView)?; let out_of_line = "this is a longer string for out-of-line storage"; let array = VarBinViewArray::from_iter_str(["hello", "world", out_of_line]).into_array(); @@ -2561,27 +2613,38 @@ mod tests { #[rstest] #[case::utf8( multi_buffer_varbinview(DType::Utf8(Nullability::NonNullable)), - DataType::Utf8View + VarBinExportLayout::VarBin, + DataType::Utf8 )] #[case::binary( multi_buffer_varbinview(DType::Binary(Nullability::NonNullable)), + VarBinExportLayout::VarBin, DataType::Binary )] + #[case::utf8_view( + multi_buffer_varbinview(DType::Utf8(Nullability::NonNullable)), + VarBinExportLayout::VarBinView, + DataType::Utf8View + )] + #[case::binary_view( + multi_buffer_varbinview(DType::Binary(Nullability::NonNullable)), + VarBinExportLayout::VarBinView, + DataType::BinaryView + )] #[crate::test] async fn test_export_varbinview_multiple_variadic_buffers( #[case] fixture: (ArrayRef, [usize; 2]), + #[case] layout: VarBinExportLayout, #[case] expected_data_type: DataType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + let mut ctx = cuda_ctx_with_varbin_layout(layout)?; let (array, expected_data_buffer_lengths) = fixture; let mut exported = array.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - let is_binary = expected_data_type == DataType::Binary; assert_eq!(field, Field::new("", expected_data_type, false)); - if is_binary { + if layout == VarBinExportLayout::VarBin { assert_binary_layout( &exported.array.array, 3, @@ -2981,8 +3044,15 @@ mod tests { .slice(1..4)?; let mut exported = utf8.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - assert_eq!(field, Field::new("", DataType::Utf8View, true)); - assert_varbinview_shape(&exported.array.array, 3, 1)?; + assert_eq!(field, Field::new("", DataType::Utf8, true)); + let sliced_out_of_line = "this out-of-line value remains in the slice"; + assert_binary_layout( + &exported.array.array, + 3, + 1, + &[0, 5, 5, i32::try_from(5 + sliced_out_of_line.len())?], + ["hello", sliced_out_of_line].concat().as_bytes(), + )?; assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA); let private_data = unsafe { &*exported.array.array.private_data.cast::() }; @@ -3389,7 +3459,7 @@ mod tests { Ok(()) } - // Check nullable string-view exports include Arrow null bitmaps. + // Check nullable variable-length exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_varbinview() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -3402,7 +3472,7 @@ mod tests { Some("this is a longer string for out-of-line storage"), ]) .into_array(), - 4, + 3, 1, &mut ctx, ) @@ -3503,7 +3573,7 @@ mod tests { assert_eq!(nested_primitive_child.n_children, 0); let string_child = unsafe { &*nested_children[1] }; - assert_eq!(string_child.n_buffers, 4); + assert_eq!(string_child.n_buffers, 3); assert_eq!(string_child.n_children, 0); let string_buffers = unsafe { std::slice::from_raw_parts( @@ -3514,7 +3584,6 @@ mod tests { assert!(string_buffers[0].is_null()); assert!(!string_buffers[1].is_null()); assert!(!string_buffers[2].is_null()); - assert!(!string_buffers[3].is_null()); unsafe { release_exported_array(&raw mut device_array.array) }; Ok(()) @@ -3623,7 +3692,7 @@ mod tests { Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Utf8View, false), + Field::new("c", DataType::Utf8, false), ]) ); assert_eq!(exported.array.array.length, 5); @@ -3654,7 +3723,7 @@ mod tests { "nested", DataType::Struct(Fields::from(vec![ Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Utf8View, false), + Field::new("c", DataType::Utf8, false), ])), false, ), @@ -3749,23 +3818,35 @@ mod tests { } #[crate::test] - async fn test_export_varbinview_with_schema_uses_utf8_view_layout() -> VortexResult<()> { + async fn test_export_varbinview_with_schema_uses_utf8_layout() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let japanese = "こんにちは"; let long_emoji = "🦀 and 🚀 make this string out-of-line"; - let array = VarBinViewArray::from_iter_str(["", "hello", "é", "🦀", japanese, long_emoji]) - .into_array(); + let values = ["", "hello", "é", "🦀", japanese, long_emoji]; + let array = VarBinViewArray::from_iter_str(values).into_array(); let mut exported = array.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - assert_eq!(field, Field::new("", DataType::Utf8View, false)); - assert_varbinview_layout( + assert_eq!(field, Field::new("", DataType::Utf8, false)); + let mut expected_offsets = Vec::with_capacity(values.len() + 1); + expected_offsets.push(0i32); + for value in values { + expected_offsets.push( + expected_offsets + .last() + .copied() + .vortex_expect("offsets is non-empty") + + i32::try_from(value.len())?, + ); + } + assert_binary_layout( &exported.array.array, 6, 0, - &[japanese.len() + long_emoji.len()], + &expected_offsets, + values.concat().as_bytes(), )?; assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index aff582159b9..60fd2622cc5 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -58,6 +58,7 @@ use vortex_arrow::ArrowSessionExt; use crate::CudaBufferExt; use crate::CudaExecutionCtx; +use crate::VarBinExportLayout; mod arrow_c_abi { #![allow(dead_code)] @@ -814,10 +815,15 @@ fn arrow_device_export_field( .arrow() .to_arrow_field(name.as_ref(), dtype)?; - let data_type = match dtype { - DType::Binary(_) => DataType::Binary, - DType::Decimal(decimal_dtype, _) => arrow_device_export_decimal_data_type(*decimal_dtype), - DType::Struct(struct_dtype, _) => { + let data_type = match (ctx.cuda_session().varbin_export_layout(), dtype) { + (VarBinExportLayout::VarBin, DType::Utf8(_)) => DataType::Utf8, + (VarBinExportLayout::VarBin, DType::Binary(_)) => DataType::Binary, + (VarBinExportLayout::VarBinView, DType::Utf8(_)) => DataType::Utf8View, + (VarBinExportLayout::VarBinView, DType::Binary(_)) => DataType::BinaryView, + (_, DType::Decimal(decimal_dtype, _)) => { + arrow_device_export_decimal_data_type(*decimal_dtype) + } + (_, DType::Struct(struct_dtype, _)) => { DataType::Struct(arrow_device_export_struct_fields(struct_dtype, ctx)?.into()) } _ => return Ok(field), diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index eef0db049b8..fe94f82967c 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -17,14 +17,15 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::varbin::VarBinArrayExt; -use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex::array::arrays::varbinview::build_views::build_views; +use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; use vortex::array::match_each_integer_ptype; use vortex::array::match_each_unsigned_integer_ptype; +use vortex::array::validity::Validity; use vortex::buffer::Alignment; -use vortex::buffer::Buffer; +use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; @@ -38,6 +39,15 @@ use crate::CudaDeviceBuffer; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; +/// Device-resident offset-based result of FSST decompression. +pub(crate) struct FSSTVarBin { + pub(crate) dtype: DType, + pub(crate) len: usize, + pub(crate) offsets: BufferHandle, + pub(crate) values: BufferHandle, + pub(crate) validity: Validity, +} + /// CUDA decoder for FSST. #[derive(Debug)] pub(crate) struct FSSTExecutor; @@ -61,16 +71,15 @@ impl CudaExecute for FSSTExecutor { let dtype = fsst.dtype().clone(); let validity = fsst.codes().validity()?; - if fsst.is_empty() || validity.definitely_all_null() { - let empty = unsafe { - VarBinViewArray::new_unchecked( - Buffer::::zeroed(fsst.len()), - Arc::from([]), - dtype, - validity, - ) - }; - return Ok(Canonical::VarBinView(empty)); + if fsst.is_empty() { + return Ok(Canonical::empty(&dtype)); + } + + if validity.definitely_all_null() { + let views = ctx.copy_to_device(vec![0i128; fsst.len()])?.await?; + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })); } let lens = fsst @@ -105,6 +114,133 @@ impl CudaExecute for FSSTExecutor { } } +/// Decode FSST directly into Arrow-compatible i32 offsets and contiguous values on device. +pub(crate) async fn decode_fsst_varbin( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = fsst.dtype().clone(); + let validity = fsst.codes().validity()?; + let len = fsst.len(); + let lens = fsst + .uncompressed_lengths() + .clone() + .execute::(ctx.execution_ctx())?; + let codes_offsets = fsst + .codes() + .offsets() + .clone() + .execute::(ctx.execution_ctx())?; + + let output_offsets = match_each_integer_ptype!(lens.ptype(), |P| { + let mut offsets = Vec::with_capacity(lens.len() + 1); + let mut acc = 0u64; + offsets.push(0i32); + #[allow(clippy::unnecessary_cast)] + for &length in lens.as_slice::

() { + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("FSST uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("FSST decoded size overflow"))?; + offsets + .push(i32::try_from(acc).map_err(|_| { + vortex_err!("FSST decoded size exceeds Arrow i32 offset range") + })?); + } + VortexResult::Ok(offsets) + })?; + let total_size = usize::try_from( + *output_offsets + .last() + .vortex_expect("output_offsets has at least one entry"), + )?; + + if total_size == 0 { + let offsets = ctx.copy_to_device(output_offsets)?.await?; + let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); + let values = BufferHandle::new_device(allocation.slice(0..0)); + return Ok(FSSTVarBin { + dtype, + len, + offsets, + values, + validity, + }); + } + + match_each_unsigned_integer_ptype!(codes_offsets.ptype().to_unsigned(), |U| { + decode_fsst_varbin_typed::(fsst, codes_offsets, output_offsets, total_size, ctx).await + }) +} + +async fn decode_fsst_varbin_typed( + fsst: FSSTArray, + codes_offsets: PrimitiveArray, + output_offsets: Vec, + total_size: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + U: NativePType + DeviceRepr + Send + Sync + 'static, +{ + let dtype = fsst.dtype().clone(); + let validity = fsst.codes().validity()?; + let len = fsst.len(); + let len_u64 = len as u64; + let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); + let symbol_lengths = fsst.symbol_lengths().clone(); + let codes_bytes_handle = fsst.codes_bytes_handle().clone(); + let PrimitiveDataParts { + buffer: codes_offsets_buffer, + .. + } = codes_offsets.into_data_parts(); + let (.., validity_bits) = validity + .clone() + .execute_mask(len, ctx.execution_ctx())? + .into_bit_buffer() + .sliced() + .into_inner(); + + let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( + ctx.copy_to_device(symbols_u64)?, + ctx.copy_to_device(symbol_lengths)?, + ctx.copy_to_device(output_offsets)?, + ctx.copy_to_device(validity_bits.to_vec())?, + ctx.ensure_on_device(codes_bytes_handle), + ctx.ensure_on_device(codes_offsets_buffer), + )?; + + let output = ctx.device_alloc::(total_size)?; + let codes_bytes_view = codes_bytes.cuda_view::()?; + let codes_offsets_view = codes_offsets.cuda_view::()?; + let symbols_view = symbols.cuda_view::()?; + let symbol_lengths_view = symbol_lengths.cuda_view::()?; + let output_offsets_view = output_offsets.cuda_view::()?; + let validity_view = validity_device.cuda_view::()?; + let ptype = U::PTYPE.to_string(); + let cuda_function = ctx.load_function_with_suffixes("fsst", &["varbin", &ptype])?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&codes_bytes_view) + .arg(&codes_offsets_view) + .arg(&symbols_view) + .arg(&symbol_lengths_view) + .arg(&output_offsets_view) + .arg(&validity_view) + .arg(&output) + .arg(&len_u64); + })?; + + Ok(FSSTVarBin { + dtype, + len, + offsets: output_offsets, + values: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))), + validity, + }) +} + async fn decode_fsst( fsst: FSSTArray, codes_offsets: PrimitiveArray, @@ -126,6 +262,13 @@ where ) .vortex_expect("total_size fits in usize"); + if total_size == 0 { + let views = ctx.copy_to_device(vec![0i128; num_strings])?.await?; + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })); + } + let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); let symbol_lengths = fsst.symbol_lengths().clone(); let codes_bytes_handle = fsst.codes_bytes_handle().clone(); @@ -153,6 +296,9 @@ where // The kernel checks store alignment relative to the base via // `out_pos % N`, so the base must satisfy the widest store (u128 → 16). let device_output = ctx.device_alloc::(total_size)?; + let device_views = (total_size <= MAX_BUFFER_LEN) + .then(|| ctx.device_alloc::(num_strings)) + .transpose()?; let (output_base_ptr, _) = device_output.device_ptr(ctx.stream()); assert_eq!( output_base_ptr % 16, @@ -168,6 +314,7 @@ where let validity_view = validity_device.cuda_view::()?; let cuda_function = ctx.load_function("fsst", &[U::PTYPE])?; + let null_views = 0u64; ctx.launch_kernel(&cuda_function, num_strings, |args| { args.arg(&codes_bytes_view) .arg(&codes_offsets_view) @@ -175,10 +322,25 @@ where .arg(&symbol_lengths_view) .arg(&output_offsets_view) .arg(&validity_view) - .arg(&device_output) - .arg(&num_strings_u64); + .arg(&device_output); + if let Some(device_views) = device_views.as_ref() { + args.arg(device_views); + } else { + args.arg(&null_views); + } + args.arg(&num_strings_u64); })?; + if let Some(device_views) = device_views { + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); + let bytes = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_output))); + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) + })); + } + + // BinaryView offsets are u32. Retain the host rollover path for decoded heaps + // that need multiple backing buffers; ordinary batches stay entirely on-device. let host_bytes = CudaDeviceBuffer::new(device_output) .copy_to_host(Alignment::new(1))? .await?; @@ -200,10 +362,13 @@ where #[cfg(test)] mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; use rstest::rstest; use vortex::array::IntoArray; use vortex::array::arrays::VarBinArray; use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::encodings::fsst::fsst_compress; @@ -213,34 +378,73 @@ mod tests { use super::*; use crate::CanonicalCudaExt; + use crate::arrow::DeviceArrayExt; + use crate::arrow::release_device_array; + use crate::arrow::release_schema; use crate::session::CudaSession; + use crate::session::VarBinExportLayout; + + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + + fn assert_device_resident(canonical: &Canonical) { + let varbinview = canonical.as_varbinview(); + assert!(varbinview.views_handle().is_on_device()); + assert!( + varbinview + .data_buffers() + .iter() + .all(BufferHandle::is_on_device) + ); + } #[rstest] - #[case::non_null( + #[case::binary_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex fsst test string"[..])], + DType::Binary(Nullability::NonNullable), + )] + #[case::utf8_non_null( vec![Some(&b"the quick brown fox"[..]), Some(&b"jumps over the lazy dog"[..]), Some(&b"hello world"[..]), Some(&b"vortex fsst test string"[..])], - Nullability::NonNullable, + DType::Utf8(Nullability::NonNullable), )] - #[case::partial_nulls( + #[case::utf8_inline_boundary( + vec![Some(&b""[..]), + Some(&b"123456789012"[..]), + Some(&b"1234567890123"[..]), + Some(&b"this is another outlined value"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_partial_nulls( vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], - Nullability::Nullable, + DType::Utf8(Nullability::Nullable), + )] + #[case::binary_all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), )] - #[case::all_nulls( + #[case::binary_all_nulls( vec![None, None, None, None, None], - Nullability::Nullable, + DType::Binary(Nullability::Nullable), )] #[crate::test] async fn test_cuda_fsst_decompression_roundtrip( #[case] strings: Vec>, - #[case] nullability: Nullability, + #[case] dtype: DType, ) -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - let varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)).into_array(); + let varbin = VarBinArray::from_iter(strings, dtype.clone()).into_array(); let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; let fsst_array = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); @@ -248,12 +452,102 @@ mod tests { let gpu_result = FSSTExecutor .execute(fsst_array.clone(), &mut cuda_ctx) .await - .vortex_expect("GPU decompression failed") - .into_host() - .await? - .into_array(); + .vortex_expect("GPU decompression failed"); + assert_eq!(gpu_result.dtype(), &dtype); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(fsst_array, host_result, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_fsst_direct_varbin_output() -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: [&[u8]; 3] = [ + b"", + b"short", + b"this value is stored directly in the values buffer", + ]; + let varbin = VarBinArray::from_iter( + values.into_iter().map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?; + + let output = decode_fsst_varbin(fsst, &mut cuda_ctx).await?; + assert_eq!(output.dtype, DType::Utf8(Nullability::NonNullable)); + assert_eq!(output.len, values.len()); + assert!(output.offsets.is_on_device()); + assert!(output.values.is_on_device()); + + let offsets = Buffer::::from_byte_buffer(output.offsets.try_to_host()?.await?); + assert_eq!( + offsets.as_slice(), + &[0, 0, 5, i32::try_from(5 + values[2].len())?,] + ); + assert_eq!( + output.values.try_to_host()?.await?.as_ref(), + values.concat() + ); + Ok(()) + } + + #[rstest] + #[case::binary( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Binary, + 3 + )] + #[case::utf8( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Utf8, + 3 + )] + #[case::binary_view( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::BinaryView, + 4 + )] + #[case::utf8_view( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::Utf8View, + 4 + )] + #[crate::test] + async fn test_cuda_fsst_arrow_export_uses_dtype_layout( + #[case] dtype: DType, + #[case] layout: VarBinExportLayout, + #[case] expected_data_type: DataType, + #[case] expected_n_buffers: i64, + ) -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(layout)?; + let values = [ + Some(&b"short"[..]), + Some(&b"this value is stored out of line"[..]), + ]; + let varbin = VarBinArray::from_iter(values, dtype).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst_array = + fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); - assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); + let mut exported = fsst_array + .export_device_array_with_schema(&mut cuda_ctx) + .await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, false) + ); + assert_eq!(exported.array.array.n_buffers, expected_n_buffers); + + release_device_array(&mut exported.array); + release_schema(&mut exported.schema); Ok(()) } @@ -271,12 +565,11 @@ mod tests { let gpu_result = FSSTExecutor .execute(fsst_array.clone(), &mut cuda_ctx) .await - .vortex_expect("GPU decompression failed") - .into_host() - .await? - .into_array(); + .vortex_expect("GPU decompression failed"); + assert_device_resident(&gpu_result); - assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(fsst_array, host_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index b97cba13412..6ad37e7ead1 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -21,6 +21,8 @@ pub(crate) use date_time_parts::DateTimePartsExecutor; pub(crate) use decimal_byte_parts::DecimalBytePartsExecutor; pub(crate) use for_::FoRExecutor; pub(crate) use fsst::FSSTExecutor; +pub(crate) use fsst::FSSTVarBin; +pub(crate) use fsst::decode_fsst_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 07cfe40b06f..5817e15436b 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -62,6 +62,7 @@ pub use pooled_read_at::PooledFileReadAt; pub use pooled_read_at::PooledObjectStoreReadAt; pub use session::CudaSession; pub use session::CudaSessionExt; +pub use session::VarBinExportLayout; pub use stream::VortexCudaStream; pub use stream_pool::VortexCudaStreamPool; use vortex::array::ArrayVTable; diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index df3e3c13d55..a5410db0c99 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -30,6 +30,16 @@ use crate::stream_pool::VortexCudaStreamPool; /// Default maximum number of streams in the pool. const DEFAULT_STREAM_POOL_CAPACITY: usize = 4; +/// Arrow Device layout used when exporting variable-length UTF-8 and binary arrays. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum VarBinExportLayout { + /// Offset-based Arrow `Utf8`/`Binary` with one contiguous values buffer. + #[default] + VarBin, + /// Arrow `Utf8View`/`BinaryView` with 16-byte views and variadic data buffers. + VarBinView, +} + /// CUDA session for GPU accelerated execution. /// /// Maintains a registry of CUDA kernel implementations for array encodings. @@ -39,6 +49,7 @@ pub struct CudaSession { context: Arc, kernels: Arc>, export_device_array: Arc, + varbin_export_layout: VarBinExportLayout, kernel_loader: Arc, stream_pool: Arc, pinned_buffer_pool: Arc, @@ -65,11 +76,23 @@ impl CudaSession { kernels: Arc::new(DashMap::default()), kernel_loader: Arc::new(KernelLoader::new()), export_device_array: Arc::new(CanonicalDeviceArrayExport), + varbin_export_layout: VarBinExportLayout::default(), stream_pool, pinned_buffer_pool, } } + /// Selects the Arrow Device layout for variable-length UTF-8 and binary exports. + pub fn with_varbin_export_layout(mut self, layout: VarBinExportLayout) -> Self { + self.varbin_export_layout = layout; + self + } + + /// Returns the Arrow Device layout used for variable-length UTF-8 and binary exports. + pub fn varbin_export_layout(&self) -> VarBinExportLayout { + self.varbin_export_layout + } + /// Creates a default CUDA session using device 0, with all GPU array kernels preloaded. /// /// Unlike [`Default::default`], this returns an error instead of panicking when CUDA cannot be From 30da3923454dc56dbb37e6966090c5007661579e Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 13:23:34 +0000 Subject: [PATCH 02/13] fix(cuda): apply varbin export layout to list element schema fields arrow_device_export_field patched Utf8/Binary only at the top level of a dtype and recursed into structs, but List fell through to the plain to_arrow_field mapping. With the offset-based default layout, a schema derived from the dtype alone (e.g. for a chunked list of strings) declared Utf8View elements while the data was exported as 3-buffer Utf8, so consumers read offsets as views. Rebuild list element fields recursively like struct fields. Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 49 ++++++++++++++++++++++++++++++ vortex-cuda/src/arrow/mod.rs | 34 +++++++++++++-------- 2 files changed, 70 insertions(+), 13 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 4537b4dda31..be09fb0452b 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -2397,6 +2397,55 @@ mod tests { Ok(()) } + // Regression test: when the export schema is derived from the dtype alone (the top-level + // array is not one of the concretely-handled encodings), list element fields must still + // reflect the session's varbin export layout, or consumers would read the offset-based + // element data through a view-typed schema. + #[rstest] + #[case::varbin(VarBinExportLayout::VarBin, DataType::Utf8, 3)] + #[case::varbin_view(VarBinExportLayout::VarBinView, DataType::Utf8View, 4)] + #[crate::test] + async fn test_export_chunked_list_utf8_element_matches_schema( + #[case] layout: VarBinExportLayout, + #[case] expected_element_type: DataType, + #[case] expected_element_n_buffers: i64, + ) -> VortexResult<()> { + use vortex::array::arrays::ChunkedArray; + + let mut ctx = cuda_ctx_with_varbin_layout(layout)?; + + let elements = VarBinViewArray::from_iter_str([ + "hello", + "world", + "this is a longer string for out-of-line storage", + ]) + .into_array(); + let chunk = ListArray::try_new( + elements, + PrimitiveArray::from_iter([0i32, 2, 3]).into_array(), + Validity::NonNullable, + )? + .into_array(); + let dtype = chunk.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk], dtype)?.into_array(); + + let mut exported = chunked.export_device_array_with_schema(&mut ctx).await?; + + let field = Field::try_from(&exported.schema)?; + let DataType::List(element_field) = field.data_type() else { + vortex_bail!("expected List schema, got {:?}", field.data_type()); + }; + assert_eq!(element_field.data_type(), &expected_element_type); + + assert_eq!(exported.array.array.n_children, 1); + let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; + let element_array = unsafe { &*children[0] }; + assert_eq!(element_array.n_buffers, expected_element_n_buffers); + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + #[crate::test] async fn test_export_host_contiguous_list_view() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 60fd2622cc5..071e22085c4 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -815,19 +815,27 @@ fn arrow_device_export_field( .arrow() .to_arrow_field(name.as_ref(), dtype)?; - let data_type = match (ctx.cuda_session().varbin_export_layout(), dtype) { - (VarBinExportLayout::VarBin, DType::Utf8(_)) => DataType::Utf8, - (VarBinExportLayout::VarBin, DType::Binary(_)) => DataType::Binary, - (VarBinExportLayout::VarBinView, DType::Utf8(_)) => DataType::Utf8View, - (VarBinExportLayout::VarBinView, DType::Binary(_)) => DataType::BinaryView, - (_, DType::Decimal(decimal_dtype, _)) => { - arrow_device_export_decimal_data_type(*decimal_dtype) - } - (_, DType::Struct(struct_dtype, _)) => { - DataType::Struct(arrow_device_export_struct_fields(struct_dtype, ctx)?.into()) - } - _ => return Ok(field), - }; + let data_type = + match (ctx.cuda_session().varbin_export_layout(), dtype) { + (VarBinExportLayout::VarBin, DType::Utf8(_)) => DataType::Utf8, + (VarBinExportLayout::VarBin, DType::Binary(_)) => DataType::Binary, + (VarBinExportLayout::VarBinView, DType::Utf8(_)) => DataType::Utf8View, + (VarBinExportLayout::VarBinView, DType::Binary(_)) => DataType::BinaryView, + (_, DType::Decimal(decimal_dtype, _)) => { + arrow_device_export_decimal_data_type(*decimal_dtype) + } + (_, DType::Struct(struct_dtype, _)) => { + DataType::Struct(arrow_device_export_struct_fields(struct_dtype, ctx)?.into()) + } + // List elements are exported through the same layout-dependent paths as top-level + // arrays, so their fields must be rebuilt recursively. Without this, an element such + // as Utf8 would keep `to_arrow_field`'s Utf8View mapping while the data is exported + // with the offset-based layout. + (_, DType::List(element_dtype, _)) => DataType::List(Arc::new( + arrow_device_export_field(Field::LIST_FIELD_DEFAULT_NAME, element_dtype, ctx)?, + )), + _ => return Ok(field), + }; Ok( Field::new(field.name().clone(), data_type, field.is_nullable()) From 2fa8f3a3ebbfdf5a15003344104d58191ff1a360 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 13:32:36 +0000 Subject: [PATCH 03/13] fix(cuda): fall back to view layout for FSST exports over i32 offset range Under the offset-based default layout, an FSST array whose decoded size exceeds i32::MAX failed to export, while the previous Utf8View default handled such heaps via the multi-buffer host rollover. Route oversized arrays to the view layout instead: the schema derivation and the data export share fsst_varbin_offsets_fit, so a given array always resolves to the same layout. Verified end-to-end with a fabricated 2.1 GiB decode on a GH200: the export produces a Utf8View schema with views, two rolled-over data buffers, and variadic sizes. Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 165 ++++++++++++++++------- vortex-cuda/src/arrow/mod.rs | 17 +++ vortex-cuda/src/kernel/encodings/fsst.rs | 55 ++++++++ vortex-cuda/src/kernel/encodings/mod.rs | 1 + 4 files changed, 186 insertions(+), 52 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index be09fb0452b..d788497a658 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -83,6 +83,7 @@ use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; use crate::kernel::FSSTVarBin; use crate::kernel::decode_fsst_varbin; +use crate::kernel::fsst_varbin_offsets_fit; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -227,7 +228,17 @@ fn export_array( }; let array = match array.try_downcast::() { Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { - return export_fsst_varbin(fsst, ctx).await; + if fsst_varbin_offsets_fit(&fsst, ctx.execution_ctx())? { + return export_fsst_varbin(fsst, ctx).await; + } + // The decoded heap exceeds Arrow's i32 offset range, so decode to views and + // export the view layout, matching the per-array schema fallback in + // `arrow_device_export_field_for_array`. + let Canonical::VarBinView(varbinview) = fsst.into_array().execute_cuda(ctx).await? + else { + vortex_bail!("FSST decode must produce a VarBinView canonical"); + }; + return export_varbinview(varbinview, ctx).await; } Ok(fsst) => fsst.into_array(), Err(array) => array, @@ -318,57 +329,7 @@ fn export_canonical( if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin { return export_varbin(varbinview, ctx).await; } - - let len = varbinview.len(); - let VarBinViewDataParts { - views, - buffers: data_buffers, - validity, - .. - } = varbinview.into_data_parts(); - - let (validity_buffer, null_count) = - export_arrow_validity_buffer(validity, len, 0, ctx).await?; - - let views = ctx.ensure_on_device(views).await?; - let mut buffers = Vec::with_capacity(data_buffers.len() + 3); - buffers.push(validity_buffer); - buffers.push(Some(views)); - for buffer in data_buffers.iter() { - buffers.push(Some(ctx.ensure_on_device(buffer.clone()).await?)); - } - // Nanoarrow's Utf8View/BinaryView C layout stores the variadic data buffer sizes - // as the final buffer slot, after the null bitmap, views, and data buffers. - let variadic_buffer_sizes = data_buffers - .iter() - .map(|buffer| i64::try_from(buffer.len())) - .collect::, _>>()?; - buffers.push(Some( - ctx.ensure_on_device(BufferHandle::new_host( - Buffer::from(variadic_buffer_sizes).into_byte_buffer(), - )) - .await?, - )); - - let n_buffers = i64::try_from(buffers.len())?; - let mut private_data = PrivateData::new(buffers, vec![], ctx)?; - let sync_event = private_data.sync_event(); - let arrow_array = ArrowArray { - length: len as i64, - null_count, - offset: 0, - // Arrow Utf8View/BinaryView layout: optional null bitmap, views, data buffers, - // and trailing variadic buffer sizes. - n_buffers, - buffers: private_data.buffer_ptrs.as_mut_ptr(), - n_children: 0, - children: ptr::null_mut(), - release: Some(release_array), - dictionary: ptr::null_mut(), - private_data: Box::into_raw(private_data).cast(), - }; - - Ok((arrow_array, sync_event)) + export_varbinview(varbinview, ctx).await } c => vortex_bail!("unsupported Arrow Device export for {} array", c.dtype()), } @@ -553,6 +514,62 @@ where Ok(BufferHandle::new_device(Arc::new(output_device))) } +/// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. +async fn export_varbinview( + varbinview: VarBinViewArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let len = varbinview.len(); + let VarBinViewDataParts { + views, + buffers: data_buffers, + validity, + .. + } = varbinview.into_data_parts(); + + let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; + + let views = ctx.ensure_on_device(views).await?; + let mut buffers = Vec::with_capacity(data_buffers.len() + 3); + buffers.push(validity_buffer); + buffers.push(Some(views)); + for buffer in data_buffers.iter() { + buffers.push(Some(ctx.ensure_on_device(buffer.clone()).await?)); + } + // Nanoarrow's Utf8View/BinaryView C layout stores the variadic data buffer sizes + // as the final buffer slot, after the null bitmap, views, and data buffers. + let variadic_buffer_sizes = data_buffers + .iter() + .map(|buffer| i64::try_from(buffer.len())) + .collect::, _>>()?; + buffers.push(Some( + ctx.ensure_on_device(BufferHandle::new_host( + Buffer::from(variadic_buffer_sizes).into_byte_buffer(), + )) + .await?, + )); + + let n_buffers = i64::try_from(buffers.len())?; + let mut private_data = PrivateData::new(buffers, vec![], ctx)?; + let sync_event = private_data.sync_event(); + let arrow_array = ArrowArray { + length: len as i64, + null_count, + offset: 0, + // Arrow Utf8View/BinaryView layout: optional null bitmap, views, data buffers, + // and trailing variadic buffer sizes. + n_buffers, + buffers: private_data.buffer_ptrs.as_mut_ptr(), + n_children: 0, + children: ptr::null_mut(), + release: Some(release_array), + dictionary: ptr::null_mut(), + private_data: Box::into_raw(private_data).cast(), + }; + + Ok((arrow_array, sync_event)) +} + /// Export Vortex binary views as an Arrow Device array with standard `Utf8`/`Binary` layout. async fn export_varbin( varbinview: VarBinViewArray, @@ -1442,6 +1459,7 @@ mod tests { use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::TemporalArray; + use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveArrayExt; use vortex::array::arrays::varbinview::BinaryView; @@ -1459,6 +1477,10 @@ mod tests { use vortex::dtype::PType; use vortex::dtype::half::f16; use vortex::dtype::i256; + use vortex::encodings::fsst::FSST; + use vortex::encodings::fsst::FSSTArrayExt; + use vortex::encodings::fsst::fsst_compress; + use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -1471,6 +1493,7 @@ mod tests { use crate::arrow::ArrowDeviceArray; use crate::arrow::DeviceArrayExt; use crate::arrow::PrivateData; + use crate::arrow::arrow_schema_for_array; use crate::arrow::canonical::export_arrow_validity_buffer; use crate::arrow::canonical::repack_arrow_validity_buffer; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; @@ -2446,6 +2469,44 @@ mod tests { Ok(()) } + // An FSST array whose decoded size exceeds Arrow's i32 offset range cannot use the + // offset-based layout, so its export schema must fall back to the view layout that + // `export_array` produces for it. + #[crate::test] + async fn test_oversized_fsst_schema_falls_back_to_view_layout() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + + let varbin = VarBinArray::from_iter( + [Some(&b"short"[..]), Some(&b"another value"[..])], + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&varbin, ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, ctx.execution_ctx())?; + let schema = arrow_schema_for_array(&fsst.clone().into_array(), &mut ctx)?; + assert_eq!( + Field::try_from(&schema)?, + Field::new("", DataType::Utf8, false) + ); + + // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. + let oversized = FSST::try_new( + DType::Utf8(Nullability::NonNullable), + fsst.symbols().clone(), + fsst.symbol_lengths().clone(), + fsst.codes(), + PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), + ctx.execution_ctx(), + )? + .into_array(); + let schema = arrow_schema_for_array(&oversized, &mut ctx)?; + assert_eq!( + Field::try_from(&schema)?, + Field::new("", DataType::Utf8View, false) + ); + Ok(()) + } + #[crate::test] async fn test_export_host_contiguous_list_view() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 071e22085c4..a66b96a4e6b 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -48,6 +48,7 @@ use vortex::dtype::DecimalDType; use vortex::dtype::DecimalType; use vortex::dtype::PType; use vortex::dtype::StructFields; +use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; @@ -59,6 +60,7 @@ use vortex_arrow::ArrowSessionExt; use crate::CudaBufferExt; use crate::CudaExecutionCtx; use crate::VarBinExportLayout; +use crate::kernel::fsst_varbin_offsets_fit; mod arrow_c_abi { #![allow(dead_code)] @@ -786,6 +788,21 @@ fn arrow_device_export_field_for_array( return Ok(Field::new_list(name, element, array.dtype().is_nullable())); } + if let Ok(fsst) = array.clone().try_downcast::() { + // An FSST array whose decoded size exceeds Arrow's i32 offset range cannot use the + // offset-based layout; `export_array` decodes it to views instead, so the schema must + // also fall back to the view layout for this array. + if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin + && !fsst_varbin_offsets_fit(&fsst, ctx.execution_ctx())? + { + return ctx + .execution_ctx() + .session() + .arrow() + .to_arrow_field(name, array.dtype()); + } + } + arrow_device_export_field(name, &arrow_device_export_dtype(array.dtype()), ctx) } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index fe94f82967c..496dafec2f9 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -13,6 +13,7 @@ use cudarc::driver::PushKernelArg; use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; +use vortex::array::ExecutionCtx; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; @@ -114,6 +115,35 @@ impl CudaExecute for FSSTExecutor { } } +/// Returns whether the array's total decoded size fits Arrow's i32 offset range, i.e. whether +/// [`decode_fsst_varbin`] can represent it. Oversized arrays must use the view layout, whose +/// host rollover path splits the decoded heap across multiple data buffers. +/// +/// The schema and data export paths both call this, so a given array always resolves to the +/// same layout. +pub(crate) fn fsst_varbin_offsets_fit( + fsst: &FSSTArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lens = fsst + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + let total_size = match_each_integer_ptype!(lens.ptype(), |P| { + let mut acc = 0u64; + #[allow(clippy::unnecessary_cast)] + for &length in lens.as_slice::

() { + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("FSST uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("FSST decoded size overflow"))?; + } + VortexResult::Ok(acc) + })?; + Ok(i32::try_from(total_size).is_ok()) +} + /// Decode FSST directly into Arrow-compatible i32 offsets and contiguous values on device. pub(crate) async fn decode_fsst_varbin( fsst: FSSTArray, @@ -461,6 +491,31 @@ mod tests { Ok(()) } + #[crate::test] + async fn test_fsst_varbin_offsets_fit() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let varbin = VarBinArray::from_iter( + [Some(&b"short"[..]), Some(&b"another value"[..])], + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&varbin, &mut ctx)?; + let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?; + assert!(fsst_varbin_offsets_fit(&fsst, &mut ctx)?); + + // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. + let oversized = FSST::try_new( + DType::Utf8(Nullability::NonNullable), + fsst.symbols().clone(), + fsst.symbol_lengths().clone(), + fsst.codes(), + PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), + &mut ctx, + )?; + assert!(!fsst_varbin_offsets_fit(&oversized, &mut ctx)?); + Ok(()) + } + #[crate::test] async fn test_cuda_fsst_direct_varbin_output() -> VortexResult<()> { let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index 6ad37e7ead1..b8c65d184a6 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -23,6 +23,7 @@ pub(crate) use for_::FoRExecutor; pub(crate) use fsst::FSSTExecutor; pub(crate) use fsst::FSSTVarBin; pub(crate) use fsst::decode_fsst_varbin; +pub(crate) use fsst::fsst_varbin_offsets_fit; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; From e278c5d4679d1a0be65e8338eedb396fe78c73f5 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 13:35:32 +0000 Subject: [PATCH 04/13] chore(cuda): pair FSST varbin decode with the wide-store alignment assertion Assert the varbin output base is 16-byte aligned like the view decode path does, since both share Scratch::drain's alignment-gated stores. Also refresh the stale kernel comment about output offset widths and note that the offset-based FSST export bypasses CudaDispatchMode. Signed-off-by: Alexander Droste --- vortex-cuda/kernels/src/fsst.cu | 4 +++- vortex-cuda/src/arrow/canonical.rs | 2 ++ vortex-cuda/src/kernel/encodings/fsst.rs | 9 +++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index fb726eb26df..5cb07fa88eb 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -53,7 +53,9 @@ // the next add (≤ 8 bytes) within the 24-byte capacity. // // `codes_offsets` is templated over the four unsigned integer widths -// (u8/u16/u32/u64). `output_offsets` is uint64_t. +// (u8/u16/u32/u64). `output_offsets` is uint64_t for the view kernels +// (`fsst_*`, which also take an optional views pointer) and int32_t Arrow +// varbin offsets for the bytes-only varbin kernels (`fsst_varbin_*`). // 24-byte scratch buffer split across three u64 lanes. `cursor` is the // number of bytes currently buffered and the next-push offset. diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index d788497a658..39689b6fb80 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -226,6 +226,8 @@ fn export_array( Ok(list_view) => return export_list_view(list_view, ctx).await, Err(array) => array, }; + // The offset-based FSST export always uses the standalone varbin kernel; + // `CudaDispatchMode` only governs `execute_cuda`'s fused-vs-standalone planning. let array = match array.try_downcast::() { Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { if fsst_varbin_offsets_fit(&fsst, ctx.execution_ctx())? { diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 496dafec2f9..2b3afc9a682 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -241,7 +241,16 @@ where ctx.ensure_on_device(codes_offsets_buffer), )?; + // The kernel gates store widths on `out_pos % N` relative to the base, so the base must + // satisfy the widest store (u128 → 16). let output = ctx.device_alloc::(total_size)?; + let (output_base_ptr, _) = output.device_ptr(ctx.stream()); + assert_eq!( + output_base_ptr % 16, + 0, + "output base not 16-aligned: {output_base_ptr:#x}", + ); + let codes_bytes_view = codes_bytes.cuda_view::()?; let codes_offsets_view = codes_offsets.cuda_view::()?; let symbols_view = symbols.cuda_view::()?; From d3c4445a34fc4af2ca86ce6d7fb5e7e34a821db5 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 13:55:19 +0000 Subject: [PATCH 05/13] fix(buffer): shift unaligned bits in BitBuffer::sliced instead of copying bytes BitBuffer::sliced documents that it returns a buffer with offset reset to zero, but the non-byte-aligned branch used bitwise_unary_op_copy, which preserves the source offset. Consumers that hand the raw bytes to FFI or device code indexing from bit zero (the CUDA FSST decoders) read validity shifted by the slice offset, decoding null slots and skipping valid ones. Copy through BitBufferMut::append_buffer, which performs the unaligned shift into an offset-0 buffer. Signed-off-by: Alexander Droste --- vortex-buffer/src/bit/buf.rs | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index d229329fcd2..53f04bb7915 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -461,8 +461,13 @@ impl BitBuffer { ); } - // Allocate directly rather than clone + identity op which would fail try_into_mut. - bitwise_unary_op_copy(self, |a| a) + // An unaligned offset requires shifting every bit down to the buffer start; + // `append_buffer` performs the unaligned copy into an offset-0 buffer. + let mut target = BitBufferMut::with_capacity(self.len); + target.append_buffer(self); + let target = target.freeze(); + assert_eq!(target.offset(), 0); + target } } @@ -674,6 +679,25 @@ mod tests { use crate::bit::BitBuffer; use crate::buffer; + // `sliced` must shift the bits down when the offset is not byte-aligned, not just copy + // the backing bytes: consumers hand the raw bytes to FFI/device code that indexes from + // bit zero. + #[rstest] + #[case::unaligned(1..9)] + #[case::aligned(8..17)] + #[case::empty(3..3)] + fn test_sliced_resets_offset(#[case] range: std::ops::Range) { + let bits = BitBuffer::from_iter((0..20).map(|bit_index| bit_index % 3 == 0)); + let view = bits.slice(range.clone()); + + let sliced = view.sliced(); + assert_eq!(sliced.offset(), 0); + assert_eq!(sliced.len(), view.len()); + for (bit_index, absolute_index) in range.enumerate() { + assert_eq!(sliced.value(bit_index), absolute_index % 3 == 0); + } + } + #[test] fn test_bool() { // Create a new Buffer of length 1024 where the 8th bit is set. From 0897e72d5b354147fab94a35ded31da436f5cf69 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 13:55:28 +0000 Subject: [PATCH 06/13] test(cuda): cover FSST varbin export contents, validity, slices, and nesting Exercise the direct FSST-to-Arrow-varbin path with full content assertions (offsets, values, null bitmap) across partial-null, all-null, all-empty, empty, and sliced arrays, plus FSST nested as a struct field and as dictionary values. The sliced case caught the BitBuffer::sliced offset bug fixed in the previous commit: the decoded batch dropped 'delta' and decoded a null slot instead. Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 180 +++++++++++++++++++++++++++++ 1 file changed, 180 insertions(+) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 39689b6fb80..ecd59d715c6 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -1517,6 +1517,52 @@ mod tests { CudaSession::create_execution_ctx(&session) } + // Compress values into a concrete FSST array so exports take the direct FSST path. + fn fsst_array_from( + values: &[Option<&'static [u8]>], + dtype: DType, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let varbin = VarBinArray::from_iter(values.iter().copied(), dtype).into_array(); + let compressor = fsst_train_compressor(&varbin, ctx.execution_ctx())?; + Ok(fsst_compress(&varbin, &compressor, ctx.execution_ctx())?.into_array()) + } + + // Assert an exported varbin array against the logical values: offsets are the prefix sum + // of lengths (nulls contribute zero), values are the non-null bytes concatenated, and the + // Arrow null bitmap marks exactly the null slots. + fn assert_varbin_contents( + array: &ArrowArray, + values: &[Option<&'static [u8]>], + ) -> VortexResult<()> { + let mut expected_offsets = vec![0i32]; + let mut expected_values = Vec::new(); + for value in values { + if let Some(value) = value { + expected_values.extend_from_slice(value); + } + expected_offsets.push(i32::try_from(expected_values.len())?); + } + let null_count = values.iter().filter(|value| value.is_none()).count(); + + assert_binary_layout( + array, + i64::try_from(values.len())?, + i64::try_from(null_count)?, + &expected_offsets, + &expected_values, + )?; + + if null_count > 0 { + let bitmap = private_data_buffer_bytes(array, 0)?; + for (index, value) in values.iter().enumerate() { + let bit = (bitmap.as_ref()[index / 8] >> (index % 8)) & 1; + assert_eq!(bit == 1, value.is_some(), "validity bit {index}"); + } + } + Ok(()) + } + // Assert Arrow Device metadata that consumers use before reading buffers. fn assert_device_metadata( device_array: &ArrowDeviceArray, @@ -2509,6 +2555,140 @@ mod tests { Ok(()) } + // Content coverage for the direct FSST varbin export: offsets, values, and null bitmap, + // not just schema shape. + #[rstest] + #[case::utf8_inline_and_outlined( + vec![Some(&b""[..]), + Some(&b"short"[..]), + Some(&b"this value is stored out of line in the heap"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::partial_nulls( + vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], + DType::Utf8(Nullability::Nullable), + )] + #[case::all_nulls( + vec![None, None, None, None, None], + DType::Binary(Nullability::Nullable), + )] + #[case::all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), + )] + #[case::empty(vec![], DType::Utf8(Nullability::NonNullable))] + #[crate::test] + async fn test_export_fsst_varbin_contents( + #[case] values: Vec>, + #[case] dtype: DType, + ) -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let fsst = fsst_array_from(&values, dtype.clone(), &mut ctx)?; + + let mut exported = fsst.export_device_array_with_schema(&mut ctx).await?; + let expected_data_type = if matches!(dtype, DType::Utf8(_)) { + DataType::Utf8 + } else { + DataType::Binary + }; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, dtype.is_nullable()) + ); + assert_varbin_contents(&exported.array.array, &values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // A sliced FSST array keeps its encoding, so the direct varbin export must respect the + // slice's codes offsets and validity. + #[crate::test] + async fn test_export_sliced_fsst_varbin() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"alpha"[..]), + Some(&b"this value is stored out of line in the heap"[..]), + None, + Some(&b"delta"[..]), + Some(&b"echo"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::Nullable), &mut ctx)?; + let sliced = fsst.slice(1..4)?; + assert!(sliced.as_opt::().is_some()); + + let mut exported = sliced.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", DataType::Utf8, true) + ); + assert_varbin_contents(&exported.array.array, &values[1..4])?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // FSST fields inside a struct take the direct varbin export path through the child + // recursion, and the struct schema reflects the offset-based layout. + #[crate::test] + async fn test_export_struct_with_fsst_field() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"short"[..]), + Some(&b"this value is stored out of line in the heap"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::NonNullable), &mut ctx)?; + let array = StructArray::new( + FieldNames::from_iter(["s"]), + vec![fsst], + values.len(), + Validity::NonNullable, + ) + .into_array(); + + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Schema::try_from(&exported.schema)?, + Schema::new(vec![Field::new("s", DataType::Utf8, false)]) + ); + assert_eq!(exported.array.array.n_children, 1); + let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; + assert_varbin_contents(unsafe { &*children[0] }, values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // FSST dictionary values take the direct varbin export path, and the dictionary schema + // reflects the offset-based layout. + #[crate::test] + async fn test_export_dict_with_fsst_values() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"alpha"[..]), + Some(&b"this dictionary value is stored out of line"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::NonNullable), &mut ctx)?; + let array = DictArray::try_new(PrimitiveArray::from_iter([0u8, 1, 0]).into_array(), fsst)? + .into_array(); + + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new( + "", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + false, + ) + ); + assert!(!exported.array.array.dictionary.is_null()); + let dictionary = unsafe { &*exported.array.array.dictionary }; + assert_varbin_contents(dictionary, values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + #[crate::test] async fn test_export_host_contiguous_list_view() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) From 72bbbc0e311340983fd57b64a68e39c59dca02ff Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 14:00:18 +0000 Subject: [PATCH 07/13] chore(cuda): move ChunkedArray import to the test module top Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index ecd59d715c6..e13df086d4a 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -1452,6 +1452,7 @@ mod tests { use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::BoolArray; + use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::DictArray; use vortex::array::arrays::FixedSizeListArray; @@ -2481,8 +2482,6 @@ mod tests { #[case] expected_element_type: DataType, #[case] expected_element_n_buffers: i64, ) -> VortexResult<()> { - use vortex::array::arrays::ChunkedArray; - let mut ctx = cuda_ctx_with_varbin_layout(layout)?; let elements = VarBinViewArray::from_iter_str([ From 499579dcfdaf9686e01dcd3d2158ec78f111a3d2 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 14:14:38 +0000 Subject: [PATCH 08/13] style(cuda): clang-format FSST kernel Signed-off-by: Alexander Droste --- vortex-cuda/kernels/src/fsst.cu | 38 +++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 5cb07fa88eb..081c36772d3 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -176,10 +176,9 @@ __device__ inline void fsst_write_view(const FSSTArgs(args, sid); \ + fsst_decode_string(args, sid); \ } #define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \ @@ -256,14 +254,20 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ - codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \ - validity_bits, output_views, \ + codes_bytes, \ + codes_offsets, \ + symbols, \ + symbol_lengths, \ + output_bytes, \ + output_offsets, \ + validity_bits, \ + output_views, \ }; \ FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \ } #define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \ - extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \ + extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \ const CodeOffsetT *__restrict codes_offsets, \ const uint64_t *__restrict symbols, \ const uint8_t *__restrict symbol_lengths, \ @@ -272,8 +276,14 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ - codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \ - validity_bits, nullptr, \ + codes_bytes, \ + codes_offsets, \ + symbols, \ + symbol_lengths, \ + output_bytes, \ + output_offsets, \ + validity_bits, \ + nullptr, \ }; \ FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \ } From 2325fafc2e62dfe54f1d1822d39500c2826b2146 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 14:23:20 +0000 Subject: [PATCH 09/13] fix(buffer): make unaligned BitBuffer slicing Miri-safe Realign unaligned sliced buffers through padded logical words instead of bitvec unaligned copying, preserving the documented offset-zero result without violating Stacked Borrows. Signed-off-by: Alexander Droste --- vortex-buffer/src/bit/buf.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 53f04bb7915..0a7eb0a2aae 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -461,13 +461,18 @@ impl BitBuffer { ); } - // An unaligned offset requires shifting every bit down to the buffer start; - // `append_buffer` performs the unaligned copy into an offset-0 buffer. - let mut target = BitBufferMut::with_capacity(self.len); - target.append_buffer(self); - let target = target.freeze(); - assert_eq!(target.offset(), 0); - target + // An unaligned offset requires shifting every bit down to the buffer start. + // `iter_padded` realigns the logical bits into offset-0 words without relying on + // bitvec's unaligned copy, which violates Miri's Stacked Borrows model. + let n_words = self.len.div_ceil(64); + let mut words = BufferMut::::with_capacity(n_words); + for word in self.chunks().iter_padded().take(n_words) { + words.push(word); + } + + let mut bytes = words.into_byte_buffer(); + bytes.truncate(self.len.div_ceil(8)); + Self::new(bytes.freeze(), self.len) } } From bab3cb61761a0dbccd87111ed23f7843e9a8159c Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 14:36:58 +0000 Subject: [PATCH 10/13] clean Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 38 +++++++--------- vortex-cuda/src/arrow/mod.rs | 17 -------- vortex-cuda/src/kernel/encodings/fsst.rs | 55 ------------------------ vortex-cuda/src/kernel/encodings/mod.rs | 1 - 4 files changed, 16 insertions(+), 95 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index e13df086d4a..66c50a0413c 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -83,7 +83,6 @@ use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; use crate::kernel::FSSTVarBin; use crate::kernel::decode_fsst_varbin; -use crate::kernel::fsst_varbin_offsets_fit; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -230,17 +229,7 @@ fn export_array( // `CudaDispatchMode` only governs `execute_cuda`'s fused-vs-standalone planning. let array = match array.try_downcast::() { Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { - if fsst_varbin_offsets_fit(&fsst, ctx.execution_ctx())? { - return export_fsst_varbin(fsst, ctx).await; - } - // The decoded heap exceeds Arrow's i32 offset range, so decode to views and - // export the view layout, matching the per-array schema fallback in - // `arrow_device_export_field_for_array`. - let Canonical::VarBinView(varbinview) = fsst.into_array().execute_cuda(ctx).await? - else { - vortex_bail!("FSST decode must produce a VarBinView canonical"); - }; - return export_varbinview(varbinview, ctx).await; + return export_fsst_varbin(fsst, ctx).await; } Ok(fsst) => fsst.into_array(), Err(array) => array, @@ -2516,11 +2505,10 @@ mod tests { Ok(()) } - // An FSST array whose decoded size exceeds Arrow's i32 offset range cannot use the - // offset-based layout, so its export schema must fall back to the view layout that - // `export_array` produces for it. + // Standard Arrow Utf8/Binary uses i32 offsets. Oversized FSST exports keep that stable schema + // and return a clear error rather than changing layout based on batch contents. #[crate::test] - async fn test_oversized_fsst_schema_falls_back_to_view_layout() -> VortexResult<()> { + async fn test_oversized_fsst_varbin_export_errors() -> VortexResult<()> { let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; let varbin = VarBinArray::from_iter( @@ -2530,11 +2518,6 @@ mod tests { .into_array(); let compressor = fsst_train_compressor(&varbin, ctx.execution_ctx())?; let fsst = fsst_compress(&varbin, &compressor, ctx.execution_ctx())?; - let schema = arrow_schema_for_array(&fsst.clone().into_array(), &mut ctx)?; - assert_eq!( - Field::try_from(&schema)?, - Field::new("", DataType::Utf8, false) - ); // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. let oversized = FSST::try_new( @@ -2549,7 +2532,18 @@ mod tests { let schema = arrow_schema_for_array(&oversized, &mut ctx)?; assert_eq!( Field::try_from(&schema)?, - Field::new("", DataType::Utf8View, false) + Field::new("", DataType::Utf8, false) + ); + + let error = oversized + .export_device_array(&mut ctx) + .await + .err() + .vortex_expect("oversized FSST varbin export must fail"); + assert!( + error + .to_string() + .contains("FSST decoded size exceeds Arrow i32 offset range") ); Ok(()) } diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index a66b96a4e6b..071e22085c4 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -48,7 +48,6 @@ use vortex::dtype::DecimalDType; use vortex::dtype::DecimalType; use vortex::dtype::PType; use vortex::dtype::StructFields; -use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; @@ -60,7 +59,6 @@ use vortex_arrow::ArrowSessionExt; use crate::CudaBufferExt; use crate::CudaExecutionCtx; use crate::VarBinExportLayout; -use crate::kernel::fsst_varbin_offsets_fit; mod arrow_c_abi { #![allow(dead_code)] @@ -788,21 +786,6 @@ fn arrow_device_export_field_for_array( return Ok(Field::new_list(name, element, array.dtype().is_nullable())); } - if let Ok(fsst) = array.clone().try_downcast::() { - // An FSST array whose decoded size exceeds Arrow's i32 offset range cannot use the - // offset-based layout; `export_array` decodes it to views instead, so the schema must - // also fall back to the view layout for this array. - if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin - && !fsst_varbin_offsets_fit(&fsst, ctx.execution_ctx())? - { - return ctx - .execution_ctx() - .session() - .arrow() - .to_arrow_field(name, array.dtype()); - } - } - arrow_device_export_field(name, &arrow_device_export_dtype(array.dtype()), ctx) } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 2b3afc9a682..a74c3df5354 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -13,7 +13,6 @@ use cudarc::driver::PushKernelArg; use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; -use vortex::array::ExecutionCtx; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; @@ -115,35 +114,6 @@ impl CudaExecute for FSSTExecutor { } } -/// Returns whether the array's total decoded size fits Arrow's i32 offset range, i.e. whether -/// [`decode_fsst_varbin`] can represent it. Oversized arrays must use the view layout, whose -/// host rollover path splits the decoded heap across multiple data buffers. -/// -/// The schema and data export paths both call this, so a given array always resolves to the -/// same layout. -pub(crate) fn fsst_varbin_offsets_fit( - fsst: &FSSTArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let lens = fsst - .uncompressed_lengths() - .clone() - .execute::(ctx)?; - let total_size = match_each_integer_ptype!(lens.ptype(), |P| { - let mut acc = 0u64; - #[allow(clippy::unnecessary_cast)] - for &length in lens.as_slice::

() { - let length = u64::try_from(length as i128) - .map_err(|_| vortex_err!("FSST uncompressed length cannot be negative"))?; - acc = acc - .checked_add(length) - .ok_or_else(|| vortex_err!("FSST decoded size overflow"))?; - } - VortexResult::Ok(acc) - })?; - Ok(i32::try_from(total_size).is_ok()) -} - /// Decode FSST directly into Arrow-compatible i32 offsets and contiguous values on device. pub(crate) async fn decode_fsst_varbin( fsst: FSSTArray, @@ -500,31 +470,6 @@ mod tests { Ok(()) } - #[crate::test] - async fn test_fsst_varbin_offsets_fit() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); - let varbin = VarBinArray::from_iter( - [Some(&b"short"[..]), Some(&b"another value"[..])], - DType::Utf8(Nullability::NonNullable), - ) - .into_array(); - let compressor = fsst_train_compressor(&varbin, &mut ctx)?; - let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?; - assert!(fsst_varbin_offsets_fit(&fsst, &mut ctx)?); - - // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. - let oversized = FSST::try_new( - DType::Utf8(Nullability::NonNullable), - fsst.symbols().clone(), - fsst.symbol_lengths().clone(), - fsst.codes(), - PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), - &mut ctx, - )?; - assert!(!fsst_varbin_offsets_fit(&oversized, &mut ctx)?); - Ok(()) - } - #[crate::test] async fn test_cuda_fsst_direct_varbin_output() -> VortexResult<()> { let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index b8c65d184a6..6ad37e7ead1 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -23,7 +23,6 @@ pub(crate) use for_::FoRExecutor; pub(crate) use fsst::FSSTExecutor; pub(crate) use fsst::FSSTVarBin; pub(crate) use fsst::decode_fsst_varbin; -pub(crate) use fsst::fsst_varbin_offsets_fit; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; From 80266a2c930556ace5fb2a96363c97e9939100c7 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 14:39:46 +0000 Subject: [PATCH 11/13] test(buffer): cover word-boundary-crossing slice in BitBuffer::sliced Signed-off-by: Alexander Droste --- vortex-buffer/src/bit/buf.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 0a7eb0a2aae..51ea8472f62 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -691,8 +691,9 @@ mod tests { #[case::unaligned(1..9)] #[case::aligned(8..17)] #[case::empty(3..3)] + #[case::crosses_word_boundary(60..70)] fn test_sliced_resets_offset(#[case] range: std::ops::Range) { - let bits = BitBuffer::from_iter((0..20).map(|bit_index| bit_index % 3 == 0)); + let bits = BitBuffer::from_iter((0..80).map(|bit_index| bit_index % 3 == 0)); let view = bits.slice(range.clone()); let sliced = view.sliced(); From 35b029581b60f39e1de9da263735396fbe29930c Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 15:40:28 +0000 Subject: [PATCH 12/13] fix(cuda): pass FSST validity bit offset Signed-off-by: Alexander Droste --- vortex-buffer/src/bit/buf.rs | 34 +----------- vortex-cuda/kernels/src/fsst.cu | 9 +++- vortex-cuda/src/kernel/encodings/fsst.rs | 66 +++++++++++++++++++----- 3 files changed, 62 insertions(+), 47 deletions(-) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 51ea8472f62..d229329fcd2 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -461,18 +461,8 @@ impl BitBuffer { ); } - // An unaligned offset requires shifting every bit down to the buffer start. - // `iter_padded` realigns the logical bits into offset-0 words without relying on - // bitvec's unaligned copy, which violates Miri's Stacked Borrows model. - let n_words = self.len.div_ceil(64); - let mut words = BufferMut::::with_capacity(n_words); - for word in self.chunks().iter_padded().take(n_words) { - words.push(word); - } - - let mut bytes = words.into_byte_buffer(); - bytes.truncate(self.len.div_ceil(8)); - Self::new(bytes.freeze(), self.len) + // Allocate directly rather than clone + identity op which would fail try_into_mut. + bitwise_unary_op_copy(self, |a| a) } } @@ -684,26 +674,6 @@ mod tests { use crate::bit::BitBuffer; use crate::buffer; - // `sliced` must shift the bits down when the offset is not byte-aligned, not just copy - // the backing bytes: consumers hand the raw bytes to FFI/device code that indexes from - // bit zero. - #[rstest] - #[case::unaligned(1..9)] - #[case::aligned(8..17)] - #[case::empty(3..3)] - #[case::crosses_word_boundary(60..70)] - fn test_sliced_resets_offset(#[case] range: std::ops::Range) { - let bits = BitBuffer::from_iter((0..80).map(|bit_index| bit_index % 3 == 0)); - let view = bits.slice(range.clone()); - - let sliced = view.sliced(); - assert_eq!(sliced.offset(), 0); - assert_eq!(sliced.len(), view.len()); - for (bit_index, absolute_index) in range.enumerate() { - assert_eq!(sliced.value(bit_index), absolute_index % 3 == 0); - } - } - #[test] fn test_bool() { // Create a new Buffer of length 1024 where the 8th bit is set. diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 081c36772d3..22e8790a479 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -148,6 +148,8 @@ struct FSSTArgs { const OutputOffsetT *__restrict output_offsets; // Validity of each string. const uint8_t *__restrict validity_bits; + // Bit offset of string zero within the first validity byte (0..7). + uint64_t validity_bit_offset; // Optional output views, one 16-byte uint4 per string. A null pointer // requests bytes-only decoding for heaps that need the host rollover path. uint4 *__restrict output_views; @@ -184,7 +186,8 @@ __device__ inline void fsst_write_view(const FSSTArgs __device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { - if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) { + const uint64_t validity_index = sid + args.validity_bit_offset; + if (((args.validity_bits[validity_index >> 3] >> (validity_index & 7u)) & 1u) == 0u) { if (args.output_views != nullptr) { args.output_views[sid] = make_uint4(0, 0, 0, 0); } @@ -250,6 +253,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ @@ -283,6 +289,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs VortexResult<(u64, ByteBuffer)> { + let bits = validity + .clone() + .execute_mask(len, ctx.execution_ctx())? + .into_bit_buffer(); + let (bit_offset, bit_len, bytes) = bits.into_inner(); + debug_assert_eq!(bit_len, len); + let byte_len = (bit_offset + bit_len).div_ceil(8); + Ok((bit_offset as u64, bytes.slice(0..byte_len))) +} + /// CUDA decoder for FSST. #[derive(Debug)] pub(crate) struct FSSTExecutor; @@ -195,18 +215,13 @@ where buffer: codes_offsets_buffer, .. } = codes_offsets.into_data_parts(); - let (.., validity_bits) = validity - .clone() - .execute_mask(len, ctx.execution_ctx())? - .into_bit_buffer() - .sliced() - .into_inner(); + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx)?; let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(symbols_u64)?, ctx.copy_to_device(symbol_lengths)?, ctx.copy_to_device(output_offsets)?, - ctx.copy_to_device(validity_bits.to_vec())?, + ctx.copy_to_device(validity_bits)?, ctx.ensure_on_device(codes_bytes_handle), ctx.ensure_on_device(codes_offsets_buffer), )?; @@ -237,6 +252,7 @@ where .arg(&symbol_lengths_view) .arg(&output_offsets_view) .arg(&validity_view) + .arg(&validity_bit_offset) .arg(&output) .arg(&len_u64); })?; @@ -286,18 +302,13 @@ where .. } = codes_offsets.into_data_parts(); - let (.., validity_bits) = validity - .clone() - .execute_mask(num_strings, ctx.execution_ctx())? - .into_bit_buffer() - .sliced() - .into_inner(); + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx)?; let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(symbols_u64)?, ctx.copy_to_device(symbol_lengths)?, ctx.copy_to_device(output_offsets)?, - ctx.copy_to_device(validity_bits.to_vec())?, + ctx.copy_to_device(validity_bits)?, ctx.ensure_on_device(codes_bytes_handle), ctx.ensure_on_device(codes_offsets_buffer), )?; @@ -331,6 +342,7 @@ where .arg(&symbol_lengths_view) .arg(&output_offsets_view) .arg(&validity_view) + .arg(&validity_bit_offset) .arg(&device_output); if let Some(device_views) = device_views.as_ref() { args.arg(device_views); @@ -470,6 +482,32 @@ mod tests { Ok(()) } + /// Verifies that the view kernel applies a sliced validity bitmap's nonzero bit offset. + #[crate::test] + async fn test_cuda_fsst_decompression_sliced_validity() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = [ + Some(&b"before"[..]), + None, + Some(&b"gamma"[..]), + None, + Some(&b"after"[..]), + ]; + let varbin = + VarBinArray::from_iter(values, DType::Utf8(Nullability::Nullable)).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); + let sliced = fsst.slice(1..4)?; + + let gpu_result = FSSTExecutor.execute(sliced.clone(), &mut cuda_ctx).await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + #[crate::test] async fn test_cuda_fsst_direct_varbin_output() -> VortexResult<()> { let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; From 2cc117e73d5dfcf8628f42b905bbcbf22dec3560 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 16 Jul 2026 15:41:52 +0000 Subject: [PATCH 13/13] docs(cuda): explain FSST view fast path Signed-off-by: Alexander Droste --- vortex-cuda/src/kernel/encodings/fsst.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index c7bc5037ce9..de46d0b8fa1 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -352,6 +352,9 @@ where args.arg(&num_strings_u64); })?; + // Fast path: the decoded heap fits in one BinaryView backing buffer, so the kernel wrote + // views directly and both buffers can remain on-device. Larger heaps use the host rollover + // path below to split decoded bytes across multiple backing buffers. if let Some(device_views) = device_views { let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); let bytes = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_output)));