fix(vortex-row): decimal sort keys are not memcmp-comparable across chunks

vortex-array/src/arrays/decimal/utils.rs

@@ -5,6 +5,9 @@ use itertools::Itertools;
 use itertools::MinMaxResult;
 use vortex_buffer::Buffer;
 use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_mask::Mask;
 
 use crate::arrays::DecimalArray;
 use crate::arrays::decimal::DecimalArrayExt;
@@ -28,6 +31,63 @@ pub(crate) fn widened_buffer<W: NativeDecimalType>(array: &DecimalArray) -> Buff
     })
 }
 
+/// Return the array's unscaled values converted to exactly `W`, whatever the array's
+/// storage type. Zero-copy when the array is already stored at `W`.
+///
+/// Widening is lossless. Narrowing fails for any *valid* value that does not fit `W`.
+/// Null slots may hold arbitrary bytes and never fail; their contents in the returned
+/// buffer are likewise arbitrary and must not be read.
+pub fn converted_buffer<W: NativeDecimalType>(
+    array: &DecimalArray,
+    validity: &Mask,
+) -> VortexResult<Buffer<W>> {
+    // Widening can never fail, so it needs no validation pass.
+    if array.values_type() <= W::DECIMAL_TYPE {
+        return Ok(widened_buffer(array));
+    }
+    match_each_decimal_value_type!(array.values_type(), |T| {
+        let src = array.buffer::<T>();
+        match validity {
+            Mask::AllTrue(_) => {
+                // Keeping the overflow scan branchless and vectorizable. Only on overflow
+                // do we rescan for diagnostics.
+                let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none());
+                if any_overflow {
+                    let (i, v) = src
+                        .iter()
+                        .enumerate()
+                        .find(|&(_, v)| W::from(*v).is_none())
+                        .vortex_expect("overflow scan found an overflowing value");
+                    vortex_bail!(
+                        "decimal value {v} at index {i} does not fit {}",
+                        W::DECIMAL_TYPE
+                    );
+                }
+            }
+            Mask::AllFalse(_) => return Ok(Buffer::zeroed(src.len())),
+            Mask::Values(values) => {
+                let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none());
+                if any_overflow {
+                    for (i, v) in src.iter().enumerate() {
+                        if values.value(i) && W::from(*v).is_none() {
+                            vortex_bail!(
+                                "decimal value {v} at index {i} does not fit {}",
+                                W::DECIMAL_TYPE
+                            );
+                        }
+                    }
+                }
+            }
+        }
+        // The convert pass is infallible: every valid value fits, and out-of-range null-slot
+        // garbage is mapped to the default value. Callers must still ignore null slots.
+        Ok(src
+            .iter()
+            .map(|v| W::from(*v).unwrap_or_default())
+            .collect())
+    })
+}
+
 macro_rules! try_downcast {
     ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{
         use crate::dtype::BigCast;

vortex-row/src/codec.rs

@@ -33,10 +33,13 @@ use vortex_array::arrays::NullArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
+use vortex_array::arrays::decimal::DecimalArrayExt;
+use vortex_array::arrays::decimal::converted_buffer;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
 use vortex_array::arrays::struct_::StructArrayExt;
 use vortex_array::dtype::DType;
+use vortex_array::dtype::DecimalDType;
 use vortex_array::dtype::DecimalType;
 use vortex_array::dtype::NativePType;
 use vortex_array::dtype::half::f16;
@@ -183,7 +186,7 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult<RowWidth> {
             ptype.byte_width(),
         )))),
         DType::Decimal(dt, _) => {
-            let vt = DecimalType::smallest_decimal_value_type(dt);
+            let vt = decimal_key_type(dt);
             if matches!(vt, DecimalType::I256) {
                 vortex_bail!("row encoding for Decimal256 is not yet implemented");
             }
@@ -385,10 +388,19 @@ fn add_size_primitive(arr: &PrimitiveArray, sizes: &mut [u32]) {
 }
 
 fn add_size_decimal(arr: &DecimalArray, sizes: &mut [u32]) {
-    let width = byte_width_u32(arr.values_type().byte_width());
+    let width = byte_width_u32(decimal_key_type(&arr.decimal_dtype()).byte_width());
     add_size_const(sizes, encoded_size_for_fixed(width));
 }
 
+/// The decimal type every chunk of a decimal column encodes its keys at.
+///
+/// Derived from the declared decimal dtype rather than the chunk's physical `values_type`,
+/// so keys from differently compressed chunks stay memcmp-comparable. Both the size plan
+/// ([`row_width_for_dtype`]) and the encoder derive their width from here.
+fn decimal_key_type(dt: &DecimalDType) -> DecimalType {
+    DecimalType::smallest_decimal_value_type(dt)
+}
+
 fn add_size_varbinview(
     arr: &VarBinViewArray,
     sizes: &mut [u32],
@@ -634,7 +646,7 @@ fn encode_decimal(
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<()> {
     let mask = arr.as_ref().validity()?.execute_mask(arr.len(), ctx)?;
-    match arr.values_type() {
+    match decimal_key_type(&arr.decimal_dtype()) {
         DecimalType::I8 => {
             encode_decimal_typed::<i8>(arr, &mask, field, row_offsets, col_offset, out)
         }
@@ -654,7 +666,6 @@ fn encode_decimal(
             vortex_bail!("row encoding for Decimal256 is not yet implemented")
         }
     }
-    Ok(())
 }
 
 fn encode_decimal_typed<T>(
@@ -664,14 +675,15 @@ fn encode_decimal_typed<T>(
     row_offsets: &[u32],
     col_offset: &mut [u32],
     out: &mut [u8],
-) where
+) -> VortexResult<()>
+where
     T: vortex_array::dtype::NativeDecimalType + RowEncode,
 {
     let non_null = field.non_null_sentinel();
     let null = field.null_sentinel();
     let value_bytes = size_of::<T>();
     let total = encoded_size_for_fixed(byte_width_u32(value_bytes));
-    let slice = arr.buffer::<T>();
+    let slice = converted_buffer::<T>(arr, mask)?;
     for i in 0..slice.len() {
         let pos = (row_offsets[i] + col_offset[i]) as usize;
         if mask.value(i) {
@@ -685,6 +697,7 @@ fn encode_decimal_typed<T>(
         }
         col_offset[i] += total;
     }
+    Ok(())
 }
 
 fn encode_varbinview(

vortex-row/src/tests.rs

@@ -10,16 +10,20 @@ use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
 use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::DecimalArray;
 use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::ListViewArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::arrays::listview::ListViewArrayExt;
 use vortex_array::arrays::listview::ListViewArraySlotsExt;
+use vortex_array::dtype::DecimalDType;
 use vortex_array::dtype::Nullability;
 use vortex_array::extension::datetime::Date;
 use vortex_array::extension::datetime::TimeUnit;
+use vortex_array::validity::Validity;
+use vortex_buffer::buffer;
 use vortex_error::VortexResult;
 
 use crate::RowEncoder;
@@ -614,3 +618,77 @@ fn reject_list_dtype_early() {
         "expected error mentioning List, got: {err}"
     );
 }
+
+/// Chunks of one decimal column can compress to different physical value widths. The key
+/// width must come from the declared dtype, not the chunk's `values_type`, otherwise keys
+/// from different chunks are not memcmp-comparable.
+#[rstest]
+#[case::ascending(false)]
+#[case::descending(true)]
+fn decimal_keys_comparable_across_chunk_widths(#[case] descending: bool) -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let dtype = DecimalDType::new(7, 5);
+    let field = RowSortField::new(descending, true);
+
+    // One logical DECIMAL(7, 5) column whose chunks compressed to different physical widths.
+    let chunk_i8 = DecimalArray::new(buffer![91i8, -5], dtype, Validity::NonNullable).into_array();
+    let chunk_i16 =
+        DecimalArray::new(buffer![484i16, 300], dtype, Validity::NonNullable).into_array();
+    let values: [i32; 4] = [91, -5, 484, 300];
+
+    let keys_i8 = collect_row_bytes(&convert_columns(&[chunk_i8], &[field], &mut ctx)?);
+    let keys_i16 = collect_row_bytes(&convert_columns(&[chunk_i16], &[field], &mut ctx)?);
+    let keys = [keys_i8, keys_i16].concat();
+
+    for (i, vi) in values.iter().enumerate() {
+        for (j, vj) in values.iter().enumerate() {
+            let expected = if descending { vj.cmp(vi) } else { vi.cmp(vj) };
+            assert_eq!(
+                keys[i].cmp(&keys[j]),
+                expected,
+                "keys for {vi} and {vj} do not compare like the values"
+            );
+        }
+    }
+    Ok(())
+}
+
+/// A null slot's backing value is unspecified and might not fit the dtype-derived key width;
+/// encoding must ignore it rather than report a spurious overflow.
+#[test]
+fn decimal_null_slot_garbage_does_not_error() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let dtype = DecimalDType::new(7, 5);
+    let field = RowSortField::new(false, true);
+
+    let chunk = DecimalArray::new(
+        buffer![484i64, 10_000_000_000_000, 91],
+        dtype,
+        Validity::from_iter([true, false, true]),
+    )
+    .into_array();
+
+    let keys = collect_row_bytes(&convert_columns(&[chunk], &[field], &mut ctx)?);
+    assert!(keys[1] < keys[2], "null must sort before non-nulls");
+    assert!(keys[2] < keys[0], "91 must sort before 484");
+    Ok(())
+}
+
+/// A valid value too large for the dtype-derived key width must fail loudly instead of
+/// silently encoding a corrupt key.
+#[test]
+fn decimal_value_not_fitting_key_width_errors() {
+    let mut ctx = array_session().create_execution_ctx();
+    let dtype = DecimalDType::new(7, 5);
+    let field = RowSortField::new(false, true);
+
+    let chunk = DecimalArray::new(buffer![10_000_000_000_000i64], dtype, Validity::NonNullable)
+        .into_array();
+
+    let err = convert_columns(&[chunk], &[field], &mut ctx)
+        .expect_err("a valid value wider than the key width must be rejected");
+    assert!(
+        err.to_string().contains("does not fit"),
+        "expected a does-not-fit error, got: {err}"
+    );
+}