From 4016ab4e7190f9997f874cb90ac80e7acf0efa73 Mon Sep 17 00:00:00 2001 From: rapour Date: Thu, 20 Aug 2026 11:08:37 +0330 Subject: [PATCH 1/2] feat: elias-fano encoding plus reduce and kernel computes Signed-off-by: rapour --- Cargo.lock | 21 + Cargo.toml | 2 + encodings/elias-fano/Cargo.toml | 35 + .../goldenfiles/elias_fano.metadata | 2 + encodings/elias-fano/src/array.rs | 722 +++++++++ encodings/elias-fano/src/compress.rs | 390 +++++ encodings/elias-fano/src/compute/cast.rs | 45 + encodings/elias-fano/src/compute/compare.rs | 109 ++ encodings/elias-fano/src/compute/filter.rs | 73 + encodings/elias-fano/src/compute/is_sorted.rs | 43 + encodings/elias-fano/src/compute/min_max.rs | 61 + encodings/elias-fano/src/compute/mod.rs | 10 + encodings/elias-fano/src/compute/slice.rs | 29 + encodings/elias-fano/src/compute/take.rs | 132 ++ encodings/elias-fano/src/cursor.rs | 577 +++++++ encodings/elias-fano/src/kernel.rs | 26 + encodings/elias-fano/src/lib.rs | 71 + encodings/elias-fano/src/params.rs | 172 ++ encodings/elias-fano/src/rules.rs | 15 + encodings/elias-fano/src/tests.rs | 1443 +++++++++++++++++ vortex-buffer/src/bit/buf.rs | 39 + vortex-buffer/src/bit/select.rs | 309 +++- vortex-file/Cargo.toml | 1 + vortex-file/src/lib.rs | 1 + vortex/Cargo.toml | 1 + vortex/src/editions/mod.rs | 4 +- vortex/src/editions/unstable/mod.rs | 2 + vortex/src/editions/unstable/v2026_08.rs | 21 + vortex/src/lib.rs | 5 + 29 files changed, 4305 insertions(+), 56 deletions(-) create mode 100644 encodings/elias-fano/Cargo.toml create mode 100644 encodings/elias-fano/goldenfiles/elias_fano.metadata create mode 100644 encodings/elias-fano/src/array.rs create mode 100644 encodings/elias-fano/src/compress.rs create mode 100644 encodings/elias-fano/src/compute/cast.rs create mode 100644 encodings/elias-fano/src/compute/compare.rs create mode 100644 encodings/elias-fano/src/compute/filter.rs create mode 100644 encodings/elias-fano/src/compute/is_sorted.rs create mode 100644 encodings/elias-fano/src/compute/min_max.rs create mode 100644 encodings/elias-fano/src/compute/mod.rs create mode 100644 encodings/elias-fano/src/compute/slice.rs create mode 100644 encodings/elias-fano/src/compute/take.rs create mode 100644 encodings/elias-fano/src/cursor.rs create mode 100644 encodings/elias-fano/src/kernel.rs create mode 100644 encodings/elias-fano/src/lib.rs create mode 100644 encodings/elias-fano/src/params.rs create mode 100644 encodings/elias-fano/src/rules.rs create mode 100644 encodings/elias-fano/src/tests.rs create mode 100644 vortex/src/editions/unstable/v2026_08.rs diff --git a/Cargo.lock b/Cargo.lock index a92a0f5be59..db6126e0a81 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9542,6 +9542,7 @@ dependencies = [ "vortex-datetime-parts", "vortex-decimal-byte-parts", "vortex-edition", + "vortex-elias-fano", "vortex-error", "vortex-fastlanes", "vortex-file", @@ -10067,6 +10068,25 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-elias-fano" +version = "0.1.0" +dependencies = [ + "fastlanes", + "lending-iterator", + "num-traits", + "prost 0.14.4", + "rstest", + "smallvec", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-mask", + "vortex-proto", + "vortex-session", +] + [[package]] name = "vortex-error" version = "0.1.0" @@ -10151,6 +10171,7 @@ dependencies = [ "vortex-datetime-parts", "vortex-decimal-byte-parts", "vortex-edition", + "vortex-elias-fano", "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..09abc62b64f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,7 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", + "encodings/elias-fano", # Benchmarks "benchmarks/bench-support", "benchmarks/lance-bench", @@ -307,6 +308,7 @@ vortex-datafusion = { version = "0.1.0", path = "./vortex-datafusion", default-f vortex-datetime-parts = { version = "0.1.0", path = "./encodings/datetime-parts", default-features = false } vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte-parts", default-features = false } vortex-edition = { version = "0.1.0", path = "./vortex-edition", default-features = false } +vortex-elias-fano = { version = "0.1.0", path = "./encodings/elias-fano", default-features = false } vortex-error = { version = "0.1.0", path = "./vortex-error", default-features = false } vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false } vortex-file = { version = "0.1.0", path = "./vortex-file", default-features = false } diff --git a/encodings/elias-fano/Cargo.toml b/encodings/elias-fano/Cargo.toml new file mode 100644 index 00000000000..5769017a97d --- /dev/null +++ b/encodings/elias-fano/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "vortex-elias-fano" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex Elias-Fano encoded array for monotonic integer sequences" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +fastlanes = { workspace = true } +lending-iterator = { workspace = true } +num-traits = { workspace = true } +prost = { workspace = true } +smallvec = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } +vortex-mask = { workspace = true } +vortex-proto = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[lints] +workspace = true diff --git a/encodings/elias-fano/goldenfiles/elias_fano.metadata b/encodings/elias-fano/goldenfiles/elias_fano.metadata new file mode 100644 index 00000000000..47396431f86 --- /dev/null +++ b/encodings/elias-fano/goldenfiles/elias_fano.metadata @@ -0,0 +1,2 @@ + + ÿÿÿÿÿÿÿÿÿ þÿÿÿÿÿÿÿÿÿ ÿÿÿÿÿÿÿÿÿ(ÿÿÿÿÿÿÿÿÿ0ÿÿÿÿÿÿÿÿÿ \ No newline at end of file diff --git a/encodings/elias-fano/src/array.rs b/encodings/elias-fano/src/array.rs new file mode 100644 index 00000000000..d6114b808ba --- /dev/null +++ b/encodings/elias-fano/src/array.rs @@ -0,0 +1,722 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; + +use prost::Message; +use smallvec::smallvec; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::array_slots; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision as StatPrecision; +use vortex_array::expr::stats::Stat; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::serde::ArrayChildren; +use vortex_array::stats::StatsSet; +use vortex_array::validity::Validity; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityVTable; +use vortex_buffer::Alignment; +use vortex_buffer::BitBuffer; +use vortex_buffer::ByteBuffer; +use vortex_buffer::read_u64_le; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::compress::elias_fano_decompress; +use crate::cursor::EliasFanoCursor; +use crate::params; +use crate::params::LOG_SAMPLING0; +use crate::params::LOG_SAMPLING1; +use crate::rules::RULES; + +/// An [`EliasFano`]-encoded Vortex array. +pub type EliasFanoArray = Array; + +/// The dtype of the bit-packed low-bits child, always `u64` whatever the array's own ptype. +/// +/// FastLanes packs `128 * bit_width` bytes per block regardless of element type, so a wide child +/// costs no space and every low-bits read monomorphises once with no runtime ptype dispatch. +pub(crate) const LOWER_DTYPE: DType = DType::Primitive(PType::U64, NonNullable); + +#[array_slots(EliasFano)] +pub struct EliasFanoSlots { + /// The low [`EliasFanoData::lower_width`] bits of every element, in element order. Normally + /// bit-packed, or a constant zero array at width zero. Its length is the *encoded* element + /// count, which after a slice exceeds the array's own; see [`EliasFanoData::first_rank`]. + #[slot(0)] + pub lower: ArrayRef, +} + +/// Wire-format metadata persisted alongside the buffers `[upper, samples]` and [`EliasFanoSlots`]. +/// +/// Only what cannot be re-derived: the seam between the two sample tables is absent, because +/// `params::num_samples0` recovers it from the universe. +#[derive(Clone, prost::Message)] +pub struct EliasFanoMetadata { + /// The value subtracted from every element before encoding. + #[prost(message, tag = "1")] + reference: Option, + /// The largest value in the *encoded* sequence, which fixes the universe. + #[prost(message, tag = "2")] + max: Option, + /// Number of low bits per element. + #[prost(uint32, tag = "3")] + lower_width: u32, + /// Length in bits of the `upper` buffer's bit array. + #[prost(uint64, tag = "4")] + upper_len: u64, + /// Rank of this array's first element within the encoded sequence. + #[prost(uint64, tag = "5")] + first_rank: u64, + /// Number of elements in the encoded sequence. Duplicates the low-bits child's length in + /// memory, but deserialization must declare a child's length before constructing it, and after + /// a slice that length is not the array's own. + #[prost(uint64, tag = "6")] + num_elements: u64, +} + +/// An Elias-Fano encoded monotonically non-decreasing integer sequence. +/// +/// Holds only what cannot be re-derived from the layout in the crate-private `params` module: the +/// two buffers, the universe bounds, and the four numbers that size it. +/// +/// Both buffers are host-resident. The upper array is read bit by bit, so there is no way to serve +/// it one entry at a time from device memory; `with_buffers` and `deserialize` copy to the host +/// once so no accessor below has to ask. +#[derive(Clone, Debug)] +pub struct EliasFanoData { + /// The unary upper array, `upper_len` bits, byte-padded. + upper: ByteBuffer, + /// The zero-sample positions followed by the one-sample positions, as little-endian `u64`s. + /// Where one table ends and the other begins is derived, not stored; see + /// [`Self::sample_bytes`]. It is read unaligned, because a deserialized buffer carries no + /// alignment guarantee. + samples: ByteBuffer, + reference: Scalar, + max: Scalar, + lower_width: u8, + upper_len: u64, + first_rank: u64, +} + +impl Display for EliasFanoData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "reference: {}, max: {}, lower_width: {}, upper_len: {}, first_rank: {}", + self.reference, self.max, self.lower_width, self.upper_len, self.first_rank + ) + } +} + +impl EliasFanoData { + /// Construct the per-array data, validating what can be checked without the child slot. + // There is one parameter per metadata field, which is what makes the two sides easy to line up. + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new( + upper: ByteBuffer, + samples: ByteBuffer, + reference: Scalar, + max: Scalar, + lower_width: u8, + upper_len: u64, + first_rank: u64, + ) -> VortexResult { + vortex_ensure!( + reference.dtype().is_int() && !reference.dtype().is_nullable(), + "Elias-Fano reference must be a non-nullable integer, got {}", + reference.dtype() + ); + vortex_ensure!( + max.dtype() == reference.dtype(), + "Elias-Fano max dtype {} does not match reference dtype {}", + max.dtype(), + reference.dtype() + ); + vortex_ensure!( + lower_width <= params::MAX_LOWER_WIDTH, + "Elias-Fano lower_width {lower_width} exceeds {}", + params::MAX_LOWER_WIDTH + ); + vortex_ensure!( + upper.len() == usize::try_from(upper_len.div_ceil(8))?, + "Elias-Fano upper buffer is {} bytes, expected {} for {upper_len} bits", + upper.len(), + upper_len.div_ceil(8) + ); + vortex_ensure!( + samples.len().is_multiple_of(size_of::()), + "Elias-Fano samples buffer of {} bytes is not a whole number of u64s", + samples.len() + ); + // The zero table comes first, so the buffer must reach at least as far as the seam for + // `sample_bytes` to be able to split there. + let span = scalar_bits(&max).wrapping_sub(scalar_bits(&reference)); + let num_samples0 = params::num_samples0(span, lower_width); + vortex_ensure!( + (samples.len() / size_of::()) as u64 >= num_samples0, + "Elias-Fano samples buffer holds {} entries, fewer than the {num_samples0} \ + zero-samples its universe implies", + samples.len() / size_of::() + ); + + Ok(Self { + upper, + samples, + reference, + max, + lower_width, + upper_len, + first_rank, + }) + } + + /// Returns the same layout, read from a different starting rank. + /// + /// A slice is nothing more than this; see [`Self::first_rank`]. + pub(crate) fn with_first_rank(mut self, first_rank: u64) -> Self { + self.first_rank = first_rank; + self + } + + /// Returns the same layout, relabelled with bounds of a different ptype. Sound only when both + /// bounds are exactly representable there, which is what leaves the span unchanged. See + /// [`CastReduce`](vortex_array::scalar_fn::fns::cast::CastReduce) for `EliasFano`. + pub(crate) fn with_bounds(mut self, reference: Scalar, max: Scalar) -> Self { + self.reference = reference; + self.max = max; + self + } + + /// The value subtracted from every element before encoding, and added back on read. + #[inline] + pub fn reference_scalar(&self) -> &Scalar { + &self.reference + } + + /// The largest value of the *encoded* sequence, fixing the universe the upper array was sized + /// for, so slicing leaves it untouched. Therefore **not** the maximum of a sliced array — use + /// [`EliasFanoCursor::access`](crate::EliasFanoCursor::access) at `len - 1` for that. + #[inline] + pub fn max_scalar(&self) -> &Scalar { + &self.max + } + + /// Number of low bits stored per element in the child slot. + #[inline] + pub fn lower_width(&self) -> u8 { + self.lower_width + } + + /// Length in bits of the upper array. + #[inline] + pub fn upper_len(&self) -> u64 { + self.upper_len + } + + /// Rank of this array's element 0 within the encoded sequence. + /// + /// Slicing cannot trim the buffers, because the sample tables hold *absolute* bit positions, so + /// a slice records where it starts and space is reclaimed on rewrite. This offsets both the + /// upper-array ranks and the low-bits child, so element `i` is rank `first_rank + i` in both. + #[inline] + pub fn first_rank(&self) -> u64 { + self.first_rank + } + + /// Number of zero-samples stored at the front of the samples buffer. + /// + /// This count is derived from the universe rather than stored, and the crate-private + /// `params::num_samples0` explains why the element count drops out of the derivation. + #[inline] + pub(crate) fn num_samples0(&self) -> u64 { + params::num_samples0(self.span(), self.lower_width) + } + + #[inline] + pub(crate) fn upper_buffer(&self) -> &ByteBuffer { + &self.upper + } + + #[inline] + pub(crate) fn samples_buffer(&self) -> &ByteBuffer { + &self.samples + } + + pub(crate) fn upper_bits(&self) -> VortexResult { + Ok(BitBuffer::new( + self.upper.clone().aligned(Alignment::none()), + usize::try_from(self.upper_len)?, + )) + } + + /// The zero-sample and one-sample tables, still as raw little-endian bytes. + /// + /// The two share a buffer; the seam is recomputed here from the universe alone, which buys back + /// a metadata field for a shift run once per cursor rather than per element. Deserialized + /// buffers carry no alignment guarantee, so entries are read one at a time with + /// [`read_sample`]. + pub(crate) fn sample_bytes(&self) -> VortexResult<(&[u8], &[u8])> { + let bytes = self.samples.as_slice(); + let num_samples0 = self.num_samples0(); + // Every path that builds an `EliasFanoData` goes through `try_new`, which proves the buffer + // reaches the seam. This re-derives rather than trusting that, because the seam is computed + // from the universe on each call and a raise is cheaper to reason about than an assertion. + usize::try_from(num_samples0) + .ok() + .and_then(|entries| entries.checked_mul(size_of::())) + .and_then(|seam| bytes.split_at_checked(seam)) + .ok_or_else(|| { + vortex_err!( + "Elias-Fano samples buffer of {} bytes is too short for the {num_samples0} \ + zero-samples its universe implies", + bytes.len() + ) + }) + } + + /// The reference value as a sign-extended 64-bit pattern. + /// + /// Encoding works in this domain throughout: `element = + /// value_bits.wrapping_sub(reference_bits)` and back. Sign-extend, wrap, truncate is exactly + /// two's complement, so one `u64` path serves every integer ptype, signed or not. + #[inline] + pub(crate) fn reference_bits(&self) -> u64 { + scalar_bits(&self.reference) + } + + /// The span of the encoded universe: `max - reference`, so the universe is `span + 1` values. + #[inline] + pub(crate) fn span(&self) -> u64 { + scalar_bits(&self.max).wrapping_sub(scalar_bits(&self.reference)) + } +} + +#[inline] +pub(crate) fn read_sample(table: &[u8], idx: usize) -> u64 { + read_u64_le(&table[idx * 8..][..8]) +} + +/// The two's-complement bit pattern of an integer scalar, sign-extended to 64 bits. +// The widening is what sign-extends, and it is a no-op only in the `u64` arm the macro also +// expands to, which is the arm the lint sees. +#[expect(clippy::unnecessary_cast)] +pub(crate) fn scalar_bits(scalar: &Scalar) -> u64 { + let pvalue = scalar + .as_primitive() + .pvalue() + .vortex_expect("Elias-Fano bounds are non-null integers"); + match_each_integer_ptype!(pvalue.ptype(), |P| { + pvalue + .cast::

() + .vortex_expect("pvalue is already of this ptype") as u64 + }) +} + +pub(crate) fn scalar_from_bits(dtype: &DType, bits: u64) -> VortexResult { + let value = match_each_integer_ptype!(dtype.as_ptype(), |P| { + ScalarValue::Primitive(PValue::from(bits as P)) + }); + Scalar::try_new(dtype.clone(), Some(value)) +} + +impl ArrayHash for EliasFanoData { + fn array_hash(&self, state: &mut H, accuracy: EqMode) { + self.reference.hash(state); + self.max.hash(state); + self.lower_width.hash(state); + self.upper_len.hash(state); + self.first_rank.hash(state); + self.upper.array_hash(state, accuracy); + self.samples.array_hash(state, accuracy); + } +} + +impl ArrayEq for EliasFanoData { + fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { + self.reference == other.reference + && self.max == other.max + && self.lower_width == other.lower_width + && self.upper_len == other.upper_len + && self.first_rank == other.first_rank + && self.upper.array_eq(&other.upper, accuracy) + && self.samples.array_eq(&other.samples, accuracy) + } +} + +impl VTable for EliasFano { + type TypedArrayData = EliasFanoData; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.elias_fano"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let lower = EliasFanoSlotsView::from_slots(slots).lower; + validate_parts(data, lower, dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 2 + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + match idx { + 0 => BufferHandle::new_host(array.upper_buffer().clone()), + 1 => BufferHandle::new_host(array.samples_buffer().clone()), + _ => vortex_panic!("EliasFanoArray buffer index {idx} out of bounds"), + } + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + match idx { + 0 => Some("upper".to_string()), + 1 => Some("samples".to_string()), + _ => None, + } + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 2, + "Expected 2 buffers, got {}", + buffers.len() + ); + let previous = array.data(); + // Back through `try_new` rather than assigning the fields: the constructor is the only + // place that holds the buffer invariants, so a replacement buffer that a caller reported as + // the right length still has to satisfy them here. + let data = EliasFanoData::try_new( + buffers[0].try_to_host_sync()?, + buffers[1].try_to_host_sync()?, + previous.reference.clone(), + previous.max.clone(), + previous.lower_width, + previous.upper_len, + previous.first_rank, + )?; + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + EliasFanoSlots::NAMES[idx].to_string() + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + EliasFanoMetadata { + reference: Some(ScalarValue::to_proto(array.reference_scalar().value())), + max: Some(ScalarValue::to_proto(array.max_scalar().value())), + lower_width: u32::from(array.lower_width()), + upper_len: array.upper_len(), + first_rank: array.first_rank(), + num_elements: array.lower().len() as u64, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 2, + "EliasFanoArray expects 2 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == 1, + "EliasFanoArray expects 1 child, got {}", + children.len() + ); + let metadata = EliasFanoMetadata::decode(metadata)?; + + let bound = |value: Option<&vortex_proto::scalar::ScalarValue>, what: &str| { + let value = value.ok_or_else(|| vortex_err!("Elias-Fano {what} is required"))?; + Scalar::from_proto_value(value, dtype, session) + }; + + let lower = children.get( + EliasFanoSlots::LOWER, + &LOWER_DTYPE, + usize::try_from(metadata.num_elements)?, + )?; + + let data = EliasFanoData::try_new( + buffers[0].try_to_host_sync()?, + buffers[1].try_to_host_sync()?, + bound(metadata.reference.as_ref(), "reference")?, + bound(metadata.max.as_ref(), "max")?, + u8::try_from(metadata.lower_width).map_err(|_| { + vortex_err!("Elias-Fano lower_width {} > 255", metadata.lower_width) + })?, + metadata.upper_len, + metadata.first_rank, + )?; + + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data) + .with_slots(smallvec![Some(lower)])) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + elias_fano_decompress(&array, ctx)?.into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for EliasFano { + fn scalar_at( + array: ArrayView<'_, EliasFano>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + EliasFanoCursor::try_new(array, ctx)?.access(index) + } +} + +impl ValidityVTable for EliasFano { + fn validity(_array: ArrayView<'_, EliasFano>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +/// Elias-Fano encoding for monotonically non-decreasing integer sequences. +#[derive(Clone, Debug)] +pub struct EliasFano; + +impl EliasFano { + /// Assemble an Elias-Fano array from encoded parts. + /// + /// Prefer [`elias_fano_encode`](crate::elias_fano_encode) unless you already hold a layout, + /// which is the case for a slice, a cast, or a rewrite of the low-bits child. + pub fn try_new( + data: EliasFanoData, + lower: ArrayRef, + len: usize, + ) -> VortexResult { + let dtype = data.reference_scalar().dtype().clone(); + let slots: ArraySlots = smallvec![Some(lower)]; + Array::try_from_parts(ArrayParts::new(EliasFano, dtype, len, data).with_slots(slots)) + .map(|array| array.with_stats_set(Self::stats())) + } + + /// Statistics that hold for every Elias-Fano array by construction. + /// + /// `IsSorted` is required, not a nicety: [`ListArray::new`](vortex_array::arrays::ListArray) + /// refuses offsets that do not report it. `IsStrictSorted` is absent because repeated offsets + /// are legal — an empty list contributes two identical ones — so it must be computed. + pub(crate) fn stats() -> StatsSet { + // SAFETY: a single stat cannot be duplicated. + unsafe { + StatsSet::new_unchecked(smallvec![( + Stat::IsSorted, + StatPrecision::Exact(true.into()), + )]) + } + } +} + +fn validate_parts( + data: &EliasFanoData, + lower: &ArrayRef, + dtype: &DType, + len: usize, +) -> VortexResult<()> { + vortex_ensure!( + dtype.is_int(), + "Elias-Fano requires an integer dtype, got {dtype}" + ); + vortex_ensure!( + !dtype.is_nullable(), + "Elias-Fano requires a non-nullable dtype, got {dtype}" + ); + vortex_ensure!( + data.reference_scalar().dtype() == dtype, + "Elias-Fano reference dtype {} does not match array dtype {dtype}", + data.reference_scalar().dtype() + ); + // Any integer array of the right width is acceptable here, not just `BitPacked`: a file + // roundtrip can hand the slot back wrapped (for example as `Patched(BitPacked)`), and a + // rewrite may replace it outright. + vortex_ensure!( + lower.dtype() == &LOWER_DTYPE, + "Elias-Fano low-bits child must be {LOWER_DTYPE}, got {}", + lower.dtype() + ); + // A bit-packed slot's width is metadata, so this is free to check and is the only part of the + // low bits checkable at all. Narrower is legal — a rewrite may repack tighter. Wider is not: + // the reader ORs the low bits in under `lower_width`, so anything above bleeds into the high + // part. + if let Some(packed) = lower.as_opt::() { + vortex_ensure!( + packed.bit_width() <= data.lower_width(), + "Elias-Fano low-bits child is packed at {} bits, above the {} the layout allows", + packed.bit_width(), + data.lower_width() + ); + } + + let num_elements = lower.len() as u64; + let end = data + .first_rank() + .checked_add(len as u64) + .ok_or_else(|| vortex_err!("Elias-Fano slice bounds overflow"))?; + vortex_ensure!( + end <= num_elements, + "Elias-Fano slice of {len} from rank {} exceeds the {num_elements} encoded elements", + data.first_rank() + ); + + if num_elements > 0 { + let expected_width = params::lower_width(data.span(), num_elements as usize); + vortex_ensure!( + data.lower_width() == expected_width, + "Elias-Fano lower_width {} does not match the {expected_width} implied by span {} \ + over {num_elements} elements", + data.lower_width(), + data.span() + ); + let expected_upper_len = + params::upper_len(data.span(), num_elements as usize, expected_width)?; + vortex_ensure!( + data.upper_len() == expected_upper_len, + "Elias-Fano upper_len {} does not match the {expected_upper_len} implied by span {} \ + over {num_elements} elements", + data.upper_len(), + data.span() + ); + + // Both tables sample from rank 1 upward, so their sizes follow from the layout and a reader + // never has to bounds-check a lookup. The zero count is also what `sample_bytes` splits on, + // derived on both sides, so this catches an encoder that disagrees about the seam. + let expected_samples0 = params::num_samples0(data.span(), expected_width); + debug_assert_eq!( + expected_samples0, + (params::num_zeros(expected_upper_len, num_elements as usize) - 1) >> LOG_SAMPLING0, + "the two derivations of the zero-sample count must agree" + ); + let expected_samples1 = (num_elements - 1) >> LOG_SAMPLING1; + let num_samples = (data.samples_buffer().len() / size_of::()) as u64; + vortex_ensure!( + num_samples == expected_samples0 + expected_samples1, + "Elias-Fano holds {num_samples} samples, expected {}", + expected_samples0 + expected_samples1 + ); + + // A sample is fed straight to `BitBuffer::select_range` as a window start, which asserts + // rather than raises past the end, so check every one — there are only `n / 256 + zeros / + // 512`. Each is pinned above by the upper array's length and below by the rank it stands + // for, which gives the strict increase a sampled search relies on. + let (samples0, samples1) = data.sample_bytes()?; + for (table, name, log_sampling, floor) in [ + (samples0, "zero", LOG_SAMPLING0, 0), + (samples1, "one", LOG_SAMPLING1, 1), + ] { + let mut previous = None; + for index in 0..table.len() / size_of::() { + let sample = read_sample(table, index); + let minimum = (((index + 1) as u64) << log_sampling) + floor; + vortex_ensure!( + (minimum..expected_upper_len).contains(&sample), + "Elias-Fano {name}-sample {index} points to bit {sample}, outside the \ + {minimum}..{expected_upper_len} its rank allows" + ); + vortex_ensure!( + previous.is_none_or(|previous| previous < sample), + "Elias-Fano {name}-samples are not strictly increasing at index {index}" + ); + previous = Some(sample); + } + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_array::test_harness::check_metadata; + + use super::*; + + #[cfg_attr(miri, ignore)] + #[test] + fn test_elias_fano_metadata() { + check_metadata( + "elias_fano.metadata", + &EliasFanoMetadata { + reference: Some((&ScalarValue::from(i64::MIN)).into()), + max: Some((&ScalarValue::from(i64::MAX)).into()), + lower_width: u32::from(u8::MAX), + upper_len: u64::MAX, + first_rank: u64::MAX, + num_elements: u64::MAX, + } + .encode_to_vec(), + ); + } +} diff --git a/encodings/elias-fano/src/compress.rs b/encodings/elias-fano/src/compress.rs new file mode 100644 index 00000000000..44d6373956c --- /dev/null +++ b/encodings/elias-fano/src/compress.rs @@ -0,0 +1,390 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Building an Elias-Fano array from a sorted primitive array, and taking it apart again. +//! +//! See [`crate::params`] for the layout both directions read and write. + +use std::iter; + +use lending_iterator::prelude::LendingIterator; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::match_each_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::Alignment; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BitIndexIterator; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked; + +use crate::EliasFano; +use crate::EliasFanoArray; +use crate::EliasFanoData; +use crate::array::EliasFanoArraySlotsExt; +use crate::array::scalar_from_bits; +use crate::cursor::position_of_rank; +use crate::params; + +/// Encode a sorted, non-nullable integer array with Elias-Fano. +/// +/// The input must be monotonically non-decreasing; duplicates are fine and cost one set bit each. +/// Nulls are rejected, because a null has no position in an ordering and the layout has nowhere to +/// put one. `ctx` is taken for symmetry with the other integer encoders; encoding does not use it. +// Values widen into the 64-bit element domain here, a no-op in the `u64` arm the lint sees. +#[expect(clippy::unnecessary_cast)] +pub fn elias_fano_encode( + array: ArrayView<'_, Primitive>, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + let dtype = array.dtype().clone(); + vortex_ensure!( + dtype.is_int(), + "Elias-Fano requires an integer dtype, got {dtype}" + ); + vortex_ensure!( + !dtype.is_nullable(), + "Elias-Fano requires a non-nullable dtype, got {dtype}" + ); + + let n = array.len(); + if n == 0 { + return empty(&dtype); + } + + // Work in sign-extended 64-bit patterns throughout; see `EliasFanoData::reference_bits`. + let (reference_bits, max_bits) = match_each_integer_ptype!(array.ptype(), |P| { + let values = array.as_slice::

(); + (values[0] as u64, values[n - 1] as u64) + }); + + let span = max_bits.wrapping_sub(reference_bits); + let lower_width = params::lower_width(span, n); + let upper_len = params::upper_len(span, n, lower_width)?; + let lower_mask = params::lower_mask(lower_width); + + let mut upper = UpperBuilder::new(usize::try_from(upper_len)?); + let mut lower = BufferMut::::with_capacity(n); + + match_each_integer_ptype!(array.ptype(), |P| { + let values = array.as_slice::

(); + // Check monotonicity in the value domain: an element is a modular difference, so on + // unsorted input the differences can still come out non-decreasing after wrapping. + let mut previous = values[0]; + for (index, &value) in values.iter().enumerate() { + vortex_ensure!( + value >= previous, + "Elias-Fano requires a non-decreasing sequence, but the value at index {index} \ + is below its predecessor" + ); + previous = value; + + let element = (value as u64).wrapping_sub(reference_bits); + let rank = index as u64; + upper.push(rank, (element >> lower_width) + rank + 1); + lower.push(element & lower_mask); + } + }); + + let (upper_bytes, sample_bytes) = upper.finish(n as u64, upper_len)?; + + let lower = pack_lower(lower.freeze(), lower_width, n)?; + + let data = EliasFanoData::try_new( + upper_bytes, + sample_bytes, + scalar_from_bits(&dtype, reference_bits)?, + scalar_from_bits(&dtype, max_bits)?, + lower_width, + upper_len, + 0, + )?; + EliasFano::try_new(data, lower, n) +} + +pub(crate) fn elias_fano_decompress( + array: &EliasFanoArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let ptype = array.dtype().as_ptype(); + if len == 0 { + return Ok(match_each_integer_ptype!(ptype, |P| { + PrimitiveArray::empty::

(NonNullable) + })); + } + + let first_rank = array.first_rank(); + let upper = array.upper_bits()?; + let upper_len = usize::try_from(array.upper_len())?; + let (_, samples1) = array.sample_bytes()?; + + // Trim the upper array to the window holding exactly our elements' set bits, so the walk below + // needs no per-element bound check and no early exit. Two sampled selects buy that. + let start = position_of_rank(&upper, samples1, upper_len, first_rank)?; + let end = position_of_rank(&upper, samples1, upper_len, first_rank + len as u64 - 1)? + 1; + let window = upper.slice(start..end); + + let lower_width = array.lower_width(); + let fold = Fold { + start, + first_rank, + reference_bits: array.reference_bits(), + lower_width, + lower_mask: params::lower_mask(lower_width), + }; + + Ok(match_each_integer_ptype!(ptype, |P| { + PrimitiveArray::new( + fold.decode::

