From aa080444b14c5bcf3b26abd925d133a02b476b70 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Mon, 6 Jul 2026 17:42:22 +0100 Subject: [PATCH] feat(vortex-onpair): port to onpair 0.1.0 API; add LIKE, list_contains, and take kernels - Port the onpair encoding to the onpair 0.1.0 public API (raw column parts, validated dictionaries, MaxDictBits config). - Add compressed-domain LIKE pushdown backed by onpair prefix/substring search, plus list_contains and take compute kernels. - onpair is temporarily a git dependency pinned to the 0.1.0 merge commit; it will switch to the crates.io release before this PR is marked ready for review. Signed-off-by: Francesco Gargiulo Co-Authored-By: Claude Fable 5 --- Cargo.lock | 5 +- Cargo.toml | 2 +- encodings/experimental/onpair/README.md | 18 +- .../experimental/onpair/benches/decode.rs | 53 +-- .../onpair/goldenfiles/onpair.metadata | 2 +- encodings/experimental/onpair/src/array.rs | 71 ++-- .../experimental/onpair/src/canonical.rs | 33 +- encodings/experimental/onpair/src/compress.rs | 98 ++---- .../experimental/onpair/src/compute/cast.rs | 108 +++++- .../onpair/src/compute/compare.rs | 123 ++++++- .../experimental/onpair/src/compute/filter.rs | 1 - .../experimental/onpair/src/compute/like.rs | 325 ++++++++++++++++++ .../onpair/src/compute/list_contains.rs | 280 +++++++++++++++ .../experimental/onpair/src/compute/mod.rs | 3 + .../experimental/onpair/src/compute/slice.rs | 1 - .../experimental/onpair/src/compute/take.rs | 114 ++++++ encodings/experimental/onpair/src/kernel.rs | 17 +- encodings/experimental/onpair/src/lib.rs | 9 +- encodings/experimental/onpair/src/ops.rs | 20 +- encodings/experimental/onpair/src/tests.rs | 2 - vortex-btrblocks/src/schemes/string/onpair.rs | 3 +- 21 files changed, 1062 insertions(+), 226 deletions(-) create mode 100644 encodings/experimental/onpair/src/compute/like.rs create mode 100644 encodings/experimental/onpair/src/compute/list_contains.rs create mode 100644 encodings/experimental/onpair/src/compute/take.rs diff --git a/Cargo.lock b/Cargo.lock index 459b58dded4..1fb0973c99e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6564,9 +6564,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?rev=a31dfd507f5e4f07f104a9fd11346dd2d6a20f9b#a31dfd507f5e4f07f104a9fd11346dd2d6a20f9b" dependencies = [ "hashbrown 0.16.1", "rand 0.9.4", diff --git a/Cargo.toml b/Cargo.toml index e781ec222d1..868038b5d6d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -199,7 +199,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", rev = "a31dfd507f5e4f07f104a9fd11346dd2d6a20f9b", version = "0.1.0" } 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..0dc91a11f20 100644 --- a/encodings/experimental/onpair/README.md +++ b/encodings/experimental/onpair/README.md @@ -10,15 +10,19 @@ 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`, `take`, byte length, +constant equality / inequality, simple constant `LIKE` patterns, and constant +string-list membership (`IN (...)`). Unsupported operators or pattern shapes +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 +30,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 58fd0332ac3..7e4b1316e7a 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. @@ -103,9 +98,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,26 +121,18 @@ 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 { + pub fn new(dict_bytes: BufferHandle, dict_size: u32, len: usize) -> Self { Self { dict_bytes, - bits, + dict_size, len, } } @@ -160,7 +146,7 @@ impl OnPairData { } pub fn bits(&self) -> u32 { - self.bits + u32::from(onpair::code_bits_for_num_tokens(self.dict_size as usize)) } pub fn dict_bytes(&self) -> &ByteBuffer { @@ -176,9 +162,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 +175,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 +185,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 +205,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,7 +213,6 @@ impl OnPair { codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> VortexResult { validate_parts( &dtype, @@ -234,10 +220,11 @@ impl OnPair { &codes, &codes_offsets, &uncompressed_lengths, - bits, )?; let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); + let dict_size = u32::try_from(dict_offsets.len().saturating_sub(1)) + .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; + let data = OnPairData::new(dict_bytes, dict_size, len); let slots = OnPairSlots { dict_offsets, codes, @@ -251,7 +238,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 +246,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 +273,11 @@ fn validate_parts( codes: &ArrayRef, codes_offsets: &ArrayRef, uncompressed_lengths: &ArrayRef, - bits: u32, ) -> 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"); @@ -338,7 +325,6 @@ impl VTable for OnPair { s.codes, s.codes_offsets, s.uncompressed_lengths, - data.bits, )?; if s.uncompressed_lengths.len() != len { vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array"); @@ -395,7 +381,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 +452,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, diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 20a058e7f90..92640aa71c6 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::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,6 +25,7 @@ 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; @@ -66,10 +67,9 @@ pub(crate) fn onpair_decode_views( // bound the contiguous run of `codes` belonging to the rows present in // this array: `slice` keeps the full `codes` child and only narrows // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), - // while `filter` rebuilds both children so the window is the whole stream. - // OnPair has no `TakeExecute`, so a reordering take is served from the - // canonical `VarBinView` and never reaches this path. We only need those - // two boundaries, so point-look them up rather than decoding every offset. + // while `filter`/`take` rebuild both children so the window is the whole + // stream. We only need those two boundaries, so point-look them up rather + // than decoding every offset. let codes_offsets = array.codes_offsets(); let code_start = code_boundary_at(codes_offsets, 0, ctx)?; let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; @@ -90,19 +90,18 @@ 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)?; + 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(), - ); + let mut out_bytes = ByteBufferMut::with_capacity(total_size + onpair::DECODE_PADDING); + // SAFETY: `out_bytes` has capacity for `total_size + DECODE_PADDING`, and + // `total_size` is the decoded length from Vortex's uncompressed_lengths + // child. + let written = + unsafe { onpair::decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) }; debug_assert_eq!(written, total_size); - // SAFETY: `decompress_into` initialised exactly `written` bytes of the + // SAFETY: `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 adf9e73461c..d2fa4cda8f5 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; @@ -39,19 +37,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); @@ -75,7 +60,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()) { @@ -95,24 +83,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, ) } @@ -125,57 +121,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..ede85120220 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,72 @@ 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 onpair = onpair_compress( + &array, + DEFAULT_DICT12_CONFIG, + &mut SESSION.create_execution_ctx(), + )?; + test_cast_conformance(&onpair.into_array()); + Ok(()) } } diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index 939a1a2fb42..63b58331d2b 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_ensure; +use vortex_error::vortex_err; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::decode::collect_widened; impl CompareKernel for OnPair { fn compare( @@ -28,29 +33,78 @@ 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 offsets = collect_widened::(lhs.codes_offsets(), ctx)?; + let code_start = offsets[0] as usize; + let code_end = offsets[lhs.len()] as usize; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= lhs.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + lhs.codes().len() + ); + for window in offsets[..=lhs.len()].windows(2) { + vortex_ensure!( + window[0] <= window[1], + "OnPair codes_offsets must be nondecreasing" + ); + } + let codes = collect_widened::(&lhs.codes().slice(code_start..code_end)?, ctx)?; + + let negated = operator == CompareOperator::NotEq; + BitBuffer::collect_bool(lhs.len(), |i| { + let start = offsets[i] as usize - code_start; + let end = offsets[i + 1] as usize - code_start; + search::equals(&codes[start..end], &query) != negated + }) + }; + Ok(Some( BoolArray::new( buffer, @@ -142,4 +196,35 @@ 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(()) + } } 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/like.rs b/encodings/experimental/onpair/src/compute/like.rs new file mode 100644 index 00000000000..781902f0096 --- /dev/null +++ b/encodings/experimental/onpair/src/compute/like.rs @@ -0,0 +1,325 @@ +// 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; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::scalar_fn::fns::like::LikeKernel; +use vortex_array::scalar_fn::fns::like::LikeOptions; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; +use crate::decode::collect_widened; + +enum SearchPattern { + Exact(Vec), + Prefix(Vec), + Contains(Vec), +} + +enum PreparedPattern { + Exact(Vec), + Prefix(search::PrefixQuery), + Contains(search::ContainsTable), +} + +impl PreparedPattern { + fn matches(&self, codes: &[u16]) -> bool { + match self { + Self::Exact(query) => search::equals(codes, query), + Self::Prefix(query) => search::starts_with(codes, query), + Self::Contains(table) => search::contains(codes, table), + } + } +} + +impl LikeKernel for OnPair { + fn like( + array: ArrayView<'_, Self>, + pattern: &ArrayRef, + options: LikeOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if options.case_insensitive { + return Ok(None); + } + + let Some(pattern_scalar) = pattern.as_constant() else { + return Ok(None); + }; + let pattern_bytes: &[u8] = if let Some(s) = pattern_scalar.as_utf8_opt() { + let Some(v) = s.value() else { + return Ok(None); + }; + v.as_ref() + } else if let Some(b) = pattern_scalar.as_binary_opt() { + let Some(v) = b.value() else { + return Ok(None); + }; + v + } else { + return Ok(None); + }; + let Some(search_pattern) = classify_like_pattern(pattern_bytes) else { + return Ok(None); + }; + + 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 offsets = collect_widened::(array.codes_offsets(), ctx)?; + let code_start = offsets[0] as usize; + let code_end = offsets[array.len()] as usize; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); + for window in offsets[..=array.len()].windows(2) { + vortex_ensure!( + window[0] <= window[1], + "OnPair codes_offsets must be nondecreasing" + ); + } + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + + let prepared = match search_pattern { + SearchPattern::Exact(needle) => PreparedPattern::Exact(search::tokenize(&needle, dict)), + SearchPattern::Prefix(prefix) => { + PreparedPattern::Prefix(search::PrefixQuery::new(&prefix, dict)) + } + SearchPattern::Contains(pattern) => { + if pattern.len() > u8::MAX as usize { + return Ok(None); + } + PreparedPattern::Contains(search::ContainsTable::new(&pattern, dict)) + } + }; + + let negated = options.negated; + let result = BitBuffer::collect_bool(array.len(), |i| { + let start = offsets[i] as usize - code_start; + let end = offsets[i + 1] as usize - code_start; + let matched = prepared.matches(&codes[start..end]); + if negated { !matched } else { matched } + }); + + let validity = array + .array() + .validity()? + .union_nullability(pattern_scalar.dtype().nullability()); + + Ok(Some(BoolArray::new(result, validity).into_array())) + } +} + +fn classify_like_pattern(pattern: &[u8]) -> Option { + let mut literal = Vec::with_capacity(pattern.len()); + let mut wildcards = Vec::new(); + let mut i = 0; + while i < pattern.len() { + match pattern[i] { + b'\\' => { + i += 1; + if i < pattern.len() { + literal.push(pattern[i]); + } else { + literal.push(b'\\'); + } + } + b'%' => wildcards.push(literal.len()), + b'_' => return None, + b => literal.push(b), + } + i += 1; + } + + match wildcards.as_slice() { + [] => Some(SearchPattern::Exact(literal)), + [end] if *end == literal.len() => Some(SearchPattern::Prefix(literal)), + [0, end] if *end == literal.len() => Some(SearchPattern::Contains(literal)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar_fn::fns::like::Like; + use vortex_array::scalar_fn::fns::like::LikeKernel; + use vortex_array::scalar_fn::fns::like::LikeOptions; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::OnPair; + use crate::OnPairArray; + 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 + }); + + fn make_onpair(strings: &[Option<&str>], nullability: Nullability) -> OnPairArray { + let array = + VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)).into_array(); + onpair_compress( + &array, + DEFAULT_DICT12_CONFIG, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + } + + fn run_like(array: OnPairArray, pattern: &str, opts: LikeOptions) -> VortexResult { + let len = array.len(); + let arr = array.into_array(); + let pattern = ConstantArray::new(pattern, len).into_array(); + let result = Like + .try_new_array(len, opts, [arr, pattern])? + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + Ok(result.into_bool()) + } + + fn like(array: OnPairArray, pattern: &str) -> VortexResult { + run_like(array, pattern, LikeOptions::default()) + } + + #[test] + fn test_like_exact() -> VortexResult<()> { + let onpair = make_onpair( + &[Some("alpha"), Some("alphabet"), Some("beta"), Some("alpha")], + Nullability::NonNullable, + ); + let result = like(onpair, "alpha")?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, false, true]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_like_prefix() -> VortexResult<()> { + let onpair = make_onpair( + &[Some("http://a"), Some("ftp://b"), Some("http://c")], + Nullability::NonNullable, + ); + let result = like(onpair, "http%")?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_like_contains_with_nulls() -> VortexResult<()> { + let onpair = make_onpair( + &[ + Some("hello world"), + None, + Some("goodbye"), + Some("say hello"), + ], + Nullability::Nullable, + ); + let result = like(onpair, "%hello%")?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([Some(true), None, Some(false), Some(true)]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_not_like_contains() -> VortexResult<()> { + let onpair = make_onpair( + &[Some("foobar_sdf"), Some("sdf_start"), Some("nothing")], + Nullability::NonNullable, + ); + let opts = LikeOptions { + negated: true, + case_insensitive: false, + }; + let result = run_like(onpair, "%sdf%", opts)?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([false, false, true]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_like_kernel_handles_contains() -> VortexResult<()> { + let onpair = make_onpair( + &[Some("hello world"), Some("goodbye")], + Nullability::NonNullable, + ); + let pattern = ConstantArray::new("%world%", onpair.len()).into_array(); + let result = ::like( + onpair.as_view(), + &pattern, + LikeOptions::default(), + &mut SESSION.create_execution_ctx(), + )?; + assert!( + result.is_some(), + "OnPair LikeKernel should handle %literal%" + ); + Ok(()) + } + + #[test] + fn test_like_kernel_rejects_overlong_contains() -> VortexResult<()> { + let onpair = make_onpair( + &[Some("hello world"), Some("goodbye")], + Nullability::NonNullable, + ); + let pattern = format!("%{}%", "x".repeat(usize::from(u8::MAX) + 1)); + let pattern = ConstantArray::new(pattern.as_str(), onpair.len()).into_array(); + let result = ::like( + onpair.as_view(), + &pattern, + LikeOptions::default(), + &mut SESSION.create_execution_ctx(), + )?; + assert!( + result.is_none(), + "overlong contains should fall back instead of panicking" + ); + Ok(()) + } +} diff --git a/encodings/experimental/onpair/src/compute/list_contains.rs b/encodings/experimental/onpair/src/compute/list_contains.rs new file mode 100644 index 00000000000..2999fa61df2 --- /dev/null +++ b/encodings/experimental/onpair/src/compute/list_contains.rs @@ -0,0 +1,280 @@ +// 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; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; +use crate::decode::collect_widened; + +impl ListContainsElementKernel for OnPair { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(list_scalar) = list.as_constant() else { + return Ok(None); + }; + + let list_scalar = list_scalar.as_list(); + let Some(elements) = list_scalar.elements() else { + return Ok(Some( + ConstantArray::new(null_bool(list, element), element.len()).into_array(), + )); + }; + if !matches!( + list_scalar.element_dtype(), + DType::Utf8(_) | DType::Binary(_) + ) { + return Ok(None); + } + + let needles = elements.iter().filter_map(scalar_bytes).collect::>(); + if needles.is_empty() { + return Ok(Some( + BoolArray::new( + BitBuffer::new_unset(element.len()), + element + .array() + .validity()? + .union_nullability(list.dtype().nullability()), + ) + .into_array(), + )); + } + + let dict_offsets = collect_widened::(element.dict_offsets(), ctx)?; + let dict = CompactDictionaryView::validate( + element.dict_bytes().as_slice(), + dict_offsets.as_slice(), + ) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + + let mut queries = needles + .iter() + .map(|bytes| search::tokenize(bytes, dict)) + .collect::>(); + queries.sort(); + queries.dedup(); + + let offsets = collect_widened::(element.codes_offsets(), ctx)?; + let code_start = offsets[0] as usize; + let code_end = offsets[element.len()] as usize; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= element.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + element.codes().len() + ); + for window in offsets[..=element.len()].windows(2) { + vortex_ensure!( + window[0] <= window[1], + "OnPair codes_offsets must be nondecreasing" + ); + } + let codes = collect_widened::(&element.codes().slice(code_start..code_end)?, ctx)?; + + let result = BitBuffer::collect_bool(element.len(), |i| { + let start = offsets[i] as usize - code_start; + let end = offsets[i + 1] as usize - code_start; + let row = &codes[start..end]; + queries.iter().any(|query| search::equals(row, query)) + }); + + Ok(Some( + BoolArray::new( + result, + element + .array() + .validity()? + .union_nullability(list.dtype().nullability()), + ) + .into_array(), + )) + } +} + +fn scalar_bytes(scalar: &Scalar) -> Option<&[u8]> { + match scalar.dtype() { + DType::Utf8(_) => scalar.as_utf8().value().map(|value| value.as_bytes()), + DType::Binary(_) => scalar.as_binary().value().map(|value| value.as_slice()), + _ => None, + } +} + +fn null_bool(list: &ArrayRef, element: ArrayView<'_, OnPair>) -> Scalar { + Scalar::null(DType::Bool( + list.dtype().nullability() | element.dtype().nullability(), + )) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::expr::list_contains; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::OnPair; + use crate::OnPairArray; + 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 + }); + + fn make_onpair(strings: &[Option<&str>], nullability: Nullability) -> OnPairArray { + let array = + VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)).into_array(); + onpair_compress( + &array, + DEFAULT_DICT12_CONFIG, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + } + + fn string_list(values: Vec, nullability: Nullability) -> Scalar { + Scalar::list( + Arc::new(DType::Utf8(Nullability::Nullable)), + values, + nullability, + ) + } + + #[test] + fn test_list_contains_string_literals() -> VortexResult<()> { + let onpair = make_onpair( + &[ + Some("alpha"), + None, + Some("beta"), + Some("gamma"), + Some("alpha"), + ], + Nullability::Nullable, + ); + let list = string_list( + vec![ + Scalar::utf8("alpha", Nullability::Nullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + Scalar::utf8("gamma", Nullability::Nullable), + ], + Nullability::Nullable, + ); + + let expr = list_contains(lit(list), root()); + let result = onpair + .into_array() + .apply(&expr)? + .execute::(&mut SESSION.create_execution_ctx())?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([Some(true), None, Some(false), Some(true), Some(true)]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_list_contains_empty_list() -> VortexResult<()> { + let onpair = make_onpair(&[Some("alpha"), Some("beta")], Nullability::NonNullable); + let list = string_list(Vec::new(), Nullability::Nullable); + + let expr = list_contains(lit(list), root()); + let result = onpair + .into_array() + .apply(&expr)? + .execute::(&mut SESSION.create_execution_ctx())?; + assert_arrays_eq!( + &result, + &BoolArray::from_iter([Some(false), Some(false)]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_list_contains_null_list() -> VortexResult<()> { + let onpair = make_onpair(&[Some("alpha"), Some("beta")], Nullability::NonNullable); + let list = Scalar::null(DType::List( + Arc::new(DType::Utf8(Nullability::Nullable)), + Nullability::Nullable, + )); + let pattern = ConstantArray::new(list, onpair.len()).into_array(); + let result = ::list_contains( + &pattern, + onpair.as_view(), + &mut SESSION.create_execution_ctx(), + )? + .expect("constant null list is handled by the OnPair kernel") + .execute::(&mut SESSION.create_execution_ctx())?; + + assert_arrays_eq!( + &result, + &BoolArray::from_iter([None::, None]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_list_contains_non_constant_list_falls_back() -> VortexResult<()> { + let onpair = make_onpair(&[Some("alpha"), Some("beta")], Nullability::NonNullable); + let elements = VarBinArray::from_iter( + [Some("alpha"), Some("beta")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let offsets = PrimitiveArray::from_iter([0u32, 1, 2]).into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); + + let result = ::list_contains( + &list, + onpair.as_view(), + &mut SESSION.create_execution_ctx(), + )?; + + assert!(result.is_none()); + Ok(()) + } +} diff --git a/encodings/experimental/onpair/src/compute/mod.rs b/encodings/experimental/onpair/src/compute/mod.rs index 4ad5f48f578..d96ebe7ff6b 100644 --- a/encodings/experimental/onpair/src/compute/mod.rs +++ b/encodings/experimental/onpair/src/compute/mod.rs @@ -5,4 +5,7 @@ mod byte_length; mod cast; mod compare; mod filter; +mod like; +mod list_contains; mod slice; +mod take; 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/compute/take.rs b/encodings/experimental/onpair/src/compute/take.rs new file mode 100644 index 00000000000..8b45d06e89d --- /dev/null +++ b/encodings/experimental/onpair/src/compute/take.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::OnPair; +use crate::OnPairArrayExt; +use crate::OnPairArraySlotsExt; + +impl TakeExecute for OnPair { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let codes = unsafe { + ListArray::new_unchecked( + array.codes().clone(), + array.codes_offsets().clone(), + Validity::NonNullable, + ) + }; + let taken_codes_ref = ::take(codes.as_view(), indices, ctx)? + .vortex_expect("List take kernel always returns Some"); + let taken_codes = taken_codes_ref + .try_downcast::() + .ok() + .vortex_expect("take for OnPair codes must return list array"); + + let lengths = array + .uncompressed_lengths() + .take(indices.clone())? + .fill_null(Scalar::zero_value(array.uncompressed_lengths().dtype()))?; + let validity = array.array_validity().take(indices)?; + + Ok(Some( + unsafe { + OnPair::new_unchecked( + array + .dtype() + .clone() + .union_nullability(indices.dtype().nullability()), + array.dict_bytes_handle().clone(), + array.dict_offsets().clone(), + taken_codes.elements().clone(), + taken_codes.offsets().clone(), + lengths, + validity, + ) + } + .into_array(), + )) + } +} + +#[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::take::test_take_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( + ["hello world", "testing onpair", "compression test", "data array", "vortex encoding"].map(Some), + DType::Utf8(Nullability::NonNullable), + ))] + #[case(VarBinArray::from_iter( + [Some("hello"), None, Some("world"), Some("test"), None], + DType::Utf8(Nullability::Nullable), + ))] + #[case(VarBinArray::from_iter( + ["single element"].map(Some), + DType::Utf8(Nullability::NonNullable), + ))] + fn test_take_onpair_conformance(#[case] varbin: VarBinArray) -> VortexResult<()> { + let array = varbin.into_array(); + let onpair = onpair_compress( + &array, + DEFAULT_DICT12_CONFIG, + &mut SESSION.create_execution_ctx(), + )?; + test_take_conformance(&onpair.into_array()); + Ok(()) + } +} diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index 8863d750a72..8a7bf5def69 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::Filter; +use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; @@ -10,18 +12,31 @@ 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_array::scalar_fn::fns::like::Like; +use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; 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(Dict.id(), OnPair, TakeExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel( ByteLength.id(), OnPair, ByteLengthExecuteAdaptor(OnPair), ); + kernels.register_execute_parent_kernel(Like.id(), OnPair, LikeExecuteAdaptor(OnPair)); + kernels.register_execute_parent_kernel( + ListContains.id(), + OnPair, + ListContainsElementExecuteAdaptor(OnPair), + ); } 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..c9da30a32a9 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; @@ -33,12 +33,9 @@ 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(), - }; + 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 +45,13 @@ 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()); + let mut buf: Vec = Vec::with_capacity(len + onpair::DECODE_PADDING); + // SAFETY: `buf` has capacity for `len + DECODE_PADDING`, and `len` is + // the decoded length read from Vortex's uncompressed_lengths child. + let written = + unsafe { onpair::decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) }; debug_assert_eq!(written, len); - // SAFETY: `decompress_into` initialised `written` bytes of the spare + // SAFETY: `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..c8acd15d0ed 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -48,7 +48,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 +314,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 7bbe480d059..b19996b6bba 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()) }