diff --git a/rust/lance-index/src/vector/bq/layered.rs b/rust/lance-index/src/vector/bq/layered.rs index b6dffd1b82c..651135fe7d7 100644 --- a/rust/lance-index/src/vector/bq/layered.rs +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -293,7 +293,8 @@ pub enum EntryColumns { 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 + /// attaches the store's rows of the others, views of the store for a + /// whole plane or partition and copies for gathered rows, so the scored /// batch is the one an `All` entry holds. Only an index whose small /// columns are resident keeps such entries. Codes, diff --git a/rust/lance-index/src/vector/bq/layered_stats.rs b/rust/lance-index/src/vector/bq/layered_stats.rs index 1a34f3ace97..6d568e93c64 100644 --- a/rust/lance-index/src/vector/bq/layered_stats.rs +++ b/rust/lance-index/src/vector/bq/layered_stats.rs @@ -318,18 +318,42 @@ layered_lazy_counters! { /// 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. + /// Batches assembled with an index's resident rows, copied or shared: + /// 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. + /// Bytes of the resident rows those batches hold, copied or shared, + /// each row at its columns' fixed widths: what a full entry would + /// have kept in the cache. resident_attach_bytes, - /// Time spent assembling them, copies included. + /// Time spent assembling them once their rows are resolved: the + /// copies (or views) of the resident rows, the other columns' clones + /// and the batch. Summed over every attach, including those of + /// partitions prepared at once, so it is CPU time rather than time + /// on a query's critical path. resident_attach_ns, + /// Of those batches, whole partitions and planes in which every + /// resident column shares the store's buffers. + resident_attach_shares, + /// Whole partitions and planes that copied a resident column: reads + /// that become cache entries (full entries, or partitions of a native + /// index with a graph), partitions streamed out of the index, and + /// columns a view cannot cover. + resident_attach_whole_copies, + /// Sparse gathers, whose rows are always copied. + resident_attach_gathers, + /// Bytes the whole copies and the gathers copied. + resident_attach_copied_bytes, + /// Storages of native partitions built from their code-only entry + /// (`LANCE_RQ_ENTRY_COLUMNS=codes`), once per read of the partition. + code_only_partition_builds, + /// Time those builds took: the attach of the resident rows and the + /// storage's construction, without a fallback load of the store. + code_only_partition_build_ns, /// Bytes of packed row columns that reads of plane-row IVF_RQ files /// unpacked into the column layout's fields. plane_rows_unpack_bytes, diff --git a/rust/lance-index/src/vector/bq/partition_codes.rs b/rust/lance-index/src/vector/bq/partition_codes.rs index 56a8ebf8959..e76238e9768 100644 --- a/rust/lance-index/src/vector/bq/partition_codes.rs +++ b/rust/lance-index/src/vector/bq/partition_codes.rs @@ -5,9 +5,10 @@ //! //! 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 +//! (`EntryColumns::Codes`): every read attaches views of the resident rows +//! and builds the storage for that read alone, 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`]): //! diff --git a/rust/lance-index/src/vector/bq/resident.rs b/rust/lance-index/src/vector/bq/resident.rs index 16ae0daa3b0..a52364d5b25 100644 --- a/rust/lance-index/src/vector/bq/resident.rs +++ b/rust/lance-index/src/vector/bq/resident.rs @@ -8,11 +8,17 @@ //! a few bytes per row: row ids and factors. With those columns read once //! 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 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 +//! 8, and a layered one 3 instead of 13. Reads attach the resident 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. +//! resident rows again: a whole partition or plane gets views of the store's +//! buffers, each of exactly its rows' bytes. Gathered rows, whole reads that +//! become cache entries, and partitions streamed out of the index get +//! copies. A view keeps its store column's whole allocation alive until it +//! drops, so views live within a read and never +//! reach a cache entry; [`resident_store_views`] counts the live ones and +//! [`shares_resident_store`] tells whether a batch holds store memory. //! //! An index loads the store when it opens, so that no read waits for it. //! The store is an entry of the index cache ([`ResidentColumnsKey`]), @@ -24,10 +30,16 @@ use std::borrow::Cow; use std::collections::HashMap; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, LazyLock, Mutex, Weak}; use std::time::Instant; -use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt64Array, new_empty_array}; +use arrow::array::ArrayData; +use arrow::buffer::Buffer; +use arrow_array::{ + Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt64Array, make_array, new_empty_array, +}; use arrow_schema::{Field, Schema, SchemaRef}; use arrow_select::take::take; use futures::TryStreamExt; @@ -49,7 +61,9 @@ use super::storage::{ RABIT_EX_CODE_COLUMN, }; use crate::vector::exact_buffers::{exact_array, exact_batch}; -use crate::vector::storage::{IndexFileKey, ResidentLifetime, WeakRegistry, shared_by_key}; +use crate::vector::storage::{ + IndexFileKey, ResidentAttach, ResidentLifetime, WeakRegistry, shared_by_key, +}; /// Columns that reads always fetch from the file: the codes, which hold most /// of its bytes, and the estimator bounds, which only some scans read. @@ -119,6 +133,24 @@ pub(crate) enum ResidentLoadTrigger { Read, } +/// The file rows a read attaches the resident rows of. +#[derive(Debug, Clone)] +pub(crate) enum ResidentRows<'a> { + /// Every file row of a whole partition or plane. + Range(Range), + /// The sorted, unique file rows of a sparse gather. + Rows(&'a UInt64Array), +} + +impl ResidentRows<'_> { + fn len(&self) -> usize { + match self { + Self::Range(range) => range.len(), + Self::Rows(rows) => rows.len(), + } + } +} + /// One index file's store, shared by every handle of the file and by its /// index cache entry, and loaded by the first index of the file to open (or, /// as a fallback, by the first read). @@ -198,6 +230,105 @@ pub fn resident_store_is_live(file: &IndexFileKey) -> bool { .is_some_and(|slot| slot.is_loaded()) } +/// Live views of resident store rows in the process; see +/// [`resident_store_views`]. +static LIVE_STORE_VIEWS: AtomicU64 = AtomicU64::new(0); + +/// Live views of resident store rows in the process: buffers that per-read +/// batches of whole partitions and planes share with a store +/// ([`ResidentAttach::Share`]), one per resident column of each. A gauge, +/// not reset by `layered_stats::snapshot_and_reset`. Each view keeps its +/// store column's whole allocation alive, even after the store's eviction, +/// so it returns to 0 once the reads that hold views have dropped them: +/// a nonzero value after every query is done means a view leaked, such as +/// into a cache entry. +pub fn resident_store_views() -> u64 { + LIVE_STORE_VIEWS.load(Ordering::Relaxed) +} + +/// The live views of one resident store's rows, which outlives the store +/// while any of them does; see [`ResidentColumns::store_views`]. +#[derive(Debug, Clone, Default)] +pub struct ResidentStoreViews(Arc); + +impl ResidentStoreViews { + /// Views of the store's rows alive now: one per resident column of each + /// batch that shares the store. + pub fn live(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } +} + +/// Rows of a resident store column that a view shares. It holds the +/// column's buffer, so the column's whole allocation lives until the last +/// view of it drops, and it counts itself in [`resident_store_views`] and in +/// its store's [`ResidentStoreViews`] while it lives. +struct StoreRows { + rows: Buffer, + store_views: ResidentStoreViews, +} + +impl StoreRows { + fn new(rows: Buffer, store_views: &ResidentStoreViews) -> Self { + LIVE_STORE_VIEWS.fetch_add(1, Ordering::Relaxed); + store_views.0.fetch_add(1, Ordering::Relaxed); + Self { + rows, + store_views: store_views.clone(), + } + } +} + +impl AsRef<[u8]> for StoreRows { + fn as_ref(&self) -> &[u8] { + self.rows.as_slice() + } +} + +impl Drop for StoreRows { + fn drop(&mut self) { + LIVE_STORE_VIEWS.fetch_sub(1, Ordering::Relaxed); + self.store_views.0.fetch_sub(1, Ordering::Relaxed); + } +} + +/// Whether any buffer of `batch` lies in the memory of a loaded resident +/// store alive in the process: a view of the store's rows, or a slice of a +/// store column. A cache entry must never hold such a buffer, since it +/// would keep the store column's whole allocation alive while charged its +/// rows alone; the loaders of cache entries check it in debug builds. It +/// knows the stores of the process's registry, which every index open +/// binds, and not those of storages built with a store of their own. +pub fn shares_resident_store(batch: &RecordBatch) -> bool { + let slots = live_slots(); + let stores: Vec<&ResidentColumnStore> = + slots.iter().filter_map(|slot| slot.store.get()).collect(); + !stores.is_empty() + && batch + .columns() + .iter() + .any(|column| array_in_stores(&column.to_data(), &stores)) +} + +/// Whether a buffer of `data` (or of its children) lies in the memory of +/// one of `stores`. +fn array_in_stores(data: &ArrayData, stores: &[&ResidentColumnStore]) -> bool { + data.buffers() + .iter() + .chain(data.nulls().map(|nulls| nulls.inner().inner())) + .any(|buffer| stores.iter().any(|store| store.holds(buffer))) + || data + .child_data() + .iter() + .any(|child| array_in_stores(child, stores)) +} + +/// The allocation `buffer` lies in, as a range of addresses. +fn allocation(buffer: &Buffer) -> Range { + let start = buffer.data_ptr().as_ptr() as usize; + start..start + buffer.capacity() +} + /// Cache key of the resident store of an IVF_RQ index file. An index keeps /// it in its namespace of the index cache without a fragment reuse segment, /// so that a new fragment reuse index, which changes the namespace of the @@ -403,6 +534,13 @@ impl ResidentColumns { self.slot.store.get().map(|store| store.bytes) } + /// The live views of the store's rows (see [`resident_store_views`]), + /// `None` until it has loaded. The count outlives the store: views keep + /// the store's memory, not the store. + pub fn store_views(&self) -> Option { + self.slot.store.get().map(|store| store.views.clone()) + } + /// The store of `reader`'s file, loaded on first use: see /// [`ResidentSlot::get_or_load`]. Cache the loaded allocation without /// making subsequent reads depend on whether it stays in the cache. @@ -424,6 +562,8 @@ pub(crate) struct ResidentColumnStore { num_rows: u64, /// See [`ResidentColumns::loaded_bytes`]. bytes: u64, + /// See [`ResidentColumns::store_views`]. + views: ResidentStoreViews, } /// Names the columns instead of printing hundreds of megabytes of values. @@ -525,6 +665,7 @@ impl ResidentColumnStore { columns, num_rows, bytes, + views: ResidentStoreViews::default(), }) } @@ -538,35 +679,76 @@ 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. + /// `schema`'s columns at the file rows `rows`: this store's rows of the + /// columns it keeps 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, bit for bit and byte for byte. `mode` says how a whole + /// partition or plane ([`ResidentRows::Range`]) takes the store's rows: + /// views of its buffers ([`ResidentColumn::share_rows`]) or copies; + /// sparse rows ([`ResidentRows::Rows`]) are always copied + /// ([`ResidentColumn::copy_rows`]). A batch with views must never be + /// cached. The mode has no default, so every caller states it. pub(crate) fn attach( &self, schema: SchemaRef, file_batch: Option<&RecordBatch>, - rows: &UInt64Array, + rows: ResidentRows<'_>, + mode: ResidentAttach, ) -> Result { - let started = Instant::now(); + let num_rows = rows.len(); if let Some(batch) = file_batch - && batch.num_rows() != rows.len() + && batch.num_rows() != num_rows { return Err(Error::internal(format!( - "resident columns attach to {} rows of a batch of {} rows", - rows.len(), + "resident columns attach to {num_rows} rows of a batch of {} rows", batch.num_rows() ))); } - let mut copied_bytes = 0; + if let ResidentRows::Range(range) = &rows + && (range.start > range.end || range.end as u64 > self.num_rows) + { + return Err(Error::invalid_input(format!( + "resident rows {range:?} are not within the {} rows of the file", + self.num_rows + ))); + } + // Resolve the rows before the timer: a copy of a whole range copies + // the rows of a `UInt64Array` built here, as callers built it before + // reads could share the store, so `resident_attach_ns` times the + // same work in `copy` mode. + let whole = matches!(rows, ResidentRows::Range(_)); + // What each resident column takes: a range to view, or rows to copy. + let range_rows; + let taken = match (rows, mode) { + (ResidentRows::Range(range), ResidentAttach::Share) => ResidentRows::Range(range), + (ResidentRows::Range(range), ResidentAttach::Copy) => { + range_rows = UInt64Array::from_iter_values(range.start as u64..range.end as u64); + ResidentRows::Rows(&range_rows) + } + (ResidentRows::Rows(rows), _) => ResidentRows::Rows(rows), + }; + let started = Instant::now(); + let mut attached_bytes = 0u64; + let mut copied_bytes = 0u64; + let mut every_column_shared = true; 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 + let (array, shared) = match &taken { + ResidentRows::Range(range) => { + column.share_rows(range.clone(), &self.views)? + } + ResidentRows::Rows(rows) => (column.copy_rows(rows)?, false), + }; + let bytes = column.row_bytes(num_rows); + attached_bytes += bytes; + if !shared { + copied_bytes += bytes; + every_column_shared = false; + } + array } None => file_batch .and_then(|batch| batch.column_by_name(field.name())) @@ -578,16 +760,36 @@ impl ResidentColumnStore { let batch = RecordBatch::try_new_with_options( schema, columns, - &RecordBatchOptions::new().with_row_count(Some(rows.len())), + &RecordBatchOptions::new().with_row_count(Some(num_rows)), )?; 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_rows.add(num_rows as u64); + stats.resident_attach_bytes.add(attached_bytes); + stats.resident_attach_copied_bytes.add(copied_bytes); + match (whole, every_column_shared) { + (true, true) => stats.resident_attach_shares.incr(), + (true, false) => stats.resident_attach_whole_copies.incr(), + (false, _) => stats.resident_attach_gathers.incr(), + } stats.resident_attach_ns.add_elapsed(started); Ok(batch) } + /// Whether `buffer`'s bytes lie in the allocation of one of this + /// store's columns. + fn holds(&self, buffer: &Buffer) -> bool { + let start = buffer.as_ptr() as usize; + let bytes = start..start + buffer.len(); + !bytes.is_empty() + && self.columns.values().any(|column| { + column.values.to_data().buffers().iter().any(|own| { + let own = allocation(own); + bytes.start < own.end && own.start < bytes.end + }) + }) + } + /// 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. @@ -684,11 +886,71 @@ pub(crate) struct ResidentColumn { } impl ResidentColumn { + /// Bytes of `rows` rows of the column's values, which have a fixed width + /// (see [`resident_width`]). + fn row_bytes(&self, rows: usize) -> u64 { + let width = self + .values + .data_type() + .primitive_width() + .unwrap_or_default(); + (rows * width) as u64 + } + + /// The file rows `range` as a view of the column's buffer: an array + /// whose one buffer is exactly those rows' bytes, which it shares with + /// the column rather than copies, so it is charged what a copy is + /// (an `ArrayRef::slice` would be charged the whole column). The view + /// keeps the column's whole allocation alive until it drops, and counts + /// itself in [`resident_store_views`] and in `store_views` meanwhile. + /// Returns whether the array is a view: a column a view cannot cover + /// (with a validity buffer, or more than one buffer) is copied as + /// [`Self::copy_rows`] copies, and so is a buffer Arrow refuses for the + /// column's type. An empty range is an empty array, which is no copy. + pub(crate) fn share_rows( + &self, + range: Range, + store_views: &ResidentStoreViews, + ) -> Result<(ArrayRef, bool)> { + if range.start > range.end || range.end > self.values.len() { + return Err(Error::invalid_input(format!( + "resident rows {range:?} are not within the {} rows of the file", + self.values.len() + ))); + } + if range.is_empty() { + return Ok((new_empty_array(self.values.data_type()), true)); + } + let data = self.values.to_data(); + if let (None, [values], Some(width)) = ( + data.nulls(), + data.buffers(), + data.data_type().primitive_width(), + ) { + // The store's buffers come from `exact_batch`, aligned for their + // type, so the rows' first byte keeps that alignment. + let rows = values + .slice_with_length((data.offset() + range.start) * width, range.len() * width); + let view = Buffer::from(bytes::Bytes::from_owner(StoreRows::new(rows, store_views))); + if let Ok(view) = ArrayData::builder(data.data_type().clone()) + .len(range.len()) + .add_buffer(view) + .build() + { + return Ok((make_array(view), true)); + } + } + let rows = UInt64Array::from_iter_values(range.start as u64..range.end as u64); + Ok((self.copy_rows(&rows)?, false)) + } + /// A copy of the ascending file rows `rows` in a buffer of exactly their /// values, as a whole read of the file returns them (see /// `read_projected`), so that a cache entry holding it is charged the - /// same bytes. A slice would instead keep, and be charged, the whole - /// column. + /// same bytes, and so that it never keeps the store's memory alive: what + /// a sparse gather, a whole read that becomes a cache entry and every + /// attach under [`ResidentAttach::Copy`] take. A slice would instead + /// keep, and be charged, the whole column. pub(crate) fn copy_rows(&self, rows: &UInt64Array) -> Result { let offsets = rows.values(); // The page split below relies on the order. @@ -810,6 +1072,7 @@ mod tests { columns: HashMap::from([("factor".to_string(), column)]), num_rows: rows as u64, bytes: 4 * rows as u64, + views: ResidentStoreViews::default(), } } @@ -951,6 +1214,7 @@ mod tests { columns: map, num_rows: rows as u64, bytes, + views: ResidentStoreViews::default(), } } @@ -1075,48 +1339,251 @@ 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. + /// Serializes the tests that assert exact deltas of the attach counters + /// and the view gauges, which only these tests move in this crate. + static ATTACH_TEST_LOCK: Mutex<()> = Mutex::new(()); + + /// A store of [`column`], 100 rows of `factor`, live in the registry + /// under a key of its own, as an index open binds it. + fn live_column_store(test: &str, column: ResidentColumn) -> Arc { + slot_with( + &test_file(test), + ResidentColumnStore { + columns: HashMap::from([("factor".to_string(), column)]), + num_rows: 100, + bytes: 400, + views: ResidentStoreViews::default(), + }, + ) + } + + /// The attach counters, in a fixed order: calls, rows, bytes, shares, + /// whole copies, gathers and copied bytes. + fn attach_counters() -> [u64; 7] { + let counters = layered_stats::counters(); + [ + counters.resident_attach_calls.get(), + counters.resident_attach_rows.get(), + counters.resident_attach_bytes.get(), + counters.resident_attach_shares.get(), + counters.resident_attach_whole_copies.get(), + counters.resident_attach_gathers.get(), + counters.resident_attach_copied_bytes.get(), + ] + } + + fn counter_deltas(before: [u64; 7]) -> [u64; 7] { + let after = attach_counters(); + std::array::from_fn(|counter| after[counter] - before[counter]) + } + + /// A view of a range of rows shares the store column's buffer, at the + /// range's first row, and is charged exactly its rows' bytes; it lives, + /// readable, after the column drops, and the gauges count it until it + /// drops. An empty range is an empty array, a range past the file an + /// error. + #[test] + fn share_rows_views_the_store_without_copying() { + let _serial = ATTACH_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let column = column(); + let store_values = column + .values + .as_primitive::() + .values() + .as_ptr(); + let store_views = ResidentStoreViews::default(); + let baseline = resident_store_views(); + let mut held = Vec::new(); + for range in [10..30, 30..60, 0..100, 99..100] { + let context = format!("{range:?}"); + let (view, shared) = column.share_rows(range.clone(), &store_views).unwrap(); + assert!(shared, "{context}"); + let rows: Vec = (range.start as u64..range.end as u64).collect(); + assert_eq!(values(&view), as_values(&rows), "{context}"); + let view_values = view.as_primitive::().values().as_ptr(); + assert_eq!( + view_values, + store_values.wrapping_add(range.start), + "{context}" + ); + assert!(view_values.is_aligned(), "{context}"); + assert!(view.nulls().is_none(), "{context}"); + assert_eq!(view.get_buffer_memory_size(), range.len() * 4, "{context}"); + assert_eq!( + view.as_ref().deep_size_of_children(&mut Context::new()), + range.len() * 4, + "{context}" + ); + held.push(view); + assert_eq!(store_views.live(), held.len() as u64, "{context}"); + assert_eq!( + resident_store_views(), + baseline + held.len() as u64, + "{context}" + ); + } + let (empty, _) = column.share_rows(100..100, &store_views).unwrap(); + assert_eq!(empty.len(), 0); + let error = column.share_rows(90..101, &store_views).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert_eq!(store_views.live(), held.len() as u64); + + drop(column); + let every_row: Vec = (0..100).collect(); + assert_eq!(values(&held[2]), as_values(&every_row)); + drop(held); + assert_eq!(store_views.live(), 0); + assert_eq!(resident_store_views(), baseline); + } + + /// A column with a validity buffer cannot be viewed as one buffer, so + /// `share_rows` copies it, and the attach counts a whole copy. Arrow + /// keeps a validity buffer only where a value is null, so the column + /// has one null, past the rows read. #[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, + fn share_rows_copies_a_column_with_nulls() { + let _serial = ATTACH_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let factors: Vec> = (0..100) + .map(|row| (row != 50).then_some(row as f32)) + .collect(); + let column = ResidentColumn { + values: Arc::new(Float32Array::from(factors)), + page_ends: vec![40, 40, 100], }; + assert!(column.values.to_data().nulls().is_some()); + let store_views = ResidentStoreViews::default(); + let (copied, shared) = column.share_rows(10..30, &store_views).unwrap(); + assert!(!shared); + let rows: Vec = (10..30).collect(); + assert_eq!(values(&copied), as_values(&rows)); + assert_eq!(store_views.live(), 0); + + let slot = live_column_store("share-rows-nulls", column); + let store = slot.store.get().unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "factor", + DataType::Float32, + true, + )])); + let before = attach_counters(); + let batch = store + .attach( + schema, + None, + ResidentRows::Range(10..30), + ResidentAttach::Share, + ) + .unwrap(); + // One call of 20 rows of 4 bytes, a whole copy of them. + assert_eq!(counter_deltas(before), [1, 20, 80, 0, 1, 0, 80]); + assert!(!shares_resident_store(&batch)); + assert_eq!(values(batch.column(0)), as_values(&rows)); + } + + /// Attaching the resident columns to a batch of the other columns at + /// the same rows lays the columns out in the schema's order, takes the + /// resident rows as the mode says and shares the others' buffers. Under + /// `share`, a whole range views the store, charged its rows alone; a + /// copy, under `copy` or of sparse rows, is a buffer of exactly its + /// rows. Either way the batch is the same, and the store's cache entry + /// is charged the same. A batch of other rows, a column the store does + /// not keep missing from it, or a range past the file is an error. + #[rstest::rstest] + fn attach_shares_or_copies_resident_rows( + #[values(ResidentAttach::Share, ResidentAttach::Copy)] mode: ResidentAttach, + #[values(false, true)] sparse: bool, + ) { + let _serial = ATTACH_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let slot = live_column_store(&format!("attach-{mode}-{sparse}"), column()); + let store = slot.store.get().unwrap(); + let entry_bytes = ResidentColumnsEntry::new(slot.clone()).deep_size_of(); 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 rows: Vec = if sparse { + vec![5, 39, 40, 99] + } else { + (5..45).collect() + }; + let file_rows = UInt64Array::from(rows.clone()); + let resident_rows = || { + if sparse { + ResidentRows::Rows(&file_rows) + } else { + ResidentRows::Range(5..45) + } + }; + let codes: ArrayRef = Arc::new(arrow_array::UInt8Array::from_iter_values( + (0..rows.len()).map(|row| row as u8), + )); 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 shared = mode == ResidentAttach::Share && !sparse; + + let before = attach_counters(); let batch = store - .attach( - schema.clone(), - Some(&file_batch), - &UInt64Array::from(rows.to_vec()), - ) + .attach(schema.clone(), Some(&file_batch), resident_rows(), mode) .unwrap(); - assert!(calls() > before); + let bytes = 4 * rows.len() as u64; + let copied = if shared { 0 } else { bytes }; + let (shares, whole_copies, gathers) = match (sparse, shared) { + (true, _) => (0, 0, 1), + (false, true) => (1, 0, 0), + (false, false) => (0, 1, 0), + }; + assert_eq!( + counter_deltas(before), + [ + 1, + rows.len() as u64, + bytes, + shares, + whole_copies, + gathers, + copied + ] + ); 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_eq!(shares_resident_store(&batch), shared); + assert_eq!(store.views.live(), u64::from(shared)); + assert_eq!(batch.column(1).get_buffer_memory_size(), rows.len() * 4); + assert_eq!( + ResidentColumnsEntry::new(slot.clone()).deep_size_of(), + entry_bytes + ); + + // The other mode attaches the same batch, of the same bytes. + let other = match mode { + ResidentAttach::Share => ResidentAttach::Copy, + ResidentAttach::Copy => ResidentAttach::Share, + }; + let twin = store + .attach(schema.clone(), Some(&file_batch), resident_rows(), other) + .unwrap(); + assert_eq!(twin, batch); + assert_eq!( + twin.column(1).get_buffer_memory_size(), + batch.column(1).get_buffer_memory_size() + ); + drop((batch, twin)); + assert_eq!(store.views.live(), 0); + + let short = file_batch.slice(0, 2); + for (batch, rows) in [(Some(&short), resident_rows()), (None, resident_rows())] { + let error = store.attach(schema.clone(), batch, rows, mode).unwrap_err(); assert!(matches!(error, Error::Internal { .. }), "{error}"); } + let error = store + .attach(schema, None, ResidentRows::Range(90..101), mode) + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); } /// Rows across pages are copied into one buffer of exactly their diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 16ba0024e3c..29f026ff296 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -11,7 +11,9 @@ use crate::vector::bq::layered::{ use crate::vector::bq::layered_stats; use crate::vector::bq::partition_codes::{PartitionCodes, PartitionCodesKey}; use crate::vector::bq::plane_rows::{PACKED_COLUMNS, PlaneRowsSpec}; -use crate::vector::bq::resident::{ResidentColumnStore, ResidentLoadTrigger, is_resident}; +use crate::vector::bq::resident::{ + ResidentColumnStore, ResidentLoadTrigger, ResidentRows, is_resident, +}; use crate::vector::bq::storage::{RQRowLayout, normalize_entry_codes}; use crate::vector::bq::transform::ERROR_FACTORS_COLUMN; use crate::vector::quantizer::{QuantizationMetadata, QuantizationType, QuantizerStorage}; @@ -69,8 +71,9 @@ use super::quantizer::{Quantizer, QuantizerMetadata}; use super::{ApproxMode, DISTANCE_TYPE_KEY}; pub use crate::vector::bq::resident::{ - ResidentColumns, ResidentColumnsEntry, ResidentColumnsKey, resident_columns_bytes, - resident_store_charge, resident_store_count, resident_store_is_live, + ResidentColumns, ResidentColumnsEntry, ResidentColumnsKey, ResidentStoreViews, + resident_columns_bytes, resident_store_charge, resident_store_count, resident_store_is_live, + resident_store_views, shares_resident_store, }; /// Coalesce source-index reads independently of the scoring vector batch size. @@ -1077,11 +1080,12 @@ fn resident_lifetime_from(value: Option<&str>) -> Result { /// 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. +/// reads fetch from the file, and every read attaches the resident rows: +/// views of the store for whole partitions and planes, copies for gathered +/// rows (see [`ResidentAttach`]), 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. @@ -1112,6 +1116,41 @@ fn entry_columns_from(value: Option<&str>) -> Result { } } +/// How a read of an IVF_RQ index with resident small columns attaches the +/// store's rows to a batch of a whole partition or plane that is never +/// cached: share (default) or copy. Gathered rows, whole reads that +/// become cache entries, and partitions streamed out of the index to its +/// caller are always copied. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum ResidentAttach { + /// A whole partition or plane gets views of the store's buffers, each + /// of exactly its rows' bytes, so the read copies nothing and the batch + /// is charged what a copy is. A view keeps its store column's whole + /// allocation alive until it drops, so no view may reach a cache entry + /// or outlive the read that made it. + #[default] + Share, + /// Every attach copies the store's rows into buffers of their own, as + /// before the store was shared. + Copy, +} + +impl ResidentAttach { + /// The mode's name in diagnostics: `share` or `copy`. + pub fn as_str(self) -> &'static str { + match self { + Self::Share => "share", + Self::Copy => "copy", + } + } +} + +impl std::fmt::Display for ResidentAttach { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + /// Parse a `0`/`1` (or `false`/`true`) environment flag. fn parse_flag(name: &str, value: &str) -> Result { match value.trim() { @@ -1160,14 +1199,23 @@ pub struct LayeredExLayout { pub row_bytes: [usize; 2], } -/// 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), - } +/// `batch` as the cache entry of a plane. An entry holds buffers of its own: +/// one viewing the resident store would keep the store column's whole +/// allocation alive from the cache, charged its rows alone, so every loader +/// of a plane entry checks it in debug builds. +fn plane_cache_entry(batch: RecordBatch) -> PlaneBatch { + debug_assert!( + !shares_resident_store(&batch), + "a plane cache entry holds memory of a resident store" + ); + PlaneBatch(batch) +} + +/// The file rows of the sorted offsets `rows` of a partition whose rows +/// start at file row `start`. +fn partition_file_rows(start: usize, rows: &[u32]) -> UInt64Array { + let start = start as u64; + UInt64Array::from_iter_values(rows.iter().map(|&row| start + u64::from(row))) } /// How stage-2 survivors index a [`GatheredEx`]'s batches. @@ -1242,6 +1290,8 @@ pub struct IvfQuantizationStorage { /// What the index asks its cache entries to hold; see /// [`Self::entry_columns`]. entry_columns: EntryColumns, + /// See [`Self::resident_attach`]. + resident_attach: ResidentAttach, /// 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`]. @@ -1372,6 +1422,7 @@ impl IvfQuantizationStorage { resident_columns: ResidentColumns::default(), resident_columns_enabled: false, entry_columns: EntryColumns::default(), + resident_attach: ResidentAttach::default(), packed_metadata: OnceLock::new(), attach_schemas: AttachSchemas::default(), index_file: None, @@ -1460,6 +1511,7 @@ impl IvfQuantizationStorage { resident_columns: ResidentColumns::default(), resident_columns_enabled: false, entry_columns: EntryColumns::default(), + resident_attach: ResidentAttach::default(), packed_metadata: OnceLock::new(), attach_schemas: AttachSchemas::default(), index_file: None, @@ -1614,6 +1666,24 @@ impl IvfQuantizationStorage { self.entry_columns.resolve(self.resident_columns_enabled) } + /// Attach the resident rows to the batches of whole partitions and + /// planes that reads build and never cache as `resident_attach`; see + /// [`Self::resident_attach`]. This lets tests compare both ownership modes. + pub fn with_resident_attach(mut self, resident_attach: ResidentAttach) -> Self { + self.resident_attach = resident_attach; + self + } + + /// How reads attach the resident rows to a batch of a whole partition + /// or plane that is never cached: views of the store + /// ([`ResidentAttach::Share`]) unless the index set + /// [`ResidentAttach::Copy`] with [`Self::with_resident_attach`]. + /// Gathered rows, whole reads that become cache entries, and partitions + /// streamed out of the index are always copied. + pub fn resident_attach(&self) -> ResidentAttach { + self.resident_attach + } + /// The cache key of plane `plane` of `partition` under this storage's /// bounds placement and entry columns. pub fn plane_key(&self, partition: usize, plane: u8) -> PlaneKey { @@ -1906,16 +1976,22 @@ impl IvfQuantizationStorage { } /// 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 index loaded the store when it opened; a - /// storage built outside an index open loads it on its first build, as a - /// fallback, and adds the load's I/O to `load_stats`. + /// code-only entry: the columns the resident store keeps are the store's + /// rows, attached as `attach` says, 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. A read's own storage takes + /// [`Self::resident_attach`] (views of the store by default); one handed + /// out of the index, which may outlive the index and its lease on the + /// store, takes [`ResidentAttach::Copy`]. The storage is built for one + /// read and never cached: views of the store must not reach the cache. + /// The index loaded the store when it opened; a storage built outside an + /// index open loads it on its first build, as a fallback, and adds the + /// load's I/O to `load_stats`. Counted in `code_only_partition_builds`. pub async fn partition_from_codes( &self, part_id: usize, codes: &PartitionCodes, + attach: ResidentAttach, load_stats: Option<&IoStats>, ) -> Result { let range = self.ivf.row_range(part_id); @@ -1934,47 +2010,64 @@ impl IvfQuantizationStorage { .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) + // As a read of no rows, which loads no store. + let store = if range.is_empty() { + None } 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 store = self.resident_store(&schema, load_stats).await?; + Some(store.ok_or_else(|| { + Error::internal(format!( + "the code-only entry of partition {part_id} needs the resident store" + )) + })?) + }; + // The build alone, without a fallback load of the store. + let started = Instant::now(); + let batch = match store { + None => RecordBatch::new_empty(schema), + Some(store) => store.attach(schema, Some(entry), ResidentRows::Range(range), attach)?, }; let (batch, remapper) = self.remap_partition_batch(batch).await?; - Q::Storage::try_from_batch_with_remapper( + let storage = Q::Storage::try_from_batch_with_remapper( batch, self.entry_metadata()?, self.distance_type, remapper, - ) + )?; + let stats = layered_stats::counters(); + stats.code_only_partition_builds.incr(); + stats.code_only_partition_build_ns.add_elapsed(started); + Ok(storage) } /// 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`]), - /// which the index loaded when it opened, so a hit reads nothing from the - /// file. + /// Every read attaches the resident columns as `attach` says + /// ([`Self::partition_from_codes`]); the index loaded them when it + /// opened, so a hit reads nothing from the file. The entry holds the + /// codes alone; the storage, which views the store's rows under + /// [`ResidentAttach::Share`], is built for this read and never cached. pub async fn load_partition_cached( &self, part_id: usize, cache: &WeakLanceCache, write_cache: bool, + attach: ResidentAttach, 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 + let codes = self + .read_partition_codes(part_id, io_stats.as_ref()) + .await?; + debug_assert!( + !shares_resident_store(&codes.0), + "a code-only partition entry holds memory of a resident store" + ); + Ok(codes) }) .await? } else if let Some(codes) = cache.get_with_key(&key).await { @@ -1986,7 +2079,7 @@ impl IvfQuantizationStorage { (Arc::new(codes), false) }; let storage = self - .partition_from_codes(part_id, &codes, io_stats.as_ref()) + .partition_from_codes(part_id, &codes, attach, io_stats.as_ref()) .await?; Ok((storage, hit)) } @@ -2010,7 +2103,7 @@ impl IvfQuantizationStorage { for &plane in &planes { cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { - Ok(PlaneBatch( + Ok(plane_cache_entry( self.read_plane_entry(part_id, plane, None).await?, )) }) @@ -2031,7 +2124,7 @@ impl IvfQuantizationStorage { } cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { - Ok(PlaneBatch( + Ok(plane_cache_entry( self.read_plane_entry(part_id, plane, None).await?, )) }) @@ -2150,7 +2243,7 @@ impl IvfQuantizationStorage { if !cache.plane_admission_gated() { return cache .get_or_insert_with_key(key, || async { - Ok(PlaneBatch( + Ok(plane_cache_entry( self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) @@ -2165,7 +2258,7 @@ impl IvfQuantizationStorage { let batch = if sign_resident { cache .get_or_insert_with_key_hit(key, || async { - Ok(PlaneBatch( + Ok(plane_cache_entry( self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) @@ -2279,12 +2372,16 @@ impl IvfQuantizationStorage { /// 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 index loaded the store - /// when it opened; a storage built outside an index open loads it on its - /// first attach, as a fallback, and adds the load's I/O to `load_stats`. + /// entry ([`EntryColumns::Codes`]) gets 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`]): a whole plane as [`Self::resident_attach`] + /// says (views of the store by default), selected rows as copies. The + /// batch is for this read alone and must never be cached, since a view + /// keeps its store column's whole allocation alive. The index loaded the + /// store when it opened; a storage built outside an index open loads it + /// on its first attach, as a fallback, and adds the load's I/O to + /// `load_stats`. pub async fn attach_resident( &self, part_id: usize, @@ -2314,8 +2411,13 @@ impl IvfQuantizationStorage { // entry. return Ok(entry.clone()); }; + let file_rows = rows.map(|rows| partition_file_rows(range.start, rows)); + let rows = match &file_rows { + Some(file_rows) => ResidentRows::Rows(file_rows), + None => ResidentRows::Range(range), + }; store - .attach(schema, Some(entry), &partition_file_rows(&range, rows)) + .attach(schema, Some(entry), rows, self.resident_attach) .map_err(|error| { Error::internal(format!( "attaching the resident columns of partition {part_id} plane {plane}: {error}" @@ -2675,7 +2777,7 @@ impl IvfQuantizationStorage { if promote { cache .get_or_insert_with_key(self.plane_key(part_id, plane), || async { - Ok(PlaneBatch( + Ok(plane_cache_entry( self.read_plane_entry(part_id, plane, io_stats.clone()) .await?, )) @@ -2928,7 +3030,10 @@ impl IvfQuantizationStorage { /// either every row or the sorted partition offsets `rows`, as `params` /// selects them. Columns `store` keeps are copies of its rows, and only /// the others are read from the file, through `reader`, so the batch is - /// the one a read of every column from the file returns. + /// the one a read of every column from the file returns. Whole reads + /// become cache entries (plane entries holding every column, and + /// partition entries), so they copy too: a view of the store would keep + /// its whole column alive from the cache. async fn read_with_resident_columns( &self, store: &ResidentColumnStore, @@ -2977,11 +3082,14 @@ impl IvfQuantizationStorage { concat_batches(&file_schema, batches.iter())? }) }; - store.attach( - schema, - file_batch.as_ref(), - &partition_file_rows(&range, rows), - ) + let file_rows = rows.map(|rows| partition_file_rows(range.start, rows)); + let rows = match &file_rows { + Some(file_rows) => ResidentRows::Rows(file_rows), + None => ResidentRows::Range(range), + }; + // Always copied: a whole read becomes a plane entry or a partition + // entry (`IVFPartitionKey`), which must not hold views of the store. + store.attach(schema, file_batch.as_ref(), rows, ResidentAttach::Copy) } /// `columns` of the file at every row of the partition rows `range`, diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 1c532c3572f..72671124c61 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -66,10 +66,10 @@ use lance_index::vector::quantizer::{ use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ IndexFileKey, OriginLatencyClass, PlaneAccessTracker, QueryResidual, QueryScratch, - QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, ResidentColumns, + QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, ResidentAttach, ResidentColumns, ResidentColumnsSetting, ResidentStoreSize, VectorStore, entry_columns_setting, origin_latency_setting, resident_columns_setting, resident_lifetime_setting, - resident_store_fits, sign_bounds_setting, + resident_store_fits, shares_resident_store, sign_bounds_setting, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -940,6 +940,12 @@ pub struct PartitionEntry { pub storage: Q::Storage, partition_rows: OnceLock>, partition_rows_accounted: AtomicBool, + /// Whether [`IVFIndex::cache_partition_rows`] may cache the entry whole. + /// An entry built for one read, from a code-only entry or of a layered + /// partition's planes, sets it false: its resident columns may be views + /// of the resident store, charged their rows' bytes (which the store's + /// cache entry already charges) yet keeping the whole store column + /// alive, so such an entry must never be cached. cache_whole_partition: bool, /// Memoized size of the immutable parts (this struct, the sub-index and the /// quantized storage): every query that prepares this partition needs its @@ -980,6 +986,15 @@ impl PartitionEntry { immutable + self.coverage_bytes.get().copied().unwrap_or_default() } + /// Whether the entry's storage holds memory of a resident store, which + /// no cached entry may (see [`shares_resident_store`]); the loaders of + /// partition entries check it in debug builds. + fn shares_resident_store(&self) -> bool { + self.storage + .to_batches() + .is_ok_and(|mut batches| batches.any(|batch| shares_resident_store(&batch))) + } + fn partition_rows(&self) -> Arc { let rows = self .partition_rows @@ -1206,6 +1221,10 @@ impl IVFIndex { if partition.cache_whole_partition && !partition.partition_rows_accounted.load(Ordering::Acquire) { + debug_assert!( + !partition.shares_resident_store(), + "partition {partition_id} would be cached with memory of a resident store" + ); let cache_key = IVFPartitionKey::::new(partition_id); if index_cache .insert_with_key(&cache_key, partition.clone()) @@ -2005,6 +2024,21 @@ impl IVFIndex { entry_columns_setting() } + /// How reads attach the resident rows to the whole partitions and + /// planes they never cache; see + /// [`IvfQuantizationStorage::resident_attach`]. + pub fn resident_attach(&self) -> ResidentAttach { + self.storage.resident_attach() + } + + /// Replace how reads attach the resident rows, so tests can run both + /// modes in one process. + #[cfg(test)] + pub(crate) fn with_resident_attach_for_test(mut self, resident_attach: ResidentAttach) -> Self { + self.storage = self.storage.with_resident_attach(resident_attach); + self + } + /// 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 @@ -2079,6 +2113,20 @@ impl IVFIndex { partition_id: usize, write_cache: bool, metrics: &dyn MetricsCollector, + ) -> Result>> { + self.load_partition_attached(partition_id, write_cache, self.resident_attach(), metrics) + .await + } + + /// [`Self::load_partition`], with the resident rows of a code-only + /// partition attached as `attach` says; see + /// [`IvfQuantizationStorage::partition_from_codes`]. + async fn load_partition_attached( + &self, + partition_id: usize, + write_cache: bool, + attach: ResidentAttach, + metrics: &dyn MetricsCollector, ) -> Result>> { let _stage_timer = IndexTimer::new(metrics, IndexTiming::PartitionLoad); if partition_id >= self.ivf.num_partitions() { @@ -2091,7 +2139,7 @@ impl IVFIndex { if self.caches_partition_codes() { return self - .load_code_only_partition(partition_id, write_cache, metrics) + .load_code_only_partition(partition_id, write_cache, attach, metrics) .await; } @@ -2103,8 +2151,14 @@ impl IVFIndex { .get_or_insert_with_key_hit(cache_key, || async { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); - self.load_partition_entry(partition_id, metrics.io_stats()) - .await + let entry = self + .load_partition_entry(partition_id, metrics.io_stats()) + .await?; + debug_assert!( + !entry.shares_resident_store(), + "partition {partition_id} would be cached with memory of a resident store" + ); + Ok(entry) }) .await; match &result { @@ -2130,12 +2184,14 @@ 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`]. + /// from it, with the resident rows attached as `attach` says, is built + /// again on every read and never cached; see + /// [`IvfQuantizationStorage::load_partition_cached`]. async fn load_code_only_partition( &self, partition_id: usize, write_cache: bool, + attach: ResidentAttach, metrics: &dyn MetricsCollector, ) -> Result>> { let io_stats = metrics.io_stats(); @@ -2145,6 +2201,7 @@ impl IVFIndex { partition_id, &self.index_cache, write_cache, + attach, io_stats.clone(), ) .await?; @@ -2786,8 +2843,14 @@ impl IVFIndex { let key = IVFPartitionKey::::new(follower_id); self.index_cache .get_or_insert_with_key(key, || async move { - self.materialize_prewarm_partition(follower_id, batches) - .await + let entry = self + .materialize_prewarm_partition(follower_id, batches) + .await?; + debug_assert!( + !entry.shares_resident_store(), + "prewarmed partition {follower_id} holds memory of a resident store" + ); + Ok(entry) }) .await .map(|_| ()) @@ -2804,8 +2867,14 @@ impl IVFIndex { if let Some(error) = first_error { return Err(error); } - self.materialize_prewarm_partition(partition_id, leader_batches) - .await + let entry = self + .materialize_prewarm_partition(partition_id, leader_batches) + .await?; + debug_assert!( + !entry.shares_resident_store(), + "prewarmed partition {partition_id} holds memory of a resident store" + ); + Ok(entry) }) .await?; partition_id = if was_cached { @@ -2861,6 +2930,10 @@ impl IVFIndex { } for (partition, batches) in run.zip(batches) { let codes = self.storage.partition_codes_from_batches(batches)?; + debug_assert!( + !shares_resident_store(&codes.0), + "prewarmed codes of partition {partition} hold memory of a resident store" + ); self.index_cache .get_or_insert_with_key( PartitionCodesKey { partition }, @@ -3839,7 +3912,12 @@ impl VectorIndex for IVFInd with_vector: bool, metrics: &dyn MetricsCollector, ) -> Result { - let partition = self.load_partition(partition_id, false, metrics).await?; + // The stream leaves the index: its caller may hold the batches after + // the index, and its lease on the resident store, are gone, where a + // view of the store would keep a whole store column alive, uncharged. + let partition = self + .load_partition_attached(partition_id, false, ResidentAttach::Copy, metrics) + .await?; let store = &partition.storage; let schema = if with_vector { store.schema().clone() @@ -8275,10 +8353,12 @@ mod tests { use lance_index::vector::bq::partition_codes::{PartitionCodes, PartitionCodesKey}; use lance_index::vector::bq::storage::RABIT_CODE_COLUMN; use lance_index::vector::storage::{ - GatherPlan, IndexFileKey, OriginLatencyClass, PlaneSource, ResidentColumnsEntry, - ResidentColumnsKey, ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, - entry_columns_setting, origin_latency_setting, resident_columns_setting, - resident_lifetime_setting, resident_store_is_live, sign_bounds_setting, + GatherPlan, IndexFileKey, OriginLatencyClass, PlaneSource, ResidentAttach, + ResidentColumns, ResidentColumnsEntry, ResidentColumnsKey, ResidentColumnsSetting, + ResidentLifetime, ResidentStoreSize, ResidentStoreViews, entry_columns_setting, + origin_latency_setting, resident_columns_setting, resident_lifetime_setting, + resident_store_is_live, resident_store_views, shares_resident_store, + sign_bounds_setting, }; use lance_index::vector::{ApproxMode, VECTOR_RESULT_SCHEMA}; use lance_io::ReadBatchParams; @@ -9018,6 +9098,8 @@ mod tests { let column = open(columns.clone()).await; let row = open(rows.clone()).await; assert!(!lazy_index(&row).resident_columns_enabled()); + let full_column_entries = + lazy_index(&column).entry_columns() == EntryColumns::All; for key in &keys { for nprobes in [1, 4, LAZY_PARTITIONS] { for k in [1, 10, 100, LAZY_ROWS + 1] { @@ -9029,13 +9111,17 @@ mod tests { // A cascade scores a partition's first stage at // high precision only while its high plane is // resident, so it compares the layouts where - // both caches hold the same planes: a - // high-latency column index caches code-only - // entries, which the plane-row index cannot. + // both caches hold the same planes. The low + // twin starts from prewarm; the high twin + // inherits each layout's lazy-read cache + // history. Resident column indexes also use + // code-only keys absent from the plane-row + // cache, including the low twin under `on`. if layered && precision == RQPrecision::Full && case.is_multiple_of(4) - && (class == OriginLatencyClass::Low + && ((class == OriginLatencyClass::Low + && full_column_entries) || cache == LazyTestCache::Origin) { query.rq_cascade_factor = Some(4); @@ -9687,10 +9773,12 @@ mod tests { } /// Open the index of `dataset` over the dataset's caches, with its - /// small columns resident or read from the file. + /// small columns resident or read from the file, attached to whole + /// partitions and planes as `attach` says. async fn open_with_resident_columns( dataset: &Dataset, resident: bool, + attach: ResidentAttach, ) -> Arc { let index = dataset.load_indices().await.unwrap()[0].clone(); let ivf = IvfRq::try_new( @@ -9707,7 +9795,8 @@ mod tests { .unwrap() .with_resident_columns_for_test(resident) .await - .unwrap(); + .unwrap() + .with_resident_attach_for_test(attach); Arc::new(ivf) } @@ -9850,10 +9939,15 @@ mod tests { .await .unwrap(); let load = layered_stats::snapshot_and_reset(); - assert!(load.resident_columns_load_requests > 0, "{load:?}"); - assert_eq!( - requests(&first), - load.resident_columns_load_requests + file_columns + assert_eq!(load.resident_columns_read_loads, 1, "{load:?}"); + // The read made the load's requests besides its file columns'. + // Other tests' opens load stores of their own meanwhile, so the + // process-wide count only bounds this load's from above. + let load_requests = requests(&first).checked_sub(file_columns).unwrap(); + assert!(load_requests > 0, "{load:?}"); + assert!( + load.resident_columns_load_requests >= load_requests, + "{load:?}" ); for (storage, expected) in [(&file, columns), (&resident, file_columns)] { @@ -9893,9 +9987,9 @@ mod tests { assert!(sparse[1] < sparse[0], "plane={plane} {sparse:?}"); } } - // Only the first read loaded the store. + // Only the first read loaded the store: no read loaded one since. let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + assert_eq!(stats.resident_columns_read_loads, 0, "{stats:?}"); } /// First reads issued at once through a storage built outside an index @@ -9952,9 +10046,12 @@ mod tests { for read in reads { read.await.unwrap().unwrap(); } + // One read load, which only this storage could make; other tests' + // opens load stores of their own meanwhile, so the process-wide + // bytes only bound this load's from below. let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_bytes, bytes, "{stats:?}"); assert_eq!(stats.resident_columns_read_loads, 1, "{stats:?}"); + assert!(stats.resident_columns_bytes >= bytes, "{stats:?}"); // The arrays hold at least their values. assert!(stats.resident_columns_alloc_bytes >= bytes, "{stats:?}"); assert!(stats.resident_columns_load_requests > 0, "{stats:?}"); @@ -9968,13 +10065,14 @@ mod tests { .unwrap(); shared.load_partition(partition, None).await.unwrap(); let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + assert_eq!(stats.resident_columns_read_loads, 0, "{stats:?}"); } /// Searches return the same batches whether the small columns are - /// resident or read from the origin file: every precision, the - /// cascade and the lazy scan of a layered index, and a native index in - /// both approximation modes. + /// resident or read from the origin file, and whether reads view or + /// copy the resident rows: every precision, the cascade and the lazy + /// scan of a layered index, and a native index in both approximation + /// modes. #[rstest] #[case::native(false)] #[case::layered(true)] @@ -9987,8 +10085,9 @@ mod tests { let vectors = batch["vector"].as_fixed_size_list(); let (dataset, _, _) = open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; layered_stats::snapshot_and_reset(); - let file = open_with_resident_columns(&dataset, false).await; - let resident = open_with_resident_columns(&dataset, true).await; + let file = open_with_resident_columns(&dataset, false, ResidentAttach::Share).await; + let shared = open_with_resident_columns(&dataset, true, ResidentAttach::Share).await; + let copied = open_with_resident_columns(&dataset, true, ResidentAttach::Copy).await; for key in [vectors.value(0), vectors.value(777)] { let queries = if layered { sign_bounds_test_queries(&key) @@ -10005,19 +10104,34 @@ mod tests { }; for (label, query) in queries { let mut results = Vec::new(); - for index in [&file, &resident] { + for index in [&file, &shared, &copied] { let result = search_global(index, &query, Arc::new(NoFilter)).await; results.push(result_bits(&result.unwrap())); } let context = format!("layered={layered} {label}"); assert!(!results[0].0.is_empty(), "{context}"); - assert_eq!(results[1], results[0], "{context}"); + assert_eq!(results[1], results[0], "{context} share"); + assert_eq!(results[2], results[0], "{context} copy"); } } + // Each resident twin holds a store of its own, which other tests + // cannot load, so its bytes are checked on the twin. + let bytes = lazy_index(&shared).resident_columns_bytes(); + for (index, attach) in [ + (&shared, ResidentAttach::Share), + (&copied, ResidentAttach::Copy), + ] { + let ivf = lazy_index(index); + assert_eq!(ivf.resident_attach(), attach); + assert_eq!( + ivf.storage.resident_columns().loaded_bytes(), + Some(bytes), + "{attach}" + ); + } let stats = layered_stats::snapshot_and_reset(); - let bytes = lazy_index(&resident).resident_columns_bytes(); - assert_eq!(stats.resident_columns_bytes, bytes, "{stats:?}"); + assert!(stats.resident_columns_bytes >= 2 * bytes, "{stats:?}"); let file_store = lazy_index(&file).storage.resident_columns(); assert_eq!(file_store.loaded_bytes(), None); } @@ -11440,21 +11554,25 @@ mod tests { /// 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. + /// small columns resident or read from the file, asking for + /// `entry_columns` and attaching the resident rows as `attach` says + /// (`None` keeps what the environment sets); 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, + attach: Option, ) -> ( 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; + let index = + open_entry_columns_index(&dataset, class, resident, entry_columns, attach).await; if cache.is_warmed() { index.prewarm().await.unwrap(); } @@ -11468,9 +11586,10 @@ mod tests { class: OriginLatencyClass, resident: bool, entry_columns: EntryColumns, + attach: Option, ) -> Arc { let (store, index_dir) = index_files(dataset).await; - let ivf = open_rq_index(dataset, store, index_dir) + let mut ivf = open_rq_index(dataset, store, index_dir) .await .with_origin_latency_for_test(class) .unwrap() @@ -11478,6 +11597,9 @@ mod tests { .await .unwrap() .with_entry_columns_for_test(entry_columns); + if let Some(attach) = attach { + ivf = ivf.with_resident_attach_for_test(attach); + } Arc::new(ivf) } @@ -11511,12 +11633,17 @@ mod tests { /// 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. + /// file version, whether whole reads view or copy the resident rows. + /// Whole partitions and planes view the store under `share` alone, + /// and gathered rows never do. 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) { + async fn test_code_only_entries_batch_remapping( + #[case] version: LanceFileVersion, + #[values(ResidentAttach::Copy, ResidentAttach::Share)] attach: ResidentAttach, + ) { use lance_index::scalar::BatchRowIdRemapper; const ROW_ID_OFFSET: u64 = 1 << 40; @@ -11558,7 +11685,7 @@ mod tests { assert!(expected.len() < unmapped.len()); for expected_hit in [false, true] { let (actual, hit) = codes - .load_partition_cached(0, &weak, true, None) + .load_partition_cached(0, &weak, true, attach, None) .await .unwrap(); assert_eq!(hit, expected_hit); @@ -11569,7 +11696,7 @@ mod tests { } // The code cache and shared resident columns keep source rows. let (reloaded, hit) = original - .load_partition_cached(0, &weak, true, None) + .load_partition_cached(0, &weak, true, attach, None) .await .unwrap(); assert!(hit); @@ -11588,6 +11715,7 @@ mod tests { async fn test_code_only_entries_match_file_reads( #[case] layered: bool, #[case] version: LanceFileVersion, + #[values(ResidentAttach::Share, ResidentAttach::Copy)] attach: ResidentAttach, ) { let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); @@ -11608,8 +11736,18 @@ mod tests { 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); + // A store of the process's registry, as an index open binds, + // so that `shares_resident_store` knows it. + let store_file = IndexFileKey::new( + &format!("code-only-{layered}-{version:?}-{page_rows:?}-{attach}"), + "memory", + "resident/storage.lance", + ); + let codes = rq_storage_over(source, &reader, true) + .with_entry_columns(EntryColumns::Codes) + .with_resident_attach(attach) + .with_resident_columns(ResidentColumns::shared(&store_file)); + let shares = attach == ResidentAttach::Share; assert_eq!(resident.entry_columns(), EntryColumns::All); assert_eq!(codes.entry_columns(), EntryColumns::Codes); // Code-only entries need the resident store. @@ -11629,15 +11767,26 @@ mod tests { [RABIT_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_COLUMN], "{context}" ); + layered_stats::snapshot_and_reset(); let actual = codes - .partition_from_codes(partition, &entry, None) + .partition_from_codes(partition, &entry, attach, None) .await .unwrap(); + let built = layered_stats::snapshot_and_reset(); + assert!(built.code_only_partition_builds >= 1, "{context} {built:?}"); assert_eq!( actual.to_batches().unwrap().collect::>(), expected.to_batches().unwrap().collect::>(), "{context}" ); + assert_eq!( + actual + .to_batches() + .unwrap() + .any(|batch| shares_resident_store(&batch)), + shares && size > 0, + "{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. @@ -11676,6 +11825,13 @@ mod tests { .unwrap(); assert_eq!(whole, full, "{context}"); assert_eq!(whole.deep_size_of(), full.deep_size_of(), "{context}"); + // The bounds plane keeps no resident column. + let attaches = whole.num_columns() > entry.num_columns(); + assert_eq!( + shares_resident_store(&whole), + shares && attaches && size > 0, + "{context}" + ); // The entry is charged what the plane is but the // copies of the resident rows. assert_eq!( @@ -11701,6 +11857,11 @@ mod tests { .await .unwrap(); assert_eq!(actual, expected, "{context} rows={}", selected.len()); + assert!( + !shares_resident_store(&actual), + "{context} rows={}", + selected.len() + ); } if !rows.is_empty() && rows.len() < size { let error = codes @@ -11756,18 +11917,21 @@ mod tests { } } } + // Only code-only entries attach the resident rows, so only they + // run both attach modes. let variants = [ - (EntryColumns::All, false), - (EntryColumns::All, true), - (EntryColumns::Codes, false), - (EntryColumns::Codes, true), + (EntryColumns::All, false, ResidentAttach::Share), + (EntryColumns::All, true, ResidentAttach::Share), + (EntryColumns::Codes, false, ResidentAttach::Share), + (EntryColumns::Codes, true, ResidentAttach::Share), + (EntryColumns::Codes, true, ResidentAttach::Copy), ]; 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 { + for (entry_columns, resident, attach) in variants { let context = format!( - "layered={layered} cache={cache:?} entries={entry_columns} resident={resident}" + "layered={layered} cache={cache:?} entries={entry_columns} resident={resident} attach={attach}" ); layered_stats::snapshot_and_reset(); let (_dataset, index, _) = open_entry_columns_test_index( @@ -11776,6 +11940,7 @@ mod tests { OriginLatencyClass::High, resident, entry_columns, + Some(attach), ) .await; let ivf = lazy_index(&index); @@ -11877,6 +12042,343 @@ mod tests { } } + /// Whether every fixed-width column of `batch`, the row ids and + /// factors a resident store keeps, holds exactly its values: no + /// validity buffer, and a values buffer that allocates its rows' + /// bytes and no more, as the index cache charges an entry. A batch + /// of no rows, such as an empty partition's prewarmed entry, which + /// may keep what the prewarm's read decoded, is not checked. + fn fixed_width_columns_exact(batch: &RecordBatch) -> bool { + batch.num_rows() == 0 + || batch.columns().iter().all(|column| { + let data = column.to_data(); + match (column.data_type().primitive_width(), data.buffers()) { + (Some(width), [values]) => { + data.nulls().is_none() && values.capacity() == data.len() * width + } + _ => true, + } + }) + } + + /// Wait until no view of the store that `store_views` counts is + /// alive, as the tasks a search spawned drop what they hold. + async fn wait_until_views_dropped(store_views: &ResidentStoreViews) -> u64 { + for _ in 0..DRAIN_POLLS { + if store_views.live() == 0 { + return 0; + } + tokio::time::sleep(DRAIN_POLL_INTERVAL).await; + } + store_views.live() + } + + /// A search of `index` per query of `queries`, sparse gathers of + /// every partition's ex planes of a layered index, from the origin + /// file and from its cached planes, and every partition's stream + /// with and without vectors: every read path that can load a cache + /// entry. + async fn read_every_way( + index: &Arc, + queries: &[(String, Query, Arc)], + label: &str, + ) { + let ivf = lazy_index(index); + for (name, query, filter) in queries { + let result = search_global(index, query, filter.clone()).await; + + result.unwrap_or_else(|error| panic!("{label} {name}: {error}")); + } + if ivf.is_layered_rq() { + for partition in 0..LAZY_PARTITIONS { + let rows = resident_test_rows(ivf.storage.partition_size(partition)); + if rows.is_empty() { + continue; + } + for source in [PlaneSource::Sparse, PlaneSource::Resident] { + let plan = GatherPlan { + planes: [source; 2], + origin_whole: [false; 2], + }; + let gathered = ivf + .storage + .gather_ex_rows(partition, &rows, plan, None, &ivf.index_cache, None) + .await + .unwrap(); + // Gathered rows are copies. + if source == PlaneSource::Sparse { + for batch in [&gathered.high, &gathered.low] { + assert!(!shares_resident_store(batch), "{label} {partition}"); + } + } + } + } + } + + for partition in 0..LAZY_PARTITIONS { + for with_vector in [false, true] { + let batches = index + .partition_reader(partition, with_vector, &NoOpMetricsCollector) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!( + rows, + ivf.storage.partition_size(partition), + "{label} {partition} with_vector={with_vector}" + ); + // The stream leaves the index, so it holds copies, + // whatever reads attach. + assert!( + !batches.iter().any(shares_resident_store), + "{label} {partition} with_vector={with_vector}" + ); + } + } + } + + /// Probe every key an entry of `ivf` can take in its index cache, + /// which cannot be walked, and check that each entry found holds no + /// memory of a resident store and its fixed-width columns exactly + /// their values; returns how many entries it found. + async fn check_cached_entries(ivf: &IvfRq, label: &str) -> usize { + let mut probed = 0; + for partition in 0..LAZY_PARTITIONS { + let mut entries: Vec<(String, Vec)> = Vec::new(); + let key = IVFPartitionKey::::new(partition); + if let Some(entry) = ivf.index_cache.get_resident_with_key(&key).await { + let batches = entry.storage.to_batches().unwrap().collect(); + entries.push((format!("partition {partition}"), batches)); + } + let key = PartitionCodesKey { partition }; + if let Some(codes) = ivf.index_cache.get_resident_with_key(&key).await { + entries.push((format!("codes {partition}"), vec![codes.0.clone()])); + } + if ivf.is_layered_rq() { + for plane in [0, 1, 2, SIGN_BOUNDS_PLANE] { + let key = ivf.storage.plane_key(partition, plane); + if let Some(entry) = ivf.index_cache.get_resident_with_key(&key).await { + let name = format!("plane {partition}/{plane}"); + entries.push((name, vec![entry.0.clone()])); + } + } + } + for (name, batches) in entries { + probed += 1; + for batch in &batches { + assert!(!shares_resident_store(batch), "{label} {name}"); + assert!(fixed_width_columns_exact(batch), "{label} {name}"); + } + } + } + probed + } + + /// No cache entry holds memory of the resident store, whatever the + /// entries hold and however reads attach the resident rows: after a + /// prewarm, and again after the cache was cleared so that the reads + /// load every entry, eager, lazy, cascade and prefiltered searches, + /// sparse gathers and partition streams with and without vectors + /// leave every partition, code-only and plane entry the index cache + /// holds in buffers of its own, its fixed-width columns of exactly + /// their values, and every view the reads took has dropped. A read + /// of a code-only entry views the store under `share` alone, and a + /// partition stream, which leaves the index, never does. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_cached_entries_never_hold_resident_memory( + #[case] layered: bool, + #[values(EntryColumns::All, EntryColumns::Codes)] entry_columns: EntryColumns, + #[values(ResidentAttach::Share, ResidentAttach::Copy)] attach: ResidentAttach, + ) { + let _serial = LAZY_TEST_LOCK.lock().await; + if resident_columns_setting().unwrap() == ResidentColumnsSetting::Off { + return; + } + let label = format!("layered={layered} entries={entry_columns} attach={attach}"); + 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 (dataset, _) = open_lazy_test_dataset(dir.as_str(), LazyTestCache::Resident).await; + let index_meta = dataset.load_indices().await.unwrap()[0].clone(); + let (store, index_dir) = index_files(&dataset).await; + // Open as a session that keeps the store does, so that the store + // is the process's and `shares_resident_store` knows it. + let file_cache = dataset.index_cache.for_index(&index_meta.uuid, None); + let ivf = IvfRq::try_new( + store, + index_dir, + index_meta.uuid, + None, + &dataset.metadata_cache, + file_cache.clone(), + index_meta.file_size_map(), + IvfOpenContext { + resident_columns: ResidentColumnsSetting::On, + file_cache: Some(file_cache), + ..Default::default() + }, + ) + .await + .unwrap() + .with_entry_columns_for_test(entry_columns) + .with_resident_attach_for_test(attach); + assert!(ivf.resident_columns_enabled(), "{label}"); + let index: Arc = Arc::new(ivf); + let ivf = lazy_index(&index); + assert_eq!(ivf.entry_columns(), entry_columns, "{label}"); + let store_views = ivf + .storage + .resident_columns() + .store_views() + .expect("the open loaded the store"); + + let mut queries: Vec<(String, Query, Arc)> = Vec::new(); + if layered { + for (name, query) in sign_bounds_test_queries(&key) { + queries.push((name, query, Arc::new(NoFilter))); + } + } else { + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + let mut query = lazy_test_query(key.clone(), 100, 8); + query.approx_mode = approx_mode; + let name = format!("approx={approx_mode:?}"); + queries.push((name, query, Arc::new(NoFilter))); + } + } + // A prefiltered query caches the whole partitions it probes + // (`cache_partition_rows`) where the index caches whole entries. + for (name, filter) in lazy_test_filters() { + if name == "sparse" || name == "dense" { + let query = lazy_test_query(key.clone(), 100, LAZY_PARTITIONS); + queries.push((format!("filter={name}"), query, filter)); + } + } + + index.prewarm().await.unwrap(); + read_every_way(&index, &queries, &label).await; + let warmed = check_cached_entries(ivf, &label).await; + assert!(warmed >= LAZY_PARTITIONS, "{label} {warmed}"); + // The reads load every entry they need into an empty cache. + dataset.index_cache.clear().await; + read_every_way(&index, &queries, &label).await; + let loaded = check_cached_entries(ivf, &label).await; + assert!(loaded > 0, "{label}"); + assert_eq!(wait_until_views_dropped(&store_views).await, 0, "{label}"); + + // A read of a whole partition or plane, held: it views the store + // exactly where the index caches code-only entries and shares. + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 1) + .unwrap(); + let held = if layered { + ivf.storage + .load_plane(partition, 0, &ivf.index_cache, None) + .await + .unwrap() + } else { + let entry = ivf + .load_partition(partition, true, &NoOpMetricsCollector) + .await + .unwrap(); + entry.storage.to_batches().unwrap().next().unwrap() + }; + let shares = entry_columns == EntryColumns::Codes && attach == ResidentAttach::Share; + assert_eq!(shares_resident_store(&held), shares, "{label}"); + assert_eq!(store_views.live() > 0, shares, "{label}"); + drop(held); + assert_eq!(store_views.live(), 0, "{label}"); + } + + /// A view keeps its store column's memory alive, not the store: a + /// native partition built for one read from its code-only entry + /// still holds the store's rows, readable and equal to a new read's, + /// after its index dropped and the cache evicted the store, which is + /// then no longer live while the partition's views are. Once the + /// partition drops they are gone. The reopen loads the store again + /// when it opens, and its searches load nothing and find what the + /// first index found. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_shared_rows_outlive_store_eviction() { + let _serial = LAZY_TEST_LOCK.lock().await; + if resident_columns_setting().unwrap() != ResidentColumnsSetting::Auto + || entry_columns_setting().unwrap() != EntryColumns::Codes + || resident_lifetime_setting().unwrap() != ResidentLifetime::Index + { + return; + } + 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(0), 10, 8); + let session = Session::new( + LAZY_LARGE_CACHE_BYTES, + LAZY_METADATA_CACHE_BYTES, + Arc::new(ObjectStoreRegistry::default()), + ); + let dataset = open_with_session(dir.as_str(), session).await; + let index_meta = dataset.load_indices().await.unwrap()[0].clone(); + let file = index_file_key(&dataset, &index_meta); + let open = + || dataset.open_vector_index("vector", &index_meta.uuid, &NoOpMetricsCollector); + + let index = open().await.unwrap(); + let ivf = lazy_index(&index); + assert_eq!(ivf.entry_columns(), EntryColumns::Codes); + let expected = search_global(&index, &query, Arc::new(NoFilter)) + .await + .unwrap(); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 1) + .unwrap(); + // Built for this read from the cached code-only entry, not cached. + let held = ivf + .load_partition(partition, true, &NoOpMetricsCollector) + .await + .unwrap(); + let held_batches: Vec = held.storage.to_batches().unwrap().collect(); + assert!(held_batches.iter().any(shares_resident_store)); + let store_views = ivf.storage.resident_columns().store_views().unwrap(); + let views = store_views.live(); + assert!(views > 0); + + drop(index); + dataset.index_cache.clear().await; + assert!(wait_until_store_freed(&file).await); + assert_eq!(store_views.live(), views); + assert!(resident_store_views() >= views); + + layered_stats::snapshot_and_reset(); + let reopened = open().await.unwrap(); + let opened = layered_stats::snapshot_and_reset(); + assert!(opened.resident_columns_loads >= 1, "{opened:?}"); + assert_eq!(opened.resident_columns_read_loads, 0, "{opened:?}"); + let ivf = lazy_index(&reopened); + let bytes = ivf.resident_columns_bytes(); + assert_eq!(ivf.storage.resident_columns().loaded_bytes(), Some(bytes)); + let fresh = ivf + .load_partition(partition, true, &NoOpMetricsCollector) + .await + .unwrap(); + let fresh_batches: Vec = fresh.storage.to_batches().unwrap().collect(); + assert_eq!(held_batches, fresh_batches); + drop((fresh, fresh_batches)); + + drop((held, held_batches)); + assert_eq!(store_views.live(), 0); + let result = search_global(&reopened, &query, Arc::new(NoFilter)) + .await + .unwrap(); + assert_eq!(result_bits(&result), result_bits(&expected)); + let searched = layered_stats::snapshot_and_reset(); + assert_eq!(searched.resident_columns_read_loads, 0, "{searched:?}"); + } + /// 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, @@ -11924,8 +12426,11 @@ mod tests { .unwrap(); } } + // The storage was built outside an index open, so its loads are + // read loads, which other tests' opens never count. let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_loads, 0, "{stats:?}"); + assert_eq!(stats.resident_columns_read_loads, 0, "{stats:?}"); + assert_eq!(codes.resident_columns().loaded_bytes(), None); let reads: Vec<_> = (0..RESIDENT_TEST_CONCURRENT_READS) .map(|read| { @@ -11941,7 +12446,13 @@ mod tests { .map(|plane| (plane.num_rows(), true)) } else { codes - .load_partition_cached(partition, &cache, true, None) + .load_partition_cached( + partition, + &cache, + true, + codes.resident_attach(), + None, + ) .await .map(|(storage, hit)| (storage.len(), hit)) } @@ -11953,11 +12464,10 @@ mod tests { assert!(hit); } let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!(stats.resident_columns_read_loads, 1, "{stats:?}"); assert_eq!( - stats.resident_columns_bytes, - source.resident_columns_bytes(), - "{stats:?}" + codes.resident_columns().loaded_bytes(), + Some(source.resident_columns_bytes()) ); assert!(stats.resident_attach_calls > 0, "{stats:?}"); @@ -11980,7 +12490,13 @@ mod tests { } else { let stats = IoStats::new(); let (_, was_hit) = codes - .load_partition_cached(partition, &cache, true, Some(stats.clone())) + .load_partition_cached( + partition, + &cache, + true, + codes.resident_attach(), + Some(stats.clone()), + ) .await .unwrap(); assert_eq!(was_hit, hit); @@ -11989,7 +12505,7 @@ mod tests { } } let stats = layered_stats::snapshot_and_reset(); - assert_eq!(stats.resident_columns_loads, 0, "{stats:?}"); + assert_eq!(stats.resident_columns_read_loads, 0, "{stats:?}"); } /// Code-only and full entries of one index persist under keys of @@ -12031,6 +12547,7 @@ mod tests { OriginLatencyClass::High, true, entry_columns, + None, ) .await; assert_eq!(lazy_index(&index).entry_columns(), entry_columns); @@ -12153,6 +12670,7 @@ mod tests { OriginLatencyClass::High, true, entry_columns, + None, ) .await; let tiered = tiered.unwrap(); @@ -12226,6 +12744,7 @@ mod tests { OriginLatencyClass::High, true, entry_columns, + None, ) .await; let ivf = lazy_index(&index); @@ -12428,12 +12947,24 @@ mod tests { .unwrap() } - /// [`open_rq_index`], as if the index's origin were of `class`. The - /// twins these tests open stand for an index on a cloud object store - /// that keeps its small columns resident, in a store of its own, and - /// a local one that reads them from the file, whatever the - /// environment: residency does not follow the class, so the twins - /// pin it. + /// Whether a twin [`open_with_origin_latency`] opens as if its origin + /// were of `class` keeps its small columns resident, in a store of + /// its own: as `LANCE_RQ_RESIDENT_COLUMNS` sets it when it is `on` + /// or `off`, and under `auto`, which leaves it to the index, as the + /// twins stand for an index on a cloud object store that keeps them + /// and a local one that reads them from the file. Residency does not + /// follow the class, so the twins pin it. + fn twin_keeps_resident_store(class: OriginLatencyClass) -> bool { + match resident_columns_setting().unwrap() { + ResidentColumnsSetting::Off => false, + ResidentColumnsSetting::On => true, + ResidentColumnsSetting::Auto => class == OriginLatencyClass::High, + } + } + + /// [`open_rq_index`], as if the index's origin were of `class`, with + /// its small columns resident where [`twin_keeps_resident_store`] + /// says. async fn open_with_origin_latency( dataset: &Dataset, store: Arc, @@ -12444,7 +12975,7 @@ mod tests { .await .with_origin_latency_for_test(class) .unwrap() - .with_resident_columns_for_test(class == OriginLatencyClass::High) + .with_resident_columns_for_test(twin_keeps_resident_store(class)) .await .unwrap(); Arc::new(ivf)