From 1b7a72f39b1d1b576f515e10525bca375db8e268 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Mon, 6 Jul 2026 17:42:22 +0100 Subject: [PATCH 01/13] feat(vortex-onpair): port to onpair 0.1.0 API Port the OnPair encoding to the 0.1.0 public API, including raw column parts, validated compact dictionaries, and MaxDictBits configuration. Pin the latest upstream validation fix, add compressed-domain equality pushdown, and strengthen file-borne dictionary and code validation. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Fable 5 --- Cargo.lock | 5 +- Cargo.toml | 2 +- encodings/experimental/onpair/README.md | 17 +- .../experimental/onpair/benches/decode.rs | 53 ++--- .../onpair/goldenfiles/onpair.metadata | 2 +- encodings/experimental/onpair/src/array.rs | 102 +++++----- .../experimental/onpair/src/canonical.rs | 36 ++-- encodings/experimental/onpair/src/compress.rs | 98 +++------- .../experimental/onpair/src/compute/cast.rs | 105 ++++++++-- .../onpair/src/compute/compare.rs | 181 ++++++++++++++++-- .../experimental/onpair/src/compute/filter.rs | 1 - .../experimental/onpair/src/compute/slice.rs | 1 - encodings/experimental/onpair/src/decode.rs | 85 ++++++++ encodings/experimental/onpair/src/kernel.rs | 4 +- encodings/experimental/onpair/src/lib.rs | 9 +- encodings/experimental/onpair/src/ops.rs | 36 ++-- encodings/experimental/onpair/src/tests.rs | 28 ++- vortex-btrblocks/src/schemes/string/onpair.rs | 3 +- 18 files changed, 539 insertions(+), 229 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ca59b40e9dc..2e28df315f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6451,9 +6451,8 @@ checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" [[package]] name = "onpair" -version = "0.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7f54ca0079320b4b43a7502822b8a2e8a2c713adb7fc65ae4e65f940921a7eb" +version = "0.1.0" +source = "git+https://github.com/spiraldb/onpair.git?rev=28df115c799674c663adedc8f81620ac82c7a890#28df115c799674c663adedc8f81620ac82c7a890" dependencies = [ "hashbrown 0.16.1", "rand 0.9.4", diff --git a/Cargo.toml b/Cargo.toml index 12e91c7ef20..29145f8b4e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -192,7 +192,7 @@ num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } once_cell = "1.21" oneshot = { version = "0.2.0", features = ["async"] } -onpair = { version = "0.0.4" } +onpair = { git = "https://github.com/spiraldb/onpair.git", rev = "28df115c799674c663adedc8f81620ac82c7a890" } opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" opentelemetry_sdk = "0.32.0" diff --git a/encodings/experimental/onpair/README.md b/encodings/experimental/onpair/README.md index 9628c006201..c2419dd0c36 100644 --- a/encodings/experimental/onpair/README.md +++ b/encodings/experimental/onpair/README.md @@ -10,15 +10,18 @@ cascading-compressor support on every integer child. ## Compute -Like the FSST encoding, this crate provides `cast` and `filter` -pushdown. Other operators fall back to ordinary decompression. +Like the FSST encoding, this crate pushes down common operations over the +encoded representation. It supports `cast`, `filter`, byte length, and +constant equality / inequality. Unsupported operators fall back to ordinary +decompression. ## Default Configuration -The default training preset is **dict-12**: 12 bits per token, -dictionary capped at 4 096 entries. Token codes are stored as a -`PrimitiveArray`; downstream `FastLanes::BitPacking` losslessly -narrows the child to exactly `bits`-bit codes on disk. +The default training preset is **dict-12**: the OnPair trainer may build a +dictionary with up to 4 096 tokens. After compression, the runtime code width is +derived from the actual dictionary size. Vortex stores token codes as an integer +child array; downstream integer compression may narrow or bit-pack that child +independently of the OnPair metadata. ## Layout @@ -26,7 +29,7 @@ narrows the child to exactly `bits`-bit codes on disk. padded with `MAX_TOKEN_SIZE` trailing zero bytes so the over-copy decoder can read 16 bytes past the last token. - Slot 0 — `dict_offsets`: `PrimitiveArray`, len `dict_size + 1`. -- Slot 1 — `codes`: `PrimitiveArray`, length `total_tokens`. +- Slot 1 — `codes`: integer `PrimitiveArray`, length `total_tokens`. - Slot 2 — `codes_offsets`: `PrimitiveArray`, length `num_rows + 1`. - Slot 3 — `uncompressed_lengths`: integer `PrimitiveArray`, length `num_rows`. diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index ff4158cda32..e1055bc1eb3 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -3,7 +3,7 @@ // //! Decode-path microbenchmarks for the OnPair Vortex array. //! -//! * `decompress_into` — the upstream `onpair::decompress_into` decoder hot +//! * `decode_into` — the upstream `onpair::decode_into` decoder hot //! loop, fed by pre-materialised [`DecodeInputs`]. Measures the inner loop //! only (no child `execute`, no allocation). //! * `canonicalize_to_varbinview` — the full Vortex @@ -28,7 +28,7 @@ use std::mem::MaybeUninit; use std::sync::LazyLock; use divan::Bencher; -use onpair::Parts; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -38,12 +38,12 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; +use vortex_array::buffer::BufferHandle; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_buffer::Buffer; -use vortex_buffer::ByteBuffer; use vortex_mask::Mask; use vortex_onpair::DEFAULT_DICT12_CONFIG; use vortex_onpair::OnPair; @@ -51,30 +51,32 @@ use vortex_onpair::OnPairArray; use vortex_onpair::OnPairArraySlotsExt; /// Host-resident decode inputs, materialised once so the decode-loop benchmark -/// measures only `onpair::decompress_into` (not child `execute`/allocation). +/// measures only `onpair::decode_into` (not child `execute`/allocation). struct DecodeInputs { - dict_bytes: ByteBuffer, + dict_bytes: BufferHandle, dict_offsets: Buffer, codes: Buffer, - bits: u32, } impl DecodeInputs { - fn as_parts(&self) -> Parts<'_> { - Parts { - dict_bytes: self.dict_bytes.as_slice(), - dict_offsets: self.dict_offsets.as_slice(), - bits: self.bits, - codes: self.codes.as_slice(), + fn dict(&self) -> CompactDictionaryView<'_> { + // SAFETY: `materialise` validates this borrowed dictionary once after + // widening offsets. The benchmark then keeps both buffers immutable. + unsafe { + CompactDictionaryView::new_unchecked( + self.dict_bytes.as_host().as_slice(), + self.dict_offsets.as_slice(), + ) } } - fn decompressed_len(&self) -> usize { - onpair::decompressed_len(self.as_parts()) + fn decoded_len(&self) -> usize { + onpair::decoded_len(self.codes.as_slice(), self.dict()) } - fn decompress_into(&self, out: &mut [MaybeUninit]) -> usize { - onpair::decompress_into(self.as_parts(), out) + fn decode_into(&self, out: &mut [MaybeUninit]) -> usize { + // SAFETY: callers allocate `decoded_len + DECODE_PADDING` bytes. + unsafe { onpair::decode_into(self.codes.as_slice(), self.dict(), out) } } } use vortex_onpair::onpair_compress; @@ -175,13 +177,16 @@ fn widen(arr: &ArrayRef, ctx: &mut ExecutionCtx) -> Buffer { fn materialise(arr: &OnPairArray, ctx: &mut ExecutionCtx) -> (DecodeInputs, usize) { let view = arr.as_view(); + let dict_offsets = widen::(view.dict_offsets(), ctx); + let dict_bytes = view.dict_bytes_handle().clone(); + CompactDictionaryView::validate(dict_bytes.as_host().as_slice(), dict_offsets.as_slice()) + .expect("valid OnPair dictionary"); let inputs = DecodeInputs { - dict_bytes: view.dict_bytes().clone(), - dict_offsets: widen::(view.dict_offsets(), ctx), + dict_bytes, + dict_offsets, codes: widen::(view.codes(), ctx), - bits: view.bits(), }; - let total = inputs.decompressed_len(); + let total = inputs.decoded_len(); (inputs, total) } @@ -194,16 +199,16 @@ const CASES: &[(Shape, usize)] = &[ ]; /// Raw decode loop time, excluding child `execute` and the output allocation. -/// Hits `onpair::decompress_into` directly. +/// Hits `onpair::decode_into` directly. #[divan::bench(args = CASES)] -fn decompress_into_bench(bencher: Bencher, case: (Shape, usize)) { +fn decode_into_bench(bencher: Bencher, case: (Shape, usize)) { let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; let arr = compress(n, shape, &mut ctx); let (inputs, total) = materialise(&arr, &mut ctx); bencher.bench_local(|| { - let mut out: Vec = Vec::with_capacity(total); - let written = inputs.decompress_into(out.spare_capacity_mut()); + let mut out: Vec = Vec::with_capacity(total + onpair::DECODE_PADDING); + let written = inputs.decode_into(out.spare_capacity_mut()); unsafe { out.set_len(written) }; divan::black_box(out); }); diff --git a/encodings/experimental/onpair/goldenfiles/onpair.metadata b/encodings/experimental/onpair/goldenfiles/onpair.metadata index e96baf1a0ab..80b7fcdef71 100644 --- a/encodings/experimental/onpair/goldenfiles/onpair.metadata +++ b/encodings/experimental/onpair/goldenfiles/onpair.metadata @@ -1 +1 @@ - € €è(08 \ No newline at end of file +€ €è(08 \ No newline at end of file diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index 519243b100e..95287ef2114 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -52,9 +52,8 @@ pub type OnPairArray = Array; /// /// On disk the layout is FSST-shape: /// -/// * Buffer 0 — `dict_bytes`: the dictionary blob built by the OnPair trainer, -/// padded with `onpair::MAX_TOKEN_SIZE` trailing zero -/// bytes so the over-copy decoder can read 16 bytes past the last token. +/// * Buffer 0 — `dict_bytes`: the read-padded dictionary blob built by the +/// OnPair trainer. /// * Slots — see [`OnPairSlots`]. /// /// The four integer slot children flow through the standard `compress_child` @@ -66,12 +65,8 @@ pub struct OnPairMetadata { /// Width of the per-row primitive `uncompressed_lengths` child. #[prost(enumeration = "PType", tag = "1")] pub uncompressed_lengths_ptype: i32, - /// Bits-per-token the column was compressed with (9..=16). Every value - /// in the `codes` child only uses its low `bits` bits. - #[prost(uint32, tag = "2")] - pub bits: u32, /// Number of dictionary tokens. `dict_offsets` has length `dict_size + 1`. - /// Bounded by `2^bits ≤ 2^16 = 65_536`, so `u32` is comfortably wide. + /// Bounded by 65_536, so `u32` is comfortably wide. #[prost(uint32, tag = "3")] pub dict_size: u32, /// Total number of tokens across all rows. `codes` has this length; @@ -83,7 +78,7 @@ pub struct OnPairMetadata { #[prost(enumeration = "PType", tag = "5")] pub dict_offsets_ptype: i32, /// PType of the `codes` slot child (typically U16, may be narrowed to U8 - /// when `bits <= 8`). + /// when the dictionary has at most 256 tokens). #[prost(enumeration = "PType", tag = "6")] pub codes_ptype: i32, /// PType of the `codes_offsets` slot child. @@ -92,6 +87,7 @@ pub struct OnPairMetadata { } impl OnPairMetadata { + /// Decode the recorded [`PType`] of the `uncompressed_lengths` slot child. pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult { PType::try_from(self.uncompressed_lengths_ptype) .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype)) @@ -103,9 +99,8 @@ pub struct OnPairSlots { /// `PrimitiveArray`, length `dict_size + 1`. Cascading compressor may /// narrow the ptype to U16/U8. pub dict_offsets: ArrayRef, - /// `PrimitiveArray`. Each value only uses its low `bits` bits; - /// downstream `FastLanes::BitPacking` losslessly shrinks the child to - /// exactly `bits`-bit codes on disk. + /// Primitive integer token codes. Downstream integer compression may + /// narrow or bit-pack this child independently of the OnPair metadata. pub codes: ArrayRef, /// `PrimitiveArray`, length `num_rows + 1`. FoR / RunEnd / etc. apply /// naturally via the cascading compressor. @@ -127,46 +122,45 @@ pub struct OnPairSlots { pub struct OnPairData { /// The dictionary blob (buffer 0). /// - /// INVARIANT: this buffer must be over-padded past its logical end - /// (`dict_offsets.last()`) by the decoder's fixed token read width, - /// `onpair::MAX_TOKEN_SIZE`. The over-copy decoder reads - /// every dictionary entry with one fixed-width load and then advances the - /// cursor by the token's true length, so the load for the final, shortest - /// token over-reads past the logical end of the dictionary. This is the - /// same over-read the decoder accounts for on the final few codes; the - /// trailing padding absorbs it so that any entry can be read in bounds. - /// `onpair_compress` establishes this padding (see `parts_to_children`); - /// the over-copy decoder lives in the `onpair` crate. + /// INVARIANT: this buffer is an OnPair compact dictionary byte buffer, + /// including its trailing read padding. dict_bytes: BufferHandle, - bits: u32, + dict_size: u32, len: usize, } impl OnPairData { - pub fn new(dict_bytes: BufferHandle, bits: u32, len: usize) -> Self { + /// Build [`OnPairData`] from the dictionary blob, its token count, and the + /// number of rows. + pub fn new(dict_bytes: BufferHandle, dict_size: u32, len: usize) -> Self { Self { dict_bytes, - bits, + dict_size, len, } } + /// Number of rows in the array. pub fn len(&self) -> usize { self.len } + /// Whether the array has zero rows. pub fn is_empty(&self) -> bool { self.len == 0 } + /// Bits per token code, derived from the dictionary size (not stored). pub fn bits(&self) -> u32 { - self.bits + u32::from(onpair::code_bits_for_num_tokens(self.dict_size as usize)) } + /// The dictionary blob as a host byte buffer. pub fn dict_bytes(&self) -> &ByteBuffer { self.dict_bytes.as_host() } + /// The [`BufferHandle`] holding the dictionary blob (buffer 0). pub fn dict_bytes_handle(&self) -> &BufferHandle { &self.dict_bytes } @@ -176,9 +170,10 @@ impl Display for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( f, - "len: {}, bits: {}, dict_bytes_len: {}", + "len: {}, dict_size: {}, bits: {}, dict_bytes_len: {}", self.len, - self.bits, + self.dict_size, + self.bits(), self.dict_bytes.len() ) } @@ -188,7 +183,8 @@ impl Debug for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("OnPairData") .field("len", &self.len) - .field("bits", &self.bits) + .field("dict_size", &self.dict_size) + .field("bits", &self.bits()) .field("dict_bytes_len", &self.dict_bytes.len()) .finish() } @@ -197,13 +193,13 @@ impl Debug for OnPairData { impl ArrayHash for OnPairData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.dict_bytes.as_host().array_hash(state, accuracy); - state.write_u32(self.bits); + state.write_u32(self.dict_size); } } impl ArrayEq for OnPairData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.bits == other.bits + self.dict_size == other.dict_size && self .dict_bytes .as_host() @@ -217,7 +213,6 @@ pub struct OnPair; impl OnPair { /// Build an [`OnPairArray`] from already-materialised parts. - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] pub fn try_new( dtype: DType, dict_bytes: BufferHandle, @@ -226,18 +221,16 @@ impl OnPair { codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> VortexResult { - validate_parts( + let dict_size = validate_parts( &dtype, &dict_offsets, &codes, &codes_offsets, &uncompressed_lengths, - bits, )?; let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); + let data = OnPairData::new(dict_bytes, dict_size, len); let slots = OnPairSlots { dict_offsets, codes, @@ -251,7 +244,6 @@ impl OnPair { }) } - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] pub(crate) unsafe fn new_unchecked( dtype: DType, dict_bytes: BufferHandle, @@ -260,10 +252,13 @@ impl OnPair { codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> OnPairArray { let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); + let dict_size = match u32::try_from(dict_offsets.len().saturating_sub(1)) { + Ok(dict_size) => dict_size, + Err(_) => vortex_panic!("OnPair dict_size exceeds u32"), + }; + let data = OnPairData::new(dict_bytes, dict_size, len); let slots = OnPairSlots { dict_offsets, codes, @@ -284,13 +279,11 @@ fn validate_parts( codes: &ArrayRef, codes_offsets: &ArrayRef, uncompressed_lengths: &ArrayRef, - bits: u32, -) -> VortexResult<()> { +) -> VortexResult { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), "OnPair arrays must be Binary or Utf8, found {dtype}" ); - vortex_ensure!((9..=16).contains(&bits), "bits {bits} out of range [9, 16]"); if !dict_offsets.dtype().is_int() || dict_offsets.dtype().is_nullable() { vortex_bail!(InvalidArgument: "dict_offsets must be non-nullable integer"); @@ -304,6 +297,7 @@ fn validate_parts( if !uncompressed_lengths.dtype().is_int() || uncompressed_lengths.dtype().is_nullable() { vortex_bail!(InvalidArgument: "uncompressed_lengths must be non-nullable integer"); } + let dict_size = dictionary_size_from_offsets(dict_offsets.len())?; if codes_offsets.len() != uncompressed_lengths.len() + 1 { vortex_bail!(InvalidArgument: "codes_offsets.len ({}) != uncompressed_lengths.len + 1 ({})", @@ -311,7 +305,15 @@ fn validate_parts( uncompressed_lengths.len() + 1 ); } - Ok(()) + Ok(dict_size) +} + +fn dictionary_size_from_offsets(dict_offsets_len: usize) -> VortexResult { + let dict_size = dict_offsets_len + .checked_sub(1) + .ok_or_else(|| vortex_err!(InvalidArgument: "OnPair dict_offsets must not be empty"))?; + u32::try_from(dict_size) + .map_err(|_| vortex_err!(InvalidArgument: "OnPair dictionary size exceeds u32")) } impl VTable for OnPair { @@ -332,14 +334,19 @@ impl VTable for OnPair { slots: &[Option], ) -> VortexResult<()> { let s = OnPairSlotsView::from_slots(slots); - validate_parts( + let dict_size = validate_parts( dtype, s.dict_offsets, s.codes, s.codes_offsets, s.uncompressed_lengths, - data.bits, )?; + if data.dict_size != dict_size { + vortex_bail!(InvalidArgument: + "OnPairData dict_size {} does not match dict_offsets size {dict_size}", + data.dict_size + ); + } if s.uncompressed_lengths.len() != len { vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array"); } @@ -395,7 +402,6 @@ impl VTable for OnPair { Ok(Some( OnPairMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - bits: array.bits(), dict_size, total_tokens, dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(), @@ -467,7 +473,7 @@ impl VTable for OnPair { other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), }; - let data = OnPairData::new(buffers[0].clone(), metadata.bits, len); + let data = OnPairData::new(buffers[0].clone(), metadata.dict_size, len); let slots = OnPairSlots { dict_offsets, codes, @@ -534,6 +540,8 @@ impl ValidityVTable for OnPair { /// Convenience methods on top of the macro-generated [`OnPairArraySlotsExt`]. pub trait OnPairArrayExt: OnPairArraySlotsExt { + /// The array's [`Validity`], derived from the optional validity child and + /// the outer dtype's nullability. fn array_validity(&self) -> Validity { child_to_validity( self.as_ref().slots()[OnPairSlots::VALIDITY].as_ref(), diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 20a058e7f90..748a9a6dd7d 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -2,14 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors // //! Convert an [`OnPairArray`] to its canonical `VarBinViewArray` by handing -//! the materialised parts to `onpair::decompress_into`. +//! the materialised parts to `onpair::try_decode_into`. //! //! [`OnPairArray`]: crate::OnPairArray use std::sync::Arc; use num_traits::AsPrimitive; -use onpair::Parts; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -25,11 +25,13 @@ use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::OnPair; use crate::OnPairArraySlotsExt; use crate::decode::code_boundary_at; use crate::decode::collect_widened; +use crate::decode::validate_codes; pub(super) fn canonicalize_onpair( array: ArrayView<'_, OnPair>, @@ -90,19 +92,27 @@ pub(crate) fn onpair_decode_views( // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; + // The codes child is file-borne: reject out-of-range codes here so the + // decoder's panicking bounds check never fires on corrupt data. + validate_codes(codes.as_slice(), dict_offsets.len().saturating_sub(1))?; + let dict = + CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; - let mut out_bytes = ByteBufferMut::with_capacity(total_size); - let written = onpair::decompress_into( - Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }, - out_bytes.spare_capacity_mut(), + // `try_decode_into` derives its write bound from the buffer itself, so it + // is sound even when the file-borne `uncompressed_lengths` understate the + // real decoded size; the extra DECODE_PADDING merely keeps it on the + // all-over-copy fast path. + let mut out_bytes = ByteBufferMut::with_capacity(total_size + onpair::DECODE_PADDING); + let written = onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) + .map_err(|_| { + vortex_err!("OnPair codes decode to more bytes than uncompressed_lengths records") + })?; + vortex_ensure!( + written == total_size, + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" ); - debug_assert_eq!(written, total_size); - // SAFETY: `decompress_into` initialised exactly `written` bytes of the + // SAFETY: `try_decode_into` initialised exactly `written` bytes of the // spare capacity reserved above. unsafe { out_bytes.set_len(written) }; diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index 3616ec7e717..dbc27f908c6 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -8,11 +8,9 @@ use onpair::Offset; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; -use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -42,19 +40,6 @@ where { let len = array.len(); let mask = array.validity()?.execute_mask(len, ctx)?; - if mask.all_false() { - return OnPair::try_new( - array.dtype().clone(), - BufferHandle::new_host(ByteBuffer::empty()), - ConstantArray::new(0, len).into_array(), - ConstantArray::new(0u16, len).into_array(), - ConstantArray::new(0u32, len + 1).into_array(), - ConstantArray::new(0i32, len).into_array(), - Validity::AllInvalid, - 9, - ); - } - let mut flat: Vec = Vec::with_capacity(len * 16); let mut offsets: Vec = Vec::with_capacity(len + 1); let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); @@ -78,7 +63,10 @@ where } } AllOr::None => { - unreachable!("all_false() should have been caught earlier"); + offsets.resize(len + 1, O::from_usize(0)); + for _ in 0..len { + uncompressed_lengths.push(0); + } } AllOr::Some(validity) => { for (view, valid) in views.iter().zip(validity.iter()) { @@ -98,24 +86,32 @@ where let column = onpair::compress(&flat, &offsets, config) .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?; - let bits = column.bits; - let dict_bytes = dict_bytes_to_buffer(column.dict_bytes); - let codes_offsets = - build_codes_offsets(&column.codes, &column.dict_offsets, &offsets)?.into_array(); - let codes = Buffer::from(column.codes).into_array(); - let dict_offsets = Buffer::from(column.dict_offsets).into_array(); + 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::>>()?, + ) + .into_array(); + let codes = Buffer::from(codes).into_array(); + let dict_offsets = Buffer::from(dict_offsets).into_array(); let uncompressed_lengths = uncompressed_lengths.into_array(); OnPair::try_new( array.dtype().clone(), - dict_bytes, + dict_bytes_to_buffer(dict_bytes), dict_offsets, codes, codes_offsets, uncompressed_lengths, array.validity()?, - bits, ) } @@ -128,57 +124,13 @@ fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] } } -/// Lift compressed dictionary bytes into the Vortex buffer slot. fn dict_bytes_to_buffer(dict_bytes: Vec) -> BufferHandle { - // Pad the dictionary blob with MAX_TOKEN_SIZE zero bytes so the - // over-copy decoder can issue a fixed 16-byte load for every token - // without risking an OOB read on the last entry. - // // Align dict_bytes to 8 bytes so the segment that ultimately holds the - // OnPair tree starts at an 8-aligned in-memory address. Without this - // anchor, the per-buffer padding the serializer inserts is only - // *relative* to the segment start; if the segment lands at a u8-aligned - // heap address, downstream `PrimitiveArray::deserialize` panics - // with `Misaligned buffer cannot be used to build PrimitiveArray of u32`. - let mut padded = ByteBufferMut::with_capacity_aligned( - dict_bytes.len() + onpair::MAX_TOKEN_SIZE, - Alignment::new(8), - ); - padded.extend_from_slice(&dict_bytes); - unsafe { padded.push_n_unchecked(0, dict_bytes.len() + onpair::MAX_TOKEN_SIZE - padded.len()) }; - BufferHandle::new_host(padded.freeze()) -} - -/// Reconstruct the per-row `codes_offsets` from the flat `codes`, the -/// dictionary `dict_offsets` (token byte lengths) and the per-row decoded byte -/// boundaries. Returns `nrows + 1` cumulative code counts (`u32`). -// TODO(joe): can we compute this while compressing the array, yes but a worse API. -fn build_codes_offsets( - codes: &[u16], - dict_offsets: &[u32], - row_byte_offsets: &[O], -) -> VortexResult> { - let nrows = row_byte_offsets.len() - 1; - let mut codes_offsets = BufferMut::with_capacity(nrows + 1); - codes_offsets.push(0u32); - let mut decoded_bytes: u64 = 0; - let mut code_idx: usize = 0; - for r in 0..nrows { - let target = row_byte_offsets[r + 1] - .to_usize() - .ok_or_else(|| vortex_err!("OnPair row byte offset does not fit usize"))? - as u64; - while decoded_bytes < target { - let code = codes[code_idx] as usize; - decoded_bytes += u64::from(dict_offsets[code + 1] - dict_offsets[code]); - code_idx += 1; - } - codes_offsets.push( - u32::try_from(code_idx) - .map_err(|_| vortex_err!("OnPair: code boundary {code_idx} does not fit u32"))?, - ); - } - Ok(codes_offsets.freeze()) + // OnPair tree starts at an 8-aligned in-memory address. Without this anchor, + // downstream primitive children may deserialize from a misaligned segment. + let mut aligned = ByteBufferMut::with_capacity_aligned(dict_bytes.len(), Alignment::new(8)); + aligned.extend_from_slice(&dict_bytes); + BufferHandle::new_host(aligned.freeze()) } /// Compress any [`ArrayRef`] whose canonical form is a string array, by first diff --git a/encodings/experimental/onpair/src/compute/cast.rs b/encodings/experimental/onpair/src/compute/cast.rs index 93e1fdd8f8a..667a6cb5c57 100644 --- a/encodings/experimental/onpair/src/compute/cast.rs +++ b/encodings/experimental/onpair/src/compute/cast.rs @@ -3,16 +3,38 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::dtype::DType; +use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use crate::OnPair; use crate::OnPairArraySlotsExt; -/// Cast between `Utf8` and `Binary` (or adjust nullability) without touching -/// any of the encoded payload — we only rewrap into a new outer DType. +/// Adjust nullability without touching any encoded payload — we only rewrap +/// into a new outer DType. +fn build_with_validity( + array: ArrayView<'_, OnPair>, + dtype: &DType, + validity: Validity, +) -> ArrayRef { + unsafe { + OnPair::new_unchecked( + dtype.clone(), + array.dict_bytes_handle().clone(), + array.dict_offsets().clone(), + array.codes().clone(), + array.codes_offsets().clone(), + array.uncompressed_lengths().clone(), + validity, + ) + } + .into_array() +} + impl CastReduce for OnPair { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { if !array.dtype().eq_ignore_nullability(dtype) { @@ -24,20 +46,69 @@ impl CastReduce for OnPair { else { return Ok(None); }; - Ok(Some( - unsafe { - OnPair::new_unchecked( - dtype.clone(), - array.dict_bytes_handle().clone(), - array.dict_offsets().clone(), - array.codes().clone(), - array.codes_offsets().clone(), - array.uncompressed_lengths().clone(), - new_validity, - array.bits(), - ) - } - .into_array(), - )) + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +impl CastKernel for OnPair { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !array.dtype().eq_ignore_nullability(dtype) { + return Ok(None); + } + let new_validity = array.array().validity()?.cast_nullability( + dtype.nullability(), + array.array().len(), + ctx, + )?; + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinArray; + use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::compress::DEFAULT_DICT12_CONFIG; + use crate::compress::onpair_compress; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[rstest] + #[case(VarBinArray::from_iter( + vec![Some("hello"), Some("world"), Some("hello world")], + DType::Utf8(Nullability::NonNullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("foo"), None, Some("bar"), Some("foobar")], + DType::Utf8(Nullability::Nullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("test")], + DType::Utf8(Nullability::NonNullable) + ))] + fn test_cast_onpair_conformance(#[case] array: VarBinArray) -> VortexResult<()> { + let array = array.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let onpair = onpair_compress(&array, DEFAULT_DICT12_CONFIG, &mut ctx)?; + test_cast_conformance(&onpair.into_array(), &mut ctx); + Ok(()) } } diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index 939a1a2fb42..dcddc79a370 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use onpair::CompactDictionaryView; +use onpair::search; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -14,9 +16,12 @@ use vortex_array::scalar_fn::fns::binary::CompareKernel; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::decode::collect_codes_window; +use crate::decode::collect_widened; impl CompareKernel for OnPair { fn compare( @@ -28,29 +33,56 @@ impl CompareKernel for OnPair { let Some(constant) = rhs.as_constant() else { return Ok(None); }; - let is_empty = match constant.dtype() { - DType::Utf8(_) => constant.as_utf8().is_empty(), - DType::Binary(_) => constant.as_binary().is_empty(), + let needle = match constant.dtype() { + DType::Utf8(_) => constant + .as_utf8() + .value() + .map(|value| value.as_bytes().to_vec()), + DType::Binary(_) => constant + .as_binary() + .value() + .map(|value| value.as_slice().to_vec()), _ => return Ok(None), }; - if is_empty != Some(true) { + let Some(needle) = needle else { return Ok(None); - } - - let lengths = lhs.uncompressed_lengths(); - let buffer = match operator { - // every value is greater than an empty string - CompareOperator::Gte => BitBuffer::new_set(lhs.len()), - // no value is less than an empty string - CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), - _ => lengths - .binary( - ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) - .into_array(), - operator.into(), - )? - .execute(ctx)?, }; + + let buffer = if needle.is_empty() { + let lengths = lhs.uncompressed_lengths(); + match operator { + // every value is greater than an empty string + CompareOperator::Gte => BitBuffer::new_set(lhs.len()), + // no value is less than an empty string + CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), + _ => lengths + .binary( + ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) + .into_array(), + operator.into(), + )? + .execute(ctx)?, + } + } else { + if !matches!(operator, CompareOperator::Eq | CompareOperator::NotEq) { + return Ok(None); + } + + let dict_offsets = collect_widened::(lhs.dict_offsets(), ctx)?; + let dict = CompactDictionaryView::validate( + lhs.dict_bytes().as_slice(), + dict_offsets.as_slice(), + ) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + let query = search::tokenize(&needle, dict); + let window = collect_codes_window(lhs, ctx)?; + + let negated = operator == CompareOperator::NotEq; + BitBuffer::collect_bool(lhs.len(), |i| { + search::equals(window.row(i), &query) != negated + }) + }; + Ok(Some( BoolArray::new( buffer, @@ -69,7 +101,6 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; @@ -78,14 +109,22 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::binary::CompareKernel; + use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::OnPair; use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; - static SESSION: LazyLock = LazyLock::new(array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[cfg_attr(miri, ignore)] #[rstest] @@ -142,4 +181,104 @@ mod tests { ); Ok(()) } + + #[cfg_attr(miri, ignore)] + #[test] + fn compare_nonempty_string_nullable() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), None, Some("world"), Some("hello")], + DType::Utf8(Nullability::Nullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = arr + .binary(rhs.clone(), Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq, + &BoolArray::from_iter([Some(true), None, Some(false), Some(true)]), + &mut ctx + ); + + let neq = arr + .binary(rhs, Operator::NotEq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &neq, + &BoolArray::from_iter([Some(false), None, Some(true), Some(false)]), + &mut ctx + ); + Ok(()) + } + + /// Call `CompareKernel::compare` directly and verify it returns `Some` + /// (i.e. the kernel handles the constant needle rather than silently + /// falling back to canonical decompression), and that ordering against a + /// non-empty needle declines — code sequences are not order-preserving. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_handles_constant_needle() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), Some("world"), Some("hello"), Some("he")], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = + ::compare(arr.as_view(), &rhs, CompareOperator::Eq, &mut ctx)? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!( + eq, + BoolArray::from_iter([true, false, true, false]), + &mut ctx + ); + + let lt = + ::compare(arr.as_view(), &rhs, CompareOperator::Lt, &mut ctx)?; + assert!( + lt.is_none(), + "ordering compare against a non-empty needle must fall back to canonical" + ); + Ok(()) + } + + /// The equality path resolves row windows relative to `codes_offsets[0]`, + /// which is nonzero for a sliced array. Exercise the kernel directly on a + /// slice so the windowed row arithmetic is covered. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_on_sliced_array() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [ + Some("aardvark"), + Some("hello"), + Some("world"), + Some("hello"), + Some("zebra"), + ], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let sliced = arr.slice(1..4)?; + assert!(sliced.is::(), "slice dropped OnPair encoding"); + let sliced = sliced + .try_downcast::() + .map_err(|_| vortex_err!("sliced array was not OnPair"))?; + + let rhs = ConstantArray::new("hello", sliced.len()).into_array(); + let eq = ::compare( + sliced.as_view(), + &rhs, + CompareOperator::Eq, + &mut ctx, + )? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!(eq, BoolArray::from_iter([true, false, true]), &mut ctx); + Ok(()) + } } diff --git a/encodings/experimental/onpair/src/compute/filter.rs b/encodings/experimental/onpair/src/compute/filter.rs index c26f3eeeacd..623fde4e6eb 100644 --- a/encodings/experimental/onpair/src/compute/filter.rs +++ b/encodings/experimental/onpair/src/compute/filter.rs @@ -69,7 +69,6 @@ impl FilterKernel for OnPair { filtered_codes.offsets().clone(), uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/experimental/onpair/src/compute/slice.rs b/encodings/experimental/onpair/src/compute/slice.rs index fcfebf413bf..861c6661ed2 100644 --- a/encodings/experimental/onpair/src/compute/slice.rs +++ b/encodings/experimental/onpair/src/compute/slice.rs @@ -34,7 +34,6 @@ impl SliceReduce for OnPair { codes_offsets, uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs index fd0b02b0284..e0b21db7ce1 100644 --- a/encodings/experimental/onpair/src/decode.rs +++ b/encodings/experimental/onpair/src/decode.rs @@ -5,6 +5,7 @@ //! `onpair` decoder consumes. use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; @@ -12,8 +13,12 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use crate::OnPair; +use crate::OnPairArraySlotsExt; + /// Canonicalise a slot child to the decoder's native primitive width. pub(crate) fn collect_widened( arr: &ArrayRef, @@ -41,3 +46,83 @@ pub(crate) fn code_boundary_at( .as_::() .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) } + +/// Ensure every code indexes the dictionary (`code < num_tokens`). +/// +/// The slot children are file-borne, so a malformed `codes` child is a +/// recoverable error, not a bug: checking up front keeps the upstream +/// decoder/search primitives — which back their bounds checks with panics — +/// off file-borne data they would panic on. +pub(crate) fn validate_codes(codes: &[u16], num_tokens: usize) -> VortexResult<()> { + // `fold(max)` instead of `Iterator::max`: the latter's last-max-wins + // semantics defeat autovectorization, turning this scan scalar. + let max = codes.iter().fold(0u16, |acc, &c| acc.max(c)); + vortex_ensure!( + codes.is_empty() || (max as usize) < num_tokens, + "OnPair code {max} out of range for dictionary of {num_tokens} tokens" + ); + Ok(()) +} + +/// A validated, materialised window over an array's `codes`: the widened +/// per-row `codes_offsets` boundaries plus the codes they bound. +/// +/// `slice` keeps the full `codes` child and only narrows `codes_offsets`, so +/// for a sliced array the window starts at `offsets[0] > 0`; [`row`] resolves +/// row indices relative to that start. Built once per query by the +/// compressed-domain compare kernel. +/// +/// [`row`]: CodesWindow::row +pub(crate) struct CodesWindow { + offsets: Buffer, + codes: Buffer, + 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] + } +} + +/// Materialise and validate the [`CodesWindow`] for every row of `array`: +/// offsets must be nondecreasing and end within the `codes` child, and every +/// code must index the dictionary (see [`validate_codes`]). +pub(crate) fn collect_codes_window( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let offsets = collect_widened::(array.codes_offsets(), ctx)?; + vortex_ensure!( + offsets.len() == len + 1, + "OnPair codes_offsets has {} entries, expected len + 1 = {}", + offsets.len(), + len + 1 + ); + vortex_ensure!( + offsets.is_sorted(), + "OnPair codes_offsets must be nondecreasing" + ); + let code_start = offsets[0] as usize; + let code_end = offsets[len] as usize; + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + validate_codes( + codes.as_slice(), + array.dict_offsets().len().saturating_sub(1), + )?; + Ok(CodesWindow { + offsets, + codes, + code_start, + }) +} diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index 8863d750a72..e6ae93248b5 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -10,13 +10,15 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; +use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use vortex_session::VortexSession; use crate::OnPair; -// TODO: implement ListExecute & TakeExecute for OnPair pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), OnPair, CastExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Filter.id(), OnPair, FilterExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel( diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index 11b22f63bc1..e6aedf24183 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -2,10 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Vortex string array backed by the [OnPair][onpair] short-string -//! compression library, with `cast` and `filter` pushdown. +//! compression library, with pushdown for common string operations. //! -//! The default training preset is `dict-12` (12 bits per token, dictionary -//! capped at 4 096 entries). See [`onpair_compress`] for the entry point and +//! The default training preset is `dict-12`: the trainer may build a dictionary +//! with up to 4 096 tokens, while runtime code width is derived from the actual +//! dictionary size. See [`onpair_compress`] for the entry point and //! [`OnPairArray`] for the resulting array type. //! //! [onpair]: https://arxiv.org/abs/2508.02280 @@ -23,9 +24,9 @@ mod tests; pub use array::*; pub use compress::*; -pub use onpair::Bits; pub use onpair::Config; pub use onpair::Error as OnPairError; +pub use onpair::MaxDictBits; pub use onpair::Threshold; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index a6e097bbfd4..fce1dadf3db 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use onpair::Parts; +use onpair::CompactDictionaryView; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::varbin::varbin_scalar; @@ -9,12 +9,14 @@ use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::OnPair; use crate::OnPairArraySlotsExt; use crate::decode::code_boundary_at; use crate::decode::collect_widened; +use crate::decode::validate_codes; impl OperationsVTable for OnPair { fn scalar_at( @@ -33,12 +35,12 @@ impl OperationsVTable for OnPair { let codes = collect_widened::(&array.codes().slice(row_start..row_end)?, ctx)?; let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - let parts = Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }; + // The codes child is file-borne: reject out-of-range codes here so the + // decoder's panicking bounds check never fires on corrupt data. + validate_codes(codes.as_slice(), dict_offsets.len().saturating_sub(1))?; + let dict = + CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; // The per-row decoded length is recorded in the `uncompressed_lengths` // child, so read it directly instead of asking the decoder to compute it. @@ -48,10 +50,22 @@ impl OperationsVTable for OnPair { .as_primitive() .as_::() .ok_or_else(|| vortex_err!("OnPair uncompressed_lengths[{index}] is null"))?; - let mut buf: Vec = Vec::with_capacity(len); - let written = onpair::decompress_into(parts, buf.spare_capacity_mut()); - debug_assert_eq!(written, len); - // SAFETY: `decompress_into` initialised `written` bytes of the spare + // `try_decode_into` derives its write bound from the buffer itself, so + // it is sound even when the file-borne `uncompressed_lengths` child + // understates the row's real decoded size; the extra DECODE_PADDING + // merely keeps it on the all-over-copy fast path. + let mut buf: Vec = Vec::with_capacity(len + onpair::DECODE_PADDING); + let written = onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) + .map_err(|_| { + vortex_err!( + "OnPair row {index} decodes to more bytes than uncompressed_lengths records" + ) + })?; + vortex_ensure!( + written == len, + "OnPair row {index} decoded to {written} bytes but uncompressed_lengths records {len}" + ); + // SAFETY: `try_decode_into` initialised `written` bytes of the spare // capacity reserved above. unsafe { buf.set_len(written) }; Ok(varbin_scalar(ByteBuffer::from(buf), array.dtype())) diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index 4f2124025f2..c2acfe9f02c 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -3,6 +3,7 @@ use std::sync::LazyLock; +use onpair::CompactDictionaryView; use prost::Message; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -41,6 +42,31 @@ fn sample_input() -> VarBinArray { ) } +#[test] +fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> { + let num_tokens = 100_000usize; + let mut tokens = Vec::with_capacity(num_tokens); + tokens.extend((u8::MIN..=u8::MAX).map(|byte| vec![byte])); + let additional_tokens = u32::try_from(num_tokens - tokens.len())?; + tokens.extend((0..additional_tokens).map(|value| { + let bytes = value.to_be_bytes(); + bytes[1..].to_vec() + })); + tokens.sort_unstable(); + + let mut dict_bytes = Vec::new(); + let mut dict_offsets = Vec::with_capacity(num_tokens + 1); + dict_offsets.push(0u32); + for token in tokens { + dict_bytes.extend_from_slice(&token); + dict_offsets.push(u32::try_from(dict_bytes.len())?); + } + dict_bytes.resize(dict_bytes.len() + onpair::MAX_TOKEN_SIZE, 0); + + assert!(CompactDictionaryView::validate(&dict_bytes, &dict_offsets).is_err()); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_metadata_golden() { @@ -48,7 +74,6 @@ fn test_onpair_metadata_golden() { "onpair.metadata", &OnPairMetadata { uncompressed_lengths_ptype: PType::I32 as i32, - bits: 12, dict_size: 4096, total_tokens: 128_000, dict_offsets_ptype: PType::U32 as i32, @@ -315,7 +340,6 @@ fn narrow_codes_offsets(arr: &crate::OnPairArray, target: PType) -> crate::OnPai narrowed_array, view.uncompressed_lengths().clone(), view.array_validity(), - view.bits(), ) } } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index f0a8966089c..eb2b6680abf 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -49,7 +49,7 @@ impl Scheme for OnPairScheme { /// 4 primitive slot children flow through the cascading compressor: /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 → - /// `FastLanes::BitPacked` to exactly `bits` = 12 by default), + /// usually `FastLanes::BitPacked` after scheme selection), /// `codes_offsets` (u32 → `FoR`), `uncompressed_lengths` (i32 → narrow /// + `FoR`). Validity stays untouched. fn num_children(&self) -> usize { @@ -116,7 +116,6 @@ impl Scheme for OnPairScheme { codes_offsets, uncompressed_lengths, onpair_array.array_validity(), - onpair_array.bits(), )? .into_array()) } From ef8b7e923afac8988d441305041efc825c81b2b4 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 10:09:58 +0100 Subject: [PATCH 02/13] test(vortex-onpair): cover all-null compression Signed-off-by: Francesco Gargiulo --- encodings/experimental/onpair/src/tests.rs | 29 ++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index c2acfe9f02c..4f3e0fa286e 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -11,6 +11,7 @@ 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::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -245,6 +246,34 @@ fn test_onpair_empty() -> vortex_error::VortexResult<()> { Ok(()) } +/// All-null input has no training values, but still produces a valid OnPair +/// array backed by an empty code stream and `len + 1` zero boundaries. +#[cfg_attr(miri, ignore)] +#[test] +fn test_onpair_all_null() -> vortex_error::VortexResult<()> { + let input = VarBinArray::from_iter( + [None::<&str>, None, None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&input, DEFAULT_DICT12_CONFIG, &mut ctx)?; + + assert!(arr.codes().is_empty()); + let codes_offsets = arr + .codes_offsets() + .clone() + .execute::(&mut ctx)?; + assert_eq!(codes_offsets.as_slice::(), &[0, 0, 0, 0]); + let uncompressed_lengths = arr + .uncompressed_lengths() + .clone() + .execute::(&mut ctx)?; + assert_eq!(uncompressed_lengths.as_slice::(), &[0, 0, 0]); + assert_arrays_eq!(arr.into_array(), input, &mut ctx); + Ok(()) +} + /// Filter must share the dictionary — never recompress (this is the /// regression cause on TPC-H Q22 SF=10). Exercise both selectivities /// and check that the result is bit-exact and still an OnPairArray. From 87adef86fedec7c315a17b1e38d85986355182ee Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 11:59:40 +0100 Subject: [PATCH 03/13] fix(vortex-onpair): remove redundant validate_codes pre-check OnPair's own decoder already bounds-checks every code in-loop before copying (a near-free, predicted-never-taken branch); vortex-onpair's validate_codes was a second full pass over the codes buffer ahead of decode, doubling memory traffic on a hot path for no additional soundness (out-of-range codes still surface, as a panic instead of a VortexResult error, matching the behavior prior to the onpair 0.1.0 port). Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- .../experimental/onpair/src/canonical.rs | 4 --- encodings/experimental/onpair/src/decode.rs | 28 +++---------------- encodings/experimental/onpair/src/ops.rs | 4 --- 3 files changed, 4 insertions(+), 32 deletions(-) diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 748a9a6dd7d..9fa2a8ed84f 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -31,7 +31,6 @@ use crate::OnPair; use crate::OnPairArraySlotsExt; use crate::decode::code_boundary_at; use crate::decode::collect_widened; -use crate::decode::validate_codes; pub(super) fn canonicalize_onpair( array: ArrayView<'_, OnPair>, @@ -92,9 +91,6 @@ pub(crate) fn onpair_decode_views( // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - // The codes child is file-borne: reject out-of-range codes here so the - // decoder's panicking bounds check never fires on corrupt data. - validate_codes(codes.as_slice(), dict_offsets.len().saturating_sub(1))?; let dict = CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs index e0b21db7ce1..1edad930fab 100644 --- a/encodings/experimental/onpair/src/decode.rs +++ b/encodings/experimental/onpair/src/decode.rs @@ -47,23 +47,6 @@ pub(crate) fn code_boundary_at( .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) } -/// Ensure every code indexes the dictionary (`code < num_tokens`). -/// -/// The slot children are file-borne, so a malformed `codes` child is a -/// recoverable error, not a bug: checking up front keeps the upstream -/// decoder/search primitives — which back their bounds checks with panics — -/// off file-borne data they would panic on. -pub(crate) fn validate_codes(codes: &[u16], num_tokens: usize) -> VortexResult<()> { - // `fold(max)` instead of `Iterator::max`: the latter's last-max-wins - // semantics defeat autovectorization, turning this scan scalar. - let max = codes.iter().fold(0u16, |acc, &c| acc.max(c)); - vortex_ensure!( - codes.is_empty() || (max as usize) < num_tokens, - "OnPair code {max} out of range for dictionary of {num_tokens} tokens" - ); - Ok(()) -} - /// A validated, materialised window over an array's `codes`: the widened /// per-row `codes_offsets` boundaries plus the codes they bound. /// @@ -88,9 +71,10 @@ impl CodesWindow { } } -/// Materialise and validate the [`CodesWindow`] for every row of `array`: -/// offsets must be nondecreasing and end within the `codes` child, and every -/// code must index the dictionary (see [`validate_codes`]). +/// Materialise the [`CodesWindow`] for every row of `array`: offsets must be +/// nondecreasing and end within the `codes` child. Codes themselves are +/// trusted to the upstream decoder/search primitives, which bounds-check them +/// in-loop and panic on a malformed value. pub(crate) fn collect_codes_window( array: ArrayView<'_, OnPair>, ctx: &mut ExecutionCtx, @@ -116,10 +100,6 @@ pub(crate) fn collect_codes_window( array.codes().len() ); let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - validate_codes( - codes.as_slice(), - array.dict_offsets().len().saturating_sub(1), - )?; Ok(CodesWindow { offsets, codes, diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index fce1dadf3db..969c7ec15eb 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -16,7 +16,6 @@ use crate::OnPair; use crate::OnPairArraySlotsExt; use crate::decode::code_boundary_at; use crate::decode::collect_widened; -use crate::decode::validate_codes; impl OperationsVTable for OnPair { fn scalar_at( @@ -35,9 +34,6 @@ impl OperationsVTable for OnPair { let codes = collect_widened::(&array.codes().slice(row_start..row_end)?, ctx)?; let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - // The codes child is file-borne: reject out-of-range codes here so the - // decoder's panicking bounds check never fires on corrupt data. - validate_codes(codes.as_slice(), dict_offsets.len().saturating_sub(1))?; let dict = CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; From ab329f7facb2186f54ab27ba07a96dc074e1e381 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 12:09:39 +0100 Subject: [PATCH 04/13] docs(vortex-onpair): clarify codes_ptype U8 narrowing condition The codes slot narrows to U8 only when dict_size is exactly 256, which is also the minimum size for a valid (complete) OnPair dictionary. The prior "at most 256 tokens" phrasing implied a range below 256 that can never occur. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/src/array.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index 95287ef2114..b683035eb0b 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -77,8 +77,10 @@ pub struct OnPairMetadata { /// narrowed to U16/U8 by the cascading compressor when values fit). #[prost(enumeration = "PType", tag = "5")] pub dict_offsets_ptype: i32, - /// PType of the `codes` slot child (typically U16, may be narrowed to U8 - /// when the dictionary has at most 256 tokens). + /// PType of the `codes` slot child (typically U16; U8 only when the + /// dictionary has exactly 256 tokens — its minimum, since a valid OnPair + /// dictionary must contain all 256 single-byte tokens, so codes then fit + /// in 8 bits). #[prost(enumeration = "PType", tag = "6")] pub codes_ptype: i32, /// PType of the `codes_offsets` slot child. From f838da3f6ba8a939af9847c22f6f41998d6eb57b Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 12:12:32 +0100 Subject: [PATCH 05/13] docs(vortex-onpair): clarify bits() is minimum addressing width bits() returns ceil(log2(dict_size)), the minimum code width to address the dictionary, not the physical per-code storage width. The codes child is stored as a U8/U16 primitive array and may be re-encoded by the cascading compressor, so "Bits per token code" was misleading. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/src/array.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index b683035eb0b..a3bfa3df34e 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -152,7 +152,11 @@ impl OnPairData { self.len == 0 } - /// Bits per token code, derived from the dictionary size (not stored). + /// Minimum bits needed to address the dictionary (`ceil(log2(dict_size))`), + /// derived from the dictionary size rather than stored. This is a logical + /// lower bound, not the physical width of the `codes` child — that child is + /// stored as a U8/U16 primitive array and may be independently re-encoded by + /// the cascading compressor. pub fn bits(&self) -> u32 { u32::from(onpair::code_bits_for_num_tokens(self.dict_size as usize)) } From 3bb9dc73e06ec21d5f07b7534c32267d4776d005 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 12:17:32 +0100 Subject: [PATCH 06/13] docs(vortex-onpair): simplify decode buffer padding comments The previous comments described try_decode_into's safe, buffer-derived bound and fast-path internals, which are implementation-specific and would not hold if the decode path switches to unsafe decode_into. Keep only the durable point: pad the output for the decoder's per-token over-copy and verify the exact size afterward. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/src/canonical.rs | 6 ++---- encodings/experimental/onpair/src/ops.rs | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 9fa2a8ed84f..231e8e81d0f 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -95,10 +95,8 @@ pub(crate) fn onpair_decode_views( CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; - // `try_decode_into` derives its write bound from the buffer itself, so it - // is sound even when the file-borne `uncompressed_lengths` understate the - // real decoded size; the extra DECODE_PADDING merely keeps it on the - // all-over-copy fast path. + // Pad the output with DECODE_PADDING to absorb the decoder's fixed + // per-token over-copy; the exact decoded size is checked below. let mut out_bytes = ByteBufferMut::with_capacity(total_size + onpair::DECODE_PADDING); let written = onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) .map_err(|_| { diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index 969c7ec15eb..9ead98be6cc 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -46,10 +46,8 @@ impl OperationsVTable for OnPair { .as_primitive() .as_::() .ok_or_else(|| vortex_err!("OnPair uncompressed_lengths[{index}] is null"))?; - // `try_decode_into` derives its write bound from the buffer itself, so - // it is sound even when the file-borne `uncompressed_lengths` child - // understates the row's real decoded size; the extra DECODE_PADDING - // merely keeps it on the all-over-copy fast path. + // Pad the row buffer with DECODE_PADDING to absorb the decoder's fixed + // per-token over-copy; the exact decoded size is checked below. let mut buf: Vec = Vec::with_capacity(len + onpair::DECODE_PADDING); let written = onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) .map_err(|_| { From e5705e807191471bb9ba41082f75a0b868fd1070 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 14:19:15 +0100 Subject: [PATCH 07/13] 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 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 Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/src/array.rs | 5 +- encodings/experimental/onpair/src/compress.rs | 68 ++++++++++++++++--- encodings/experimental/onpair/src/decode.rs | 19 ++++-- encodings/experimental/onpair/src/tests.rs | 65 ++++++++++++++++++ 4 files changed, 137 insertions(+), 20 deletions(-) diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index a3bfa3df34e..a502bb20761 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/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`, length `num_rows + 1`. FoR / RunEnd / etc. apply - /// naturally via the cascading compressor. + /// `PrimitiveArray` (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. diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index dbc27f908c6..6f4bf389273 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/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::>>()?, - ) - .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) -> 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(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::>(), + ) + .into_array() + } else { + Buffer::from( + row_offsets + .iter() + .map(|&o| u64::try_from(o.to_usize()).vortex_expect("token count fits u64")) + .collect::>(), + ) + .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::(ctx)?; onpair_compress_varbinview::(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::(&[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::(&[0, 3, 8], 7); + assert_eq!( + wide.dtype(), + &DType::Primitive(PType::U64, Nullability::NonNullable) + ); + } +} diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs index 1edad930fab..cc24f4669f4 100644 --- a/encodings/experimental/onpair/src/decode.rs +++ b/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,7 +58,7 @@ pub(crate) fn code_boundary_at( /// /// [`row`]: CodesWindow::row pub(crate) struct CodesWindow { - offsets: Buffer, + offsets: Buffer, codes: Buffer, code_start: usize, } @@ -65,9 +66,13 @@ pub(crate) struct CodesWindow { 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 { let len = array.len(); - let offsets = collect_widened::(array.codes_offsets(), ctx)?; + let offsets = collect_widened::(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 {}", diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index 4f3e0fa286e..7a43740a1f6 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/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::(&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::(&mut ctx)?; + let wide_decoded = wide + .clone() + .into_array() + .execute::(&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::(&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<()> { From d43343e5268eaf0efe44012d08dd00fde38c0374 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 14:34:46 +0100 Subject: [PATCH 08/13] chore(deps): keep a version alongside the onpair git pin The onpair dep was a bare git rev with no `version`, which fails the publish dry-run CI job (Cargo requires a version to package/publish). Add `version = "0.1.0"` alongside the pin; Cargo only checks a version is present, so local builds still resolve the pinned commit. Matches the tpchgen pattern below. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 29145f8b4e4..c9d5efbc3dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -192,7 +192,7 @@ num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } once_cell = "1.21" oneshot = { version = "0.2.0", features = ["async"] } -onpair = { git = "https://github.com/spiraldb/onpair.git", rev = "28df115c799674c663adedc8f81620ac82c7a890" } +onpair = { version = "0.1.0", git = "https://github.com/spiraldb/onpair.git", rev = "28df115c799674c663adedc8f81620ac82c7a890" } opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" opentelemetry_sdk = "0.32.0" From 752e76f2c3d4aa57a302537f84cafeaf404f1a0f Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Thu, 16 Jul 2026 15:24:38 +0100 Subject: [PATCH 09/13] docs(vortex-onpair): fix redundant rustdoc link Signed-off-by: Francesco Gargiulo --- encodings/experimental/onpair/src/decode.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs index cc24f4669f4..96add100e36 100644 --- a/encodings/experimental/onpair/src/decode.rs +++ b/encodings/experimental/onpair/src/decode.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors // -//! Helpers for turning [`OnPair`](super::OnPair) slot children into the inputs the upstream +//! Helpers for turning [`OnPair`] slot children into the inputs the upstream //! `onpair` decoder consumes. use vortex_array::ArrayRef; From fc335b48c93b923f9aa0ee51bfb3b02d1da1bc98 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Fri, 17 Jul 2026 13:08:09 +0100 Subject: [PATCH 10/13] refactor(vortex-onpair): slim OnPair metadata and return Constant for all-null Drop derived state and redundant metadata from the OnPair representation: - Remove `dict_size`/`bits()` from `OnPairData`; the dictionary size is recoverable from `dict_offsets.len() - 1`, so storing it (and validating it against the child) was redundant. - Rename the `total_tokens` metadata field to `codes_len`, which is what it actually records: the length of the `codes` slot child (a sliced array may retain codes outside its visible row range). - Have `onpair_compress` return `ArrayRef` and emit a `ConstantArray` of nulls for all-null input rather than a degenerate empty-code-stream OnPair array. Callers (btrblocks, benches, tests) downcast to `OnPair` as needed. - Drop the generic `Offset` parameter from compression; flatten to `u64` byte offsets unconditionally and size the flat buffer from the actual view lengths. `codes_offsets` width selection keeps the u32/u64 adaptivity. Decode paths (`canonical`, `ops`) no longer over-allocate `DECODE_PADDING` and `vortex_panic!` on a decoded-length mismatch instead of returning an error: the length is an internal invariant, not attacker-controlled input, so a mismatch is a bug rather than a recoverable condition. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/README.md | 20 ++- .../experimental/onpair/benches/decode.rs | 2 + encodings/experimental/onpair/src/array.rs | 109 +++++------------ .../experimental/onpair/src/canonical.rs | 28 +++-- encodings/experimental/onpair/src/compress.rs | 114 ++++++------------ .../onpair/src/compute/compare.rs | 4 +- encodings/experimental/onpair/src/lib.rs | 6 +- encodings/experimental/onpair/src/ops.rs | 29 +++-- encodings/experimental/onpair/src/tests.rs | 43 +++---- .../experimental/onpair/tests/big_data.rs | 5 +- vortex-btrblocks/src/schemes/string/onpair.rs | 5 +- 11 files changed, 139 insertions(+), 226 deletions(-) diff --git a/encodings/experimental/onpair/README.md b/encodings/experimental/onpair/README.md index c2419dd0c36..f19697a412d 100644 --- a/encodings/experimental/onpair/README.md +++ b/encodings/experimental/onpair/README.md @@ -17,22 +17,18 @@ decompression. ## Default Configuration -The default training preset is **dict-12**: the OnPair trainer may build a -dictionary with up to 4 096 tokens. After compression, the runtime code width is -derived from the actual dictionary size. Vortex stores token codes as an integer -child array; downstream integer compression may narrow or bit-pack that child -independently of the OnPair metadata. +The default training configuration uses OnPair's default dictionary budget and +a fixed seed. Vortex stores token codes as an integer child array; downstream +integer compression may narrow or bit-pack that child independently. ## Layout - Buffer 0 — `dict_bytes`: dictionary blob built by the OnPair trainer, - padded with `MAX_TOKEN_SIZE` trailing zero bytes so the over-copy - decoder can read 16 bytes past the last token. -- Slot 0 — `dict_offsets`: `PrimitiveArray`, len `dict_size + 1`. -- Slot 1 — `codes`: integer `PrimitiveArray`, length `total_tokens`. -- Slot 2 — `codes_offsets`: `PrimitiveArray`, length `num_rows + 1`. -- Slot 3 — `uncompressed_lengths`: integer `PrimitiveArray`, length - `num_rows`. + including the read padding required by the decoder. +- Slot 0 — `dict_offsets`: integer child, len `dict_size + 1`. +- Slot 1 — `codes`: integer child, length `total_tokens`. +- Slot 2 — `codes_offsets`: integer child, length `num_rows + 1`. +- Slot 3 — `uncompressed_lengths`: integer child, length `num_rows`. - Slot 4 — optional validity child. All four integer slot children flow through the standard cascading diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index e1055bc1eb3..8a0f54bc2d8 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -164,6 +164,8 @@ fn compress(n: usize, shape: Shape, ctx: &mut ExecutionCtx) -> OnPairArray { ); onpair_compress(varbin.as_array(), DEFAULT_DICT12_CONFIG, ctx) .unwrap_or_else(|e| panic!("onpair_compress failed: {e}")) + .try_downcast::() + .unwrap_or_else(|array| panic!("expected OnPair array, got {}", array.encoding_id())) } /// Canonicalise a slot child to the decoder's native primitive width. diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index a502bb20761..4f3503ba213 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -66,21 +66,17 @@ pub struct OnPairMetadata { #[prost(enumeration = "PType", tag = "1")] pub uncompressed_lengths_ptype: i32, /// Number of dictionary tokens. `dict_offsets` has length `dict_size + 1`. - /// Bounded by 65_536, so `u32` is comfortably wide. #[prost(uint32, tag = "3")] pub dict_size: u32, - /// Total number of tokens across all rows. `codes` has this length; - /// `codes_offsets.last() == total_tokens`. + /// Length of the `codes` slot child. A sliced array may retain codes that + /// fall outside its visible row range. #[prost(uint64, tag = "4")] - pub total_tokens: u64, + pub codes_len: u64, /// PType of the `dict_offsets` slot child (defaults to U32, may be /// narrowed to U16/U8 by the cascading compressor when values fit). #[prost(enumeration = "PType", tag = "5")] pub dict_offsets_ptype: i32, - /// PType of the `codes` slot child (typically U16; U8 only when the - /// dictionary has exactly 256 tokens — its minimum, since a valid OnPair - /// dictionary must contain all 256 single-byte tokens, so codes then fit - /// in 8 bits). + /// PType of the `codes` slot child. #[prost(enumeration = "PType", tag = "6")] pub codes_ptype: i32, /// PType of the `codes_offsets` slot child. @@ -98,18 +94,17 @@ impl OnPairMetadata { #[array_slots(OnPair)] pub struct OnPairSlots { - /// `PrimitiveArray`, length `dict_size + 1`. Cascading compressor may - /// narrow the ptype to U16/U8. + /// Primitive integer dictionary offsets, length `dict_size + 1`. The + /// cascading compressor may re-encode this child independently. pub dict_offsets: ArrayRef, /// Primitive integer token codes. Downstream integer compression may /// narrow or bit-pack this child independently of the OnPair metadata. pub codes: ArrayRef, - /// `PrimitiveArray` (or `u64` when a chunk exceeds `u32::MAX` tokens), - /// length `num_rows + 1`. FoR / RunEnd / etc. apply naturally via the - /// cascading compressor. + /// Primitive integer row offsets into `codes`, length `num_rows + 1`. The + /// cascading compressor may re-encode this child independently. pub codes_offsets: ArrayRef, - /// Integer `PrimitiveArray`, length `num_rows`. Used to size the canonical - /// output buffer. + /// Integer decoded-length child, length `num_rows`. Used to size the + /// canonical output buffer. pub uncompressed_lengths: ArrayRef, /// Optional validity child for the outer string column. pub validity: Option, @@ -128,19 +123,13 @@ pub struct OnPairData { /// INVARIANT: this buffer is an OnPair compact dictionary byte buffer, /// including its trailing read padding. dict_bytes: BufferHandle, - dict_size: u32, len: usize, } impl OnPairData { - /// Build [`OnPairData`] from the dictionary blob, its token count, and the - /// number of rows. - pub fn new(dict_bytes: BufferHandle, dict_size: u32, len: usize) -> Self { - Self { - dict_bytes, - dict_size, - len, - } + /// Build [`OnPairData`] from the dictionary blob and the number of rows. + pub fn new(dict_bytes: BufferHandle, len: usize) -> Self { + Self { dict_bytes, len } } /// Number of rows in the array. @@ -153,15 +142,6 @@ impl OnPairData { self.len == 0 } - /// Minimum bits needed to address the dictionary (`ceil(log2(dict_size))`), - /// derived from the dictionary size rather than stored. This is a logical - /// lower bound, not the physical width of the `codes` child — that child is - /// stored as a U8/U16 primitive array and may be independently re-encoded by - /// the cascading compressor. - pub fn bits(&self) -> u32 { - u32::from(onpair::code_bits_for_num_tokens(self.dict_size as usize)) - } - /// The dictionary blob as a host byte buffer. pub fn dict_bytes(&self) -> &ByteBuffer { self.dict_bytes.as_host() @@ -177,10 +157,8 @@ impl Display for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( f, - "len: {}, dict_size: {}, bits: {}, dict_bytes_len: {}", + "len: {}, dict_bytes_len: {}", self.len, - self.dict_size, - self.bits(), self.dict_bytes.len() ) } @@ -190,8 +168,6 @@ impl Debug for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("OnPairData") .field("len", &self.len) - .field("dict_size", &self.dict_size) - .field("bits", &self.bits()) .field("dict_bytes_len", &self.dict_bytes.len()) .finish() } @@ -200,17 +176,14 @@ impl Debug for OnPairData { impl ArrayHash for OnPairData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.dict_bytes.as_host().array_hash(state, accuracy); - state.write_u32(self.dict_size); } } impl ArrayEq for OnPairData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.dict_size == other.dict_size - && self - .dict_bytes - .as_host() - .array_eq(other.dict_bytes.as_host(), accuracy) + self.dict_bytes + .as_host() + .array_eq(other.dict_bytes.as_host(), accuracy) } } @@ -229,7 +202,7 @@ impl OnPair { uncompressed_lengths: ArrayRef, validity: Validity, ) -> VortexResult { - let dict_size = validate_parts( + validate_parts( &dtype, &dict_offsets, &codes, @@ -237,7 +210,7 @@ impl OnPair { &uncompressed_lengths, )?; let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, dict_size, len); + let data = OnPairData::new(dict_bytes, len); let slots = OnPairSlots { dict_offsets, codes, @@ -261,11 +234,7 @@ impl OnPair { validity: Validity, ) -> OnPairArray { let len = uncompressed_lengths.len(); - let dict_size = match u32::try_from(dict_offsets.len().saturating_sub(1)) { - Ok(dict_size) => dict_size, - Err(_) => vortex_panic!("OnPair dict_size exceeds u32"), - }; - let data = OnPairData::new(dict_bytes, dict_size, len); + let data = OnPairData::new(dict_bytes, len); let slots = OnPairSlots { dict_offsets, codes, @@ -286,7 +255,7 @@ fn validate_parts( codes: &ArrayRef, codes_offsets: &ArrayRef, uncompressed_lengths: &ArrayRef, -) -> VortexResult { +) -> VortexResult<()> { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), "OnPair arrays must be Binary or Utf8, found {dtype}" @@ -304,7 +273,6 @@ fn validate_parts( if !uncompressed_lengths.dtype().is_int() || uncompressed_lengths.dtype().is_nullable() { vortex_bail!(InvalidArgument: "uncompressed_lengths must be non-nullable integer"); } - let dict_size = dictionary_size_from_offsets(dict_offsets.len())?; if codes_offsets.len() != uncompressed_lengths.len() + 1 { vortex_bail!(InvalidArgument: "codes_offsets.len ({}) != uncompressed_lengths.len + 1 ({})", @@ -312,15 +280,7 @@ fn validate_parts( uncompressed_lengths.len() + 1 ); } - Ok(dict_size) -} - -fn dictionary_size_from_offsets(dict_offsets_len: usize) -> VortexResult { - let dict_size = dict_offsets_len - .checked_sub(1) - .ok_or_else(|| vortex_err!(InvalidArgument: "OnPair dict_offsets must not be empty"))?; - u32::try_from(dict_size) - .map_err(|_| vortex_err!(InvalidArgument: "OnPair dictionary size exceeds u32")) + Ok(()) } impl VTable for OnPair { @@ -341,19 +301,13 @@ impl VTable for OnPair { slots: &[Option], ) -> VortexResult<()> { let s = OnPairSlotsView::from_slots(slots); - let dict_size = validate_parts( + validate_parts( dtype, s.dict_offsets, s.codes, s.codes_offsets, s.uncompressed_lengths, )?; - if data.dict_size != dict_size { - vortex_bail!(InvalidArgument: - "OnPairData dict_size {} does not match dict_offsets size {dict_size}", - data.dict_size - ); - } if s.uncompressed_lengths.len() != len { vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array"); } @@ -405,12 +359,12 @@ impl VTable for OnPair { ) -> VortexResult>> { let dict_size = u32::try_from(array.dict_offsets().len().saturating_sub(1)) .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; - let total_tokens = array.codes().len() as u64; + let codes_len = array.codes().len() as u64; Ok(Some( OnPairMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), dict_size, - total_tokens, + codes_len, dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(), codes_ptype: array.codes().dtype().as_ptype().into(), codes_offsets_ptype: array.codes_offsets().dtype().as_ptype().into(), @@ -434,12 +388,11 @@ impl VTable for OnPair { let metadata = OnPairMetadata::decode(metadata)?; let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?; - // Slot children. We pass `usize::MAX` for slots whose length we - // don't know up front (`dict_offsets` and `codes`). `codes_offsets` - // has known length `len + 1`. + // Slot children do not persist their own lengths, so metadata records + // the dictionary and code-stream sizes needed to deserialize them. let dict_offsets_len = metadata.dict_size as usize + 1; - let total_tokens = usize::try_from(metadata.total_tokens) - .map_err(|_| vortex_err!("total_tokens {} overflows usize", metadata.total_tokens))?; + let codes_len = usize::try_from(metadata.codes_len) + .map_err(|_| vortex_err!("codes_len {} overflows usize", metadata.codes_len))?; // The cascading compressor may have narrowed any of these integer // children to a tighter ptype; the recorded ptype tells the framework // exactly which dtype to materialise as. @@ -462,7 +415,7 @@ impl VTable for OnPair { let codes = children.get( 1, &DType::Primitive(codes_ptype, Nullability::NonNullable), - total_tokens, + codes_len, )?; let codes_offsets = children.get( 2, @@ -480,7 +433,7 @@ impl VTable for OnPair { other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), }; - let data = OnPairData::new(buffers[0].clone(), metadata.dict_size, len); + let data = OnPairData::new(buffers[0].clone(), len); let slots = OnPairSlots { dict_offsets, codes, diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 231e8e81d0f..530e68c0bae 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -26,6 +26,7 @@ use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; @@ -95,19 +96,20 @@ pub(crate) fn onpair_decode_views( CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; - // Pad the output with DECODE_PADDING to absorb the decoder's fixed - // per-token over-copy; the exact decoded size is checked below. - let mut out_bytes = ByteBufferMut::with_capacity(total_size + onpair::DECODE_PADDING); - let written = onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) - .map_err(|_| { - vortex_err!("OnPair codes decode to more bytes than uncompressed_lengths records") - })?; - vortex_ensure!( - written == total_size, - "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" - ); - // SAFETY: `try_decode_into` initialised exactly `written` bytes of the - // spare capacity reserved above. + let mut out_bytes = ByteBufferMut::with_capacity(total_size); + let written = + match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) { + Ok(written) => written, + Err(_) => { + vortex_panic!("OnPair codes decode to more bytes than uncompressed_lengths records") + } + }; + if written != total_size { + vortex_panic!( + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" + ); + } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; match_each_integer_ptype!(lengths.ptype(), |P| { diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index 6f4bf389273..69170616385 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -4,13 +4,14 @@ //! Train + compress entry points for the OnPair encoding. use onpair::Config; -use onpair::Offset; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; +use vortex_array::scalar::Scalar; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -22,7 +23,6 @@ use vortex_error::vortex_err; use vortex_mask::AllOr; use crate::OnPair; -use crate::OnPairArray; /// Default OnPair training configuration: 12-bit codes ("dict-12"). pub const DEFAULT_DICT12_CONFIG: Config = Config { @@ -30,21 +30,32 @@ pub const DEFAULT_DICT12_CONFIG: Config = Config { ..onpair::DEFAULT_CONFIG }; -fn onpair_compress_varbinview( - array: VarBinViewArray, +/// Compress any [`ArrayRef`] whose canonical form is a string array. +/// +/// All-null inputs are returned as a [`ConstantArray`]. +pub fn onpair_compress( + array: &ArrayRef, config: Config, ctx: &mut ExecutionCtx, -) -> VortexResult -where - O: Offset, -{ +) -> VortexResult { + let array = array.clone().execute::(ctx)?; let len = array.len(); - let mask = array.validity()?.execute_mask(len, ctx)?; - let mut flat: Vec = Vec::with_capacity(len * 16); - let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - offsets.push(O::from_usize(0)); + let validity = array.validity()?; + let mask = validity.execute_mask(len, ctx)?; + if matches!(mask.bit_buffer(), AllOr::None) { + // CascadingCompressor handles this earlier, but direct callers can reach it. + return Ok(ConstantArray::new(Scalar::null(array.dtype().clone()), len).into_array()); + } + let views = array.views(); + let flat_bytes: usize = views.iter().map(|v| v.len() as usize).sum(); + + // TODO(francesco): we flatten because onpair training needs a contiguous `(bytes, offsets)` + // pair. Allowing onpair to train on a slice-of-slices would let us skip this copy. + let mut flat: Vec = Vec::with_capacity(flat_bytes); + let mut offsets: Vec = Vec::with_capacity(len + 1); + let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); + offsets.push(0); let buffers = array .data_buffers() .as_ref() @@ -57,27 +68,22 @@ where for view in views { let bytes = view_bytes(view, &buffers); flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); } } - AllOr::None => { - offsets.resize(len + 1, O::from_usize(0)); - for _ in 0..len { - uncompressed_lengths.push(0); - } - } + AllOr::None => unreachable!("all-null input handled above"), AllOr::Some(validity) => { for (view, valid) in views.iter().zip(validity.iter()) { if valid { let bytes = view_bytes(view, &buffers); flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); } else { - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths.push(0); } } @@ -88,7 +94,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 = codes_offsets_array(&row_offsets, u32::MAX as usize); + let codes_offsets = codes_offsets_array(&row_offsets); let codes = Buffer::from(codes).into_array(); let dict_offsets = Buffer::from(dict_offsets).into_array(); @@ -101,8 +107,9 @@ where codes, codes_offsets, uncompressed_lengths, - array.validity()?, + validity, ) + .map(IntoArray::into_array) } fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] { @@ -128,64 +135,19 @@ fn dict_bytes_to_buffer(dict_bytes: Vec) -> BufferHandle { /// `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(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 { +/// single chunk carries more than `u32::MAX` tokens, matching the `u64` byte +/// offsets accepted at compression. +fn codes_offsets_array(row_offsets: &[u64]) -> ArrayRef { + let total_tokens = row_offsets.last().copied().unwrap_or(0); + if u32::try_from(total_tokens).is_ok() { Buffer::from( row_offsets .iter() - .map(|&o| u32::try_from(o.to_usize()).vortex_expect("code boundary fits u32")) + .map(|&o| u32::try_from(o).vortex_expect("code boundary fits u32")) .collect::>(), ) .into_array() } else { - Buffer::from( - row_offsets - .iter() - .map(|&o| u64::try_from(o.to_usize()).vortex_expect("token count fits u64")) - .collect::>(), - ) - .into_array() - } -} - -/// Compress any [`ArrayRef`] whose canonical form is a string array, by first -/// canonicalising to `VarBinViewArray`. -pub fn onpair_compress( - array: &ArrayRef, - config: Config, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let view = array.clone().execute::(ctx)?; - onpair_compress_varbinview::(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::(&[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::(&[0, 3, 8], 7); - assert_eq!( - wide.dtype(), - &DType::Primitive(PType::U64, Nullability::NonNullable) - ); + Buffer::from(row_offsets.to_vec()).into_array() } } diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index dcddc79a370..d6aa35a5dcf 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -225,7 +225,9 @@ mod tests { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)? + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; let rhs = ConstantArray::new("hello", arr.len()).into_array(); let eq = diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index e6aedf24183..4c2ae4529cc 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -4,10 +4,8 @@ //! Vortex string array backed by the [OnPair][onpair] short-string //! compression library, with pushdown for common string operations. //! -//! The default training preset is `dict-12`: the trainer may build a dictionary -//! with up to 4 096 tokens, while runtime code width is derived from the actual -//! dictionary size. See [`onpair_compress`] for the entry point and -//! [`OnPairArray`] for the resulting array type. +//! See [`onpair_compress`] for the compression entry point and [`OnPairArray`] +//! for the encoded representation of non-null inputs. //! //! [onpair]: https://arxiv.org/abs/2508.02280 diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index 9ead98be6cc..64fc967b096 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -9,8 +9,8 @@ use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; @@ -46,21 +46,20 @@ impl OperationsVTable for OnPair { .as_primitive() .as_::() .ok_or_else(|| vortex_err!("OnPair uncompressed_lengths[{index}] is null"))?; - // Pad the row buffer with DECODE_PADDING to absorb the decoder's fixed - // per-token over-copy; the exact decoded size is checked below. - let mut buf: Vec = Vec::with_capacity(len + onpair::DECODE_PADDING); - let written = onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) - .map_err(|_| { - vortex_err!( + let mut buf: Vec = Vec::with_capacity(len); + let written = + match onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) { + Ok(written) => written, + Err(_) => vortex_panic!( "OnPair row {index} decodes to more bytes than uncompressed_lengths records" - ) - })?; - vortex_ensure!( - written == len, - "OnPair row {index} decoded to {written} bytes but uncompressed_lengths records {len}" - ); - // SAFETY: `try_decode_into` initialised `written` bytes of the spare - // capacity reserved above. + ), + }; + if written != len { + vortex_panic!( + "OnPair row {index} decoded to {written} bytes but uncompressed_lengths records {len}" + ); + } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { buf.set_len(written) }; Ok(varbin_scalar(ByteBuffer::from(buf), array.dtype())) } diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index 7a43740a1f6..ba2709ddf43 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -8,6 +8,7 @@ use prost::Message; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; @@ -34,6 +35,17 @@ use crate::compress::onpair_compress; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +fn compress_onpair( + array: &vortex_array::ArrayRef, + ctx: &mut vortex_array::ExecutionCtx, +) -> vortex_error::VortexResult { + onpair_compress(array, DEFAULT_DICT12_CONFIG, ctx)? + .try_downcast::() + .map_err(|array| { + vortex_error::vortex_err!("expected OnPair array, got {}", array.encoding_id()) + }) +} + fn sample_input() -> VarBinArray { VarBinArray::from_iter( [ @@ -80,7 +92,7 @@ fn test_onpair_metadata_golden() { &OnPairMetadata { uncompressed_lengths_ptype: PType::I32 as i32, dict_size: 4096, - total_tokens: 128_000, + codes_len: 128_000, dict_offsets_ptype: PType::U32 as i32, codes_ptype: PType::U16 as i32, codes_offsets_ptype: PType::U32 as i32, @@ -127,11 +139,7 @@ fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> { #[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, - )?; + let narrow = compress_onpair(&sample_input().into_array(), &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. @@ -311,8 +319,7 @@ fn test_onpair_empty() -> vortex_error::VortexResult<()> { Ok(()) } -/// All-null input has no training values, but still produces a valid OnPair -/// array backed by an empty code stream and `len + 1` zero boundaries. +/// All-null input is already represented optimally by a constant null array. #[cfg_attr(miri, ignore)] #[test] fn test_onpair_all_null() -> vortex_error::VortexResult<()> { @@ -324,18 +331,8 @@ fn test_onpair_all_null() -> vortex_error::VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let arr = onpair_compress(&input, DEFAULT_DICT12_CONFIG, &mut ctx)?; - assert!(arr.codes().is_empty()); - let codes_offsets = arr - .codes_offsets() - .clone() - .execute::(&mut ctx)?; - assert_eq!(codes_offsets.as_slice::(), &[0, 0, 0, 0]); - let uncompressed_lengths = arr - .uncompressed_lengths() - .clone() - .execute::(&mut ctx)?; - assert_eq!(uncompressed_lengths.as_slice::(), &[0, 0, 0]); - assert_arrays_eq!(arr.into_array(), input, &mut ctx); + assert!(arr.is::()); + assert_arrays_eq!(arr, input, &mut ctx); Ok(()) } @@ -354,7 +351,7 @@ fn test_onpair_filter_shares_dict() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let dict_bytes_before = arr.dict_bytes().clone(); let dict_offsets_len_before = arr.dict_offsets().len(); @@ -457,7 +454,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() -> vortex_error::VortexR DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; // Force `codes_offsets` to u16 so the panicking pre-fix // `as_slice::()` would fire. @@ -513,7 +510,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u8() -> vortex_error::VortexRe DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let arr = narrow_codes_offsets(&arr, PType::U8); assert_eq!(arr.as_view().codes_offsets().dtype().as_ptype(), PType::U8); diff --git a/encodings/experimental/onpair/tests/big_data.rs b/encodings/experimental/onpair/tests/big_data.rs index e7f74fdd782..15f50bb3322 100644 --- a/encodings/experimental/onpair/tests/big_data.rs +++ b/encodings/experimental/onpair/tests/big_data.rs @@ -76,10 +76,9 @@ fn smoke_100k_rows() -> vortex_error::VortexResult<()> { let t0 = Instant::now(); let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; let compress_elapsed = t0.elapsed(); - let bits = arr.bits(); eprintln!( - "compressed {} rows ({} raw bytes) in {:?}, bits={}", - n, raw_bytes, compress_elapsed, bits + "compressed {} rows ({} raw bytes) in {:?}", + n, raw_bytes, compress_elapsed ); let arr_ref = arr.into_array(); diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index eb2b6680abf..8c7b0561502 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -73,7 +73,10 @@ impl Scheme for OnPairScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let utf8 = data.array_as_varbinview().into_owned(); - let onpair_array = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; + let encoded = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; + let Some(onpair_array) = encoded.as_opt::() else { + return Ok(encoded); + }; let dict_offsets = compress_offsets_child( compressor, From b099b5b3d25c574ad747585563b93d0049876b56 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Fri, 17 Jul 2026 16:30:08 +0100 Subject: [PATCH 11/13] perf(vortex-onpair): memoize widened dict_offsets and validate dictionary once Widen the dict_offsets child to u32 and validate the compact dictionary lazily, memoizing both in an Arc on OnPairData so the cost is paid at most once per dictionary rather than per operation. Thread OnPairData through new_unchecked so slice/filter/cast share the cell, seed it from the trainer's conformant offsets at compress time, and drop the cell on buffer replacement. Add a test covering lazy rejection of a corrupt dictionary on first use, including on a derived slice. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- encodings/experimental/onpair/src/array.rs | 92 ++++++++++++++++++- .../experimental/onpair/src/canonical.rs | 8 +- encodings/experimental/onpair/src/compress.rs | 15 ++- .../experimental/onpair/src/compute/cast.rs | 2 +- .../onpair/src/compute/compare.rs | 11 +-- .../experimental/onpair/src/compute/filter.rs | 2 +- .../experimental/onpair/src/compute/slice.rs | 2 +- encodings/experimental/onpair/src/ops.rs | 7 +- encodings/experimental/onpair/src/tests.rs | 40 +++++++- 9 files changed, 147 insertions(+), 32 deletions(-) diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index 4f3503ba213..111269ac484 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -5,7 +5,10 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::sync::Arc; +use std::sync::OnceLock; +use onpair::CompactDictionaryView; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -32,6 +35,7 @@ use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -43,6 +47,7 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_onpair; use crate::canonical::onpair_decode_views; +use crate::decode::collect_widened; use crate::rules::RULES; /// An [`OnPair`]-encoded Vortex array. @@ -124,12 +129,52 @@ pub struct OnPairData { /// including its trailing read padding. dict_bytes: BufferHandle, len: usize, + /// The `dict_offsets` child widened to `u32`, memoized on first use so the + /// child is decompressed and the dictionary content is validated at most + /// once per dictionary — never per operation. (`dict_bytes` needs no such + /// cache: it is a raw buffer, so access is already zero-cost.) + /// + /// INVARIANT: once populated, the offsets passed + /// [`CompactDictionaryView::validate`] against `dict_bytes` (or came from + /// the trainer via [`init_dict_offsets`](Self::init_dict_offsets)), and + /// they are the widened values of the array's `dict_offsets` child. The + /// `Arc` cell is shared only between arrays with identical `dict_bytes` + /// and logically identical `dict_offsets` (slice / filter / cast keep + /// both). + dict_offsets: Arc>>, } impl OnPairData { /// Build [`OnPairData`] from the dictionary blob and the number of rows. pub fn new(dict_bytes: BufferHandle, len: usize) -> Self { - Self { dict_bytes, len } + Self { + dict_bytes, + len, + dict_offsets: Arc::new(OnceLock::new()), + } + } + + /// Seed the widened-offsets cell with dictionary offsets that are already + /// known to be conformant, so the first operation skips validation. + /// + /// This is the crate-internal trust mint (the moral equivalent of + /// [`onpair::CompactDictionary::new_unchecked`]): compression seeds it + /// with the trainer's offsets, which are conformant by construction. + /// + /// # Safety + /// `(self.dict_bytes, offsets)` must satisfy every [`onpair`] compact + /// dictionary invariant (i.e. [`CompactDictionaryView::validate`] would + /// succeed on them), and `offsets` must be the widened values of the + /// array's `dict_offsets` child. [`dict_view`] relies on this to build + /// unchecked dictionary views. + pub(crate) unsafe fn init_dict_offsets(&self, offsets: Buffer) { + debug_assert!( + CompactDictionaryView::validate(self.dict_bytes().as_slice(), offsets.as_slice()) + .is_ok(), + "init_dict_offsets called with a non-conformant dictionary" + ); + // A benign race can only ever install another conformant value. + drop(self.dict_offsets.set(offsets)); } /// Number of rows in the array. @@ -153,6 +198,35 @@ impl OnPairData { } } +/// A conformant [`CompactDictionaryView`] over `array`'s dictionary. +/// +/// The first call per dictionary widens the `dict_offsets` child to `u32` and +/// validates the dictionary content; both results are memoized in +/// [`OnPairData`], so subsequent calls — including on arrays derived by +/// slice / filter / cast, which share the cell — pay neither cost again. +pub(crate) fn dict_view<'a>( + array: ArrayView<'a, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let data = array.data(); + let offsets = match data.dict_offsets.get() { + Some(offsets) => offsets, + None => { + let widened = collect_widened::(array.dict_offsets(), ctx)?; + CompactDictionaryView::validate(data.dict_bytes().as_slice(), widened.as_slice()) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + data.dict_offsets.get_or_init(|| widened) + } + }; + // SAFETY: the cell only ever holds offsets that satisfy the compact + // dictionary invariants against this `dict_bytes` (validated above, or + // guaranteed by `init_dict_offsets`'s contract), and both buffers are + // immutable. + Ok(unsafe { + CompactDictionaryView::new_unchecked(data.dict_bytes().as_slice(), offsets.as_slice()) + }) +} + impl Display for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( @@ -224,9 +298,17 @@ impl OnPair { }) } + /// Build an [`OnPairArray`] without validation, carrying `data` — and with + /// it the memoized widened `dict_offsets` — from an existing array. + /// + /// # Safety + /// The parts must satisfy the same invariants [`try_new`](Self::try_new) + /// checks. If `data`'s widened-offsets cell is populated (or shared with a + /// live array), `dict_offsets` must hold the same logical offsets the cell + /// was built from. pub(crate) unsafe fn new_unchecked( dtype: DType, - dict_bytes: BufferHandle, + mut data: OnPairData, dict_offsets: ArrayRef, codes: ArrayRef, codes_offsets: ArrayRef, @@ -234,7 +316,7 @@ impl OnPair { validity: Validity, ) -> OnPairArray { let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, len); + data.len = len; let slots = OnPairSlots { dict_offsets, codes, @@ -347,6 +429,10 @@ impl VTable for OnPair { ); let mut data = array.data().clone(); data.dict_bytes = buffers[0].clone(); + // The replacement blob may differ from the one the memoized offsets + // were validated against, so drop the (shared) cell rather than + // carry a claim we can no longer prove. + data.dict_offsets = Arc::new(OnceLock::new()); Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 530e68c0bae..f5e55467eff 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -9,7 +9,6 @@ use std::sync::Arc; use num_traits::AsPrimitive; -use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -25,11 +24,11 @@ use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::code_boundary_at; use crate::decode::collect_widened; @@ -91,10 +90,7 @@ pub(crate) fn onpair_decode_views( // contiguous decoder walks `codes` in order and never reads the per-row // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - let dict = - CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) - .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + let dict = dict_view(array, ctx)?; let mut out_bytes = ByteBufferMut::with_capacity(total_size); let written = diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index 69170616385..a2e63f2c8a2 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -96,20 +96,25 @@ pub fn onpair_compress( let (dict_bytes, dict_offsets) = dict.into_raw(); let codes_offsets = codes_offsets_array(&row_offsets); let codes = Buffer::from(codes).into_array(); - let dict_offsets = Buffer::from(dict_offsets).into_array(); + // The `dict_offsets` child and the memoized widened-offsets cell share + // this buffer, so seeding below costs no copy. + let dict_offsets = Buffer::from(dict_offsets); let uncompressed_lengths = uncompressed_lengths.into_array(); - OnPair::try_new( + let encoded = OnPair::try_new( array.dtype().clone(), dict_bytes_to_buffer(dict_bytes), - dict_offsets, + dict_offsets.clone().into_array(), codes, codes_offsets, uncompressed_lengths, validity, - ) - .map(IntoArray::into_array) + )?; + // SAFETY: the trainer's dictionary is conformant by construction, and + // `dict_offsets` is exactly the u32 child attached above. + unsafe { encoded.init_dict_offsets(dict_offsets) }; + Ok(encoded.into_array()) } fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] { diff --git a/encodings/experimental/onpair/src/compute/cast.rs b/encodings/experimental/onpair/src/compute/cast.rs index 667a6cb5c57..c31893f82fd 100644 --- a/encodings/experimental/onpair/src/compute/cast.rs +++ b/encodings/experimental/onpair/src/compute/cast.rs @@ -24,7 +24,7 @@ fn build_with_validity( unsafe { OnPair::new_unchecked( dtype.clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), array.codes().clone(), array.codes_offsets().clone(), diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index d6aa35a5dcf..bb10fcfd6aa 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use onpair::CompactDictionaryView; use onpair::search; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -16,12 +15,11 @@ use vortex_array::scalar_fn::fns::binary::CompareKernel; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::collect_codes_window; -use crate::decode::collect_widened; impl CompareKernel for OnPair { fn compare( @@ -68,12 +66,7 @@ impl CompareKernel for OnPair { return Ok(None); } - let dict_offsets = collect_widened::(lhs.dict_offsets(), ctx)?; - let dict = CompactDictionaryView::validate( - lhs.dict_bytes().as_slice(), - dict_offsets.as_slice(), - ) - .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + let dict = dict_view(lhs, ctx)?; let query = search::tokenize(&needle, dict); let window = collect_codes_window(lhs, ctx)?; diff --git a/encodings/experimental/onpair/src/compute/filter.rs b/encodings/experimental/onpair/src/compute/filter.rs index 623fde4e6eb..1543bff136a 100644 --- a/encodings/experimental/onpair/src/compute/filter.rs +++ b/encodings/experimental/onpair/src/compute/filter.rs @@ -63,7 +63,7 @@ impl FilterKernel for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), filtered_codes.elements().clone(), filtered_codes.offsets().clone(), diff --git a/encodings/experimental/onpair/src/compute/slice.rs b/encodings/experimental/onpair/src/compute/slice.rs index 861c6661ed2..49d840f2029 100644 --- a/encodings/experimental/onpair/src/compute/slice.rs +++ b/encodings/experimental/onpair/src/compute/slice.rs @@ -28,7 +28,7 @@ impl SliceReduce for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), array.codes().clone(), codes_offsets, diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index 64fc967b096..728e5a0e6f2 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use onpair::CompactDictionaryView; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::varbin::varbin_scalar; @@ -14,6 +13,7 @@ use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::code_boundary_at; use crate::decode::collect_widened; @@ -33,10 +33,7 @@ impl OperationsVTable for OnPair { let row_end = code_boundary_at(codes_offsets, index + 1, ctx)?; let codes = collect_widened::(&array.codes().slice(row_start..row_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - let dict = - CompactDictionaryView::validate(array.dict_bytes().as_slice(), dict_offsets.as_slice()) - .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + let dict = dict_view(array, ctx)?; // The per-row decoded length is recorded in the `uncompressed_lengths` // child, so read it directly instead of asking the decoder to compute it. diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index ba2709ddf43..20aa7b4e061 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -15,6 +15,7 @@ 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::buffer::BufferHandle; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -84,6 +85,43 @@ fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> Ok(()) } +/// Dictionary content is validated lazily — construction stays lightweight, +/// and a corrupt dictionary is rejected by the first operation that decodes +/// or searches through it, including on derived (sliced) arrays. +#[cfg_attr(miri, ignore)] +#[test] +fn test_corrupt_dictionary_rejected_on_first_use() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let arr = compress_onpair(&sample_input().into_array(), &mut ctx)?; + let view = arr.as_view(); + + // Chop one byte off the blob so the trailing read-padding invariant fails. + let truncated = view.dict_bytes().slice(0..view.dict_bytes().len() - 1); + let corrupt = OnPair::try_new( + view.dtype().clone(), + BufferHandle::new_host(truncated), + view.dict_offsets().clone(), + view.codes().clone(), + view.codes_offsets().clone(), + view.uncompressed_lengths().clone(), + view.array_validity(), + )?; + + let sliced = corrupt.clone().into_array().slice(1..3)?; + // Canonical decode, random access, and ops on a derived slice all pass + // through the same validation door. + assert!( + corrupt + .clone() + .into_array() + .execute::(&mut ctx) + .is_err() + ); + assert!(corrupt.into_array().execute_scalar(0, &mut ctx).is_err()); + assert!(sliced.execute::(&mut ctx).is_err()); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_metadata_golden() { @@ -425,7 +463,7 @@ fn narrow_codes_offsets(arr: &crate::OnPairArray, target: PType) -> crate::OnPai unsafe { OnPair::new_unchecked( view.dtype().clone(), - view.dict_bytes_handle().clone(), + view.data().clone(), view.dict_offsets().clone(), view.codes().clone(), narrowed_array, From a0a4ee3588cd82461a67a928451407b714d1e215 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Fri, 17 Jul 2026 16:44:59 +0100 Subject: [PATCH 12/13] chore(deps): use released onpair 0.1.1 from crates.io instead of git pin The onpair dependency was a git source pinned to commit 28df115, which goes against the workspace guideline of depending on published crates. onpair 0.1.1 is now on crates.io and its v0.1.1 tag is the exact commit that was pinned, so this is behavior-preserving: the resolved package is byte-identical. Drop the git source and version placeholder in favor of a plain registry dependency. Verified: cargo build -p vortex-onpair and all 36 vortex-onpair tests pass against the crates.io 0.1.1 package. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 5 +++-- Cargo.toml | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2e28df315f0..c77172e3ea1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6451,8 +6451,9 @@ checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" [[package]] name = "onpair" -version = "0.1.0" -source = "git+https://github.com/spiraldb/onpair.git?rev=28df115c799674c663adedc8f81620ac82c7a890#28df115c799674c663adedc8f81620ac82c7a890" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf79077d8decdf5714242f1813f844f23c80a0cf009aba52881effb4c36ecbe" dependencies = [ "hashbrown 0.16.1", "rand 0.9.4", diff --git a/Cargo.toml b/Cargo.toml index c9d5efbc3dd..68be0c50653 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -192,7 +192,7 @@ num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } once_cell = "1.21" oneshot = { version = "0.2.0", features = ["async"] } -onpair = { version = "0.1.0", git = "https://github.com/spiraldb/onpair.git", rev = "28df115c799674c663adedc8f81620ac82c7a890" } +onpair = "0.1.1" opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" opentelemetry_sdk = "0.32.0" From 76a54b37c6aa9a9f75fb6c4d884b87276af131ae Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Fri, 17 Jul 2026 17:24:57 +0100 Subject: [PATCH 13/13] DCO Remediation Commit for Francesco Gargiulo I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: ab329f7facb2186f54ab27ba07a96dc074e1e381 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: f838da3f6ba8a939af9847c22f6f41998d6eb57b I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: 3bb9dc73e06ec21d5f07b7534c32267d4776d005 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: e5705e807191471bb9ba41082f75a0b868fd1070 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: d43343e5268eaf0efe44012d08dd00fde38c0374 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: fc335b48c93b923f9aa0ee51bfb3b02d1da1bc98 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: b099b5b3d25c574ad747585563b93d0049876b56 I, Francesco Gargiulo , hereby add my Signed-off-by to this commit: a0a4ee3588cd82461a67a928451407b714d1e215 Signed-off-by: Francesco Gargiulo