Unsound `sort_run_downcasted` constructs invalid `RunArray` on empty array or zero limit

[Manishearth]

Note

This finding was identified during an agentic unsafe Rust code review performed by Gemini AI, followed by human review and verification.

The Issue

In arrow-ord, the internal run-end encoded array sorting function sort_run_downcasted constructs a sorted RunArray. When the logical length of the output run array (output_len) is computed to be 0 (which occurs if the input array has logical length 0, or if limit is Some(0)), the function proceeds to call sort_run_inner. For empty arrays, sort_run_inner calculates a physical length of 1 and calls consume_runs once with a run length of 0. Consequently, new_run_ends_builder appends a single value 0.

arrow-rs/arrow-ord/src/sort.rs

Lines 673 to 728 in 7505005

	fn sort_run_downcasted<R: RunEndIndexType>(
	values: &dyn Array,
	options: Option<SortOptions>,
	limit: Option<usize>,
	) -> Result<ArrayRef, ArrowError> {
	let run_array = values.as_any().downcast_ref::<RunArray<R>>().unwrap();
	// Determine the length of output run array.
	let output_len = if let Some(limit) = limit {
	limit.min(run_array.len())
	} else {
	run_array.len()
	};
	let run_ends = run_array.run_ends();
	let mut new_run_ends_builder = BufferBuilder::<R::Native>::new(run_ends.len());
	let mut new_run_end: usize = 0;
	let mut new_physical_len: usize = 0;
	let consume_runs = |run_length, _| {
	new_run_end += run_length;
	new_physical_len += 1;
	new_run_ends_builder.append(R::Native::from_usize(new_run_end).unwrap());
	};
	let (values_indices, run_values) = sort_run_inner(run_array, options, output_len, consume_runs);
	let new_run_ends = unsafe {
	// Safety:
	// The function builds a valid run_ends array and hence need not be validated.
	ArrayDataBuilder::new(R::DATA_TYPE)
	.len(new_physical_len)
	.add_buffer(new_run_ends_builder.finish())
	.build_unchecked()
	};
	// slice the sorted value indices based on limit.
	let new_values_indices: PrimitiveArray<UInt32Type> = values_indices
	.slice(0, new_run_ends.len())
	.into_data()
	.into();
	let new_values = take(&run_values, &new_values_indices, None)?;
	// Build sorted run array
	let builder = ArrayDataBuilder::new(run_array.data_type().clone())
	.len(new_run_end)
	.add_child_data(new_run_ends)
	.add_child_data(new_values.into_data());
	let array_data: RunArray<R> = unsafe {
	// Safety:
	// This function builds a valid run array and hence can skip validation.
	builder.build_unchecked().into()
	};
	Ok(Arc::new(array_data))

The safety invariant of RunArray requires that its run_ends array contains strictly increasing positive integers greater than zero. By calling build_unchecked() with a buffer containing [0] and converting it into a RunArray, the code bypasses validation and violates the unsafe safety contract of RunEndBuffer::new_unchecked. Because safe public sorting APIs (sort and sort_limit) can be invoked by safe callers and result in violating an unsafe contract (RunEndBuffer::new_unchecked requiring strictly positive run_ends > 0), this constitutes a soundness vulnerability.

When validating the resulting ArrayData with validate_full(), it panics due to the invalid run_ends values:

Minimal Reproduction (Structural Invariant Violation)

use arrow_array::types::{Int16Type, Int32Type};
use arrow_array::builder::PrimitiveRunBuilder;
use arrow_array::Array;
use arrow_ord::sort::sort_limit;

fn main() {
    let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
    builder.extend([Some(10), Some(20)]);
    let run_array = builder.finish();
    let sorted = sort_limit(&run_array, None, Some(0)).unwrap();

    // Validate the sorted ArrayData, panicking if it is invalid.
    let data = sorted.into_data();
    data.validate_full().unwrap();
}

thread 'main' panicked at src/bin/repro1.rs:14:26:
called `Result::unwrap()` on an `Err` value: InvalidArgumentError("The values in run_ends array should be strictly positive. Found value 0 at index 0 that does not match the criteria.")

Miri Undefined Behavior Demonstration

The caller operates strictly within expected type contracts: the returned sorted array from safe API sort_limit is a RunArray with logical length 0. Obtaining its physical index via get_start_physical_index() is valid. However, because sort_limit violated RunEndBuffer::new_unchecked by constructing an invalid run_ends = [0] buffer, physical_idx evaluates to 0.

When downstream code consumes this RunArray and uses value_unchecked(physical_idx) on a 0-length sliced values array (which is sound assuming valid RunArray invariants hold), Miri immediately halts due to an out-of-bounds pointer read on an unallocated address:

use arrow_array::types::{Int16Type, Int32Type};
use arrow_array::builder::PrimitiveRunBuilder;
use arrow_array::Array;
use arrow_ord::sort::sort_limit;

fn main() {
    let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
    builder.extend([Some(10), Some(20)]);
    let run_array = builder.finish();
    let sorted = sort_limit(&run_array, None, Some(0)).unwrap();

    let sorted_run_array = sorted.as_any().downcast_ref::<arrow_array::RunArray<Int16Type>>().unwrap();
    
    // get_start_physical_index() relies on run_ends > 0 invariant, returning physical_idx = 0
    let physical_idx = sorted_run_array.run_ends().get_start_physical_index();
    let zero_values = sorted_run_array.values().slice(0, 0);
    let typed_values = zero_values.as_any().downcast_ref::<arrow_array::Int32Array>().unwrap();

    unsafe {
        // `physical_idx` is supposed to be a valid index. 
        // Because sorted_run_array was constructed unsoundly with run_ends = [0],
        // value_unchecked(0) executes unchecked pointer arithmetic on a 0-length allocation.
        let val = typed_values.value_unchecked(physical_idx);
        println!("val = {}", val);
    }
}

error: Undefined Behavior: `assume` called with `false`
   --> /usr/local/google/home/manishearth/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/arrow-array-57.3.1/src/array/primitive_array.rs:761:19
    |
761 |         unsafe { *self.values.get_unchecked(i) }
    |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Undefined Behavior occurred here

Suggested Fix

At the beginning of sort_run_downcasted, explicitly check if output_len == 0. If so, immediately return an empty RunArray without performing run-length calculations or calling build_unchecked().

---

Note

The full audit report below also contains additional minor findings (such as missing safety comments or undocumented FFI assumptions) that are probably worth fixing as well but not the primary goal of this issue. The audit report has not been human-reviewed, it may contain misleading claims.

Full Gemini Codebase Audit Report Appendix

Unsafe Rust Review: arrow_ord (v57)

Overall Safety Assessment

arrow_ord contains a moderate amount of unsafe code, mostly used for performance optimization (such as bypassing bounds checks via value_unchecked and get_unchecked or skipping array validation via build_unchecked).

While the majority of unsafe operations are soundly structured, we identified a critical unsoundness issue in the run-end encoded array sorting (sort_run_downcasted) where sorting an empty array or sorting with limit = Some(0) constructs an invalid RunArray that violates its safety invariants, leading to a safety contract violation in downstream safe code.

Additionally, safety documentation is generally sparse or missing for many internal unsafe blocks and trait implementations, and we identified a logic bug that causes integer underflow inside an unsafe block when sorting sliced empty run arrays.

Critical Findings

Unsoundness in sort_run_downcasted with Empty Array or Zero Limit 🔴 🚨

• Priority: 🔴 High
• Threat Vector: 🚨 Untrusted Input
• Bug Type: Out-of-Bounds Access

In src/sort.rs at line 673, the sort_run_downcasted function sorts a RunArray. When the logical length of the output run array (output_len) is computed to be 0 (which happens if the input array has logical length 0, or if limit is Some(0)), the function proceeds to call sort_run_inner.

In sort_run_inner, since the input array logical length is 0, it determines that physical_len is 1 (due to end_physical_index - start_physical_index + 1 where both indices default to 0 for empty arrays). It then enters a loop to calculate the run lengths and calls consume_runs once with a run length of 0.

The closure consume_runs (at line 693) appends the run lengths to new_run_ends_builder:

    let consume_runs = |run_length, _| {
        new_run_end += run_length;
        new_physical_len += 1;
        new_run_ends_builder.append(R::Native::from_usize(new_run_end).unwrap());
    };

As a result, new_run_ends_builder appends a single value 0. Then, new_run_ends is built containing the buffer [0] of length 1. Finally, the sorted RunArray is built:

    // Build sorted run array
    let builder = ArrayDataBuilder::new(run_array.data_type().clone())
        .len(new_run_end) // new_run_end is 0
        .add_child_data(new_run_ends)
        .add_child_data(new_values.into_data());
    let array_data: RunArray<R> = unsafe {
        // Safety:
        //  This function builds a valid run array and hence can skip validation.
        builder.build_unchecked().into()
    };

The safety invariant of a RunArray requires that its run_ends child array does not contain null values, has strictly increasing positive integers, and that the last value matches the logical length of the RunArray if it is non-empty. For an empty RunArray (logical length 0), the run_ends must be empty.

By using build_unchecked() and passing new_run_ends containing [0], we bypass validation and construct an invalid ArrayData representing the RunArray. When .into() is called, it invokes From<ArrayData> for RunArray<R>, which is safe code. However, this implementation internally calls:

        let run_ends = unsafe {
            let scalar = child.buffers()[0].clone().into();
            RunEndBuffer::new_unchecked(scalar, data.offset(), data.len())
        };

RunEndBuffer::new_unchecked requires that the run ends buffer contains values greater than zero. Since the buffer contains [0], this safety contract is violated. Because a safe function (sort or sort_limit containing this logic) can be invoked by safe callers and results in the violation of an unsafe contract (new_unchecked), this function is unsound.

Remediation: sort_run_downcasted should check if output_len == 0 at the very beginning, and if so, return an empty RunArray immediately without proceeding to run-length calculations and building invalid array data.

---

Fishy Findings

Incorrect Safety Comment in sort_run_downcasted 🟡 ⚠️

• Priority: 🟡 Low

• Threat Vector: ⚠️ Accidental Misuse

• Bug Type: Invalid Safety Comment

In src/sort.rs at line 723, the unsafe block has the following safety comment:

    let array_data: RunArray<R> = unsafe {
        // Safety:
        //  This function builds a valid run array and hence can skip validation.
        builder.build_unchecked().into()
    };

As shown in the Critical Findings section, this statement is false when output_len is 0, because the function builds an invalid run array containing [0] in its run_ends buffer.

Underflow inside Unsafe Block in sort_run_inner 🟡 🚨

• Priority: 🟡 Low
• Threat Vector: 🚨 Untrusted Input
• Bug Type: Integer Underflow

In src/sort.rs at line 795, when logical_length == 0 but the underlying array has non-zero physical elements (e.g., sorting a sliced array with logical length 0), start_physical_index and end_physical_index default to 0. This causes the loop to process physical index 0. Inside the unsafe block, it evaluates:

run_ends.get_unchecked(physical_index).as_usize() - run_array.offset()

Since physical_index is 0, it reads the first run end (e.g. 3). But run_array.offset() is 5. This evaluates to 3 - 5, which underflows. In debug mode, this triggers a panic. In release mode, it wraps around to usize::MAX - 1. While it does not directly cause UB here because new_run_length is capped to remaining_len (which is 0), executing wrapping calculations or panicking inside unsafe blocks due to incorrect state representation is extremely risky and indicates a logic bug in boundary handling.

---

Missing Safety Comments

src/cmp.rs

ArrayOrd::value_unchecked Implementations 🔴

The implementations of value_unchecked in ArrayOrd for &BooleanArray (line 509) and &[T] (line 529) lack safety comments explaining how their preconditions are discharged.

Proposed Comments: For &BooleanArray (line 509):

    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
        // SAFETY: The caller guarantees that `idx < self.len()`.
        // `BooleanArray::value_unchecked` has the same precondition.
        unsafe { BooleanArray::value_unchecked(self, idx) }
    }

