From 930da9c63b466e0009e72833f74d7bd95f5ab172 Mon Sep 17 00:00:00 2001 From: Thor Date: Mon, 31 Aug 2026 14:58:48 -0500 Subject: [PATCH 1/7] PoC: IndexedLayout Adds an indexed layout to the Vortex layout crate based on https://github.com/vortex-data/vortex/issues/9024 --- vortex-layout/Cargo.toml | 1 + vortex-layout/src/layouts/indexed/index.rs | 192 ++++++ vortex-layout/src/layouts/indexed/mod.rs | 366 +++++++++++ vortex-layout/src/layouts/indexed/reader.rs | 284 +++++++++ vortex-layout/src/layouts/indexed/session.rs | 62 ++ vortex-layout/src/layouts/indexed/tests.rs | 626 +++++++++++++++++++ vortex-layout/src/layouts/indexed/writer.rs | 236 +++++++ vortex-layout/src/layouts/mod.rs | 1 + vortex-layout/src/session.rs | 2 + 9 files changed, 1770 insertions(+) create mode 100644 vortex-layout/src/layouts/indexed/index.rs create mode 100644 vortex-layout/src/layouts/indexed/mod.rs create mode 100644 vortex-layout/src/layouts/indexed/reader.rs create mode 100644 vortex-layout/src/layouts/indexed/session.rs create mode 100644 vortex-layout/src/layouts/indexed/tests.rs create mode 100644 vortex-layout/src/layouts/indexed/writer.rs diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 5d4a62d86e7..ba095231cde 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -33,6 +33,7 @@ parking_lot = { workspace = true } paste = { workspace = true } pin-project-lite = { workspace = true } prost = { workspace = true } +roaring = { workspace = true } rustc-hash = { workspace = true } sketches-ddsketch = { workspace = true } termtree = { workspace = true } diff --git a/vortex-layout/src/layouts/indexed/index.rs b/vortex-layout/src/layouts/indexed/index.rs new file mode 100644 index 00000000000..dc6d4fe9b3b --- /dev/null +++ b/vortex-layout/src/layouts/indexed/index.rs @@ -0,0 +1,192 @@ +//! The pluggable index-kind contract: what a kind must implement to be built at write time and +//! probed at read time. + +use std::fmt::Debug; +use std::ops::Range; +use std::sync::Arc; + +use roaring::RoaringBitmap; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::Expression; +use vortex_array::stream::SendableArrayStream; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::Id; + +/// Stable registry id of an index kind, e.g. `vortex.idx.reverse_index`. +pub type IndexId = Id; + +/// Shared handle to a registered index kind. +pub type IndexVTableRef = Arc; + +/// A pluggable index kind. +/// +/// Mirrors the layout `VTable` machinery: implementations are registered in an +/// [`IndexSession`](crate::layouts::indexed::session::IndexSession) under a stable string id, +/// which is what gets written into the layout metadata. A reader that does not have the kind +/// registered drops the index child and reads the data child directly. +pub trait IndexVTable: 'static + Send + Sync + Debug { + /// Stable string id, e.g. `vortex.idx.reverse_index`. + fn id(&self) -> IndexId; + + /// Whether this kind can build an index over values of `dtype`. + fn supports_dtype(&self, dtype: &DType) -> bool; + + /// Construct a builder for the write path. + /// + /// `data_block_len` is the data child's repartition block size when known. Kinds that emit + /// block-granular locators should default their block length to it so pruned blocks line up + /// with chunk and segment boundaries. + fn builder( + &self, + dtype: &DType, + options: &[u8], + data_block_len: Option, + session: &VortexSession, + ) -> VortexResult>; + + /// Decide whether this index can serve `expr`, a single conjunct scoped to the data child's + /// dtype. + /// + /// `None` means "no claim" and is always safe: the scan falls back to the data child. + fn plan( + &self, + expr: &BoundExpression, + dtype: &DType, + options: &[u8], + ) -> VortexResult>; +} + +/// Accumulates index content while the data stream is written. +pub trait IndexBuilder: Send { + /// Chunks arrive in stream order with their absolute row offset within this layout. + fn push( + &mut self, + chunk: &ArrayRef, + row_offset: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()>; + + /// Emit the index content as an array stream, to be written through a child layout strategy. + /// + /// Returns the final serialized options alongside it, so builders can record normalization + /// choices or block sizes discovered during the build. + /// + /// `None` declines: nothing worth keeping was built, so no index child and no spec are written, + /// and the wrapper collapses to the plain data layout if every builder declines. This is the + /// only point at which size can be judged — a builder is constructed before the first chunk + /// arrives, so row count and cardinality are not knowable earlier. Declining is always safe: an + /// absent index reads exactly like an unregistered one. + fn finish(self: Box) -> VortexResult)>>; + + /// Bytes currently buffered, reported up through the write context's buffered-bytes tracker. + fn buffered_bytes(&self) -> u64; +} + +/// What a probe result means and how precisely it locates. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum IndexExactness { + /// True bits are exactly the matching rows, so the probe may serve `filter_evaluation` + /// directly and skip decoding the data child for that conjunct. + Exact, + /// False bits are proven non-matching; true bits are "maybe". Serves `pruning_evaluation` + /// only, and the real predicate re-checks the survivors. + Superset, +} + +/// Where an index located its matches, in the data child's row space. +/// +/// Both variants are roaring bitmaps: postings are naturally set-like, intersect and union +/// cheaply, and expand into a [`Mask`] by walking runs in sorted order. +/// +/// Roaring bitmaps are `u32`-keyed, which caps a single layout at `u32::MAX` rows. That is far +/// above any practical Vortex file, and [`RowLocator::mask_for`] errors rather than truncating if +/// it is ever exceeded. +#[derive(Clone, Debug)] +pub enum RowLocator { + /// Row positions local to the data child. + Rows(RoaringBitmap), + /// Ids of fixed `block_len`-row blocks, expanded by broadcasting each block's bit across its + /// rows — the same shape as the zoned reader's per-zone expansion. + Blocks { block_len: u64, ids: RoaringBitmap }, +} + +impl RowLocator { + /// An empty locator: nothing matches, so everything prunes. + pub fn empty_rows() -> Self { + RowLocator::Rows(RoaringBitmap::new()) + } + + /// Expand this locator into a mask covering `row_range` of the data child. + /// + /// The returned mask has length `row_range.len()` and is *not* intersected with any input + /// mask; callers do that. + pub fn mask_for(&self, row_range: &Range) -> VortexResult { + let len = usize::try_from(row_range.end - row_range.start)?; + let mut bits = BitBufferMut::with_capacity(len); + + match self { + // Walk the set bits in ascending order, emitting the false run before each one. The + // bitmap is sorted, so this is a single linear pass with no random access. + RowLocator::Rows(rows) => { + let start = u32::try_from(row_range.start)?; + let end = u32::try_from(row_range.end)?; + let mut pos = row_range.start; + for row in rows.range(start..end) { + let row = u64::from(row); + bits.append_n(false, usize::try_from(row - pos)?); + bits.append_n(true, 1); + pos = row + 1; + } + bits.append_n(false, usize::try_from(row_range.end - pos)?); + } + // Broadcast each block's bit across the rows it covers, clipped to `row_range`. + RowLocator::Blocks { block_len, ids } => { + let mut row = row_range.start; + while row < row_range.end { + let block = row / block_len; + let block_end = ((block + 1) * block_len).min(row_range.end); + let hit = ids.contains(u32::try_from(block)?); + bits.append_n(hit, usize::try_from(block_end - row)?); + row = block_end; + } + } + } + + Ok(Mask::from(bits.freeze())) + } +} + +/// How an index intends to answer one expression. +/// +/// The probe runs `filter` as an ordinary scan over the index child — inheriting its zone maps, +/// lazy segment IO and compression — then hands the surviving index rows to `resolve`, which folds +/// them into a locator over the data child's rows. +pub struct IndexQueryPlan { + /// Whether the resulting mask is exact or a superset. + pub exactness: IndexExactness, + /// Predicate over the index child's dtype, selecting the posting rows this query needs. + /// + /// Unbound: the index child's dtype is only known once its layout child is materialized, so + /// the reader binds this against the index child's dtype right before scanning. + pub filter: Expression, + /// Folds the selected posting rows into a locator over the data child's row space. + pub resolve: Arc, +} + +/// Post-processes probed index rows into a [`RowLocator`]. +pub trait IndexResolve: 'static + Send + Sync { + /// `postings` are the index-child rows that survived [`IndexQueryPlan::filter`], projected in + /// the index child's own schema. + fn resolve( + &self, + postings: &ArrayRef, + data_row_count: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult; +} diff --git a/vortex-layout/src/layouts/indexed/mod.rs b/vortex-layout/src/layouts/indexed/mod.rs new file mode 100644 index 00000000000..ab1f5256803 --- /dev/null +++ b/vortex-layout/src/layouts/indexed/mod.rs @@ -0,0 +1,366 @@ +//! A prototype of the `vortex.indexed` layout proposed in +//! [vortex-data/vortex#9024](https://github.com/vortex-data/vortex/issues/9024). +//! +//! A `vortex.indexed` layout wraps a data layout with zero or more *locating indexes* held as +//! auxiliary children. Writers build indexes through the pluggable [`IndexVTable`] registry while +//! streaming; readers probe those indexes to prune (or outright answer) filter predicates, and +//! fall back to a plain scan of the data child whenever an index is missing, unknown, or has no +//! claim on the expression. +//! +//! Indexes are optional at every stage. A builder that finds nothing worth keeping declines once +//! the stream is drained — see [`IndexBuilder::finish`] — and if every builder declines, no wrapper +//! is written at all and the file carries the plain data layout. +//! +//! # Shape +//! +//! ```text +//! vortex.indexed +//! ├── child 0: data Transparent("data") +//! ├── child 1: index #0 Auxiliary("index:") +//! └── child n: index #n Auxiliary("index:") +//! ``` +//! +//! Index content is written through an ordinary layout strategy, so it is chunked, zone-mapped, +//! and compressed by the same machinery as data. Probing an index is therefore just a pruned scan +//! over the index child: a sorted key column's zone map narrows the probe to a handful of zones. +//! +//! # What ships here +//! +//! The generic wrapper only: [`Indexed`], [`writer::IndexedStrategy`] and [`reader::IndexedReader`], +//! plus the [`IndexVTable`] contract that index kinds implement. Concrete kinds are registered into +//! an [`session::IndexSession`] — see the `vortex-reverse-index` crate for a worked example, a +//! minimal equality index over integer columns. + +pub mod index; +pub(crate) mod reader; +pub mod session; +pub mod writer; + +use std::sync::Arc; + +use prost::Message; +use vortex_array::DeserializeMetadata; +use vortex_array::SerializeMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::proto::dtype as pb; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +pub use self::index::IndexBuilder; +pub use self::index::IndexExactness; +pub use self::index::IndexId; +pub use self::index::IndexQueryPlan; +pub use self::index::IndexResolve; +pub use self::index::IndexVTable; +pub use self::index::IndexVTableRef; +pub use self::index::RowLocator; +pub use self::session::IndexSession; +pub use self::session::IndexSessionExt; +pub use self::writer::IndexConfig; +pub use self::writer::IndexedStrategy; +use crate::Layout; +use crate::LayoutChildType; +use crate::LayoutDeserializeArgs; +use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::LayoutRef; +use crate::VTable; +use crate::children::layout_children; +use crate::layouts::indexed::reader::IndexedReader; +use crate::segments::SegmentSource; + +/// Registry id of this layout encoding. +pub const INDEXED_LAYOUT_ID: &str = "vortex.indexed"; + +/// Leading byte of the serialized metadata, so the protobuf can be re-shaped later. +const INDEXED_METADATA_VERSION: u8 = 1; + +/// Layout vtable for the `vortex.indexed` layout. +#[derive(Clone, Debug)] +pub struct Indexed; + +/// One index attached to the data child. +#[derive(Clone, Debug)] +pub struct IndexSpec { + id: IndexId, + options: Arc<[u8]>, + index_dtype: DType, + /// The resolved kind, or `None` when it is not registered in this session. An unresolved spec + /// is inert: its child is never probed and reads fall through to the data child. + vtable: Option, +} + +impl IndexSpec { + /// Create a spec for an index that was just built by `vtable`. + pub fn new(vtable: IndexVTableRef, options: Vec, index_dtype: DType) -> Self { + Self { + id: vtable.id(), + options: options.into(), + index_dtype, + vtable: Some(vtable), + } + } + + /// The registry id of this index's kind. + pub fn id(&self) -> IndexId { + self.id + } + + /// The kind-defined, self-versioned options blob. + pub fn options(&self) -> &[u8] { + &self.options + } + + /// The dtype of this index's layout child. + pub fn index_dtype(&self) -> &DType { + &self.index_dtype + } + + /// The resolved kind, or `None` if it is not registered in this session. + pub fn vtable(&self) -> Option<&IndexVTableRef> { + self.vtable.as_ref() + } +} + +/// Layout-specific data for the [`IndexedLayout`]. +/// +/// Child 0 is the data, sharing this layout's dtype and row space. Children 1.. are index content, +/// one per entry in [`IndexedData::indexes`]. +#[derive(Clone, Debug)] +pub struct IndexedData { + indexes: Arc<[IndexSpec]>, +} + +/// A layout that attaches locating indexes to a data child. +pub type IndexedLayout = Layout; + +impl IndexedLayout { + /// Assemble an indexed layout from a data child and one layout child per index spec. + pub fn try_new( + data: LayoutRef, + index_layouts: Vec, + indexes: Vec, + ) -> VortexResult { + vortex_ensure!( + index_layouts.len() == indexes.len(), + "IndexedLayout got {} index children for {} specs", + index_layouts.len(), + indexes.len() + ); + for (layout, spec) in index_layouts.iter().zip(&indexes) { + vortex_ensure!( + layout.dtype() == &spec.index_dtype, + "Index child dtype {} does not match spec dtype {} for {}", + layout.dtype(), + spec.index_dtype, + spec.id + ); + } + + let dtype = data.dtype().clone(); + let row_count = data.row_count(); + let mut children = Vec::with_capacity(1 + index_layouts.len()); + children.push(data); + children.extend(index_layouts); + + Ok(LayoutParts::new( + Indexed, + dtype, + row_count, + Vec::new(), + layout_children(children), + IndexedData { + indexes: indexes.into(), + }, + ) + .into_typed()) + } + + /// The indexes attached to the data child. + pub fn indexes(&self) -> &Arc<[IndexSpec]> { + &self.indexes + } +} + +impl VTable for Indexed { + type LayoutData = IndexedData; + type Metadata = IndexedMetadata; + + fn id(&self) -> LayoutId { + static ID: CachedId = CachedId::new(INDEXED_LAYOUT_ID); + *ID + } + + fn metadata(layout: &Layout) -> Self::Metadata { + IndexedMetadata { + indexes: layout + .indexes + .iter() + .map(|spec| IndexSpecProto { + id: spec.id.to_string(), + options: spec.options.to_vec(), + index_dtype: Some( + pb::DType::try_from(&spec.index_dtype) + .vortex_expect("index child dtype should be serializable"), + ), + }) + .collect::>() + .into(), + } + } + + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &IndexedMetadata, + ) -> VortexResult { + vortex_ensure!( + args.children.nchildren() == 1 + metadata.indexes.len(), + "IndexedLayout expects {} children (data + {} indexes), got {}", + 1 + metadata.indexes.len(), + metadata.indexes.len(), + args.children.nchildren() + ); + + let registry = args.session.indexes(); + let indexes = metadata + .indexes + .iter() + .map(|spec| { + let index_dtype = spec + .index_dtype + .as_ref() + .map(|dtype| DType::from_proto(dtype, args.session)) + .transpose()? + .ok_or_else(|| vortex_err!("Index spec {} is missing its dtype", spec.id))?; + let id = IndexId::from(spec.id.as_str()); + + // An unknown kind degrades to an inert spec rather than failing the read: the + // child stays addressable so child counts and dtypes still line up, but nothing + // ever probes it. + Ok(IndexSpec { + id, + options: spec.options.as_slice().into(), + index_dtype, + vtable: registry.find(&id), + }) + }) + .collect::>>()?; + + args.children.child(0, args.dtype)?; + for (idx, spec) in indexes.iter().enumerate() { + args.children.child(idx + 1, &spec.index_dtype)?; + } + + Ok(IndexedData { + indexes: indexes.into(), + }) + } + + fn child_dtype(layout: &Layout, slot: usize) -> VortexResult { + match slot { + 0 => Ok(layout.dtype().clone()), + _ => { + let Some(spec) = layout.indexes.get(slot - 1) else { + vortex_bail!("Invalid child index: {}", slot); + }; + Ok(spec.index_dtype.clone()) + } + } + } + + fn child_type(layout: &Layout, slot: usize) -> LayoutChildType { + match slot { + 0 => LayoutChildType::Transparent("data".into()), + _ => match layout.indexes.get(slot - 1) { + Some(spec) => LayoutChildType::Auxiliary(format!("index:{}", spec.id).into()), + None => vortex_panic!("Invalid child index: {}", slot), + }, + } + } + + fn new_reader( + layout: &Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Ok(Arc::new(IndexedReader::try_new( + layout.clone(), + name, + segment_source, + session.clone(), + ctx.clone(), + )?)) + } +} + +/// Serialized indexed-layout metadata: one entry per index child, in child order. +#[derive(Debug, Clone, PartialEq)] +pub struct IndexedMetadata { + indexes: Arc<[IndexSpecProto]>, +} + +#[derive(Clone, PartialEq, Message)] +struct IndexedMetadataProto { + #[prost(message, repeated, tag = "1")] + indexes: Vec, +} + +#[derive(Clone, PartialEq, Message)] +struct IndexSpecProto { + /// Registry id of the index kind, e.g. `vortex.idx.reverse_index`. + #[prost(string, tag = "1")] + id: String, + /// Kind-defined, self-versioned options. + #[prost(bytes = "vec", tag = "2")] + options: Vec, + /// The dtype of the index child. Layout nodes carry no dtype of their own — it flows top-down + /// during deserialization — so an auxiliary child's dtype has to be recorded here. + #[prost(message, optional, tag = "3")] + index_dtype: Option, +} + +impl SerializeMetadata for IndexedMetadata { + fn serialize(self) -> Vec { + let proto = IndexedMetadataProto { + indexes: self.indexes.to_vec(), + }; + let mut metadata = vec![INDEXED_METADATA_VERSION]; + metadata.extend(proto.encode_to_vec()); + metadata + } +} + +impl DeserializeMetadata for IndexedMetadata { + type Output = Self; + + fn deserialize(metadata: &[u8]) -> VortexResult { + let Some((&version, proto_bytes)) = metadata.split_first() else { + vortex_bail!("Indexed metadata missing protobuf version"); + }; + vortex_ensure!( + version == INDEXED_METADATA_VERSION, + "Unsupported indexed metadata version: {}", + version + ); + + let proto = IndexedMetadataProto::decode(proto_bytes) + .map_err(|err| vortex_err!("Failed to decode indexed metadata: {err}"))?; + Ok(Self { + indexes: proto.indexes.into(), + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/layouts/indexed/reader.rs b/vortex-layout/src/layouts/indexed/reader.rs new file mode 100644 index 00000000000..92969ce4309 --- /dev/null +++ b/vortex-layout/src/layouts/indexed/reader.rs @@ -0,0 +1,284 @@ +//! Read-time probing: answer or prune a conjunct from an index child, else defer to the data +//! child. + +use std::any::Any; +use std::ops::BitAnd; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use tracing::trace; +use vortex_array::ArrayRef; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::expr::BoundExpression; +use vortex_array::stream::ArrayStreamExt; +use vortex_error::SharedVortexResult; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::dash_map::Entry; + +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::LazyReaderChildren; +use crate::RowSplits; +use crate::SplitRange; +use crate::layouts::indexed::IndexedLayout; +use crate::layouts::indexed::index::IndexExactness; +use crate::layouts::indexed::index::IndexQueryPlan; +use crate::layouts::indexed::index::RowLocator; +use crate::scan::scan_builder::ScanBuilder; +use crate::segments::SegmentSource; + +/// One probe result, shared by every split that needs it. +type SharedProbe = Shared>>>; + +/// A reader for the [`crate::layouts::indexed::Indexed`] layout. +/// +/// Probes happen once per expression per file: the shared future is cached, and each split slices +/// its own row range out of the resulting locator rather than re-probing. +pub struct IndexedReader { + layout: IndexedLayout, + name: Arc, + lazy_children: Arc, + session: VortexSession, + /// Cached probes keyed by expression. `None` means no index claimed the expression, so the + /// lookup is not retried. + probes: DashMap>, +} + +#[derive(Clone)] +struct CachedProbe { + exactness: IndexExactness, + locator: SharedProbe, +} + +impl IndexedReader { + pub(crate) fn try_new( + layout: IndexedLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ctx: LayoutReaderContext, + ) -> VortexResult { + let mut dtypes = Vec::with_capacity(1 + layout.indexes().len()); + let mut names = Vec::with_capacity(1 + layout.indexes().len()); + dtypes.push(layout.dtype().clone()); + names.push(Arc::clone(&name)); + for spec in layout.indexes().iter() { + dtypes.push(spec.index_dtype().clone()); + names.push(format!("{}.index:{}", name, spec.id()).into()); + } + + let lazy_children = Arc::new(LazyReaderChildren::new( + Arc::clone(layout.children()), + dtypes, + names, + segment_source, + session.clone(), + ctx, + )); + + Ok(Self { + layout, + name, + lazy_children, + session, + probes: DashMap::default(), + }) + } + + fn data_child(&self) -> VortexResult<&LayoutReaderRef> { + self.lazy_children.get(0) + } + + /// Find the first index kind with a claim on `expr` and start (or reuse) its probe. + /// + /// One probe per expression per file: every split slices its own row range out of the shared + /// locator rather than re-probing. The vacant-entry insert holds the shard lock across + /// planning so two splits racing on the same expression cannot both issue the probe's IO; + /// planning only touches the child readers, never this map, so it cannot re-enter. + fn probe(&self, expr: &BoundExpression) -> VortexResult> { + if let Some(cached) = self.probes.get(expr) { + return Ok(cached.value().clone()); + } + + match self.probes.entry(expr.clone()) { + Entry::Occupied(entry) => Ok(entry.get().clone()), + Entry::Vacant(entry) => { + let probe = self.plan_probe(expr)?; + entry.insert(probe.clone()); + Ok(probe) + } + } + } + + fn plan_probe(&self, expr: &BoundExpression) -> VortexResult> { + for (idx, spec) in self.layout.indexes().iter().enumerate() { + // Unregistered kinds are inert: their child is never read. + let Some(vtable) = spec.vtable() else { + trace!(index = %spec.id(), "index kind not registered, skipping"); + continue; + }; + + let Some(plan) = vtable.plan(expr, self.layout.dtype(), spec.options())? else { + continue; + }; + + trace!(index = %spec.id(), %expr, filter = %plan.filter, "index claimed expression"); + + let index_reader = Arc::clone(self.lazy_children.get(idx + 1)?); + let exactness = plan.exactness; + let locator = probe_index( + index_reader, + plan, + self.layout.row_count(), + self.session.clone(), + )?; + + return Ok(Some(CachedProbe { exactness, locator })); + } + + Ok(None) + } +} + +/// Run a plan's filter as a real scan over the index child, then fold the surviving posting rows +/// into a locator. +/// +/// Going through [`ScanBuilder`] rather than calling the reader's evaluations directly is what +/// makes the probe cheap. The scan loop splits the index child at its natural chunk boundaries and +/// prunes each split before projecting it, so the sorted key column's zone map narrows the probe to +/// the few chunks that can hold the query's keys and no other posting bytes are ever fetched. +/// Evaluating the whole index child in one call instead would decode every posting list in it. +fn probe_index( + index_reader: LayoutReaderRef, + plan: IndexQueryPlan, + data_row_count: u64, + session: VortexSession, +) -> VortexResult { + // The index child's dtype is only known once its layout child is materialized, so the plan's + // filter is bound here rather than by the index kind that produced it. + let bound_filter = plan.filter.bind(index_reader.dtype())?; + let postings = ScanBuilder::new(session.clone(), index_reader) + .with_filter(bound_filter) + .into_array_stream()?; + + let resolve = Arc::clone(&plan.resolve); + Ok(async move { + // Only the rows matching the plan's key predicate survive, one per query term, so + // collecting them into a single array is cheap regardless of index size. + let postings: ArrayRef = postings.read_all().await?; + let mut ctx = session.create_execution_ctx(); + let locator = resolve.resolve(&postings, data_row_count, &mut ctx)?; + Ok(Arc::new(locator)) + } + .map_err(Arc::new) + .boxed() + .shared()) +} + +impl LayoutReader for IndexedReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn dtype(&self) -> &DType { + self.layout.dtype() + } + + fn row_count(&self) -> u64 { + self.layout.row_count() + } + + fn register_splits( + &self, + field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + self.data_child()? + .register_splits(field_mask, split_range, splits) + } + + fn pruning_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: Mask, + ) -> VortexResult { + let data_eval = self + .data_child()? + .pruning_evaluation(row_range, expr, mask.clone())?; + + let Some(probe) = self.probe(expr)? else { + return Ok(data_eval); + }; + + let row_range = row_range.clone(); + let name = Arc::clone(&self.name); + let expr = expr.clone(); + + Ok(MaskFuture::new(mask.len(), async move { + let locator = probe.locator.await?; + let mut result = mask.bitand(&locator.mask_for(&row_range)?); + + // Only bother the data child if the index left anything alive. + if !result.all_false() { + result = result.bitand(&data_eval.await?); + } + + trace!(%name, %expr, density = result.density(), "index pruning evaluation"); + Ok(result) + })) + } + + fn filter_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + // An exact index answers the conjunct outright, so the data child is never decoded for it. + // A superset index can only prune, and the data child re-checks the real predicate. Either + // way this reuses the cached probe, so a superset conjunct costs no extra IO here. + if let Some(probe) = self + .probe(expr)? + .filter(|probe| probe.exactness == IndexExactness::Exact) + { + let row_range = row_range.clone(); + let len = mask.len(); + return Ok(MaskFuture::new(len, async move { + let locator = probe.locator.await?; + let index_mask = locator.mask_for(&row_range)?; + // Post-condition: the result must be intersected with the input mask. + Ok(mask.await?.bitand(&index_mask)) + })); + } + + self.data_child()?.filter_evaluation(row_range, expr, mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + self.data_child()? + .projection_evaluation(row_range, expr, mask) + } +} diff --git a/vortex-layout/src/layouts/indexed/session.rs b/vortex-layout/src/layouts/indexed/session.rs new file mode 100644 index 00000000000..3be4a7a7b14 --- /dev/null +++ b/vortex-layout/src/layouts/indexed/session.rs @@ -0,0 +1,62 @@ +//! Session registry of index kinds. + +use std::any::Any; + +use vortex_session::ArcSwapMap; +use vortex_session::SessionExt; +use vortex_session::SessionGuard; +use vortex_session::SessionVar; +use vortex_session::registry::Id; + +use crate::layouts::indexed::index::IndexId; +use crate::layouts::indexed::index::IndexVTableRef; + +/// Registry of index kinds, keyed by their stable [`IndexId`]. +type IndexRegistry = ArcSwapMap; + +/// Session state holding the registered index kinds. +/// +/// Empty by default: index kinds live in their own crates and are registered explicitly, the same +/// way layout encodings are. A kind that is not registered is simply ignored on read — its child is +/// dropped from consideration and the data child answers the query directly. +#[derive(Clone, Debug, Default)] +pub struct IndexSession { + registry: IndexRegistry, +} + +impl IndexSession { + /// Register an index kind, replacing any existing kind with the same id. + pub fn register(&self, index: IndexVTableRef) { + self.registry.insert(index.id(), index); + } + + /// Find a registered index kind by id. + pub fn find(&self, id: &IndexId) -> Option { + self.registry.get(id) + } + + /// The underlying registry. + pub fn registry(&self) -> &IndexRegistry { + &self.registry + } +} + +impl SessionVar for IndexSession { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +/// Extension trait for reaching the index-kind registry from a session. +pub trait IndexSessionExt: SessionExt { + /// Returns the index-kind registry. + fn indexes(&self) -> SessionGuard<'_, IndexSession> { + self.get::() + } +} + +impl IndexSessionExt for S {} diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs new file mode 100644 index 00000000000..3f544a0e51e --- /dev/null +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -0,0 +1,626 @@ +//! End-to-end tests for the generic wrapper. +//! +//! Concrete index kinds live in their own crates (see `vortex-reverse-index`), so these exercise +//! the machinery through a test-only [`exact_value::ExactValueIndex`] instead. + +use std::sync::Arc; + +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::expr::eq; +use vortex_array::expr::like; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::stream::ArrayStreamExt; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::INDEXED_LAYOUT_ID; +use super::IndexConfig; +use super::IndexSessionExt; +use super::IndexedStrategy; +use crate::LayoutChildType; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::indexed::tests::exact_value::DecliningIndex; +use crate::layouts::indexed::tests::exact_value::ExactValueIndex; +use crate::layouts::repartition::RepartitionStrategy; +use crate::layouts::repartition::RepartitionWriterOptions; +use crate::scan::scan_builder::ScanBuilder; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialArrayStreamExt; +use crate::test::new_session; + +/// The only index kind these tests attach, standing in for a real plugin. +fn exact_configs() -> Vec { + vec![IndexConfig::with_defaults(ExactValueIndex::new_ref())] +} + +/// Small enough that a 12-row file spans three blocks, making the row/block granularity +/// difference visible in a single assertion. +const BLOCK_LEN: usize = 4; + +/// Rows 1 and 9 contain "needle"; nothing else does. With `BLOCK_LEN` of 4 they land in +/// blocks 0 and 2, leaving block 1 prunable. +const ROWS: [&str; 12] = [ + "alpha", + "a needle here", + "beta", + "gamma", + "delta", + "epsilon", + "zeta", + "eta", + "theta", + "needle again", + "iota", + "kappa", +]; + +/// A session knowing the index kinds this test suite ships. +/// +/// Deliberately not a shared global session: sessions clone by sharing one `Arc`, so registering +/// into a shared session would leak between tests, and +/// [`unregistered_index_kind_falls_back_to_the_data_child`] depends on two sessions with different +/// index registries. +fn session_with_exact_index() -> VortexSession { + let session = new_session(); + session.indexes().register(ExactValueIndex::new_ref()); + session +} + +fn text_column() -> ArrayRef { + VarBinViewArray::from_iter_str(ROWS).into_array() +} + +/// A write strategy that attaches `configs` to the text column. +/// +/// The indexed wrapper sits directly above repartitioning — the same slot `ZonedStrategy` occupies +/// — so it sees whole chunks in row order and knows the data child's block size. +fn strategy(configs: Vec) -> IndexedStrategy { + let data = RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: BLOCK_LEN, + block_size_target: None, + canonicalize: false, + }, + ); + IndexedStrategy::new(data, FlatLayoutStrategy::default(), configs) + .with_data_block_len(BLOCK_LEN as u64) +} + +/// Writes the text column, returning the resulting layout and the segments backing it. +async fn write( + session: &VortexSession, + configs: Vec, +) -> VortexResult<(LayoutRef, Arc)> { + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + + let layout = strategy(configs) + .write_stream( + ctx.into(), + Arc::::clone(&segments), + text_column().to_array_stream().sequenced(ptr), + eof, + session, + ) + .await?; + Ok((layout, segments)) +} + +fn text_reader( + session: &VortexSession, + layout: &LayoutRef, + segments: Arc, +) -> VortexResult { + layout.new_reader("text".into(), segments, session, &Default::default()) +} + +/// The pruning mask the text reader produces for `LIKE '%needle%'`. +/// +/// Goes straight at the indexed reader rather than through a scan, so the mask under test is +/// unambiguously the one the index produced. +async fn prune_mask( + session: &VortexSession, + layout: &LayoutRef, + segments: Arc, + needle: &str, +) -> VortexResult { + let reader = text_reader(session, layout, segments)?; + let row_count = reader.row_count(); + let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype())?; + + reader + .pruning_evaluation( + &(0..row_count), + &filter, + Mask::new_true(usize::try_from(row_count)?), + )? + .await +} + +/// The mask the text reader produces for `text == value`. +/// +/// An `Exact` plan serves `filter_evaluation` directly, so the result is the index's own answer +/// rather than the data child's, intersected with `input`. +async fn exact_mask( + session: &VortexSession, + layout: &LayoutRef, + segments: Arc, + value: &str, + input: MaskFuture, +) -> VortexResult { + let reader = text_reader(session, layout, segments)?; + let row_count = reader.row_count(); + let filter = eq(root(), lit(value)).bind(reader.dtype())?; + + reader + .filter_evaluation(&(0..row_count), &filter, input)? + .await +} + +async fn scan_matching( + session: &VortexSession, + layout: &LayoutRef, + segments: Arc, + needle: &str, +) -> VortexResult> { + let reader = text_reader(session, layout, segments)?; + let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype())?; + + let text = ScanBuilder::new(session.clone(), reader) + .with_filter(filter) + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = session.create_execution_ctx(); + let text = text.execute::(&mut ctx)?; + Ok((0..text.len()) + .map(|idx| String::from_utf8_lossy(text.bytes_at(idx).as_slice()).into_owned()) + .collect()) +} + +#[tokio::test] +async fn unregistered_index_kind_falls_back_to_the_data_child() -> VortexResult<()> { + // Written by a session that knows the exact value index... + let (layout, segments) = write( + &session_with_exact_index(), + vec![IndexConfig::with_defaults(ExactValueIndex::new_ref())], + ) + .await?; + + // ...and read by one that does not. The spec goes inert, nothing probes the index child, and + // the data child answers everything — indexes are strictly optional accelerators. + let read_session = new_session(); + + assert!( + prune_mask(&read_session, &layout, Arc::clone(&segments), "needle") + .await? + .all_true() + ); + assert_eq!( + scan_matching(&read_session, &layout, segments, "needle").await?, + vec![ROWS[1].to_string(), ROWS[9].to_string()], + ); + Ok(()) +} + +#[tokio::test] +async fn exact_index_answers_the_filter_itself() -> VortexResult<()> { + let session = session_with_exact_index(); + let (layout, segments) = write( + &session, + vec![IndexConfig::with_defaults(ExactValueIndex::new_ref())], + ) + .await?; + + let mask = exact_mask( + &session, + &layout, + Arc::clone(&segments), + ROWS[2], + MaskFuture::new_true(ROWS.len()), + ) + .await?; + assert_eq!(mask, Mask::from_iter((0..ROWS.len()).map(|row| row == 2))); + + // The post-condition is that the result is intersected with the input mask, so an input that + // excludes the match must yield nothing. + let excluded = exact_mask( + &session, + &layout, + segments, + ROWS[2], + MaskFuture::ready(Mask::from_iter((0..ROWS.len()).map(|row| row != 2))), + ) + .await?; + assert!(excluded.all_false()); + + Ok(()) +} + +#[tokio::test] +async fn layout_carries_one_auxiliary_child_per_index() -> VortexResult<()> { + let session = new_session(); + let (layout, _segments) = write(&session, exact_configs()).await?; + + assert_eq!(layout.encoding_id().as_str(), INDEXED_LAYOUT_ID); + + // The edge types are load-bearing beyond this crate: anything attributing segments to a role + // walks for the nearest `Auxiliary` edge, so an index child hung off a transparent edge would + // silently be counted as data. + assert_eq!( + (0..layout.nslots()) + .filter_map(|slot| layout.slot_type(slot)) + .collect::>(), + vec![ + LayoutChildType::Transparent("data".into()), + LayoutChildType::Auxiliary("index:test.idx.exact_value".into()), + ], + ); + + // Index content is an ordinary layout tree, so it inherits chunking and zone maps for free. + let index_child = layout + .slot(1)? + .ok_or_else(|| vortex_error::vortex_err!("an exact-value index was configured"))?; + assert_eq!( + index_child + .dtype() + .as_struct_fields() + .names() + .iter() + .map(|name| name.to_string()) + .collect::>(), + vec!["key".to_string(), "postings".to_string()], + ); + assert!(index_child.row_count() > 0); + Ok(()) +} + +/// A builder that finds nothing worth keeping must leave no trace at all: no index child, no spec, +/// and no wrapper, so the file reads exactly as if no index had been configured. +#[tokio::test] +async fn every_builder_declining_writes_no_wrapper() -> VortexResult<()> { + let session = new_session(); + let (layout, segments) = write( + &session, + vec![IndexConfig::with_defaults(DecliningIndex::new_ref())], + ) + .await?; + + assert_ne!(layout.encoding_id().as_str(), INDEXED_LAYOUT_ID); + assert_eq!( + scan_matching(&session, &layout, segments, "needle").await?, + vec![ROWS[1].to_string(), ROWS[9].to_string()], + ); + Ok(()) +} + +/// One index declining must not disturb the others: the wrapper survives with exactly the children +/// that were actually built, and the declining kind leaves no spec behind to probe. +#[tokio::test] +async fn one_builder_declining_leaves_the_others_intact() -> VortexResult<()> { + let session = session_with_exact_index(); + let (layout, segments) = write( + &session, + vec![ + IndexConfig::with_defaults(DecliningIndex::new_ref()), + IndexConfig::with_defaults(ExactValueIndex::new_ref()), + ], + ) + .await?; + + assert_eq!(layout.encoding_id().as_str(), INDEXED_LAYOUT_ID); + assert_eq!( + layout + .child_names() + .map(|name| name.to_string()) + .collect::>(), + vec!["data".to_string(), "index:test.idx.exact_value".to_string()], + ); + + // The surviving index still answers, so the spec ordering did not shift out from under it. + let mask = exact_mask( + &session, + &layout, + segments, + ROWS[2], + MaskFuture::new_true(ROWS.len()), + ) + .await?; + assert_eq!(mask, Mask::from_iter((0..ROWS.len()).map(|row| row == 2))); + Ok(()) +} + +/// A test-only sorted value index, present to exercise the [`super::IndexExactness::Exact`] path +/// that a real posting-list index kind (such as an n-gram index) would rarely reach for equality +/// queries. +/// +/// One row per distinct string, sorted, with a roaring posting list of the rows holding it. That +/// makes equality answerable outright, so `filter_evaluation` returns the index's mask and the +/// data child is never decoded for that conjunct. +mod exact_value { + use std::collections::BTreeMap; + use std::sync::Arc; + + use roaring::RoaringBitmap; + use vortex_array::ArrayRef; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::arrays::varbinview::VarBinViewArrayExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::StructFields; + use vortex_array::expr::BoundExpression; + use vortex_array::expr::col; + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::scalar_fn::fns::binary::Binary; + use vortex_array::scalar_fn::fns::literal::Literal; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::stream::ArrayStreamExt; + use vortex_array::stream::SendableArrayStream; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_session::VortexSession; + use vortex_session::registry::CachedId; + + use crate::layouts::indexed::IndexBuilder; + use crate::layouts::indexed::IndexExactness; + use crate::layouts::indexed::IndexId; + use crate::layouts::indexed::IndexQueryPlan; + use crate::layouts::indexed::IndexResolve; + use crate::layouts::indexed::IndexVTable; + use crate::layouts::indexed::IndexVTableRef; + use crate::layouts::indexed::RowLocator; + + pub const EXACT_VALUE_ID: &str = "test.idx.exact_value"; + pub const DECLINING_ID: &str = "test.idx.declining"; + const KEY_FIELD: &str = "key"; + const POSTINGS_FIELD: &str = "postings"; + + fn index_fields() -> StructFields { + let names: FieldNames = vec![KEY_FIELD, POSTINGS_FIELD].into(); + StructFields::new( + names, + vec![DType::Utf8(NonNullable), DType::Binary(NonNullable)], + ) + } + + #[derive(Debug)] + pub struct ExactValueIndex; + + impl ExactValueIndex { + pub fn new_ref() -> IndexVTableRef { + Arc::new(Self) + } + } + + impl IndexVTable for ExactValueIndex { + fn id(&self) -> IndexId { + static ID: CachedId = CachedId::new(EXACT_VALUE_ID); + *ID + } + + fn supports_dtype(&self, dtype: &DType) -> bool { + matches!(dtype, DType::Utf8(_)) + } + + fn builder( + &self, + _dtype: &DType, + _options: &[u8], + _data_block_len: Option, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Box::new(Builder { + postings: BTreeMap::new(), + })) + } + + fn plan( + &self, + expr: &BoundExpression, + _dtype: &DType, + _options: &[u8], + ) -> VortexResult> { + // Only ` == `. + if !expr.is::() || *expr.as_::() != Operator::Eq { + return Ok(None); + } + if !expr.child(0).is_root() || !expr.child(1).is::() { + return Ok(None); + } + let Some(value) = expr.child(1).as_::().as_utf8().value() else { + return Ok(None); + }; + let value = value.to_string(); + + Ok(Some(IndexQueryPlan { + exactness: IndexExactness::Exact, + filter: eq(col(KEY_FIELD), lit(value.clone())), + resolve: Arc::new(Resolve { value }), + })) + } + } + + struct Builder { + /// Sorted by construction, which is what gives the key column a useful zone map. + postings: BTreeMap, + } + + impl IndexBuilder for Builder { + fn push( + &mut self, + chunk: &ArrayRef, + row_offset: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let values = chunk.clone().execute::(ctx)?; + let validity = values + .varbinview_validity() + .execute_mask(values.len(), ctx)?; + + for idx in 0..values.len() { + if !validity.value(idx) { + continue; + } + let value = String::from_utf8_lossy(values.bytes_at(idx).as_slice()).into_owned(); + self.postings + .entry(value) + .or_default() + .insert(u32::try_from(row_offset + idx as u64)?); + } + Ok(()) + } + + fn finish(self: Box) -> VortexResult)>> { + let mut keys = Vec::with_capacity(self.postings.len()); + let mut lists = Vec::with_capacity(self.postings.len()); + for (key, bitmap) in self.postings { + let mut buffer = Vec::with_capacity(bitmap.serialized_size()); + bitmap + .serialize_into(&mut buffer) + .map_err(|err| vortex_err!("Failed to serialize postings: {err}"))?; + keys.push(key); + lists.push(buffer); + } + + let len = keys.len(); + let array = StructArray::try_new_with_dtype( + vec![ + VarBinViewArray::from_iter_str(keys).into_array(), + VarBinViewArray::from_iter_bin(lists).into_array(), + ], + index_fields(), + len, + Validity::NonNullable, + )?; + + Ok(Some((array.into_array().to_array_stream().boxed(), vec![]))) + } + + fn buffered_bytes(&self) -> u64 { + self.postings + .values() + .map(|bitmap| bitmap.serialized_size() as u64) + .sum() + } + } + + /// A kind that always declines at `finish`. + /// + /// Standing in for "the index would not be worth its bytes", so the decline paths are testable + /// without a fixture large enough to trip a real threshold. + #[derive(Debug)] + pub struct DecliningIndex; + + impl DecliningIndex { + pub fn new_ref() -> IndexVTableRef { + Arc::new(Self) + } + } + + impl IndexVTable for DecliningIndex { + fn id(&self) -> IndexId { + static ID: CachedId = CachedId::new(DECLINING_ID); + *ID + } + + fn supports_dtype(&self, dtype: &DType) -> bool { + matches!(dtype, DType::Utf8(_)) + } + + fn builder( + &self, + _dtype: &DType, + _options: &[u8], + _data_block_len: Option, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Box::new(DecliningBuilder)) + } + + fn plan( + &self, + _expr: &BoundExpression, + _dtype: &DType, + _options: &[u8], + ) -> VortexResult> { + Ok(None) + } + } + + struct DecliningBuilder; + + impl IndexBuilder for DecliningBuilder { + fn push( + &mut self, + _chunk: &ArrayRef, + _row_offset: u64, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + Ok(()) + } + + fn finish(self: Box) -> VortexResult)>> { + Ok(None) + } + + fn buffered_bytes(&self) -> u64 { + 0 + } + } + + struct Resolve { + value: String, + } + + impl IndexResolve for Resolve { + fn resolve( + &self, + postings: &ArrayRef, + _data_row_count: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let entries = postings.clone().execute::(ctx)?; + let keys = entries + .unmasked_field_by_name(KEY_FIELD)? + .clone() + .execute::(ctx)?; + let lists = entries + .unmasked_field_by_name(POSTINGS_FIELD)? + .clone() + .execute::(ctx)?; + + for idx in 0..keys.len() { + if keys.bytes_at(idx).as_slice() != self.value.as_bytes() { + continue; + } + let bitmap = RoaringBitmap::deserialize_from(lists.bytes_at(idx).as_slice()) + .map_err(|err| vortex_err!("Failed to deserialize postings: {err}"))?; + return Ok(RowLocator::Rows(bitmap)); + } + + Ok(RowLocator::empty_rows()) + } + } +} diff --git a/vortex-layout/src/layouts/indexed/writer.rs b/vortex-layout/src/layouts/indexed/writer.rs new file mode 100644 index 00000000000..912af030ce7 --- /dev/null +++ b/vortex-layout/src/layouts/indexed/writer.rs @@ -0,0 +1,236 @@ +//! Write-time assembly: forward chunks to the data child while feeding index builders, then write +//! each index's content after all data segments. + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use parking_lot::Mutex; +use tracing::trace; +use vortex_array::ArrayRef; +use vortex_array::VortexSessionExecute; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::BufferedBytesReservation; +use crate::BufferedBytesTracker; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::LayoutWriterContext; +use crate::layouts::indexed::IndexSpec; +use crate::layouts::indexed::IndexedLayout; +use crate::layouts::indexed::index::IndexBuilder; +use crate::layouts::indexed::index::IndexVTableRef; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialArrayStreamExt; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// An index to attach to a column, as configured on the write side. +#[derive(Clone, Debug)] +pub struct IndexConfig { + vtable: IndexVTableRef, + options: Vec, +} + +impl IndexConfig { + /// Configure `vtable` with kind-defined options. + pub fn new(vtable: IndexVTableRef, options: Vec) -> Self { + Self { vtable, options } + } + + /// Configure `vtable` with its default options. + pub fn with_defaults(vtable: IndexVTableRef) -> Self { + Self::new(vtable, Vec::new()) + } +} + +/// Wraps a data-child strategy with one or more index builders. +/// +/// Sits in the same slot as `ZonedStrategy`, meaning above the repartition step, so it knows the +/// data child's row block size and can hand it to block-granular index kinds, making their blocks +/// line up with the data child's chunks by default. +pub struct IndexedStrategy { + data: Arc, + index: Arc, + configs: Arc<[IndexConfig]>, + data_block_len: Option, +} + +impl IndexedStrategy { + /// Create a strategy writing data through `data` and each index's content through `index`. + pub fn new( + data: D, + index: I, + configs: Vec, + ) -> Self { + Self { + data: Arc::new(data), + index: Arc::new(index), + configs: configs.into(), + data_block_len: None, + } + } + + /// Tell block-granular index kinds the data child's row block size, so pruned blocks align + /// with chunk and segment boundaries. + pub fn with_data_block_len(mut self, data_block_len: u64) -> Self { + self.data_block_len = Some(data_block_len); + self + } +} + +/// Builders plus the running row offset, shared between the stream-mapping closure and the +/// finishing code. Index building is globally stateful, so pushes stay sequential in stream order. +struct BuilderState { + builders: Vec<(IndexVTableRef, Box)>, + row_offset: u64, + /// Reservation reflecting the builders' current combined `buffered_bytes()`. + /// + /// Builders report a running total rather than a per-chunk delta, so each push replaces this + /// reservation (drop the old, reserve the new total) instead of accumulating one reservation + /// per chunk the way `BufferedStrategy` does for known-size chunks. + buffered: Option, +} + +impl BuilderState { + fn push( + &mut self, + chunk: &ArrayRef, + session: &VortexSession, + tracker: &BufferedBytesTracker, + ) -> VortexResult<()> { + let mut ctx = session.create_execution_ctx(); + for (_, builder) in &mut self.builders { + builder.push(chunk, self.row_offset, &mut ctx)?; + } + self.row_offset += chunk.len() as u64; + + let total: u64 = self + .builders + .iter() + .map(|(_, builder)| builder.buffered_bytes()) + .sum(); + self.buffered = Some(tracker.reserve(total)); + Ok(()) + } +} + +#[async_trait] +impl LayoutStrategy for IndexedStrategy { + async fn write_stream( + &self, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + + let mut builders = Vec::with_capacity(self.configs.len()); + for config in self.configs.iter() { + if !config.vtable.supports_dtype(&dtype) { + continue; + } + let builder = + config + .vtable + .builder(&dtype, &config.options, self.data_block_len, session)?; + builders.push((Arc::clone(&config.vtable), builder)); + } + + // Nothing to index for this dtype: don't emit a wrapper at all, so readers see the plain + // data layout. + if builders.is_empty() { + return self + .data + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + let state = Arc::new(Mutex::new(BuilderState { + builders, + row_offset: 0, + buffered: None, + })); + + let feed_state = Arc::clone(&state); + let feed_session = session.clone(); + let feed_tracker = ctx.buffered_bytes_tracker().clone(); + let stream = SequentialStreamAdapter::new( + dtype, + stream.map(move |item| { + let (sequence_id, chunk) = item?; + feed_state + .lock() + .push(&chunk, &feed_session, &feed_tracker)?; + Ok((sequence_id, chunk)) + }), + ) + .sendable(); + + // Data segments come first, so a reader that ignores indexes keeps its locality and a + // streaming writer never has to seek back. + let data_eof = eof.split_off(); + let data_layout = self + .data + .write_stream( + ctx.clone(), + Arc::clone(&segment_sink), + stream, + data_eof, + session, + ) + .await?; + + // The stream is drained, so every builder has seen every chunk. Bytes now move from being + // buffered in memory to being written out below, so the reservation is released here + // rather than left to drop at the end of the function. + let builders = { + let mut state = state.lock(); + state.buffered.take(); + std::mem::take(&mut state.builders) + }; + + let mut index_layouts = Vec::with_capacity(builders.len()); + let mut specs = Vec::with_capacity(builders.len()); + for (vtable, builder) in builders { + // A builder that found nothing worth keeping leaves no trace: no child, no spec, and no + // sequence pointer, since the splits below are what allocate one. + let Some((content, options)) = builder.finish()? else { + trace!(index = %vtable.id(), "index builder declined, writing no child"); + continue; + }; + let index_dtype = content.dtype().clone(); + + // Each index child gets its own (stream pointer, eof) pair, all ordered after the data + // segments. + let content_ptr = eof.split_off(); + let child_eof = eof.split_off(); + let layout = self + .index + .write_stream( + ctx.clone(), + Arc::clone(&segment_sink), + content.sequenced(content_ptr), + child_eof, + session, + ) + .await?; + + specs.push(IndexSpec::new(vtable, options, index_dtype)); + index_layouts.push(layout); + } + + // Every builder declined, so there is nothing to wrap. The data layout is already written + // and stands on its own, so hand it back as if no index had been configured. + if index_layouts.is_empty() { + return Ok(data_layout); + } + + Ok(IndexedLayout::try_new(data_layout, index_layouts, specs)?.into_layout()) + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 47fa31aa3d9..d2a8dd71ca9 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -16,6 +16,7 @@ pub mod dict; pub mod file_stats; pub mod flat; pub(crate) mod foreign; +pub mod indexed; pub mod list; pub(crate) mod partitioned; pub mod repartition; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 0cf2234c09c..f349d9d0ea4 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -14,6 +14,7 @@ use crate::LayoutEncodingRef; use crate::layouts::chunked::Chunked; use crate::layouts::dict::Dict; use crate::layouts::flat::Flat; +use crate::layouts::indexed::Indexed; use crate::layouts::list::List; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; @@ -62,6 +63,7 @@ impl Default for LayoutSession { this.register(&LegacyStats as &dyn LayoutEncoding); this.register(&Dict as &dyn LayoutEncoding); this.register(&List as &dyn LayoutEncoding); + this.register(&Indexed as &dyn LayoutEncoding); this } } From f80bc954c3dd67eea3dd36a6e4ce5dd281db129d Mon Sep 17 00:00:00 2001 From: Thor Date: Mon, 31 Aug 2026 14:59:51 -0500 Subject: [PATCH 2/7] vortex-reverse-index: An example index Intended as an example implementation of a indexed-layout index. --- vortex-reverse-index/Cargo.toml | 33 +++ vortex-reverse-index/src/lib.rs | 244 +++++++++++++++++++++ vortex-reverse-index/src/tests.rs | 350 ++++++++++++++++++++++++++++++ 3 files changed, 627 insertions(+) create mode 100644 vortex-reverse-index/Cargo.toml create mode 100644 vortex-reverse-index/src/lib.rs create mode 100644 vortex-reverse-index/src/tests.rs diff --git a/vortex-reverse-index/Cargo.toml b/vortex-reverse-index/Cargo.toml new file mode 100644 index 00000000000..ab2436aec95 --- /dev/null +++ b/vortex-reverse-index/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "vortex-reverse-index" +description = "A simple value -> row reverse index over any column, demonstrating vortex-layout's indexed layout IndexVTable" +authors.workspace = true +categories.workspace = true +edition.workspace = true +homepage.workspace = true +include.workspace = true +keywords.workspace = true +license.workspace = true +readme.workspace = true +repository.workspace = true +rust-version.workspace = true +version.workspace = true + +[dependencies] +roaring = { workspace = true } +vortex-array = { workspace = true } +vortex-error = { workspace = true } +vortex-layout = { workspace = true } +vortex-session = { workspace = true } +vortex-utils = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } +vortex-buffer = { workspace = true } +vortex-edition = { workspace = true } +vortex-file = { workspace = true } +vortex-io = { workspace = true, features = ["tokio"] } +vortex-mask = { workspace = true } + +[lints] +workspace = true diff --git a/vortex-reverse-index/src/lib.rs b/vortex-reverse-index/src/lib.rs new file mode 100644 index 00000000000..3922cd4a9d9 --- /dev/null +++ b/vortex-reverse-index/src/lib.rs @@ -0,0 +1,244 @@ +//! A simple `value -> rows` reverse index over any column, demonstrating [`vortex_layout`]'s +//! `vortex.indexed` layout [`IndexVTable`] contract with a minimal concrete index kind. +//! +//! This is deliberately narrow: it supports only equality (`column == literal`), answered exactly +//! via a sorted `key -> postings` table, the same shape as the indexed layout's own test-only +//! `exact_value` index but generalized to any dtype instead of being fixed to `Utf8`. It exists as +//! a worked, non-test example of a concrete index kind for the `vortex.indexed` layout prototype +//! described in [vortex-data/vortex#9024](https://github.com/vortex-data/vortex/issues/9024). + +use std::sync::Arc; + +use roaring::RoaringBitmap; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::builder_with_capacity; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::col; +use vortex_array::expr::eq; +use vortex_array::expr::lit; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::search_sorted::SearchSorted; +use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::stream::ArrayStreamExt; +use vortex_array::stream::SendableArrayStream; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_layout::layouts::indexed::IndexBuilder; +use vortex_layout::layouts::indexed::IndexExactness; +use vortex_layout::layouts::indexed::IndexId; +use vortex_layout::layouts::indexed::IndexQueryPlan; +use vortex_layout::layouts::indexed::IndexResolve; +use vortex_layout::layouts::indexed::IndexVTable; +use vortex_layout::layouts::indexed::IndexVTableRef; +use vortex_layout::layouts::indexed::RowLocator; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; +use vortex_utils::aliases::hash_map::HashMap; + +#[cfg(test)] +mod tests; + +/// Stable registry id of this index kind. +pub const REVERSE_INDEX_ID: &str = "vortex.idx.reverse_index"; + +const KEY_FIELD: &str = "key"; +const POSTINGS_FIELD: &str = "postings"; + +/// A value -> rows reverse index over a single column of any dtype. +/// +/// One row per distinct value, sorted by key, with a roaring posting list of the rows holding it. +/// Sorting the key column gives it a useful zone map, so probing the index is a pruned scan rather +/// than a full decode. Equality is answered exactly: [`IndexVTable::plan`] only claims +/// `column == literal`, so the probe's mask is the answer, not just a filter to re-check. +#[derive(Debug)] +pub struct ReverseIndex; + +impl ReverseIndex { + /// A shared handle to this index kind, ready to register into an + /// [`IndexSession`](vortex_layout::layouts::indexed::IndexSession). + pub fn new_ref() -> IndexVTableRef { + Arc::new(Self) + } +} + +impl IndexVTable for ReverseIndex { + fn id(&self) -> IndexId { + static ID: CachedId = CachedId::new(REVERSE_INDEX_ID); + *ID + } + + fn supports_dtype(&self, dtype: &DType) -> bool { + // Every other dtype decodes to a `Scalar` and has a canonical `ArrayBuilder`; `Union` and + // `Variant` do not yet, so decline rather than panic building their key column. + !matches!(dtype, DType::Union(..) | DType::Variant(_)) + } + + fn builder( + &self, + dtype: &DType, + _options: &[u8], + _data_block_len: Option, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Box::new(Builder { + dtype: dtype.clone(), + postings: HashMap::new(), + })) + } + + fn plan( + &self, + expr: &BoundExpression, + dtype: &DType, + _options: &[u8], + ) -> VortexResult> { + // Only ` == `. + if !expr.is::() || *expr.as_::() != Operator::Eq { + return Ok(None); + } + if !expr.child(0).is_root() || !expr.child(1).is::() { + return Ok(None); + } + + let target = expr.child(1).as_::(); + if !target.dtype().eq_ignore_nullability(dtype) { + // Binding rejects mismatched dtypes for ordinary comparisons, but exempts `Extension` + // dtypes from that check; decline rather than guess at cross-type equality. + return Ok(None); + } + if target.is_null() { + // Null literal: `column == NULL` never matches under normal equality, and this index + // does not model that, so decline rather than claim it incorrectly. + return Ok(None); + } + let target = target.clone(); + + Ok(Some(IndexQueryPlan { + exactness: IndexExactness::Exact, + filter: eq(col(KEY_FIELD), lit(target.clone())), + resolve: Arc::new(Resolve { target }), + })) + } +} + +fn index_fields(key_dtype: DType) -> StructFields { + let names: FieldNames = vec![KEY_FIELD, POSTINGS_FIELD].into(); + StructFields::new(names, vec![key_dtype, DType::Binary(NonNullable)]) +} + +struct Builder { + dtype: DType, + /// Deduplicated by scalar equality; sorted into key order in `finish`, which is what gives + /// the key column a useful zone map. + postings: HashMap, +} + +impl IndexBuilder for Builder { + fn push( + &mut self, + chunk: &ArrayRef, + row_offset: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + for idx in 0..chunk.len() { + let value = chunk.execute_scalar(idx, ctx)?; + if value.is_null() { + continue; + } + self.postings + .entry(value) + .or_default() + .insert(u32::try_from(row_offset + idx as u64)?); + } + Ok(()) + } + + fn finish(self: Box) -> VortexResult)>> { + let Builder { dtype, postings } = *self; + + let mut entries: Vec<(Scalar, RoaringBitmap)> = postings.into_iter().collect(); + entries.sort_by(|(a, _), (b, _)| { + a.partial_cmp(b) + .vortex_expect("keys were all decoded from the same column, so they share a dtype") + }); + + let key_dtype = dtype.as_nonnullable(); + let mut key_builder = builder_with_capacity(&key_dtype, entries.len()); + let mut lists = Vec::with_capacity(entries.len()); + for (key, bitmap) in &entries { + // Keys are never null (`push` skips them), but may carry the source column's nullable + // dtype; the key column itself is non-nullable, so normalize before appending. + key_builder.append_scalar(&key.cast(&key_dtype)?)?; + let mut buffer = Vec::with_capacity(bitmap.serialized_size()); + bitmap + .serialize_into(&mut buffer) + .map_err(|err| vortex_err!("Failed to serialize postings: {err}"))?; + lists.push(buffer); + } + + let len = entries.len(); + let array = StructArray::try_new_with_dtype( + vec![ + key_builder.finish(), + VarBinViewArray::from_iter_bin(lists).into_array(), + ], + index_fields(key_dtype), + len, + Validity::NonNullable, + )?; + + Ok(Some((array.into_array().to_array_stream().boxed(), vec![]))) + } + + fn buffered_bytes(&self) -> u64 { + self.postings + .values() + .map(|bitmap| bitmap.serialized_size() as u64) + .sum() + } +} + +struct Resolve { + target: Scalar, +} + +impl IndexResolve for Resolve { + fn resolve( + &self, + postings: &ArrayRef, + _data_row_count: u64, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let entries = postings.clone().execute::(ctx)?; + let keys = entries.unmasked_field_by_name(KEY_FIELD)?; + let lists = entries + .unmasked_field_by_name(POSTINGS_FIELD)? + .clone() + .execute::(ctx)?; + + let Some(idx) = keys + .search_sorted(&self.target, SearchSortedSide::Left)? + .to_found() + else { + return Ok(RowLocator::empty_rows()); + }; + + let bitmap = RoaringBitmap::deserialize_from(lists.bytes_at(idx).as_slice()) + .map_err(|err| vortex_err!("Failed to deserialize postings: {err}"))?; + Ok(RowLocator::Rows(bitmap)) + } +} diff --git a/vortex-reverse-index/src/tests.rs b/vortex-reverse-index/src/tests.rs new file mode 100644 index 00000000000..eee07001f8c --- /dev/null +++ b/vortex-reverse-index/src/tests.rs @@ -0,0 +1,350 @@ +//! End-to-end tests: write a column indexed by [`ReverseIndex`], then probe it. + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::FieldPath; +use vortex_array::expr::col; +use vortex_array::expr::eq; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::stream::ArrayStreamExt; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::VortexFile; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::session::RuntimeSession; +use vortex_layout::LayoutChildType; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::indexed::INDEXED_LAYOUT_ID; +use vortex_layout::layouts::indexed::IndexConfig; +use vortex_layout::layouts::indexed::IndexSessionExt; +use vortex_layout::layouts::indexed::IndexedStrategy; +use vortex_layout::layouts::repartition::RepartitionStrategy; +use vortex_layout::layouts::repartition::RepartitionWriterOptions; +use vortex_layout::session::LayoutSession; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::ReverseIndex; + +/// Small enough that a 12-row file spans three blocks, making the row/block granularity +/// difference visible in a single assertion. +const BLOCK_LEN: usize = 4; + +const VALUE_FIELD: &str = "value"; + +/// `20` appears at rows 1 and 9; nothing else repeats. With `BLOCK_LEN` of 4 those land in blocks +/// 0 and 2, leaving block 1 prunable. `999` never appears, to exercise the "claimed but no match" +/// path distinctly from "no index claimed this at all". +const VALUES: [i32; 12] = [10, 20, 30, 40, 50, 60, 70, 80, 90, 20, 100, 110]; + +fn reverse_index_configs() -> Vec { + vec![IndexConfig::with_defaults(ReverseIndex::new_ref())] +} + +/// The array/layout encodings these tests need to write. +/// +/// The default Vortex file writer only permits array/layout ids covered by the session's enabled +/// editions, but those first-party declarations live in the `vortex` facade crate, which this +/// out-of-tree crate deliberately does not depend on. Declaring and enabling a tiny test-only +/// edition here is the local equivalent. +const TEST_EDITION: EditionId = EditionId::new("vortex-reverse-index-test", 2026, 1, 0); + +static TEST_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: TEST_EDITION, + min_vortex_version: None, + }, + added: &[ + EditionMember::array(&"vortex.struct"), + EditionMember::array(&"vortex.primitive"), + // The index child's postings column is serialized roaring bitmaps, written as varbinview. + EditionMember::array(&"vortex.varbinview"), + EditionMember::layout(&"vortex.struct"), + EditionMember::layout(&"vortex.chunked"), + EditionMember::layout(&"vortex.flat"), + EditionMember::layout(&INDEXED_LAYOUT_ID), + ], +}; + +/// A session knowing the `vortex.indexed` layout and the reverse index. +/// +/// The `vortex.indexed` layout is registered by default in every [`LayoutSession`], the same as +/// any other built-in layout, so only the edition declaration needs setting up here. +/// +/// Deliberately not a shared global session: sessions clone by sharing one `Arc`, so registering +/// into a shared session would leak between tests, and +/// [`unregistered_index_kind_falls_back_to_the_data_child`] depends on two sessions with different +/// index registries. +fn session() -> VortexSession { + let session = array_session() + .with::() + .with::() + .with::(); + session + .register_edition(&TEST_EDITION_DECLARATION) + .expect("test edition declaration should be valid"); + session + .enable_edition(TEST_EDITION) + .expect("test edition was just registered"); + session +} + +/// A session that additionally knows the reverse index. +fn session_with_reverse_index() -> VortexSession { + let session = session(); + session.indexes().register(ReverseIndex::new_ref()); + session +} + +fn value_column() -> VortexResult { + let values: PrimitiveArray = VALUES.into_iter().collect(); + struct_column(values.into_array()) +} + +fn struct_column(value_column: ArrayRef) -> VortexResult { + Ok(StructArray::from_fields([(VALUE_FIELD, value_column)].as_slice())?.into_array()) +} + +/// A write strategy that attaches `configs` to the value column. +/// +/// The indexed wrapper sits directly above repartitioning — the same slot `ZonedStrategy` +/// occupies — so it sees whole chunks in row order and knows the data child's block size. +fn strategy(configs: Vec) -> Arc { + let data = RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: BLOCK_LEN, + block_size_target: None, + canonicalize: false, + }, + ); + let indexed = IndexedStrategy::new(data, FlatLayoutStrategy::default(), configs) + .with_data_block_len(BLOCK_LEN as u64); + + WriteStrategyBuilder::default() + .with_row_block_size(BLOCK_LEN) + .with_field_writer(FieldPath::from_name(VALUE_FIELD), Arc::new(indexed)) + .build() +} + +async fn write_file( + session: &VortexSession, + configs: Vec, +) -> VortexResult { + write_file_with_column(session, configs, value_column()?).await +} + +async fn write_file_with_column( + session: &VortexSession, + configs: Vec, + column: ArrayRef, +) -> VortexResult { + let mut bytes = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(strategy(configs)) + .write(&mut bytes, column.to_array_stream()) + .await?; + Ok(bytes.freeze()) +} + +fn value_reader(file: &VortexFile) -> VortexResult> { + let value_layout = file + .footer() + .layout() + .slot(1)? + .vortex_expect("root struct always has the value column"); + value_layout.new_reader( + VALUE_FIELD.into(), + file.segment_source(), + file.session(), + &Default::default(), + ) +} + +/// The mask the value column's reader produces for `value == target`. +/// +/// An `Exact` plan serves `filter_evaluation` directly, so the result is the index's own answer +/// rather than the data child's, intersected with `input`. +async fn exact_mask(file: &VortexFile, target: i32, input: MaskFuture) -> VortexResult { + let reader = value_reader(file)?; + let row_count = reader.row_count(); + let filter = eq(root(), lit(target)).bind(reader.dtype())?; + + reader + .filter_evaluation(&(0..row_count), &filter, input)? + .await +} + +async fn scan_matching(file: &VortexFile, target: i32) -> VortexResult> { + let filter = eq(col(VALUE_FIELD), lit(target)).bind(file.dtype())?; + let result = file + .scan()? + .with_filter(filter) + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = file.session().create_execution_ctx(); + let values = result + .execute::(&mut ctx)? + .unmasked_field_by_name(VALUE_FIELD)? + .clone() + .execute::(&mut ctx)?; + + Ok(values.as_slice::().to_vec()) +} + +#[tokio::test] +async fn unregistered_index_kind_falls_back_to_the_data_child() -> VortexResult<()> { + // Written by a session that knows the reverse index... + let bytes = write_file(&session_with_reverse_index(), reverse_index_configs()).await?; + + // ...and read by one that does not. The spec goes inert, nothing probes the index child, and + // the data child answers everything — indexes are strictly optional accelerators. + let file = session().open_options().open_buffer(bytes)?; + + assert_eq!(scan_matching(&file, 20).await?, vec![20, 20]); + Ok(()) +} + +#[tokio::test] +async fn exact_index_answers_the_filter_itself() -> VortexResult<()> { + let session = session_with_reverse_index(); + let bytes = write_file(&session, reverse_index_configs()).await?; + let file = session.open_options().open_buffer(bytes)?; + + let mask = exact_mask(&file, 20, MaskFuture::new_true(VALUES.len())).await?; + assert_eq!( + mask, + Mask::from_iter((0..VALUES.len()).map(|row| VALUES[row] == 20)) + ); + + // The post-condition is that the result is intersected with the input mask, so an input that + // excludes both matches must yield nothing. + let excluded = exact_mask( + &file, + 20, + MaskFuture::ready(Mask::from_iter( + (0..VALUES.len()).map(|row| VALUES[row] != 20), + )), + ) + .await?; + assert!(excluded.all_false()); + + Ok(()) +} + +#[tokio::test] +async fn absent_value_resolves_to_no_matches() -> VortexResult<()> { + let session = session_with_reverse_index(); + let bytes = write_file(&session, reverse_index_configs()).await?; + let file = session.open_options().open_buffer(bytes)?; + + // The index claims the expression (it is an equality over an integer column) but finds no + // posting for 999, distinct from "no index claimed this at all". + let mask = exact_mask(&file, 999, MaskFuture::new_true(VALUES.len())).await?; + assert!(mask.all_false()); + assert_eq!(scan_matching(&file, 999).await?, Vec::::new()); + Ok(()) +} + +#[tokio::test] +async fn layout_carries_one_auxiliary_child_per_index() -> VortexResult<()> { + let session = session(); + let bytes = write_file(&session, reverse_index_configs()).await?; + let file = session.open_options().open_buffer(bytes)?; + + let value_layout = file + .footer() + .layout() + .slot(1)? + .vortex_expect("root struct always has the value column"); + assert_eq!(value_layout.encoding_id().as_str(), INDEXED_LAYOUT_ID); + + assert_eq!( + (0..value_layout.nslots()) + .filter_map(|slot| value_layout.slot_type(slot)) + .collect::>(), + vec![ + LayoutChildType::Transparent("data".into()), + LayoutChildType::Auxiliary(format!("index:{}", crate::REVERSE_INDEX_ID).into()), + ], + ); + + // Index content is an ordinary layout tree, so it inherits chunking and zone maps for free. + let index_child = value_layout + .slot(1)? + .vortex_expect("a reverse index was configured"); + assert_eq!( + index_child + .dtype() + .as_struct_fields() + .names() + .iter() + .map(|name| name.to_string()) + .collect::>(), + vec!["key".to_string(), "postings".to_string()], + ); + // 11 distinct values across 12 rows (20 repeats once). + assert_eq!(index_child.row_count(), 11); + Ok(()) +} + +/// The reverse index is not limited to integer columns: it decodes keys as +/// [`Scalar`](vortex_array::scalar::Scalar)s, so any dtype the layout can carry works the same +/// way. +#[tokio::test] +async fn string_valued_column_is_indexed_exactly() -> VortexResult<()> { + const STRING_VALUES: [&str; 6] = ["b", "a", "b", "c", "a", "d"]; + + let session = session_with_reverse_index(); + let values = VarBinViewArray::from_iter_str(STRING_VALUES); + let bytes = write_file_with_column( + &session, + reverse_index_configs(), + struct_column(values.into_array())?, + ) + .await?; + let file = session.open_options().open_buffer(bytes)?; + + let reader = value_reader(&file)?; + let row_count = reader.row_count(); + let filter = eq(root(), lit("b")).bind(reader.dtype())?; + let mask = reader + .filter_evaluation( + &(0..row_count), + &filter, + MaskFuture::new_true(STRING_VALUES.len()), + )? + .await?; + + assert_eq!( + mask, + Mask::from_iter(STRING_VALUES.iter().map(|value| *value == "b")) + ); + Ok(()) +} From df008c60d3b46d04b8855fb998bc2cf8f880249d Mon Sep 17 00:00:00 2001 From: Thor Date: Mon, 31 Aug 2026 15:00:35 -0500 Subject: [PATCH 3/7] Datafusion tests with an indexed layout --- vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/tests/indexed_layout.rs | 339 ++++++++++++++++++ vortex-datafusion/src/tests/mod.rs | 1 + 3 files changed, 341 insertions(+) create mode 100644 vortex-datafusion/src/tests/indexed_layout.rs diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 49fb22d4f59..9c22ddd4bec 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -48,6 +48,7 @@ rstest = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "fs"] } url = { workspace = true } +vortex-reverse-index = { workspace = true } [lints] workspace = true diff --git a/vortex-datafusion/src/tests/indexed_layout.rs b/vortex-datafusion/src/tests/indexed_layout.rs new file mode 100644 index 00000000000..adc030cd9fc --- /dev/null +++ b/vortex-datafusion/src/tests/indexed_layout.rs @@ -0,0 +1,339 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Test that DataFusion can query a file whose column uses the `vortex.indexed` layout. + +use std::sync::Arc; + +use arrow_array::record_batch; +use datafusion::arrow::array::RecordBatch; +use datafusion::assert_batches_sorted_eq; +use datafusion::physical_plan::collect; +use datafusion_expr::col; +use datafusion_expr::lit; +use datafusion_physical_plan::metrics::MetricsSet; +use rstest::rstest; +use vortex::VortexSessionDefault; +use vortex::dtype::FieldPath; +use vortex::editions::Edition; +use vortex::editions::EditionDeclaration; +use vortex::editions::EditionId; +use vortex::editions::EditionMember; +use vortex::editions::EditionSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::io::VortexWrite; +use vortex::io::object_store::ObjectStoreWrite; +use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex::layout::layouts::indexed::INDEXED_LAYOUT_ID; +use vortex::layout::layouts::indexed::IndexConfig; +use vortex::layout::layouts::indexed::IndexSessionExt; +use vortex::layout::layouts::indexed::IndexedStrategy; +use vortex::layout::layouts::repartition::RepartitionStrategy; +use vortex::layout::layouts::repartition::RepartitionWriterOptions; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; +use vortex_reverse_index::ReverseIndex; + +use crate::VortexFormatFactory; +use crate::VortexTableOptions; +use crate::common_tests::TestSessionContext; +use crate::metrics::VortexMetricsFinder; + +const VALUE_FIELD: &str = "value"; +/// Small enough that the 12-row batch spans three blocks under the indexed layout, leaving the +/// index a real pruning decision to make instead of degenerating to a single block. +const BLOCK_LEN: usize = 4; + +/// The default session doesn't enable `vortex.indexed` for writing — it's a layout prototype, not +/// part of any frozen `core` edition — so this test registers a tiny edition just for it. +const INDEXED_TEST_EDITION: EditionId = + EditionId::new("vortex-datafusion-indexed-test", 2026, 1, 0); + +static INDEXED_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: INDEXED_TEST_EDITION, + min_vortex_version: None, + }, + added: &[EditionMember::layout(&INDEXED_LAYOUT_ID)], +}; + +/// A session that can write and read the `vortex.indexed` layout, with a [`ReverseIndex`] +/// registered as an index kind. +fn session_with_indexed_layout() -> anyhow::Result { + let session = VortexSession::default(); + session.register_edition(&INDEXED_TEST_DECLARATION)?; + session.enable_edition(INDEXED_TEST_EDITION)?; + session.indexes().register(ReverseIndex::new_ref()); + Ok(session) +} + +/// A write strategy that attaches a [`ReverseIndex`] to the `value` field, chunked into blocks of +/// `block_len` rows. +fn indexed_write_strategy(block_len: usize) -> Arc { + let data = RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: block_len, + block_size_target: None, + canonicalize: false, + }, + ); + let indexed = IndexedStrategy::new( + data, + FlatLayoutStrategy::default(), + vec![IndexConfig::with_defaults(ReverseIndex::new_ref())], + ) + .with_data_block_len(block_len as u64); + + WriteStrategyBuilder::default() + .with_row_block_size(block_len) + .with_field_writer(FieldPath::from_name(VALUE_FIELD), Arc::new(indexed)) + .build() +} + +/// `20` repeats at rows 1 and 9; `999` never appears. With [`BLOCK_LEN`] of 4 those land in +/// blocks 0 and 2, leaving block 1 fully prunable by an exact equality index. +fn test_batch() -> anyhow::Result { + Ok(record_batch!(( + "value", + Int32, + vec![ + Some(10), + Some(20), + Some(30), + Some(40), + Some(50), + Some(60), + Some(70), + Some(80), + Some(90), + Some(20), + Some(100), + Some(110) + ] + ))?) +} + +async fn write_indexed_batch( + ctx: &TestSessionContext, + session: &VortexSession, + path: &str, + batch: &RecordBatch, + block_len: usize, +) -> anyhow::Result<()> { + let array = session + .arrow() + .from_arrow_record_batch(batch.clone(), &batch.schema())?; + let mut write = ObjectStoreWrite::new(Arc::clone(&ctx.store), &path.into()).await?; + session + .write_options() + .with_strategy(indexed_write_strategy(block_len)) + .write(&mut write, array.to_array_stream()) + .await?; + write.shutdown().await?; + Ok(()) +} + +/// `20` repeats at rows 1 and 9; `999` never appears. An equality filter on `value` is exactly +/// what [`ReverseIndex::plan`](vortex_layout::layouts::indexed::IndexVTable::plan) claims, so this +/// exercises both "claimed and found" and "claimed but no match" through DataFusion's predicate +/// pushdown into the indexed column. +#[rstest] +#[tokio::test] +async fn test_query_over_indexed_column( + #[values(false, true)] projection_pushdown: bool, +) -> anyhow::Result<()> { + let session = session_with_indexed_layout()?; + + let opts = VortexTableOptions { + projection_pushdown, + ..Default::default() + }; + let factory = Arc::new(VortexFormatFactory::new_with_options(session.clone(), opts)); + let ctx = TestSessionContext::new_with_factory(factory); + + let batch = test_batch()?; + write_indexed_batch(&ctx, &session, "files/indexed.vortex", &batch, BLOCK_LEN).await?; + + let schema = batch.schema(); + let provider = ctx + .table_provider("indexed_tbl", "/files/", schema.as_ref().clone()) + .await?; + let table = ctx.session.read_table(provider)?; + + let matches = table + .clone() + .filter(col(VALUE_FIELD).eq(lit(20)))? + .collect() + .await?; + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 20 |", + "| 20 |", + "+-------+", + ], + &matches + ); + + let absent = table + .filter(col(VALUE_FIELD).eq(lit(999)))? + .collect() + .await?; + assert!(absent.iter().all(|batch| batch.num_rows() == 0)); + + Ok(()) +} + +/// Rows per block for [`pruning_test_batch`], and the number of blocks it spans. +const PRUNING_BLOCK_LEN: usize = 100; +const PRUNING_BLOCK_COUNT: usize = 5; +const PAYLOAD_FIELD: &str = "payload"; + +/// `value` is clustered by block: block `k` (rows `k * PRUNING_BLOCK_LEN` to +/// `(k + 1) * PRUNING_BLOCK_LEN - 1`) holds nothing but the constant `k + 1`. Filtering on a +/// single value therefore claims exactly one block as a match and leaves every other block fully +/// prunable, unlike [`test_batch`]'s handful of rows, where the pruned savings are too small to +/// stand out over the index's own storage overhead. +/// +/// `payload` carries the row index and is neither indexed nor filtered on — selecting it instead +/// of `value` means the data child only needs the rows the index's mask actually claims, rather +/// than every row `value` itself requires to be decoded and re-checked against the filter. +fn pruning_test_batch() -> anyhow::Result { + let row_count = PRUNING_BLOCK_LEN * PRUNING_BLOCK_COUNT; + let mut values: Vec> = Vec::with_capacity(row_count); + for block in 0..PRUNING_BLOCK_COUNT { + let value = i32::try_from(block)? + 1; + values.extend(std::iter::repeat_n(Some(value), PRUNING_BLOCK_LEN)); + } + let payload: Vec> = (0..row_count) + .map(|row| i32::try_from(row).map(Some)) + .collect::>()?; + + Ok(record_batch!( + ("value", Int32, values), + ("payload", Int32, payload) + )?) +} + +/// Total bytes read from storage across every Vortex-backed data source in the plan, per +/// [`InstrumentedReadAt`](vortex::io::VortexReadAt)'s `vortex.io.read.total_size` counter. +fn total_bytes_read(metrics_sets: &[MetricsSet]) -> usize { + metrics_sets + .iter() + .filter_map(|set| set.sum_by_name("vortex.io.read.total_size")) + .map(|value| value.as_usize()) + .sum() +} + +/// Runs `value = ` against a freshly written copy of `batch`, executing the physical plan +/// directly (rather than through [`DataFrame::collect`]) so the same plan instance can be +/// inspected for metrics afterward. +/// +/// [`DataFrame::collect`]: datafusion::dataframe::DataFrame::collect +async fn run_equality_filter( + write_session: &VortexSession, + read_session: VortexSession, + projection_pushdown: bool, + batch: &RecordBatch, + block_len: usize, + target: i32, +) -> anyhow::Result<(Vec, usize)> { + let opts = VortexTableOptions { + projection_pushdown, + ..Default::default() + }; + let factory = Arc::new(VortexFormatFactory::new_with_options(read_session, opts)); + let ctx = TestSessionContext::new_with_factory(factory); + + write_indexed_batch( + &ctx, + write_session, + "files/indexed.vortex", + batch, + block_len, + ) + .await?; + + let schema = batch.schema(); + let provider = ctx + .table_provider("indexed_tbl", "/files/", schema.as_ref().clone()) + .await?; + ctx.session.register_table("indexed_tbl", provider)?; + + let df = ctx + .session + .sql(&format!( + "SELECT {PAYLOAD_FIELD} FROM indexed_tbl WHERE {VALUE_FIELD} = {target}" + )) + .await?; + let physical_plan = ctx + .session + .state() + .create_physical_plan(df.logical_plan()) + .await?; + let results = collect(Arc::clone(&physical_plan), ctx.session.task_ctx()).await?; + let bytes_read = total_bytes_read(&VortexMetricsFinder::find_all(physical_plan.as_ref())); + + Ok((results, bytes_read)) +} + +/// Correctness alone can't distinguish "the index answered the filter exactly" from "the index +/// was silently ignored and a full scan happened to get the right answer anyway" — both produce +/// identical query results. This test tells them apart by comparing bytes read from storage for +/// the same file and filter, once with [`ReverseIndex`] registered for reading and once without +/// (which forces the documented fallback: `IndexedReader::plan_probe` treats an unregistered +/// index kind's spec as inert). [`pruning_test_batch`] puts a single value in each block, so a +/// real exact-index probe lets the scan skip every block but one, while the unregistered run must +/// decode all of them to filter. +#[rstest] +#[tokio::test] +async fn test_index_avoids_reading_pruned_blocks( + #[values(false, true)] projection_pushdown: bool, +) -> anyhow::Result<()> { + let write_session = session_with_indexed_layout()?; + let batch = pruning_test_batch()?; + // Block 0, not a middle block: object_store's read coalescing merges nearby byte ranges into + // one physical read, so a "hole" in the middle of the file still gets pulled in as padding. A + // match confined to the first block is a genuine prefix, so skipping the rest of the file + // actually shrinks the range requested. + let target = 1; + + let (with_index_rows, with_index_bytes) = run_equality_filter( + &write_session, + session_with_indexed_layout()?, + projection_pushdown, + &batch, + PRUNING_BLOCK_LEN, + target, + ) + .await?; + let (without_index_rows, without_index_bytes) = run_equality_filter( + &write_session, + VortexSession::default(), + projection_pushdown, + &batch, + PRUNING_BLOCK_LEN, + target, + ) + .await?; + + assert_eq!(with_index_rows, without_index_rows); + let matched_rows: usize = with_index_rows.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(matched_rows, PRUNING_BLOCK_LEN); + + assert!( + with_index_bytes < without_index_bytes, + "expected the registered index to skip the prunable blocks and read fewer bytes than the \ + unregistered fallback, got {with_index_bytes} (indexed) vs {without_index_bytes} \ + (fallback)" + ); + + Ok(()) +} diff --git a/vortex-datafusion/src/tests/mod.rs b/vortex-datafusion/src/tests/mod.rs index e3421f63c67..da4b9f597be 100644 --- a/vortex-datafusion/src/tests/mod.rs +++ b/vortex-datafusion/src/tests/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod indexed_layout; mod nested_projection; mod schema_evolution; From f81f922c4bbdcfc44b4846a1f3207d1c15fbd56f Mon Sep 17 00:00:00 2001 From: Thor Date: Mon, 31 Aug 2026 15:00:54 -0500 Subject: [PATCH 4/7] Cargo.toml --- Cargo.lock | 20 ++++++++++++++++++++ Cargo.toml | 2 ++ 2 files changed, 22 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..83d62046cd1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10904,6 +10904,7 @@ dependencies = [ "url", "vortex", "vortex-arrow", + "vortex-reverse-index", "vortex-utils", ] @@ -11254,6 +11255,7 @@ dependencies = [ "paste", "pin-project-lite", "prost 0.14.4", + "roaring", "rstest", "rustc-hash", "sketches-ddsketch", @@ -11420,6 +11422,24 @@ dependencies = [ "vortex-python-abi", ] +[[package]] +name = "vortex-reverse-index" +version = "0.1.0" +dependencies = [ + "roaring", + "tokio", + "vortex-array", + "vortex-buffer", + "vortex-edition", + "vortex-error", + "vortex-file", + "vortex-io", + "vortex-layout", + "vortex-mask", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-row" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index ae164db13c2..7205427feeb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,7 @@ members = [ "benchmarks/random-access-bench", "benchmarks/string-bench", "vortex-spatial", + "vortex-reverse-index", ] exclude = ["java/testfiles", "wasm-test"] resolver = "3" @@ -323,6 +324,7 @@ vortex-onpair = { version = "0.1.0", path = "./encodings/onpair", default-featur vortex-parquet-variant = { version = "0.1.0", path = "./encodings/parquet-variant" } vortex-pco = { version = "0.1.0", path = "./encodings/pco", default-features = false } vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features = false } +vortex-reverse-index = { version = "0.1.0", path = "./vortex-reverse-index", default-features = false } vortex-row = { version = "0.1.0", path = "./vortex-row", default-features = false } vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false } vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false } From aaf40a24cab906651ca262e213ebdc8265c6306e Mon Sep 17 00:00:00 2001 From: Thor Date: Tue, 1 Sep 2026 12:14:39 -0500 Subject: [PATCH 5/7] claude: Move vortex-reverse-index into vortex-layout examples vortex-reverse-index was a standalone worked example of the vortex.indexed IndexVTable contract; move it to vortex-layout/examples/reverse_index, following the plain cargo-example convention used by vortex-ffi/examples instead of shipping it as its own workspace crate. vortex-datafusion could no longer dev-depend on it as a library once it became a cargo example, so its DataFusion integration test moves along with it (examples/reverse_index/datafusion_tests.rs), rebuilt against vortex-datafusion's public API and a manually assembled session instead of the vortex facade crate, which vortex-layout cannot depend on. Signed-off-by: "Thor" --- Cargo.lock | 30 +-- Cargo.toml | 2 - vortex-datafusion/Cargo.toml | 1 - vortex-datafusion/src/tests/mod.rs | 1 - vortex-layout/Cargo.toml | 17 +- .../reverse_index/datafusion_tests.rs | 207 ++++++++++++++---- .../examples/reverse_index/main.rs | 150 +++++++++++++ .../examples/reverse_index}/tests.rs | 4 +- vortex-layout/src/layouts/indexed/mod.rs | 4 +- vortex-layout/src/layouts/indexed/tests.rs | 5 +- vortex-reverse-index/Cargo.toml | 33 --- 11 files changed, 344 insertions(+), 110 deletions(-) rename vortex-datafusion/src/tests/indexed_layout.rs => vortex-layout/examples/reverse_index/datafusion_tests.rs (60%) rename vortex-reverse-index/src/lib.rs => vortex-layout/examples/reverse_index/main.rs (62%) rename {vortex-reverse-index/src => vortex-layout/examples/reverse_index}/tests.rs (98%) delete mode 100644 vortex-reverse-index/Cargo.toml diff --git a/Cargo.lock b/Cargo.lock index 83d62046cd1..fa0f6f40b64 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10904,7 +10904,6 @@ dependencies = [ "url", "vortex", "vortex-arrow", - "vortex-reverse-index", "vortex-utils", ] @@ -11237,6 +11236,7 @@ dependencies = [ name = "vortex-layout" version = "0.1.0" dependencies = [ + "anyhow", "arcref", "arrow-array 59.2.0", "arrow-schema 59.2.0", @@ -11244,12 +11244,18 @@ dependencies = [ "async-trait", "bit-vec", "codspeed-divan-compat", + "datafusion 55.0.0", + "datafusion-catalog 55.0.0", + "datafusion-common 55.0.0", + "datafusion-expr 55.0.0", + "datafusion-physical-plan 55.0.0", "flatbuffers", "futures", "insta", "itertools 0.14.0", "kanal", "moka", + "object_store", "once_cell", "parking_lot", "paste", @@ -11263,12 +11269,16 @@ dependencies = [ "termtree", "tokio", "tracing", + "url", "uuid", "vortex-array", "vortex-arrow", "vortex-btrblocks", "vortex-buffer", + "vortex-datafusion", + "vortex-edition", "vortex-error", + "vortex-file", "vortex-flatbuffers", "vortex-io", "vortex-mask", @@ -11422,24 +11432,6 @@ dependencies = [ "vortex-python-abi", ] -[[package]] -name = "vortex-reverse-index" -version = "0.1.0" -dependencies = [ - "roaring", - "tokio", - "vortex-array", - "vortex-buffer", - "vortex-edition", - "vortex-error", - "vortex-file", - "vortex-io", - "vortex-layout", - "vortex-mask", - "vortex-session", - "vortex-utils", -] - [[package]] name = "vortex-row" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 7205427feeb..ae164db13c2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,7 +71,6 @@ members = [ "benchmarks/random-access-bench", "benchmarks/string-bench", "vortex-spatial", - "vortex-reverse-index", ] exclude = ["java/testfiles", "wasm-test"] resolver = "3" @@ -324,7 +323,6 @@ vortex-onpair = { version = "0.1.0", path = "./encodings/onpair", default-featur vortex-parquet-variant = { version = "0.1.0", path = "./encodings/parquet-variant" } vortex-pco = { version = "0.1.0", path = "./encodings/pco", default-features = false } vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features = false } -vortex-reverse-index = { version = "0.1.0", path = "./vortex-reverse-index", default-features = false } vortex-row = { version = "0.1.0", path = "./vortex-row", default-features = false } vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false } vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false } diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 9c22ddd4bec..49fb22d4f59 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -48,7 +48,6 @@ rstest = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "fs"] } url = { workspace = true } -vortex-reverse-index = { workspace = true } [lints] workspace = true diff --git a/vortex-datafusion/src/tests/mod.rs b/vortex-datafusion/src/tests/mod.rs index da4b9f597be..e3421f63c67 100644 --- a/vortex-datafusion/src/tests/mod.rs +++ b/vortex-datafusion/src/tests/mod.rs @@ -1,6 +1,5 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod indexed_layout; mod nested_projection; mod schema_evolution; diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index ba095231cde..14fb53e2fe8 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -55,14 +55,27 @@ vortex-session = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] +anyhow = { workspace = true } +arrow-array = { workspace = true } +datafusion = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-common = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-physical-plan = { workspace = true } divan = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } +object_store = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } -tokio = { workspace = true, features = ["rt", "macros"] } +tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros", "time"] } +url = { workspace = true } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } -vortex-io = { path = "../vortex-io", features = ["tokio"] } +vortex-buffer = { workspace = true } +vortex-datafusion = { workspace = true } +vortex-edition = { workspace = true } +vortex-file = { workspace = true } +vortex-io = { path = "../vortex-io", features = ["tokio", "object_store"] } [features] default = ["wasm-bindgen"] diff --git a/vortex-datafusion/src/tests/indexed_layout.rs b/vortex-layout/examples/reverse_index/datafusion_tests.rs similarity index 60% rename from vortex-datafusion/src/tests/indexed_layout.rs rename to vortex-layout/examples/reverse_index/datafusion_tests.rs index adc030cd9fc..0c6122045bd 100644 --- a/vortex-datafusion/src/tests/indexed_layout.rs +++ b/vortex-layout/examples/reverse_index/datafusion_tests.rs @@ -1,75 +1,136 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Test that DataFusion can query a file whose column uses the `vortex.indexed` layout. +//! Test that DataFusion can query a file whose column uses the `vortex.indexed` layout, with +//! [`ReverseIndex`] answering an equality filter. +//! +//! This lives alongside the [`ReverseIndex`] example rather than in `vortex-datafusion` because +//! `vortex-layout` sits below `vortex-datafusion` in the dependency graph: `vortex-datafusion` +//! cannot depend on an example that lives in `vortex-layout`, but an example's `dev-dependencies` +//! may freely depend on `vortex-datafusion`. use std::sync::Arc; use arrow_array::record_batch; use datafusion::arrow::array::RecordBatch; use datafusion::assert_batches_sorted_eq; +use datafusion::datasource::provider::DefaultTableFactory; +use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::collect; +use datafusion::prelude::SessionContext; +use datafusion_catalog::TableProvider; +use datafusion_common::DFSchema; +use datafusion_common::GetExt; +use datafusion_expr::CreateExternalTable; use datafusion_expr::col; use datafusion_expr::lit; use datafusion_physical_plan::metrics::MetricsSet; +use object_store::ObjectStore; +use object_store::memory::InMemory; use rstest::rstest; -use vortex::VortexSessionDefault; -use vortex::dtype::FieldPath; -use vortex::editions::Edition; -use vortex::editions::EditionDeclaration; -use vortex::editions::EditionId; -use vortex::editions::EditionMember; -use vortex::editions::EditionSessionExt; -use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; -use vortex::io::VortexWrite; -use vortex::io::object_store::ObjectStoreWrite; -use vortex::layout::LayoutStrategy; -use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; -use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; -use vortex::layout::layouts::indexed::INDEXED_LAYOUT_ID; -use vortex::layout::layouts::indexed::IndexConfig; -use vortex::layout::layouts::indexed::IndexSessionExt; -use vortex::layout::layouts::indexed::IndexedStrategy; -use vortex::layout::layouts::repartition::RepartitionStrategy; -use vortex::layout::layouts::repartition::RepartitionWriterOptions; -use vortex::session::VortexSession; +use url::Url; +use vortex_array::array_session; +use vortex_array::dtype::FieldPath; use vortex_arrow::ArrowSessionExt; -use vortex_reverse_index::ReverseIndex; - -use crate::VortexFormatFactory; -use crate::VortexTableOptions; -use crate::common_tests::TestSessionContext; -use crate::metrics::VortexMetricsFinder; +use vortex_datafusion::VortexFormatFactory; +use vortex_datafusion::VortexTableOptions; +use vortex_datafusion::metrics::VortexMetricsFinder; +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::VortexWrite; +use vortex_io::object_store::ObjectStoreWrite; +use vortex_io::session::RuntimeSession; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::indexed::INDEXED_LAYOUT_ID; +use vortex_layout::layouts::indexed::IndexConfig; +use vortex_layout::layouts::indexed::IndexSessionExt; +use vortex_layout::layouts::indexed::IndexedStrategy; +use vortex_layout::layouts::repartition::RepartitionStrategy; +use vortex_layout::layouts::repartition::RepartitionWriterOptions; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; + +use crate::ReverseIndex; const VALUE_FIELD: &str = "value"; /// Small enough that the 12-row batch spans three blocks under the indexed layout, leaving the /// index a real pruning decision to make instead of degenerating to a single block. const BLOCK_LEN: usize = 4; -/// The default session doesn't enable `vortex.indexed` for writing — it's a layout prototype, not -/// part of any frozen `core` edition — so this test registers a tiny edition just for it. +/// The array/layout encodings this test needs to write, converted from Arrow via +/// [`vortex_arrow::ArrowSessionExt`]. +/// +/// The default Vortex file writer only permits array/layout ids covered by the session's enabled +/// editions, but those first-party declarations live in the `vortex` facade crate, which +/// `vortex-layout` cannot depend on. Declaring and enabling a tiny edition here is the local +/// equivalent. const INDEXED_TEST_EDITION: EditionId = - EditionId::new("vortex-datafusion-indexed-test", 2026, 1, 0); + EditionId::new("vortex-layout-reverse-index-datafusion-test", 2026, 1, 0); static INDEXED_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { edition: Edition { id: INDEXED_TEST_EDITION, min_vortex_version: None, }, - added: &[EditionMember::layout(&INDEXED_LAYOUT_ID)], + added: &[ + EditionMember::array(&"vortex.struct"), + EditionMember::array(&"vortex.primitive"), + // The index child's postings column is serialized roaring bitmaps, written as varbinview. + EditionMember::array(&"vortex.varbinview"), + // `payload`'s sequential row-index values compress with the sequence encoding; `value`'s + // single repeated value per block compresses with the constant encoding. + EditionMember::array(&"vortex.sequence"), + EditionMember::array(&"vortex.constant"), + EditionMember::layout(&"vortex.struct"), + EditionMember::layout(&"vortex.chunked"), + EditionMember::layout(&"vortex.flat"), + EditionMember::layout(&"vortex.zoned"), + EditionMember::layout(&INDEXED_LAYOUT_ID), + // Zone-map stats over each chunk need min/max/null-count, computed as aggregates during + // writing. + EditionMember::aggregate(&"vortex.min"), + EditionMember::aggregate(&"vortex.max"), + EditionMember::aggregate(&"vortex.null_count"), + ], }; /// A session that can write and read the `vortex.indexed` layout, with a [`ReverseIndex`] /// registered as an index kind. +/// +/// [`array_session`] already bundles everything [`vortex::VortexSessionDefault::default`] would +/// (arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, and memory); +/// this only adds the layout, runtime, and edition state that live in higher-level crates. fn session_with_indexed_layout() -> anyhow::Result { - let session = VortexSession::default(); + let session = array_session() + .with::() + .with::() + .with::(); + vortex_arrow::initialize(&session); + vortex_sequence::initialize(&session); session.register_edition(&INDEXED_TEST_DECLARATION)?; session.enable_edition(INDEXED_TEST_EDITION)?; session.indexes().register(ReverseIndex::new_ref()); Ok(session) } +/// A session that can read the `vortex.indexed` layout, but without a [`ReverseIndex`] +/// registered — the fallback path `IndexedReader::plan_probe` takes when an index kind's spec +/// goes unclaimed. +fn session_without_reverse_index() -> VortexSession { + let session = array_session() + .with::() + .with::() + .with::(); + vortex_arrow::initialize(&session); + vortex_sequence::initialize(&session); + session +} + /// A write strategy that attaches a [`ReverseIndex`] to the `value` field, chunked into blocks of /// `block_len` rows. fn indexed_write_strategy(block_len: usize) -> Arc { @@ -118,6 +179,64 @@ fn test_batch() -> anyhow::Result { ))?) } +/// A minimal DataFusion harness over an in-memory [`ObjectStore`], built only from +/// `vortex-datafusion`'s public API (the crate's own richer `TestSessionContext` is `#[cfg(test)]` +/// only, so it isn't visible outside `vortex-datafusion` itself). +struct TestSessionContext { + store: Arc, + session: SessionContext, +} + +impl TestSessionContext { + fn new_with_factory(factory: Arc) -> Self { + let store = Arc::new(InMemory::new()); + let mut session_state_builder = SessionStateBuilder::new() + .with_default_features() + .with_table_factory( + factory.get_ext().to_uppercase(), + Arc::new(DefaultTableFactory::new()), + ) + .with_object_store( + &Url::try_from("file://").unwrap(), + Arc::::clone(&store), + ); + + if let Some(file_formats) = session_state_builder.file_formats() { + file_formats.push(factory as _); + } + + let session = + SessionContext::new_with_state(session_state_builder.build()).enable_url_table(); + + Self { store, session } + } + + async fn table_provider( + &self, + name: &str, + location: impl Into, + schema: S, + ) -> anyhow::Result> + where + DFSchema: TryFrom, + anyhow::Error: From<>::Error>, + { + let factory = self.session.table_factory("VORTEX").unwrap(); + + let cmd = CreateExternalTable::builder( + name, + location.into(), + "vortex", + DFSchema::try_from(schema)?.into(), + ) + .build(); + + let table = factory.create(&self.session.state(), &cmd).await?; + + Ok(table) + } +} + async fn write_indexed_batch( ctx: &TestSessionContext, session: &VortexSession, @@ -149,10 +268,8 @@ async fn test_query_over_indexed_column( ) -> anyhow::Result<()> { let session = session_with_indexed_layout()?; - let opts = VortexTableOptions { - projection_pushdown, - ..Default::default() - }; + let mut opts = VortexTableOptions::default(); + opts.projection_pushdown = projection_pushdown; let factory = Arc::new(VortexFormatFactory::new_with_options(session.clone(), opts)); let ctx = TestSessionContext::new_with_factory(factory); @@ -223,7 +340,7 @@ fn pruning_test_batch() -> anyhow::Result { } /// Total bytes read from storage across every Vortex-backed data source in the plan, per -/// [`InstrumentedReadAt`](vortex::io::VortexReadAt)'s `vortex.io.read.total_size` counter. +/// `InstrumentedReadAt`'s `vortex.io.read.total_size` counter. fn total_bytes_read(metrics_sets: &[MetricsSet]) -> usize { metrics_sets .iter() @@ -245,10 +362,8 @@ async fn run_equality_filter( block_len: usize, target: i32, ) -> anyhow::Result<(Vec, usize)> { - let opts = VortexTableOptions { - projection_pushdown, - ..Default::default() - }; + let mut opts = VortexTableOptions::default(); + opts.projection_pushdown = projection_pushdown; let factory = Arc::new(VortexFormatFactory::new_with_options(read_session, opts)); let ctx = TestSessionContext::new_with_factory(factory); @@ -316,7 +431,7 @@ async fn test_index_avoids_reading_pruned_blocks( .await?; let (without_index_rows, without_index_bytes) = run_equality_filter( &write_session, - VortexSession::default(), + session_without_reverse_index(), projection_pushdown, &batch, PRUNING_BLOCK_LEN, diff --git a/vortex-reverse-index/src/lib.rs b/vortex-layout/examples/reverse_index/main.rs similarity index 62% rename from vortex-reverse-index/src/lib.rs rename to vortex-layout/examples/reverse_index/main.rs index 3922cd4a9d9..19968cb13d6 100644 --- a/vortex-reverse-index/src/lib.rs +++ b/vortex-layout/examples/reverse_index/main.rs @@ -1,3 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect( + clippy::use_debug, + reason = "example output favors a quick debug print over a formatted display" +)] //! A simple `value -> rows` reverse index over any column, demonstrating [`vortex_layout`]'s //! `vortex.indexed` layout [`IndexVTable`] contract with a minimal concrete index kind. //! @@ -6,6 +13,12 @@ //! `exact_value` index but generalized to any dtype instead of being fixed to `Utf8`. It exists as //! a worked, non-test example of a concrete index kind for the `vortex.indexed` layout prototype //! described in [vortex-data/vortex#9024](https://github.com/vortex-data/vortex/issues/9024). +//! +//! Run it with: +//! +//! ```ignore +//! cargo run -p vortex-layout --example reverse_index +//! ``` use std::sync::Arc; @@ -13,12 +26,16 @@ use roaring::RoaringBitmap; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; +use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; @@ -34,21 +51,44 @@ use vortex_array::search_sorted::SearchSortedSide; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_array::validity::Validity; +use vortex_buffer::ByteBufferMut; +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::session::RuntimeSession; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::indexed::INDEXED_LAYOUT_ID; use vortex_layout::layouts::indexed::IndexBuilder; +use vortex_layout::layouts::indexed::IndexConfig; use vortex_layout::layouts::indexed::IndexExactness; use vortex_layout::layouts::indexed::IndexId; use vortex_layout::layouts::indexed::IndexQueryPlan; use vortex_layout::layouts::indexed::IndexResolve; +use vortex_layout::layouts::indexed::IndexSessionExt; use vortex_layout::layouts::indexed::IndexVTable; use vortex_layout::layouts::indexed::IndexVTableRef; +use vortex_layout::layouts::indexed::IndexedStrategy; use vortex_layout::layouts::indexed::RowLocator; +use vortex_layout::layouts::repartition::RepartitionStrategy; +use vortex_layout::layouts::repartition::RepartitionWriterOptions; +use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_map::HashMap; +#[cfg(test)] +mod datafusion_tests; #[cfg(test)] mod tests; @@ -242,3 +282,113 @@ impl IndexResolve for Resolve { Ok(RowLocator::Rows(bitmap)) } } + +/// `20` repeats at rows 1 and 9; nothing else does, and `999` never appears. +const VALUES: [i32; 12] = [10, 20, 30, 40, 50, 60, 70, 80, 90, 20, 100, 110]; +const VALUE_FIELD: &str = "value"; + +/// The array/layout encodings this example needs to write. +/// +/// The default Vortex file writer only permits array/layout ids covered by the session's enabled +/// editions, but those first-party declarations live in the `vortex` facade crate, which this +/// example deliberately does not depend on (see the module doc comment). Declaring and enabling a +/// tiny edition here is the local equivalent. +const DEMO_EDITION: EditionId = EditionId::new("vortex-layout-reverse-index-example", 2026, 1, 0); + +static DEMO_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: DEMO_EDITION, + min_vortex_version: None, + }, + added: &[ + EditionMember::array(&"vortex.struct"), + EditionMember::array(&"vortex.primitive"), + // The index child's postings column is serialized roaring bitmaps, written as varbinview. + EditionMember::array(&"vortex.varbinview"), + EditionMember::layout(&"vortex.struct"), + EditionMember::layout(&"vortex.chunked"), + EditionMember::layout(&"vortex.flat"), + EditionMember::layout(&INDEXED_LAYOUT_ID), + ], +}; + +/// A session knowing the `vortex.indexed` layout and the reverse index. +fn demo_session() -> VortexResult { + let session = array_session() + .with::() + .with::() + .with::(); + session + .register_edition(&DEMO_EDITION_DECLARATION) + .map_err(|err| vortex_err!("{err}"))?; + session + .enable_edition(DEMO_EDITION) + .map_err(|err| vortex_err!("{err}"))?; + session.indexes().register(ReverseIndex::new_ref()); + Ok(session) +} + +/// A write strategy that attaches a [`ReverseIndex`] to the `value` field, chunked into blocks of +/// 4 rows so the 12-row demo column spans three blocks. +fn demo_write_strategy() -> Arc { + const BLOCK_LEN: usize = 4; + let data = RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: BLOCK_LEN, + block_size_target: None, + canonicalize: false, + }, + ); + let indexed = IndexedStrategy::new( + data, + FlatLayoutStrategy::default(), + vec![IndexConfig::with_defaults(ReverseIndex::new_ref())], + ) + .with_data_block_len(BLOCK_LEN as u64); + + WriteStrategyBuilder::default() + .with_row_block_size(BLOCK_LEN) + .with_field_writer(FieldPath::from_name(VALUE_FIELD), Arc::new(indexed)) + .build() +} + +#[tokio::main] +async fn main() -> VortexResult<()> { + let session = demo_session()?; + + let values: PrimitiveArray = VALUES.into_iter().collect(); + let column = + StructArray::from_fields([(VALUE_FIELD, values.into_array())].as_slice())?.into_array(); + + let mut bytes = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(demo_write_strategy()) + .write(&mut bytes, column.to_array_stream()) + .await?; + let bytes = bytes.freeze(); + + let file = session.open_options().open_buffer(bytes)?; + let filter = eq(col(VALUE_FIELD), lit(20)).bind(file.dtype())?; + let result = file + .scan()? + .with_filter(filter) + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = file.session().create_execution_ctx(); + let matches = result + .execute::(&mut ctx)? + .unmasked_field_by_name(VALUE_FIELD)? + .clone() + .execute::(&mut ctx)?; + + println!( + "rows matching value == 20, answered by the {REVERSE_INDEX_ID} index: {:?}", + matches.as_slice::() + ); + Ok(()) +} diff --git a/vortex-reverse-index/src/tests.rs b/vortex-layout/examples/reverse_index/tests.rs similarity index 98% rename from vortex-reverse-index/src/tests.rs rename to vortex-layout/examples/reverse_index/tests.rs index eee07001f8c..18d1b01a3ec 100644 --- a/vortex-reverse-index/src/tests.rs +++ b/vortex-layout/examples/reverse_index/tests.rs @@ -67,8 +67,8 @@ fn reverse_index_configs() -> Vec { /// /// The default Vortex file writer only permits array/layout ids covered by the session's enabled /// editions, but those first-party declarations live in the `vortex` facade crate, which this -/// out-of-tree crate deliberately does not depend on. Declaring and enabling a tiny test-only -/// edition here is the local equivalent. +/// example deliberately does not depend on. Declaring and enabling a tiny test-only edition here +/// is the local equivalent. const TEST_EDITION: EditionId = EditionId::new("vortex-reverse-index-test", 2026, 1, 0); static TEST_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { diff --git a/vortex-layout/src/layouts/indexed/mod.rs b/vortex-layout/src/layouts/indexed/mod.rs index ab1f5256803..1aa5b11e05a 100644 --- a/vortex-layout/src/layouts/indexed/mod.rs +++ b/vortex-layout/src/layouts/indexed/mod.rs @@ -28,8 +28,8 @@ //! //! The generic wrapper only: [`Indexed`], [`writer::IndexedStrategy`] and [`reader::IndexedReader`], //! plus the [`IndexVTable`] contract that index kinds implement. Concrete kinds are registered into -//! an [`session::IndexSession`] — see the `vortex-reverse-index` crate for a worked example, a -//! minimal equality index over integer columns. +//! an [`session::IndexSession`] — see the `reverse_index` example (`examples/reverse_index/`) for +//! a worked example, a minimal equality index over any column. pub mod index; pub(crate) mod reader; diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs index 3f544a0e51e..483a3cdb04f 100644 --- a/vortex-layout/src/layouts/indexed/tests.rs +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -1,7 +1,8 @@ //! End-to-end tests for the generic wrapper. //! -//! Concrete index kinds live in their own crates (see `vortex-reverse-index`), so these exercise -//! the machinery through a test-only [`exact_value::ExactValueIndex`] instead. +//! Concrete index kinds live elsewhere (see the `reverse_index` example under +//! `examples/reverse_index/`), so these exercise the machinery through a test-only +//! [`exact_value::ExactValueIndex`] instead. use std::sync::Arc; diff --git a/vortex-reverse-index/Cargo.toml b/vortex-reverse-index/Cargo.toml deleted file mode 100644 index ab2436aec95..00000000000 --- a/vortex-reverse-index/Cargo.toml +++ /dev/null @@ -1,33 +0,0 @@ -[package] -name = "vortex-reverse-index" -description = "A simple value -> row reverse index over any column, demonstrating vortex-layout's indexed layout IndexVTable" -authors.workspace = true -categories.workspace = true -edition.workspace = true -homepage.workspace = true -include.workspace = true -keywords.workspace = true -license.workspace = true -readme.workspace = true -repository.workspace = true -rust-version.workspace = true -version.workspace = true - -[dependencies] -roaring = { workspace = true } -vortex-array = { workspace = true } -vortex-error = { workspace = true } -vortex-layout = { workspace = true } -vortex-session = { workspace = true } -vortex-utils = { workspace = true } - -[dev-dependencies] -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } -vortex-buffer = { workspace = true } -vortex-edition = { workspace = true } -vortex-file = { workspace = true } -vortex-io = { workspace = true, features = ["tokio"] } -vortex-mask = { workspace = true } - -[lints] -workspace = true From 9a91972a36fb435f4e5549edd330279b6a0e71b4 Mon Sep 17 00:00:00 2001 From: Thor Date: Tue, 1 Sep 2026 15:02:06 -0500 Subject: [PATCH 6/7] claude: Combine sibling index claims instead of using only the first IndexedReader::plan_probe returned on the first spec that claimed an expression, regardless of exactness. A Superset claim from an earlier spec could shadow a later, more precise Exact claim, and multiple Superset claims on the same expression never got combined. CachedProbe is now Exact(locator) or Superset(Vec): plan_probe scans every spec, preferring any Exact claim outright and otherwise collecting all Superset claims. pruning_evaluation intersects every Superset locator (short-circuiting once nothing survives), and filter_evaluation only fast-paths on an Exact claim. Signed-off-by: "Thor" --- vortex-layout/src/layouts/indexed/reader.rs | 62 +++++-- vortex-layout/src/layouts/indexed/tests.rs | 195 ++++++++++++++++++++ 2 files changed, 241 insertions(+), 16 deletions(-) diff --git a/vortex-layout/src/layouts/indexed/reader.rs b/vortex-layout/src/layouts/indexed/reader.rs index 92969ce4309..7cbca735189 100644 --- a/vortex-layout/src/layouts/indexed/reader.rs +++ b/vortex-layout/src/layouts/indexed/reader.rs @@ -45,7 +45,9 @@ type SharedProbe = Shared> /// A reader for the [`crate::layouts::indexed::Indexed`] layout. /// /// Probes happen once per expression per file: the shared future is cached, and each split slices -/// its own row range out of the resulting locator rather than re-probing. +/// its own row range out of the resulting locator rather than re-probing. When more than one spec +/// claims the same expression, the first `Exact` claim wins outright; failing that, every claiming +/// spec's locator is kept and intersected at evaluation time. pub struct IndexedReader { layout: IndexedLayout, name: Arc, @@ -56,10 +58,16 @@ pub struct IndexedReader { probes: DashMap>, } +/// One cached probe result for an expression, combining every claiming spec's index. +/// +/// `Exact` holds a single spec's locator: an exact claim already fully answers the expression, so +/// no other spec's claim on it, exact or not, needs combining with it. `Superset` holds every +/// claiming spec's locator, since each one only narrows what's proven non-matching, and +/// [`IndexedReader::pruning_evaluation`] intersects them all. #[derive(Clone)] -struct CachedProbe { - exactness: IndexExactness, - locator: SharedProbe, +enum CachedProbe { + Exact(SharedProbe), + Superset(Vec), } impl IndexedReader { @@ -123,6 +131,8 @@ impl IndexedReader { } fn plan_probe(&self, expr: &BoundExpression) -> VortexResult> { + let mut supersets = Vec::new(); + for (idx, spec) in self.layout.indexes().iter().enumerate() { // Unregistered kinds are inert: their child is never read. let Some(vtable) = spec.vtable() else { @@ -145,10 +155,18 @@ impl IndexedReader { self.session.clone(), )?; - return Ok(Some(CachedProbe { exactness, locator })); + if exactness == IndexExactness::Exact { + // Already the best possible answer: no other spec's claim on this expression, + // exact or not, can sharpen it or needs combining with it. + return Ok(Some(CachedProbe::Exact(locator))); + } + supersets.push(locator); } - Ok(None) + if supersets.is_empty() { + return Ok(None); + } + Ok(Some(CachedProbe::Superset(supersets))) } } @@ -233,8 +251,22 @@ impl LayoutReader for IndexedReader { let expr = expr.clone(); Ok(MaskFuture::new(mask.len(), async move { - let locator = probe.locator.await?; - let mut result = mask.bitand(&locator.mask_for(&row_range)?); + let locators: &[SharedProbe] = match &probe { + CachedProbe::Exact(locator) => std::slice::from_ref(locator), + CachedProbe::Superset(locators) => locators, + }; + + // Every claiming spec's mask only narrows what's proven non-matching, so intersect + // them all; stop as soon as nothing is left alive, rather than awaiting a probe whose + // answer can no longer change the result. + let mut result = mask; + for locator in locators { + if result.all_false() { + break; + } + let locator = locator.clone().await?; + result = result.bitand(&locator.mask_for(&row_range)?); + } // Only bother the data child if the index left anything alive. if !result.all_false() { @@ -252,17 +284,15 @@ impl LayoutReader for IndexedReader { expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { - // An exact index answers the conjunct outright, so the data child is never decoded for it. - // A superset index can only prune, and the data child re-checks the real predicate. Either - // way this reuses the cached probe, so a superset conjunct costs no extra IO here. - if let Some(probe) = self - .probe(expr)? - .filter(|probe| probe.exactness == IndexExactness::Exact) - { + // Only an exact claim answers the conjunct outright, so the data child is never decoded + // for it. A superset claim, however many specs contributed to it, can only prune, so the + // real predicate always re-checks through the data child. Either way this reuses the + // cached probe, so a superset conjunct costs no extra IO here. + if let Some(CachedProbe::Exact(locator)) = self.probe(expr)? { let row_range = row_range.clone(); let len = mask.len(); return Ok(MaskFuture::new(len, async move { - let locator = probe.locator.await?; + let locator = locator.await?; let index_mask = locator.mask_for(&row_range)?; // Post-condition: the result must be intersected with the input mask. Ok(mask.await?.bitand(&index_mask)) diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs index 483a3cdb04f..335c8bd0c43 100644 --- a/vortex-layout/src/layouts/indexed/tests.rs +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -6,6 +6,7 @@ use std::sync::Arc; +use roaring::RoaringBitmap; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -32,6 +33,7 @@ use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::indexed::tests::exact_value::DecliningIndex; use crate::layouts::indexed::tests::exact_value::ExactValueIndex; +use crate::layouts::indexed::tests::fixed_superset::FixedSupersetIndex; use crate::layouts::repartition::RepartitionStrategy; use crate::layouts::repartition::RepartitionWriterOptions; use crate::scan::scan_builder::ScanBuilder; @@ -346,6 +348,79 @@ async fn one_builder_declining_leaves_the_others_intact() -> VortexResult<()> { Ok(()) } +/// Two `Superset` claims on the same conjunct must combine, not pick one and drop the other: +/// neither `{1, 2, 9}` nor `{2, 9, 10}` alone leaves only `{2, 9}` standing, only their +/// intersection does. +#[tokio::test] +async fn multiple_superset_claims_intersect() -> VortexResult<()> { + let session = new_session(); + let (layout, segments) = write( + &session, + vec![ + IndexConfig::with_defaults(FixedSupersetIndex::new_ref( + "test.idx.fixed_a", + RoaringBitmap::from_iter([1u32, 2, 9]), + )), + IndexConfig::with_defaults(FixedSupersetIndex::new_ref( + "test.idx.fixed_b", + RoaringBitmap::from_iter([2u32, 9, 10]), + )), + ], + ) + .await?; + + let reader = text_reader(&session, &layout, segments)?; + let row_count = reader.row_count(); + // `FixedSupersetIndex` claims unconditionally, so any bound Utf8 conjunct exercises it. + let filter = eq(root(), lit("irrelevant")).bind(reader.dtype())?; + let mask = reader + .pruning_evaluation( + &(0..row_count), + &filter, + Mask::new_true(usize::try_from(row_count)?), + )? + .await?; + + assert_eq!( + mask, + Mask::from_iter((0..ROWS.len()).map(|row| row == 2 || row == 9)) + ); + Ok(()) +} + +/// An `Exact` claim must discard an earlier, misleading `Superset` claim on the same conjunct +/// rather than intersect with it: the empty superset here would prune away row 2 if it were kept +/// around, but the exact claim that follows it proves row 2 is the real, correct answer. +#[tokio::test] +async fn exact_claim_discards_a_preceding_superset_claim() -> VortexResult<()> { + let session = session_with_exact_index(); + let (layout, segments) = write( + &session, + vec![ + IndexConfig::with_defaults(FixedSupersetIndex::new_ref( + "test.idx.fixed_empty", + RoaringBitmap::new(), + )), + IndexConfig::with_defaults(ExactValueIndex::new_ref()), + ], + ) + .await?; + + let reader = text_reader(&session, &layout, segments)?; + let row_count = reader.row_count(); + let filter = eq(root(), lit(ROWS[2])).bind(reader.dtype())?; + let mask = reader + .pruning_evaluation( + &(0..row_count), + &filter, + Mask::new_true(usize::try_from(row_count)?), + )? + .await?; + + assert_eq!(mask, Mask::from_iter((0..ROWS.len()).map(|row| row == 2))); + Ok(()) +} + /// A test-only sorted value index, present to exercise the [`super::IndexExactness::Exact`] path /// that a real posting-list index kind (such as an n-gram index) would rarely reach for equality /// queries. @@ -625,3 +700,123 @@ mod exact_value { } } } + +/// A test-only index kind that always claims any expression with `Superset` exactness and answers +/// with a fixed locator supplied at construction, ignoring both the expression and the data. +/// +/// Real superset-only kinds (bloom filters, n-gram indexes) derive their locator from the data; +/// this one is a stand-in that hands a test exact, known masks to combine, so its assertions are +/// about the reader's sibling-combination logic rather than any kind's own indexing correctness. +mod fixed_superset { + use std::sync::Arc; + + use roaring::RoaringBitmap; + use vortex_array::ArrayRef; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::expr::BoundExpression; + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::stream::ArrayStreamExt; + use vortex_array::stream::SendableArrayStream; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::layouts::indexed::IndexBuilder; + use crate::layouts::indexed::IndexExactness; + use crate::layouts::indexed::IndexId; + use crate::layouts::indexed::IndexQueryPlan; + use crate::layouts::indexed::IndexResolve; + use crate::layouts::indexed::IndexVTable; + use crate::layouts::indexed::IndexVTableRef; + use crate::layouts::indexed::RowLocator; + + #[derive(Debug)] + pub struct FixedSupersetIndex { + id: &'static str, + rows: RoaringBitmap, + } + + impl FixedSupersetIndex { + pub fn new_ref(id: &'static str, rows: RoaringBitmap) -> IndexVTableRef { + Arc::new(Self { id, rows }) + } + } + + impl IndexVTable for FixedSupersetIndex { + fn id(&self) -> IndexId { + IndexId::from(self.id) + } + + fn supports_dtype(&self, dtype: &DType) -> bool { + matches!(dtype, DType::Utf8(_)) + } + + fn builder( + &self, + _dtype: &DType, + _options: &[u8], + _data_block_len: Option, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Box::new(Builder)) + } + + fn plan( + &self, + _expr: &BoundExpression, + _dtype: &DType, + _options: &[u8], + ) -> VortexResult> { + Ok(Some(IndexQueryPlan { + exactness: IndexExactness::Superset, + filter: eq(root(), lit(0i32)), + resolve: Arc::new(Resolve { + rows: self.rows.clone(), + }), + })) + } + } + + /// Writes one dummy row so the layout has real content for `plan`'s filter to select; the + /// value itself is never inspected. + struct Builder; + + impl IndexBuilder for Builder { + fn push( + &mut self, + _chunk: &ArrayRef, + _row_offset: u64, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + Ok(()) + } + + fn finish(self: Box) -> VortexResult)>> { + let array = PrimitiveArray::from_iter([0i32]).into_array(); + Ok(Some((array.to_array_stream().boxed(), vec![]))) + } + + fn buffered_bytes(&self) -> u64 { + 0 + } + } + + struct Resolve { + rows: RoaringBitmap, + } + + impl IndexResolve for Resolve { + fn resolve( + &self, + _postings: &ArrayRef, + _data_row_count: u64, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(RowLocator::Rows(self.rows.clone())) + } + } +} From b28f8e3f2571fd0a2aa5e0d392111ead5e75282d Mon Sep 17 00:00:00 2001 From: Thor Date: Wed, 2 Sep 2026 12:10:44 -0500 Subject: [PATCH 7/7] Added comment about nesting the indexed layout --- vortex-layout/src/layouts/indexed/mod.rs | 27 ++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/vortex-layout/src/layouts/indexed/mod.rs b/vortex-layout/src/layouts/indexed/mod.rs index 1aa5b11e05a..5f0b7ff78cf 100644 --- a/vortex-layout/src/layouts/indexed/mod.rs +++ b/vortex-layout/src/layouts/indexed/mod.rs @@ -24,6 +24,33 @@ //! and compressed by the same machinery as data. Probing an index is therefore just a pruned scan //! over the index child: a sorted key column's zone map narrows the probe to a handful of zones. //! +//! # Nesting +//! +//! Both child slots hold an ordinary [`LayoutRef`], so nesting is not special-cased anywhere in +//! this module — it falls out of every child being read and written through the generic +//! [`crate::LayoutReader`] and [`crate::LayoutStrategy`] traits: +//! +//! - **Data child.** Slot 0 can be any layout — chunked, zoned, struct, flat, or another +//! `vortex.indexed` — since [`writer::IndexedStrategy`] writes it through a plain +//! `Arc` and [`reader::IndexedReader`] reads it back through the generic +//! reader trait rather than assuming a concrete kind. A data child that is itself wrapped in +//! `vortex.indexed` prunes with its own indexes when the outer wrapper delegates a +//! `pruning_evaluation`/`filter_evaluation` call to it, so a struct layout whose fields are +//! independently indexed, or two stacked `vortex.indexed` layers with different index kinds, +//! both work without this layout knowing about it. +//! - **Index child.** Index content is written through the same kind of `Arc` +//! and probed by running an ordinary scan (`ScanBuilder`) over it, so it too can be chunked, +//! zone-mapped, or itself wrapped in `vortex.indexed`. Wrapping an index child in another index +//! (for example, a small index over keys that point into a large posting list) prunes the probe +//! scan the same way it would prune a plain data scan. +//! - **`vortex.indexed` as someone else's child.** Since it is an ordinary [`LayoutRef`], this +//! layout is not restricted to the top of a file — it can equally be a struct field, one chunk +//! of a chunked layout, or the child of any other layout that composes over children. +//! +//! The one thing that does not nest is the [`IndexSession`] registry itself: every level of +//! nesting is resolved against the same session passed down from the reader, so an index kind +//! unregistered at any level degrades to inert there, independent of the other levels. +//! //! # What ships here //! //! The generic wrapper only: [`Indexed`], [`writer::IndexedStrategy`] and [`reader::IndexedReader`],