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..a61bcfa605e --- /dev/null +++ b/encodings/elias-fano/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "vortex-elias-fano" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex Elias-Fano and unary integer encodings" +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..f9d88567b9d --- /dev/null +++ b/encodings/elias-fano/benches/unary.rs @@ -0,0 +1,171 @@ +// 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. +//! +//! 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)`, 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)] + +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]; + +/// 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: usize = LEN; + +/// 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; + +/// 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); + 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) + }), + GAP + LEN / 2, + ), + } +} + +fn primitive(shape: Shape) -> PrimitiveArray { + PrimitiveArray::from_iter(counts(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 unset 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..array.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); + let start = array.len() / 2; + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| { + for index in start..start + 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); + 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(); + 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..1c2f041a6bc --- /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..c9c1401a9fd --- /dev/null +++ b/encodings/elias-fano/src/lib.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Elias-Fano and unary encodings for integers. +//! +//! [`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; + +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..b745158568d --- /dev/null +++ b/encodings/elias-fano/src/unary/array.rs @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +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::match_each_unsigned_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +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::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::bits::for_each_unset_index; +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 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) num_ones: u64, +} + +/// The parts of a [`Unary`] array: its bits. +#[derive(Clone, Debug)] +pub struct UnaryData { + bits: BitBuffer, +} + +impl Display for UnaryData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "bits: {}", self.bits.len()) + } +} + +impl ArrayHash for UnaryData { + fn array_hash(&self, state: &mut H, accuracy: EqMode) { + self.bits.array_hash(state, accuracy); + } +} + +impl ArrayEq for UnaryData { + fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { + self.bits.array_eq(&other.bits, accuracy) + } +} + +/// 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(num_ones) + .and_then(|bits| usize::try_from(bits).ok()) +} + +/// 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. 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() + ); + Ok(Self { + bits: BitBuffer::new(buffer, bits), + }) + } + + /// Check the parts against an array of `len` values of `dtype`. + /// + /// 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}"); + }; + vortex_ensure!( + 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.bits.is_empty() || !self.bits.value(self.bits.len() - 1), + "UnaryArray bits must end with an unset 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}"); + } + Ok(()) + } + + pub(super) fn bits(&self) -> &BitBuffer { + &self.bits + } + + /// 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 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 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 + .slice(start..) + .select_unset(0) + .ok_or_else(missing)?; + Ok(u64::try_from(value)?) + } +} + +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_eq!(slots.len(), 0); + 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()?, + num_ones_of(array)?, + 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 { + num_ones: num_ones_of(array)?, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + 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.num_ones, 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 every value within the ptype. + 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-nullable unsigned integers, in any order. +/// +/// 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 `[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. +#[derive(Clone, Debug)] +pub struct Unary; + +impl Unary { + /// 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, ptype: PType, len: usize) -> VortexResult { + let dtype = DType::Primitive(ptype, NonNullable); + // `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())) + } +} 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 new file mode 100644 index 00000000000..a7e826a9a61 --- /dev/null +++ b/encodings/elias-fano/src/unary/compress.rs @@ -0,0 +1,78 @@ +// 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::BitBufferMut; +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; +/// - 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); + }; + if !ptype.is_unsigned_int() { + return Ok(None); + } + let len = array.len(); + + match_each_unsigned_integer_ptype!(ptype, |T| { + let values = array.as_slice::(); + let Some(num_ones) = values + .iter() + .try_fold(0u64, |sum, &value| sum.checked_add(value.as_())) + else { + return Ok(None); + }; + let Some(bits) = bit_len(len, num_ones) else { + return Ok(None); + }; + if len + .checked_mul(ptype.bit_width()) + .is_none_or(|plain| bits > plain) + { + return Ok(None); + } + + // 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 unset bits. +pub(super) fn unary_decode(array: &UnaryArray) -> PrimitiveArray { + let len = array.len(); + + // `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); + 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 new file mode 100644 index 00000000000..71f6f4d09f9 --- /dev/null +++ b/encodings/elias-fano/src/unary/mod.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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; + +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..b80394736dd --- /dev/null +++ b/encodings/elias-fano/src/unary/tests.rs @@ -0,0 +1,400 @@ +// 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::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::bits::for_each_unset_index; +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` 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); + (0..len).map(|_| rng.random_range(0..=max)).collect() +} + +/// 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(&[2, 0, 3])?; + assert_eq!( + array.bits(), + &BitBuffer::from(vec![true, true, false, false, true, true, true, false]) + ); + Ok(()) +} + +#[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::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::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)? + .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::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!( + 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 + 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(()) +} + +#[test] +fn test_try_new_rejects() -> VortexResult<()> { + let bits = |bits: &[bool]| BitBuffer::from(bits); + let rejected = [ + ("signed", Unary::try_new(bits(&[false]), PType::I64, 1)), + ( + "nonzero offset", + Unary::try_new(bits(&[true, true, true, false]).slice(3..), PType::U64, 1), + ), + ( + "fewer bits than values", + Unary::try_new(bits(&[false, false]), PType::U64, 3), + ), + ( + "last bit set", + Unary::try_new(bits(&[false, true]), PType::U64, 1), + ), + ( + "empty with bits", + Unary::try_new(bits(&[true]), PType::U64, 0), + ), + // `[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(&[false, false, false]), PType::U64, 2), + ), + ( + "an extra set bit", + 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 { + 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 = |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) + .and_then(Array::try_from_parts) + }; + + let one = slice::from_ref(&bits); + + 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(1), + one, + &[primitive(PType::U64, &[1]).into_array()] + ) + .is_err() + ); + // The buffer holds 8 bits, short of the 102 these claim. + assert!(deserialize(&metadata(100), one, &no_children).is_err()); + // 2 + `u64::MAX` bits overflow the count. + 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 = small(1000, 7, 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(&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) }?; + assert!(rebuilt.is::()); + assert!(rebuilt.array_eq(&array, EqMode::Value)); + Ok(()) +} + +#[test] +fn test_array_eq_hash() -> VortexResult<()> { + 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)); + + // `[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(()) +} + +#[test] +fn test_nbytes() -> VortexResult<()> { + let values = small(1000, 7, 9); + let num_ones: u64 = values.iter().sum(); + assert_eq!( + encode(&values)?.into_array().nbytes(), + (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(()) +} + +#[cfg_attr(miri, ignore)] +#[test] +fn test_unary_metadata() { + check_metadata( + "unary.metadata", + &UnaryMetadata { num_ones: u64::MAX }.encode_to_vec(), + ); +} + +#[rstest] +#[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(); + 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_lazy_slice() -> VortexResult<()> { + 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(()) +}