From 9d3bd7d461d3596a589c8d37da40b6fab7813e8b Mon Sep 17 00:00:00 2001 From: rapour Date: Fri, 18 Sep 2026 13:51:46 +0330 Subject: [PATCH] feat: elias-fano encoding module ef and basic vortex vtable traits Signed-off-by: rapour --- Cargo.lock | 18 + Cargo.toml | 2 + encodings/elias-fano/Cargo.toml | 38 + encodings/elias-fano/benches/elias_fano.rs | 154 ++++ .../goldenfiles/elias_fano.metadata | 2 + encodings/elias-fano/src/access.rs | 94 +++ encodings/elias-fano/src/array.rs | 646 ++++++++++++++++ encodings/elias-fano/src/compress.rs | 274 +++++++ encodings/elias-fano/src/compute/is_sorted.rs | 42 ++ encodings/elias-fano/src/compute/mod.rs | 5 + encodings/elias-fano/src/compute/slice.rs | 29 + encodings/elias-fano/src/ef/decode.rs | 112 +++ encodings/elias-fano/src/ef/encode.rs | 80 ++ encodings/elias-fano/src/ef/mod.rs | 285 +++++++ encodings/elias-fano/src/ef/params.rs | 106 +++ encodings/elias-fano/src/ef/read.rs | 134 ++++ encodings/elias-fano/src/ef/select.rs | 218 ++++++ encodings/elias-fano/src/ef/tests.rs | 587 +++++++++++++++ encodings/elias-fano/src/ef/upper.rs | 187 +++++ encodings/elias-fano/src/ef/validate.rs | 105 +++ encodings/elias-fano/src/lib.rs | 77 ++ encodings/elias-fano/src/lower.rs | 44 ++ encodings/elias-fano/src/rules.rs | 12 + encodings/elias-fano/src/tests.rs | 712 ++++++++++++++++++ 24 files changed, 3963 insertions(+) create mode 100644 encodings/elias-fano/Cargo.toml create mode 100644 encodings/elias-fano/benches/elias_fano.rs create mode 100644 encodings/elias-fano/goldenfiles/elias_fano.metadata create mode 100644 encodings/elias-fano/src/access.rs 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/is_sorted.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/ef/decode.rs create mode 100644 encodings/elias-fano/src/ef/encode.rs create mode 100644 encodings/elias-fano/src/ef/mod.rs create mode 100644 encodings/elias-fano/src/ef/params.rs create mode 100644 encodings/elias-fano/src/ef/read.rs create mode 100644 encodings/elias-fano/src/ef/select.rs create mode 100644 encodings/elias-fano/src/ef/tests.rs create mode 100644 encodings/elias-fano/src/ef/upper.rs create mode 100644 encodings/elias-fano/src/ef/validate.rs create mode 100644 encodings/elias-fano/src/lib.rs create mode 100644 encodings/elias-fano/src/lower.rs create mode 100644 encodings/elias-fano/src/rules.rs create mode 100644 encodings/elias-fano/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 275db7ef8ee..1aaaf1104fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11040,6 +11040,24 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-elias-fano" +version = "0.1.0" +dependencies = [ + "codspeed-divan-compat", + "lending-iterator", + "num-traits", + "prost 0.14.4", + "rstest", + "smallvec", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-mask", + "vortex-session", +] + [[package]] name = "vortex-error" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 6a94cd1eb56..098dc69b7f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,6 +61,7 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", + "encodings/elias-fano", # Benchmarks "benchmarks/bench-support", "benchmarks/lance-bench", @@ -312,6 +313,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..bc36247e8aa --- /dev/null +++ b/encodings/elias-fano/Cargo.toml @@ -0,0 +1,38 @@ +[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] +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-session = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +rstest = { workspace = true } +vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[[bench]] +name = "elias_fano" +harness = false + +[lints] +workspace = true diff --git a/encodings/elias-fano/benches/elias_fano.rs b/encodings/elias-fano/benches/elias_fano.rs new file mode 100644 index 00000000000..13b4be8af48 --- /dev/null +++ b/encodings/elias-fano/benches/elias_fano.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +// +//! Microbenchmarks for the Elias-Fano array's read and write paths. +//! +//! `scalar_at` is the per-element path and reads the low-bits child once per probe, so it is what a +//! change to that path has to be judged against. `encode` and `decode_bulk` cover the two batch +//! paths. +//! +//! Three shapes, because the layout behaves differently in each: +//! +//! * `Sparse` — a wide universe, so `lower_width` is large and nearly every read touches the child. +//! * `Dense` — a universe no wider than the row count, so `lower_width` is zero and the low-bits +//! child is never read at all. The difference against `Sparse` is the child's whole cost. +//! * `Duplicates` — few distinct values over a wide universe, so each occupied high-part bucket is +//! deep. Random data almost never produces this. + +#![allow( + clippy::cast_possible_truncation, + clippy::expect_used, + clippy::tests_outside_test_module, + clippy::unwrap_used +)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_elias_fano::EliasFanoArray; +use vortex_elias_fano::elias_fano_encode; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + vortex_elias_fano::initialize(&session); + session +}); + +/// Deterministic xorshift, so a run 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 + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Shape { + Sparse, + Dense, + Duplicates, +} + +/// The universe every shape draws from, wide enough that `Sparse` gets ten low bits at 2^20 rows. +const UNIVERSE: u64 = 1 << 30; + +/// Distinct values in the `Duplicates` shape: at 2^20 rows that is ~64 rows per value, so an +/// occupied bucket is deep enough for `search_bucket` to bisect rather than walk. +const LEVELS: u64 = 1 << 14; + +fn values(n: usize, shape: Shape) -> Vec { + let mut rng = Rng(0x5EED_1234_ABCD_0001); + let mut out: Vec = match shape { + Shape::Sparse => (0..n).map(|_| rng.below(UNIVERSE)).collect(), + // Universe == row count, which drives `lower_width` to zero. + Shape::Dense => (0..n).map(|_| rng.below(n as u64)).collect(), + Shape::Duplicates => { + let step = UNIVERSE / LEVELS; + (0..n).map(|_| rng.below(LEVELS) * step).collect() + } + }; + out.sort_unstable(); + out +} + +fn encoded(n: usize, shape: Shape) -> EliasFanoArray { + let array = PrimitiveArray::from_iter(values(n, shape)); + let mut ctx = SESSION.create_execution_ctx(); + elias_fano_encode(array.as_ref().as_::(), &mut ctx).expect("encode") +} + +/// Probe count held fixed across shapes and row counts, so the reported figure is comparable. +const PROBES: usize = 4096; + +const CASES: &[(Shape, usize)] = &[ + (Shape::Sparse, 1 << 16), + (Shape::Sparse, 1 << 20), + (Shape::Dense, 1 << 20), + (Shape::Duplicates, 1 << 20), +]; + +fn random_indices(n: usize) -> Vec { + let mut rng = Rng(0xA11C_E000_0000_0001); + (0..PROBES).map(|_| rng.below(n as u64) as usize).collect() +} + +/// Point lookups through `OperationsVTable::scalar_at`: one sampled `select1` and one low-bits read +/// apiece, in a random order so nothing about locality is being measured by accident. +#[divan::bench(args = CASES)] +fn scalar_at(bencher: Bencher, case: (Shape, usize)) { + let (shape, n) = case; + let array: ArrayRef = encoded(n, shape).into_array(); + let indices = random_indices(n); + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| { + for &index in &indices { + divan::black_box(array.execute_scalar(index, &mut ctx).unwrap()); + } + }); +} + +/// Whole-array decode, which walks the upper array once and reads the low bits a FastLanes block at +/// a time rather than one element at a time. +#[divan::bench(args = CASES)] +fn decode_bulk(bencher: Bencher, case: (Shape, usize)) { + let (shape, n) = case; + let array = encoded(n, shape); + bencher + .with_inputs(|| (array.clone().into_array(), SESSION.create_execution_ctx())) + .bench_local_values(|(array, mut ctx)| { + divan::black_box(array.execute::(&mut ctx).unwrap()); + }); +} + +#[divan::bench(args = CASES)] +fn encode(bencher: Bencher, case: (Shape, usize)) { + let (shape, n) = case; + let array = PrimitiveArray::from_iter(values(n, shape)); + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| { + divan::black_box( + elias_fano_encode(array.as_ref().as_::(), &mut ctx).unwrap(), + ); + }); +} + +fn main() { + divan::main(); +} 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/access.rs b/encodings/elias-fano/src/access.rs new file mode 100644 index 00000000000..67845d0d172 --- /dev/null +++ b/encodings/elias-fano/src/access.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Binding the codec's random access to Vortex: describing an array's buffers as a layout, +//! supplying the low bits out of the child array, and turning elements back into scalars. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::scalar::Scalar; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::EliasFano; +use crate::array::EliasFanoSlotsView; +use crate::array::scalar_from_bits; +use crate::ef; +use crate::ef::LowBits; +use crate::malformed; + +/// The low-bits child, as a source the codec can pull from. +/// +/// Built per call rather than held: it borrows the execution context. +pub(crate) struct LowerSource<'a, 'c> { + pub(crate) lower: &'a ArrayRef, + pub(crate) ctx: &'c mut ExecutionCtx, +} + +impl LowBits for LowerSource<'_, '_> { + type Error = VortexError; + + /// One low part, through the child's own `scalar_at`, so a rewritten slot needs no case here. + fn get(&mut self, rank: u64) -> VortexResult { + self.lower + .execute_scalar(usize::try_from(rank)?, self.ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| vortex_err!("Elias-Fano low-bits child holds no value at rank {rank}")) + } +} + +/// Carry a read's failure into Vortex's error type. A function rather than a `From` impl, for the +/// reason [`crate::malformed`] gives. +pub(crate) fn read_error(error: ef::ReadError) -> VortexError { + match error { + ef::ReadError::Malformed(error) => malformed(error), + ef::ReadError::LowBits(error) => error, + } +} + +/// The layout an array describes, borrowed from its buffers. +/// +/// The upper array is taken as raw bytes rather than a `BitBuffer`, which would strip alignment and +/// can reallocate to carry the three numbers the codec wants. +pub(crate) fn layout<'a>(array: ArrayView<'a, EliasFano>) -> VortexResult> { + let data = array.data(); + let (_, samples1) = data.sample_bytes()?; + let upper = ef::Bits::new( + data.upper_buffer().as_slice(), + 0, + usize::try_from(data.upper_len())?, + ); + Ok(ef::Layout::new( + upper, + samples1, + data.lower_width(), + data.first_rank(), + array.len(), + )) +} + +/// The value at logical `index`. +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 reference_bits = array.data().reference_bits(); + let layout = layout(array)?; + + // 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 mut source = LowerSource { lower, ctx }; + + let element = ef::element_at(layout, index, &mut source).map_err(read_error)?; + scalar_from_bits(array.dtype(), reference_bits.wrapping_add(element)) +} diff --git a/encodings/elias-fano/src/array.rs b/encodings/elias-fano/src/array.rs new file mode 100644 index 00000000000..edf4978f953 --- /dev/null +++ b/encodings/elias-fano/src/array.rs @@ -0,0 +1,646 @@ +// 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::proto::scalar::ScalarValue as ProtoScalarValue; +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::ByteBuffer; +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::access::access_at; +use crate::compress::elias_fano_decompress; +use crate::ef; +use crate::malformed; +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 +/// `ef::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::ef`] module: the two +/// buffers, the universe bounds, and the three 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 <= ef::MAX_LOWER_WIDTH, + "Elias-Fano lower_width {lower_width} exceeds {}", + ef::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 = ef::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 + } + + /// 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 — read + /// the element 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 [`ef::num_samples0`] + /// explains why the element count drops out of the derivation. + #[inline] + pub(crate) fn num_samples0(&self) -> u64 { + ef::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 + } + + /// 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 read rather than per element. Deserialized + /// buffers carry no alignment guarantee, so entries are read one at a time with + /// [`ef::read_sample`]. + pub(crate) fn sample_bytes(&self) -> VortexResult<(&[u8], &[u8])> { + let bytes = self.samples.as_slice(); + let num_samples0 = self.num_samples0(); + // `try_new` already proved the buffer reaches the seam; this raises rather than asserting. + 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)) + } +} + +/// 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, so a replacement buffer still + // has to satisfy the constructor's invariants. + 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<&ProtoScalarValue>, 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 { + // No probe state: like every other encoding in the tree, this stays on the `scalar_at` path + // that `probe_scalar` defaults to. + type ProbeState = (); + + /// One sampled `select1` for the high part, and one read of the low-bits child. + fn scalar_at( + array: ArrayView<'_, EliasFano>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + access_at(array, index, ctx) + } +} + +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() + ); + + // Everything below is codec geometry, re-derived from `(span, n)` alone: the widths, the upper + // length, both sample counts, and every sample's range and order. + ef::validate_layout( + data.span(), + num_elements as usize, + data.lower_width(), + data.upper_len(), + data.samples_buffer().as_slice(), + ) + .map_err(malformed) +} + +#[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..dab25e1fea1 --- /dev/null +++ b/encodings/elias-fano/src/compress.rs @@ -0,0 +1,274 @@ +// 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::ef`] for the layout both directions read and write. + +use std::mem::MaybeUninit; + +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::BitBufferMut; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; +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::ef; +use crate::lower::materialise; +use crate::lower::readable_in_place; +use crate::malformed; + +/// 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: the layout has nowhere to put one. `ctx` is taken for symmetry with the +/// other integer encoders and goes unused. +// 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 encoded = match_each_integer_ptype!(array.ptype(), |P| { + let values = array.as_slice::

