diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 060e85a4d67..5cf754dfa3a 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1362,7 +1362,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?; } } @@ -1406,7 +1406,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?; } @@ -1432,12 +1432,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(); @@ -1545,7 +1540,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( @@ -1596,7 +1591,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(); @@ -2277,8 +2272,10 @@ pub fn document_input( #[cfg(test)] mod tests { use super::*; + 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}; @@ -2311,6 +2308,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, } @@ -3492,6 +3500,126 @@ mod tests { Ok(()) } + #[tokio::test] + 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(), + index_dir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + + 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(true) + .stem(false) + .max_token_length(Some(6)) + .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 (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(()) + } + + #[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, + block_size: InvertedIndexParams::default().block_size, + }, + ) + .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)] diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index a1217d06ce7..e04d6d1924a 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -6964,12 +6964,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(_) => { @@ -11760,6 +11761,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 85b14a0321d..c799ea07f06 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -578,6 +578,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, } @@ -674,7 +675,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, @@ -846,6 +847,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() @@ -875,7 +877,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, })); } @@ -1139,24 +1141,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(()); }; @@ -1207,10 +1198,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, @@ -1905,8 +1901,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, @@ -1933,7 +1930,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() { @@ -1946,8 +1946,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 { @@ -3260,7 +3260,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); } } @@ -4089,6 +4093,91 @@ mod tests { ); } + #[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()),