From cc50dfd4e417f0374c72db84dbdfe96750e45164 Mon Sep 17 00:00:00 2001 From: Ecthlion_zyy <48782306+Ecthlion@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:52:08 +0800 Subject: [PATCH 1/2] fix: skip empty string docs in FTS stats --- .../src/scalar/inverted/builder.rs | 119 ++++++++++++++++-- 1 file changed, 110 insertions(+), 9 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index e3c2569f774..577e6e8767f 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1360,7 +1360,7 @@ impl IndexWorker { .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); for (doc, row_id) in docs { - self.process_document(row_id, DocumentSource::Text(doc), false) + self.process_document(row_id, DocumentSource::Text(doc)) .await?; } } @@ -1404,7 +1404,7 @@ impl IndexWorker { continue; }; - self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), true) + self.process_document(*row_id, DocumentSource::StringList(doc.as_ref())) .await?; } @@ -1430,12 +1430,7 @@ impl IndexWorker { doc } - async fn process_document( - &mut self, - row_id: u64, - document: DocumentSource<'_>, - skip_empty_document: bool, - ) -> Result<()> { + async fn process_document(&mut self, row_id: u64, document: DocumentSource<'_>) -> Result<()> { let with_position = self.has_position(); let builder_was_empty = self.builder.docs.is_empty(); let old_temporary_memory_size = self.temporary_memory_size(); @@ -1541,7 +1536,7 @@ impl IndexWorker { self.builder.tokens.memory_size() as u64, ); - if skip_empty_document && token_num == 0 { + if token_num == 0 { self.last_token_count = 0; self.trim_temporary_buffers(); self.adjust_tracked_memory_size( @@ -2189,6 +2184,7 @@ pub fn document_input( #[cfg(test)] mod tests { use super::*; + use crate::Index; use crate::metrics::NoOpMetricsCollector; use crate::progress::IndexBuildProgress; use crate::scalar::{IndexFile, IndexReader, IndexWriter, ScalarIndex}; @@ -2223,6 +2219,17 @@ mod tests { RecordBatch::try_new(schema, vec![docs, row_ids]).unwrap() } + fn make_doc_batch_from_docs(docs: Vec>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("doc", DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let num_rows = docs.len(); + let docs = Arc::new(StringArray::from(docs)); + let row_ids = Arc::new(UInt64Array::from_iter_values(0..num_rows as u64)); + RecordBatch::try_new(schema, vec![docs, row_ids]).unwrap() + } + struct FailingListObjectStore { inner: InMemory, } @@ -3456,6 +3463,100 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_empty_string_documents_are_skipped_in_corpus_stats() -> Result<()> { + 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_from_docs(vec![Some(""), Some(" "), None, Some("hello")]); + let stream = RecordBatchStreamAdapter::new(batch.schema(), stream::iter(vec![Ok(batch)])); + let params = + InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) + .with_position(false) + .remove_stop_words(false) + .stem(false) + .max_token_length(None) + .num_workers(1); + + let mut builder = InvertedIndexBuilder::new(params); + builder + .update(Box::pin(stream), store.as_ref(), None) + .await?; + + let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; + let stats = index.bm25_stats_for_terms(&["hello".to_string()]).await?; + assert_eq!(stats, (1, 1, vec![1])); + + Ok(()) + } + + #[tokio::test] + async fn test_all_empty_string_documents_build_empty_index() -> Result<()> { + 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_from_docs(vec![Some(""), Some(" "), None]); + let stream = RecordBatchStreamAdapter::new(batch.schema(), stream::iter(vec![Ok(batch)])); + let params = + InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) + .with_position(false) + .remove_stop_words(false) + .stem(false) + .max_token_length(None) + .num_workers(1); + + let mut builder = InvertedIndexBuilder::new(params); + builder + .update(Box::pin(stream), store.as_ref(), None) + .await?; + + let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; + assert!(index.partitions.is_empty()); + let statistics = index.statistics()?; + assert_eq!(statistics["num_tokens"], 0); + assert_eq!(statistics["num_docs"], 0); + + Ok(()) + } + + #[tokio::test] + async fn test_all_empty_string_documents_do_not_create_tail_partition() -> Result<()> { + let tokenizer = InvertedIndexParams::default().build()?; + let store = Arc::new(CountingStore::new()); + let id_alloc = Arc::new(AtomicU64::new(0)); + let mut worker = IndexWorker::new( + tokenizer, + store, + id_alloc, + IndexWorkerConfig { + with_position: false, + format_version: InvertedListFormatVersion::V1, + fragment_mask: None, + token_set_format: TokenSetFormat::default(), + worker_memory_limit_bytes: u64::MAX, + }, + ) + .await?; + + worker + .process_batch(make_doc_batch_from_docs(vec![Some(""), Some(" "), None])) + .await?; + let output = worker.finish().await?; + + assert!(output.partitions.is_empty()); + assert!(output.tail_partition.is_none()); + + Ok(()) + } + lance_testing::define_stage_event_progress!(RecordingProgress, IndexBuildProgress, Result<()>); #[derive(Debug, Default)] From acbe39583fdef917ec956f8feae483841b1d74d0 Mon Sep 17 00:00:00 2001 From: Ecthlion_zyy <48782306+Ecthlion@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:56:18 +0800 Subject: [PATCH 2/2] fix: align zero-token policy across FTS paths --- .../src/scalar/inverted/builder.rs | 40 +++++- rust/lance-index/src/scalar/inverted/index.rs | 65 +++++++-- rust/lance/src/dataset/mem_wal/index/fts.rs | 129 +++++++++++++++--- 3 files changed, 199 insertions(+), 35 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 0955aeba78e..75748ae4336 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1480,7 +1480,7 @@ impl IndexWorker { new_posting_memory_size as i64 - old_posting_memory_size as i64; } Self::apply_delta(&mut self.memory_size, posting_memory_delta); - } else if token_num > 0 { + } else { self.token_ids.sort_unstable(); let mut iter = self.token_ids.iter(); let mut current = *iter.next().unwrap(); @@ -2081,6 +2081,7 @@ mod tests { use crate::Index; use crate::metrics::NoOpMetricsCollector; use crate::progress::IndexBuildProgress; + use crate::scalar::inverted::{MemBM25Scorer, Scorer}; use crate::scalar::{IndexFile, IndexReader, IndexWriter, ScalarIndex}; use arrow_array::{RecordBatch, StringArray, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; @@ -3300,7 +3301,7 @@ mod tests { } #[tokio::test] - async fn test_empty_string_documents_are_skipped_in_corpus_stats() -> Result<()> { + async fn test_zero_token_string_documents_are_skipped_in_corpus_stats() -> Result<()> { let index_dir = TempDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -3308,14 +3309,21 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - let batch = make_doc_batch_from_docs(vec![Some(""), Some(" "), None, Some("hello")]); + let batch = make_doc_batch_from_docs(vec![ + Some(""), + Some(" "), + Some("the"), + Some("overlength"), + None, + Some("hello"), + ]); let stream = RecordBatchStreamAdapter::new(batch.schema(), stream::iter(vec![Ok(batch)])); let params = InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) .with_position(false) - .remove_stop_words(false) + .remove_stop_words(true) .stem(false) - .max_token_length(None) + .max_token_length(Some(6)) .num_workers(1); let mut builder = InvertedIndexBuilder::new(params); @@ -3324,8 +3332,26 @@ mod tests { .await?; let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; - let stats = index.bm25_stats_for_terms(&["hello".to_string()]).await?; - assert_eq!(stats, (1, 1, vec![1])); + let (total_tokens, num_docs, token_docs) = + index.bm25_stats_for_terms(&["hello".to_string()]).await?; + assert_eq!(total_tokens, 1); + assert_eq!(num_docs, 1); + assert_eq!(token_docs, vec![1]); + + let actual_scorer = MemBM25Scorer::new( + total_tokens, + num_docs, + HashMap::from([("hello".to_string(), token_docs[0])]), + ); + let expected_scorer = MemBM25Scorer::new(1, 1, HashMap::from([("hello".to_string(), 1)])); + assert_eq!( + actual_scorer.avg_doc_length(), + expected_scorer.avg_doc_length() + ); + assert_eq!( + actual_scorer.query_weight("hello"), + expected_scorer.query_weight("hello") + ); Ok(()) } diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 26b0f0256b0..07bbf04dbd5 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -5881,12 +5881,13 @@ async fn tokenize_and_count( .extend(std::iter::repeat_n(0, query_tokens.len())); let Some(doc) = doc else { - append_counts(*row_id, 0, &temp_query_token_counts); continue; }; let all_tokens = count_text(doc, &mut temp_query_token_counts); - append_counts(*row_id, all_tokens, &temp_query_token_counts); + if all_tokens > 0 { + append_counts(*row_id, all_tokens, &temp_query_token_counts); + } } } DataType::List(_) => { @@ -6078,12 +6079,6 @@ fn flat_bm25_score( for _ in 0..counted_input.num_rows() { let num_tokens_in_doc = all_token_counts_iter.next().expect_ok()?; let row_id = row_ids_iter.next().expect_ok()?; - if num_tokens_in_doc == 0 { - for _ in query_tokens { - query_token_counts_iter.next().expect_ok()?; - } - continue; - } let doc_norm = K1 * (1.0 - B + B * num_tokens_in_doc as f32 / scorer.avg_doc_length()); let mut score = 0.0; for token in query_tokens { @@ -9469,6 +9464,60 @@ mod tests { ); } + #[tokio::test] + async fn flat_bm25_skips_zero_token_documents_from_corpus_stats() { + let schema = Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("text", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from(vec![0_u64, 1, 2, 3, 4, 5])) as ArrayRef, + Arc::new(StringArray::from(vec![ + Some(""), + Some(" "), + Some("the"), + Some("overlength"), + None, + Some("hello"), + ])) as ArrayRef, + ], + ) + .unwrap(); + let params = InvertedIndexParams::new("whitespace".to_string(), Language::English) + .remove_stop_words(true) + .stem(false) + .max_token_length(Some(6)); + let query_tokens = Arc::new(Tokens::new(vec!["hello".to_string()], DocType::Text)); + + let counted_input = tokenize_and_count( + stream::iter(vec![Ok(batch)]), + params.build().unwrap(), + query_tokens.clone(), + 1, + None, + ) + .await + .unwrap(); + + assert_eq!(counted_input.num_rows(), 1); + assert_eq!( + counted_input[ROW_ID].as_primitive::().values(), + &[5] + ); + let scorer = initialize_scorer(None, query_tokens.as_ref(), &counted_input); + let expected_scorer = MemBM25Scorer::new(1, 1, HashMap::from([("hello".to_string(), 1)])); + assert_eq!(scorer.total_tokens, 1); + assert_eq!(scorer.num_docs(), 1); + assert_eq!(scorer.num_docs_containing_token("hello"), 1); + assert_eq!(scorer.avg_doc_length(), expected_scorer.avg_doc_length()); + assert_eq!( + scorer.query_weight("hello"), + expected_scorer.query_weight("hello") + ); + } + #[tokio::test] async fn flat_bm25_search_uses_full_document_length_for_normalization() { let schema = Arc::new(Schema::new(vec![ diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 61c797db041..a5bb3198108 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -543,6 +543,7 @@ struct BatchMeta { batch_position: usize, row_offset: u64, /// `doc_lengths[i]` is the token count of the row at `row_offset + i`. + /// Zero entries preserve row-position alignment but are not documents. doc_lengths: Vec, rows: u32, } @@ -639,7 +640,7 @@ struct Snapshot { /// visible_count` for any snapshot the writer has stored (each publish /// appends one entry and bumps `visible_count`). batches: BatchLog, - /// `Σ batches[i].rows` for `i < visible_count`. + /// Number of non-zero-token documents for `i < visible_count`. cumulative_doc_count: u64, /// `Σ batches[i].doc_lengths.iter().sum()` for `i < visible_count`. cumulative_total_tokens: u64, @@ -811,6 +812,7 @@ impl TailIndex { term_builders: FxHashMap, BatchTermBuilder>, with_position: bool, ) { + let doc_count = doc_lengths.iter().filter(|&&len| len > 0).count() as u64; let mut cache = self .writer_term_cache .lock() @@ -840,7 +842,7 @@ impl TailIndex { self.snapshot.store(Arc::new(Snapshot { visible_count: cur.visible_count + 1, batches: cur.batches.pushed(new_meta), - cumulative_doc_count: cur.cumulative_doc_count + rows as u64, + cumulative_doc_count: cur.cumulative_doc_count + doc_count, cumulative_total_tokens: cur.cumulative_total_tokens + total_tokens, })); } @@ -1086,24 +1088,13 @@ impl FtsMemIndex { fn insert_batch(&self, batch: &RecordBatch, row_offset: u64) -> Result<()> { let st = self.state.load_full(); - let batch_position = st.tail.next_position(); let Some(col_idx) = batch .schema() .column_with_name(&self.column_name) .map(|(idx, _)| idx) else { - // Column missing: nothing to index, but publish an empty batch so - // the tail's visibility counters keep up with the writer. - st.tail.append_batch( - batch_position, - row_offset, - batch.num_rows() as u32, - vec![0; batch.num_rows()], - 0, - FxHashMap::default(), - self.params.has_positions(), - ); + // A missing column has no searchable documents. return Ok(()); }; @@ -1155,10 +1146,15 @@ impl FtsMemIndex { total_tokens += doc_token_count as u64; } + if total_tokens == 0 { + return Ok(()); + } + // Drop the tokenizer guard before publishing so we don't hold it // across the snapshot install. drop(tok_guard); + let batch_position = st.tail.next_position(); st.tail.append_batch( batch_position, row_offset, @@ -1752,8 +1748,9 @@ impl FtsMemIndex { /// Export the in-memory FTS index to an `InnerBuilder` ready to be /// written to disk. /// - /// Doc row positions are kept in insert order to match the forward-written - /// flush data file 1:1. `total_rows` is used only to validate positions. + /// Documents are kept in insert order and retain their positions in the + /// forward-written flush data file; zero-token rows are omitted. + /// `total_rows` is used only to validate positions. pub fn to_index_builder( &self, partition_id: u64, @@ -1779,7 +1776,10 @@ impl FtsMemIndex { let tail_snap = st.tail.snapshot(); for batch in tail_snap.batches.iter().take(tail_snap.visible_count) { for i in 0..batch.rows as usize { - all_docs.push((batch.row_offset + i as u64, batch.doc_lengths[i])); + let num_tokens = batch.doc_lengths[i]; + if num_tokens > 0 { + all_docs.push((batch.row_offset + i as u64, num_tokens)); + } } } if all_docs.is_empty() { @@ -1791,8 +1791,8 @@ impl FtsMemIndex { )); } - // Step 2: assign doc_ids in ascending insert-position order, so the - // stored row positions line up 1:1 with the forward-written data file. + // Step 2: assign doc_ids in ascending insert-position order while + // preserving each document's position in the forward-written data file. let mut entries: Vec<(u64, u32)> = Vec::with_capacity(all_docs.len()); for (original, num_tokens) in &all_docs { if *original >= total_rows_u64 { @@ -2970,7 +2970,11 @@ impl Partition { for batch in snap.batches.iter().take(snap.visible_count) { for i in 0..batch.rows as usize { let rp = batch.row_offset + i as u64; - let doc_id = docs.append(rp, batch.doc_lengths[i]); + let num_tokens = batch.doc_lengths[i]; + if num_tokens == 0 { + continue; + } + let doc_id = docs.append(rp, num_tokens); pos_to_doc.insert(rp, doc_id); } } @@ -3652,6 +3656,91 @@ mod tests { assert!(entries.is_empty()); } + #[test] + fn test_zero_token_documents_are_skipped_across_memwal_paths() { + let params = + InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) + .remove_stop_words(true) + .stem(false) + .max_token_length(Some(6)); + let schema = create_test_schema(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec![ + Some(""), + Some(" "), + Some("the"), + Some("overlength"), + None, + Some("hello"), + ])), + ], + ) + .unwrap(); + let index = FtsMemIndex::with_params(1, "description".to_string(), params.clone()); + index.insert(&batch, 0).unwrap(); + + assert_eq!(index.doc_count(), 1); + let st = index.state.load_full(); + let tail_snap = st.tail.snapshot(); + let tokens = vec!["hello".to_string()]; + let tail_scorer = build_scorer(&st, &tail_snap, &tokens, true); + let expected_scorer = MemBM25Scorer::new(1, 1, HashMap::from([("hello".to_string(), 1)])); + assert_eq!(tail_scorer.total_tokens, 1); + assert_eq!(tail_scorer.num_docs(), 1); + assert_eq!(tail_scorer.num_docs_containing_token("hello"), 1); + assert_eq!( + tail_scorer.avg_doc_length(), + expected_scorer.avg_doc_length() + ); + assert_eq!( + tail_scorer.query_weight("hello"), + expected_scorer.query_weight("hello") + ); + let tail_results = index.search("hello"); + assert_eq!(rows(tail_results.clone()), vec![5]); + let tail_score = tail_results[0].score; + assert!(!index.to_index_builder(0, 6).unwrap().is_empty()); + + index.flush(); + let st = index.state.load_full(); + assert_eq!(st.partitions.len(), 1); + assert_eq!( + st.partitions[0] + .docs + .iter() + .map(|(row_id, num_tokens)| (*row_id, *num_tokens)) + .collect::>(), + vec![(5, 1)] + ); + let frozen_scorer = build_scorer(&st, &st.tail.snapshot(), &tokens, true); + assert_eq!(frozen_scorer.total_tokens, 1); + assert_eq!(frozen_scorer.num_docs(), 1); + assert_eq!(frozen_scorer.num_docs_containing_token("hello"), 1); + assert_eq!( + frozen_scorer.avg_doc_length(), + expected_scorer.avg_doc_length() + ); + assert_eq!( + frozen_scorer.query_weight("hello"), + expected_scorer.query_weight("hello") + ); + let frozen_results = index.search("hello"); + assert_eq!(rows(frozen_results.clone()), vec![5]); + assert!((frozen_results[0].score - tail_score).abs() < f32::EPSILON); + + let all_zero_batch = batch.slice(0, 5); + let all_zero_index = FtsMemIndex::with_params(1, "description".to_string(), params); + all_zero_index.insert(&all_zero_batch, 0).unwrap(); + assert!(all_zero_index.is_empty()); + assert_eq!(all_zero_index.doc_count(), 0); + assert!(all_zero_index.to_index_builder(0, 5).unwrap().is_empty()); + all_zero_index.flush(); + assert!(all_zero_index.state.load().partitions.is_empty()); + } + fn create_phrase_test_batch(schema: &ArrowSchema) -> RecordBatch { RecordBatch::try_new( Arc::new(schema.clone()),