From bfd3004fd160c73020594189dee832beeaa49228 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 11 Sep 2026 14:39:09 +0100 Subject: [PATCH] perf(array): reuse probe state in primitive and compressed arrays Signed-off-by: Joe Isaacs --- Cargo.lock | 3 + encodings/fastlanes/src/rle/mod.rs | 2 + encodings/fastlanes/src/rle/probe.rs | 210 ++++++++++++++ encodings/fastlanes/src/rle/probe/tests.rs | 96 +++++++ .../fastlanes/src/rle/vtable/operations.rs | 14 +- encodings/pco/Cargo.toml | 6 + encodings/pco/benches/probe.md | 122 ++++++++ encodings/pco/benches/probe.rs | 217 ++++++++++++++ encodings/pco/examples/probe.rs | 40 +++ encodings/pco/src/array.rs | 15 +- encodings/pco/src/lib.rs | 1 + encodings/pco/src/probe.rs | 153 ++++++++++ encodings/pco/src/probe/tests.rs | 268 ++++++++++++++++++ encodings/runend/src/lib.rs | 1 + encodings/runend/src/ops.rs | 13 + encodings/runend/src/probe.rs | 47 +++ encodings/runend/src/probe/tests.rs | 89 ++++++ vortex-array/PROBE_DESIGN.md | 32 ++- .../src/arrays/primitive/vtable/operations.rs | 57 +++- 19 files changed, 1378 insertions(+), 8 deletions(-) create mode 100644 encodings/fastlanes/src/rle/probe.rs create mode 100644 encodings/fastlanes/src/rle/probe/tests.rs create mode 100644 encodings/pco/benches/probe.md create mode 100644 encodings/pco/benches/probe.rs create mode 100644 encodings/pco/examples/probe.rs create mode 100644 encodings/pco/src/probe.rs create mode 100644 encodings/pco/src/probe/tests.rs create mode 100644 encodings/runend/src/probe.rs create mode 100644 encodings/runend/src/probe/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 54801630a5b..6f25bfeedf3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11359,6 +11359,7 @@ dependencies = [ name = "vortex-pco" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "pco", "prost 0.14.4", "rstest", @@ -11366,7 +11367,9 @@ dependencies = [ "vortex-arrow", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-mask", + "vortex-runend", "vortex-session", ] diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 25610742c19..37a2db13395 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -10,6 +10,8 @@ pub use array::RLESlots; mod compute; mod kernel; +mod probe; + mod vtable; pub use vtable::RLE; pub use vtable::RLEArray; diff --git a/encodings/fastlanes/src/rle/probe.rs b/encodings/fastlanes/src/rle/probe.rs new file mode 100644 index 00000000000..da8231d4bab --- /dev/null +++ b/encodings/fastlanes/src/rle/probe.rs @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! RLE probes route through the indices, offsets, and values slots. +//! +//! A slot that already holds a materialized primitive array is read directly from its typed +//! buffer. Any other child, including one with a lazy validity expression, is read through a +//! child probe that keeps its own state, so nesting works to any depth. + +use num_traits::ToPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::ProbeChildren; +use vortex_array::ProbeCtx; +use vortex_array::arrays::Bool; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::match_each_native_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::FL_CHUNK_SIZE; +use crate::RLE; +use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; +use crate::rle::RLESlots; + +/// State for repeated RLE probes: one reader per child slot plus the slice's base value offset. +#[derive(Default)] +pub struct RleProbeState<'a> { + children: Option>, +} + +struct Children<'a> { + indices: Child<'a>, + offsets: Child<'a>, + values: Child<'a>, + base: usize, +} + +enum Child<'a> { + /// A materialized primitive array with its validity resolved once. + Primitive { + view: ArrayView<'a, Primitive>, + validity: Option, + }, + /// A slot read through the context's retained child probe. + Probe(usize), +} + +impl<'a> Child<'a> { + fn new( + slots: &'a [Option], + slot: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let child = slots[slot] + .as_ref() + .ok_or_else(|| vortex_err!("RLE slot {slot} is missing"))?; + let Some(view) = child.as_opt::() else { + return Ok(Self::Probe(slot)); + }; + let validity = match PrimitiveArrayExt::validity(&view) { + Validity::NonNullable | Validity::AllValid => None, + // A lazy validity expression may fail on rows this probe never requests, so the + // primitive probe evaluates it per row instead. + Validity::Array(array) if !array.is::() => { + return Ok(Self::Probe(slot)); + } + validity => Some(validity.execute_mask(view.len(), ctx)?), + }; + Ok(Self::Primitive { view, validity }) + } + + /// Read an unsigned routing value, or `None` when the row is null. + fn index_at( + &mut self, + index: usize, + children: &mut ProbeChildren<'a>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + match self { + Self::Primitive { view, validity } => { + vortex_ensure!(index < view.len(), OutOfBounds: index, 0, view.len()); + if validity.as_ref().is_some_and(|mask| !mask.value(index)) { + return Ok(None); + } + match_each_unsigned_integer_ptype!(view.ptype(), |T| { + view.as_slice::()[index] + .to_usize() + .map(Some) + .ok_or_else(|| vortex_err!("RLE index does not fit usize")) + }) + } + Self::Probe(slot) => Ok(children + .child(*slot)? + .scalar_at(index, ctx)? + .as_primitive() + .as_::()), + } + } + + /// Read a value, tagged with the RLE array's dtype. + fn value_at( + &mut self, + index: usize, + array: ArrayView<'_, RLE>, + children: &mut ProbeChildren<'a>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + match self { + Self::Primitive { view, .. } => { + vortex_ensure!(index < view.len(), OutOfBounds: index, 0, view.len()); + Ok(match_each_native_ptype!(view.ptype(), |T| { + Scalar::primitive(view.as_slice::()[index], array.dtype().nullability()) + })) + } + Self::Probe(slot) => Scalar::try_new( + array.dtype().clone(), + children.child(*slot)?.scalar_at(index, ctx)?.into_value(), + ), + } + } +} + +pub(crate) fn scalar_at<'a>( + array: ArrayView<'a, RLE>, + index: usize, + probe: Option<&mut ProbeCtx<'a, RleProbeState<'a>>>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let logical_index = array.offset() + index; + let chunk = logical_index / FL_CHUNK_SIZE; + let Some(probe) = probe else { + // The index scalar supplies RLE's nullness and routing value in one lookup. + let code = array.indices().execute_scalar(logical_index, ctx)?; + let Some(code) = code.as_primitive().as_::() else { + return Ok(Scalar::null(array.dtype().clone())); + }; + let offset = if chunk == 0 { + 0 + } else { + let offsets = array.values_idx_offsets(); + read_offset(offsets.execute_scalar(chunk, ctx)?)? + .checked_sub(read_offset(offsets.execute_scalar(0, ctx)?)?) + .ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))? + }; + let value_index = value_index(offset, code)?; + return Scalar::try_new( + array.dtype().clone(), + array + .values() + .execute_scalar(value_index, ctx)? + .into_value(), + ); + }; + + let (state, probes) = probe.parts(); + let children = match &mut state.children { + Some(children) => children, + slot @ None => { + // Borrow from the source slots, whose lifetime is independent of the temporary view. + let slots = array.slots(); + let mut offsets = Child::new(slots, RLESlots::VALUES_IDX_OFFSETS, ctx)?; + let base = offsets + .index_at(0, probes, ctx)? + .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))?; + slot.insert(Children { + indices: Child::new(slots, RLESlots::INDICES, ctx)?, + offsets, + values: Child::new(slots, RLESlots::VALUES, ctx)?, + base, + }) + } + }; + let Some(code) = children.indices.index_at(logical_index, probes, ctx)? else { + return Ok(Scalar::null(array.dtype().clone())); + }; + let offset = children + .offsets + .index_at(chunk, probes, ctx)? + .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))? + .checked_sub(children.base) + .ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))?; + children + .values + .value_at(value_index(offset, code)?, array, probes, ctx) +} + +fn read_offset(scalar: Scalar) -> VortexResult { + scalar + .as_primitive() + .as_::() + .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize")) +} + +fn value_index(offset: usize, code: usize) -> VortexResult { + offset + .checked_add(code) + .ok_or_else(|| vortex_err!("RLE value index overflow")) +} + +#[cfg(test)] +mod tests; diff --git a/encodings/fastlanes/src/rle/probe/tests.rs b/encodings/fastlanes/src/rle/probe/tests.rs new file mode 100644 index 00000000000..a5c75c99ba6 --- /dev/null +++ b/encodings/fastlanes/src/rle/probe/tests.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::ProbeUsage; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::RLE; +use crate::RLEData; + +#[rstest] +#[case(ProbeUsage::Once, false)] +#[case(ProbeUsage::Repeated, false)] +#[case(ProbeUsage::Once, true)] +#[case(ProbeUsage::Repeated, true)] +fn random_access_across_chunks_and_nulls( + #[case] usage: ProbeUsage, + #[case] sliced: bool, +) -> VortexResult<()> { + let mut ctx = crate::test::SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..8192u32).map(|i| (i % 11 != 0).then_some(i / 16))); + let encoded = RLEData::encode(input.as_view(), &mut ctx)?.into_array(); + let range = if sliced { 1777..7333 } else { 0..8192 }; + let source = if sliced { + encoded + .slice(range.clone())? + .execute::(&mut ctx)? + } else { + encoded + }; + assert!(source.is::()); + let input = input.slice(range)?; + let indices = [0u32, 1, 1023, 1024, 2048, 2047, 4097, 11, 33, 17, 0]; + let mut actual = builder_with_capacity_in(source.dtype(), indices.len(), ctx.allocator()); + let mut probe = source.probe(usage); + for index in indices { + actual.append_scalar(&probe.scalar_at(index as usize, &mut ctx)?)?; + } + assert_arrays_eq!( + actual.finish(), + input.take(PrimitiveArray::from_iter(indices).into_array())?, + &mut ctx + ); + assert!(probe.scalar_at(source.len(), &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn primitive_slots_use_direct_readers() -> VortexResult<()> { + let mut ctx = crate::test::SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter((0..4096u32).map(|i| i / 16)); + let encoded = RLEData::encode(input.as_view(), &mut ctx)?; + for slot in [ + super::RLESlots::INDICES, + super::RLESlots::VALUES_IDX_OFFSETS, + super::RLESlots::VALUES, + ] { + let reader = super::Child::new(encoded.as_view().slots(), slot, &mut ctx)?; + assert!(matches!(reader, super::Child::Primitive { .. })); + } + Ok(()) +} + +#[test] +fn lazy_validity_does_not_evaluate_unrequested_rows() -> VortexResult<()> { + let mut ctx = crate::test::SESSION.create_execution_ctx(); + let numerators = PrimitiveArray::from_iter(vec![1u32; 1024]).into_array(); + let denominators = + PrimitiveArray::from_iter((0..1024).map(|i| u32::from(i != 1023))).into_array(); + let validity = numerators + .binary(denominators, Operator::Div)? + .binary(numerators, Operator::Eq)?; + let array = RLE::try_new( + PrimitiveArray::from_iter([42u32]).into_array(), + PrimitiveArray::new(vec![0u16; 1024], Validity::Array(validity)).into_array(), + PrimitiveArray::from_iter([0u64]).into_array(), + 0, + 1024, + )? + .into_array(); + let expected = array.execute_scalar(0, &mut ctx)?; + let mut probe = array.probe(ProbeUsage::Repeated); + assert_eq!(probe.scalar_at(0, &mut ctx)?, expected); + assert!(probe.scalar_at(1023, &mut ctx).is_err()); + Ok(()) +} diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 8045fc550cc..ade2f11cdb1 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -3,6 +3,7 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::ProbeCtx; use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; @@ -12,9 +13,20 @@ use super::RLE; use crate::FL_CHUNK_SIZE; use crate::rle::RLEArrayExt; use crate::rle::RLEArraySlotsExt; +use crate::rle::probe; +use crate::rle::probe::RleProbeState; impl OperationsVTable for RLE { - type ProbeState<'a> = (); + type ProbeState<'a> = RleProbeState<'a>; + + fn probe_scalar<'a>( + array: ArrayView<'a, RLE>, + index: usize, + probe: Option<&mut ProbeCtx<'a, Self::ProbeState<'a>>>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + probe::scalar_at(array, index, probe, ctx) + } fn scalar_at( array: ArrayView<'_, RLE>, diff --git a/encodings/pco/Cargo.toml b/encodings/pco/Cargo.toml index 1683d4875e9..eefbd9422b1 100644 --- a/encodings/pco/Cargo.toml +++ b/encodings/pco/Cargo.toml @@ -26,7 +26,13 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } +vortex-fastlanes = { workspace = true } +vortex-runend = { workspace = true } +[[bench]] +name = "probe" +harness = false diff --git a/encodings/pco/benches/probe.md b/encodings/pco/benches/probe.md new file mode 100644 index 00000000000..17fabda3b5e --- /dev/null +++ b/encodings/pco/benches/probe.md @@ -0,0 +1,122 @@ +# Scalar probe benchmarks + +Compare unchanged `execute_scalar` with `ArrayProbe::scalar_at` in the same binary. The +[example](../examples/probe.rs) demonstrates both `Once` and `Repeated` access. The +[design note](../../../vortex-array/PROBE_DESIGN.md) describes the caller and vtable APIs. + +| Encoding | Work retained across lookups | +|---|---| +| Primitive | Materialized validity, or a child probe for lazy validity at requested rows. | +| FastLanes RLE | Typed readers for materialized primitive slots, slot IDs for other children, and the slice base offset. | +| RunEnd | The context's ends probe across binary-search comparisons and lookups, and its values probe for selected runs, including nullness. | +| PCO | Validity, prefix ranks for non-null positions, page boundaries, and the last decoded page. | + +Other encodings use the existing scalar path. Encodings request child probes by slot from +`ProbeCtx`; local state does not store child probes manually. Each child gets its own +context, initialized once and reused recursively. + +Construction allocates nothing. `Once` initializes no retained state, though scalar +execution can still allocate decoding buffers. On this ARM64 build, combined contexts +occupy 48 bytes for Primitive, 32 for RunEnd, 128 for PCO, and 136 for RLE. With the +128-byte inline capacity, RLE spills once on first repeated use. The other contexts fit +inline. Requesting a child allocates a slot table; PCO allocates indexes and page buffers +as needed. No child table is allocated for RLE's direct primitive readers. + +## Run + +```bash +RUSTC_WRAPPER= cargo bench -p vortex-pco --bench probe -- --sample-count 100 --min-time 0.1 +RUSTC_WRAPPER= cargo bench -p vortex-pco --bench probe -- '\(1, (false|true), false\)' --sample-count 100 --sample-size 64 --min-time 0.2 +RUSTC_WRAPPER= cargo run -p vortex-pco --example probe +``` + +`RUSTC_WRAPPER=` bypasses this environment's sccache permission failure. No runtime feature +toggles were set. For these measurements, the binary was first built with +`RUSTC_WRAPPER= cargo bench -p vortex-pco --bench probe --no-run`, then invoked directly as +`target/release/deps/probe-944ff408a6b9fe8f --bench` with the corresponding arguments above. + +Cases are `(access_count, nullable, scattered)`. The leaf fixture has 16,384 `u32` rows +with values `row / 16`; nullable cases mark every 11th row null. PCO uses compression level +8 and 1,024 non-null values per page. Deterministic random indices are clustered within +256 logical rows starting at 4,096, or scattered over the whole array. + +The stacked fixture is `RunEnd(ends=PCO, values=PCO)`: 4,096 runs of length four, covering +16,384 logical rows. Ends are `4, 8, ...`; values are run numbers, with every 11th run null +in nullable cases. Both children use PCO level 8 with 1,024 non-null values per page. + +Timings include probe construction, first-use preparation, lookups, and destruction. +Compression, index generation, and execution-context creation are excluded equally. +One-access cases use `Once`; larger groups use `Repeated`. The `repeated_first` cases +measure the first lookup using `Repeated` separately. + +## Results + +Measured on 2026-09-11, ARM64 macOS, Rust 1.98.0, using the repository's bench profile. +Runs were serial, without concurrent builds or tests. Values are medians per complete +group of accesses. Single-access results use 64 independent probe lifetimes per sample +and are normalized per lookup; larger cases use the first command above. + +| Encoding | Accesses | Nullable | Pattern | `execute_scalar` | Probe | +|---|---:|---|---|---:|---:| +| RLE | 1 | no | single | 99.31 ns | 105.8 ns | +| RLE | 1 | yes | single | 179.3 ns | 137.7 ns | +| RLE | 64 | no | clustered | 6.458 µs | 1.374 µs | +| RLE | 64 | no | scattered | 6.374 µs | 1.384 µs | +| RLE | 64 | yes | clustered | 10.95 µs | 1.541 µs | +| RLE | 64 | yes | scattered | 10.29 µs | 1.509 µs | +| RLE | 1,024 | no | clustered | 101 µs | 20.91 µs | +| RLE | 1,024 | no | scattered | 100.9 µs | 20.91 µs | +| RLE | 1,024 | yes | clustered | 168 µs | 21.16 µs | +| RLE | 1,024 | yes | scattered | 164.8 µs | 21.16 µs | +| PCO | 1 | no | single | 5.363 µs | 5.323 µs | +| PCO | 1 | yes | single | 4.404 µs | 4.418 µs | +| PCO | 64 | no | clustered | 330.9 µs | 6.208 µs | +| PCO | 64 | no | scattered | 331.2 µs | 309.1 µs | +| PCO | 64 | yes | clustered | 269.3 µs | 5.708 µs | +| PCO | 64 | yes | scattered | 241.6 µs | 192.6 µs | +| PCO | 1,024 | no | clustered | 5.419 ms | 21.08 µs | +| PCO | 1,024 | no | scattered | 5.412 ms | 4.91 ms | +| PCO | 1,024 | yes | clustered | 4.127 ms | 31.62 µs | +| PCO | 1,024 | yes | scattered | 4.056 ms | 3.161 ms | +| RunEnd(PCO, PCO) | 1 | no | single | 25.97 µs | 25.18 µs | +| RunEnd(PCO, PCO) | 1 | yes | single | 53.68 µs | 53.18 µs | +| RunEnd(PCO, PCO) | 64 | no | clustered | 2.065 ms | 245.5 µs | +| RunEnd(PCO, PCO) | 64 | no | scattered | 2.033 ms | 334.7 µs | +| RunEnd(PCO, PCO) | 64 | yes | clustered | 4.01 ms | 262.7 µs | +| RunEnd(PCO, PCO) | 64 | yes | scattered | 4.044 ms | 446.4 µs | +| RunEnd(PCO, PCO) | 1,024 | no | clustered | 34.26 ms | 3.999 ms | +| RunEnd(PCO, PCO) | 1,024 | no | scattered | 33.22 ms | 5.487 ms | +| RunEnd(PCO, PCO) | 1,024 | yes | clustered | 66.57 ms | 4.12 ms | +| RunEnd(PCO, PCO) | 1,024 | yes | scattered | 64.74 ms | 7.044 ms | + +For 1,024 clustered non-nullable reads, this run shows approximately 4.8× faster RLE, +257× faster PCO, and 8.6× faster RunEnd-over-PCO access. RLE avoids scalar dispatch for +materialized primitive slots. PCO amortizes decoding across reads of the same page. +RunEnd preserves its children's preparation and avoids a separate search for nullness. + +PCO retains one page. Scattered access often misses that cache, so the improvement is +much smaller: 5.412 ms → 4.91 ms for 1,024 non-nullable reads here. In stacked RunEnd, +the ends search crosses page boundaries and still causes decodes; the values child +retains its own independent page. These local measurements do not establish a general +speedup for arbitrary access patterns or encoding trees. + +One-off non-nullable RLE adds about 6.5 ns in the batched measurement. Nullable RLE's +one-off path combines routing and nullness in one index lookup. PCO and stacked RunEnd +one-off timings are close to their baselines; small differences should be treated as +measurement variation rather than a reliable wrapper speedup. + +The first `Repeated` lookup, including teardown, measured separately in the main run: + +| Encoding | Non-nullable | Nullable | +|---|---:|---:| +| RLE | 65.87 ns | 163.5 ns | +| PCO | 5.124 µs | 4.082 µs | + +A deeper regression fixture, `RunEnd(PCO, RunEnd(PCO, PCO))`, uses three single-page PCO +leaves and verifies exactly three state initializations and three decodes across repeated +lookups, with one drop per state. New root probes prepare independent state, and `Once` +initializes none. This checks reuse independently of timing. + +Raw local output: `/private/tmp/vortex-probe-impl-bench.log` and +`/private/tmp/vortex-probe-impl-bench-once.log`. Context sizes were measured from the same +source and are implementation details, not ABI guarantees. diff --git a/encodings/pco/benches/probe.rs b/encodings/pco/benches/probe.rs new file mode 100644 index 00000000000..3c9a9731d5a --- /dev/null +++ b/encodings/pco/benches/probe.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar access comparisons, including preparation and teardown for every group of lookups. + +use std::hint::black_box; +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::ProbeUsage; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_fastlanes::RLEData; +use vortex_pco::Pco; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +const LEN: usize = 16_384; +// Number of accesses, nullable, scattered (otherwise clustered within 256 rows). +const CASES: &[(usize, bool, bool)] = &[ + (1, false, false), + (1, true, false), + (64, false, false), + (64, true, false), + (64, false, true), + (64, true, true), + (1024, false, false), + (1024, true, false), + (1024, false, true), + (1024, true, true), +]; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + vortex_runend::initialize(&session); + session +}); + +fn input(nullable: bool) -> PrimitiveArray { + let validity = if nullable { + Validity::from_iter((0..LEN).map(|i| i % 11 != 0)) + } else { + Validity::NonNullable + }; + PrimitiveArray::new( + (0..LEN) + .map(|i| u32::try_from(i / 16).vortex_expect("fixture values fit u32")) + .collect::>(), + validity, + ) +} + +fn indices(count: usize, scattered: bool) -> Vec { + let span = if scattered { LEN } else { 256 }; + let base = if scattered { 0 } else { 4096 }; + let mut seed = 42u64; + (0..count) + .map(|_| { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + base + ((seed >> 32) as usize % span) + }) + .collect() +} + +fn rle(nullable: bool) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + RLEData::encode(input(nullable).as_view(), &mut ctx) + .vortex_expect("RLE compression") + .into_array() +} + +fn pco(nullable: bool) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + Pco::from_primitive(input(nullable).as_view(), 8, 1024, &mut ctx) + .vortex_expect("PCO compression") + .into_array() +} + +fn runend_pco(nullable: bool) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let runs = u32::try_from(LEN / 4).vortex_expect("run count fits u32"); + let ends = PrimitiveArray::from_iter((1..=runs).map(|run| run * 4)); + let values = PrimitiveArray::new( + (0..runs).collect::>(), + if nullable { + Validity::from_iter((0..runs).map(|run| run % 11 != 0)) + } else { + Validity::NonNullable + }, + ); + RunEnd::try_new( + Pco::from_primitive(ends.as_view(), 8, 1024, &mut ctx) + .vortex_expect("PCO ends compression") + .into_array(), + Pco::from_primitive(values.as_view(), 8, 1024, &mut ctx) + .vortex_expect("PCO values compression") + .into_array(), + &mut ctx, + ) + .vortex_expect("RunEnd construction") + .into_array() +} + +fn execute_scalar(bencher: Bencher, array: ArrayRef, indices: &[usize]) { + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_refs(|ctx| { + for &index in indices { + black_box( + array + .execute_scalar(black_box(index), ctx) + .vortex_expect("scalar access"), + ); + } + }); +} + +fn probe(bencher: Bencher, array: ArrayRef, indices: &[usize], usage: ProbeUsage) { + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_refs(|ctx| { + let mut probe = array.probe(usage); + for &index in indices { + black_box( + probe + .scalar_at(black_box(index), ctx) + .vortex_expect("probe access"), + ); + } + }); +} + +#[divan::bench(args = CASES)] +fn rle_probe(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + probe( + bencher, + rle(nullable), + &indices(count, scattered), + if count == 1 { + ProbeUsage::Once + } else { + ProbeUsage::Repeated + }, + ); +} + +#[divan::bench(args = CASES)] +fn pco_probe(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + probe( + bencher, + pco(nullable), + &indices(count, scattered), + if count == 1 { + ProbeUsage::Once + } else { + ProbeUsage::Repeated + }, + ); +} + +#[divan::bench(args = [false, true])] +fn rle_repeated_first(bencher: Bencher, nullable: bool) { + probe( + bencher, + rle(nullable), + &indices(1, false), + ProbeUsage::Repeated, + ); +} + +#[divan::bench(args = [false, true])] +fn pco_repeated_first(bencher: Bencher, nullable: bool) { + probe( + bencher, + pco(nullable), + &indices(1, false), + ProbeUsage::Repeated, + ); +} + +#[divan::bench(args = CASES)] +fn rle_execute_scalar(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + execute_scalar(bencher, rle(nullable), &indices(count, scattered)); +} + +#[divan::bench(args = CASES)] +fn pco_execute_scalar(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + execute_scalar(bencher, pco(nullable), &indices(count, scattered)); +} + +#[divan::bench(args = CASES)] +fn runend_pco_execute_scalar(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + execute_scalar(bencher, runend_pco(nullable), &indices(count, scattered)); +} + +#[divan::bench(args = CASES)] +fn runend_pco_probe(bencher: Bencher, (count, nullable, scattered): (usize, bool, bool)) { + probe( + bencher, + runend_pco(nullable), + &indices(count, scattered), + if count == 1 { + ProbeUsage::Once + } else { + ProbeUsage::Repeated + }, + ); +} diff --git a/encodings/pco/examples/probe.rs b/encodings/pco/examples/probe.rs new file mode 100644 index 00000000000..518cf99acac --- /dev/null +++ b/encodings/pco/examples/probe.rs @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Run with `cargo run -p vortex-pco --example probe`. + +use vortex_array::IntoArray; +use vortex_array::ProbeUsage; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_error::VortexResult; +use vortex_fastlanes::RLEData; +use vortex_pco::Pco; +use vortex_runend::RunEnd; + +fn main() -> VortexResult<()> { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + vortex_runend::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let values = + PrimitiveArray::from_option_iter((0..4096u32).map(|i| (i % 11 != 0).then_some(i / 16))); + let rle = RLEData::encode(values.as_view(), &mut ctx)?.into_array(); + let pco = Pco::from_primitive(values.as_view(), 3, 1024, &mut ctx)?.into_array(); + let ends = PrimitiveArray::from_iter((1..=4096u32).map(|run| run * 4)); + let ends = Pco::from_primitive(ends.as_view(), 3, 1024, &mut ctx)?.into_array(); + let runend_pco = RunEnd::try_new(ends, pco.clone(), &mut ctx)?.into_array(); + + for (name, array) in [("RLE", rle), ("PCO", pco), ("RunEnd(PCO, PCO)", runend_pco)] { + // Once passes no retained state to the encoding. + let scalar = array.probe(ProbeUsage::Once).scalar_at(17, &mut ctx)?; + println!("{name} single lookup: {scalar}"); + + // RLE and RunEnd retain child probes; each PCO child keeps its own decoded page. + let mut probe = array.probe(ProbeUsage::Repeated); + for index in [17, 33, 22, 1025, 17] { + println!("{name}[{index}] = {}", probe.scalar_at(index, &mut ctx)?); + } + } + Ok(()) +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 48939f9efe8..74a69fe0975 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -29,6 +29,7 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::ProbeCtx; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::Primitive; @@ -778,7 +779,19 @@ impl ValidityVTable for Pco { } impl OperationsVTable for Pco { - type ProbeState<'a> = (); + type ProbeState<'a> = crate::probe::PcoProbeState; + + fn probe_scalar<'a>( + array: ArrayView<'a, Pco>, + index: usize, + probe: Option<&mut ProbeCtx<'a, Self::ProbeState<'a>>>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + match probe { + Some(probe) => crate::probe::scalar_at(array, index, probe.state_mut(), ctx), + None => array.array().execute_scalar(index, ctx), + } + } fn scalar_at( array: ArrayView<'_, Pco>, diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 601bd0826bd..70d71abf7d2 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -20,6 +20,7 @@ mod array; mod compute; +mod probe; mod rules; mod slice; diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs new file mode 100644 index 00000000000..5002b324b36 --- /dev/null +++ b/encodings/pco/src/probe.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! PCO probes retain one decoded page and index the compacted, non-null value positions. + +use std::ops::Range; + +use pco::data_types::Number; +use pco::data_types::NumberType; +use pco::match_number_enum; +use pco::wrapped::FileDecompressor; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::array::number_type_from_ptype; +use crate::array::vortex_err_from_pco; + +const RANK_STRIDE: usize = 512; + +/// State retained by repeated PCO probes. Construction performs no allocation. +/// +/// The mask and rank index use unsliced logical rows; page boundaries use compacted value +/// positions. Only the most recently accessed page is retained, bounding decoded storage. +#[derive(Default)] +pub struct PcoProbeState { + validity: Option, + rank: Vec, + pages: Vec, + decoded: Option<(Range, PrimitiveArray)>, + #[cfg(test)] + decoded_pages: usize, + #[cfg(test)] + tracking: tests::TrackedProbeState, +} + +struct Page { + values: Range, + chunk: usize, +} + +pub(crate) fn scalar_at( + array: ArrayView<'_, Pco>, + index: usize, + state: &mut PcoProbeState, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = match &mut state.validity { + Some(mask) => mask, + slot @ None => slot.insert( + array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?, + ), + }; + let logical_index = array.slice_start() + index; + if !mask.value(logical_index) { + return Ok(Scalar::null(array.dtype().clone())); + } + + let value_index = match mask { + Mask::AllTrue(_) => logical_index, + Mask::AllFalse(_) => unreachable!("null rows returned above"), + Mask::Values(values) => { + let bits = values.bit_buffer(); + if state.rank.is_empty() { + state.rank.reserve(bits.len().div_ceil(RANK_STRIDE)); + let mut count = 0; + for start in (0..bits.len()).step_by(RANK_STRIDE) { + state.rank.push(count); + count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len())); + } + } + let block = logical_index / RANK_STRIDE; + state.rank[block] + bits.count_range(block * RANK_STRIDE, logical_index) + } + }; + + let (range, values) = match &state.decoded { + Some(decoded) if decoded.0.contains(&value_index) => decoded, + _ => { + if state.pages.is_empty() { + let mut start = 0; + for (chunk, metadata) in array.metadata.chunks.iter().enumerate() { + for page in &metadata.pages { + let end = start + page.n_values as usize; + state.pages.push(Page { + values: start..end, + chunk, + }); + start = end; + } + } + } + let page_index = state + .pages + .partition_point(|page| page.values.end <= value_index); + let page = state + .pages + .get(page_index) + .ok_or_else(|| vortex_err!("Missing PCO page for value {value_index}"))?; + let decoded = match_number_enum!( + number_type_from_ptype(array.dtype().as_ptype()), + NumberType => { decode_page::(array, page, array.pages[page_index].as_slice(), ctx)? } + ); + #[cfg(test)] + { + state.decoded_pages += 1; + state.tracking.record_decode(); + } + state.decoded.insert((page.values.clone(), decoded)) + } + }; + Ok(match_each_native_ptype!(values.ptype(), |T| { + Scalar::primitive( + values.as_slice::()[value_index - range.start], + array.dtype().nullability(), + ) + })) +} + +fn decode_page( + array: ArrayView<'_, Pco>, + page: &Page, + buffer: &[u8], + ctx: &mut ExecutionCtx, +) -> VortexResult { + let (file, _) = + FileDecompressor::new(array.metadata.header.as_slice()).map_err(vortex_err_from_pco)?; + let (mut chunk, _) = file + .chunk_decompressor::(array.chunk_metas[page.chunk].as_ref()) + .map_err(vortex_err_from_pco)?; + let mut decoder = chunk + .page_decompressor(buffer, page.values.len()) + .map_err(vortex_err_from_pco)?; + let mut values = BufferMut::::zeroed_in(page.values.len(), ctx.allocator().clone()); + decoder.read(&mut values).map_err(vortex_err_from_pco)?; + Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) +} + +#[cfg(test)] +mod tests; diff --git a/encodings/pco/src/probe/tests.rs b/encodings/pco/src/probe/tests.rs new file mode 100644 index 00000000000..4a7c30d1a17 --- /dev/null +++ b/encodings/pco/src/probe/tests.rs @@ -0,0 +1,268 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cell::Cell; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::ProbeUsage; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_runend::RunEnd; + +use super::PcoProbeState; +use crate::Pco; + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +struct ProbeCounts { + initialized: usize, + decoded: usize, + dropped: usize, +} + +thread_local! { + static COUNTS: Cell = Cell::default(); +} + +pub(super) struct TrackedProbeState; + +impl Default for TrackedProbeState { + fn default() -> Self { + let counts = COUNTS.get(); + COUNTS.set(ProbeCounts { + initialized: counts.initialized + 1, + ..counts + }); + Self + } +} + +impl TrackedProbeState { + pub(super) fn record_decode(&self) { + let counts = COUNTS.get(); + COUNTS.set(ProbeCounts { + decoded: counts.decoded + 1, + ..counts + }); + } +} + +impl Drop for TrackedProbeState { + fn drop(&mut self) { + let counts = COUNTS.get(); + COUNTS.set(ProbeCounts { + dropped: counts.dropped + 1, + ..counts + }); + } +} + +fn stacked_runend( + runs: u32, + page_size: usize, + nested: bool, + nullable: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let ends = PrimitiveArray::from_iter((1..=runs).map(|run| run * 4)); + let ends = Pco::from_primitive(ends.as_view(), 3, page_size, ctx)?.into_array(); + let values = if nested { + stacked_runend(runs / 4, page_size, false, nullable, ctx)? + } else { + let values = PrimitiveArray::new( + (0..runs).collect::>(), + if nullable { + Validity::from_iter((0..runs).map(|run| run % 11 != 0)) + } else { + Validity::NonNullable + }, + ); + Pco::from_primitive(values.as_view(), 3, page_size, ctx)?.into_array() + }; + Ok(RunEnd::try_new(ends, values, ctx)?.into_array()) +} + +#[test] +fn stacked_runend_reuses_each_child_state_and_drops_it_once() -> VortexResult<()> { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + let mut ctx = session.create_execution_ctx(); + // RunEnd(PCO, RunEnd(PCO, PCO)): three leaves, each fitting in one page. + let array = stacked_runend(256, 512, true, false, &mut ctx)?; + COUNTS.set(ProbeCounts::default()); + for lifetime in 1..=2 { + let mut probe = array.probe(ProbeUsage::Repeated); + assert_eq!(COUNTS.get().initialized, 3 * (lifetime - 1)); + assert!(probe.scalar_at(array.len(), &mut ctx).is_err()); + assert_eq!(COUNTS.get().initialized, 3 * (lifetime - 1)); + for index in [1, 5, 127, 255, 511, 1023, 0, 1] { + assert_eq!( + probe.scalar_at(index, &mut ctx)?, + u32::try_from(index / 16)?.into() + ); + assert_eq!( + COUNTS.get(), + ProbeCounts { + initialized: 3 * lifetime, + decoded: 3 * lifetime, + dropped: 3 * (lifetime - 1), + } + ); + } + drop(probe); + assert_eq!(COUNTS.get().dropped, 3 * lifetime); + } + let before = COUNTS.get(); + let mut once = array.probe(ProbeUsage::Once); + for index in [1, 17, 511] { + assert_eq!( + once.scalar_at(index, &mut ctx)?, + u32::try_from(index / 16)?.into() + ); + } + assert_eq!(COUNTS.get(), before); + Ok(()) +} + +#[rstest] +fn stacked_runend_random_access( + #[values(ProbeUsage::Once, ProbeUsage::Repeated)] usage: ProbeUsage, + #[values(false, true)] nested: bool, + #[values(false, true)] nullable: bool, + #[values(false, true)] sliced: bool, +) -> VortexResult<()> { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let encoded = stacked_runend(4096, 128, nested, nullable, &mut ctx)?; + let range = if sliced { 777..15333 } else { 0..16384 }; + let source = if sliced { + encoded + .slice(range.clone())? + .execute::(&mut ctx)? + } else { + encoded + }; + assert!(source.is::()); + let indices = [ + 0, + 1, + 3, + 4, + 15, + 16, + 44, + 176, + 511, + 512, + 1023, + 1024, + 4097, + source.len() - 1, + 0, + ]; + let mut actual = builder_with_capacity_in(source.dtype(), indices.len(), ctx.allocator()); + let mut probe = source.probe(usage); + for &index in &indices { + actual.append_scalar(&probe.scalar_at(index, &mut ctx)?)?; + } + let expected = indices + .map(|index| u32::try_from((range.start + index) / if nested { 16 } else { 4 })) + .into_iter() + .collect::, _>>()?; + let validity = if nullable { + Validity::from_iter(expected.iter().map(|value| value % 11 != 0)) + } else { + Validity::NonNullable + }; + let expected = PrimitiveArray::new(expected, validity); + assert_arrays_eq!(actual.finish(), expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case(ProbeUsage::Once, false)] +#[case(ProbeUsage::Repeated, false)] +#[case(ProbeUsage::Once, true)] +#[case(ProbeUsage::Repeated, true)] +fn sliced_nullable_random_access( + #[case] usage: ProbeUsage, + #[case] sliced: bool, +) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..4096i32).map(|i| (i % 7 != 0).then_some(i * 19))); + let encoded = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?.into_array(); + let range = if sliced { 777..3333 } else { 0..4096 }; + let source = encoded.slice(range.clone())?; + assert!(source.is::()); + let input = input.slice(range)?; + let indices = [0u32, 1, 127, 128, 512, 2048, 129, 2, 1024, 0]; + let mut actual = builder_with_capacity_in(source.dtype(), indices.len(), ctx.allocator()); + let mut probe = source.probe(usage); + for index in indices { + actual.append_scalar(&probe.scalar_at(index as usize, &mut ctx)?)?; + } + assert_arrays_eq!( + actual.finish(), + input.take(PrimitiveArray::from_iter(indices).into_array())?, + &mut ctx + ); + assert!(probe.scalar_at(source.len(), &mut ctx).is_err()); + Ok(()) +} + +#[rstest] +#[case(PrimitiveArray::from_iter([1u16, 9, 32768, 65535]))] +#[case(PrimitiveArray::from_iter([i64::MIN, -1, 0, i64::MAX]))] +#[case(PrimitiveArray::from_iter([1.25f64, -2.5, 0.0, f64::INFINITY]))] +fn preserves_physical_type(#[case] input: PrimitiveArray) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let encoded = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?.into_array(); + let mut probe = encoded.probe(ProbeUsage::Repeated); + let mut actual = builder_with_capacity_in(input.dtype(), input.len(), ctx.allocator()); + for i in 0..input.len() { + actual.append_scalar(&probe.scalar_at(i, &mut ctx)?)?; + } + assert_arrays_eq!(actual.finish(), input, &mut ctx); + Ok(()) +} + +#[test] +fn decodes_once_per_cached_page_and_evicts() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..4096i32); + let encoded = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?; + let mut state = PcoProbeState::default(); + for index in [1, 5, 2, 100, 7] { + assert_eq!( + super::scalar_at(encoded.as_view(), index, &mut state, &mut ctx)?, + i32::try_from(index)?.into() + ); + } + assert_eq!(state.decoded_pages, 1); + super::scalar_at(encoded.as_view(), 2048, &mut state, &mut ctx)?; + assert_eq!(state.decoded_pages, 2); + super::scalar_at(encoded.as_view(), 1, &mut state, &mut ctx)?; + assert_eq!(state.decoded_pages, 3); + Ok(()) +} + +#[test] +fn all_null_access_does_not_decode() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let input = PrimitiveArray::new(vec![0i32; 128], Validity::AllInvalid); + let encoded = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?; + let mut state = PcoProbeState::default(); + let result = super::scalar_at(encoded.as_view(), 42, &mut state, &mut ctx)?; + assert!(result.is_null()); + assert_eq!(state.decoded_pages, 0); + assert!(state.rank.is_empty() && state.pages.is_empty()); + Ok(()) +} diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index b991609c19c..e1e5719a864 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -15,6 +15,7 @@ pub mod decompress_bool; mod iter; mod kernel; pub mod ops; +mod probe; mod rules; #[cfg(test)] #[cfg(not(codspeed))] diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 14acef6c3e1..d947e228366 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -4,6 +4,7 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::ProbeCtx; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; use vortex_array::search_sorted::SearchResult; @@ -20,6 +21,18 @@ use crate::array::RunEndArraySlotsExt; impl OperationsVTable for RunEnd { type ProbeState<'a> = (); + fn probe_scalar<'a>( + array: ArrayView<'a, RunEnd>, + index: usize, + probe: Option<&mut ProbeCtx<'a, Self::ProbeState<'a>>>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + match probe { + Some(probe) => crate::probe::scalar_at(array, index, probe, ctx), + None => array.array().execute_scalar(index, ctx), + } + } + fn scalar_at( array: ArrayView<'_, RunEnd>, index: usize, diff --git a/encodings/runend/src/probe.rs b/encodings/runend/src/probe.rs new file mode 100644 index 00000000000..23a685cbaf7 --- /dev/null +++ b/encodings/runend/src/probe.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! RunEnd probes retain their routing and value child probes across lookups. + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::ProbeCtx; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::RunEnd; +use crate::RunEndArrayExt; +use crate::RunEndSlots; + +pub(crate) fn scalar_at<'a>( + array: ArrayView<'a, RunEnd>, + index: usize, + probe: &mut ProbeCtx<'a, ()>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let logical_index = array + .offset() + .checked_add(index) + .ok_or_else(|| vortex_err!("RunEnd logical index overflow"))?; + // Search for the first end strictly greater than the logical index. Every comparison + // uses the same ends probe, preserving child preparation within and between searches. + let ends = probe.child(RunEndSlots::ENDS)?; + let mut left = 0; + let mut right = ends.array().len(); + while left < right { + let mid = left + (right - left) / 2; + let end = usize::try_from(&ends.scalar_at(mid, ctx)?)?; + if end <= logical_index { + left = mid + 1; + } else { + right = mid; + } + } + // The selected value supplies nullness too; probing expanded RunEnd validity would + // repeat the search and lose the values child's retained state. + probe.child(RunEndSlots::VALUES)?.scalar_at(left, ctx) +} + +#[cfg(test)] +mod tests; diff --git a/encodings/runend/src/probe/tests.rs b/encodings/runend/src/probe/tests.rs new file mode 100644 index 00000000000..8b82f80c666 --- /dev/null +++ b/encodings/runend/src/probe/tests.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::ProbeUsage; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::RunEnd; +use crate::tests::SESSION; + +#[rstest] +fn sliced_strings_and_null_runs( + #[values(ProbeUsage::Once, ProbeUsage::Repeated)] usage: ProbeUsage, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + buffer![2u8, 5, 9].into_array(), + VarBinViewArray::from_iter( + [Some("first"), None, Some("last")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(), + 1, + 7, + &mut ctx, + )? + .into_array(); + let mut probe = array.probe(usage); + let mut actual = builder_with_capacity_in(array.dtype(), 5, ctx.allocator()); + for index in [6, 0, 1, 3, 4] { + actual.append_scalar(&probe.scalar_at(index, &mut ctx)?)?; + } + assert_arrays_eq!( + actual.finish(), + VarBinViewArray::from_iter( + [Some("last"), Some("first"), None, None, Some("last")], + DType::Utf8(Nullability::Nullable) + ), + &mut ctx + ); + assert!(probe.scalar_at(array.len(), &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn empty_probe_checks_bounds() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + Buffer::::empty().into_array(), + Buffer::::empty().into_array(), + &mut ctx, + )? + .into_array(); + for usage in [ProbeUsage::Once, ProbeUsage::Repeated] { + assert!(array.probe(usage).scalar_at(0, &mut ctx).is_err()); + } + Ok(()) +} + +#[test] +fn lazy_value_validity_only_evaluates_requested_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let numerators = buffer![1u32, 1].into_array(); + let validity = numerators + .binary(buffer![1u32, 0].into_array(), Operator::Div)? + .binary(numerators, Operator::Eq)?; + let values = PrimitiveArray::new(buffer![42u32, 99], Validity::Array(validity)); + let array = + RunEnd::try_new(buffer![4u32, 8].into_array(), values.into_array(), &mut ctx)?.into_array(); + let expected = array.execute_scalar(0, &mut ctx)?; + let mut probe = array.probe(ProbeUsage::Repeated); + assert_eq!(probe.scalar_at(0, &mut ctx)?, expected); + assert_eq!(probe.scalar_at(3, &mut ctx)?, expected); + assert!(probe.scalar_at(4, &mut ctx).is_err()); + Ok(()) +} diff --git a/vortex-array/PROBE_DESIGN.md b/vortex-array/PROBE_DESIGN.md index f3cad7a5a08..e31647e62a6 100644 --- a/vortex-array/PROBE_DESIGN.md +++ b/vortex-array/PROBE_DESIGN.md @@ -40,9 +40,9 @@ fn probe_scalar<'a>( } ``` -All encodings initially select `ProbeState<'a> = ()` and use this default, including -primitive, PCO, FastLanes RLE, RunEnd, ScalarFn, and Zstd. The API change provides no -encoding-specific speedup on its own; optimized hooks are a separate follow-up. +Encodings can select `ProbeState<'a> = ()` and use this default without implementing a +hook. Primitive, PCO, FastLanes RLE, and RunEnd override it; other encodings, including +ScalarFn and Zstd, retain their existing scalar execution path. An encoding opting into preparation chooses its own concrete `ProbeState`. `Default` should be cheap and avoid allocation or execution. Fallible preparation belongs in the @@ -104,6 +104,10 @@ provides 128 inline bytes aligned to 16 bytes and a heap fallback for larger or contexts. The child-cache header counts toward that capacity. Child tables, indexes, and decoded buffers may allocate separately when an encoding needs them. +On the measured ARM64 build, the combined contexts occupy 48 bytes for Primitive, 32 for +RunEnd, 128 for PCO, and 136 for RLE. RLE therefore spills once on first repeated use; +the other contexts fit inline. These sizes are implementation details, not ABI guarantees. + ```text ArrayProbe::scalar_at -> existing DynArrayData::probe_scalar dispatch @@ -126,5 +130,23 @@ encoding state is not required to implement those traits. Tests cover inline and spilled storage, alignment, moves, borrowed state, exactly-once initialization and destruction, default scalar behavior, bounds and nulls, lazy child creation, repeated slot reuse, source binding, independent slots/contexts, and simultaneous -local-state and child access. Encoding follow-ups should test recursive cache reuse and -benchmark complete probe lifetimes, including preparation and teardown. +local-state and child access. + +## Encoding implementations + +| Encoding | Local state | Child access | +|---|---|---| +| Primitive | Prepared validity mask or a marker for lazy validity. | Lazy validity uses the validity slot's retained probe and evaluates only requested rows. | +| FastLanes RLE | Borrowed primitive readers with prepared validity, or slot IDs for encoded children; the slice base offset. | Materialized primitive slots are read directly. Other slots use the context's child probes. | +| RunEnd | Unit state. | Reuses the ends probe across binary-search comparisons and lookups, then probes the selected value, including its nullness. | +| PCO | Validity, non-null prefix ranks, page boundaries, and one decoded page. | Reads compressed buffers directly; no child probe is needed. | + +For `RunEnd(ends=PCO, values=PCO)`, the RunEnd context owns both child probes and each +PCO context retains its own page. This applies recursively: the regression fixture +`RunEnd(PCO, RunEnd(PCO, PCO))` verifies one initialization and decode per single-page +PCO leaf across repeated reads, and one drop per state. Independent root probes prepare +independent state. Page changes can still require decoding; PCO retains only one page. + +The [runnable example](../encodings/pco/examples/probe.rs) covers both usage modes and +stacked encodings. [Paired benchmarks](../encodings/pco/benches/probe.md) include complete +probe lifetimes, including preparation and teardown, alongside unchanged `execute_scalar`. diff --git a/vortex-array/src/arrays/primitive/vtable/operations.rs b/vortex-array/src/arrays/primitive/vtable/operations.rs index 7cc7e5a47a6..0f469d62634 100644 --- a/vortex-array/src/arrays/primitive/vtable/operations.rs +++ b/vortex-array/src/arrays/primitive/vtable/operations.rs @@ -2,16 +2,71 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::ExecutionCtx; +use crate::ProbeCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::arrays::Bool; use crate::arrays::Primitive; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::primitive::array::PrimitiveSlots; use crate::match_each_native_ptype; use crate::scalar::Scalar; +/// State for repeated primitive probes: the validity, resolved once. +/// +/// Encodings that route through primitive children reach this state via their child probes, +/// so a nullable child costs one bit read per lookup rather than a validity execution. +#[derive(Default)] +pub struct PrimitiveProbeState { + validity: Option, +} + +enum PreparedValidity { + Mask(Mask), + Lazy, +} + impl OperationsVTable for Primitive { - type ProbeState<'a> = (); + type ProbeState<'a> = PrimitiveProbeState; + + fn probe_scalar<'a>( + array: ArrayView<'a, Primitive>, + index: usize, + probe: Option<&mut ProbeCtx<'a, Self::ProbeState<'a>>>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(probe) = probe else { + return array.array().execute_scalar(index, ctx); + }; + let (state, children) = probe.parts(); + let validity = match &mut state.validity { + Some(validity) => validity, + slot @ None => slot.insert(match array.slots()[PrimitiveSlots::VALIDITY].as_ref() { + // A lazy validity expression may fail on rows this probe never requests. + Some(child) if !child.is::() => PreparedValidity::Lazy, + _ => PreparedValidity::Mask( + PrimitiveArrayExt::validity(&array).execute_mask(array.len(), ctx)?, + ), + }), + }; + let valid = match validity { + PreparedValidity::Mask(mask) => mask.value(index), + PreparedValidity::Lazy => children + .child(PrimitiveSlots::VALIDITY)? + .scalar_at(index, ctx)? + .as_bool() + .value() + .ok_or_else(|| vortex_err!("validity value at index {index} is null"))?, + }; + if !valid { + return Ok(Scalar::null(array.dtype().clone())); + } + Self::scalar_at(array, index, ctx) + } fn scalar_at( array: ArrayView<'_, Primitive>,