(); + ensure_non_decreasing(values)?; + ef::encode( + values + .iter() + .map(|&value| (value as u64).wrapping_sub(reference_bits)), + span, + ) + }) + .map_err(crate::unrepresentable)?; + + // Every buffer adopts its `Vec`'s allocation rather than copying; see `Buffer: From>`. + let lower = pack_lower(Buffer::from(encoded.lower), encoded.lower_width, n)?; + + let data = EliasFanoData::try_new( + ByteBuffer::from(encoded.upper), + Buffer::from(encoded.samples).into_byte_buffer(), + scalar_from_bits(&dtype, reference_bits)?, + scalar_from_bits(&dtype, max_bits)?, + encoded.lower_width, + encoded.upper_len, + 0, + )?; + EliasFano::try_new(data, lower, n) +} + +/// Reject a sequence that decreases anywhere. +/// +/// Checked in the **value** domain, before the reference is subtracted: an element is a modular +/// difference, so unsorted input can still yield non-decreasing elements after wrapping, and +/// [`ef::encode`] would build a layout no reader could fault. +fn ensure_non_decreasing(values: &[P]) -> VortexResult<()> { + if values.is_sorted() { + return Ok(()); + } + let index = values + .windows(2) + .position(|pair| pair[1] < pair[0]) + .map_or(0, |first| first + 1); + vortex_bail!( + "Elias-Fano requires a non-decreasing sequence, but the value at index {index} is below \ + its predecessor" + ) +} + +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 bits = ef::Bits::new( + array.upper_buffer().as_slice(), + 0, + 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 = ef::position_of_rank(bits, samples1, first_rank).map_err(malformed)?; + let end = + ef::position_of_rank(bits, samples1, first_rank + len as u64 - 1).map_err(malformed)? + 1; + let words = ef::window_words(bits, start, end); + + Ok(match_each_integer_ptype!(ptype, |P| { + PrimitiveArray::new( + decode::

(array, &words, start, ctx)?, + Validity::NonNullable, + ) + })) +} + +/// Decode `len` elements into the column's own width, feeding [`ef::Decoder`] the low bits one +/// FastLanes block at a time. +/// +/// Decides only the shape of the low-bits child: whether its packed bytes can be unpacked a block +/// at a time, or have to be materialised first. +fn decode( + array: &EliasFanoArray, + words: &[u64], + start: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + u64: AsPrimitive