For &[T] (line 529):

    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
        // SAFETY: The caller guarantees that `idx < self.len()`.
        // By `slice::get_unchecked`'s contract, `idx` must be in bounds of the slice, which is `0..self.len()`.
        // Thus, the precondition implies the safety contract of `get_unchecked`.
        unsafe { *self.get_unchecked(idx) }
    }

Calls to value_unchecked in apply_op 🔴

The calls to value_unchecked in apply_op at lines 445, 456, and 460 lack safety comments.

Proposed Comments: At line 445:

            // SAFETY: `collect_bool` calls the closure with `idx` in `0..l.len()`.
            // Since `l.len() == r.len()`, `idx` is also in `0..r.len()`.
            // This satisfies the precondition of `value_unchecked`.
            collect_bool(l.len(), neg, |idx| unsafe {
                op(l.value_unchecked(idx), r.value_unchecked(idx))
            })

At line 456:

            // SAFETY: `collect_bool` calls the closure with `idx` in `0..r.len()`.
            // This satisfies the precondition of `r.value_unchecked`.
            collect_bool(r.len(), neg, |idx| op(v, unsafe { r.value_unchecked(idx) }))

At line 460:

            // SAFETY: `collect_bool` calls the closure with `idx` in `0..l.len()`.
            // This satisfies the precondition of `l.value_unchecked`.
            collect_bool(l.len(), neg, |idx| op(unsafe { l.value_unchecked(idx) }, v))

