diff --git a/vortex-cuda/kernels/src/dict.cu b/vortex-cuda/kernels/src/dict.cu index 3422980c65e..a456b476821 100644 --- a/vortex-cuda/kernels/src/dict.cu +++ b/vortex-cuda/kernels/src/dict.cu @@ -19,12 +19,51 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, (block_start + elements_per_block < codes_len) ? (block_start + elements_per_block) : codes_len; for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { - IndexT code = codes[idx]; + const IndexT code = codes[idx]; output[idx] = values[code]; } } -// Macro to generate dict kernels for all value/index type combinations +// Nullable dictionary values require gathering validity through the row codes: +// +// row_validity[i] = values_validity[codes[i]] +// +// Vortex represents this gather lazily as `Dict`; null codes are carried by the resulting +// boolean array's own validity. Consequently, this kernel may materialize validity for a +// dictionary whose logical values are strings or another non-boolean type. +// +// Vortex booleans are bit-packed, so fixed-width `output[i] = values[codes[i]]` semantics do not +// apply. Assigning each complete output byte to one thread avoids races between individual bit +// writes. +template +__device__ void dict_bool_kernel(const IndexT *const __restrict codes, + uint64_t codes_len, + const uint8_t *const __restrict values, + uint64_t values_bit_offset, + uint8_t *const __restrict output) { + const uint64_t output_len = (codes_len + 7) / 8; + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = + (block_start + elements_per_block < output_len) ? (block_start + elements_per_block) : output_len; + + for (uint64_t output_idx = block_start + threadIdx.x; output_idx < block_end; output_idx += blockDim.x) { + const uint64_t row_start = output_idx * 8; + uint8_t packed = 0; +#pragma unroll + for (uint32_t bit = 0; bit < 8; ++bit) { + const uint64_t row = row_start + bit; + if (row < codes_len) { + const uint64_t value_idx = values_bit_offset + static_cast(codes[row]); + const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; + packed |= static_cast(value << bit); + } + } + output[output_idx] = packed; + } +} + +// Macro to generate dict kernels for all fixed-width value/index type combinations. #define GENERATE_DICT_KERNEL(value_suffix, ValueType, index_suffix, IndexType) \ extern "C" __global__ void dict_##value_suffix##_##index_suffix( \ const IndexType *const __restrict codes, \ @@ -41,5 +80,21 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, GENERATE_DICT_KERNEL(value_suffix, ValueType, u32, uint32_t) \ GENERATE_DICT_KERNEL(value_suffix, ValueType, u64, uint64_t) -// Generate for all native ptypes & decimal values +#define GENERATE_DICT_BOOL_KERNEL(index_suffix, IndexType) \ + extern "C" __global__ void dict_bool_##index_suffix(const IndexType *const __restrict codes, \ + uint64_t codes_len, \ + const uint8_t *const __restrict values, \ + uint64_t values_bit_offset, \ + uint8_t *const __restrict output) { \ + dict_bool_kernel(codes, codes_len, values, values_bit_offset, output); \ + } + +// Generate fixed-width kernels for all native ptypes and decimal values. FOR_EACH_NUMERIC(GENERATE_DICT_FOR_ALL_INDICES) + +// Boolean values use a different physical layout and launch unit, but still dispatch over the +// same unsigned dictionary-code types. +GENERATE_DICT_BOOL_KERNEL(u8, uint8_t) +GENERATE_DICT_BOOL_KERNEL(u16, uint16_t) +GENERATE_DICT_BOOL_KERNEL(u32, uint32_t) +GENERATE_DICT_BOOL_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index d5fa1c26682..af141bfe548 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -371,12 +371,14 @@ async fn export_dict( ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { let len = array.len(); + let validity = array.validity()?; let parts = array.into_parts(); - let PrimitiveDataParts { - buffer, validity, .. - } = export_dictionary_codes(parts.codes, ctx).await?; + let PrimitiveDataParts { buffer, .. } = export_dictionary_codes(parts.codes, ctx).await?; let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; let codes_buffer = ctx.ensure_on_device(buffer).await?; + // Arrow permits null dictionary values, so preserve the child's validity bitmap. The outer + // bitmap independently marks each row whose code selects a null dictionary value, ensuring + // consumers that require non-null dictionary keys do not lose the logical nulls. let (dictionary, _) = export_array(parts.values, ctx).await?; let mut private_data = PrivateData::new_with_dictionary( @@ -1837,6 +1839,35 @@ mod tests { Ok(()) } + #[crate::test] + async fn test_export_dictionary_propagates_value_nulls_to_codes() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let array = DictArray::try_new( + PrimitiveArray::from_iter([0u8, 1, 2, 1]).into_array(), + VarBinViewArray::from_iter_nullable_str([ + Some("alpha"), + None, + Some("a dictionary value stored out-of-line"), + ]) + .into_array(), + )? + .into_array(); + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + + assert_eq!(exported.array.array.null_count, 2); + assert_eq!( + private_data_buffer_bytes(&exported.array.array, 0)?.as_ref(), + &[0b0000_0101] + ); + let dictionary = unsafe { &*exported.array.array.dictionary }; + assert_eq!(dictionary.null_count, 1); + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + #[crate::test] async fn test_export_struct_preserves_dictionary_child() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -1900,7 +1931,11 @@ mod tests { true, ) ); - assert_eq!(exported.array.array.null_count, 0); + assert_eq!(exported.array.array.null_count, 1); + assert_eq!( + private_data_buffer_bytes(&exported.array.array, 0)?.as_ref(), + &[0b0000_0101] + ); assert_eq!( private_data_buffer_i16_values(&exported.array.array, 1)?, [0, 1, 0] @@ -2691,7 +2726,11 @@ mod tests { let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; let elements = unsafe { &*children[0] }; - assert_eq!(elements.null_count, 0); + assert_eq!(elements.null_count, 2); + assert_eq!( + private_data_buffer_bytes(elements, 0)?.as_ref(), + &[0b0001_0101] + ); assert_eq!( private_data_buffer_i16_values(elements, 1)?, [0, 1, 0, 1, 2] diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 394de699f49..840a3cdbbac 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -10,11 +10,13 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::Dict; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::decimal::DecimalDataParts; use vortex::array::arrays::dict::DictArraySlotsExt; use vortex::array::arrays::primitive::PrimitiveDataParts; @@ -29,6 +31,7 @@ use vortex::dtype::NativePType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -55,6 +58,8 @@ impl CudaExecute for DictExecutor { let values_dtype = dict_array.values().dtype().clone(); match &values_dtype { + // Nullable dictionary values expose their logical validity as a lazy `Dict`. + DType::Bool(..) => execute_dict_bool(dict_array, ctx).await, DType::Decimal(..) => execute_dict_decimal(dict_array, ctx).await, DType::Primitive(..) => execute_dict_prim(dict_array, ctx).await, DType::Utf8(..) | DType::Binary(..) => execute_dict_varbinview(dict_array, ctx).await, @@ -86,6 +91,80 @@ async fn execute_dict_prim(dict: DictArray, ctx: &mut CudaExecutionCtx) -> Vorte }) } +/// Gather bit-packed boolean values through dictionary codes. +/// +/// This path is especially important for dictionary validity. When dictionary values are +/// nullable, `Dict::validity()` represents `take(values_validity, codes)` lazily as a `Dict`. +/// Materializing that validity on the GPU therefore requires a boolean dictionary gather even +/// when the user-visible array contains strings or another non-boolean type. +async fn execute_dict_bool(dict: DictArray, ctx: &mut CudaExecutionCtx) -> VortexResult { + let values = dict.values().clone().execute_cuda(ctx).await?.into_bool(); + let codes = dict + .codes() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + let codes_ptype = codes.ptype(); + + match_each_integer_ptype!(codes_ptype, |I| { + execute_dict_bool_typed::(values, codes, ctx).await + }) +} + +async fn execute_dict_bool_typed( + values: BoolArray, + codes: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!(!codes.is_empty(), "cannot CUDA-decode an empty dictionary"); + let codes_len = codes.len(); + + let values_len = values.len(); + let values_validity = values.validity()?; + let BoolDataParts { + bits: values_buffer, + meta: values_meta, + } = values.into_data().into_parts(values_len); + let output_validity = values_validity.take(&codes.clone().into_array())?; + let PrimitiveDataParts { + buffer: codes_buffer, + .. + } = codes.into_data_parts(); + + let values_device = ctx.ensure_on_device(values_buffer).await?; + let codes_device = ctx.ensure_on_device(codes_buffer).await?; + + // Each CUDA thread owns complete output bytes, avoiding races between threads writing + // different bits in the same byte. The kernel handles the final partial byte explicitly. + let output_bytes = codes_len.div_ceil(8); + let output_slice = ctx.device_alloc::(output_bytes)?; + let output_device = CudaDeviceBuffer::new(output_slice); + + let values_view = values_device.cuda_view::()?; + let codes_view = codes_device.cuda_view::()?; + let output_view = output_device.as_view::(); + let codes_len_u64 = codes_len as u64; + let values_offset_u64 = values_meta.offset() as u64; + + let codes_ptype = I::PTYPE.to_string(); + let kernel_function = ctx.load_function_with_suffixes("dict", &["bool", &codes_ptype])?; + ctx.launch_kernel(&kernel_function, output_bytes, |args| { + args.arg(&codes_view) + .arg(&codes_len_u64) + .arg(&values_view) + .arg(&values_offset_u64) + .arg(&output_view); + })?; + + Ok(Canonical::Bool(BoolArray::new_handle( + BufferHandle::new_device(Arc::new(output_device)), + 0, + codes_len, + output_validity, + ))) +} + async fn execute_dict_prim_typed( values: PrimitiveArray, codes: PrimitiveArray, @@ -298,6 +377,7 @@ async fn execute_dict_varbinview( #[cfg(test)] mod tests { use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; @@ -323,6 +403,43 @@ mod tests { )) } + #[crate::test] + async fn test_cuda_dict_bool_gathers_packed_validity_bits() -> 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"); + + // Slicing leaves the dictionary values at a non-zero bit offset. Thirteen codes also + // exercise a final partial output byte. + let values = BoolArray::from_iter([ + false, true, false, true, false, true, true, false, true, false, + ]) + .into_array() + .slice(3..8)?; + let codes = PrimitiveArray::new( + Buffer::from(vec![0u8, 1, 2, 3, 4, 3, 2, 1, 0, 4, 1, 3, 0]), + NonNullable, + ); + let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?.into_array(); + + let codes_handle = cuda_ctx + .ensure_on_device(codes.buffer_handle().clone()) + .await?; + let device_codes = + PrimitiveArray::from_buffer_handle(codes_handle, codes.ptype(), codes.validity()?); + let dict = DictArray::try_new(device_codes.into_array(), values)?.into_array(); + + let actual = DictExecutor + .execute(dict, &mut cuda_ctx) + .await? + .into_host() + .await? + .into_bool(); + + assert_arrays_eq!(actual.into_array(), expected, &mut ctx); + Ok(()) + } + #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx();