Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
147 changes: 137 additions & 10 deletions rust/lance-index/src/scalar/inverted/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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?;
}
}
Expand Down Expand Up @@ -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?;
}

Expand All @@ -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();
Expand Down Expand Up @@ -1545,7 +1540,7 @@ impl IndexWorker {
self.builder.tokens.memory_size() as u64,
);

if skip_empty_document && token_num == 0 {
if token_num == 0 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we apply this zero-token policy consistently across the other FTS paths? The flat/unindexed Utf8 path still appends every row and adds counted_input.num_rows() to num_docs, while MemWAL keeps zero-length documents and flushes them into the DocSet. A whitespace-only, stop-word-only, or overlength document is therefore excluded after a normal index build but counted while it is unindexed or flushed through MemWAL, so avgdl/IDF and potentially top-k ordering can change as the same data moves between these paths. Please update those paths as part of this change and add cross-path coverage.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. I understood the consistency issue at a high level, but I wasn’t sure about the best way to implement the policy across the flat/unindexed and MemWAL paths.

This follow-up commit was implemented mostly with assistance from GPT-5.6-sol. It now excludes zero-token documents from the flat counted input, MemWAL corpus statistics, frozen DocSets, and the flush builder. The added coverage uses whitespace-only, stop-word-only, overlength, null, and normal documents, and checks that num_docs, avgdl, and the IDF/query weight match the normal-build baseline through the flat path and the MemWAL tail-to-frozen/flush transitions. It also checks that the MemWAL score remains unchanged across freeze.

Thanks again for the detailed review—it clarified the invariant that needs to hold across the full FTS lifecycle.

self.last_token_count = 0;
self.trim_temporary_buffers();
self.adjust_tracked_memory_size(
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -2282,8 +2277,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};
Expand Down Expand Up @@ -2316,6 +2313,17 @@ mod tests {
RecordBatch::try_new(schema, vec![docs, row_ids]).unwrap()
}

fn make_doc_batch_from_docs(docs: Vec<Option<&str>>) -> 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,
}
Expand Down Expand Up @@ -3493,6 +3501,125 @@ 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,
},
)
.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)]
Expand Down
65 changes: 57 additions & 8 deletions rust/lance-index/src/scalar/inverted/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6784,12 +6784,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(_) => {
Expand Down Expand Up @@ -6981,12 +6982,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 {
Expand Down Expand Up @@ -11129,6 +11124,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();
Comment on lines +11129 to +11147

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the target area and nearby test context.
sed -n '9440,9525p' rust/lance-index/src/scalar/inverted/index.rs

# Find existing uses of record_batch!() in Rust tests for repo-local style.
rg -n --glob '*.rs' 'record_batch!\(' rust | head -n 80

# Check the Arrow crate version in the workspace lockfile / manifests.
rg -n --glob 'Cargo.lock' 'name = "(arrow-array|arrow-ar|arrow)"|version = ' Cargo.lock
rg -n --glob 'Cargo.toml' 'arrow-array|arrow' Cargo.toml rust/**/Cargo.toml

Repository: lance-format/lance

Length of output: 39859


Use record_batch!() for this test fixture.

Replace the manual Schema / RecordBatch::try_new setup with arrow_array::record_batch!() and keep the nullable string values.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 9469 - 9487,
Replace the manual Schema and RecordBatch::try_new fixture setup with
arrow_array::record_batch!(), preserving the ROW_ID_FIELD values and nullable
text values in the test data.

Source: Coding guidelines

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::<UInt64Type>().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![
Expand Down
Loading
Loading