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..22e8790a479 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. @@ -123,13 +125,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 +145,57 @@ 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; + // 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; }; -template -__device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { - if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) { +// 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; } - OffT in_pos = args.codes_offsets[sid]; - const OffT in_end = args.codes_offsets[sid + 1]; + 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) { + 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); + } + return; + } + + 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,25 +226,38 @@ __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, \ + uint64_t validity_bit_offset, \ uint8_t *__restrict output_bytes, \ + uint4 *__restrict output_views, \ uint64_t num_strings) { \ - const FSSTArgs args = { \ + const FSSTArgs args = { \ codes_bytes, \ codes_offsets, \ symbols, \ @@ -207,20 +265,42 @@ __device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t s output_bytes, \ output_offsets, \ validity_bits, \ + validity_bit_offset, \ + 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, \ + uint64_t validity_bit_offset, \ + 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, \ + validity_bit_offset, \ + 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..66c50a0413c 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,15 @@ 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 => { + 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,60 +317,10 @@ 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(); - 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()), } @@ -541,8 +505,64 @@ 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 `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, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { @@ -558,16 +578,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, @@ -1393,6 +1441,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; @@ -1402,6 +1451,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; @@ -1419,6 +1469,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; @@ -1431,11 +1485,13 @@ 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; 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 +1501,58 @@ 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) + } + + // 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, @@ -1812,7 +1920,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 +1941,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 +2008,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 +2296,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(); @@ -2345,6 +2458,230 @@ 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<()> { + 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(()) + } + + // 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_varbin_export_errors() -> 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())?; + + // 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::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(()) + } + + // 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()) @@ -2561,27 +2898,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 +3329,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 +3744,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 +3757,7 @@ mod tests { Some("this is a longer string for out-of-line storage"), ]) .into_array(), - 4, + 3, 1, &mut ctx, ) @@ -3503,7 +3858,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 +3869,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 +3977,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 +4008,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 +4103,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..071e22085c4 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,14 +815,27 @@ 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, _) => { - 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()) diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index eef0db049b8..de46d0b8fa1 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -17,14 +17,16 @@ 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::buffer::ByteBuffer; +use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; @@ -38,6 +40,34 @@ 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, +} + +/// Returns validity backing bytes and the bit offset of the first row for the FSST kernels. +/// +/// `BitBuffer` slices normalize whole-byte offsets but can retain a sub-byte offset. Passing that +/// offset to CUDA lets the kernels address sliced validity directly without repacking it on host. +fn cuda_validity( + validity: &Validity, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> 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; @@ -61,16 +91,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 +134,138 @@ 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_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)?, + ctx.ensure_on_device(codes_bytes_handle), + 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::()?; + 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(&validity_bit_offset) + .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 +287,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(); @@ -134,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), )?; @@ -153,6 +316,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 +334,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 +342,29 @@ where .arg(&symbol_lengths_view) .arg(&output_offsets_view) .arg(&validity_view) - .arg(&device_output) - .arg(&num_strings_u64); + .arg(&validity_bit_offset) + .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); })?; + // 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))); + 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 +386,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 +402,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"[..])], - Nullability::NonNullable, + DType::Binary(Nullability::NonNullable), )] - #[case::partial_nulls( + #[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"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[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 +476,128 @@ 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(()) + } + + /// 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)?; + 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 +615,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