, +{ + let len = array.len(); + let first_rank = array.first_rank(); + let reference_bits = array.reference_bits(); + let lower_width = array.lower_width(); + + let mut decoder = + ef::Decoder::new(words, start, first_rank, len, lower_width).map_err(malformed)?; + + let mut values = BufferMut::

::zeroed(len); + + if lower_width == 0 { + // Nothing is stored, so do not execute the slot just to read `len` zeros. + segment(&mut decoder, &mut values, reference_bits, None); + } else { + let first = usize::try_from(first_rank)?; + let window_lower = array.lower().slice(first..first + len)?; + + if let Some(packed) = readable_in_place(&window_lower) { + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = packed.unpacked_chunks::(&mut scratch)?; + if let Some(initial) = chunks.initial() { + segment(&mut decoder, &mut values, reference_bits, Some(initial)); + } + if decoder.remaining() > 0 { + let mut full = chunks.full_chunks(); + while let Some(chunk) = full.next() { + segment(&mut decoder, &mut values, reference_bits, Some(chunk)); + } + } + if decoder.remaining() > 0 + && let Some(trailer) = chunks.trailer() + { + segment(&mut decoder, &mut values, reference_bits, Some(trailer)); + } + } else { + // The slot is patched, device-resident, or some other encoding after a rewrite. + let dense = materialise(window_lower, ctx)?; + segment( + &mut decoder, + &mut values, + reference_bits, + Some(dense.as_slice()), + ); + } + } + + decoder.finish().map_err(malformed)?; + Ok(values.freeze()) +} + +/// One run of low parts, folded into the column's own width. +fn segment( + decoder: &mut ef::Decoder<'_>, + values: &mut BufferMut

, + reference_bits: u64, + lows: Option<&[u64]>, +) where + u64: AsPrimitive

, +{ + let out = values.as_mut_slice(); + decoder.segment(lows, |index, element| { + // 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. + out[index] = reference_bits.wrapping_add(element).as_(); + }); +} + +fn pack_lower(lower: Buffer, lower_width: u8, n: usize) -> VortexResult { + if lower_width == 0 { + // Nothing to store, and a constant array 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, representable rather than rejected so an empty chunk needs no +/// special handling upstream. +/// +/// The bounds go unused, there being 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/is_sorted.rs b/encodings/elias-fano/src/compute/is_sorted.rs new file mode 100644 index 00000000000..452d4d642fd --- /dev/null +++ b/encodings/elias-fano/src/compute/is_sorted.rs @@ -0,0 +1,42 @@ +// 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 declined, because duplicates are legal and finding whether any are present +/// means a full scan of the low bits. Returning `None` leaves that to the generic path. +#[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/mod.rs b/encodings/elias-fano/src/compute/mod.rs new file mode 100644 index 00000000000..4224a9920f1 --- /dev/null +++ b/encodings/elias-fano/src/compute/mod.rs @@ -0,0 +1,5 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(crate) mod is_sorted; +mod slice; 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/ef/decode.rs b/encodings/elias-fano/src/ef/decode.rs new file mode 100644 index 00000000000..03cf5e9204a --- /dev/null +++ b/encodings/elias-fano/src/ef/decode.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bulk decode, one run at a time, so a caller can supply low bits a block at a time rather than +//! materialising the whole sequence. + +use super::Malformed; +use super::Ones; +use super::element_of; +use super::high_of; +use super::lower_mask; + +/// Walks the set bits of an upper window, pairing each with a low part. +pub struct Decoder<'a> { + ones: Ones<'a>, + /// Absolute bit position the window starts at, which its set-bit indices are relative to. + start: usize, + first_rank: u64, + len: usize, + lower_width: u8, + lower_mask: u64, + /// How many elements have been written so far. + rank: usize, +} + +impl<'a> Decoder<'a> { + /// Open a decoder over `words`, the window's bits as whole `u64`s at a zero bit offset. + /// + /// Checks what the per-element loop relies on once rather than `len` times: `len` set bits in + /// the window, and `start > first_rank` so every position exceeds its own rank. + pub fn new( + words: &'a [u64], + start: usize, + first_rank: u64, + len: usize, + lower_width: u8, + ) -> Result { + let found = words.iter().map(|word| word.count_ones() as usize).sum(); + if found != len { + return Err(Malformed::SetBitCount { + expected: len, + found, + }); + } + if start as u64 <= first_rank { + return Err(Malformed::PositionAtOrBelowRank { + rank: first_rank, + position: start as u64, + }); + } + Ok(Self { + ones: Ones::new(words), + start, + first_rank, + len, + lower_width, + lower_mask: lower_mask(lower_width), + rank: 0, + }) + } + + /// How many elements are still to be written. A caller feeding runs stops at zero. + pub fn remaining(&self) -> usize { + self.len - self.rank + } + + /// Fold one run of consecutive low parts, calling `out(index, element)` for each element. + /// + /// `lows` of `None` is the `lower_width == 0` layout, where the run covers whatever is left. + /// + /// Cannot fail: an uninvertible position leaves the count short for [`Self::finish`]. Inlined + /// because out of line, a zero-width layout has no unpacking to hide the call behind and + /// decodes at half the speed. + #[inline] + pub fn segment(&mut self, lows: Option<&[u64]>, mut out: impl FnMut(usize, u64)) { + let remaining = self.len - self.rank; + let take = lows.map_or(remaining, |lows| lows.len().min(remaining)); + + for offset in 0..take { + let index = self.rank + offset; + let Some(position) = self.ones.next() else { + self.rank = index; + return; + }; + let Some(high) = high_of( + (self.start + position) as u64, + self.first_rank + index as u64, + ) else { + self.rank = index; + return; + }; + // A supplier may hand back bits above `lower_width`, which would bleed into the high + // part. + let low = lows.map_or(0, |lows| lows[offset] & self.lower_mask); + out(index, element_of(high, low, self.lower_width)); + } + + self.rank += take; + } + + /// Close the decode, refusing a layout that produced fewer elements than it claimed: either a + /// position that could not be inverted, or runs of low bits that ran short. + pub fn finish(self) -> Result<(), Malformed> { + if self.rank != self.len { + return Err(Malformed::SetBitCount { + expected: self.len, + found: self.rank, + }); + } + Ok(()) + } +} diff --git a/encodings/elias-fano/src/ef/encode.rs b/encodings/elias-fano/src/ef/encode.rs new file mode 100644 index 00000000000..a0d38ad1982 --- /dev/null +++ b/encodings/elias-fano/src/ef/encode.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Building the layout from a non-decreasing sequence of elements. + +use super::Error; +use super::UpperBuilder; +use super::lower_mask; +use super::lower_width; +use super::position_of; +use super::upper_len; + +/// The buffers an encoded sequence occupies. +/// +/// Every field is a plain `Vec`, which a caller with its own buffer type can adopt rather than copy. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Encoded { + /// The upper bit array, LSB-first, padded to whole bytes. + pub upper: Vec, + /// Zero-samples followed by one-samples, sharing one table. The seam is not recorded: a reader + /// recovers it from the universe with [`num_samples0`](super::num_samples0). + pub samples: Vec, + /// The low [`lower_width`](Self::lower_width) bits of every element, in element order. Empty of + /// meaning at width zero, where every entry is `0`. + pub lower: Vec, + /// Low bits per element. + pub lower_width: u8, + /// Length of the upper array in bits, which is not `upper.len() * 8`. + pub upper_len: u64, +} + +/// Encode a non-decreasing sequence of elements spanning `0..=span`. +/// +/// An *element* is a value with the sequence's reference already subtracted, so `span` is the last +/// element. +/// +/// A caller holding values must check monotonicity **in the value domain** before subtracting: an +/// element is a modular difference, so an unsorted sequence can still yield non-decreasing elements +/// after wrapping, which nothing here can detect afterwards. +/// +/// # Panics +/// +/// Panics in debug builds if `elements` is empty, or if it is not non-decreasing, or if its last +/// element is not `span`. A release build produces a layout no reader will accept. +pub fn encode(elements: impl ExactSizeIterator, span: u64) -> Result { + let n = elements.len(); + debug_assert!(n > 0, "the empty sequence has no layout to build"); + + let lower_width = lower_width(span, n); + let upper_len = upper_len(span, n, lower_width)?; + let mask = lower_mask(lower_width); + + // `upper_len` has already refused a length `usize` cannot address. + let bits = usize::try_from(upper_len).map_err(|_| Error::UpperLenTooLarge { upper_len })?; + let mut upper = UpperBuilder::new(bits); + let mut lower = Vec::with_capacity(n); + let mut previous = 0u64; + + for (index, element) in elements.enumerate() { + debug_assert!(element >= previous, "elements must be non-decreasing"); + debug_assert!(element <= span, "element {element} exceeds the span {span}"); + previous = element; + + let rank = index as u64; + upper.push(rank, position_of(element, rank, lower_width)); + lower.push(element & mask); + } + + debug_assert_eq!(lower.len(), n, "ExactSizeIterator yielded the wrong count"); + debug_assert_eq!(previous, span, "the last element is the span"); + + let (upper, samples) = upper.finish(n as u64, upper_len); + Ok(Encoded { + upper, + samples, + lower, + lower_width, + upper_len, + }) +} diff --git a/encodings/elias-fano/src/ef/mod.rs b/encodings/elias-fano/src/ef/mod.rs new file mode 100644 index 00000000000..d210a0a4d9a --- /dev/null +++ b/encodings/elias-fano/src/ef/mod.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A self-contained Elias-Fano codec, depending on nothing beyond `std`. +//! +//! [`encode()`] builds the layout from a non-decreasing sequence of elements, [`element_at`] reads +//! one back, [`Decoder`] turns a whole window back into elements, and [`validate_layout`] checks +//! that stored geometry describes the universe it claims. +//! +//! Each element splits into an `l`-bit low part and a high part `element >> l`, set as one bit at +//! position `high + rank + 1` of a bit array. The `+ rank` keeps positions distinct when elements +//! share a high part, so reading element `i` is a `select1` and the inverse is +//! `high = position - rank - 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. [`LOG_SAMPLING1`] and [`LOG_SAMPLING0`] bound the scans over +//! that array. +//! +//! The low parts live outside these buffers, supplied through [`LowBits`]: they compress better +//! under a dedicated encoding than anything here would manage. +//! +//! Samples are written native-endian and read back with `from_le_bytes`, so a root building this +//! for a big-endian target wants `#![cfg(target_endian = "little")]`. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::fmt::Result as FmtResult; + +mod decode; +mod encode; +mod params; +mod read; +mod select; +mod upper; +mod validate; + +#[cfg(test)] +mod tests; + +pub use decode::Decoder; +pub use encode::Encoded; +pub use encode::encode; +pub use params::LOG_SAMPLING0; +pub use params::LOG_SAMPLING1; +pub use params::MAX_LOWER_WIDTH; +pub use params::lower_mask; +pub use params::lower_width; +pub use params::num_samples0; +pub use params::num_samples1; +pub use params::num_zeros; +pub use params::upper_len; +pub use read::Layout; +pub use read::LowBits; +pub use read::element_at; +pub use read::position_of_rank; +pub use select::Bits; +pub use select::select_range; +pub use select::select_zero_range; +pub use select::window_words; +pub use upper::Ones; +pub use upper::UpperBuilder; +pub use upper::element_of; +pub use upper::high_of; +pub use upper::position_of; +pub use upper::read_sample; +pub use upper::sampled_select; +pub use validate::validate_layout; + +/// A sequence whose layout cannot be represented, i.e. bad arguments to the encoder. +/// +/// Both variants need a universe of very nearly `2^64`, so neither is reachable from a sequence +/// held in memory. The same geometry is re-derived from untrusted metadata, where they are. +#[non_exhaustive] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Error { + /// The upper array's length overflows a `u64`. + UpperLenOverflow { + /// Number of elements. + n: usize, + /// The universe's width, `max - reference`. + span: u64, + /// Low bits per element. + lower_width: u8, + }, + /// The upper array is longer than `usize` can address. + UpperLenTooLarge { + /// The length in bits that does not fit. + upper_len: u64, + }, +} + +impl Display for Error { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + match self { + Self::UpperLenOverflow { + n, + span, + lower_width, + } => write!( + f, + "upper array overflows: n {n}, span {span}, lower_width {lower_width}" + ), + Self::UpperLenTooLarge { upper_len } => { + write!(f, "upper array of {upper_len} bits does not fit in memory") + } + } + } +} + +impl core::error::Error for Error {} + +/// Which of the two sample tables a fault was found in. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Table { + /// The table sampling unset bits, one per `1 << LOG_SAMPLING0`. + Zero, + /// The table sampling set bits, one per `1 << LOG_SAMPLING1`. + One, +} + +impl Display for Table { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.write_str(match self { + Self::Zero => "zero", + Self::One => "one", + }) + } +} + +/// Buffers that do not describe a valid Elias-Fano sequence. +/// +/// [`Error`] is about arguments the encoder cannot serve; this is about a layout handed to a +/// reader. Nothing here builds one, but the upper array's contents are never checked against the +/// elements, so a corrupt file can. +#[non_exhaustive] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Malformed { + /// The stored geometry describes a layout that cannot exist at all. + Unrepresentable(Error), + /// The upper array runs out before the element of this rank. + NoElementOfRank { + /// The rank that has no set bit. + rank: u64, + }, + /// A set bit sits at or below its own rank, so [`high_of`] cannot invert it. + PositionAtOrBelowRank { + /// The element's rank. + rank: u64, + /// The bit it was found at. + position: u64, + }, + /// A window holds a different number of set bits than it has elements. + SetBitCount { + /// How many the layout calls for. + expected: usize, + /// How many were found. + found: usize, + }, + /// The stored low-bits width disagrees with the one the universe implies. + LowerWidth { + /// The width the universe implies. + expected: u8, + /// The width stored. + found: u8, + }, + /// The stored upper length disagrees with the one the universe implies. + UpperLen { + /// The length the universe implies. + expected: u64, + /// The length stored. + found: u64, + }, + /// The samples buffer holds a different number of entries than the layout calls for. + SampleCount { + /// How many the layout calls for. + expected: u64, + /// How many are stored. + found: u64, + }, + /// A sample points outside the range its own rank allows. + SampleOutOfRange { + /// Which table it is in. + table: Table, + /// Its index within that table. + index: usize, + /// The position it points to. + sample: u64, + /// The lowest position its rank permits. + min: u64, + /// One past the highest. + max: u64, + }, + /// Samples within one table are not strictly increasing, which a sampled search relies on. + SamplesNotIncreasing { + /// Which table. + table: Table, + /// The index at which the order breaks. + index: usize, + }, +} + +impl Display for Malformed { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + match self { + Self::Unrepresentable(error) => write!(f, "{error}"), + Self::NoElementOfRank { rank } => { + write!(f, "upper array holds no element of rank {rank}") + } + Self::PositionAtOrBelowRank { rank, position } => write!( + f, + "upper array is malformed: the element of rank {rank} sits at bit {position}, at \ + or below its own rank" + ), + Self::SetBitCount { expected, found } => write!( + f, + "upper array is malformed: expected exactly {expected} set bits above their own \ + ranks, found {found}" + ), + Self::LowerWidth { expected, found } => write!( + f, + "lower_width {found} does not match the {expected} its universe implies" + ), + Self::UpperLen { expected, found } => write!( + f, + "upper_len {found} does not match the {expected} its universe implies" + ), + Self::SampleCount { expected, found } => { + write!(f, "holds {found} samples, expected {expected}") + } + Self::SampleOutOfRange { + table, + index, + sample, + min, + max, + } => write!( + f, + "{table}-sample {index} points to bit {sample}, outside the {min}..{max} its rank \ + allows" + ), + Self::SamplesNotIncreasing { table, index } => { + write!( + f, + "{table}-samples are not strictly increasing at index {index}" + ) + } + } + } +} + +impl core::error::Error for Malformed {} + +impl From for Malformed { + fn from(error: Error) -> Self { + Self::Unrepresentable(error) + } +} + +/// A read that either met a malformed layout or could not obtain its low bits. +/// +/// Generic over the supplier's error, so a host's own type survives the round trip rather than +/// being flattened into a string. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReadError { + /// The layout does not describe a valid sequence. + Malformed(Malformed), + /// The low-bits source could not supply a value. + LowBits(E), +} + +impl From for ReadError { + fn from(error: Malformed) -> Self { + Self::Malformed(error) + } +} + +impl Display for ReadError { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + match self { + Self::Malformed(error) => write!(f, "{error}"), + Self::LowBits(error) => write!(f, "low bits unavailable: {error}"), + } + } +} + +impl core::error::Error for ReadError {} diff --git a/encodings/elias-fano/src/ef/params.rs b/encodings/elias-fano/src/ef/params.rs new file mode 100644 index 00000000000..1952ce523b8 --- /dev/null +++ b/encodings/elias-fano/src/ef/params.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The layout's geometry, as pure arithmetic over `(span, n)`. See the [module docs](super) for the +//! layout these size. +//! +//! An encoder stores the results and [`validate_layout`](super::validate_layout) re-derives them, +//! refusing a layout that disagrees. + +use super::Error; + +/// 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 — eight words, a window +/// short enough for [`select_zero_range`](super::select_zero_range) to walk without vectorising. +/// [`LOG_SAMPLING1`] is sized the same way. +pub const LOG_SAMPLING0: usize = 9; + +/// One one-sample is stored per `1 << LOG_SAMPLING1` set bits of the upper array. +pub const LOG_SAMPLING1: usize = 8; + +/// The widest low part we will store. +/// +/// `l == 64` would leave no high part at all. Only a single element spanning the whole `u64` range +/// reaches the clamp. +pub 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 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 fn upper_len(span: u64, n: usize, lower_width: u8) -> Result { + let buckets = span >> lower_width; + let upper_len = (n as u64) + .checked_add(buckets) + .and_then(|v| v.checked_add(2)) + .ok_or(Error::UpperLenOverflow { + n, + span, + lower_width, + })?; + if usize::try_from(upper_len).is_err() { + return Err(Error::UpperLenTooLarge { upper_len }); + } + Ok(upper_len) +} + +/// The number of unset bits in an upper array of `upper_len` bits holding `n` elements. +/// +/// No read path calls this. It states the identity [`num_samples0`] must agree with, which +/// [`validate_layout`](super::validate_layout) asserts. +#[inline] +pub 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. A reader therefore splits the shared samples buffer into +/// its two tables without the seam being stored. +#[inline] +pub 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 +} + +/// The number of one-samples the layout calls for. +/// +/// Rank 0 is never sampled — the first set bit is where a reader starts anyway — so the samples are +/// counted over ranks `1..n`. +#[inline] +pub fn num_samples1(n: usize) -> u64 { + (n as u64).saturating_sub(1) >> LOG_SAMPLING1 +} + +/// The mask keeping the low `lower_width` bits of an element. +#[inline] +pub fn lower_mask(lower_width: u8) -> u64 { + if lower_width == 0 { + 0 + } else { + u64::MAX >> (64 - u32::from(lower_width)) + } +} diff --git a/encodings/elias-fano/src/ef/read.rs b/encodings/elias-fano/src/ef/read.rs new file mode 100644 index 00000000000..e2b346ec106 --- /dev/null +++ b/encodings/elias-fano/src/ef/read.rs @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Random access over an encoded sequence, in `O(1)`. +//! +//! Reading element `i` is one sampled `select1` for its high part and one read of its low part. + +use super::Bits; +use super::LOG_SAMPLING1; +use super::Malformed; +use super::ReadError; +use super::element_of; +use super::high_of; +use super::lower_mask; +use super::sampled_select; + +/// The low bits of each element, supplied by the caller. +/// +/// Asked for one rank at a time, as a read decides which it needs, so a caller keeps its own +/// storage rather than flattening it first. A supplier need not mask what it returns: the reader +/// keeps only `lower_width` bits. +pub trait LowBits { + /// What supplying a low part can fail with. Use [`Infallible`](core::convert::Infallible) for a + /// source that cannot, such as a slice already in memory. + type Error; + + /// The low part of the element at absolute rank `rank`. + fn get(&mut self, rank: u64) -> Result; +} + +/// The borrowed parts of an encoded sequence a read needs. +/// +/// Every field is a slice the caller already holds, so describing a layout copies nothing. +/// `first_rank` and `len` narrow it to one window of a longer encoded sequence, the buffers +/// unchanged. +#[derive(Clone, Copy, Debug)] +pub struct Layout<'a> { + upper: Bits<'a>, + samples1: &'a [u8], + lower_width: u8, + first_rank: u64, + len: usize, +} + +impl<'a> Layout<'a> { + /// Describe a window of an encoded sequence. + /// + /// `upper` is the whole upper array, not the window's part of it: the sample table holds + /// absolute positions. `first_rank` is where this window starts within the encoded sequence and + /// `len` how many elements it covers. + pub fn new( + upper: Bits<'a>, + samples1: &'a [u8], + lower_width: u8, + first_rank: u64, + len: usize, + ) -> Self { + Self { + upper, + samples1, + lower_width, + first_rank, + len, + } + } + + /// How many elements this layout covers. + pub fn len(&self) -> usize { + self.len + } + + /// Whether this layout covers no elements at all. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// 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. + fn position_of_rank(&self, rank: u64) -> Result { + position_of_rank(self.upper, self.samples1, rank) + } + + /// The low bits at `rank`, masked. + /// + /// A supplier may hand back bits above `lower_width`, which would bleed into the high part. At + /// width zero nothing is stored, so the supplier is not asked at all. + fn low_at(&self, rank: u64, low: &mut L) -> Result> { + if self.lower_width == 0 { + return Ok(0); + } + low.get(rank) + .map(|bits| bits & lower_mask(self.lower_width)) + .map_err(ReadError::LowBits) + } + + /// The element seated at absolute `rank`, given the bit position of its set bit. + fn element_at_position( + &self, + position: usize, + rank: u64, + low: &mut L, + ) -> Result> { + let high = high_of(position as u64, rank).ok_or(Malformed::PositionAtOrBelowRank { + rank, + position: position as u64, + })?; + Ok(element_of(high, self.low_at(rank, low)?, self.lower_width)) + } +} + +/// The bit position of the set bit belonging to absolute rank `rank`, taking the upper array and +/// its sample table directly, for a caller that has no [`Layout`] to hand. +pub fn position_of_rank(upper: Bits<'_>, samples1: &[u8], rank: u64) -> Result { + let end = upper.len(); + sampled_select(upper, samples1, LOG_SAMPLING1, rank, end, false) + .ok_or(Malformed::NoElementOfRank { rank }) +} + +/// The element at logical `index`, in one sampled `select1`. +/// +/// # Panics +/// +/// Panics if `index` is at or beyond the layout's length. +pub fn element_at( + layout: Layout<'_>, + index: usize, + low: &mut L, +) -> Result> { + assert!(index < layout.len, "index {index} is out of bounds"); + let rank = layout.first_rank + index as u64; + let position = layout.position_of_rank(rank)?; + layout.element_at_position(position, rank, low) +} diff --git a/encodings/elias-fano/src/ef/select.rs b/encodings/elias-fano/src/ef/select.rs new file mode 100644 index 00000000000..247cbe8db6b --- /dev/null +++ b/encodings/elias-fano/src/ef/select.rs @@ -0,0 +1,218 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Windowed rank/select over the upper bit array, taking a half-open bit range so a sample can give +//! the scan its lower bound. +//! +//! Reads the backing bytes directly, so a lookup allocates nothing. [`LOG_SAMPLING0`] and +//! [`LOG_SAMPLING1`] cap a window at 512 unset or 256 set bits — eight to sixteen words — which a +//! scalar walk covers without vectorising. +//! +//! [`LOG_SAMPLING0`]: super::LOG_SAMPLING0 +//! [`LOG_SAMPLING1`]: super::LOG_SAMPLING1 + +/// A window of bits: the backing bytes, the bit offset of bit zero within them, and a length. +#[derive(Clone, Copy, Debug)] +pub struct Bits<'a> { + bytes: &'a [u8], + offset: usize, + len: usize, +} + +impl<'a> Bits<'a> { + /// A window of `len` bits beginning `offset` bits into `bytes`. + /// + /// # Panics + /// + /// Panics if `bytes` is too short to hold `offset + len` bits. + #[inline] + pub fn new(bytes: &'a [u8], offset: usize, len: usize) -> Self { + assert!( + offset.saturating_add(len) <= bytes.len().saturating_mul(8), + "{} bytes cannot hold {len} bits at offset {offset}", + bytes.len() + ); + Self { bytes, offset, len } + } + + /// The number of bits in the window. + #[inline] + pub fn len(&self) -> usize { + self.len + } + + /// Whether the window holds no bits at all. + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +/// Returns the position of the `nth` set bit within `[start, end)` of `bits`, relative to `start`, +/// or `None` if the range holds `nth` set bits or fewer. +/// +/// # Panics +/// +/// Panics if `start > end` or `end > bits.len()`. +#[inline] +pub fn select_range(bits: Bits<'_>, start: usize, end: usize, nth: usize) -> Option { + select_in_range(bits, start, end, nth, false) +} + +/// Returns the position of the `nth` unset bit within `[start, end)` of `bits`, relative to +/// `start`, or `None` if the range holds `nth` unset bits or fewer. +/// +/// The complement of [`select_range`]: a bucket boundary in the upper array is a zero, so this is +/// what turns a high part into a rank. +/// +/// # Panics +/// +/// Panics if `start > end` or `end > bits.len()`. +#[inline] +pub fn select_zero_range(bits: Bits<'_>, start: usize, end: usize, nth: usize) -> Option { + select_in_range(bits, start, end, nth, true) +} + +/// The bits of `[start, end)` as whole `u64` words at a zero bit offset, so a decode shifts +/// once here rather than per word. +/// +/// The last word is zero-padded above the window's end, so counting set bits across the result +/// counts exactly the window's. At roughly two bits per element the allocation is `n / 32` words — +/// a few KB per million rows. +/// +/// # Panics +/// +/// Panics if `start > end` or `end > bits.len()`. +pub fn window_words(bits: Bits<'_>, start: usize, end: usize) -> Vec { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= bits.len, "end {end} exceeds len {}", bits.len); + + let total = end - start; + let base = bits.offset + start; + let mut words = Vec::with_capacity(total.div_ceil(u64::BITS as usize)); + let mut pos = 0usize; + while pos < total { + let width = (total - pos).min(64); + words.push(load_bits(bits.bytes, base + pos, width)); + pos += width; + } + words +} + +/// Shared walk behind [`select_range`] (`zeros == false`) and [`select_zero_range`] +/// (`zeros == true`). +#[inline] +fn select_in_range( + bits: Bits<'_>, + start: usize, + end: usize, + nth: usize, + zeros: bool, +) -> Option { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= bits.len, "end {end} exceeds len {}", bits.len); + + // The window begins `offset + start` bits into the backing bytes. + let bytes = bits.bytes; + let base = bits.offset + start; + let total = end - start; + + let mut remaining = nth; + let mut pos = 0usize; + while pos < total { + let width = (total - pos).min(64); + let mut word = load_bits(bytes, base + pos, width); + if zeros { + // Complementing turns the padding above `width` into ones, so mask it off again. + word = mask_to(!word, width); + } + let count = word.count_ones() as usize; + if remaining < count { + return Some(pos + select_in_word(word, remaining)); + } + remaining -= count; + pos += width; + } + None +} + +/// Reads `width` bits starting at absolute bit `at`, returned in the low bits with the rest zero. +/// +/// An unaligned 64 bits straddles nine bytes: the common path takes the first eight as one +/// little-endian word and folds the ninth in across the shift. Within nine bytes of the end the +/// slower path assembles whatever is there, since the bytes it cannot read land above `width`. +#[inline] +fn load_bits(bytes: &[u8], at: usize, width: usize) -> u64 { + debug_assert!(width > 0 && width <= 64, "width {width} out of range"); + + let first = at / 8; + let shift = at % 8; + + let head = bytes.get(first..).and_then(<[u8]>::first_chunk::<8>); + let word = match (head, bytes.get(first + 8)) { + (Some(chunk), Some(&straddle)) => { + let lo = u64::from_le_bytes(*chunk); + // Shifting a `u64` by 64 is undefined, so the aligned case needs its own arm. + if shift == 0 { + lo + } else { + (lo >> shift) | (u64::from(straddle) << (64 - shift)) + } + } + _ => { + let mut raw = 0u128; + for (i, &byte) in bytes.iter().skip(first).take(9).enumerate() { + raw |= u128::from(byte) << (8 * i); + } + (raw >> shift) as u64 + } + }; + mask_to(word, width) +} + +/// Clears every bit at or above `width`. +#[inline] +fn mask_to(word: u64, width: usize) -> u64 { + if width == 64 { + word + } else { + word & ((1u64 << width) - 1) + } +} + +/// Returns the index of the `nth` set bit of `word`, which must hold more than `nth` set bits. +/// +/// Narrows a byte at a time first, so [`select_in_byte`] loops at most seven times. +#[inline] +fn select_in_word(word: u64, nth: usize) -> usize { + debug_assert!( + nth < word.count_ones() as usize, + "rank {nth} is not present in the word" + ); + + let mut remaining = nth; + let mut rest = word; + let mut shift = 0usize; + loop { + let byte = (rest & 0xFF) as u8; + let count = byte.count_ones() as usize; + if remaining < count { + return shift + select_in_byte(byte, remaining); + } + remaining -= count; + rest >>= 8; + shift += 8; + } +} + +/// Returns the index of the `nth` set bit of `byte`. +/// +/// `byte & (byte - 1)` clears the lowest set bit, so `nth` of those leave the target lowest. +#[inline] +fn select_in_byte(byte: u8, nth: usize) -> usize { + let mut rest = byte; + for _ in 0..nth { + rest &= rest - 1; + } + rest.trailing_zeros() as usize +} diff --git a/encodings/elias-fano/src/ef/tests.rs b/encodings/elias-fano/src/ef/tests.rs new file mode 100644 index 00000000000..066caf5b34a --- /dev/null +++ b/encodings/elias-fano/src/ef/tests.rs @@ -0,0 +1,587 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::convert::Infallible; + +use rstest::rstest; + +use super::*; + +// ── Geometry ─────────────────────────────────────────────────────────── + +#[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, +) -> Result<(), Error> { + 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) -> Result<(), Error> { + 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); +} + +/// One sample per `1 << LOG_SAMPLING1` set bits above the first, and none for an empty sequence. +#[test] +fn test_num_samples1() { + assert_eq!(num_samples1(0), 0); + assert_eq!(num_samples1(1), 0); + assert_eq!(num_samples1(1 << LOG_SAMPLING1), 0); + assert_eq!(num_samples1((1 << LOG_SAMPLING1) + 1), 1); + assert_eq!(num_samples1(1 << (LOG_SAMPLING1 + 1)), 1); + assert_eq!(num_samples1((1 << (LOG_SAMPLING1 + 1)) + 1), 2); +} + +// ── Select ───────────────────────────────────────────────────────────── + +/// A pattern near 50% density, deliberately not byte-periodic. +fn mixed_bytes(len: usize) -> Vec { + (0..len) + .map(|i| (i as u8).wrapping_mul(37) ^ 0x5A) + .collect() +} + +/// The bit at `index` of a window beginning `offset` bits into `bytes`, read one bit at a time. +fn bit_at(bytes: &[u8], offset: usize, index: usize) -> bool { + let at = offset + index; + (bytes[at / 8] >> (at % 8)) & 1 == 1 +} + +/// Positions of the bits equal to `want`, as the reference answer. +fn naive_positions( + bytes: &[u8], + offset: usize, + start: usize, + end: usize, + want: bool, +) -> Vec { + (start..end) + .filter(|&i| bit_at(bytes, offset, i) == want) + .map(|i| i - start) + .collect() +} + +/// Both variants must agree with the bit-at-a-time reference across the whole rank range, and +/// both must report `None` one past the last rank. +fn check_against_naive(bytes: &[u8], offset: usize, len: usize, start: usize, end: usize) { + let bits = Bits::new(bytes, offset, len); + for (want, select) in [ + ( + true, + select_range as fn(Bits<'_>, usize, usize, usize) -> Option, + ), + ( + false, + select_zero_range as fn(Bits<'_>, usize, usize, usize) -> Option, + ), + ] { + let expected = naive_positions(bytes, offset, start, end, want); + for (nth, &expected_pos) in expected.iter().enumerate() { + assert_eq!( + select(bits, start, end, nth), + Some(expected_pos), + "want={want} offset={offset} start={start} end={end} nth={nth}" + ); + } + assert_eq!( + select(bits, start, end, expected.len()), + None, + "want={want} offset={offset} start={start} end={end} past-the-end rank" + ); + } +} + +#[rstest] +#[case(0, 0, 128)] +#[case(3, 0, 100)] +#[case(7, 0, 50)] +#[case(0, 0, 1)] +#[case(0, 0, 64)] +#[case(1, 0, 64)] +#[case(0, 0, 65)] +#[case(3, 0, 256)] +#[case(0, 0, 512)] +#[case(0, 0, 513)] +#[case(5, 0, 1024)] +// Windows that start partway in, which is the shape a sample lower bound produces. +#[case(0, 1, 128)] +#[case(0, 63, 128)] +#[case(0, 64, 200)] +#[case(0, 65, 200)] +#[case(3, 70, 300)] +#[case(5, 511, 1024)] +// Sub-word windows, where the only word's valid width is what the zero path must respect. +#[case(1, 0, 1)] +#[case(1, 2, 8)] +#[case(4, 3, 6)] +#[case(7, 0, 1)] +#[case(0, 9, 17)] +#[case(2, 71, 71)] +fn select_agrees_with_naive(#[case] offset: usize, #[case] start: usize, #[case] end: usize) { + let bytes = mixed_bytes((offset + end).div_ceil(8) + 1); + check_against_naive(&bytes, offset, end, start, end); +} + +/// 50% density is exactly where confusing ones with zeros is least visible; these make it +/// obvious. +#[rstest] +#[case::all_zero(0x00)] +#[case::all_one(0xFF)] +#[case::sparse(0x01)] +#[case::dense(0xFE)] +fn select_uniform_density(#[case] fill: u8) { + for (offset, start, end) in [ + (0usize, 0usize, 8usize), + (0, 0, 128), + (3, 2, 5), + (5, 7, 130), + (1, 200, 517), + ] { + let bytes = vec![fill; (offset + end).div_ceil(8) + 1]; + check_against_naive(&bytes, offset, end, start, end); + } +} + +#[test] +fn select_degenerate_buffers() { + // All ones: no zero to find, at any rank. + let ones = vec![0xFFu8; 17]; + let ones = Bits::new(&ones, 0, 128); + assert_eq!(select_zero_range(ones, 0, 128, 0), None); + assert_eq!(select_range(ones, 0, 128, 127), Some(127)); + + // All zeros: the nth zero is at position n, and no set bit exists. + let zeros = vec![0x00u8; 17]; + let zeros = Bits::new(&zeros, 0, 128); + for nth in 0..128 { + assert_eq!( + select_zero_range(zeros, 0, 128, nth), + Some(nth), + "nth={nth}" + ); + } + assert_eq!(select_zero_range(zeros, 0, 128, 128), None); + assert_eq!(select_range(zeros, 0, 128, 0), None); +} + +/// `window_words` has to reproduce the window bit for bit, and zero-pad above it so that counting +/// set bits over the result counts exactly the window's — which is what the decoder relies on. +#[rstest] +#[case(0, 0, 128)] +#[case(3, 5, 130)] +#[case(7, 1, 64)] +#[case(0, 63, 65)] +#[case(2, 0, 1)] +#[case(5, 100, 100)] +#[case(1, 7, 1000)] +fn window_words_reproduces_the_window( + #[case] offset: usize, + #[case] start: usize, + #[case] end: usize, +) { + let bytes = mixed_bytes((offset + end).div_ceil(8) + 1); + let bits = Bits::new(&bytes, offset, end); + let words = window_words(bits, start, end); + + assert_eq!(words.len(), (end - start).div_ceil(64), "word count"); + for index in 0..end - start { + let expected = bit_at(&bytes, offset, start + index); + let actual = (words[index / 64] >> (index % 64)) & 1 == 1; + assert_eq!(actual, expected, "bit {index}"); + } + + let ones: u32 = words.iter().map(|word| word.count_ones()).sum(); + let expected = naive_positions(&bytes, offset, start, end, true).len(); + assert_eq!(ones as usize, expected, "set bits"); +} + +/// An empty window holds nothing, whatever it is asked for. +#[test] +fn select_empty_window() { + let bytes = mixed_bytes(16); + let bits = Bits::new(&bytes, 0, 128); + assert_eq!(select_range(bits, 64, 64, 0), None); + assert_eq!(select_zero_range(bits, 64, 64, 0), None); +} + +/// The counts a window reports must match a bit-at-a-time count, and the first absent rank is the +/// one at that count. +#[test] +fn select_count_agrees_with_naive_count() { + let bytes = mixed_bytes(300); + let bits = Bits::new(&bytes, 3, 2000); + for (start, end) in [(0usize, 2000usize), (5, 1999), (7, 8), (64, 583)] { + let ones = naive_positions(&bytes, 3, start, end, true).len(); + let zeros = (end - start) - ones; + if let Some(last) = ones.checked_sub(1) { + assert!(select_range(bits, start, end, last).is_some()); + } + assert_eq!(select_range(bits, start, end, ones), None); + if let Some(last) = zeros.checked_sub(1) { + assert!(select_zero_range(bits, start, end, last).is_some()); + } + assert_eq!(select_zero_range(bits, start, end, zeros), None); + } +} + +/// A window has to fit in the bytes backing it, whatever the offset. +#[test] +#[should_panic(expected = "cannot hold")] +fn bits_rejects_a_window_past_the_end() { + let bytes = mixed_bytes(2); + let _ = Bits::new(&bytes, 4, 13); +} + +// ── The upper array ──────────────────────────────────────────────────── + +/// `position_of` and `high_of` have to invert each other for every rank and width, since the +/// encoder uses one and all three readers use the other. +#[rstest] +#[case(0, 0, 0)] +#[case(1, 0, 0)] +#[case(0, 7, 3)] +#[case(1_000_000, 4095, 10)] +#[case(u64::MAX, 0, 63)] +#[case(u32::MAX as u64, 1023, 17)] +fn position_and_high_invert(#[case] element: u64, #[case] rank: u64, #[case] lower_width: u8) { + let position = position_of(element, rank, lower_width); + assert_eq!( + high_of(position, rank), + Some(element >> lower_width), + "element={element} rank={rank} lower_width={lower_width}" + ); +} + +/// A position at or below its own rank is what a corrupt upper buffer looks like, and cannot be +/// inverted. +#[test] +fn high_of_rejects_a_position_below_its_rank() { + assert_eq!(high_of(0, 0), None); + assert_eq!(high_of(5, 5), None); + assert_eq!(high_of(5, 9), None); + assert_eq!(high_of(6, 5), Some(0)); +} + +/// Build the upper array for `elements` the way the encoder does. +/// +/// Returns the bits, the shared sample buffer, and the array's length in bits. +fn build_upper(elements: &[u64], lower_width: u8) -> (Vec, Vec, usize) { + let n = elements.len(); + let span = elements[n - 1]; + let upper_len = upper_len(span, n, lower_width).expect("representable"); + let mut builder = UpperBuilder::new(upper_len as usize); + for (index, &element) in elements.iter().enumerate() { + let rank = index as u64; + builder.push(rank, position_of(element, rank, lower_width)); + } + let (bits, samples) = builder.finish(n as u64, upper_len); + (bits, samples, upper_len as usize) +} + +/// A spread wide enough to fill both sample tables: 2000 elements over a span of ~6000 gives +/// `lower_width` 1, five zero-samples and seven one-samples. +fn spread() -> Vec { + (0..2000u64).map(|i| i * 3).collect() +} + +/// Both tables have to hold exactly what `num_samples0` and `num_samples1` predict, because a +/// reader splits the shared buffer at the seam those two imply rather than at a stored offset. +#[test] +fn upper_builder_fills_both_sample_tables() { + let elements = spread(); + let (_, samples, _) = build_upper(&elements, 1); + + let span = elements[elements.len() - 1]; + let zeros = num_samples0(span, 1); + let ones = num_samples1(elements.len()); + assert!(zeros > 0 && ones > 0, "the case must exercise both tables"); + assert_eq!(samples.len() as u64, zeros + ones, "total samples"); +} + +/// Every element's bit has to be recoverable through the one-sample table, and its high part +/// through `high_of` — the round trip a point lookup makes. +#[test] +fn sampled_select_recovers_every_element() { + let elements = spread(); + let lower_width = 1u8; + let (bytes, samples, upper_len) = build_upper(&elements, lower_width); + let bits = Bits::new(&bytes, 0, upper_len); + + let span = elements[elements.len() - 1]; + let seam = num_samples0(span, lower_width) as usize; + let samples1: Vec = samples[seam..] + .iter() + .flat_map(|s| s.to_le_bytes()) + .collect(); + + for (index, &element) in elements.iter().enumerate() { + let rank = index as u64; + let position = sampled_select(bits, &samples1, LOG_SAMPLING1, rank, upper_len, false) + .unwrap_or_else(|| panic!("no set bit for rank {rank}")); + assert_eq!(position as u64, position_of(element, rank, lower_width)); + assert_eq!(high_of(position as u64, rank), Some(element >> lower_width)); + } +} + +/// The zero-sample table answers the other query: `select0(high) - high` is the number of elements +/// whose high part is below `high`, with no rank directory stored. +#[test] +fn sampled_select_zero_counts_elements_below_a_bucket() { + let elements = spread(); + let lower_width = 1u8; + let (bytes, samples, upper_len) = build_upper(&elements, lower_width); + let bits = Bits::new(&bytes, 0, upper_len); + + let span = elements[elements.len() - 1]; + let seam = num_samples0(span, lower_width) as usize; + let samples0: Vec = samples[..seam] + .iter() + .flat_map(|s| s.to_le_bytes()) + .collect(); + + for high in [0u64, 1, 2, 500, 1000, 2998] { + let position = sampled_select(bits, &samples0, LOG_SAMPLING0, high, upper_len, true) + .unwrap_or_else(|| panic!("no bucket boundary for high part {high}")); + let expected = elements + .iter() + .take_while(|&&element| (element >> lower_width) < high) + .count() as u64; + assert_eq!(position as u64 - high, expected, "high={high}"); + } +} + +/// `Ones` has to walk exactly the positions the builder set, in order. +#[test] +fn ones_walks_every_set_bit() { + let elements = spread(); + let lower_width = 1u8; + let (bytes, _, upper_len) = build_upper(&elements, lower_width); + + // The decoder materialises the window as whole words first; do the same by hand. + let words: Vec = bytes + .chunks(8) + .map(|chunk| { + let mut buf = [0u8; 8]; + buf[..chunk.len()].copy_from_slice(chunk); + u64::from_le_bytes(buf) + }) + .collect(); + + let mut ones = Ones::new(&words); + for (index, &element) in elements.iter().enumerate() { + let rank = index as u64; + assert_eq!( + ones.next().map(|p| p as u64), + Some(position_of(element, rank, lower_width)), + "rank {rank}" + ); + } + assert_eq!(ones.next(), None, "no set bits past the last element"); + assert!(upper_len > 0); +} + +// ── The codec, end to end ────────────────────────────────────────────── + +/// A low-bits source backed by a plain slice. +/// +/// Six lines, and the only thing an embedder has to supply. A host reading its own bit-packed +/// column is not much longer. +struct VecLows(Vec); + +impl LowBits for VecLows { + type Error = Infallible; + + fn get(&mut self, rank: u64) -> Result { + Ok(self.0[rank as usize]) + } +} + +/// Everything a reader needs, held together so a test can keep the borrows alive. +struct RoundTrip { + upper: Vec, + samples: Vec, + seam: usize, + lows: VecLows, + lower_width: u8, + upper_len: usize, + len: usize, + span: u64, +} + +impl RoundTrip { + fn encode(elements: &[u64]) -> Self { + let span = elements[elements.len() - 1]; + let encoded = encode(elements.iter().copied(), span).expect("representable"); + + // The two sample tables share one buffer and the seam is derived, never stored. + let samples: Vec = encoded + .samples + .iter() + .flat_map(|sample| sample.to_le_bytes()) + .collect(); + let seam = num_samples0(span, encoded.lower_width) as usize * size_of::(); + + Self { + upper: encoded.upper, + samples, + seam, + lows: VecLows(encoded.lower), + lower_width: encoded.lower_width, + upper_len: encoded.upper_len as usize, + len: elements.len(), + span, + } + } + + fn layout(&self) -> Layout<'_> { + let (_, samples1) = self.samples.split_at(self.seam); + Layout::new( + Bits::new(&self.upper, 0, self.upper_len), + samples1, + self.lower_width, + 0, + self.len, + ) + } +} + +fn sparse() -> Vec { + (0..200u64).map(|i| i * 977).collect() +} + +fn dense() -> Vec { + (0..500u64).collect() +} + +/// Fifty distinct values spread over a wide universe, each repeated twenty times, so `lower_width` +/// stays positive and every occupied bucket is deep enough to send `next_geq` into its bisection. +fn duplicates() -> Vec { + (0..50u64) + .flat_map(|value| std::iter::repeat_n(value * 4096, 20)) + .collect() +} + +/// Encode, then read every element back through the public API alone. +/// +/// This is the test the module exists to make possible. If it cannot be written without reaching +/// past that API, this is a bag of helpers rather than a codec. +#[rstest] +#[case::sparse(sparse())] +#[case::dense(dense())] +#[case::duplicates(duplicates())] +#[case::all_equal(vec![7u64; 64])] +#[case::single(vec![0u64])] +#[case::single_wide(vec![u64::MAX])] +fn codec_reads_back_every_element(#[case] elements: Vec) { + let round_trip = RoundTrip::encode(&elements); + let mut lows = VecLows(round_trip.lows.0.clone()); + let layout = round_trip.layout(); + + for (index, &expected) in elements.iter().enumerate() { + assert_eq!( + element_at(layout, index, &mut lows), + Ok(expected), + "index {index}" + ); + } + + // And once more in reverse. A read holds no state, so the order cannot matter — which is the + // claim being pinned. + for (index, &expected) in elements.iter().enumerate().rev() { + assert_eq!( + element_at(layout, index, &mut lows), + Ok(expected), + "reverse index {index}" + ); + } +} + +/// What the encoder writes, the validator must accept. +#[rstest] +#[case::sparse(sparse())] +#[case::dense(dense())] +#[case::duplicates(duplicates())] +#[case::single_wide(vec![u64::MAX])] +fn encoder_output_validates(#[case] elements: Vec) { + let round_trip = RoundTrip::encode(&elements); + assert_eq!( + validate_layout( + round_trip.span, + round_trip.len, + round_trip.lower_width, + round_trip.upper_len as u64, + &round_trip.samples, + ), + Ok(()) + ); +} + +/// The bulk decoder has to reproduce the sequence from the same buffers the reader walks. +#[rstest] +#[case::sparse(sparse())] +#[case::dense(dense())] +#[case::duplicates(duplicates())] +fn decoder_reproduces_the_sequence(#[case] elements: Vec) { + let round_trip = RoundTrip::encode(&elements); + let (_, samples1) = round_trip.samples.split_at(round_trip.seam); + let bits = Bits::new(&round_trip.upper, 0, round_trip.upper_len); + let last = elements.len() as u64 - 1; + + // Trim to the window holding exactly these elements' set bits, which is what a caller does and + // what lets the walk below run without a per-element bound check. + let start = position_of_rank(bits, samples1, 0).expect("first element"); + let end = position_of_rank(bits, samples1, last).expect("last element") + 1; + let words = window_words(bits, start, end); + + let mut decoder = Decoder::new(&words, start, 0, elements.len(), round_trip.lower_width) + .expect("well-formed"); + let mut decoded = vec![0u64; elements.len()]; + decoder.segment(Some(&round_trip.lows.0), |index, element| { + decoded[index] = element; + }); + assert_eq!(decoder.finish(), Ok(())); + assert_eq!(decoded, elements); +} diff --git a/encodings/elias-fano/src/ef/upper.rs b/encodings/elias-fano/src/ef/upper.rs new file mode 100644 index 00000000000..39265794afc --- /dev/null +++ b/encodings/elias-fano/src/ef/upper.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The upper bit array: one set bit per element at [`position_of`], one unset bit per bucket +//! boundary, with [`high_of`] inverting it. + +use super::Bits; +use super::LOG_SAMPLING0; +use super::LOG_SAMPLING1; +use super::select_range; +use super::select_zero_range; + +/// The bit position claimed by the element of rank `rank`. +/// +/// The `+ rank` keeps positions distinct when elements share a high part; the `+ 1` is the sentinel +/// that aligns the unset bits with the high parts. +#[inline] +pub fn position_of(element: u64, rank: u64, lower_width: u8) -> u64 { + (element >> lower_width) + rank + 1 +} + +/// The high part of the element of rank `rank` sitting at `position`, inverting [`position_of`]. +/// +/// `None` when the position is at or below its own rank. No array this crate builds can produce +/// that, but the upper buffer's contents are never validated, so a corrupt file can. +#[inline] +pub fn high_of(position: u64, rank: u64) -> Option { + position.checked_sub(rank + 1) +} + +/// The element with high part `high` and low part `low`. +#[inline] +pub fn element_of(high: u64, low: u64, lower_width: u8) -> u64 { + (high << lower_width) | low +} + +/// One entry of a sample table, stored as a raw little-endian `u64`. +/// +/// Deserialized buffers carry no alignment guarantee, hence byte-at-a-time rather than a cast. +/// +/// # Panics +/// +/// Panics if `table` holds fewer than `index + 1` entries. +#[inline] +pub fn read_sample(table: &[u8], index: usize) -> u64 { + let mut bytes = [0u8; 8]; + bytes.copy_from_slice(&table[index * 8..][..8]); + u64::from_le_bytes(bytes) +} + +/// The bit position of the `target`-th set bit of `bits`, or of its `target`-th unset bit when +/// `zeros`, found through the sample table that brackets it. +/// +/// `table` holds one position per `1 << log_sampling` bits of the kind being counted, from rank 1 +/// upward; rank 0 is never stored, since the search starts at bit zero anyway. The sample lower- +/// bounds the scan, so it costs the sampling rate rather than the length of the array. +/// +/// Returns the absolute position, not one relative to the window. `inline(always)` because both +/// callers pass `log_sampling` and `zeros` as constants, which measurably do not fold otherwise. +#[inline(always)] +#[allow(clippy::inline_always)] +pub fn sampled_select( + bits: Bits<'_>, + table: &[u8], + log_sampling: usize, + target: u64, + end: usize, + zeros: bool, +) -> Option { + let sample = (target >> log_sampling) as usize; + let start = if sample == 0 { + 0 + } else { + usize::try_from(read_sample(table, sample - 1)).ok()? + }; + let nth = usize::try_from(target - ((sample as u64) << log_sampling)).ok()?; + let offset = if zeros { + select_zero_range(bits, start, end, nth) + } else { + select_range(bits, start, end, nth) + }?; + Some(start + offset) +} + +/// A walk over the set bits of a window, a `u64` word at a time. +/// +/// Costs a `trailing_zeros` and a clear-lowest-bit per element. Takes whole words, so an unaligned +/// window goes through [`window_words`](super::window_words) first. +pub 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> { + /// Walk the set bits of `words`, in order. + pub fn new(words: &'a [u64]) -> Self { + Self { + words, + word: 0, + current: words.first().copied().unwrap_or(0), + } + } +} + +impl Iterator for Ones<'_> { + type Item = usize; + + /// The next set bit's index within the window, or `None` once the words run out. + #[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) + } +} + +/// Builds the upper array and both sample tables together, in one pass over the elements. +/// +/// A zero-sample is the position of a sampled *unset* bit, and the unset runs are only known as the +/// set bits bounding them are written, so the tables cannot be built in a later pass. +pub struct UpperBuilder { + bits: Vec, + len: usize, + samples0: Vec, + samples1: Vec, + /// 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 { + /// An all-unset array of `upper_len` bits, with both sample tables empty. + pub fn new(upper_len: usize) -> Self { + Self { + bits: vec![0u8; upper_len.div_ceil(8)], + len: upper_len, + samples0: Vec::new(), + samples1: Vec::new(), + next_zero_sample: 1 << LOG_SAMPLING0, + } + } + + /// Record the element of rank `rank` as a set bit at `position`. + /// + /// Must be called with strictly increasing `rank` and `position`. + pub fn push(&mut self, rank: u64, position: u64) { + self.sample_zeros_below(position, rank); + + debug_assert!(position < self.len as u64, "position out of bounds"); + self.bits[(position / 8) as usize] |= 1 << (position % 8); + + if rank > 0 && rank.is_multiple_of(1 << LOG_SAMPLING1) { + self.samples1.push(position); + } + } + + /// Emit a zero-sample for every sampled unset rank below `position`. + /// + /// An unset bit in this run has exactly `ones` set bits before it, so its rank is + /// `position - ones` — inverted here to get the position back. + fn sample_zeros_below(&mut self, position: u64, ones: u64) { + while self.next_zero_sample + ones < position { + self.samples0.push(self.next_zero_sample + ones); + self.next_zero_sample += 1 << LOG_SAMPLING0; + } + } + + /// Close the array, returning the upper bytes and both sample tables in one buffer, zeros + /// first. The seam is not returned: a reader recomputes it from the universe with + /// [`num_samples0`](super::num_samples0). + pub fn finish(mut self, n: u64, upper_len: u64) -> (Vec, Vec) { + // The trailing unset bits past the last element, which `push` never reached: the bucket + // boundaries above the maximum element's high part, plus the guard zero. + self.sample_zeros_below(upper_len, n); + + self.samples0.extend_from_slice(&self.samples1); + (self.bits, self.samples0) + } +} diff --git a/encodings/elias-fano/src/ef/validate.rs b/encodings/elias-fano/src/ef/validate.rs new file mode 100644 index 00000000000..8a884889f30 --- /dev/null +++ b/encodings/elias-fano/src/ef/validate.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Checking a stored layout against the universe it claims. +//! +//! Every stored number is recomputed from `(span, n)` alone and a layout that disagrees is refused, +//! which is what lets the sample tables carry no seam and every sampled search skip its bounds +//! checks. +//! +//! The upper array's contents are *not* checked: that costs a pass over the whole bit array, and a +//! reader raises rather than underflowing on a position it cannot invert. + +use super::LOG_SAMPLING0; +use super::LOG_SAMPLING1; +use super::Malformed; +use super::Table; +use super::lower_width; +use super::num_samples0; +use super::num_samples1; +use super::num_zeros; +use super::read_sample; +use super::upper_len; + +/// Check that a stored layout describes the universe it claims. +/// +/// `samples` is the shared table, zero-samples first; the seam between the two is derived here +/// rather than stored, so a buffer too short to reach it fails as a count mismatch. +/// +/// An empty sequence is accepted unconditionally, having no geometry to check. +pub fn validate_layout( + span: u64, + n: usize, + stored_lower_width: u8, + stored_upper_len: u64, + samples: &[u8], +) -> Result<(), Malformed> { + if n == 0 { + return Ok(()); + } + + let expected_width = lower_width(span, n); + if stored_lower_width != expected_width { + return Err(Malformed::LowerWidth { + expected: expected_width, + found: stored_lower_width, + }); + } + + let expected_upper_len = upper_len(span, n, expected_width)?; + if stored_upper_len != expected_upper_len { + return Err(Malformed::UpperLen { + expected: expected_upper_len, + found: stored_upper_len, + }); + } + + // Both tables sample from rank 1 upward, so their sizes follow from the layout and a reader + // never has to bounds-check a lookup. + let expected_samples0 = num_samples0(span, expected_width); + debug_assert_eq!( + expected_samples0, + (num_zeros(expected_upper_len, n) - 1) >> LOG_SAMPLING0, + "the two derivations of the zero-sample count must agree" + ); + let expected_samples1 = num_samples1(n); + let expected = expected_samples0 + expected_samples1; + let found = (samples.len() / size_of::()) as u64; + if found != expected { + return Err(Malformed::SampleCount { expected, found }); + } + + // The zero table comes first. The count check above already proves the buffer reaches the seam. + let seam = (expected_samples0 as usize) * size_of::(); + let (samples0, samples1) = samples.split_at(seam); + + // A sample is fed straight to `select_range` as a window start, which asserts rather than + // raises past the end, so every one is checked — 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. + for (table, name, log_sampling, floor) in [ + (samples0, Table::Zero, LOG_SAMPLING0, 0), + (samples1, Table::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; + if !(minimum..expected_upper_len).contains(&sample) { + return Err(Malformed::SampleOutOfRange { + table: name, + index, + sample, + min: minimum, + max: expected_upper_len, + }); + } + if previous.is_some_and(|previous| previous >= sample) { + return Err(Malformed::SamplesNotIncreasing { table: name, index }); + } + previous = Some(sample); + } + } + + Ok(()) +} diff --git a/encodings/elias-fano/src/lib.rs b/encodings/elias-fano/src/lib.rs new file mode 100644 index 00000000000..3577d0a3911 --- /dev/null +++ b/encodings/elias-fano/src/lib.rs @@ -0,0 +1,77 @@ +// 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 in constant time. Against bit-packing at `ceil(log2(u))` bits the saving +//! is `log2(n)`, so it widens with row count. +//! +//! Inputs must be non-decreasing and non-nullable; duplicates are fine. See [`elias_fano_encode`] +//! for the compression entry point and [`initialize`] to register the encoding in a session. +//! +//! [`ef`] holds the codec; everything else here binds it to Vortex's array model. +//! +//! The sampled select index follows Vigna's [broadword][] construction, with the two-table sampling +//! scheme and its parameters after [`rise-rs`][] (MIT); [`vers`][] was consulted as a further +//! reference. No code is taken from either. +//! +//! [broadword]: https://vigna.di.unimi.it/ftp/papers/Broadword.pdf +//! [`rise-rs`]: https://github.com/AngeloSav/rise-rs +//! [`vers`]: https://github.com/Cydhra/vers + +mod access; +mod array; +mod compress; +mod compute; +pub mod ef; +mod lower; +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; +use vortex_array::ArrayVTable; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::session::ArraySessionExt; +use vortex_error::VortexError; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +/// The codec's own errors, carried into Vortex's. +/// +/// Functions rather than `From` impls: once [`ef`] is its own crate both types are foreign here and +/// the orphan rule refuses the impl. +pub(crate) fn unrepresentable(error: ef::Error) -> VortexError { + vortex_err!("Elias-Fano {error}") +} + +pub(crate) fn malformed(error: ef::Malformed) -> VortexError { + vortex_err!("Elias-Fano {error}") +} + +/// Initialize the Elias-Fano encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(EliasFano); + + // Answered from the layout rather than the data: the encoding only decodes correctly for a + // non-decreasing sequence, so sortedness is a precondition rather than a measurement. + session.aggregate_fns().register_aggregate_kernel( + EliasFano.id(), + Some(IsSorted.id()), + &compute::is_sorted::EliasFanoIsSortedKernel, + ); +} + +#[cfg(test)] +mod tests; diff --git a/encodings/elias-fano/src/lower.rs b/encodings/elias-fano/src/lower.rs new file mode 100644 index 00000000000..ab1a0c7933a --- /dev/null +++ b/encodings/elias-fano/src/lower.rs @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reading the low-bits child in bulk, which is an ordinary FastLanes `BitPacked` array. +//! +//! Only the bulk decode comes through here; a per-element read goes through the child's own +//! `scalar_at` in [`crate::access`]. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; + +/// The low-bits child as a `BitPacked` view, if its bytes can be read where they lie: +/// FastLanes-packed, unpatched, host-resident and `u64`-aligned. +/// +/// `unpacked_chunks` reinterprets those bytes as `&[u64]` without a check of its own, so the four +/// conditions are stated here. `None` means the low bits must be materialised. +pub(crate) fn readable_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::()?; + (packed.patches().is_none() + && packed + .packed() + .as_host_opt() + .is_some_and(|buffer| buffer.is_aligned(Alignment::of::()))) + .then_some(packed) +} + +/// The low bits of an already-windowed child, materialised when its bytes cannot be read in place. +/// +/// The child spans the whole encoded sequence, so a caller windows it to the range being folded +/// first. +pub(crate) fn materialise(windowed: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + Ok(windowed + .execute::(ctx)? + .into_buffer::()) +} diff --git a/encodings/elias-fano/src/rules.rs b/encodings/elias-fano/src/rules.rs new file mode 100644 index 00000000000..b54a1a3807a --- /dev/null +++ b/encodings/elias-fano/src/rules.rs @@ -0,0 +1,12 @@ +// 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 crate::EliasFano; + +/// Reductions an Elias-Fano array can absorb from its parent without reading a buffer: slicing, +/// which costs one metadata field. +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new(&[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..b92328f9cfb --- /dev/null +++ b/encodings/elias-fano/src/tests.rs @@ -0,0 +1,712 @@ +// 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::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +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_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::EliasFano; +use crate::EliasFanoArray; +use crate::EliasFanoArraySlotsExt; +use crate::EliasFanoData; +use crate::ef; +use crate::elias_fano_encode; + +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 `scalar_at`, must match `expected`. +/// +/// Probed in a random order and then in sequence. The path is stateless, so the two must give the +/// same answers — which is what the shuffle is there to catch. +fn check_access(array: &EliasFanoArray, expected: &[Scalar], seed: u64) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for index in Rng(seed).permutation(expected.len()) { + assert_eq!( + array.execute_scalar(index, &mut ctx)?, + expected[index], + "index {index}" + ); + } + for (index, want) in expected.iter().enumerate() { + assert_eq!( + &array.execute_scalar(index, &mut ctx)?, + want, + "sequential index {index}" + ); + } + Ok(()) +} + +/// 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)?; + 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(()) +} + +// ── 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)?; + 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(()) +} + +/// `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(()) +} + +#[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 per-element read 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)?; + + // And sliced, which is the one place two bases have to be reconciled: the slice's own rank + // offset, and the offset the slot carries. A mishandled base comes back as a wrong element. + 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)?; + 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) & ef::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)?; + 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 = ef::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 per-element read. +#[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 = !ef::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)?; + 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 ───────────────────────────────────────── + +// ── Take and filter pushdown ──────────────────────────────────────────── + +// ── 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(()) +} + +// ── 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`. +#[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 `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 per-element read 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" + ); + assert!( + rebuilt.execute_scalar(0, &mut ctx).is_err(), + "scalar_at into a malformed upper array must raise" + ); + Ok(()) +}