feat[array]: dict array allows signed codes (#6336)

To support the current `take` compute function allowing signed indices
we allow the dict array to also store signed codes`.
This is a backwards compatible change.

We could add another array take that is only used for compute/?

---------

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-array/src/arrays/dict/arbitrary.rs

@@ -46,10 +46,26 @@ impl ArbitraryDictArray {
 
         // Choose a random PType at least as wide as the minimum
         let valid_ptypes: &[PType] = match min_codes_ptype {
-            PType::U8 => &[PType::U8, PType::U16, PType::U32, PType::U64],
-            PType::U16 => &[PType::U16, PType::U32, PType::U64],
-            PType::U32 => &[PType::U32, PType::U64],
-            PType::U64 => &[PType::U64],
+            PType::U8 => &[
+                PType::U8,
+                PType::U16,
+                PType::U32,
+                PType::U64,
+                PType::I8,
+                PType::I16,
+                PType::I32,
+                PType::I64,
+            ],
+            PType::U16 => &[
+                PType::U16,
+                PType::U32,
+                PType::U64,
+                PType::I16,
+                PType::I32,
+                PType::I64,
+            ],
+            PType::U32 => &[PType::U32, PType::U64, PType::I32, PType::I64],
+            PType::U64 => &[PType::U64, PType::I64],
             _ => unreachable!(),
         };
         let codes_ptype = *u.choose(valid_ptypes)?;
@@ -61,6 +77,10 @@ impl ArbitraryDictArray {
             PType::U16 => random_codes::<u16>(u, codes_len, values_len, codes_nullable)?,
             PType::U32 => random_codes::<u32>(u, codes_len, values_len, codes_nullable)?,
             PType::U64 => random_codes::<u64>(u, codes_len, values_len, codes_nullable)?,
+            PType::I8 => random_codes::<i8>(u, codes_len, values_len, codes_nullable)?,
+            PType::I16 => random_codes::<i16>(u, codes_len, values_len, codes_nullable)?,
+            PType::I32 => random_codes::<i32>(u, codes_len, values_len, codes_nullable)?,
+            PType::I64 => random_codes::<i64>(u, codes_len, values_len, codes_nullable)?,
             _ => unreachable!(),
         };
 
@@ -71,7 +91,7 @@ impl ArbitraryDictArray {
     }
 }
 
-/// Generate random codes for a DictArray with a specific unsigned integer type.
+/// Generate random codes for a DictArray with a specific integer type.
 fn random_codes<T>(
     u: &mut Unstructured,
     len: usize,

vortex-array/src/arrays/dict/array.rs

@@ -103,18 +103,18 @@ impl DictArray {
 
     /// Build a new `DictArray` from its components, `codes` and `values`.
     ///
-    /// The codes must be unsigned integers, and may be nullable. Values can be any type, and
-    /// may also be nullable. This mirrors the nullability of the Arrow `DictionaryArray`.
+    /// The codes must be integers, and may be nullable. Values can be any
+    /// type, and may also be nullable. This mirrors the nullability of the Arrow `DictionaryArray`.
     ///
     /// # Errors
     ///
-    /// The `codes` **must** be unsigned integers, and the maximum code must be less than the length
+    /// The `codes` **must** be integers, and the maximum code must be less than the length
     /// of the `values` array. Otherwise, this constructor returns an error.
     ///
     /// It is an error to provide a nullable `codes` with non-nullable `values`.
     pub fn try_new(codes: ArrayRef, values: ArrayRef) -> VortexResult<Self> {
-        if !codes.dtype().is_unsigned_int() {
-            vortex_bail!(MismatchedTypes: "unsigned int", codes.dtype());
+        if !codes.dtype().is_int() {
+            vortex_bail!(MismatchedTypes: "int", codes.dtype());
         }
 
         Ok(unsafe { Self::new_unchecked(codes, values) })
@@ -194,7 +194,11 @@ impl DictArray {
         match codes_validity.bit_buffer() {
             AllOr::All => {
                 match_each_integer_ptype!(codes_primitive.ptype(), |P| {
-                    #[allow(clippy::cast_possible_truncation)]
+                    #[allow(
+                        clippy::cast_possible_truncation,
+                        clippy::cast_sign_loss,
+                        reason = "codes are non-negative indices; a negative signed code would wrap to a large usize and panic on the bounds-checked array index"
+                    )]
                     for &code in codes_primitive.as_slice::<P>().iter() {
                         values_vec[code as usize] = referenced_value;
                     }
@@ -205,7 +209,11 @@ impl DictArray {
                 match_each_integer_ptype!(codes_primitive.ptype(), |P| {
                     let codes = codes_primitive.as_slice::<P>();
 
-                    #[allow(clippy::cast_possible_truncation)]
+                    #[allow(
+                        clippy::cast_possible_truncation,
+                        clippy::cast_sign_loss,
+                        reason = "codes are non-negative indices; a negative signed code would wrap to a large usize and panic on the bounds-checked array index"
+                    )]
                     buf.set_indices().for_each(|idx| {
                         values_vec[codes[idx] as usize] = referenced_value;
                     })

vortex-array/src/builders/dict/mod.rs

@@ -58,6 +58,9 @@ pub fn dict_encoder(array: &dyn Array, constraints: &DictConstraints) -> Box<dyn
     dict_builder
 }
 
+/// Encode an array as a `DictArray` subject to the given constraints.
+///
+/// Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility.
 pub fn dict_encode_with_constraints(
     array: &dyn Array,
     constraints: &DictConstraints,

vortex-btrblocks/src/compressor/float/dictionary.rs

@@ -2,6 +2,8 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 //! Float-specific dictionary encoding implementation.
+//!
+//! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility.
 
 use vortex_array::IntoArray;
 use vortex_array::arrays::DictArray;

vortex-btrblocks/src/compressor/integer/dictionary.rs

@@ -2,6 +2,8 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 //! Dictionary compressor that reuses the unique values in the `IntegerStats`.
+//!
+//! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility.
 
 use vortex_array::IntoArray;
 use vortex_array::arrays::DictArray;

vortex-cuda/src/kernel/arrays/dict.rs

@@ -21,8 +21,8 @@ use vortex_dtype::DType;
 use vortex_dtype::NativeDecimalType;
 use vortex_dtype::NativePType;
 use vortex_dtype::match_each_decimal_value_type;
+use vortex_dtype::match_each_integer_ptype;
 use vortex_dtype::match_each_native_simd_ptype;
-use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -75,7 +75,7 @@ async fn execute_dict_prim(dict: DictArray, ctx: &mut CudaExecutionCtx) -> Vorte
 
     // Dispatch based on both value type and code type
     match_each_native_simd_ptype!(values_ptype, |V| {
-        match_each_unsigned_integer_ptype!(codes_ptype, |I| {
+        match_each_integer_ptype!(codes_ptype, |I| {
             execute_dict_prim_typed::<V, I>(values_prim, codes_prim, ctx).await
         })
     })
@@ -165,7 +165,7 @@ async fn execute_dict_decimal(
     let decimal_type = values_decimal.values_type();
 
     match_each_decimal_value_type!(decimal_type, |V| {
-        match_each_unsigned_integer_ptype!(codes_ptype, |C| {
+        match_each_integer_ptype!(codes_ptype, |C| {
             execute_dict_decimal_typed::<V, C>(values_decimal, codes_prim, dtype, ctx).await
         })
     })

vortex-layout/src/layouts/dict/reader.rs

@@ -223,7 +223,7 @@ impl LayoutReader for DictReader {
             let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?;
 
             // SAFETY: Layout was validated at write time.
-            //  * The codes dtype is guaranteed to be an unsigned integer type from the layout
+            //  * The codes dtype is guaranteed to be an integer type from the layout
             //  * The codes child reader ensures the correct dtype.
             //  * The layout stores `all_values_referenced` and if this is malicious then it must
             //    only affect correctness not memory safety.

vortex-layout/src/layouts/dict/writer.rs

@@ -66,6 +66,8 @@ pub struct DictLayoutConstraints {
     /// The codes dtype is determined upfront from this constraint:
     /// - [`PType::U8`] when max_len <= 255
     /// - [`PType::U16`] when max_len > 255
+    ///
+    /// Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility.
     pub max_len: u16,
 }