apply_op_vectored Unsafe Block 🔴

The unsafe block inside apply_op_vectored (line 475) lacks safety comments for its get_unchecked and value_unchecked calls.

Proposed Comment:

    // SAFETY:
    // - `collect_bool` calls the closure with `idx` in `0..l_v.len()`.
    //   Since `l_v.len() == r_v.len()`, `idx` is in bounds for both slices, making `get_unchecked` safe.
    // - `l_v` and `r_v` are indices derived from `AnyDictionaryArray::normalized_keys()`.
    //   By the invariants of `AnyDictionaryArray`, the normalized keys are valid indices into the values array.
    //   Therefore `l_idx < l.len()` and `r_idx < r.len()`, satisfying the preconditions of `value_unchecked`.
    collect_bool(l_v.len(), neg, |idx| unsafe {
        let l_idx = *l_v.get_unchecked(idx);
        let r_idx = *r_v.get_unchecked(idx);
        op(l.value_unchecked(l_idx), r.value_unchecked(r_idx))
    })

---

src/sort.rs

partition_validity_scan Unsafe Block 🔴

The unsafe block in partition_validity_scan (line 219) lacks a safety comment explaining why calling set_len on valid and nulls is sound.

Proposed Comment:

    // SAFETY:
    // - `valid` and `nulls` are allocated with capacity `len - null_count` and `null_count` respectively.
    // - By the invariants of `NullBuffer` (which `bitmap` is), the number of set bits is exactly `len - null_count`,
    //   and the number of unset bits is exactly `null_count`.
    // - `set_indices_u32()` yields exactly one index per set bit, so the first loop writes exactly `len - null_count`
    //   elements, completely initializing the spare capacity of `valid`.
    // - Similarly, the negated buffer yields exactly `null_count` elements, completely initializing `nulls`.
    // - Therefore, calling `set_len` with these lengths is sound as all elements in the active range are initialized.
    unsafe {
        ...
    }

