From 5602cd02f117e3bd869ecdbb9a9fefb1ec6a5dba Mon Sep 17 00:00:00 2001
From: Connor Tsui <connor.tsui20@gmail.com>
Date: Fri, 23 Oct 2026 15:57:58 -0400
Subject: [PATCH] fix(array): avoid truncating dictionary fill codes

The dictionary fill_null kernel converted the index of the fill value into the codes type with an unchecked cast. A dictionary can hold more values than its codes type can index, so an unreferenced fill value could truncate to the wrong code, or panic for signed codes. The kernel now checks the index against the range of the codes and declines when it does not fit, so the generic path fills the decoded array. Non-nullable codes need no fill code, so they keep the fast path.

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
---
 .../src/arrays/dict/compute/fill_null.rs      | 76 ++++++++++++++-----
 1 file changed, 56 insertions(+), 20 deletions(-)

diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs
index 077ee865bda..09ba3bba970 100644
--- a/vortex-array/src/arrays/dict/compute/fill_null.rs
+++ b/vortex-array/src/arrays/dict/compute/fill_null.rs
@@ -16,9 +16,7 @@ use crate::arrays::bool::BoolArrayExt;
 use crate::arrays::dict::DictArrayExt;
 use crate::arrays::dict::DictArraySlotsExt;
 use crate::builtins::ArrayBuiltins;
-use crate::match_each_integer_ptype;
 use crate::scalar::Scalar;
-use crate::scalar::ScalarValue;
 use crate::scalar_fn::fns::fill_null::FillNullKernel;
 use crate::scalar_fn::fns::operators::Operator;
 
@@ -54,26 +52,27 @@ impl FillNullKernel for Dict {
             ));
         };
 
-        // Now we rewrite the nullable codes to point at the fill value.
+        // A dictionary can hold more values than its codes type can index, so the index of an
+        // unreferenced fill value can exceed the range of the codes. Non-nullable codes need no
+        // fill code. For nullable codes, decline in that case, so that the generic path fills the
+        // decoded array. The first match has the smallest index, so no later match fits either.
         let codes = array.codes();
+        let codes = if codes.dtype().is_nullable() {
+            let codes_dtype = codes.dtype().as_nonnullable();
+            let fill_index = existing_fill_value_index as u64;
+            if fill_index > codes_dtype.as_ptype().max_value_as_u64() {
+                return Ok(None);
+            }
 
-        // Cast the index to the correct unsigned integer type matching the codes' ptype.
-        let codes_ptype = codes.dtype().as_ptype();
-
-        #[expect(
-            clippy::cast_possible_truncation,
-            reason = "The existing index must be representable by the existing ptype"
-        )]
-        let fill_scalar_value = match_each_integer_ptype!(codes_ptype, |P| {
-            ScalarValue::from(existing_fill_value_index as P)
-        });
-
-        // Fill nulls in both the codes and the values. Note that the precondition of this function
-        // states that the fill value is non-null, so we do not have to worry about the nullability.
-        let codes = codes.clone().fill_null(Scalar::try_new(
-            codes.dtype().as_nonnullable(),
-            Some(fill_scalar_value),
-        )?)?;
+            codes
+                .clone()
+                .fill_null(Scalar::from(fill_index).cast(&codes_dtype)?)?
+        } else {
+            codes.clone()
+        };
+
+        // Fill the nulls of the values. Note that the precondition of this function states that the
+        // fill value is non-null, so we do not have to worry about the nullability.
         let values = array.values().clone().fill_null(fill_value.clone())?;
 
         // SAFETY: invariants are still satisfied after patching nulls.
@@ -91,9 +90,11 @@ impl FillNullKernel for Dict {
 mod tests {
     use std::sync::LazyLock;
 
+    use rstest::rstest;
     use vortex_buffer::BitBuffer;
     use vortex_buffer::buffer;
     use vortex_error::VortexExpect;
+    use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
     use crate::IntoArray;
@@ -108,6 +109,41 @@ mod tests {
 
     static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
 
+    #[rstest]
+    #[case::u8_max(PrimitiveArray::from_option_iter([Some(0u8), None]), 255)]
+    #[case::u8_overflow(PrimitiveArray::from_option_iter([Some(0u8), None]), 256)]
+    #[case::u8_truncation(PrimitiveArray::from_option_iter([Some(0u8), None]), 300)]
+    #[case::i8_max(PrimitiveArray::from_option_iter([Some(0i8), None]), 127)]
+    #[case::i8_overflow(PrimitiveArray::from_option_iter([Some(0i8), None]), 128)]
+    fn fill_unreferenced_value(
+        #[case] codes: PrimitiveArray,
+        #[case] fill_value: i32,
+    ) -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let dict = DictArray::try_new(
+            codes.into_array(),
+            PrimitiveArray::from_iter(0..=fill_value).into_array(),
+        )?;
+
+        let filled = dict.into_array().fill_null(fill_value)?;
+        assert_arrays_eq!(filled, PrimitiveArray::from_iter([0, fill_value]), &mut ctx);
+
+        Ok(())
+    }
+
+    #[test]
+    fn non_nullable_codes_fill_unreferenced_value() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values =
+            PrimitiveArray::from_option_iter((0..=300).map(|value| (value != 1).then_some(value)));
+        let dict = DictArray::try_new(buffer![0u8, 1].into_array(), values.into_array())?;
+
+        let filled = dict.into_array().fill_null(300)?;
+        assert_arrays_eq!(filled, PrimitiveArray::from_iter([0, 300]), &mut ctx);
+
+        Ok(())
+    }
+
     #[test]
     fn nullable_codes_fill_in_values() {
         let mut ctx = SESSION.create_execution_ctx();