fix(cuda): preserve dictionary value nulls (#8774)

Export logical dictionary validity through Arrow Device arrays. Add a
packed-boolean dictionary gather so device-resident validity remains on
the GPU.

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

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<bool>`; 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 <typename IndexT>
+__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<uint64_t>(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<uint64_t>(codes[row]);
+                const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1;
+                packed |= static_cast<uint8_t>(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<IndexType>(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)

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]

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<bool>`.
+            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<bool>`.
+/// 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<Canonical> {
+    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::<I>(values, codes, ctx).await
+    })
+}
+
+async fn execute_dict_bool_typed<I: DeviceRepr + NativePType>(
+    values: BoolArray,
+    codes: PrimitiveArray,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<Canonical> {
+    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::<u8>(output_bytes)?;
+    let output_device = CudaDeviceBuffer::new(output_slice);
+
+    let values_view = values_device.cuda_view::<u8>()?;
+    let codes_view = codes_device.cuda_view::<I>()?;
+    let output_view = output_device.as_view::<u8>();
+    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<V: DeviceRepr + NativePType, I: DeviceRepr + NativePType>(
     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();