From b14f48c9c92d69ad6bcae4f03098833344ac0278 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Mon, 17 Aug 2026 12:01:52 +0200 Subject: [PATCH 1/4] feat(fts): add combined_fields (BM25F) cross-field search Score several text columns as one virtual field (Lucene's `CombinedFieldQuery` / BM25F blend) instead of the per-field max fusion `MultiMatch` does. Adds the query type with serde and JSON parsing, the `CombinedFieldsBM25Scorer`, the indexed scan and the planner and execution nodes that drive it. BM25F blends per-column term frequencies and document lengths into one `tf'`/`dl'` per row, so the scan is row-granular by construction. Two consequences shape the design: Row granularity, not document granularity. An inverted index may hold one document per list element and report `_doc_index` coordinates. BM25F cannot use such an index: it joins the target columns on the row address, and element coordinates of different columns have no correspondence to pair them on. `combined_fields` therefore declares itself row-granular everywhere the granularity plumbing asks, and rejects a target column that can only supply element documents. Corpus statistics must match that granularity. Releases before #7656 indexed each `List` element as its own document, so those files report element-scoped `docCount`/`docFreq` while the scan accumulates by row. Mixing the two domains corrupts `idf'` and `avgdl'`, shifting an old index's top-k relative to the same data reindexed on a current build. Hence `bm25_row_stats_for_terms`, which counts distinct rows, delegating to the document-granular path on V3 where one row owns one document. A cross-field score is complete only when every target column's index holds the row, because `dl'` sums each column's length and a row absent from a column's `DocSet` contributes 0. This commit therefore requires every target column to cover every scanned fragment and refuses the query otherwise, naming the uncovered fragments and the columns to reindex. Scoring the rows no index covers is the next commit. The indexed scan reads every posting up front, then scores the union of the query terms' postings and keeps a bounded top-k. Every candidate is scored, so the result is exact by construction, and candidates are visited in ascending row-id order, which makes the top-k deterministic under ties. MAXSCORE pruning and read pruning are both follow-ups. --- rust/lance-index/src/scalar/inverted.rs | 9 +- .../src/scalar/inverted/combined.rs | 91 +++ .../src/scalar/inverted/combined/cursor.rs | 67 ++ .../src/scalar/inverted/combined/search.rs | 179 +++++ .../src/scalar/inverted/combined/stats.rs | 82 ++ .../src/scalar/inverted/compound.rs | 15 + .../src/scalar/inverted/documents.rs | 347 +++++++- rust/lance-index/src/scalar/inverted/index.rs | 32 +- .../src/scalar/inverted/index/doc_set.rs | 123 ++- .../src/scalar/inverted/index/partition.rs | 50 ++ .../src/scalar/inverted/index/search.rs | 65 ++ .../index/tests/format_and_builder.rs | 61 ++ .../src/scalar/inverted/index/tests/stats.rs | 86 ++ .../lance-index/src/scalar/inverted/parser.rs | 104 ++- rust/lance-index/src/scalar/inverted/query.rs | 418 ++++++++++ .../lance-index/src/scalar/inverted/scorer.rs | 169 ++++ .../src/scalar/inverted/tokenizer.rs | 75 ++ rust/lance-index/src/traits.rs | 9 + rust/lance/src/dataset/mem_wal/index/fts.rs | 10 + .../src/dataset/mem_wal/scanner/fts_search.rs | 7 + rust/lance/src/dataset/scanner.rs | 321 +++++++- rust/lance/src/dataset/tests/dataset_index.rs | 6 + rust/lance/src/index/prefilter.rs | 27 + rust/lance/src/index/scalar/inverted.rs | 39 + rust/lance/src/io/exec/fts.rs | 745 ++++++++++++++++-- rust/lance/src/io/exec/utils.rs | 73 +- 26 files changed, 3083 insertions(+), 127 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/combined.rs create mode 100644 rust/lance-index/src/scalar/inverted/combined/cursor.rs create mode 100644 rust/lance-index/src/scalar/inverted/combined/search.rs create mode 100644 rust/lance-index/src/scalar/inverted/combined/stats.rs diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 474c9125c93..7fef58b1874 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -3,6 +3,7 @@ pub mod builder; mod cache_codec; +mod combined; mod compound; mod cross_column; mod documents; @@ -11,6 +12,8 @@ mod impact; mod index; mod iter; pub mod json; +/// Brute-force scoring reference for tests and benches. Never built normally; see +/// the module docs for the gating. pub mod parser; pub mod query; mod scorer; @@ -23,6 +26,10 @@ use std::sync::{Arc, LazyLock}; use arrow_schema::{DataType, Field}; use async_trait::async_trait; pub use builder::InvertedIndexBuilder; +pub use combined::{ + CombinedFieldColumn, build_combined_bm25_scorer, combined_fields_search, + validate_combined_tokenizers, +}; pub use compound::{ compound_search, compound_search_prepared_match, compound_search_prepared_match_with_score_floor, compound_search_with_base_scorer, @@ -35,7 +42,7 @@ use datafusion::execution::SendableRecordBatchStream; pub use index::*; use lance_core::{Result, cache::LanceCache}; pub use lance_tokenizer::Language; -pub use scorer::{MemBM25Scorer, Scorer}; +pub use scorer::{CombinedFieldsBM25Scorer, MemBM25Scorer, Scorer}; pub use tokenizer::*; use crate::scalar::inverted::query::{FtsSearchParams, Tokens, uses_fuzzy_expansion}; diff --git a/rust/lance-index/src/scalar/inverted/combined.rs b/rust/lance-index/src/scalar/inverted/combined.rs new file mode 100644 index 00000000000..26c5d9cc42b --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Cross-field BM25F scoring (`combined_fields`). +//! +//! The target columns are treated as one virtual field so that term statistics +//! are blended across fields instead of scored independently (contrast the +//! field-centric `best_fields` fusion of a `MultiMatch` query). The blend rules +//! are Lucene's `CombinedFieldQuery`: +//! +//! ```text +//! tf'(t, d) = Σ_f w_f · tf_f(t, d) (weighted sum) +//! dl'(d) = Σ_f w_f · dl_f(d) (weighted sum) +//! docFreq'(t) = max_f docFreq_f(t) (max) +//! docCount' = max_f docCount_f (max) +//! sumTotalTermFreq' = Σ_f w_f · sumTotalTermFreq_f (weighted sum) +//! avgdl' = sumTotalTermFreq' / docCount' +//! score(t, d) = idf'(t) · (k1 + 1) · tf' / (tf' + k1·(1 - b + b·dl'/avgdl')) +//! ``` +//! +//! Scoring uses exact (non-quantized) document lengths and a shared tokenizer +//! across columns. Every candidate in the union of the query terms' postings is +//! scored; see [`combined_fields_search`]. + +mod cursor; +mod search; +mod stats; + +use std::sync::Arc; + +use lance_core::{Error, Result}; + +pub use search::combined_fields_search; +pub use stats::build_combined_bm25_scorer; + +use super::index::InvertedIndex; +use super::query::Tokens; + +/// One target column of a `combined_fields` query: its per-column weight and +/// the opened FTS segments (one per committed segment; usually a single one). +pub struct CombinedFieldColumn { + /// Column name, used only for error messages. + pub column: String, + /// Per-column BM25F weight `w_f` (`>= 1`, validated at query construction). + pub weight: f32, + /// Opened inverted-index segments for this column. + pub indices: Vec>, +} + +/// Deduplicate the query tokens into the unique terms of the virtual field, +/// preserving first-seen order. Duplicate terms collapse to one (mirrors the +/// per-column `load_posting_lists` dedup), so each term is scored once. +fn unique_terms(tokens: &Tokens) -> Vec { + let mut terms = Vec::with_capacity(tokens.len()); + let mut seen = std::collections::HashSet::new(); + for token in tokens { + if seen.insert(token.as_str()) { + terms.push(token.clone()); + } + } + terms +} + +/// Reject a `combined_fields` query whose target columns do not share an +/// identical index/tokenizer configuration. BM25F is only well-defined when the +/// fields tokenize the same way, so mixing configurations is an error rather +/// than a silently wrong score. The error lists the offending columns. +pub fn validate_combined_tokenizers(columns: &[CombinedFieldColumn]) -> Result<()> { + // A column with no index has no tokenizer to disagree with, so it is skipped. + let mut indexed = columns + .iter() + .filter_map(|column| Some((column, column.indices.first()?))); + let Some((reference, reference_index)) = indexed.next() else { + return Ok(()); + }; + // Compare only the tokenization-affecting params: two columns may differ in + // storage/layout knobs (e.g. `with_position`) yet still tokenize identically, + // which is all BM25F requires. + let offending: Vec<&str> = indexed + .filter(|(_, index)| !reference_index.params().same_tokenization(index.params())) + .map(|(column, _)| column.column.as_str()) + .collect(); + if !offending.is_empty() { + return Err(Error::invalid_input(format!( + "combined_fields requires every target column to share the same tokenizer/index \ + configuration; column(s) {:?} differ from column '{}'", + offending, reference.column + ))); + } + Ok(()) +} diff --git a/rust/lance-index/src/scalar/inverted/combined/cursor.rs b/rust/lance-index/src/scalar/inverted/combined/cursor.rs new file mode 100644 index 00000000000..62e8af0f250 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/cursor.rs @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Per-term cross-column postings for `combined_fields`: one term's postings +//! merged across every target column into the shared row-id space, and the +//! loaded posting sources they are built from. + +use std::collections::HashMap; +use std::sync::Arc; + +use lance_select::RowAddrMask; + +use super::super::documents::AddressKeyedDocuments; +use super::super::index::{PostingList, live_posting_rows}; +use super::super::scorer::CombinedFieldsBM25Scorer; + +/// One query term's postings, merged across every target column/partition into +/// the shared row-id space. +/// +/// Entries are unique row ids sorted ascending, each carrying the blended term +/// frequency `tf'(t, d) = Σ_f w_f · freq_f(t, d)`. +pub(super) struct CombinedTermPostings { + pub(super) idf: f32, + pub(super) postings: Vec<(u64, f32)>, +} + +impl CombinedTermPostings { + /// `tf'` for `row_id`, or 0 when the term does not occur in the document. + #[inline] + pub(super) fn tf_prime(&self, row_id: u64) -> f32 { + match self.postings.binary_search_by_key(&row_id, |(id, _)| *id) { + Ok(idx) => self.postings[idx].1, + Err(_) => 0.0, + } + } +} + +/// A `(column, index, partition)` posting source loaded for one term. +pub(super) struct LoadedSource { + pub(super) weight: f32, + pub(super) docs: AddressKeyedDocuments, + pub(super) is_legacy: bool, + pub(super) posting: PostingList, +} + +/// Merge every source's postings for `term` into the shared row-id space, +/// accumulating `tf'` in the canonical order. +pub(super) fn build_term_postings( + term: &str, + sources: Vec, + mask: &Arc, + scorer: &CombinedFieldsBM25Scorer, +) -> CombinedTermPostings { + let mut acc: HashMap = HashMap::new(); + for source in &sources { + for (row_id, freq) in live_posting_rows(&source.posting, &source.docs, source.is_legacy) { + if !mask.selected(row_id) { + continue; + } + *acc.entry(row_id).or_insert(0.0) += source.weight * freq as f32; + } + } + let idf = scorer.query_weight(term); + let mut postings: Vec<(u64, f32)> = acc.into_iter().collect(); + postings.sort_unstable_by_key(|(row_id, _)| *row_id); + CombinedTermPostings { idf, postings } +} diff --git a/rust/lance-index/src/scalar/inverted/combined/search.rs b/rust/lance-index/src/scalar/inverted/combined/search.rs new file mode 100644 index 00000000000..ceae5ede3c8 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/search.rs @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The indexed `combined_fields` entry point: load every term's postings across +//! the target columns, merge them into the shared row-id space, and score every +//! candidate. + +use std::cmp::Reverse; +use std::collections::{BTreeSet, BinaryHeap}; +use std::sync::Arc; + +use lance_core::Result; +use lance_core::utils::tokio::spawn_cpu; + +use super::super::documents::AddressKeyedDocuments; +use super::super::query::{FtsSearchParams, Operator, Tokens}; +use super::super::scorer::CombinedFieldsBM25Scorer; +use super::cursor::{CombinedTermPostings, LoadedSource, build_term_postings}; +use super::{CombinedFieldColumn, unique_terms}; +use crate::metrics::MetricsCollector; +use crate::prefilter::PreFilter; +use crate::vector::graph::OrderedFloat; + +/// A scored candidate ordered the way callers read results: `score DESC, row_id +/// ASC`, the same order the single-column path imposes in +/// `classify_wand_exactness_certificate`. +/// +/// [`ScoredDoc`](super::super::builder::ScoredDoc) compares on score alone, which +/// is not enough for a bounded heap: among equal scores it evicts whichever row +/// the heap happens to hold at the bottom, so rows that belong in the top-k are +/// dropped and no later sort can bring them back. A fully covered plan returns +/// this search's output directly, with no `SortExec` above it, so the order and +/// the membership both have to be settled here. +/// +/// `row_id` is stored reversed so that the derived lexicographic ordering ranks a +/// higher row id lower. The heap's smallest element is then the lowest score with +/// the highest row id, which is exactly the candidate a full heap should evict. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +struct RankedDoc { + score: OrderedFloat, + row_id: Reverse, +} + +impl RankedDoc { + fn new(row_id: u64, score: f32) -> Self { + Self { + score: OrderedFloat(score), + row_id: Reverse(row_id), + } + } +} + +/// Exact cross-field BM25F search over the target columns. +/// +/// Loads each query term's postings across every column and partition, then +/// scores the union of those postings and keeps a bounded top-k. Every candidate +/// is scored; there is no pruning, so the result is exact by construction. +/// Results come back ordered by `score DESC, row_id ASC`, which is a total order, +/// so the same data always yields the same top-k even when scores tie. +/// +/// `operator` applies across the virtual field: `And` keeps only docs where +/// every query term appears in at least one column; `Or` keeps docs matching +/// any term. Per-column `boost` is folded into `tf'`. +pub async fn combined_fields_search( + columns: &[CombinedFieldColumn], + tokens: &Tokens, + params: &FtsSearchParams, + operator: Operator, + scorer: &CombinedFieldsBM25Scorer, + prefilter: Arc, + metrics: &dyn MetricsCollector, +) -> Result<(Vec, Vec)> { + let terms = unique_terms(tokens); + let limit = params.limit.unwrap_or(usize::MAX); + if terms.is_empty() || limit == 0 { + return Ok((Vec::new(), Vec::new())); + } + + let mask = prefilter.mask(); + let require_all_terms = operator == Operator::And; + + // Load every term's postings across all columns, in the canonical + // column → index → partition order. The length sources are collected in that + // same order so `dl'` sums in the exact scan's order too (float addition is + // order-sensitive; matching the order keeps every score bit-identical). + let mut loaded: Vec> = (0..terms.len()).map(|_| Vec::new()).collect(); + let mut length_sources: Vec<(f32, AddressKeyedDocuments)> = Vec::new(); + for column in columns { + let weight = column.weight; + for index in &column.indices { + for partition in &index.partitions { + let docs = partition.docs.address_keyed().await?; + let is_legacy = partition.is_legacy(); + for (term_index, term) in terms.iter().enumerate() { + let Some(token_id) = partition.tokens.get(term) else { + continue; + }; + let posting = partition + .inverted_list + .posting_list(token_id, false, metrics) + .await?; + loaded[term_index].push(LoadedSource { + weight, + docs: docs.clone(), + is_legacy, + posting, + }); + } + length_sources.push((weight, docs)); + } + } + } + // Everything past the loads is uninterruptible CPU work: building and sorting + // a per-term `HashMap`, then the whole scoring loop with no await. Offload it + // so a large query cannot hold a DataFusion + // worker past a stream drop or task cancellation, matching how the + // single-column `InvertedIndex::bm25_search` dispatches its per-partition + // scoring and how `flat_combined_fields_search_stream` dispatches its own + // scoring loop. The `'static` closure clones the borrowed `scorer` (a handful + // of per-term statistics) and moves everything else in. + let scorer = Arc::new(scorer.clone()); + let top = spawn_cpu(move || { + let dl_prime = |row_id: u64| -> f32 { + length_sources + .iter() + .map(|(weight, docs)| weight * docs.doc_length_at(row_id) as f32) + .sum() + }; + let terms: Vec = terms + .iter() + .zip(loaded) + .map(|(term, sources)| build_term_postings(term, sources, &mask, scorer.as_ref())) + .collect(); + + // Score every candidate: the union of the terms' postings for `Or`, and the + // same union filtered to the documents holding every term for `And`. A row + // absent from a term contributes no `tf'`, so it is skipped rather than + // scored as zero. + // + // Ties are settled by [`RankedDoc`], not left to the heap: it orders on + // `(score, row_id)` as a whole, so both which rows survive the k-th score and + // the order they come back in are fixed by the data alone. + let candidates: BTreeSet = terms + .iter() + .flat_map(|term| term.postings.iter().map(|(row_id, _)| *row_id)) + .collect(); + let mut top: BinaryHeap> = BinaryHeap::new(); + for row_id in candidates { + let dl = dl_prime(row_id); + let mut score = 0.0f32; + let mut missing_term = false; + for term in &terms { + let tf = term.tf_prime(row_id); + if tf <= 0.0 { + missing_term = true; + continue; + } + score += term.idf * scorer.doc_weight(tf, dl); + } + if require_all_terms && missing_term { + continue; + } + top.push(Reverse(RankedDoc::new(row_id, score))); + if top.len() > limit { + top.pop(); + } + } + Result::Ok(top) + }) + .await?; + + // Ascending in `Reverse` is descending in `RankedDoc`, i.e. best + // first: highest score, and within a score the lowest row id. + Ok(top + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id.0, doc.score.0)) + .unzip()) +} diff --git a/rust/lance-index/src/scalar/inverted/combined/stats.rs b/rust/lance-index/src/scalar/inverted/combined/stats.rs new file mode 100644 index 00000000000..11fba5d2ad0 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/stats.rs @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Corpus statistics for `combined_fields`: the cross-column BM25F scorer +//! build. + +use std::collections::HashMap; + +use lance_core::Result; + +use super::super::query::Tokens; +use super::super::scorer::CombinedFieldsBM25Scorer; +use super::{CombinedFieldColumn, unique_terms}; +use crate::metrics::MetricsCollector; + +/// Fold the target columns' segment statistics into a single [`CombinedFieldsBM25Scorer`]. +/// +/// Generalizes [`build_global_bm25_scorer`](super::super::build_global_bm25_scorer) +/// (which folds segments of one column) to fold across columns too, with +/// per-column weights and the BM25F blend: +/// `docCount'`/`docFreq'` take the max across columns while `sumTotalTermFreq'` +/// is the weighted sum. The query terms are deduplicated the same way the scan +/// deduplicates them, so the resulting per-term `docFreq'` covers exactly the +/// terms [`combined_fields_search`](super::combined_fields_search) scores. +/// +/// `metrics`, when provided, receives the per-token posting-metadata cache +/// lookups this fold triggers, exactly as +/// [`build_global_bm25_scorer`](super::super::build_global_bm25_scorer) +/// reports them on the single-column path. Without it a cold cross-field query +/// undercounts `index_cache_misses` by one lookup per (term, partition, column). +pub async fn build_combined_bm25_scorer( + columns: &[CombinedFieldColumn], + tokens: &Tokens, + metrics: Option<&dyn MetricsCollector>, +) -> Result { + let terms = unique_terms(tokens); + let mut doc_count = 0usize; + let mut sum_total_term_freq = 0f64; + let mut doc_freq: HashMap = terms.iter().map(|t| (t.clone(), 0)).collect(); + + for column in columns { + let mut column_num_docs = 0usize; + let mut column_total_tokens = 0u64; + let mut column_doc_freq = vec![0usize; terms.len()]; + { + for index in &column.indices { + // Row granularity, not document granularity: `combined_fields_search` + // blends every posting a row owns into one `tf'` and every document + // length it owns into one `dl'`, and released V1/V2 indexes may hold + // one document per list element. See + // [`InvertedIndex::bm25_row_stats_for_terms`]. + let (total_tokens, num_docs, token_docs) = + index.bm25_row_stats_for_terms(&terms, metrics).await?; + column_total_tokens += total_tokens; + column_num_docs += num_docs; + for (slot, df) in token_docs.into_iter().enumerate() { + column_doc_freq[slot] += df; + } + } + } + + doc_count = doc_count.max(column_num_docs); + sum_total_term_freq += column.weight as f64 * column_total_tokens as f64; + for (term, df) in terms.iter().zip(column_doc_freq) { + let entry = doc_freq + .get_mut(term) + .expect("doc_freq initialized for every term"); + *entry = (*entry).max(df); + } + } + + let avg_doc_length = if doc_count > 0 { + (sum_total_term_freq / doc_count as f64) as f32 + } else { + 0.0 + }; + Ok(CombinedFieldsBM25Scorer::new( + doc_count, + avg_doc_length, + doc_freq, + )) +} diff --git a/rust/lance-index/src/scalar/inverted/compound.rs b/rust/lance-index/src/scalar/inverted/compound.rs index 7bbba53ee52..e37302bffa0 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -686,6 +686,7 @@ impl CompoundScorerPlan { .map(|query| Self::from_query(query, num_leaves)) .collect::>>()?, }), + FtsQuery::CombinedFields(_) => Err(combined_fields_unsupported()), } } @@ -3769,10 +3770,24 @@ pub(super) fn collect_leaf_queries(query: &FtsQuery, leaves: &mut Vec collect_leaf_queries(child, leaves)?; } } + FtsQuery::CombinedFields(_) => return Err(combined_fields_unsupported()), } Ok(()) } +/// A `combined_fields` node reaching the compound scorer is a planner bug. +/// +/// Every leaf here draws its postings from one column's index and scores them +/// with a `MemBM25Scorer`, whereas BM25F blends `tf'`/`dl'`/`docFreq'` across +/// several columns before scoring. The planner keeps such trees out via +/// `supports_compound_scorer`, so this is an explicit error. +fn combined_fields_unsupported() -> Error { + Error::not_supported( + "the compound FTS scorer cannot score a combined_fields (BM25F) node: its statistics \ + are blended across columns and do not fit the single-index leaf protocol", + ) +} + struct PreparedLeaf { query: Arc, params: Arc, diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index f7631355b8b..a53ea4a1fe0 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -8,6 +8,7 @@ //! never has to infer which value a numeric slot represents. use std::borrow::Cow; +use std::ops::Range; use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering as AtomicOrdering}; use std::sync::{Arc, OnceLock, Weak}; @@ -28,7 +29,7 @@ use crate::FtsPrewarmDocumentStatus; use crate::scalar::{IndexReader, IndexStore, RowIdRemapper}; use super::index::{ - DocSet, NUM_TOKEN_COL, dequantize_doc_length, doc_index_storage_column, + DocSet, NUM_TOKEN_COL, count_row_id_runs, dequantize_doc_length, doc_index_storage_column, document_coordinate_rank, quantize_doc_length, }; @@ -301,6 +302,18 @@ impl AddressDocIdLookup { left } + /// The positions of every live DocId whose address falls in `start..=end`. + fn positions_in_address_range( + &self, + projection: &ResidentAddressProjection, + start: u64, + end: u64, + ) -> Range { + let first = self.partition_point(projection, |address| address < start); + let after_last = self.partition_point(projection, |address| address <= end); + first..after_last + } + fn insert_address_range( &self, projection: &ResidentAddressProjection, @@ -308,9 +321,7 @@ impl AddressDocIdLookup { end: u64, selected: &mut RoaringBitmap, ) { - let first = self.partition_point(projection, |address| address < start); - let after_last = self.partition_point(projection, |address| address <= end); - for position in first..after_last { + for position in self.positions_in_address_range(projection, start, end) { selected.insert(self.doc_id_at(position)); } } @@ -971,6 +982,28 @@ impl ResidentAddressProjection { .unwrap_or_else(|| (0..self.len() as u32).collect()) } + /// True iff every slot is live and addresses ascend strictly with DocId. + /// + /// DocId order then equals address order and no two documents share a row, + /// so a contiguous DocId range spans a contiguous address interval. Row + /// granularity relies on the second half: cross-field statistics count + /// distinct rows, so an ascending partition can answer + /// [`AddressKeyedDocuments::num_distinct_rows`] without walking it. + /// + /// The ordering half is [`OrderedRowAddressProjection`], whose verdict this + /// projection already caches for the cross-column scorer. Density is the + /// extra condition: an ordered projection may still hold dead slots, and + /// those report [`RowAddress::TOMBSTONE_ROW`], which the block-skip cursor + /// cannot tell from an exhausted cursor. A remapper is attached whenever the + /// dataset carries a fragment reuse index, even when it retains every + /// document, so liveness has to be decided on cardinality. + fn dense_and_strictly_ascending(&self) -> bool { + match self.try_ordered_row_addresses() { + Ok(ordered) => !ordered.has_sparse_live_docs(), + Err(_) => false, + } + } + async fn doc_ids_by_address(&self) -> Result> { self.projection .doc_ids_by_address @@ -1186,6 +1219,7 @@ impl PartitionDocumentStore { prewarm_complete: true, scoring_ready: true, reverse_lookup_ready: true, + ascending_addresses_ready: true, projection_resident: true, }, Self::Modern(docs) => docs.prewarm_status(), @@ -1198,6 +1232,128 @@ impl PartitionDocumentStore { Self::Modern(docs) => docs.load_build_docset().await, } } + + /// This partition's documents keyed by row address; see + /// [`AddressKeyedDocuments`]. + pub(crate) async fn address_keyed(&self) -> Result { + match self { + Self::Legacy(docs) => Ok(AddressKeyedDocuments::from_docset(docs.clone())), + Self::Modern(docs) => docs.address_keyed().await, + } + } +} + +/// One partition's documents keyed by row address instead of [`DocId`]. +/// +/// A scan that spans several columns can only join them on the row address: +/// each column has its own index, its own partitioning and its own DocId space. +/// Such a scan needs both directions: `DocId -> address` to place a posting, +/// and `address -> length` to read a candidate row's contribution from a column +/// it may hold no posting for at all, which rules out carrying a DocId through +/// the merge instead. +/// +/// The modern variant materializes nothing per query: it clones the partition's +/// cached lengths, its resident address projection and the address-sorted DocId +/// lookup, all of which are per-partition state that [`PartitionDocuments::prewarm`] +/// fills and that outlives the query. It does keep the whole address column +/// resident, unlike the single-column path which resolves addresses only for the +/// final top-k, because every candidate row needs a length from every column. +#[derive(Debug, Clone)] +pub(super) struct AddressKeyedDocuments(AddressKeyedSource); + +#[derive(Debug, Clone)] +enum AddressKeyedSource { + /// Legacy partitions already own a complete row-address-keyed [`DocSet`]. + Legacy(Arc), + Modern { + projection: ResidentAddressProjection, + doc_ids_by_address: Arc, + lengths: Arc, + /// Memoized per partition; see + /// [`ResidentAddressProjection::dense_and_strictly_ascending`]. + strictly_ascending: bool, + }, +} + +impl AddressKeyedDocuments { + /// View a legacy partition's complete [`DocSet`], which is already keyed by + /// row address. + pub(crate) fn from_docset(docs: Arc) -> Self { + Self(AddressKeyedSource::Legacy(docs)) + } + + /// Number of documents, counting the dead slots a remapped partition keeps + /// so its DocIds stay aligned with the posting lists. + pub(crate) fn len(&self) -> usize { + match &self.0 { + AddressKeyedSource::Legacy(docs) => docs.len(), + AddressKeyedSource::Modern { lengths, .. } => lengths.len(), + } + } + + /// Row address of `doc_id`, or [`RowAddress::TOMBSTONE_ROW`] when the slot is + /// not live. + #[inline] + pub(crate) fn row_address(&self, doc_id: u32) -> u64 { + match &self.0 { + AddressKeyedSource::Legacy(docs) => docs.row_id(doc_id), + AddressKeyedSource::Modern { projection, .. } => projection + .address(DocId::new(doc_id)) + .unwrap_or(RowAddress::TOMBSTONE_ROW), + } + } + + /// Total length of the row at `address`: the sum over every document the row + /// owns in this partition, or 0 when the partition holds none (an empty or + /// null field, a row outside the partition, or a dead slot). + /// + /// Released V1/V2 list indexes indexed each `List` element as its own + /// document, so one row can own a whole run of documents and every one of + /// them contributes. + #[inline] + pub(crate) fn doc_length_at(&self, address: u64) -> u64 { + match &self.0 { + AddressKeyedSource::Legacy(docs) => docs.doc_length_by_row_id(address), + AddressKeyedSource::Modern { + projection, + doc_ids_by_address, + lengths, + .. + } => doc_ids_by_address + .positions_in_address_range(projection, address, address) + .map(|position| { + u64::from(lengths.exact(DocId::new(doc_ids_by_address.doc_id_at(position)))) + }) + .sum(), + } + } + + /// Number of distinct row addresses the live documents cover. + /// + /// Equal to [`Self::len`] whenever each row owns a single live document. + /// Row-granularity corpus statistics need this count; see + /// [`Self::doc_length_at`] for why the two can differ. + pub(crate) fn num_distinct_rows(&self) -> usize { + match &self.0 { + AddressKeyedSource::Legacy(docs) => docs.num_distinct_rows(), + AddressKeyedSource::Modern { + projection, + doc_ids_by_address, + strictly_ascending, + .. + } => { + if *strictly_ascending { + return self.len(); + } + // The lookup lists the live DocIds in address order, so the + // documents of one row form a contiguous run. + count_row_id_runs( + (0..doc_ids_by_address.len(projection)) + .map(|position| doc_ids_by_address.address_at(projection, position)), + ) + } + } + } } impl std::fmt::Debug for PartitionDocuments { @@ -1319,6 +1475,11 @@ impl PartitionDocuments { .projection .get() .is_some_and(|projection| projection.doc_ids_by_address.initialized()), + ascending_addresses_ready: self.projection.get().is_some_and(|projection| { + CachedRowAddressOrder::from_raw( + projection.ordered_validation.load(AtomicOrdering::Acquire), + ) != CachedRowAddressOrder::Unknown + }), projection_resident: self.projection_resident(), } } @@ -1807,6 +1968,39 @@ impl PartitionDocuments { self.num_docs.saturating_mul(std::mem::size_of::()) } + /// This partition's documents keyed by row address; see + /// [`AddressKeyedDocuments`]. + /// + /// Every part is a per-partition cache, so a warm partition answers with + /// four `Arc` clones and no IO, CPU-pool work or per-query allocation. A cold + /// partition reads the two columns concurrently, since it needs both whatever + /// the outcome. + pub(crate) async fn address_keyed(&self) -> Result { + let (lengths, projection) = futures::try_join!(self.lengths(), self.address_projection())?; + let doc_ids_by_address = projection.doc_ids_by_address().await?; + let strictly_ascending = self.strictly_ascending_addresses(&projection).await?; + Ok(AddressKeyedDocuments(AddressKeyedSource::Modern { + projection, + doc_ids_by_address, + lengths, + strictly_ascending, + })) + } + + /// Answer to [`ResidentAddressProjection::dense_and_strictly_ascending`], + /// memoized by the projection's own ordering cache. The validation scan is + /// O(num_docs), so it belongs to prewarmed partition state. + async fn strictly_ascending_addresses( + &self, + projection: &ResidentAddressProjection, + ) -> Result { + if projection.cached_row_address_order() != CachedRowAddressOrder::Unknown { + return Ok(projection.dense_and_strictly_ascending()); + } + let projection = projection.clone(); + spawn_cpu(move || Result::Ok(projection.dense_and_strictly_ascending())).await + } + /// Materialize the build-side table for rewrite/update operations. pub(crate) async fn load_build_docset(&self) -> Result { DocSet::load(self.reader().await?, false, self.remapper.clone()).await @@ -1861,10 +2055,9 @@ impl PartitionDocuments { Result::Ok(()) }) .await?; - self.address_projection() - .await? - .doc_ids_by_address() - .await?; + let projection = self.address_projection().await?; + projection.doc_ids_by_address().await?; + self.strictly_ascending_addresses(&projection).await?; Result::Ok(()) }) .await?; @@ -2721,6 +2914,46 @@ mod tests { assert!(!projection.doc_ids_by_address.initialized()); } + /// A remapped partition that kept every document must still pass the gate. + /// That is the common case: a fragment reuse index attaches a remapper on + /// every index load, whether or not it removes anything. + /// + /// A dead slot stores address 0, so `dead_first_slot` keeps the stored + /// addresses ascending and is rejected by the liveness half alone. + #[rstest::rstest] + #[case::no_remapper_ascending(vec![10, 20, 30], None, true)] + #[case::no_remapper_ties(vec![10, 20, 20], None, false)] + #[case::remapper_retains_every_document(vec![10, 20, 30], Some(vec![]), true)] + #[case::remapper_rewrites_in_order(vec![10, 20, 30], Some(vec![(20, Some(25))]), true)] + #[case::remapper_ties_addresses(vec![10, 20, 30], Some(vec![(20, Some(10))]), false)] + #[case::remapper_reverses_addresses(vec![10, 20, 30], Some(vec![(30, Some(5))]), false)] + #[case::dead_first_slot(vec![10, 20, 30], Some(vec![(10, None)]), false)] + #[case::dead_middle_slot(vec![10, 20, 30], Some(vec![(20, None)]), false)] + fn ascending_address_gate_requires_live_slots_and_strict_order( + #[case] raw: Vec, + #[case] remapping: Option)>>, + #[case] expected: bool, + ) { + let raw = Arc::new(UInt64Array::from(raw)); + let remapper = remapping.map(|entries| TestRemapper { + mapping: entries.into_iter().collect(), + }); + let projection = Arc::new( + VersionAddressProjection::try_new( + raw.as_ref(), + raw.len(), + remapper + .as_ref() + .map(|remapper| remapper as &dyn RowIdRemapper), + "docs", + ) + .expect("valid projection"), + ); + let projection = projection.resident(Some(raw)).unwrap(); + + assert_eq!(projection.dense_and_strictly_ascending(), expected); + } + #[test] fn doc_lengths_validate_shape_total_and_memory() { let mismatch = DocLengths::try_new(ScalarBuffer::from(vec![2, 3]), 3, None, false, "docs") @@ -2871,6 +3104,104 @@ mod tests { .await .is_some() ); + assert_eq!( + CachedRowAddressOrder::from_raw( + documents + .projection + .get() + .unwrap() + .ordered_validation + .load(AtomicOrdering::Acquire) + ), + CachedRowAddressOrder::Ordered + ); + } + + /// Prewarm owns the O(num_docs) address scan, so the first cross-field query + /// against a prewarmed partition spends no CPU deciding the fast path. + #[rstest::rstest] + #[case::retains_every_document(vec![], true)] + #[case::deletes_a_document(vec![(20, None)], false)] + #[tokio::test] + async fn prewarm_answers_the_ascending_gate_for_a_remapped_partition( + #[case] remapping: Vec<(u64, Option)>, + #[case] strictly_ascending: bool, + ) { + let (_directory, store, cache) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + let remapper: Arc = Arc::new(TestRemapper { + mapping: remapping.into_iter().collect(), + }); + let documents = open_documents(store, path, cache.as_ref(), Some(remapper)) + .await + .unwrap(); + + documents.prewarm().await.unwrap(); + + assert!(documents.query_ready()); + assert_ne!( + CachedRowAddressOrder::from_raw( + documents + .projection + .get() + .unwrap() + .ordered_validation + .load(AtomicOrdering::Acquire) + ), + CachedRowAddressOrder::Unknown, + "prewarm must leave the projection's ordering verdict cached" + ); + let keyed = documents.address_keyed().await.unwrap(); + assert_eq!(keyed.row_address(0), 10); + assert_eq!(keyed.doc_length_at(30), 5); + if strictly_ascending { + assert_eq!(keyed.num_distinct_rows(), 3); + assert_eq!(keyed.row_address(1), 20); + assert_eq!(keyed.doc_length_at(20), 3); + } else { + assert_eq!(keyed.num_distinct_rows(), 2); + assert_eq!(keyed.row_address(1), RowAddress::TOMBSTONE_ROW); + assert_eq!(keyed.doc_length_at(20), 0); + } + } + + /// A partition holding no documents satisfies the cross-field fast-path gate + /// vacuously, and both document representations must say so: the gate is a + /// per-partition decision, so disagreeing here would make an empty legacy + /// partition alone force a whole query onto the full-read fallback. + #[tokio::test] + async fn the_ascending_gate_agrees_on_an_empty_partition() { + let (_directory, store, cache) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(Vec::::new()), + UInt32Array::from(Vec::::new()), + Some("0"), + ) + .await; + let modern = open_documents(store, path, cache.as_ref(), None) + .await + .unwrap() + .address_keyed() + .await + .unwrap(); + + let legacy = AddressKeyedDocuments::from_docset(Arc::new(DocSet::default())); + + for (label, keyed) in [("modern", &modern), ("legacy", &legacy)] { + assert_eq!(keyed.len(), 0, "{label}"); + assert_eq!(keyed.num_distinct_rows(), 0, "{label}"); + } } #[tokio::test] diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 660339fefd3..ab66b97467e 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -51,7 +51,7 @@ use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::utils::tracing::{IO_TYPE_LOAD_SCALAR_PART, TRACE_IO_EVENTS}; use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; use lance_select::{RowAddrMask, RowAddrTreeMap}; -use roaring::RoaringBitmap; +use roaring::{RoaringBitmap, RoaringTreemap}; use std::sync::LazyLock; use tokio::{ sync::{Mutex, OnceCell}, @@ -60,7 +60,8 @@ use tokio::{ use tracing::{debug, info, instrument, warn}; use super::documents::{ - DocId, DocLengths, DocVisibility, PartitionDocumentStore, PartitionDocuments, + AddressKeyedDocuments, DocId, DocLengths, DocVisibility, PartitionDocumentStore, + PartitionDocuments, }; use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder}; use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder}; @@ -126,5 +127,32 @@ use prewarm::*; pub(super) use search_candidates::*; pub use token_set::*; +/// Walk `posting` as `(row_id, frequency)` over the rows that still exist, +/// translating each posting's key into a row id. +/// +/// Compressed postings key on a partition-local doc id, so they go through +/// `docs`; legacy plain postings key on the row id directly. +/// +/// A remapped partition keeps a deleted document's slot so the posting lists stay +/// aligned with its DocIds, and `row_address` answers `TOMBSTONE_ROW` for it. +/// Those postings are dropped here: a default prefilter mask is an empty block +/// list, which selects them, and `doc_length_at` reports 0 for them, which would +/// hand them the largest `doc_weight` there is. The single-column `wand` cursor +/// guards the same way. +pub(in crate::scalar::inverted) fn live_posting_rows<'a>( + posting: &'a PostingList, + docs: &'a AddressKeyedDocuments, + is_legacy: bool, +) -> impl Iterator + 'a { + posting.iter().filter_map(move |(posting_doc_id, freq, _)| { + let row_id = if is_legacy { + posting_doc_id + } else { + docs.row_address(posting_doc_id as u32) + }; + (row_id != RowAddress::TOMBSTONE_ROW).then_some((row_id, freq)) + }) +} + #[cfg(test)] mod tests; diff --git a/rust/lance-index/src/scalar/inverted/index/doc_set.rs b/rust/lance-index/src/scalar/inverted/index/doc_set.rs index 4a77e00de67..d9ad17fa125 100644 --- a/rust/lance-index/src/scalar/inverted/index/doc_set.rs +++ b/rust/lance-index/src/scalar/inverted/index/doc_set.rs @@ -237,6 +237,9 @@ pub struct DocSet { // partitions never set the flag and keep exact scoring. pub(super) scoring_quantized: bool, pub(super) norms: Arc>>, + // Memoized answer to `row_ids_strictly_ascending`. Filled on first query use; + // build-time appends happen before any query touches it. + row_ids_ascending: Arc>, } impl DeepSizeOf for DocSet { @@ -253,6 +256,21 @@ impl DeepSizeOf for DocSet { } } +/// Count the runs of equal values in a sequence sorted by row id, i.e. the +/// number of distinct row ids it contains. Row addresses sort the same way, so +/// the address-keyed document view shares it. +pub(in crate::scalar::inverted) fn count_row_id_runs(row_ids: impl Iterator) -> usize { + let mut distinct = 0; + let mut previous = None; + for row_id in row_ids { + if previous != Some(row_id) { + distinct += 1; + previous = Some(row_id); + } + } + distinct +} + impl DocSet { pub(crate) fn with_coordinate_rank(coordinate_rank: usize) -> Self { Self { @@ -307,13 +325,66 @@ impl DocSet { self.doc_indices.len() } + /// True iff `row_id(doc_id)` is strictly increasing in `doc_id`, i.e. the + /// per-doc `row_ids` array is sorted with no duplicates. This is the exact + /// condition the combined_fields read-pruning fast path requires: doc-id + /// order equals row-id order (so a posting block's row-id span is an + /// interval, letting a row-id seek skip whole blocks) and every row owns a + /// single doc (no list-element or legacy list multiplicity, so a term + /// contributes at most one posting per partition per row). + /// + /// This checks ascension only, not liveness. A tombstoned slot holds + /// [`RowAddress::TOMBSTONE_ROW`] (`u64::MAX`), which breaks ascension wherever + /// another slot follows it but not when it is the last one, so a set whose + /// final document is tombstoned still reports `true`. Only the build-side sets + /// (`PartitionDocuments::load_build_docset`) tombstone in place, and + /// `combined_fields_search` forces the full-read fallback for legacy + /// partitions regardless of this answer. + /// + /// Computed once per loaded set and memoized (shared across clones). An empty + /// set is vacuously ascending, matching the modern side's + /// `ResidentAddressProjection::dense_and_strictly_ascending`. + pub fn row_ids_strictly_ascending(&self) -> bool { + *self + .row_ids_ascending + .get_or_init(|| self.row_ids.windows(2).all(|w| w[0] < w[1])) + } + + /// Number of distinct row ids this set covers, which row-granularity BM25F + /// statistics need. + /// + /// Equal to [`Self::len`] whenever each row owns a single document, and + /// strictly smaller on an `inv`-keyed set where a row owns a run of documents + /// (see [`Self::doc_ids`]). + /// + /// Linear in the number of documents, over an already-loaded set: both `inv` + /// and the legacy `row_ids` array are sorted by row id, so the documents of + /// one row form a contiguous run. Tombstoned (frag-reuse deleted) documents + /// are absent from `inv` and therefore uncounted, matching + /// [`Self::doc_length_by_row_id`], which reports 0 for them. + pub fn num_distinct_rows(&self) -> usize { + // Memoized and shared with the combined_fields fast-path check, so the + // one-document-per-row case answers without a scan after the first query. + if self.row_ids_strictly_ascending() { + return self.len(); + } + if !self.inv.is_empty() { + return count_row_id_runs(self.inv.iter().map(|entry| entry.0)); + } + count_row_id_runs(self.row_ids.iter().copied()) + } + /// Resolve a `row_id` to every `doc_id` it owns. /// /// Row-document indexes map each row to a single document. Element-document /// indexes (and older list indexes) can map one row to several documents, /// so a single `row_id` may own multiple `doc_id`s sharing that key in `inv`. /// The prefilter path (`flat_search`) walks an allow-list of row_ids and - /// must evaluate all legacy documents for that row. + /// must evaluate all of those documents for that row. + /// + /// Only the `inv` branch answers with several. A no-metadata legacy set + /// holds one document per row, so the branch below resolves it with a single + /// binary search. pub fn doc_ids(&self, row_id: u64) -> impl Iterator + '_ { if self.inv.is_empty() { // in legacy format, the row id is doc id (one document per row) @@ -455,6 +526,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), } } @@ -519,6 +591,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }); } @@ -563,6 +636,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }); } @@ -585,6 +659,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }) } @@ -602,6 +677,7 @@ impl DocSet { .map(|_| Vec::with_capacity(len)) .collect(); self.invalidate_norms(); + self.invalidate_row_ids_ascending(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { match mapping.get(row_id) { @@ -681,6 +757,41 @@ impl DocSet { .unwrap_or(0) } + /// Total document length of `row_id`: the sum of `num_tokens` over every + /// document the row owns, or `0` when the row is absent (e.g. an empty or + /// null field that was never indexed). + /// + /// Lets cross-field BM25F read a candidate's per-column length with a + /// targeted lookup instead of scanning the whole `DocSet`. The result is + /// exact against a full `iter()` scan filtered to the row: every document + /// the row owns contributes, whether the row is resolved through the sorted + /// `inv` index or, on a legacy set, through the run of equal `row_ids`. + /// + /// A row owns several documents only on the `inv` path (see [`Self::doc_ids`]). + /// A no-metadata legacy set cannot reach that state, because the writer of that + /// era keyed its documents by row id in a map, so the elements of one row + /// collapsed before they were written. The legacy branch below sums a run + /// anyway, so the two paths answer alike whatever they are handed. + #[inline] + pub fn doc_length_by_row_id(&self, row_id: u64) -> u64 { + if self.inv.is_empty() { + // Legacy: row id == doc id, `row_ids` is sorted; duplicate row ids + // (one per indexed list element) form a contiguous run. + let lo = self.row_ids.partition_point(|&id| id < row_id); + let hi = self.row_ids.partition_point(|&id| id <= row_id); + (lo..hi).map(|doc_id| self.num_tokens[doc_id] as u64).sum() + } else { + // Compressed: `inv` is sorted by row id and holds one entry per + // document owned by the row. + let lo = self.inv.partition_point(|entry| entry.0 < row_id); + let hi = self.inv.partition_point(|entry| entry.0 <= row_id); + self.inv[lo..hi] + .iter() + .map(|entry| self.num_tokens[entry.1 as usize] as u64) + .sum() + } + } + // append a document to the doc set // returns the doc_id (the number of documents before appending) pub fn append(&mut self, row_id: u64, num_tokens: u32) -> u32 { @@ -688,6 +799,7 @@ impl DocSet { self.num_tokens.push(num_tokens); self.total_tokens += num_tokens as u64; self.invalidate_norms(); + self.invalidate_row_ids_ascending(); self.row_ids.len() as u32 - 1 } @@ -714,6 +826,7 @@ impl DocSet { } self.total_tokens += num_tokens as u64; self.invalidate_norms(); + self.invalidate_row_ids_ascending(); Ok(self.row_ids.len() as u32 - 1) } @@ -725,6 +838,14 @@ impl DocSet { } } + // Drop the memoized ascending-row_ids answer after a mutation; it + // recomputes on the next query use. + fn invalidate_row_ids_ascending(&mut self) { + if self.row_ids_ascending.get().is_some() { + self.row_ids_ascending = Arc::new(std::sync::OnceLock::new()); + } + } + pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.memory_size() diff --git a/rust/lance-index/src/scalar/inverted/index/partition.rs b/rust/lance-index/src/scalar/inverted/index/partition.rs index b2b26e85825..d680028fe0a 100644 --- a/rust/lance-index/src/scalar/inverted/index/partition.rs +++ b/rust/lance-index/src/scalar/inverted/index/partition.rs @@ -692,6 +692,56 @@ impl InvertedPartition { self.inverted_list.is_legacy_layout() } + /// This partition's `(num_rows, per_term_row_freq)` counted over distinct + /// row ids rather than documents. See + /// [`InvertedIndex::bm25_row_stats_for_terms`] for why the distinction + /// exists and when this runs. + /// + /// A partition whose documents map one-to-one onto rows takes the same + /// single-metadata-row `posting_len_for_token` lookup the document-granularity + /// statistics use, so only a partition that really indexed one document per + /// list element pays for reading posting lists. + pub(super) async fn row_stats_for_terms( + &self, + terms: &[String], + metrics: Option<&dyn MetricsCollector>, + ) -> Result<(usize, Vec)> { + let docs = self.docs.address_keyed().await?; + let num_rows = docs.num_distinct_rows(); + let one_document_per_row = num_rows == docs.len(); + let is_legacy = self.is_legacy(); + let mut row_freqs = Vec::with_capacity(terms.len()); + for term in terms { + let Some(token_id) = self.tokens.get(term) else { + row_freqs.push(0); + continue; + }; + if one_document_per_row { + row_freqs.push( + self.inverted_list + .posting_len_for_token(token_id, metrics) + .await?, + ); + continue; + } + // Deduplicating needs the postings themselves; `posting_len_for_token` + // only knows how many documents there are. The read is cached, and a + // partition that reaches this branch forces the full-read fallback in + // `combined_fields_search` anyway (legacy layout or non-ascending + // row_ids), so nothing that would otherwise have been pruned is read. + let posting = self + .inverted_list + .posting_list(token_id, false, metrics.unwrap_or(&NoOpMetricsCollector)) + .await?; + let mut rows = RoaringTreemap::new(); + for (row_id, _) in live_posting_rows(&posting, &docs, is_legacy) { + rows.insert(row_id); + } + row_freqs.push(rows.len() as usize); + } + Ok((num_rows, row_freqs)) + } + pub async fn load( store: Arc, id: u64, diff --git a/rust/lance-index/src/scalar/inverted/index/search.rs b/rust/lance-index/src/scalar/inverted/index/search.rs index 78a533d7acb..be3bd144ecd 100644 --- a/rust/lance-index/src/scalar/inverted/index/search.rs +++ b/rust/lance-index/src/scalar/inverted/index/search.rs @@ -116,6 +116,10 @@ impl InvertedIndex { Ok(MemBM25Scorer::new(total_tokens, num_docs, token_docs)) } + /// Corpus statistics at document granularity: `(total_tokens, docCount, + /// per_term_docFreq)`, where a document is one entry of a partition's + /// [`DocSet`]. This is what the single-column scoring path pairs with, since + /// wand scores one posting per document. pub async fn bm25_stats_for_terms( &self, terms: &[String], @@ -179,6 +183,67 @@ impl InvertedIndex { Ok(Some((total_tokens, num_docs, token_docs))) } + /// Corpus statistics at row granularity: `(total_tokens, docCount, + /// per_term_docFreq)` counted over distinct row ids. + /// + /// Cross-field BM25F scores rows, not documents: `combined_fields_search` + /// sums all of a row's postings into one `tf'` and all of its document + /// lengths into one `dl'`. Released V1/V2 indexes indexed each + /// `List` element as its own document, so pairing that row-granular + /// `tf'`/`dl'` with [`Self::bm25_stats_for_terms`] would measure it against a + /// corpus of a different size: `idf'` and `avgdl'` would describe elements + /// while the frequencies describe rows, shifting an old index's top-k relative + /// to the same data reindexed on a current build. + /// + /// V3 indexes always map one row to one document, so they delegate to + /// [`Self::bm25_stats_for_terms`] and keep its IO profile: no `DocSet` load, + /// one posting-metadata row per term and partition. `total_tokens` is + /// granularity-independent, so it comes from the shared + /// `aggregate_corpus_stats` cache either way. + /// + /// `metrics`, when provided, receives the posting-metadata cache lookups these + /// statistics trigger, so both granularities report comparable + /// `index_cache_hits`/`index_cache_misses`. + pub async fn bm25_row_stats_for_terms( + &self, + terms: &[String], + metrics: Option<&dyn MetricsCollector>, + ) -> Result<(u64, usize, Vec)> { + if !matches!( + self.format_version, + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2 + ) { + return self.bm25_stats_for_terms(terms, metrics).await; + } + + let (total_tokens, _) = self.aggregate_corpus_stats().await?; + let io_parallelism = self.store.io_parallelism(); + let futures = self + .partitions + .iter() + .map(|partition| { + let partition = partition.clone(); + async move { partition.row_stats_for_terms(terms, metrics).await } + }) + .collect::>(); + let per_partition: Vec<(usize, Vec)> = stream::iter(futures) + .buffer_unordered(io_parallelism) + .try_collect() + .await?; + + let mut num_rows = 0usize; + let mut row_freqs = vec![0usize; terms.len()]; + for (partition_rows, partition_freqs) in per_partition { + // Partitions hold disjoint row sets, so per-partition distinct counts + // add up to the index's distinct count. + num_rows += partition_rows; + for (total, partition_freq) in row_freqs.iter_mut().zip(partition_freqs) { + *total += partition_freq; + } + } + Ok((total_tokens, num_rows, row_freqs)) + } + /// Aggregate immutable per-partition corpus statistics. New modern files /// read both values from the already-opened docs footer; older partitioned /// files scan `_num_tokens` once as a compatibility fallback. diff --git a/rust/lance-index/src/scalar/inverted/index/tests/format_and_builder.rs b/rust/lance-index/src/scalar/inverted/index/tests/format_and_builder.rs index fa6716e9c16..b0d625b44b2 100644 --- a/rust/lance-index/src/scalar/inverted/index/tests/format_and_builder.rs +++ b/rust/lance-index/src/scalar/inverted/index/tests/format_and_builder.rs @@ -457,6 +457,67 @@ fn test_cached_num_tokens_uses_supplied_total_and_full_stays_owned() { assert_eq!(full.row_id(1), 20); } +#[test] +fn test_combined_fields_append_invalidates_row_ids_ascending_cache() { + let row_ids = UInt64Array::from(vec![10, 20, 30]); + let num_tokens = UInt32Array::from(vec![1, 1, 1]); + let mut docs = DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap(); + + // Memoize the ascending answer for the current (sorted) row_ids. + assert!(docs.row_ids_strictly_ascending()); + + // Appending a smaller row_id breaks the invariant; the cached answer + // must be dropped so the recomputed value reflects the mutation. + docs.append(5, 1); + assert!(!docs.row_ids_strictly_ascending()); + + // The element-document append is the same mutation and owes the same + // invalidation. A set built by `with_coordinate_rank` is shared by its + // clones, so a stale answer would outlive the builder that produced it. + let mut docs = DocSet::with_coordinate_rank(1); + for (doc_id, row_id) in [10u64, 20, 30].into_iter().enumerate() { + assert_eq!( + docs.append_with_doc_index(row_id, 1, &[doc_id as u32]) + .unwrap(), + doc_id as u32 + ); + } + assert!(docs.row_ids_strictly_ascending()); + + docs.append_with_doc_index(5, 1, &[0]).unwrap(); + assert!(!docs.row_ids_strictly_ascending()); +} + +#[test] +fn test_doc_length_by_row_id_matches_scan() { + // Compressed layout: row ids are stored in doc-id order and resolved + // through `inv`. The targeted lookup must equal an `iter()` scan that + // sums the matching row's lengths, and return 0 for an absent row. + let row_ids = UInt64Array::from(vec![30, 10, 20]); + let num_tokens = UInt32Array::from(vec![8, 3, 5]); + let docs = DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap(); + for target in [10u64, 20, 30, 99] { + let expected: u64 = docs + .iter() + .filter(|(id, _)| **id == target) + .map(|(_, nt)| *nt as u64) + .sum(); + assert_eq!(docs.doc_length_by_row_id(target), expected, "row {target}"); + } + assert_eq!(docs.doc_length_by_row_id(10), 3); + assert_eq!(docs.doc_length_by_row_id(99), 0); + + // Legacy layout: row id == doc id, row ids are sorted, and a row indexed + // as several list documents forms a contiguous run whose lengths sum. + let legacy_row_ids = UInt64Array::from(vec![10, 10, 20]); + let legacy_num_tokens = UInt32Array::from(vec![3, 4, 5]); + let legacy = DocSet::from_columns(&legacy_row_ids, &legacy_num_tokens, true, None).unwrap(); + assert!(legacy.inv.is_empty()); + assert_eq!(legacy.doc_length_by_row_id(10), 7); + assert_eq!(legacy.doc_length_by_row_id(20), 5); + assert_eq!(legacy.doc_length_by_row_id(99), 0); +} + #[test] fn test_posting_builder_writes_impacts_for_supported_block_sizes() { for block_size in [128, 256] { diff --git a/rust/lance-index/src/scalar/inverted/index/tests/stats.rs b/rust/lance-index/src/scalar/inverted/index/tests/stats.rs index 3519c56547e..5c376d78364 100644 --- a/rust/lance-index/src/scalar/inverted/index/tests/stats.rs +++ b/rust/lance-index/src/scalar/inverted/index/tests/stats.rs @@ -331,6 +331,50 @@ async fn test_bm25_stats_for_terms_reuses_posting_metadata_cache() { ); } +/// Row-granularity statistics must not cost the posting file any extra IO +/// when a partition's documents already map one-to-one onto rows, which is +/// every V3 index and every V1/V2 index that is not a legacy list index. The +/// deduplicating branch reads posting lists, so it has to stay behind the +/// duplicate-row-id check rather than run for every legacy-format index. +#[tokio::test] +async fn test_bm25_row_stats_for_terms_keeps_the_lazy_posting_metadata_path() { + let (index, counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; + // A V3 index would take the delegating early return, which proves nothing + // about the legacy branch this test exists for. + assert!( + matches!( + index.format_version(), + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2 + ), + "expected a legacy format version, got {:?}", + index.format_version(), + ); + + let terms = ["t0".to_string()]; + let documents = index.bm25_stats_for_terms(&terms, None).await.unwrap(); + assert_eq!(documents, (100, 100, vec![1])); + let metadata_rows = counter.metadata_rows_read(); + let rows = counter.rows_read(); + assert_eq!(metadata_rows, 1); + + assert_eq!( + index.bm25_row_stats_for_terms(&terms, None).await.unwrap(), + documents, + "one document per row leaves the two granularities identical", + ); + assert_eq!( + counter.metadata_rows_read() - metadata_rows, + 1, + "row statistics should read one metadata row per (term, partition)", + ); + assert_eq!( + counter.rows_read() - rows, + 1, + "row statistics must not read the posting list itself (got {} extra rows)", + counter.rows_read() - rows, + ); +} + #[tokio::test] async fn test_bm25_stats_for_terms_records_metadata_cache_stats() { let cache = LanceCache::with_capacity(1024 * 1024); @@ -360,6 +404,48 @@ async fn test_bm25_stats_for_terms_records_metadata_cache_stats() { assert_eq!(warm.index_cache_hits(), terms.len()); } +/// Row-granularity statistics go through `InvertedPartition::row_stats_for_terms` +/// on a V1/V2 index, and its posting-metadata lookups must reach the caller's +/// collector too: the cross-field scorer build is the only consumer, and it +/// runs under an `ExecutionPlan` whose cache counters would otherwise miss +/// them entirely. +#[tokio::test] +async fn test_bm25_row_stats_for_terms_records_metadata_cache_stats() { + let cache = LanceCache::with_capacity(1024 * 1024); + let (index, _counter, _tmpdir) = load_counted_v2_index(100, cache.clone()).await; + // A V3 index would delegate to the document-granularity path, which + // already has its own coverage. + assert!( + matches!( + index.format_version(), + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2 + ), + "expected a legacy format version, got {:?}", + index.format_version(), + ); + + let terms = ["t0".to_string(), "t1".to_string(), "t2".to_string()]; + let cold = LocalMetricsCollector::default(); + let cold_stats = index + .bm25_row_stats_for_terms(&terms, Some(&cold)) + .await + .unwrap(); + assert_eq!(cold_stats, (100, 100, vec![1, 1, 1])); + assert_eq!(cold.index_cache_misses(), terms.len()); + assert_eq!(cold.index_cache_hits(), 0); + + let warm = LocalMetricsCollector::default(); + assert_eq!( + index + .bm25_row_stats_for_terms(&terms, Some(&warm)) + .await + .unwrap(), + cold_stats, + ); + assert_eq!(warm.index_cache_misses(), 0); + assert_eq!(warm.index_cache_hits(), terms.len()); +} + #[tokio::test] async fn test_aggregate_corpus_stats_reuses_cached_value() { let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; diff --git a/rust/lance-index/src/scalar/inverted/parser.rs b/rust/lance-index/src/scalar/inverted/parser.rs index 03f39e66ad0..8993522b517 100644 --- a/rust/lance-index/src/scalar/inverted/parser.rs +++ b/rust/lance-index/src/scalar/inverted/parser.rs @@ -3,7 +3,8 @@ use super::DocumentGranularity; use super::query::{ - BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery, + BooleanQuery, BoostQuery, CombinedFieldsQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, + Operator, PhraseQuery, }; use lance_core::{Error, Result}; use serde_json::Value; @@ -128,6 +129,54 @@ impl JsonParser for MultiMatchQuery { } } +impl JsonParser for CombinedFieldsQuery { + fn from_json(value: &Value) -> Result { + let terms = value["query"] + .as_str() + .ok_or_else(|| Error::invalid_input("missing query in combined_fields query"))? + .to_string(); + let columns = value["columns"] + .as_array() + .ok_or_else(|| Error::invalid_input("missing columns in combined_fields query"))? + .iter() + .map(|v| { + v.as_str().map(String::from).ok_or_else(|| { + Error::invalid_input( + "columns must be an array of strings in combined_fields query", + ) + }) + }) + .collect::>>()?; + + let query = Self::try_new(terms, columns)?; + + let query = match value.get("boost") { + Some(Value::Array(boosts)) => { + let boosts = boosts + .iter() + .map(|v| { + v.as_f64().map(|f| f as f32).ok_or_else(|| { + Error::invalid_input( + "boost must be an array of numbers in combined_fields query", + ) + }) + }) + .collect::>>()?; + query.try_with_boosts(boosts)? + } + _ => query, + }; + + let operator = value["operator"] + .as_str() + .map(Operator::try_from) + .transpose()? + .unwrap_or_default(); + + Ok(query.with_operator(operator)) + } +} + impl JsonParser for BooleanQuery { fn from_json(value: &Value) -> Result { let mut clauses = Vec::new(); @@ -171,6 +220,9 @@ fn from_json_value(value: &Value) -> Result { "phrase" => Ok(FtsQuery::Phrase(PhraseQuery::from_json(query_val)?)), "boost" => Ok(FtsQuery::Boost(BoostQuery::from_json(query_val)?)), "multi_match" => Ok(FtsQuery::MultiMatch(MultiMatchQuery::from_json(query_val)?)), + "combined_fields" => Ok(FtsQuery::CombinedFields(CombinedFieldsQuery::from_json( + query_val, + )?)), "boolean" => Ok(FtsQuery::Boolean(BooleanQuery::from_json(query_val)?)), _ => Err(Error::invalid_input(format!( "unknown fts query type: {}", @@ -344,4 +396,54 @@ mod tests { ])); assert_eq!(fts_query, expected_query); } + + #[test] + fn test_from_json_combined_fields() { + let json = r#" + { + "combined_fields": { + "query": "hello world", + "columns": ["title", "body"], + "boost": [2.0, 1.0], + "operator": "and" + } + }"#; + let fts_query = from_json(json).unwrap(); + let expected = CombinedFieldsQuery::try_new( + "hello world".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap() + .try_with_boosts(vec![2.0, 1.0]) + .unwrap() + .with_operator(Operator::And); + assert_eq!(fts_query, FtsQuery::CombinedFields(expected)); + } + + #[test] + fn test_from_json_combined_fields_defaults_and_validation() { + // boost + operator omitted: weights default to 1.0, operator to Or. + let json = r#"{ "combined_fields": { "query": "hi", "columns": ["a", "b"] } }"#; + let FtsQuery::CombinedFields(query) = from_json(json).unwrap() else { + panic!("expected combined_fields query"); + }; + assert_eq!( + query.weighted_columns().collect::>(), + vec![("a", 1.0), ("b", 1.0)] + ); + assert_eq!(query.operator(), Operator::Or); + + // A weight below 1 is rejected. + let json = r#"{ "combined_fields": { "query": "hi", "columns": ["a", "b"], "boost": [0.5, 1.0] } }"#; + assert!(from_json(json).is_err()); + + // So is a weight above the upper bound that keeps the blended length + // finite. + let json = r#"{ "combined_fields": { "query": "hi", "columns": ["a", "b"], "boost": [1.0, 1e30] } }"#; + assert!(from_json(json).is_err()); + + // Empty columns are rejected. + let json = r#"{ "combined_fields": { "query": "hi", "columns": [] } }"#; + assert!(from_json(json).is_err()); + } } diff --git a/rust/lance-index/src/scalar/inverted/query.rs b/rust/lance-index/src/scalar/inverted/query.rs index 9952c36cc55..ebd2ff740a5 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -164,6 +164,7 @@ pub enum FtsQuery { // compound queries Boost(BoostQuery), MultiMatch(MultiMatchQuery), + CombinedFields(CombinedFieldsQuery), Boolean(BooleanQuery), } @@ -178,6 +179,7 @@ impl std::fmt::Display for FtsQuery { query.positive, query.negative, query.negative_boost ), Self::MultiMatch(query) => write!(f, "MultiMatch({:?})", query), + Self::CombinedFields(query) => write!(f, "CombinedFields({:?})", query), Self::Boolean(query) => { write!( f, @@ -206,6 +208,7 @@ impl FtsQueryNode for FtsQuery { } columns } + Self::CombinedFields(query) => query.columns(), Self::Boolean(query) => query.columns(), } } @@ -218,6 +221,7 @@ impl FtsQuery { Self::Phrase(query) => format!("\"{}\"", query.terms), // Phrase queries are quoted Self::Boost(query) => query.positive.query(), Self::MultiMatch(query) => query.match_queries[0].terms.clone(), + Self::CombinedFields(query) => query.terms().to_string(), Self::Boolean(_) => { // Bool queries don't have a single query string, they are composed of multiple queries String::new() @@ -233,6 +237,9 @@ impl FtsQuery { query.positive.is_missing_column() || query.negative.is_missing_column() } Self::MultiMatch(query) => query.match_queries.iter().any(|q| q.column.is_none()), + // `try_new` rejects an empty column list, so the target columns of a + // combined_fields query are always known. + Self::CombinedFields(_) => false, Self::Boolean(query) => { query.must.iter().any(|q| q.is_missing_column()) || query.should.iter().any(|q| q.is_missing_column()) @@ -262,6 +269,11 @@ impl FtsQuery { .collect(); Self::MultiMatch(MultiMatchQuery { match_queries }) } + Self::CombinedFields(query) => { + // combined_fields carries all target columns (and their per-column + // boosts) at construction, so a single-column override is a no-op. + Self::CombinedFields(query) + } Self::Boolean(query) => { let must = query .must @@ -312,6 +324,12 @@ impl From for FtsQuery { } } +impl From for FtsQuery { + fn from(query: CombinedFieldsQuery) -> Self { + Self::CombinedFields(query) + } +} + impl From for FtsQuery { fn from(query: BooleanQuery) -> Self { Self::Boolean(query) @@ -631,6 +649,202 @@ impl FtsQueryNode for MultiMatchQuery { } } +/// A cross-field (BM25F) full-text query, exposing Elasticsearch's +/// `combined_fields` semantics: the target columns are treated as one virtual +/// field so term statistics (document frequency, term frequency, and document +/// length) are blended across fields rather than scored independently. +/// +/// This contrasts with [`MultiMatchQuery`], which fans out into one +/// [`MatchQuery`] per column and fuses the per-field scores by taking the +/// maximum (Elasticsearch `best_fields`). A `CombinedFieldsQuery` stays a single +/// node and is scored once over the blended statistics. +/// +/// Per-column weights follow Lucene's `CombinedFieldQuery`; see +/// [`Self::try_with_boosts`] for the accepted range. +/// +/// Each column is stored together with its weight and every field is private, so +/// the pairing cannot go out of sync: [`Self::try_new`] and +/// [`Self::try_with_boosts`] are the only constructors and no accessor hands out a +/// `&mut`. +#[derive(Debug, Clone, PartialEq)] +pub struct CombinedFieldsQuery { + /// The columns combined into the virtual field, each paired with its BM25F + /// weight `w_f`. Non-empty, column names unique, weights in + /// `[MIN_BOOST, MAX_BOOST]`. + weighted_columns: Vec<(String, f32)>, + /// The query string, tokenized once and matched against every target column. + terms: String, + /// How to combine terms: `And` (all terms must match) or `Or` (default). + operator: Operator, +} + +impl CombinedFieldsQuery { + /// Minimum allowed per-column weight (Lucene `CombinedFieldQuery` constraint). + const MIN_BOOST: f32 = 1.0; + + /// Maximum allowed per-column weight. + /// + /// The blended length `dl'(d) = Σ_f w_f · dl_f(d)` and the identically shaped + /// `tf'` are accumulated in `f32`, and a per-column document length is a `u32` + /// token count. One term of that sum is therefore at most + /// `2^20 · (2^32 - 1) < 2^52`, so with `f32::MAX ≈ 2^128` it would take more + /// than `2^76` maximally long columns to reach infinity. The column count is + /// bounded by the schema, many orders of magnitude below that, so an accepted + /// weight cannot turn a score into `Inf`/`NaN`. + const MAX_BOOST: f32 = 1_048_576.0; // 2^20 + + /// Create a combined-fields query over `columns`, with every weight + /// defaulting to `1.0` and the `Or` operator. + /// + /// Returns an error if `columns` is empty or contains duplicates. + pub fn try_new(terms: String, columns: Vec) -> Result { + if columns.is_empty() { + return Err(Error::invalid_input( + "Cannot create CombinedFieldsQuery with no columns".to_string(), + )); + } + // A duplicate column would double-count its postings and inflate + // sum_total_term_freq, skewing the blended BM25F statistics. + let mut seen = HashSet::with_capacity(columns.len()); + for column in &columns { + if !seen.insert(column.as_str()) { + return Err(Error::invalid_input(format!( + "CombinedFieldsQuery columns must be unique, but '{}' is duplicated", + column + ))); + } + } + Ok(Self { + weighted_columns: columns + .into_iter() + .map(|column| (column, Self::MIN_BOOST)) + .collect(), + terms, + operator: Operator::Or, + }) + } + + /// Set per-column weights, positionally aligned with the columns passed to + /// [`Self::try_new`]. + /// + /// Returns an error if the number of boosts does not match the number of + /// columns, or if any weight falls outside `[1, 2^20]`. The lower bound + /// mirrors Lucene's `CombinedFieldQuery`, which requires `weight >= 1` so the + /// combined length norm stays additive; the upper bound keeps the blended + /// length finite in `f32`. `NaN` and infinities are outside the range and so + /// rejected too. + pub fn try_with_boosts(mut self, boosts: Vec) -> Result { + if boosts.len() != self.weighted_columns.len() { + return Err(Error::invalid_input(format!( + "The number of boosts ({}) must match the number of columns ({})", + boosts.len(), + self.weighted_columns.len() + ))); + } + for ((column, _), &boost) in self.weighted_columns.iter().zip(&boosts) { + if !(Self::MIN_BOOST..=Self::MAX_BOOST).contains(&boost) { + return Err(Error::invalid_input(format!( + "combined_fields boost for column '{}' must be a finite value in [{}, {}], got {}", + column, + Self::MIN_BOOST, + Self::MAX_BOOST, + boost + ))); + } + } + for ((_, weight), boost) in self.weighted_columns.iter_mut().zip(boosts) { + *weight = boost; + } + Ok(self) + } + + /// Set the operator used to combine terms. + pub fn with_operator(mut self, operator: Operator) -> Self { + self.operator = operator; + self + } + + /// The query string, tokenized once and matched against every target column. + pub fn terms(&self) -> &str { + &self.terms + } + + /// How the query terms are combined. + pub fn operator(&self) -> Operator { + self.operator + } + + /// The target columns in query order. + pub fn column_names(&self) -> impl ExactSizeIterator { + self.weighted_columns + .iter() + .map(|(column, _)| column.as_str()) + } + + /// Each target column paired with its BM25F weight, in query order. + /// + /// Iterating the pairs is the only way to read the weights, so an execution + /// path cannot pair a column with the wrong weight or silently drop a column + /// that has no weight. + pub fn weighted_columns(&self) -> impl ExactSizeIterator { + self.weighted_columns + .iter() + .map(|(column, weight)| (column.as_str(), *weight)) + } +} + +impl Serialize for CombinedFieldsQuery { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let mut map = serializer.serialize_map(Some(4))?; + map.serialize_entry("query", &self.terms)?; + map.serialize_entry("columns", &self.column_names().collect::>())?; + map.serialize_entry( + "boost", + &self + .weighted_columns() + .map(|(_, weight)| weight) + .collect::>(), + )?; + map.serialize_entry("operator", &self.operator)?; + map.end() + } +} + +impl<'de> Deserialize<'de> for CombinedFieldsQuery { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + struct CombinedFieldsQueryData { + query: String, + columns: Vec, + boost: Option>, + #[serde(default)] + operator: Operator, + } + + let data = CombinedFieldsQueryData::deserialize(deserializer)?; + let query = Self::try_new(data.query, data.columns).map_err(serde::de::Error::custom)?; + let query = match data.boost { + Some(boosts) => query + .try_with_boosts(boosts) + .map_err(serde::de::Error::custom)?, + None => query, + }; + Ok(query.with_operator(data.operator)) + } +} + +impl FtsQueryNode for CombinedFieldsQuery { + fn columns(&self) -> HashSet { + self.column_names().map(String::from).collect() + } +} + pub enum Occur { /// The clause may match and contributes its score when it does. Should, @@ -984,6 +1198,11 @@ pub fn fill_fts_query_column( .collect(); Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries })) } + FtsQuery::CombinedFields(combined_query) => { + // combined_fields carries its target columns (and per-column boosts) + // at construction, so there is nothing to fill or replace. + Ok(FtsQuery::CombinedFields(combined_query.clone())) + } FtsQuery::Boolean(bool_query) => { let must = bool_query .must @@ -1318,4 +1537,203 @@ mod tests { let query = MatchQuery::new("hello".to_string()); assert!(BooleanMatchPlan::try_build(&FtsQuery::Match(query)).is_none()); } + + #[test] + fn test_combined_fields_query_serde() { + use super::*; + use serde_json::json; + + let query = CombinedFieldsQuery::try_new( + "hello world".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap() + .try_with_boosts(vec![2.0, 1.0]) + .unwrap() + .with_operator(Operator::And); + + // Serializes with multi_match-style keys, plus a round-tripped operator. + let serialized = serde_json::to_value(&query).unwrap(); + let expected = json!({ + "query": "hello world", + "columns": ["title", "body"], + "boost": [2.0, 1.0], + "operator": "And", + }); + assert_eq!(serialized, expected); + + // The wire format is a contract, so pin the exact bytes (key order + // included), not just the equivalent `Value`. + assert_eq!( + serde_json::to_string(&query).unwrap(), + r#"{"query":"hello world","columns":["title","body"],"boost":[2.0,1.0],"operator":"And"}"# + ); + + let deserialized: CombinedFieldsQuery = serde_json::from_value(serialized).unwrap(); + assert_eq!(deserialized, query); + + // Round-trips as a wrapped FtsQuery variant under the "combined_fields" tag. + let wrapped = FtsQuery::CombinedFields(query); + let value = serde_json::to_value(&wrapped).unwrap(); + assert!(value.get("combined_fields").is_some()); + let round_trip: FtsQuery = serde_json::from_value(value).unwrap(); + assert_eq!(round_trip, wrapped); + } + + #[test] + fn test_combined_fields_query_defaults() { + use super::*; + use serde_json::json; + + // Omitting boost + operator defaults weights to 1.0 and the operator to Or. + let value = json!({ + "query": "hello", + "columns": ["title", "body"], + }); + let query: CombinedFieldsQuery = serde_json::from_value(value).unwrap(); + assert_eq!(query.terms(), "hello"); + assert_eq!( + query.weighted_columns().collect::>(), + vec![("title", 1.0), ("body", 1.0)] + ); + assert_eq!(query.operator(), Operator::Or); + } + + #[test] + fn test_combined_fields_query_validation() { + use super::*; + + // Assert the result is an invalid-input error whose message names the + // rejected cause, not just that it failed. + let assert_invalid_input = |result: Result, needle: &str| { + let err = result.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + assert!( + err.to_string().contains(needle), + "error {err:?} should mention {needle:?}" + ); + }; + + // Empty columns are rejected. + assert_invalid_input( + CombinedFieldsQuery::try_new("hello".to_string(), vec![]), + "no columns", + ); + + // Duplicate columns are rejected (they would double-count postings). + assert_invalid_input( + CombinedFieldsQuery::try_new( + "hello".to_string(), + vec!["title".to_string(), "title".to_string()], + ), + "duplicated", + ); + + let query = CombinedFieldsQuery::try_new( + "hello".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap(); + + // A boost count that does not match the column count is rejected in both + // directions, and the message names both lengths. + assert_invalid_input( + query.clone().try_with_boosts(vec![1.0]), + "number of boosts (1) must match the number of columns (2)", + ); + assert_invalid_input( + query.clone().try_with_boosts(vec![1.0, 1.0, 1.0]), + "number of boosts (3) must match the number of columns (2)", + ); + + // Weights outside [1, 2^20] are rejected: below 1 breaks Lucene's + // additive length norm, above 2^20 could overflow the blended length. + // NaN and the infinities fall outside the range as well. + let out_of_range = "must be a finite value in [1, 1048576]"; + for boosts in [ + vec![1.0, 0.5], + vec![0.0, 1.0], + vec![f32::NAN, 1.0], + vec![f32::INFINITY, 1.0], + vec![f32::NEG_INFINITY, 1.0], + vec![1.0, f32::MAX], + vec![1.0, 1_048_577.0], + ] { + assert_invalid_input(query.clone().try_with_boosts(boosts), out_of_range); + } + // The message names the offending column and its value. + assert_invalid_input(query.clone().try_with_boosts(vec![1.0, 0.5]), "'body'"); + assert_invalid_input(query.clone().try_with_boosts(vec![1.0, 0.5]), "got 0.5"); + + // Fractional weights >= 1 are accepted, and so is the upper bound itself. + assert!(query.clone().try_with_boosts(vec![1.5, 2.0]).is_ok()); + assert!(query.try_with_boosts(vec![1.0, 1_048_576.0]).is_ok()); + } + + /// The columns and their weights cannot be desynced by any operation the type + /// permits, which is what the execution boundary relies on when it iterates + /// [`CombinedFieldsQuery::weighted_columns`]. + /// + /// Field privacy carries the structural half and the compiler enforces it, so + /// the only thing left to check at runtime is that `try_new` and + /// `try_with_boosts`, the sole mutation paths, keep the pairs aligned. + #[test] + fn test_combined_fields_query_pairing_is_stable() { + use super::*; + + let query = CombinedFieldsQuery::try_new( + "hello".to_string(), + vec!["title".to_string(), "body".to_string(), "tags".to_string()], + ) + .unwrap(); + // Defaults: one weight per column, in query order. + assert_eq!( + query.weighted_columns().collect::>(), + vec![("title", 1.0), ("body", 1.0), ("tags", 1.0)] + ); + + let query = query.try_with_boosts(vec![3.0, 1.0, 2.5]).unwrap(); + assert_eq!( + query.weighted_columns().collect::>(), + vec![("title", 3.0), ("body", 1.0), ("tags", 2.5)] + ); + // The name-only view and the pair view agree on order and length, so a + // consumer that reads either sees the same columns. + assert_eq!( + query.column_names().collect::>(), + query + .weighted_columns() + .map(|(column, _)| column) + .collect::>() + ); + assert_eq!(query.column_names().len(), 3); + // The `FtsQueryNode` set view still reports the target columns, and is not + // shadowed by the inherent accessors. + assert_eq!( + FtsQueryNode::columns(&query), + HashSet::from(["title".to_string(), "body".to_string(), "tags".to_string()]) + ); + + // A rejected boost list leaves no partially updated query behind: the + // builder consumes `self`, so the caller keeps the validated value. + let rejected = query.clone().try_with_boosts(vec![1.0, f32::MAX, 1.0]); + assert!(rejected.is_err()); + assert_eq!( + query.weighted_columns().collect::>(), + vec![("title", 3.0), ("body", 1.0), ("tags", 2.5)] + ); + + // Clone/PartialEq compare the pairs, so a differing weight is a differing + // query even when the column lists match. + assert_eq!(query.clone(), query); + let reweighted = query.clone().try_with_boosts(vec![3.0, 1.0, 2.0]).unwrap(); + assert_ne!(reweighted, query); + // Debug renders the columns with their weights. + let debug = format!("{:?}", query); + assert!(debug.contains("title"), "unexpected Debug output: {debug}"); + assert!(debug.contains("2.5"), "unexpected Debug output: {debug}"); + } } diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 3523387936b..7db1ea36ca6 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -239,6 +239,94 @@ pub fn idf(token_docs: usize, num_docs: usize) -> f32 { ((num_docs - token_docs as f32 + 0.5) / (token_docs as f32 + 0.5) + 1.0).ln() } +/// BM25F scorer over a virtual field formed by combining several columns +/// (Lucene `CombinedFieldQuery` / Elasticsearch `combined_fields`). +/// +/// Where [`MemBM25Scorer`] / `IndexBM25Scorer` score one column against that +/// column's own statistics, this scorer holds statistics blended across the +/// target columns per the BM25F rules: +/// - `doc_count` = `max_f docCount_f` +/// - `avg_doc_length` = `sumTotalTermFreq' / doc_count`, where +/// `sumTotalTermFreq' = Σ_f w_f · sumTotalTermFreq_f` +/// - `doc_freq(t)` = `max_f docFreq_f(t)` +/// +/// The caller supplies the blended term frequency `tf' = Σ_f w_f · tf_f(t, d)` +/// and blended document length `dl' = Σ_f w_f · dl_f(d)` per candidate document; +/// this type turns those into an IDF weight and a BM25 document weight. Lance's +/// `(k1 + 1)` numerator is kept (a constant factor vs. Lucene that preserves +/// ranking). +#[derive(Debug, Clone)] +pub struct CombinedFieldsBM25Scorer { + doc_count: usize, + avg_doc_length: f32, + doc_freq: HashMap, +} + +impl CombinedFieldsBM25Scorer { + pub fn new(doc_count: usize, avg_doc_length: f32, doc_freq: HashMap) -> Self { + Self { + doc_count, + avg_doc_length, + doc_freq, + } + } + + pub fn doc_count(&self) -> usize { + self.doc_count + } + + pub fn avg_doc_length(&self) -> f32 { + self.avg_doc_length + } + + /// Blended IDF for `term` over the virtual field. Returns `0.0` for a term + /// absent from every target column (`docFreq' == 0`), and for the inconsistent + /// statistics (`docFreq' > docCount'`) that drive `idf` negative or, once + /// `docFreq'` is large enough for the ratio to round to `-1`, to `-inf`. A + /// `-inf` weight times a zero document weight is `NaN`; see [`Self::doc_weight`] + /// for why no score may be non-finite. + pub fn query_weight(&self, term: &str) -> f32 { + match self.doc_freq.get(term).copied() { + Some(token_docs) if token_docs > 0 => { + let idf = idf(token_docs, self.doc_count); + if idf.is_finite() && idf > 0.0 { + idf + } else { + 0.0 + } + } + _ => 0.0, + } + } + + /// BM25 term contribution for a document, given the blended term frequency + /// `tf'` and blended document length `dl'`. + /// + /// The result is always finite and within `[0, BM25_DOC_WEIGHT_UPPER_BOUND]`, + /// whatever reaches it. BM25 saturates the `tf'` factor at `K1 + 1` in exact + /// arithmetic, but the f32 evaluation can land above it: once `doc_norm` is + /// small relative to `ulp(tf')` the rounded `tf' + doc_norm` collapses back to + /// `tf'`, and `fl(fl((K1 + 1) · tf') / tf')` then rounds up. Large per-column + /// boosts can also push `tf'`, `dl'`, or `avgdl'` to infinity, and `Inf / Inf` + /// is `NaN`; one `NaN` contribution would poison a whole query's scores. The + /// clamp makes the ceiling hold by construction rather than by trusting the + /// query-level validation of the boosts. + pub fn doc_weight(&self, tf_prime: f32, dl_prime: f32) -> f32 { + if self.avg_doc_length <= 0.0 { + return 0.0; + } + let doc_norm = K1 * (1.0 - B + B * dl_prime / self.avg_doc_length); + let weight = (K1 + 1.0) * tf_prime / (tf_prime + doc_norm); + if weight.is_nan() { + // `Inf / Inf`, or a `NaN` that came in through `tf'` / `dl'` / + // `avgdl'`. There is no meaningful weight to report, and `clamp` + // would propagate the `NaN`. + return 0.0; + } + weight.clamp(0.0, BM25_DOC_WEIGHT_UPPER_BOUND) + } +} + #[cfg(test)] mod tests { use super::*; @@ -252,4 +340,85 @@ mod tests { assert!(doc_weight > K1 + 1.0); assert!(scorer.doc_weight_upper_bound().unwrap() >= doc_weight); } + + #[test] + fn test_combined_scorer_query_weight_matches_blended_idf() { + let doc_freq = HashMap::from([("rare".to_string(), 2), ("common".to_string(), 800)]); + let scorer = CombinedFieldsBM25Scorer::new(1000, 12.0, doc_freq); + + // IDF uses the blended docFreq' over the blended docCount'. + assert_eq!(scorer.query_weight("rare"), idf(2, 1000)); + assert_eq!(scorer.query_weight("common"), idf(800, 1000)); + // A rarer term outweighs a common one. + assert!(scorer.query_weight("rare") > scorer.query_weight("common")); + // A term absent from every field contributes nothing. + assert_eq!(scorer.query_weight("missing"), 0.0); + } + + /// No corpus statistics and no blended `(tf', dl')` may produce a non-finite + /// term score. A `NaN` would order arbitrarily in the top-k heap, and a weight + /// above [`BM25_DOC_WEIGHT_UPPER_BOUND`] would break any pruning bound derived + /// from that ceiling. + #[test] + fn test_combined_scorer_stays_finite_for_extreme_statistics() { + // `docFreq' > docCount'` drives `idf` negative and then to `-inf`. + let doc_freq = HashMap::from([ + ("ok".to_string(), 1), + ("over".to_string(), 2_000), + ("way_over".to_string(), usize::MAX), + ]); + let blends = [ + (1.0f32, 8.0f32), + (0.0, 0.0), + (f32::MAX, f32::MAX), + (f32::INFINITY, f32::INFINITY), + (f32::INFINITY, 8.0), + (8.0, f32::INFINITY), + (f32::NAN, f32::NAN), + (-1.0, -1.0), + ]; + for avg_doc_length in [10.0f32, 0.0, f32::MAX, f32::INFINITY, f32::NAN, -1.0] { + let scorer = CombinedFieldsBM25Scorer::new(1000, avg_doc_length, doc_freq.clone()); + for term in ["ok", "over", "way_over", "missing"] { + let query_weight = scorer.query_weight(term); + assert!( + query_weight.is_finite() && query_weight >= 0.0, + "query_weight({term}) = {query_weight:e}" + ); + for (tf_prime, dl_prime) in blends { + let score = query_weight * scorer.doc_weight(tf_prime, dl_prime); + assert!( + score.is_finite() && score >= 0.0, + "score for {term} at tf'={tf_prime:e} dl'={dl_prime:e} \ + avgdl'={avg_doc_length:e} was {score:e}" + ); + } + } + } + // The unclamped reference: `idf` really does leave the usable range for + // these statistics, so the clamp in `query_weight` is important. In the + // `-inf` case multiplied by a zero `doc_weight` it would yield `NaN`. + assert!(idf(2_000, 1000) < 0.0); + assert_eq!(idf(usize::MAX, 1000), f32::NEG_INFINITY); + } + + #[test] + fn test_combined_scorer_doc_weight_saturates_and_penalizes_length() { + let scorer = CombinedFieldsBM25Scorer::new(1000, 10.0, HashMap::new()); + + // Matches the BM25 formula (with Lance's (k1 + 1) numerator). + let expected = { + let doc_norm = K1 * (1.0 - B + B * 20.0 / 10.0); + (K1 + 1.0) * 3.0 / (3.0 + doc_norm) + }; + assert!((scorer.doc_weight(3.0, 20.0) - expected).abs() < 1e-6); + // More term frequency scores higher, saturating below (k1 + 1). + assert!(scorer.doc_weight(5.0, 20.0) > scorer.doc_weight(1.0, 20.0)); + assert!(scorer.doc_weight(1000.0, 20.0) < K1 + 1.0); + // A longer document is penalized for the same term frequency. + assert!(scorer.doc_weight(3.0, 40.0) < scorer.doc_weight(3.0, 5.0)); + // A degenerate (empty) corpus scores zero rather than dividing by zero. + let empty = CombinedFieldsBM25Scorer::new(0, 0.0, HashMap::new()); + assert_eq!(empty.doc_weight(3.0, 20.0), 0.0); + } } diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index dcc5b5c3ac7..0234f3f5605 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -733,6 +733,65 @@ impl InvertedIndexParams { Ok(self) } + /// Whether `self` and `other` tokenize text identically. + /// + /// Compares only the fields that affect how documents are turned into tokens + /// (tokenizer, language, casing, stemming, stop words, n-gram bounds, + /// code-analyzer splitting, ...). Storage/layout and build-time fields + /// (`with_position`, `block_size`, worker/memory/format settings) are ignored: + /// two columns can differ there yet still be safe to combine in a + /// `combined_fields` (BM25F) query. Used by cross-field FTS validation. + pub(crate) fn same_tokenization(&self, other: &Self) -> bool { + // Destructure exhaustively so that adding a new field forces an explicit + // decision here about whether it affects tokenization. + let Self { + lance_tokenizer, + base_tokenizer, + language, + max_token_length, + lower_case, + stem, + remove_stop_words, + custom_stop_words, + ascii_folding, + min_ngram_length, + max_ngram_length, + prefix_only, + // Code-analyzer tokenization options: change the emitted token stream. + split_identifiers, + split_on_numerics, + preserve_original, + index_operators, + // Not tokenization-affecting (storage/layout/build-time only): + with_position: _, + block_size: _, + memory_limit_mb: _, + num_workers: _, + format_version: _, + // Decides what a document is, not how its text is tokenized. Cross-field + // FTS checks it separately because it needs a different answer: the + // blend is per row, so a list-element index is rejected outright rather + // than merely required to match its siblings. + document_granularity: _, + } = self; + lance_tokenizer == &other.lance_tokenizer + && base_tokenizer == &other.base_tokenizer + && language == &other.language + && max_token_length == &other.max_token_length + && lower_case == &other.lower_case + && stem == &other.stem + && remove_stop_words == &other.remove_stop_words + && custom_stop_words == &other.custom_stop_words + && ascii_folding == &other.ascii_folding + && min_ngram_length == &other.min_ngram_length + && max_ngram_length == &other.max_ngram_length + && prefix_only == &other.prefix_only + && split_identifiers == &other.split_identifiers + && split_on_numerics == &other.split_on_numerics + && preserve_original == &other.preserve_original + && index_operators == &other.index_operators + } + pub fn lance_tokenizer(mut self, lance_tokenizer: String) -> Self { self.lance_tokenizer = Some(lance_tokenizer); self @@ -1172,6 +1231,22 @@ mod tests { assert!(json.get("format_version").is_none()); } + #[test] + fn test_same_tokenization_ignores_storage_only_fields() { + let base = InvertedIndexParams::new("simple".to_string(), Language::English); + + // Storage/layout-only differences keep the same tokenization. + assert!(base.same_tokenization(&base.clone().with_position(true))); + + // Tokenizer-affecting differences are detected. + assert!(!base.same_tokenization(&base.clone().stem(!base.stem))); + assert!(!base.same_tokenization(&base.clone().lower_case(!base.lower_case))); + assert!(!base.same_tokenization(&InvertedIndexParams::new( + "whitespace".to_string(), + Language::English, + ))); + } + #[test] fn test_memory_limit_serde_accepts_legacy_worker_field_name() { let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); diff --git a/rust/lance-index/src/traits.rs b/rust/lance-index/src/traits.rs index f5441bd80f0..1ddb8cd9418 100644 --- a/rust/lance-index/src/traits.rs +++ b/rust/lance-index/src/traits.rs @@ -246,6 +246,9 @@ impl Display for FtsPrewarmPartitionStatus { if !self.documents.reverse_lookup_ready { missing.push("reverse document lookup"); } + if !self.documents.ascending_addresses_ready { + missing.push("memoized ascending-address gate"); + } if !self.documents.projection_resident { missing.push("resident row-address projection"); } @@ -278,6 +281,11 @@ pub struct FtsPrewarmDocumentStatus { pub prewarm_complete: bool, pub scoring_ready: bool, pub reverse_lookup_ready: bool, + /// Whether the memoized "are the stored addresses strictly ascending?" + /// answer is populated. Cross-field search reads it on every query, so a + /// prewarmed partition that left it empty would still pay an O(num_docs) + /// address scan on its first query. + pub ascending_addresses_ready: bool, pub projection_resident: bool, } @@ -286,6 +294,7 @@ impl FtsPrewarmDocumentStatus { self.prewarm_complete && self.scoring_ready && self.reverse_lookup_ready + && self.ascending_addresses_ready && self.projection_resident } } diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 50ab8cc3b36..b027144dca9 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1633,6 +1633,11 @@ impl FtsMemIndex { visit(index, query, terms)?; } } + FtsQuery::CombinedFields(_) => { + return Err(Error::invalid_input( + "residual compound FTS does not support combined_fields (BM25F) leaves", + )); + } } Ok(()) } @@ -1703,6 +1708,11 @@ impl FtsMemIndex { visit(index, query, scorer, leaves)?; } } + FtsQuery::CombinedFields(_) => { + return Err(Error::invalid_input( + "residual compound FTS does not support combined_fields (BM25F) leaves", + )); + } } Ok(()) } diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 12f58db6354..d55f229995b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -118,6 +118,11 @@ fn requested_query_document_granularity( } Ok(()) } + // BM25F blends the target columns per row, so combined_fields is + // row-granular by construction and carries no granularity field. + // Merging Row still catches a tree that mixes it with a + // list-element leaf. + IndexFtsQuery::CombinedFields(_) => merge(current, Some(DocumentGranularity::Row)), } } @@ -156,6 +161,8 @@ fn set_query_document_granularity( child.document_granularity = Some(document_granularity); } } + // Row-granular by construction, and it carries no field to set. + IndexFtsQuery::CombinedFields(_) => {} } } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index fe6985fdfb7..5e0d9e6d028 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -72,12 +72,12 @@ use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::expression::PlannerIndexExt; use lance_index::scalar::expression::ScalarIndexExpr; use lance_index::scalar::inverted::query::{ - FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, PhraseQuery, - fill_fts_query_column, + CombinedFieldsQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, + PhraseQuery, fill_fts_query_column, }; use lance_index::scalar::inverted::{ - DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity, INVERTED_INDEX_VERSION_V2, - INVERTED_INDEX_VERSION_V3, SCORE_COL, SCORE_FIELD, fts_schema, + CombinedFieldsBM25Scorer, DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity, + INVERTED_INDEX_VERSION_V2, INVERTED_INDEX_VERSION_V3, SCORE_COL, SCORE_FIELD, fts_schema, }; use lance_index::scalar::registry::VALUE_COLUMN_NAME; use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, DIST_COL, Query}; @@ -103,7 +103,7 @@ use crate::index::DatasetIndexInternalExt; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, - resolve_fts_field, resolve_query_document_granularity, + resolve_fts_field, resolve_query_document_granularity, validate_combined_fields_target_column, }; use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ @@ -113,9 +113,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, CombinedFieldsQueryExec, CompoundQueryExec, CrossColumnCompoundQueryExec, + FlatMatchFilterExec, FlatMatchQueryExec, FtsDocumentExec, HybridCompoundQueryExec, + MatchQueryExec, PhraseQueryExec, SharedFtsScorer, }; use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; @@ -210,6 +210,13 @@ fn collect_fts_columns_in_order(query: &FtsQuery) -> Vec { visit(child, columns, seen); } } + FtsQuery::CombinedFields(query) => { + for column in query.column_names() { + if seen.insert(column.to_string()) { + columns.push(column.to_string()); + } + } + } } } @@ -240,7 +247,8 @@ fn collect_phrase_columns(query: &FtsQuery, columns: &mut HashSet) { collect_phrase_columns(child, columns); } } - FtsQuery::Match(_) | FtsQuery::MultiMatch(_) => {} + // combined_fields is term-based; it never needs positions. + FtsQuery::Match(_) | FtsQuery::MultiMatch(_) | FtsQuery::CombinedFields(_) => {} } } @@ -262,6 +270,10 @@ fn supports_compound_scorer(query: &FtsQuery) -> bool { fn supports_shape(query: &FtsQuery) -> bool { match query { FtsQuery::Match(_) | FtsQuery::Phrase(_) | FtsQuery::MultiMatch(_) => true, + // BM25F blends term statistics across columns, so it cannot be a leaf + // of the single-index compound scorer. Such trees keep the + // union/sort plan built by `plan_combined_fields_query`. + FtsQuery::CombinedFields(_) => false, FtsQuery::Boolean(query) => { (!query.should.is_empty() || !query.must.is_empty()) && query @@ -305,6 +317,9 @@ fn supports_indexed_stats_residual_compound(query: &FtsQuery) -> bool { .chain(&query.must) .chain(&query.must_not) .all(supports_indexed_stats_residual_compound), + // The compound scorer scores each leaf from one column's index, which a + // cross-field blend cannot be expressed as. + FtsQuery::CombinedFields(_) => false, } } @@ -368,6 +383,11 @@ fn validate_fts_query_contract(query: &FtsQuery) -> Result<()> { match query { FtsQuery::Match(query) => validate_multiplier("MatchQuery boost", query.boost), FtsQuery::Phrase(_) => Ok(()), + // Nothing to re-check: the weights are private and every constructor, + // including `Deserialize`, goes through `try_with_boosts`, which enforces a + // finite value in `[MIN_BOOST, MAX_BOOST]`, stricter than the non-negative + // contract checked here. + FtsQuery::CombinedFields(_) => Ok(()), FtsQuery::Boost(query) => { validate_multiplier("BoostQuery negative_boost", query.negative_boost)?; validate_fts_query_contract(&query.positive)?; @@ -407,7 +427,9 @@ fn normalize_fts_zero_boosts(query: &mut FtsQuery) { match query { FtsQuery::Match(query) => normalize_zero(&mut query.boost), - FtsQuery::Phrase(_) => {} + // combined_fields weights are validated into `[1.0, 2^20]`, so no zero + // to normalize. + FtsQuery::Phrase(_) | FtsQuery::CombinedFields(_) => {} FtsQuery::Boost(query) => { normalize_zero(&mut query.negative_boost); normalize_fts_zero_boosts(&mut query.positive); @@ -443,7 +465,8 @@ fn apply_dataset_planner_auto_fuzziness_compatibility_gate(query: &mut FtsQuery) FtsQuery::Match(query) => { query.fuzziness.get_or_insert(0); } - FtsQuery::Phrase(_) => {} + // combined_fields matches terms exactly; it has no fuzziness setting. + FtsQuery::Phrase(_) | FtsQuery::CombinedFields(_) => {} FtsQuery::Boost(query) => { apply_dataset_planner_auto_fuzziness_compatibility_gate(&mut query.positive); apply_dataset_planner_auto_fuzziness_compatibility_gate(&mut query.negative); @@ -3947,6 +3970,19 @@ impl Scanner { ) .await } + FtsQuery::CombinedFields(combined) => { + // A doc's fragment must be covered by every target column's index + // for the prefilter to be exact, mirroring MultiMatch. + for column in combined.column_names() { + if !self + .fragments_covered_by_fts_leaf(column, DocumentGranularity::Row, accum) + .await? + { + return Ok(false); + } + } + Ok(true) + } FtsQuery::Boolean(bool_query) => { for query in bool_query .must @@ -4083,6 +4119,17 @@ impl Scanner { } Ok(()) } + FtsQuery::CombinedFields(query) => { + // BM25F sums each target column's contribution for one row, + // so combined_fields is row-granular by construction. A + // target column that only has a list-element index was + // already rejected by + // `resolve_fts_query_document_granularity`. + for column in query.column_names() { + add_leaf(schema, Some(column), Some(DocumentGranularity::Row), state)?; + } + Ok(()) + } } } @@ -4171,6 +4218,15 @@ impl Scanner { } Ok(FtsQuery::MultiMatch(query)) } + FtsQuery::CombinedFields(query) => { + // There is nothing to resolve: a combined_fields query carries no + // granularity because BM25F only makes sense over row documents. + // Reject a target column that cannot supply them. + for column in query.column_names() { + validate_combined_fields_target_column(self.dataset.as_ref(), column).await?; + } + Ok(FtsQuery::CombinedFields(query)) + } } } @@ -4210,6 +4266,8 @@ impl Scanner { .get_or_insert(document_granularity); } } + // Row-granular by construction, and it carries no field to fill in. + FtsQuery::CombinedFields(_) => {} } } @@ -4236,6 +4294,7 @@ impl Scanner { .document_granularity .is_some_and(DocumentGranularity::is_list_element) }), + FtsQuery::CombinedFields(_) => false, } } @@ -4678,27 +4737,16 @@ impl Scanner { fts_node, schema, )?); - let sort_exprs = [ - PhysicalSortExpr { - expr: expressions::col(SCORE_COL, fts_node.schema().as_ref())?, - options: SortOptions { - descending: true, - nulls_first: false, - }, - }, - PhysicalSortExpr { - expr: expressions::col(ROW_ID, fts_node.schema().as_ref())?, - options: SortOptions { - descending: false, - nulls_first: false, - }, - }, - ]; + let sort_exprs = Self::fts_score_sort_exprs(fts_node.schema().as_ref())?; // `params.limit` is the recursive planning contract. Compound // parents pass `None` when they require every candidate. Arc::new(SortExec::new(sort_exprs.into(), fts_node).with_fetch(params.limit)) } + FtsQuery::CombinedFields(query) => { + self.plan_combined_fields_query(query, params, prefilter_source) + .await? + } FtsQuery::Boolean(query) => { // TODO: rewrite the query for better performance @@ -5258,6 +5306,221 @@ impl Scanner { Ok(flat_match_plan) } + /// Sort keys for an FTS plan whose scores come from more than one source, or + /// from a source that does not sort. + /// + /// `row_id ASC` is a required second key, not a nicety: a score-only sort over + /// concurrently read sources breaks ties by arrival order, so equal-scoring + /// documents would come back in a different order run to run and + /// `limit`/`offset` pagination could skip and repeat rows. `combined_fields` + /// needs it for the same reason: `combined_fields_search` visits candidates in + /// ascending row-id order, so its top-k is reproducible, but it does not order + /// ties itself (`ScoredDoc` compares on score alone), so this sort is what + /// decides which of two equal-scoring rows comes first. + fn fts_score_sort_exprs(schema: &ArrowSchema) -> Result<[PhysicalSortExpr; 2]> { + Ok([ + PhysicalSortExpr { + expr: expressions::col(SCORE_COL, schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: expressions::col(ROW_ID, schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }, + ]) + } + + /// Plan a cross-field (BM25F) `combined_fields` query. + /// + /// Unlike a single-column match, coverage here is per column: a BM25F score is + /// only complete when every target column's index covers the row's fragment, + /// because `dl'` sums each column's document length and the per-column lookup + /// (`AddressKeyedDocuments::doc_length_at` over that partition's `DocLengths`) + /// returns 0 for a row that column's index holds no document for. So the + /// fragments the index scan must not touch are the union of the per-column + /// unindexed sets, not their intersection: a fragment indexed for `title` but + /// not `body` would otherwise be emitted with a partial `tf'`/`dl'`. + /// + /// The same union absorbs the fragments whose index entries a newer data + /// overlay made stale, per target column, so the indexed scan neither returns a + /// pre-overlay hit nor hides a new one. See [`Self::fts_overlay_plan`]. + /// + /// A fragment any target column's index does not fully cover cannot be scored, + /// so the query is refused rather than answered from partial statistics. + async fn plan_combined_fields_query( + &self, + query: &CombinedFieldsQuery, + params: &FtsSearchParams, + prefilter_source: &PreFilterSource, + ) -> Result> { + let target_fragments: &[Fragment] = self + .fragments + .as_deref() + .unwrap_or_else(|| self.dataset.fragments()); + // An explicitly empty fragment list selects no rows. The prefilter would + // already reject every row, but without this the fully covered branch below + // still opens the segments, builds the scorer and runs a full scan to return + // nothing. `plan_match_query` and `plan_phrase_query` short-circuit the same way. + if self.fragments.as_ref().is_some_and(Vec::is_empty) { + return Ok(Arc::new(EmptyExec::new(fts_schema( + DocumentGranularity::Row, + )))); + } + + let mut uncovered = RoaringBitmap::new(); + let mut any_indexed = false; + for column in query.column_names() { + let column_uncovered: RoaringBitmap = match self + .dataset + .load_scalar_index( + IndexCriteria::default() + .for_column(column) + .supports_fts() + .with_fts_document_granularity(DocumentGranularity::Row), + ) + .await? + { + Some(index) => { + any_indexed = true; + self.dataset + .unindexed_fragments(&index.name) + .await? + .iter() + .map(|fragment| fragment.id as u32) + .collect() + } + // No index on this column at all, so no fragment is covered for it. + None => target_fragments + .iter() + .map(|fragment| fragment.id as u32) + .collect(), + }; + + // Fragments this column's index holds documents for, but whose entries a + // newer data overlay made stale. The index cannot score them from its own + // statistics, so they count as uncovered. + // + // Coverage is decided per fragment, not per row: a BM25F score needs every + // target column's `tf_f`/`dl_f`, so a row must be scored wholly from the + // index or not at all, and the fragment is the granularity the indexed + // scan's prefilter restriction works at. + match self + .fts_overlay_plan(column, DocumentGranularity::Row, target_fragments) + .await? + { + FtsOverlayPlan::Unchanged(_) => {} + FtsOverlayPlan::RowLevel { stale_rows, .. } => { + uncovered.extend(stale_rows.keys().copied()); + } + // A legacy segment reports no fragment coverage, so no target fragment + // can be proven free of stale entries and no row can be named as one. + FtsOverlayPlan::FullScan => { + uncovered.extend(target_fragments.iter().map(|fragment| fragment.id as u32)); + } + } + + uncovered |= &column_uncovered; + } + + // BM25F needs one shared tokenizer configuration across the target columns + // (`validate_combined_tokenizers`), and it is read off an index. Error out + // when no target column has one rather than silently falling back to a + // default tokenizer that may not match the data. + if !any_indexed { + return Err(Error::invalid_input(format!( + "combined_fields requires an inverted index on at least one of {:?}", + query.column_names().collect::>() + ))); + } + + let uncovered_fragments: Vec = target_fragments + .iter() + .filter(|fragment| uncovered.contains(fragment.id as u32)) + .cloned() + .collect(); + // The complement: fragments every target column indexes, and so the only + // ones the indexed scan can score completely. + let covered: RoaringBitmap = target_fragments + .iter() + .map(|fragment| fragment.id as u32) + .filter(|fragment_id| !uncovered.contains(*fragment_id)) + .collect(); + + // The only construction site for the indexed side, so the options every + // shape needs cannot be set on one path and forgotten on another. + // + // `covered` restricts the scan to the fragments every target column indexes + // and no overlay made stale; `None` means there is nothing to restrict. + // Segments stay unrestricted either way: the fragment restriction already + // keeps a stale row out of the results, and dropping a segment would drop + // its documents from the corpus statistics too. + // + // `shared_scorer` is set only when a flat sibling exists: that side alone + // sees the rows no index covers, so it publishes the corpus statistics both + // children must score against. + let index_exec = |covered: Option, + shared_scorer: Option>>| + -> Arc { + let mut exec = CombinedFieldsQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + ) + .with_external_mask(self.external_row_mask.clone()); + if let Some(covered) = covered { + exec = exec.with_covered_fragments(covered); + } + if let Some(shared_scorer) = shared_scorer { + exec = exec.with_shared_scorer(shared_scorer); + } + Arc::new(exec) + }; + + // Every target column covers every target fragment, with no overlay-stale + // entries anywhere: one unified scan that already emits merged hits sorted + // by score with the top-k limit applied, so no union/sort is needed, and + // there are no fragments to restrict the prefilter to. + if uncovered_fragments.is_empty() { + return Ok(index_exec(None, None)); + } + // `fast_search` is index-only by contract, but must still drop the + // partially covered fragments so no partial score is emitted. When no + // fragment is covered the answer is definitionally empty, and the index + // exec would instead fail on the target column that has no segments. + if self.fast_search { + if covered.is_empty() { + return Ok(Arc::new(EmptyExec::new(fts_schema( + DocumentGranularity::Row, + )))); + } + return Ok(index_exec(Some(covered), None)); + } + + // The flat scan that would score the rows no index covers is not part of + // this change, so a partial plan would silently drop them. Refuse instead, + // and name what is missing. + Err(Error::invalid_input(format!( + "combined_fields requires every target column to be indexed over every \ + scanned fragment, but {} of {} fragments are not fully covered. Optimize \ + the indexes on {} and retry.", + uncovered_fragments.len(), + target_fragments.len(), + query + .columns() + .iter() + .cloned() + .collect::>() + .join(", "), + ))) + } + // ANN/KNN search execution node with optional prefilter #[async_recursion] async fn vector_search( @@ -7457,7 +7720,7 @@ mod test { fn boost_bits(query: &FtsQuery) -> Vec { match query { FtsQuery::Match(query) => vec![query.boost.to_bits()], - FtsQuery::Phrase(_) => Vec::new(), + FtsQuery::Phrase(_) | FtsQuery::CombinedFields(_) => Vec::new(), FtsQuery::Boost(query) => std::iter::once(query.negative_boost.to_bits()) .chain(boost_bits(&query.positive)) .chain(boost_bits(&query.negative)) @@ -7514,7 +7777,7 @@ mod test { fn collect_fuzziness(query: &FtsQuery, values: &mut Vec>) { match query { FtsQuery::Match(query) => values.push(query.fuzziness), - FtsQuery::Phrase(_) => {} + FtsQuery::Phrase(_) | FtsQuery::CombinedFields(_) => {} FtsQuery::Boost(query) => { collect_fuzziness(&query.positive, values); collect_fuzziness(&query.negative, values); diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 2d7c48c72a1..8fc0e47cc79 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1453,6 +1453,12 @@ fn independent_compound_fts_oracle<'a>( } result } + // This oracle scores each leaf against its own column's statistics, + // which is exactly what BM25F does not do. combined_fields never + // reaches the cross-column compound path. + FtsQuery::CombinedFields(_) => { + unreachable!("combined_fields is not a cross-column compound leaf") + } } }) } diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 78dcbd3ad43..0030c06c7e9 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -106,6 +106,33 @@ impl DatasetPreFilter { } } + /// Like [`Self::new_with_filter_future`], but restricted to an explicit + /// fragment set instead of the union of the indices' fragment bitmaps. + /// + /// `combined_fields` needs the intersection of its target columns' coverage: a + /// cross-field BM25F score is only complete when every column's index holds the + /// row, so the union [`Self::new`] derives would admit rows that only some + /// columns cover. Routing through [`Self::create_restricted_deletion_mask`] + /// keeps the restriction in whichever id space the indices actually store, + /// stable row ids or row addresses, which a fragment-address block list cannot + /// do. + pub(crate) fn new_restricted_to_fragments( + 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..2449c19283d 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -636,6 +636,45 @@ pub(crate) async fn resolve_query_document_granularity( Ok(resolved) } +/// Check that one `combined_fields` target column can take part in a BM25F +/// blend, and resolve its path. +/// +/// Cross-field scoring joins the target columns on the row address and sums their +/// per-row term frequencies and document lengths, so it needs row documents. A +/// column indexed only at list-element granularity cannot supply them: that index +/// stores several documents per row, identified by element coordinates a +/// cross-column scan cannot pair up between columns and the combined result schema +/// cannot report. A column carrying both granularities is fine; the row index is +/// the one used. +/// +/// Any field shape a row-document index accepts is allowed, list nesting included: +/// the flat sibling scan reaches such a leaf through +/// [`flatten_fts_document_column`], which uses the same traversal a single-column +/// match drives through [`FtsDocument`]. +pub(crate) async fn validate_combined_fields_target_column( + dataset: &Dataset, + column: &str, +) -> Result<()> { + let indices = indexed_fts_document_granularities(dataset, column).await?; + if !indices.is_empty() + && indices + .iter() + .all(|(_, granularity)| granularity.is_list_element()) + { + let indexed = indices + .iter() + .map(|(name, granularity)| format!("'{name}' ({granularity:?})")) + .collect::>() + .join(", "); + return Err(Error::not_supported(format!( + "combined_fields (BM25F) scores whole rows and needs a Row document granularity FTS \ + index on every target column, but '{column}' only has: {indexed}" + ))); + } + resolve_fts_field(dataset.schema(), column, DocumentGranularity::Row)?; + Ok(()) +} + pub(crate) fn fts_document_schema(coordinate_rank: usize) -> Arc { let mut fields = vec![ ArrowField::new(VALUE_COLUMN_NAME, DataType::Utf8, false), diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 3483a476d52..5b0846bc89e 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -41,7 +41,9 @@ use lance_table::format::IndexMetadata; use rustc_hash::FxHashSet; use super::PreFilterSource; -use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; +use super::utils::{ + IndexMetrics, PreFilterMasks, build_prefilter, build_prefilter_restricted_to_fragments, +}; use crate::dataset::mem_wal::index::{QueryLocalFtsIndex, QueryLocalFtsStats}; use crate::index::scalar::inverted::{ ResolvedFtsField, fts_document_schema, load_segment_details, load_segments, @@ -53,17 +55,18 @@ use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::builder::document_input; use lance_index::scalar::inverted::document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}; use lance_index::scalar::inverted::query::{ - BoostQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, PhraseQuery, Tokens, - collect_query_tokens, has_query_token, uses_fuzzy_expansion, + BoostQuery, CombinedFieldsQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, + PhraseQuery, Tokens, collect_query_tokens, has_query_token, uses_fuzzy_expansion, }; use lance_index::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_index::scalar::inverted::{ - DOC_INDEX_COL, DocumentGranularity, FTS_SCHEMA, FlatBm25SearchOptions, InvertedIndex, - MemBM25Scorer, PreparedBm25Query, SCORE_COL, Scorer, build_global_bm25_scorer, compound_search, + CombinedFieldColumn, CombinedFieldsBM25Scorer, DOC_INDEX_COL, DocumentGranularity, FTS_SCHEMA, + FlatBm25SearchOptions, InvertedIndex, MemBM25Scorer, PreparedBm25Query, SCORE_COL, Scorer, + build_combined_bm25_scorer, build_global_bm25_scorer, combined_fields_search, compound_search, compound_search_prepared_match, compound_search_prepared_match_with_score_floor, compound_search_with_base_scorer, cross_column_compound_search, exclusive_scaled_score_floor, flat_bm25_search_stream_with_options_and_scorer, fts_schema, materialized_compound_top_k, - prepare_bm25_query, + prepare_bm25_query, validate_combined_tokenizers, }; use lance_index::{prefilter::PreFilter, scalar::inverted::query::BooleanQuery}; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; @@ -581,6 +584,10 @@ fn count_fts_leaves(query: &FtsQuery) -> usize { count_fts_leaves(&query.positive) + count_fts_leaves(&query.negative) } FtsQuery::MultiMatch(query) => query.match_queries.len(), + // Unreachable while `supports_compound_scorer` rejects BM25F. Counted as + // the per-column posting scans anyway, so the arm stays honest if that + // changes. + FtsQuery::CombinedFields(query) => query.column_names().count(), FtsQuery::Boolean(query) => query .should .iter() @@ -628,6 +635,14 @@ fn compound_leaf_columns(query: &FtsQuery) -> Result> { visit(query, columns)?; } } + // Unreachable while `supports_compound_scorer` rejects BM25F: a + // combined_fields leaf blends statistics across its columns, so it + // cannot name one column per leaf the way this list requires. + FtsQuery::CombinedFields(_) => { + return Err(Error::invalid_input( + "cross-column compound FTS cannot take a combined_fields leaf".to_string(), + )); + } } Ok(()) } @@ -640,7 +655,8 @@ fn compound_leaf_columns(query: &FtsQuery) -> Result> { fn compound_query_uses_fuzzy_expansion(query: &FtsQuery) -> bool { match query { FtsQuery::Match(query) => uses_fuzzy_expansion(query.fuzziness), - FtsQuery::Phrase(_) => false, + // combined_fields matches terms exactly; it has no fuzziness setting. + FtsQuery::Phrase(_) | FtsQuery::CombinedFields(_) => false, FtsQuery::Boost(query) => { compound_query_uses_fuzzy_expansion(&query.positive) || compound_query_uses_fuzzy_expansion(&query.negative) @@ -2092,6 +2108,10 @@ fn tokenize_compound_query(query: &FtsQuery, index: &InvertedIndex) -> Tokenized }); } } + // Unreachable: only `CompoundQueryExec` tokenizes this way, and + // `supports_compound_scorer` rejects BM25F at any depth. The + // combined_fields execs snapshot their own tokens. + FtsQuery::CombinedFields(_) => {} FtsQuery::Boolean(query) => { for query in query .should @@ -2185,6 +2205,13 @@ fn tokenize_cross_column_compound_query( visit(query, indices, leaves)?; } } + // Unreachable while `supports_compound_scorer` rejects BM25F; see + // `compound_leaf_columns`. + FtsQuery::CombinedFields(_) => { + return Err(Error::invalid_input( + "cross-column compound FTS cannot take a combined_fields leaf".to_string(), + )); + } } Ok(()) } @@ -2194,24 +2221,28 @@ fn tokenize_cross_column_compound_query( Ok(TokenizedCompoundQuery(leaves)) } -type SharedScorerResult = std::result::Result, Arc>; +type SharedScorerResult = std::result::Result, Arc>; /// Coordinates BM25 corpus statistics between the indexed and flat branches /// of a mixed search. The flat branch extends the indexed statistics with the /// unindexed documents, then publishes the resulting corpus-wide scorer. +/// +/// Generic over the scorer so the single-column path can share a +/// [`MemBM25Scorer`] and `combined_fields` a [`CombinedFieldsBM25Scorer`]; the +/// coordination is the same either way. #[derive(Debug)] -pub(crate) struct SharedFtsScorer { - sender: tokio::sync::watch::Sender>, +pub(crate) struct SharedFtsScorer { + sender: tokio::sync::watch::Sender>>, } -impl SharedFtsScorer { +impl SharedFtsScorer { pub(crate) fn new() -> Self { let (sender, _) = tokio::sync::watch::channel(None); Self { sender } } - fn publish(&self, scorer: MemBM25Scorer) { - self.sender.send_replace(Some(Ok(Arc::new(scorer)))); + fn publish(&self, scorer: Arc) { + self.sender.send_replace(Some(Ok(scorer))); } fn publish_error(&self, error: &DataFusionError) { @@ -2219,7 +2250,7 @@ impl SharedFtsScorer { .send_replace(Some(Err(Arc::from(error.to_string())))); } - async fn wait(&self) -> DataFusionResult> { + async fn wait(&self) -> DataFusionResult> { let mut receiver = self.sender.subscribe(); loop { let result = receiver.borrow_and_update().clone(); @@ -2236,20 +2267,20 @@ impl SharedFtsScorer { } } -struct SharedFtsScorerProducer { - scorer: Arc, +struct SharedFtsScorerProducer { + scorer: Arc>, completed: bool, } -impl SharedFtsScorerProducer { - fn new(scorer: Arc) -> Self { +impl SharedFtsScorerProducer { + fn new(scorer: Arc>) -> Self { Self { scorer, completed: false, } } - fn publish(mut self, scorer: MemBM25Scorer) { + fn publish(mut self, scorer: Arc) { self.scorer.publish(scorer); self.completed = true; } @@ -2260,7 +2291,7 @@ impl SharedFtsScorerProducer { } } -impl Drop for SharedFtsScorerProducer { +impl Drop for SharedFtsScorerProducer { fn drop(&mut self) { if !self.completed { self.scorer.sender.send_replace(Some(Err(Arc::from( @@ -2998,6 +3029,426 @@ impl ExecutionPlan for MatchQueryExec { } } +/// Cross-field BM25F full-text search (`combined_fields`). +/// +/// Unlike a `MultiMatch` plan (one [`MatchQueryExec`] per column fused by a +/// `max` aggregate, i.e. `best_fields`), this is a single node: it opens every +/// target column's segments, blends their corpus statistics, and runs one merged +/// scan so term statistics are shared across fields. Per-column boosts are baked +/// into the blended term frequency, so there is no post-scan boost multiplier or +/// `AggregateExec(max)`. Emits `(ROW_ID, SCORE)` in [`FTS_SCHEMA`]. +#[derive(Debug)] +pub struct CombinedFieldsQueryExec { + dataset: Arc, + query: CombinedFieldsQuery, + /// Tokens `execute()` actually produced, for `analyze_plan`. One snapshot + /// covers every target column because `validate_combined_tokenizers` + /// requires them to share a tokenizer, so the scan tokenizes once. + tokenized_query: Arc>, + params: FtsSearchParams, + prefilter_source: PreFilterSource, + /// When set, `execute()` skips `build_combined_bm25_scorer` and threads this + /// blended scorer down to the merged scan (distributed corpus-global stats). + base_scorer: Option>, + /// Waits for the flat sibling's blended scorer; see + /// [`Self::with_shared_scorer`]. + shared_scorer: Option>>, + /// When set, restrict this scan to exactly these fragments: the ones every + /// target column's index covers. See [`Self::with_covered_fragments`]. + covered_fragments: Option, + /// Optional external row-address mask ANDead into the prefilter so only + /// masked rows are scored (see [`MatchQueryExec::with_external_mask`]). + external_mask: Option>, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl CombinedFieldsQueryExec { + pub fn new( + dataset: Arc, + query: CombinedFieldsQuery, + params: FtsSearchParams, + prefilter_source: PreFilterSource, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(FTS_SCHEMA.clone()), + Partitioning::RoundRobinBatch(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + dataset, + query, + tokenized_query: Arc::new(OnceLock::new()), + params, + prefilter_source, + base_scorer: None, + shared_scorer: None, + covered_fragments: None, + external_mask: None, + properties, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// See [`MatchQueryExec::with_external_mask`]. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + + /// Restrict this scan to `fragments`, the ones every target column's index + /// covers. + /// + /// A BM25F score is only complete when every target column's index covers the + /// row's fragment: `dl'` sums each column's document length, and the per-column + /// lookup returns 0 for a row that column's index holds no document for, so a + /// fragment indexed for some target columns but not others would score with a + /// partial `tf'`/`dl'`. The planner routes those fragments to the flat scan and + /// passes the remainder here. + /// + /// This has to be an allow list rather than a fragment block list: the inverted + /// index trains with `_rowid`, which is a logical stable row id when the dataset + /// uses stable row ids, so a fragment-address block list would match nothing and + /// both sides of the union would emit the same row. + pub fn with_covered_fragments(mut self, fragments: roaring::RoaringBitmap) -> Self { + self.covered_fragments = Some(fragments); + self + } + + /// Override the blended BM25F scorer used by `execute()`. When set, the + /// local `build_combined_bm25_scorer` call is skipped. Mirrors + /// [`MatchQueryExec::with_base_scorer`] for distributed queries that + /// aggregate cross-column corpus statistics out-of-band. + pub fn with_base_scorer(mut self, scorer: Arc) -> Self { + self.base_scorer = Some(scorer); + self + } + + /// Score against the blended scorer the flat sibling publishes instead of + /// building one from this side's index statistics alone. + /// + /// Only the flat scan sees the rows no index covers, so only it can produce + /// the statistics that describe the whole scanned corpus. Without this the two + /// children of a mixed plan score against different `docCount'`/`docFreq'`/ + /// `avgdl'`, which makes their scores incomparable and the union's sort wrong. + pub(crate) fn with_shared_scorer( + mut self, + scorer: Arc>, + ) -> Self { + self.shared_scorer = Some(scorer); + self + } + + pub fn query(&self) -> &CombinedFieldsQuery { + &self.query + } + + pub fn params(&self) -> &FtsSearchParams { + &self.params + } + + pub fn prefilter_source(&self) -> &PreFilterSource { + &self.prefilter_source + } + + fn clone_with_prefilter_source(&self, prefilter_source: PreFilterSource) -> Self { + Self { + dataset: self.dataset.clone(), + query: self.query.clone(), + tokenized_query: self.tokenized_query.clone(), + params: self.params.clone(), + prefilter_source, + base_scorer: self.base_scorer.clone(), + shared_scorer: self.shared_scorer.clone(), + covered_fragments: self.covered_fragments.clone(), + external_mask: self.external_mask.clone(), + properties: self.properties.clone(), + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for CombinedFieldsQueryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + fmt_combined_fields( + "CombinedFieldsQuery", + &self.query, + &self.tokenized_query, + t, + f, + ) + } +} + +/// The plan display every `combined_fields` exec uses: the label, the target +/// columns and the query text, plus the tokens `execute()` produced once it has +/// run. +fn fmt_combined_fields( + label: &str, + query: &CombinedFieldsQuery, + tokenized_query: &OnceLock, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, +) -> std::fmt::Result { + let columns = query.column_names().collect::>().join(", "); + // The one-line forms label with `: ` and separate fields with `, `; the tree + // render puts every field on its own line. + let (after_label, separator) = match t { + DisplayFormatType::TreeRender => ("\n", "\n"), + DisplayFormatType::Default | DisplayFormatType::Verbose => (": ", ", "), + }; + write!( + f, + "{label}{after_label}columns=[{columns}]{separator}query={}", + query.terms() + )?; + fmt_tokenized_query(tokenized_query, separator, f) +} + +/// What [`open_combined_fields_scan`] hands back to the `combined_fields` exec. +struct CombinedFieldsScan { + columns: Vec, + /// Every segment opened across all target columns, for the indexed side's + /// prefilter. + segments: Vec, + tokens: Tokens, +} + +/// Open every target column's segments, pair each with its boost, and tokenize +/// the query once. +/// +/// Always every committed segment at row granularity: BM25F blends the target +/// columns per row, and query planning rejects a column that only has a +/// list-element index. The query stores each column together with its weight, so +/// no column can be dropped here for want of one. +/// +/// Both sides count the opened segments toward parts searched: the flat side needs +/// them for the tokenizer and the blended corpus statistics, the indexed side for +/// scoring. +async fn open_combined_fields_scan( + dataset: &Dataset, + query: &CombinedFieldsQuery, + tokenized_query: &OnceLock, + metrics: &FtsIndexMetrics, +) -> DataFusionResult { + let mut columns = Vec::with_capacity(query.column_names().len()); + let mut all_segments = Vec::new(); + for (column, weight) in query.weighted_columns() { + let segments = Some( + FtsSegmentSelection::AllCommitted + .resolve( + dataset, + column, + DocumentGranularity::Row, + &metrics.segment_bind_duration, + ) + .await?, + ); + let indices = match segments { + Some(segments) => { + let _details = load_segment_details(dataset, column, &segments).await?; + let indices = + open_fts_segments(dataset, column, &segments, &metrics.index_metrics).await?; + all_segments.extend(segments.iter().cloned()); + indices + } + None => Vec::new(), + }; + columns.push(CombinedFieldColumn { + column: column.to_string(), + weight, + indices, + }); + } + validate_combined_tokenizers(&columns)?; + metrics.record_parts_searched( + columns + .iter() + .flat_map(|column| &column.indices) + .map(|index| index.partition_count()) + .sum(), + ); + + // Fields share a tokenizer (validated above), so tokenize once. + let first_index = columns + .iter() + .find_map(|column| column.indices.first()) + .ok_or_else(|| { + // Not a user error: segment resolution already failed for any column + // without an index, and `CombinedFieldsQuery::try_new` rejects an empty + // column list, so reaching this means one of those two guarantees broke. + DataFusionError::Internal( + "combined_fields query reached execution with no target columns".to_string(), + ) + })?; + let mut tokenizer = first_index.tokenizer(); + let tokens = collect_query_tokens(query.terms(), &mut tokenizer); + record_tokenized_query(tokenized_query, &tokens); + + Ok(CombinedFieldsScan { + columns, + segments: all_segments, + tokens, + }) +} + +impl ExecutionPlan for CombinedFieldsQueryExec { + fn name(&self) -> &str { + "CombinedFieldsQueryExec" + } + + fn children(&self) -> Vec<&Arc> { + self.prefilter_source.execution_plan().into_iter().collect() + } + + fn required_input_distribution(&self) -> Vec { + self.children() + .iter() + .map(|_| Distribution::SinglePartition) + .collect() + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + let expected = self.children().len(); + if children.len() != expected { + return Err(DataFusionError::Internal(format!( + "CombinedFieldsQueryExec expected {expected} children, got {}", + children.len() + ))); + } + let prefilter_source = match children.pop() { + Some(source) => self.prefilter_source.with_execution_plan(source)?, + None => PreFilterSource::None, + }; + Ok(Arc::new(self.clone_with_prefilter_source(prefilter_source))) + } + + #[instrument(name = "combined_fields_query_exec", level = "debug", skip_all)] + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let query = self.query.clone(); + let tokenized_query = self.tokenized_query.clone(); + let params = self.params.clone(); + let ds = self.dataset.clone(); + let prefilter_source = self.prefilter_source.clone(); + let preset_base_scorer = self.base_scorer.clone(); + let shared_scorer = self.shared_scorer.clone(); + let covered_fragments = self.covered_fragments.clone(); + let external_mask = self.external_mask.clone(); + let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); + let stream = stream::once(async move { + let _timer = metrics.baseline_metrics.elapsed_compute().timer(); + + let CombinedFieldsScan { + columns, + segments: all_segments, + tokens, + .. + } = open_combined_fields_scan(&ds, &query, &tokenized_query, metrics.as_ref()).await?; + + // With mixed coverage the planner supplies the fragments every target + // column indexes, and the scan is restricted to exactly those: the + // rest go to the flat sibling. Otherwise the prefilter spans the union + // of the target columns' segments. + let mut pre_filter = match covered_fragments { + Some(covered) => build_prefilter_restricted_to_fragments( + context.clone(), + partition, + &prefilter_source, + ds, + covered, + external_mask, + )?, + None => build_prefilter( + context.clone(), + partition, + &prefilter_source, + ds, + &all_segments, + PreFilterMasks { + overlay_block: None, + external_mask, + }, + )?, + }; + let deleted_fragments = columns.iter().flat_map(|column| &column.indices).fold( + roaring::RoaringBitmap::new(), + |mut deleted, index| { + deleted |= index.deleted_fragments().clone(); + deleted + }, + ); + if !deleted_fragments.is_empty() { + Arc::get_mut(&mut pre_filter) + .expect("prefilter just created") + .set_deleted_fragments(deleted_fragments); + } + let scorer = match (preset_base_scorer, shared_scorer) { + (Some(scorer), _) => scorer, + // An injected scorer describes the corpus the whole plan scores + // against; wait for it rather than folding this side's statistics. + (None, Some(shared_scorer)) => shared_scorer.wait().await?, + (None, None) => { + let scorer_start = std::time::Instant::now(); + let scorer = Arc::new( + build_combined_bm25_scorer(&columns, &tokens, Some(metrics.as_ref())) + .boxed() + .await?, + ); + metrics.record_scorer_build(scorer_start.elapsed()); + scorer + } + }; + + pre_filter.wait_for_ready().await?; + let (doc_ids, scores) = combined_fields_search( + &columns, + &tokens, + ¶ms, + query.operator(), + scorer.as_ref(), + pre_filter, + metrics.as_ref(), + ) + .await?; + metrics.baseline_metrics.record_output(doc_ids.len()); + + let batch = RecordBatch::try_new( + FTS_SCHEMA.clone(), + vec![ + Arc::new(UInt64Array::from(doc_ids)), + Arc::new(Float32Array::from(scores)), + ], + )?; + Ok::<_, DataFusionError>(batch) + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream.stream_in_current_span().boxed(), + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn supports_limit_pushdown(&self) -> bool { + false + } +} + /// Filters the input according to a match query's token operator. #[derive(Debug)] pub struct FlatMatchFilterExec { @@ -3696,11 +4147,10 @@ impl ExecutionPlan for FlatMatchQueryExec { let document_column = self.document_column.clone(); let phrase_slop = self.params.phrase_slop; - // CPU time accumulator passed into `flat_bm25_search_stream_with_metrics` - // so it can attribute the spawn_cpu tokenize work and synchronous - // scoring back onto this node's `elapsed_compute`. Sharing the same - // `Time` handle that's already inside the FtsIndexMetrics avoids - // registering a duplicate metric. + // Lets `flat_bm25_search_stream_with_metrics` attribute the spawn_cpu + // tokenize work and synchronous scoring back onto this node's + // `elapsed_compute`. Reusing the handle already inside `FtsIndexMetrics` + // avoids registering a duplicate metric. let elapsed_compute = metrics.baseline_metrics.elapsed_compute().clone(); let column = query.column.ok_or(DataFusionError::Execution(format!( @@ -3782,7 +4232,7 @@ impl ExecutionPlan for FlatMatchQueryExec { match result { Ok((stream, scorer)) => { if let Some(producer) = shared_scorer_producer { - producer.publish(scorer); + producer.publish(Arc::new(scorer)); } Ok(stream) } @@ -3796,9 +4246,9 @@ impl ExecutionPlan for FlatMatchQueryExec { }) .try_flatten() .map(move |batch| { - // record_poll records output_rows, output_bytes, and output_batches - // on the shared BaselineMetrics — same pattern DataFusion's own - // FilterExec uses inside its hand-written poll_next. + // Records output_rows, output_bytes, and output_batches on the shared + // BaselineMetrics, as DataFusion's own FilterExec does in its + // hand-written poll_next. let poll = metrics_clone .baseline_metrics .record_poll(std::task::Poll::Ready(Some(batch))); @@ -4816,17 +5266,17 @@ mod tests { use lance_core::{ROW_ID, utils::address::RowAddress}; use lance_datafusion::datagen::DatafusionDatagenExt; use lance_datafusion::exec::{ExecutionStatsCallback, ExecutionSummaryCounts}; - use lance_datafusion::utils::PARTITIONS_SEARCHED_METRIC; + use lance_datafusion::utils::{INDEX_CACHE_HITS_METRIC, PARTITIONS_SEARCHED_METRIC}; use lance_datagen::{BatchCount, ByteCount, RowCount}; use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::query::{ - BooleanQuery, BoostQuery, FtsQuery, FtsSearchParams, MatchQuery, Occur, Operator, - PhraseQuery, collect_query_tokens, has_query_token, + BooleanQuery, BoostQuery, CombinedFieldsQuery, FtsQuery, FtsSearchParams, MatchQuery, + Occur, Operator, PhraseQuery, collect_query_tokens, has_query_token, }; use lance_index::scalar::inverted::{ - DocumentGranularity, FTS_SCHEMA, InvertedIndex, Language, SCORE_COL, - build_global_bm25_scorer, prepare_bm25_query, + CombinedFieldColumn, DocumentGranularity, FTS_SCHEMA, InvertedIndex, Language, SCORE_COL, + build_combined_bm25_scorer, build_global_bm25_scorer, prepare_bm25_query, }; use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams}; use lance_index::{IndexCriteria, IndexType}; @@ -4843,11 +5293,12 @@ mod tests { }; use super::{ - BoolSlot, BoostQueryExec, CompoundQueryExec, CrossColumnCompoundQueryExec, - FTS_SEGMENT_BIND_DURATION_METRIC, FlatMatchFilterExec, FlatMatchQueryExec, MatchQueryExec, - PhraseQueryExec, WAND_TIE_COMPLETION_BUDGET, WandExactnessCertificate, - build_boolean_query_children, classify_wand_exactness_certificate, default_text_tokenizer, - open_fts_segments, tokenizer_for_match_query, + BoolSlot, BoostQueryExec, CombinedFieldsQueryExec, CompoundQueryExec, + CrossColumnCompoundQueryExec, FTS_SEGMENT_BIND_DURATION_METRIC, FlatMatchFilterExec, + FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec, WAND_TIE_COMPLETION_BUDGET, + WandExactnessCertificate, build_boolean_query_children, + classify_wand_exactness_certificate, default_text_tokenizer, open_fts_segments, + tokenizer_for_match_query, }; use crate::io::exec::utils::IndexMetrics; use datafusion::physical_plan::empty::EmptyExec; @@ -5020,33 +5471,47 @@ mod tests { .ascii_folding(false) } - async fn create_tokenized_query_fixture(with_unindexed_append: bool) -> Dataset { - let mut dataset = lance_datagen::gen_batch() - .col( - "text", - lance_datagen::array::cycle_utf8_literals(&["first and second"]), - ) + /// Builds a dataset where every column in `columns` holds "first and second" + /// and carries its own inverted index. `combined_fields` needs an FTS index + /// per target column, so the columns are indexed one at a time rather than + /// as a single multi-column index. + /// + /// With `with_unindexed_append`, extra rows land outside the indices, which + /// forces the mixed plan where an indexed exec and its flat sibling both run. + async fn create_tokenized_query_fixture( + columns: &[&str], + with_unindexed_append: bool, + ) -> Dataset { + let corpus = || { + columns + .iter() + .fold(lance_datagen::gen_batch(), |batch, column| { + batch.col( + *column, + lance_datagen::array::cycle_utf8_literals(&["first and second"]), + ) + }) + }; + + let mut dataset = corpus() .into_ram_dataset(FragmentCount::from(1), FragmentRowCount::from(2)) .await .unwrap(); - dataset - .create_index( - &["text"], - IndexType::Inverted, - None, - &tokenized_query_index_params(), - true, - ) - .await - .unwrap(); + for column in columns { + dataset + .create_index( + &[*column], + IndexType::Inverted, + None, + &tokenized_query_index_params(), + true, + ) + .await + .unwrap(); + } if with_unindexed_append { - let appended = lance_datagen::gen_batch() - .col( - "text", - lance_datagen::array::cycle_utf8_literals(&["first and second"]), - ) - .into_reader_rows(RowCount::from(2), BatchCount::from(1)); + let appended = corpus().into_reader_rows(RowCount::from(2), BatchCount::from(1)); dataset.append(appended, None).await.unwrap(); } dataset @@ -5168,7 +5633,7 @@ mod tests { #[tokio::test] async fn shared_fts_scorer_reports_cancelled_producer() { - let scorer = Arc::new(super::SharedFtsScorer::new()); + let scorer = Arc::new(super::SharedFtsScorer::::new()); let producer = super::SharedFtsScorerProducer::new(scorer.clone()); drop(producer); @@ -5439,7 +5904,7 @@ mod tests { #[tokio::test] async fn test_analyze_plan_shows_indexed_and_flat_match_tokens() { - let dataset = create_tokenized_query_fixture(true).await; + let dataset = create_tokenized_query_fixture(&["text"], true).await; let query = MatchQuery::new("FIRST and SECOND".to_string()) .with_column(Some("text".to_string())) .with_operator(Operator::And); @@ -5468,7 +5933,7 @@ mod tests { #[tokio::test] async fn test_analyze_plan_shows_indexed_and_flat_phrase_tokens() { - let dataset = create_tokenized_query_fixture(true).await; + let dataset = create_tokenized_query_fixture(&["text"], true).await; let query = PhraseQuery::new("FIRST and SECOND".to_string()).with_column(Some("text".to_string())); let mut scanner = dataset.scan(); @@ -5490,7 +5955,7 @@ mod tests { #[tokio::test] async fn test_analyze_plan_shows_compound_leaf_tokens() { - let dataset = create_tokenized_query_fixture(false).await; + let dataset = create_tokenized_query_fixture(&["text"], false).await; let query = BooleanQuery::new([ ( Occur::Should, @@ -5591,6 +6056,138 @@ mod tests { ); } + /// The cross-field scorer build reads one posting-metadata row per + /// (term, column, partition), and both `combined_fields` execs must report + /// those to their query's index-cache counters. Unreported, `EXPLAIN ANALYZE` + /// on a cold cross-field query undercounts `index_cache_misses`, so its + /// `index_cache_hit_ratio` is not comparable with `match`'s on the same data. + #[tokio::test] + async fn test_combined_fields_scorer_build_reports_index_cache_metrics() { + let mut dataset = lance_datagen::gen_batch() + .col( + "title", + lance_datagen::array::cycle_utf8_literals(&[ + "hello world", + "lance search", + "hello lance", + ]), + ) + .col( + "body", + lance_datagen::array::cycle_utf8_literals(&["lance", "hello", "search hello"]), + ) + .into_ram_dataset(FragmentCount::from(1), FragmentRowCount::from(9)) + .await + .unwrap(); + for column in ["title", "body"] { + dataset + .create_index( + &[column], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + } + let dataset = Arc::new(dataset); + + let query = CombinedFieldsQuery::try_new( + "hello lance".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap(); + let params = FtsSearchParams::default().with_limit(Some(10)); + + // Open the same segments the execs will, both to size the expected lookup + // count and to hand the preset runs a scorer identical to the one they + // would have built. + let mut columns = Vec::new(); + for (column, weight) in query.weighted_columns() { + let segments = crate::index::scalar::inverted::load_segments( + &dataset, + column, + DocumentGranularity::Row, + ) + .await + .unwrap() + .expect("FTS index just created"); + let metrics_set = ExecutionPlanMetricsSet::new(); + let indices = open_fts_segments( + &dataset, + column, + &segments, + &IndexMetrics::new(&metrics_set, 0), + ) + .await + .unwrap(); + columns.push(CombinedFieldColumn { + column: column.to_string(), + weight, + indices, + }); + } + let mut tokenizer = columns[0].indices[0].tokenizer(); + let tokens = collect_query_tokens(query.terms(), &mut tokenizer); + let mut terms = (0..tokens.len()) + .map(|index| tokens.get_token(index).to_string()) + .collect::>(); + terms.sort_unstable(); + terms.dedup(); + let partitions: usize = columns + .iter() + .flat_map(|column| &column.indices) + .map(|index| index.partition_count()) + .sum(); + // One posting-metadata row per (term, column, partition). + let expected_lookups = terms.len() * partitions; + assert!(expected_lookups > 0); + let scorer = Arc::new( + build_combined_bm25_scorer(&columns, &tokens, None) + .await + .unwrap(), + ); + + // Warm the index cache first: a cold run's counts depend on which entry + // each lookup happens to load, while on a warm cache every lookup is a + // hit and the two runs below differ only in the scorer build. + let warmup = CombinedFieldsQueryExec::new( + dataset.clone(), + query.clone(), + params.clone(), + PreFilterSource::None, + ); + let expected = execute_results(&warmup).await.unwrap(); + assert!(!expected.is_empty()); + + let built = CombinedFieldsQueryExec::new( + dataset.clone(), + query.clone(), + params.clone(), + PreFilterSource::None, + ); + assert_eq!(execute_results(&built).await.unwrap(), expected); + let preset = CombinedFieldsQueryExec::new( + dataset.clone(), + query.clone(), + params.clone(), + PreFilterSource::None, + ) + .with_base_scorer(scorer.clone()); + assert_eq!( + execute_results(&preset).await.unwrap(), + expected, + "the preset scorer must match the one the exec builds", + ); + assert_eq!( + metric_value(&built, INDEX_CACHE_HITS_METRIC) + - metric_value(&preset, INDEX_CACHE_HITS_METRIC), + expected_lookups, + "CombinedFieldsQueryExec must report the scorer build's cache lookups", + ); + } + #[tokio::test] async fn test_match_query_exec_segment_selection() { let (dataset, segments, fragment_ids) = create_segment_selection_fixture().await; @@ -5790,6 +6387,24 @@ mod tests { ); } + #[tokio::test] + async fn test_combined_fields_exec_opens_all_committed_segments() { + let (dataset, _segments, fragment_ids) = create_segment_selection_fixture().await; + let query = CombinedFieldsQuery::try_new("quick".to_string(), vec!["text".to_string()]) + .unwrap() + .try_with_boosts(vec![2.0]) + .unwrap(); + let params = FtsSearchParams::default().with_limit(Some(20)); + + // One segment per fragment, so a scan that opens every committed segment + // returns every fragment's matching row. + let exec = CombinedFieldsQueryExec::new(dataset, query, params, PreFilterSource::None); + assert_eq!( + execute_row_ids(&exec).await.unwrap(), + expected_row_ids(&fragment_ids) + ); + } + #[tokio::test] async fn test_phrase_query_exec_segment_selection() { let (dataset, segments, fragment_ids) = create_segment_selection_fixture().await; diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index 7b99d732e55..a84fc04adb2 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; @@ -365,14 +366,18 @@ fn shared_prefilter_future( .boxed() } -pub(crate) fn build_prefilter( +/// Resolve a prefilter source into the future that yields its mask, ANDing in +/// the external row-address mask when the scan carries one. +/// +/// The external mask restricts index-side scoring to the masked rows (mirroring +/// the ANN path). It is independent of `overlay_block`, which the prefilter +/// applies separately to drop index entries staled by a data overlay. +fn prefilter_mask_future( context: Arc, partition: usize, prefilter_source: &PreFilterSource, - ds: Arc, - index_meta: &[IndexMetadata], - masks: PreFilterMasks, -) -> Result> { + external_mask: Option>, +) -> Result>>>> { let mut shared_filter = None; let prefilter_loader = match &prefilter_source { PreFilterSource::FilteredRowIds(src_node) => { @@ -407,12 +412,8 @@ pub(crate) fn build_prefilter( } PreFilterSource::None => None, }; - // Combine the external row-address mask (logical AND) with whatever the - // filter produced, so an FTS prefilter restricts BM25 scoring to masked rows - // (mirrors the ANN path). Independent of `overlay_block`, which the prefilter - // applies separately to drop index entries staled by a data overlay. - let mut prefilter = if let Some(shared_filter) = shared_filter { - let shared_filter = match masks.external_mask { + 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(), @@ -421,22 +422,54 @@ pub(crate) fn build_prefilter( .boxed(), None => shared_filter, }; - DatasetPreFilter::new_with_filter_future(ds, index_meta, Some(shared_filter)) - } else { - let prefilter_loader = match masks.external_mask { - Some(mask) => { - Some(Box::new(MaskAndLoader::new(mask, prefilter_loader)) as Box) - } - None => prefilter_loader, - }; - DatasetPreFilter::new(ds, index_meta, prefilter_loader) + return Ok(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(prefilter_loader.map(|loader| { + async move { loader.load().await.map(Arc::new) } + .in_current_span() + .boxed() + })) +} + +pub(crate) fn build_prefilter( + context: Arc, + partition: usize, + prefilter_source: &PreFilterSource, + ds: Arc, + index_meta: &[IndexMetadata], + masks: PreFilterMasks, +) -> Result> { + let filter = prefilter_mask_future(context, partition, prefilter_source, masks.external_mask)?; + let mut prefilter = DatasetPreFilter::new_with_filter_future(ds, index_meta, filter); if let Some(overlay_block) = masks.overlay_block { prefilter = prefilter.with_overlay_block(overlay_block); } Ok(Arc::new(prefilter)) } +/// Build a prefilter restricted to `fragments` rather than to the union of +/// `index_meta`'s fragment bitmaps. See +/// [`DatasetPreFilter::new_restricted_to_fragments`]. +pub(crate) fn build_prefilter_restricted_to_fragments( + context: Arc, + partition: usize, + prefilter_source: &PreFilterSource, + ds: Arc, + fragments: RoaringBitmap, + external_mask: Option>, +) -> Result> { + let filter = prefilter_mask_future(context, partition, prefilter_source, external_mask)?; + Ok(Arc::new(DatasetPreFilter::new_restricted_to_fragments( + ds, fragments, filter, + ))) +} + // Utility to convert an input (containing row ids) into a prefilter pub(crate) struct FilteredRowIdsToPrefilter(pub SendableRecordBatchStream); From 18a58b1732b83a2feec67388429b7beb23f07d99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Mon, 17 Aug 2026 12:03:55 +0200 Subject: [PATCH 2/4] test(fts): cover combined_fields end to end Dataset-level coverage for BM25F, checked against an independent brute-force BM25F reference (`lance_index::scalar::inverted::oracle`) that re-derives every statistic from the raw text, so it shares no code with the scan it checks. Each case asserts exact scores rather than just a hit set, because a wrong corpus size still returns the right rows in almost the right order. That is what pins down the parts easy to get subtly wrong: the per-column `w_f` factors, which are invisible at unit weights; ties, where the score-then-row ordering has to be deterministic across runs; and top-k across every k, where the pruning must agree with an exhaustive scan. Also covers the released-format fixtures (V1 and V2) so the row-granularity statistics path runs against real files rather than synthetic ones, nulls and empty strings, and the refusal paths: no index on any target column, and `fast_search` without full coverage. --- rust/lance-index/Cargo.toml | 3 + rust/lance-index/src/scalar/inverted.rs | 2 + .../src/scalar/inverted/combined.rs | 2 + .../src/scalar/inverted/combined/cursor.rs | 87 ++ .../src/scalar/inverted/combined/search.rs | 258 +++++ .../src/scalar/inverted/combined/stats.rs | 126 +++ .../src/scalar/inverted/combined/testing.rs | 303 ++++++ .../lance-index/src/scalar/inverted/oracle.rs | 143 +++ rust/lance/Cargo.toml | 1 + .../tests/dataset_fts_combined_fields.rs | 952 ++++++++++++++++++ rust/lance/src/dataset/tests/dataset_index.rs | 155 ++- rust/lance/src/dataset/tests/mod.rs | 1 + rust/lance/tests/query/inverted.rs | 73 +- 13 files changed, 2061 insertions(+), 45 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/combined/testing.rs create mode 100644 rust/lance-index/src/scalar/inverted/oracle.rs create mode 100644 rust/lance/src/dataset/tests/dataset_fts_combined_fields.rs diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index ab8ad36783c..32c53f1f5f3 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -96,6 +96,9 @@ jieba-rs = ["tokenizer-jieba"] lindera = ["tokenizer-lindera"] tokenizer-lindera = ["lance-tokenizer/tokenizer-lindera"] tokenizer-jieba = ["dep:jieba-rs", "lance-tokenizer/tokenizer-jieba"] +# Exposes `scalar::inverted::oracle`, the brute-force scoring reference, to +# downstream test and bench targets. Enable it as a dev-dependency feature only. +test-oracle = [] [build-dependencies] prost-build.workspace = true diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 7fef58b1874..f14fbf41e8b 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -14,6 +14,8 @@ mod iter; pub mod json; /// Brute-force scoring reference for tests and benches. Never built normally; see /// the module docs for the gating. +#[cfg(any(test, feature = "test-oracle"))] +pub mod oracle; pub mod parser; pub mod query; mod scorer; diff --git a/rust/lance-index/src/scalar/inverted/combined.rs b/rust/lance-index/src/scalar/inverted/combined.rs index 26c5d9cc42b..9de77185ce7 100644 --- a/rust/lance-index/src/scalar/inverted/combined.rs +++ b/rust/lance-index/src/scalar/inverted/combined.rs @@ -25,6 +25,8 @@ mod cursor; mod search; mod stats; +#[cfg(test)] +mod testing; use std::sync::Arc; diff --git a/rust/lance-index/src/scalar/inverted/combined/cursor.rs b/rust/lance-index/src/scalar/inverted/combined/cursor.rs index 62e8af0f250..33f538364f2 100644 --- a/rust/lance-index/src/scalar/inverted/combined/cursor.rs +++ b/rust/lance-index/src/scalar/inverted/combined/cursor.rs @@ -65,3 +65,90 @@ pub(super) fn build_term_postings( postings.sort_unstable_by_key(|(row_id, _)| *row_id); CombinedTermPostings { idf, postings } } + +#[cfg(test)] +mod tests { + use super::super::testing::{compressed_list, modern_identity_docs}; + use super::*; + use lance_core::utils::address::RowAddress; + use lance_select::RowAddrTreeMap; + + #[tokio::test] + async fn test_build_term_postings_merges_legacy_and_compressed() { + // A legacy (Plain, row-id-keyed, list-multiplicity) + // source and a compressed source merge into one ordered `tf'` stream, + // masked rows dropped and contributions summed in column order. + use super::super::super::index::PlainPostingList; + use arrow::buffer::ScalarBuffer; + + let scorer = CombinedFieldsBM25Scorer::new(1000, 12.0, HashMap::new()); + // Legacy column (weight 2): the posting keys directly on row ids; row 20 + // appears twice (list multiplicity), so its contributions sum. + let legacy = PostingList::Plain(PlainPostingList::new( + ScalarBuffer::from(vec![10u64, 20, 20, 30]), + ScalarBuffer::from(vec![1.0f32, 2.0, 3.0, 1.0]), + Some(0.0), + None, + )); + // Compressed column (weight 1): doc id 20 maps through the modern + // projection (the only representation a compressed posting is loaded + // alongside) to row 20; row 42 is blocked by the mask below. + let compressed = PostingList::Compressed(compressed_list(&[(20, 5), (42, 7)])); + let docs = modern_identity_docs(&vec![1u32; 64], &[]).await; + let sources = vec![ + LoadedSource { + weight: 2.0, + docs: docs.clone(), + is_legacy: true, + posting: legacy, + }, + LoadedSource { + weight: 1.0, + docs, + is_legacy: false, + posting: compressed, + }, + ]; + let mask = Arc::new(RowAddrMask::all_rows().also_block(RowAddrTreeMap::from_iter([42u64]))); + let term = build_term_postings("t", sources, &mask, &scorer); + + // row 10: 2*1 = 2; row 20: 2*2 + 2*3 + 1*5 = 15; row 30: 2*1 = 2. + // Row 42 is masked out entirely. + assert_eq!(term.postings, vec![(10, 2.0), (20, 15.0), (30, 2.0)]); + } + + #[tokio::test] + async fn test_build_term_postings_skips_tombstoned_addresses() { + // A remapped partition keeps a deleted document's DocId slot so the + // posting lists stay aligned and answers `TOMBSTONE_ROW` for its address. + // Nothing else stops that address: a default mask is an empty block list, + // which selects it, and `doc_length_at(TOMBSTONE_ROW) == 0` would give it + // the largest `doc_weight` there is, so it must be dropped at the source. + const DEAD_ROW: u64 = 20; + let scorer = CombinedFieldsBM25Scorer::new(1000, 12.0, HashMap::new()); + let docs = modern_identity_docs(&[4u32; 40], &[DEAD_ROW]).await; + assert_eq!( + docs.row_address(DEAD_ROW as u32), + RowAddress::TOMBSTONE_ROW, + "the deleted document must keep its slot as a tombstone" + ); + assert_eq!(docs.doc_length_at(RowAddress::TOMBSTONE_ROW), 0); + let sources = vec![LoadedSource { + weight: 2.0, + docs, + is_legacy: false, + posting: PostingList::Compressed(compressed_list(&[ + (10, 1), + (DEAD_ROW as u32, 7), + (30, 3), + ])), + }]; + let term = build_term_postings("t", sources, &Arc::new(RowAddrMask::default()), &scorer); + assert_eq!( + term.postings, + vec![(10, 2.0), (30, 6.0)], + "the tombstoned address must never be accumulated, and the live rows \ + must keep their exact contributions" + ); + } +} diff --git a/rust/lance-index/src/scalar/inverted/combined/search.rs b/rust/lance-index/src/scalar/inverted/combined/search.rs index ceae5ede3c8..f9650ffd5bc 100644 --- a/rust/lance-index/src/scalar/inverted/combined/search.rs +++ b/rust/lance-index/src/scalar/inverted/combined/search.rs @@ -177,3 +177,261 @@ pub async fn combined_fields_search( .map(|Reverse(doc)| (doc.row_id.0, doc.score.0)) .unzip()) } + +#[cfg(test)] +mod tests { + use super::super::super::index::InvertedListFormatVersion; + use super::super::super::scorer::idf; + use super::super::super::tokenizer::document_tokenizer::DocType; + use super::super::stats::build_combined_bm25_scorer; + use super::super::testing::{ + ElementRows, as_row_documents, combined_columns, combined_top_k, element_document_index, + }; + use super::*; + use crate::metrics::NoOpMetricsCollector; + use crate::prefilter::NoFilter; + use rstest::rstest; + + /// `title` is a `List` where row 0 holds ten `"alpha"` elements, row 1 + /// one `"beta"`, and rows 2..10 one `"gamma"`; `body` holds one `"zzz"` per + /// row. Row 0 matches `"alpha"` ten times over, so it must win + /// `"alpha beta"` at `limit = 1`. + /// + /// With document-granularity statistics it does not: `docCount'` counts the + /// 19 title elements instead of the 10 rows and `docFreq'("alpha")` counts 10 + /// element postings instead of 1 row, which collapses `alpha`'s `idf'` and + /// inflates `beta`'s enough to invert the ranking (row 1 at ~2.2984569). Row + /// granularity restores row 0 at ~3.1963050, bit-identical to the same data + /// indexed one document per row. + #[rstest] + #[case::v1(InvertedListFormatVersion::V1)] + #[case::v2(InvertedListFormatVersion::V2)] + #[tokio::test] + async fn test_combined_fields_scores_legacy_list_elements_by_row( + #[case] format_version: InvertedListFormatVersion, + ) { + let vocab = ["alpha", "beta", "gamma", "zzz"]; + let mut title: ElementRows = vec![vec![vec!["alpha"]; 10], vec![vec!["beta"]]]; + title.extend((2..10).map(|_| vec![vec!["gamma"]])); + let body: ElementRows = (0..10).map(|_| vec![vec!["zzz"]]).collect(); + + let (title_index, _title_dir) = + element_document_index(format_version, &vocab, &title).await; + let (body_index, _body_dir) = element_document_index(format_version, &vocab, &body).await; + + // Precondition: the fixture really is element-per-document, or the test + // covers nothing. + let title_docs = title_index.partitions[0] + .docs + .address_keyed() + .await + .unwrap(); + assert_eq!(title_docs.len(), 19, "one document per title list element"); + assert_eq!(title_docs.num_distinct_rows(), 10); + + // The single-column path reads document granularity; only the cross-field + // path reads row granularity. + let terms = ["alpha".to_owned(), "beta".to_owned()]; + assert_eq!( + title_index + .bm25_stats_for_terms(&terms, None) + .await + .unwrap(), + (19, 19, vec![10, 1]), + "document-granularity statistics must keep counting elements", + ); + assert_eq!( + title_index + .bm25_row_stats_for_terms(&terms, None) + .await + .unwrap(), + (19, 10, vec![1, 1]), + "row-granularity statistics must count distinct row ids", + ); + + let legacy = combined_top_k( + &combined_columns(vec![title_index, body_index]), + &["alpha", "beta"], + 1, + ) + .await; + assert_eq!(legacy.len(), 1); + assert_eq!(legacy[0].0, 0, "row 0 matches `alpha` ten times over"); + assert!( + (legacy[0].1 - 3.196_305).abs() < 1e-5, + "unexpected score {}", + legacy[0].1 + ); + + // Reindexing the same data one document per row must agree bit for bit. + let (row_title, _row_title_dir) = element_document_index( + InvertedListFormatVersion::V3, + &vocab, + &as_row_documents(&title), + ) + .await; + let (row_body, _row_body_dir) = element_document_index( + InvertedListFormatVersion::V3, + &vocab, + &as_row_documents(&body), + ) + .await; + let rebuilt = combined_top_k( + &combined_columns(vec![row_title, row_body]), + &["alpha", "beta"], + 1, + ) + .await; + assert_eq!( + legacy + .iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>(), + rebuilt + .iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>(), + "legacy index must score like the same data reindexed per row", + ); + } + + /// The single-column path stays at document granularity. It reads + /// [`InvertedIndex::bm25_stats_for_terms`] via `build_global_bm25_scorer` and + /// wand scores one posting per document, so a legacy list index is + /// self-consistent there: `docCount` and `docFreq` count elements, `tf` is one + /// element's frequency and `dl` one element's length. V1/V2 are released + /// stable formats, so these `(row_id, score)` bits pin that behavior; moving + /// it is a separate compatibility decision. + /// + /// The element domain also shows through in the results: row 0's ten `"alpha"` + /// documents surface as ten separate hits at the same row id. + #[rstest] + #[case::v1(InvertedListFormatVersion::V1)] + #[case::v2(InvertedListFormatVersion::V2)] + #[tokio::test] + async fn test_single_column_search_keeps_element_granularity( + #[case] format_version: InvertedListFormatVersion, + ) { + let vocab = ["alpha", "beta", "gamma"]; + let mut title: ElementRows = vec![vec![vec!["alpha"]; 10], vec![vec!["beta"]]]; + title.extend((2..10).map(|_| vec![vec!["gamma"]])); + let (index, _dir) = element_document_index(format_version, &vocab, &title).await; + + let terms = ["alpha".to_owned(), "beta".to_owned()]; + assert_eq!( + index.bm25_stats_for_terms(&terms, None).await.unwrap(), + (19, 19, vec![10, 1]), + "the single-column path's statistics stay at document granularity", + ); + + let tokens = Arc::new(Tokens::new(terms.to_vec(), DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(4))); + let scorer = crate::scalar::inverted::build_global_bm25_scorer( + std::slice::from_ref(&index), + &tokens, + ¶ms, + None, + ) + .await + .unwrap(); + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + Some(&scorer), + ) + .await + .unwrap(); + // `beta` (row 1) leads on its inflated element-level idf, and row 0 + // repeats once per matching element: the element domain showing through + // in the hits themselves. + assert_eq!(row_ids, vec![1, 0, 0, 0]); + // Every element is one token long and `avgdl` is the element average + // (19 tokens / 19 documents), so `doc_weight` is exactly 1 and each score + // is the bare element-granularity `idf`: `ln(1 + (19 - df + 0.5)/(df + 0.5))` + // with df = 1 for `beta` and df = 10 for `alpha`. + assert_eq!( + scores.iter().map(|s| s.to_bits()).collect::>(), + vec![0x4025_c6f0, 0x3f24_f495, 0x3f24_f495, 0x3f24_f495], + ); + assert_eq!(scores[0].to_bits(), idf(1, 19).to_bits()); + assert_eq!(scores[1].to_bits(), idf(10, 19).to_bits()); + } + + /// One legacy element-per-document column and one modern row-per-document + /// column in the same query. `docCount'` and `docFreq'` are a `max_f` across + /// columns, so a column left at element granularity poisons the blend even + /// when the other column is fine. The whole query must score as if the legacy + /// column had been reindexed per row. + #[tokio::test] + async fn test_combined_fields_mixes_legacy_and_modern_columns() { + let vocab = ["alpha", "beta"]; + let legacy_rows: ElementRows = vec![ + vec![vec!["alpha"]; 6], + vec![vec!["beta"]], + vec![vec!["beta"]], + vec![vec!["beta"]], + ]; + let modern_rows: ElementRows = (0..4).map(|_| vec![vec!["alpha", "beta"]]).collect(); + + let (legacy, _legacy_dir) = + element_document_index(InvertedListFormatVersion::V2, &vocab, &legacy_rows).await; + let (modern, _modern_dir) = + element_document_index(InvertedListFormatVersion::V3, &vocab, &modern_rows).await; + + let terms = ["alpha".to_owned(), "beta".to_owned()]; + assert_eq!( + legacy.bm25_stats_for_terms(&terms, None).await.unwrap(), + (9, 9, vec![6, 3]), + "the fixture must still be element-per-document", + ); + assert_eq!( + legacy.bm25_row_stats_for_terms(&terms, None).await.unwrap(), + (9, 4, vec![1, 3]), + ); + assert_eq!( + modern.bm25_row_stats_for_terms(&terms, None).await.unwrap(), + modern.bm25_stats_for_terms(&terms, None).await.unwrap(), + "a row-per-document index must be untouched", + ); + + let columns = combined_columns(vec![legacy, modern.clone()]); + let mixed = combined_top_k(&columns, &["alpha", "beta"], 4).await; + // `docCount'` is the 4 rows, not the legacy column's 9 elements. + let scorer = build_combined_bm25_scorer( + &columns, + &Tokens::new(vec!["alpha".to_owned(), "beta".to_owned()], DocType::Text), + None, + ) + .await + .unwrap(); + assert_eq!(scorer.doc_count(), 4); + + let (rebuilt_legacy, _rebuilt_dir) = element_document_index( + InvertedListFormatVersion::V3, + &vocab, + &as_row_documents(&legacy_rows), + ) + .await; + let rebuilt = combined_top_k( + &combined_columns(vec![rebuilt_legacy, modern]), + &["alpha", "beta"], + 4, + ) + .await; + assert_eq!(mixed.len(), 4); + assert_eq!( + mixed + .iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>(), + rebuilt + .iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>(), + ); + } +} diff --git a/rust/lance-index/src/scalar/inverted/combined/stats.rs b/rust/lance-index/src/scalar/inverted/combined/stats.rs index 11fba5d2ad0..8be1f095d52 100644 --- a/rust/lance-index/src/scalar/inverted/combined/stats.rs +++ b/rust/lance-index/src/scalar/inverted/combined/stats.rs @@ -80,3 +80,129 @@ pub async fn build_combined_bm25_scorer( doc_freq, )) } + +#[cfg(test)] +mod tests { + use super::super::super::index::InvertedListFormatVersion; + use super::super::super::tokenizer::document_tokenizer::DocType; + use super::super::testing::{ElementRows, combined_columns, element_document_index}; + use super::*; + use crate::metrics::LocalMetricsCollector; + use rstest::rstest; + + /// The scorer build must report its index cache lookups to the query's + /// collector, the way the single-column `build_global_bm25_scorer` does; see + /// [`build_combined_bm25_scorer`] for what goes wrong otherwise. + /// + /// The two format versions cover both statistics arms: + /// `bm25_row_stats_for_terms` delegates to the document-granularity path on + /// V3, and takes `InvertedPartition::row_stats_for_terms` on V1/V2, where this + /// element-per-document fixture makes it read the posting lists themselves. + #[rstest] + #[case::modern(InvertedListFormatVersion::V3)] + #[case::legacy_elements(InvertedListFormatVersion::V2)] + #[tokio::test] + async fn test_combined_scorer_build_reports_index_cache_lookups( + #[case] format_version: InvertedListFormatVersion, + ) { + let vocab = ["alpha", "beta"]; + // Row 0 owns two elements, so the legacy fixture is element-per-document + // and its row-granularity statistics take the deduplicating branch. + let rows: ElementRows = vec![vec![vec!["alpha"], vec!["beta"]], vec![vec!["alpha"]]]; + let (first, _first_dir) = element_document_index(format_version, &vocab, &rows).await; + let (second, _second_dir) = element_document_index(format_version, &vocab, &rows).await; + let columns = combined_columns(vec![first, second]); + let tokens = Tokens::new(vec!["alpha".to_owned(), "beta".to_owned()], DocType::Text); + // One lookup per (term, column, partition), over 2 terms and the single + // partition each fixture index holds. + let expected_lookups = 2 * columns.len(); + + let indexed = LocalMetricsCollector::default(); + build_combined_bm25_scorer(&columns, &tokens, Some(&indexed)) + .await + .unwrap(); + assert_eq!( + indexed.index_cache_hits() + indexed.index_cache_misses(), + expected_lookups, + "the scorer build must report every cache lookup it makes", + ); + } + + /// A legacy index whose documents already map one-to-one onto rows (a plain + /// `Utf8` column, or a list column with a single element per row) must be + /// completely unaffected: the row-granularity statistics equal the + /// document-granularity ones, so no score can move. + #[rstest] + #[case::v1(InvertedListFormatVersion::V1)] + #[case::v2(InvertedListFormatVersion::V2)] + #[case::v3(InvertedListFormatVersion::V3)] + #[tokio::test] + async fn test_row_stats_match_document_stats_without_list_multiplicity( + #[case] format_version: InvertedListFormatVersion, + ) { + let vocab = ["alpha", "beta"]; + // Row 0: a multi-token `Utf8` document. Row 1: a single-element list. + // Row 2: an empty list, so the row owns no document at all. + let rows: ElementRows = vec![ + vec![vec!["alpha", "beta", "alpha"]], + vec![vec!["beta"]], + vec![Vec::new()], + vec![vec!["alpha"]], + ]; + let (index, _dir) = element_document_index(format_version, &vocab, &rows).await; + + let docs = index.partitions[0].docs.address_keyed().await.unwrap(); + assert_eq!(docs.len(), 3, "the empty list must not become a document"); + assert_eq!(docs.num_distinct_rows(), docs.len()); + let terms = ["alpha".to_owned(), "beta".to_owned(), "absent".to_owned()]; + let documents = index.bm25_stats_for_terms(&terms, None).await.unwrap(); + assert_eq!(documents, (5, 3, vec![2, 2, 0])); + assert_eq!( + index.bm25_row_stats_for_terms(&terms, None).await.unwrap(), + documents + ); + } + + /// Rows a legacy list index never indexed (a null or empty list, a list of + /// empty strings) own no document, so they must not count toward + /// `docCount_f` or `docFreq_f`. This is the same rule the indexed and flat + /// sides already follow (`AddressKeyedDocuments::doc_length_at` reports 0 for them and + /// `FlatFieldStats::fold_row` skips a column with `dl_f == 0`), applied to + /// distinct-row counting. + #[rstest] + #[case::v1(InvertedListFormatVersion::V1)] + #[case::v2(InvertedListFormatVersion::V2)] + #[tokio::test] + async fn test_row_stats_skip_rows_a_legacy_list_index_never_indexed( + #[case] format_version: InvertedListFormatVersion, + ) { + let vocab = ["alpha", "beta"]; + let rows: ElementRows = vec![ + // Several elements, one of them empty. + vec![vec!["alpha"], Vec::new(), vec!["alpha"], vec!["beta"]], + // A null or empty list. + Vec::new(), + vec![vec!["beta"]], + // A list of nothing but empty strings. + vec![Vec::new(), Vec::new()], + ]; + let (index, _dir) = element_document_index(format_version, &vocab, &rows).await; + + let docs = index.partitions[0].docs.address_keyed().await.unwrap(); + assert_eq!(docs.len(), 4, "only rows 0 and 2 own documents"); + assert_eq!(docs.num_distinct_rows(), 2); + for unindexed in [1u64, 3] { + assert_eq!(docs.doc_length_at(unindexed), 0); + } + + let terms = ["alpha".to_owned(), "beta".to_owned()]; + assert_eq!( + index.bm25_stats_for_terms(&terms, None).await.unwrap(), + (4, 4, vec![2, 2]), + ); + assert_eq!( + index.bm25_row_stats_for_terms(&terms, None).await.unwrap(), + (4, 2, vec![1, 2]), + ); + } +} diff --git a/rust/lance-index/src/scalar/inverted/combined/testing.rs b/rust/lance-index/src/scalar/inverted/combined/testing.rs new file mode 100644 index 00000000000..e6f83fa60d6 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/testing.rs @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Shared test fixtures for the `combined_fields` submodules: term cursors, +//! document views, index and flat-scan corpora, and the exact reference scans +//! the tests compare against. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::{ArrayRef, RecordBatch, UInt32Array, UInt64Array}; +use lance_core::Result; +use lance_core::cache::{LanceCache, WeakLanceCache}; +use lance_core::utils::tempfile::TempObjDir; +use lance_io::object_store::ObjectStore; +use lance_select::RowAddrTreeMap; +use roaring::RoaringTreemap; + +use super::super::builder::{BLOCK_SIZE, InnerBuilder, PositionRecorder}; +use super::super::documents::{AddressKeyedDocuments, PartitionDocuments}; +use super::super::encoding::{ + MAX_POSTING_BLOCK_SIZE, compress_posting_list_with_tail_codec_and_block_size, +}; +use super::super::index::{ + CompressedPostingList, FTS_FORMAT_VERSION_KEY, InvertedIndex, InvertedListFormatVersion, + METADATA_FILE, NUM_TOKEN_COL, POSTING_BLOCK_SIZE_KEY, POSTING_TAIL_CODEC_KEY, + PostingListBuilder, PostingTailCodec, TOKEN_SET_FORMAT_KEY, TokenSetFormat, +}; +use super::super::query::{FtsSearchParams, Operator, Tokens}; +use super::super::tokenizer::document_tokenizer::DocType; +use super::super::tokenizer::{InvertedIndexParams, LEGACY_BLOCK_SIZE}; +use super::{CombinedFieldColumn, build_combined_bm25_scorer, combined_fields_search}; +use crate::metrics::NoOpMetricsCollector; +use crate::prefilter::NoFilter; +use crate::scalar::lance_format::LanceIndexStore; +use crate::scalar::{IndexStore, RowIdRemapper}; + +pub(super) fn compressed_list(postings: &[(u32, u32)]) -> CompressedPostingList { + let doc_ids: Vec = postings.iter().map(|(d, _)| *d).collect(); + let freqs: Vec = postings.iter().map(|(_, f)| *f).collect(); + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + freqs.iter(), + std::iter::repeat(0.0f32), + PostingTailCodec::VarintDelta, + BLOCK_SIZE, + ) + .unwrap(); + CompressedPostingList::new( + blocks, + 1.0, + doc_ids.len() as u32, + PostingTailCodec::VarintDelta, + BLOCK_SIZE, + None, + None, + ) +} + +#[derive(Debug)] +struct DeletedRows(Vec); + +impl RowIdRemapper for DeletedRows { + fn remap_row_id(&self, row_id: u64) -> Option { + (!self.0.contains(&row_id)).then_some(row_id) + } + + fn remap_row_addrs_tree_map(&self, _: &RowAddrTreeMap) -> RowAddrTreeMap { + unreachable!("document fixtures only remap single row ids") + } + + fn remap_row_ids_roaring_tree_map(&self, _: &RoaringTreemap) -> RoaringTreemap { + unreachable!("document fixtures only remap single row ids") + } + + fn remap_row_ids_record_batch(&self, _: RecordBatch, _: usize) -> Result { + unreachable!("document fixtures only remap single row ids") + } +} + +/// Modern identity documents, written to a real store and loaded through +/// [`PartitionDocuments::address_keyed`] so the fixture is the representation +/// production scores against rather than a stand-in. +/// +/// `dead_rows` are addresses a fragment-reuse remapper has deleted: the slot +/// survives (posting lists key on its DocId positionally) but +/// `row_address` answers [`RowAddress::TOMBSTONE_ROW`] for it. Every part of +/// the returned view is resident, so the temporary store is dropped here. +pub(super) async fn modern_identity_docs( + num_tokens: &[u32], + dead_rows: &[u64], +) -> AddressKeyedDocuments { + const DOCS_PATH: &str = "combined_fixture_docs.lance"; + let tmpdir = TempObjDir::default(); + let cache = Arc::new(LanceCache::no_cache()); + let store: Arc = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + cache.clone(), + )); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new(lance_core::ROW_ID, arrow_schema::DataType::UInt64, false), + arrow_schema::Field::new(NUM_TOKEN_COL, arrow_schema::DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from( + (0..num_tokens.len() as u64).collect::>(), + )) as ArrayRef, + Arc::new(UInt32Array::from(num_tokens.to_vec())) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = store.new_index_file(DOCS_PATH, schema).await.unwrap(); + writer.write_record_batch(batch).await.unwrap(); + writer.finish().await.unwrap(); + + let reader = store.open_index_file(DOCS_PATH).await.unwrap(); + let remapper: Option> = (!dead_rows.is_empty()) + .then(|| Arc::new(DeletedRows(dead_rows.to_vec())) as Arc); + PartitionDocuments::try_new( + store.clone(), + DOCS_PATH.to_owned(), + 0, + WeakLanceCache::from(cache.as_ref()), + reader.as_ref(), + remapper, + false, + ) + .unwrap() + .address_keyed() + .await + .unwrap() +} + +// Legacy element-per-document corpus statistics. +// +// Released V1/V2 indexes indexed every `List` element as its own document, +// so one row there owns a run of documents. `combined_fields` scores at row level, +// so `docCount'` / `docFreq'` must be row level too or `idf'` and `avgdl'` describe +// a different corpus than the frequencies they divide. The current builder writes +// one document per row, so these fixtures write the legacy partition files +// directly, and assert that they really are element-per-document so the coverage +// cannot lapse. + +/// A fixture corpus: `rows[row][element]` is one indexed element's tokens. +pub(super) type ElementRows<'a> = Vec>>; + +/// Build a one-partition FTS index whose documents are the individual +/// elements of `rows`, each carrying its row's id (row `r` has row id `r`). +/// `vocab` fixes the token order; tokens absent from `rows` are left out of +/// the token set, as the builder leaves them out. Elements that tokenize to +/// nothing are skipped, also as the builder does. +pub(super) async fn element_document_index( + format_version: InvertedListFormatVersion, + vocab: &[&str], + rows: &ElementRows<'_>, +) -> (Arc, TempObjDir) { + let vocab: Vec<&str> = vocab + .iter() + .copied() + .filter(|token| rows.iter().flatten().any(|element| element.contains(token))) + .collect(); + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let block_size = match format_version { + InvertedListFormatVersion::V3 => MAX_POSTING_BLOCK_SIZE, + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2 => LEGACY_BLOCK_SIZE, + }; + let posting_tail_codec = format_version.posting_tail_codec(); + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + format_version, + block_size, + ); + let mut postings: Vec = vocab + .iter() + .map(|token| { + builder.tokens.add((*token).to_owned()); + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + posting_tail_codec, + block_size, + ) + }) + .collect(); + let mut doc_id = 0u32; + for (row, elements) in rows.iter().enumerate() { + for element in elements { + if element.is_empty() { + continue; + } + for (token_id, token) in vocab.iter().enumerate() { + let freq = element.iter().filter(|t| *t == token).count() as u32; + if freq > 0 { + postings[token_id].add(doc_id, PositionRecorder::Count(freq)); + } + } + builder.docs.append(row as u64, element.len() as u32); + doc_id += 1; + } + } + builder.set_posting_lists(postings); + builder.write(store.as_ref()).await.unwrap(); + + let params = InvertedIndexParams::default() + .block_size(block_size) + .unwrap(); + let metadata = HashMap::from([ + ( + "partitions".to_owned(), + serde_json::to_string(&vec![0u64]).unwrap(), + ), + ("params".to_owned(), serde_json::to_string(¶ms).unwrap()), + ( + TOKEN_SET_FORMAT_KEY.to_owned(), + TokenSetFormat::default().to_string(), + ), + ( + POSTING_TAIL_CODEC_KEY.to_owned(), + posting_tail_codec.as_str().to_owned(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), + ]); + let mut writer = store + .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) + .await + .unwrap(); + writer.finish_with_metadata(metadata).await.unwrap(); + + let index = InvertedIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + (index, tmpdir) +} + +/// The same corpus as one document per row: a row's elements joined into a +/// single document, which is what the current builder writes for a +/// `List` column. +pub(super) fn as_row_documents<'a>(rows: &ElementRows<'a>) -> ElementRows<'a> { + rows.iter() + .map(|elements| { + let joined: Vec<&str> = elements.iter().flatten().copied().collect(); + if joined.is_empty() { + Vec::new() + } else { + vec![joined] + } + }) + .collect() +} + +pub(super) fn combined_columns(indices: Vec>) -> Vec { + indices + .into_iter() + .enumerate() + .map(|(slot, index)| CombinedFieldColumn { + column: format!("col{slot}"), + weight: 1.0, + indices: vec![index], + }) + .collect() +} + +/// Run a `combined_fields` top-k over `columns`, returning `(row_id, score)`. +pub(super) async fn combined_top_k( + columns: &[CombinedFieldColumn], + terms: &[&str], + limit: usize, +) -> Vec<(u64, f32)> { + let tokens = Tokens::new( + terms.iter().map(|t| (*t).to_owned()).collect(), + DocType::Text, + ); + let scorer = build_combined_bm25_scorer(columns, &tokens, None) + .await + .unwrap(); + let params = FtsSearchParams::new().with_limit(Some(limit)); + let (row_ids, scores) = combined_fields_search( + columns, + &tokens, + ¶ms, + Operator::Or, + &scorer, + Arc::new(NoFilter), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + row_ids.into_iter().zip(scores).collect() +} diff --git a/rust/lance-index/src/scalar/inverted/oracle.rs b/rust/lance-index/src/scalar/inverted/oracle.rs new file mode 100644 index 00000000000..3977d71bba0 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/oracle.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Brute-force BM25F reference for the `combined_fields` tests and bench. +//! +//! Re-derives every statistic from the raw text on each call, so it shares no code with +//! the scan it checks. `cfg(test)` here and the `test-oracle` feature downstream keep it +//! out of normal builds. + +use std::collections::HashSet; + +const K1: f32 = 1.2; +const B: f32 = 0.75; + +/// Splits on whitespace only, while the `simple` tokenizer splits on every +/// non-alphanumeric character. Keep test corpora alphanumeric or scores diverge. +fn tokenize(text: &str) -> Vec { + text.split_whitespace() + .map(|word| word.to_lowercase()) + .collect() +} + +/// Exact BM25F: `docFreq'`/`docCount'` are the max across fields, `tf'`/`dl'`/ +/// `sumTotalTermFreq'` weighted sums, and the `(k1 + 1)` numerator is Lance's. `None` +/// for a document that does not match. `columns` pairs each field's weight with its +/// text, one entry per document. +/// +/// A `""` value stands for anything that tokenizes to nothing (NULL, empty string, +/// empty list): absent from that column's statistics, so it must not raise +/// `docCount_f`/`docFreq_f`. List columns are their elements joined by a space. +#[cfg_attr(coverage, coverage(off))] +pub fn brute_force_bm25f( + columns: &[(f32, Vec<&str>)], + query: &str, + require_all_terms: bool, +) -> Vec> { + let mut seen = HashSet::new(); + let terms: Vec = tokenize(query) + .into_iter() + .filter(|term| seen.insert(term.clone())) + .collect(); + let num_docs = columns[0].1.len(); + let tokenized: Vec>> = columns + .iter() + .map(|(_, texts)| texts.iter().map(|text| tokenize(text)).collect()) + .collect(); + + // Zero-token documents are absent from a column's `DocSet`, and the flat scan's + // fold skips a column whose `dl_f == 0`, so they cannot raise `docCount_f` either. + let doc_count = tokenized + .iter() + .map(|column| column.iter().filter(|doc| !doc.is_empty()).count()) + .max() + .unwrap_or(0); + let mut sum_total_term_freq = 0f64; + let mut doc_freq = vec![0usize; terms.len()]; + for (column_index, (weight, _)) in columns.iter().enumerate() { + let total_tokens: usize = tokenized[column_index].iter().map(|doc| doc.len()).sum(); + sum_total_term_freq += *weight as f64 * total_tokens as f64; + for (term_index, term) in terms.iter().enumerate() { + let df = tokenized[column_index] + .iter() + .filter(|doc| doc.contains(term)) + .count(); + doc_freq[term_index] = doc_freq[term_index].max(df); + } + } + let avgdl = if doc_count == 0 { + 0.0 + } else { + (sum_total_term_freq / doc_count as f64) as f32 + }; + let idf: Vec = doc_freq + .iter() + .map(|&df| { + if df == 0 { + 0.0 + } else { + ((doc_count as f32 - df as f32 + 0.5) / (df as f32 + 0.5) + 1.0).ln() + } + }) + .collect(); + + (0..num_docs) + .map(|doc| { + let mut tf = vec![0f32; terms.len()]; + let mut dl = 0f32; + for (column_index, (weight, _)) in columns.iter().enumerate() { + dl += weight * tokenized[column_index][doc].len() as f32; + for (term_index, term) in terms.iter().enumerate() { + let count = tokenized[column_index][doc] + .iter() + .filter(|token| *token == term) + .count(); + tf[term_index] += weight * count as f32; + } + } + let matched = if require_all_terms { + tf.iter().all(|&freq| freq > 0.0) + } else { + tf.iter().any(|&freq| freq > 0.0) + }; + if !matched { + return None; + } + let mut score = 0.0; + for term_index in 0..terms.len() { + if tf[term_index] <= 0.0 { + continue; + } + let doc_norm = K1 * (1.0 - B + B * dl / avgdl); + score += + idf[term_index] * (K1 + 1.0) * tf[term_index] / (tf[term_index] + doc_norm); + } + Some(score) + }) + .collect() +} + +/// The ids [`brute_force_bm25f`] matched. +#[cfg_attr(coverage, coverage(off))] +pub fn brute_force_ids(scores: &[Option]) -> HashSet { + (0..scores.len()) + .filter(|&doc| scores[doc].is_some()) + .map(|doc| doc as i32) + .collect() +} + +/// The `k` highest-scoring ids, descending, id as tiebreak. +#[cfg_attr(coverage, coverage(off))] +pub fn brute_force_top_k(scores: &[Option], k: usize) -> Vec { + let mut ranked: Vec<(i32, f32)> = scores + .iter() + .enumerate() + .filter_map(|(doc, score)| score.map(|score| (doc as i32, score))) + .collect(); + ranked.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then(a.0.cmp(&b.0)) + }); + ranked.into_iter().take(k).map(|(doc, _)| doc).collect() +} diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 417d41a5fe1..4a0e3123701 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -117,6 +117,7 @@ all_asserts.workspace = true mock_instant.workspace = true lance-testing = { workspace = true } lance-io = { workspace = true, features = ["test-util"] } +lance-index = { workspace = true, features = ["test-oracle"] } tracing-subscriber = { version = "0.3.17", features = ["env-filter"] } env_logger.workspace = true tempfile.workspace = true diff --git a/rust/lance/src/dataset/tests/dataset_fts_combined_fields.rs b/rust/lance/src/dataset/tests/dataset_fts_combined_fields.rs new file mode 100644 index 00000000000..914f9fa5055 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_fts_combined_fields.rs @@ -0,0 +1,952 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::{HashMap, HashSet}; +use std::vec; + +use crate::Dataset; + +use crate::dataset::write::{WriteMode, WriteParams}; +use crate::index::DatasetIndexExt; +use crate::utils::test::copy_test_data_to_tmp; +use arrow::array::AsArray; +use arrow_array::RecordBatch; +use arrow_array::record_batch; +use arrow_array::{ + RecordBatchIterator, + types::{Float32Type, Int32Type, Int64Type}, +}; +use lance_core::utils::tempfile::TempStrDir; +use lance_index::IndexType; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::scalar::inverted::{ + Language, + query::{MatchQuery, Operator}, + tokenizer::InvertedIndexParams, +}; + +use lance_index::scalar::inverted::builder::BLOCK_SIZE; +use lance_index::scalar::inverted::oracle::{brute_force_bm25f, brute_force_ids}; +use lance_index::scalar::inverted::query::{CombinedFieldsQuery, FtsQuery, MultiMatchQuery}; +use rstest::rstest; + +/// Whitespace-tokenizing index params, so `brute_force_bm25f` can mirror them. +fn combined_fields_test_params() -> InvertedIndexParams { + InvertedIndexParams::new("simple".to_string(), Language::English) + .lower_case(true) + .stem(false) + .remove_stop_words(false) + .ascii_folding(false) + .max_token_length(None) +} + +/// The standard `id`/`title`/`body` batch these tests score over. +fn combined_fields_batch(ids: Vec, titles: Vec<&str>, bodies: Vec<&str>) -> RecordBatch { + record_batch!( + ("id", Int32, ids), + ("title", Utf8, titles), + ("body", Utf8, bodies) + ) + .unwrap() +} + +/// Write one batch to `uri`, creating the dataset or appending to it. +/// +/// `write_params` of `None` takes the write defaults: one fragment, row addresses +/// rather than stable row ids. Pass a [`WriteParams`] to pick a fragment layout +/// (`max_rows_per_file`), append instead of create (`mode`), or switch row-id +/// scheme (`enable_stable_row_ids`). +async fn write_fts_dataset( + uri: &str, + batch: RecordBatch, + write_params: Option, +) -> Dataset { + let schema = batch.schema(); + Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + uri, + write_params, + ) + .await + .unwrap() +} + +/// Append one batch to an existing dataset. The appended fragments carry no +/// index, which is what routes their rows to the flat scan. +/// +/// `max_rows_per_file` of `None` takes the write default, i.e. one fragment for +/// the whole batch; `Some(n)` spreads it over several. +async fn append_fts_dataset( + uri: &str, + batch: RecordBatch, + max_rows_per_file: Option, +) -> Dataset { + let mut params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + }; + if let Some(max_rows_per_file) = max_rows_per_file { + params.max_rows_per_file = max_rows_per_file; + } + write_fts_dataset(uri, batch, Some(params)).await +} + +/// Build one inverted index per column in `columns`, at the dataset's current +/// version. Call sites list the columns explicitly because which ones carry an +/// index, and at which version, is what the coverage-skew tests vary. +async fn create_inverted_indices( + dataset: &mut Dataset, + columns: &[&str], + params: &InvertedIndexParams, +) { + for &column in columns { + dataset + .create_index(&[column], IndexType::Inverted, None, params, true) + .await + .unwrap(); + } +} + +/// Write `titles`/`bodies` as the standard `id`/`title`/`body` batch, ids running +/// `0..titles.len()`, and index both text columns with [`combined_fields_test_params`]. +/// `max_rows_per_file` picks the fragment layout, `None` taking the default of one. +/// +/// The returned [`TempStrDir`] owns the dataset directory, so the caller has to keep it +/// bound while it reads the dataset; it is also what [`append_fts_dataset`] appends to. +async fn indexed_two_column_dataset( + titles: &[&str], + bodies: &[&str], + max_rows_per_file: Option, +) -> (TempStrDir, Dataset) { + let ids = (0..titles.len() as i32).collect(); + let batch = combined_fields_batch(ids, titles.to_vec(), bodies.to_vec()); + let write_params = max_rows_per_file.map(|max_rows_per_file| WriteParams { + max_rows_per_file, + ..Default::default() + }); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, write_params).await; + let params = combined_fields_test_params(); + create_inverted_indices(&mut dataset, &["title", "body"], ¶ms).await; + (test_uri, dataset) +} + +fn combined_query(terms: &str, operator: Operator) -> FtsQuery { + combined_query_with_boosts(terms, operator, None) +} + +/// `boosts` of `None` leaves every weight at the default `1.0`. Explicit weights +/// exercise the `w_f` factors in `tf'`/`dl'`, which are invisible at 1.0: an +/// implementation that dropped `weight` still scores correctly with unit weights. +fn combined_query_with_boosts( + terms: &str, + operator: Operator, + boosts: Option>, +) -> FtsQuery { + combined_query_over(&["title", "body"], terms, operator, boosts) +} + +/// Like [`combined_query_with_boosts`] but over an explicit column list, for the +/// datasets that are not the two-column `title`/`body` shape. +fn combined_query_over( + columns: &[&str], + terms: &str, + operator: Operator, + boosts: Option>, +) -> FtsQuery { + let query = CombinedFieldsQuery::try_new( + terms.to_string(), + columns.iter().map(|c| c.to_string()).collect(), + ) + .unwrap(); + let query = match boosts { + Some(boosts) => query.try_with_boosts(boosts).unwrap(), + None => query, + }; + FtsQuery::CombinedFields(query.with_operator(operator)) +} + +/// Run an unlimited full-text query and return the matched `id`s in result order. +async fn fts_result_ids(dataset: &Dataset, query: FtsQuery) -> Vec { + fts_result_id_scores(dataset, query, None) + .await + .into_iter() + .map(|(id, _)| id) + .collect() +} + +/// Run a full-text query and return `(id, score)` pairs in result order. +/// +/// `limit` is the query's top-k: `Some(k)` arms the pruning path, `None` takes the +/// exact scan over every match. +async fn fts_result_id_scores( + dataset: &Dataset, + query: FtsQuery, + limit: Option, +) -> Vec<(i32, f32)> { + let batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query).limit(limit)) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_primitive::(); + let scores = batch["_score"].as_primitive::(); + (0..batch.num_rows()) + .map(|i| (ids.value(i), scores.value(i))) + .collect() +} + +/// The scan's ids as a set. A set compare on its own would hide a duplicate +/// emission, so the count is pinned here. +fn unique_ids(ids: Vec) -> HashSet { + let set: HashSet = ids.iter().copied().collect(); + assert_eq!(set.len(), ids.len(), "duplicate row ids emitted: {ids:?}"); + set +} + +/// Assert a scan's `(id, score)` results are exactly the documents +/// [`brute_force_bm25f`] matched, with scores equal to within 1e-3. +/// +/// `context` names the invariant the call site is pinning down and is appended to +/// every failure message, so a failure says which one broke. +fn assert_matches_brute_force(actual: &[(i32, f32)], expected: &[Option], context: &str) { + let expected_ids = brute_force_ids(expected); + let actual_ids: HashSet = actual.iter().map(|(id, _)| *id).collect(); + // A set compare would hide a duplicate emission, so pin the row count too. + assert_eq!( + actual.len(), + expected_ids.len(), + "duplicate row emitted ({context}): {actual:?}" + ); + assert_eq!(actual_ids, expected_ids, "matched ids differ ({context})"); + for (id, score) in actual { + let want = expected[*id as usize].expect("scan returned an unmatched doc"); + assert!( + (score - want).abs() < 1e-3, + "score mismatch for id {id} ({context}): scan={score}, brute force={want}" + ); + } +} + +/// Assert `actual` is a valid exact top-`k` of `expected`: the right hit count, the exact +/// ranking's scores, each returned doc carrying its own score, no duplicate or +/// below-cutoff id. Identity is membership rather than an exact id set because ties at +/// the k-th score make the winning set ambiguous. +fn assert_topk_matches_brute_force( + actual: &[(i32, f32)], + expected: &[Option], + k: usize, + context: &str, +) { + let mut expected_ranked: Vec = expected.iter().filter_map(|score| *score).collect(); + expected_ranked.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)); + // Not an early return: an all-empty comparison satisfies everything below. + assert!( + !expected_ranked.is_empty(), + "the oracle matched nothing, so there is nothing to compare ({context})" + ); + + assert_eq!( + actual.len(), + expected_ranked.len().min(k), + "hit count mismatch for k={k} ({context})" + ); + + let mut actual_scores: Vec = actual.iter().map(|(_, score)| *score).collect(); + actual_scores.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)); + for (got, want) in actual_scores.iter().zip(&expected_ranked) { + assert!( + (got - want).abs() < 1e-3, + "top-{k} score mismatch ({context}): pruned={actual_scores:?} exact={expected_ranked:?}" + ); + } + + // Per id: sorted scores alone accept ids and scores paired up wrongly. + for (id, score) in actual { + let want = expected[*id as usize].expect("scan returned an unmatched doc"); + assert!( + (score - want).abs() < 1e-3, + "score mismatch for id {id} at k={k} ({context}): scan={score}, brute force={want}" + ); + } + + let returned_ids = unique_ids(actual.iter().map(|(id, _)| *id).collect()); + let cutoff = expected_ranked[actual.len() - 1]; + for id in &returned_ids { + let want = expected[*id as usize].expect("scan returned an unmatched doc"); + assert!( + want >= cutoff - 1e-3, + "id {id} scores {want}, below the top-{k} cutoff {cutoff} ({context})" + ); + } +} + +#[tokio::test] +async fn test_fts_combined_fields_cross_field_and() { + // combined_fields (BM25F) treats the target columns as one virtual field, so an + // AND query matches when each term appears in at least one field. best_fields + // (MultiMatch) evaluates AND per field, so it only matches a single field that + // contains every term. + let params = InvertedIndexParams::default(); + // row 0: the two terms are split across the fields; + // row 1: both terms live in a single field; + // row 2: only one of the two terms appears anywhere. + let batch = combined_fields_batch( + vec![0, 1, 2], + vec!["john", "john smith", "john"], + vec!["smith", "foo", "alice"], + ); + let test_uri = TempStrDir::default(); + // Spread the rows across fragments so the merged scan exercises multi-fragment + // row-id resolution. + let mut dataset = write_fts_dataset( + &test_uri, + batch, + Some(WriteParams { + max_rows_per_file: 1, + ..Default::default() + }), + ) + .await; + create_inverted_indices(&mut dataset, &["title", "body"], ¶ms).await; + + let columns = vec!["title".to_string(), "body".to_string()]; + let combined = |op| combined_query("john smith", op); + let multi = |op| { + FtsQuery::MultiMatch( + MultiMatchQuery::try_new("john smith".to_string(), columns.clone()) + .unwrap() + .with_operator(op), + ) + }; + + // AND over the virtual field: row 0 (john|title + smith|body) and row 1 + // (both terms in title) match; row 2 (no "smith" anywhere) does not. + assert_eq!( + unique_ids(fts_result_ids(&dataset, combined(Operator::And)).await), + HashSet::from([0, 1]) + ); + // best_fields AND matches only row 1, where a single field holds both terms. + assert_eq!( + unique_ids(fts_result_ids(&dataset, multi(Operator::And)).await), + HashSet::from([1]) + ); + + // OR matches any doc containing either term: every row has "john". + assert_eq!( + unique_ids(fts_result_ids(&dataset, combined(Operator::Or)).await), + HashSet::from([0, 1, 2]) + ); + assert_eq!( + unique_ids(fts_result_ids(&dataset, multi(Operator::Or)).await), + HashSet::from([0, 1, 2]) + ); +} + +#[tokio::test] +async fn test_fts_combined_fields_boost_ranking() { + // Per-column boosts move a document up the ranking in the BM25F direction: + // boosting the field a term lives in counts that term more. Stemming and + // stop words are disabled so the filler tokens contribute to document length + // ("other" is an English stop word). + let params = InvertedIndexParams::new("simple".to_string(), Language::English) + .stem(false) + .remove_stop_words(false); + // id 0 has the term only in `title`; id 1 has it only in the shorter `body`. + let batch = combined_fields_batch( + vec![0, 1], + vec!["lance", "other"], + vec!["other other other", "lance"], + ); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, None).await; + create_inverted_indices(&mut dataset, &["title", "body"], ¶ms).await; + + let combined = + |boosts: Vec| combined_query_with_boosts("lance", Operator::Or, Some(boosts)); + // Compare by score rather than result row-order (FTS batch order is not a + // guaranteed ranking; only the scores are). + let scores = |ids: Vec<(i32, f32)>| ids.into_iter().collect::>(); + + // Equal weights: the shorter document (id 1, term in the 1-token body) wins. + let equal = scores(fts_result_id_scores(&dataset, combined(vec![1.0, 1.0]), None).await); + assert!( + equal[&1] > equal[&0], + "equal weights: {equal:?} should rank id 1 above id 0" + ); + // Boosting `title` 3x lifts id 0 (term in title) above id 1. + let boosted = scores(fts_result_id_scores(&dataset, combined(vec![3.0, 1.0]), None).await); + assert!( + boosted[&0] > boosted[&1], + "title-boosted: {boosted:?} should rank id 0 above id 1" + ); +} + +#[tokio::test] +async fn test_fts_combined_fields_nulls() { + // NULL edge cases: a doc null in one field is still scored via the other; a + // doc null in every field never matches; an entirely-null column is a no-op. + let params = InvertedIndexParams::default(); + let batch = record_batch!( + ("id", Int32, [0, 1, 2, 3]), + ("title", Utf8, [Some("lance"), None, None, Some("lance")]), + ("body", Utf8, [None, Some("lance"), None, Some("lance")]), + ("empty", Utf8, [None::<&str>, None, None, None]) + ) + .unwrap(); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, None).await; + create_inverted_indices(&mut dataset, &["title", "body", "empty"], ¶ms).await; + + // Null in one field (0, 1) is scored via the other; null in both (2) never + // matches; present in both (3) matches once. + let query = combined_query("lance", Operator::Or); + assert_eq!( + unique_ids(fts_result_ids(&dataset, query).await), + HashSet::from([0, 1, 3]) + ); + + // An entirely-null column contributes nothing (docCount' uses the max), so + // combining it with `title` matches exactly the `title` hits. + let with_empty = combined_query_over(&["title", "empty"], "lance", Operator::Or, None); + assert_eq!( + unique_ids(fts_result_ids(&dataset, with_empty).await), + HashSet::from([0, 3]) + ); +} + +#[tokio::test] +async fn test_fts_combined_fields_matches_brute_force_bm25f() { + // Validate the combined-fields scan against an independent brute-force BM25F + // reference (the primary correctness oracle). Stemming and stop words are + // disabled so the reference can tokenize by whitespace. + let titles = vec!["aa bb", "aa", "cc dd", "aa aa bb", "cc"]; + let bodies = vec!["cc", "bb cc dd", "aa", "dd", "aa bb"]; + let (_test_uri, dataset) = indexed_two_column_dataset(&titles, &bodies, None).await; + + let weights = [2.0f32, 1.0f32]; + let query = combined_query_with_boosts("aa bb", Operator::Or, Some(weights.to_vec())); + + let expected = brute_force_bm25f( + &[(weights[0], titles.clone()), (weights[1], bodies.clone())], + "aa bb", + false, + ); + let actual = fts_result_id_scores(&dataset, query, None).await; + assert_matches_brute_force(&actual, &expected, "the fully-indexed two-column scan"); +} + +/// `combined_fields` against a real released-format index read from disk rather +/// than a synthetic fixture, so the row-granularity statistics path +/// ([`lance_index::scalar::inverted::index::InvertedIndex::bm25_row_stats_for_terms`]) +/// is exercised on actual V1 and V2 files. +/// +/// The checked-in fixtures index a plain `Utf8` column, so every row owns exactly +/// one document and they cannot reproduce the list-element multiplicity the +/// row-granularity path exists for (that lives in `combined/search.rs`'s unit tests). +/// What they do pin is that a released-format index scores against the +/// brute-force BM25F reference. +#[rstest] +#[case::v1("v3.0.1/fts_v1", 1)] +#[case::v2("v4.0.1/fts_v2", 2)] +#[tokio::test] +async fn test_fts_combined_fields_on_released_format_fixture( + #[case] fixture_path: &str, + #[case] expected_version: i32, +) { + let test_dir = copy_test_data_to_tmp(fixture_path).unwrap(); + let dataset = Dataset::open(&test_dir.path_str()).await.unwrap(); + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].index_version, expected_version); + + // `test_data/{v3.0.1,v4.0.1}/datagen.py` writes 300 rows of + // "lance database compatibility shared" for id % 3 == 0 and + // "database lance compatibility shared" otherwise. + const NUM_ROWS: usize = 300; + let texts: Vec<&str> = (0..NUM_ROWS) + .map(|id| { + if id % 3 == 0 { + "lance database compatibility shared" + } else { + "database lance compatibility shared" + } + }) + .collect(); + let expected = brute_force_bm25f(&[(1.0, texts)], "lance compatibility", false); + + let query = combined_query_over(&["text"], "lance compatibility", Operator::Or, None); + let batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query)) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + let scores = batch + .column_by_name("_score") + .unwrap() + .as_primitive::(); + assert_eq!(batch.num_rows(), NUM_ROWS); + for i in 0..batch.num_rows() { + let id = ids.value(i) as usize; + let want = expected[id].expect("every fixture row carries both query terms"); + assert!( + (scores.value(i) - want).abs() < 1e-3, + "score mismatch for id {id}: scan={}, brute force={want}", + scores.value(i), + ); + } +} + +/// Drive the top-k path end-to-end through the scanner and confirm the pruned +/// top-k equals the exact brute-force BM25F top-k for OR and AND across every k. +/// +/// Only a limited query arms MAXSCORE pruning; an unlimited one takes the exact +/// scan. The corpus is skewed, mixing a rare high-idf term ("zeta") with a common +/// low-idf term ("beta"), so MAXSCORE makes "beta" non-essential. The OR query +/// additionally carries an absent term ("missingterm", idf' == 0) to check the +/// clamped-ceiling handling. +/// +/// `max_rows_per_file` picks the partition layout, so the multi-partition case also +/// exercises the cursors' cross-partition row-id merge. +#[rstest] +#[case::single_partition(None)] +#[case::multi_partition(Some(7))] +#[tokio::test] +async fn test_fts_combined_fields_topk_matches_brute_force( + #[case] max_rows_per_file: Option, +) { + let n = 40usize; + let titles: Vec = (0..n) + .map(|i| { + if i % 9 == 0 { + "zeta gamma".to_string() + } else { + "gamma".to_string() + } + }) + .collect(); + let bodies: Vec = (0..n) + .map(|i| { + // "beta" fills every body with a growing count (common, low idf); + // one body also carries the rare "zeta" so a doc can hold it in + // either field. + let beta = vec!["beta"; 1 + i % 3].join(" "); + if i == 3 { format!("{beta} zeta") } else { beta } + }) + .collect(); + + let title_refs: Vec<&str> = titles.iter().map(|s| s.as_str()).collect(); + let body_refs: Vec<&str> = bodies.iter().map(|s| s.as_str()).collect(); + let (_test_uri, dataset) = + indexed_two_column_dataset(&title_refs, &body_refs, max_rows_per_file).await; + + let weights = [2.0f32, 1.0f32]; + + for operator in [Operator::Or, Operator::And] { + let require_all = operator == Operator::And; + // The absent term is only valid for OR (AND with an absent term matches + // nothing); keep it out of the AND case. + let query_str = if require_all { + "zeta beta" + } else { + "zeta beta missingterm" + }; + let expected = brute_force_bm25f( + &[ + (weights[0], title_refs.clone()), + (weights[1], body_refs.clone()), + ], + query_str, + require_all, + ); + for k in [1usize, 2, 3, 5, 10, 50] { + let query = combined_query_with_boosts(query_str, operator, Some(weights.to_vec())); + let actual = fts_result_id_scores(&dataset, query, Some(k as i64)).await; + assert_topk_matches_brute_force(&actual, &expected, k, &format!("op={operator:?}")); + } + } +} + +/// Top-k recall where a term's postings span several [`BLOCK_SIZE`] blocks, so the +/// cursors decode and merge across block boundaries. `beta` sits under `docCount'` +/// (at equality its ceiling is 0 and pruning is trivial) with fewer `zeta` +/// documents than `beta` has blocks. +/// +/// `alpha` has no such margin and covers the same paths with pruning that cannot +/// fire. `max_rows_per_file` varies fragments, not partitions. +#[rstest] +#[case::single_fragment(None)] +#[case::multi_fragment(Some(BLOCK_SIZE))] +#[tokio::test] +async fn test_fts_combined_fields_topk_matches_brute_force_across_blocks( + #[case] max_rows_per_file: Option, +) { + const CORPUS_BLOCKS: usize = 16; + let num_docs = CORPUS_BLOCKS * BLOCK_SIZE; + + // Also derive the expected match counts, so corpus and expectations cannot drift. + // Three `zeta` probes leave most of `beta`'s ten blocks untouched, and the first is + // early so the threshold rises before much is decoded. + let has_zeta = |doc: usize| [3, 703, 1403].contains(&doc); + let has_beta = |doc: usize| doc % 5 >= 2; + let has_alpha = |doc: usize| doc.is_multiple_of(7); + // A subset, so its cursor merges two sources and no match count changes. + let has_beta_in_title = |doc: usize| doc % 5 == 3; + let count = |carries: &dyn Fn(usize) -> bool| (0..num_docs).filter(|&d| carries(d)).count(); + + // `gamma` and `delta` are filler: no query uses them, they just vary `dl'`. + let titles: Vec = (0..num_docs) + .map(|doc| { + let mut title = if has_zeta(doc) { + "zeta gamma".to_string() + } else { + "gamma".to_string() + }; + if has_beta_in_title(doc) { + title.push_str(" beta"); + } + title + }) + .collect(); + let bodies: Vec = (0..num_docs) + .map(|doc| { + let mut body = if has_beta(doc) { + vec!["beta"; 1 + doc % 3].join(" ") + } else { + "delta".to_string() + }; + if has_alpha(doc) { + body.push_str(" alpha"); + } + body + }) + .collect(); + + // `beta` multi-block with a real ceiling, and rarer terms above it by ceiling. + let beta_docs = count(&has_beta); + let zeta_docs = count(&has_zeta); + let alpha_docs = count(&has_alpha); + assert!( + beta_docs > 4 * BLOCK_SIZE && beta_docs < num_docs, + "beta must span several blocks without covering the corpus, got {beta_docs} of {num_docs}" + ); + assert!( + zeta_docs < beta_docs / BLOCK_SIZE, + "zeta ({zeta_docs} docs) must stay under beta's block count ({})", + beta_docs / BLOCK_SIZE + ); + assert!( + 2 * alpha_docs < beta_docs && 2 * zeta_docs < beta_docs, + "the rare terms must stay well below beta ({beta_docs}): zeta={zeta_docs} alpha={alpha_docs}" + ); + let beta_title_docs = count(&has_beta_in_title); + assert!( + beta_title_docs > 0 && count(&|doc| has_beta_in_title(doc) && !has_beta(doc)) == 0, + "beta needs a second source in `title`, drawn from its own documents, got {beta_title_docs}" + ); + + let title_refs: Vec<&str> = titles.iter().map(|title| title.as_str()).collect(); + let body_refs: Vec<&str> = bodies.iter().map(|body| body.as_str()).collect(); + let (_test_uri, dataset) = + indexed_two_column_dataset(&title_refs, &body_refs, max_rows_per_file).await; + + let weights = vec![2.0f32, 1.0f32]; + + for (terms, operator, want_matches) in [ + ( + "zeta beta", + Operator::Or, + count(&|doc| has_zeta(doc) || has_beta(doc)), + ), + ( + "alpha beta", + Operator::Or, + count(&|doc| has_alpha(doc) || has_beta(doc)), + ), + ( + "zeta beta", + Operator::And, + count(&|doc| has_zeta(doc) && has_beta(doc)), + ), + ] { + let expected = brute_force_bm25f( + &[ + (weights[0], title_refs.clone()), + (weights[1], body_refs.clone()), + ], + terms, + operator == Operator::And, + ); + // The oracle has to agree with the predicates the corpus was built from. + assert_eq!( + brute_force_ids(&expected).len(), + want_matches, + "corpus no longer matches as intended for {terms:?} {operator:?}" + ); + assert!(want_matches > 0, "no matches for {terms:?} {operator:?}"); + // Below, at, and above the match count, to cover the exhausted-cursor path. + for k in [1usize, 10, 100, num_docs] { + let query = combined_query_with_boosts(terms, operator, Some(weights.clone())); + let actual = fts_result_id_scores(&dataset, query, Some(k as i64)).await; + assert_topk_matches_brute_force( + &actual, + &expected, + k, + &format!("terms={terms:?} op={operator:?}"), + ); + } + } +} + +#[tokio::test] +async fn test_fts_combined_fields_tokenizer_validation() { + // combined_fields accepts columns that differ only in storage-only params + // (e.g. with_position) but rejects columns configured with different + // tokenizers. + let with_pos = InvertedIndexParams::new("simple".to_string(), Language::English) + .with_position(true) + .stem(false) + .remove_stop_words(false); + let without_pos = InvertedIndexParams::new("simple".to_string(), Language::English) + .with_position(false) + .stem(false) + .remove_stop_words(false); + let whitespace = InvertedIndexParams::new("whitespace".to_string(), Language::English) + .stem(false) + .remove_stop_words(false); + + let batch = record_batch!( + ("id", Int32, [0, 1]), + ("title", Utf8, ["aa", "bb"]), + ("body", Utf8, ["bb", "aa"]), + ("alt", Utf8, ["aa", "aa"]) + ) + .unwrap(); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, None).await; + // One index per column, each with its own tokenizer configuration, which is + // what the query-time validation reads back. + create_inverted_indices(&mut dataset, &["title"], &with_pos).await; + create_inverted_indices(&mut dataset, &["body"], &without_pos).await; + create_inverted_indices(&mut dataset, &["alt"], &whitespace).await; + + // title (with positions) and body (without) tokenize identically: accepted. + let accepted = combined_query("aa", Operator::Or); + assert_eq!( + unique_ids(fts_result_ids(&dataset, accepted).await), + HashSet::from([0, 1]) + ); + + // title (simple) and alt (whitespace) use different tokenizers: rejected. + let rejected = combined_query_over(&["title", "alt"], "aa", Operator::Or, None); + let result = dataset + .scan() + .full_text_search(FullTextSearchQuery::new_query(rejected)) + .unwrap() + .try_into_batch() + .await; + let message = result + .expect_err("expected a tokenizer-mismatch error") + .to_string(); + assert!( + message.contains("combined_fields") && message.contains("tokenizer"), + "unexpected error: {message}" + ); +} + +#[tokio::test] +async fn test_fts_combined_fields_concatenation_identity() { + // With integer weights, BM25F over (title^w_t, body^w_b) is identical to plain + // BM25 over a single column that concatenates title repeated w_t times with body + // repeated w_b times, as long as every title token also appears in body so the + // max-based docFreq' blend matches the concatenated union. This cross-checks the + // combined scan against Lance's own single-field BM25, an independent code path, + // catching spec-interpretation errors. + let params = combined_fields_test_params(); + let titles = ["cat", "dog", "bird", "cat dog"]; + let bodies = ["cat dog", "dog bird cat", "bird cat", "cat dog bird"]; + let (w_title, w_body) = (2usize, 1usize); + let concat: Vec = titles + .iter() + .zip(&bodies) + .map(|(title, body)| { + let mut parts: Vec<&str> = Vec::with_capacity(w_title + w_body); + for _ in 0..w_title { + parts.push(title); + } + for _ in 0..w_body { + parts.push(body); + } + parts.join(" ") + }) + .collect(); + + let concat_refs: Vec<&str> = concat.iter().map(|s| s.as_str()).collect(); + let batch = record_batch!( + ("id", Int32, (0..titles.len() as i32).collect::>()), + ("title", Utf8, titles.to_vec()), + ("body", Utf8, bodies.to_vec()), + ("concat", Utf8, concat_refs) + ) + .unwrap(); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, None).await; + create_inverted_indices(&mut dataset, &["title", "body", "concat"], ¶ms).await; + + let combined = combined_query_with_boosts( + "cat dog", + Operator::Or, + Some(vec![w_title as f32, w_body as f32]), + ); + let plain = FtsQuery::Match( + MatchQuery::new("cat dog".to_string()).with_column(Some("concat".to_string())), + ); + + let combined_scores: HashMap = fts_result_id_scores(&dataset, combined, None) + .await + .into_iter() + .collect(); + let plain_scores: HashMap = fts_result_id_scores(&dataset, plain, None) + .await + .into_iter() + .collect(); + + assert_eq!( + combined_scores.keys().copied().collect::>(), + plain_scores.keys().copied().collect::>(), + ); + for (id, combined_score) in &combined_scores { + let plain_score = plain_scores[id]; + assert!( + (combined_score - plain_score).abs() < 1e-3, + "id {id}: combined={combined_score}, concatenated single-field={plain_score}" + ); + } +} + +#[tokio::test] +async fn test_fts_combined_fields_fast_search_without_full_coverage_is_empty() { + // `fast_search` is index-only. When a target column has no index at all no + // fragment is fully covered, so the answer is definitionally empty, and the + // indexed exec must not be built, because it requires every target column to + // have segments. Mirrors `plan_match_query`, which returns `EmptyExec` here. + let params = combined_fields_test_params(); + let batch = combined_fields_batch(vec![0, 1], vec!["aa", "bb"], vec!["bb", "aa"]); + let test_uri = TempStrDir::default(); + let mut dataset = write_fts_dataset(&test_uri, batch, None).await; + create_inverted_indices(&mut dataset, &["title"], ¶ms).await; + + let batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(combined_query( + "aa", + Operator::Or, + ))) + .unwrap() + .fast_search() + .try_into_batch() + .await + .expect("fast_search without full coverage should be empty, not an error"); + assert_eq!(batch.num_rows(), 0); +} + +#[tokio::test] +async fn test_fts_combined_fields_fast_search_skips_uncovered_fragments() { + // `fast_search` is index-only by contract, so appended rows stay invisible. It + // must still exclude a fragment that not every target column indexes, otherwise + // those rows come back with a partial `tf'`/`dl'`. + let params = combined_fields_test_params(); + let (test_uri, _indexed) = indexed_two_column_dataset(&["aa"], &["aa"], None).await; + + let second = combined_fields_batch(vec![1], vec!["aa"], vec!["aa"]); + let mut dataset = append_fts_dataset(&test_uri, second, None).await; + // `title` now covers both fragments, `body` only the first. + create_inverted_indices(&mut dataset, &["title"], ¶ms).await; + + let batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(combined_query( + "aa", + Operator::Or, + ))) + .unwrap() + .fast_search() + .try_into_batch() + .await + .unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec(); + assert_eq!( + ids, + vec![0], + "fast_search returned a row whose cross-field data is only partly indexed" + ); +} + +#[tokio::test] +async fn test_fts_combined_fields_empty_fragment_list_is_empty() { + // An explicitly empty fragment list selects no rows. The fully indexed plan + // carries no fragment restriction, so without an explicit short circuit the + // indexed scan answers from every fragment the index holds. + let (_test_uri, dataset) = indexed_two_column_dataset(&["aa", "bb"], &["bb", "aa"], None).await; + + let mut scan = dataset.scan(); + scan.project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(combined_query( + "aa", + Operator::Or, + ))) + .unwrap() + .with_fragments(vec![]); + let plan = scan.explain_plan(false).await.unwrap(); + assert!(plan.contains("EmptyExec"), "unexpected plan: {plan}"); + assert_eq!(scan.try_into_batch().await.unwrap().num_rows(), 0); +} + +#[tokio::test] +async fn test_fts_combined_fields_requires_an_index() { + // BM25F reads its shared tokenizer configuration off an index, so a query + // whose target columns are all unindexed is rejected rather than scored with a + // default tokenizer that may not match how the data would be indexed. + let batch = combined_fields_batch(vec![0, 1], vec!["aa", "bb"], vec!["bb", "aa"]); + let test_uri = TempStrDir::default(); + let dataset = write_fts_dataset(&test_uri, batch, None).await; + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(combined_query( + "aa", + Operator::Or, + ))) + .unwrap() + .try_into_batch() + .await; + let message = result.expect_err("expected an error").to_string(); + assert!( + message.contains("combined_fields") && message.contains("inverted index"), + "unexpected error: {message}" + ); +} diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 8fc0e47cc79..1f925c1df45 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -62,7 +62,7 @@ use datafusion::common::{assert_contains, assert_not_contains}; use futures::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_arrow::json::ARROW_JSON_EXT_NAME; -use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; +use lance_index::scalar::inverted::query::{CombinedFieldsQuery, FtsQuery, MultiMatchQuery}; use lance_table::format::BasePath; use lance_testing::datagen::generate_random_array; use rand::Rng; @@ -1285,6 +1285,56 @@ async fn create_fragmented_fts_index_with_groups( assert_eq!(segments.len(), expected_segments); } +/// Three fragments over two indexed text columns, each column indexed as one +/// segment per fragment. +/// +/// Rows 3 and 4 carry the same text in both columns, so any query that scores +/// them symmetrically produces an exact score tie and exposes tie ordering. +async fn compound_fts_dataset() -> Dataset { + let batch = arrow_array::record_batch!( + ( + "title", + Utf8, + [ + "common", + "common filler filler filler filler filler filler filler", + "irrelevant", + "common tie", + "common tie", + "irrelevant" + ] + ), + ( + "body", + Utf8, + [ + "penalty", + "special", + "common", + "neutral", + "neutral", + "common filler filler filler penalty" + ] + ) + ) + .unwrap(); + let schema = batch.schema(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + "memory://", + Some(WriteParams { + max_rows_per_file: 2, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + create_fragmented_fts_index(&mut dataset, "title", false).await; + create_fragmented_fts_index(&mut dataset, "body", false).await; + dataset +} + fn compound_multimatch_query() -> FtsQuery { MultiMatchQuery::try_new( "common".to_owned(), @@ -1296,6 +1346,17 @@ fn compound_multimatch_query() -> FtsQuery { .into() } +fn compound_combined_fields_query() -> FtsQuery { + CombinedFieldsQuery::try_new( + "common".to_owned(), + vec!["title".to_owned(), "body".to_owned()], + ) + .unwrap() + .try_with_boosts(vec![10.0, 1.0]) + .unwrap() + .into() +} + fn compound_match_query(term: &str, column: &str, boost: f32) -> FtsQuery { MatchQuery::new(term.to_owned()) .with_column(Some(column.to_owned())) @@ -3213,47 +3274,7 @@ async fn test_boolean_must_scores_sum_across_execution_paths() { #[tokio::test] async fn test_nested_multimatch_limit_propagation() { - let batch = arrow_array::record_batch!( - ( - "title", - Utf8, - [ - "common", - "common filler filler filler filler filler filler filler", - "irrelevant", - "common tie", - "common tie", - "irrelevant" - ] - ), - ( - "body", - Utf8, - [ - "penalty", - "special", - "common", - "neutral", - "neutral", - "common filler filler filler penalty" - ] - ) - ) - .unwrap(); - let schema = batch.schema(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), - "memory://", - Some(WriteParams { - max_rows_per_file: 2, - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!(dataset.get_fragments().len(), 3); - create_fragmented_fts_index(&mut dataset, "title", false).await; - create_fragmented_fts_index(&mut dataset, "body", false).await; + let dataset = compound_fts_dataset().await; let must_query: FtsQuery = BooleanQuery::new([ (Occur::Must, compound_multimatch_query()), @@ -3335,6 +3356,54 @@ async fn test_nested_multimatch_limit_propagation() { .await; } +/// `combined_fields` is planned recursively like every other FTS node, so it owes +/// a compound parent the same contract a nested `MultiMatch` does (see +/// `test_nested_multimatch_limit_propagation`): `FtsSearchParams::limit` decides +/// completeness, and the ambient scanner limit must not reach the child. A child +/// that applied the scanner limit itself would drop candidates before the parent +/// finished composing scores. +#[tokio::test] +async fn test_nested_combined_fields_limit_propagation() { + let dataset = compound_fts_dataset().await; + + let must_query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_combined_fields_query()), + ( + Occur::Should, + compound_match_query("special", "body", 100.0), + ), + ]) + .into(); + let must_results = compound_fts_results(&dataset, must_query.clone(), None).await; + assert!( + must_results + .windows(2) + .any(|rows| rows[0].1 == rows[1].1 && rows[0].0 < rows[1].0), + "the exhaustive result should include a deterministic score tie" + ); + assert_compound_fts_top_k(&dataset, must_query, 2).await; + + let should_query: FtsQuery = BooleanQuery::new([ + (Occur::Should, compound_combined_fields_query()), + ( + Occur::Should, + compound_match_query("special", "body", 100.0), + ), + ]) + .into(); + assert_compound_fts_top_k(&dataset, should_query, 2).await; + + let boost_query: FtsQuery = BoostQuery::new( + compound_combined_fields_query(), + compound_match_query("penalty", "body", 100.0), + Some(1.0), + ) + .into(); + assert_compound_fts_top_k(&dataset, boost_query, 2).await; + + assert_compound_fts_top_k(&dataset, compound_combined_fields_query(), 1).await; +} + #[tokio::test] async fn test_same_column_compound_scorer_is_exact_and_bounded() { let batch = arrow_array::record_batch!(( @@ -3608,7 +3677,7 @@ async fn test_compound_tie_uses_resolved_row_id() { assert_eq!(exhaustive.len(), 384); } -fn nested_fts_batch( +pub(super) fn nested_fts_batch( ids: Vec, a_values: Vec>, b_values: Vec>, diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs index 1204f352966..b26ae565eee 100644 --- a/rust/lance/src/dataset/tests/mod.rs +++ b/rust/lance/src/dataset/tests/mod.rs @@ -6,6 +6,7 @@ mod data_file_part; mod dataset_aggregate; mod dataset_common; mod dataset_concurrency_store; +mod dataset_fts_combined_fields; #[cfg(feature = "geo")] mod dataset_geo; mod dataset_index; diff --git a/rust/lance/tests/query/inverted.rs b/rust/lance/tests/query/inverted.rs index b1db7fd6f2f..e9d10824cc5 100644 --- a/rust/lance/tests/query/inverted.rs +++ b/rust/lance/tests/query/inverted.rs @@ -22,8 +22,8 @@ use lance_index::metrics::NoOpMetricsCollector; use lance_index::optimize::OptimizeOptions; use lance_index::prefilter::NoFilter; use lance_index::scalar::inverted::query::{ - BooleanQuery, BoostQuery, FtsQuery, FtsSearchParams, MatchQuery, MultiMatchQuery, Occur, - Operator, PhraseQuery, collect_query_tokens, + BooleanQuery, BoostQuery, CombinedFieldsQuery, FtsQuery, FtsSearchParams, MatchQuery, + MultiMatchQuery, Occur, Operator, PhraseQuery, collect_query_tokens, }; use lance_index::scalar::inverted::{DocumentGranularity, Language}; use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams}; @@ -85,6 +85,16 @@ fn row_match(column: &str, terms: &str) -> FullTextSearchQuery { FullTextSearchQuery::new_query(FtsQuery::Match(row_match_node(column, terms))) } +fn combined_fields(terms: &str, columns: &[&str]) -> FullTextSearchQuery { + FullTextSearchQuery::new_query(FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + terms.to_string(), + columns.iter().map(|column| column.to_string()).collect(), + ) + .unwrap(), + )) +} + fn row_phrase(column: &str, terms: &str) -> FullTextSearchQuery { FullTextSearchQuery::new_query(FtsQuery::Phrase( PhraseQuery::new(terms.to_string()) @@ -479,6 +489,20 @@ async fn test_element_document_fts_flat_indexed_and_mixed() { "{err}" ); + // BM25F sums each target column's contribution for one row and reports no + // element coordinates, so a column that only has a list-element index has to + // be rejected rather than silently collapsed to a row score. + let mut element_only_combined = ds.scan(); + element_only_combined + .full_text_search(combined_fields("alpha", &["tags"])) + .unwrap(); + let err = element_only_combined.try_into_batch().await.unwrap_err(); + assert!( + err.to_string().contains("combined_fields") + && err.to_string().contains("Row document granularity"), + "{err}" + ); + ds.create_index( &["tags"], IndexType::Inverted, @@ -488,6 +512,17 @@ async fn test_element_document_fts_flat_indexed_and_mixed() { ) .await .unwrap(); + + // With a row index alongside the list-element one, combined_fields picks the + // row index and scores whole rows, without a `_doc_index` column. + let combined = run_fts(&ds, combined_fields("alpha", &["tags"]), None).await; + assert_eq!( + combined["id"] + .as_primitive::() + .values(), + &[0] + ); + assert!(combined.column_by_name("_doc_index").is_none()); let names = ds .load_indices() .await @@ -969,6 +1004,40 @@ async fn test_element_document_nested_lists_use_deepest_boundary() { assert_eq!(row.fields, elements.fields); assert_eq!(row.fields, vec![content.id]); assert_ne!(row.fields, vec![docs.children[0].id]); + + // A cross-field scan resolves the same path through the row index, and its + // flat sibling walks two list levels down to the leaf. Rows 2 and 3 repeat + // rows 0 and 1 in a fragment no index covers, so both sides contribute. + let combined_indexed = run_fts(&ds, combined_fields("alpha", &[path]), None).await; + assert_eq!( + combined_indexed["id"] + .as_primitive::() + .values(), + &[0, 1] + ); + assert!(combined_indexed.column_by_name("_doc_index").is_none()); + + let appended = RecordBatch::try_from_iter(vec![ + ("id", Arc::new(Int32Array::from(vec![2, 3])) as ArrayRef), + ("groups", batch.column_by_name("groups").unwrap().clone()), + ]) + .unwrap(); + let ds = InsertBuilder::new(Arc::new(ds)) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute(vec![appended]) + .await + .unwrap(); + + let combined_mixed = run_fts(&ds, combined_fields("alpha", &[path]), None).await; + assert_eq!( + combined_mixed["id"] + .as_primitive::() + .values(), + &[0, 1, 2, 3] + ); } #[tokio::test] From 6ef87509f3a8a4ca6c158c68e28cc49a53cb009b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Mon, 17 Aug 2026 12:07:44 +0200 Subject: [PATCH 3/4] feat(fts): score unindexed and partially indexed fragments in combined_fields The previous commit refuses a `combined_fields` query whose target columns do not all cover every scanned fragment, so a default full-text search fails on any dataset with rows appended since the indexes were built. `MatchQuery` already unions in a flat scan for its unindexed fragments; this does the same for BM25F. Coverage is per column here, which makes it more than a copy of the single-column path. `dl'` sums each column's document length and a row absent from a column's `DocSet` contributes 0, so a fragment indexed for `title` but not `body` cannot be scored from the index at all. The indexed scan is therefore restricted to the intersection of per-column coverage and everything else goes to the flat scan, rather than splitting on the union. Both sides then score against one shared corpus. The flat side alone sees the rows no index covers, so it measures their contribution and publishes the blend; the indexed side waits for it instead of folding only its own `docCount'`/`docFreq'`/`avgdl'`. Without that, a row reached through either path would rank differently depending on which side happened to score it. Data overlays are handled by measuring rather than patching. When a target column carries an overlay-stale index entry, folding the flat row into the index statistics would double count it against the entry it replaces, and the flat scan cannot subtract what it replaced. So the corpus is measured from current data instead: every target fragment is scanned, every row folded into every column, and the index statistics left out. That costs a full scan of the target columns, so it stays confined to the stale case. `fast_search` is unchanged, being index-only by contract. --- rust/lance-index/src/scalar/inverted.rs | 4 +- .../src/scalar/inverted/combined.rs | 4 +- .../src/scalar/inverted/combined/flat.rs | 246 +++++++ .../src/scalar/inverted/combined/search.rs | 3 +- .../src/scalar/inverted/combined/stats.rs | 278 +++++++- .../src/scalar/inverted/combined/testing.rs | 64 +- rust/lance-index/src/scalar/inverted/index.rs | 4 +- .../src/scalar/inverted/index/flat_search.rs | 235 ++++++- rust/lance/src/dataset/scanner.rs | 268 +++++++- rust/lance/src/index/scalar/inverted.rs | 65 ++ rust/lance/src/io/exec/fts.rs | 625 ++++++++++++++++-- 11 files changed, 1704 insertions(+), 92 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/combined/flat.rs diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index f14fbf41e8b..4b0b7e657bb 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -29,8 +29,8 @@ use arrow_schema::{DataType, Field}; use async_trait::async_trait; pub use builder::InvertedIndexBuilder; pub use combined::{ - CombinedFieldColumn, build_combined_bm25_scorer, combined_fields_search, - validate_combined_tokenizers, + CombinedCorpusStats, CombinedFieldColumn, FlatFieldStats, build_combined_bm25_scorer, + combined_fields_search, flat_combined_fields_search_stream, validate_combined_tokenizers, }; pub use compound::{ compound_search, compound_search_prepared_match, diff --git a/rust/lance-index/src/scalar/inverted/combined.rs b/rust/lance-index/src/scalar/inverted/combined.rs index 9de77185ce7..9c1810b1f90 100644 --- a/rust/lance-index/src/scalar/inverted/combined.rs +++ b/rust/lance-index/src/scalar/inverted/combined.rs @@ -23,6 +23,7 @@ //! scored; see [`combined_fields_search`]. mod cursor; +mod flat; mod search; mod stats; #[cfg(test)] @@ -32,8 +33,9 @@ use std::sync::Arc; use lance_core::{Error, Result}; +pub use flat::flat_combined_fields_search_stream; pub use search::combined_fields_search; -pub use stats::build_combined_bm25_scorer; +pub use stats::{CombinedCorpusStats, FlatFieldStats, build_combined_bm25_scorer}; use super::index::InvertedIndex; use super::query::Tokens; diff --git a/rust/lance-index/src/scalar/inverted/combined/flat.rs b/rust/lance-index/src/scalar/inverted/combined/flat.rs new file mode 100644 index 00000000000..9cb2c13c250 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/flat.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The unindexed `combined_fields` plan: blend the scanned column values into +//! `dl'`/`tf'` and score the rows no target column's index covers. + +use std::sync::Arc; + +use arrow::array::{Float32Builder, UInt64Builder}; +use arrow_array::{ArrayRef, RecordBatch}; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::metrics::Time; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{FutureExt, stream}; +use lance_core::Error; +use lance_core::error::DataFusionResult; +use lance_core::utils::tokio::spawn_cpu; +use lance_select::RowAddrMask; + +use super::super::index::{BlendedRows, FTS_SCHEMA, slice_into_batches, tokenize_and_blend_multi}; +use super::super::query::{Operator, Tokens}; +use super::super::scorer::CombinedFieldsBM25Scorer; +use super::super::tokenizer::document_tokenizer::LanceTokenizer; +use super::stats::{CombinedCorpusStats, FlatFieldStats, build_combined_bm25_scorer}; +use super::{CombinedFieldColumn, unique_terms}; +use crate::metrics::MetricsCollector; + +/// Exact cross-field BM25F search over rows that no index covers. +/// +/// The indexed [`combined_fields_search`](super::combined_fields_search) can +/// only score a row whose fragment +/// every target column's index covers; otherwise `dl'` and `tf'` would be missing +/// a column's contribution. This scores the remaining rows straight from their +/// column values, so the two plans together cover the whole dataset. +/// +/// `input` must carry `_rowid` plus every target column, in `columns` order. +/// Query terms are deduplicated exactly as +/// [`combined_fields_search`](super::combined_fields_search) does, and +/// `operator` applies across the virtual field: `And` requires every term to +/// appear in at least one target column. +/// +/// The blended scorer folds these rows' own per-column statistics in through +/// [`FlatFieldStats`](super::FlatFieldStats), so `avgdl'` and `idf'` reflect the +/// scanned rows rather than only the indexed ones. `stats_masks` keeps that fold +/// from double counting rows a column's index already holds. `metrics` covers +/// only that scorer build's index reads; the plan already accounts for the input +/// stream's IO. +/// +/// `emit_mask`, when set, selects the rows to emit. The caller reads `input` +/// unfiltered in that case, so the fold sees rows the query filtered out and a folded +/// row's contribution does not move with the filter. +/// +/// `flat_covers_whole_corpus` drops the index statistics in favour of this scan's own; +/// see [`CombinedCorpusStats::FlatOnly`]. +/// +/// Returns the scorer alongside the stream. The whole input is consumed before +/// the first output batch, so by then the blend describes the entire scanned +/// corpus and an indexed sibling can score against the same statistics. +#[allow(clippy::too_many_arguments)] +pub async fn flat_combined_fields_search_stream( + input: SendableRecordBatchStream, + columns: &[CombinedFieldColumn], + doc_col_indices: Vec, + stats_masks: &[Arc], + emit_mask: Option>, + flat_covers_whole_corpus: bool, + tokens: &Tokens, + tokenizer: Box, + operator: Operator, + target_batch_size: usize, + elapsed_compute: Option