diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 7b48ae44ad3..7b55e874d73 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -32,6 +32,7 @@ use lance_select::RowSetOps; use object_store::path::Path; use roaring::RoaringBitmap; use smallvec::SmallVec; +use std::any::Any; use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; @@ -145,6 +146,17 @@ fn merge_all_tail_partitions( Ok(merged_builders) } +#[derive(Clone)] +struct BuildLifetimeGuard(Arc); + +impl Debug for BuildLifetimeGuard { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BuildLifetimeGuard") + .field("strong_count", &Arc::strong_count(&self.0)) + .finish_non_exhaustive() + } +} + #[derive(Debug)] pub struct InvertedIndexBuilder { params: InvertedIndexParams, @@ -157,6 +169,7 @@ pub struct InvertedIndexBuilder { src_store: Option>, progress: Arc, deleted_fragments: RoaringBitmap, + build_lifetime: Option, } impl InvertedIndexBuilder { @@ -201,6 +214,7 @@ impl InvertedIndexBuilder { posting_tail_codec: format_version.posting_tail_codec(), progress: noop_progress(), deleted_fragments, + build_lifetime: None, } } @@ -230,6 +244,14 @@ impl InvertedIndexBuilder { self } + /// Retain a caller-owned resource guard until all asynchronously spawned + /// build work, including CPU compression tasks, has finished. + #[doc(hidden)] + pub fn with_build_lifetime(mut self, guard: Arc) -> Self { + self.build_lifetime = Some(BuildLifetimeGuard(guard)); + self + } + pub async fn update( &mut self, new_data: SendableRecordBatchStream, @@ -344,12 +366,14 @@ impl InvertedIndexBuilder { // upload's whole-request timeout keeps running while // its task waits to be polled. The builder is moved // in and handed back so ownership survives the hop. - accumulated = spawn_cpu(move || { - accumulated.merge_from(partition_builder)?; - Result::Ok(accumulated) - }) - .await?; - merged = Some(accumulated); + let build_lifetime = self.build_lifetime.clone(); + let (next_accumulated, _completed_build_lifetime) = + spawn_cpu(move || { + accumulated.merge_from(partition_builder)?; + Result::Ok((accumulated, build_lifetime)) + }) + .await?; + merged = Some(next_accumulated); } } None => merged = Some(partition_builder), @@ -371,7 +395,11 @@ impl InvertedIndexBuilder { let partition_id = self.next_partition_id() | self.fragment_mask.unwrap_or(0); builder.set_id(partition_id); let files = builder - .write_to(dest_store, self.partition_write_target()) + .write_to_with_lifetime( + dest_store, + self.partition_write_target(), + self.build_lifetime.clone(), + ) .await?; self.new_partitions.push(partition_id); Ok(files) @@ -417,9 +445,13 @@ impl InvertedIndexBuilder { let next_id = self.next_partition_id(); let id_alloc = Arc::new(AtomicU64::new(next_id)); let tokenized_count = Arc::new(AtomicU64::new(0)); + let build_lifetime = self.build_lifetime.clone(); let (sender, receiver) = async_channel::bounded(num_workers); let dest_store = dest_store.clone_arc(); - let mut index_tasks = Vec::with_capacity(num_workers); + // JoinSet aborts all still-running workers when update_index is + // cancelled or returns early. Plain JoinHandles detach on drop and + // would let query-time residual builders outlive their resource permit. + let mut index_tasks = tokio::task::JoinSet::new(); for _ in 0..num_workers { let tokenizer = tokenizer.clone(); let receiver: async_channel::Receiver = receiver.clone(); @@ -427,9 +459,16 @@ impl InvertedIndexBuilder { let id_alloc = id_alloc.clone(); let progress = self.progress.clone(); let tokenized_count = tokenized_count.clone(); - index_tasks.push(tokio::task::spawn(async move { - let mut worker = - IndexWorker::new(tokenizer, dest_store, id_alloc, worker_config).await?; + let build_lifetime = build_lifetime.clone(); + index_tasks.spawn(async move { + let mut worker = IndexWorker::new( + tokenizer, + dest_store, + id_alloc, + worker_config, + build_lifetime, + ) + .await?; while let Ok(batch) = receiver.recv().await { let num_rows = batch.num_rows(); worker.process_batch(batch).await?; @@ -441,7 +480,7 @@ impl InvertedIndexBuilder { .await?; } worker.finish().await - })); + }); } let index_build = async { @@ -485,18 +524,30 @@ impl InvertedIndexBuilder { let start = std::time::Instant::now(); let mut tail_partitions = Vec::new(); let mut files = Vec::new(); - for index_task in index_tasks { - let output = index_task.await??; + while let Some(index_task) = index_tasks.join_next().await { + let output = index_task??; self.new_partitions.extend(output.partitions); files.extend(output.files); if let Some(tail_partition) = output.tail_partition { tail_partitions.push(tail_partition); } } - let merged_tail_partitions = spawn_cpu(move || { - merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes) - }) - .await?; + let merged_tail_partitions = if tail_partitions.len() <= 1 { + // A single worker can produce at most one tail. Keep its + // potentially large builder owned by this cancellable future; + // handing it to spawn_cpu would detach up to the full worker + // memory budget when the outer query is cancelled. + merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes)? + } else { + let build_lifetime = build_lifetime.clone(); + let (merged, _completed_build_lifetime) = spawn_cpu(move || { + let merged = + merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes)?; + Result::Ok((merged, build_lifetime)) + }) + .await?; + merged + }; // Tail partitions hold most of the data when workers rarely hit the // flush threshold; writing them one at a time serializes the // posting-list compression of nearly the whole index behind a @@ -505,9 +556,16 @@ impl InvertedIndexBuilder { let mut tail_writes = futures::stream::iter(merged_tail_partitions.into_iter().map(|mut builder| { let dest_store = dest_store.clone(); + let build_lifetime = build_lifetime.clone(); async move { let partition_id = builder.id(); - let files = builder.write_to(dest_store.as_ref(), write_target).await?; + let files = builder + .write_to_with_lifetime( + dest_store.as_ref(), + write_target, + build_lifetime, + ) + .await?; Result::Ok((partition_id, files)) } })) @@ -543,7 +601,11 @@ impl InvertedIndexBuilder { builder.remap(mapping).await?; files.extend( builder - .write_to(dest_store, self.partition_write_target()) + .write_to_with_lifetime( + dest_store, + self.partition_write_target(), + self.build_lifetime.clone(), + ) .await?, ); } @@ -1109,11 +1171,25 @@ impl InnerBuilder { &mut self, store: &dyn IndexStore, target: PartitionWriteTarget, + ) -> Result> { + self.write_to_with_lifetime(store, target, None).await + } + + async fn write_to_with_lifetime( + &mut self, + store: &dyn IndexStore, + target: PartitionWriteTarget, + build_lifetime: Option, ) -> Result> { let docs = Arc::new(std::mem::take(&mut self.docs)); let files = vec![ - self.write_posting_lists(store, docs.clone(), &target.posting_path(self.id)) - .await?, + self.write_posting_lists( + store, + docs.clone(), + &target.posting_path(self.id), + build_lifetime, + ) + .await?, self.write_tokens(store, &target.token_path(self.id)) .await?, self.write_docs(store, docs, &target.doc_path(self.id)) @@ -1128,6 +1204,7 @@ impl InnerBuilder { store: &dyn IndexStore, docs: Arc, path: &str, + build_lifetime: Option, ) -> Result { let id = self.id; let mut writer = store @@ -1168,7 +1245,9 @@ impl InnerBuilder { // accumulates enough data to flush an encoded column, the flush also needs // that pool. The parked producer and the starved consumer would then wait // on each other forever. - let producer = tokio::spawn(async move { + let mut producer = tokio::task::JoinSet::new(); + producer.spawn(async move { + let _build_lifetime = build_lifetime.clone(); let mut batch_builder = PostingListBatchBuilder::new( schema_for_batches, with_position, @@ -1179,31 +1258,33 @@ impl InnerBuilder { let mut encode_elapsed = Duration::ZERO; loop { let docs_for_batches = docs_for_batches.clone(); + let build_lifetime = build_lifetime.clone(); let encode_started = Instant::now(); // Build the next batch on the CPU pool. The builder and the // remaining posting lists are moved in and handed back so state // persists across batches. - let (next_builder, next_posting_lists, batch) = spawn_cpu(move || { - let mut batch_builder = batch_builder; - let mut posting_lists = posting_lists; - let mut batch = None; - for posting_list in posting_lists.by_ref() { - posting_list.append_to_batch_with_docs( - &docs_for_batches, - &mut batch_builder, - format_version, - )?; - if batch_builder.len() >= batch_rows { + let (next_builder, next_posting_lists, batch, _completed_batch_lifetime) = + spawn_cpu(move || { + let mut batch_builder = batch_builder; + let mut posting_lists = posting_lists; + let mut batch = None; + for posting_list in posting_lists.by_ref() { + posting_list.append_to_batch_with_docs( + &docs_for_batches, + &mut batch_builder, + format_version, + )?; + if batch_builder.len() >= batch_rows { + batch = Some(batch_builder.finish()?); + break; + } + } + if batch.is_none() && !batch_builder.is_empty() { batch = Some(batch_builder.finish()?); - break; } - } - if batch.is_none() && !batch_builder.is_empty() { - batch = Some(batch_builder.finish()?); - } - Result::Ok((batch_builder, posting_lists, batch)) - }) - .await?; + Result::Ok((batch_builder, posting_lists, batch, build_lifetime)) + }) + .await?; encode_elapsed += encode_started.elapsed(); batch_builder = next_builder; posting_lists = next_posting_lists; @@ -1229,13 +1310,20 @@ impl InnerBuilder { write_elapsed += write_started.elapsed(); if let Err(err) = result { drop(rx); - // Wait for producer to stop; preserve the write error as the primary failure. - let _ = producer.await; + // Cancel production promptly, then wait for its async task to + // stop while preserving the writer error as the primary one. + // Any in-flight CPU batch retains the build lifetime itself. + producer.abort_all(); + let _ = producer.join_next().await; return Err(err); } } drop(rx); - let encode_elapsed = producer.await??; + let producer_result = producer + .join_next() + .await + .ok_or_else(|| Error::internal("posting producer exited without a result"))?; + let encode_elapsed = producer_result??; let finish_started = Instant::now(); let file = writer.finish().await?; write_elapsed += finish_started.elapsed(); @@ -1325,6 +1413,7 @@ struct IndexWorker { token_ids: Vec, last_token_count: usize, coordinate_rank: usize, + build_lifetime: Option, } struct TailPartition { @@ -1393,6 +1482,7 @@ impl IndexWorker { dest_store: Arc, id_alloc: Arc, config: IndexWorkerConfig, + build_lifetime: Option, ) -> Result { let schema = inverted_list_schema_for_version_with_block_size( config.with_position, @@ -1426,6 +1516,7 @@ impl IndexWorker { token_ids: Vec::new(), last_token_count: 0, coordinate_rank: config.coordinate_rank, + build_lifetime, }) } @@ -1809,7 +1900,11 @@ impl IndexWorker { PartitionWriteTarget::Final }; let files = builder - .write_to(self.dest_store.as_ref(), target) + .write_to_with_lifetime( + self.dest_store.as_ref(), + target, + self.build_lifetime.clone(), + ) .await .map_err(|err| { Error::execution(format!( @@ -2914,6 +3009,140 @@ mod tests { write_count: Arc, } + #[derive(Clone, Copy, Debug)] + enum GatedWriteBehavior { + Pending, + Fail, + } + + #[derive(Clone, Debug)] + struct GatedWriteStore { + entered: Arc, + behavior: GatedWriteBehavior, + } + + impl GatedWriteStore { + fn new(behavior: GatedWriteBehavior) -> Self { + Self { + entered: Arc::new(tokio::sync::Notify::new()), + behavior, + } + } + } + + impl DeepSizeOf for GatedWriteStore { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + #[derive(Debug)] + struct GatedWriter { + entered: Arc, + behavior: GatedWriteBehavior, + } + + #[async_trait] + impl IndexWriter for GatedWriter { + async fn write_record_batch(&mut self, _batch: RecordBatch) -> Result { + self.entered.notify_one(); + match self.behavior { + GatedWriteBehavior::Pending => std::future::pending().await, + GatedWriteBehavior::Fail => Err(Error::internal("injected posting write failure")), + } + } + + async fn add_global_buffer(&mut self, _data: Bytes) -> Result { + Ok(1) + } + + async fn finish(&mut self) -> Result { + Err(Error::internal("gated writer unexpectedly finished")) + } + + async fn finish_with_metadata( + &mut self, + _metadata: HashMap, + ) -> Result { + Err(Error::internal("gated writer unexpectedly finished")) + } + } + + #[async_trait] + impl IndexStore for GatedWriteStore { + fn as_any(&self) -> &dyn Any { + self + } + + fn clone_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn io_parallelism(&self) -> usize { + 1 + } + + fn with_io_priority(&self, _io_priority: u64) -> Arc { + self.clone_arc() + } + + async fn new_index_file( + &self, + _name: &str, + _schema: Arc, + ) -> Result> { + Ok(Box::new(GatedWriter { + entered: self.entered.clone(), + behavior: self.behavior, + })) + } + + async fn open_index_file(&self, _name: &str) -> Result> { + Err(Error::not_supported( + "GatedWriteStore does not support reading", + )) + } + + async fn copy_index_file( + &self, + _name: &str, + _dest_store: &dyn IndexStore, + ) -> Result { + Err(Error::not_supported( + "GatedWriteStore does not support copying", + )) + } + + async fn rename_index_file(&self, _name: &str, _new_name: &str) -> Result { + Err(Error::not_supported( + "GatedWriteStore does not support renaming", + )) + } + + async fn delete_index_file(&self, _name: &str) -> Result<()> { + Err(Error::not_supported( + "GatedWriteStore does not support deleting", + )) + } + + async fn list_files_with_sizes(&self) -> Result> { + Ok(vec![]) + } + } + + #[derive(Debug)] + struct ObservedBuildLifetime { + stopped: Arc, + resource: Option, + } + + impl Drop for ObservedBuildLifetime { + fn drop(&mut self) { + drop(self.resource.take()); + self.stopped.notify_one(); + } + } + #[async_trait] impl IndexWriter for CountingWriter { async fn write_record_batch(&mut self, _batch: RecordBatch) -> Result { @@ -3022,13 +3251,97 @@ mod tests { let store = CountingStore::new(); let docs = Arc::new(std::mem::take(&mut builder.docs)); builder - .write_posting_lists(&store, docs, &posting_file_path(0)) + .write_posting_lists(&store, docs, &posting_file_path(0), None) .await?; assert_eq!(store.write_count(), 1); Ok(()) } + fn posting_builder_with_multiple_batches() -> (InnerBuilder, Arc) { + let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); + let doc_id = builder.docs.append(0, 1); + for _ in 0..(*LANCE_FTS_POSTING_BATCH_ROWS * 3) { + let mut posting_list = PostingListBuilder::new(false); + posting_list.add(doc_id, PositionRecorder::Count(1)); + builder.posting_lists.push(posting_list); + } + let docs = Arc::new(std::mem::take(&mut builder.docs)); + (builder, docs) + } + + async fn observed_build_lifetime() -> ( + BuildLifetimeGuard, + Arc, + Arc, + ) { + let resource = Arc::new(tokio::sync::Semaphore::new(1)); + let permit = resource + .clone() + .acquire_owned() + .await + .expect("build lifetime semaphore closed"); + let stopped = Arc::new(tokio::sync::Notify::new()); + let lifetime = Arc::new(ObservedBuildLifetime { + stopped: stopped.clone(), + resource: Some(permit), + }); + (BuildLifetimeGuard(lifetime), resource, stopped) + } + + #[tokio::test] + async fn test_write_posting_lists_cancellation_stops_producer() { + let (mut builder, docs) = posting_builder_with_multiple_batches(); + let store = GatedWriteStore::new(GatedWriteBehavior::Pending); + let entered = store.entered.clone(); + let (build_lifetime, resource, stopped) = observed_build_lifetime().await; + let write_task = tokio::spawn(async move { + builder + .write_posting_lists(&store, docs, &posting_file_path(0), Some(build_lifetime)) + .await + }); + + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .expect("posting writer did not receive a batch"); + let stopped = stopped.notified(); + tokio::pin!(stopped); + write_task.abort(); + let join_error = write_task + .await + .expect_err("cancelled posting write unexpectedly completed"); + assert!(join_error.is_cancelled()); + tokio::time::timeout(Duration::from_secs(5), stopped) + .await + .expect("posting producer retained the build lifetime after cancellation"); + assert!( + resource.try_acquire_owned().is_ok(), + "build resource should be reusable after producer cancellation" + ); + } + + #[tokio::test] + async fn test_write_posting_lists_preserves_writer_error_during_cleanup() { + let (mut builder, docs) = posting_builder_with_multiple_batches(); + let store = GatedWriteStore::new(GatedWriteBehavior::Fail); + let (build_lifetime, resource, stopped) = observed_build_lifetime().await; + let error = builder + .write_posting_lists(&store, docs, &posting_file_path(0), Some(build_lifetime)) + .await + .expect_err("injected posting write failure should surface"); + assert!( + error.to_string().contains("injected posting write failure"), + "unexpected posting write error: {error}" + ); + tokio::time::timeout(Duration::from_secs(5), stopped.notified()) + .await + .expect("posting producer retained the build lifetime after writer failure"); + assert!( + resource.try_acquire_owned().is_ok(), + "build resource should be reusable after writer failure" + ); + } + async fn write_partition_file_marker( store: &dyn IndexStore, path: &str, @@ -3665,6 +3978,7 @@ mod tests { block_size: params.block_size, coordinate_rank: 0, }, + None, ) .await?; worker1 @@ -3690,6 +4004,7 @@ mod tests { block_size: params.block_size, coordinate_rank: 0, }, + None, ) .await?; worker2 @@ -4074,6 +4389,7 @@ mod tests { block_size: InvertedIndexParams::default().block_size, coordinate_rank: 0, }, + None, ) .await?; @@ -4108,6 +4424,60 @@ mod tests { } } + #[derive(Debug)] + struct PendingWorkerProgress { + entered: tokio::sync::Notify, + stopped: tokio::sync::Notify, + resource: std::sync::Mutex>, + } + + impl PendingWorkerProgress { + fn new(resource: tokio::sync::OwnedSemaphorePermit) -> Self { + Self { + entered: tokio::sync::Notify::new(), + stopped: tokio::sync::Notify::new(), + resource: std::sync::Mutex::new(Some(resource)), + } + } + } + + struct PendingWorkerGuard<'a> { + stopped: &'a tokio::sync::Notify, + _resource: tokio::sync::OwnedSemaphorePermit, + } + + impl Drop for PendingWorkerGuard<'_> { + fn drop(&mut self) { + self.stopped.notify_one(); + } + } + + #[async_trait] + impl IndexBuildProgress for PendingWorkerProgress { + async fn stage_start(&self, _stage: &str, _total: Option, _unit: &str) -> Result<()> { + Ok(()) + } + + async fn stage_progress(&self, _stage: &str, _completed: u64) -> Result<()> { + let resource = self + .resource + .lock() + .expect("pending worker resource lock poisoned") + .take() + .expect("stage_progress called more than once"); + let _guard = PendingWorkerGuard { + stopped: &self.stopped, + _resource: resource, + }; + self.entered.notify_one(); + std::future::pending().await + } + + async fn stage_complete(&self, _stage: &str) -> Result<()> { + Ok(()) + } + } + #[tokio::test] async fn test_builder_reports_progress_stages() -> Result<()> { let index_dir = TempDir::default(); @@ -4378,6 +4748,7 @@ mod tests { block_size: InvertedIndexParams::default().block_size, coordinate_rank: 0, }, + None, ) .await?; @@ -4411,6 +4782,7 @@ mod tests { block_size: InvertedIndexParams::default().block_size, coordinate_rank: 0, }, + None, ) .await?; @@ -4451,6 +4823,7 @@ mod tests { block_size: InvertedIndexParams::default().block_size, coordinate_rank: 0, }, + None, ) .await?; worker @@ -4896,6 +5269,55 @@ mod tests { ); } + #[tokio::test] + async fn test_update_index_cancellation_aborts_workers() { + let index_dir = TempDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + index_dir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + let batch = make_doc_batch("hello world", 0); + let stream = + RecordBatchStreamAdapter::new(batch.schema(), stream::iter(std::iter::once(Ok(batch)))); + let stream = Box::pin(stream); + + let worker_resource = Arc::new(tokio::sync::Semaphore::new(1)); + let worker_permit = worker_resource + .clone() + .acquire_owned() + .await + .expect("worker resource semaphore closed"); + let progress = Arc::new(PendingWorkerProgress::new(worker_permit)); + let mut builder = InvertedIndexBuilder::new( + InvertedIndexParams::default() + .memory_limit_mb(1) + .num_workers(1), + ) + .with_progress(progress.clone()); + + let update_task = + tokio::spawn(async move { builder.update_index(stream, store.as_ref()).await }); + tokio::time::timeout(Duration::from_secs(5), progress.entered.notified()) + .await + .expect("worker did not reach its observable pending point"); + let stopped = progress.stopped.notified(); + tokio::pin!(stopped); + + update_task.abort(); + let join_error = update_task + .await + .expect_err("cancelled update task unexpectedly completed"); + assert!(join_error.is_cancelled()); + tokio::time::timeout(Duration::from_secs(5), stopped) + .await + .expect("worker remained detached after update cancellation"); + assert!( + worker_resource.try_acquire_owned().is_ok(), + "worker resource should be reusable after cancellation" + ); + } + #[tokio::test] async fn test_new_index_has_empty_deleted_fragments() { let index_dir = TempDir::default(); diff --git a/rust/lance-index/src/scalar/inverted/index/inverted_index.rs b/rust/lance-index/src/scalar/inverted/index/inverted_index.rs index 32208eea61f..67da4a9117d 100644 --- a/rust/lance-index/src/scalar/inverted/index/inverted_index.rs +++ b/rust/lance-index/src/scalar/inverted/index/inverted_index.rs @@ -19,7 +19,7 @@ impl InvertedPrewarmState { pub struct InvertedIndex { pub(super) params: InvertedIndexParams, pub(super) store: Arc, - pub(super) tokenizer: Box, + pub(super) tokenizer: Arc, pub(super) token_set_format: TokenSetFormat, pub(super) format_version: InvertedListFormatVersion, pub(crate) partitions: Vec>, @@ -47,7 +47,7 @@ impl Debug for InvertedIndex { impl DeepSizeOf for InvertedIndex { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.partitions.deep_size_of_children(context) + self.params.deep_size_of_children(context) + self.partitions.deep_size_of_children(context) } } @@ -123,6 +123,13 @@ impl InvertedIndex { } pub fn tokenizer(&self) -> Box { + self.tokenizer.box_clone() + } + + /// Return the immutable analyzer shared by this index. Callers still clone + /// it before tokenization because token streams require mutable state. + #[doc(hidden)] + pub fn shared_tokenizer(&self) -> Arc { self.tokenizer.clone() } @@ -201,6 +208,20 @@ impl InvertedIndex { } impl InvertedIndex { + /// Materialize the lazily owned query state used by cache weight accounting. + /// + /// This does not prewarm posting payloads into the supplied index cache. + /// It only fills state owned by this `InvertedIndex`, making its + /// [`DeepSizeOf`] value stable before a long-lived cache admits it. + #[doc(hidden)] + pub async fn materialize_cache_weight(&self) -> Result<()> { + for partition in &self.partitions { + partition.inverted_list.ensure_metadata_loaded().await?; + partition.docs.prewarm().await?; + } + Ok(()) + } + async fn load_legacy_index( store: Arc, frag_reuse_index: Option>, @@ -245,7 +266,7 @@ impl InvertedIndex { let inverted_list = invert_list_fut.await??; let docs = docs_fut.await??; - let tokenizer = tokenizer_config.build()?; + let tokenizer = Arc::from(tokenizer_config.build()?); Ok(Arc::new(Self { params: tokenizer_config, @@ -322,6 +343,34 @@ impl InvertedIndex { frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> + where + Self: Sized, + { + Self::load_inner(store, frag_reuse_index, index_cache, None).await + } + + /// Load an immutable segment while sharing an already validated analyzer. + /// This avoids retaining one language model or custom stop-word set per + /// cached residual fragment. + #[doc(hidden)] + pub async fn load_with_shared_tokenizer( + store: Arc, + frag_reuse_index: Option>, + index_cache: &LanceCache, + tokenizer: Arc, + ) -> Result> + where + Self: Sized, + { + Self::load_inner(store, frag_reuse_index, index_cache, Some(tokenizer)).await + } + + async fn load_inner( + store: Arc, + frag_reuse_index: Option>, + index_cache: &LanceCache, + shared_tokenizer: Option>, + ) -> Result> where Self: Sized, { @@ -408,7 +457,11 @@ impl InvertedIndex { DocumentGranularity::ListElement }; - let tokenizer = params.build()?; + let tokenizer = if let Some(tokenizer) = shared_tokenizer { + tokenizer + } else { + Arc::from(params.build()?) + }; Ok(Arc::new(Self { params, store, @@ -422,10 +475,11 @@ impl InvertedIndex { deleted_fragments, })) } - Err(_) => { + Err(_) if shared_tokenizer.is_none() => { // old index format Self::load_legacy_index(store, frag_reuse_index, index_cache).await } + Err(error) => Err(error), } } } @@ -638,7 +692,7 @@ impl InvertedIndex { /// Search docs match the input text. async fn do_search(&self, text: &str) -> Result { let params = FtsSearchParams::new(); - let mut tokenizer = self.tokenizer.clone(); + let mut tokenizer = self.tokenizer(); let tokens = collect_query_tokens(text, &mut tokenizer); let (doc_ids, _) = self diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index fd9b9294d7b..24810f3be9c 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use lance_core::{Error, Result}; +use lance_core::{ + Error, Result, + deepsize::{Context, DeepSizeOf}, +}; use serde::{Deserialize, Deserializer, Serialize}; use std::{env, path::PathBuf}; @@ -235,6 +238,14 @@ pub struct InvertedIndexParams { pub(crate) format_version: Option, } +impl DeepSizeOf for InvertedIndexParams { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.lance_tokenizer.deep_size_of_children(context) + + self.base_tokenizer.deep_size_of_children(context) + + self.custom_stop_words.deep_size_of_children(context) + } +} + // Unknown fields must remain ignored because these params are persisted across Lance versions. #[derive(Debug, Deserialize)] struct RawInvertedIndexParams { @@ -615,6 +626,15 @@ impl Default for InvertedIndexParams { } impl InvertedIndexParams { + /// Whether this analyzer loads an opaque external language model whose + /// retained heap cannot currently be measured by `DeepSizeOf`. + #[doc(hidden)] + pub fn uses_external_language_model(&self) -> bool { + self.base_tokenizer.starts_with("lindera/") + || self.base_tokenizer.starts_with("jieba/") + || self.base_tokenizer == "jieba" + } + /// Create a new `InvertedIndexParams` with the given base tokenizer and language. /// /// The `base_tokenizer` can be one of the following: @@ -1147,7 +1167,7 @@ mod tests { use crate::pbold::inverted_index_details::DocumentGranularity as PbDocumentGranularity; use super::{DocumentGranularity, InvertedIndexParams, InvertedListFormatVersion}; - use lance_core::Error; + use lance_core::{Error, deepsize::DeepSizeOf}; use lance_tokenizer::{Language, TokenStream}; use rstest::rstest; use serde_json::json; @@ -1676,6 +1696,16 @@ mod tests { assert_eq!(tokens, vec!["the".to_string(), "data".to_string()]); } + #[test] + fn params_deep_size_charges_dynamic_tokenizer_configuration() { + let params = InvertedIndexParams::default().custom_stop_words(Some(vec![ + "a deliberately heap allocated custom stop word".repeat(8), + "another custom stop word".repeat(8), + ])); + let empty = InvertedIndexParams::default(); + assert!(params.deep_size_of() > empty.deep_size_of() + 256); + } + #[rstest] #[case::icu("icu")] #[case::icu_split("icu/split")] diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index ef4ea65de26..0b7a3c797de 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -91,7 +91,7 @@ pub use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::{Fragment, IndexMetadata}; use prost::Message; use roaring::RoaringBitmap; -use tracing::{Span, info_span, instrument}; +use tracing::{Span, info_span, instrument, warn}; use uuid::Uuid; use super::Dataset; @@ -103,7 +103,8 @@ use crate::dataset::utils::SchemaAdapter; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segment_params, load_segments, - normalize_inverted_details, resolve_fts_field, resolve_query_document_granularity, + normalize_inverted_details, residual::ResidualFtsSpec, resolve_fts_field, + resolve_query_document_granularity, }; use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ @@ -114,9 +115,9 @@ use crate::io::exec::filtered_read::{ FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, }; use crate::io::exec::fts::{ - BoostQueryExec, CompoundQueryExec, CrossColumnCompoundQueryExec, FlatMatchFilterExec, - FlatMatchQueryExec, FtsDocumentExec, HybridCompoundQueryExec, MatchQueryExec, PhraseQueryExec, - SharedFtsScorer, + BoostQueryExec, CachedResidualCompoundQueryExec, CompoundQueryExec, + CrossColumnCompoundQueryExec, FlatMatchFilterExec, FlatMatchQueryExec, FtsDocumentExec, + HybridCompoundQueryExec, MatchQueryExec, PhraseQueryExec, SharedFtsScorer, }; use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; @@ -285,13 +286,13 @@ fn supports_compound_scorer(query: &FtsQuery) -> bool { !columns.is_empty() && (!matches!(query, FtsQuery::MultiMatch(_)) || columns.len() == 1) } +/// The residual segment MVP avoids fuzzy vocabulary expansion while its cache +/// admission and invalidation contract is being established. Exact matches, +/// phrase queries, and their compound parents use the standard posting scorer. fn supports_exact_residual_compound(query: &FtsQuery) -> bool { match query { FtsQuery::Match(query) => query.fuzziness == Some(0), - // MemWAL phrase matching currently collapses tokenizer position gaps. - // Keep phrase queries on the established fallback until it can retain - // those gaps exactly (notably when stop words are configured). - FtsQuery::Phrase(_) => false, + FtsQuery::Phrase(_) => true, FtsQuery::Boost(query) => { supports_exact_residual_compound(&query.positive) && supports_exact_residual_compound(&query.negative) @@ -344,17 +345,14 @@ fn has_exact_hybrid_fts_coverage( indexed | residual == target } -fn has_compatible_hybrid_physical_segments( - params: &[InvertedIndexParams], +fn compatible_live_physical_params<'a>( + params: &'a [InvertedIndexParams], details: &[InvertedIndexDetails], has_deleted_fragments: &[bool], -) -> bool { - let Some(first) = params.first() else { - return false; - }; - params.len() == details.len() +) -> Option<&'a InvertedIndexParams> { + let first = params.first()?; + (params.len() == details.len() && params.len() == has_deleted_fragments.len() - && first.posting_block_size() == 128 && params.iter().all(|params| params == first) && details.iter().all(|details| { matches!( @@ -362,7 +360,28 @@ fn has_compatible_hybrid_physical_segments( Some(INVERTED_INDEX_VERSION_V2 | INVERTED_INDEX_VERSION_V3) ) }) - && has_deleted_fragments.iter().all(|has_deleted| !has_deleted) + && has_deleted_fragments.iter().all(|has_deleted| !has_deleted)) + .then_some(first) +} + +fn has_compatible_hybrid_physical_segments( + params: &[InvertedIndexParams], + details: &[InvertedIndexDetails], + has_deleted_fragments: &[bool], +) -> bool { + compatible_live_physical_params(params, details, has_deleted_fragments) + .is_some_and(|params| params.posting_block_size() == 128) +} + +fn can_use_query_local_residual_fallback( + compatible_hybrid_segments: bool, + residual_fragment_count: usize, + target_fragment_count: usize, + requires_positions: bool, +) -> bool { + compatible_hybrid_segments + && !requires_positions + && residual_fragment_count < target_fragment_count } fn validate_fts_query_contract(query: &FtsQuery) -> Result<()> { @@ -4312,7 +4331,7 @@ impl Scanner { && document_granularity == DocumentGranularity::Row && target_fragments .iter() - .all(|fragment| fragment.deletion_file.is_none()) + .all(|fragment| fragment.deletion_file.is_none() && fragment.overlays.is_empty()) && supports_exact_residual_compound(query); let segment_groups = futures::future::try_join_all(columns.into_iter().map(|column| { @@ -4339,8 +4358,7 @@ impl Scanner { let unindexed_fragments = self.retain_target_fragments(unindexed_fragments); if !unindexed_fragments.is_empty() && (!self.fast_search || unindexed_fragments.len() == target_fragments.len()) - && !(allow_exact_residual - && unindexed_fragments.len() < target_fragments.len()) + && !allow_exact_residual { // Flat and posting-backed leaves do not share a document // domain, so preserve the exact fallback for partial index @@ -4363,6 +4381,9 @@ impl Scanner { } FtsOverlayPlan::RowLevel { .. } | FtsOverlayPlan::FullScan => return Ok(None), }; + let mut cached_residual_params = None; + let mut query_local_residual_compatible = false; + let mut validated_segment_details = None; if allow_exact_residual && !unindexed_fragments.is_empty() { if !has_exact_hybrid_fts_coverage( &segments, @@ -4372,19 +4393,24 @@ impl Scanner { return Ok(None); } if segments.is_empty() { - return Err(Error::internal( - "hybrid compound FTS requires one indexed segment", - )); + // Current logical metadata omits custom stop words, so an + // all-residual snapshot cannot prove the complete tokenizer + // identity needed to build a compatible posting segment. + return Ok(None); } // Preserve the established semantic mismatch error before - // applying the narrower physical fast-path gate. - load_segment_details(&self.dataset, &column, &segments).await?; + // applying the physical fast-path gates. + let details = + load_segment_details(&self.dataset, &column, &segments).await?; if !has_append_only_indexed_field_history(&self.dataset, &segments).await { // Logical coverage can prune a same-id field rewrite // while the physical segment still contributes the // obsolete document to BM25 corpus statistics. return Ok(None); } + // `load_segment_params` reads root params JSON, whose + // format_version is intentionally not serialized. Use the + // normalized physical details for the legacy boundary. let physical_details = futures::future::try_join_all( segments.iter().map(|segment| { load_physical_fts_details(&self.dataset, &column, segment) @@ -4423,18 +4449,37 @@ impl Scanner { }), ) .await?; - if !has_compatible_hybrid_physical_segments( + let Some(params) = compatible_live_physical_params( &segment_params, &physical_details, &has_deleted_fragments, - ) { - // Larger posting blocks quantize document lengths, and - // retired physical documents remain in corpus stats. - // Either would make the two arms incomparable. + ) else { + // Retired physical documents remain in corpus stats, and + // incompatible tokenizers cannot share one residual writer. return Ok(None); - } + }; + query_local_residual_compatible = + has_compatible_hybrid_physical_segments( + &segment_params, + &physical_details, + &has_deleted_fragments, + ); + let format_version = InvertedIndexParams::try_from( + physical_details.first().ok_or_else(|| { + Error::internal( + "cached residual FTS requires one physical segment detail" + .to_string(), + ) + })?, + )? + .resolved_format_version(); + cached_residual_params = Some(params.clone().format_version(format_version)); + validated_segment_details = Some(details); } + let mut residual_spec = None; + let mut use_query_local_residual = false; + if cross_column { let details = futures::future::try_join_all( segments.iter().map(|segment| { @@ -4459,17 +4504,76 @@ impl Scanner { if !all_modern { return Ok(None); } - } else if phrase_columns.contains(&column) { - let details = load_segment_details(&self.dataset, &column, &segments).await?; - if !details.with_position { - return Err(Error::invalid_input( - "position is not found but required for phrase queries, try recreating the index with position" - .to_string(), - )); + } else { + if phrase_columns.contains(&column) { + let details = if let Some(details) = validated_segment_details.take() { + details + } else { + load_segment_details(&self.dataset, &column, &segments).await? + }; + if !details.with_position { + return Err(Error::invalid_input( + "position is not found but required for phrase queries, try recreating the index with position" + .to_string(), + )); + } + } + if allow_exact_residual && !unindexed_fragments.is_empty() { + let Some(params) = cached_residual_params else { + return Ok(None); + }; + let resolved = resolve_fts_field( + self.dataset.schema(), + &column, + document_granularity, + )?; + let committed_segment_uuid = segments + .first() + .ok_or_else(|| { + Error::internal( + "cached residual FTS lost its validated committed segment" + .to_string(), + ) + })? + .uuid; + match ResidualFtsSpec::try_new( + &self.dataset, + &index, + &column, + committed_segment_uuid, + &unindexed_fragments, + resolved, + params, + )? { + Ok(spec) => { + residual_spec = Some(spec); + use_query_local_residual = + can_use_query_local_residual_fallback( + query_local_residual_compatible, + unindexed_fragments.len(), + target_fragments.len(), + phrase_columns.contains(&column), + ); + } + Err(reason) => { + tracing::debug!( + column, + reason, + "cached residual FTS segment admission rejected" + ); + return Ok(None); + } + } } } - Ok(Some((column, segments, unindexed_fragments))) + Ok(Some(( + column, + segments, + unindexed_fragments, + residual_spec, + use_query_local_residual, + ))) } })) .await?; @@ -4478,11 +4582,11 @@ impl Scanner { }; if !cross_column { - let (column, segments, unindexed_fragments) = + let (column, segments, unindexed_fragments, residual_spec, use_query_local_residual) = segment_groups.into_iter().next().ok_or_else(|| { Error::internal("compound scorer requires one column".to_string()) })?; - if allow_exact_residual && !unindexed_fragments.is_empty() { + let hybrid_fallback = if use_query_local_residual { let resolved = resolve_fts_field(self.dataset.schema(), &column, document_granularity)?; let scan_column = if resolved.has_lists() { @@ -4506,25 +4610,49 @@ impl Scanner { None, ) .await?; - return Ok(Some(Arc::new(HybridCompoundQueryExec::new( + Some(Arc::new(HybridCompoundQueryExec::new( self.dataset.clone(), query.clone(), params.clone(), column, - segments, + segments.clone(), plan, - )))); - } - return Ok(Some(Arc::new( - CompoundQueryExec::new_with_segments( + )) as Arc) + } else { + None + }; + if let Some(residual_spec) = residual_spec { + let fallback = if let Some(hybrid_fallback) = hybrid_fallback { + hybrid_fallback + } else { + Box::pin(self.plan_fts_inner( + query, + params, + filter_plan, + prefilter_source, + false, + )) + .await? + }; + return Ok(Some(Arc::new(CachedResidualCompoundQueryExec::new( self.dataset.clone(), query.clone(), params.clone(), prefilter_source.clone(), segments, - ) - .with_external_mask(self.external_row_mask.clone()), - ))); + residual_spec, + fallback, + )))); + } + let exec = CompoundQueryExec::new_with_segments( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + segments, + ) + .with_external_mask(self.external_row_mask.clone()); + return Ok(Some(Arc::new(exec))); } let mut coverage_groups = segment_groups.iter(); @@ -4551,7 +4679,7 @@ impl Scanner { } let segment_groups = segment_groups .into_iter() - .map(|(column, segments, _)| (column, segments)) + .map(|(column, segments, _, _, _)| (column, segments)) .collect(); let exec = CrossColumnCompoundQueryExec::new_with_segments( self.dataset.clone(), @@ -4570,9 +4698,22 @@ impl Scanner { params: &FtsSearchParams, filter_plan: &ExprFilterPlan, prefilter_source: &PreFilterSource, + ) -> Result> { + self.plan_fts_inner(query, params, filter_plan, prefilter_source, true) + .await + } + + async fn plan_fts_inner( + &self, + query: &FtsQuery, + params: &FtsSearchParams, + filter_plan: &ExprFilterPlan, + prefilter_source: &PreFilterSource, + allow_compound: bool, ) -> Result> { let document_granularity = self.fts_document_granularity(query)?; - if !document_granularity.is_list_element() + if allow_compound + && !document_granularity.is_list_element() && supports_compound_scorer(query) && let Some(plan) = self .plan_compound_scorer( @@ -4604,17 +4745,19 @@ impl Scanner { // the documents that are not in the top-k results of the positive query, // but in the final top-k results. let unlimited_params = params.clone().with_limit(None); - let positive_exec = Box::pin(self.plan_fts( + let positive_exec = Box::pin(self.plan_fts_inner( &query.positive, &unlimited_params, filter_plan, prefilter_source, + allow_compound, )); - let negative_exec = Box::pin(self.plan_fts( + let negative_exec = Box::pin(self.plan_fts_inner( &query.negative, &unlimited_params, filter_plan, prefilter_source, + allow_compound, )); let (positive_exec, negative_exec) = futures::future::try_join(positive_exec, negative_exec).await?; @@ -4645,7 +4788,7 @@ impl Scanner { .map(|(match_query, field_prefilter_source)| { let unlimited_params = &unlimited_params; async move { - if can_use_bounded_compound { + if allow_compound && can_use_bounded_compound { let child_query = FtsQuery::Match(match_query.clone()); if let Some(plan) = self .plan_compound_scorer( @@ -4732,11 +4875,12 @@ impl Scanner { let mut should = Vec::with_capacity(query.should.len()); for subquery in &query.should { should.push( - Box::pin(self.plan_fts( + Box::pin(self.plan_fts_inner( subquery, &unlimited_params, filter_plan, prefilter_source, + allow_compound, )) .await?, ); @@ -4744,11 +4888,12 @@ impl Scanner { let mut must = Vec::with_capacity(query.must.len()); for subquery in &query.must { must.push( - Box::pin(self.plan_fts( + Box::pin(self.plan_fts_inner( subquery, &unlimited_params, filter_plan, prefilter_source, + allow_compound, )) .await?, ); @@ -4756,11 +4901,12 @@ impl Scanner { let mut must_not = Vec::with_capacity(query.must_not.len()); for subquery in &query.must_not { must_not.push( - Box::pin(self.plan_fts( + Box::pin(self.plan_fts_inner( subquery, &unlimited_params, filter_plan, prefilter_source, + allow_compound, )) .await?, ); @@ -7392,6 +7538,7 @@ mod test { use lance_index::scalar::inverted::query::{ BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, PhraseQuery, }; + use lance_index::scalar::inverted::{INVERTED_INDEX_VERSION_V1, InvertedListFormatVersion}; use lance_index::vector::hnsw::builder::HnswBuildParams; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; @@ -7644,6 +7791,27 @@ mod test { &[InvertedIndexDetails::default()], &[false] )); + assert!( + compatible_live_physical_params( + &[params.clone().block_size(256).unwrap()], + std::slice::from_ref(&modern_details), + &[false] + ) + .is_some(), + "cached standard segments support the committed 256-doc encoding" + ); + assert!( + compatible_live_physical_params( + std::slice::from_ref(¶ms), + &[InvertedIndexDetails { + posting_format_version: Some(1), + ..Default::default() + }], + &[false] + ) + .is_none(), + "cached residual segments must not extend the legacy v1 format" + ); assert!(!has_compatible_hybrid_physical_segments( &[params], &[modern_details], @@ -7651,6 +7819,138 @@ mod test { )); } + #[tokio::test] + async fn cached_residual_rejects_physically_written_v1_segment() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + let params = InvertedIndexParams::default() + .with_position(true) + .remove_stop_words(false) + .format_version(InvertedListFormatVersion::V1); + test_ds + .dataset + .create_index(&["s"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + test_ds.append_new_data().await.unwrap(); + + let segments = load_segments(&test_ds.dataset, "s", DocumentGranularity::Row) + .await + .unwrap() + .unwrap(); + let details = load_physical_fts_details(&test_ds.dataset, "s", &segments[0]) + .await + .unwrap(); + assert_eq!( + details.posting_format_version, + Some(INVERTED_INDEX_VERSION_V1), + "the fixture must exercise a physically written v1 segment" + ); + + let query: FtsQuery = BooleanQuery::new([( + Occur::Should, + MatchQuery::new("409".to_string()) + .with_column(Some("s".to_string())) + .into(), + )]) + .into(); + let mut scan = test_ds.dataset.scan(); + scan.full_text_search(FullTextSearchQuery::new_query(query).limit(Some(10))) + .unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("CachedResidualCompoundFtsScorer"), + "v1 postings must retain the exact fallback:\n{plan}" + ); + assert_eq!(scan.try_into_batch().await.unwrap().num_rows(), 1); + } + + #[test] + fn test_cached_residual_scorer_accepts_exact_compounds_only() { + let exact: FtsQuery = BooleanQuery::new([ + ( + Occur::Must, + MatchQuery::new("alpha".to_string()) + .with_column(Some("body".to_string())) + .into(), + ), + ( + Occur::Should, + PhraseQuery::new("beta gamma".to_string()) + .with_column(Some("body".to_string())) + .into(), + ), + ]) + .into(); + let fuzzy: FtsQuery = BooleanQuery::new([( + Occur::Must, + MatchQuery::new("alpha".to_string()) + .with_column(Some("body".to_string())) + .with_fuzziness(Some(1)) + .into(), + )]) + .into(); + let automatic_fuzzy: FtsQuery = BooleanQuery::new([( + Occur::Must, + MatchQuery::new("alpha".to_string()) + .with_column(Some("body".to_string())) + .with_fuzziness(None) + .into(), + )]) + .into(); + + assert!(supports_exact_residual_compound(&exact)); + assert!(!supports_exact_residual_compound(&fuzzy)); + assert!(!supports_exact_residual_compound(&automatic_fuzzy)); + } + + #[test] + fn test_cached_residual_requires_exact_disjoint_coverage() { + fn segment(coverage: Option) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + fields: vec![0], + covering_fields: Vec::new(), + name: "fts".to_string(), + dataset_version: 1, + fragment_bitmap: coverage, + index_details: None, + index_version: 3, + created_at: None, + base_id: None, + files: None, + } + } + let targets = [Fragment::new(1), Fragment::new(2)]; + let all_residual = [Fragment::new(1), Fragment::new(2)]; + assert!(has_exact_hybrid_fts_coverage(&[], &all_residual, &targets,)); + assert!(has_exact_hybrid_fts_coverage( + &[segment(Some(RoaringBitmap::new()))], + &all_residual, + &targets, + )); + + assert!(!has_exact_hybrid_fts_coverage( + &[ + segment(Some(RoaringBitmap::from_iter([1]))), + segment(Some(RoaringBitmap::from_iter([1]))), + ], + &[Fragment::new(2)], + &targets, + )); + assert!(!has_exact_hybrid_fts_coverage( + &[segment(Some(RoaringBitmap::from_iter([1, 99])))], + &[Fragment::new(2)], + &targets, + )); + assert!(!has_exact_hybrid_fts_coverage( + &[segment(None)], + &all_residual, + &targets, + )); + } + #[test] fn test_collect_phrase_columns_traverses_prohibited_subtrees() { let phrase = @@ -8146,6 +8446,164 @@ mod test { ); } + #[tokio::test] + async fn cached_residual_compound_fts_matches_exact_flat_fallback() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_fts_index().await.unwrap(); + test_ds.append_data_with_range(0, 10).await.unwrap(); + test_ds.append_data_with_range(10, 20).await.unwrap(); + + let query = || { + FullTextSearchQuery::new_query(FtsQuery::Boolean(BooleanQuery::new([ + ( + Occur::Should, + MatchQuery::new("4".to_string()) + .with_column(Some("s".to_string())) + .into(), + ), + ( + Occur::Should, + PhraseQuery::new("s 9".to_string()) + .with_column(Some("s".to_string())) + .into(), + ), + ]))) + .limit(Some(20)) + }; + + let mut cached_scan = test_ds.dataset.scan(); + cached_scan.full_text_search(query()).unwrap().with_row_id(); + let logical_index = test_ds + .dataset + .load_scalar_index( + IndexCriteria::default() + .for_column("s") + .supports_fts() + .with_fts_document_granularity(DocumentGranularity::Row), + ) + .await + .unwrap() + .unwrap(); + let residual_cache = test_ds + .dataset + .index_cache + .for_index(&logical_index.uuid, None) + .with_key_prefix("residual-fts"); + let committed_segments = load_segments(&test_ds.dataset, "s", DocumentGranularity::Row) + .await + .unwrap() + .unwrap(); + for segment in committed_segments { + test_ds + .dataset + .open_scalar_index("s", &segment.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + } + let before = residual_cache.stats().await; + let mut coalesced_scan = test_ds.dataset.scan(); + coalesced_scan + .full_text_search(query()) + .unwrap() + .with_row_id(); + let (plan, coalesced_plan) = tokio::join!( + cached_scan.explain_plan(true), + coalesced_scan.explain_plan(true) + ); + let plan = plan.unwrap(); + let coalesced_plan = coalesced_plan.unwrap(); + assert!( + plan.contains("CachedResidualCompoundFtsScorer"), + "expected cached residual scorer, got:\n{plan}" + ); + assert!(coalesced_plan.contains("CachedResidualCompoundFtsScorer")); + let after_explain = residual_cache.stats().await; + assert_eq!(after_explain.misses, before.misses); + + // First use only records the exact working set and runs the fallback. + cached_scan.try_into_batch().await.unwrap(); + let after_seen = residual_cache.stats().await; + assert!(after_seen.misses > after_explain.misses); + + // Repeated concurrent uses build each standard immutable segment once; + // same-key cache single-flight coalesces the peer execution. + let mut coalesced_peer = test_ds.dataset.scan(); + coalesced_peer + .full_text_search(query()) + .unwrap() + .with_row_id(); + let (built, peer) = tokio::join!( + coalesced_scan.try_into_batch(), + coalesced_peer.try_into_batch() + ); + let built = built.unwrap(); + assert_eq!(batch_row_ids(&built), batch_row_ids(&peer.unwrap())); + let after_cold_build = residual_cache.stats().await; + assert!(after_cold_build.misses > after_seen.misses); + + let mut reused_scan = test_ds.dataset.scan(); + reused_scan.full_text_search(query()).unwrap(); + let reused_plan = reused_scan.explain_plan(true).await.unwrap(); + assert!(reused_plan.contains("CachedResidualCompoundFtsScorer")); + let cached = reused_scan.try_into_batch().await.unwrap(); + let after_reuse = residual_cache.stats().await; + assert_eq!(after_reuse.misses, after_cold_build.misses); + assert!(after_reuse.hits > after_cold_build.hits); + assert_eq!(batch_row_ids(&cached), batch_row_ids(&built)); + + let mut all_rows_scan = test_ds.dataset.scan(); + all_rows_scan.with_row_id(); + let all_rows = batch_row_ids(&all_rows_scan.try_into_batch().await.unwrap()); + let mut fallback_scan = test_ds.dataset.scan(); + fallback_scan + .full_text_search(query()) + .unwrap() + .with_row_id() + .with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + all_rows, + ))); + let fallback_plan = fallback_scan.explain_plan(true).await.unwrap(); + assert!( + !fallback_plan.contains("CachedResidualCompoundFtsScorer"), + "external masks must retain the established fallback, got:\n{fallback_plan}" + ); + let fallback = fallback_scan.try_into_batch().await.unwrap(); + + assert_eq!( + batch_row_ids(&cached), + batch_row_ids(&fallback), + "cached residual scorer changed ordered top-k membership" + ); + assert_eq!( + cached + .column_by_name(SCORE_COL) + .unwrap() + .as_primitive::() + .values(), + fallback + .column_by_name(SCORE_COL) + .unwrap() + .as_primitive::() + .values(), + "cached residual scorer changed BM25 scores" + ); + + test_ds.append_data_with_range(20, 30).await.unwrap(); + let mut appended_scan = test_ds.dataset.scan(); + appended_scan.full_text_search(query()).unwrap(); + let appended_plan = appended_scan.explain_plan(true).await.unwrap(); + assert!(appended_plan.contains("CachedResidualCompoundFtsScorer")); + appended_scan.try_into_batch().await.unwrap(); + let mut appended_reuse_scan = test_ds.dataset.scan(); + appended_reuse_scan.full_text_search(query()).unwrap(); + appended_reuse_scan.try_into_batch().await.unwrap(); + let after_append = residual_cache.stats().await; + assert!(after_append.misses > after_reuse.misses); + assert!(after_append.hits > after_reuse.hits); + } + /// A cross-column boolean query plans into CrossColumnCompoundFtsScorer, /// which is a different exec from the same-column CompoundFtsScorer and /// builds its own prefilter, so it needs the mask threaded separately. diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index a0b240d2380..14af2b6fed0 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2609,7 +2609,19 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { ) .await .unwrap(); - create_fragmented_fts_index(&mut dataset, "text", true).await; + let physical_params = InvertedIndexParams::default() + .with_position(true) + .custom_stop_words(Some(vec!["unused-custom-stop-word".to_string()])); + dataset + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + &physical_params, + true, + ) + .await + .unwrap(); let appended = arrow_array::record_batch!(("text", Utf8, ["fresh alpha"]), ("id", Int32, [2])).unwrap(); @@ -2637,19 +2649,47 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { exact_scanner.limit(Some(2), None).unwrap(); let exact_plan = exact_scanner.explain_plan(false).await.unwrap(); assert!( - exact_plan.contains("HybridCompoundFtsScorer"), - "exact partial coverage should build one query-local residual index:\n{exact_plan}" - ); - assert!( - !exact_plan.contains("FlatMatchQuery"), - "hybrid compound scoring must not scan the residual once per leaf:\n{exact_plan}" + exact_plan.contains("CachedResidualCompoundFtsScorer"), + "exact partial coverage should build one cached immutable residual segment:\n{exact_plan}" ); + // The lazy node intentionally carries the exact fallback as a child. A + // one-off cold query uses it without paying the posting-build cost. let exact = exact_scanner.try_into_batch().await.unwrap(); assert_eq!( exact["id"].as_primitive::().values(), &[0, 2], "exact search should include the appended hit" ); + let mut residual_row_scan = dataset.scan(); + residual_row_scan.with_row_id().filter("id = 2").unwrap(); + let residual_row_id = residual_row_scan.try_into_batch().await.unwrap()[ROW_ID] + .as_primitive::() + .value(0); + let built_results = compound_fts_results(&dataset, query.clone(), Some(2)).await; + assert!( + built_results + .iter() + .any(|(row_id, _)| *row_id == residual_row_id), + "the cached compound prefilter must include the residual fragment" + ); + let reused_results = compound_fts_results(&dataset, query.clone(), Some(2)).await; + assert_eq!(reused_results, built_results); + + let appended = + arrow_array::record_batch!(("text", Utf8, ["new noise"]), ("id", Int32, [3])).unwrap(); + let schema = appended.schema(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + // The new working set is first observed without a build. On the next use, + // the old fragment is reused and only the newly appended fragment loads. + compound_fts_results(&dataset, query.clone(), Some(2)).await; + let appended_results = compound_fts_results(&dataset, query.clone(), Some(2)).await; + assert_eq!(appended_results, built_results); let empty_terms_query: FtsQuery = BooleanQuery::new([ (Occur::Must, compound_match_query("", "text", 1.0)), @@ -2658,8 +2698,8 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { .into(); let empty_terms_plan = compound_fts_plan(&dataset, empty_terms_query.clone(), 2).await; assert!( - empty_terms_plan.contains("HybridCompoundFtsScorer"), - "the empty analyzed-term case must exercise the hybrid short circuit:\n{empty_terms_plan}" + empty_terms_plan.contains("CachedResidualCompoundFtsScorer"), + "the warm residual cache should retain the compound short circuit:\n{empty_terms_plan}" ); let empty_results = compound_fts_results(&dataset, empty_terms_query, Some(2)).await; assert!(empty_results.is_empty()); @@ -2675,7 +2715,8 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { filtered_scanner.limit(Some(2), None).unwrap(); let filtered_plan = filtered_scanner.explain_plan(false).await.unwrap(); assert!( - !filtered_plan.contains("HybridCompoundFtsScorer"), + !filtered_plan.contains("CachedResidualCompoundFtsScorer") + && !filtered_plan.contains("HybridCompoundFtsScorer"), "prefiltered residual scoring must retain the exact fallback:\n{filtered_plan}" ); @@ -2691,8 +2732,8 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { .into(); let phrase_plan = compound_fts_plan(&dataset, phrase_query, 2).await; assert!( - !phrase_plan.contains("HybridCompoundFtsScorer"), - "phrase position gaps are not yet supported by the residual index:\n{phrase_plan}" + phrase_plan.contains("CachedResidualCompoundFtsScorer"), + "standard residual postings should preserve phrase positions:\n{phrase_plan}" ); let mut fast_scanner = dataset.scan(); @@ -2743,7 +2784,7 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { } #[tokio::test] -async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { +async fn test_partial_compound_cached_residual_matches_rebuilt_scores_and_ties() { let initial = arrow_array::record_batch!( ("text", Utf8, ["fresh alpha", "blocked fresh alpha"]), ("id", Int32, [0, 1]) @@ -2823,12 +2864,12 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { assert_eq!( scored_row_bits(&partial_boost), scored_row_bits(&rebuilt_boost), - "hybrid Boost scores must be bit-identical to a rebuilt index" + "cached residual Boost scores must be bit-identical to a rebuilt index" ); assert_eq!( scored_row_bits(&partial_multimatch), scored_row_bits(&rebuilt_multimatch), - "hybrid MultiMatch scores must be bit-identical to a rebuilt index" + "cached residual MultiMatch scores must be bit-identical to a rebuilt index" ); } diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 78dcbd3ad43..cc9ec5739dc 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -106,6 +106,40 @@ impl DatasetPreFilter { } } + /// Build a prefilter for an exact, caller-proven fragment coverage set. + /// + /// This is used when some searchable segments are query-time immutable + /// segments and therefore do not have fabricated [`IndexMetadata`]. + pub(crate) fn new_with_fragment_bitmap( + dataset: Arc, + fragments: RoaringBitmap, + filter: Option>, + ) -> Self { + let filter = filter.map(|filter| { + async move { filter.load().await.map(Arc::new) } + .in_current_span() + .boxed() + }); + Self::new_with_fragment_bitmap_and_filter_future(dataset, fragments, filter) + } + + pub(crate) fn new_with_fragment_bitmap_and_filter_future( + dataset: Arc, + fragments: RoaringBitmap, + filter: Option>>>, + ) -> Self { + let deleted_ids = Self::create_restricted_deletion_mask(dataset, fragments) + .map(SharedPrerequisite::spawn); + let filtered_ids = filter.map(SharedPrerequisite::spawn); + Self { + deleted_ids, + filtered_ids, + deleted_fragments: None, + overlay_block: None, + final_mask: Mutex::new(OnceCell::new()), + } + } + #[instrument(level = "debug", skip_all)] async fn do_create_deletion_mask( dataset: Arc, diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index c7c521050b4..ab4c74a34ed 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -38,6 +38,8 @@ use crate::{ index::{DatasetIndexExt, scalar::fetch_index_details}, }; +pub(crate) mod residual; + #[derive(Debug, Clone)] enum FtsTraversal { Text, diff --git a/rust/lance/src/index/scalar/inverted/residual.rs b/rust/lance/src/index/scalar/inverted/residual.rs new file mode 100644 index 00000000000..a09478b8960 --- /dev/null +++ b/rust/lance/src/index/scalar/inverted/residual.rs @@ -0,0 +1,1782 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Query-time immutable FTS segments for append-only residual fragments. +//! +//! This mirrors Lucene's near-real-time model: newly appended fragments are +//! converted into the same immutable posting format as committed FTS segments. +//! The segment, rather than a query result, is cached and can therefore serve +//! arbitrary exact compound queries until the fragment or index configuration +//! changes. + +use std::{ + borrow::Cow, + collections::HashSet, + fmt::Display, + ops::Range, + sync::{ + Arc, LazyLock, Mutex, Weak, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, +}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; +use lance_core::{ + Result, + cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache}, + deepsize::{Context, DeepSizeOf}, +}; +use lance_index::{ + progress::noop_progress, + scalar::{ + TrainingCriteria, TrainingOrdering, + inverted::{InvertedIndex, InvertedIndexBuilder, InvertedIndexParams}, + lance_format::LanceIndexStore, + }, +}; +use lance_io::object_store::ObjectStore as LanceObjectStore; +use lance_table::format::{Fragment, IndexMetadata}; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OsObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, + PutResult, RenameOptions, Result as OsResult, UploadPart, memory::InMemory, path::Path, +}; +use tokio::sync::Semaphore; +use uuid::Uuid; + +use super::ResolvedFtsField; +use crate::{Dataset, index::DatasetIndexInternalExt, index::scalar::load_fts_training_data}; + +/// Bound query-time indexing without serializing unrelated datasets behind one +/// long build. Exact and overlapping working sets are coalesced by the claim +/// registry below before any fragment data is scanned. +static RESIDUAL_FTS_BUILD_SEMAPHORE: LazyLock> = LazyLock::new(|| { + let permits = std::thread::available_parallelism() + .map(|cpus| cpus.get().div_ceil(4)) + .unwrap_or(1) + .clamp(1, 2); + Arc::new(Semaphore::new(permits)) +}); + +pub(crate) const MAX_RESIDUAL_FTS_FRAGMENTS: usize = 16; +pub(crate) const MAX_RESIDUAL_FTS_ROWS: usize = 1_000_000; +const MAX_RESIDUAL_FTS_BUILD_MEMORY_MB: u64 = 256; +const MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES: usize = + (MAX_RESIDUAL_FTS_BUILD_MEMORY_MB as usize) << 20; +const MAX_RESIDUAL_FTS_SERIALIZED_BYTES: usize = 1 << 30; +const MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES: usize = 2 << 30; +const RESIDUAL_FTS_FRAGMENT_BUILD_CONCURRENCY: usize = 2; + +#[derive(Debug)] +struct BudgetedMemoryStore { + inner: InMemory, + uploaded_bytes: Arc, + budget_exceeded: Arc, + max_uploaded_bytes: usize, + build_lifetime: Option>, +} + +impl BudgetedMemoryStore { + #[cfg(test)] + fn new(max_uploaded_bytes: usize) -> Self { + Self::with_counter( + Arc::new(AtomicUsize::new(0)), + Arc::new(AtomicBool::new(false)), + max_uploaded_bytes, + ) + } + + fn with_counter( + uploaded_bytes: Arc, + budget_exceeded: Arc, + max_uploaded_bytes: usize, + ) -> Self { + Self { + inner: InMemory::new(), + uploaded_bytes, + budget_exceeded, + max_uploaded_bytes, + build_lifetime: None, + } + } + + fn with_build_lifetime(mut self, build_lifetime: Weak) -> Self { + self.build_lifetime = Some(build_lifetime); + self + } + + fn reserve(&self, bytes: usize) -> OsResult<()> { + reserve_upload_bytes( + &self.uploaded_bytes, + &self.budget_exceeded, + self.max_uploaded_bytes, + bytes, + ) + } + + async fn reserve_source(&self, source: &Path) -> OsResult<()> { + let source_size = self.inner.head(source).await?.size; + let source_size = + usize::try_from(source_size).map_err(|_| object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: format!( + "source object {source} has size {source_size}, which does not fit usize" + ) + .into(), + })?; + self.reserve(source_size) + } +} + +fn reserve_upload_bytes( + uploaded_bytes: &AtomicUsize, + budget_exceeded: &AtomicBool, + max_uploaded_bytes: usize, + bytes: usize, +) -> OsResult<()> { + uploaded_bytes + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + current + .checked_add(bytes) + .filter(|next| *next <= max_uploaded_bytes) + }) + .map(|_| ()) + .map_err(|_| { + budget_exceeded.store(true, Ordering::Relaxed); + object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: format!( + "residual FTS build exceeded the {max_uploaded_bytes} byte resident budget" + ) + .into(), + } + }) +} + +impl Display for BudgetedMemoryStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BudgetedMemoryStore") + } +} + +#[async_trait] +impl OsObjectStore for BudgetedMemoryStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> OsResult { + self.reserve(payload.content_length())?; + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OsResult> { + let inner = self.inner.put_multipart_opts(location, opts).await?; + let build_lifetime = match &self.build_lifetime { + Some(build_lifetime) => { + Some( + build_lifetime + .upgrade() + .ok_or_else(|| object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: "residual FTS build lifetime ended before multipart creation" + .into(), + })?, + ) + } + None => None, + }; + Ok(Box::new(BudgetedMultipartUpload { + inner: Some(inner), + uploaded_bytes: self.uploaded_bytes.clone(), + budget_exceeded: self.budget_exceeded.clone(), + max_uploaded_bytes: self.max_uploaded_bytes, + part_bytes: 0, + build_lifetime, + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OsResult { + self.inner.get_opts(location, options).await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OsResult> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OsResult>, + ) -> BoxStream<'static, OsResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OsResult> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OsResult> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OsResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OsResult<()> { + // `InMemory::copy_opts` creates another logical object and is used by + // `merge_index_files` to publish staged partitions. Charge the source + // before the copy so this path cannot bypass the group-wide hard cap. + self.reserve_source(from).await?; + self.inner.copy_opts(from, to, opts).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OsResult<()> { + // `InMemory` inherits ObjectStore's copy-then-delete rename. Its Bytes + // payload is shared, but the operation is not an atomic map move, so + // conservatively charge the source just like copy. + self.reserve_source(from).await?; + self.inner.rename_opts(from, to, opts).await + } +} + +#[derive(Debug)] +struct BudgetedMultipartUpload { + inner: Option>, + uploaded_bytes: Arc, + budget_exceeded: Arc, + max_uploaded_bytes: usize, + part_bytes: usize, + // Keep this field last: implicit struct drop must destroy the inner upload + // and its retained parts before releasing group admission. + build_lifetime: Option>, +} + +struct TransientBudgetReservation { + uploaded_bytes: Arc, + bytes: usize, +} + +impl Drop for TransientBudgetReservation { + fn drop(&mut self) { + self.uploaded_bytes.fetch_sub(self.bytes, Ordering::Relaxed); + } +} + +#[async_trait] +impl MultipartUpload for BudgetedMultipartUpload { + fn put_part(&mut self, payload: PutPayload) -> UploadPart { + let Some(inner) = self.inner.as_mut() else { + return Box::pin(async { + Err(object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: "residual FTS multipart upload is already completed".into(), + }) + }); + }; + let payload_bytes = payload.content_length(); + let Some(part_bytes) = self.part_bytes.checked_add(payload_bytes) else { + self.budget_exceeded.store(true, Ordering::Relaxed); + return Box::pin(async { + Err(object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: "residual FTS multipart byte count overflowed".into(), + }) + }); + }; + if let Err(error) = reserve_upload_bytes( + &self.uploaded_bytes, + &self.budget_exceeded, + self.max_uploaded_bytes, + payload_bytes, + ) { + return Box::pin(async move { Err(error) }); + } + self.part_bytes = part_bytes; + inner.put_part(payload) + } + + async fn complete(&mut self) -> OsResult { + let Some(inner) = self.inner.as_mut() else { + return Err(object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: "residual FTS multipart upload is already completed".into(), + }); + }; + // InMemory concatenates all parts into a newly allocated buffer before + // the upload handle releases them. Reserve that full-copy peak across + // every concurrently building residual fragment. On success, drop the + // inner upload (and its retained parts) before releasing the transient + // reservation for the final concatenated object. + reserve_upload_bytes( + &self.uploaded_bytes, + &self.budget_exceeded, + self.max_uploaded_bytes, + self.part_bytes, + )?; + let _transient_reservation = TransientBudgetReservation { + uploaded_bytes: self.uploaded_bytes.clone(), + bytes: self.part_bytes, + }; + let result = inner.complete().await; + if result.is_ok() { + drop(self.inner.take()); + drop(self.build_lifetime.take()); + } + result + } + + async fn abort(&mut self) -> OsResult<()> { + let Some(inner) = self.inner.as_mut() else { + return Err(object_store::Error::Generic { + store: "BudgetedMemoryStore", + source: "residual FTS multipart upload is already completed".into(), + }); + }; + let result = inner.abort().await; + if result.is_ok() { + drop(self.inner.take()); + drop(self.build_lifetime.take()); + } + result + } +} + +/// Stable identity for one query-time FTS segment. +/// +/// The dataset URI is already a namespace of `DSIndexCache`. The remaining +/// fields deliberately omit dataset version so an unchanged fragment is reused +/// after another append. The complete serialized fragment metadata invalidates +/// data rewrites, overlays, deletions, and row-id metadata changes. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub(crate) struct ResidualFtsFragmentKey { + pub store_identity: Arc, + pub index_uuid: Uuid, + pub index_version: i32, + pub fragment_id: u64, + pub fragment_fingerprint: Arc<[u8]>, + pub field_id: i32, + pub canonical_path: Arc, + pub field_type: Arc, + pub params_fingerprint: Arc<[u8]>, +} + +impl CacheKey for ResidualFtsFragmentKey { + type ValueType = CachedResidualFtsEntry; + + fn key(&self) -> Cow<'_, str> { + Cow::Owned(format!( + "residual-fts/{}/{}/{}/{}", + self.index_uuid, self.index_version, self.field_id, self.fragment_id + )) + } + + fn type_name() -> &'static str { + "CachedResidualFtsEntry" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.index.fts-residual-fragment-key", 2) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_str(&self.store_identity); + builder.write_fixed_bytes(self.index_uuid.as_bytes()); + builder.write_i32(self.index_version); + builder.write_u64(self.fragment_id); + builder.write_bytes(&self.fragment_fingerprint); + builder.write_i32(self.field_id); + builder.write_str(&self.canonical_path); + builder.write_str(&self.field_type); + builder.write_bytes(&self.params_fingerprint); + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +struct ResidualFtsGroupKey { + store_identity: Arc, + index_uuid: Uuid, + index_version: i32, + members: Arc<[ResidualFtsFragmentKey]>, +} + +#[derive(Default)] +struct ActiveResidualFtsBuilds { + groups: HashSet, + fragments: HashSet, +} + +static ACTIVE_RESIDUAL_FTS_BUILDS: LazyLock> = + LazyLock::new(|| Mutex::new(ActiveResidualFtsBuilds::default())); + +#[derive(Debug)] +struct ResidualFtsBuildClaim { + group: ResidualFtsGroupKey, + fragments: Vec, +} + +impl ResidualFtsBuildClaim { + fn try_acquire( + group: &ResidualFtsGroupKey, + fragments: &[ResidualFtsFragmentKey], + ) -> Option { + let mut active = ACTIVE_RESIDUAL_FTS_BUILDS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if active.groups.contains(group) + || fragments + .iter() + .any(|fragment| active.fragments.contains(fragment)) + { + return None; + } + active.groups.insert(group.clone()); + active.fragments.extend(fragments.iter().cloned()); + Some(Self { + group: group.clone(), + fragments: fragments.to_vec(), + }) + } +} + +impl Drop for ResidualFtsBuildClaim { + fn drop(&mut self) { + let mut active = ACTIVE_RESIDUAL_FTS_BUILDS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + active.groups.remove(&self.group); + for fragment in &self.fragments { + active.fragments.remove(fragment); + } + } +} + +#[derive(Debug)] +struct ResidualFtsBuildLifetime { + _permit: tokio::sync::OwnedSemaphorePermit, + _claim: ResidualFtsBuildClaim, +} + +struct ResidualFtsFragmentBuildLifetime { + _group: Arc, + _working_set: TransientBudgetReservation, +} + +impl CacheKey for ResidualFtsGroupKey { + type ValueType = ResidualFtsGroupState; + + fn key(&self) -> Cow<'_, str> { + Cow::Owned(format!( + "residual-fts-group/{}/{}/{}", + self.index_uuid, + self.index_version, + self.members.len() + )) + } + + fn type_name() -> &'static str { + "ResidualFtsGroupState" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.index.fts-residual-group-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_str(&self.store_identity); + builder.write_fixed_bytes(self.index_uuid.as_bytes()); + builder.write_i32(self.index_version); + builder.write_u64(self.members.len() as u64); + for member in self.members.iter() { + member.write_key(builder); + } + } +} + +#[derive(Debug)] +enum ResidualFtsGroupState { + Seen, + Rejected, +} + +impl DeepSizeOf for ResidualFtsGroupState { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + 0 + } +} + +#[derive(Clone, Debug)] +struct ResidualFtsRetentionProbeKey(ResidualFtsGroupKey); + +impl CacheKey for ResidualFtsRetentionProbeKey { + type ValueType = ResidualFtsRetentionProbe; + + fn key(&self) -> Cow<'_, str> { + Cow::Owned(format!("retention-probe/{}", self.0.key())) + } + + fn type_name() -> &'static str { + "ResidualFtsRetentionProbe" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.index.fts-residual-retention-probe-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + self.0.write_key(builder); + } +} + +#[derive(Debug)] +struct ResidualFtsRetentionProbe; + +impl DeepSizeOf for ResidualFtsRetentionProbe { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + 0 + } +} + +#[derive(Debug)] +pub(crate) enum CachedResidualFtsEntry { + Ready(CachedResidualFtsSegment), +} + +impl DeepSizeOf for CachedResidualFtsEntry { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + match self { + Self::Ready(segment) => segment.deep_size_of_children(context), + } + } +} + +/// A standard immutable FTS segment plus the memory-store bytes backing it. +#[derive(Debug)] +pub(crate) struct CachedResidualFtsSegment { + index: Arc, + resident_bytes: usize, +} + +impl CachedResidualFtsSegment { + pub fn index(&self) -> Arc { + self.index.clone() + } + + pub fn resident_bytes(&self) -> usize { + self.resident_bytes + } +} + +impl DeepSizeOf for CachedResidualFtsSegment { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + self.resident_bytes + } +} + +#[derive(Debug)] +pub(crate) struct LoadedResidualFtsSegments { + pub segments: Vec>, + pub fragment_bitmap: roaring::RoaringBitmap, +} + +#[derive(Debug)] +pub(crate) enum ResidualFtsAdmission { + Eligible(LoadedResidualFtsSegments), + Deferred(&'static str), + Rejected(&'static str), +} + +/// Everything execution needs to probe or construct one exact residual group. +/// Construction is metadata-only; it never scans fragment data. +#[derive(Clone, Debug)] +pub(crate) struct ResidualFtsSpec { + logical_index: IndexMetadata, + column: Arc, + committed_segment_uuid: Uuid, + fragments: Arc<[Fragment]>, + resolved: ResolvedFtsField, + params: InvertedIndexParams, + keys: Arc<[ResidualFtsFragmentKey]>, + group_key: ResidualFtsGroupKey, + fragment_bitmap: roaring::RoaringBitmap, +} + +fn fragment_fingerprint(fragment: &Fragment) -> Result> { + Ok(Arc::from(serde_json::to_vec(fragment)?)) +} + +fn params_fingerprint(params: &InvertedIndexParams) -> Result> { + Ok(Arc::from(serde_json::to_vec(¶ms.to_training_json()?)?)) +} + +fn residual_key( + dataset: &Dataset, + index: &IndexMetadata, + fragment: &Fragment, + resolved: &ResolvedFtsField, + params_fingerprint: Arc<[u8]>, +) -> Result { + let field = dataset + .schema() + .field_by_id(resolved.final_field_id) + .ok_or_else(|| { + lance_core::Error::invalid_input(format!( + "residual FTS field id {} is missing from the dataset schema", + resolved.final_field_id + )) + })?; + Ok(ResidualFtsFragmentKey { + store_identity: Arc::from(dataset.object_store.store_prefix.as_str()), + index_uuid: index.uuid, + index_version: index.index_version, + fragment_id: fragment.id, + fragment_fingerprint: fragment_fingerprint(fragment)?, + field_id: resolved.final_field_id, + canonical_path: Arc::from(resolved.canonical_path.as_str()), + field_type: Arc::from(field.data_type().to_string()), + params_fingerprint, + }) +} + +impl ResidualFtsSpec { + pub(crate) fn try_new( + dataset: &Dataset, + logical_index: &IndexMetadata, + column: &str, + committed_segment_uuid: Uuid, + fragments: &[Fragment], + resolved: ResolvedFtsField, + params: InvertedIndexParams, + ) -> Result> { + if params.get_document_granularity().is_list_element() { + return Ok(Err("non-row document granularity")); + } + // The model tokenizers are opaque to DeepSizeOf. Until they expose an + // exact retained size, do not let a residual Arc keep an uncharged + // model alive after the committed index is evicted. + if params.uses_external_language_model() { + return Ok(Err( + "external language-model tokenizer is not cache-accountable", + )); + } + if fragments.is_empty() { + return Ok(Err("empty residual fragment group")); + } + if fragments.len() > MAX_RESIDUAL_FTS_FRAGMENTS { + return Ok(Err("too many residual fragments")); + } + if fragments + .iter() + .any(|fragment| !fragment.overlays.is_empty()) + { + return Ok(Err("residual fragment has overlays")); + } + if fragments + .iter() + .any(|fragment| fragment.deletion_file.is_some()) + { + return Ok(Err("residual fragment has deletions")); + } + let Some(total_rows) = fragments.iter().try_fold(0_usize, |total, fragment| { + fragment + .physical_rows + .and_then(|rows| total.checked_add(rows)) + }) else { + return Ok(Err("unknown or overflowing residual row count")); + }; + if total_rows > MAX_RESIDUAL_FTS_ROWS { + return Ok(Err("too many residual rows")); + } + let fragment_bitmap = fragments + .iter() + .map(|fragment| u32::try_from(fragment.id)) + .collect::>() + .map_err(|_| { + lance_core::Error::invalid_input( + "residual FTS fragment id does not fit u32".to_string(), + ) + })?; + + let params_fingerprint = params_fingerprint(¶ms)?; + let keys = fragments + .iter() + .map(|fragment| { + residual_key( + dataset, + logical_index, + fragment, + &resolved, + params_fingerprint.clone(), + ) + }) + .collect::>>()?; + let keys: Arc<[ResidualFtsFragmentKey]> = Arc::from(keys); + let group_key = ResidualFtsGroupKey { + store_identity: Arc::from(dataset.object_store.store_prefix.as_str()), + index_uuid: logical_index.uuid, + index_version: logical_index.index_version, + members: keys.clone(), + }; + Ok(Ok(Self { + logical_index: logical_index.clone(), + column: Arc::from(column), + committed_segment_uuid, + fragments: Arc::from(fragments.to_vec()), + resolved, + params, + keys, + group_key, + fragment_bitmap, + })) + } + + pub(crate) fn fragment_bitmap(&self) -> roaring::RoaringBitmap { + self.fragment_bitmap.clone() + } +} + +async fn serialized_store_size(object_store: &LanceObjectStore, index_dir: &Path) -> Result { + let mut entries = object_store.read_dir_all(index_dir, None); + let mut total = 0_u64; + while let Some(entry) = entries.next().await { + total = total.checked_add(entry?.size).ok_or_else(|| { + lance_core::Error::io(format!( + "residual FTS serialized byte count overflowed for {index_dir}" + )) + })?; + } + usize::try_from(total).map_err(|_| { + lance_core::Error::io(format!( + "residual FTS serialized size {total} does not fit usize for {index_dir}" + )) + }) +} + +async fn build_residual_segment( + dataset: &Dataset, + fragment: Fragment, + resolved: &ResolvedFtsField, + params: InvertedIndexParams, + shared_tokenizer: Arc, + build_lifetime: Arc, + group_accounted_bytes: Arc, + group_budget_exceeded: Arc, +) -> Result { + let fragment_id = u32::try_from(fragment.id).map_err(|_| { + lance_core::Error::invalid_input(format!( + "residual FTS fragment id {} does not fit u32", + fragment.id + )) + })?; + let stream = load_fts_training_data( + dataset, + resolved, + &TrainingCriteria::new(TrainingOrdering::None).with_row_id(), + Some(vec![fragment]), + true, + None, + ) + .await?; + + // Bound the builder's working set in addition to the row admission limit. + // The serialized store is checked before the value can enter the cache. + let params = params + .memory_limit_mb(MAX_RESIDUAL_FTS_BUILD_MEMORY_MB) + .num_workers(1); + // The builder's per-worker memory is temporary but overlaps the immutable + // store as files are emitted. Reserve it from the same group hard cap and + // let detached CPU work retain the reservation after query cancellation. + reserve_upload_bytes( + group_accounted_bytes.as_ref(), + group_budget_exceeded.as_ref(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES, + )?; + let weak_build_lifetime = Arc::downgrade(&build_lifetime); + let fragment_build_lifetime = Arc::new(ResidualFtsFragmentBuildLifetime { + _group: build_lifetime, + _working_set: TransientBudgetReservation { + uploaded_bytes: group_accounted_bytes.clone(), + bytes: MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES, + }, + }); + let mut object_store = LanceObjectStore::memory(); + object_store.inner = Arc::new( + BudgetedMemoryStore::with_counter( + group_accounted_bytes.clone(), + group_budget_exceeded.clone(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + ) + .with_build_lifetime(weak_build_lifetime), + ); + let object_store = Arc::new(object_store); + let index_dir = Path::from(format!("residual-{fragment_id}")); + let private_cache = Arc::new(LanceCache::no_cache()); + let store = Arc::new(LanceIndexStore::new( + object_store.clone(), + index_dir.clone(), + private_cache.clone(), + )); + let mut builder = + InvertedIndexBuilder::new_with_fragment_mask(params, Some(u64::from(fragment_id) << 32)) + .with_build_lifetime(fragment_build_lifetime.clone()); + let update_result = builder.update(stream, store.as_ref(), None).await; + drop(builder); + drop(fragment_build_lifetime); + update_result?; + lance_index::scalar::inverted::builder::merge_index_files( + object_store.as_ref(), + &index_dir, + store.clone(), + noop_progress(), + ) + .await?; + let serialized_bytes = serialized_store_size(object_store.as_ref(), &index_dir).await?; + if serialized_bytes > MAX_RESIDUAL_FTS_SERIALIZED_BYTES { + group_budget_exceeded.store(true, Ordering::Relaxed); + return Err(lance_core::Error::io(format!( + "residual FTS fragment {fragment_id} produced {serialized_bytes} serialized bytes, exceeding the {} byte build budget", + MAX_RESIDUAL_FTS_SERIALIZED_BYTES + ))); + } + let index = InvertedIndex::load_with_shared_tokenizer( + store, + None, + private_cache.as_ref(), + shared_tokenizer, + ) + .await?; + // The private readers use a no-cache backend, so decoded posting groups + // cannot accumulate outside this value's cache weight. Materialize the + // lazy metadata and document state that the index itself does retain before + // DSIndexCache computes its fixed admission weight. + index.materialize_cache_weight().await?; + // `InvertedIndex` intentionally does not charge its object store. The + // store is private to this cache value, so account its serialized files + // explicitly in addition to decoded partition state. Cache this traversal + // before admission so DSIndexCache and aggregate group admission can use + // the fixed weight in O(1). + let mut size_context = Context::default(); + let resident_bytes = index + .deep_size_of_children(&mut size_context) + .checked_add(serialized_bytes) + .ok_or_else(|| lance_core::Error::io("residual FTS resident byte count overflowed"))?; + // Store writes have already charged the serialized backing bytes. Charge + // the retained decoded/index state before returning the segment so two + // concurrent fragment materializations cannot overshoot the group cap and + // defer the rejection until after both large values are resident. + let decoded_bytes = resident_bytes + .checked_sub(serialized_bytes) + .ok_or_else(|| { + lance_core::Error::io("residual FTS decoded resident byte count underflowed") + })?; + reserve_upload_bytes( + group_accounted_bytes.as_ref(), + group_budget_exceeded.as_ref(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + decoded_bytes, + )?; + Ok(CachedResidualFtsSegment { + index, + resident_bytes, + }) +} + +/// Load or build one immutable posting segment for each append-only fragment. +/// +/// Deterministically unsupported shapes return `Rejected`; callers should keep +/// the existing exact flat fallback. I/O/build errors are returned so callers +/// can log them and fall back without poisoning the single-flight cache. +pub(crate) async fn load_residual_fts_segments( + dataset: &Dataset, + spec: &ResidualFtsSpec, +) -> Result { + let cache = dataset + .index_cache + .for_index(&spec.logical_index.uuid, None) + .with_key_prefix("residual-fts"); + // Usage and negative-admission state must not compete with the posting + // entries whose aggregate pressure it records. + let group_cache = dataset + .metadata_cache + .with_key_prefix("residual-fts-groups"); + if matches!( + group_cache.get_with_key(&spec.group_key).await.as_deref(), + Some(ResidualFtsGroupState::Rejected) + ) { + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds cache admission capacity", + )); + } + + // A completely warm group can bypass the usage marker and build permit, + // but not aggregate admission. Different append snapshots can warm + // individual members separately, so the combined group must still fit the + // same resident hard cap before it is searched. + let mut warm_entries = Vec::with_capacity(spec.keys.len()); + for key in spec.keys.iter() { + let Some(entry) = cache.get_with_key(key).await else { + warm_entries.clear(); + break; + }; + warm_entries.push(entry); + } + if warm_entries.len() == spec.keys.len() { + let warm_resident_bytes = checked_resident_sum(warm_entries.iter().map(|entry| { + let CachedResidualFtsEntry::Ready(segment) = entry.as_ref(); + segment.resident_bytes() + })); + if !matches!(warm_resident_bytes, Some(bytes) if bytes <= MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES) + { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds cache admission capacity", + )); + } + // A concurrent exact-group owner may have rejected admission while + // this task was probing members. Never let a fully warm posting set + // bypass that durable negative decision. + if matches!( + group_cache.get_with_key(&spec.group_key).await.as_deref(), + Some(ResidualFtsGroupState::Rejected) + ) { + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds cache admission capacity", + )); + } + return Ok(ResidualFtsAdmission::Eligible(loaded_segments( + warm_entries, + spec.fragment_bitmap(), + ))); + } + + match group_cache.get_with_key(&spec.group_key).await { + Some(state) if matches!(state.as_ref(), ResidualFtsGroupState::Rejected) => { + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds cache admission capacity", + )); + } + Some(_) => {} + None => { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Seen)) + .await; + if group_cache.get_with_key(&spec.group_key).await.is_none() { + return Ok(ResidualFtsAdmission::Rejected( + "residual metadata cache cannot retain the working-set marker", + )); + } + return Ok(ResidualFtsAdmission::Deferred( + "first use records residual working set", + )); + } + } + + // The independent group state can be available while the index cache is + // disabled. Probe posting-entry retention before scanning any fragment. + let retention_probe_key = ResidualFtsRetentionProbeKey(spec.group_key.clone()); + cache + .insert_with_key(&retention_probe_key, Arc::new(ResidualFtsRetentionProbe)) + .await; + if cache.get_with_key(&retention_probe_key).await.is_none() { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual index cache cannot retain posting entries", + )); + } + + // Claim the complete working set before building anything. This prevents + // concurrent invocations (including overlapping append snapshots) from + // each winning different fragment loaders under independent byte budgets. + let mut cached_entries = Vec::with_capacity(spec.keys.len()); + for key in spec.keys.iter() { + match cache.get_with_key(key).await { + Some(entry) => cached_entries.push(Some(entry)), + None => cached_entries.push(None), + } + } + let Some(build_claim) = ResidualFtsBuildClaim::try_acquire(&spec.group_key, &spec.keys) else { + return Ok(ResidualFtsAdmission::Deferred( + "residual working set is already building", + )); + }; + // Busy means exact fallback, not head-of-line blocking across datasets. + let Ok(build_permit) = RESIDUAL_FTS_BUILD_SEMAPHORE.clone().try_acquire_owned() else { + return Ok(ResidualFtsAdmission::Deferred( + "residual build resources are busy", + )); + }; + // Claim and permit share one lifetime. Detached builder CPU output keeps + // both until its heavy state is destroyed, so cancellation cannot admit a + // retry of the same group under an independent resident counter. + let build_lifetime = Arc::new(ResidualFtsBuildLifetime { + _permit: build_permit, + _claim: build_claim, + }); + + // A different append snapshot may have committed one of the missing keys + // immediately before this claim. Re-probe while the overlapping-fragment + // claim is held so the baseline and local build list are exact. + let mut missing = Vec::new(); + for ((key, fragment), entry) in spec + .keys + .iter() + .zip(spec.fragments.iter()) + .zip(cached_entries.iter_mut()) + { + if entry.is_none() { + *entry = cache.get_with_key(key).await; + } + if entry.is_none() { + missing.push((key.clone(), fragment.clone())); + } + } + + let baseline_resident_bytes = + checked_resident_sum(cached_entries.iter().flatten().map(|entry| { + let CachedResidualFtsEntry::Ready(segment) = entry.as_ref(); + segment.resident_bytes() + })); + let Some(baseline_resident_bytes) = baseline_resident_bytes else { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working-set byte count overflowed", + )); + }; + if baseline_resident_bytes > MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds query-time build byte budget", + )); + } + if missing.is_empty() { + let retained = cached_entries.into_iter().flatten().collect::>(); + return Ok(ResidualFtsAdmission::Eligible(loaded_segments( + retained, + spec.fragment_bitmap(), + ))); + } + + let committed_index = dataset + .open_scalar_index( + spec.column.as_ref(), + &spec.committed_segment_uuid, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await?; + let committed_index = committed_index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + lance_core::Error::internal(format!( + "residual FTS source segment {} is not an inverted index", + spec.committed_segment_uuid + )) + })?; + let shared_tokenizer = committed_index.shared_tokenizer(); + // Count fixed cached resident weight before any new allocation. A retry + // after cancellation therefore cannot build the group piecemeal around the + // invocation-local upload cap. + let group_accounted_bytes = Arc::new(AtomicUsize::new(baseline_resident_bytes)); + let group_budget_exceeded = Arc::new(AtomicBool::new(false)); + + let built_entries = futures::stream::iter(missing.into_iter().map(|(key, fragment)| { + let resolved = spec.resolved.clone(); + let params = spec.params.clone(); + let shared_tokenizer = shared_tokenizer.clone(); + let build_lifetime = build_lifetime.clone(); + let group_accounted_bytes = group_accounted_bytes.clone(); + let group_budget_exceeded = group_budget_exceeded.clone(); + async move { + build_residual_segment( + dataset, + fragment, + &resolved, + params, + shared_tokenizer, + build_lifetime, + group_accounted_bytes, + group_budget_exceeded, + ) + .await + .map(|entry| (key, CachedResidualFtsEntry::Ready(entry))) + } + })) + .buffered(RESIDUAL_FTS_FRAGMENT_BUILD_CONCURRENCY) + .try_collect::>() + .await; + let built_entries = match built_entries { + Ok(entries) => entries, + Err(_) if group_budget_exceeded.load(Ordering::Relaxed) => { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds query-time build byte budget", + )); + } + Err(error) => return Err(error), + }; + + let group_resident_bytes = + built_entries + .iter() + .try_fold(baseline_resident_bytes, |total, (_, entry)| { + let CachedResidualFtsEntry::Ready(segment) = entry; + total.checked_add(segment.resident_bytes()) + }); + if !matches!(group_resident_bytes, Some(bytes) if bytes <= MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES) + { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds query-time build byte budget", + )); + } + + // Publish only after every missing segment was built and the aggregate + // fixed weight passed admission. Cancellation during a retry cannot use + // partial prior builds to escape the group cap because those entries are + // included in the next baseline above. + for (key, entry) in built_entries { + cache.insert_with_key(&key, Arc::new(entry)).await; + } + + // Verify the whole working set after all insertions. Per-entry success is + // insufficient for a sharded cache: later siblings may already have + // evicted an earlier one from the same shard. + let mut retained = Vec::with_capacity(spec.keys.len()); + for key in spec.keys.iter() { + let Some(entry) = cache.get_with_key(key).await else { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual working set exceeds cache admission capacity", + )); + }; + retained.push(entry); + } + if !matches!( + group_cache.get_with_key(&spec.group_key).await.as_deref(), + Some(ResidualFtsGroupState::Seen) + ) { + group_cache + .insert_with_key(&spec.group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + return Ok(ResidualFtsAdmission::Rejected( + "residual metadata cache could not retain the working-set marker", + )); + } + + Ok(ResidualFtsAdmission::Eligible(loaded_segments( + retained, + spec.fragment_bitmap(), + ))) +} + +fn checked_resident_sum(bytes: impl IntoIterator) -> Option { + bytes + .into_iter() + .try_fold(0_usize, |total, bytes| total.checked_add(bytes)) +} + +fn loaded_segments( + entries: Vec>, + fragment_bitmap: roaring::RoaringBitmap, +) -> LoadedResidualFtsSegments { + let mut loaded = LoadedResidualFtsSegments { + segments: Vec::with_capacity(entries.len()), + fragment_bitmap, + }; + for entry in entries { + let CachedResidualFtsEntry::Ready(segment) = entry.as_ref(); + loaded.segments.push(segment.index()); + } + loaded +} + +#[cfg(test)] +mod tests { + use std::{sync::atomic::AtomicUsize, time::Duration}; + + use lance_core::cache::{CacheNamespace, KeyBuilder, recommended_cache_shards}; + + use super::*; + use crate::dataset::DEFAULT_INDEX_CACHE_SIZE; + + fn internal_key(key: &ResidualFtsFragmentKey) -> [u8; 16] { + let mut builder = KeyBuilder::new( + CacheNamespace::root(), + ResidualFtsFragmentKey::stable_type_id(), + ResidualFtsFragmentKey::schema(), + ); + key.write_key(&mut builder); + builder.finish().into_bytes() + } + + fn key() -> ResidualFtsFragmentKey { + ResidualFtsFragmentKey { + store_identity: Arc::from("memory"), + index_uuid: Uuid::nil(), + index_version: 3, + fragment_id: 7, + fragment_fingerprint: Arc::from(&b"fragment-a"[..]), + field_id: 2, + canonical_path: Arc::from("body"), + field_type: Arc::from("Utf8"), + params_fingerprint: Arc::from(&b"params-a"[..]), + } + } + + fn group_key_with_uuid(uuid: u128) -> ResidualFtsGroupKey { + let mut member = key(); + member.index_uuid = Uuid::from_u128(uuid); + ResidualFtsGroupKey { + store_identity: member.store_identity.clone(), + index_uuid: member.index_uuid, + index_version: member.index_version, + members: Arc::from([member]), + } + } + + async fn test_build_lifetime( + uuid: u128, + ) -> ( + Arc, + Arc, + ResidualFtsGroupKey, + Vec, + ) { + let group = group_key_with_uuid(uuid); + let fragments = group.members.to_vec(); + let claim = ResidualFtsBuildClaim::try_acquire(&group, &fragments) + .expect("build should own the working set"); + let resources = Arc::new(Semaphore::new(1)); + let permit = resources + .clone() + .acquire_owned() + .await + .expect("build resource semaphore closed"); + ( + Arc::new(ResidualFtsBuildLifetime { + _permit: permit, + _claim: claim, + }), + resources, + group, + fragments, + ) + } + + struct PressureKey(&'static str); + + impl CacheKey for PressureKey { + type ValueType = Vec; + + fn key(&self) -> Cow<'_, str> { + Cow::Borrowed(self.0) + } + + fn type_name() -> &'static str { + "ResidualFtsPressureValue" + } + } + + #[derive(Debug)] + struct DropObservedUpload { + accounted_bytes: Arc, + accounted_at_drop: Arc, + } + + #[derive(Debug)] + struct BlockingAbortUpload { + abort_started: Arc, + release_abort: Arc, + } + + #[async_trait] + impl MultipartUpload for BlockingAbortUpload { + fn put_part(&mut self, _payload: PutPayload) -> UploadPart { + Box::pin(async { Ok(()) }) + } + + async fn complete(&mut self) -> OsResult { + Ok(PutResult { + e_tag: None, + version: None, + }) + } + + async fn abort(&mut self) -> OsResult<()> { + self.abort_started.notify_one(); + self.release_abort.notified().await; + Ok(()) + } + } + + impl Drop for DropObservedUpload { + fn drop(&mut self) { + self.accounted_at_drop.store( + self.accounted_bytes.load(Ordering::SeqCst), + Ordering::SeqCst, + ); + } + } + + #[async_trait] + impl MultipartUpload for DropObservedUpload { + fn put_part(&mut self, _payload: PutPayload) -> UploadPart { + Box::pin(async { Ok(()) }) + } + + async fn complete(&mut self) -> OsResult { + Ok(PutResult { + e_tag: None, + version: None, + }) + } + + async fn abort(&mut self) -> OsResult<()> { + Ok(()) + } + } + + #[test] + fn residual_key_reuses_identical_fragment() { + assert_eq!(internal_key(&key()), internal_key(&key())); + } + + #[test] + fn cancelled_group_claim_can_be_retried() { + let group = group_key_with_uuid(1); + let fragments = group.members.to_vec(); + let claim = ResidualFtsBuildClaim::try_acquire(&group, &fragments) + .expect("first build should own the working set"); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_none()); + + drop(claim); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_some()); + } + + #[tokio::test] + async fn concurrent_group_builds_have_one_owner() { + let start = Arc::new(tokio::sync::Barrier::new(3)); + let attempted = Arc::new(tokio::sync::Barrier::new(3)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let start = start.clone(); + let attempted = attempted.clone(); + tasks.push(tokio::spawn(async move { + let group = group_key_with_uuid(2); + let fragments = group.members.to_vec(); + start.wait().await; + let claim = ResidualFtsBuildClaim::try_acquire(&group, &fragments); + attempted.wait().await; + claim.is_some() + })); + } + start.wait().await; + attempted.wait().await; + let acquired = futures::future::join_all(tasks) + .await + .into_iter() + .filter(|result| *result.as_ref().expect("claim task panicked")) + .count(); + assert_eq!(acquired, 1); + } + + #[tokio::test] + async fn detached_work_retains_group_claim_and_permit() { + let (lifetime, resources, group, fragments) = test_build_lifetime(5).await; + let detached_work = lifetime.clone(); + drop(lifetime); + + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_none()); + assert!(resources.clone().try_acquire_owned().is_err()); + + drop(detached_work); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_some()); + assert!(resources.try_acquire_owned().is_ok()); + } + + #[test] + fn overlapping_snapshots_cannot_build_different_missing_members() { + let first_group = group_key_with_uuid(3); + let mut second_member = first_group.members[0].clone(); + second_member.fragment_id += 1; + second_member.fragment_fingerprint = Arc::from(&b"fragment-b"[..]); + let second_group = ResidualFtsGroupKey { + store_identity: first_group.store_identity.clone(), + index_uuid: first_group.index_uuid, + index_version: first_group.index_version, + members: Arc::from([first_group.members[0].clone(), second_member]), + }; + + let _first_claim = + ResidualFtsBuildClaim::try_acquire(&first_group, first_group.members.as_ref()) + .expect("first snapshot should own its working set"); + assert!( + ResidualFtsBuildClaim::try_acquire(&second_group, second_group.members.as_ref()) + .is_none(), + "overlapping append snapshots must not build under independent budgets" + ); + } + + #[test] + fn cached_group_baseline_consumes_shared_build_budget() { + let uploaded = AtomicUsize::new(MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES - 4); + let exceeded = AtomicBool::new(false); + reserve_upload_bytes( + &uploaded, + &exceeded, + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + 4, + ) + .unwrap(); + assert!( + reserve_upload_bytes( + &uploaded, + &exceeded, + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + 1, + ) + .is_err() + ); + assert!(exceeded.load(Ordering::Relaxed)); + } + + #[test] + fn concurrent_builder_working_sets_consume_resident_budget() { + let accounted_bytes = Arc::new(AtomicUsize::new(0)); + let exceeded = Arc::new(AtomicBool::new(false)); + let mut reservations = Vec::new(); + for _ in 0..RESIDUAL_FTS_FRAGMENT_BUILD_CONCURRENCY { + reserve_upload_bytes( + accounted_bytes.as_ref(), + exceeded.as_ref(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES, + ) + .unwrap(); + reservations.push(TransientBudgetReservation { + uploaded_bytes: accounted_bytes.clone(), + bytes: MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES, + }); + } + assert_eq!( + accounted_bytes.load(Ordering::Relaxed), + MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES * RESIDUAL_FTS_FRAGMENT_BUILD_CONCURRENCY + ); + let remaining = + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES - accounted_bytes.load(Ordering::Relaxed); + reserve_upload_bytes( + accounted_bytes.as_ref(), + exceeded.as_ref(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + remaining, + ) + .unwrap(); + assert!( + reserve_upload_bytes( + accounted_bytes.as_ref(), + exceeded.as_ref(), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES, + 1, + ) + .is_err() + ); + + drop(reservations); + assert_eq!( + accounted_bytes.load(Ordering::Relaxed), + MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES + - MAX_RESIDUAL_FTS_BUILD_MEMORY_BYTES * RESIDUAL_FTS_FRAGMENT_BUILD_CONCURRENCY + ); + } + + #[test] + fn separately_warmed_segments_still_obey_combined_resident_cap() { + let first_snapshot = checked_resident_sum([MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES / 2]); + let second_snapshot = checked_resident_sum([MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES / 2 + 1]); + assert!( + matches!(first_snapshot, Some(bytes) if bytes <= MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES) + ); + assert!( + matches!(second_snapshot, Some(bytes) if bytes <= MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES) + ); + + let combined = checked_resident_sum([ + first_snapshot.expect("first snapshot byte count overflowed"), + second_snapshot.expect("second snapshot byte count overflowed"), + ]); + assert!(!matches!(combined, Some(bytes) if bytes <= MAX_RESIDUAL_FTS_GROUP_RESIDENT_BYTES)); + } + + #[test] + fn default_cache_admits_large_residual_fixture() { + const FIXTURE_ROWS: usize = 100_000; + const FIXTURE_SERIALIZED_BYTES: usize = 396 << 20; + + let shards = recommended_cache_shards(DEFAULT_INDEX_CACHE_SIZE); + assert!(FIXTURE_ROWS <= MAX_RESIDUAL_FTS_ROWS); + assert!(FIXTURE_SERIALIZED_BYTES < DEFAULT_INDEX_CACHE_SIZE / shards); + } + + #[test] + fn residual_key_invalidates_every_semantic_identity() { + let original = key(); + let original_key = internal_key(&original); + let variants = [ + ResidualFtsFragmentKey { + store_identity: Arc::from("memory-other"), + ..original.clone() + }, + ResidualFtsFragmentKey { + index_uuid: Uuid::from_u128(1), + ..original.clone() + }, + ResidualFtsFragmentKey { + index_version: 4, + ..original.clone() + }, + ResidualFtsFragmentKey { + fragment_id: 8, + ..original.clone() + }, + ResidualFtsFragmentKey { + fragment_fingerprint: Arc::from(&b"fragment-b"[..]), + ..original.clone() + }, + ResidualFtsFragmentKey { + field_id: 3, + ..original.clone() + }, + ResidualFtsFragmentKey { + canonical_path: Arc::from("renamed"), + ..original.clone() + }, + ResidualFtsFragmentKey { + field_type: Arc::from("LargeUtf8"), + ..original.clone() + }, + ResidualFtsFragmentKey { + params_fingerprint: Arc::from(&b"params-b"[..]), + ..original + }, + ]; + for variant in variants { + assert_ne!(original_key, internal_key(&variant)); + } + } + + #[tokio::test] + async fn budgeted_memory_store_rejects_growth_before_allocation() { + use object_store::ObjectStoreExt; + + let store = BudgetedMemoryStore::new(4); + let error = store + .put(&Path::from("too-large"), PutPayload::from_static(b"12345")) + .await + .unwrap_err(); + assert!(error.to_string().contains("4 byte resident budget")); + assert!(store.list(None).next().await.is_none()); + } + + #[tokio::test] + async fn multi_fragment_builds_share_one_upload_budget() { + use object_store::ObjectStoreExt; + + let uploaded = Arc::new(AtomicUsize::new(0)); + let exceeded = Arc::new(AtomicBool::new(false)); + let first = BudgetedMemoryStore::with_counter(uploaded.clone(), exceeded.clone(), 7); + let second = BudgetedMemoryStore::with_counter(uploaded, exceeded.clone(), 7); + first + .put(&Path::from("first"), PutPayload::from_static(b"1234")) + .await + .unwrap(); + let error = second + .put(&Path::from("second"), PutPayload::from_static(b"5678")) + .await + .unwrap_err(); + assert!(error.to_string().contains("7 byte resident budget")); + assert!(exceeded.load(Ordering::Relaxed)); + assert!(second.list(None).next().await.is_none()); + } + + #[tokio::test] + async fn multipart_completion_reserves_full_copy_peak() { + use object_store::ObjectStoreExt; + + let store = BudgetedMemoryStore::new(7); + let destination = Path::from("multipart"); + let mut upload = store.put_multipart(&destination).await.unwrap(); + upload + .put_part(PutPayload::from_static(b"1234")) + .await + .unwrap(); + + let error = upload.complete().await.unwrap_err(); + assert!(error.to_string().contains("7 byte resident budget")); + assert!(store.head(&destination).await.is_err()); + } + + #[tokio::test] + async fn multipart_success_drops_parts_before_copy_reservation() { + let (build_lifetime, resources, group, fragments) = test_build_lifetime(7).await; + let accounted_bytes = Arc::new(AtomicUsize::new(0)); + let exceeded = Arc::new(AtomicBool::new(false)); + let accounted_at_drop = Arc::new(AtomicUsize::new(0)); + let mut upload = BudgetedMultipartUpload { + inner: Some(Box::new(DropObservedUpload { + accounted_bytes: accounted_bytes.clone(), + accounted_at_drop: accounted_at_drop.clone(), + })), + uploaded_bytes: accounted_bytes.clone(), + budget_exceeded: exceeded, + max_uploaded_bytes: 8, + part_bytes: 0, + build_lifetime: Some(build_lifetime.clone()), + }; + drop(build_lifetime); + + upload + .put_part(PutPayload::from_static(b"1234")) + .await + .unwrap(); + upload.complete().await.unwrap(); + assert!(upload.build_lifetime.is_none()); + assert_eq!(accounted_at_drop.load(Ordering::SeqCst), 8); + assert_eq!(accounted_bytes.load(Ordering::SeqCst), 4); + assert!(upload.complete().await.is_err()); + assert!( + upload + .put_part(PutPayload::from_static(b"5678")) + .await + .is_err() + ); + assert_eq!(accounted_bytes.load(Ordering::SeqCst), 4); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_some()); + assert!(resources.try_acquire_owned().is_ok()); + } + + #[tokio::test] + async fn detached_multipart_abort_retains_group_admission() { + let (build_lifetime, resources, group, fragments) = test_build_lifetime(6).await; + let abort_started = Arc::new(tokio::sync::Notify::new()); + let release_abort = Arc::new(tokio::sync::Notify::new()); + let mut upload = BudgetedMultipartUpload { + inner: Some(Box::new(BlockingAbortUpload { + abort_started: abort_started.clone(), + release_abort: release_abort.clone(), + })), + uploaded_bytes: Arc::new(AtomicUsize::new(0)), + budget_exceeded: Arc::new(AtomicBool::new(false)), + max_uploaded_bytes: 8, + part_bytes: 0, + build_lifetime: Some(build_lifetime.clone()), + }; + drop(build_lifetime); + + let abort_task = tokio::spawn(async move { + let result = upload.abort().await; + (result, upload) + }); + tokio::time::timeout(Duration::from_secs(5), abort_started.notified()) + .await + .expect("detached multipart abort did not start"); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_none()); + assert!(resources.clone().try_acquire_owned().is_err()); + + release_abort.notify_one(); + let (result, upload) = abort_task.await.expect("abort task panicked"); + result.unwrap(); + assert!(upload.build_lifetime.is_none()); + assert!(ResidualFtsBuildClaim::try_acquire(&group, &fragments).is_some()); + assert!(resources.try_acquire_owned().is_ok()); + } + + #[tokio::test] + async fn object_copy_cannot_bypass_upload_budget() { + use object_store::ObjectStoreExt; + + let store = BudgetedMemoryStore::new(7); + let source = Path::from("source"); + let destination = Path::from("destination"); + store + .put(&source, PutPayload::from_static(b"1234")) + .await + .unwrap(); + + let error = store.copy(&source, &destination).await.unwrap_err(); + assert!(error.to_string().contains("7 byte resident budget")); + assert!(store.head(&source).await.is_ok()); + assert!(store.head(&destination).await.is_err()); + } + + #[tokio::test] + async fn object_rename_is_conservatively_charged() { + use object_store::ObjectStoreExt; + + let store = BudgetedMemoryStore::new(7); + let source = Path::from("source"); + let destination = Path::from("destination"); + store + .put(&source, PutPayload::from_static(b"1234")) + .await + .unwrap(); + + let error = store.rename(&source, &destination).await.unwrap_err(); + assert!(error.to_string().contains("7 byte resident budget")); + assert!(store.head(&source).await.is_ok()); + assert!(store.head(&destination).await.is_err()); + } + + #[tokio::test] + async fn aggregate_rejection_is_independent_of_posting_eviction() { + let posting_cache = LanceCache::with_capacity(192); + let group_cache = LanceCache::with_capacity(4096); + let group_key = group_key_with_uuid(4); + group_cache + .insert_with_key(&group_key, Arc::new(ResidualFtsGroupState::Rejected)) + .await; + + // Two individually admissible values exceed aggregate capacity and + // evict one another, but cannot evict the independent negative state. + posting_cache + .insert_with_key(&PressureKey("a"), Arc::new(vec![0; 128])) + .await; + posting_cache + .insert_with_key(&PressureKey("b"), Arc::new(vec![0; 128])) + .await; + assert!(posting_cache.size_bytes().await <= 192); + let loader_runs = AtomicUsize::new(0); + if !matches!( + group_cache.get_with_key(&group_key).await.as_deref(), + Some(ResidualFtsGroupState::Rejected) + ) { + loader_runs.fetch_add(1, Ordering::Relaxed); + } + assert_eq!(loader_runs.load(Ordering::Relaxed), 0); + } +} diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 7c6e68ccc00..9e6e5941cbd 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -40,10 +40,13 @@ use lance_select::RowAddrMask; use lance_table::format::IndexMetadata; use super::PreFilterSource; -use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; +use super::utils::{ + IndexMetrics, PreFilterMasks, build_prefilter, build_prefilter_with_fragment_bitmap, +}; use crate::dataset::mem_wal::index::QueryLocalFtsIndex; use crate::index::scalar::inverted::{ ResolvedFtsField, fts_document_schema, load_segment_details, load_segments, + residual::{ResidualFtsAdmission, ResidualFtsSpec, load_residual_fts_segments}, transform_fts_document_stream, }; use crate::{Dataset, index::DatasetIndexInternalExt}; @@ -673,6 +676,9 @@ pub struct CompoundQueryExec { /// subset. prepared_match: Option>, segment_selection: FtsSegmentSelection, + /// Query-time immutable posting segments for append-only fragments. + residual_indices: Arc<[Arc]>, + residual_fragment_bitmap: roaring::RoaringBitmap, /// Caller-supplied row-address mask, intersected into the prefilter so the /// compound scorer ranks only surviving rows (see /// [`MatchQueryExec::with_external_mask`]). @@ -730,6 +736,8 @@ impl CompoundQueryExec { base_scorer: None, prepared_match: None, segment_selection, + residual_indices: Arc::from([]), + residual_fragment_bitmap: roaring::RoaringBitmap::new(), external_mask: None, properties: Arc::new(PlanProperties::new( EquivalenceProperties::new(FTS_SCHEMA.clone()), @@ -769,6 +777,22 @@ impl CompoundQueryExec { self } + /// Add cached immutable posting segments covering append-only fragments. + pub(crate) fn with_cached_residual_segments( + mut self, + indices: Vec>, + fragment_bitmap: roaring::RoaringBitmap, + ) -> Self { + self.residual_indices = Arc::from(indices); + self.residual_fragment_bitmap = fragment_bitmap; + self + } + + fn with_metrics(mut self, metrics: ExecutionPlanMetricsSet) -> Self { + self.metrics = metrics; + self + } + pub fn dataset(&self) -> &Arc { &self.dataset } @@ -815,6 +839,162 @@ async fn index_query_local_residual_batch( .await } +/// Lazily activates immutable residual postings after the working set has +/// demonstrated reuse. Cold, busy, rejected, and failed builds execute the +/// supplied exact fallback instead. +#[derive(Debug)] +pub(crate) struct CachedResidualCompoundQueryExec { + dataset: Arc, + query: FtsQuery, + params: FtsSearchParams, + prefilter_source: PreFilterSource, + segments: Arc<[IndexMetadata]>, + residual_spec: ResidualFtsSpec, + fallback: Arc, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl CachedResidualCompoundQueryExec { + pub(crate) fn new( + dataset: Arc, + query: FtsQuery, + params: FtsSearchParams, + prefilter_source: PreFilterSource, + segments: Vec, + residual_spec: ResidualFtsSpec, + fallback: Arc, + ) -> Self { + Self { + dataset, + query, + params, + prefilter_source, + segments: Arc::from(segments), + residual_spec, + fallback, + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(FTS_SCHEMA.clone()), + Partitioning::RoundRobinBatch(1), + EmissionType::Final, + Boundedness::Bounded, + )), + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for CachedResidualCompoundQueryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "CachedResidualCompoundFtsScorer: state=lazy, query={}", + self.query + ), + DisplayFormatType::TreeRender => write!( + f, + "CachedResidualCompoundFtsScorer\nstate=lazy\nquery={}", + self.query + ), + } + } +} + +impl ExecutionPlan for CachedResidualCompoundQueryExec { + fn name(&self) -> &str { + "CachedResidualCompoundQueryExec" + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.fallback] + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "cached residual compound FTS expected one fallback child, got {}", + children.len() + ))); + } + let fallback = children.pop().ok_or_else(|| { + DataFusionError::Internal( + "cached residual compound FTS lost its fallback child".to_string(), + ) + })?; + Ok(Arc::new(Self::new( + self.dataset.clone(), + self.query.clone(), + self.params.clone(), + self.prefilter_source.clone(), + self.segments.to_vec(), + self.residual_spec.clone(), + fallback, + ))) + } + + #[instrument(name = "cached_residual_compound_fts_exec", level = "debug", skip_all)] + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let dataset = self.dataset.clone(); + let query = self.query.clone(); + let params = self.params.clone(); + let prefilter_source = self.prefilter_source.clone(); + let segments = self.segments.clone(); + let residual_spec = self.residual_spec.clone(); + let fallback = self.fallback.clone(); + let metrics_set = self.metrics.clone(); + let context_for_choice = context.clone(); + + let chosen = async move { + match load_residual_fts_segments(&dataset, &residual_spec).await { + Ok(ResidualFtsAdmission::Eligible(loaded)) => CompoundQueryExec::new_with_segments( + dataset, + query, + params, + prefilter_source, + segments.to_vec(), + ) + .with_cached_residual_segments(loaded.segments, loaded.fragment_bitmap) + .with_metrics(metrics_set) + .execute(partition, context_for_choice), + Ok(ResidualFtsAdmission::Deferred(reason)) + | Ok(ResidualFtsAdmission::Rejected(reason)) => { + tracing::debug!(reason, "using exact residual FTS fallback"); + fallback.execute(partition, context_for_choice) + } + Err(error) => { + tracing::warn!(error = %error, "cached residual FTS build failed; using exact fallback"); + fallback.execute(partition, context_for_choice) + } + } + }; + let stream = stream::once(chosen).try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream.boxed(), + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn properties(&self) -> &Arc { + &self.properties + } +} + /// Exact compound FTS over committed postings plus an append-only residual /// scan. The residual documents are tokenized once into query-local postings, /// rather than once for every compound leaf. @@ -1151,11 +1331,29 @@ impl DisplayAs for CompoundQueryExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "CompoundFtsScorer: query={}", self.query)?; + if self.residual_indices.is_empty() { + write!(f, "CompoundFtsScorer: query={}", self.query)?; + } else { + write!( + f, + "CachedResidualCompoundFtsScorer: query={}, residual_segments={}", + self.query, + self.residual_indices.len() + )?; + } fmt_tokenized_compound_query(&self.tokenized_query, ", ", f) } DisplayFormatType::TreeRender => { - write!(f, "CompoundFtsScorer\nquery={}", self.query)?; + if self.residual_indices.is_empty() { + write!(f, "CompoundFtsScorer\nquery={}", self.query)?; + } else { + write!( + f, + "CachedResidualCompoundFtsScorer\nquery={}\nresidual_segments={}", + self.query, + self.residual_indices.len() + )?; + } fmt_tokenized_compound_query(&self.tokenized_query, "\n", f) } } @@ -1207,6 +1405,8 @@ impl ExecutionPlan for CompoundQueryExec { base_scorer: self.base_scorer.clone(), prepared_match: self.prepared_match.clone(), segment_selection: self.segment_selection.clone(), + residual_indices: self.residual_indices.clone(), + residual_fragment_bitmap: self.residual_fragment_bitmap.clone(), external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), @@ -1227,6 +1427,8 @@ impl ExecutionPlan for CompoundQueryExec { let preset_base_scorer = self.base_scorer.clone(); let preset_prepared_match = self.prepared_match.clone(); let segment_selection = self.segment_selection.clone(); + let residual_indices = self.residual_indices.clone(); + let residual_fragment_bitmap = self.residual_fragment_bitmap.clone(); let external_mask = self.external_mask.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); @@ -1266,24 +1468,52 @@ impl ExecutionPlan for CompoundQueryExec { } else { true }; - let _details = load_segment_details(&dataset, column, &segments).await?; - let indices = + if !segments.is_empty() { + let _details = load_segment_details(&dataset, column, &segments).await?; + } else if residual_indices.is_empty() { + return Err(DataFusionError::Execution(format!( + "FTS index for column {column} has no searchable segments" + ))); + } + let mut indices = open_fts_segments(&dataset, column, &segments, &metrics.index_metrics).await?; + indices.extend(residual_indices.iter().cloned()); if let Some(first_index) = indices.first() { tokenized_query .get_or_init(|| tokenize_compound_query(&query, first_index.as_ref())); } - let mut prefilter = build_prefilter( - context, - partition, - &prefilter_source, - dataset, - &segments, - PreFilterMasks { - overlay_block: None, + let mut prefilter = if residual_fragment_bitmap.is_empty() { + build_prefilter( + context, + partition, + &prefilter_source, + dataset, + &segments, + PreFilterMasks { + overlay_block: None, + external_mask, + }, + )? + } else { + let mut coverage = residual_fragment_bitmap; + for segment in segments.iter() { + let Some(bitmap) = &segment.fragment_bitmap else { + return Err(DataFusionError::Execution( + "cached residual FTS requires exact committed fragment coverage" + .to_string(), + )); + }; + coverage |= bitmap; + } + build_prefilter_with_fragment_bitmap( + context, + partition, + &prefilter_source, + dataset, + coverage, external_mask, - }, - )?; + )? + }; let deleted_fragments = indices .iter() diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index fb36f612742..7b59ef58f64 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -10,6 +10,7 @@ use lance_index::metrics::MetricsCollector; use lance_io::scheduler::{IoStats, ScanScheduler, ScanStats}; use lance_table::format::IndexMetadata; use pin_project::pin_project; +use roaring::RoaringBitmap; use std::collections::HashMap; use std::future::Future; use std::pin::Pin; @@ -437,6 +438,79 @@ pub(crate) fn build_prefilter( Ok(Arc::new(prefilter)) } +pub(crate) fn build_prefilter_with_fragment_bitmap( + context: Arc, + partition: usize, + prefilter_source: &PreFilterSource, + ds: Arc, + fragments: RoaringBitmap, + external_mask: Option>, +) -> Result> { + let mut shared_filter = None; + let prefilter_loader = match prefilter_source { + PreFilterSource::FilteredRowIds(src_node) => { + if let Some(shared) = src_node.downcast_ref::() { + shared_filter = Some(shared_prefilter_future( + shared.materialization.clone(), + shared.source.clone(), + false, + context, + partition, + )); + None + } else { + let stream = src_node.execute(partition, context)?; + Some(Box::new(FilteredRowIdsToPrefilter(stream)) as Box) + } + } + PreFilterSource::ScalarIndexQuery(src_node) => { + if let Some(shared) = src_node.downcast_ref::() { + shared_filter = Some(shared_prefilter_future( + shared.materialization.clone(), + shared.source.clone(), + true, + context, + partition, + )); + None + } else { + let stream = src_node.execute(partition, context)?; + Some(Box::new(SelectionVectorToPrefilter(stream)) as Box) + } + } + PreFilterSource::None => None, + }; + if let Some(shared_filter) = shared_filter { + let shared_filter = match external_mask { + Some(mask) => async move { + Ok(Arc::new( + mask.as_ref().clone() & shared_filter.await?.as_ref().clone(), + )) + } + .boxed(), + None => shared_filter, + }; + return Ok(Arc::new( + DatasetPreFilter::new_with_fragment_bitmap_and_filter_future( + ds, + fragments, + Some(shared_filter), + ), + )); + } + let prefilter_loader = match external_mask { + Some(mask) => { + Some(Box::new(MaskAndLoader::new(mask, prefilter_loader)) as Box) + } + None => prefilter_loader, + }; + Ok(Arc::new(DatasetPreFilter::new_with_fragment_bitmap( + ds, + fragments, + prefilter_loader, + ))) +} + // Utility to convert an input (containing row ids) into a prefilter pub(crate) struct FilteredRowIdsToPrefilter(pub SendableRecordBatchStream);