From d926f8824517cd2228637ad12f8ed4a985292c8e Mon Sep 17 00:00:00 2001 From: rapour Date: Mon, 19 Oct 2026 15:30:44 +0330 Subject: [PATCH 1/2] feat: non-indexed unary encoding Signed-off-by: rapour --- .github/workflows/codspeed.yml | 2 +- Cargo.lock | 16 + Cargo.toml | 2 + encodings/elias-fano/Cargo.toml | 35 ++ encodings/elias-fano/benches/unary.rs | 154 +++++++ .../elias-fano/goldenfiles/unary.metadata | 1 + encodings/elias-fano/src/lib.rs | 26 ++ encodings/elias-fano/src/unary/array.rs | 400 +++++++++++++++++ encodings/elias-fano/src/unary/compress.rs | 89 ++++ encodings/elias-fano/src/unary/mod.rs | 14 + encodings/elias-fano/src/unary/tests.rs | 403 ++++++++++++++++++ 11 files changed, 1141 insertions(+), 1 deletion(-) create mode 100644 encodings/elias-fano/Cargo.toml create mode 100644 encodings/elias-fano/benches/unary.rs create mode 100644 encodings/elias-fano/goldenfiles/unary.metadata create mode 100644 encodings/elias-fano/src/lib.rs create mode 100644 encodings/elias-fano/src/unary/array.rs create mode 100644 encodings/elias-fano/src/unary/compress.rs create mode 100644 encodings/elias-fano/src/unary/mod.rs create mode 100644 encodings/elias-fano/src/unary/tests.rs diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 23f86dde9cd..4d689b925a7 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -58,7 +58,7 @@ jobs: - { shard: 5, name: "Core & encodings", packages: "vortex-buffer vortex-error vortex-mask vortex-compute vortex-file vortex-alp vortex-bytebool vortex-datetime-parts" } - { shard: 6, name: "Main library & storage", packages: "vortex vortex-btrblocks vortex-compressor vortex-row" } - { shard: 7, name: "FastLanes & decimals", packages: "vortex-decimal-byte-parts vortex-fastlanes", features: "--features _test-harness" } - - { shard: 8, name: "Encodings & layout", packages: "vortex-pco vortex-runend vortex-sequence vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } + - { shard: 8, name: "Encodings & layout", packages: "vortex-elias-fano vortex-pco vortex-runend vortex-sequence vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } - { shard: 9, name: "Strings, tensor & spatial", packages: "vortex-fsst vortex-tensor vortex-spatial", features: "--features _test-harness" } name: "Benchmark with Codspeed (Shard #${{ matrix.shard }} - ${{ matrix.name }})" timeout-minutes: 30 diff --git a/Cargo.lock b/Cargo.lock index 38e2760c516..48646286090 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11906,6 +11906,22 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-elias-fano" +version = "0.1.0" +dependencies = [ + "codspeed-divan-compat", + "mimalloc", + "num-traits", + "prost 0.14.4", + "rand 0.10.3", + "rstest", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-session", +] + [[package]] name = "vortex-error" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index a551b0f397b..b540f6500f7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,6 +61,7 @@ members = [ "encodings/parquet-variant", "encodings/onpair", "encodings/uuid", + "encodings/elias-fano", # Benchmarks "benchmarks/bench-support", "benchmarks/lance-bench", @@ -317,6 +318,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..99a44eacd6c --- /dev/null +++ b/encodings/elias-fano/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "vortex-elias-fano" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex encodings for non-decreasing integer sequences" +edition = { 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] +num-traits = { workspace = true } +prost = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +mimalloc = { workspace = true } +rand = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } + +[[bench]] +name = "unary" +harness = false + +[lints] +workspace = true diff --git a/encodings/elias-fano/benches/unary.rs b/encodings/elias-fano/benches/unary.rs new file mode 100644 index 00000000000..cc7845c38f4 --- /dev/null +++ b/encodings/elias-fano/benches/unary.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Baselines for the naive `Unary` encoding, so each later optimisation can be measured +//! against them. Keep the names and inputs stable: CodSpeed compares benchmarks by name. +//! +//! Four shapes, all small enough for the encoder's size rule: +//! +//! * `Dense`: values drawn from `[0, LEN)`, about one unset bit per row. +//! * `Spread`: values drawn from `[0, 2 * LEN)`, the densest an Elias-Fano high part gets. +//! * `Runs`: `LEN / 64` values, 64 rows each, so set bits come in long runs. +//! * `Gap`: dense halves either side of one jump of `GAP` values, a long run of unset bits that +//! every scan crossing it has to walk. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +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::unary_encode; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_elias_fano::initialize(&session); + session +}); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Shape { + Dense, + Spread, + Runs, + Gap, +} + +const SHAPES: &[Shape] = &[Shape::Dense, Shape::Spread, Shape::Runs, Shape::Gap]; + +/// Rows per array, small enough that a whole decode stays inside the 1 ms per-iteration budget. +const LEN: usize = 1 << 16; + +/// The jump in `Gap`, in values. +const GAP: u64 = 4 * LEN as u64; + +/// Rows per value in `Runs`. +const ROWS_PER_VALUE: usize = 64; + +/// Probes per iteration. Each is a linear scan from the first bit, so 256 keep an iteration near +/// the budget, as in `vortex-array`'s probe benchmark. +const PROBES: usize = 256; + +fn values(shape: Shape) -> Vec { + let mut rng = StdRng::seed_from_u64(0); + let len = LEN as u64; + let mut out: Vec = match shape { + Shape::Dense => (0..LEN).map(|_| rng.random_range(0..len)).collect(), + Shape::Spread => (0..LEN).map(|_| rng.random_range(0..2 * len)).collect(), + Shape::Runs => (0..LEN).map(|row| (row / ROWS_PER_VALUE) as u64).collect(), + Shape::Gap => (0..LEN) + .map(|row| { + let offset = if row < LEN / 2 { 0 } else { GAP }; + offset + rng.random_range(0..len / 2) + }) + .collect(), + }; + out.sort_unstable(); + out +} + +fn primitive(shape: Shape) -> PrimitiveArray { + PrimitiveArray::from_iter(values(shape)) +} + +fn encoded(shape: Shape) -> ArrayRef { + unary_encode(primitive(shape).as_ref().as_::()) + .unwrap() + .unwrap() + .into_array() +} + +#[divan::bench(args = SHAPES)] +fn encode(bencher: Bencher, shape: Shape) { + let array = primitive(shape); + bencher.bench_local(|| { + divan::black_box(unary_encode(array.as_ref().as_::()).unwrap()); + }); +} + +/// The whole array, in one walk over the set bits. +#[divan::bench(args = SHAPES)] +fn decode(bencher: Bencher, shape: Shape) { + let array = encoded(shape); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_local_values(|(array, mut ctx)| { + divan::black_box(array.execute::(&mut ctx).unwrap()); + }); +} + +/// Point reads in a random order, each a scan from the first bit. +#[divan::bench(args = SHAPES)] +fn scalar_at_random(bencher: Bencher, shape: Shape) { + let array = encoded(shape); + let mut rng = StdRng::seed_from_u64(1); + let indices: Vec = (0..PROBES).map(|_| rng.random_range(0..LEN)).collect(); + 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()); + } + }); +} + +/// Consecutive point reads from the middle, which a cursor could answer from the last position. +#[divan::bench(args = SHAPES)] +fn scalar_at_sequential(bencher: Bencher, shape: Shape) { + let array = encoded(shape); + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| { + for index in LEN / 2..LEN / 2 + PROBES { + divan::black_box(array.execute_scalar(index, &mut ctx).unwrap()); + } + }); +} + +/// The middle half, through the default slice path, which decodes the whole array first. +#[divan::bench(args = SHAPES)] +fn slice_decode(bencher: Bencher, shape: Shape) { + let array = encoded(shape); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_local_values(|(array, mut ctx)| { + let sliced = array.slice(LEN / 4..3 * LEN / 4).unwrap(); + divan::black_box(sliced.execute::(&mut ctx).unwrap()); + }); +} + +fn main() { + divan::main(); +} diff --git a/encodings/elias-fano/goldenfiles/unary.metadata b/encodings/elias-fano/goldenfiles/unary.metadata new file mode 100644 index 00000000000..d3da94c6217 --- /dev/null +++ b/encodings/elias-fano/goldenfiles/unary.metadata @@ -0,0 +1 @@ +ÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿ \ No newline at end of file diff --git a/encodings/elias-fano/src/lib.rs b/encodings/elias-fano/src/lib.rs new file mode 100644 index 00000000000..175f15ae391 --- /dev/null +++ b/encodings/elias-fano/src/lib.rs @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Encodings for non-decreasing integer sequences. +//! +//! [`Unary`] stores a non-decreasing, non-nullable unsigned column as one bit array, writing +//! each value's gap from the previous one in unary. It takes `len + (max - min)` bits, so it suits +//! dense sequences, where values rise by about one per row or less. It is also meant to hold the +//! high bits of the Elias-Fano encoding this crate is being built towards. +//! +//! None of these encodings is in a file edition or registered by `vortex-file`, so the default +//! writers reject them for now. [`initialize`] registers them in a session. + +mod unary; + +pub use unary::Unary; +pub use unary::UnaryArray; +pub use unary::UnaryData; +pub use unary::unary_encode; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Register this crate's encodings in `session`. +pub fn initialize(session: &VortexSession) { + session.arrays().register(Unary); +} diff --git a/encodings/elias-fano/src/unary/array.rs b/encodings/elias-fano/src/unary/array.rs new file mode 100644 index 00000000000..072b47c5dc9 --- /dev/null +++ b/encodings/elias-fano/src/unary/array.rs @@ -0,0 +1,400 @@ +// 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 num_traits::AsPrimitive; +use prost::Message; +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::buffer::BufferHandle; +use vortex_array::dtype::DType; +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::match_each_unsigned_integer_ptype; +use vortex_array::scalar::Scalar; +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::BitBuffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::compress::unary_decode; + +/// A [`Unary`]-encoded Vortex array. +pub type UnaryArray = Array; + +/// Metadata stored alongside the `bits` buffer. +#[derive(Clone, prost::Message)] +pub(super) struct UnaryMetadata { + /// The first value. + #[prost(uint64, tag = "1")] + pub(super) base: u64, + /// The last value. + #[prost(uint64, tag = "2")] + pub(super) max: u64, +} + +/// The parts of a [`Unary`] array: its bits, and its first and last values. +#[derive(Clone, Debug)] +pub struct UnaryData { + bits: BitBuffer, + base: u64, + max: u64, +} + +impl Display for UnaryData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "base: {}, max: {}", self.base, self.max) + } +} + +impl ArrayHash for UnaryData { + fn array_hash(&self, state: &mut H, accuracy: EqMode) { + self.base.hash(state); + self.max.hash(state); + self.bits.array_hash(state, accuracy); + } +} + +impl ArrayEq for UnaryData { + fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { + self.base == other.base + && self.max == other.max + && self.bits.array_eq(&other.bits, accuracy) + } +} + +/// The number of bits `len` values spanning `span` take, if a `usize` can count them. +pub(super) fn bit_len(len: usize, span: u64) -> Option { + u64::try_from(len) + .ok()? + .checked_add(span) + .and_then(|bits| usize::try_from(bits).ok()) +} + +fn too_many_bits(len: usize, span: u64) -> VortexError { + vortex_err!("UnaryArray of {len} values spanning {span} needs more bits than a usize counts") +} + +impl UnaryData { + /// Rebuild the parts from a stored buffer. The buffer's size is checked first, because + /// [`BitBuffer::new`] panics on one too short for its bits. + fn from_buffer( + buffer: ByteBuffer, + base: u64, + max: u64, + dtype: &DType, + len: usize, + ) -> VortexResult { + let span = max + .checked_sub(base) + .ok_or_else(|| vortex_err!("UnaryArray base {base} exceeds max {max}"))?; + let bits = bit_len(len, span).ok_or_else(|| too_many_bits(len, span))?; + vortex_ensure!( + buffer.len().checked_mul(8).is_some_and(|held| held >= bits), + "UnaryArray needs {bits} bits, but its buffer holds {} bytes", + buffer.len() + ); + let data = Self { + bits: BitBuffer::new(buffer, bits), + base, + max, + }; + data.validate(dtype, len)?; + Ok(data) + } + + /// Check the parts against an array of `len` values of `dtype`. + /// + /// Besides the metadata and the first and last bits, this counts the set bits, one pass over + /// the buffer. Any bits that pass then describe a valid array: with exactly `len` set bits + /// among `len + span`, the `k`-th sits in `[k, k + span]`, so every value lies within + /// `[base, max]`. The readers rely on this and check no value themselves. + fn validate(&self, dtype: &DType, len: usize) -> VortexResult<()> { + let DType::Primitive(ptype, NonNullable) = dtype else { + vortex_bail!("UnaryArray requires a non-nullable unsigned integer dtype, got {dtype}"); + }; + vortex_ensure!( + ptype.is_unsigned_int(), + "UnaryArray requires a non-nullable unsigned integer dtype, got {dtype}" + ); + vortex_ensure!( + self.base <= self.max, + "UnaryArray base {} exceeds max {}", + self.base, + self.max + ); + vortex_ensure!( + self.max <= ptype.max_value_as_u64(), + "UnaryArray max {} does not fit {ptype}", + self.max + ); + let span = self.span(); + let bits = bit_len(len, span).ok_or_else(|| too_many_bits(len, span))?; + vortex_ensure!( + self.bits.offset() == 0, + "UnaryArray bits must start at offset 0, got {}", + self.bits.offset() + ); + vortex_ensure_eq!(self.bits.len(), bits); + if len == 0 { + vortex_ensure!( + span == 0, + "An empty UnaryArray must have its base equal to its max, got a span of {span}" + ); + } else { + // Each value ends in a set bit, so these make `base` and `max` the real first and last + // values. + vortex_ensure!( + self.bits.value(0) && self.bits.value(bits - 1), + "UnaryArray bits must start and end with a set bit" + ); + } + // The only linear check, so it runs last. + vortex_ensure_eq!(self.bits.true_count(), len); + Ok(()) + } + + pub(super) fn bits(&self) -> &BitBuffer { + &self.bits + } + + /// The first value. + pub(super) fn base(&self) -> u64 { + self.base + } + + /// The last value. + pub(super) fn max(&self) -> u64 { + self.max + } + + /// `max - base`, which is also the number of unset bits. + pub(super) fn span(&self) -> u64 { + self.max - self.base + } + + /// The value at `index`, found by scanning for its set bit from the start. + pub(super) fn value_at(&self, index: usize) -> VortexResult { + let position = self + .bits + .select(index) + .ok_or_else(|| vortex_err!("UnaryArray has no set bit for index {index}"))?; + // The `index`-th set bit has `index` set bits before it, so this cannot underflow, and + // `validate` keeps it within `span`, so the sum cannot pass `max`. + let delta = u64::try_from(position - index)?; + Ok(self.base + delta) + } +} + +impl VTable for Unary { + type TypedArrayData = UnaryData; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.unary"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.is_empty(), + "UnaryArray has no children, got {}", + slots.len() + ); + data.validate(dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 1 + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + match idx { + 0 => BufferHandle::new_host(array.bits().inner().clone()), + _ => vortex_panic!("UnaryArray buffer index {idx} out of bounds"), + } + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + match idx { + 0 => Some("bits".to_string()), + _ => None, + } + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure_eq!(buffers.len(), 1); + let data = UnaryData::from_buffer( + buffers[0].try_to_host_sync()?, + array.base(), + array.max(), + array.dtype(), + array.len(), + )?; + Ok(ArrayParts::new( + self.clone(), + array.dtype().clone(), + array.len(), + data, + ArraySlots::new(), + )) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + UnaryMetadata { + base: array.base(), + max: array.max(), + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "UnaryArray expects 1 buffer, got {}", + buffers.len() + ); + vortex_ensure!( + children.is_empty(), + "UnaryArray expects 0 children, got {}", + children.len() + ); + let metadata = UnaryMetadata::decode(metadata)?; + let data = UnaryData::from_buffer( + buffers[0].try_to_host_sync()?, + metadata.base, + metadata.max, + dtype, + len, + )?; + Ok(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + ArraySlots::new(), + )) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + vortex_panic!("UnaryArray has no slots, got index {idx}") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(unary_decode(&array).into_array())) + } +} + +impl OperationsVTable for Unary { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Unary>, + index: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + let value = array.value_at(index)?; + // `validate` keeps `max` within the ptype and every value within `max`. + Ok(match_each_unsigned_integer_ptype!( + array.dtype().as_ptype(), + |T| { Scalar::primitive::(value.as_(), NonNullable) } + )) + } +} + +impl ValidityVTable for Unary { + fn validity(_array: ArrayView<'_, Unary>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +/// Unary encoding for non-decreasing, non-nullable unsigned integers. +/// +/// The values are stored as one bit array. Each value's gap from the one before it is written in +/// unary, as that many unset bits followed by a set bit. The first value is the array's `base`, so +/// its gap is zero. The set bit of the value at index `i` therefore sits at `(value - base) + i`, +/// and the array takes `len + (max - base)` bits. +/// +/// For `[3, 5, 5, 9]` the base is 3 and the bits are `1 001 1 00001`. +/// +/// Reads are deliberately simple for now: `scalar_at` scans from the first bit, and a slice +/// decodes the whole array. +#[derive(Clone, Debug)] +pub struct Unary; + +impl Unary { + /// Assemble an array of `len` values from its bits and its first and last values. + /// + /// Prefer [`unary_encode`](crate::unary_encode) unless you already hold the bits. + pub fn try_new( + bits: BitBuffer, + base: u64, + max: u64, + ptype: PType, + len: usize, + ) -> VortexResult { + let dtype = DType::Primitive(ptype, NonNullable); + let data = UnaryData { bits, base, max }; + Array::try_from_parts(ArrayParts::new(Unary, dtype, len, data, ArraySlots::new())) + .map(|array| array.with_stats_set(Self::stats())) + } + + /// Statistics that hold for every Unary array by construction. + fn stats() -> StatsSet { + StatsSet::of(Stat::IsSorted, Precision::Exact(true.into())) + } +} diff --git a/encodings/elias-fano/src/unary/compress.rs b/encodings/elias-fano/src/unary/compress.rs new file mode 100644 index 00000000000..95a07d537ba --- /dev/null +++ b/encodings/elias-fano/src/unary/compress.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; +use vortex_array::ArrayView; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use super::array::Unary; +use super::array::UnaryArray; +use super::array::bit_len; + +/// Encode `array` as a [`Unary`] array. +/// +/// Returns `None` when the values cannot, or should not, be stored this way: +/// - the dtype is not a non-nullable unsigned integer; +/// - a value is below the one before it; +/// - the encoding would take more bits than the plain array's `len × bit width`. This keeps sparse +/// values from growing past their input, and bounds what the encoder allocates. +pub fn unary_encode(array: ArrayView<'_, Primitive>) -> VortexResult> { + let &DType::Primitive(ptype, NonNullable) = array.dtype() else { + return Ok(None); + }; + if !ptype.is_unsigned_int() { + return Ok(None); + } + let len = array.len(); + if len == 0 { + return Unary::try_new(BitBuffer::empty(), 0, 0, ptype, 0).map(Some); + } + + match_each_unsigned_integer_ptype!(ptype, |T| { + let values = array.as_slice::(); + if !values.is_sorted() { + return Ok(None); + } + let base: u64 = values[0].as_(); + let max: u64 = values[len - 1].as_(); + let Some(bits) = bit_len(len, max - base) else { + return Ok(None); + }; + if len + .checked_mul(ptype.bit_width()) + .is_none_or(|plain| bits > plain) + { + return Ok(None); + } + + // Each position is below `bits`, which fits a `usize`, so the conversion is exact. + let positions = values.iter().enumerate().map(|(index, &value)| { + AsPrimitive::::as_(AsPrimitive::::as_(value) - base) + index + }); + Unary::try_new( + BitBuffer::from_indices(bits, positions), + base, + max, + ptype, + len, + ) + .map(Some) + }) +} + +/// Decode every value with one walk over the set bits. +pub(super) fn unary_decode(array: &UnaryArray) -> PrimitiveArray { + let len = array.len(); + let bits = array.bits(); + let base = array.base(); + + // `validate` counted exactly `len` set bits, which keeps every value within `[base, max]`, so + // the loop has nothing to check. + match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { + let mut values = BufferMut::::with_capacity(len); + let mut index = 0; + bits.for_each_set_index(|position| { + let value = base + AsPrimitive::::as_(position - index); + values.push(value.as_()); + index += 1; + }); + PrimitiveArray::new(values, Validity::NonNullable) + }) +} diff --git a/encodings/elias-fano/src/unary/mod.rs b/encodings/elias-fano/src/unary/mod.rs new file mode 100644 index 00000000000..aafb9966ab2 --- /dev/null +++ b/encodings/elias-fano/src/unary/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`Unary`] encoding: the array, its encoder, and the decode behind `execute`. + +mod array; +mod compress; +#[cfg(test)] +mod tests; + +pub use array::Unary; +pub use array::UnaryArray; +pub use array::UnaryData; +pub use compress::unary_encode; diff --git a/encodings/elias-fano/src/unary/tests.rs b/encodings/elias-fano/src/unary/tests.rs new file mode 100644 index 00000000000..b05dcf20f46 --- /dev/null +++ b/encodings/elias-fano/src/unary/tests.rs @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hash::DefaultHasher; +use std::hash::Hasher; +use std::slice; +use std::sync::LazyLock; + +use num_traits::AsPrimitive; +use prost::Message; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rstest::rstest; +use vortex_array::Array; +use vortex_array::ArrayContext; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::EqMode; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::compute::conformance::consistency::test_array_consistency; +use vortex_array::dtype::DType; +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::match_each_unsigned_integer_ptype; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::test_harness::check_metadata; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_buffer::BitBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use super::Unary; +use super::UnaryArray; +use super::array::UnaryMetadata; +use super::unary_encode; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +/// `values` as a primitive array of `ptype`, which must hold every one of them. +fn primitive(ptype: PType, values: &[u64]) -> PrimitiveArray { + match_each_unsigned_integer_ptype!(ptype, |T| { + PrimitiveArray::from_iter(values.iter().map(|&value| AsPrimitive::::as_(value))) + }) +} + +fn encode_as(ptype: PType, values: &[u64]) -> VortexResult> { + let array = primitive(ptype, values); + unary_encode(array.as_ref().as_::()) +} + +fn encode(values: &[u64]) -> VortexResult { + encode_as(PType::U64, values)? + .ok_or_else(|| vortex_err!("{} values did not encode", values.len())) +} + +/// `len` sorted values drawn from `0..=span`, with duplicates wherever they fall. +fn sorted(len: usize, span: u64, seed: u64) -> Vec { + let mut rng = StdRng::seed_from_u64(seed); + let mut values: Vec = (0..len).map(|_| rng.random_range(0..=span)).collect(); + values.sort_unstable(); + values +} + +/// Every value read back through `execute_scalar`, in a shuffled order and then in sequence. +fn check_access(array: &ArrayRef, expected: &PrimitiveArray, seed: u64) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut order: Vec = (0..array.len()).collect(); + order.shuffle(&mut StdRng::seed_from_u64(seed)); + for index in order.into_iter().chain(0..array.len()) { + assert_eq!( + array.execute_scalar(index, &mut ctx)?, + expected.execute_scalar(index, &mut ctx)?, + "index {index}" + ); + } + Ok(()) +} + +fn serde_roundtrip(array: &ArrayRef) -> VortexResult { + 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()); + } + SerializedArray::try_from(concat.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + ) +} + +fn hash_of(array: &ArrayRef) -> u64 { + let mut hasher = DefaultHasher::new(); + array.array_hash(&mut hasher, EqMode::Value); + hasher.finish() +} + +#[test] +fn test_layout_example() -> VortexResult<()> { + let array = encode(&[3, 5, 5, 9])?; + assert_eq!(array.base(), 3); + assert_eq!(array.max(), 9); + assert_eq!( + array.bits(), + &BitBuffer::from(vec![ + true, false, false, true, true, false, false, false, false, true + ]) + ); + Ok(()) +} + +#[rstest] +#[case::empty(PType::U64, vec![])] +#[case::single(PType::U64, vec![42])] +#[case::all_equal(PType::U64, vec![7; 300])] +#[case::consecutive(PType::U64, (0..300).collect())] +#[case::dense(PType::U64, sorted(500, 500, 1))] +#[case::spread(PType::U64, sorted(500, 1000, 2))] +#[case::runs(PType::U64, (0..500).map(|i| i / 64 * 3).collect())] +#[case::near_max( + PType::U64, + sorted(500, 1000, 3).into_iter().map(|v| u64::MAX - 1000 + v).collect() +)] +#[case::u8(PType::U8, sorted(300, 255, 4))] +#[case::u16(PType::U16, (0..500).map(|i| 60_000 + 2 * i).collect())] +#[case::u32( + PType::U32, + sorted(500, 1000, 5).into_iter().map(|v| u64::from(u32::MAX) - 1000 + v).collect() +)] +fn test_roundtrip(#[case] ptype: PType, #[case] values: Vec) -> VortexResult<()> { + let expected = primitive(ptype, &values); + let encoded = encode_as(ptype, &values)? + .ok_or_else(|| vortex_err!("{} values did not encode", values.len()))? + .into_array(); + assert_eq!(encoded.dtype(), &DType::Primitive(ptype, NonNullable)); + assert_arrays_eq!(encoded, expected, &mut SESSION.create_execution_ctx()); + check_access(&encoded, &expected, 0x5EED) +} + +#[test] +fn test_encode_returns_none() -> VortexResult<()> { + let inputs = [ + PrimitiveArray::from_iter([3u64, 1, 2]), + PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid), + PrimitiveArray::from_iter([1i64, 2, 3]), + PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]), + // 2 + 1,000 bits, against the plain array's 128. + PrimitiveArray::from_iter([0u64, 1000]), + ]; + for input in inputs { + assert!( + unary_encode(input.as_ref().as_::())?.is_none(), + "{} encoded", + input.dtype() + ); + } + Ok(()) +} + +#[rstest] +#[case(PType::U8)] +#[case(PType::U16)] +#[case(PType::U32)] +#[case(PType::U64)] +fn test_size_rule_boundary(#[case] ptype: PType) -> VortexResult<()> { + // Four values take `4 + span` bits, against the plain array's `4 × width`. + let span = 4 * u64::try_from(ptype.bit_width())? - 4; + assert!(encode_as(ptype, &[0, 0, 0, span])?.is_some()); + assert!(encode_as(ptype, &[0, 0, 0, span + 1])?.is_none()); + Ok(()) +} + +#[test] +fn test_try_new_rejects() -> VortexResult<()> { + let bits = |bits: &[bool]| BitBuffer::from(bits); + let rejected = [ + ( + "base above max", + Unary::try_new(bits(&[true]), 5, 3, PType::U64, 1), + ), + ( + "max above the ptype", + Unary::try_new(bits(&[true]), 0, 300, PType::U8, 1), + ), + ("signed", Unary::try_new(bits(&[true]), 0, 0, PType::I64, 1)), + ( + "nonzero offset", + Unary::try_new( + bits(&[false, false, false, true]).slice(3..), + 0, + 0, + PType::U64, + 1, + ), + ), + ( + "wrong bit length", + Unary::try_new(bits(&[true, true]), 0, 0, PType::U64, 1), + ), + ( + "first bit unset", + Unary::try_new(bits(&[false, true, true]), 0, 1, PType::U64, 2), + ), + ( + "last bit unset", + Unary::try_new(bits(&[true, true, false]), 0, 1, PType::U64, 2), + ), + ( + "empty with a span", + Unary::try_new(BitBuffer::new_unset(3), 0, 3, PType::U64, 0), + ), + // `[0, 0, 1]` takes the bits `11 01`, and `[0, 2]` takes `1 001`. Clearing or setting bit 1 + // keeps both ends set, so only the count of set bits catches them. + ( + "a set bit missing", + Unary::try_new(bits(&[true, false, false, true]), 0, 1, PType::U64, 3), + ), + ( + "an extra set bit", + Unary::try_new(bits(&[true, true, false, true]), 0, 2, PType::U64, 2), + ), + ]; + for (case, result) in rejected { + assert!(result.is_err(), "{case} was accepted"); + } + + // The vtable refuses children however the parts are put together. + let array = encode(&[1, 2])?; + let child = primitive(PType::U64, &[1]).into_array(); + assert!( + Array::try_from_parts(ArrayParts::new( + Unary, + array.dtype().clone(), + array.len(), + array.data().clone(), + [Some(child)].into_iter().collect(), + )) + .is_err() + ); + Ok(()) +} + +#[test] +fn test_deserialize_rejects() { + let dtype = DType::Primitive(PType::U64, NonNullable); + let metadata = |base: u64, max: u64| UnaryMetadata { base, max }.encode_to_vec(); + // `[0, 1]` takes the bits `1 01`. + let bits = BufferHandle::new_host(BitBuffer::from(vec![true, false, true]).inner().clone()); + let no_children: [ArrayRef; 0] = []; + let deserialize = |metadata: &[u8], buffers: &[BufferHandle], children: &[ArrayRef]| { + Unary.deserialize(&dtype, 2, metadata, buffers, &children, &SESSION) + }; + + let one = slice::from_ref(&bits); + + assert!(deserialize(&metadata(0, 1), one, &no_children).is_ok()); + assert!(deserialize(&metadata(0, 1), &[], &no_children).is_err()); + assert!(deserialize(&metadata(0, 1), &[bits.clone(), bits.clone()], &no_children).is_err()); + assert!( + deserialize( + &metadata(0, 1), + one, + &[primitive(PType::U64, &[1]).into_array()] + ) + .is_err() + ); + // The buffer holds 8 bits, short of the 102 these claim. + assert!(deserialize(&metadata(0, 100), one, &no_children).is_err()); + // 2 + `u64::MAX` bits overflow the count. + assert!(deserialize(&metadata(0, u64::MAX), one, &no_children).is_err()); + // `1 11` keeps both ends set, but holds three set bits for two values. + let extra = BufferHandle::new_host(BitBuffer::from(vec![true, true, true]).inner().clone()); + assert!(deserialize(&metadata(0, 1), slice::from_ref(&extra), &no_children).is_err()); +} + +#[test] +fn test_serde_roundtrip() -> VortexResult<()> { + let values = sorted(1000, 1500, 6); + let decoded = serde_roundtrip(&encode(&values)?.into_array())?; + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + primitive(PType::U64, &values), + &mut SESSION.create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn test_with_buffers_roundtrip() -> VortexResult<()> { + let array = encode(&sorted(1000, 1500, 7))?.into_array(); + let buffers = array.buffer_handles(); + // SAFETY: the replacement buffers are the array's own, so they hold the same values. + let rebuilt = unsafe { array.clone().with_buffers(buffers) }?; + assert!(rebuilt.is::()); + assert!(rebuilt.array_eq(&array, EqMode::Value)); + Ok(()) +} + +#[test] +fn test_array_eq_hash() -> VortexResult<()> { + let values: Vec = (0..500).map(|i| 2 * i).collect(); + let array = encode(&values)?.into_array(); + let same = encode(&values)?.into_array(); + assert!(array.array_eq(&same, EqMode::Value)); + assert_eq!(hash_of(&array), hash_of(&same)); + + let shifted: Vec = values.iter().map(|value| value + 1).collect(); + // The same base and max, with different bits. + let mut moved = values; + moved[1] = 3; + for other in [encode(&shifted)?, encode(&moved)?] { + assert!(!array.array_eq(&other.into_array(), EqMode::Value)); + } + Ok(()) +} + +#[test] +fn test_nbytes() -> VortexResult<()> { + let values = sorted(1000, 1500, 9); + let span = values[999] - values[0]; + assert_eq!( + encode(&values)?.into_array().nbytes(), + (1000 + span).div_ceil(8) + ); + Ok(()) +} + +#[cfg_attr(miri, ignore)] +#[test] +fn test_unary_metadata() { + check_metadata( + "unary.metadata", + &UnaryMetadata { + base: u64::MAX, + max: u64::MAX, + } + .encode_to_vec(), + ); +} + +#[rstest] +#[case::dense(sorted(1000, 1000, 10))] +#[case::spread(sorted(1000, 2000, 11))] +#[case::runs((0..1000).map(|i| i / 64).collect())] +#[case::single(vec![5])] +fn test_unary_consistency(#[case] values: Vec) -> VortexResult<()> { + let array = encode(&values)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + test_array_consistency(&array, &mut ctx); + if values.len() > 2 { + test_array_consistency(&array.slice(1..values.len() - 1)?, &mut ctx); + } + Ok(()) +} + +#[test] +fn test_is_sorted_stat() -> VortexResult<()> { + for values in [vec![], vec![4], sorted(500, 500, 12)] { + let array = encode(&values)?.into_array(); + assert_eq!( + array + .statistics() + .with_typed_stats_set(|stats| stats.get_as::(Stat::IsSorted)), + Precision::Exact(true), + "IsSorted over {} values", + values.len() + ); + } + Ok(()) +} + +#[test] +fn test_lazy_slice() -> VortexResult<()> { + let values = sorted(1000, 1500, 13); + let sliced = encode(&values)?.into_array().slice(100..900)?; + let expected = primitive(PType::U64, &values[100..900]); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); + check_access(&sliced, &expected, 0x51CE) +} From e96311a9ece9f80445dd23bfbd5e3844b8ad6353 Mon Sep 17 00:00:00 2001 From: rapour Date: Thu, 22 Oct 2026 21:42:42 +0330 Subject: [PATCH 2/2] chore: switch from delta unary to unary and take implicit bucketing out Signed-off-by: rapour --- encodings/elias-fano/Cargo.toml | 2 +- encodings/elias-fano/benches/unary.rs | 73 +++--- .../elias-fano/goldenfiles/unary.metadata | 2 +- encodings/elias-fano/src/lib.rs | 11 +- encodings/elias-fano/src/unary/array.rs | 237 +++++++----------- encodings/elias-fano/src/unary/bits.rs | 27 ++ encodings/elias-fano/src/unary/compress.rs | 59 ++--- encodings/elias-fano/src/unary/mod.rs | 4 +- encodings/elias-fano/src/unary/tests.rs | 221 ++++++++-------- 9 files changed, 305 insertions(+), 331 deletions(-) create mode 100644 encodings/elias-fano/src/unary/bits.rs diff --git a/encodings/elias-fano/Cargo.toml b/encodings/elias-fano/Cargo.toml index 99a44eacd6c..a61bcfa605e 100644 --- a/encodings/elias-fano/Cargo.toml +++ b/encodings/elias-fano/Cargo.toml @@ -2,7 +2,7 @@ name = "vortex-elias-fano" authors = { workspace = true } categories = { workspace = true } -description = "Vortex encodings for non-decreasing integer sequences" +description = "Vortex Elias-Fano and unary integer encodings" edition = { workspace = true } include = { workspace = true } keywords = { workspace = true } diff --git a/encodings/elias-fano/benches/unary.rs b/encodings/elias-fano/benches/unary.rs index cc7845c38f4..f9d88567b9d 100644 --- a/encodings/elias-fano/benches/unary.rs +++ b/encodings/elias-fano/benches/unary.rs @@ -4,12 +4,14 @@ //! Baselines for the naive `Unary` encoding, so each later optimisation can be measured //! against them. Keep the names and inputs stable: CodSpeed compares benchmarks by name. //! -//! Four shapes, all small enough for the encoder's size rule: +//! Each shape is what an Elias-Fano high part will hold for `LEN` sorted values: one count per +//! bucket of how many values fall in it. All four are small enough for the encoder's size rule: //! -//! * `Dense`: values drawn from `[0, LEN)`, about one unset bit per row. -//! * `Spread`: values drawn from `[0, 2 * LEN)`, the densest an Elias-Fano high part gets. -//! * `Runs`: `LEN / 64` values, 64 rows each, so set bits come in long runs. -//! * `Gap`: dense halves either side of one jump of `GAP` values, a long run of unset bits that +//! * `Dense`: values drawn from `[0, LEN)`, so about one value per bucket. +//! * `Spread`: values drawn from `[0, 2 * LEN)`, about half a value per bucket, the sparsest an +//! Elias-Fano high part gets. +//! * `Runs`: `LEN / 32` values, 32 rows each, so every count is 32 and set bits come in long runs. +//! * `Gap`: dense halves either side of one jump of `GAP` values, a long run of empty buckets that //! every scan crossing it has to walk. #![expect(clippy::unwrap_used)] @@ -48,39 +50,50 @@ enum Shape { const SHAPES: &[Shape] = &[Shape::Dense, Shape::Spread, Shape::Runs, Shape::Gap]; -/// Rows per array, small enough that a whole decode stays inside the 1 ms per-iteration budget. +/// Sorted values behind each shape, few enough that a whole decode stays inside the 1 ms +/// per-iteration budget. const LEN: usize = 1 << 16; /// The jump in `Gap`, in values. -const GAP: u64 = 4 * LEN as u64; +const GAP: usize = LEN; -/// Rows per value in `Runs`. -const ROWS_PER_VALUE: usize = 64; +/// Rows per value in `Runs`. At 64, each count would take 65 bits, past the size rule. +const ROWS_PER_VALUE: usize = 32; /// Probes per iteration. Each is a linear scan from the first bit, so 256 keep an iteration near /// the budget, as in `vortex-array`'s probe benchmark. const PROBES: usize = 256; -fn values(shape: Shape) -> Vec { +/// How many of `values` fall in each of `buckets` buckets, one value wide. +fn bucket_counts(values: impl Iterator, buckets: usize) -> Vec { + let mut counts = vec![0; buckets]; + for value in values { + counts[value] += 1; + } + counts +} + +fn counts(shape: Shape) -> Vec { let mut rng = StdRng::seed_from_u64(0); - let len = LEN as u64; - let mut out: Vec = match shape { - Shape::Dense => (0..LEN).map(|_| rng.random_range(0..len)).collect(), - Shape::Spread => (0..LEN).map(|_| rng.random_range(0..2 * len)).collect(), - Shape::Runs => (0..LEN).map(|row| (row / ROWS_PER_VALUE) as u64).collect(), - Shape::Gap => (0..LEN) - .map(|row| { + match shape { + Shape::Dense => bucket_counts((0..LEN).map(|_| rng.random_range(0..LEN)), LEN), + Shape::Spread => bucket_counts((0..LEN).map(|_| rng.random_range(0..2 * LEN)), 2 * LEN), + Shape::Runs => bucket_counts( + (0..LEN).map(|row| row / ROWS_PER_VALUE), + LEN / ROWS_PER_VALUE, + ), + Shape::Gap => bucket_counts( + (0..LEN).map(|row| { let offset = if row < LEN / 2 { 0 } else { GAP }; - offset + rng.random_range(0..len / 2) - }) - .collect(), - }; - out.sort_unstable(); - out + offset + rng.random_range(0..LEN / 2) + }), + GAP + LEN / 2, + ), + } } fn primitive(shape: Shape) -> PrimitiveArray { - PrimitiveArray::from_iter(values(shape)) + PrimitiveArray::from_iter(counts(shape)) } fn encoded(shape: Shape) -> ArrayRef { @@ -98,7 +111,7 @@ fn encode(bencher: Bencher, shape: Shape) { }); } -/// The whole array, in one walk over the set bits. +/// The whole array, in one walk over the unset bits. #[divan::bench(args = SHAPES)] fn decode(bencher: Bencher, shape: Shape) { let array = encoded(shape); @@ -114,7 +127,9 @@ fn decode(bencher: Bencher, shape: Shape) { fn scalar_at_random(bencher: Bencher, shape: Shape) { let array = encoded(shape); let mut rng = StdRng::seed_from_u64(1); - let indices: Vec = (0..PROBES).map(|_| rng.random_range(0..LEN)).collect(); + let indices: Vec = (0..PROBES) + .map(|_| rng.random_range(0..array.len())) + .collect(); bencher .with_inputs(|| SESSION.create_execution_ctx()) .bench_local_values(|mut ctx| { @@ -128,10 +143,11 @@ fn scalar_at_random(bencher: Bencher, shape: Shape) { #[divan::bench(args = SHAPES)] fn scalar_at_sequential(bencher: Bencher, shape: Shape) { let array = encoded(shape); + let start = array.len() / 2; bencher .with_inputs(|| SESSION.create_execution_ctx()) .bench_local_values(|mut ctx| { - for index in LEN / 2..LEN / 2 + PROBES { + for index in start..start + PROBES { divan::black_box(array.execute_scalar(index, &mut ctx).unwrap()); } }); @@ -141,10 +157,11 @@ fn scalar_at_sequential(bencher: Bencher, shape: Shape) { #[divan::bench(args = SHAPES)] fn slice_decode(bencher: Bencher, shape: Shape) { let array = encoded(shape); + let len = array.len(); bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .bench_local_values(|(array, mut ctx)| { - let sliced = array.slice(LEN / 4..3 * LEN / 4).unwrap(); + let sliced = array.slice(len / 4..3 * len / 4).unwrap(); divan::black_box(sliced.execute::(&mut ctx).unwrap()); }); } diff --git a/encodings/elias-fano/goldenfiles/unary.metadata b/encodings/elias-fano/goldenfiles/unary.metadata index d3da94c6217..1c2f041a6bc 100644 --- a/encodings/elias-fano/goldenfiles/unary.metadata +++ b/encodings/elias-fano/goldenfiles/unary.metadata @@ -1 +1 @@ -ÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿÿ \ No newline at end of file +ÿÿÿÿÿÿÿÿÿ \ No newline at end of file diff --git a/encodings/elias-fano/src/lib.rs b/encodings/elias-fano/src/lib.rs index 175f15ae391..c9c1401a9fd 100644 --- a/encodings/elias-fano/src/lib.rs +++ b/encodings/elias-fano/src/lib.rs @@ -1,15 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Encodings for non-decreasing integer sequences. +//! Elias-Fano and unary encodings for integers. //! -//! [`Unary`] stores a non-decreasing, non-nullable unsigned column as one bit array, writing -//! each value's gap from the previous one in unary. It takes `len + (max - min)` bits, so it suits -//! dense sequences, where values rise by about one per row or less. It is also meant to hold the -//! high bits of the Elias-Fano encoding this crate is being built towards. +//! [`Unary`] stores a non-nullable unsigned column as one bit array, writing each value in unary. +//! It takes `len + num_ones` bits, where `num_ones` is the sum of the values, so it suits small +//! values in any order. It is also meant to hold the bucket counts of the high bits of the +//! Elias-Fano encoding this crate is being built towards. //! //! None of these encodings is in a file edition or registered by `vortex-file`, so the default //! writers reject them for now. [`initialize`] registers them in a session. +#![deny(missing_docs)] mod unary; diff --git a/encodings/elias-fano/src/unary/array.rs b/encodings/elias-fano/src/unary/array.rs index 072b47c5dc9..b745158568d 100644 --- a/encodings/elias-fano/src/unary/array.rs +++ b/encodings/elias-fano/src/unary/array.rs @@ -3,7 +3,6 @@ use std::fmt::Display; use std::fmt::Formatter; -use std::hash::Hash; use std::hash::Hasher; use num_traits::AsPrimitive; @@ -24,19 +23,15 @@ 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; -use vortex_array::expr::stats::Stat; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; 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::BitBuffer; use vortex_buffer::ByteBuffer; -use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -46,6 +41,7 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use super::bits::for_each_unset_index; use super::compress::unary_decode; /// A [`Unary`]-encoded Vortex array. @@ -54,90 +50,75 @@ pub type UnaryArray = Array; /// Metadata stored alongside the `bits` buffer. #[derive(Clone, prost::Message)] pub(super) struct UnaryMetadata { - /// The first value. + /// The number of set bits, which is also the sum of the values. With the array's length it + /// gives the exact bit length, which the buffer's bytes alone do not. #[prost(uint64, tag = "1")] - pub(super) base: u64, - /// The last value. - #[prost(uint64, tag = "2")] - pub(super) max: u64, + pub(super) num_ones: u64, } -/// The parts of a [`Unary`] array: its bits, and its first and last values. +/// The parts of a [`Unary`] array: its bits. #[derive(Clone, Debug)] pub struct UnaryData { bits: BitBuffer, - base: u64, - max: u64, } impl Display for UnaryData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "base: {}, max: {}", self.base, self.max) + write!(f, "bits: {}", self.bits.len()) } } impl ArrayHash for UnaryData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { - self.base.hash(state); - self.max.hash(state); self.bits.array_hash(state, accuracy); } } impl ArrayEq for UnaryData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.base == other.base - && self.max == other.max - && self.bits.array_eq(&other.bits, accuracy) + self.bits.array_eq(&other.bits, accuracy) } } -/// The number of bits `len` values spanning `span` take, if a `usize` can count them. -pub(super) fn bit_len(len: usize, span: u64) -> Option { +/// The number of bits `len` values take with `num_ones` set bits, if a `usize` can count them. +pub(super) fn bit_len(len: usize, num_ones: u64) -> Option { u64::try_from(len) .ok()? - .checked_add(span) + .checked_add(num_ones) .and_then(|bits| usize::try_from(bits).ok()) } -fn too_many_bits(len: usize, span: u64) -> VortexError { - vortex_err!("UnaryArray of {len} values spanning {span} needs more bits than a usize counts") +/// The number of set bits in `array`, which `validate` counted as the bits past its length. +fn num_ones_of(array: ArrayView<'_, Unary>) -> VortexResult { + Ok(u64::try_from(array.bits().len() - array.len())?) } impl UnaryData { - /// Rebuild the parts from a stored buffer. The buffer's size is checked first, because - /// [`BitBuffer::new`] panics on one too short for its bits. - fn from_buffer( - buffer: ByteBuffer, - base: u64, - max: u64, - dtype: &DType, - len: usize, - ) -> VortexResult { - let span = max - .checked_sub(base) - .ok_or_else(|| vortex_err!("UnaryArray base {base} exceeds max {max}"))?; - let bits = bit_len(len, span).ok_or_else(|| too_many_bits(len, span))?; + /// Rebuild the parts from a stored buffer. Only the buffer's size is checked here, because + /// [`BitBuffer::new`] panics on one too short for its bits. [`Array::try_from_parts`] validates + /// the rest. + fn from_buffer(buffer: ByteBuffer, num_ones: u64, len: usize) -> VortexResult { + let bits = bit_len(len, num_ones).ok_or_else(|| { + vortex_err!( + "UnaryArray of {len} values with {num_ones} set bits overflows a usize bit count" + ) + })?; vortex_ensure!( buffer.len().checked_mul(8).is_some_and(|held| held >= bits), "UnaryArray needs {bits} bits, but its buffer holds {} bytes", buffer.len() ); - let data = Self { + Ok(Self { bits: BitBuffer::new(buffer, bits), - base, - max, - }; - data.validate(dtype, len)?; - Ok(data) + }) } /// Check the parts against an array of `len` values of `dtype`. /// - /// Besides the metadata and the first and last bits, this counts the set bits, one pass over - /// the buffer. Any bits that pass then describe a valid array: with exactly `len` set bits - /// among `len + span`, the `k`-th sits in `[k, k + span]`, so every value lies within - /// `[base, max]`. The readers rely on this and check no value themselves. + /// Besides the dtype, the bit length and the last bit, this counts the set bits, one pass over + /// the buffer. When the values could sum past the ptype's maximum, it also walks the unset bits + /// once to check that each value fits. Any bits that pass then describe a valid array, and the + /// readers rely on this and check no value themselves. fn validate(&self, dtype: &DType, len: usize) -> VortexResult<()> { let DType::Primitive(ptype, NonNullable) = dtype else { vortex_bail!("UnaryArray requires a non-nullable unsigned integer dtype, got {dtype}"); @@ -146,40 +127,29 @@ impl UnaryData { ptype.is_unsigned_int(), "UnaryArray requires a non-nullable unsigned integer dtype, got {dtype}" ); + vortex_ensure_eq!(self.bits.offset(), 0); + let num_ones = self.bits.len().checked_sub(len).ok_or_else(|| { + vortex_err!( + "UnaryArray of {len} values has only {} bits", + self.bits.len() + ) + })?; + // An unset last bit closes the last value, and leaves an empty array no bits at all. vortex_ensure!( - self.base <= self.max, - "UnaryArray base {} exceeds max {}", - self.base, - self.max - ); - vortex_ensure!( - self.max <= ptype.max_value_as_u64(), - "UnaryArray max {} does not fit {ptype}", - self.max - ); - let span = self.span(); - let bits = bit_len(len, span).ok_or_else(|| too_many_bits(len, span))?; - vortex_ensure!( - self.bits.offset() == 0, - "UnaryArray bits must start at offset 0, got {}", - self.bits.offset() + self.bits.is_empty() || !self.bits.value(self.bits.len() - 1), + "UnaryArray bits must end with an unset bit" ); - vortex_ensure_eq!(self.bits.len(), bits); - if len == 0 { - vortex_ensure!( - span == 0, - "An empty UnaryArray must have its base equal to its max, got a span of {span}" - ); - } else { - // Each value ends in a set bit, so these make `base` and `max` the real first and last - // values. - vortex_ensure!( - self.bits.value(0) && self.bits.value(bits - 1), - "UnaryArray bits must start and end with a set bit" - ); + // The linear checks run last. `num_ones` set bits among `len + num_ones` leave exactly + // `len` unset, one closing each value. + vortex_ensure_eq!(self.bits.true_count(), num_ones); + let max = ptype.max_value_as_u64(); + if u64::try_from(num_ones)? > max { + let mut fits = true; + self.for_each_value(|value| { + fits &= u64::try_from(value).is_ok_and(|value| value <= max); + }); + vortex_ensure!(fits, "UnaryArray holds a value that does not fit {ptype}"); } - // The only linear check, so it runs last. - vortex_ensure_eq!(self.bits.true_count(), len); Ok(()) } @@ -187,31 +157,31 @@ impl UnaryData { &self.bits } - /// The first value. - pub(super) fn base(&self) -> u64 { - self.base + /// Call `f` with every value, in order, from one walk over the unset bits. + pub(super) fn for_each_value(&self, mut f: impl FnMut(usize)) { + let mut start = 0; + for_each_unset_index(&self.bits, |end| { + f(end - start); + start = end + 1; + }); } - /// The last value. - pub(super) fn max(&self) -> u64 { - self.max - } - - /// `max - base`, which is also the number of unset bits. - pub(super) fn span(&self) -> u64 { - self.max - self.base - } - - /// The value at `index`, found by scanning for its set bit from the start. + /// The value at `index`, found by scanning from the first bit to the unset bit before it, then + /// on to its own closing bit. pub(super) fn value_at(&self, index: usize) -> VortexResult { - let position = self + let missing = || vortex_err!("UnaryArray has no unset bit for index {index}"); + let start = match index.checked_sub(1) { + Some(previous) => self.bits.select_unset(previous).ok_or_else(missing)? + 1, + None => 0, + }; + // The value is the run of set bits from `start` to the next unset bit, so this continues + // from `start` rather than scanning from the first bit again. + let value = self .bits - .select(index) - .ok_or_else(|| vortex_err!("UnaryArray has no set bit for index {index}"))?; - // The `index`-th set bit has `index` set bits before it, so this cannot underflow, and - // `validate` keeps it within `span`, so the sum cannot pass `max`. - let delta = u64::try_from(position - index)?; - Ok(self.base + delta) + .slice(start..) + .select_unset(0) + .ok_or_else(missing)?; + Ok(u64::try_from(value)?) } } @@ -233,11 +203,7 @@ impl VTable for Unary { len: usize, slots: &[Option], ) -> VortexResult<()> { - vortex_ensure!( - slots.is_empty(), - "UnaryArray has no children, got {}", - slots.len() - ); + vortex_ensure_eq!(slots.len(), 0); data.validate(dtype, len) } @@ -267,9 +233,7 @@ impl VTable for Unary { vortex_ensure_eq!(buffers.len(), 1); let data = UnaryData::from_buffer( buffers[0].try_to_host_sync()?, - array.base(), - array.max(), - array.dtype(), + num_ones_of(array)?, array.len(), )?; Ok(ArrayParts::new( @@ -287,8 +251,7 @@ impl VTable for Unary { ) -> VortexResult>> { Ok(Some( UnaryMetadata { - base: array.base(), - max: array.max(), + num_ones: num_ones_of(array)?, } .encode_to_vec(), )) @@ -303,24 +266,10 @@ impl VTable for Unary { children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "UnaryArray expects 1 buffer, got {}", - buffers.len() - ); - vortex_ensure!( - children.is_empty(), - "UnaryArray expects 0 children, got {}", - children.len() - ); + vortex_ensure_eq!(buffers.len(), 1); + vortex_ensure_eq!(children.len(), 0); let metadata = UnaryMetadata::decode(metadata)?; - let data = UnaryData::from_buffer( - buffers[0].try_to_host_sync()?, - metadata.base, - metadata.max, - dtype, - len, - )?; + let data = UnaryData::from_buffer(buffers[0].try_to_host_sync()?, metadata.num_ones, len)?; Ok(ArrayParts::new( self.clone(), dtype.clone(), @@ -348,7 +297,7 @@ impl OperationsVTable for Unary { _ctx: &mut ExecutionCtx, ) -> VortexResult { let value = array.value_at(index)?; - // `validate` keeps `max` within the ptype and every value within `max`. + // `validate` keeps every value within the ptype. Ok(match_each_unsigned_integer_ptype!( array.dtype().as_ptype(), |T| { Scalar::primitive::(value.as_(), NonNullable) } @@ -362,14 +311,14 @@ impl ValidityVTable for Unary { } } -/// Unary encoding for non-decreasing, non-nullable unsigned integers. +/// Unary encoding for non-nullable unsigned integers, in any order. /// -/// The values are stored as one bit array. Each value's gap from the one before it is written in -/// unary, as that many unset bits followed by a set bit. The first value is the array's `base`, so -/// its gap is zero. The set bit of the value at index `i` therefore sits at `(value - base) + i`, -/// and the array takes `len + (max - base)` bits. +/// The values are stored as one bit array. Each value `v` is written in unary, as `v` set bits +/// followed by an unset bit that closes it, so the array takes `len + num_ones` bits, where +/// `num_ones`, the number of set bits, is the sum of the values. The value at index `k` is the +/// number of set bits between the `k`-th unset bit and the one before it. /// -/// For `[3, 5, 5, 9]` the base is 3 and the bits are `1 001 1 00001`. +/// For `[2, 0, 3]` the bits are `110 0 1110`. /// /// Reads are deliberately simple for now: `scalar_at` scans from the first bit, and a slice /// decodes the whole array. @@ -377,24 +326,16 @@ impl ValidityVTable for Unary { pub struct Unary; impl Unary { - /// Assemble an array of `len` values from its bits and its first and last values. + /// Assemble an array of `len` values from its bits. /// /// Prefer [`unary_encode`](crate::unary_encode) unless you already hold the bits. - pub fn try_new( - bits: BitBuffer, - base: u64, - max: u64, - ptype: PType, - len: usize, - ) -> VortexResult { + pub fn try_new(bits: BitBuffer, ptype: PType, len: usize) -> VortexResult { let dtype = DType::Primitive(ptype, NonNullable); - let data = UnaryData { bits, base, max }; + // `shrink_offset` also trims the bytes past the last bit, which `buffer` would otherwise + // export with the array. + let data = UnaryData { + bits: bits.shrink_offset(), + }; Array::try_from_parts(ArrayParts::new(Unary, dtype, len, data, ArraySlots::new())) - .map(|array| array.with_stats_set(Self::stats())) - } - - /// Statistics that hold for every Unary array by construction. - fn stats() -> StatsSet { - StatsSet::of(Stat::IsSorted, Precision::Exact(true.into())) } } diff --git a/encodings/elias-fano/src/unary/bits.rs b/encodings/elias-fano/src/unary/bits.rs new file mode 100644 index 00000000000..e0ec80bda4c --- /dev/null +++ b/encodings/elias-fano/src/unary/bits.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A naive walk over unset bits, which [`BitBuffer`] only offers for set bits. + +use vortex_buffer::BitBuffer; + +/// Call `f` with the index of every unset bit, in ascending order. +pub(super) fn for_each_unset_index(bits: &BitBuffer, mut f: impl FnMut(usize)) { + let len = bits.len(); + let mut start = 0; + // The remainder word comes last, zero padded even when it holds no bits, so its padding comes + // out set once inverted and has to be masked off. + for chunk in bits.chunks().iter_padded() { + let held = len - start; + let mut word = if held < 64 { + !chunk & ((1 << held) - 1) + } else { + !chunk + }; + while word != 0 { + f(start + word.trailing_zeros() as usize); + word &= word - 1; + } + start += 64; + } +} diff --git a/encodings/elias-fano/src/unary/compress.rs b/encodings/elias-fano/src/unary/compress.rs index 95a07d537ba..a7e826a9a61 100644 --- a/encodings/elias-fano/src/unary/compress.rs +++ b/encodings/elias-fano/src/unary/compress.rs @@ -9,7 +9,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::validity::Validity; -use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -21,9 +21,10 @@ use super::array::bit_len; /// /// Returns `None` when the values cannot, or should not, be stored this way: /// - the dtype is not a non-nullable unsigned integer; -/// - a value is below the one before it; -/// - the encoding would take more bits than the plain array's `len × bit width`. This keeps sparse -/// values from growing past their input, and bounds what the encoder allocates. +/// - the values sum past `u64::MAX`; +/// - the encoding would take more bits than the plain array's `len × bit width`, which happens +/// once the values average more than `bit width - 1`. This keeps large values from growing past +/// their input, and bounds what the encoder allocates. pub fn unary_encode(array: ArrayView<'_, Primitive>) -> VortexResult> { let &DType::Primitive(ptype, NonNullable) = array.dtype() else { return Ok(None); @@ -32,18 +33,16 @@ pub fn unary_encode(array: ArrayView<'_, Primitive>) -> VortexResult(); - if !values.is_sorted() { + let Some(num_ones) = values + .iter() + .try_fold(0u64, |sum, &value| sum.checked_add(value.as_())) + else { return Ok(None); - } - let base: u64 = values[0].as_(); - let max: u64 = values[len - 1].as_(); - let Some(bits) = bit_len(len, max - base) else { + }; + let Some(bits) = bit_len(len, num_ones) else { return Ok(None); }; if len @@ -53,37 +52,27 @@ pub fn unary_encode(array: ArrayView<'_, Primitive>) -> VortexResult::as_(AsPrimitive::::as_(value) - base) + index - }); - Unary::try_new( - BitBuffer::from_indices(bits, positions), - base, - max, - ptype, - len, - ) - .map(Some) + // Each value is at most `num_ones`, below `bits`, which fits a `usize`, so the conversion + // is exact. + let mut buffer = BitBufferMut::new_set(bits); + let mut end = 0; + for &value in values { + end += AsPrimitive::::as_(value); + buffer.unset(end); + end += 1; + } + Unary::try_new(buffer.freeze(), ptype, len).map(Some) }) } -/// Decode every value with one walk over the set bits. +/// Decode every value with one walk over the unset bits. pub(super) fn unary_decode(array: &UnaryArray) -> PrimitiveArray { let len = array.len(); - let bits = array.bits(); - let base = array.base(); - // `validate` counted exactly `len` set bits, which keeps every value within `[base, max]`, so - // the loop has nothing to check. + // `validate` keeps every value within the ptype, so the loop has nothing to check. match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { let mut values = BufferMut::::with_capacity(len); - let mut index = 0; - bits.for_each_set_index(|position| { - let value = base + AsPrimitive::::as_(position - index); - values.push(value.as_()); - index += 1; - }); + array.for_each_value(|value| values.push(AsPrimitive::::as_(value))); PrimitiveArray::new(values, Validity::NonNullable) }) } diff --git a/encodings/elias-fano/src/unary/mod.rs b/encodings/elias-fano/src/unary/mod.rs index aafb9966ab2..71f6f4d09f9 100644 --- a/encodings/elias-fano/src/unary/mod.rs +++ b/encodings/elias-fano/src/unary/mod.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`Unary`] encoding: the array, its encoder, and the decode behind `execute`. +//! The [`Unary`] encoding: the array, its encoder, the decode behind `execute`, and the walk over +//! unset bits that decoding and validation use. mod array; +mod bits; mod compress; #[cfg(test)] mod tests; diff --git a/encodings/elias-fano/src/unary/tests.rs b/encodings/elias-fano/src/unary/tests.rs index b05dcf20f46..b80394736dd 100644 --- a/encodings/elias-fano/src/unary/tests.rs +++ b/encodings/elias-fano/src/unary/tests.rs @@ -30,9 +30,6 @@ use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DType; 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::match_each_unsigned_integer_ptype; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; @@ -50,6 +47,7 @@ use vortex_session::registry::ReadContext; use super::Unary; use super::UnaryArray; use super::array::UnaryMetadata; +use super::bits::for_each_unset_index; use super::unary_encode; static SESSION: LazyLock = LazyLock::new(|| { @@ -75,12 +73,10 @@ fn encode(values: &[u64]) -> VortexResult { .ok_or_else(|| vortex_err!("{} values did not encode", values.len())) } -/// `len` sorted values drawn from `0..=span`, with duplicates wherever they fall. -fn sorted(len: usize, span: u64, seed: u64) -> Vec { +/// `len` values drawn from `0..=max`, in no particular order. +fn small(len: usize, max: u64, seed: u64) -> Vec { let mut rng = StdRng::seed_from_u64(seed); - let mut values: Vec = (0..len).map(|_| rng.random_range(0..=span)).collect(); - values.sort_unstable(); - values + (0..len).map(|_| rng.random_range(0..=max)).collect() } /// Every value read back through `execute_scalar`, in a shuffled order and then in sequence. @@ -120,14 +116,10 @@ fn hash_of(array: &ArrayRef) -> u64 { #[test] fn test_layout_example() -> VortexResult<()> { - let array = encode(&[3, 5, 5, 9])?; - assert_eq!(array.base(), 3); - assert_eq!(array.max(), 9); + let array = encode(&[2, 0, 3])?; assert_eq!( array.bits(), - &BitBuffer::from(vec![ - true, false, false, true, true, false, false, false, false, true - ]) + &BitBuffer::from(vec![true, true, false, false, true, true, true, false]) ); Ok(()) } @@ -135,21 +127,18 @@ fn test_layout_example() -> VortexResult<()> { #[rstest] #[case::empty(PType::U64, vec![])] #[case::single(PType::U64, vec![42])] +#[case::all_zeros(PType::U64, vec![0; 300])] #[case::all_equal(PType::U64, vec![7; 300])] -#[case::consecutive(PType::U64, (0..300).collect())] -#[case::dense(PType::U64, sorted(500, 500, 1))] -#[case::spread(PType::U64, sorted(500, 1000, 2))] +#[case::unsorted(PType::U64, small(500, 3, 1))] +#[case::decreasing(PType::U64, (0..60).rev().collect())] #[case::runs(PType::U64, (0..500).map(|i| i / 64 * 3).collect())] -#[case::near_max( - PType::U64, - sorted(500, 1000, 3).into_iter().map(|v| u64::MAX - 1000 + v).collect() -)] -#[case::u8(PType::U8, sorted(300, 255, 4))] -#[case::u16(PType::U16, (0..500).map(|i| 60_000 + 2 * i).collect())] -#[case::u32( - PType::U32, - sorted(500, 1000, 5).into_iter().map(|v| u64::from(u32::MAX) - 1000 + v).collect() -)] +#[case::u8(PType::U8, small(300, 7, 2))] +#[case::u8_max(PType::U8, (0..41).map(|i| if i == 20 { 255 } else { 0 }).collect())] +#[case::u16(PType::U16, small(500, 20, 3))] +#[case::u32(PType::U32, small(500, 40, 4))] +// A value of 1,000 that starts 4 bits into a byte, so reading it scans whole 512-bit chunks from an +// offset. +#[case::long_value(PType::U16, [vec![3, 1000], vec![0; 68]].concat())] fn test_roundtrip(#[case] ptype: PType, #[case] values: Vec) -> VortexResult<()> { let expected = primitive(ptype, &values); let encoded = encode_as(ptype, &values)? @@ -163,12 +152,13 @@ fn test_roundtrip(#[case] ptype: PType, #[case] values: Vec) -> VortexResul #[test] fn test_encode_returns_none() -> VortexResult<()> { let inputs = [ - PrimitiveArray::from_iter([3u64, 1, 2]), PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid), PrimitiveArray::from_iter([1i64, 2, 3]), PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]), // 2 + 1,000 bits, against the plain array's 128. PrimitiveArray::from_iter([0u64, 1000]), + // The sum overflows a `u64`. + PrimitiveArray::from_iter([u64::MAX, 1]), ]; for input in inputs { assert!( @@ -186,10 +176,10 @@ fn test_encode_returns_none() -> VortexResult<()> { #[case(PType::U32)] #[case(PType::U64)] fn test_size_rule_boundary(#[case] ptype: PType) -> VortexResult<()> { - // Four values take `4 + span` bits, against the plain array's `4 × width`. - let span = 4 * u64::try_from(ptype.bit_width())? - 4; - assert!(encode_as(ptype, &[0, 0, 0, span])?.is_some()); - assert!(encode_as(ptype, &[0, 0, 0, span + 1])?.is_none()); + // Four values take `4 + num_ones` bits, against the plain array's `4 × width`. + let num_ones = 4 * u64::try_from(ptype.bit_width())? - 4; + assert!(encode_as(ptype, &[0, 0, 0, num_ones])?.is_some()); + assert!(encode_as(ptype, &[0, 0, 0, num_ones + 1])?.is_none()); Ok(()) } @@ -197,50 +187,37 @@ fn test_size_rule_boundary(#[case] ptype: PType) -> VortexResult<()> { fn test_try_new_rejects() -> VortexResult<()> { let bits = |bits: &[bool]| BitBuffer::from(bits); let rejected = [ - ( - "base above max", - Unary::try_new(bits(&[true]), 5, 3, PType::U64, 1), - ), - ( - "max above the ptype", - Unary::try_new(bits(&[true]), 0, 300, PType::U8, 1), - ), - ("signed", Unary::try_new(bits(&[true]), 0, 0, PType::I64, 1)), + ("signed", Unary::try_new(bits(&[false]), PType::I64, 1)), ( "nonzero offset", - Unary::try_new( - bits(&[false, false, false, true]).slice(3..), - 0, - 0, - PType::U64, - 1, - ), - ), - ( - "wrong bit length", - Unary::try_new(bits(&[true, true]), 0, 0, PType::U64, 1), + Unary::try_new(bits(&[true, true, true, false]).slice(3..), PType::U64, 1), ), ( - "first bit unset", - Unary::try_new(bits(&[false, true, true]), 0, 1, PType::U64, 2), + "fewer bits than values", + Unary::try_new(bits(&[false, false]), PType::U64, 3), ), ( - "last bit unset", - Unary::try_new(bits(&[true, true, false]), 0, 1, PType::U64, 2), + "last bit set", + Unary::try_new(bits(&[false, true]), PType::U64, 1), ), ( - "empty with a span", - Unary::try_new(BitBuffer::new_unset(3), 0, 3, PType::U64, 0), + "empty with bits", + Unary::try_new(bits(&[true]), PType::U64, 0), ), - // `[0, 0, 1]` takes the bits `11 01`, and `[0, 2]` takes `1 001`. Clearing or setting bit 1 - // keeps both ends set, so only the count of set bits catches them. + // `[1, 0]` takes the bits `10 0`. Clearing bit 0 or setting bit 1 keeps the last bit unset, + // so only the count of set bits catches them. ( "a set bit missing", - Unary::try_new(bits(&[true, false, false, true]), 0, 1, PType::U64, 3), + Unary::try_new(bits(&[false, false, false]), PType::U64, 2), ), ( "an extra set bit", - Unary::try_new(bits(&[true, true, false, true]), 0, 2, PType::U64, 2), + Unary::try_new(bits(&[true, true, false]), PType::U64, 2), + ), + // 256 set bits and one unset: a single value of 256. + ( + "a value above the ptype", + Unary::try_new((0..257).map(|bit| bit < 256).collect(), PType::U8, 1), ), ]; for (case, result) in rejected { @@ -266,39 +243,42 @@ fn test_try_new_rejects() -> VortexResult<()> { #[test] fn test_deserialize_rejects() { let dtype = DType::Primitive(PType::U64, NonNullable); - let metadata = |base: u64, max: u64| UnaryMetadata { base, max }.encode_to_vec(); - // `[0, 1]` takes the bits `1 01`. - let bits = BufferHandle::new_host(BitBuffer::from(vec![true, false, true]).inner().clone()); + let metadata = |num_ones: u64| UnaryMetadata { num_ones }.encode_to_vec(); + // `[1, 0]` takes the bits `10 0`. + let bits = BufferHandle::new_host(BitBuffer::from(vec![true, false, false]).inner().clone()); let no_children: [ArrayRef; 0] = []; + // `try_from_parts` validates the parts, as the array plugin does after `deserialize`. let deserialize = |metadata: &[u8], buffers: &[BufferHandle], children: &[ArrayRef]| { - Unary.deserialize(&dtype, 2, metadata, buffers, &children, &SESSION) + Unary + .deserialize(&dtype, 2, metadata, buffers, &children, &SESSION) + .and_then(Array::try_from_parts) }; let one = slice::from_ref(&bits); - assert!(deserialize(&metadata(0, 1), one, &no_children).is_ok()); - assert!(deserialize(&metadata(0, 1), &[], &no_children).is_err()); - assert!(deserialize(&metadata(0, 1), &[bits.clone(), bits.clone()], &no_children).is_err()); + assert!(deserialize(&metadata(1), one, &no_children).is_ok()); + assert!(deserialize(&metadata(1), &[], &no_children).is_err()); + assert!(deserialize(&metadata(1), &[bits.clone(), bits.clone()], &no_children).is_err()); assert!( deserialize( - &metadata(0, 1), + &metadata(1), one, &[primitive(PType::U64, &[1]).into_array()] ) .is_err() ); // The buffer holds 8 bits, short of the 102 these claim. - assert!(deserialize(&metadata(0, 100), one, &no_children).is_err()); + assert!(deserialize(&metadata(100), one, &no_children).is_err()); // 2 + `u64::MAX` bits overflow the count. - assert!(deserialize(&metadata(0, u64::MAX), one, &no_children).is_err()); - // `1 11` keeps both ends set, but holds three set bits for two values. - let extra = BufferHandle::new_host(BitBuffer::from(vec![true, true, true]).inner().clone()); - assert!(deserialize(&metadata(0, 1), slice::from_ref(&extra), &no_children).is_err()); + assert!(deserialize(&metadata(u64::MAX), one, &no_children).is_err()); + // `11 0` keeps the last bit unset, but holds two set bits where the metadata says one. + let extra = BufferHandle::new_host(BitBuffer::from(vec![true, true, false]).inner().clone()); + assert!(deserialize(&metadata(1), slice::from_ref(&extra), &no_children).is_err()); } #[test] fn test_serde_roundtrip() -> VortexResult<()> { - let values = sorted(1000, 1500, 6); + let values = small(1000, 7, 6); let decoded = serde_roundtrip(&encode(&values)?.into_array())?; assert!(decoded.is::()); assert_arrays_eq!( @@ -311,7 +291,7 @@ fn test_serde_roundtrip() -> VortexResult<()> { #[test] fn test_with_buffers_roundtrip() -> VortexResult<()> { - let array = encode(&sorted(1000, 1500, 7))?.into_array(); + let array = encode(&small(1000, 7, 7))?.into_array(); let buffers = array.buffer_handles(); // SAFETY: the replacement buffers are the array's own, so they hold the same values. let rebuilt = unsafe { array.clone().with_buffers(buffers) }?; @@ -322,17 +302,13 @@ fn test_with_buffers_roundtrip() -> VortexResult<()> { #[test] fn test_array_eq_hash() -> VortexResult<()> { - let values: Vec = (0..500).map(|i| 2 * i).collect(); - let array = encode(&values)?.into_array(); - let same = encode(&values)?.into_array(); + let array = encode(&[1, 0, 2])?.into_array(); + let same = encode(&[1, 0, 2])?.into_array(); assert!(array.array_eq(&same, EqMode::Value)); assert_eq!(hash_of(&array), hash_of(&same)); - let shifted: Vec = values.iter().map(|value| value + 1).collect(); - // The same base and max, with different bits. - let mut moved = values; - moved[1] = 3; - for other in [encode(&shifted)?, encode(&moved)?] { + // `[0, 1, 2]` has the same length and sum, with different bits. + for other in [encode(&[0, 1, 2])?, encode(&[1, 0, 3])?] { assert!(!array.array_eq(&other.into_array(), EqMode::Value)); } Ok(()) @@ -340,11 +316,20 @@ fn test_array_eq_hash() -> VortexResult<()> { #[test] fn test_nbytes() -> VortexResult<()> { - let values = sorted(1000, 1500, 9); - let span = values[999] - values[0]; + let values = small(1000, 7, 9); + let num_ones: u64 = values.iter().sum(); assert_eq!( encode(&values)?.into_array().nbytes(), - (1000 + span).div_ceil(8) + (1000 + num_ones).div_ceil(8) + ); + + // The value 2, as the first 3 bits of a 125-byte buffer, keeps only the byte it uses. + let prefix: BitBuffer = (0..1000).map(|bit| bit < 2).collect(); + assert_eq!( + Unary::try_new(prefix.slice(..3), PType::U64, 1)? + .into_array() + .nbytes(), + 1 ); Ok(()) } @@ -354,18 +339,15 @@ fn test_nbytes() -> VortexResult<()> { fn test_unary_metadata() { check_metadata( "unary.metadata", - &UnaryMetadata { - base: u64::MAX, - max: u64::MAX, - } - .encode_to_vec(), + &UnaryMetadata { num_ones: u64::MAX }.encode_to_vec(), ); } #[rstest] -#[case::dense(sorted(1000, 1000, 10))] -#[case::spread(sorted(1000, 2000, 11))] +#[case::small(small(1000, 3, 10))] +#[case::wide(small(1000, 40, 11))] #[case::runs((0..1000).map(|i| i / 64).collect())] +#[case::zeros(vec![0; 100])] #[case::single(vec![5])] fn test_unary_consistency(#[case] values: Vec) -> VortexResult<()> { let array = encode(&values)?.into_array(); @@ -377,27 +359,42 @@ fn test_unary_consistency(#[case] values: Vec) -> VortexResult<()> { Ok(()) } -#[test] -fn test_is_sorted_stat() -> VortexResult<()> { - for values in [vec![], vec![4], sorted(500, 500, 12)] { - let array = encode(&values)?.into_array(); - assert_eq!( - array - .statistics() - .with_typed_stats_set(|stats| stats.get_as::(Stat::IsSorted)), - Precision::Exact(true), - "IsSorted over {} values", - values.len() - ); - } - Ok(()) -} - #[test] fn test_lazy_slice() -> VortexResult<()> { - let values = sorted(1000, 1500, 13); + let values = small(1000, 7, 13); let sliced = encode(&values)?.into_array().slice(100..900)?; let expected = primitive(PType::U64, &values[100..900]); assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); check_access(&sliced, &expected, 0x51CE) } + +#[rstest] +#[case(0)] +#[case(1)] +#[case(63)] +#[case(64)] +#[case(65)] +#[case(128)] +#[case(200)] +fn test_for_each_unset_index_agrees_with_naive(#[case] len: usize) -> VortexResult<()> { + let mut rng = StdRng::seed_from_u64(u64::try_from(len)?); + let random: BitBuffer = (0..len).map(|_| rng.random_bool(0.5)).collect(); + // The same kind of bits, starting 3 bits into their buffer. + let offset: BitBuffer = (0..len + 3).map(|_| rng.random_bool(0.5)).collect(); + for bits in [ + random, + offset.slice(3..), + BitBuffer::new_set(len), + BitBuffer::new_unset(len), + ] { + let expected: Vec = bits + .iter() + .enumerate() + .filter_map(|(index, set)| (!set).then_some(index)) + .collect(); + let mut found = Vec::new(); + for_each_unset_index(&bits, |index| found.push(index)); + assert_eq!(found, expected); + } + Ok(()) +}