sort_bytes Unsafe Block 🔴

The unsafe block in sort_bytes (line 357) lacks safety comments for value_unchecked and std::ptr::read_unaligned.

Proposed Comment:

        // SAFETY:
        // - `idx` is from `value_indices`, which only contains indices `< values.len()` (guaranteed by `partition_validity`).
        //   This satisfies the precondition of `values.value_unchecked`.
        // - `slice.len() >= 4` ensures that `slice.as_ptr()` points to at least 4 valid initialized bytes in a single allocation.
        //   Therefore, reading a `u32` using `read_unaligned` is safe and within bounds.
        .map(|idx| unsafe {
            let slice: &[u8] = values.value_unchecked(idx as usize).as_ref();
            let len = slice.len() as u64;
            let prefix = if slice.len() >= 4 {
                let raw = std::ptr::read_unaligned(slice.as_ptr() as *const u32);
                u32::from_be(raw)
            ...

cmp_bytes Unsafe Comparator Block 🔴

The unsafe block in cmp_bytes (line 386) lacks safety comments.

Proposed Comment:

    // SAFETY:
    // - `ia` and `ib` are `idx` values from `valids`, which are copied from `value_indices`.
    // - As established, `value_indices` only contains indices `< values.len()`.
    // - This satisfies the precondition of `values.value_unchecked`.
    let cmp_bytes = |a: &(u32, u32, u64), b: &(u32, u32, u64)| unsafe {
        let (ia, pa, la) = *a;
        let (ib, pb, lb) = *b;
        ...
        let a_bytes: &[u8] = values.value_unchecked(ia as usize).as_ref();
        let b_bytes: &[u8] = values.value_unchecked(ib as usize).as_ref();
        a_bytes.cmp(b_bytes)
    };

sort_byte_view Unsafe Value Accesses 🔴

The unsafe block inside the mixed comparator in sort_byte_view (lines 494, 495) lacks safety comments.

Proposed Comment:

            // SAFETY:
            // - `a.0` and `b.0` are indices from `valids`, which are copied from `value_indices`.
            // - `value_indices` only contains indices `< values.len()`.
            // - This satisfies the precondition of `values.value_unchecked`.
            let full_a: &[u8] = unsafe { values.value_unchecked(a.0 as usize).as_ref() };
            let full_b: &[u8] = unsafe { values.value_unchecked(b.0 as usize).as_ref() };

[Jefffrey]

thanks for identifying this, have raised a PR to fix it