From bc2cc33043bed29c63e3b9f452cd5f84bd6f83e6 Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Sat, 10 Oct 2026 23:08:30 +0800 Subject: [PATCH] feat(index): add raw and code-only RaBitQ cache entries --- rust/lance-index/src/vector/bq.rs | 2 + rust/lance-index/src/vector/bq/layered.rs | 183 ++- .../src/vector/bq/layered_stats.rs | 35 +- .../src/vector/bq/partition_codes.rs | 227 ++++ rust/lance-index/src/vector/bq/plane_cache.rs | 953 +++++++++++---- rust/lance-index/src/vector/bq/raw_body.rs | 609 +++++++++ rust/lance-index/src/vector/bq/resident.rs | 110 +- rust/lance-index/src/vector/bq/storage.rs | 264 +++- rust/lance-index/src/vector/storage.rs | 708 +++++++++-- .../src/index/vector/ivf/partition_serde.rs | 143 ++- rust/lance/src/index/vector/ivf/v2.rs | 1087 ++++++++++++++++- .../src/index/vector/ivf/v2/lazy_full.rs | 13 +- 12 files changed, 3947 insertions(+), 387 deletions(-) create mode 100644 rust/lance-index/src/vector/bq/partition_codes.rs create mode 100644 rust/lance-index/src/vector/bq/raw_body.rs diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index e1fb3927d83..14cc540be61 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -25,8 +25,10 @@ pub mod layered_stats; #[cfg(test)] mod lazy_tests; pub(crate) mod pairwise; +pub mod partition_codes; mod plane_cache; pub mod prune; +pub mod raw_body; pub mod resident; pub mod rotation; pub mod storage; diff --git a/rust/lance-index/src/vector/bq/layered.rs b/rust/lance-index/src/vector/bq/layered.rs index 4e16d79dcd3..6979d55de62 100644 --- a/rust/lance-index/src/vector/bq/layered.rs +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -242,6 +242,10 @@ pub const SIGN_BOUNDS_PLANE: u8 = 3; /// plane over the high plane over the low plane. const SIGN_PLANE_MEMORY_PRIORITY: u8 = 3; +/// Appended to the key of a sign, high or low plane entry that holds its +/// codes alone ([`EntryColumns::Codes`]). +const CODE_ONLY_KEY_SUFFIX: &str = "-codes"; + /// Where a layered index's cache keeps the estimator bounds columns /// ([`HIGH_BOUNDS_COLUMN`], [`FULL_BOUNDS_COLUMN`]). The placement decides /// what the sign plane entry holds, so it is part of that entry's key. @@ -274,6 +278,47 @@ impl std::fmt::Display for SignBounds { } } +/// Which columns an IVF_RQ index's cache entries hold: a layered +/// partition's plane entries and a native flat partition's entry. The +/// composition decides what an entry holds, so it is part of its key. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum EntryColumns { + /// Every column a read of the plane or partition returns. + #[default] + All, + /// Only the columns the index's resident store does not keep: the codes + /// and the estimator bounds (see [`plane_entry_columns`]). Every read + /// attaches copies of the store's rows of the others, so the scored + /// batch is the one an `All` entry holds. Only an index whose small + /// columns are resident keeps such entries. + Codes, +} + +impl EntryColumns { + /// The spelling of `LANCE_RQ_ENTRY_COLUMNS`: `all` or `codes`. + pub fn as_str(self) -> &'static str { + match self { + Self::All => "all", + Self::Codes => "codes", + } + } + + /// What an index that asks for `self` keeps in its entries: `Codes` + /// needs the resident store, so without it the entries hold `All`. + pub fn resolve(self, resident: bool) -> Self { + match self { + Self::Codes if resident => Self::Codes, + _ => Self::All, + } + } +} + +impl std::fmt::Display for EntryColumns { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + /// Plane keys live under the immutable index UUID namespace. pub struct PlaneKey { pub partition: usize, @@ -281,17 +326,26 @@ pub struct PlaneKey { /// The index's bounds placement. Only the sign plane's key depends on /// it: the other planes hold the same columns under either placement. pub sign_bounds: SignBounds, + /// What the index's entries hold. The keys of planes 0-2 depend on it; + /// the bounds plane holds file columns alone either way. + pub entry_columns: EntryColumns, } impl CacheKey for PlaneKey { type ValueType = PlaneBatch; fn key(&self) -> Cow<'_, str> { - // Persisted entries outlive the process, so each sign-plane - // composition needs its own key. The sign plane with bounds keeps - // the key its entries had before the bounds plane existed. + // Persisted entries outlive the process, so each composition of an + // entry needs its own key. The sign plane with bounds keeps the key + // its entries had before the bounds plane existed, and a full entry + // the key it had before code-only entries existed. + let codes = if self.entry_columns == EntryColumns::Codes && self.plane < SIGN_BOUNDS_PLANE { + CODE_ONLY_KEY_SUFFIX + } else { + "" + }; if self.plane == 0 && self.sign_bounds == SignBounds::Lazy { - format!("{}:0-bounds", self.partition).into() + format!("{}:0-bounds{codes}", self.partition).into() } else { - format!("{}:{}", self.partition, self.plane).into() + format!("{}:{}{codes}", self.partition, self.plane).into() } } fn type_name() -> &'static str { @@ -357,6 +411,23 @@ pub fn plane_columns(plane: u8, sign_bounds: SignBounds) -> &'static [&'static s } } +/// Columns a plane entry of `plane` holds under `entry_columns`, in +/// [`plane_columns`] order: every column of the plane, or under +/// [`EntryColumns::Codes`] only those a resident store does not keep. +pub fn plane_entry_columns( + plane: u8, + sign_bounds: SignBounds, + entry_columns: EntryColumns, +) -> Vec<&'static str> { + let columns = plane_columns(plane, sign_bounds).iter().copied(); + match entry_columns { + EntryColumns::All => columns.collect(), + EntryColumns::Codes => columns + .filter(|name| super::resident::is_file_column(name)) + .collect(), + } +} + /// A quantized candidate tied to an immutable index and its physical partition. /// /// Distributed executors may perform a global distance cut on these values, then @@ -728,17 +799,20 @@ mod tests { #[test] fn plane_key_codec_tags_plane_and_priority() { - for (plane, priority) in [(0, 3), (1, 2), (2, 1), (SIGN_BOUNDS_PLANE, 2)] { - let codec = PlaneKey { - partition: 7, - plane, - sign_bounds: SignBounds::Lazy, + for entry_columns in [EntryColumns::All, EntryColumns::Codes] { + for (plane, priority) in [(0, 3), (1, 2), (2, 1), (SIGN_BOUNDS_PLANE, 2)] { + let codec = PlaneKey { + partition: 7, + plane, + sign_bounds: SignBounds::Lazy, + entry_columns, + } + .codec_for_key() + .unwrap(); + assert_eq!(codec.plane_tag(), Some(plane)); + assert_eq!(codec.memory_priority(), priority); + assert!(codec.supports_row_selection()); } - .codec_for_key() - .unwrap(); - assert_eq!(codec.plane_tag(), Some(plane)); - assert_eq!(codec.memory_priority(), priority); - assert!(codec.supports_row_selection()); } assert_eq!(PlaneKey::codec().unwrap().plane_tag(), None); } @@ -752,6 +826,7 @@ mod tests { partition: 7, plane, sign_bounds, + entry_columns: EntryColumns::All, }; plane_key.key().into_owned() }; @@ -786,6 +861,84 @@ mod tests { assert!(plane_columns(SIGN_BOUNDS_PLANE + 1, SignBounds::Lazy).is_empty()); } + /// Code-only entries of the sign, high and low planes take their own + /// keys; every other key is the one its entries had before, so full + /// entries persisted by an earlier build keep being found. + #[test] + fn plane_key_separates_entry_columns() { + let key = |plane, sign_bounds, entry_columns| { + PlaneKey { + partition: 7, + plane, + sign_bounds, + entry_columns, + } + .key() + .into_owned() + }; + let full_keys = [ + (0, SignBounds::Lazy, "7:0-bounds"), + (0, SignBounds::Eager, "7:0"), + (1, SignBounds::Lazy, "7:1"), + (1, SignBounds::Eager, "7:1"), + (2, SignBounds::Lazy, "7:2"), + (2, SignBounds::Eager, "7:2"), + (SIGN_BOUNDS_PLANE, SignBounds::Lazy, "7:3"), + (SIGN_BOUNDS_PLANE, SignBounds::Eager, "7:3"), + ]; + for (plane, sign_bounds, full) in full_keys { + assert_eq!(key(plane, sign_bounds, EntryColumns::All), full); + let codes = key(plane, sign_bounds, EntryColumns::Codes); + if plane == SIGN_BOUNDS_PLANE { + assert_eq!(codes, full); + } else { + assert_eq!(codes, format!("{full}-codes")); + } + } + } + + /// A code-only entry holds the plane's code and bounds columns, in the + /// plane's order; the bounds plane holds the same columns either way. + #[test] + fn plane_entry_columns_keep_file_columns() { + for sign_bounds in [SignBounds::Lazy, SignBounds::Eager] { + for plane in [0, 1, 2, SIGN_BOUNDS_PLANE] { + assert_eq!( + plane_entry_columns(plane, sign_bounds, EntryColumns::All), + plane_columns(plane, sign_bounds) + ); + } + } + let codes = + |plane, sign_bounds| plane_entry_columns(plane, sign_bounds, EntryColumns::Codes); + assert_eq!(codes(0, SignBounds::Lazy), [RABIT_CODE_COLUMN]); + assert_eq!( + codes(0, SignBounds::Eager), + [RABIT_CODE_COLUMN, HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN] + ); + assert_eq!(codes(1, SignBounds::Lazy), [RABIT_BLOCKED_EX_CODE_COLUMN]); + assert_eq!( + codes(2, SignBounds::Eager), + [RABIT_BLOCKED_EX_CODE_LO_COLUMN] + ); + assert_eq!( + codes(SIGN_BOUNDS_PLANE, SignBounds::Lazy), + [HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN] + ); + assert!(codes(SIGN_BOUNDS_PLANE, SignBounds::Eager).is_empty()); + } + + #[test] + fn entry_columns_need_a_resident_store() { + assert_eq!(EntryColumns::Codes.resolve(true), EntryColumns::Codes); + assert_eq!(EntryColumns::Codes.resolve(false), EntryColumns::All); + assert_eq!(EntryColumns::All.resolve(true), EntryColumns::All); + assert_eq!(EntryColumns::All.resolve(false), EntryColumns::All); + assert_eq!(EntryColumns::default(), EntryColumns::All); + assert_eq!(EntryColumns::Codes.to_string(), "codes"); + assert_eq!(EntryColumns::All.to_string(), "all"); + } + /// A layered store needs the full bounds only to prune full precision on /// a file without error factors. Without the high bounds, a High scan of /// a layered store scores every row instead of pruning with the error diff --git a/rust/lance-index/src/vector/bq/layered_stats.rs b/rust/lance-index/src/vector/bq/layered_stats.rs index 4806a9d5d6e..1bfb1ee4d93 100644 --- a/rust/lance-index/src/vector/bq/layered_stats.rs +++ b/rust/lance-index/src/vector/bq/layered_stats.rs @@ -6,10 +6,12 @@ //! The counters only advance while a query runs the lazy scan (or is checked //! for it), so production queries with the scan disabled touch at most the //! `ineligible_disabled` counter. The `resident_columns_*`, -//! `resident_store_evictions` counters are the -//! exception: they advance when an IVF_RQ index, layered or not, opens with -//! or loads its resident columns. Readers take deltas with -//! [`snapshot_and_reset`]. +//! `resident_attach_*`, `resident_store_evictions` +//! counters are the exception: they advance when an IVF_RQ index, layered or +//! not, opens with, loads or reads through its resident columns. So is +//! `storage_construct_repacks`, which advances when any RaBitQ storage is +//! built from codes it must rewrite. +//! Readers take deltas with [`snapshot_and_reset`]. use std::sync::atomic::{AtomicU64, Ordering}; @@ -52,6 +54,13 @@ impl Counter { fn take(&self) -> u64 { self.0.swap(0, Ordering::Relaxed) } + + /// The count so far, without resetting it: tests that run beside others + /// can only bound the growth of a process-wide counter from below. + #[cfg(test)] + pub(crate) fn get(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } } /// A cumulative counter split by probe-rank bucket (see [`rank_bucket`]). @@ -299,6 +308,24 @@ layered_lazy_counters! { /// Resident store entries the index cache dropped: evicted while /// idle, cleared, or refused at admission. resident_store_evictions, + /// RaBitQ storages, layered or not, whose construction rewrote the + /// stored codes: packed sign codes stored unpacked, or repacked + /// legacy sequential ex codes into the blocked layout. A storage + /// built from codes stored packed and blocked, as the index builder + /// writes them and cache entries keep them, rewrites nothing. + storage_construct_repacks, + /// Batches assembled with copies of an index's resident rows: every + /// read of a plane or partition from the file that takes its small + /// columns from the resident store, and every read of a code-only + /// cache entry (`LANCE_RQ_ENTRY_COLUMNS=codes`), hits included. + resident_attach_calls, + /// Rows those batches hold. + resident_attach_rows, + /// Bytes of the resident rows copied into them: what a full entry + /// would have kept in the cache. + resident_attach_bytes, + /// Time spent assembling them, copies included. + resident_attach_ns, } ranked { /// Probes gathered lazily. diff --git a/rust/lance-index/src/vector/bq/partition_codes.rs b/rust/lance-index/src/vector/bq/partition_codes.rs new file mode 100644 index 00000000000..56a8ebf8959 --- /dev/null +++ b/rust/lance-index/src/vector/bq/partition_codes.rs @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Code-only cache entries of native IVF_RQ partitions. +//! +//! A native IVF_RQ index with a flat sub-index whose small columns are +//! resident caches a partition as its code columns alone +//! (`EntryColumns::Codes`): every read attaches copies of the resident rows +//! and builds the storage, so the cache holds no second copy of the row ids +//! and factors the store already keeps. The codes are kept packed and +//! blocked (see `normalize_entry_codes`), and persisted as a raw fixed-width +//! body (see [`super::raw_body`]): +//! +//! ```text +//! v1: [raw body] +//! ``` + +use std::borrow::Cow; + +use arrow_array::RecordBatch; +use lance_core::cache::{ + CacheCodec, CacheCodecImpl, CacheEntryReader, CacheEntryWriter, CacheKey, CacheKeySchema, + KeyBuilder, +}; +use lance_core::deepsize::{Context, DeepSizeOf}; +use lance_core::{Error, Result}; + +use super::raw_body::{read_raw_batch, write_raw_batch}; + +/// Body version of the entries this build writes: a raw body. +const RAW_BODY_VERSION: u32 = 1; + +/// The code columns of a native IVF_RQ partition: the sign codes, packed, +/// and the ex codes, blocked, at the partition's rows. +#[derive(Debug)] +pub struct PartitionCodes(pub RecordBatch); + +impl DeepSizeOf for PartitionCodes { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.0.deep_size_of_children(context) + } +} + +impl CacheCodecImpl for PartitionCodes { + const TYPE_ID: &'static str = "lance.vector.rq.partition-codes"; + const CURRENT_VERSION: u32 = RAW_BODY_VERSION; + + fn serialize(&self, writer: &mut CacheEntryWriter<'_>) -> Result<()> { + write_raw_batch(writer, &self.0) + } + + fn deserialize(reader: &mut CacheEntryReader<'_>) -> Result { + match reader.version() { + RAW_BODY_VERSION => Ok(Self(read_raw_batch(reader)?)), + version => Err(Error::invalid_input(format!( + "unsupported partition codes cache version {version}" + ))), + } + } +} + +/// Key of a native IVF_RQ partition's code-only entry, under the index's +/// namespace of the index cache. Its entries are a type of their own, so an +/// index that caches whole partitions never finds one, nor the reverse. +pub struct PartitionCodesKey { + pub partition: usize, +} + +impl CacheKey for PartitionCodesKey { + type ValueType = PartitionCodes; + + fn key(&self) -> Cow<'_, str> { + format!("ivf-codes-{}", self.partition).into() + } + + fn type_name() -> &'static str { + "IVFPartitionCodes" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.index.ivf-partition-codes-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_u64(self.partition as u64); + } + + /// No plane tag and the lowest memory priority, as a whole native + /// partition's entry. + fn codec() -> Option { + Some(CacheCodec::from_impl::()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array, UInt8Array}; + use bytes::Bytes; + use lance_arrow::FixedSizeListArrayExt; + use lance_core::cache::{CacheDecode, CacheMissReason}; + + use crate::vector::bq::raw_body::RAW_BODY_KIND; + use crate::vector::bq::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_CODE_COLUMN}; + + /// Envelope bytes ahead of an entry's body. + fn body_offset() -> usize { + 4 + 1 + 2 + PartitionCodes::TYPE_ID.len() + 4 + } + + fn codes(rows: usize, width: usize) -> ArrayRef { + Arc::new( + FixedSizeListArray::try_new_from_values( + UInt8Array::from( + (0..rows * width) + .map(|v| (v % 251) as u8) + .collect::>(), + ), + width as i32, + ) + .unwrap(), + ) + } + + fn partition_codes(rows: usize) -> RecordBatch { + RecordBatch::try_from_iter([ + (RABIT_CODE_COLUMN, codes(rows, 16)), + (RABIT_BLOCKED_EX_CODE_COLUMN, codes(rows, 48)), + ]) + .unwrap() + } + + fn encode(batch: &RecordBatch) -> Bytes { + let codec = PartitionCodesKey { partition: 0 }.codec_for_key().unwrap(); + let mut encoded = Vec::new(); + codec + .serialize( + &(Arc::new(PartitionCodes(batch.clone())) as Arc), + &mut encoded, + ) + .unwrap(); + Bytes::from(encoded) + } + + fn decode(encoded: &Bytes) -> CacheDecode { + let codec = PartitionCodesKey { partition: 0 }.codec_for_key().unwrap(); + match codec.deserialize(encoded) { + CacheDecode::Hit(value) => { + CacheDecode::Hit(value.downcast::().unwrap().0.clone()) + } + CacheDecode::Miss(reason) => CacheDecode::Miss(reason), + } + } + + /// Codes round trip through a raw body of exactly their values, beyond + /// the schema and padding: no validity bitmap. Decoded entries are + /// charged their values. + #[test] + fn partition_codes_round_trip_through_raw_bodies() { + for rows in [0, 1, 33, 300] { + let original = partition_codes(rows); + let encoded = encode(&original); + let body = &encoded[body_offset()..]; + assert_eq!(body[0], RAW_BODY_KIND, "rows={rows}"); + let schema_len = u32::from_le_bytes(body[1..5].try_into().unwrap()) as usize; + let values = [16 * rows, 48 * rows] + .map(|bytes| bytes.next_multiple_of(8)) + .iter() + .sum::(); + assert_eq!(body.len(), 1 + 4 + schema_len + 8 + values, "rows={rows}"); + let CacheDecode::Hit(decoded) = decode(&encoded) else { + panic!("rows={rows}: decode missed") + }; + assert_eq!(decoded, original, "rows={rows}"); + assert_eq!( + PartitionCodes(decoded).deep_size_of(), + PartitionCodes(original).deep_size_of(), + "rows={rows}" + ); + } + } + + /// A body this build cannot read is a miss, never a panic. + #[test] + fn newer_or_corrupt_partition_codes_miss() { + let encoded = encode(&partition_codes(40)).to_vec(); + let version_at = body_offset() - 4; + let mut newer = encoded.clone(); + newer[version_at..body_offset()] + .copy_from_slice(&(PartitionCodes::CURRENT_VERSION + 1).to_le_bytes()); + assert!(matches!( + decode(&Bytes::from(newer)), + CacheDecode::Miss(CacheMissReason::VersionTooNew) + )); + for len in [body_offset(), body_offset() + 1, encoded.len() - 1] { + assert!( + matches!( + decode(&Bytes::from(encoded[..len].to_vec())), + CacheDecode::Miss(_) + ), + "truncated to {len} bytes" + ); + } + } + + /// Entries of other partitions, and whole-partition entries, have other + /// keys; the codec tags no plane and keeps the lowest priority. + #[test] + fn partition_codes_key_and_codec() { + let key = |partition| PartitionCodesKey { partition }; + assert_ne!(key(1).key(), key(2).key()); + assert_eq!(PartitionCodesKey::type_name(), "IVFPartitionCodes"); + let codec = key(1).codec_for_key().unwrap(); + assert_eq!(codec.type_id(), PartitionCodes::TYPE_ID); + assert_eq!(codec.plane_tag(), None); + assert_eq!(codec.memory_priority(), 0); + assert!(!codec.supports_row_selection()); + let factor: ArrayRef = Arc::new(Float32Array::from(vec![0.5f32; 3])); + let other = RecordBatch::try_from_iter([("factor", factor)]).unwrap(); + let CacheDecode::Hit(decoded) = decode(&encode(&other)) else { + panic!("decode missed") + }; + assert_eq!(decoded, other); + } +} diff --git a/rust/lance-index/src/vector/bq/plane_cache.rs b/rust/lance-index/src/vector/bq/plane_cache.rs index a2e2fb02e94..217ba0d5ad2 100644 --- a/rust/lance-index/src/vector/bq/plane_cache.rs +++ b/rust/lance-index/src/vector/bq/plane_cache.rs @@ -1,69 +1,153 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Row-addressable persistent ex planes. Sign codes retain their transposed -//! IPC body, and so do the bounds columns, which scans read whole. +//! Persistent plane entries. The ex planes (high and low) are +//! row-addressable: fixed-stride rows of codes, each followed by its +//! factors, so a sparse gather reads only its rows. The sign plane and the +//! bounds plane are read whole and keep a raw fixed-width body (see +//! [`super::raw_body`]); sign codes keep their partition-local transposition. +//! +//! Body layouts by entry version: +//! +//! ```text +//! v1 : Arrow IPC section, every plane +//! v2 sign : [plane: u8 = 0][Arrow IPC section], the bounds plane too +//! v2 ex plane : [plane: u8][rows: u64 LE][width: u32 LE][rows] +//! v3 sign : [plane: u8 = 0, or 3 for the bounds plane][raw body] +//! v3 ex plane : [plane: u8][factor_columns: u8][rows: u64 LE][width: u32 LE][rows] +//! ex-plane row : [codes: width bytes][add: f32 LE][scale: f32 LE] 2 factor columns +//! [codes: width bytes] 0 factor columns +//! ``` +//! +//! v2 ex-plane rows always carry both factors. A code-only entry +//! (`EntryColumns::Codes`) holds the file columns of its plane alone: its ex +//! plane rows carry no factor columns, and its sign plane only the codes (and +//! the bounds under eager bounds). use std::ops::Range; use std::sync::Arc; -use arrow_array::{Array, Float32Array, RecordBatch, UInt8Array, cast::AsArray, types::UInt8Type}; +use arrow_array::{ + Array, FixedSizeListArray, Float32Array, RecordBatch, UInt8Array, + cast::AsArray, + types::{Float32Type, UInt8Type}, +}; use arrow_schema::{DataType, Field, Schema}; use lance_arrow::FixedSizeListArrayExt; use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter, CacheRangeReader}; use lance_core::{Error, Result}; -use super::layered::{PlaneBatch, SignBounds, plane_columns}; +use super::layered::{PlaneBatch, SIGN_BOUNDS_PLANE, SignBounds, plane_columns}; +use super::raw_body::{read_raw_batch, write_raw_batch}; use super::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN}; -const HEADER_BYTES: usize = 13; +/// Body layouts of the plane entry versions this build reads. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BodyVersion { + /// Every plane as an Arrow IPC section. + Ipc = 1, + /// Ex planes as rows with both factors behind a 13-byte header; the sign + /// and bounds planes as a tagged Arrow IPC section. + RowsIpcSign = 2, + /// Ex planes as rows with 0 or 2 factor columns behind a 14-byte header; + /// the sign and bounds planes as a tagged raw body. + RowsRawSign = 3, +} + +impl BodyVersion { + fn of(version: u32) -> Result { + match version { + 1 => Ok(Self::Ipc), + 2 => Ok(Self::RowsIpcSign), + 3 => Ok(Self::RowsRawSign), + _ => Err(Error::invalid_input(format!( + "unsupported plane cache version {version}" + ))), + } + } + + /// Bytes of the header ahead of an ex plane's rows. + fn ex_header_bytes(self) -> Result { + match self { + Self::Ipc => Err(Error::invalid_input( + "plane cache version has no row directory", + )), + Self::RowsIpcSign => Ok(V2_HEADER_BYTES), + Self::RowsRawSign => Ok(V3_HEADER_BYTES), + } + } +} + +const V2_HEADER_BYTES: usize = 13; +const V3_HEADER_BYTES: usize = 14; +/// Factor columns of an ex-plane row with its add and scale factors. +const EX_FACTOR_COLUMNS: u8 = 2; +const FACTOR_BYTES: usize = size_of::(); const COALESCE_GAP_BYTES: usize = 4096; struct Header { plane: u8, + /// 0 or [`EX_FACTOR_COLUMNS`]. + factor_columns: u8, rows: usize, width: usize, + /// Bytes of the header itself; the rows follow it. + len: usize, } + impl Header { - fn parse(bytes: &[u8]) -> Result { - if bytes.len() < HEADER_BYTES || !matches!(bytes[0], 1 | 2) { - return Err(Error::invalid_input("invalid ex-plane cache header")); + fn parse(bytes: &[u8], version: BodyVersion) -> Result { + let invalid = || Error::invalid_input("invalid ex-plane cache header"); + let len = version.ex_header_bytes()?; + let header = bytes.get(..len).ok_or_else(invalid)?; + // v3 adds the factor column count after the plane; v2 rows always + // carry both factors. + let (factor_columns, fields) = match version { + BodyVersion::RowsRawSign => (header[1], &header[2..]), + _ => (EX_FACTOR_COLUMNS, &header[1..]), + }; + let plane = header[0]; + if !matches!(plane, 1 | 2) || !matches!(factor_columns, 0 | EX_FACTOR_COLUMNS) { + return Err(invalid()); } - let rows = usize::try_from(u64::from_le_bytes(bytes[1..9].try_into().unwrap())) + let (rows, width) = fields.split_at(size_of::()); + let rows = usize::try_from(u64::from_le_bytes(rows.try_into().map_err(|_| invalid())?)) .map_err(|_| Error::invalid_input("plane row count overflow"))?; - let width = u32::from_le_bytes(bytes[9..13].try_into().unwrap()) as usize; + let width = u32::from_le_bytes(width.try_into().map_err(|_| invalid())?) as usize; if width == 0 || width > i32::MAX as usize { return Err(Error::invalid_input("invalid cached code width")); } Ok(Self { - plane: bytes[0], + plane, + factor_columns, rows, width, + len, }) } + fn stride(&self) -> usize { - self.width + 8 + self.width + FACTOR_BYTES * self.factor_columns as usize } + fn batch(&self, codes: Vec, adds: Vec, scales: Vec) -> Result { // Ex planes hold the same columns under either bounds placement. let names = plane_columns(self.plane, SignBounds::default()); - let codes = arrow_array::FixedSizeListArray::try_new_from_values( - UInt8Array::from(codes), - self.width as i32, - )?; - let fields = vec![ - Field::new(names[0], codes.data_type().clone(), true), - Field::new(names[1], DataType::Float32, true), - Field::new(names[2], DataType::Float32, true), - ]; + let codes = + FixedSizeListArray::try_new_from_values(UInt8Array::from(codes), self.width as i32)?; + let mut fields = vec![Field::new(names[0], codes.data_type().clone(), true)]; + let mut columns: Vec> = vec![Arc::new(codes)]; + if self.factor_columns == EX_FACTOR_COLUMNS { + for (name, values) in names[1..].iter().zip([adds, scales]) { + fields.push(Field::new(*name, DataType::Float32, true)); + columns.push(Arc::new(Float32Array::from(values))); + } + } Ok(PlaneBatch(RecordBatch::try_new( Arc::new(Schema::new(fields)), - vec![ - Arc::new(codes), - Arc::new(Float32Array::from(adds)), - Arc::new(Float32Array::from(scales)), - ], + columns, )?)) } + fn decode(&self, bytes: &[u8]) -> Result { let size = self .rows @@ -72,84 +156,193 @@ impl Header { if bytes.len() != size { return Err(Error::invalid_input("truncated ex-plane cache body")); } + let factor_rows = if self.factor_columns == EX_FACTOR_COLUMNS { + self.rows + } else { + 0 + }; let mut codes = Vec::with_capacity(self.rows * self.width); - let mut adds = Vec::with_capacity(self.rows); - let mut scales = Vec::with_capacity(self.rows); + let mut adds = Vec::with_capacity(factor_rows); + let mut scales = Vec::with_capacity(factor_rows); for row in bytes.chunks_exact(self.stride()) { self.append(row, &mut codes, &mut adds, &mut scales); } self.batch(codes, adds, scales) } + + /// Append one row of [`Self::stride`] bytes. fn append(&self, row: &[u8], codes: &mut Vec, adds: &mut Vec, scales: &mut Vec) { - codes.extend_from_slice(&row[..self.width]); - adds.push(f32::from_le_bytes( - row[self.width..self.width + 4].try_into().unwrap(), - )); - scales.push(f32::from_le_bytes( - row[self.width + 4..self.width + 8].try_into().unwrap(), - )); + let (row_codes, factors) = row.split_at(self.width); + codes.extend_from_slice(row_codes); + if self.factor_columns == EX_FACTOR_COLUMNS { + let (add, scale) = factors.split_at(FACTOR_BYTES); + adds.push(le_f32(add)); + scales.push(le_f32(scale)); + } } } -impl CacheCodecImpl for PlaneBatch { - const TYPE_ID: &'static str = "lance.vector.rq.plane"; - const CURRENT_VERSION: u32 = 2; - const SUPPORTS_ROW_SELECTION: bool = true; +/// The little-endian `f32` at the start of `bytes`. +fn le_f32(bytes: &[u8]) -> f32 { + let mut value = [0; FACTOR_BYTES]; + value.copy_from_slice(&bytes[..FACTOR_BYTES]); + f32::from_le_bytes(value) +} - fn serialize(&self, writer: &mut CacheEntryWriter<'_>) -> Result<()> { - let plane = if self - .0 - .column_by_name(RABIT_BLOCKED_EX_CODE_LO_COLUMN) - .is_some() - { - 2 - } else if self - .0 - .column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN) - .is_some() - { - 1 - } else { - 0 - }; - writer.write_u8(plane)?; - if plane == 0 { - return writer.write_ipc(&self.0); - } +/// The ex plane (1 or 2) whose codes `batch` holds, or `None` for the sign +/// and bounds planes. +fn ex_plane(batch: &RecordBatch) -> Option { + if batch + .column_by_name(RABIT_BLOCKED_EX_CODE_LO_COLUMN) + .is_some() + { + Some(2) + } else if batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() { + Some(1) + } else { + None + } +} + +/// The plane tag of a sign or bounds plane batch: [`SIGN_BOUNDS_PLANE`] for +/// exactly the bounds plane's columns, 0 otherwise. +fn whole_plane(batch: &RecordBatch) -> u8 { + let bounds = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy); + let schema = batch.schema_ref(); + if schema.fields().len() == bounds.len() + && schema + .fields() + .iter() + .zip(bounds) + .all(|(field, name)| field.name() == name) + { + SIGN_BOUNDS_PLANE + } else { + 0 + } +} + +/// A float factor column of an ex-plane batch, which must hold no nulls. +fn factor_values<'a>(batch: &'a RecordBatch, name: &str) -> Result> { + let Some(column) = batch.column_by_name(name) else { + return Ok(None); + }; + let values = column + .as_primitive_opt::() + .filter(|values| values.null_count() == 0) + .ok_or_else(|| { + Error::invalid_input(format!( + "cached ex-plane factor column {name} must be non-null float32, got {}", + column.data_type() + )) + })?; + let values: &[f32] = values.values(); + Ok(Some(values)) +} + +impl PlaneBatch { + /// Write an ex plane as rows. Its factor columns are both present (2) or + /// both absent (0). + fn serialize_ex_plane(&self, plane: u8, writer: &mut CacheEntryWriter<'_>) -> Result<()> { let names = plane_columns(plane, SignBounds::default()); - let codes = self.0[names[0]].as_fixed_size_list(); - let adds = self.0[names[1]].as_primitive::(); - let scales = self.0[names[2]].as_primitive::(); - if codes.null_count() != 0 || adds.null_count() != 0 || scales.null_count() != 0 { - return Err(Error::invalid_input("null cached ex-plane row")); - } + let batch = &self.0; + let codes = batch + .column_by_name(names[0]) + .and_then(|codes| codes.as_fixed_size_list_opt()) + .filter(|codes| codes.null_count() == 0) + .ok_or_else(|| { + Error::invalid_input(format!( + "cached ex plane {plane} needs a non-null fixed-size list column {}", + names[0] + )) + })?; + let values = codes + .values() + .as_primitive_opt::() + .filter(|values| values.null_count() == 0) + .ok_or_else(|| { + Error::invalid_input(format!( + "cached ex-plane codes {} must be non-null uint8, got {}", + names[0], + codes.value_type() + )) + })? + .values(); + let factors = match ( + factor_values(batch, names[1])?, + factor_values(batch, names[2])?, + ) { + (Some(adds), Some(scales)) => Some((adds, scales)), + (None, None) => None, + _ => { + return Err(Error::invalid_input(format!( + "cached ex plane {plane} has only one of its factor columns {} and {}", + names[1], names[2] + ))); + } + }; let width = codes.value_length() as usize; - let values = codes.values().as_primitive::().values(); + let width_u32 = u32::try_from(width) + .ok() + .filter(|&width| width > 0) + .ok_or_else(|| { + Error::invalid_input(format!("invalid cached ex-plane code width {width}")) + })?; + writer.write_u8(plane)?; + writer.write_u8(if factors.is_some() { + EX_FACTOR_COLUMNS + } else { + 0 + })?; let writer = writer.raw_writer(); - writer.write_all(&(self.0.num_rows() as u64).to_le_bytes())?; - writer.write_all(&(width as u32).to_le_bytes())?; - for row in 0..self.0.num_rows() { + writer.write_all(&(batch.num_rows() as u64).to_le_bytes())?; + writer.write_all(&width_u32.to_le_bytes())?; + for row in 0..batch.num_rows() { writer.write_all(&values[row * width..(row + 1) * width])?; - writer.write_all(&adds.value(row).to_le_bytes())?; - writer.write_all(&scales.value(row).to_le_bytes())?; + if let Some((adds, scales)) = factors { + writer.write_all(&adds[row].to_le_bytes())?; + writer.write_all(&scales[row].to_le_bytes())?; + } } Ok(()) } - fn deserialize(reader: &mut CacheEntryReader<'_>) -> Result { - if reader.version() == 1 { - return Ok(Self(reader.read_ipc()?)); - } - if reader.version() != Self::CURRENT_VERSION { - return Err(Error::invalid_input("unsupported plane cache version")); +} + +impl CacheCodecImpl for PlaneBatch { + const TYPE_ID: &'static str = "lance.vector.rq.plane"; + const CURRENT_VERSION: u32 = BodyVersion::RowsRawSign as u32; + const SUPPORTS_ROW_SELECTION: bool = true; + + fn serialize(&self, writer: &mut CacheEntryWriter<'_>) -> Result<()> { + match ex_plane(&self.0) { + Some(plane) => self.serialize_ex_plane(plane, writer), + None => { + writer.write_u8(whole_plane(&self.0))?; + write_raw_batch(writer, &self.0) + } } + } + + fn deserialize(reader: &mut CacheEntryReader<'_>) -> Result { + let version = BodyVersion::of(reader.version())?; let body = reader.body(); - if body.first() == Some(&0) { - reader.read_u8()?; - return Ok(Self(reader.read_ipc()?)); + match (version, body.first().copied()) { + (BodyVersion::Ipc, _) => Ok(Self(reader.read_ipc()?)), + (BodyVersion::RowsIpcSign, Some(0)) => { + reader.read_u8()?; + Ok(Self(reader.read_ipc()?)) + } + (BodyVersion::RowsRawSign, Some(0 | SIGN_BOUNDS_PLANE)) => { + reader.read_u8()?; + Ok(Self(read_raw_batch(reader)?)) + } + _ => { + let header = Header::parse(&body, version)?; + header.decode(&body[header.len..]) + } } - let header = Header::parse(&body)?; - header.decode(&body[HEADER_BYTES..]) } + fn deserialize_rows( reader: &dyn CacheRangeReader, offset: usize, @@ -157,7 +350,6 @@ impl CacheCodecImpl for PlaneBatch { rows: &[u32], ) -> Result { let header = read_row_header(reader, offset, version)?; - let width = header.width; let stride = header.stride(); let mut codes = Vec::new(); let mut adds = Vec::new(); @@ -172,7 +364,7 @@ impl CacheCodecImpl for PlaneBatch { for &row in &rows[run] { let local = (row as usize - first) * stride; header.append( - &bytes[local..local + width + 8], + &bytes[local..local + stride], &mut codes, &mut adds, &mut scales, @@ -181,6 +373,7 @@ impl CacheCodecImpl for PlaneBatch { } header.batch(codes, adds, scales) } + fn plan_row_ranges( reader: &dyn CacheRangeReader, offset: usize, @@ -198,12 +391,11 @@ impl CacheCodecImpl for PlaneBatch { } fn read_row_header(reader: &dyn CacheRangeReader, offset: usize, version: u32) -> Result
{ - if version != PlaneBatch::CURRENT_VERSION { - return Err(Error::invalid_input( - "plane cache version has no row directory", - )); - } - Header::parse(&reader.read_range(offset..offset + HEADER_BYTES)?) + let version = BodyVersion::of(version)?; + let end = offset + .checked_add(version.ex_header_bytes()?) + .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; + Header::parse(&reader.read_range(offset..end)?, version) } /// Payload byte ranges for sorted, unique `rows`, each paired with the span of @@ -220,7 +412,7 @@ fn row_runs( } let stride = header.stride(); let base = offset - .checked_add(HEADER_BYTES) + .checked_add(header.len) .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; let byte_offset = |row: usize| { row.checked_mul(stride) @@ -246,87 +438,164 @@ fn row_runs( #[cfg(test)] mod tests { use super::*; - use crate::vector::bq::layered::{PlaneKey, SIGN_BOUNDS_PLANE}; + use crate::vector::bq::layered::{ + EntryColumns, FULL_BOUNDS_COLUMN, HIGH_BOUNDS_COLUMN, PlaneKey, plane_entry_columns, + }; + use crate::vector::bq::raw_body::RAW_BODY_KIND; + use crate::vector::bq::storage::RABIT_CODE_COLUMN; + use arrow_array::{ArrayRef, UInt64Array}; use bytes::Bytes; use lance_arrow::RecordBatchExt; - use lance_core::cache::{CacheCodec, CacheDecode, CacheKey}; - use std::cell::Cell; + use lance_core::cache::{CacheCodec, CacheDecode, CacheKey, CacheMissReason}; + use std::cell::{Cell, RefCell}; - /// Bounds plane entries keep an IPC body, as sign planes do: they round - /// trip whole and have no row directory, so a row read misses. - #[test] - fn bounds_plane_round_trips_through_ipc() { - const ROWS: usize = 70; - let names = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy); - let bounds = |offset: f32| { - let values: Vec = (0..ROWS * 3).map(|v| v as f32 + offset).collect(); - arrow_array::FixedSizeListArray::try_new_from_values(Float32Array::from(values), 3) - .unwrap() - }; - let (high, full) = (bounds(0.5), bounds(0.25)); - let original = RecordBatch::try_new( - Arc::new(Schema::new(vec![ - Field::new(names[0], high.data_type().clone(), true), - Field::new(names[1], full.data_type().clone(), true), - ])), - vec![Arc::new(high), Arc::new(full)], - ) - .unwrap(); - let codec = PlaneKey { - partition: 0, - plane: SIGN_BOUNDS_PLANE, - sign_bounds: SignBounds::Lazy, - } - .codec_for_key() - .unwrap(); + /// Ex-plane code widths, including one that is not a multiple of 8. + const WIDTHS: [usize; 4] = [56, 128, 512, 1024]; + + /// Envelope bytes ahead of a plane entry's body. + fn body_offset() -> usize { + 4 + 1 + 2 + PlaneBatch::TYPE_ID.len() + 4 + } + + fn encode(codec: &CacheCodec, batch: &RecordBatch) -> Bytes { let mut encoded = Vec::new(); codec .serialize( - &(Arc::new(PlaneBatch(original.clone())) as Arc), + &(Arc::new(PlaneBatch(batch.clone())) as Arc), &mut encoded, ) .unwrap(); - let encoded = Bytes::from(encoded); - let CacheDecode::Hit(decoded) = codec.deserialize(&encoded) else { - panic!("bounds plane decode failed") - }; - assert_eq!(decoded.downcast::().unwrap().0, original); + Bytes::from(encoded) + } + + fn decode(codec: &CacheCodec, encoded: &Bytes) -> CacheDecode { + match codec.deserialize(encoded) { + CacheDecode::Hit(decoded) => { + CacheDecode::Hit(decoded.downcast::().unwrap().0.clone()) + } + CacheDecode::Miss(reason) => CacheDecode::Miss(reason), + } + } + + fn decode_rows(codec: &CacheCodec, encoded: &Bytes, rows: &[u32]) -> CacheDecode { let read = |range: Range| -> Result { Ok(encoded.slice(range)) }; - assert!(matches!( - codec.deserialize_rows(&read, &[0, 5]), - CacheDecode::Miss(_) - )); - assert!(codec.plan_rows(&read, &[0, 5]).is_none()); + match codec.deserialize_rows(&read, rows) { + CacheDecode::Hit(decoded) => { + CacheDecode::Hit(decoded.downcast::().unwrap().0.clone()) + } + CacheDecode::Miss(reason) => CacheDecode::Miss(reason), + } + } + + fn ex_plane_batch(plane: u8, rows: usize, width: usize, factor_columns: u8) -> RecordBatch { + Header { + plane, + factor_columns, + rows, + width, + len: V3_HEADER_BYTES, + } + .batch( + (0..rows * width).map(|v| (v % 251) as u8).collect(), + (0..rows).map(|v| v as f32).collect(), + (0..rows).map(|v| v as f32 * 0.5 + 2.0).collect(), + ) + .unwrap() + .0 + } + + /// A sign or bounds plane batch with the columns `names`. + fn whole_plane_batch(names: &[&str], rows: usize) -> RecordBatch { + let columns = names + .iter() + .map(|name| { + let array: ArrayRef = match *name { + lance_core::ROW_ID => Arc::new(UInt64Array::from( + (0..rows as u64).map(|v| v * 3).collect::>(), + )), + RABIT_CODE_COLUMN => Arc::new( + FixedSizeListArray::try_new_from_values( + UInt8Array::from( + (0..rows * 16).map(|v| (v % 253) as u8).collect::>(), + ), + 16, + ) + .unwrap(), + ), + HIGH_BOUNDS_COLUMN | FULL_BOUNDS_COLUMN => Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from( + (0..rows * 3).map(|v| v as f32 - 1.5).collect::>(), + ), + 3, + ) + .unwrap(), + ), + _ => Arc::new(Float32Array::from( + (0..rows).map(|v| v as f32 + 0.25).collect::>(), + )), + }; + (*name, array) + }) + .collect::>(); + RecordBatch::try_from_iter(columns).unwrap() + } + + #[test] + fn ex_planes_round_trip_whole_and_by_rows() { + const ROWS: usize = 700; + let codec = CacheCodec::from_impl::(); + for plane in [1, 2] { + for width in WIDTHS { + for factor_columns in [EX_FACTOR_COLUMNS, 0] { + let context = format!("plane {plane} width {width} factors {factor_columns}"); + let original = ex_plane_batch(plane, ROWS, width, factor_columns); + let encoded = encode(&codec, &original); + let body = &encoded[body_offset()..]; + assert_eq!(&body[..2], &[plane, factor_columns], "{context}"); + let stride = width + FACTOR_BYTES * factor_columns as usize; + assert_eq!(body.len(), V3_HEADER_BYTES + ROWS * stride, "{context}"); + let CacheDecode::Hit(full) = decode(&codec, &encoded) else { + panic!("whole decode failed: {context}") + }; + assert_eq!(full, original, "{context}"); + for rows in [ + vec![], + vec![0], + vec![ROWS as u32 - 1], + vec![0, 1, 350, 351, 699], + (0..ROWS as u32).step_by(3).collect(), + (0..ROWS as u32).collect(), + ] { + let CacheDecode::Hit(selected) = decode_rows(&codec, &encoded, &rows) + else { + panic!("row decode failed: {context} {rows:?}") + }; + assert_eq!( + selected, + original + .take(&arrow_array::UInt32Array::from(rows.clone())) + .unwrap(), + "{context} {rows:?}" + ); + } + for rows in [vec![ROWS as u32], vec![1, 1], vec![2, 1]] { + assert!( + matches!(decode_rows(&codec, &encoded, &rows), CacheDecode::Miss(_)), + "{context} {rows:?}" + ); + } + } + } + } } #[test] fn sparse_plane_cache_reads_match_full_decode() { for plane in [1, 2] { - let header = Header { - plane, - rows: 4096, - width: 128, - }; - let original = header - .batch( - (0..4096 * 128).map(|v| (v % 251) as u8).collect(), - (0..4096).map(|v| v as f32).collect(), - vec![2.; 4096], - ) - .unwrap(); + let original = ex_plane_batch(plane, 4096, 128, EX_FACTOR_COLUMNS); let codec = CacheCodec::from_impl::(); - let mut encoded = Vec::new(); - codec - .serialize( - &(Arc::new(original) as Arc), - &mut encoded, - ) - .unwrap(); - let encoded = Bytes::from(encoded); - let CacheDecode::Hit(full) = codec.deserialize(&encoded) else { - panic!("full decode failed") - }; - let full = full.downcast::().unwrap(); + let encoded = encode(&codec, &original); let bytes_read = Cell::new(0); let read = |range: std::ops::Range| -> Result { bytes_read.set(bytes_read.get() + range.len()); @@ -338,83 +607,327 @@ mod tests { }; assert_eq!( selected.downcast::().unwrap().0, - full.0 + original .take(&arrow_array::UInt32Array::from(rows.to_vec())) .unwrap() ); assert!(bytes_read.get() < encoded.len() / 10); - assert!(matches!( - codec.deserialize_rows(&read, &[4096]), - CacheDecode::Miss(_) - )); - assert!(matches!( - codec.deserialize_rows(&read, &[1, 1]), - CacheDecode::Miss(_) - )); } } #[test] fn row_plan_matches_ranges_read_by_decode() { const ROWS: usize = 4096; - const WIDTH: usize = 56; - let header = Header { - plane: 1, - rows: ROWS, - width: WIDTH, + for width in WIDTHS { + for factor_columns in [EX_FACTOR_COLUMNS, 0] { + let original = ex_plane_batch(1, ROWS, width, factor_columns); + let codec = CacheCodec::from_impl::(); + let encoded = encode(&codec, &original); + let requested = RefCell::new(Vec::new()); + let read = |range: Range| -> Result { + requested.borrow_mut().push(range.clone()); + Ok(encoded.slice(range)) + }; + let stride = width + FACTOR_BYTES * factor_columns as usize; + let gap = COALESCE_GAP_BYTES / stride; + let row_sets: Vec> = vec![ + vec![], + vec![0], + vec![ROWS as u32 - 1], + (0..ROWS as u32).collect(), + (0..ROWS as u32).step_by(gap + 1).collect(), + (0..ROWS as u32).step_by(gap + 2).collect(), + vec![1, 2, 3, 900, 901, 3000, 4095], + ]; + for rows in row_sets { + requested.borrow_mut().clear(); + let CacheDecode::Hit(_) = codec.deserialize_rows(&read, &rows) else { + panic!("range decode failed for {rows:?}") + }; + // The envelope and body header come first; the rest are row reads. + let decoded_rows = requested.borrow()[3..].to_vec(); + let plan = codec.plan_rows(&read, &rows).expect("plane codec plans"); + assert_eq!(plan, decoded_rows); + let expected_runs = if rows.is_empty() { + 0 + } else { + 1 + rows + .windows(2) + .filter(|w| (w[1] - w[0] - 1) as usize > gap) + .count() + }; + assert_eq!(plan.len(), expected_runs, "width {width} {rows:?}"); + } + assert!(codec.plan_rows(&read, &[5, 5]).is_none()); + assert!(codec.plan_rows(&read, &[ROWS as u32]).is_none()); + } + } + } + + /// Sign and bounds plane entries hold a raw body: they round trip whole, + /// with their schema, and have no row directory, so a row read misses. + #[test] + fn sign_and_bounds_planes_round_trip_through_raw_bodies() { + let sign_lazy = plane_columns(0, SignBounds::Lazy); + let sign_eager = plane_columns(0, SignBounds::Eager); + let bounds = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy); + let cases: [(&[&str], u8, SignBounds); 5] = [ + (sign_lazy, 0, SignBounds::Lazy), + (sign_eager, 0, SignBounds::Eager), + (bounds, SIGN_BOUNDS_PLANE, SignBounds::Lazy), + // Code-only sign entries: the file columns alone. + (&[RABIT_CODE_COLUMN], 0, SignBounds::Lazy), + ( + &[RABIT_CODE_COLUMN, HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN], + 0, + SignBounds::Eager, + ), + ]; + for (names, tag, sign_bounds) in cases { + for rows in [0, 70] { + let codec = PlaneKey { + partition: 0, + plane: tag, + sign_bounds, + entry_columns: EntryColumns::default(), + } + .codec_for_key() + .unwrap(); + let original = whole_plane_batch(names, rows); + let encoded = encode(&codec, &original); + let body = &encoded[body_offset()..]; + assert_eq!(&body[..2], &[tag, RAW_BODY_KIND], "{names:?}"); + let CacheDecode::Hit(decoded) = decode(&codec, &encoded) else { + panic!("{names:?} decode failed") + }; + assert_eq!(decoded, original, "{names:?}"); + assert_eq!(decoded.schema(), original.schema(), "{names:?}"); + assert!(matches!( + decode_rows(&codec, &encoded, &[0]), + CacheDecode::Miss(_) + )); + let read = |range: Range| -> Result { Ok(encoded.slice(range)) }; + assert!(codec.plan_rows(&read, &[0]).is_none()); + } + } + } + + /// Code-only entries (`EntryColumns::Codes`) hold the file columns of + /// their plane: an ex plane's rows carry no factor columns, and the sign + /// and bounds planes keep raw bodies, whole and by rows. + #[test] + fn code_only_entries_round_trip() { + const ROWS: usize = 90; + for sign_bounds in [SignBounds::Lazy, SignBounds::Eager] { + for plane in [0, 1, 2, SIGN_BOUNDS_PLANE] { + let names = plane_entry_columns(plane, sign_bounds, EntryColumns::Codes); + if names.is_empty() { + continue; + } + let codec = PlaneKey { + partition: 0, + plane, + sign_bounds, + entry_columns: EntryColumns::Codes, + } + .codec_for_key() + .unwrap(); + let context = format!("plane {plane} {sign_bounds:?} {names:?}"); + let original = match plane { + 1 | 2 => ex_plane_batch(plane, ROWS, 128, 0), + _ => whole_plane_batch(&names, ROWS), + }; + let schema = original.schema(); + let columns: Vec<&str> = schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect(); + assert_eq!(columns, names, "{context}"); + let encoded = encode(&codec, &original); + let body = &encoded[body_offset()..]; + match plane { + 1 | 2 => { + assert_eq!(&body[..2], &[plane, 0], "{context}"); + assert_eq!(body.len(), V3_HEADER_BYTES + ROWS * 128, "{context}"); + let rows = [0, 7, 8, 89]; + let CacheDecode::Hit(selected) = decode_rows(&codec, &encoded, &rows) + else { + panic!("{context}: row decode missed") + }; + assert_eq!( + selected, + original + .take(&arrow_array::UInt32Array::from(rows.to_vec())) + .unwrap(), + "{context}" + ); + } + _ => assert_eq!(body[1], RAW_BODY_KIND, "{context}"), + } + let CacheDecode::Hit(decoded) = decode(&codec, &encoded) else { + panic!("{context}: decode missed") + }; + assert_eq!(decoded, original, "{context}"); + } + } + } + + type AnyValue = Arc; + + /// A v1 body: an IPC section for every plane. + fn write_v1(any: &AnyValue, writer: &mut CacheEntryWriter<'_>) -> Result<()> { + writer.write_ipc(&any.downcast_ref::().unwrap().0) + } + + /// A v2 body: a tagged IPC section for the sign and bounds planes, rows + /// with both factors behind a 13-byte header for the ex planes. + fn write_v2(any: &AnyValue, writer: &mut CacheEntryWriter<'_>) -> Result<()> { + let batch = &any.downcast_ref::().unwrap().0; + let Some(plane) = ex_plane(batch) else { + writer.write_u8(0)?; + return writer.write_ipc(batch); }; - let original = header - .batch( - (0..ROWS * WIDTH).map(|v| (v % 253) as u8).collect(), - (0..ROWS).map(|v| v as f32).collect(), - vec![0.5; ROWS], - ) - .unwrap(); - let codec = CacheCodec::from_impl::(); - let mut encoded = Vec::new(); - codec - .serialize( - &(Arc::new(original) as Arc), - &mut encoded, - ) - .unwrap(); - let encoded = Bytes::from(encoded); - let requested = std::cell::RefCell::new(Vec::new()); - let read = |range: Range| -> Result { - requested.borrow_mut().push(range.clone()); - Ok(encoded.slice(range)) + let names = plane_columns(plane, SignBounds::default()); + let codes = batch[names[0]].as_fixed_size_list(); + let width = codes.value_length() as usize; + let values = codes.values().as_primitive::().values(); + let adds = batch[names[1]].as_primitive::(); + let scales = batch[names[2]].as_primitive::(); + writer.write_u8(plane)?; + let writer = writer.raw_writer(); + writer.write_all(&(batch.num_rows() as u64).to_le_bytes())?; + writer.write_all(&(width as u32).to_le_bytes())?; + for row in 0..batch.num_rows() { + writer.write_all(&values[row * width..(row + 1) * width])?; + writer.write_all(&adds.value(row).to_le_bytes())?; + writer.write_all(&scales.value(row).to_le_bytes())?; + } + Ok(()) + } + + fn write_only(_: &mut CacheEntryReader<'_>) -> Result { + Err(Error::invalid_input("legacy test codecs only serialize")) + } + + /// A codec that writes the bodies of an earlier plane cache version. + fn legacy_codec(version: u32) -> CacheCodec { + let serialize = match version { + 1 => write_v1, + 2 => write_v2, + _ => unreachable!("no legacy plane cache version {version}"), }; - let stride = WIDTH + 8; - let gap = COALESCE_GAP_BYTES / stride; - let row_sets: Vec> = vec![ - vec![], - vec![0], - vec![ROWS as u32 - 1], - (0..ROWS as u32).collect(), - (0..ROWS as u32).step_by(gap + 1).collect(), - (0..ROWS as u32).step_by(gap + 2).collect(), - vec![1, 2, 3, 900, 901, 3000, 4095], - ]; - for rows in row_sets { - requested.borrow_mut().clear(); - let CacheDecode::Hit(_) = codec.deserialize_rows(&read, &rows) else { - panic!("range decode failed for {rows:?}") - }; - // The envelope and body header come first; the rest are row reads. - let decoded_rows = requested.borrow()[3..].to_vec(); - let plan = codec.plan_rows(&read, &rows).expect("plane codec plans"); - assert_eq!(plan, decoded_rows); - let expected_runs = if rows.is_empty() { - 0 + CacheCodec::new(PlaneBatch::TYPE_ID, version, serialize, write_only) + } + + #[test] + fn earlier_versions_still_decode() { + let codec = CacheCodec::from_impl::(); + let sign = whole_plane_batch(plane_columns(0, SignBounds::Eager), 40); + let bounds = whole_plane_batch(plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy), 40); + let ex = ex_plane_batch(2, 300, 56, EX_FACTOR_COLUMNS); + for version in [1, 2] { + let legacy = legacy_codec(version); + for original in [&sign, &bounds, &ex] { + let encoded = encode(&legacy, original); + let CacheDecode::Hit(decoded) = decode(&codec, &encoded) else { + panic!("v{version} decode failed") + }; + assert_eq!(&decoded, original, "v{version}"); + } + let encoded = encode(&legacy, &ex); + let rows = [0, 5, 299]; + let selected = decode_rows(&codec, &encoded, &rows); + if version == 1 { + assert!(matches!(selected, CacheDecode::Miss(_))); } else { - 1 + rows - .windows(2) - .filter(|w| (w[1] - w[0] - 1) as usize > gap) - .count() - }; - assert_eq!(plan.len(), expected_runs); + let CacheDecode::Hit(selected) = selected else { + panic!("v2 row decode failed") + }; + assert_eq!( + selected, + ex.take(&arrow_array::UInt32Array::from(rows.to_vec())) + .unwrap() + ); + } + } + } + + #[test] + fn mismatched_factor_columns_are_errors() { + let ex = ex_plane_batch(1, 10, 56, EX_FACTOR_COLUMNS); + let names = plane_columns(1, SignBounds::default()); + let only_add = ex.project(&[0, 1]).unwrap(); + let only_scale = ex.project(&[0, 2]).unwrap(); + let (codes, adds, scales) = ( + ex[names[0]].clone(), + ex[names[1]].clone(), + ex[names[2]].clone(), + ); + let float_codes: ArrayRef = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.5f32; 10 * 4]), 4) + .unwrap(), + ); + let wrong_codes = RecordBatch::try_from_iter([ + (names[0], float_codes), + (names[1], adds), + (names[2], scales.clone()), + ]) + .unwrap(); + let integer_adds: ArrayRef = Arc::new(UInt64Array::from((0..10u64).collect::>())); + let wrong_factor = RecordBatch::try_from_iter([ + (names[0], codes), + (names[1], integer_adds), + (names[2], scales), + ]) + .unwrap(); + for batch in [only_add, only_scale, wrong_codes, wrong_factor] { + let mut buf = Vec::new(); + let result = PlaneBatch(batch).serialize(&mut CacheEntryWriter::new(&mut buf)); + assert!(result.is_err()); + } + + // A header that claims one factor column misses rather than panics. + let codec = CacheCodec::from_impl::(); + let mut encoded = encode(&codec, &ex).to_vec(); + encoded[body_offset() + 1] = 1; + let encoded = Bytes::from(encoded); + assert!(matches!( + decode(&codec, &encoded), + CacheDecode::Miss(CacheMissReason::BodyError) + )); + assert!(matches!( + decode_rows(&codec, &encoded, &[0]), + CacheDecode::Miss(_) + )); + } + + #[test] + fn newer_version_is_a_miss() { + let codec = CacheCodec::from_impl::(); + for original in [ + ex_plane_batch(1, 20, 128, EX_FACTOR_COLUMNS), + whole_plane_batch(plane_columns(0, SignBounds::Lazy), 20), + ] { + let mut encoded = encode(&codec, &original).to_vec(); + let version_at = body_offset() - 4; + assert_eq!( + u32::from_le_bytes(encoded[version_at..body_offset()].try_into().unwrap()), + PlaneBatch::CURRENT_VERSION + ); + encoded[version_at..body_offset()] + .copy_from_slice(&(PlaneBatch::CURRENT_VERSION + 1).to_le_bytes()); + let encoded = Bytes::from(encoded); + assert!(matches!( + decode(&codec, &encoded), + CacheDecode::Miss(CacheMissReason::VersionTooNew) + )); + assert!(matches!( + decode_rows(&codec, &encoded, &[0]), + CacheDecode::Miss(_) + )); + let read = |range: Range| -> Result { Ok(encoded.slice(range)) }; + assert!(codec.plan_rows(&read, &[0]).is_none()); } - assert!(codec.plan_rows(&read, &[5, 5]).is_none()); - assert!(codec.plan_rows(&read, &[ROWS as u32]).is_none()); } } diff --git a/rust/lance-index/src/vector/bq/raw_body.rs b/rust/lance-index/src/vector/bq/raw_body.rs new file mode 100644 index 00000000000..0a730d2abe0 --- /dev/null +++ b/rust/lance-index/src/vector/bq/raw_body.rs @@ -0,0 +1,609 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Raw fixed-width bodies for RaBitQ cache entries. +//! +//! Arrow IPC writes a validity bitmap for every array, all-valid ones +//! included, so a persisted `FixedSizeList` column costs `n / 8` +//! bytes per row beyond its values: 12.5-13% of a sign plane or a native +//! partition. RaBitQ cache entries hold only non-null fixed-width columns, +//! so they are written as their schema and value buffers instead: +//! +//! ```text +//! [kind: u8 = 1] +//! [schema_len: u32 LE][Arrow IPC schema message] +//! [rows: u64 LE] +//! per column, in schema order: +//! [values, little-endian][zero padding to a multiple of 8 bytes] +//! ``` +//! +//! A primitive column contributes its values and a `FixedSizeList` of a +//! primitive its child values. The schema message keeps the names, types, +//! nullability and metadata, so a decoded batch has the source's schema. +//! Any other batch (a column with nulls or of another type, or one written +//! on a big-endian host) is written as `[kind: u8 = 0]` followed by an +//! Arrow IPC section, so every batch round trips. +//! +//! Decoding copies the values into buffers aligned to their type, sized +//! exactly to them, as the ex-plane codec does (see `plane_cache`). + +use std::sync::Arc; + +use arrow::array::ArrayData; +use arrow::buffer::Buffer; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, RecordBatch, RecordBatchOptions, cast::AsArray, make_array, +}; +use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions}; +use arrow_schema::{DataType, Schema}; +use lance_core::cache::{CacheEntryReader, CacheEntryWriter}; +use lance_core::{Error, Result}; + +/// Body kind of an Arrow IPC section, the fallback for any batch. +const IPC_BODY_KIND: u8 = 0; +/// Body kind of a raw fixed-width body. +pub const RAW_BODY_KIND: u8 = 1; +/// Every column's values are zero-padded to a multiple of this many bytes. +const VALUE_PADDING: usize = 8; +const ZEROS: [u8; VALUE_PADDING] = [0; VALUE_PADDING]; +/// Name reported by decode errors, which have no file path. +const BODY_NAME: &str = "raw cache body"; + +/// How one column's values are stored in a raw body. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct RawColumn { + /// Bytes of one primitive value. + value_width: usize, + /// Primitive values per row: 1, or the list size of a `FixedSizeList`. + values_per_row: usize, +} + +impl RawColumn { + /// The layout of a column of `data_type`, or `None` if a raw body cannot + /// hold it: only integer and floating-point values, alone or in a + /// `FixedSizeList`, qualify. + fn of(data_type: &DataType) -> Option { + match data_type { + DataType::FixedSizeList(item, size) => Some(Self { + value_width: raw_value_width(item.data_type())?, + values_per_row: usize::try_from(*size).ok()?, + }), + data_type => Some(Self { + value_width: raw_value_width(data_type)?, + values_per_row: 1, + }), + } + } + + /// Bytes of the values of `rows` rows, before padding. + fn value_bytes(&self, rows: usize) -> Option { + rows.checked_mul(self.values_per_row)? + .checked_mul(self.value_width) + } +} + +fn raw_value_width(data_type: &DataType) -> Option { + if data_type.is_integer() || data_type.is_floating() { + data_type.primitive_width() + } else { + None + } +} + +fn padding(len: usize) -> usize { + (VALUE_PADDING - len % VALUE_PADDING) % VALUE_PADDING +} + +/// Write `batch` as the rest of the entry body: a raw fixed-width body when +/// every column qualifies (see the module docs), an Arrow IPC section +/// otherwise. Nothing may be written after it. +pub fn write_raw_batch(writer: &mut CacheEntryWriter<'_>, batch: &RecordBatch) -> Result<()> { + let Some(values) = raw_values(batch) else { + writer.write_u8(IPC_BODY_KIND)?; + return writer.write_ipc(batch); + }; + writer.write_u8(RAW_BODY_KIND)?; + let schema = schema_message(batch.schema_ref()); + let schema_len = u32::try_from(schema.len()).map_err(|_| { + Error::invalid_input(format!( + "raw cache body schema message has {} bytes, more than a u32 length", + schema.len() + )) + })?; + let out = writer.raw_writer(); + out.write_all(&schema_len.to_le_bytes())?; + out.write_all(&schema)?; + out.write_all(&(batch.num_rows() as u64).to_le_bytes())?; + for column in values { + out.write_all(column.as_slice())?; + out.write_all(&ZEROS[..padding(column.len())])?; + } + Ok(()) +} + +/// Read a body written by [`write_raw_batch`]: the rest of the entry body. +pub fn read_raw_batch(reader: &mut CacheEntryReader<'_>) -> Result { + match reader.read_u8()? { + IPC_BODY_KIND => reader.read_ipc(), + RAW_BODY_KIND => decode_raw(&reader.body()), + kind => Err(Error::corrupt_file_named( + BODY_NAME, + format!("unknown body kind {kind}"), + )), + } +} + +/// The value bytes of every column of `batch`, or `None` if some column +/// cannot go in a raw body. Raw bodies are little-endian, which is the +/// in-memory layout only on a little-endian host. +fn raw_values(batch: &RecordBatch) -> Option> { + if cfg!(target_endian = "big") { + return None; + } + batch + .columns() + .iter() + .map(|column| column_values(column.as_ref())) + .collect() +} + +fn column_values(array: &dyn Array) -> Option { + let layout = RawColumn::of(array.data_type())?; + if array.null_count() != 0 { + return None; + } + let values = match array.data_type() { + DataType::FixedSizeList(..) => { + let values = array.as_fixed_size_list_opt()?.values(); + // A list's child holds exactly its rows' values once sliced. + if values.len() != array.len().checked_mul(layout.values_per_row)? { + return None; + } + values.to_data() + } + _ => array.to_data(), + }; + if values.null_count() != 0 { + return None; + } + let buffer = values.buffers().first()?; + Some(buffer.slice_with_length( + values.offset() * layout.value_width, + values.len() * layout.value_width, + )) +} + +fn schema_message(schema: &Schema) -> Vec { + IpcDataGenerator::default() + .schema_to_bytes_with_dictionary_tracker( + schema, + &mut DictionaryTracker::new(false), + &IpcWriteOptions::default(), + ) + .ipc_message +} + +fn parse_schema(message: &[u8]) -> Result { + let corrupt = |reason: String| Error::corrupt_file_named(BODY_NAME, reason); + let message = arrow_ipc::root_as_message(message) + .map_err(|e| corrupt(format!("invalid schema message: {e}")))?; + let schema = message + .header_as_schema() + .ok_or_else(|| corrupt("schema message holds no schema".to_string()))?; + // `fb_to_schema` assumes the field list is present. + if schema.fields().is_none() { + return Err(corrupt("schema message has no field list".to_string())); + } + Ok(arrow_ipc::convert::fb_to_schema(schema)) +} + +/// A bounds-checked cursor over a raw body. +struct BodyCursor<'a> { + bytes: &'a [u8], + pos: usize, +} + +impl<'a> BodyCursor<'a> { + fn take(&mut self, len: usize, what: &str) -> Result<&'a [u8]> { + let end = self + .pos + .checked_add(len) + .filter(|&end| end <= self.bytes.len()) + .ok_or_else(|| { + Error::corrupt_file_named( + BODY_NAME, + format!( + "truncated at {what}: needs {len} bytes at offset {}, body has {}", + self.pos, + self.bytes.len() + ), + ) + })?; + let taken = &self.bytes[self.pos..end]; + self.pos = end; + Ok(taken) + } + + fn take_array(&mut self, what: &str) -> Result<[u8; N]> { + let mut array = [0; N]; + array.copy_from_slice(self.take(N, what)?); + Ok(array) + } +} + +fn decode_raw(body: &[u8]) -> Result { + let mut cursor = BodyCursor { + bytes: body, + pos: 0, + }; + let schema_len = u32::from_le_bytes(cursor.take_array("schema length")?) as usize; + let schema = Arc::new(parse_schema(cursor.take(schema_len, "schema")?)?); + let rows = u64::from_le_bytes(cursor.take_array("row count")?); + let rows = usize::try_from(rows).map_err(|_| { + Error::corrupt_file_named(BODY_NAME, format!("row count {rows} overflows usize")) + })?; + let mut columns = Vec::with_capacity(schema.fields().len()); + for field in schema.fields() { + let layout = RawColumn::of(field.data_type()).ok_or_else(|| { + Error::corrupt_file_named( + BODY_NAME, + format!( + "column {} of type {} is not fixed-width", + field.name(), + field.data_type() + ), + ) + })?; + let len = layout.value_bytes(rows).ok_or_else(|| { + Error::corrupt_file_named( + BODY_NAME, + format!("column {} of {rows} rows overflows usize", field.name()), + ) + })?; + let values = cursor.take(len, field.name())?; + cursor.take(padding(len), field.name())?; + columns.push(decode_column(field.data_type(), layout, rows, values)?); + } + if cursor.pos != body.len() { + return Err(Error::corrupt_file_named( + BODY_NAME, + format!( + "{} trailing bytes after the last column", + body.len() - cursor.pos + ), + )); + } + Ok(RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(rows)), + )?) +} + +fn decode_column( + data_type: &DataType, + layout: RawColumn, + rows: usize, + values: &[u8], +) -> Result { + // `values` holds `rows * values_per_row` values, as `value_bytes` checked. + let len = rows * layout.values_per_row; + match data_type { + DataType::FixedSizeList(item, size) => { + let child = decode_primitive(item.data_type(), layout.value_width, len, values)?; + Ok(Arc::new(FixedSizeListArray::try_new( + item.clone(), + *size, + child, + None, + )?)) + } + data_type => decode_primitive(data_type, layout.value_width, len, values), + } +} + +fn decode_primitive( + data_type: &DataType, + value_width: usize, + len: usize, + values: &[u8], +) -> Result { + let buffer = match value_width { + 1 => Buffer::from_vec(values.to_vec()), + 2 => Buffer::from_vec(le_values(values, u16::from_le_bytes)), + 4 => Buffer::from_vec(le_values(values, u32::from_le_bytes)), + 8 => Buffer::from_vec(le_values(values, u64::from_le_bytes)), + _ => { + return Err(Error::corrupt_file_named( + BODY_NAME, + format!("values of type {data_type} are {value_width} bytes wide"), + )); + } + }; + let data = ArrayData::builder(data_type.clone()) + .len(len) + .add_buffer(buffer) + .build()?; + Ok(make_array(data)) +} + +/// `bytes` as little-endian `N`-byte values, in a vector of exactly their +/// count, aligned for `T`. +fn le_values(bytes: &[u8], from_le_bytes: fn([u8; N]) -> T) -> Vec { + bytes + .as_chunks::() + .0 + .iter() + .copied() + .map(from_le_bytes) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::{Float32Array, StringArray, UInt8Array, UInt64Array}; + use arrow_schema::Field; + use bytes::Bytes; + use lance_arrow::FixedSizeListArrayExt; + use lance_core::deepsize::DeepSizeOf; + use std::collections::HashMap; + + use crate::vector::bq::layered::{ + FULL_BOUNDS_COLUMN, HIGH_ADD_FACTORS_COLUMN, HIGH_BOUNDS_COLUMN, HIGH_SCALE_FACTORS_COLUMN, + SIGN_BOUNDS_PLANE, SignBounds, plane_columns, + }; + use crate::vector::bq::storage::{ + RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN, RABIT_CODE_COLUMN, + }; + use crate::vector::bq::transform::{ + ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + }; + + /// Sign code bytes of a 1024-dimensional index, as on MS MARCO. + const SIGN_CODE_BYTES: i32 = 128; + /// Blocked ex-code bytes of a 1024-dimensional RQ7 (6 ex bits) index. + const EX_CODE_BYTES: i32 = 768; + const LOW_CODE_BYTES: i32 = 256; + const BOUNDS_VALUES: i32 = 3; + + /// A column of an RQ cache entry with deterministic, row-dependent values. + fn column(name: &str, rows: usize) -> (Field, ArrayRef) { + let codes = |width: i32| -> ArrayRef { + let values = (0..rows * width as usize) + .map(|v| (v % 251) as u8) + .collect::>(); + Arc::new( + FixedSizeListArray::try_new_from_values(UInt8Array::from(values), width).unwrap(), + ) + }; + let array: ArrayRef = match name { + lance_core::ROW_ID => Arc::new(UInt64Array::from( + (0..rows as u64) + .map(|v| v * 7 + (1 << 40)) + .collect::>(), + )), + RABIT_CODE_COLUMN => codes(SIGN_CODE_BYTES), + RABIT_BLOCKED_EX_CODE_COLUMN => codes(EX_CODE_BYTES), + RABIT_BLOCKED_EX_CODE_LO_COLUMN => codes(LOW_CODE_BYTES), + HIGH_BOUNDS_COLUMN | FULL_BOUNDS_COLUMN => { + let values = (0..rows * BOUNDS_VALUES as usize) + .map(|v| v as f32 * 0.5 - 3.0) + .collect::>(); + Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(values), + BOUNDS_VALUES, + ) + .unwrap(), + ) + } + _ => Arc::new(Float32Array::from( + (0..rows).map(|v| v as f32 * 0.25 + 1.0).collect::>(), + )), + }; + (Field::new(name, array.data_type().clone(), true), array) + } + + fn batch(names: &[&str], rows: usize) -> RecordBatch { + let (fields, columns): (Vec<_>, Vec<_>) = + names.iter().map(|name| column(name, rows)).unzip(); + let metadata = HashMap::from([("lance:rq".to_string(), "entry".to_string())]); + RecordBatch::try_new( + Arc::new(Schema::new_with_metadata(fields, metadata)), + columns, + ) + .unwrap() + } + + /// The column sets of every sign, bounds and native entry: full entries + /// and code-only ones (the file columns), under both bounds placements. + fn entry_schemas() -> Vec<(&'static str, Vec<&'static str>)> { + let native = vec![ + lance_core::ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + ]; + let layered = vec![ + lance_core::ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + HIGH_ADD_FACTORS_COLUMN, + HIGH_SCALE_FACTORS_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + HIGH_BOUNDS_COLUMN, + FULL_BOUNDS_COLUMN, + ]; + let is_code = |name: &&str| { + [ + RABIT_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + HIGH_BOUNDS_COLUMN, + FULL_BOUNDS_COLUMN, + ] + .contains(name) + }; + let code_only = |names: &[&'static str]| -> Vec<&'static str> { + names.iter().copied().filter(is_code).collect() + }; + let sign_lazy = plane_columns(0, SignBounds::Lazy).to_vec(); + let sign_eager = plane_columns(0, SignBounds::Eager).to_vec(); + let bounds = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy).to_vec(); + vec![ + ("sign lazy codes", code_only(&sign_lazy)), + ("sign eager codes", code_only(&sign_eager)), + ("sign lazy", sign_lazy), + ("sign eager", sign_eager), + ("bounds", bounds), + ("native codes", code_only(&native)), + ("native", native), + ("layered partition codes", code_only(&layered)), + ("layered partition", layered), + ] + } + + fn encode(batch: &RecordBatch) -> Bytes { + let mut buf = Vec::new(); + write_raw_batch(&mut CacheEntryWriter::new(&mut buf), batch).unwrap(); + Bytes::from(buf) + } + + fn decode(bytes: &Bytes) -> Result { + read_raw_batch(&mut CacheEntryReader::new(bytes, 0, 1)) + } + + /// Bytes one row of `schema` holds in a raw body: its logical width. + fn row_bytes(schema: &Schema) -> Vec { + schema + .fields() + .iter() + .map(|field| { + let layout = RawColumn::of(field.data_type()).unwrap(); + layout.value_width * layout.values_per_row + }) + .collect() + } + + #[test] + fn entry_batches_round_trip_with_schema_and_size() { + for (name, columns) in entry_schemas() { + for rows in [0, 1, 7, 300] { + let source = batch(&columns, rows); + let encoded = encode(&source); + assert_eq!(encoded[0], RAW_BODY_KIND, "{name} {rows}"); + let decoded = decode(&encoded).unwrap(); + assert_eq!(decoded, source, "{name} {rows}"); + assert_eq!(decoded.schema(), source.schema(), "{name} {rows}"); + assert_eq!( + decoded.deep_size_of(), + source.deep_size_of(), + "{name} {rows}" + ); + } + } + } + + #[test] + fn body_holds_schema_and_logical_row_bytes_only() { + for (name, columns) in entry_schemas() { + let rows = 301; + let source = batch(&columns, rows); + let encoded = encode(&source); + let schema_len = u32::from_le_bytes(encoded[1..5].try_into().unwrap()) as usize; + assert_eq!( + schema_len, + schema_message(source.schema_ref()).len(), + "{name}" + ); + let row_bytes = row_bytes(source.schema_ref()); + let values: usize = row_bytes.iter().map(|bytes| rows * bytes).sum(); + let padding: usize = row_bytes.iter().map(|bytes| padding(rows * bytes)).sum(); + assert!(padding < VALUE_PADDING * row_bytes.len(), "{name}"); + // Kind, schema length, schema, row count, then the values: no + // validity bitmap, which IPC adds per column. + assert_eq!( + encoded.len(), + 1 + 4 + schema_len + 8 + values + padding, + "{name}" + ); + } + } + + #[test] + fn sliced_columns_write_their_rows_only() { + let columns = plane_columns(0, SignBounds::Eager); + let sliced = batch(columns, 50).slice(7, 30); + let encoded = encode(&sliced); + assert_eq!(decode(&encoded).unwrap(), sliced); + assert_eq!(encoded.len(), encode(&batch(columns, 30)).len()); + } + + #[test] + fn other_batches_fall_back_to_ipc() { + let with_nulls = RecordBatch::try_from_iter([( + ADD_FACTORS_COLUMN, + Arc::new(Float32Array::from(vec![Some(1.0), None, Some(3.0)])) as ArrayRef, + )]) + .unwrap(); + let null_child = RecordBatch::try_from_iter([( + HIGH_BOUNDS_COLUMN, + Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![Some(1.0), None, Some(3.0)]), + BOUNDS_VALUES, + ) + .unwrap(), + ) as ArrayRef, + )]) + .unwrap(); + let variable_width = RecordBatch::try_from_iter([ + ( + lance_core::ROW_ID, + Arc::new(UInt64Array::from(vec![1u64, 2])) as ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec!["a", "bc"])) as ArrayRef, + ), + ]) + .unwrap(); + for source in [with_nulls, null_child, variable_width] { + let encoded = encode(&source); + assert_eq!(encoded[0], IPC_BODY_KIND); + assert_eq!(decode(&encoded).unwrap(), source); + } + } + + #[test] + fn truncated_or_corrupt_bodies_are_errors() { + let source = batch(plane_columns(0, SignBounds::Eager), 5); + let encoded = encode(&source); + for len in 0..encoded.len() { + assert!(decode(&encoded.slice(..len)).is_err(), "truncated to {len}"); + } + let mut trailing = encoded.to_vec(); + trailing.push(0); + let err = decode(&Bytes::from(trailing)).unwrap_err().to_string(); + assert!(err.contains("trailing"), "{err}"); + let mut unknown_kind = encoded.to_vec(); + unknown_kind[0] = 7; + let err = decode(&Bytes::from(unknown_kind)).unwrap_err().to_string(); + assert!(err.contains("unknown body kind 7"), "{err}"); + let mut bad_schema = encoded.to_vec(); + bad_schema[5..13].fill(0xFF); + assert!(decode(&Bytes::from(bad_schema)).is_err()); + } +} diff --git a/rust/lance-index/src/vector/bq/resident.rs b/rust/lance-index/src/vector/bq/resident.rs index c60df58d8c1..d66fef0cb72 100644 --- a/rust/lance-index/src/vector/bq/resident.rs +++ b/rust/lance-index/src/vector/bq/resident.rs @@ -9,7 +9,10 @@ //! per index, later reads fetch only the codes (and bounds) from the file: //! on an object store a native partition miss takes 2 requests instead of //! 8, and a layered one 3 instead of 13. Reads attach copies of the resident -//! rows, so cache entries hold, and are charged, what a file read returns. +//! rows, so a batch holds, and is charged, what a file read returns. Cache +//! entries hold that batch, or only the columns reads fetch from the file +//! (`EntryColumns::Codes`), and every read of such an entry attaches the +//! resident rows again. //! //! The store is an entry of the index cache ([`ResidentColumnsKey`]), //! charged what its arrays allocate and evictable under ordinary cache @@ -24,8 +27,8 @@ use std::sync::{Arc, LazyLock, Mutex, Weak}; use std::time::Instant; use arrow::compute::concat_batches; -use arrow_array::{Array, ArrayRef, UInt64Array, new_empty_array}; -use arrow_schema::{Field, Schema}; +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt64Array, new_empty_array}; +use arrow_schema::{Field, Schema, SchemaRef}; use arrow_select::concat::concat; use arrow_select::take::take; use futures::TryStreamExt; @@ -73,6 +76,13 @@ pub(crate) fn is_resident(field: &Field) -> bool { resident_width(field).is_some() } +/// Whether reads always fetch column `name` of an IVF_RQ storage file from +/// the file: the columns a code-only cache entry holds. Every other column +/// of a plane or partition is a fixed-width one the store keeps. +pub(crate) fn is_file_column(name: &str) -> bool { + FILE_COLUMNS.contains(&name) +} + /// Bytes of values the resident store holds for an IVF_RQ storage file with /// `schema` and `num_rows` rows: every row of every column it keeps, at the /// column's fixed width. It reads nothing, so an open can size the store @@ -503,6 +513,56 @@ impl ResidentColumnStore { self.num_rows } + /// `schema`'s columns at the ascending file rows `rows`: copies of this + /// store's rows of the columns it keeps (see [`ResidentColumn::copy_rows`]) + /// and, by name, the other columns of `file_batch`, which holds them at + /// the same rows, read from the file or a code-only cache entry. The + /// batch is the one a read of every column from the file returns. + pub(crate) fn attach( + &self, + schema: SchemaRef, + file_batch: Option<&RecordBatch>, + rows: &UInt64Array, + ) -> Result { + let started = Instant::now(); + if let Some(batch) = file_batch + && batch.num_rows() != rows.len() + { + return Err(Error::internal(format!( + "resident columns attach to {} rows of a batch of {} rows", + rows.len(), + batch.num_rows() + ))); + } + let mut copied_bytes = 0; + let mut columns = Vec::with_capacity(schema.fields().len()); + for field in schema.fields() { + let column = match self.column(field.name()) { + Some(column) => { + let copy = column.copy_rows(rows)?; + copied_bytes += copy.get_buffer_memory_size() as u64; + copy + } + None => file_batch + .and_then(|batch| batch.column_by_name(field.name())) + .cloned() + .ok_or_else(|| Error::internal(format!("unread column {}", field.name())))?, + }; + columns.push(column); + } + let batch = RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(rows.len())), + )?; + let stats = layered_stats::counters(); + stats.resident_attach_calls.incr(); + stats.resident_attach_rows.add(rows.len() as u64); + stats.resident_attach_bytes.add(copied_bytes); + stats.resident_attach_ns.add_elapsed(started); + Ok(batch) + } + /// Heap memory the store holds: its arrays' buffers, which the /// `resident_columns_alloc_bytes` of its load count, plus its pages and /// its map of columns. @@ -986,6 +1046,50 @@ mod tests { } } + /// Attaching the resident columns to a batch of the other columns at + /// the same rows lays the columns out in the schema's order, copies the + /// resident rows and shares the others' buffers; a batch of other rows + /// or without a column the store does not keep is an error. + #[test] + fn attach_copies_resident_rows_and_takes_file_columns() { + let store = ResidentColumnStore { + columns: HashMap::from([("factor".to_string(), column())]), + num_rows: 100, + bytes: 400, + }; + assert!(is_file_column(RABIT_CODE_COLUMN)); + assert!(!is_file_column("factor")); + let schema = Arc::new(Schema::new(vec![ + Field::new(RABIT_CODE_COLUMN, DataType::UInt8, true), + Field::new("factor", DataType::Float32, true), + ])); + let rows = [5u64, 39, 40, 99]; + let codes: ArrayRef = Arc::new(arrow_array::UInt8Array::from(vec![1u8, 2, 3, 4])); + let file_batch = RecordBatch::try_from_iter([(RABIT_CODE_COLUMN, codes.clone())]).unwrap(); + let calls = || layered_stats::counters().resident_attach_calls.get(); + let before = calls(); + let batch = store + .attach( + schema.clone(), + Some(&file_batch), + &UInt64Array::from(rows.to_vec()), + ) + .unwrap(); + assert!(calls() > before); + assert_eq!(batch.schema(), schema); + assert_eq!(values(batch.column(1)), as_values(&rows)); + assert_eq!( + batch.column(0).to_data().buffers()[0].as_ptr(), + codes.to_data().buffers()[0].as_ptr() + ); + for (batch, rows) in [(Some(&file_batch), vec![5, 39]), (None, rows.to_vec())] { + let error = store + .attach(schema.clone(), batch, &UInt64Array::from(rows)) + .unwrap_err(); + assert!(matches!(error, Error::Internal { .. }), "{error}"); + } + } + #[test] fn copy_rows_across_pages_matches_a_file_read() { let column = column(); diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index efbbc3d1ccc..05b16cd1571 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -14,7 +14,8 @@ use std::sync::{ use arrow::array::AsArray; use arrow::datatypes::{Float16Type, Float32Type, Float64Type, UInt8Type, UInt64Type}; use arrow_array::{ - Array, FixedSizeListArray, Float32Array, RecordBatch, UInt8Array, UInt32Array, UInt64Array, + Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, UInt8Array, UInt32Array, + UInt64Array, }; use arrow_schema::{DataType, Field, SchemaRef}; use async_trait::async_trait; @@ -58,7 +59,7 @@ use crate::vector::bq::transform::{ SCALE_FACTORS_COLUMN, }; use crate::vector::bq::{ - RQRotationType, rabit_binary_code_bytes, rabit_ex_bits, rabit_ex_code_bytes, + RQRotationType, layered_stats, rabit_binary_code_bytes, rabit_ex_bits, rabit_ex_code_bytes, validate_rq_num_bits, }; use crate::vector::graph::{OrderedFloat, OrderedNode}; @@ -633,8 +634,17 @@ impl RabitQuantizationStorage { Rebuild the index." ))); } - let row_ids = batch[ROW_ID].as_primitive::().clone(); - let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); + // Columns are looked up by name and checked, so that a batch missing + // one, such as a code-only cache entry read without its resident + // columns, is an error rather than a panic. + let row_ids = required_column(&batch, ROW_ID)? + .as_primitive_opt::() + .ok_or_else(|| invalid_column_type(ROW_ID, "uint64"))? + .clone(); + let codes = required_column(&batch, RABIT_CODE_COLUMN)? + .as_fixed_size_list_opt() + .ok_or_else(|| invalid_column_type(RABIT_CODE_COLUMN, "a fixed-size list"))? + .clone(); // `rotated_dim() == 0` means the metadata never recorded a code // dimension, so the width check below could only ever report that the // column needs 0 bytes. Reject up front with the real cause. @@ -658,18 +668,21 @@ impl RabitQuantizationStorage { expected_code_bytes ))); } - let add_factors = batch[ADD_FACTORS_COLUMN] - .as_primitive::() - .clone(); - let scale_factors = batch[SCALE_FACTORS_COLUMN] - .as_primitive::() - .clone(); + let factors = |name: &str| -> Result { + Ok(required_column(&batch, name)? + .as_primitive_opt::() + .ok_or_else(|| invalid_column_type(name, "float32"))? + .clone()) + }; + let add_factors = factors(ADD_FACTORS_COLUMN)?; + let scale_factors = factors(SCALE_FACTORS_COLUMN)?; // Full precision preserves native binary pruning regardless of how // the ex codes are stored. Projected prefixes remove these factors and // use their own estimator-difference bounds instead. let error_factors = batch .column_by_name(ERROR_FACTORS_COLUMN) - .map(|factors| factors.as_primitive::().clone()); + .map(|_| factors(ERROR_FACTORS_COLUMN)) + .transpose()?; let ex_bits = rabit_ex_bits(metadata.num_bits)?; let mut batch = batch; let mut ex_codes = None; @@ -727,6 +740,15 @@ impl RabitQuantizationStorage { "RabitQ low plane requires layered metadata", )); } + // Codes a file stores unpacked or in the legacy sequential ex layout + // are rewritten below on every construction, including each one + // from a cache entry that kept them as stored. + let repacks_ex = ex_bits != 0 + && scope == PlaneScope::All + && repacks_sequential_ex_codes(&batch, rotated_dim, ex_bits); + if !metadata.packed || repacks_ex { + layered_stats::counters().storage_construct_repacks.incr(); + } if ex_bits != 0 && scope == PlaneScope::SignOnlyForFull { // The ex planes are gathered per query; see `ExRows`. } else if ex_bits != 0 { @@ -2842,6 +2864,23 @@ fn maybe_pack_ex_codes( } } +/// Whether legacy sequential ex-code rows of `seq_codes`' width are already +/// in the blocked layout, so they are used as stored. +fn sequential_rows_are_blocked(seq_codes: &FixedSizeListArray, dim: usize, ex_bits: u8) -> bool { + sequential_matches_blocked(ex_bits) + && seq_codes.value_length() as usize == blocked_ex_code_bytes(dim, ex_bits) +} + +/// Whether loading the ex codes of `batch` repacks rows: it holds only the +/// legacy sequential column, in a layout the blocked one does not match. +fn repacks_sequential_ex_codes(batch: &RecordBatch, dim: usize, ex_bits: u8) -> bool { + batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_none() + && batch + .column_by_name(RABIT_EX_CODE_COLUMN) + .and_then(|codes| codes.as_fixed_size_list_opt()) + .is_some_and(|codes| !sequential_rows_are_blocked(codes, dim, ex_bits)) +} + /// Bring legacy sequential ex codes into the blocked kernel layout: rows are /// repacked, except for the widths whose layouts agree byte-for-byte (then /// the column is used as stored). @@ -2850,9 +2889,7 @@ fn blocked_ex_codes_from_sequential( dim: usize, ex_bits: u8, ) -> Result { - if sequential_matches_blocked(ex_bits) - && seq_codes.value_length() as usize == blocked_ex_code_bytes(dim, ex_bits) - { + if sequential_rows_are_blocked(seq_codes, dim, ex_bits) { return Ok(seq_codes.clone()); } let seq_code_len = seq_codes.value_length() as usize; @@ -2871,6 +2908,56 @@ fn blocked_ex_codes_from_sequential( )?) } +/// Column `name` of a batch a storage is built from. +fn required_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef> { + batch.column_by_name(name).ok_or_else(|| { + Error::invalid_input(format!( + "RabitQ storage batch missing column {name}; it has {:?}", + batch + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + )) + }) +} + +fn invalid_column_type(name: &str, expected: &str) -> Error { + Error::invalid_input(format!("RabitQ storage column {name} must be {expected}")) +} + +/// The code columns of a code-only cache entry in the form storages use: +/// sign codes packed and ex codes in the blocked layout, rewritten here once +/// when the entry is read from the file rather than on every construction +/// from it (see `storage_construct_repacks`). A storage built from such an +/// entry takes `metadata` with `packed` set. Other columns are kept, and a +/// file whose codes are stored packed and blocked, as the index builder +/// writes them, is kept whole. +pub(crate) fn normalize_entry_codes( + batch: RecordBatch, + metadata: &RabitQuantizationMetadata, +) -> Result { + let mut batch = batch; + if !metadata.packed + && let Some(codes) = batch.column_by_name(RABIT_CODE_COLUMN) + { + let codes = codes + .as_fixed_size_list_opt() + .ok_or_else(|| invalid_column_type(RABIT_CODE_COLUMN, "a fixed-size list"))?; + let packed = Arc::new(pack_codes(codes)); + batch = batch.replace_column_by_name(RABIT_CODE_COLUMN, packed)?; + } + let ex_bits = rabit_ex_bits(metadata.num_bits)?; + if ex_bits != 0 + && batch.column_by_name(RABIT_EX_CODE_COLUMN).is_some() + && batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_none() + { + batch = load_blocked_ex_codes(batch, metadata.rotated_dim(), metadata.num_bits)?.0; + } + Ok(batch) +} + /// Load the ex-code planes of an index batch. A layered `1 + 4 + 4` batch /// carries [`RABIT_BLOCKED_EX_CODE_LO_COLUMN`]; then both planes are 4-bit /// blocked rows and the metadata must say 9 bits. Otherwise this is @@ -5931,6 +6018,155 @@ mod tests { assert!(stored_batch.column_by_name(ERROR_FACTORS_COLUMN).is_some()); } + /// A storage built from codes stored packed and blocked, as the index + /// builder writes them and cache entries keep them, uses them as stored. + /// Unpacked sign codes and sequential ex codes are rewritten, and + /// counted, on every construction. + #[test] + fn test_storage_construct_does_not_repack() { + const ROWS: usize = 50; + const CONSTRUCTIONS: u64 = 3; + // 6 ex bits: the sequential and blocked row layouts differ. + const NUM_BITS: u8 = 7; + let codes = make_test_codes(ROWS, 64); + let code_dim = codes.value_length() as usize * 8; + let ex_bits = rabit_ex_bits(NUM_BITS).unwrap(); + let mut metadata = make_test_metadata(code_dim); + metadata.num_bits = NUM_BITS; + assert!(!metadata.packed); + let file_batch = + make_test_batch_with_ex(codes, make_test_ex_codes(ROWS, code_dim, NUM_BITS)); + assert!(repacks_sequential_ex_codes(&file_batch, code_dim, ex_bits)); + + let repacks = || layered_stats::counters().storage_construct_repacks.get(); + let before = repacks(); + let construct = |batch: &RecordBatch, metadata: &RabitQuantizationMetadata| { + RabitQuantizationStorage::try_from_batch( + batch.clone(), + metadata, + DistanceType::L2, + None, + ) + .unwrap() + }; + for _ in 1..CONSTRUCTIONS { + construct(&file_batch, &metadata); + } + let storage = construct(&file_batch, &metadata); + // Other tests construct storages concurrently: only a lower bound holds. + assert!(repacks() - before >= CONSTRUCTIONS); + + let cached = storage.to_batches().unwrap().next().unwrap(); + assert!(storage.metadata().packed); + assert!(!repacks_sequential_ex_codes(&cached, code_dim, ex_bits)); + let rebuilt = construct(&cached, storage.metadata()); + let rebuilt_batch = rebuilt.to_batches().unwrap().next().unwrap(); + assert_eq!(rebuilt_batch, cached); + for column in [RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN] { + let values_ptr = |batch: &RecordBatch| { + batch[column] + .as_fixed_size_list() + .values() + .to_data() + .buffers()[0] + .as_ptr() + }; + assert_eq!( + values_ptr(&rebuilt_batch), + values_ptr(&cached), + "{column} was rewritten" + ); + } + + // A code-only cache entry keeps the codes normalized once, when it + // is read from the file: every storage built from it, with packed + // metadata, holds the storage built from the file's codes and uses + // the entry's codes as they are. + let values_ptr = |batch: &RecordBatch, column: &str| { + batch[column] + .as_fixed_size_list() + .values() + .to_data() + .buffers()[0] + .as_ptr() + }; + let normalized = normalize_entry_codes(file_batch, &metadata).unwrap(); + assert!(!repacks_sequential_ex_codes(&normalized, code_dim, ex_bits)); + let mut packed = metadata; + packed.packed = true; + for _ in 0..CONSTRUCTIONS { + let from_entry = construct(&normalized, &packed); + let entry_batch = from_entry.to_batches().unwrap().next().unwrap(); + assert_eq!(entry_batch, cached); + for column in [RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN] { + assert_eq!( + values_ptr(&entry_batch, column), + values_ptr(&normalized, column), + "{column} was rewritten" + ); + } + } + // Codes stored packed and blocked are kept as they are. + let kept = normalize_entry_codes(cached.clone(), &packed).unwrap(); + for column in [RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN] { + assert_eq!( + values_ptr(&kept, column), + values_ptr(&cached, column), + "{column}" + ); + } + } + + /// A batch missing a column a storage needs, as a code-only cache entry + /// that was not attached to its resident columns would, or holding one + /// of the wrong type, is an error. + #[test] + fn test_storage_batch_missing_columns_is_an_error() { + let codes = make_test_codes(50, 64); + let metadata = make_test_metadata(codes.value_length() as usize * 8); + let batch = make_test_batch(codes); + for column in [ + ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ] { + let error = RabitQuantizationStorage::try_from_batch( + batch.drop_column(column).unwrap(), + &metadata, + DistanceType::L2, + None, + ) + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains(column), "{error}"); + } + let (fields, columns): (Vec, Vec) = batch + .schema() + .fields() + .iter() + .zip(batch.columns()) + .map(|(field, column)| { + if field.name() == ADD_FACTORS_COLUMN { + let integers: ArrayRef = Arc::new(UInt64Array::from(vec![0u64; column.len()])); + ( + Field::new(ADD_FACTORS_COLUMN, DataType::UInt64, true), + integers, + ) + } else { + (field.as_ref().clone(), column.clone()) + } + }) + .unzip(); + let wrong_type = + RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns).unwrap(); + let error = + RabitQuantizationStorage::try_from_batch(wrong_type, &metadata, DistanceType::L2, None) + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains(ADD_FACTORS_COLUMN), "{error}"); + } + #[test] fn test_try_from_batch_accepts_missing_error_factors_for_compatibility() { let original_codes = make_test_codes(50, 64); diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 7025245a3b0..410cc8fe060 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -4,11 +4,16 @@ //! Vector Storage, holding (quantized) vectors and providing distance calculation. use crate::vector::bq::ex_dot::blocked_ex_code_bytes; -use crate::vector::bq::layered::{PlaneBatch, PlaneKey, RQLayout, SIGN_BOUNDS_PLANE, SignBounds}; +use crate::vector::bq::layered::{ + EntryColumns, PlaneBatch, PlaneKey, RQLayout, SIGN_BOUNDS_PLANE, SignBounds, plane_columns, + plane_entry_columns, +}; use crate::vector::bq::layered_stats; +use crate::vector::bq::partition_codes::{PartitionCodes, PartitionCodesKey}; use crate::vector::bq::resident::{ResidentColumnStore, is_resident}; +use crate::vector::bq::storage::normalize_entry_codes; use crate::vector::bq::transform::ERROR_FACTORS_COLUMN; -use crate::vector::quantizer::QuantizerStorage; +use crate::vector::quantizer::{QuantizationMetadata, QuantizerStorage}; use arrow::compute::concat_batches; use arrow_array::{ Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, UInt64Array, @@ -22,7 +27,7 @@ use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, ROW_ID, Result}; use lance_encoding::decoder::FilterExpression; -use lance_file::reader::FileReader; +use lance_file::reader::{FileReader, ReaderProjection}; use lance_index_core::remapping::{BatchRowIdRemapper, remap_row_ids_preserving_layout_async}; use lance_io::ReadBatchParams; use lance_io::object_store::{DEFAULT_CLOUD_IO_PARALLELISM, ObjectStore}; @@ -1274,6 +1279,45 @@ fn resident_lifetime_from(value: Option<&str>) -> Result { } } +/// What an IVF_RQ index's cache entries hold while its small columns are +/// resident (see [`RESIDENT_COLUMNS_ENV`]): `codes` (default) or `all`. With +/// `codes`, a layered index's sign, high and low plane entries and a native +/// flat index's partition entries ([`PartitionCodes`]) hold only the columns +/// reads fetch from the file, and every read attaches copies of the resident +/// rows, so the cache keeps no second copy of the row ids and factors. An +/// index without a resident store keeps `all`, whatever the setting (see +/// [`EntryColumns::resolve`]). Results are the same either way. Read once +/// per process; resolved when an index opens. +pub const ENTRY_COLUMNS_ENV: &str = "LANCE_RQ_ENTRY_COLUMNS"; + +/// The setting of [`ENTRY_COLUMNS_ENV`] when it is unset. +pub const DEFAULT_ENTRY_COLUMNS: EntryColumns = EntryColumns::Codes; + +/// [`ENTRY_COLUMNS_ENV`], read once per process. +static ENTRY_COLUMNS: LazyLock> = LazyLock::new(|| { + entry_columns_from(std::env::var(ENTRY_COLUMNS_ENV).ok().as_deref()) + .map_err(|err| err.to_string()) +}); + +/// The setting of [`ENTRY_COLUMNS_ENV`]. The variable is read once per +/// process; an invalid value fails here and every IVF_RQ index open. +pub fn entry_columns_setting() -> Result { + ENTRY_COLUMNS.clone().map_err(Error::invalid_input) +} + +fn entry_columns_from(value: Option<&str>) -> Result { + let Some(value) = value else { + return Ok(DEFAULT_ENTRY_COLUMNS); + }; + match value.trim() { + "codes" => Ok(EntryColumns::Codes), + "all" => Ok(EntryColumns::All), + _ => Err(Error::invalid_input(format!( + "{ENTRY_COLUMNS_ENV}={value:?} is invalid, expected codes or all" + ))), + } +} + /// Parse a `0`/`1` (or `false`/`true`) environment flag. fn parse_flag(name: &str, value: &str) -> Result { match value.trim() { @@ -1719,6 +1763,16 @@ fn sparse_read_extent(rows: &[u32], row_bytes: usize) -> (usize, usize) { (runs, pages) } +/// The file rows of the partition rows `range`: all of them, or the sorted +/// partition offsets `rows`. +fn partition_file_rows(range: &std::ops::Range, rows: Option<&[u32]>) -> UInt64Array { + let start = range.start as u64; + match rows { + Some(rows) => UInt64Array::from_iter_values(rows.iter().map(|&row| start + u64::from(row))), + None => UInt64Array::from_iter_values(start..range.end as u64), + } +} + /// How stage-2 survivors index a [`GatheredEx`]'s batches. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ExIndex { @@ -1752,6 +1806,22 @@ struct FetchedPlane { promotion: Option, } +/// Planes of a layered partition that have a cache entry: sign, high, low +/// and [`SIGN_BOUNDS_PLANE`]. +const ENTRY_PLANES: usize = SIGN_BOUNDS_PLANE as usize + 1; + +/// The schemas a storage attaches code-only entries to, built once per +/// storage rather than on every read of an entry. +#[derive(Debug, Default)] +struct AttachSchemas { + /// A plane's projection of the file, by bounds placement (lazy, eager) + /// and plane. + planes: [[OnceLock; ENTRY_PLANES]; 2], + /// A native partition's columns, which a storage built from its + /// code-only entry holds. + partition: OnceLock, +} + /// Loader to load partitioned PQ storage from disk. #[derive(Debug)] pub struct IvfQuantizationStorage { @@ -1775,6 +1845,17 @@ pub struct IvfQuantizationStorage { resident_columns: ResidentColumns, /// See [`Self::resident_columns_enabled`]. resident_columns_enabled: bool, + /// What the index asks its cache entries to hold; see + /// [`Self::entry_columns`]. + entry_columns: EntryColumns, + /// The file's metadata with its sign codes marked packed, which + /// storages built from code-only entries take, built on first use; `None` + /// when the file stores them packed. See [`Self::entry_metadata`]. + packed_metadata: OnceLock>, + /// The schemas code-only entries are attached to, built on first use: + /// every plane's projection of the file, by bounds placement, and a + /// native partition's; see [`Self::plane_schema`]. + attach_schemas: AttachSchemas, /// The file the storage reads, when bound to the runtime handles every /// open of it shares; see [`Self::with_index_file`]. `None` keeps the /// far gather pools to the plane access tracker. @@ -1873,6 +1954,9 @@ impl IvfQuantizationStorage { sign_bounds: SignBounds::default(), resident_columns: ResidentColumns::default(), resident_columns_enabled: false, + entry_columns: EntryColumns::default(), + packed_metadata: OnceLock::new(), + attach_schemas: AttachSchemas::default(), index_file: None, }) } @@ -1911,6 +1995,9 @@ impl IvfQuantizationStorage { sign_bounds: SignBounds::default(), resident_columns: ResidentColumns::default(), resident_columns_enabled: false, + entry_columns: EntryColumns::default(), + packed_metadata: OnceLock::new(), + attach_schemas: AttachSchemas::default(), index_file: None, } } @@ -2012,7 +2099,8 @@ impl IvfQuantizationStorage { } /// The size of the resident store of this storage's file, loaded or - /// not, and what an index cache charges for it. It reads nothing. + /// not, and what an index cache charges for it; zero for a plane-row + /// file. It reads nothing. pub fn resident_store_size(&self) -> Result { let bytes = self.resident_columns_bytes(); if bytes == 0 { @@ -2024,13 +2112,67 @@ impl IvfQuantizationStorage { }) } + /// Ask the cache entries to hold `entry_columns`, resolved when the + /// index opened ([`ENTRY_COLUMNS_ENV`]); see [`Self::entry_columns`]. + pub fn with_entry_columns(mut self, entry_columns: EntryColumns) -> Self { + self.entry_columns = entry_columns; + self + } + + /// What the cache entries of this storage hold: [`EntryColumns::Codes`] + /// when the index asked for it and reads take the small columns from + /// the resident store ([`Self::resident_columns_enabled`]), which every + /// read of a code-only entry needs; [`EntryColumns::All`] otherwise. + pub fn entry_columns(&self) -> EntryColumns { + self.entry_columns.resolve(self.resident_columns_enabled) + } + /// The cache key of plane `plane` of `partition` under this storage's - /// bounds placement. + /// bounds placement and entry columns. pub fn plane_key(&self, partition: usize, plane: u8) -> PlaneKey { PlaneKey { partition, plane, sign_bounds: self.sign_bounds, + entry_columns: self.entry_columns(), + } + } + + /// The metadata a storage built from this storage's cache entries takes. + /// Code-only entries keep their codes packed and blocked (see + /// [`Self::normalize_entry`]), so a storage built from them takes the + /// file's metadata with its sign codes marked packed and rewrites nothing; + /// full entries keep the codes as the file stores them. + fn entry_metadata(&self) -> Result<&Q::Metadata> { + if self.entry_columns() == EntryColumns::All { + return Ok(&self.metadata); + } + let packed = match self.packed_metadata.get() { + Some(packed) => packed, + None => { + let packed = match self.quantizer()? { + Quantizer::Rabit(rq) if !rq.metadata_ref().packed => { + let quantizer = Q::try_from(Quantizer::Rabit(rq))?; + Some(quantizer.metadata(Some(QuantizationMetadata { + transposed: true, + ..Default::default() + }))) + } + _ => None, + }; + self.packed_metadata.get_or_init(|| packed) + } + }; + Ok(packed.as_ref().unwrap_or(&self.metadata)) + } + + /// `batch`, the columns of a code-only entry read from the file, with its + /// codes as the entry keeps them: packed and blocked, once here rather + /// than on every construction of a storage from the entry. + fn normalize_entry(&self, batch: RecordBatch) -> Result { + match self.quantizer()? { + Quantizer::Rabit(rq) => normalize_entry_codes(batch, rq.metadata_ref()), + _ => Ok(batch), } } @@ -2116,6 +2258,23 @@ impl IvfQuantizationStorage { self.ivf.num_partitions() } + /// Translate assembled rows without changing cached encoded columns. + async fn remap_partition_batch( + &self, + batch: RecordBatch, + ) -> Result<(RecordBatch, Option>)> { + if let Some(remapping) = &self.batch_remapper { + lance_index_core::remapping::check_batch_remapping_entry()?; + let row_id_idx = batch.schema().index_of(ROW_ID)?; + let (batch, remapper) = + remap_row_ids_preserving_layout_async(remapping.as_ref(), batch, row_id_idx) + .await?; + Ok((batch, Some(remapper))) + } else { + Ok((batch, self.frag_reuse_index.clone())) + } + } + /// Load a partition's quantization storage, optionally measuring the exact /// I/O it performs into `io_stats`. /// @@ -2154,29 +2313,178 @@ impl IvfQuantizationStorage { .await?; concat_batches(&schema, batches.iter())? }; - if let Some(remapping) = &self.batch_remapper { - // Tagged asynchronous path. - lance_index_core::remapping::check_batch_remapping_entry()?; - let row_id_idx = batch.schema().index_of(ROW_ID)?; - let (batch, remapper) = - remap_row_ids_preserving_layout_async(remapping.as_ref(), batch, row_id_idx) - .await?; - return Q::Storage::try_from_batch_with_remapper( - batch, - self.metadata(), - self.distance_type, - Some(remapper), - ); - } - // Legacy synchronous remapping path. + let (batch, remapper) = self.remap_partition_batch(batch).await?; Q::Storage::try_from_batch_with_remapper( batch, self.metadata(), self.distance_type, - self.frag_reuse_index.clone(), + remapper, + ) + } + + /// The file columns a native partition's code-only entry + /// ([`PartitionCodes`]) holds, in the file's order: those the resident + /// store does not keep, the codes. + fn partition_code_columns(&self) -> Vec { + self.schema() + .fields() + .iter() + .filter(|field| !is_resident(field)) + .map(|field| field.name().clone()) + .collect() + } + + /// [`Self::partition_code_columns`] as a projection of the file. + pub fn partition_codes_projection(&self) -> Result { + let columns = self.partition_code_columns(); + let columns: Vec<&str> = columns.iter().map(String::as_str).collect(); + lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + &columns, + ) + } + + /// Read the code-only entry of native partition `part_id` from the file, + /// through a reader that also records into `io_stats`: its code columns + /// alone, which an object store serves in one request each. + pub async fn read_partition_codes( + &self, + part_id: usize, + io_stats: Option<&IoStats>, + ) -> Result { + let columns = self.partition_code_columns(); + let columns: Vec<&str> = columns.iter().map(String::as_str).collect(); + let batch = self + .read_file_columns(self.ivf.row_range(part_id), &columns, io_stats) + .await?; + Ok(PartitionCodes(self.normalize_entry(batch)?)) + } + + /// The code-only entry of a native partition from `batches`, its rows of + /// the [`Self::partition_codes_projection`] columns that a prewarm read + /// of several partitions decoded: copied into buffers of its own, so the + /// entry does not keep the whole read. + pub fn partition_codes_from_batches( + &self, + batches: Vec, + ) -> Result { + Ok(PartitionCodes( + self.normalize_entry(compact_prewarm_batches(batches)?)?, + )) + } + + /// The columns a storage built from a native partition's code-only + /// entry of `entry_schema` holds: the file's in its order, then the + /// entry's columns the file lacks. Ex codes the file stores sequentially + /// are kept blocked, where a storage built from the file lays them out + /// (see `load_blocked_ex_codes`). + fn partition_schema(&self, entry_schema: &arrow_schema::Schema) -> SchemaRef { + let file_schema = self.schema(); + let fields: Vec<_> = file_schema + .fields() + .iter() + .filter(|field| { + is_resident(field) || entry_schema.field_with_name(field.name()).is_ok() + }) + .chain( + entry_schema + .fields() + .iter() + .filter(|field| file_schema.field_with_name(field.name()).is_err()), + ) + .cloned() + .collect(); + Arc::new(arrow_schema::Schema::new_with_metadata( + fields, + file_schema.metadata().clone(), + )) + } + + /// The storage of native partition `part_id` built from `codes`, its + /// code-only entry: the columns the resident store keeps are copies of + /// its rows, the others the entry's, in the order a storage built from + /// a read of the file holds them, so the storage is that one, bit for + /// bit and byte for byte. The first build loads the store, whose I/O is + /// added to `load_stats`. + pub async fn partition_from_codes( + &self, + part_id: usize, + codes: &PartitionCodes, + load_stats: Option<&IoStats>, + ) -> Result { + let range = self.ivf.row_range(part_id); + let entry = &codes.0; + if entry.num_rows() != range.len() { + return Err(Error::internal(format!( + "code-only entry of partition {part_id} holds {} rows of its {}", + entry.num_rows(), + range.len() + ))); + } + // The entries of a file all hold the same columns, so the schema + // built for the first one serves every other. + let schema = self + .attach_schemas + .partition + .get_or_init(|| self.partition_schema(&entry.schema())) + .clone(); + let batch = if range.is_empty() { + // As a read of no rows, which loads no store. + RecordBatch::new_empty(schema) + } else { + let store = self + .resident_store(&schema, load_stats) + .await? + .ok_or_else(|| { + Error::internal(format!( + "the code-only entry of partition {part_id} needs the resident store" + )) + })?; + store.attach(schema, Some(entry), &partition_file_rows(&range, None))? + }; + let (batch, remapper) = self.remap_partition_batch(batch).await?; + Q::Storage::try_from_batch_with_remapper( + batch, + self.entry_metadata()?, + self.distance_type, + remapper, ) } + /// The storage of native partition `part_id` read through its code-only + /// entry in `cache` ([`PartitionCodesKey`]), loaded from the file on a + /// miss and admitted when `write_cache`, and whether the entry was a hit. + /// Every read attaches the resident columns ([`Self::partition_from_codes`]), + /// so a hit reads nothing from the file once the store is loaded. + pub async fn load_partition_cached( + &self, + part_id: usize, + cache: &WeakLanceCache, + write_cache: bool, + io_stats: Option, + ) -> Result<(Q::Storage, bool)> { + let key = PartitionCodesKey { partition: part_id }; + let (codes, hit) = if write_cache { + cache + .get_or_insert_with_key_hit(key, || async { + self.read_partition_codes(part_id, io_stats.as_ref()).await + }) + .await? + } else if let Some(codes) = cache.get_with_key(&key).await { + (codes, true) + } else { + let codes = self + .read_partition_codes(part_id, io_stats.as_ref()) + .await?; + (Arc::new(codes), false) + }; + let storage = self + .partition_from_codes(part_id, &codes, io_stats.as_ref()) + .await?; + Ok((storage, hit)) + } + /// Warm every plane entry a full-precision scan reads. Backend admission /// enforces its byte budget. /// @@ -2197,7 +2505,7 @@ impl IvfQuantizationStorage { cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { Ok(PlaneBatch( - self.read_plane(part_id, plane, None, None).await?, + self.read_plane_entry(part_id, plane, None).await?, )) }) .await?; @@ -2218,7 +2526,7 @@ impl IvfQuantizationStorage { cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { Ok(PlaneBatch( - self.read_plane(part_id, plane, None, None).await?, + self.read_plane_entry(part_id, plane, None).await?, )) }) .await?; @@ -2272,7 +2580,7 @@ impl IvfQuantizationStorage { return self.load_partition(part_id, io_stats).await; } let after_sign = self.planes_after_sign(precision); - let load_plane = |plane| self.load_plane_entry(part_id, plane, cache, io_stats.clone()); + let load_plane = |plane| self.load_plane(part_id, plane, cache, io_stats.clone()); let planes = if sequential_plane_loads()? || cache.plane_admission_gated() { // Admit the sign dependency first. The other reads can then // overlap without changing admission policy or assembled column order. @@ -2290,23 +2598,14 @@ impl IvfQuantizationStorage { let mut fields = Vec::new(); let mut columns = Vec::new(); for batch in planes { - fields.extend(batch.0.schema().fields().iter().cloned()); - columns.extend(batch.0.columns().iter().cloned()); + fields.extend(batch.schema().fields().iter().cloned()); + columns.extend(batch.columns().iter().cloned()); } let batch = RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns)?; - let (batch, remapper) = if let Some(remapping) = &self.batch_remapper { - lance_index_core::remapping::check_batch_remapping_entry()?; - let row_id_idx = batch.schema().index_of(ROW_ID)?; - let (batch, remapper) = - remap_row_ids_preserving_layout_async(remapping.as_ref(), batch, row_id_idx) - .await?; - (batch, Some(remapper)) - } else { - (batch, self.frag_reuse_index.clone()) - }; + let (batch, remapper) = self.remap_partition_batch(batch).await?; Q::Storage::try_from_batch_at_precision( batch, - self.metadata(), + self.entry_metadata()?, self.distance_type, remapper, precision, @@ -2318,7 +2617,10 @@ impl IvfQuantizationStorage { self.frag_reuse_index.is_none() && self.batch_remapper.is_none() } - /// Load one whole plane entry of a layered partition. + /// Load one whole plane entry of a layered partition, which holds the + /// plane or, under [`EntryColumns::Codes`], its file columns alone (see + /// [`Self::read_plane_entry`]); [`Self::load_plane`] attaches the + /// resident columns. /// /// A backend that gates lower planes on their sign plane admits a lower /// plane only while the partition's sign plane is resident; otherwise the @@ -2337,7 +2639,7 @@ impl IvfQuantizationStorage { return cache .get_or_insert_with_key(key, || async { Ok(PlaneBatch( - self.read_plane(part_id, plane, None, io_stats.clone()) + self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) }) @@ -2352,7 +2654,7 @@ impl IvfQuantizationStorage { cache .get_or_insert_with_key_hit(key, || async { Ok(PlaneBatch( - self.read_plane(part_id, plane, None, io_stats.clone()) + self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) }) @@ -2364,13 +2666,130 @@ impl IvfQuantizationStorage { // An oversized sign plane falls back to partition streaming; do // not admit smaller ex entries without their sign dependency. Arc::new(PlaneBatch( - self.read_plane(part_id, plane, None, io_stats.clone()) + self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) }; Ok(batch) } + /// One whole plane of a layered partition, read through the cache: its + /// entry ([`Self::load_plane_entry`]) with the resident columns attached + /// ([`Self::attach_resident`]), so the batch is the plane a read of the + /// file returns whatever the entry holds. + pub async fn load_plane( + &self, + part_id: usize, + plane: u8, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result { + let entry = self + .load_plane_entry(part_id, plane, cache, io_stats.clone()) + .await?; + self.attach_resident(part_id, plane, &entry.0, None, io_stats.as_ref()) + .await + } + + /// What the cache entry of plane `plane` of `part_id` holds, read from + /// the file: the whole plane ([`Self::read_plane`]) or, under + /// [`EntryColumns::Codes`], only its file columns + /// ([`plane_entry_columns`]), with the sign codes packed. Every loader + /// of a plane entry reads it here. + pub async fn read_plane_entry( + &self, + part_id: usize, + plane: u8, + io_stats: Option, + ) -> Result { + let entry_columns = self.entry_columns(); + if entry_columns == EntryColumns::All { + return self.read_plane(part_id, plane, None, io_stats).await; + } + let columns = plane_entry_columns(plane, self.sign_bounds, entry_columns); + if columns.is_empty() { + return Err(Error::invalid_input(format!( + "a layered partition has no plane {plane} with {} bounds", + self.sign_bounds + ))); + } + let range = self.ivf.row_range(part_id); + let batch = self + .read_file_columns(range, &columns, io_stats.as_ref()) + .await?; + self.normalize_entry(batch) + } + + /// Plane `plane`'s projection of the file under this storage's bounds + /// placement, as a read of the plane returns it. + fn plane_schema(&self, plane: u8) -> Result { + let build = || -> Result { + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + plane_columns(plane, self.sign_bounds), + )?; + Ok(Arc::new(arrow_schema::Schema::from( + projection.schema.as_ref(), + ))) + }; + let placement = usize::from(self.sign_bounds == SignBounds::Eager); + let Some(slot) = self.attach_schemas.planes[placement].get(usize::from(plane)) else { + return build(); + }; + if let Some(schema) = slot.get() { + return Ok(schema.clone()); + } + let schema = build()?; + Ok(slot.get_or_init(|| schema).clone()) + } + + /// Plane `plane` of partition `part_id` at the sorted partition offsets + /// `rows` (every row when `None`), assembled from `entry`, a cache entry + /// of the plane at those rows. A full entry is the plane. A code-only + /// entry ([`EntryColumns::Codes`]) gets copies of the resident store's + /// rows at the rows' file offsets, in the plane's column order, so the + /// batch is the one, bit for bit and byte for byte, that a read of the + /// file returns (see [`Self::read_plane`]). The first attach loads the + /// store, whose I/O is added to `load_stats`. + pub async fn attach_resident( + &self, + part_id: usize, + plane: u8, + entry: &RecordBatch, + rows: Option<&[u32]>, + load_stats: Option<&IoStats>, + ) -> Result { + if self.entry_columns() == EntryColumns::All { + return Ok(entry.clone()); + } + let range = self.ivf.row_range(part_id); + let expected_rows = rows.map_or(range.len(), <[u32]>::len); + if entry.num_rows() != expected_rows { + return Err(Error::internal(format!( + "code-only entry of partition {part_id} plane {plane} holds {} rows, the read needs {expected_rows}", + entry.num_rows() + ))); + } + let schema = self.plane_schema(plane)?; + if expected_rows == 0 { + // As a read of no rows, which loads no store. + return Ok(RecordBatch::new_empty(schema)); + } + let Some(store) = self.resident_store(&schema, load_stats).await? else { + // A plane without resident columns, the bounds plane, is its + // entry. + return Ok(entry.clone()); + }; + store + .attach(schema, Some(entry), &partition_file_rows(&range, rows)) + .map_err(|error| { + Error::internal(format!( + "attaching the resident columns of partition {part_id} plane {plane}: {error}" + )) + }) + } + /// Ex-plane geometry of a layered index. pub fn layered_ex_layout(&self) -> Result { let Quantizer::Rabit(rq) = self.quantizer()? else { @@ -2407,12 +2826,8 @@ impl IvfQuantizationStorage { "the lazy full-precision scan requires an index without a row-id remapper", )); } - let sign = self.load_plane_entry(part_id, 0, cache, io_stats).await?; - Q::Storage::try_from_sign_plane_for_full( - sign.0.clone(), - self.metadata(), - self.distance_type, - ) + let sign = self.load_plane(part_id, 0, cache, io_stats).await?; + Q::Storage::try_from_sign_plane_for_full(sign, self.entry_metadata()?, self.distance_type) } /// Fetch the high and low ex-plane rows at the sorted, unique partition @@ -2522,11 +2937,14 @@ impl IvfQuantizationStorage { let started = Instant::now(); let key = self.plane_key(part_id, plane); if source == PlaneSource::Resident - && let Some(batch) = cache.get_resident_with_key(&key).await + && let Some(entry) = cache.get_resident_with_key(&key).await { + let batch = self + .attach_resident(part_id, plane, &entry.0, None, io_stats.as_ref()) + .await?; stats.fetch_resident_ns.add_elapsed(started); return Ok(FetchedPlane { - batch: batch.0.clone(), + batch, source: PlaneSource::Resident, from_origin: false, promotion: None, @@ -2536,7 +2954,18 @@ impl IvfQuantizationStorage { let persisted = cache.get_rows_with_key(&key, rows).await; if persisted.is_some() || !origin_whole || cache.plane_admission_gated() { let (batch, from_origin) = match persisted { - Some(batch) => (batch.0.clone(), false), + Some(entry) => ( + self.attach_resident( + part_id, + plane, + &entry.0, + Some(rows), + io_stats.as_ref(), + ) + .await?, + false, + ), + // An origin read returns every column of the plane. None => ( self.read_origin_rows(part_id, plane, rows, config, io_stats.as_ref()) .await?, @@ -2563,12 +2992,10 @@ impl IvfQuantizationStorage { // it cached. stats.origin_whole_fallbacks.incr(); } - let batch = self - .load_plane_entry(part_id, plane, cache, io_stats) - .await?; + let batch = self.load_plane(part_id, plane, cache, io_stats).await?; stats.fetch_whole_ns.add_elapsed(started); Ok(FetchedPlane { - batch: batch.0.clone(), + batch, source: PlaneSource::Whole, from_origin: false, promotion: None, @@ -2650,6 +3077,7 @@ impl IvfQuantizationStorage { )); } let indices = arrow_array::UInt32Array::from(rows.clone()); + let code_only = self.entry_columns() == EntryColumns::Codes; // The planes are independent. Overlap their cache/origin reads, // while preserving plane order when assembling the full row schema. let planes = std::iter::once(0).chain(self.planes_after_sign(RQPrecision::Full)); @@ -2661,7 +3089,7 @@ impl IvfQuantizationStorage { let key = self.plane_key(part_id, plane); let resident = cache.get_resident_with_key(&key).await; let was_resident = resident.is_some(); - // Sign codes and bounds keep an IPC body, which has no row + // Sign codes and bounds keep a raw body, which has no row // directory, so their persisted entries are read whole. let row_addressable = plane != 0 && plane != SIGN_BOUNDS_PLANE; let cached = if !row_addressable && resident.is_none() { @@ -2674,30 +3102,64 @@ impl IvfQuantizationStorage { } else { None }; + // An entry's rows, attached to the resident columns when the + // entry holds its codes alone. + let attach = |entry: RecordBatch| { + let io_stats = io_stats.clone(); + async move { + self.attach_resident( + part_id, + plane, + &entry, + Some(rows.as_slice()), + io_stats.as_ref(), + ) + .await + } + }; let batch = if plane == 0 { // Sign codes are transposed across rows. Gather their physical // offsets directly, then pack only the selected rows. let raw = if let Some(value) = cached { value.0.clone() } else { - self.read_plane(part_id, plane, None, io_stats.clone()) + self.read_plane_entry(part_id, plane, io_stats.clone()) .await? }; let codes = raw .column_by_name(RABIT_CODE_COLUMN) .ok_or_else(|| Error::invalid_input("missing sign codes"))?; let selected_codes = take_packed_codes(codes.as_fixed_size_list(), rows)?; - // The transposed sign column cannot be gathered with Arrow - // take. Avoid copying it only to replace that copy immediately. - let field = raw.schema().field_with_name(RABIT_CODE_COLUMN)?.clone(); - raw.drop_column(RABIT_CODE_COLUMN)? - .take(indices)? - .try_with_column(field, Arc::new(selected_codes))? + if code_only { + // The entry's other columns, the bounds, at the rows. + let columns = raw + .schema() + .fields() + .iter() + .zip(raw.columns()) + .map(|(field, column)| { + if field.name() == RABIT_CODE_COLUMN { + Ok(Arc::new(selected_codes.clone()) as ArrayRef) + } else { + Ok(arrow_select::take::take(column, indices, None)?) + } + }) + .collect::>>()?; + attach(RecordBatch::try_new(raw.schema(), columns)?).await? + } else { + // The transposed sign column cannot be gathered with Arrow + // take. Avoid copying it only to replace that copy immediately. + let field = raw.schema().field_with_name(RABIT_CODE_COLUMN)?.clone(); + raw.drop_column(RABIT_CODE_COLUMN)? + .take(indices)? + .try_with_column(field, Arc::new(selected_codes))? + } } else if let Some(value) = cached { - value.0.take(indices)? + attach(value.0.take(indices)?).await? } else if let Some(value) = selected_cached { - value.0.clone() + attach(value.0.clone()).await? } else { + // An origin read returns every column of the plane. self.read_plane(part_id, plane, Some(rows.clone()), io_stats.clone()) .await? }; @@ -2735,8 +3197,8 @@ impl IvfQuantizationStorage { if promote { cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { - Ok(super::bq::layered::PlaneBatch( - self.read_plane(part_id, plane, None, io_stats.clone()) + Ok(PlaneBatch( + self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) }) @@ -2756,7 +3218,7 @@ impl IvfQuantizationStorage { let batch = RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns)?; Q::Storage::try_from_batch_at_precision( batch, - self.metadata(), + self.entry_metadata()?, self.distance_type, self.frag_reuse_index.clone(), RQPrecision::Full, @@ -2992,26 +3454,44 @@ impl IvfQuantizationStorage { .await?; Some(concat_batches(&file_schema, batches.iter())?) }; - let start = range.start as u64; - let file_rows = match rows { - Some(rows) => { - UInt64Array::from_iter_values(rows.iter().map(|&row| start + u64::from(row))) - } - None => UInt64Array::from_iter_values(start..range.end as u64), - }; - let mut columns = Vec::with_capacity(schema.fields().len()); - for field in schema.fields() { - let column = match store.column(field.name()) { - Some(column) => column.copy_rows(&file_rows)?, - None => file_batch - .as_ref() - .and_then(|batch| batch.column_by_name(field.name())) - .cloned() - .ok_or_else(|| Error::internal(format!("unread column {}", field.name())))?, - }; - columns.push(column); + store.attach( + schema, + file_batch.as_ref(), + &partition_file_rows(&range, rows), + ) + } + + /// `columns` of the file at every row of the partition rows `range`, + /// read through a reader that also records into `io_stats`: what a + /// code-only entry holds before its codes are normalized. + async fn read_file_columns( + &self, + range: std::ops::Range, + columns: &[&str], + io_stats: Option<&IoStats>, + ) -> Result { + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + columns, + )?; + let schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); + if range.is_empty() { + return Ok(RecordBatch::new_empty(schema)); } - Ok(RecordBatch::try_new(schema, columns)?) + let batches = self + .file_reader(io_stats, None) + .read_stream_projected( + ReadBatchParams::Range(range), + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + Ok(concat_batches(&schema, batches.iter())?) } /// Read only `columns` of the given rows; the batch keeps file schema order. @@ -3334,16 +3814,17 @@ impl IvfQuantizationStorage { #[cfg(test)] mod tests { use super::{ - DenseGatherMode, DenseToEager, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, - LAZY_DENSE_BYTES_FRACTION_ENV, LAZY_DENSE_ENV, LAZY_DENSE_TO_EAGER_ENV, - LAZY_EAGER_BEFORE_FULL_ENV, LAZY_FAR_INFLIGHT_ENV, LAZY_FAR_WINDOW_ENV, LAZY_FULL_ENV, - LAZY_INLINE_ROWS_ENV, LAZY_MAX_RUNS_ENV, LAZY_ORIGIN_GAP_BYTES_ENV, - LAZY_ORIGIN_MAX_RUNS_ENV, LAZY_PROMOTE_ENV, LAZY_PROMOTE_INFLIGHT_ENV, LAZY_WINDOW_ENV, - LayeredLazyConfig, LazyFarPermits, LazyOriginGap, LazyPromotion, ORIGIN_LATENCY_ENV, - OriginLatencyClass, PlaneSource, QueryScratchCapacity, QueryScratchPool, - RESIDENT_COLUMNS_ENV, RESIDENT_LIFETIME_ENV, ResidentColumnsSetting, ResidentLifetime, - ResidentStoreSize, SEQUENTIAL_PLANE_LOADS_ENV, SIGN_BOUNDS_ENV, SignBounds, - compact_prewarm_batches, origin_latency_from, origin_latency_setting, + DEFAULT_ENTRY_COLUMNS, DenseGatherMode, DenseToEager, ENTRY_COLUMNS_ENV, EntryColumns, + HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, LAZY_DENSE_BYTES_FRACTION_ENV, + LAZY_DENSE_ENV, LAZY_DENSE_TO_EAGER_ENV, LAZY_EAGER_BEFORE_FULL_ENV, LAZY_FAR_INFLIGHT_ENV, + LAZY_FAR_WINDOW_ENV, LAZY_FULL_ENV, LAZY_INLINE_ROWS_ENV, LAZY_MAX_RUNS_ENV, + LAZY_ORIGIN_GAP_BYTES_ENV, LAZY_ORIGIN_MAX_RUNS_ENV, LAZY_PROMOTE_ENV, + LAZY_PROMOTE_INFLIGHT_ENV, LAZY_WINDOW_ENV, LayeredLazyConfig, LazyFarPermits, + LazyOriginGap, LazyPromotion, ORIGIN_LATENCY_ENV, OriginLatencyClass, PlaneSource, + QueryScratchCapacity, QueryScratchPool, RESIDENT_COLUMNS_ENV, RESIDENT_LIFETIME_ENV, + ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, SEQUENTIAL_PLANE_LOADS_ENV, + SIGN_BOUNDS_ENV, SignBounds, compact_prewarm_batches, entry_columns_from, + entry_columns_setting, origin_latency_from, origin_latency_setting, origin_reads_whole_plane, plan_plane_gather, resident_columns_from, resident_columns_setting, resident_lifetime_from, resident_lifetime_setting, resident_store_fits, sequential_plane_loads, sequential_plane_loads_from, sign_bounds_from, @@ -4127,6 +4608,41 @@ mod tests { } } + /// `codes` by default; an index keeps code-only entries only with a + /// resident store. + #[test] + fn test_entry_columns_knob() { + let process = std::env::var(ENTRY_COLUMNS_ENV).ok(); + assert_eq!( + entry_columns_setting().ok(), + entry_columns_from(process.as_deref()).ok() + ); + assert_eq!(DEFAULT_ENTRY_COLUMNS, EntryColumns::Codes); + assert_eq!(entry_columns_from(None).unwrap(), EntryColumns::Codes); + assert_eq!( + entry_columns_from(Some(" all ")).unwrap(), + EntryColumns::All + ); + for setting in [EntryColumns::All, EntryColumns::Codes] { + assert_eq!(entry_columns_from(Some(setting.as_str())).unwrap(), setting); + } + for (setting, resident, expected) in [ + (EntryColumns::Codes, true, EntryColumns::Codes), + (EntryColumns::Codes, false, EntryColumns::All), + (EntryColumns::All, true, EntryColumns::All), + (EntryColumns::All, false, EntryColumns::All), + ] { + assert_eq!(setting.resolve(resident), expected, "{setting} {resident}"); + } + for value in ["code", "ALL", "1", ""] { + let error = entry_columns_from(Some(value)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(ENTRY_COLUMNS_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + } + /// `off` never keeps a store; `on` keeps any store with columns, however /// large; `auto` keeps one whose charge fits the entry limit of the /// cache's largest admissible entry, whatever its values' diff --git a/rust/lance/src/index/vector/ivf/partition_serde.rs b/rust/lance/src/index/vector/ivf/partition_serde.rs index b1230078b17..e4dbf2c5268 100644 --- a/rust/lance/src/index/vector/ivf/partition_serde.rs +++ b/rust/lance/src/index/vector/ivf/partition_serde.rs @@ -8,6 +8,10 @@ //! Arrow IPC sections in a fixed, version-keyed order: the sub-index, then any //! quantizer-specific arrays (PQ codebook, RabitQ Matrix rotation), then the //! quantizer storage batches. Sections decode zero-copy via [`lance_arrow::ipc`]. +//! +//! RabitQ entries from version 2 on write their storage batch as a raw +//! fixed-width body instead (see [`lance_index::vector::bq::raw_body`]), which +//! has no validity bitmaps; decoding copies it. use std::sync::Arc; @@ -17,6 +21,7 @@ use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter}; use lance_core::{Error, Result}; use lance_index::vector::bq::RQRotationType; use lance_index::vector::bq::builder::RabitQuantizer; +use lance_index::vector::bq::raw_body::{read_raw_batch, write_raw_batch}; use lance_index::vector::bq::storage::{RabitQuantizationMetadata, RabitQueryEstimator}; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatMetadata, FlatQuantizer}; use lance_index::vector::pq::ProductQuantizer; @@ -407,11 +412,16 @@ impl CacheCodecImpl for PartitionEntry { // RabitQ // --------------------------------------------------------------------------- -impl CacheCodecImpl for PartitionEntry { - const TYPE_ID: &'static str = "lance.vector.ivf.PartitionEntry.Rabit"; - const CURRENT_VERSION: u32 = 1; +/// RabitQ entry version whose storage batch is an Arrow IPC section. +const RABIT_IPC_STORAGE_VERSION: u32 = 1; +/// RabitQ entry version whose storage batch is a raw fixed-width body. +const RABIT_RAW_STORAGE_VERSION: u32 = 2; - fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { +impl PartitionEntry { + /// Write the sections every RabitQ entry version puts ahead of the + /// storage batch: the header, the sub-index and, for Matrix rotation, the + /// rotation matrix. + fn write_sections_before_storage(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { let metadata = self.storage.metadata(); let header = RabitPartitionHeader { layered: metadata.layered, @@ -436,10 +446,23 @@ impl CacheCodecImpl for PartitionEntry { })?; w.write_ipc(&fsl_to_batch(mat, "rotate_mat")?)?; } + Ok(()) + } +} - w.write_ipc_batches(self.storage.to_batches()?)?; +impl CacheCodecImpl for PartitionEntry { + const TYPE_ID: &'static str = "lance.vector.ivf.PartitionEntry.Rabit"; + const CURRENT_VERSION: u32 = RABIT_RAW_STORAGE_VERSION; - Ok(()) + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + self.write_sections_before_storage(w)?; + let mut batches = self.storage.to_batches()?; + let (Some(storage_batch), None) = (batches.next(), batches.next()) else { + return Err(Error::io( + "RabitQ partition storage must be exactly one batch".to_string(), + )); + }; + write_raw_batch(w, &storage_batch) } fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { @@ -456,7 +479,15 @@ impl CacheCodecImpl for PartitionEntry { None }; - let storage_batch = read_single_storage_batch(r)?; + let storage_batch = match r.version() { + RABIT_IPC_STORAGE_VERSION => read_single_storage_batch(r)?, + RABIT_RAW_STORAGE_VERSION => read_raw_batch(r)?, + version => { + return Err(Error::io(format!( + "unsupported RabitQ partition entry version {version}" + ))); + } + }; let index = S::load(sub_index_batch)?; // Read the proto enum accessor before moving fields out of `header`. @@ -502,6 +533,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema}; use half::f16; use lance_arrow::FixedSizeListArrayExt; + use lance_index::vector::bq::raw_body::RAW_BODY_KIND; use lance_index::vector::bq::storage::RABIT_CODE_COLUMN; use lance_index::vector::bq::transform::{ADD_FACTORS_COLUMN, SCALE_FACTORS_COLUMN}; use lance_index::vector::bq::{RQRotationType, builder::RabitQuantizer}; @@ -996,6 +1028,7 @@ mod tests { orig_codes.values().as_primitive::().values(), rest_codes.values().as_primitive::().values(), ); + assert_eq!(rest_batch, orig_batch); } #[test] @@ -1073,6 +1106,102 @@ mod tests { orig_mat.values().as_primitive::().values(), rest_mat.values().as_primitive::().values(), ); + assert_eq!(storage_batch(&restored), storage_batch(&entry)); + } + + fn storage_batch(entry: &PartitionEntry) -> RecordBatch { + entry.storage.to_batches().unwrap().next().unwrap() + } + + /// A v1 RabitQ body: the storage batch as an Arrow IPC section. + fn ser_rabit_v1_body(entry: &PartitionEntry) -> Vec { + let mut buf = Vec::new(); + let mut w = CacheEntryWriter::new(&mut buf); + entry.write_sections_before_storage(&mut w).unwrap(); + w.write_ipc_batches(entry.storage.to_batches().unwrap()) + .unwrap(); + buf + } + + /// Bodies written before the storage batch became a raw body still + /// decode, under either rotation. + #[test] + fn test_rabitq_v1_bodies_still_decode() { + for rotation_type in [RQRotationType::Fast, RQRotationType::Matrix] { + let storage = make_rabit_storage( + 40, + 32, + DistanceType::L2, + rotation_type, + RabitQueryEstimator::ResidualQuery, + ); + let entry = + PartitionEntry::::new(FlatIndex::default(), storage); + let v1 = bytes::Bytes::from(ser_rabit_v1_body(&entry)); + let restored = PartitionEntry::::deserialize( + &mut CacheEntryReader::new(&v1, 0, RABIT_IPC_STORAGE_VERSION), + ) + .unwrap(); + assert_eq!(storage_batch(&restored), storage_batch(&entry)); + assert_eq!( + restored.storage.metadata().rotation_type, + entry.storage.metadata().rotation_type + ); + assert_eq!( + restored.storage.metadata().rotate_mat, + entry.storage.metadata().rotate_mat + ); + } + } + + /// Bytes of one row of a fixed-width column: its logical width. + fn logical_row_bytes(data_type: &DataType) -> usize { + match data_type { + DataType::FixedSizeList(item, size) => { + *size as usize * item.data_type().primitive_width().unwrap() + } + data_type => data_type.primitive_width().unwrap(), + } + } + + /// The storage section of a v2 body holds the storage schema and each + /// column's values, 8-byte padded: none of the validity bitmaps an IPC + /// section writes. + #[test] + fn test_rabitq_storage_section_has_no_bitmap() { + const ROWS: usize = 41; + const VALUE_PADDING: usize = 8; + for rotation_type in [RQRotationType::Fast, RQRotationType::Matrix] { + let storage = make_rabit_storage( + ROWS, + 64, + DistanceType::L2, + rotation_type, + RabitQueryEstimator::ResidualQuery, + ); + let entry = + PartitionEntry::::new(FlatIndex::default(), storage); + let body = ser_body(&entry); + let mut prefix = Vec::new(); + entry + .write_sections_before_storage(&mut CacheEntryWriter::new(&mut prefix)) + .unwrap(); + assert_eq!(&body[..prefix.len()], &prefix[..]); + let section = &body[prefix.len()..]; + assert_eq!(section[0], RAW_BODY_KIND); + let schema_len = u32::from_le_bytes(section[1..5].try_into().unwrap()) as usize; + let values: usize = storage_batch(&entry) + .schema() + .fields() + .iter() + .map(|field| { + (ROWS * logical_row_bytes(field.data_type())).next_multiple_of(VALUE_PADDING) + }) + .sum(); + // Kind, schema length, schema, row count, then the values. + assert_eq!(section.len(), 1 + 4 + schema_len + 8 + values); + assert!(body.len() < ser_rabit_v1_body(&entry).len()); + } } /// SQ storage (a multi-batch IPC section) must decode zero-copy through the diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 6e71fc6711e..1e1e41f6679 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -50,8 +50,9 @@ use lance_index::prefilter::NoFilter; use lance_index::vector::VectorIndexCacheEntry; use lance_index::vector::bq::builder::RabitQuantizer; use lance_index::vector::bq::ex_dot::{blocked_ex_code_bytes, padded_query_len}; -use lance_index::vector::bq::layered::SignBounds; +use lance_index::vector::bq::layered::{EntryColumns, SignBounds}; use lance_index::vector::bq::layered_stats; +use lance_index::vector::bq::partition_codes::PartitionCodesKey; use lance_index::vector::bq::rabit_ex_bits; use lance_index::vector::bq::storage::{RabitQueryEstimator, SEGMENT_NUM_CODES}; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex, FlatQuantizer}; @@ -66,8 +67,8 @@ use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ IndexFileKey, LayeredLazyConfig, OriginLatencyClass, PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, ResidentColumns, - ResidentColumnsSetting, VectorStore, origin_latency_setting, resident_columns_setting, - resident_lifetime_setting, resident_store_fits, sign_bounds_setting, + ResidentColumnsSetting, VectorStore, entry_columns_setting, origin_latency_setting, + resident_columns_setting, resident_lifetime_setting, resident_store_fits, sign_bounds_setting, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -1724,6 +1725,7 @@ impl IVFIndex { .with_origin_latency(origin_latency) .with_resident_columns(resident_store) .with_resident_columns_enabled(resident_columns) + .with_entry_columns(Self::entry_columns_at_open()?) .with_index_file(index_file); // Cache file metadata so reconstructions from IvfIndexState can skip @@ -1988,6 +1990,50 @@ impl IVFIndex { self } + /// The entry columns an opening index asks its storage for: an IVF_RQ + /// index reads (and validates) `LANCE_RQ_ENTRY_COLUMNS`, and its storage + /// resolves them against residency (see + /// [`IvfQuantizationStorage::entry_columns`]); any other index keeps + /// full entries. + fn entry_columns_at_open() -> Result { + if Q::quantization_type() != QuantizationType::Rabit { + return Ok(EntryColumns::All); + } + entry_columns_setting() + } + + /// What the index's cache entries hold, resolved when it opened; see + /// `LANCE_RQ_ENTRY_COLUMNS`. [`EntryColumns::Codes`] only for an index + /// whose small columns are resident: a layered index's sign, high and + /// low plane entries, or a native flat index's partition entries, which + /// are then [`PartitionCodes`](lance_index::vector::bq::partition_codes::PartitionCodes). + /// A native index with a graph caches whole partitions. + pub fn entry_columns(&self) -> EntryColumns { + if self.layered_rq || self.caches_partition_codes() { + self.storage.entry_columns() + } else { + EntryColumns::All + } + } + + /// Whether the index caches its partitions as code-only entries: a + /// native IVF_RQ index with a flat sub-index, which holds nothing but + /// the storage, and code-only entries (see [`Self::entry_columns`]). + fn caches_partition_codes(&self) -> bool { + Q::quantization_type() == QuantizationType::Rabit + && !self.layered_rq + && S::name() == ::name() + && self.storage.entry_columns() == EntryColumns::Codes + } + + /// Replace the entry columns the index asks for, so tests can open both + /// in one process; the storage still resolves them against residency. + #[cfg(test)] + pub(crate) fn with_entry_columns_for_test(mut self, entry_columns: EntryColumns) -> Self { + self.storage = self.storage.with_entry_columns(entry_columns); + self + } + /// Where a newly opened index's cache keeps the bounds columns. Only /// layered indexes read (and validate) `LANCE_RQ_SIGN_BOUNDS`. fn sign_bounds_at_open(layered_rq: bool) -> Result { @@ -2042,6 +2088,12 @@ impl IVFIndex { ))); } + if self.caches_partition_codes() { + return self + .load_code_only_partition(partition_id, write_cache, metrics) + .await; + } + let cache_key = IVFPartitionKey::::new(partition_id); if write_cache { @@ -2075,6 +2127,40 @@ impl IVFIndex { } } + /// A native flat partition read through its code-only entry, which is + /// cached (admitted when `write_cache`) while the partition entry built + /// from it, with copies of the resident rows, is built again on every + /// read; see [`IvfQuantizationStorage::load_partition_cached`]. + async fn load_code_only_partition( + &self, + partition_id: usize, + write_cache: bool, + metrics: &dyn MetricsCollector, + ) -> Result>> { + let io_stats = metrics.io_stats(); + let (storage, hit) = self + .storage + .load_partition_cached( + partition_id, + &self.index_cache, + write_cache, + io_stats.clone(), + ) + .await?; + if hit { + metrics.record_index_cache_hit(); + } else { + metrics.record_index_cache_miss(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); + metrics.record_part_load(); + } + // A flat sub-index holds no rows, so building it reads nothing. + let index = self.load_sub_index(partition_id, io_stats).await?; + let mut entry = PartitionEntry::new(index, storage); + entry.cache_whole_partition = false; + Ok(Arc::new(entry)) + } + /// Cut candidates across all probed partitions before fetching lower planes. pub async fn quantized_candidates( &self, @@ -2624,6 +2710,9 @@ impl IVFIndex { Q::Metadata: 'static, Q::Storage: 'static, { + if self.caches_partition_codes() { + return self.prewarm_partition_codes_window(partitions).await; + } let index_schema = Arc::new(match &self.read_projection { Some(projection) => projection.schema.as_ref().into(), None => self.reader.schema().as_ref().into(), @@ -2730,6 +2819,58 @@ impl IVFIndex { Ok(()) } + /// [`Self::prewarm_partition_window`] of an index that caches its + /// partitions as code-only entries: one read of the window's code + /// columns, cut into the entries of its partitions. The flat sub-index + /// holds nothing to warm, and the resident store loads on first use. + async fn prewarm_partition_codes_window(&self, partitions: Range) -> Result<()> { + let projection = self.storage.partition_codes_projection()?; + let schema: arrow_schema::SchemaRef = Arc::new(projection.schema.as_ref().into()); + let mut partition_id = partitions.start; + while partition_id < partitions.end { + let cached = |partition| async move { + let key = PartitionCodesKey { partition }; + self.index_cache.get_with_key(&key).await.is_some() + }; + if cached(partition_id).await { + partition_id += 1; + continue; + } + // Stop at a cached hole, as the whole-partition prewarm does. + let mut run_end = partition_id + 1; + while run_end < partitions.end && !cached(run_end).await { + run_end += 1; + } + let run = partition_id..run_end; + let batches = read_partition_window_batches( + self.storage.reader(), + Some(&projection), + &schema, + self.storage.ivf(), + run.clone(), + None, + ) + .await?; + if batches.len() != run.len() { + return Err(Error::internal(format!( + "IVF prewarm run {run:?} produced {} partitions of codes", + batches.len() + ))); + } + for (partition, batches) in run.zip(batches) { + let codes = self.storage.partition_codes_from_batches(batches)?; + self.index_cache + .get_or_insert_with_key( + PartitionCodesKey { partition }, + || async move { Ok(codes) }, + ) + .await?; + } + partition_id = run_end; + } + Ok(()) + } + pub async fn load_partition_storage( &self, partition_id: usize, @@ -3883,6 +4024,7 @@ async fn reconstruct_typed( let storage = storage .with_resident_columns(resident_columns) .with_resident_columns_enabled(resident) + .with_entry_columns(IVFIndex::::entry_columns_at_open()?) .with_index_file(index_file); let rq_search_cache = IVFIndex::::rq_search_cache_from_state(state, &storage)?; @@ -7490,7 +7632,7 @@ mod tests { #[tokio::test] async fn test_layered_plane_promotion_survives_query_reconstruction() { - use lance_index::vector::bq::layered::{PlaneKey, SignBounds}; + use lance_index::vector::bq::layered::{EntryColumns, PlaneKey, SignBounds}; let dir = TempStrDir::default(); let (mut dataset, vectors) = @@ -7532,10 +7674,13 @@ mod tests { for partition in 0..4 { for plane in [1, 2] { // Ex plane keys do not depend on the bounds placement. + // A local index keeps no resident store, so its entries + // hold every column. let key = PlaneKey { partition, plane, sign_bounds: SignBounds::default(), + entry_columns: EntryColumns::All, }; let resident = cache.get_resident_with_key(&key).await; assert_eq!( @@ -8105,7 +8250,9 @@ mod tests { use std::time::Duration; use arrow::compute::concat_batches; + use arrow_array::UInt32Array; use bytes::Bytes; + use lance_arrow::RecordBatchExt; use lance_core::cache::{ CacheCodec, CacheEntry, CacheTier, InternalCacheKey, QuickCacheBackend, }; @@ -8113,20 +8260,23 @@ mod tests { use lance_file::version::LanceFileVersion; use lance_file::writer::FileWriterOptions; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; + use lance_index::metrics::LocalMetricsCollector; use lance_index::vector::bq::layered::{ - FULL_BOUNDS_COLUMN, PlaneBatch, RQPrecision, SIGN_BOUNDS_PLANE, SignBounds, - plane_columns, + EntryColumns, FULL_BOUNDS_COLUMN, PlaneBatch, RQPrecision, SIGN_BOUNDS_PLANE, + SignBounds, plane_columns, plane_entry_columns, }; use lance_index::vector::bq::layered_stats::{self, LayeredLazyStats, RANK_BUCKETS}; + use lance_index::vector::bq::partition_codes::{PartitionCodes, PartitionCodesKey}; use lance_index::vector::bq::storage::{ - RabitPruneStatsSnapshot, rabit_prune_stats_snapshot, + RABIT_CODE_COLUMN, RabitPruneStatsSnapshot, rabit_prune_stats_snapshot, }; use lance_index::vector::storage::{ - DenseGatherMode, DenseToEager, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, - LayeredLazyConfig, LazyOriginGap, LazyPromotion, OriginLatencyClass, - ResidentColumnsKey, ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, - origin_latency_setting, resident_columns_setting, resident_lifetime_setting, - resident_store_is_live, sign_bounds_setting, + DenseGatherMode, DenseToEager, GatherPlan, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, + IndexFileKey, LayeredLazyConfig, LazyOriginGap, LazyPromotion, OriginLatencyClass, + PlaneSource, ResidentColumnsKey, ResidentColumnsSetting, ResidentLifetime, + ResidentStoreSize, entry_columns_setting, origin_latency_setting, + resident_columns_setting, resident_lifetime_setting, resident_store_is_live, + sign_bounds_setting, }; use lance_index::vector::{ApproxMode, PartitionSearchControl, VECTOR_RESULT_SCHEMA}; use lance_io::ReadBatchParams; @@ -8843,6 +8993,31 @@ mod tests { Arc, Option>, ) { + let (dataset, tiered) = open_lazy_test_dataset(uri, cache).await; + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let index = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + if cache.is_warmed() { + index.prewarm().await.unwrap(); + } + (dataset, index, tiered) + } + + impl LazyTestCache { + /// Whether the tests warm an index opened over this cache. + fn is_warmed(self) -> bool { + !matches!(self, Self::Origin | Self::ColdUngated | Self::ColdGated) + } + } + + /// The dataset at `uri` over a new session whose index cache is + /// `cache`, and the tiered backend of that cache, if it is one. + async fn open_lazy_test_dataset( + uri: &str, + cache: LazyTestCache, + ) -> (Dataset, Option>) { let registry = Arc::new(ObjectStoreRegistry::default()); let (session, tiered) = match cache { LazyTestCache::Origin => { @@ -8887,18 +9062,7 @@ mod tests { .load() .await .unwrap(); - let uuid = dataset.load_indices().await.unwrap()[0].uuid; - let index = dataset - .open_vector_index("vector", &uuid, &NoOpMetricsCollector) - .await - .unwrap(); - if !matches!( - cache, - LazyTestCache::Origin | LazyTestCache::ColdUngated | LazyTestCache::ColdGated - ) { - index.prewarm().await.unwrap(); - } - (dataset, index, tiered) + (dataset, tiered) } fn lazy_index(index: &Arc) -> &IvfRq { @@ -9602,6 +9766,7 @@ mod tests { let setting = origin_latency_setting().unwrap(); let lazy_setting = LayeredLazyConfig::from_env().unwrap(); let resident_setting = resident_columns_setting().unwrap(); + let entry_setting = entry_columns_setting().unwrap(); let vectors = batch["vector"].as_fixed_size_list(); let queries = [ lazy_test_query(vectors.value(0), 10, 8), @@ -9661,6 +9826,12 @@ mod tests { resident_setting.resolve(expected), "{context}" ); + // Code-only entries follow residency, so the hint. + assert_eq!( + ivf.entry_columns(), + entry_setting.resolve(ivf.resident_columns_enabled()), + "{context}" + ); assert_eq!( ivf.lazy_origin_gap_bytes(), lazy_setting.origin_gap.resolve(expected, block_size), @@ -9703,6 +9874,7 @@ mod tests { .downcast_ref::() .expect("IVF_PQ index"); assert_eq!(ivf_pq.origin_latency(), pq_class, "hint={hint:?} {path}"); + assert_eq!(ivf_pq.entry_columns(), EntryColumns::All); } } wait_for_promotions().await; @@ -12170,6 +12342,855 @@ mod tests { assert_eq!(results[0], results[1]); } + /// Open the index of the dataset at `uri` over a new session whose + /// index cache is `cache`, as if its origin were of `class`, with its + /// small columns resident or read from the file and asking for + /// `entry_columns`; warm it when the tests warm that cache. + async fn open_entry_columns_test_index( + uri: &str, + cache: LazyTestCache, + class: OriginLatencyClass, + resident: bool, + entry_columns: EntryColumns, + ) -> ( + Dataset, + Arc, + Option>, + ) { + let (dataset, tiered) = open_lazy_test_dataset(uri, cache).await; + let index = open_entry_columns_index(&dataset, class, resident, entry_columns).await; + if cache.is_warmed() { + index.prewarm().await.unwrap(); + } + (dataset, index, tiered) + } + + /// Open the index of `dataset` over the dataset's caches, as + /// [`open_entry_columns_test_index`] does. + async fn open_entry_columns_index( + dataset: &Dataset, + class: OriginLatencyClass, + resident: bool, + entry_columns: EntryColumns, + ) -> Arc { + let (store, index_dir) = index_files(dataset).await; + let ivf = open_rq_index(dataset, store, index_dir) + .await + .with_origin_latency_for_test(class) + .unwrap() + .with_resident_columns_for_test(resident) + .with_entry_columns_for_test(entry_columns); + Arc::new(ivf) + } + + /// Bytes the columns of `full` that `entry` does not hold add to the + /// charge of `full`: each array and its buffers. + fn bytes_beyond(full: &RecordBatch, entry: &RecordBatch) -> usize { + full.schema() + .fields() + .iter() + .zip(full.columns()) + .filter(|(field, _)| entry.column_by_name(field.name()).is_none()) + .map(|(_, column)| { + column.deep_size_of_children(&mut lance_core::deepsize::Context::new()) + }) + .sum() + } + + fn column_names(batch: &RecordBatch) -> Vec { + batch + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() + } + + /// A code-only entry holds the file columns of its plane or native + /// partition, and is charged the bytes of a read of the file but the + /// resident columns'. Attaching the resident columns to it, whole or + /// at selected rows, gives the batch a read of the file gives, with or + /// without the resident store, bit for bit, and whole planes and + /// partitions byte for byte: every partition, empty ones included, in + /// the writer's pages and in pages that partitions straddle, in either + /// file version. An entry of other rows is an error. + #[rstest] + #[case(LanceFileVersion::V2_0)] + #[case(LanceFileVersion::V2_2)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_batch_remapping(#[case] version: LanceFileVersion) { + use lance_index::scalar::BatchRowIdRemapper; + + const ROW_ID_OFFSET: u64 = 1 << 40; + #[derive(Debug)] + struct TestBatchRemapper; + + #[async_trait::async_trait] + impl BatchRowIdRemapper for TestBatchRemapper { + async fn remap_row_ids( + &self, + row_ids: &[u64], + ) -> lance_core::Result>> { + Ok(row_ids + .iter() + .map(|&id| (!id.is_multiple_of(2)).then_some(id + ROW_ID_OFFSET)) + .collect()) + } + } + + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, false).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + let reader = rewrite_rq_storage(source, version, None).await; + let file = rq_storage_over(source, &reader, false) + .with_row_id_remapping(Arc::new(TestBatchRemapper)); + let codes = rq_storage_over(source, &reader, true) + .with_entry_columns(EntryColumns::Codes) + .with_row_id_remapping(Arc::new(TestBatchRemapper)); + let original = + rq_storage_over(source, &reader, true).with_entry_columns(EntryColumns::Codes); + let cache = lance_core::cache::LanceCache::with_capacity(16 * 1024 * 1024); + let weak = lance_core::cache::WeakLanceCache::from(&cache); + let expected = file.load_partition(0, None).await.unwrap(); + let unmapped = original.load_partition(0, None).await.unwrap(); + assert!(!expected.is_empty()); + assert!(expected.len() < unmapped.len()); + for expected_hit in [false, true] { + let (actual, hit) = codes + .load_partition_cached(0, &weak, true, None) + .await + .unwrap(); + assert_eq!(hit, expected_hit); + assert_eq!( + actual.to_batches().unwrap().collect::>(), + expected.to_batches().unwrap().collect::>(), + ); + } + // The code cache and shared resident columns keep source rows. + let (reloaded, hit) = original + .load_partition_cached(0, &weak, true, None) + .await + .unwrap(); + assert!(hit); + assert_eq!( + reloaded.to_batches().unwrap().collect::>(), + unmapped.to_batches().unwrap().collect::>(), + ); + } + + #[rstest] + #[case::native_v2_0(false, LanceFileVersion::V2_0)] + #[case::native_v2_2(false, LanceFileVersion::V2_2)] + #[case::layered_v2_0(true, LanceFileVersion::V2_0)] + #[case::layered_v2_2(true, LanceFileVersion::V2_2)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_match_file_reads( + #[case] layered: bool, + #[case] version: LanceFileVersion, + ) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + let sign_bounds = source.sign_bounds(); + // A native file has the sign plane's columns too. + let mut planes = vec![0]; + if layered { + planes.extend([1, 2]); + if sign_bounds == SignBounds::Lazy { + planes.push(SIGN_BOUNDS_PLANE); + } + } + for page_rows in [None, Some(RESIDENT_TEST_PAGE_ROWS)] { + let reader = rewrite_rq_storage(source, version, page_rows).await; + let file = rq_storage_over(source, &reader, false); + let resident = rq_storage_over(source, &reader, true); + let codes = + rq_storage_over(source, &reader, true).with_entry_columns(EntryColumns::Codes); + assert_eq!(resident.entry_columns(), EntryColumns::All); + assert_eq!(codes.entry_columns(), EntryColumns::Codes); + // Code-only entries need the resident store. + let unresident = + rq_storage_over(source, &reader, false).with_entry_columns(EntryColumns::Codes); + assert_eq!(unresident.entry_columns(), EntryColumns::All); + for partition in 0..source.num_partitions() { + let context = format!( + "layered={layered} {version:?} pages={page_rows:?} partition={partition}" + ); + let size = source.partition_size(partition); + if !layered { + let expected = file.load_partition(partition, None).await.unwrap(); + let entry = codes.read_partition_codes(partition, None).await.unwrap(); + assert_eq!( + column_names(&entry.0), + [RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN], + "{context}" + ); + let actual = codes + .partition_from_codes(partition, &entry, None) + .await + .unwrap(); + assert_eq!( + actual.to_batches().unwrap().collect::>(), + expected.to_batches().unwrap().collect::>(), + "{context}" + ); + assert_eq!(actual.deep_size_of(), expected.deep_size_of(), "{context}"); + // The entry is charged what the partition built from + // it is but the copies of the resident rows. + let actual_batch = actual.to_batches().unwrap().next().unwrap(); + assert_eq!( + entry.0.deep_size_of(), + actual_batch.deep_size_of() - bytes_beyond(&actual_batch, &entry.0), + "{context}" + ); + } + let rows = resident_test_rows(size); + let every_row: Vec = (0..size as u32).collect(); + for &plane in &planes { + let context = format!("{context} plane={plane}"); + let full = file.read_plane(partition, plane, None, None).await.unwrap(); + assert_eq!( + resident + .read_plane(partition, plane, None, None) + .await + .unwrap(), + full, + "{context}" + ); + let entry = codes + .read_plane_entry(partition, plane, None) + .await + .unwrap(); + assert_eq!( + column_names(&entry), + plane_entry_columns(plane, sign_bounds, EntryColumns::Codes), + "{context}" + ); + let whole = codes + .attach_resident(partition, plane, &entry, None, None) + .await + .unwrap(); + assert_eq!(whole, full, "{context}"); + assert_eq!(whole.deep_size_of(), full.deep_size_of(), "{context}"); + // The entry is charged what the plane is but the + // copies of the resident rows. + assert_eq!( + entry.deep_size_of(), + whole.deep_size_of() - bytes_beyond(&whole, &entry), + "{context}" + ); + for selected in [&rows, &every_row] { + let expected = file + .read_plane(partition, plane, Some(selected.clone()), None) + .await + .unwrap(); + let entry_rows = + entry.take(&UInt32Array::from(selected.clone())).unwrap(); + let actual = codes + .attach_resident( + partition, + plane, + &entry_rows, + Some(selected), + None, + ) + .await + .unwrap(); + assert_eq!(actual, expected, "{context} rows={}", selected.len()); + } + if !rows.is_empty() && rows.len() < size { + let error = codes + .attach_resident(partition, plane, &entry, Some(&rows), None) + .await + .unwrap_err(); + assert!( + matches!(error, lance_core::Error::Internal { .. }), + "{context}: {error}" + ); + assert!( + error + .to_string() + .contains(&format!("partition {partition} plane {plane}")), + "{context}: {error}" + ); + } + } + } + } + } + + /// Searches return the same batches, and assemble partitions of the + /// same sizes, whether the index caches code-only or full entries, + /// with its small columns resident or not: every precision, the + /// cascade and the lazy scan under every setting of the parity tests + /// of a layered index, reading its planes from the origin file or + /// from a tiered cache that serves selected rows of persisted + /// entries, and a native index in both approximation modes and + /// under every prefilter. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_search_parity(#[case] layered: bool) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let key = batch["vector"].as_fixed_size_list().value(777); + let mut queries: Vec<(String, Query, LayeredLazyConfig, Arc)> = + Vec::new(); + if layered { + for (label, query, config) in sign_bounds_test_queries(&key) { + queries.push((label, query, config, Arc::new(NoFilter))); + } + for (_, config) in lazy_test_configs() { + let query = lazy_test_query(key.clone(), 100, LAZY_PARTITIONS); + queries.push((format!("{config:?}"), query, config, Arc::new(NoFilter))); + } + } else { + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + for (filter_name, filter) in lazy_test_filters() { + let mut query = lazy_test_query(key.clone(), 100, 8); + query.approx_mode = approx_mode; + let label = format!("approx={approx_mode:?} filter={filter_name}"); + queries.push((label, query, LayeredLazyConfig::default(), filter)); + } + } + } + let variants = [ + (EntryColumns::All, false), + (EntryColumns::All, true), + (EntryColumns::Codes, false), + (EntryColumns::Codes, true), + ]; + for cache in [LazyTestCache::Origin, LazyTestCache::Ungated] { + let mut expected: Option, Vec)>> = None; + let mut expected_sizes: Option> = None; + for (entry_columns, resident) in variants { + let context = format!( + "layered={layered} cache={cache:?} entries={entry_columns} resident={resident}" + ); + layered_stats::snapshot_and_reset(); + let (_dataset, index, _) = open_entry_columns_test_index( + dir.as_str(), + cache, + OriginLatencyClass::High, + resident, + entry_columns, + ) + .await; + let ivf = lazy_index(&index); + assert_eq!( + ivf.entry_columns(), + entry_columns.resolve(resident), + "{context}" + ); + let mut results = Vec::new(); + for (label, query, config, filter) in &queries { + ivf.set_layered_lazy_config_for_test(*config); + let result = search_global(&index, query, filter.clone()).await; + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let result = + result.unwrap_or_else(|error| panic!("{context} {label}: {error}")); + results.push(result_bits(&result)); + } + assert!(results.iter().any(|(ids, _)| !ids.is_empty()), "{context}"); + wait_for_promotions().await; + let stats = layered_stats::snapshot_and_reset(); + if entry_columns.resolve(resident) == EntryColumns::Codes { + assert!(stats.resident_attach_calls > 0, "{context} {stats:?}"); + } + // The tiered cache persisted every plane, so the lazy + // scan gathers selected rows from the persistent tier and + // none from the origin file, and those rows, attached to + // the resident columns, are the rows a read of the file + // returns. + if layered && cache == LazyTestCache::Ungated { + assert_eq!(stats.origin_row_reads, 0, "{context} {stats:?}"); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 1) + .unwrap(); + let rows = resident_test_rows(ivf.storage.partition_size(partition)); + let plan = GatherPlan { + planes: [PlaneSource::Sparse; 2], + origin_whole: [false; 2], + }; + let gathered = ivf + .storage + .gather_ex_rows( + partition, + &rows, + plan, + &LayeredLazyConfig::default(), + &ivf.index_cache, + None, + ) + .await + .unwrap(); + assert_eq!(gathered.sources, [Some(PlaneSource::Sparse); 2]); + assert_eq!(gathered.origin_row_reads, 0, "{context}"); + for (plane, batch) in [(1, &gathered.high), (2, &gathered.low)] { + let expected = ivf + .storage + .read_plane(partition, plane, Some(rows.clone()), None) + .await + .unwrap(); + // A row read of a full entry decodes a schema + // without the file's metadata; the columns match. + assert_eq!( + column_names(batch), + column_names(&expected), + "{context} plane={plane}" + ); + assert_eq!( + batch.columns(), + expected.columns(), + "{context} plane={plane}" + ); + } + } + match &expected { + None => expected = Some(results), + Some(expected) => { + for ((label, ..), (actual, expected)) in + queries.iter().zip(results.iter().zip(expected)) + { + assert_eq!(actual, expected, "{context} {label}"); + } + } + } + // A partition is assembled of the same arrays from the + // origin file whatever the entries hold. + if cache == LazyTestCache::Origin { + let mut sizes = Vec::with_capacity(LAZY_PARTITIONS); + for partition in 0..LAZY_PARTITIONS { + let entry = if layered { + ivf.load_query_partition( + partition, + RQPrecision::Full, + &NoOpMetricsCollector, + ) + .await + } else { + ivf.load_partition(partition, true, &NoOpMetricsCollector) + .await + }; + sizes.push(entry.unwrap().size_bytes()); + } + match &expected_sizes { + None => expected_sizes = Some(sizes), + Some(expected) => assert_eq!(&sizes, expected, "{context}"), + } + } + } + } + } + + /// With its entries cached, a storage with code-only entries reads + /// nothing from the file; a miss reads the code columns alone, one + /// request per column in the writer's pages: 1 for a layered sign, + /// high or low plane, 2 for a native partition. Reading the entries + /// loads no store, and the first attaches, issued at once, load it + /// once. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_cut_origin_requests(#[case] layered: bool) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + let reader = rewrite_rq_storage(source, LanceFileVersion::V2_0, None).await; + // Lazy bounds whatever `LANCE_RQ_SIGN_BOUNDS` sets: a sign plane + // entry then holds the code column alone. + let codes = Arc::new( + rq_storage_over(source, &reader, true) + .with_sign_bounds(SignBounds::Lazy) + .with_entry_columns(EntryColumns::Codes), + ); + let planes: &[u8] = if layered { &[0, 1, 2] } else { &[] }; + let warm = LanceCache::with_capacity(LAZY_LARGE_CACHE_BYTES); + let requests = |stats: &IoStats| stats.snapshot().iops; + + layered_stats::snapshot_and_reset(); + for partition in 0..LAZY_PARTITIONS { + if layered { + for &plane in planes { + codes + .load_plane_entry(partition, plane, &WeakLanceCache::from(&warm), None) + .await + .unwrap(); + } + } else { + let codes = &codes; + warm.get_or_insert_with_key(PartitionCodesKey { partition }, || async move { + codes.read_partition_codes(partition, None).await + }) + .await + .unwrap(); + } + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 0, "{stats:?}"); + + let reads: Vec<_> = (0..RESIDENT_TEST_CONCURRENT_READS) + .map(|read| { + let codes = codes.clone(); + let cache = WeakLanceCache::from(&warm); + tokio::spawn(async move { + let partition = read % LAZY_PARTITIONS; + if layered { + let plane = (read % 3) as u8; + codes + .load_plane(partition, plane, &cache, None) + .await + .map(|plane| (plane.num_rows(), true)) + } else { + codes + .load_partition_cached(partition, &cache, true, None) + .await + .map(|(storage, hit)| (storage.len(), hit)) + } + }) + }) + .collect(); + for read in reads { + let (_, hit) = read.await.unwrap().unwrap(); + assert!(hit); + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!( + stats.resident_columns_bytes, + source.resident_columns_bytes(), + "{stats:?}" + ); + assert!(stats.resident_attach_calls > 0, "{stats:?}"); + + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| source.partition_size(partition) > 1) + .unwrap(); + let cold = LanceCache::with_capacity(LAZY_LARGE_CACHE_BYTES); + for (cache, hit) in [(&warm, true), (&cold, false)] { + let cache = WeakLanceCache::from(cache); + if layered { + for &plane in planes { + let stats = IoStats::new(); + codes + .load_plane(partition, plane, &cache, Some(stats.clone())) + .await + .unwrap(); + let expected = if hit { 0 } else { 1 }; + assert_eq!(requests(&stats), expected, "plane={plane} hit={hit}"); + } + } else { + let stats = IoStats::new(); + let (_, was_hit) = codes + .load_partition_cached(partition, &cache, true, Some(stats.clone())) + .await + .unwrap(); + assert_eq!(was_hit, hit); + let expected = if hit { 0 } else { 2 }; + assert_eq!(requests(&stats), expected, "hit={hit}"); + } + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 0, "{stats:?}"); + } + + /// Code-only and full entries of one index persist under keys of + /// their own: an index of either kind finds none of the other's + /// entries in a cache the other warmed, and reads its own, every one + /// of which decodes, with the same results. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_persist_apart_from_full_entries(#[case] layered: bool) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let key = batch["vector"].as_fixed_size_list().value(0); + let query = lazy_test_query(key, 100, LAZY_PARTITIONS); + let partition_entry_type = + as CacheCodecImpl>::TYPE_ID; + for order in [ + [EntryColumns::All, EntryColumns::Codes], + [EntryColumns::Codes, EntryColumns::All], + ] { + let context = format!("layered={layered} order={order:?}"); + let (dataset, tiered) = + open_lazy_test_dataset(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let persisted = |entry_columns| match (layered, entry_columns) { + (true, _) => tiered.persisted_entries(::TYPE_ID), + (false, EntryColumns::All) => tiered.persisted_entries(partition_entry_type), + (false, EntryColumns::Codes) => { + tiered.persisted_entries(::TYPE_ID) + } + }; + let mut indexes = Vec::new(); + for entry_columns in order { + let index = open_entry_columns_index( + &dataset, + OriginLatencyClass::High, + true, + entry_columns, + ) + .await; + assert_eq!(lazy_index(&index).entry_columns(), entry_columns); + indexes.push(index); + } + indexes[0].prewarm().await.unwrap(); + let first_entries = persisted(order[0]); + assert!(first_entries > 0, "{context}"); + // The second kind finds none of the first kind's entries. + let second = lazy_index(&indexes[1]); + for partition in 0..LAZY_PARTITIONS { + let tiers = if layered { + let mut tiers = Vec::new(); + for plane in [0, 1, 2] { + tiers.push( + second + .storage + .plane_tier(partition, plane, &second.index_cache) + .await, + ); + } + tiers + } else if order[1] == EntryColumns::Codes { + vec![ + second + .index_cache + .peek_tier_with_key(&PartitionCodesKey { partition }) + .await, + ] + } else { + vec![ + second + .index_cache + .peek_tier_with_key( + &IVFPartitionKey::::new(partition), + ) + .await, + ] + }; + assert!( + tiers.iter().all(|&tier| tier == CacheTier::Absent), + "{context} partition={partition} {tiers:?}" + ); + } + indexes[1].prewarm().await.unwrap(); + if layered { + // Both kinds hold the sign, high and low planes. + assert_eq!(persisted(order[1]), 2 * first_entries, "{context}"); + } else { + assert_eq!(persisted(order[1]), LAZY_PARTITIONS, "{context}"); + assert_eq!(first_entries, LAZY_PARTITIONS, "{context}"); + } + for (bytes, codec, _) in tiered.disk.lock().unwrap().values() { + assert!( + codec.deserialize(bytes).hit().is_some(), + "{context}: a persisted {} entry does not decode", + codec.type_id() + ); + } + // Each kind reads its own entries back from the persistent + // tier, as it wrote them. + tiered.ram.clear().await; + let mut results = Vec::new(); + for (index, entry_columns) in indexes.iter().zip(order) { + let ivf = lazy_index(index); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 0) + .unwrap(); + let (held, expected): (Vec, Vec<&str>) = if layered { + let entry = ivf + .index_cache + .get_with_key(&ivf.storage.plane_key(partition, 1)) + .await + .unwrap(); + let sign_bounds = ivf.sign_bounds(); + ( + column_names(&entry.0), + plane_entry_columns(1, sign_bounds, entry_columns), + ) + } else if entry_columns == EntryColumns::Codes { + let entry = ivf + .index_cache + .get_with_key(&PartitionCodesKey { partition }) + .await + .unwrap(); + ( + column_names(&entry.0), + vec![RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN], + ) + } else { + (Vec::new(), Vec::new()) + }; + assert_eq!(held, expected, "{context} {entry_columns}"); + let result = search_global(index, &query, Arc::new(NoFilter)) + .await + .unwrap(); + results.push(result_bits(&result)); + } + assert_eq!(results[0], results[1], "{context}"); + } + } + + /// A native index with code-only entries warms them, one per + /// partition and no whole partition, and a search that follows reads + /// every partition from them: no cache miss and no partition load, + /// with the results of full entries. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_native_prewarm_inserts_partition_codes() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, false).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(3), 100, 16); + let partition_entry_type = + as CacheCodecImpl>::TYPE_ID; + let mut results = Vec::new(); + for entry_columns in [EntryColumns::Codes, EntryColumns::All] { + let (_dataset, index, tiered) = open_entry_columns_test_index( + dir.as_str(), + LazyTestCache::TieredResident, + OriginLatencyClass::High, + true, + entry_columns, + ) + .await; + let tiered = tiered.unwrap(); + let codes = tiered.persisted_entries(::TYPE_ID); + let whole = tiered.persisted_entries(partition_entry_type); + match entry_columns { + EntryColumns::Codes => assert_eq!((codes, whole), (LAZY_PARTITIONS, 0)), + EntryColumns::All => assert_eq!((codes, whole), (0, LAZY_PARTITIONS)), + } + let metrics = Arc::new(LocalMetricsCollector::default()); + let (partitions, dists) = index.find_partitions(&query).unwrap(); + let probes = partitions.len(); + let batches = index + .clone() + .search_partitions( + query.clone(), + Arc::new(partitions), + Arc::new(dists), + 0, + probes, + Arc::new(NoFilter), + None, + metrics.clone(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let result = concat_batches(&VECTOR_RESULT_SCHEMA, batches.iter()).unwrap(); + assert_eq!(metrics.index_cache_misses(), 0, "{entry_columns}"); + assert!(metrics.index_cache_hits() >= probes, "{entry_columns}"); + assert_eq!( + metrics.parts_loaded.load(Ordering::Relaxed), + 0, + "{entry_columns}" + ); + results.push(result_bits(&result)); + } + assert_eq!(results[0], results[1]); + } + + /// A native index whose file stores its sign codes unpacked keeps + /// them packed in its code-only entries: every partition built from a + /// cached entry uses the entry's codes as they are, and searches + /// return the results of full entries, which pack the codes on every + /// construction. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_code_only_entries_keep_unpacked_codes_packed() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (batch, centroids) = lazy_test_data(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(reader, dir.as_str(), None).await.unwrap(); + let mut params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(LAZY_PARTITIONS, centroids).unwrap(), + RQBuildParams::new(7), + ); + params.skip_transpose(true); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + let vectors = batch["vector"].as_fixed_size_list(); + let mut results = Vec::new(); + for entry_columns in [EntryColumns::All, EntryColumns::Codes] { + let (_dataset, index, _) = open_entry_columns_test_index( + dir.as_str(), + LazyTestCache::ColdUngated, + OriginLatencyClass::High, + true, + entry_columns, + ) + .await; + let ivf = lazy_index(&index); + assert!(!ivf.storage.metadata().packed); + let mut bits = Vec::new(); + for key in [vectors.value(0), vectors.value(777)] { + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + let mut query = lazy_test_query(key.clone(), 100, 8); + query.approx_mode = approx_mode; + let result = search_global(&index, &query, Arc::new(NoFilter)) + .await + .unwrap(); + bits.push(result_bits(&result)); + } + } + results.push(bits); + if entry_columns == EntryColumns::Codes { + for partition in 0..LAZY_PARTITIONS { + if ivf.storage.partition_size(partition) == 0 { + continue; + } + let loaded = ivf + .load_partition(partition, true, &NoOpMetricsCollector) + .await + .unwrap(); + let cached = ivf + .index_cache + .get_with_key(&PartitionCodesKey { partition }) + .await + .unwrap(); + let storage_batch = loaded.storage.to_batches().unwrap().next().unwrap(); + let values = |batch: &RecordBatch| { + batch[RABIT_CODE_COLUMN] + .as_fixed_size_list() + .values() + .to_data() + .buffers()[0] + .as_ptr() + }; + assert_eq!( + values(&storage_batch), + values(&cached.0), + "partition {partition} rewrote its cached codes" + ); + } + } + } + assert_eq!(results[0], results[1]); + } + /// Gaps that coalesced origin reads are checked at, in ascending /// order so that neighbours compare a narrower gap with a wider one: /// touching ranges only, V11's S3 block size, a narrower gap for a @@ -12830,7 +13851,9 @@ mod tests { } // The layout is what the tier checks of a high-latency index see. - lay_out_ex_planes(&low_index, &tiered, layout).await; + // Each index lays out its own entries: the high-latency one's + // hold their codes alone when its small columns are resident. + lay_out_ex_planes(&high_index, &tiered, layout).await; let ivf = lazy_index(&high_index); let cache = &ivf.index_cache; for partition in 0..LAZY_PARTITIONS { @@ -13113,6 +14136,10 @@ mod tests { const QUERIES: usize = 3; const STAGING_STEPS: usize = 8; const ROW_READ_DELAY: Duration = Duration::from_millis(5); + // Rows of the dropped query: a quarter of the rows, so that the + // heap fills only after several probes, with a threshold that + // leaves the far gathers after them rows to read. + const DROPPED_QUERY_K: usize = LAZY_ROWS / 4; let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; @@ -13212,6 +14239,14 @@ mod tests { } // A query dropped while a far gather holds a permit returns it. + // A permit is seen only while its gather reads rows, for the row + // read delay: a far gather whose rows the threshold prunes holds + // it for no time. The first probe of a query with `SMALL_K` fills + // the heap, and whether its threshold leaves any far gather rows + // depends on how far staging ran ahead of scoring, so this query + // takes a `k` that leaves the far gathers rows to read. + let mut dropped = queries[0].clone(); + dropped.k = DROPPED_QUERY_K; let permits = ivf.storage.lazy_far_permits(&one_permit).unwrap(); ivf.set_layered_lazy_config_for_test(one_permit); tiered.clear().await; @@ -13221,7 +14256,7 @@ mod tests { } }; let finished = tokio::select! { - result = search_global(&index, &queries[0], filter.clone()) => Some(result), + result = search_global(&index, &dropped, filter.clone()) => Some(result), () = permit_taken => None, }; assert!( diff --git a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs index e4a5e29b017..5a39ae89442 100644 --- a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs +++ b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs @@ -55,6 +55,13 @@ //! planes read together and as many probes in flight as the eager scan //! prepares, instead of its sign plane first and its ex planes within the //! gather window. +//! +//! Plane entries hold the plane or, on an index with code-only entries +//! (`LANCE_RQ_ENTRY_COLUMNS=codes` with the small columns resident), its code +//! columns alone. The storage attaches the resident rows of the other +//! columns to every plane it serves, sign stage and gathered rows alike, so +//! the scan scores the same batches either way; a sparse origin read returns +//! every column itself. use std::collections::BinaryHeap; use std::future::Future; @@ -1091,8 +1098,10 @@ impl IVFIndex { }))) } - /// Load a whole ex plane in the background, off the query's critical - /// path, admitting it as the eager load would. + /// Load a whole ex plane entry in the background, off the query's + /// critical path, admitting it as the eager load would. The entry holds + /// the plane's codes alone on an index with code-only entries, and + /// `promotion_bytes` counts what it holds. fn spawn_lazy_promotion(self: &Arc, ticket: LazyPromotionTicket) { let index = self.clone(); tokio::spawn(async move {