feat(vortex-onpair): store codes_offsets at adaptive u32/u64 width

Compression accepts u64 byte offsets (large-binary inputs), but the
codes_offsets child was always narrowed to u32, capping a chunk at 2^32
tokens and failing compression above it. Pick the narrowest of u32/u64
that holds the largest per-row code boundary instead, so codes_offsets
scales with the u64 byte-offset capacity. The cascading compressor still
narrows the common u32 case down to u16/u8, and the width round-trips via
the existing codes_offsets_ptype metadata, so the serialized format is
unchanged.

Widen CodesWindow to Buffer<u64> to match; the u64->usize conversions are
checked (cast_possible_truncation is denied) but fold away on 64-bit.

The u64 branch cannot be reached with realistic test data (>4 GiB chunk),
so cover it two ways: a unit test drives the width selection via a
threshold parameter, and a read-path test hand-widens a small array's
codes_offsets child and asserts canonical decode and the compressed-domain
equality compare (CodesWindow) behave identically to the u32 width.

Signed-off-by: Francesco Gargiulo <francesco@spiraldb.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: gargiulofrancesco

encodings/experimental/onpair/src/array.rs

@@ -104,8 +104,9 @@ pub struct OnPairSlots {
     /// Primitive integer token codes. Downstream integer compression may
     /// narrow or bit-pack this child independently of the OnPair metadata.
     pub codes: ArrayRef,
-    /// `PrimitiveArray<u32>`, length `num_rows + 1`. FoR / RunEnd / etc. apply
-    /// naturally via the cascading compressor.
+    /// `PrimitiveArray<u32>` (or `u64` when a chunk exceeds `u32::MAX` tokens),
+    /// length `num_rows + 1`. FoR / RunEnd / etc. apply naturally via the
+    /// cascading compressor.
     pub codes_offsets: ArrayRef,
     /// Integer `PrimitiveArray`, length `num_rows`. Used to size the canonical
     /// output buffer.

encodings/experimental/onpair/src/compress.rs

@@ -88,17 +88,7 @@ where
         .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?;
     let (dict, codes, row_offsets) = column.into_raw();
     let (dict_bytes, dict_offsets) = dict.into_raw();
-    let codes_offsets = Buffer::from(
-        row_offsets
-            .into_iter()
-            .map(|o| {
-                let value = o.to_usize();
-                u32::try_from(value)
-                    .map_err(|_| vortex_err!("OnPair code boundary {value} does not fit u32"))
-            })
-            .collect::<VortexResult<Vec<_>>>()?,
-    )
-    .into_array();
+    let codes_offsets = codes_offsets_array(&row_offsets, u32::MAX as usize);
     let codes = Buffer::from(codes).into_array();
     let dict_offsets = Buffer::from(dict_offsets).into_array();
 
@@ -133,6 +123,35 @@ fn dict_bytes_to_buffer(dict_bytes: Vec<u8>) -> BufferHandle {
     BufferHandle::new_host(aligned.freeze())
 }
 
+/// Build the `codes_offsets` child from the library's per-row code boundaries,
+/// storing the narrowest of `u32`/`u64` that holds the largest boundary.
+/// `row_offsets` is non-decreasing, so its last entry is that maximum and one
+/// bound check picks the width. `u32` covers the common case (the cascading
+/// compressor narrows it further to `u16`/`u8`); `u64` engages only when a
+/// single chunk carries more than `u32_max` tokens, matching the `u64` byte
+/// offsets accepted at compression. `u32_max` is a parameter so tests can drive
+/// the `u64` branch without a multi-GiB array.
+fn codes_offsets_array<O: Offset>(row_offsets: &[O], u32_max: usize) -> ArrayRef {
+    let total_tokens = row_offsets.last().map_or(0, |&o| o.to_usize());
+    if total_tokens <= u32_max {
+        Buffer::from(
+            row_offsets
+                .iter()
+                .map(|&o| u32::try_from(o.to_usize()).vortex_expect("code boundary fits u32"))
+                .collect::<Vec<u32>>(),
+        )
+        .into_array()
+    } else {
+        Buffer::from(
+            row_offsets
+                .iter()
+                .map(|&o| u64::try_from(o.to_usize()).vortex_expect("token count fits u64"))
+                .collect::<Vec<u64>>(),
+        )
+        .into_array()
+    }
+}
+
 /// Compress any [`ArrayRef`] whose canonical form is a string array, by first
 /// canonicalising to `VarBinViewArray`.
 pub fn onpair_compress(
@@ -143,3 +162,30 @@ pub fn onpair_compress(
     let view = array.clone().execute::<VarBinViewArray>(ctx)?;
     onpair_compress_varbinview::<u64>(view, config, ctx)
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+
+    use super::codes_offsets_array;
+
+    #[test]
+    fn codes_offsets_width_selection() {
+        // Largest boundary within the threshold is stored as u32.
+        let narrow = codes_offsets_array::<u64>(&[0, 3, 7], 7);
+        assert_eq!(narrow.len(), 3);
+        assert_eq!(
+            narrow.dtype(),
+            &DType::Primitive(PType::U32, Nullability::NonNullable)
+        );
+
+        // A boundary above the threshold widens the child to u64.
+        let wide = codes_offsets_array::<u64>(&[0, 3, 8], 7);
+        assert_eq!(
+            wide.dtype(),
+            &DType::Primitive(PType::U64, Nullability::NonNullable)
+        );
+    }
+}

encodings/experimental/onpair/src/decode.rs

@@ -12,6 +12,7 @@ use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::NativePType;
 use vortex_buffer::Buffer;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
@@ -57,17 +58,21 @@ pub(crate) fn code_boundary_at(
 ///
 /// [`row`]: CodesWindow::row
 pub(crate) struct CodesWindow {
-    offsets: Buffer<u32>,
+    offsets: Buffer<u64>,
     codes: Buffer<u16>,
     code_start: usize,
 }
 
 impl CodesWindow {
     /// The codes for row `i`.
     pub(crate) fn row(&self, i: usize) -> &[u16] {
-        let start = self.offsets[i] as usize - self.code_start;
-        let end = self.offsets[i + 1] as usize - self.code_start;
-        &self.codes[start..end]
+        &self.codes[self.local(i)..self.local(i + 1)]
+    }
+
+    /// Offset `i` rebased into the window's local `codes` slice. Offsets are
+    /// bounded by `codes.len()` (a `usize`), so the conversion never truncates.
+    fn local(&self, i: usize) -> usize {
+        usize::try_from(self.offsets[i]).vortex_expect("code offset fits usize") - self.code_start
     }
 }
 
@@ -80,7 +85,7 @@ pub(crate) fn collect_codes_window(
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<CodesWindow> {
     let len = array.len();
-    let offsets = collect_widened::<u32>(array.codes_offsets(), ctx)?;
+    let offsets = collect_widened::<u64>(array.codes_offsets(), ctx)?;
     vortex_ensure!(
         offsets.len() == len + 1,
         "OnPair codes_offsets has {} entries, expected len + 1 = {}",
@@ -91,8 +96,8 @@ pub(crate) fn collect_codes_window(
         offsets.is_sorted(),
         "OnPair codes_offsets must be nondecreasing"
     );
-    let code_start = offsets[0] as usize;
-    let code_end = offsets[len] as usize;
+    let code_start = usize::try_from(offsets[0]).vortex_expect("code offset fits usize");
+    let code_end = usize::try_from(offsets[len]).vortex_expect("code offset fits usize");
     vortex_ensure!(
         code_end <= array.codes().len(),
         "OnPair codes_offsets end {} exceeds codes len {}",

encodings/experimental/onpair/src/tests.rs

@@ -7,15 +7,19 @@ use onpair::CompactDictionaryView;
 use prost::Message;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::VarBinArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::arrays::filter::FilterKernel;
 use vortex_array::assert_arrays_eq;
+use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::match_each_integer_ptype;
+use vortex_array::scalar_fn::fns::operators::Operator;
 use vortex_array::test_harness::check_metadata;
 use vortex_array::validity::Validity;
 use vortex_buffer::BufferMut;
@@ -114,6 +118,67 @@ fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> {
     Ok(())
 }
 
+/// The `u64` `codes_offsets` branch only engages past `u32::MAX` tokens (a
+/// multi-GiB chunk), so exercise the read path by widening a small array's
+/// `codes_offsets` child to `u64` by hand and asserting the decode paths
+/// (canonical and the compressed-domain equality compare, which builds a
+/// `CodesWindow`) treat it identically to the default `u32` width.
+#[cfg_attr(miri, ignore)]
+#[test]
+fn test_onpair_u64_codes_offsets() -> vortex_error::VortexResult<()> {
+    let mut ctx = SESSION.create_execution_ctx();
+    let narrow = onpair_compress(
+        &sample_input().into_array(),
+        DEFAULT_DICT12_CONFIG,
+        &mut ctx,
+    )?;
+
+    // Rebuild with only codes_offsets widened to u64; every other child is the
+    // input's, so the two arrays differ solely in codes_offsets width.
+    let wide = {
+        let view = narrow.as_view();
+        let wide_offsets = view
+            .codes_offsets()
+            .cast(DType::Primitive(PType::U64, Nullability::NonNullable))?
+            .execute::<PrimitiveArray>(&mut ctx)?
+            .into_array();
+        OnPair::try_new(
+            view.dtype().clone(),
+            view.dict_bytes_handle().clone(),
+            view.dict_offsets().clone(),
+            view.codes().clone(),
+            wide_offsets,
+            view.uncompressed_lengths().clone(),
+            view.array_validity(),
+        )?
+    };
+    assert_eq!(
+        wide.as_view().codes_offsets().dtype(),
+        &DType::Primitive(PType::U64, Nullability::NonNullable)
+    );
+
+    // Canonical decode is byte-identical across the two widths.
+    let narrow_decoded = narrow.into_array().execute::<VarBinViewArray>(&mut ctx)?;
+    let wide_decoded = wide
+        .clone()
+        .into_array()
+        .execute::<VarBinViewArray>(&mut ctx)?;
+    assert_arrays_eq!(&wide_decoded, &narrow_decoded, &mut ctx);
+
+    // Equality compare drives CodesWindow over the u64 codes_offsets.
+    let needle = ConstantArray::new("https://www.example.com/page", wide.len()).into_array();
+    let eq = wide
+        .into_array()
+        .binary(needle, Operator::Eq)?
+        .execute::<BoolArray>(&mut ctx)?;
+    assert_arrays_eq!(
+        &eq,
+        &BoolArray::from_iter([true, false, false, false, true]),
+        &mut ctx
+    );
+    Ok(())
+}
+
 #[cfg_attr(miri, ignore)]
 #[test]
 fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> {