(&window, array.lower(), len, ctx)?, + Validity::NonNullable, + ) + })) +} + +/// Reassembles elements from the two halves of the layout, in the column's own width. +struct Fold { + /// Bit position the upper window starts at, which its set-bit indices are relative to. + start: usize, + first_rank: u64, + reference_bits: u64, + lower_width: u8, + lower_mask: u64, +} + +impl Fold { + /// Decode `len` elements, taking high parts from `window`'s set bits and low parts from + /// `lower`, which is read one FastLanes block at a time. + fn decode( + &self, + window: &BitBuffer, + lower: &ArrayRef, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> + where + u64: AsPrimitive

, + { + let mut values = BufferMut::

::with_capacity(len); + let mut ones = window.set_indices(); + + if self.lower_width == 0 { + // Nothing is stored, so do not execute the slot just to read `len` zeros. + self.segment(&mut values, &mut ones, iter::repeat_n(0, len))?; + return self.finish(values, ones, len); + } + + // Window before reading: the child spans the whole encoded sequence, and a decode only ever + // wants the ranks inside its own slice. + let first = usize::try_from(self.first_rank)?; + let window_lower = lower.slice(first..first + len)?; + + // The alignment test is the cursor's, for the same reason: the unpack reinterprets the + // packed bytes as `&[u64]` unchecked, so an under-aligned buffer must take the fallback. + if let Some(packed) = window_lower.as_opt::() + && packed.patches().is_none() + && packed + .packed() + .as_host_opt() + .is_some_and(|buffer| buffer.is_aligned(Alignment::of::())) + { + let mut chunks = packed.unpacked_chunks::()?; + if let Some(initial) = chunks.initial() { + self.segment(&mut values, &mut ones, initial.iter().copied())?; + } + // A single-block child is covered by `initial` alone, and the later phases would hand + // that same block back. + if values.len() < len { + let mut full = chunks.full_chunks(); + while let Some(chunk) = full.next() { + self.segment(&mut values, &mut ones, chunk.iter().copied())?; + } + } + if values.len() < len + && let Some(trailer) = chunks.trailer() + { + self.segment(&mut values, &mut ones, trailer.iter().copied())?; + } + } else { + // The slot is patched, device-resident, or some other encoding after a rewrite. + let dense = window_lower.execute::(ctx)?; + self.segment(&mut values, &mut ones, dense.as_slice::().iter().copied())?; + } + + self.finish(values, ones, len) + } + + /// Fold one run of consecutive low parts onto the end of `values`. + fn segment>( + &self, + values: &mut BufferMut

, + ones: &mut BitIndexIterator<'_>, + lows: L, + ) -> VortexResult<()> + where + u64: AsPrimitive

, + { + for low in lows { + let rank = self.first_rank + values.len() as u64; + let position = ones.next().ok_or_else(|| { + vortex_err!("Elias-Fano upper array holds no element of rank {rank}") + })?; + let position = self.start + position; + // The inverse of the encoder's `position = (element >> lower_width) + rank + 1`. + let high = (position as u64).checked_sub(rank + 1).ok_or_else(|| { + vortex_err!( + "Elias-Fano upper array is malformed: the element of rank {rank} sits at bit \ + {position}, at or below its own rank" + ) + })?; + // The low bits are masked for the same reason as in the cursor's `lower_at`: only a + // bit-packed child's width is checkable at construction, so a patched or rewritten slot + // could otherwise carry bits above `lower_width` into the high part. + let bits = self + .reference_bits + .wrapping_add((high << self.lower_width) | (low & self.lower_mask)); + // Truncating the pattern to the column's width is exactly the two's complement result, + // signed or unsigned, because the reference was added in the same modular arithmetic. + values.push(bits.as_()); + } + Ok(()) + } + + /// The window holds exactly `len` set bits, and the child exactly `len` low parts, for any + /// array this crate builds — but `validate` does not check the upper buffer's contents, so a + /// corrupt file can hold a different number of either. + fn finish( + &self, + values: BufferMut

, + mut ones: BitIndexIterator<'_>, + len: usize, + ) -> VortexResult> { + vortex_ensure!( + values.len() == len && ones.next().is_none(), + "Elias-Fano upper array is malformed: expected exactly {len} set bits above their own \ + ranks, found {}", + values.len() + ); + Ok(values.freeze()) + } +} + +/// Builds the upper array and both sample tables together, in one pass over the elements. +/// +/// The tables have to be built here: a zero-sample is the position of a sampled *unset* bit, and +/// the unset runs are only known in order between two consecutive elements as they are written. +struct UpperBuilder { + bits: BitBufferMut, + samples0: BufferMut, + samples1: BufferMut, + /// The next unset-bit rank owed a sample. Sample 0 is never stored, for either table: the + /// sentinel puts the 0th unset bit at position 0 and the 0th set bit is the array's first, both + /// of which a reader can assume. + next_zero_sample: u64, +} + +impl UpperBuilder { + fn new(upper_len: usize) -> Self { + Self { + bits: BitBufferMut::new_unset(upper_len), + samples0: BufferMut::empty(), + samples1: BufferMut::empty(), + next_zero_sample: 1 << params::LOG_SAMPLING0, + } + } + + /// Record the element of rank `rank` as a set bit at `position`. + /// + /// Must be called with strictly increasing `rank` and `position`. + fn push(&mut self, rank: u64, position: u64) { + // Every unset bit below `position` but above the previous element's has exactly `rank` set + // bits before it, so its own rank is `its position - rank`. Emit a sample for each sampled + // rank that lands in that gap. + while self.next_zero_sample + rank < position { + self.samples0.push(self.next_zero_sample + rank); + self.next_zero_sample += 1 << params::LOG_SAMPLING0; + } + + debug_assert!(position < self.bits.len() as u64, "position out of bounds"); + self.bits.set(position as usize); + + if rank > 0 && rank.is_multiple_of(1 << params::LOG_SAMPLING1) { + self.samples1.push(position); + } + } + + /// Close the array, returning the upper bytes and both sample tables packed into one buffer. + /// + /// The seam is not returned: a reader recomputes it from the universe (see + /// [`params::num_samples0`]), and `validate_parts` fails the array if the two disagree. + fn finish(mut self, n: u64, upper_len: u64) -> VortexResult<(ByteBuffer, ByteBuffer)> { + // Sample the trailing unset bits past the last element, which `push` never reached. These + // are the bucket boundaries above the maximum element's high part, plus the guard zero. + while self.next_zero_sample + n < upper_len { + self.samples0.push(self.next_zero_sample + n); + self.next_zero_sample += 1 << params::LOG_SAMPLING0; + } + + // Both tables share one buffer, zeros first. A reader gets a single zero-copy mapping, and + // a query that reseats and then walks touches both within a few cache lines of each other. + self.samples0.extend_from_slice(self.samples1.as_slice()); + + let (_, _, upper_bytes) = self.bits.freeze().into_inner(); + Ok((upper_bytes, self.samples0.freeze().into_byte_buffer())) + } +} + +fn pack_lower(lower: Buffer, lower_width: u8, n: usize) -> VortexResult { + if lower_width == 0 { + // There is nothing to store. A constant array says so explicitly, and costs nothing on + // disk. + return Ok(ConstantArray::new(0u64, n).into_array()); + } + let lower = PrimitiveArray::new(lower, Validity::NonNullable); + // SAFETY: every value was masked to `lower_width` bits as it was pushed, so all pack losslessly + // and none needs a patch. The checked path would scan for a minimum and build a bit-width + // histogram to rediscover what the encoder already guaranteed. + Ok(unsafe { bitpack_encode_unchecked(lower, lower_width) }?.into_array()) +} + + +/// The degenerate zero-element array. +/// +/// This case is representable rather than rejected, so an empty chunk needs no special handling +/// upstream. The bounds go unused, because there is nothing to offset, and the two-bit upper array +/// holds just the sentinel and its guard. +fn empty(dtype: &DType) -> VortexResult { + let upper = BitBufferMut::new_unset(2).freeze(); + let (_, _, upper_bytes) = upper.into_inner(); + let data = EliasFanoData::try_new( + upper_bytes, + Buffer::::empty().into_byte_buffer(), + scalar_from_bits(dtype, 0)?, + scalar_from_bits(dtype, 0)?, + 0, + 2, + 0, + )?; + EliasFano::try_new( + data, + PrimitiveArray::empty::(NonNullable).into_array(), + 0, + ) +} diff --git a/encodings/elias-fano/src/compute/cast.rs b/encodings/elias-fano/src/compute/cast.rs new file mode 100644 index 00000000000..bf508a2f553 --- /dev/null +++ b/encodings/elias-fano/src/compute/cast.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_error::VortexResult; + +use crate::EliasFano; +use crate::array::EliasFanoArraySlotsExt; + +impl CastReduce for EliasFano { + /// Cast by rewriting the two universe bounds, and nothing else. + /// + /// The encoded bits never mention the ptype, so if both bounds are exactly representable in the + /// target the span and every derived quantity are unchanged and both buffers carry over. A cast + /// that cannot represent a bound returns `None` and leaves it to the generic path. + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + if !dtype.is_int() || dtype.is_nullable() { + return Ok(None); + } + if dtype == array.array().dtype() { + return Ok(Some(array.array().clone())); + } + + let data = array.data(); + let (Ok(reference), Ok(max)) = ( + data.reference_scalar().cast(dtype), + data.max_scalar().cast(dtype), + ) else { + return Ok(None); + }; + + Ok(Some( + EliasFano::try_new( + data.clone().with_bounds(reference, max), + array.lower().clone(), + array.len(), + )? + .into_array(), + )) + } +} diff --git a/encodings/elias-fano/src/compute/compare.rs b/encodings/elias-fano/src/compute/compare.rs new file mode 100644 index 00000000000..74fc9fab224 --- /dev/null +++ b/encodings/elias-fano/src/compute/compare.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Comparing an Elias-Fano array against a constant, without decoding it. +//! +//! The sequence is sorted, so the matching rows form a contiguous run — or, for `NotEq`, the +//! complement of one. Two sampled searches find its bounds and the answer is a bit-buffer fill. + +use std::ops::Range; + +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::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::validity::Validity; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; + +use crate::EliasFano; +use crate::EliasFanoCursor; + +impl CompareKernel for EliasFano { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Only a constant right-hand side reduces to a range. The adaptor has already normalised + // operand order, flipping the operator if the encoded array arrived on the right. + let Some(constant) = rhs.as_constant() else { + return Ok(None); + }; + // A null operand makes every row null rather than a range. The adaptor answers that before + // a kernel is reached, but this is a public trait impl and the cursor takes no nulls. + if constant.is_null() { + return Ok(None); + } + // The cursor probes in the array's own dtype. A literal differing only in nullability casts + // cleanly; one outside the column's range is left to the generic path, which promotes both. + let Ok(constant) = constant.cast(lhs.dtype()) else { + return Ok(None); + }; + + let len = lhs.len(); + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + + // `rank` and `rank_inclusive` bracket the run equal to the constant, and every comparison + // is one side of that pair, so no operator costs more than two searches and most cost one. + let mut cursor = EliasFanoCursor::try_new(lhs, ctx)?; + let result = match operator { + // These need only the lower bound. + CompareOperator::Lt => run(0..cursor.rank(&constant)?, len, nullability), + CompareOperator::Gte => run(cursor.rank(&constant)?..len, len, nullability), + // These need only the upper bound. + CompareOperator::Lte => run(0..cursor.rank_inclusive(&constant)?, len, nullability), + CompareOperator::Gt => run(cursor.rank_inclusive(&constant)?..len, len, nullability), + CompareOperator::Eq => { + let lo = cursor.rank(&constant)?; + run(lo..cursor.rank_inclusive(&constant)?, len, nullability) + } + // The one answer that is not a single run. + CompareOperator::NotEq => { + let lo = cursor.rank(&constant)?; + complement(lo..cursor.rank_inclusive(&constant)?, len, nullability) + } + }; + Ok(Some(result)) + } +} + +fn validity(nullability: Nullability) -> Validity { + match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::AllValid, + } +} + +/// A boolean array true exactly on `range`. An all-true or all-false answer becomes a constant, so +/// whatever consumes it can skip the array entirely. +fn run(range: Range, len: usize, nullability: Nullability) -> ArrayRef { + if range.start >= range.end { + return ConstantArray::new(Scalar::bool(false, nullability), len).into_array(); + } + if range.start == 0 && range.end == len { + return ConstantArray::new(Scalar::bool(true, nullability), len).into_array(); + } + let mut buffer = BitBufferMut::new_unset(len); + buffer.fill_range(range.start, range.end, true); + BoolArray::new(buffer.freeze(), validity(nullability)).into_array() +} + +fn complement(range: Range, len: usize, nullability: Nullability) -> ArrayRef { + if range.start >= range.end { + return ConstantArray::new(Scalar::bool(true, nullability), len).into_array(); + } + if range.start == 0 && range.end == len { + return ConstantArray::new(Scalar::bool(false, nullability), len).into_array(); + } + let mut buffer = BitBufferMut::new_set(len); + buffer.fill_range(range.start, range.end, false); + BoolArray::new(buffer.freeze(), validity(nullability)).into_array() +} diff --git a/encodings/elias-fano/src/compute/filter.rs b/encodings/elias-fano/src/compute/filter.rs new file mode 100644 index 00000000000..6acb3869c92 --- /dev/null +++ b/encodings/elias-fano/src/compute/filter.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Filtering an Elias-Fano array through the cursor: the same shape and crossover as +//! [`take`](super::take). Selected rows arrive ascending, which is the cursor's best case. + +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::filter::FilterKernel; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_integer_ptype; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::EliasFano; +use crate::EliasFanoCursor; +use crate::compute::take::BULK_DECODE_THRESHOLD; + +impl FilterKernel for EliasFano { + fn filter( + array: ArrayView<'_, Self>, + mask: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let selected = mask + .values() + .vortex_expect("FilterKernel precondition: mask is Mask::Values"); + + if selected.true_count() * BULK_DECODE_THRESHOLD > array.len() { + let decoded = array.array().clone().execute::(ctx)?; + return decoded.into_array().filter(mask.clone()).map(Some); + } + + let ptype = array.dtype().as_ptype(); + let reference_bits = array.reference_bits(); + let validity = array.validity()?.filter(mask)?; + + let mut cursor = EliasFanoCursor::try_new(array, ctx)?; + let filtered = match_each_integer_ptype!(ptype, |P| { + PrimitiveArray::new( + gather_rows::

(&mut cursor, selected.indices(), reference_bits)?, + validity, + ) + }); + Ok(Some(filtered.into_array())) + } +} + +/// The selected rows, in the column's own width. +fn gather_rows( + cursor: &mut EliasFanoCursor<'_>, + positions: &[usize], + reference_bits: u64, +) -> VortexResult> +where + u64: AsPrimitive

