GPU: fuse dictionaries with nullable codes

vortex-cuda/kernels/src/dynamic_dispatch.cu

@@ -170,7 +170,8 @@ __device__ inline void scalar_op(T *values, const struct ScalarOp &op, char *__r
         const T *dict = reinterpret_cast<const T *>(smem + op.params.dict.values_smem_byte_offset);
 #pragma unroll
         for (uint32_t i = 0; i < N; ++i) {
-            values[i] = dict[static_cast<uint32_t>(values[i])];
+            const uint32_t code = static_cast<uint32_t>(values[i]);
+            values[i] = code < op.params.dict.values_len ? dict[code] : T {};
         }
         break;
     }

vortex-cuda/kernels/src/dynamic_dispatch.h

@@ -197,6 +197,7 @@ union ScalarParams {
     /// element type.
     struct DictParams {
         uint32_t values_smem_byte_offset; // byte offset to decoded dict values in smem
+        uint32_t values_len;
     } dict;
 };
 

vortex-cuda/src/dynamic_dispatch/mod.rs

@@ -415,13 +415,14 @@ impl ScalarOp {
 
     /// Dictionary gather: use current value as index into decoded values
     /// in shared memory (populated by an earlier input stage).
-    pub fn dict(values_smem_byte_offset: u32, output_ptype: PTypeTag) -> Self {
+    pub fn dict(values_smem_byte_offset: u32, values_len: u32, output_ptype: PTypeTag) -> Self {
         Self {
             op_code: ScalarOp_ScalarOpCode_DICT,
             output_ptype,
             params: ScalarParams {
                 dict: ScalarParams_DictParams {
                     values_smem_byte_offset,
+                    values_len,
                 },
             },
         }
@@ -691,7 +692,7 @@ mod tests {
                     SourceOp::bitunpack(6, 0),
                     &[
                         ScalarOp::frame_of_ref(42, PTypeTag_PTYPE_U32),
-                        ScalarOp::dict(0, PTypeTag_PTYPE_U32),
+                        ScalarOp::dict(0, 256, PTypeTag_PTYPE_U32),
                     ],
                 ),
             ],
@@ -2519,20 +2520,38 @@ mod tests {
         Ok(())
     }
 
-    /// Dict with nullable codes must fall back to Unfused (not fused).
+    /// Dict with nullable codes fuses. Null positions may contain arbitrary
+    /// physical codes, so exercise out-of-range values as well.
     #[crate::test]
-    fn test_dict_nullable_codes_rejected() -> VortexResult<()> {
+    async fn test_dict_nullable_codes_fuses() -> VortexResult<()> {
         use vortex::buffer::buffer;
 
-        let codes = PrimitiveArray::from_option_iter([Some(0u32), None, Some(1), None, Some(2)]);
-        let values = PrimitiveArray::new(buffer![10u32, 20, 30], NonNullable);
+        let mut cpu_ctx = array_session().create_execution_ctx();
+        let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?;
+
+        let codes = PrimitiveArray::new(
+            buffer![0u32, u32::MAX, 1, 99, 2],
+            Validity::from_iter([true, false, true, false, true]),
+        );
+        let values = PrimitiveArray::from_option_iter([Some(10u32), None, Some(30)]);
         let dict = DictArray::try_new(codes.into_array(), values.into_array())?;
 
-        let plan = DispatchPlan::new(&dict.into_array(), CudaDispatchMode::Auto)?;
+        let plan = DispatchPlan::new(&dict.clone().into_array(), CudaDispatchMode::Auto)?;
         assert!(
-            matches!(plan, DispatchPlan::Unfused),
-            "Dict with nullable codes should fall back to Unfused"
+            matches!(plan, DispatchPlan::Fused(..)),
+            "Dict with nullable codes should fuse"
         );
+
+        let gpu = dict
+            .clone()
+            .into_array()
+            .execute_cuda(&mut cuda_ctx)
+            .await?
+            .into_host()
+            .await?
+            .into_array();
+
+        assert_arrays_eq!(dict, gpu, &mut cpu_ctx);
         Ok(())
     }
 

vortex-cuda/src/dynamic_dispatch/plan_builder.rs

@@ -93,11 +93,6 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool {
     }
     if id == Dict.id() {
         let arr = array.as_::<Dict>();
-        // Nullable codes could hold garbage values at null positions, causing
-        // out-of-bounds shared memory reads in the DICT gather scalar op.
-        if arr.codes().dtype().is_nullable() {
-            return false;
-        }
         // Dict codes and values may have different byte widths.
         // The kernel handles mixed widths via widening input stages,
         // but only when codes are no wider than values (the output type).
@@ -303,9 +298,9 @@ impl DispatchPlan {
     /// - **F16 primitives** are not supported (no reinterpret path in the kernel).
     /// - **ALP** is supported for f32 and f64 only (including patches).
     /// - **BitPacked** with patches is supported.
-    /// - **Dict** with nullable codes is rejected (garbage at null positions
-    ///   could OOB the DICT gather). Dict with codes wider than values is
-    ///   also rejected (load would truncate code indices).
+    /// - **Dict** with codes wider than values is rejected (load would truncate
+    ///   code indices). Nullable codes are safe because the gather bounds-checks
+    ///   their unspecified physical values and validity is propagated separately.
     /// - **RunEnd** with nullable ends is rejected (garbage values break the
     ///   binary search). RunEnd with ends wider than values is also rejected.
     /// - Validity is propagated from the root array to the output.
@@ -757,8 +752,6 @@ impl FusedPlan {
     /// Cases that require a separate kernel dispatch:
     ///
     /// - **F16 primitives** — no reinterpret path in the kernel.
-    /// - **Dict with nullable codes** — garbage at null positions could OOB
-    ///   the DICT gather in shared memory.
     /// - **Dict with codes wider than values** — `load_element<T>()` would
     ///   truncate the code indices.
     /// - **RunEnd with nullable ends** — garbage values break the binary
@@ -806,7 +799,11 @@ impl FusedPlan {
         // DICT scalar op: pass byte offset directly (C ABI uses byte offsets).
         // output_ptype is the values' ptype — DICT transforms codes → values.
         pipeline.scalar_ops.push((
-            ScalarOp::dict(values_smem_byte_offset, ptype_to_tag(values_ptype)),
+            ScalarOp::dict(
+                values_smem_byte_offset,
+                values_len,
+                ptype_to_tag(values_ptype),
+            ),
             None,
         ));
         Ok(pipeline)