, +{ + let mut values = BufferMut::

::with_capacity(positions.len()); + for &index in positions { + let bits = reference_bits.wrapping_add(cursor.access_element(index)?); + // Truncating the pattern to the column's width is exactly the two's complement result, + // signed or unsigned, because the reference was added in the same modular arithmetic. + values.push(bits.as_()); + } + Ok(values.freeze()) +} diff --git a/encodings/elias-fano/src/compute/is_sorted.rs b/encodings/elias-fano/src/compute/is_sorted.rs new file mode 100644 index 00000000000..f439332f202 --- /dev/null +++ b/encodings/elias-fano/src/compute/is_sorted.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::EliasFano; + +/// Elias-Fano-specific `is_sorted` kernel. Sortedness is a precondition of the encoding — the upper +/// array only decodes correctly for a non-decreasing sequence — so the answer needs no data. +/// +/// Strict sortedness is a different question and is declined. Duplicates are legal — an empty list +/// contributes two identical offsets — and finding out whether any are present means comparing +/// adjacent low bits, which is a full scan. Returning `None` lets the generic path do that. +#[derive(Debug)] +pub(crate) struct EliasFanoIsSortedKernel; + +impl DynAggregateKernel for EliasFanoIsSortedKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(None); + }; + if options.strict || !batch.is::() { + return Ok(None); + } + Ok(Some(IsSorted::make_partial( + batch, + true, + options.strict, + ctx, + )?)) + } +} diff --git a/encodings/elias-fano/src/compute/min_max.rs b/encodings/elias-fano/src/compute/min_max.rs new file mode 100644 index 00000000000..758ad3ecc01 --- /dev/null +++ b/encodings/elias-fano/src/compute/min_max.rs @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::min_max::MinMax; +use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::EliasFano; +use crate::EliasFanoCursor; + +/// Elias-Fano-specific min/max kernel: the sequence is sorted, so two `select1`s find the extremes. +/// +/// Deliberately not read from the `reference` and `max` metadata, which describe the *encoded* +/// universe and survive slicing, so on a sliced array they are not its extremes. +/// +/// Each end goes through a one-element slice rather than one cursor over the whole array: a cursor +/// sizes its low-bits view to the array it opens on, and materialises the slot whole if it cannot +/// be read in place. +#[derive(Debug)] +pub(crate) struct EliasFanoMinMaxKernel; + +impl DynAggregateKernel for EliasFanoMinMaxKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() { + return Ok(None); + } + let Some(array) = batch.as_opt::() else { + return Ok(None); + }; + + let struct_dtype = make_minmax_dtype(batch.dtype()); + if array.is_empty() { + return Ok(Some(Scalar::null(struct_dtype))); + } + + let last = array.len() - 1; + let min = element_at(batch, 0, ctx)?; + let max = element_at(batch, last, ctx)?; + + Ok(Some(Scalar::struct_(struct_dtype, vec![min, max]))) + } +} + +fn element_at(array: &ArrayRef, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + let one = array.slice(index..index + 1)?; + let Some(one) = one.as_opt::() else { + // Slicing normally reduces into the encoding, but it is not obliged to. + return one.execute_scalar(0, ctx); + }; + EliasFanoCursor::try_new(one, ctx)?.access(0) +} diff --git a/encodings/elias-fano/src/compute/mod.rs b/encodings/elias-fano/src/compute/mod.rs new file mode 100644 index 00000000000..e2175357537 --- /dev/null +++ b/encodings/elias-fano/src/compute/mod.rs @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod cast; +mod compare; +mod filter; +pub(crate) mod is_sorted; +pub(crate) mod min_max; +mod slice; +pub(crate) mod take; diff --git a/encodings/elias-fano/src/compute/slice.rs b/encodings/elias-fano/src/compute/slice.rs new file mode 100644 index 00000000000..c02a0ec04bc --- /dev/null +++ b/encodings/elias-fano/src/compute/slice.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_error::VortexResult; + +use crate::EliasFano; +use crate::array::EliasFanoArraySlotsExt; + +impl SliceReduce for EliasFano { + /// Slice by recording where the slice starts, leaving every buffer alone: the sample tables + /// hold *absolute* bit positions and the low-bits child is packed in 1024-element blocks, so + /// one rank offset covers both. See + /// [`EliasFanoData::first_rank`](crate::EliasFanoData::first_rank). + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + let data = array + .data() + .clone() + .with_first_rank(array.first_rank() + range.start as u64); + Ok(Some( + EliasFano::try_new(data, array.lower().clone(), range.len())?.into_array(), + )) + } +} diff --git a/encodings/elias-fano/src/compute/take.rs b/encodings/elias-fano/src/compute/take.rs new file mode 100644 index 00000000000..918f9ae4d29 --- /dev/null +++ b/encodings/elias-fano/src/compute/take.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Taking arbitrary rows out of an Elias-Fano array through the cursor. + +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::dict::TakeExecute; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::AllOr; + +use crate::EliasFano; +use crate::EliasFanoCursor; + +/// The whole array is decoded in one pass, rather than walked through the cursor, once a request +/// reaches roughly one row in `BULK_DECODE_THRESHOLD` of it. +/// +/// Same constant as `vortex-fastlanes`'s take kernel, whose reasoning transfers because the +/// low-bits child *is* bit-packed. Elias-Fano's extra sampled select per index only pushes the true +/// crossover further toward bulk, so reusing 8 errs safe. +pub(crate) const BULK_DECODE_THRESHOLD: usize = 8; + +impl TakeExecute for EliasFano { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if indices.len() * BULK_DECODE_THRESHOLD > array.len() { + let decoded = array.array().clone().execute::(ctx)?; + return decoded.into_array().take(indices.clone()).map(Some); + } + + // A null index selects nothing, so its payload is not a position: it is whatever the slot + // happened to hold, and reading or bounds-checking one turns a legal take into a failure. + // Those rows come back null either way, so the walk below just skips them. + let selected = indices.validity()?.execute_mask(indices.len(), ctx)?; + let valid: Option<&BitBuffer> = match selected.bit_buffer() { + AllOr::All => None, + // Nothing to read at all, so answer before opening a cursor — which would materialise + // a low-bits child that cannot be read in place, for reads that never happen. + AllOr::None => { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + } + AllOr::Some(valid) => Some(valid), + }; + + let ptype = array.dtype().as_ptype(); + let reference_bits = array.reference_bits(); + let taken_validity = array.validity()?.take(indices)?; + let indices = indices.clone().execute::(ctx)?; + + let mut cursor = EliasFanoCursor::try_new(array, ctx)?; + let taken = gather( + &mut cursor, + &indices, + valid, + reference_bits, + taken_validity, + ptype, + )?; + Ok(Some(taken.into_array())) + } +} + +fn gather( + cursor: &mut EliasFanoCursor<'_>, + indices: &PrimitiveArray, + valid: Option<&BitBuffer>, + reference_bits: u64, + validity: Validity, + ptype: PType, +) -> VortexResult { + Ok(match_each_integer_ptype!(ptype, |P| { + match_each_integer_ptype!(indices.ptype(), |I| { + PrimitiveArray::new( + gather_rows::(cursor, indices.as_slice::(), valid, reference_bits)?, + validity, + ) + }) + })) +} + +/// The requested rows, in the column's own width. +/// +/// `valid` is `None` when every index is a real position. Otherwise the rows it marks invalid are +/// left as zero and their index payloads are never read, which is what the caller's validity +/// already says about them. +fn gather_rows( + cursor: &mut EliasFanoCursor<'_>, + indices: &[I], + valid: Option<&BitBuffer>, + reference_bits: u64, +) -> VortexResult> +where + u64: AsPrimitive

, +{ + let mut values = BufferMut::

::with_capacity(indices.len()); + for (position, &raw) in indices.iter().enumerate() { + if valid.is_some_and(|valid| !valid.value(position)) { + values.push(P::default()); + continue; + } + // Refused rather than wrapped into range; `access_element` bounds-checks what survives. + let index = raw + .to_usize() + .ok_or_else(|| vortex_err!("Elias-Fano take index {raw} is not a position"))?; + let bits = reference_bits.wrapping_add(cursor.access_element(index)?); + // Truncating the pattern to the column's width is exactly the two's complement result, + // signed or unsigned, because the reference was added in the same modular arithmetic. + values.push(bits.as_()); + } + Ok(values.freeze()) +} diff --git a/encodings/elias-fano/src/cursor.rs b/encodings/elias-fano/src/cursor.rs new file mode 100644 index 00000000000..da06ccc772e --- /dev/null +++ b/encodings/elias-fano/src/cursor.rs @@ -0,0 +1,577 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Random access and predecessor search over an encoded sequence, all in `O(1)`: `access` reads the +//! value at an index, `rank` counts the values strictly below a probe, and `next_geq` finds the +//! first value at or above one. +//! +//! The cursor is stateful on purpose. Probes usually arrive in ascending order — merge joins, +//! intersections, scans of list offsets — so it remembers where it stopped and walks forward from +//! there when the next probe is close, touching neither sample table. See [`crate::params`] for the +//! layout these operations read. + +use fastlanes::BitPacking; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_buffer::Alignment; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; +use vortex_fastlanes::bitpack_decompress::unpack_single_primitive; + +use crate::EliasFano; +use crate::array::EliasFanoSlotsView; +use crate::array::read_sample; +use crate::array::scalar_bits; +use crate::array::scalar_from_bits; +use crate::params::LINEAR_SCAN_THRESHOLD; +use crate::params::LOG_SAMPLING0; +use crate::params::LOG_SAMPLING1; +use crate::params::lower_mask; + +/// A whole FastLanes block is bulk-unpacked once *more* than this many reads have landed inside it. +/// +/// Unpacking all 1024 values costs roughly nine single-value unpacks, so that is the crossover. +/// Mirrors the crate-private `UNPACK_CHUNK_THRESHOLD` in `vortex-fastlanes`. +const BULK_UNPACK_THRESHOLD: usize = 8; + +/// The bit position of the set bit belonging to absolute rank `rank`, as a sampled `select1`. +/// +/// `samples1` bounds the search window at `1 << LOG_SAMPLING1` ones, about one cache line. Sample 0 +/// is not stored: rank 0 is always the first set bit. Free-standing because the bulk decode needs +/// the window boundaries without a cursor's low-bits reader. +pub(crate) fn position_of_rank( + upper: &BitBuffer, + samples1: &[u8], + upper_len: usize, + rank: u64, +) -> VortexResult { + let sample = (rank >> LOG_SAMPLING1) as usize; + let start = if sample == 0 { + 0 + } else { + usize::try_from(read_sample(samples1, sample - 1))? + }; + let nth = usize::try_from(rank - ((sample as u64) << LOG_SAMPLING1))?; + upper + .select_range(start, upper_len, nth) + .map(|offset| start + offset) + .ok_or_else(|| vortex_err!("Elias-Fano upper array holds no element of rank {rank}")) +} + +/// Where a probe value falls relative to the encoded universe. +enum Bound { + /// The probe sits below the reference, so it is below every element and its rank is 0. + Below, + /// The probe sits inside `reference..=max`, and this is its offset from the reference. + Inside(u64), + /// The probe sits above the encoded maximum, so it is above every element and its rank is the + /// array's length. + Above, +} + +/// How the low bits of each element can be read. +enum LowerBits<'a> { + /// The width is zero, so there are no low bits to read and nothing is stored. + Zero, + /// The normal case, where the low bits are read straight out of the FastLanes-packed child in + /// place. + Packed { + packed: &'a [u64], + bit_width: usize, + /// The child's own sub-block offset, which `unpack_single_primitive` does not apply + /// itself (unlike `unpack_single`). Forgetting it is a silent wrong answer, not a panic. + child_offset: usize, + }, + /// The fallback for a child that cannot be read in place — patches, device memory, + /// under-aligned, or a slot some rewrite replaced. The low bits are materialised once, up + /// front. + Dense { + /// Low bits for ranks `base..base + values.len()` only, not the whole child, so a slice + /// near the end of a long sequence does not materialise everything before it. + values: Buffer, + /// The absolute rank `values[0]` holds, i.e. the array's `first_rank`. + base: u64, + }, +} + +/// Where the cursor currently sits, which is one element together with everything already known +/// about it. +struct Seat { + /// Bit position of this element's set bit in the upper array. + position: usize, + /// Absolute rank within the encoded sequence, so `first_rank` is already included. + rank: u64, + /// The element, i.e. the value minus the reference. + element: u64, +} + +/// A stateful reader over an [`EliasFanoArray`](crate::EliasFanoArray). +pub struct EliasFanoCursor<'a> { + upper: BitBuffer, + /// Position of every `1 << LOG_SAMPLING0`-th unset bit, as raw little-endian `u64`s. + samples0: &'a [u8], + /// Position of every `1 << LOG_SAMPLING1`-th set bit, as raw little-endian `u64`s. + samples1: &'a [u8], + lower: LowerBits<'a>, + /// Destination for a bulk-unpacked FastLanes block, together with which block it currently + /// holds. It is allocated on the first bulk unpack, so a single point lookup never pays for it. + scratch: Option>, + scratch_chunk: Option, + /// The block the last few reads landed in, and how many landed there. Together they drive the + /// switch over to [`BULK_UNPACK_THRESHOLD`]. + hot_chunk: Option, + hot_reads: usize, + dtype: &'a DType, + ptype: PType, + reference_bits: u64, + span: u64, + lower_width: u8, + first_rank: u64, + len: usize, + upper_len: usize, + seat: Option, + /// The last probe [`Self::next_geq_element`] answered, and the answer it gave. A merge join + /// probes the same value from both sides, so the immediate repeat is common. + /// + /// Memoising the answer is not the same as reusing the seat: with duplicates a rank must count + /// from the first occurrence, and the found element is not generally the probe. + last_answer: Option<(u64, (usize, Option))>, +} + +impl<'a> EliasFanoCursor<'a> { + /// Open a cursor over `array`. `ctx` is only used if the low-bits child has to be materialised. + pub fn try_new( + array: ArrayView<'a, EliasFano>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let data = array.data(); + let (samples0, samples1) = data.sample_bytes()?; + // The slots view borrows the array behind the `ArrayView`, which outlives the cursor; the + // `lower()` accessor would borrow the (`Copy`, stack-local) view itself. + let lower = EliasFanoSlotsView::from_slots(array.slots()).lower; + let lower = LowerBits::try_new( + lower, + data.lower_width(), + data.first_rank(), + array.len(), + ctx, + )?; + + Ok(Self { + upper: data.upper_bits()?, + samples0, + samples1, + lower, + scratch: None, + scratch_chunk: None, + hot_chunk: None, + hot_reads: 0, + dtype: array.array().dtype(), + ptype: array.array().dtype().as_ptype(), + reference_bits: data.reference_bits(), + span: data.span(), + lower_width: data.lower_width(), + first_rank: data.first_rank(), + len: array.len(), + upper_len: usize::try_from(data.upper_len())?, + seat: None, + last_answer: None, + }) + } + + #[inline] + fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The value at logical index `index`. + pub fn access(&mut self, index: usize) -> VortexResult { + let element = self.access_element(index)?; + scalar_from_bits(self.dtype, self.reference_bits.wrapping_add(element)) + } + + /// The element (value minus reference) at logical index `index`, in one sampled `select1`. + pub(crate) fn access_element(&mut self, index: usize) -> VortexResult { + if index >= self.len { + vortex_bail!(OutOfBounds: index, 0usize, self.len); + } + let rank = self.first_rank + index as u64; + + // Sequential access is the common case, and stepping to the next set bit is cheaper than + // a select from the sample table, so reuse the seat when it lines up. + match &self.seat { + Some(seat) if seat.rank == rank => {} + Some(seat) if seat.rank + 1 == rank => self.advance()?, + _ => self.seek(rank)?, + } + Ok(self.seated().element) + } + + /// The number of elements strictly less than `value`. + /// + /// This is `search_sorted`'s left bound, and needs no rank directory: see the + /// `rank1(select0(h)) == select0(h) - h` identity in the crate-private `params` module. + pub fn rank(&mut self, value: &Scalar) -> VortexResult { + Ok(self.next_geq(value)?.0) + } + + /// The number of elements at or below `value`, i.e. `search_sorted`'s right bound. With + /// [`Self::rank`] this brackets the run of elements equal to `value`. + /// + /// The count is the rank of `value`'s successor, taken in the element domain because the value + /// domain would overflow at the top of the ptype. + pub fn rank_inclusive(&mut self, value: &Scalar) -> VortexResult { + match self.locate(value)? { + Bound::Below => Ok(0), + Bound::Above => Ok(self.len), + Bound::Inside(element) => match element.checked_add(1) { + Some(successor) => Ok(self.next_geq_element(successor)?.0), + None => Ok(self.len), + }, + } + } + + /// The first value at or above `value`, together with the number of elements strictly below it. + /// + /// Returns `(len, None)` when every element is below `value`. + pub fn next_geq(&mut self, value: &Scalar) -> VortexResult<(usize, Option)> { + let (rank, element) = match self.locate(value)? { + Bound::Below => ( + 0, + (!self.is_empty()) + .then(|| self.access_element(0)) + .transpose()?, + ), + Bound::Above => (self.len, None), + Bound::Inside(element) => self.next_geq_element(element)?, + }; + let value = element + .map(|element| scalar_from_bits(self.dtype, self.reference_bits.wrapping_add(element))) + .transpose()?; + Ok((rank, value)) + } + + /// [`Self::next_geq`] in the element domain. `element` must be no greater than the span; + /// callers holding a value want [`Self::next_geq`], which classifies it first. + pub(crate) fn next_geq_element(&mut self, element: u64) -> VortexResult<(usize, Option)> { + if self.is_empty() { + return Ok((0, None)); + } + if element > self.span { + return Ok((self.len, None)); + } + if let Some((probe, answer)) = self.last_answer + && probe == element + { + return Ok(answer); + } + + let end_rank = self.first_rank + self.len as u64; + + // Stepping forward is only sound when the probe is strictly above the seated element: on + // equality the seat may sit past earlier duplicates, and a rank counts from the first. + // The walk is bounded in buckets, measurable up front, and in elements, which is what it + // pays — a bucket of duplicates is arbitrarily many elements deep. + if let Some(seat) = &self.seat + && element > seat.element + && (element >> self.lower_width) - (seat.element >> self.lower_width) + < LINEAR_SCAN_THRESHOLD + && let Some(answer) = self.walk_to(element, end_rank, LINEAR_SCAN_THRESHOLD)? + { + return Ok(self.memoise(element, answer)); + } + + let answer = self.search_bucket(element, end_rank)?; + Ok(self.memoise(element, answer)) + } + + /// Locate the first element `>= element` from a standing start. + /// + /// One sampled `select0` finds where the probe's bucket begins; everything before it is below + /// the probe. A bucket normally holds about one element, so a short walk usually ends it. + /// Otherwise the bucket is a run of near-duplicates, and a second `select0` brackets it for + /// bisection — inside a bucket every element shares a high part, so the low bits alone order it + /// and a bisection step needs no `select`. + fn search_bucket(&mut self, element: u64, end_rank: u64) -> VortexResult<(usize, Option)> { + let high = element >> self.lower_width; + let start = self.rank_of_bucket(high)?.max(self.first_rank); + if start >= end_rank { + return Ok((self.len, None)); + } + self.seek(start)?; + + if let Some(answer) = self.walk_to(element, end_rank, LINEAR_SCAN_THRESHOLD)? { + return Ok(answer); + } + + // The rank the walk stopped *on* was never compared, so the bisection includes it. It may + // also be the first of a later bucket, which the clamp turns into an empty range and leaves + // as the answer — correct, since a greater high part is already above the probe. + let low = element & lower_mask(self.lower_width); + let mut lo = self.seated().rank; + let mut hi = self.rank_of_bucket(high + 1)?.clamp(lo, end_rank); + while lo < hi { + let mid = lo + (hi - lo) / 2; + if self.lower_at(mid) < low { + lo = mid + 1; + } else { + hi = mid; + } + } + + // `lo` is the first element at or above the probe: either one inside the bucket, or the + // bucket's successor, which a greater high part already puts above the probe. + if lo >= end_rank { + return Ok((self.len, None)); + } + self.seek(lo)?; + Ok((self.relative_rank(lo), Some(self.seated().element))) + } + + /// Step forward at most `budget` times looking for the first element `>= element`, leaving the + /// cursor where it stopped. `None` means the budget ran out with everything still below the + /// probe, so stepping is no longer the way to find it. + fn walk_to( + &mut self, + element: u64, + end_rank: u64, + budget: u64, + ) -> VortexResult)>> { + for _ in 0..budget { + let seat = self.seated(); + let (found, rank) = (seat.element, seat.rank); + if found >= element { + return Ok(Some((self.relative_rank(rank), Some(found)))); + } + if rank + 1 >= end_rank { + return Ok(Some((self.len, None))); + } + self.advance()?; + } + Ok(None) + } + + fn memoise(&mut self, probe: u64, answer: (usize, Option)) -> (usize, Option) { + self.last_answer = Some((probe, answer)); + answer + } + + fn position_of(&self, rank: u64) -> VortexResult { + position_of_rank(&self.upper, self.samples1, self.upper_len, rank) + } + + /// The absolute rank of the first element whose high part is at least `high`: a sampled + /// `select0` followed by the rank identity, windowed the same way as [`position_of_rank`]. + fn rank_of_bucket(&self, high: u64) -> VortexResult { + let sample = (high >> LOG_SAMPLING0) as usize; + let start = if sample == 0 { + 0 + } else { + usize::try_from(read_sample(self.samples0, sample - 1))? + }; + let nth = usize::try_from(high - ((sample as u64) << LOG_SAMPLING0))?; + let position = self + .upper + .select_zero_range(start, self.upper_len, nth) + .map(|offset| (start + offset) as u64) + .ok_or_else(|| { + vortex_err!("Elias-Fano upper array holds no bucket boundary for high part {high}") + })?; + // rank1(select0(high)) == select0(high) - high. Checked because the upper buffer's contents + // are never validated: a corrupt one must raise rather than underflow. + position.checked_sub(high).ok_or_else(|| { + vortex_err!( + "Elias-Fano upper array is malformed: bucket boundary for high part {high} sits at \ + bit {position}, below its own rank" + ) + }) + } + + fn seek(&mut self, rank: u64) -> VortexResult<()> { + let position = self.position_of(rank)?; + self.reseat(position, rank) + } + + fn advance(&mut self) -> VortexResult<()> { + let seat = self.seated(); + let (from, rank) = (seat.position + 1, seat.rank + 1); + let offset = self + .upper + .select_range(from, self.upper_len, 0) + .ok_or_else(|| vortex_err!("Elias-Fano upper array holds no element of rank {rank}"))?; + self.reseat(from + offset, rank) + } + + fn reseat(&mut self, position: usize, rank: u64) -> VortexResult<()> { + // The inverse of the encoder's `position = (element >> lower_width) + rank + 1`. Checked + // for the same reason as in `rank_of_bucket`. + let high = (position as u64).checked_sub(rank + 1).ok_or_else(|| { + vortex_err!( + "Elias-Fano upper array is malformed: the element of rank {rank} sits at bit \ + {position}, at or below its own rank" + ) + })?; + let element = (high << self.lower_width) | self.lower_at(rank); + self.seat = Some(Seat { + position, + rank, + element, + }); + Ok(()) + } + + fn seated(&self) -> &Seat { + self.seat + .as_ref() + .vortex_expect("the cursor is seated before it is read") + } + + #[inline] + fn relative_rank(&self, rank: u64) -> usize { + rank.checked_sub(self.first_rank) + .and_then(|relative| usize::try_from(relative).ok()) + .vortex_expect("rank is within this array") + } + + /// The low bits of the element at absolute rank `rank`. + /// + /// Masked, not trusted. Only a bit-packed child's width is checkable at construction — see + /// `validate_parts` — so a patched or rewritten slot can arrive as a plain `u64` array carrying + /// bits above `lower_width`. Those would bleed into the high part in [`Self::reseat`] and + /// misorder the bisection in [`Self::search_bucket`], both silently. + fn lower_at(&mut self, rank: u64) -> u64 { + self.lower_at_unmasked(rank) & lower_mask(self.lower_width) + } + + fn lower_at_unmasked(&mut self, rank: u64) -> u64 { + // Copy the descriptor out before touching the scratch buffer: the packed slice borrows the + // array, not `self`, so this keeps the borrow checker out of the way. + let (packed, bit_width, child_offset) = match &self.lower { + LowerBits::Zero => return 0, + LowerBits::Dense { values, base } => return values[(rank - base) as usize], + LowerBits::Packed { + packed, + bit_width, + child_offset, + } => (*packed, *bit_width, *child_offset), + }; + + let index = rank as usize + child_offset; + let chunk = index / FL_CHUNK_SIZE; + let within_chunk = index % FL_CHUNK_SIZE; + + if self.scratch_chunk == Some(chunk) { + return self.scratch_slice()[within_chunk]; + } + + if self.hot_chunk == Some(chunk) { + self.hot_reads += 1; + } else { + self.hot_chunk = Some(chunk); + self.hot_reads = 1; + } + + let elems_per_chunk = 128 * bit_width / size_of::(); + if self.hot_reads > BULK_UNPACK_THRESHOLD { + let block = &packed[chunk * elems_per_chunk..][..elems_per_chunk]; + let scratch = self + .scratch + .get_or_insert_with(|| Box::new([0u64; FL_CHUNK_SIZE])); + // SAFETY: `block` is exactly `elems_per_chunk` packed values, and `scratch` is exactly + // one FastLanes block of 1024 values, which is what `unchecked_unpack` requires. + unsafe { BitPacking::unchecked_unpack(bit_width, block, scratch.as_mut_slice()) }; + self.scratch_chunk = Some(chunk); + return self.scratch_slice()[within_chunk]; + } + + // SAFETY: `packed` is `BitPackedData`'s own buffer, whose length the array's validation + // already tied to `bit_width` and a whole number of blocks, and `index` is within the + // child's length because `first_rank + len` is validated against it. + unsafe { unpack_single_primitive::(packed, bit_width, index) } + } + + fn scratch_slice(&self) -> &[u64; FL_CHUNK_SIZE] { + self.scratch + .as_deref() + .vortex_expect("the scratch block is allocated before it is read") + } + + fn locate(&self, value: &Scalar) -> VortexResult { + if value.dtype() != self.dtype { + vortex_bail!( + "Elias-Fano probe dtype {} does not match array dtype {}", + value.dtype(), + self.dtype + ); + } + if value.is_null() { + vortex_bail!("Elias-Fano cannot be probed with a null value"); + } + let bits = scalar_bits(value); + let element = bits.wrapping_sub(self.reference_bits); + if element <= self.span { + return Ok(Bound::Inside(element)); + } + // Outside the universe. The subtraction above wraps for values below the reference and + // overshoots for values above the max, so tell the two apart in the ptype's own ordering. + let ptype = self.ptype; + let reference_bits = self.reference_bits; + let below = match_each_integer_ptype!(ptype, |P| { (bits as P) < (reference_bits as P) }); + Ok(if below { Bound::Below } else { Bound::Above }) + } +} + +impl<'a> LowerBits<'a> { + /// Choose how to read the low bits for absolute ranks `first_rank..first_rank + len`. + fn try_new( + lower: &'a ArrayRef, + lower_width: u8, + first_rank: u64, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if lower_width == 0 { + return Ok(LowerBits::Zero); + } + // `as_opt`, never `as_`: with the experimental patched-array plugin enabled the slot comes + // back from a file as `Patched(BitPacked)`, and a rewrite may replace it outright. + if let Some(packed) = lower.as_opt::() + && packed.patches().is_none() + && packed + .packed() + .as_host_opt() + .is_some_and(|buffer| buffer.is_aligned(Alignment::of::())) + { + return Ok(LowerBits::Packed { + // `.data()` rather than the `Deref`, which would borrow the view rather than the + // array behind it and so not live long enough. + packed: packed.data().packed_slice::(), + bit_width: packed.bit_width() as usize, + child_offset: packed.offset() as usize, + }); + } + // Window before executing: the child spans the whole encoded sequence, and a cursor only + // ever asks for ranks inside its own slice. + let first = usize::try_from(first_rank)?; + Ok(LowerBits::Dense { + values: lower + .slice(first..first + len)? + .execute::(ctx)? + .into_buffer::(), + base: first_rank, + }) + } +} diff --git a/encodings/elias-fano/src/kernel.rs b/encodings/elias-fano/src/kernel.rs new file mode 100644 index 00000000000..ff73b372ada --- /dev/null +++ b/encodings/elias-fano/src/kernel.rs @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Registering the pushdown kernels, each letting a parent operation execute against the encoded +//! array instead of canonicalising first. An unregistered kernel is still correct, just slower, so +//! the tests assert the pushdown is actually taken. + +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; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; + +use crate::EliasFano; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), EliasFano, CompareExecuteAdaptor(EliasFano)); + kernels.register_execute_parent_kernel(Filter.id(), EliasFano, FilterExecuteAdaptor(EliasFano)); + kernels.register_execute_parent_kernel(Dict.id(), EliasFano, TakeExecuteAdaptor(EliasFano)); +} diff --git a/encodings/elias-fano/src/lib.rs b/encodings/elias-fano/src/lib.rs new file mode 100644 index 00000000000..dd62adbec83 --- /dev/null +++ b/encodings/elias-fano/src/lib.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +// Elias-Fano works in sign-extended 64-bit patterns and narrows back to the array's own width on +// the way out; see `EliasFanoData::reference_bits`. Both halves of that pair are exact. +#![expect(clippy::cast_possible_truncation)] + +//! Elias-Fano encoding for monotonically non-decreasing integer sequences. +//! +//! Stores about `log2(u / n) + 2` bits per value for `n` values over a universe of `u`, while still +//! answering random access, rank, and predecessor queries in constant time. Against bit-packing at +//! `ceil(log2(u))` bits per value the saving is `log2(n)` bits, so it widens with row count. +//! +//! Inputs must be non-decreasing and non-nullable. Duplicates are fine; anything else is refused +//! rather than silently mangled. +//! +//! See [`elias_fano_encode`] for the compression entry point, [`EliasFanoCursor`] for point +//! lookups, rank, and seeks, and [`initialize`] to register the encoding in a session. The +//! crate-private `params` module documents the bit layout; its sampled select index follows +//! Vigna's [broadword][] construction. `rise-rs` and `vers` were consulted as references; no code +//! is taken from either. +//! +//! [broadword]: https://vigna.di.unimi.it/ftp/papers/Broadword.pdf + +mod array; +mod compress; +mod compute; +mod cursor; +mod kernel; +pub(crate) mod params; +mod rules; + +pub use array::EliasFano; +pub use array::EliasFanoArray; +pub use array::EliasFanoArraySlotsExt; +pub use array::EliasFanoData; +pub use array::EliasFanoMetadata; +pub use array::EliasFanoSlots; +pub use compress::elias_fano_encode; +pub use cursor::EliasFanoCursor; +use vortex_array::ArrayVTable; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; +use vortex_array::aggregate_fn::fns::min_max::MinMax; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Initialize the Elias-Fano encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(EliasFano); + kernel::initialize(session); + + // Both answer from the layout rather than the data. + session.aggregate_fns().register_aggregate_kernel( + EliasFano.id(), + Some(MinMax.id()), + &compute::min_max::EliasFanoMinMaxKernel, + ); + session.aggregate_fns().register_aggregate_kernel( + EliasFano.id(), + Some(IsSorted.id()), + &compute::is_sorted::EliasFanoIsSortedKernel, + ); +} + +// TODO(reza): add an integer scheme in `vortex-btrblocks`, so the compressor can choose this +// encoding itself rather than it having to be applied explicitly. + +#[cfg(test)] +mod tests; diff --git a/encodings/elias-fano/src/params.rs b/encodings/elias-fano/src/params.rs new file mode 100644 index 00000000000..b268389cf7a --- /dev/null +++ b/encodings/elias-fano/src/params.rs @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The bit layout's geometry, as pure arithmetic over `(span, n)`. The encoder writes the results +//! into the metadata and `validate_parts` re-derives them, refusing an array that disagrees. +//! +//! Each element splits into an `l`-bit low part, bit-packed in a child slot, and a high part +//! `element >> l` set as one bit at position `high + index + 1`. The `+ index` keeps positions +//! distinct when elements share a high part, so reading element `i` is a `select1` and the inverse +//! is `high = position - index - 1`. The `+ 1` sentinel aligns the unset bits with the high parts, +//! giving `rank1(select0(h)) == select0(h) - h`, so one `select0` counts the elements below a high +//! part with no rank directory stored. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +/// One zero-sample is stored per `1 << LOG_SAMPLING0` unset bits of the upper array. +/// +/// The upper array is roughly 50% dense, so 512 zeros span about 512 bits — one or two 64-byte +/// chunks, the window [`BitBuffer::select_zero_range`](vortex_buffer::BitBuffer::select_zero_range) +/// is fastest over. [`LOG_SAMPLING1`] is sized the same way. +pub(crate) const LOG_SAMPLING0: usize = 9; + +/// One one-sample is stored per `1 << LOG_SAMPLING1` set bits of the upper array. +pub(crate) const LOG_SAMPLING1: usize = 8; + +/// How far ahead `next_geq` walks from where the cursor sits before reseating through the +/// zero-sample table. +/// +/// Bounded in high-part buckets, which the cursor can measure up front, and in elements, which is +/// what it pays: a bucket holding a run of duplicates is arbitrarily many elements deep, so a gap +/// of a few buckets is not a gap of a few steps. +pub(crate) const LINEAR_SCAN_THRESHOLD: u64 = 8; + +/// The widest low part we will store. +/// +/// `l == 64` would leave no high part and would ask the bit-packed child for its own full width, +/// which FastLanes does not do. Only a single element spanning the whole `u64` range reaches it. +pub(crate) const MAX_LOWER_WIDTH: u8 = 63; + +/// The number of low bits to give each element, written `l` in the literature. +/// +/// `l = floor(log2(universe / n))` balances the halves: low parts cost `l` bits each and the upper +/// array costs about `n + universe / 2^l` bits, so the total lands near `n * (l + 2)`. +pub(crate) fn lower_width(span: u64, n: usize) -> u8 { + debug_assert!(n > 0, "lower_width is undefined for an empty sequence"); + + // The universe is `span + 1` values, which is 2^64 when the span fills a u64 — hence u128. + let universe = u128::from(span) + 1; + let n = u128::from(n as u64); + if universe <= n { + // More elements than distinct values: the sequence is dense, or has many duplicates. + // Every bit is better spent on the upper array, which stays O(n) either way. + return 0; + } + let width = (universe / n).ilog2(); + u8::try_from(width).unwrap_or(u8::MAX).min(MAX_LOWER_WIDTH) +} + +/// The length in bits of the upper array, written `H` in the literature. +/// +/// One set bit per element, one unset bit per high-part bucket boundary, and `+ 2` for the sentinel +/// and a trailing guard zero, so the largest selectable zero rank `span >> lower_width` is always +/// present. Bounded at roughly `3n`, since `lower_width` keeps `(span + 1) >> lower_width < 2n`. +pub(crate) fn upper_len(span: u64, n: usize, lower_width: u8) -> VortexResult { + let buckets = span >> lower_width; + let upper_len = (n as u64) + .checked_add(buckets) + .and_then(|v| v.checked_add(2)) + .ok_or_else(|| { + vortex_error::vortex_err!( + "Elias-Fano upper array overflows: n {n}, span {span}, lower_width {lower_width}" + ) + })?; + vortex_ensure!( + usize::try_from(upper_len).is_ok(), + "Elias-Fano upper array of {upper_len} bits does not fit in memory" + ); + Ok(upper_len) +} + +/// The number of unset bits in an upper array of `upper_len` bits holding `n` elements. +/// +/// No query path calls this; it states the identity [`num_samples0`]'s derivation has to agree +/// with, and `validate_parts` asserts the two against each other. +#[inline] +pub(crate) fn num_zeros(upper_len: u64, n: usize) -> u64 { + upper_len - n as u64 +} + +/// The number of zero-samples the layout calls for. +/// +/// The unset bits are the sentinel, one terminator per bucket, and the guard zero, so the universe +/// alone fixes this count whatever `n` is. That is what lets a reader split the shared samples +/// buffer into its two tables without the seam being written into the metadata. +#[inline] +pub(crate) fn num_samples0(span: u64, lower_width: u8) -> u64 { + // Saturating because `lower_width` arrives from metadata: a corrupt zero against a full-width + // span would otherwise overflow here rather than at the buffer-length check that catches it. + ((span >> lower_width).saturating_add(1)) >> LOG_SAMPLING0 +} + +#[inline] +pub(crate) fn lower_mask(lower_width: u8) -> u64 { + if lower_width == 0 { + 0 + } else { + u64::MAX >> (64 - u32::from(lower_width)) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + // Four elements over a universe of 18: 2 low bits, 4 buckets. + #[case(17, 4, 2, 10)] + // A dense run 0..n: universe == n, so no low bits, and the upper array is 2n + 1. + #[case(999, 1000, 0, 2001)] + // Sparse: 1000 elements over a 2^20 universe wants 10 low bits, leaving 1024 buckets. + #[case((1 << 20) - 1, 1000, 10, 1000 + 1023 + 2)] + // All-equal input. Span 0 means one bucket and no low bits. + #[case(0, 100, 0, 102)] + // A single element at the very top of the u64 range: the clamp fires. + #[case(u64::MAX, 1, 63, 4)] + fn test_geometry( + #[case] span: u64, + #[case] n: usize, + #[case] expected_width: u8, + #[case] expected_upper_len: u64, + ) -> VortexResult<()> { + let width = lower_width(span, n); + assert_eq!(width, expected_width, "lower_width"); + assert_eq!(upper_len(span, n, width)?, expected_upper_len, "upper_len"); + Ok(()) + } + + /// The upper array must always be long enough to hold the position that the highest element + /// claims, and to leave at least one zero rank above the highest one a query can name. + #[rstest] + #[case(0, 1)] + #[case(1, 1)] + #[case(u64::MAX, 1)] + #[case(u64::MAX, 1024)] + #[case(1_000_000, 100_000)] + #[case(7, 8)] + #[case(255, 256)] + fn test_upper_len_leaves_room(#[case] span: u64, #[case] n: usize) -> VortexResult<()> { + let width = lower_width(span, n); + let upper_len = upper_len(span, n, width)?; + + // The last element sits at `(span >> width) + (n - 1) + 1`, which must be in bounds. + let last_position = (span >> width) + n as u64; + assert!(last_position < upper_len, "last position {last_position}"); + + // A reseat may name any zero rank up to the maximum element's high part. + let max_zero_rank = span >> width; + assert!(max_zero_rank < num_zeros(upper_len, n), "max zero rank"); + Ok(()) + } + + #[test] + fn test_lower_mask() { + assert_eq!(lower_mask(0), 0); + assert_eq!(lower_mask(1), 1); + assert_eq!(lower_mask(8), 0xFF); + assert_eq!(lower_mask(63), u64::MAX >> 1); + } +} diff --git a/encodings/elias-fano/src/rules.rs b/encodings/elias-fano/src/rules.rs new file mode 100644 index 00000000000..d3ac9f89389 --- /dev/null +++ b/encodings/elias-fano/src/rules.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; + +use crate::EliasFano; + +/// Reductions an Elias-Fano array can absorb from its parent without reading a buffer: slicing, +/// which costs one metadata field, and casting, inherited from the generic adaptor. +pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&CastReduceAdaptor(EliasFano)), + ParentRuleSet::lift(&SliceReduceAdaptor(EliasFano)), +]); diff --git a/encodings/elias-fano/src/tests.rs b/encodings/elias-fano/src/tests.rs new file mode 100644 index 00000000000..3178b466d8e --- /dev/null +++ b/encodings/elias-fano/src/tests.rs @@ -0,0 +1,1443 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::compute::conformance::cast::test_cast_conformance; +use vortex_array::compute::conformance::consistency::test_array_consistency; +use vortex_array::compute::conformance::filter::test_filter_conformance; +use vortex_array::compute::conformance::take::test_take_conformance; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::eq; +use vortex_array::expr::gt; +use vortex_array::expr::gt_eq; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::lt_eq; +use vortex_array::expr::not_eq; +use vortex_array::expr::root; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::EliasFano; +use crate::EliasFanoArray; +use crate::EliasFanoArraySlotsExt; +use crate::EliasFanoCursor; +use crate::EliasFanoData; +use crate::elias_fano_encode; +use crate::params; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + crate::initialize(&session); + session +}); + +/// Deterministic xorshift, so a failure is reproducible without a `rand` dependency. +struct Rng(u64); + +impl Rng { + fn next_u64(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + + fn below(&mut self, bound: u64) -> u64 { + self.next_u64() % bound + } + + /// Uniform over `0..=bound`, including when `bound` is `u64::MAX` and `bound + 1` overflows. + fn at_most(&mut self, bound: u64) -> u64 { + match bound.checked_add(1) { + Some(universe) => self.below(universe), + None => self.next_u64(), + } + } + + /// A random permutation of `0..len`, so index-order bugs cannot hide behind a sequential walk. + fn permutation(&mut self, len: usize) -> Vec { + let mut indices: Vec = (0..len).collect(); + for i in (1..len).rev() { + indices.swap(i, self.below(i as u64 + 1) as usize); + } + indices + } +} + +/// A sorted sequence of `n` values spread over `0..=span`, with duplicates wherever they fall. +fn sorted_values(n: usize, span: u64, seed: u64) -> Vec { + let mut rng = Rng(seed); + let mut values: Vec = (0..n).map(|_| rng.at_most(span)).collect(); + values.sort_unstable(); + values +} + +fn encode(values: &[P]) -> VortexResult { + let array = PrimitiveArray::from_iter(values.iter().copied()); + let mut ctx = SESSION.create_execution_ctx(); + elias_fano_encode(array.as_ref().as_::(), &mut ctx) +} + +/// Every element, read back through the cursor in a random order, must match `expected`. +fn check_access(array: &EliasFanoArray, expected: &[Scalar], seed: u64) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cursor = EliasFanoCursor::try_new(array.as_view(), &mut ctx)?; + for index in Rng(seed).permutation(expected.len()) { + assert_eq!(cursor.access(index)?, expected[index], "index {index}"); + } + // And once more in order, which takes the cursor's step-forward path instead of a select. + let mut cursor = EliasFanoCursor::try_new(array.as_view(), &mut ctx)?; + for (index, want) in expected.iter().enumerate() { + assert_eq!(&cursor.access(index)?, want, "sequential index {index}"); + } + Ok(()) +} + +/// `next_geq`, `rank`, and `rank_inclusive` must agree with a linear scan, for probes on and off +/// the elements. +/// +/// All three share one cursor, so they interleave the way a query does: `rank_inclusive` probes the +/// successor of the value `next_geq` just answered for, which is where a stale seat or a stale +/// memoised answer would show up. +fn check_searches( + array: &EliasFanoArray, + expected: &[Scalar], + probes: &[Scalar], +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cursor = EliasFanoCursor::try_new(array.as_view(), &mut ctx)?; + for probe in probes { + let want_rank = expected.iter().take_while(|value| *value < probe).count(); + let want_value = expected.get(want_rank).cloned(); + + let (rank, value) = cursor.next_geq(probe)?; + assert_eq!(rank, want_rank, "rank of {probe}"); + assert_eq!(value, want_value, "next_geq of {probe}"); + // Repeat, which takes the memoised answer. It has to reproduce the whole answer, not just + // the rank: the element found is generally not the probe. + assert_eq!( + cursor.next_geq(probe)?, + (want_rank, want_value), + "memoised next_geq of {probe}" + ); + // The other end of the run equal to the probe, which together with `rank` brackets it. + let want_inclusive = expected.iter().take_while(|value| *value <= probe).count(); + assert_eq!( + cursor.rank_inclusive(probe)?, + want_inclusive, + "rank_inclusive of {probe}" + ); + } + Ok(()) +} + +/// Probes covering every element, both neighbours of every element, and the universe edges. +fn probes(values: &[u64], span: u64, dtype: &DType, seed: u64) -> Vec { + let ptype = dtype.as_ptype(); + let mut raw: Vec = Vec::new(); + for &value in values { + raw.push(value); + raw.push(value.saturating_sub(1)); + raw.push(value.saturating_add(1).min(span)); + } + raw.push(0); + raw.push(span); + let mut rng = Rng(seed); + raw.extend((0..64).map(|_| rng.at_most(span))); + // Shuffled, so the cursor has to reseat backwards as well as walk forwards. + for i in (1..raw.len()).rev() { + raw.swap(i, rng.below(i as u64 + 1) as usize); + } + raw.into_iter() + .map(|value| scalar_of(ptype, value)) + .collect() +} + +/// A sorted sequence of `clusters` tight runs, each `per_cluster` long, over `0..=span`. +/// +/// `lower_width` is picked for a uniform spread, so uniformly random values leave about one element +/// in each high-part bucket however wide the universe is. Clustering is what puts many elements in +/// one bucket, which is the shape a search inside a bucket has to handle without walking it. +fn clustered_values(clusters: usize, per_cluster: usize, span: u64, seed: u64) -> Vec { + let mut rng = Rng(seed); + let mut values: Vec = Vec::with_capacity(clusters * per_cluster); + for _ in 0..clusters { + let base = rng.at_most(span); + for _ in 0..per_cluster { + // A handful of distinct values per run, so the low bits inside a bucket vary rather + // than every comparison in it landing on the same answer. + values.push(base.saturating_add(rng.below(4)).min(span)); + } + } + values.sort_unstable(); + values +} + +/// A `Scalar` of `ptype` holding `value`, which must be in range for it. +fn scalar_of(ptype: PType, value: u64) -> Scalar { + crate::array::scalar_from_bits(&DType::Primitive(ptype, NonNullable), value) + .vortex_expect("value fits the ptype") +} + +fn scalars(array: &ArrayRef) -> VortexResult> { + let mut ctx = SESSION.create_execution_ctx(); + (0..array.len()) + .map(|i| array.execute_scalar(i, &mut ctx)) + .collect() +} + +// ── Roundtrip over the shapes that change the layout ──────────────────── + +#[rstest] +// Single element, and the smallest sequences at all. +#[case::one(1, 0)] +#[case::one_sparse(1, 1 << 40)] +#[case::two(2, 1)] +// A dense run: the universe is no larger than the element count, so there are no low bits at all. +#[case::dense(1000, 999)] +// All values equal: one high-part bucket, `lower_width == 0`, and n duplicates. +#[case::all_equal(500, 0)] +// The ordinary sparse case, and one sparse enough to want many low bits. +#[case::sparse(1000, 1 << 20)] +#[case::very_sparse(1000, 1 << 50)] +// Around the FastLanes block boundary, where the low-bits child gains a partial block. +#[case::block_low(1023, 1 << 20)] +#[case::block_exact(1024, 1 << 20)] +#[case::block_high(1025, 1 << 20)] +#[case::two_blocks_low(2047, 1 << 20)] +#[case::two_blocks_exact(2048, 1 << 20)] +#[case::two_blocks_high(2049, 1 << 20)] +// Long enough that both sample tables are non-empty: one-samples need n > 256, and zero-samples +// need more than 512 unset bits, which follows from the upper array being about 2n bits. +#[case::sampled(5000, 1 << 30)] +#[case::sampled_dense(5000, 6000)] +fn test_roundtrip(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0x5EED_0001 ^ n as u64); + let expected = PrimitiveArray::from_iter(values.iter().copied()); + let encoded = encode(&values)?; + + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(encoded.len(), n); + assert_arrays_eq!(encoded, expected, &mut ctx); + + let expected_scalars = scalars(&expected.into_array())?; + check_access(&encoded, &expected_scalars, 0xC0FFEE)?; + check_searches( + &encoded, + &expected_scalars, + &probes(&values, span, encoded.dtype(), 0xBEEF), + )?; + Ok(()) +} + +/// Both sample tables must actually be populated at the sizes the roundtrip cases use, or those +/// cases would be silently testing only the unsampled path. +#[test] +fn test_sample_tables_are_exercised() -> VortexResult<()> { + let encoded = encode(&sorted_values(5000, 1 << 30, 0xDEAD))?; + let samples = encoded.samples_buffer().len() / size_of::(); + let num_samples0 = encoded.num_samples0() as usize; + assert_eq!(num_samples0, 15, "zero-samples"); + assert_eq!(samples - num_samples0, 19, "one-samples"); + Ok(()) +} + +/// The bulk decode and the per-element cursor must agree — they share no code beyond the layout. +#[rstest] +#[case(1, 0)] +#[case(300, 1 << 16)] +#[case(5000, 1 << 30)] +fn test_decode_matches_cursor(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0xABCD); + let encoded = encode(&values)?; + let mut ctx = SESSION.create_execution_ctx(); + + let decoded = encoded + .clone() + .into_array() + .execute::(&mut ctx)?; + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + for (index, &want) in decoded.as_slice::().iter().enumerate() { + let element = cursor.access_element(index)?; + assert_eq!(element + values[0], want, "index {index}"); + } + Ok(()) +} + +// ── Slicing ──────────────────────────────────────────────────────────── + +/// A slice records a rank offset and keeps the buffers whole, so every read has to apply it. The +/// starts below straddle the one-sample spacing (256) and the FastLanes block size (1024). +#[rstest] +#[case(0, 1)] +#[case(0, 3000)] +#[case(1, 2999)] +#[case(255, 300)] +#[case(256, 300)] +#[case(257, 300)] +#[case(1023, 1200)] +#[case(1024, 1200)] +#[case(1025, 1200)] +#[case(2999, 3000)] +fn test_slice(#[case] start: usize, #[case] end: usize) -> VortexResult<()> { + let values = sorted_values(3000, 1 << 24, 0xF00D); + let encoded = encode(&values)?; + let sliced = encoded.slice(start..end)?; + + // The slice must stay Elias-Fano rather than falling back to a generic `SliceArray`. + assert!( + sliced.is::(), + "slice reduced away from EliasFano" + ); + let sliced = sliced.as_::().into_owned(); + assert_eq!(sliced.first_rank(), start as u64); + // The low-bits child is deliberately *not* sliced: one rank offset serves both halves. + assert_eq!(sliced.lower().len(), values.len()); + + let expected = PrimitiveArray::from_iter(values[start..end].iter().copied()); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(sliced, expected, &mut ctx); + + let expected_scalars = scalars(&expected.into_array())?; + check_access(&sliced, &expected_scalars, 0x1234)?; + check_searches( + &sliced, + &expected_scalars, + &probes(&values, 1 << 24, sliced.dtype(), 0x5678), + )?; + Ok(()) +} + +/// Slicing twice must compose, and the second slice must not re-slice the child. +#[test] +fn test_slice_of_slice() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x9999); + let encoded = encode(&values)?; + let sliced = encoded.slice(500..1500)?.slice(200..800)?; + assert_eq!(sliced.as_::().first_rank(), 700); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!( + sliced, + PrimitiveArray::from_iter(values[700..1300].iter().copied()), + &mut ctx + ); + Ok(()) +} + +// ── Element types ────────────────────────────────────────────────────── + +/// Every integer ptype, signed and unsigned, including references at the bottom of the range where +/// the element domain wraps through the whole width. +#[test] +fn test_signed_and_unsigned() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + macro_rules! check { + ($values:expr) => {{ + let values = $values; + let expected = PrimitiveArray::from_iter(values.iter().copied()); + let encoded = elias_fano_encode(expected.as_ref().as_::(), &mut ctx)?; + assert_arrays_eq!(encoded, expected, &mut ctx); + let expected_scalars = scalars(&expected.into_array())?; + check_access(&encoded, &expected_scalars, 0x2468)?; + }}; + } + + check!([0u8, 1, 7, 200, 255]); + check!([i8::MIN, -100, 0, 100, i8::MAX]); + check!([0u16, 300, 65535]); + check!([i16::MIN, 0, i16::MAX]); + check!([0u32, 1 << 20, u32::MAX]); + check!([i32::MIN, -1, 0, 1, i32::MAX]); + check!([0u64, 1 << 40, u64::MAX]); + check!([i64::MIN, -1, 0, 1, i64::MAX]); + // Single elements at the extremes, which is where `lower_width` clamps. + check!([u64::MAX]); + check!([i64::MIN]); + Ok(()) +} + +/// `next_geq` and `rank` over signed columns, probed *inside* the universe. +/// +/// [`probes`] builds its set in the unsigned element domain, so every case that calls +/// [`check_next_geq`] drives it with a `u64` column; the signed cases above check `access` alone, +/// and [`test_probes_outside_the_universe`] only reaches the universe edges. That leaves the probe +/// classification in `locate` — which tells below from above in the ptype's own ordering, after a +/// subtraction that wraps for every value under the reference — covered for signed columns only +/// indirectly, through the two bounds `compare` asks for. +#[test] +fn test_signed_next_geq() -> VortexResult<()> { + macro_rules! check { + ($ptype:expr, $values:expr) => {{ + let values = $values; + let encoded = encode(&values)?; + let expected: Vec = values + .iter() + .map(|&v| scalar_of($ptype, v as i64 as u64)) + .collect(); + // Every element, and both its neighbours, so the probe lands on a value, between two, + // and inside a run of duplicates. + let mut probes: Vec = values + .iter() + .flat_map(|&v| [v.saturating_sub(1), v, v.saturating_add(1)]) + .map(|v| scalar_of($ptype, v as i64 as u64)) + .collect(); + check_searches(&encoded, &expected, &probes)?; + // And again descending, which reseats backwards rather than walking forwards. + probes.reverse(); + check_searches(&encoded, &expected, &probes)?; + }}; + } + + check!(PType::I8, [i8::MIN, -100, -7, -7, 0, 1, 100, i8::MAX]); + check!(PType::I16, [i16::MIN, -3000, -7, -7, 0, 9, i16::MAX]); + check!(PType::I32, [i32::MIN, -70_000, -7, -7, 0, 5000, i32::MAX]); + check!(PType::I64, [i64::MIN, -1 << 40, -7, -7, 0, 1 << 40, i64::MAX]); + // A reference above zero, so the element domain does not wrap and the negative probes below it + // all classify as `Bound::Below`. + check!(PType::I32, [5i32, 5, 900, 1_000_000, i32::MAX]); + // And one entirely below zero, where every value is a negative pattern. + check!(PType::I32, [i32::MIN, i32::MIN + 1, -900_000, -5]); + Ok(()) +} + +/// `lower_width` on either side of every native width, where a naive implementation would try to +/// bit-pack at or above the child's own width. +#[rstest] +#[case(7)] +#[case(8)] +#[case(9)] +#[case(15)] +#[case(16)] +#[case(17)] +#[case(31)] +#[case(32)] +#[case(33)] +#[case(62)] +#[case(63)] +fn test_lower_width_boundaries(#[case] width: u8) -> VortexResult<()> { + // `lower_width` is `floor(log2(universe / n))`, so `n` elements over a universe of `n << width` + // land on exactly `width`. The cap keeps that universe inside 64 bits for the widest cases. + let n = 400usize.min(1usize << (64 - u32::from(width)).min(20)); + let span = u64::try_from(((n as u128) << width) - 1)?; + let mut values = sorted_values(n, span, 0x7777 + u64::from(width)); + // Pin the extremes, so the *observed* span is the one the case asked for. + values[0] = 0; + values[n - 1] = span; + let encoded = encode(&values)?; + assert_eq!(encoded.lower_width(), width); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!( + encoded, + PrimitiveArray::from_iter(values.iter().copied()), + &mut ctx + ); + check_access( + &encoded, + &values + .iter() + .map(|&v| scalar_of(PType::U64, v)) + .collect::>(), + 0x8888, + )?; + Ok(()) +} + +// ── Degenerate inputs ────────────────────────────────────────────────── + +#[test] +fn test_empty() -> VortexResult<()> { + let encoded = encode::(&[])?; + assert_eq!(encoded.len(), 0); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(encoded, PrimitiveArray::empty::(NonNullable), &mut ctx); + Ok(()) +} + +/// A probe below every element must report rank 0 and the first element, and one above every +/// element must report rank `len` and nothing. Each probe below gets a freshly opened cursor, which +/// is not yet seated on anything, so the classification has to come from the universe bounds alone. +#[test] +fn test_probes_outside_the_universe() -> VortexResult<()> { + let values: Vec = vec![-5, 0, 10, 10, 400]; + let encoded = encode(&values)?; + let mut ctx = SESSION.create_execution_ctx(); + + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + assert_eq!(cursor.rank(&scalar_of(PType::I32, 401))?, 5); + assert_eq!(cursor.next_geq(&scalar_of(PType::I32, 401))?.1, None); + + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + let (rank, value) = cursor.next_geq(&scalar_of(PType::I32, i32::MIN as i64 as u64))?; + assert_eq!(rank, 0); + assert_eq!(value, Some(scalar_of(PType::I32, -5i32 as i64 as u64))); + + // The maximum is inside the universe, so it must be found rather than clamped away. + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + assert_eq!(cursor.rank(&scalar_of(PType::I32, 400))?, 4); + Ok(()) +} + +/// `rank_inclusive` counts the elements *at or below* the probe, which is the rank of the probe's +/// successor — and at the top of the universe the probe has none. +/// +/// [`check_searches`] drives this from every roundtrip and slice case, but its probes all sit +/// inside a universe narrower than the ptype, so the overflow guard and the two +/// outside-the-universe arms need naming here. A span filling the whole width is what makes the +/// successor overflow. +#[test] +fn test_rank_inclusive_at_the_universe_edges() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + // Duplicates at the top, so the count is not simply the length. + let encoded = encode(&[0u64, 1 << 40, u64::MAX, u64::MAX])?; + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::U64, u64::MAX))?, 4); + assert_eq!( + cursor.rank_inclusive(&scalar_of(PType::U64, u64::MAX - 1))?, + 2 + ); + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::U64, 0))?, 1); + assert_eq!(cursor.rank(&scalar_of(PType::U64, u64::MAX))?, 2); + + // The same at the top of a signed ptype, where the maximum's bit pattern is not the largest + // `u64` but the span is still the full width. + let encoded = encode(&[i64::MIN, -1, i64::MAX])?; + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + assert_eq!( + cursor.rank_inclusive(&scalar_of(PType::I64, i64::MAX as u64))?, + 3 + ); + assert_eq!( + cursor.rank_inclusive(&scalar_of(PType::I64, i64::MIN as u64))?, + 1 + ); + + // Below the reference and above the maximum, where the count comes from the universe bounds + // alone rather than from a search. + let encoded = encode(&[10i32, 20, 20, 30])?; + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::I32, 9))?, 0); + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::I32, 31))?, 4); + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::I32, 20))?, 3); + + let empty = encode::(&[])?; + let mut cursor = EliasFanoCursor::try_new(empty.as_view(), &mut ctx)?; + assert_eq!(cursor.rank_inclusive(&scalar_of(PType::U64, 7))?, 0); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 7))?, 0); + Ok(()) +} + +/// Duplicates are legal, and a rank must count from the *first* occurrence. This is the case a +/// cursor that trusts wherever it happens to be seated gets wrong. +#[test] +fn test_duplicates_report_first_occurrence() -> VortexResult<()> { + let values: Vec = vec![5, 5, 5, 5, 5, 9, 9, 12]; + let encoded = encode(&values)?; + let mut ctx = SESSION.create_execution_ctx(); + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + + // Walk past the duplicates first, so the cursor is seated in the middle of the run... + assert_eq!(cursor.access(3)?, scalar_of(PType::U64, 5)); + // ...then ask for their value, which must still answer 0. + assert_eq!(cursor.rank(&scalar_of(PType::U64, 5))?, 0); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 9))?, 5); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 12))?, 7); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 6))?, 5); + Ok(()) +} + +/// Many duplicates spread over a wide universe, so runs of equal values share a bucket while the +/// buckets themselves are sparse. +#[test] +fn test_long_duplicate_runs() -> VortexResult<()> { + let mut values: Vec = Vec::new(); + for group in 0..40u64 { + for _ in 0..50 { + values.push(group * 1000); + } + } + let encoded = encode(&values)?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!( + encoded, + PrimitiveArray::from_iter(values.iter().copied()), + &mut ctx + ); + + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + for group in 0..40u64 { + assert_eq!( + cursor.rank(&scalar_of(PType::U64, group * 1000))?, + group as usize * 50, + "group {group}" + ); + } + Ok(()) +} + +/// Probes that advance by a bucket or two while a long run of duplicates sits between them. +/// +/// This is the shape that separates the two bounds on the forward walk. The gap in buckets is +/// small enough to walk, but each bucket holds a thousand elements, so the walk has to give up on +/// its element budget and reseat instead of stepping through the run. Every answer here is also +/// reachable by a linear scan, which is what makes the reseat observable only as work not done. +#[test] +fn test_ascending_probes_across_deep_buckets() -> VortexResult<()> { + // Three runs, placed so the first two share adjacent buckets and the third is far away. The + // span keeps `lower_width` at 8, so a bucket spans 256 values. + let mut values: Vec = Vec::new(); + for &value in &[0u64, 300, 900_000] { + values.extend(std::iter::repeat_n(value, 1000)); + } + let encoded = encode(&values)?; + assert_eq!(encoded.lower_width(), 8, "case depends on the bucket width"); + + let mut ctx = SESSION.create_execution_ctx(); + let mut cursor = EliasFanoCursor::try_new(encoded.as_view(), &mut ctx)?; + + // Seat the cursor at the head of the first run, then probe forward. Every probe below sits + // within `LINEAR_SCAN_THRESHOLD` buckets of the seat but a thousand elements past it. + assert_eq!(cursor.rank(&scalar_of(PType::U64, 0))?, 0); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 1))?, 1000); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 300))?, 1000); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 301))?, 2000); + assert_eq!(cursor.rank(&scalar_of(PType::U64, 900_000))?, 2000); + + let decoded = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let expected_scalars = scalars(&decoded)?; + check_searches( + &encoded, + &expected_scalars, + &probes(&values, 900_000, encoded.dtype(), 0x7070), + ) +} + +/// Clustered data, where buckets are deep, checked against a linear scan over hundreds of probes. +/// +/// The sliced half matters as much as the whole: a search inside a bucket has to clamp the bucket +/// to the slice at both ends, and a bucket running past the slice is the case that reports "nothing +/// at or above this" rather than an element the slice does not contain. +#[rstest] +#[case::few_deep(3, 1000, 900_000)] +#[case::many_shallow(200, 25, 1 << 30)] +#[case::mixed(20, 200, 1 << 24)] +fn test_clustered_next_geq( + #[case] clusters: usize, + #[case] per_cluster: usize, + #[case] span: u64, +) -> VortexResult<()> { + let seed = 0x9E37 ^ span; + let values = clustered_values(clusters, per_cluster, span, seed); + let encoded = encode(&values)?.into_array(); + + let mut arrays = vec![encoded.clone()]; + arrays.push(encoded.slice(per_cluster / 2..values.len() - per_cluster / 2)?); + + for array in &arrays { + let array = array.as_::().into_owned(); + let expected = scalars(&array.clone().into_array())?; + check_searches( + &array, + &expected, + &probes(&values, span, array.dtype(), seed ^ 0xFEED), + )?; + } + Ok(()) +} + +#[test] +fn test_rejects_unsorted_and_nullable() { + let mut ctx = SESSION.create_execution_ctx(); + assert!(encode(&[3u64, 1, 2]).is_err()); + // Nulls have no position in an ordering, so they are refused rather than worked around. + let nullable = PrimitiveArray::from_option_iter([Some(1u64), None, Some(3)]); + assert!(elias_fano_encode(nullable.as_ref().as_::(), &mut ctx).is_err()); +} + +// ── The low-bits child in shapes a rewrite or a file roundtrip can produce ── + +/// The child does not have to be a bare `BitPacked`. A file roundtrip can hand it back wrapped, and +/// a rewrite can replace it outright, so both the bulk decode and the cursor must fall back rather +/// than downcast blindly. +#[test] +fn test_unpacked_lower_child() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x4321); + let encoded = encode(&values)?; + let mut ctx = SESSION.create_execution_ctx(); + + // Replace the bit-packed child with the plain primitive array it decodes to. + let plain = encoded + .lower() + .clone() + .execute::(&mut ctx)? + .into_array(); + let rebuilt = rebuild_with_lower(&encoded, plain)?; + + assert_arrays_eq!( + rebuilt, + PrimitiveArray::from_iter(values.iter().copied()), + &mut ctx + ); + let expected_scalars = values + .iter() + .map(|&v| scalar_of(PType::U64, v)) + .collect::>(); + check_access(&rebuilt, &expected_scalars, 0x1111)?; + check_searches( + &rebuilt, + &expected_scalars, + &probes(&values, 1 << 20, rebuilt.dtype(), 0x1112), + )?; + + // And sliced. A cursor over this slot materialises it, so it materialises only the slice's own + // range and rebases into it — the one place two bases have to be reconciled. The seek paths + // matter more here than the point lookups: a mishandled base comes back as a wrong rank. + const START: usize = 1500; + let sliced = rebuilt.into_array().slice(START..2000)?; + assert!( + sliced.is::(), + "slice reduced away from EliasFano" + ); + assert_arrays_eq!( + sliced, + PrimitiveArray::from_iter(values[START..2000].iter().copied()), + &mut ctx + ); + let sliced = sliced.as_::().into_owned(); + let sliced_scalars = values[START..2000] + .iter() + .map(|&v| scalar_of(PType::U64, v)) + .collect::>(); + check_access(&sliced, &sliced_scalars, 0x2222)?; + check_searches( + &sliced, + &sliced_scalars, + &probes(&values, 1 << 20, sliced.dtype(), 0x2223), + )?; + Ok(()) +} + +/// A child carrying a non-zero FastLanes sub-block offset. `unpack_single_primitive` does not apply +/// that offset itself, so a reader that forgets it returns wrong values with no panic anywhere. +#[test] +fn test_lower_child_with_block_offset() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x2222); + let encoded = encode(&values)?; + let width = encoded.lower_width(); + + // Rebuild the low bits with `pad` junk values in front, then slice them back off. The child now + // holds the same n values at the same ranks, but starting part-way into a block. + const PAD: usize = 5; + let mut padded: Vec = vec![0; PAD]; + let reference = values[0]; + padded.extend( + values + .iter() + .map(|&v| (v - reference) & params::lower_mask(width)), + ); + let packed = unsafe { + bitpack_encode_unchecked( + PrimitiveArray::new( + padded.into_iter().collect::>(), + Validity::NonNullable, + ), + width, + ) + }? + .into_array() + .slice(PAD..PAD + values.len())?; + assert_eq!(packed.as_::().offset(), 5); + + let rebuilt = rebuild_with_lower(&encoded, packed)?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!( + rebuilt, + PrimitiveArray::from_iter(values.iter().copied()), + &mut ctx + ); + let expected_scalars = values + .iter() + .map(|&v| scalar_of(PType::U64, v)) + .collect::>(); + check_access(&rebuilt, &expected_scalars, 0x3333)?; + check_searches( + &rebuilt, + &expected_scalars, + &probes(&values, 1 << 20, rebuilt.dtype(), 0x3334), + )?; + Ok(()) +} + +/// The low bits are OR-ed in under `lower_width`, so a child packed above that width would bleed +/// into the high part. Its width is metadata, so this is refused at construction; a child packed +/// *below* it is a legal tightening and must still be accepted. +#[test] +fn test_rejects_overwide_lower_child() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x6060); + let encoded = encode(&values)?; + let width = encoded.lower_width(); + assert!(width > 1, "case needs room on both sides of the width"); + + // Masking to the packing width keeps every repack lossless, so the only thing that varies + // between these two children is the width the layout is asked to accept. + let pack = |bit_width: u8| -> VortexResult { + let mask = params::lower_mask(bit_width); + let low: Buffer = values.iter().map(|&v| (v - values[0]) & mask).collect(); + let packed = unsafe { + bitpack_encode_unchecked(PrimitiveArray::new(low, Validity::NonNullable), bit_width) + }?; + Ok(packed.into_array()) + }; + + assert!( + rebuild_with_lower(&encoded, pack(width + 1)?).is_err(), + "a child packed wider than lower_width must be rejected" + ); + rebuild_with_lower(&encoded, pack(width - 1)?)?; + Ok(()) +} + +/// Bits above `lower_width` in the low-bits child must be masked off, not trusted. +/// +/// A bit-packed child's width is metadata, so [`test_rejects_overwide_lower_child`] refuses that at +/// construction. A patched or rewritten slot arrives as a plain `u64` array instead, where nothing +/// bounds the values at all — and a bit that survives into the high part is a wrong answer with no +/// error anywhere. Both readers are covered: the bulk decode, and the cursor, whose bucket +/// bisection compares low parts directly as well as assembling them. +#[test] +fn test_lower_child_with_junk_above_the_width() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x4949); + let encoded = encode(&values)?; + let width = encoded.lower_width(); + assert!(width > 0, "case needs low bits to mask"); + let mut ctx = SESSION.create_execution_ctx(); + + let junk = !params::lower_mask(width); + let plain: Buffer = encoded + .lower() + .clone() + .execute::(&mut ctx)? + .as_slice::() + .iter() + .map(|&low| low | junk) + .collect(); + let rebuilt = rebuild_with_lower( + &encoded, + PrimitiveArray::new(plain, Validity::NonNullable).into_array(), + )?; + + let expected = PrimitiveArray::from_iter(values.iter().copied()); + assert_arrays_eq!(rebuilt, expected, &mut ctx); + let expected_scalars = scalars(&expected.into_array())?; + check_access(&rebuilt, &expected_scalars, 0x4950)?; + check_searches( + &rebuilt, + &expected_scalars, + &probes(&values, 1 << 20, rebuilt.dtype(), 0x4951), + )?; + Ok(()) +} + +fn rebuild_with_lower(array: &EliasFanoArray, lower: ArrayRef) -> VortexResult { + let len = array.len(); + EliasFano::try_new(array.as_view().data().clone(), lower, len) +} + +// ── Statistics and conformance ───────────────────────────────────────── + +#[rstest] +#[case::empty(0, 0)] +#[case::single(1, 0)] +#[case::single_sparse(1, 1 << 40)] +#[case::pair(2, 1)] +#[case::all_equal(500, 0)] +#[case::dense(1000, 999)] +#[case::sparse(2000, 1 << 40)] +fn test_is_sorted_stat(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let encoded = encode(&sorted_values(n, span, 0xAAAA + n as u64))?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let mut arrays = vec![encoded.clone()]; + if n > 3 { + arrays.push(encoded.slice(1..n - 1)?); + } + for array in arrays { + // Present without reading a buffer, which is what `ListArray::new` requires of offsets. + assert_eq!( + array + .statistics() + .with_typed_stats_set(|stats| stats.get_as::(Stat::IsSorted)), + Precision::Exact(true), + "IsSorted over {} elements", + array.len() + ); + // Strictness is declined rather than answered, so it must come back from the generic path + // with the same answer the decoded array gives. + let decoded = array.clone().execute::(&mut ctx)?.into_array(); + assert_eq!( + array.statistics().compute_is_strict_sorted(&mut ctx), + decoded.statistics().compute_is_strict_sorted(&mut ctx), + "IsStrictSorted over {} elements", + array.len() + ); + } + Ok(()) +} + +#[rstest] +#[case::empty(0, 0)] +#[case::single(1, 0)] +#[case::single_sparse(1, 1 << 40)] +#[case::pair(2, 1)] +#[case::all_equal(500, 0)] +#[case::dense(1000, 999)] +#[case::sparse(2000, 1 << 40)] +fn test_min_max_kernel(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0xBBBB + n as u64); + let encoded = encode(&values)?.into_array(); + + // On a slice, the metadata bounds are the *encoded* universe rather than the slice's own + // extremes, so the slices are what catch a kernel reading metadata instead of elements. The + // one-element slices are also the shape `min_max` itself uses internally. + let mut arrays = vec![(encoded.clone(), values.clone())]; + if n > 3 { + for (start, end) in [(0, n - 1), (1, n), (1, n - 1), (n / 2, n / 2 + 1)] { + arrays.push((encoded.slice(start..end)?, values[start..end].to_vec())); + } + } + + let mut ctx = SESSION.create_execution_ctx(); + for (array, expected) in arrays { + // The oracle is the decoded array, so the kernel is checked against the generic path and + // not only against the input it was built from. + let decoded = array.clone().execute::(&mut ctx)?.into_array(); + let len = array.len(); + for (name, got, oracle, want) in [ + ( + "min", + array.statistics().compute_min::(&mut ctx), + decoded.statistics().compute_min::(&mut ctx), + expected.first().copied(), + ), + ( + "max", + array.statistics().compute_max::(&mut ctx), + decoded.statistics().compute_max::(&mut ctx), + expected.last().copied(), + ), + ] { + assert_eq!(got, want, "{name} over {len} elements"); + assert_eq!(got, oracle, "{name} disagrees with the generic path over {len}"); + } + } + Ok(()) +} + +/// `min_max` on a signed column, where the reference is negative and the element domain wraps +/// through the whole width on the way out. +#[test] +fn test_min_max_kernel_signed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for values in [ + vec![i32::MIN, -7, -7, 0, i32::MAX], + vec![-9i32], + vec![i32::MIN, i32::MAX], + ] { + let encoded = encode(&values)?.into_array(); + let decoded = encoded.clone().execute::(&mut ctx)?.into_array(); + assert_eq!( + encoded.statistics().compute_min::(&mut ctx), + Some(values[0]), + "min of {values:?}" + ); + assert_eq!( + encoded.statistics().compute_max::(&mut ctx), + values.last().copied(), + "max of {values:?}" + ); + assert_eq!( + encoded.statistics().compute_min::(&mut ctx), + decoded.statistics().compute_min::(&mut ctx), + "min disagrees with the generic path for {values:?}" + ); + assert_eq!( + encoded.statistics().compute_max::(&mut ctx), + decoded.statistics().compute_max::(&mut ctx), + "max disagrees with the generic path for {values:?}" + ); + } + Ok(()) +} + +/// The shared compute harnesses, over both a whole array and a slice of it. +/// +/// The slice is the half that matters: take and filter both decode arbitrary index sets, which is +/// where a mishandled `first_rank` — the one number a slice records — shows up. +#[rstest] +#[case(1, 0)] +#[case(5, 100)] +#[case(1000, 999)] +#[case(2000, 1 << 20)] +#[case(5000, 1 << 40)] +fn test_conformance(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let encoded = encode(&sorted_values(n, span, 0xCCCC + n as u64))?.into_array(); + let mut arrays = vec![encoded.clone()]; + if n > 2 { + arrays.push(encoded.slice(1..n - 1)?); + } + + let mut ctx = SESSION.create_execution_ctx(); + for array in &arrays { + test_array_consistency(array, &mut ctx); + test_take_conformance(array, &mut ctx); + test_filter_conformance(array, &mut ctx); + test_cast_conformance(array, &mut ctx); + } + Ok(()) +} + +// ── Take and filter pushdown ──────────────────────────────────────────── + +/// Take and filter along the *cursor* path, not the bulk decode. +/// +/// Both kernels hand a dense request back to the framework to decode in one pass, so only a sparse +/// one reaches the per-element cursor — the half that has to apply `first_rank` itself. A +/// conformance harness that happens to ask for many rows would exercise the bulk path and leave +/// this untested, which is why the index sets here are deliberately tiny. +#[rstest] +#[case(2000, 1 << 20)] +#[case(5000, 1 << 40)] +fn test_sparse_take_and_filter(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0x5A5A + n as u64); + let encoded = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + for array in [encoded.clone(), encoded.slice(7..n - 7)?] { + let decoded = array + .clone() + .execute::(&mut ctx)? + .into_array(); + + // Sixteen rows spread across thousands: far under `BULK_DECODE_THRESHOLD`, so every read + // goes through the cursor. + let picks: Vec = (0..16).map(|i| i * array.len() / 16).collect(); + + // Take sees them in descending order, so a cursor that only ever steps forward fails here. + let descending = PrimitiveArray::from_iter(picks.iter().rev().map(|&i| i as u64)); + let indices = descending.into_array(); + assert_arrays_eq!( + array.take(indices.clone())?, + decoded.take(indices)?, + &mut ctx + ); + + let mask = Mask::from_indices(array.len(), picks.iter().copied()); + assert_arrays_eq!( + array.filter(mask.clone())?, + decoded.filter(mask)?, + &mut ctx + ); + } + Ok(()) +} + +/// A null index selects nothing, so its payload is not a position: it is whatever the slot happened +/// to hold, and the cursor path must neither bounds-check nor read it. +/// +/// `test_take_conformance` cannot catch this. It builds its nullable indices with +/// `from_option_iter`, which writes a zero under every null — in bounds for any non-empty array, so +/// a reader that ignores validity still passes. The payloads below are the ones that do not: a +/// negative, and two far past the end. +#[test] +fn test_take_ignores_null_index_payloads() -> VortexResult<()> { + let values = sorted_values(2000, 1 << 20, 0x7A7A); + let encoded = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + // Six rows out of two thousand: well under `BULK_DECODE_THRESHOLD`, so this is the cursor path + // rather than the bulk decode, which hands the whole job to the framework. + let raw: Buffer = [7i64, i64::MAX, 1500, -1, 99_999, 3].into_iter().collect(); + let indices = PrimitiveArray::new( + raw, + Validity::from_mask(Mask::from_indices(6, [0, 2, 5]), Nullable), + ) + .into_array(); + + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([ + Some(values[7]), + None, + Some(values[1500]), + None, + None, + Some(values[3]), + ]), + &mut ctx + ); + + // All-null indices are answered before a cursor is opened, which would otherwise materialise a + // low-bits child for a read that never happens. The constant is the visible sign of that. + let all_null = PrimitiveArray::new( + [i64::MIN, -3, 88_888].into_iter().collect::>(), + Validity::AllInvalid, + ); + let taken = encoded + .take(all_null.into_array())? + .execute::(&mut ctx)?; + assert_eq!(taken.len(), 3); + let constant = taken + .as_constant() + .vortex_expect("an all-null take must reduce to a constant"); + assert!(constant.is_null()); + Ok(()) +} + +// ── Comparison pushdown ───────────────────────────────────────────────── + +/// Every comparison operator, against the decoded array's own answer. +/// +/// The pushdown resolves each one to a range through two sampled searches instead of decoding, so +/// it shares no code with the generic path — which is what makes that path a real oracle. Probes +/// cover a value that is present, one that falls between two elements, one below the minimum and +/// one above the maximum, and both ends of the array. +#[rstest] +#[case(1, 0)] +#[case(5, 100)] +#[case(1000, 999)] +#[case(2000, 1 << 20)] +#[case(5000, 1 << 40)] +fn test_compare_matches_decoded(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0xC0DE + n as u64); + let encoded = encode(&values)?.into_array(); + let mut arrays = vec![encoded.clone()]; + if n > 2 { + arrays.push(encoded.slice(1..n - 1)?); + } + + let mut probes = vec![0u64, 1, u64::MAX, span, span.saturating_add(1)]; + probes.push(values[0]); + probes.push(values[n - 1]); + probes.push(values[n / 2]); + probes.push(values[n / 2].saturating_add(1)); + probes.push(values[n / 2].saturating_sub(1)); + + let mut ctx = SESSION.create_execution_ctx(); + for array in &arrays { + // The oracle: the same array, decoded, compared by the generic path. + let decoded = array.clone().execute::(&mut ctx)?.into_array(); + for &probe in &probes { + // A nullable literal must not change which rows match, but it does make the result + // nullable — which a comparison of set bits alone would not notice. + for literal in nullabilities(Scalar::from(probe))? { + for op in COMPARISONS { + let expr = op(root(), lit(literal.clone())); + let pushed = array.clone().apply(&expr)?.execute::(&mut ctx)?; + let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + assert_eq!( + pushed.dtype(), + expected.dtype(), + "dtype for {literal} over {} elements", + array.len() + ); + assert_arrays_eq!(pushed, expected, &mut ctx); + } + } + } + } + Ok(()) +} + +/// `scalar` as itself and again with a nullable dtype, which the comparison kernels have to carry +/// into the result even though the value is never null. +fn nullabilities(scalar: Scalar) -> VortexResult<[Scalar; 2]> { + let nullable = scalar.cast(&scalar.dtype().as_nullable())?; + Ok([scalar, nullable]) +} + +/// The six comparison operators, as expression constructors. +const COMPARISONS: [fn(Expression, Expression) -> Expression; 6] = + [eq, not_eq, lt, lt_eq, gt, gt_eq]; + +/// A probe at the very top of a universe that fills a `u64`. +/// +/// Every operator but `Lt` and `Gte` needs the count of elements *at or below* the probe, which is +/// the rank of its successor — and here the probe has no successor. Taking it in the value domain, +/// or without a guard in the element domain, overflows: a panic in debug and a wrong answer in +/// release, on the one input that reaches it. +#[test] +fn test_compare_at_the_top_of_the_universe() -> VortexResult<()> { + let values = [0u64, 1 << 40, u64::MAX]; + let encoded = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let decoded = encoded + .clone() + .execute::(&mut ctx)? + .into_array(); + + for probe in [u64::MAX, u64::MAX - 1, 0] { + for op in COMPARISONS { + let expr = op(root(), lit(Scalar::from(probe))); + let pushed = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + assert_eq!(pushed.dtype(), expected.dtype(), "dtype for probe {probe}"); + assert_arrays_eq!(pushed, expected, &mut ctx); + } + } + Ok(()) +} + +/// A signed, narrow column, probed on and off its elements and at both ends of its range. +/// +/// The u64 cases above never exercise the narrowing on the way out, nor a reference below zero, +/// where the element domain wraps through the whole width. +#[test] +fn test_compare_signed_narrow_column() -> VortexResult<()> { + let values = [i32::MIN, -7, -7, 0, 1, 90, i32::MAX]; + let encoded = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let decoded = encoded + .clone() + .execute::(&mut ctx)? + .into_array(); + + for probe in [i32::MIN, i32::MIN + 1, -8, -7, -6, 0, 89, i32::MAX - 1, i32::MAX] { + for literal in nullabilities(Scalar::from(probe))? { + for op in COMPARISONS { + let expr = op(root(), lit(literal.clone())); + let pushed = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + assert_eq!(pushed.dtype(), expected.dtype(), "dtype for {literal}"); + assert_arrays_eq!(pushed, expected, &mut ctx); + } + } + } + Ok(()) +} + +/// A comparison every row satisfies must come back as a constant, not a buffer of set bits. +/// +/// This is the one externally visible sign that the kernel ran at all: the generic path decodes and +/// compares element by element, so it can only ever produce a `BoolArray`. Without it an +/// unregistered kernel would leave every other test in this file passing. +#[test] +fn test_compare_is_pushed_down() -> VortexResult<()> { + let values = sorted_values(1000, 1 << 20, 0xF00D); + let encoded = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + // Every element is >= the minimum, and none is < it. + let minimum = lit(Scalar::from(values[0])); + for (expr, expected) in [ + (gt_eq(root(), minimum.clone()), true), + (lt(root(), minimum), false), + ] { + let result = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let constant = result + .as_constant() + .vortex_expect("an all-or-nothing comparison must reduce to a constant"); + assert_eq!(constant, Scalar::from(expected)); + } + Ok(()) +} + +// ── Serialization, and use as a list column's offsets ─────────────────── + +/// Serialize, decode, and read back. This is the path a file roundtrip takes, and the only one that +/// exercises `deserialize` — including that it can size the low-bits child, which after a slice is +/// not the array's own length. +#[rstest] +#[case(0, 3000)] +#[case(700, 2100)] +fn test_serde_roundtrip(#[case] start: usize, #[case] end: usize) -> VortexResult<()> { + let values = sorted_values(3000, 1 << 24, 0xE11A); + let array = encode(&values)?.into_array().slice(start..end)?; + let dtype = array.dtype().clone(); + let len = array.len(); + + let array_ctx = ArrayContext::empty(); + let mut concat = ByteBufferMut::empty(); + for buffer in array.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + concat.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + + assert!(decoded.is::(), "decoded away from EliasFano"); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter(values[start..end].iter().copied()), + &mut ctx + ); + check_access( + &decoded.as_::().into_owned(), + &values[start..end] + .iter() + .map(|&v| scalar_of(PType::U64, v)) + .collect::>(), + 0x9A9A, + )?; + Ok(()) +} + +/// One column that a sorted integer encoding lands under: a list's offsets. +/// +/// Worth its own test not because it is what the encoding is for, but because it drives the array +/// through a parent that reads it a boundary at a time. `ListArray::new` refuses offsets that do +/// not report `IsSorted`, and `offset_at` only fast-paths a `Primitive` child — so every list +/// boundary here goes through `scalar_at`, and through the cursor underneath it. +#[test] +fn test_list_offsets() -> VortexResult<()> { + // Lengths including several empty lists, which is what makes duplicate offsets ordinary. + let lengths: Vec = (0..500u64).map(|i| (i * 7) % 5).collect(); + let mut offsets: Vec = Vec::with_capacity(lengths.len() + 1); + offsets.push(0); + for length in &lengths { + offsets.push(offsets[offsets.len() - 1] + length); + } + let total = *offsets.last().vortex_expect("at least one offset") as usize; + + let elements = PrimitiveArray::from_iter((0..total as i32).map(|i| i * 3)).into_array(); + let list = ListArray::try_new( + elements.clone(), + encode(&offsets)?.into_array(), + Validity::NonNullable, + )?; + assert_eq!(list.len(), lengths.len()); + + let mut ctx = SESSION.create_execution_ctx(); + for (index, &length) in lengths.iter().enumerate() { + let slice = list.list_elements_at(index)?; + assert_eq!(slice.len(), length as usize, "list {index} length"); + assert_arrays_eq!( + slice, + elements.slice(offsets[index] as usize..offsets[index + 1] as usize)?, + &mut ctx + ); + } + Ok(()) +} + +// ── Corrupt arrays must raise, never panic ───────────────────────────── + +/// A sample table is fed straight to `BitBuffer::select_range` as a window start, and that asserts +/// on a start past the end. So a file with the right sample *count* and garbage sample *values* has +/// to be rejected at construction, not left to panic on the first query. +#[rstest] +#[case::past_the_end(u64::MAX)] +#[case::just_past_the_end(u64::MAX - 1)] +#[case::out_of_order(0)] +fn test_rejects_corrupt_samples(#[case] poison: u64) -> VortexResult<()> { + // Long enough that both sample tables are populated, so either can be poisoned. + let encoded = encode(&sorted_values(5000, 1 << 30, 0xDEFACED))?; + let samples = encoded.samples_buffer(); + assert!(samples.len() >= 2 * size_of::(), "need two samples"); + + for index in [0usize, samples.len() / size_of::() - 1] { + let mut poisoned = samples.clone().into_mut(); + let start = index * size_of::(); + poisoned[start..start + size_of::()].copy_from_slice(&poison.to_le_bytes()); + + let data = EliasFanoData::try_new( + encoded.upper_buffer().clone(), + poisoned.freeze(), + encoded.reference_scalar().clone(), + encoded.max_scalar().clone(), + encoded.lower_width(), + encoded.upper_len(), + encoded.first_rank(), + )?; + let rebuilt = EliasFano::try_new(data, encoded.lower().clone(), encoded.len()); + assert!( + rebuilt.is_err(), + "a sample of {poison} at index {index} must be rejected" + ); + } + Ok(()) +} + +/// The upper array's *contents* are not validated — that would mean walking the whole buffer on +/// every construction — so both the bulk decode and the cursor have to raise on a malformed one +/// rather than underflow or hand back a short answer. +#[test] +fn test_corrupt_upper_array_raises() -> VortexResult<()> { + let encoded = encode(&sorted_values(600, 1 << 16, 0xBADB175))?; + let mut ctx = SESSION.create_execution_ctx(); + + // Set the sentinel at bit 0. Now the first set bit sits at its own rank, so recovering its high + // part would underflow. + let upper = encoded.upper_buffer(); + let mut poisoned = upper.clone().into_mut(); + poisoned[0] |= 1; + + let data = EliasFanoData::try_new( + poisoned.freeze(), + encoded.samples_buffer().clone(), + encoded.reference_scalar().clone(), + encoded.max_scalar().clone(), + encoded.lower_width(), + encoded.upper_len(), + encoded.first_rank(), + )?; + let rebuilt = EliasFano::try_new(data, encoded.lower().clone(), encoded.len())?; + + // Both entry points must return an error. Each recovers a high part by subtracting a rank from + // a bit position, which this input drives negative, so an unchecked subtraction would panic in + // debug and hand back wrong values in release. + assert!( + rebuilt + .clone() + .into_array() + .execute::(&mut ctx) + .is_err(), + "bulk decode of a malformed upper array must raise" + ); + let mut cursor = EliasFanoCursor::try_new(rebuilt.as_view(), &mut ctx)?; + assert!( + cursor.access(0).is_err(), + "cursor access into a malformed upper array must raise" + ); + Ok(()) +} + + diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index aac8ead42fe..25bc408b745 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -29,6 +29,7 @@ use crate::bit::ops::bitwise_binary_op_lhs_owned; use crate::bit::ops::bitwise_unary_op; use crate::bit::ops::bitwise_unary_op_copy; use crate::bit::select::bit_select; +use crate::bit::select::bit_select_zero; use crate::buffer; /// An immutable bitset stored as a packed byte buffer. @@ -428,6 +429,44 @@ impl BitBuffer { bit_select(self.buffer.as_slice(), self.offset, self.len, nth) } + /// Returns the position of the `nth` set bit within the bit range `[start, end)`, relative + /// to `start`. + /// + /// Unlike `self.slice(start..end).select(nth)`, this walks the existing backing buffer + /// directly without cloning a new [`BitBuffer`], so it stays cheap when called repeatedly + /// over many small windows — the pattern a sampled select index produces. + /// + /// Panics if `start > end` or `end > len`. + #[inline] + pub fn select_range(&self, start: usize, end: usize, nth: usize) -> Option { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= self.len, "end {end} exceeds len {}", self.len); + bit_select( + self.buffer.as_slice(), + self.offset + start, + end - start, + nth, + ) + } + + /// Returns the position of the `nth` unset bit within the bit range `[start, end)`, relative + /// to `start`. + /// + /// The complement of [`Self::select_range`]; see it for why the range form exists. + /// + /// Panics if `start > end` or `end > len`. + #[inline] + pub fn select_zero_range(&self, start: usize, end: usize, nth: usize) -> Option { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= self.len, "end {end} exceeds len {}", self.len); + bit_select_zero( + self.buffer.as_slice(), + self.offset + start, + end - start, + nth, + ) + } + /// Get the number of unset bits in the buffer. #[inline] pub fn false_count(&self) -> usize { diff --git a/vortex-buffer/src/bit/select.rs b/vortex-buffer/src/bit/select.rs index deae368ddfb..6a8c2c4b1c1 100644 --- a/vortex-buffer/src/bit/select.rs +++ b/vortex-buffer/src/bit/select.rs @@ -18,46 +18,83 @@ use crate::dispatch::CpuKernel; /// - **Scalar fallback**: 4× unrolled word scan with `count_ones`, byte-level narrowing. #[inline] pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option { + bit_select_impl(bytes, offset, len, nth, false) +} + +/// Returns the position of the `nth` *unset* bit (0-indexed) within the logical range +/// `[offset, offset + len)` of the given byte slice. +/// +/// The complement of [`bit_select`], and the same walk: every tier below is shared, because a +/// fully-valid region of `width` bits holds `width - popcount` zeros. That identity is exact, so +/// no vector load has to be complemented — only the running totals change, and the complement +/// itself happens at the final scalar narrowing step. +#[inline] +pub fn bit_select_zero(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option { + bit_select_impl(bytes, offset, len, nth, true) +} + +/// Shared implementation of [`bit_select`] (`invert == false`) and [`bit_select_zero`] +/// (`invert == true`). +/// +/// `invert` is loop-invariant at every level, so it costs nothing in the steady state: the +/// caller passes a constant, and each scan loop unswitches on it. It is a runtime parameter +/// rather than a const generic because the tiered kernels declare their `CpuKernel` statics +/// inside their own function bodies, and an item in a function body cannot name that +/// function's generic parameters (E0401). +#[inline] +fn bit_select_impl( + bytes: &[u8], + offset: usize, + len: usize, + nth: usize, + invert: bool, +) -> Option { let (head, middle, tail) = align_offset_len(bytes, offset, len); let mut remaining = nth; let mut pos = 0usize; // ── partial first byte ────────────────────────────────────────────── if let Some(head) = head { + // `align_offset_len` hands back the head already shifted down and masked to its valid + // width, so the bits above that width read as zero. Counting *ones* is therefore correct + // as-is; counting zeros would also count those padding bits, which is why the zero path + // needs the valid width in order to re-mask after complementing. + let start_len = (8 - offset % 8).min(len); + let head = selectable_byte(head, start_len, invert); let count = head.count_ones() as usize; if remaining < count { return Some(select_in_byte(head, remaining)); } remaining -= count; - let start_bit = offset % 8; - pos = (8 - start_bit).min(len); + pos = start_len; } // ── aligned middle bytes ──────────────────────────────────────────── if !middle.is_empty() { let (chunks, tail_bytes) = middle.as_chunks::<64>(); - let (rem, new_pos, chunk_idx) = scan_chunks(chunks, remaining, pos); + let (rem, new_pos, chunk_idx) = scan_chunks(chunks, remaining, pos, invert); remaining = rem; pos = new_pos; if chunk_idx < chunks.len() { - return Some(pos + select_in_chunk(&chunks[chunk_idx], remaining)); + return Some(pos + select_in_chunk(&chunks[chunk_idx], remaining, invert)); } let (words, tail_bytes) = tail_bytes.as_chunks::<8>(); - let (rem, new_pos, word_idx) = scan_words(words, remaining, pos); + let (rem, new_pos, word_idx) = scan_words(words, remaining, pos, invert); remaining = rem; pos = new_pos; if word_idx < words.len() { - let word = u64::from_le_bytes(words[word_idx]); + let word = selectable_word(u64::from_le_bytes(words[word_idx]), invert); return Some(pos + select_in_word(word, remaining)); } // Remaining aligned bytes that don't fill a full u64. for &byte in tail_bytes { + let byte = selectable_byte(byte, 8, invert); let count = byte.count_ones() as usize; if remaining < count { return Some(pos + select_in_byte(byte, remaining)); @@ -68,34 +105,73 @@ pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option } // ── partial last byte ─────────────────────────────────────────────── - if let Some(tail) = tail - && remaining < tail.count_ones() as usize - { - return Some(pos + select_in_byte(tail, remaining)); + // `pos` has now consumed the head plus every aligned middle byte — exactly the `consumed` + // that `align_offset_len` subtracted from `len` to size the tail — so `len - pos` is the + // tail's valid width. + if let Some(tail) = tail { + let tail = selectable_byte(tail, len - pos, invert); + if remaining < tail.count_ones() as usize { + return Some(pos + select_in_byte(tail, remaining)); + } } None } +/// Narrow a byte down to the bits the select should walk. +/// +/// On the ones path this is the identity. On the zeros path the byte is complemented, so its set +/// bits are the input's unset bits, and then re-masked to `valid_len` bits so the padding above +/// the valid range does not become phantom zeros. +#[inline] +fn selectable_byte(byte: u8, valid_len: usize, invert: bool) -> u8 { + if !invert { + return byte; + } + let mask = if valid_len >= 8 { + u8::MAX + } else { + (1u8 << valid_len) - 1 + }; + !byte & mask +} + +/// [`selectable_byte`] for a fully-valid word: no mask is needed, every bit counts. +#[inline] +fn selectable_word(word: u64, invert: bool) -> u64 { + if invert { !word } else { word } +} + +/// How many bits a fully-valid `width`-bit region contributes to the select, given its popcount. +#[inline] +fn selectable_count(width: usize, ones: usize, invert: bool) -> usize { + if invert { width - ones } else { ones } +} + // ── 64-byte chunk scan ────────────────────────────────────────────────── /// Scan `chunks` accumulating popcounts. Returns `(remaining, position, chunk_index)`. /// /// If `chunk_index < chunks.len()`, the target bit is inside that chunk and `remaining` /// is the rank *within* that chunk. Otherwise all chunks were consumed. -type ScanChunks = unsafe fn(&[[u8; 64]], usize, usize) -> (usize, usize, usize); +type ScanChunks = unsafe fn(&[[u8; 64]], usize, usize, bool) -> (usize, usize, usize); #[inline] -fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { +fn scan_chunks( + chunks: &[[u8; 64]], + remaining: usize, + pos: usize, + invert: bool, +) -> (usize, usize, usize) { // Scans of a couple of chunks don't amortize the dispatch indirection: call the // per-architecture unconditional kernel directly so it stays inlinable (see the // size-gating note in the CpuKernel docs). if chunks.len() <= 2 { #[cfg(target_arch = "aarch64")] - return scan_chunks_neon(chunks, remaining, pos); + return scan_chunks_neon(chunks, remaining, pos, invert); #[allow(unreachable_code)] { - return scan_chunks_scalar(chunks, remaining, pos); + return scan_chunks_scalar(chunks, remaining, pos, invert); } } @@ -117,7 +193,7 @@ fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usi }); // SAFETY: the selector only returns kernels that are safe or whose required CPU // features were probed before selection. - unsafe { KERNEL.get()(chunks, remaining, pos) } + unsafe { KERNEL.get()(chunks, remaining, pos, invert) } } #[cfg(target_arch = "aarch64")] @@ -127,6 +203,7 @@ fn scan_chunks_neon( chunks: &[[u8; 64]], mut remaining: usize, mut pos: usize, + invert: bool, ) -> (usize, usize, usize) { use std::arch::aarch64::vcntq_u8; use std::arch::aarch64::vgetq_lane_u64; @@ -139,7 +216,7 @@ fn scan_chunks_neon( let ptr = chunk.as_ptr(); // SAFETY: chunk is exactly 64 bytes split across four 128-bit NEON loads. // NEON vld1q_u8 supports unaligned access. - let total = unsafe { + let ones = unsafe { let pop_0 = vcntq_u8(vld1q_u8(ptr)); let pop_1 = vcntq_u8(vld1q_u8(ptr.add(16))); let pop_2 = vcntq_u8(vld1q_u8(ptr.add(32))); @@ -158,6 +235,7 @@ fn scan_chunks_neon( + vgetq_lane_u64::<0>(sums_3) + vgetq_lane_u64::<1>(sums_3)) as usize }; + let total = selectable_count(512, ones, invert); if remaining < total { return (remaining, pos, idx); @@ -176,6 +254,7 @@ unsafe fn scan_chunks_avx512_vpopcnt( chunks: &[[u8; 64]], mut remaining: usize, mut pos: usize, + invert: bool, ) -> (usize, usize, usize) { use std::arch::x86_64::_mm512_loadu_si512; use std::arch::x86_64::_mm512_popcnt_epi64; @@ -187,8 +266,9 @@ unsafe fn scan_chunks_avx512_vpopcnt( // SAFETY: chunk is exactly 64 bytes. `_mm512_loadu_si512` supports unaligned access. let block = unsafe { _mm512_loadu_si512(chunk.as_ptr().cast()) }; let counts = _mm512_popcnt_epi64(block); - let total = + let ones = usize::try_from(_mm512_reduce_add_epi64(counts)).vortex_expect("must fit in usize"); + let total = selectable_count(512, ones, invert); if remaining < total { return (remaining, pos, idx); @@ -206,9 +286,10 @@ fn scan_chunks_scalar( chunks: &[[u8; 64]], mut remaining: usize, mut pos: usize, + invert: bool, ) -> (usize, usize, usize) { for (idx, chunk) in chunks.iter().enumerate() { - let total = count_ones_chunk(chunk); + let total = selectable_count(512, count_ones_chunk(chunk), invert); if remaining < total { return (remaining, pos, idx); } @@ -227,15 +308,25 @@ fn scan_chunks_scalar( /// If `word_index < words.len()`, the target bit is inside that word and `remaining` /// is the rank *within* that word. Otherwise all words were consumed. #[inline] -fn scan_words(words: &[[u8; 8]], remaining: usize, pos: usize) -> (usize, usize, usize) { - scan_words_impl(words, remaining, pos) +fn scan_words( + words: &[[u8; 8]], + remaining: usize, + pos: usize, + invert: bool, +) -> (usize, usize, usize) { + scan_words_impl(words, remaining, pos, invert) } // ── Scalar word scan ──────────────────────────────────────────────────── #[inline] -fn scan_words_impl(words: &[[u8; 8]], remaining: usize, pos: usize) -> (usize, usize, usize) { - scan_words_scalar(words, remaining, pos) +fn scan_words_impl( + words: &[[u8; 8]], + remaining: usize, + pos: usize, + invert: bool, +) -> (usize, usize, usize) { + scan_words_scalar(words, remaining, pos, invert) } #[inline] @@ -243,15 +334,23 @@ fn scan_words_scalar( words: &[[u8; 8]], mut remaining: usize, mut pos: usize, + invert: bool, ) -> (usize, usize, usize) { let mut idx = 0; + let count_at = |idx: usize| { + selectable_count( + 64, + u64::from_le_bytes(words[idx]).count_ones() as usize, + invert, + ) + }; // 4× unrolled: the four independent `count_ones` calls pipeline well. while idx + 4 <= words.len() { - let count_0 = u64::from_le_bytes(words[idx]).count_ones() as usize; - let count_1 = u64::from_le_bytes(words[idx + 1]).count_ones() as usize; - let count_2 = u64::from_le_bytes(words[idx + 2]).count_ones() as usize; - let count_3 = u64::from_le_bytes(words[idx + 3]).count_ones() as usize; + let count_0 = count_at(idx); + let count_1 = count_at(idx + 1); + let count_2 = count_at(idx + 2); + let count_3 = count_at(idx + 3); let total = count_0 + count_1 + count_2 + count_3; if remaining >= total { @@ -280,8 +379,7 @@ fn scan_words_scalar( } while idx < words.len() { - let word = u64::from_le_bytes(words[idx]); - let count = word.count_ones() as usize; + let count = count_at(idx); if remaining < count { return (remaining, pos, idx); } @@ -295,11 +393,11 @@ fn scan_words_scalar( // ── In-chunk select ───────────────────────────────────────────────────── -type SelectInChunk = unsafe fn(&[u8; 64], usize) -> usize; +type SelectInChunk = unsafe fn(&[u8; 64], usize, bool) -> usize; -/// Position of the `nth` set bit inside a 64-byte chunk (0-indexed). +/// Position of the `nth` set (or, when `invert`, unset) bit inside a 64-byte chunk (0-indexed). #[inline] -fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize { +fn select_in_chunk(chunk: &[u8; 64], nth: usize, invert: bool) -> usize { static KERNEL: CpuKernel = CpuKernel::new(|| { #[cfg(target_arch = "x86_64")] { @@ -314,12 +412,12 @@ fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize { }); // SAFETY: the selector only returns kernels that are safe or whose required CPU // features were probed before selection. - unsafe { KERNEL.get()(chunk, nth) } + unsafe { KERNEL.get()(chunk, nth, invert) } } #[cfg(target_arch = "x86_64")] #[target_feature(enable = "avx512f,avx512vpopcntdq,avx512vbmi2")] -unsafe fn select_in_chunk_vbmi2(chunk: &[u8; 64], mut nth: usize) -> usize { +unsafe fn select_in_chunk_vbmi2(chunk: &[u8; 64], mut nth: usize, invert: bool) -> usize { use std::arch::x86_64::_mm512_loadu_si512; use std::arch::x86_64::_mm512_popcnt_epi64; use std::arch::x86_64::_mm512_storeu_epi64; @@ -336,10 +434,12 @@ unsafe fn select_in_chunk_vbmi2(chunk: &[u8; 64], mut nth: usize) -> usize { // SAFETY: `lane_counts` has room for all eight i64 lanes. unsafe { _mm512_storeu_epi64(lane_counts.as_mut_ptr(), counts) }; - for (idx, count) in lane_counts.into_iter().enumerate() { - let count = usize::try_from(count).vortex_expect("must fit in usize"); + for (idx, ones) in lane_counts.into_iter().enumerate() { + let ones = usize::try_from(ones).vortex_expect("must fit in usize"); + let count = selectable_count(64, ones, invert); if nth < count { - return idx * 64 + select_in_word(u64::from_le_bytes(words[idx]), nth); + let word = selectable_word(u64::from_le_bytes(words[idx]), invert); + return idx * 64 + select_in_word(word, nth); } nth -= count; } @@ -348,11 +448,11 @@ unsafe fn select_in_chunk_vbmi2(chunk: &[u8; 64], mut nth: usize) -> usize { } #[inline] -fn select_in_chunk_scalar(chunk: &[u8; 64], mut nth: usize) -> usize { +fn select_in_chunk_scalar(chunk: &[u8; 64], mut nth: usize, invert: bool) -> usize { let words = chunk.as_chunks::<8>().0; for (idx, word) in words.iter().enumerate() { - let word = u64::from_le_bytes(*word); + let word = selectable_word(u64::from_le_bytes(*word), invert); let count = word.count_ones() as usize; if nth < count { return idx * 64 + select_in_word(word, nth); @@ -483,6 +583,52 @@ mod tests { assert_eq!(bit_select(&buf, 0, 8, 2), None); } + /// Deterministic ~50% density filler, matching the pattern the original tests used. + fn mixed_bytes(total_bytes: usize) -> Vec { + (0..total_bytes) + .map(|i| ((i.wrapping_mul(0x9E) ^ 0xA5) & 0xFF) as u8) + .collect() + } + + /// Every position in `[offset, offset + len)` whose bit equals `want`, in ascending order. + fn naive_positions(buf: &[u8], offset: usize, len: usize, want: bool) -> Vec { + (0..len) + .filter(|&i| { + let phys = offset + i; + ((buf[phys / 8] >> (phys % 8)) & 1 == 1) == want + }) + .collect() + } + + /// Both select variants must agree with a bit-at-a-time reference over the whole rank range, + /// and both must report `None` one past the last rank. + fn check_against_naive(buf: &[u8], offset: usize, len: usize) { + for (want, select) in [ + ( + true, + bit_select as fn(&[u8], usize, usize, usize) -> Option, + ), + ( + false, + bit_select_zero as fn(&[u8], usize, usize, usize) -> Option, + ), + ] { + let expected = naive_positions(buf, offset, len, want); + for (nth, &expected_pos) in expected.iter().enumerate() { + assert_eq!( + select(buf, offset, len, nth), + Some(expected_pos), + "want={want} offset={offset} len={len} nth={nth}" + ); + } + assert_eq!( + select(buf, offset, len, expected.len()), + None, + "want={want} offset={offset} len={len} past-the-end rank" + ); + } + } + #[rstest] #[case(0, 128)] #[case(3, 100)] @@ -497,27 +643,80 @@ mod tests { #[case(0, 512)] #[case(0, 513)] #[case(5, 1024)] + // Head-only windows: an offset inside the first byte with a length that never leaves it, so + // `align_offset_len` produces a head and nothing else. The zero path has to know the head's + // valid width here, or it counts the padding bits above it. + #[case(1, 1)] + #[case(1, 6)] + #[case(4, 3)] + #[case(7, 1)] + // Windows whose last byte is partial, exercising the tail's valid width. + #[case(0, 9)] + #[case(0, 63)] + #[case(2, 71)] + #[case(6, 130)] + #[case(3, 517)] fn test_select_agrees_with_naive(#[case] offset: usize, #[case] len: usize) { - let total_bits = offset + len; - let total_bytes = total_bits.div_ceil(8); - // Deterministic pattern with moderate density. - let buf: Vec = (0..total_bytes) - .map(|i| ((i.wrapping_mul(0x9E) ^ 0xA5) & 0xFF) as u8) - .collect(); + check_against_naive(&mixed_bytes((offset + len).div_ceil(8)), offset, len); + } - // Collect set-bit positions naively. - let expected: Vec = (0..len) - .filter(|&i| { - let phys = offset + i; - (buf[phys / 8] >> (phys % 8)) & 1 == 1 - }) - .collect(); + /// The mixed pattern above sits near 50% density, which is exactly where confusing ones with + /// zeros is least visible. These degenerate densities make it obvious. + #[rstest] + #[case::all_zero(0x00)] + #[case::all_one(0xFF)] + #[case::sparse(0x01)] + #[case::dense(0xFE)] + fn test_select_uniform_density(#[case] fill: u8) { + for (offset, len) in [(0usize, 8usize), (0, 128), (3, 5), (5, 130), (1, 517)] { + let buf = vec![fill; (offset + len).div_ceil(8)]; + check_against_naive(&buf, offset, len); + } + } + + #[test] + fn test_select_zero_degenerate_buffers() { + // All ones: no zero to find, at any rank. + let ones = [0xFFu8; 16]; + assert_eq!(bit_select_zero(&ones, 0, 128, 0), None); + + // All zeros: the nth zero is at position n, and the nth one does not exist. + let zeros = [0x00u8; 16]; + for nth in 0..128 { + assert_eq!(bit_select_zero(&zeros, 0, 128, nth), Some(nth), "nth={nth}"); + } + assert_eq!(bit_select_zero(&zeros, 0, 128, 128), None); + assert_eq!(bit_select(&zeros, 0, 128, 0), None); + } + + /// Cross-check the zero select against the already-public counting entry points: the number of + /// zeros in a window is `len - count_ones`, and that is the first rank that must miss. + #[test] + fn test_select_zero_count_agrees_with_count_ones() { + let buf = mixed_bytes(300); + for (offset, len) in [(0usize, 2400usize), (5, 2000), (7, 1), (3, 519)] { + let zeros = len - super::super::count_ones::count_ones(&buf, offset, len); + if let Some(last) = zeros.checked_sub(1) { + assert!(bit_select_zero(&buf, offset, len, last).is_some()); + } + assert_eq!(bit_select_zero(&buf, offset, len, zeros), None); + } + } + + #[test] + fn test_select_zero_large_buffer() { + // ~64 KB buffer, spanning many 64-byte chunks so the chunk-scan tier runs. + let len = 65_536 * 8; + let buf = mixed_bytes(65_536); + let zeros = len - super::super::count_ones::count_ones(&buf, 0, len); - for (nth, &expected_pos) in expected.iter().enumerate() { + for nth in [0usize, 1, 1000, zeros / 2, zeros - 1] { + let pos = bit_select_zero(&buf, 0, len, nth).expect("rank is in bounds"); + assert_eq!(buf[pos / 8] & (1 << (pos % 8)), 0, "nth={nth} pos={pos}"); assert_eq!( - bit_select(&buf, offset, len, nth), - Some(expected_pos), - "offset={offset} len={len} nth={nth}" + super::super::count_ones::count_ones(&buf, 0, pos), + pos - nth, + "nth={nth}: rank1(select0(nth)) must be select0(nth) - nth" ); } } diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index fc406f7133e..95666e81d56 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -40,6 +40,7 @@ vortex-bytebool = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } vortex-edition = { workspace = true } +vortex-elias-fano = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["file"] } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 25760f0890e..84d9388af24 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -189,6 +189,7 @@ pub fn register_default_encodings(session: &VortexSession) { vortex_alp::initialize(session); vortex_datetime_parts::initialize(session); vortex_decimal_byte_parts::initialize(session); + vortex_elias_fano::initialize(session); vortex_fastlanes::initialize(session); vortex_runend::initialize(session); vortex_sequence::initialize(session); diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 6a2a840a500..4f364c25d61 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -33,6 +33,7 @@ vortex-cloud = { workspace = true, optional = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } vortex-edition = { workspace = true } +vortex-elias-fano = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-file = { workspace = true, optional = true, default-features = true } diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index 2079a710f58..dda8a6ae6f2 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -44,13 +44,14 @@ pub use self::unstable::UNSTABLE_2025_05_0; pub use self::unstable::UNSTABLE_2026_02_0; pub use self::unstable::UNSTABLE_2026_04_0; pub use self::unstable::UNSTABLE_2026_06_0; +pub use self::unstable::UNSTABLE_2026_08_0; /// The `core` edition enabled for writing by the default Vortex session. pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_08; /// The `unstable` edition enabled for writing by the default Vortex session when the /// `unstable_encodings` feature is selected. -pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_06_0; +pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_08_0; /// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ @@ -63,6 +64,7 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &unstable::v2026_02::DECLARATION, &unstable::v2026_04::DECLARATION, &unstable::v2026_06::DECLARATION, + &unstable::v2026_08::DECLARATION, ]; /// Register the Vortex edition declarations with the session's [`EditionSession`]. diff --git a/vortex/src/editions/unstable/mod.rs b/vortex/src/editions/unstable/mod.rs index 5544ba45c0f..7c271198e25 100644 --- a/vortex/src/editions/unstable/mod.rs +++ b/vortex/src/editions/unstable/mod.rs @@ -10,8 +10,10 @@ pub mod v2025_05; pub mod v2026_02; pub mod v2026_04; pub mod v2026_06; +pub mod v2026_08; pub use v2025_05::UNSTABLE_2025_05_0; pub use v2026_02::UNSTABLE_2026_02_0; pub use v2026_04::UNSTABLE_2026_04_0; pub use v2026_06::UNSTABLE_2026_06_0; +pub use v2026_08::UNSTABLE_2026_08_0; diff --git a/vortex/src/editions/unstable/v2026_08.rs b/vortex/src/editions/unstable/v2026_08.rs new file mode 100644 index 00000000000..454b8d71a17 --- /dev/null +++ b/vortex/src/editions/unstable/v2026_08.rs @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The August 2026 `unstable` encoding cohort. + +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; + +/// The August 2026 draft edition of the `unstable` family. +pub const UNSTABLE_2026_08_0: EditionId = EditionId::new("unstable", 2026, 8, 0); + +/// The declaration of [`UNSTABLE_2026_08_0`] and the encodings that join the family at it. +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: UNSTABLE_2026_08_0, + min_vortex_version: None, + }, + added: &[EditionMember::array(&"vortex.elias_fano")], +}; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 289500e2543..33fd28f86e5 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -255,6 +255,11 @@ pub mod encodings { pub use vortex_decimal_byte_parts::*; } + /// Elias-Fano encoding for sorted integer sequences, such as a list column's offsets. + pub mod elias_fano { + pub use vortex_elias_fano::*; + } + /// FastLanes integer encodings: bit-packing, delta, frame-of-reference, and RLE. pub mod fastlanes { pub use vortex_fastlanes::*; From 1416e5c585f5c8e27800e9ba30eb0f380f31a02b Mon Sep 17 00:00:00 2001 From: rapour Date: Fri, 21 Aug 2026 16:56:21 +0330 Subject: [PATCH 2/2] chore: run scalar_at without a cursor Signed-off-by: rapour --- Cargo.lock | 1 + encodings/elias-fano/src/array.rs | 8 +- encodings/elias-fano/src/compress.rs | 187 ++++++++++++---- encodings/elias-fano/src/cursor.rs | 176 +++++++++++---- encodings/elias-fano/src/kernel.rs | 6 +- encodings/elias-fano/src/lib.rs | 4 +- encodings/elias-fano/src/params.rs | 48 ++++ encodings/elias-fano/src/tests.rs | 142 ++++++++++-- vortex-btrblocks/Cargo.toml | 7 +- vortex-btrblocks/src/builder.rs | 5 + .../src/schemes/integer/elias_fano.rs | 206 ++++++++++++++++++ vortex-btrblocks/src/schemes/integer/mod.rs | 4 + .../schemes/integer/scheme_selection_tests.rs | 133 +++++++++++ vortex-btrblocks/src/trace_tests.rs | 7 + .../golden__unstable__list_of_int_runs.snap | 14 +- 15 files changed, 826 insertions(+), 122 deletions(-) create mode 100644 vortex-btrblocks/src/schemes/integer/elias_fano.rs diff --git a/Cargo.lock b/Cargo.lock index db6126e0a81..e5b65e525ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9769,6 +9769,7 @@ dependencies = [ "vortex-compressor", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-elias-fano", "vortex-error", "vortex-fastlanes", "vortex-fsst", diff --git a/encodings/elias-fano/src/array.rs b/encodings/elias-fano/src/array.rs index d6114b808ba..4f35b84427e 100644 --- a/encodings/elias-fano/src/array.rs +++ b/encodings/elias-fano/src/array.rs @@ -52,7 +52,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::compress::elias_fano_decompress; -use crate::cursor::EliasFanoCursor; +use crate::cursor::access_at; use crate::params; use crate::params::LOG_SAMPLING0; use crate::params::LOG_SAMPLING1; @@ -529,12 +529,16 @@ impl VTable for EliasFano { } impl OperationsVTable for EliasFano { + /// A single access, without a cursor: nothing a cursor carries survives to the next call, and + /// every batched path builds its own and keeps it. A caller making a stream of probes wants + /// [`EliasFanoCursor`](crate::EliasFanoCursor) directly, which amortises the setup and + /// remembers where it stopped. fn scalar_at( array: ArrayView<'_, EliasFano>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - EliasFanoCursor::try_new(array, ctx)?.access(index) + access_at(array, index, ctx) } } diff --git a/encodings/elias-fano/src/compress.rs b/encodings/elias-fano/src/compress.rs index 44d6373956c..7067fa7baee 100644 --- a/encodings/elias-fano/src/compress.rs +++ b/encodings/elias-fano/src/compress.rs @@ -5,8 +5,6 @@ //! //! See [`crate::params`] for the layout both directions read and write. -use std::iter; - use lending_iterator::prelude::LendingIterator; use num_traits::AsPrimitive; use vortex_array::ArrayRef; @@ -24,13 +22,11 @@ use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; -use vortex_buffer::BitIndexIterator; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_fastlanes::BitPacked; use vortex_fastlanes::BitPackedArrayExt; use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked; @@ -159,6 +155,53 @@ pub(crate) fn elias_fano_decompress( })) } +/// A resumable walk over the set bits of the upper window, a `u64` word at a time. +/// +/// This is the shape the decode loop wants and a pull iterator is not: the walk holds only two +/// values, so a caller can lift them into locals for the length of one FastLanes block and pay +/// nothing per element beyond a `trailing_zeros` and a clear-lowest-bit. The words are materialised +/// once by [`window_words`] so this can index a slice instead of driving a shifting iterator. +struct Ones<'a> { + words: &'a [u64], + /// Index of the word `current` was taken from. + word: usize, + /// The bits of that word not yet returned. + current: u64, +} + +impl<'a> Ones<'a> { + fn new(words: &'a [u64]) -> Self { + Self { + words, + word: 0, + current: words.first().copied().unwrap_or(0), + } + } + + /// The next set bit's index within the window, or `None` once the words run out. + /// + /// `decode` establishes that the window holds exactly as many set bits as there are elements, + /// so on a well-formed array this never returns `None` before the last element. + #[inline] + fn next(&mut self) -> Option { + while self.current == 0 { + self.word += 1; + self.current = *self.words.get(self.word)?; + } + let bit = self.current.trailing_zeros() as usize; + self.current &= self.current - 1; + Some(self.word * u64::BITS as usize + bit) + } +} + +/// The window's bits as whole `u64` words, shifted to a zero bit offset once. +/// +/// The upper array runs at roughly two bits per element, so this is `n / 32` words — a few KB for a +/// whole chunk, against the half-megabyte a materialised copy of the low bits would cost. +fn window_words(window: &BitBuffer) -> Buffer { + window.chunks().iter_padded().collect() +} + /// Reassembles elements from the two halves of the layout, in the column's own width. struct Fold { /// Bit position the upper window starts at, which its set-bit indices are relative to. @@ -172,6 +215,15 @@ struct Fold { impl Fold { /// Decode `len` elements, taking high parts from `window`'s set bits and low parts from /// `lower`, which is read one FastLanes block at a time. + /// + /// The two invariants the per-element loop relies on are established once here rather than + /// re-tested for every element. `validate` does not read the upper buffer's contents, so a + /// corrupt file can still violate them — it just fails once, up front, instead of `len` times: + /// + /// * the window holds exactly `len` set bits, so the walk cannot run dry; and + /// * `start > first_rank`, which gives `position >= rank + 1` for every element, because + /// the `i`-th set bit of the window sits at window index `>= i` while its rank is + /// `first_rank + i`. fn decode( &self, window: &BitBuffer, @@ -182,13 +234,29 @@ impl Fold { where u64: AsPrimitive

, { - let mut values = BufferMut::

::with_capacity(len); - let mut ones = window.set_indices(); + let ones_in_window = window.true_count(); + vortex_ensure!( + ones_in_window == len, + "Elias-Fano upper array is malformed: expected exactly {len} set bits above their own \ + ranks, found {ones_in_window}" + ); + vortex_ensure!( + self.start as u64 > self.first_rank, + "Elias-Fano upper array is malformed: the element of rank {} sits at bit {}, at or \ + below its own rank", + self.first_rank, + self.start + ); + + let words = window_words(window); + let mut ones = Ones::new(words.as_slice()); + let mut values = BufferMut::

::zeroed(len); + let mut rank = 0usize; if self.lower_width == 0 { // Nothing is stored, so do not execute the slot just to read `len` zeros. - self.segment(&mut values, &mut ones, iter::repeat_n(0, len))?; - return self.finish(values, ones, len); + self.segment(values.as_mut_slice(), &mut ones, &mut rank, None, len); + return self.finish(values, rank, len); } // Window before reading: the child spans the whole encoded sequence, and a decode only ever @@ -207,80 +275,108 @@ impl Fold { { let mut chunks = packed.unpacked_chunks::()?; if let Some(initial) = chunks.initial() { - self.segment(&mut values, &mut ones, initial.iter().copied())?; + self.segment( + values.as_mut_slice(), + &mut ones, + &mut rank, + Some(initial), + len, + ); } // A single-block child is covered by `initial` alone, and the later phases would hand // that same block back. - if values.len() < len { + if rank < len { let mut full = chunks.full_chunks(); while let Some(chunk) = full.next() { - self.segment(&mut values, &mut ones, chunk.iter().copied())?; + self.segment( + values.as_mut_slice(), + &mut ones, + &mut rank, + Some(chunk), + len, + ); } } - if values.len() < len + if rank < len && let Some(trailer) = chunks.trailer() { - self.segment(&mut values, &mut ones, trailer.iter().copied())?; + self.segment( + values.as_mut_slice(), + &mut ones, + &mut rank, + Some(trailer), + len, + ); } } else { // The slot is patched, device-resident, or some other encoding after a rewrite. let dense = window_lower.execute::(ctx)?; - self.segment(&mut values, &mut ones, dense.as_slice::().iter().copied())?; + self.segment( + values.as_mut_slice(), + &mut ones, + &mut rank, + Some(dense.as_slice::()), + len, + ); } - self.finish(values, ones, len) + self.finish(values, rank, len) } - /// Fold one run of consecutive low parts onto the end of `values`. - fn segment>( + /// Fold one run of consecutive low parts into `out`, starting at `rank`. + /// + /// `lows` of `None` is the `lower_width == 0` layout, where every low part is zero. Nothing here + /// can fail: `decode` has already established both invariants, and a walk that runs dry early + /// leaves `rank` short for `finish` to reject. + fn segment( &self, - values: &mut BufferMut

, - ones: &mut BitIndexIterator<'_>, - lows: L, - ) -> VortexResult<()> - where + out: &mut [P], + ones: &mut Ones<'_>, + rank: &mut usize, + lows: Option<&[u64]>, + len: usize, + ) where u64: AsPrimitive

, { - for low in lows { - let rank = self.first_rank + values.len() as u64; - let position = ones.next().ok_or_else(|| { - vortex_err!("Elias-Fano upper array holds no element of rank {rank}") - })?; - let position = self.start + position; - // The inverse of the encoder's `position = (element >> lower_width) + rank + 1`. - let high = (position as u64).checked_sub(rank + 1).ok_or_else(|| { - vortex_err!( - "Elias-Fano upper array is malformed: the element of rank {rank} sits at bit \ - {position}, at or below its own rank" - ) - })?; + let remaining = len - *rank; + let take = lows.map_or(remaining, |lows| lows.len().min(remaining)); + + for offset in 0..take { + let index = *rank + offset; + let Some(position) = ones.next() else { + *rank = index; + return; + }; + // The inverse of the encoder's `position = (element >> lower_width) + rank + 1`. The + // subtraction is non-negative by `decode`'s second invariant. + let high = (self.start + position) as u64 - (self.first_rank + index as u64 + 1); // The low bits are masked for the same reason as in the cursor's `lower_at`: only a // bit-packed child's width is checkable at construction, so a patched or rewritten slot // could otherwise carry bits above `lower_width` into the high part. + let low = lows.map_or(0, |lows| lows[offset] & self.lower_mask); let bits = self .reference_bits - .wrapping_add((high << self.lower_width) | (low & self.lower_mask)); + .wrapping_add((high << self.lower_width) | low); // Truncating the pattern to the column's width is exactly the two's complement result, // signed or unsigned, because the reference was added in the same modular arithmetic. - values.push(bits.as_()); + out[index] = bits.as_(); } - Ok(()) + + *rank += take; } - /// The window holds exactly `len` set bits, and the child exactly `len` low parts, for any - /// array this crate builds — but `validate` does not check the upper buffer's contents, so a - /// corrupt file can hold a different number of either. + /// The child holds exactly `len` low parts for any array this crate builds, but `validate` does + /// not check that, so a corrupt file can hand back fewer. fn finish( &self, values: BufferMut

, - mut ones: BitIndexIterator<'_>, + rank: usize, len: usize, ) -> VortexResult> { vortex_ensure!( - values.len() == len && ones.next().is_none(), + rank == len, "Elias-Fano upper array is malformed: expected exactly {len} set bits above their own \ - ranks, found {}", - values.len() + ranks, found {rank}" ); Ok(values.freeze()) } @@ -364,7 +460,6 @@ fn pack_lower(lower: Buffer, lower_width: u8, n: usize) -> VortexResult { + packed: &'a [u64], + bit_width: usize, + /// The child's own sub-block offset, which `unpack_single_primitive` does not apply itself + /// (unlike `unpack_single`). Forgetting it is a silent wrong answer, not a panic. + child_offset: usize, +} + +impl PackedLower<'_> { + /// Where in the child the element of absolute rank `rank` sits. + #[inline] + fn index_of(&self, rank: u64) -> usize { + rank as usize + self.child_offset + } + + /// One value unpacked on its own, without a scratch block. + #[inline] + fn unpack_one(&self, index: usize) -> u64 { + // SAFETY: `packed` is `BitPackedData`'s own buffer, whose length the child's validation + // already tied to `bit_width` and a whole number of blocks, and `index` is within the + // child's length because `validate_parts` ties `first_rank + len` to it. + unsafe { unpack_single_primitive::(self.packed, self.bit_width, index) } + } +} + +/// The low-bits child as a packed slice, if it can be read in place: FastLanes-packed, unpatched, +/// host-resident and `u64`-aligned. `None` means the low bits have to be materialised instead. +fn packed_in_place(lower: &ArrayRef) -> Option> { + // `as_opt`, never `as_`: with the experimental patched-array plugin enabled the slot comes back + // from a file as `Patched(BitPacked)`, and a rewrite may replace it outright. + let packed = lower.as_opt::()?; + if packed.patches().is_some() + || !packed + .packed() + .as_host_opt() + .is_some_and(|buffer| buffer.is_aligned(Alignment::of::())) + { + return None; + } + Some(PackedLower { + // `.data()` rather than the `Deref`, which would borrow the view rather than the array + // behind it and so not live long enough. + packed: packed.data().packed_slice::(), + bit_width: packed.bit_width() as usize, + child_offset: packed.offset() as usize, + }) +} + +/// The value at logical `index`, read without building a cursor. +/// +/// A cursor earns its setup back over a stream of probes — the seat, the memoised answer and the +/// bulk-unpack scratch all amortise — and every batched path builds one and reuses it. A point +/// lookup through `OperationsVTable::scalar_at` amortises none of it, so it does just the two reads +/// an access needs: one sampled `select1` for the high part, and one low-bits read. +/// +/// The two readers must agree, and share only the layout; `tests::check_access` runs them against +/// each other over every shape the suite covers. +pub(crate) fn access_at( + array: ArrayView<'_, EliasFano>, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + if index >= len { + vortex_bail!(OutOfBounds: index, 0usize, len); + } + let data = array.data(); + let lower_width = data.lower_width(); + let rank = data.first_rank() + index as u64; + + let (_, samples1) = data.sample_bytes()?; + let upper = data.upper_bits()?; + let upper_len = usize::try_from(data.upper_len())?; + let position = position_of_rank(&upper, samples1, upper_len, rank)?; + // The inverse of the encoder's `position = (element >> lower_width) + rank + 1`. Checked for + // the same reason as in `EliasFanoCursor::reseat`: the upper buffer's contents are never + // validated, so a corrupt one has to raise rather than underflow. + let high = (position as u64).checked_sub(rank + 1).ok_or_else(|| { + vortex_err!( + "Elias-Fano upper array is malformed: the element of rank {rank} sits at bit \ + {position}, at or below its own rank" + ) + })?; + + // The slots view borrows the array behind the `ArrayView`; the `lower()` accessor would borrow + // the (`Copy`, stack-local) view itself. + let lower = EliasFanoSlotsView::from_slots(array.slots()).lower; + let element = (high << lower_width) | lower_at_rank(lower, lower_width, rank, ctx)?; + scalar_from_bits(array.dtype(), data.reference_bits().wrapping_add(element)) +} + +/// The low bits of the element at absolute rank `rank`, for a reader that will only ask once. +/// +/// Masked for the same reason as [`EliasFanoCursor::lower_at`]: only a bit-packed child's width is +/// checkable at construction, so a patched or rewritten slot can carry bits above `lower_width`, +/// and those would bleed into the high part. +/// +/// The fallback windows to the single rank rather than to the whole slice as +/// [`LowerBits::try_new`] does, so one probe against a child that cannot be read in place +/// materialises one value rather than `len` of them. +fn lower_at_rank( + lower: &ArrayRef, + lower_width: u8, + rank: u64, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if lower_width == 0 { + return Ok(0); + } + let bits = match packed_in_place(lower) { + Some(packed) => packed.unpack_one(packed.index_of(rank)), + None => { + let index = usize::try_from(rank)?; + lower + .slice(index..index + 1)? + .execute::(ctx)? + .into_buffer::()[0] + } + }; + Ok(bits & lower_mask(lower_width)) +} + /// How the low bits of each element can be read. enum LowerBits<'a> { /// The width is zero, so there are no low bits to read and nothing is stored. Zero, /// The normal case, where the low bits are read straight out of the FastLanes-packed child in /// place. - Packed { - packed: &'a [u64], - bit_width: usize, - /// The child's own sub-block offset, which `unpack_single_primitive` does not apply - /// itself (unlike `unpack_single`). Forgetting it is a silent wrong answer, not a panic. - child_offset: usize, - }, + Packed(PackedLower<'a>), /// The fallback for a child that cannot be read in place — patches, device memory, /// under-aligned, or a slot some rewrite replaced. The low bits are materialised once, up /// front. @@ -459,17 +580,13 @@ impl<'a> EliasFanoCursor<'a> { fn lower_at_unmasked(&mut self, rank: u64) -> u64 { // Copy the descriptor out before touching the scratch buffer: the packed slice borrows the // array, not `self`, so this keeps the borrow checker out of the way. - let (packed, bit_width, child_offset) = match &self.lower { + let packed = match &self.lower { LowerBits::Zero => return 0, LowerBits::Dense { values, base } => return values[(rank - base) as usize], - LowerBits::Packed { - packed, - bit_width, - child_offset, - } => (*packed, *bit_width, *child_offset), + LowerBits::Packed(packed) => *packed, }; - let index = rank as usize + child_offset; + let index = packed.index_of(rank); let chunk = index / FL_CHUNK_SIZE; let within_chunk = index % FL_CHUNK_SIZE; @@ -484,23 +601,22 @@ impl<'a> EliasFanoCursor<'a> { self.hot_reads = 1; } - let elems_per_chunk = 128 * bit_width / size_of::(); + let elems_per_chunk = 128 * packed.bit_width / size_of::(); if self.hot_reads > BULK_UNPACK_THRESHOLD { - let block = &packed[chunk * elems_per_chunk..][..elems_per_chunk]; + let block = &packed.packed[chunk * elems_per_chunk..][..elems_per_chunk]; let scratch = self .scratch .get_or_insert_with(|| Box::new([0u64; FL_CHUNK_SIZE])); // SAFETY: `block` is exactly `elems_per_chunk` packed values, and `scratch` is exactly // one FastLanes block of 1024 values, which is what `unchecked_unpack` requires. - unsafe { BitPacking::unchecked_unpack(bit_width, block, scratch.as_mut_slice()) }; + unsafe { + BitPacking::unchecked_unpack(packed.bit_width, block, scratch.as_mut_slice()) + }; self.scratch_chunk = Some(chunk); return self.scratch_slice()[within_chunk]; } - // SAFETY: `packed` is `BitPackedData`'s own buffer, whose length the array's validation - // already tied to `bit_width` and a whole number of blocks, and `index` is within the - // child's length because `first_rank + len` is validated against it. - unsafe { unpack_single_primitive::(packed, bit_width, index) } + packed.unpack_one(index) } fn scratch_slice(&self) -> &[u64; FL_CHUNK_SIZE] { @@ -546,22 +662,8 @@ impl<'a> LowerBits<'a> { if lower_width == 0 { return Ok(LowerBits::Zero); } - // `as_opt`, never `as_`: with the experimental patched-array plugin enabled the slot comes - // back from a file as `Patched(BitPacked)`, and a rewrite may replace it outright. - if let Some(packed) = lower.as_opt::() - && packed.patches().is_none() - && packed - .packed() - .as_host_opt() - .is_some_and(|buffer| buffer.is_aligned(Alignment::of::())) - { - return Ok(LowerBits::Packed { - // `.data()` rather than the `Deref`, which would borrow the view rather than the - // array behind it and so not live long enough. - packed: packed.data().packed_slice::(), - bit_width: packed.bit_width() as usize, - child_offset: packed.offset() as usize, - }); + if let Some(packed) = packed_in_place(lower) { + return Ok(LowerBits::Packed(packed)); } // Window before executing: the child spans the whole encoded sequence, and a cursor only // ever asks for ranks inside its own slice. diff --git a/encodings/elias-fano/src/kernel.rs b/encodings/elias-fano/src/kernel.rs index ff73b372ada..07503676844 100644 --- a/encodings/elias-fano/src/kernel.rs +++ b/encodings/elias-fano/src/kernel.rs @@ -20,7 +20,11 @@ use crate::EliasFano; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); - kernels.register_execute_parent_kernel(Binary.id(), EliasFano, CompareExecuteAdaptor(EliasFano)); + kernels.register_execute_parent_kernel( + Binary.id(), + EliasFano, + CompareExecuteAdaptor(EliasFano), + ); kernels.register_execute_parent_kernel(Filter.id(), EliasFano, FilterExecuteAdaptor(EliasFano)); kernels.register_execute_parent_kernel(Dict.id(), EliasFano, TakeExecuteAdaptor(EliasFano)); } diff --git a/encodings/elias-fano/src/lib.rs b/encodings/elias-fano/src/lib.rs index dd62adbec83..fcb92cad6b4 100644 --- a/encodings/elias-fano/src/lib.rs +++ b/encodings/elias-fano/src/lib.rs @@ -38,6 +38,7 @@ pub use array::EliasFanoMetadata; pub use array::EliasFanoSlots; pub use compress::elias_fano_encode; pub use cursor::EliasFanoCursor; +pub use params::encoded_bit_size; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; @@ -64,8 +65,5 @@ pub fn initialize(session: &VortexSession) { ); } -// TODO(reza): add an integer scheme in `vortex-btrblocks`, so the compressor can choose this -// encoding itself rather than it having to be applied explicitly. - #[cfg(test)] mod tests; diff --git a/encodings/elias-fano/src/params.rs b/encodings/elias-fano/src/params.rs index b268389cf7a..6a6f5c42bca 100644 --- a/encodings/elias-fano/src/params.rs +++ b/encodings/elias-fano/src/params.rs @@ -100,6 +100,54 @@ pub(crate) fn num_samples0(span: u64, lower_width: u8) -> u64 { ((span >> lower_width).saturating_add(1)) >> LOG_SAMPLING0 } +/// The size of the low-bits slot when `lower_width == 0`. +/// +/// That slot is `ConstantArray::new(0u64, n)`, whose single buffer is the scalar encoded as +/// protobuf: one tag byte and one varint. Pinned by +/// `tests::test_encoded_bit_size_matches_encoder`, which fails if that encoding ever changes. +const CONSTANT_LOWER_BYTES: u64 = 2; + +/// The exact size in bits of the buffers an Elias-Fano encoding of `n` values spanning `span` +/// occupies, without building one. +/// +/// This is what the compressor prices a candidate column at. It is exact rather than asymptotic: +/// the upper array rounded up to whole bytes, both sample tables, and the FastLanes-padded +/// low-bits child. It equals `8 * array.nbytes()` for the array [`crate::elias_fano_encode`] +/// produces from the same `(span, n)`, which `tests::test_encoded_bit_size_matches_encoder` pins. +/// +/// The literature's `n * (log2(u / n) + 2)` is the asymptotic form of the same quantity. This +/// includes the constant overheads a real column of a few thousand values actually pays, which is +/// what a cost model has to compare against bit-packing. +pub fn encoded_bit_size(span: u64, n: usize) -> VortexResult { + if n == 0 { + // The degenerate array `compress::empty` builds: a two-bit upper array, and nothing else. + return Ok(8); + } + + let lower_width = lower_width(span, n); + let upper_len = upper_len(span, n, lower_width)?; + + // The upper array is byte-padded on disk. + let upper_bytes = upper_len.div_ceil(8); + + // One sample per `1 << LOG_SAMPLING0` unset bits, and one per `1 << LOG_SAMPLING1` set bits + // above the first: `UpperBuilder::push` never samples rank 0, so the one-samples are counted + // over ranks `1..n`. + let samples = num_samples0(span, lower_width) + ((n as u64 - 1) >> LOG_SAMPLING1); + let sample_bytes = samples * size_of::() as u64; + + // A zero-width low part is a `ConstantArray` of `0u64`, whose only buffer is the scalar as + // protobuf. Otherwise FastLanes packs whole 1024-element blocks, padding the tail; see + // `BitPackedArray`'s buffer length. + let lower_bytes = if lower_width == 0 { + CONSTANT_LOWER_BYTES + } else { + (n.div_ceil(1024) as u64) * 128 * u64::from(lower_width) + }; + + Ok((upper_bytes + sample_bytes + lower_bytes) * 8) +} + #[inline] pub(crate) fn lower_mask(lower_width: u8) -> u64 { if lower_width == 0 { diff --git a/encodings/elias-fano/src/tests.rs b/encodings/elias-fano/src/tests.rs index 3178b466d8e..bfad1592123 100644 --- a/encodings/elias-fano/src/tests.rs +++ b/encodings/elias-fano/src/tests.rs @@ -121,6 +121,17 @@ fn check_access(array: &EliasFanoArray, expected: &[Scalar], seed: u64) -> Vorte for (index, want) in expected.iter().enumerate() { assert_eq!(&cursor.access(index)?, want, "sequential index {index}"); } + // And through `scalar_at`, which builds no cursor at all and so shares only the layout with the + // two loops above. Every shape the suite covers therefore checks the two readers against each + // other, including the low-bits children a rewrite can leave behind. Order is irrelevant here: + // the path is stateless. + for (index, want) in expected.iter().enumerate() { + assert_eq!( + &array.execute_scalar(index, &mut ctx)?, + want, + "scalar_at index {index}" + ); + } Ok(()) } @@ -273,6 +284,47 @@ fn test_sample_tables_are_exercised() -> VortexResult<()> { Ok(()) } +/// [`params::encoded_bit_size`] is the compressor's cost model, so it has to agree with the +/// encoder exactly rather than approximately. The shapes below are the layout-relevant ones: no low +/// bits, a partial FastLanes block, an exact block, and both sample tables populated. +#[rstest] +#[case::one(1, 0)] +#[case::dense(1000, 999)] +#[case::all_equal(500, 0)] +#[case::sparse(1000, 1 << 20)] +#[case::block_low(1023, 1 << 20)] +#[case::block_exact(1024, 1 << 20)] +#[case::block_high(1025, 1 << 20)] +#[case::one_sample_exact(256, 1 << 16)] +#[case::one_sample_over(257, 1 << 16)] +#[case::sampled(5000, 1 << 30)] +#[case::sampled_dense(5000, 6000)] +fn test_encoded_bit_size_matches_encoder(#[case] n: usize, #[case] span: u64) -> VortexResult<()> { + let values = sorted_values(n, span, 0x5126_0001 ^ n as u64); + let encoded = encode(&values)?; + + // The generator draws within `span`, so the sequence's own span is what the encoder saw. + let actual_span = values[n - 1] - values[0]; + let estimated = params::encoded_bit_size(actual_span, n)?; + + assert_eq!( + estimated, + encoded.into_array().nbytes() * 8, + "n {n}, span {actual_span}" + ); + Ok(()) +} + +#[test] +fn test_encoded_bit_size_of_empty() -> VortexResult<()> { + let encoded = encode::(&[])?; + assert_eq!( + params::encoded_bit_size(0, 0)?, + encoded.into_array().nbytes() * 8 + ); + Ok(()) +} + /// The bulk decode and the per-element cursor must agree — they share no code beyond the layout. #[rstest] #[case(1, 0)] @@ -424,7 +476,10 @@ fn test_signed_next_geq() -> VortexResult<()> { check!(PType::I8, [i8::MIN, -100, -7, -7, 0, 1, 100, i8::MAX]); check!(PType::I16, [i16::MIN, -3000, -7, -7, 0, 9, i16::MAX]); check!(PType::I32, [i32::MIN, -70_000, -7, -7, 0, 5000, i32::MAX]); - check!(PType::I64, [i64::MIN, -1 << 40, -7, -7, 0, 1 << 40, i64::MAX]); + check!( + PType::I64, + [i64::MIN, -1 << 40, -7, -7, 0, 1 << 40, i64::MAX] + ); // A reference above zero, so the element domain does not wrap and the negative probes below it // all classify as `Bound::Below`. check!(PType::I32, [5i32, 5, 900, 1_000_000, i32::MAX]); @@ -908,7 +963,10 @@ fn test_is_sorted_stat(#[case] n: usize, #[case] span: u64) -> VortexResult<()> ); // Strictness is declined rather than answered, so it must come back from the generic path // with the same answer the decoded array gives. - let decoded = array.clone().execute::(&mut ctx)?.into_array(); + let decoded = array + .clone() + .execute::(&mut ctx)? + .into_array(); assert_eq!( array.statistics().compute_is_strict_sorted(&mut ctx), decoded.statistics().compute_is_strict_sorted(&mut ctx), @@ -945,7 +1003,10 @@ fn test_min_max_kernel(#[case] n: usize, #[case] span: u64) -> VortexResult<()> for (array, expected) in arrays { // The oracle is the decoded array, so the kernel is checked against the generic path and // not only against the input it was built from. - let decoded = array.clone().execute::(&mut ctx)?.into_array(); + let decoded = array + .clone() + .execute::(&mut ctx)? + .into_array(); let len = array.len(); for (name, got, oracle, want) in [ ( @@ -962,7 +1023,10 @@ fn test_min_max_kernel(#[case] n: usize, #[case] span: u64) -> VortexResult<()> ), ] { assert_eq!(got, want, "{name} over {len} elements"); - assert_eq!(got, oracle, "{name} disagrees with the generic path over {len}"); + assert_eq!( + got, oracle, + "{name} disagrees with the generic path over {len}" + ); } } Ok(()) @@ -979,7 +1043,10 @@ fn test_min_max_kernel_signed() -> VortexResult<()> { vec![i32::MIN, i32::MAX], ] { let encoded = encode(&values)?.into_array(); - let decoded = encoded.clone().execute::(&mut ctx)?.into_array(); + let decoded = encoded + .clone() + .execute::(&mut ctx)? + .into_array(); assert_eq!( encoded.statistics().compute_min::(&mut ctx), Some(values[0]), @@ -1067,11 +1134,7 @@ fn test_sparse_take_and_filter(#[case] n: usize, #[case] span: u64) -> VortexRes ); let mask = Mask::from_indices(array.len(), picks.iter().copied()); - assert_arrays_eq!( - array.filter(mask.clone())?, - decoded.filter(mask)?, - &mut ctx - ); + assert_arrays_eq!(array.filter(mask.clone())?, decoded.filter(mask)?, &mut ctx); } Ok(()) } @@ -1160,7 +1223,10 @@ fn test_compare_matches_decoded(#[case] n: usize, #[case] span: u64) -> VortexRe let mut ctx = SESSION.create_execution_ctx(); for array in &arrays { // The oracle: the same array, decoded, compared by the generic path. - let decoded = array.clone().execute::(&mut ctx)?.into_array(); + let decoded = array + .clone() + .execute::(&mut ctx)? + .into_array(); for &probe in &probes { // A nullable literal must not change which rows match, but it does make the result // nullable — which a comparison of set bits alone would not notice. @@ -1168,7 +1234,10 @@ fn test_compare_matches_decoded(#[case] n: usize, #[case] span: u64) -> VortexRe for op in COMPARISONS { let expr = op(root(), lit(literal.clone())); let pushed = array.clone().apply(&expr)?.execute::(&mut ctx)?; - let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let expected = decoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; assert_eq!( pushed.dtype(), expected.dtype(), @@ -1213,8 +1282,14 @@ fn test_compare_at_the_top_of_the_universe() -> VortexResult<()> { for probe in [u64::MAX, u64::MAX - 1, 0] { for op in COMPARISONS { let expr = op(root(), lit(Scalar::from(probe))); - let pushed = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; - let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let pushed = encoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; + let expected = decoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; assert_eq!(pushed.dtype(), expected.dtype(), "dtype for probe {probe}"); assert_arrays_eq!(pushed, expected, &mut ctx); } @@ -1236,12 +1311,28 @@ fn test_compare_signed_narrow_column() -> VortexResult<()> { .execute::(&mut ctx)? .into_array(); - for probe in [i32::MIN, i32::MIN + 1, -8, -7, -6, 0, 89, i32::MAX - 1, i32::MAX] { + for probe in [ + i32::MIN, + i32::MIN + 1, + -8, + -7, + -6, + 0, + 89, + i32::MAX - 1, + i32::MAX, + ] { for literal in nullabilities(Scalar::from(probe))? { for op in COMPARISONS { let expr = op(root(), lit(literal.clone())); - let pushed = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; - let expected = decoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let pushed = encoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; + let expected = decoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; assert_eq!(pushed.dtype(), expected.dtype(), "dtype for {literal}"); assert_arrays_eq!(pushed, expected, &mut ctx); } @@ -1267,7 +1358,10 @@ fn test_compare_is_pushed_down() -> VortexResult<()> { (gt_eq(root(), minimum.clone()), true), (lt(root(), minimum), false), ] { - let result = encoded.clone().apply(&expr)?.execute::(&mut ctx)?; + let result = encoded + .clone() + .apply(&expr)? + .execute::(&mut ctx)?; let constant = result .as_constant() .vortex_expect("an all-or-nothing comparison must reduce to a constant"); @@ -1421,9 +1515,9 @@ fn test_corrupt_upper_array_raises() -> VortexResult<()> { )?; let rebuilt = EliasFano::try_new(data, encoded.lower().clone(), encoded.len())?; - // Both entry points must return an error. Each recovers a high part by subtracting a rank from - // a bit position, which this input drives negative, so an unchecked subtraction would panic in - // debug and hand back wrong values in release. + // All three entry points must return an error. Each recovers a high part by subtracting a rank + // from a bit position, which this input drives negative, so an unchecked subtraction would + // panic in debug and hand back wrong values in release. assert!( rebuilt .clone() @@ -1437,7 +1531,9 @@ fn test_corrupt_upper_array_raises() -> VortexResult<()> { cursor.access(0).is_err(), "cursor access into a malformed upper array must raise" ); + assert!( + rebuilt.execute_scalar(0, &mut ctx).is_err(), + "scalar_at into a malformed upper array must raise" + ); Ok(()) } - - diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 4e22f042adf..e09b7f4b1df 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -26,6 +26,7 @@ vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-elias-fano = { workspace = true, optional = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } @@ -53,7 +54,11 @@ vortex-session = { workspace = true } [features] # This feature enabled unstable encodings for which we don't guarantee stability. -unstable_encodings = ["dep:vortex-onpair", "vortex-zstd?/unstable_encodings"] +unstable_encodings = [ + "dep:vortex-elias-fano", + "dep:vortex-onpair", + "vortex-zstd?/unstable_encodings", +] pco = ["dep:pco", "dep:vortex-pco"] zstd = ["dep:vortex-zstd"] diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 53cfc1d2be4..b4956296e20 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -39,6 +39,11 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ // Prefer all other schemes above delta, for now (since its slower to decompress). #[cfg(feature = "unstable_encodings")] &integer::DeltaScheme::new(1.25), + // NOTE: Elias-Fano goes last in the integer block on purpose. A point lookup costs a sampled + // select, so on a tie we would rather keep the cheaper random access of the schemes above; the + // order of this list is the tie-break order. + #[cfg(feature = "unstable_encodings")] + &integer::EliasFanoScheme::new(1.2), //////////////////////////////////////////////////////////////////////////////////////////////// // Float schemes. //////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/vortex-btrblocks/src/schemes/integer/elias_fano.rs b/vortex-btrblocks/src/schemes/integer/elias_fano.rs new file mode 100644 index 00000000000..33712fcaf31 --- /dev/null +++ b/vortex-btrblocks/src/schemes/integer/elias_fano.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Elias-Fano integer encoding for monotonically non-decreasing sequences. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; +use vortex_array::arrays::Constant; +use vortex_compressor::builtins::BinaryDictScheme; +use vortex_compressor::builtins::FloatDictScheme; +use vortex_compressor::builtins::IntDictScheme; +use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateScore; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_elias_fano::EliasFano; +use vortex_elias_fano::elias_fano_encode; +use vortex_elias_fano::encoded_bit_size; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_fastlanes::BitPacked; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::SchemeExt; +use crate::schemes::string::FSSTScheme; + +/// Elias-Fano encoding for monotonically non-decreasing integers. +/// +/// Stores each value in about `log2(u / n) + 2` bits for `n` values over a universe of `u`, against +/// bit-packing's `ceil(log2(u))`. The saving is therefore roughly `log2(n)` bits per value and it +/// widens with row count, which makes the encoding most valuable on exactly the columns that are +/// largest: sorted keys, timestamps, and a list column's offsets. +/// +/// The minimum penalized compression ratio required before Elias-Fano is selected is configurable +/// via [`EliasFanoScheme::new`]; [`EliasFanoScheme::default`] uses a ratio of `1.2`. +#[derive(Debug, Copy, Clone, PartialEq)] +pub struct EliasFanoScheme { + min_ratio: f64, +} + +impl EliasFanoScheme { + /// Creates an Elias-Fano scheme requiring `min_ratio` after the penalty before it wins. + /// + /// Pass a higher ratio to make Elias-Fano more conservative, or a lower one to select it more + /// eagerly. [`EliasFanoScheme::default`] uses a ratio of `1.2`. + pub const fn new(min_ratio: f64) -> Self { + Self { min_ratio } + } +} + +impl Default for EliasFanoScheme { + fn default() -> Self { + Self::new(1.2) + } +} + +/// Multiplicative penalty applied to Elias-Fano's estimated compression ratio. +/// +/// A point lookup costs a sampled `select1` — a pointer read plus a bounded popcount scan — where +/// bit-packing costs one unpack. Elias-Fano keeps random access, unlike Delta, so the tax is no +/// heavier than Delta's; but it is not free either, so we require a real size win rather than +/// picking Elias-Fano for a single-bit gain. +const ELIAS_FANO_PENALTY: f64 = 0.95; + +/// Minimum length before Elias-Fano is worth considering. +/// +/// Below one FastLanes block the padding of the low-bits child, the two sample tables and the guard +/// bits dominate, and the asymptotic saving has not arrived yet. +const MIN_ELIAS_FANO_LEN: usize = 1024; + +impl Scheme for EliasFanoScheme { + fn scheme_name(&self) -> &'static str { + "vortex.int.elias_fano" + } + + /// Elias-Fano has nowhere to put a null: a null has no position in an ordering, and the array + /// reports `Validity::NonNullable` unconditionally. A nullable dtype can therefore never be + /// encoded, whatever its null count, so it is rejected here rather than in the estimate. + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() && !canonical.dtype().is_nullable() + } + + /// The low-bits child is built inside `elias_fano_encode` rather than handed back through + /// `compress_child`, so it never passes the edition filter on its own. Declaring both of the + /// encodings that encoder can produce keeps `retain_allowed_encodings` honest. + fn produced_encodings(&self) -> Vec { + vec![EliasFano.id(), BitPacked.id(), Constant.id()] + } + + /// Two different reasons to decline, both about the parent's access pattern rather than the + /// data: + /// + /// - **Dictionary codes** have no monotone structure, so Elias-Fano over them just adds + /// indirection. Same exclusion FoR, Sequence and Delta declare. + /// - **FSST's children** are always fully materialised, never read through a pushdown: the + /// `like` kernel calls `codes.offsets().execute::()` and canonicalisation + /// reads `uncompressed_lengths` as a slice. Measured on this branch, an Elias-Fano bulk decode + /// costs about 9x a bit-packed unpack of the same values (172 µs vs 19 µs for 65,536 u64), so + /// trading roughly 3x space for that on a child every string operation decodes is the wrong + /// way round. The cost model prices space only; a scalar penalty cannot express a 9x decode + /// difference, so it is excluded structurally instead. + fn ancestor_exclusions(&self) -> Vec { + vec![ + AncestorExclusion { + ancestor: FSSTScheme.id(), + children: ChildSelection::All, + }, + AncestorExclusion { + ancestor: IntDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: FloatDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: StringDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: BinaryDictScheme.id(), + children: ChildSelection::One(1), + }, + ] + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // A contiguous sample of a sorted column keeps the full span but a fraction of the rows, so + // `log2(span / n)` comes out too wide and the estimate is biased against Elias-Fano. Price + // the whole array instead, as SequenceScheme does for the same reason. + if compress_ctx.is_sample() { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + if data.array_len() < MIN_ELIAS_FANO_LEN { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + + let min_ratio = self.min_ratio; + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + move |_compressor, data, best_so_far, _ctx, exec_ctx| { + let primitive = data.array_as_primitive(); + let full_width = primitive.ptype().bit_width() as f64; + let n = data.array_len(); + + let stats = data.integer_stats(exec_ctx); + if stats.null_count() > 0 { + return Ok(EstimateVerdict::Skip); + } + + // On a non-decreasing sequence the minimum is the first element, which is exactly + // the reference the encoder subtracts, so this span is the one it will see. That + // holds for signed types too: both sides work in wrapping two's complement. + let span = stats.erased().max_minus_min(); + + // The cost model is exact and costs no allocation, so price the candidate before + // paying for the O(n) sortedness scan below. + let ratio = (n as f64 * full_width) / encoded_bit_size(span, n)? as f64 + * ELIAS_FANO_PENALTY; + if ratio <= min_ratio { + return Ok(EstimateVerdict::Skip); + } + let threshold = best_so_far.and_then(EstimateScore::finite_ratio); + if threshold.is_some_and(|t| ratio <= t) { + return Ok(EstimateVerdict::Skip); + } + + // Last, because it is the only part that reads every value. The result is cached as + // `Stat::IsSorted`, which `ListArray::try_new` then reuses rather than rescanning. + if !is_sorted(data.array(), exec_ctx)? { + return Ok(EstimateVerdict::Skip); + } + + Ok(EstimateVerdict::Ratio(ratio)) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + if data.integer_stats(exec_ctx).null_count() > 0 { + vortex_bail!("Elias-Fano encoding does not support nulls"); + } + elias_fano_encode(data.array_as_primitive(), exec_ctx).map(IntoArray::into_array) + } +} diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 3aae2ae5601..c3602c6aadd 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -6,6 +6,8 @@ mod bitpacking; #[cfg(feature = "unstable_encodings")] mod delta; +#[cfg(feature = "unstable_encodings")] +mod elias_fano; mod for_; mod rle; mod runend; @@ -19,6 +21,8 @@ mod pco; pub use bitpacking::BitPackingScheme; #[cfg(feature = "unstable_encodings")] pub use delta::DeltaScheme; +#[cfg(feature = "unstable_encodings")] +pub use elias_fano::EliasFanoScheme; pub use for_::FoRScheme; #[cfg(feature = "pco")] pub use pco::PcoScheme; diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index e4227a472ec..7fa0ee816b9 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -7,6 +7,8 @@ use std::iter; use std::sync::LazyLock; use rand::Rng; +#[cfg(feature = "unstable_encodings")] +use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; @@ -195,6 +197,137 @@ fn test_delta_compressed() -> VortexResult<()> { Ok(()) } +/// A sorted, non-nullable, sparse column: Elias-Fano's habitat. Sparse enough that +/// `log2(universe / n) + 2` is well below the bit-packed width, all-unique so RunEnd and Dict skip, +/// and not an arithmetic progression so Sequence skips. +#[cfg(feature = "unstable_encodings")] +fn sorted_sparse_u32(n: usize, seed: u64) -> Vec { + let mut rng = StdRng::seed_from_u64(seed); + let mut values: Vec = (0..n).map(|_| rng.random_range(0..(1u32 << 24))).collect(); + values.sort_unstable(); + values +} + +#[cfg(feature = "unstable_encodings")] +#[test] +fn test_elias_fano_compressed() -> VortexResult<()> { + use vortex_array::assert_arrays_eq; + use vortex_elias_fano::EliasFano; + + let values = sorted_sparse_u32(4096, 21u64); + let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!( + compressed.is::(), + "expected EliasFano, got tree:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!( + compressed, + array.into_array(), + &mut SESSION.create_execution_ctx() + ); + Ok(()) +} + +/// The same values in a shuffled order. Elias-Fano needs a non-decreasing sequence, so the +/// estimate's sortedness check must reject this even though the span and density are identical. +#[cfg(feature = "unstable_encodings")] +#[test] +fn test_elias_fano_skips_unsorted() -> VortexResult<()> { + use vortex_elias_fano::EliasFano; + + let mut values = sorted_sparse_u32(4096, 22u64); + let mut rng = StdRng::seed_from_u64(23u64); + for i in (1..values.len()).rev() { + values.swap(i, rng.random_range(0..=i)); + } + let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + !compressed.is::(), + "Elias-Fano must not be selected for unsorted input, got tree:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +/// A nullable dtype with no actual nulls. Elias-Fano reports `Validity::NonNullable` +/// unconditionally, so it cannot represent this column's dtype and `matches` has to decline it +/// whatever the null count. +#[cfg(feature = "unstable_encodings")] +#[test] +fn test_elias_fano_skips_nullable_dtype() -> VortexResult<()> { + use vortex_elias_fano::EliasFano; + + let values = sorted_sparse_u32(4096, 24u64); + let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::AllValid); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + !compressed.is::(), + "Elias-Fano must not be selected for a nullable dtype, got tree:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +/// A list column's offsets are the case Elias-Fano exists for: strictly increasing, non-nullable, +/// and spanning the whole elements array. `compress_list_array` resets them to start at zero and +/// narrows the ptype before the integer schemes see them, and `ListArray::try_new` then requires +/// the compressed offsets to report `Stat::IsSorted` — which Elias-Fano answers from its layout. +#[cfg(feature = "unstable_encodings")] +#[test] +fn test_elias_fano_compresses_list_offsets() -> VortexResult<()> { + use vortex_array::arrays::List; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::list::ListArraySlotsExt; + use vortex_array::assert_arrays_eq; + use vortex_elias_fano::EliasFano; + use vortex_error::vortex_err; + + const LISTS: usize = 4096; + + let mut rng = StdRng::seed_from_u64(25u64); + let mut offsets: Vec = Vec::with_capacity(LISTS + 1); + let mut offset = 0u32; + offsets.push(offset); + for _ in 0..LISTS { + offset += rng.random_range(1..20); + offsets.push(offset); + } + let elements = PrimitiveArray::new((0..offset).collect::>(), Validity::NonNullable); + let offsets = PrimitiveArray::new(Buffer::copy_from(&offsets), Validity::NonNullable); + let list = ListArray::try_new( + elements.into_array(), + offsets.into_array(), + Validity::NonNullable, + )? + .into_array(); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress(&list, &mut SESSION.create_execution_ctx())?; + + let list_view = compressed + .as_opt::() + .ok_or_else(|| vortex_err!("expected a List root"))?; + assert!( + list_view.offsets().is::(), + "expected Elias-Fano offsets, got tree:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!(compressed, list, &mut SESSION.create_execution_ctx()); + Ok(()) +} + /// Returns true if any `Delta` array appears below an ancestor `Delta` in the tree. #[cfg(feature = "unstable_encodings")] fn has_nested_delta(array: &vortex_array::ArrayRef, under_delta: bool) -> bool { diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index d779757bc77..ecf66a762c2 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -91,6 +91,8 @@ fn trace_session() -> VortexSession { vortex_alp::initialize(&session); vortex_datetime_parts::initialize(&session); vortex_decimal_byte_parts::initialize(&session); + #[cfg(feature = "unstable_encodings")] + vortex_elias_fano::initialize(&session); vortex_fastlanes::initialize(&session); vortex_runend::initialize(&session); vortex_sequence::initialize(&session); @@ -361,6 +363,11 @@ fn trace_scan_like_on_compressed_comment() -> VortexResult<()> { // No reduce rule rewrites a like over FSST; the FSST like kernel compiles the pattern and // matches in compressed space at execution time. insta::assert_snapshot!(optimized.trace.to_string(), @""); + + // The FSST like kernel materialises its offsets child + // (`codes.offsets().execute::`), so no extra `execute_until` appears here only + // because `EliasFanoScheme` declines under an FSST ancestor for exactly that reason. If that + // exclusion is ever dropped, this trace grows an Elias-Fano decode inside the kernel. insta::assert_snapshot!(executed.trace.to_string(), @" execute_until target=AnyCanonical root=vortex.like(bool, len=4096) iter 0 current=vortex.like(bool, len=4096) builder_active=false diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap index c9554add05a..c73f0ac082f 100644 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: list(i32), len=4066, nbytes=81804 -root: vortex.list(list(i32), len=4066) nbytes=11146 +root: vortex.list(list(i32), len=4066) nbytes=6197 metadata: elements: vortex.runend(i32, len=16384) nbytes=3968 metadata: offset: 0 @@ -15,11 +15,7 @@ root: vortex.list(list(i32), len=4066) nbytes=11146 metadata: reference: -49931i32 encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 metadata: bit_width: 17, offset: 0 - offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 - metadata: bit_width: 14, offset: 0 - patch_indices: vortex.primitive(u16, len=1) nbytes=2 - metadata: ptype: u16 - patch_values: vortex.constant(u16, len=1) nbytes=4 - metadata: scalar: 16384u16 - patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 - metadata: ptype: u8 + offsets: vortex.elias_fano(u16, len=4067) nbytes=2229 + metadata: reference: 0u16, max: 16384u16, lower_width: 2, upper_len: 8165, first_rank: 0 + lower: fastlanes.bitpacked(u64, len=4067) nbytes=1024 + metadata: bit_width: 2, offset: 0