diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index 7f09c4325b7..9c427db2979 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -363,6 +363,75 @@ query_result = ds.to_table( ) ``` +### Searching Multiple Columns + +When a query should span several text columns (for example `title` and `body`), +Lance offers two strategies that differ in how they combine evidence across +fields. + +#### `best_fields` with `MultiMatchQuery` + +`MultiMatchQuery` runs the query against each column independently and keeps the +**best** field's score for each document. Each +column is scored against its *own* corpus statistics, and an optional per-column +`boosts` multiplier is applied before the max. + +```python +from lance.query import MultiMatchQuery + +query_result = ds.to_table( + full_text_query=MultiMatchQuery( + "albino elephant", + ["title", "body"], + boosts=[2.0, 1.0], # weight title matches twice as much + ) +) +``` + +This works well when the fields are interchangeable and a document is relevant +if *any single field* matches strongly. + +#### `combined_fields` (BM25F) with `CombinedFieldsQuery` + +`CombinedFieldsQuery` instead treats the target columns as **one virtual field** +and blends their statistics into a single BM25F model (the same scoring as +Lucene's `CombinedFieldQuery`). This makes the scores comparable across +fields and lets a single query term match *across* fields. + +```python +from lance.query import CombinedFieldsQuery, FullTextOperator + +query_result = ds.to_table( + full_text_query=CombinedFieldsQuery( + "john smith", + ["first_name", "last_name"], + boosts=[1.0, 1.0], + operator=FullTextOperator.AND, + ) +) +``` + +Prefer `combined_fields` when: + +- a term is **rare in one field but common in another**: blended document + frequencies keep a spurious hit in a short field from outranking a genuinely + better match, which `best_fields` cannot do; or +- a term should match **across** fields: with `operator=AND`, `"john smith"` + matches a row whose `first_name` is `"john"` and `last_name` is `"smith"`, + even though neither field alone contains both terms (`best_fields` would + require a single field to contain the whole query). + +Per-column `boosts` weight how much each field contributes to the blended term +frequency. Each weight must be `>= 1` (fractional weights are allowed); a title +boost of `2` counts every title term twice. + +!!! note "Shared tokenizer required" + + All columns in a `CombinedFieldsQuery` must share the same tokenizer / index + configuration, because BM25F is only well-defined when the fields tokenize + the same way. Mixing configurations raises an error listing the offending + columns. + ## Performance Tips ### Index Maintenance diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index b9cdc221b82..e146e60a20d 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -882,6 +882,13 @@ class PyFullTextQuery: boosts: Optional[List[float]] = None, operator: str = "OR", ) -> PyFullTextQuery: ... + @staticmethod + def combined_fields_query( + query: str, + columns: List[str], + boosts: Optional[List[float]] = None, + operator: str = "OR", + ) -> PyFullTextQuery: ... class ScanStatistics: """Statistics about a scan operation.""" diff --git a/python/python/lance/query.py b/python/python/lance/query.py index 85bbf121e6e..161af34bd59 100644 --- a/python/python/lance/query.py +++ b/python/python/lance/query.py @@ -14,6 +14,7 @@ class FullTextQueryType(Enum): MATCH_PHRASE = "match_phrase" BOOST = "boost" MULTI_MATCH = "multi_match" + COMBINED_FIELDS = "combined_fields" BOOLEAN = "boolean" @@ -227,6 +228,54 @@ def query_type(self) -> FullTextQueryType: return FullTextQueryType.MULTI_MATCH +class CombinedFieldsQuery(FullTextQuery): + def __init__( + self, + query: str, + columns: list[str], + *, + boosts: Optional[list[float]] = None, + operator: FullTextOperator = FullTextOperator.OR, + ): + """ + Combined-fields (BM25F) query for full-text search. + + Unlike :class:`MultiMatchQuery`, which scores each column independently + and keeps the best field (Elasticsearch ``best_fields``), this query + treats the target columns as a single virtual field so that term + statistics are blended across fields (Elasticsearch ``combined_fields`` + / Lucene ``CombinedFieldQuery``). This makes a term that is rare in one + field but common in another score consistently, and lets a single query + term match across fields (e.g. a first name in one column and a last + name in another). + + Parameters + ---------- + query : str + The query string to match against. + columns : list[str] + The list of columns combined into the virtual field. + boosts : list[float], optional + Per-column weights aligned with ``columns``. Each weight must be + ``>= 1`` (fractional weights allowed). If not provided, every column + defaults to ``1.0``. + operator : FullTextOperator, default OR + The operator applied across the virtual field. If ``AND``, every + term must appear in at least one column; if ``OR``, at least one + term must match. + + Notes + ----- + All target columns must share the same tokenizer/index configuration. + """ + self._inner = PyFullTextQuery.combined_fields_query( + query, columns, boosts=boosts, operator=operator.value + ) + + def query_type(self) -> FullTextQueryType: + return FullTextQueryType.COMBINED_FIELDS + + class BooleanQuery(FullTextQuery): def __init__(self, queries: list[tuple[Occur, FullTextQuery]]): """ diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index b45631ccd4f..7db58bcaf2b 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -23,6 +23,7 @@ from lance.query import ( BooleanQuery, BoostQuery, + CombinedFieldsQuery, FullTextOperator, MatchQuery, MultiMatchQuery, @@ -1873,6 +1874,65 @@ def test_fts_multi_match_query(tmp_path): ) +@pytest.mark.parametrize( + "operator,combined_ids,multi_ids", + [ + # combined_fields treats the columns as one virtual field, so AND matches + # when each term appears in any field (rows 0 and 1); best_fields AND only + # matches row 1, where a single field holds both terms. + (FullTextOperator.AND, {0, 1}, {1}), + (FullTextOperator.OR, {0, 1, 2}, {0, 1, 2}), + ], +) +def test_fts_combined_fields_query(tmp_path, operator, combined_ids, multi_ids): + data = pa.table( + { + "id": [0, 1, 2], + # row 0 splits the terms across fields, row 1 has both in one field, + # row 2 has only "john". + "title": ["john", "john smith", "john"], + "body": ["smith", "foo", "alice"], + } + ) + ds = lance.write_dataset(data, tmp_path) + ds.create_scalar_index("title", "INVERTED") + ds.create_scalar_index("body", "INVERTED") + + def ids(query): + return set(ds.to_table(full_text_query=query, columns=["id"])["id"].to_pylist()) + + assert ( + ids(CombinedFieldsQuery("john smith", ["title", "body"], operator=operator)) + == combined_ids + ) + assert ( + ids(MultiMatchQuery("john smith", ["title", "body"], operator=operator)) + == multi_ids + ) + + +def test_fts_combined_fields_boost_validation(tmp_path): + # Per-column boosts must be >= 1 (Lucene CombinedFieldQuery constraint), and + # the boost count must match the column count. + data = pa.table({"title": ["hello"], "body": ["world"]}) + ds = lance.write_dataset(data, tmp_path) + ds.create_scalar_index("title", "INVERTED") + ds.create_scalar_index("body", "INVERTED") + + with pytest.raises(ValueError, match="boost"): + CombinedFieldsQuery("hello", ["title", "body"], boosts=[0.5, 1.0]) + with pytest.raises(ValueError): + CombinedFieldsQuery("hello", ["title", "body"], boosts=[1.0]) + + # Fractional weights >= 1 are accepted. + result = ds.to_table( + full_text_query=CombinedFieldsQuery( + "hello", ["title", "body"], boosts=[1.5, 1.0] + ) + ) + assert result.num_rows == 1 + + def test_fts_boolean_query(tmp_path): data = pa.table( { diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 49c0996b8fc..df7cc0722a5 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -73,7 +73,8 @@ use lance_file::reader::FileReaderOptions; use lance_index::scalar::inverted::InvertedListFormatVersion; use lance_index::scalar::inverted::query::Occur; use lance_index::scalar::inverted::query::{ - BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Operator, PhraseQuery, + BooleanQuery, BoostQuery, CombinedFieldsQuery, FtsQuery, MatchQuery, MultiMatchQuery, Operator, + PhraseQuery, }; use lance_index::{ FtsPrewarmOptions, IndexParams, IndexType, PrewarmOptions, @@ -5343,6 +5344,31 @@ impl PyFullTextQuery { }) } + #[staticmethod] + #[pyo3(signature = (query, columns, boosts=None, operator="OR"))] + fn combined_fields_query( + query: String, + columns: Vec, + boosts: Option>, + operator: &str, + ) -> PyResult { + let q = CombinedFieldsQuery::try_new(query, columns) + .map_err(|e| PyValueError::new_err(format!("Invalid query: {}", e)))?; + let q = if let Some(boosts) = boosts { + q.try_with_boosts(boosts) + .map_err(|e| PyValueError::new_err(format!("Invalid boosts: {}", e)))? + } else { + q + }; + + let op = Operator::try_from(operator) + .map_err(|e| PyValueError::new_err(format!("Invalid operator: {}", e)))?; + + Ok(Self { + inner: q.with_operator(op).into(), + }) + } + #[staticmethod] #[pyo3(signature = (queries))] fn boolean_query(queries: Vec<(String, Self)>) -> PyResult { diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 72604243c1c..23096b63770 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 encoding; mod impact; mod index; @@ -21,11 +22,15 @@ use std::sync::Arc; 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, +}; 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}; diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index c90f96053ef..19e15a86359 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -132,6 +132,13 @@ fn merge_all_tail_partitions( if let Some(builder) = merged { merged_builders.push(builder); } + // Each merged partition concatenated several workers' ascending runs, so + // restore a global row_id order (a no-op for partitions that stayed a + // single ascending run). This runs inside `spawn_cpu`, off the async + // runtime, alongside the rest of the merge. + for builder in &mut merged_builders { + builder.sort_docs_by_row_id(); + } Ok(merged_builders) } @@ -346,6 +353,16 @@ impl InvertedIndexBuilder { ) -> Result> { let partition_id = self.next_partition_id() | self.fragment_mask.unwrap_or(0); builder.set_id(partition_id); + // A partition merged from several existing segments is a concatenation + // of their doc runs; restore a global row_id order so read pruning keeps + // working after updates (a no-op when it is already ascending). Offload + // the sort to the CPU pool so a large partition does not block the async + // runtime thread. + let mut builder = spawn_cpu(move || { + builder.sort_docs_by_row_id(); + Result::Ok(builder) + }) + .await?; let files = builder .write_to(dest_store, self.partition_write_target()) .await?; @@ -1036,6 +1053,120 @@ impl InnerBuilder { Ok(()) } + /// Reorder this partition's documents so their `row_id`s are strictly + /// ascending in `doc_id` order, applying one consistent old->new `doc_id` + /// remap to the doc set (row_ids + num_tokens) and to every posting list + /// (doc_ids, frequencies, and positions). + /// + /// `doc_id`s are internal dense indices with no external meaning: a BM25 + /// score depends only on term frequency, document length + /// (`num_tokens[doc_id]`), the average document length, the document count, + /// and the posting length, none of which change when documents are + /// relabeled with a consistent bijection. So this reorder never changes any + /// query's scores or results; it only makes each partition's `row_id`s + /// monotonic, which lets combined-field read pruning skip posting blocks by + /// `row_id` range. + /// + /// The parallel builder assigns `doc_id`s per worker, so a partition merged + /// from several workers' tails (see [`merge_from`](Self::merge_from)) is a + /// concatenation of individually ascending runs and is not globally ordered; + /// this restores the global order. It is a cheap no-op when the documents are + /// already ascending, which is the common single-worker / single-tail case. + fn sort_docs_by_row_id(&mut self) { + let num_docs = self.docs.len(); + if num_docs <= 1 { + return; + } + + // Fast path: already strictly ascending, so there is nothing to remap. + // This is the common single-worker / single-tail case, and it avoids + // the O(n log n) sort below entirely. + let mut previous_row_id: Option = None; + let already_ascending = self.docs.iter().all(|(row_id, _)| { + let ascending = previous_row_id < Some(*row_id); + previous_row_id = Some(*row_id); + ascending + }); + if already_ascending { + return; + } + + // new_to_old[new_doc_id] = old_doc_id, stable-sorted by row_id. The + // stable order keeps the result deterministic even when two documents + // share a row_id (only possible for legacy list indexes re-merged via + // `merge_existing_segments`; freshly built partitions hold exactly one + // document per row, so their row_ids are distinct and strictly + // ascending after the sort). + let mut new_to_old: Vec = (0..num_docs as u32).collect(); + new_to_old.sort_by_key(|&old_doc_id| self.docs.row_id(old_doc_id)); + + // The docs were not strictly ascending but may already be in a valid + // (stable) order, e.g. ties from a legacy list index: skip the rebuild + // when the permutation is the identity. + if new_to_old + .iter() + .enumerate() + .all(|(new_doc_id, &old_doc_id)| new_doc_id as u32 == old_doc_id) + { + return; + } + + // old_to_new[old_doc_id] = new_doc_id, the inverse permutation used to + // relabel the posting lists. + let mut old_to_new = vec![0u32; num_docs]; + for (new_doc_id, &old_doc_id) in new_to_old.iter().enumerate() { + old_to_new[old_doc_id as usize] = new_doc_id as u32; + } + + // Rebuild the doc set in the new order. `append` recomputes total_tokens + // as it goes, so the corpus statistics are preserved exactly. + let mut reordered_docs = DocSet::default(); + for &old_doc_id in &new_to_old { + reordered_docs.append( + self.docs.row_id(old_doc_id), + self.docs.num_tokens(old_doc_id), + ); + } + self.docs = reordered_docs; + + // Rebuild every posting list with remapped doc_ids. Posting-list block + // compression delta-encodes doc_ids and requires them strictly + // ascending within the list, so the remapped entries are re-sorted by + // new doc_id before they are re-added. Frequencies and positions travel + // with their entry, exactly as `merge_from` and `PostingListBuilder::remap` + // reconstruct a list. + for posting_list in &mut self.posting_lists { + if posting_list.is_empty() { + continue; + } + let mut entries: Vec<(u32, u32, Option>)> = + Vec::with_capacity(posting_list.len()); + posting_list + .for_each_entry(|old_doc_id, frequency, positions| { + entries.push((old_to_new[old_doc_id as usize], frequency, positions)); + Ok::<(), Error>(()) + }) + .expect("posting list iteration is infallible"); + // doc_ids are a bijection image, hence distinct: no ties to break. + entries.sort_unstable_by_key(|(new_doc_id, _, _)| *new_doc_id); + + let mut reordered_posting = + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + posting_list.has_positions(), + self.posting_tail_codec, + self.block_size, + ); + for (new_doc_id, frequency, positions) in entries { + let positions = match positions { + Some(positions) => PositionRecorder::Position(positions.into()), + None => PositionRecorder::Count(frequency), + }; + reordered_posting.add(new_doc_id, positions); + } + *posting_list = reordered_posting; + } + } + fn memory_size(&self) -> u64 { let posting_lists_overhead = self.posting_lists.capacity() * std::mem::size_of::(); @@ -1645,6 +1776,17 @@ impl IndexWorker { Ok(()) } + /// Flush the worker's in-memory posting lists into a new on-disk partition. + /// + /// A flushed partition inherits the input scan's row_id order. For a normal + /// ascending row_id scan each worker's subsequence is ascending, so flushed + /// partitions are already sorted and the combined_fields read-pruning fast + /// path engages. A non-monotonic input stream can yield an unsorted partition, + /// which only costs that fast path; results and scores are unaffected. + /// + /// Do not add an inline sort here: it would block the async runtime (the issue + /// just fixed in `write_new_partition`). The reorder is handled in the + /// offloaded merge path. #[instrument(level = "debug", skip_all)] async fn flush(&mut self) -> Result<()> { if self.builder.tokens.is_empty() { @@ -4308,6 +4450,265 @@ mod tests { Ok(()) } + // Merging worker tails whose row_ids interleave across the tails must leave + // the merged partition with row_ids strictly ascending in doc-id order, + // while preserving every (token -> row_id -> frequency/positions) + // association exactly (the remap is identity-preserving for scoring). Covers + // the plain (V1/Fixed32) and delta (V2/V3/VarintDelta) tail codecs, both + // 128- and 256-document blocks, and with/without positions. + #[rstest::rstest] + #[case::v1(InvertedListFormatVersion::V1, LEGACY_BLOCK_SIZE)] + #[case::v2(InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE)] + #[case::v3(InvertedListFormatVersion::V3, 256)] + fn test_merge_reorders_docs_by_row_id_ascending( + #[case] format_version: InvertedListFormatVersion, + #[case] block_size: usize, + #[values(false, true)] with_position: bool, + ) { + use std::collections::BTreeMap; + + // Append one token occurrence for a document to the pre-registered + // posting list. `positions` doubles as the frequency (its length) in the + // no-position case. + fn add_entry( + builder: &mut InnerBuilder, + doc_id: u32, + token: &str, + positions: &[u32], + with_position: bool, + ) { + let token_id = builder.tokens.get(token).expect("token pre-registered") as usize; + let recorder = if with_position { + PositionRecorder::Position(positions.iter().copied().collect()) + } else { + PositionRecorder::Count(positions.len() as u32) + }; + builder.posting_lists[token_id].add(doc_id, recorder); + } + + let make_tail = |id: u64, first_token: &str, second_token: &str| { + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + id, + with_position, + TokenSetFormat::default(), + format_version, + block_size, + ); + builder.tokens.add(first_token.to_owned()); + builder.tokens.add(second_token.to_owned()); + builder.posting_lists.resize_with(builder.tokens.len(), || { + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + with_position, + format_version.posting_tail_codec(), + block_size, + ) + }); + builder + }; + + // Each tail is internally ascending by row_id, mirroring a real worker. + // Tail A registers alpha before beta; tail B registers beta before + // alpha, so the two tails carry different local token ids and the merge + // exercises the token-id remap as well as the doc-id remap. + // + // tail A: row 10 -> alpha@[0], beta@[1] (num_tokens 2) + // row 20 -> beta@[0] (num_tokens 1) + // row 30 -> alpha@[0, 1] (num_tokens 2) + // tail B: row 5 -> beta@[0] (num_tokens 1) + // row 15 -> alpha@[0] (num_tokens 1) + // row 25 -> alpha@[0], beta@[1,2] (num_tokens 3) + let mut tail_a = make_tail(0, "alpha", "beta"); + tail_a.docs.append(10, 2); + add_entry(&mut tail_a, 0, "alpha", &[0], with_position); + add_entry(&mut tail_a, 0, "beta", &[1], with_position); + tail_a.docs.append(20, 1); + add_entry(&mut tail_a, 1, "beta", &[0], with_position); + tail_a.docs.append(30, 2); + add_entry(&mut tail_a, 2, "alpha", &[0, 1], with_position); + + let mut tail_b = make_tail(1, "beta", "alpha"); + tail_b.docs.append(5, 1); + add_entry(&mut tail_b, 0, "beta", &[0], with_position); + tail_b.docs.append(15, 1); + add_entry(&mut tail_b, 1, "alpha", &[0], with_position); + tail_b.docs.append(25, 3); + add_entry(&mut tail_b, 2, "alpha", &[0], with_position); + add_entry(&mut tail_b, 2, "beta", &[1, 2], with_position); + + // Concatenating the two tails yields doc-id order row_ids + // [10, 20, 30, 5, 15, 25], which is NOT globally ascending. + let merged = merge_all_tail_partitions( + vec![ + TailPartition { builder: tail_a }, + TailPartition { builder: tail_b }, + ], + u64::MAX, + ) + .unwrap(); + assert_eq!(merged.len(), 1, "both tails fit within the memory budget"); + let merged = &merged[0]; + + // row_ids are strictly ascending in doc-id order. + let row_ids: Vec = merged.docs.iter().map(|(row_id, _)| *row_id).collect(); + assert_eq!(row_ids, vec![5, 10, 15, 20, 25, 30]); + assert!( + row_ids.windows(2).all(|w| w[0] < w[1]), + "row_ids must be strictly ascending, got {row_ids:?}" + ); + + // num_tokens travels with each row (corpus statistics preserved). + let num_tokens_by_row: BTreeMap = + merged.docs.iter().map(|(r, n)| (*r, *n)).collect(); + assert_eq!( + num_tokens_by_row, + BTreeMap::from([(5, 1), (10, 2), (15, 1), (20, 1), (25, 3), (30, 2)]), + ); + assert_eq!(merged.docs.total_tokens_num(), 10); + + // Posting lists: doc_ids strictly ascending, and each row_id keeps its + // exact frequency (and positions when recorded). + let expected_alpha: BTreeMap)> = BTreeMap::from([ + (10, (1, vec![0])), + (15, (1, vec![0])), + (25, (1, vec![0])), + (30, (2, vec![0, 1])), + ]); + let expected_beta: BTreeMap)> = BTreeMap::from([ + (5, (1, vec![0])), + (10, (1, vec![1])), + (20, (1, vec![0])), + (25, (2, vec![1, 2])), + ]); + + for (token, expected) in [("alpha", expected_alpha), ("beta", expected_beta)] { + let token_id = merged.tokens.get(token).expect("token present") as usize; + let mut doc_ids = Vec::new(); + let mut by_row: BTreeMap)> = BTreeMap::new(); + merged.posting_lists[token_id] + .for_each_entry(|doc_id, freq, positions| { + doc_ids.push(doc_id); + let recorded_positions = if with_position { + positions.expect("positions present when with_position") + } else { + assert!(positions.is_none()); + Vec::new() + }; + by_row.insert(merged.docs.row_id(doc_id), (freq, recorded_positions)); + Ok::<(), Error>(()) + }) + .unwrap(); + + assert!( + doc_ids.windows(2).all(|w| w[0] < w[1]), + "{token} posting doc_ids must be strictly ascending, got {doc_ids:?}" + ); + if with_position { + assert_eq!(by_row, expected, "{token} row_id -> (freq, positions)"); + } else { + let actual_freqs: BTreeMap = + by_row.iter().map(|(r, (f, _))| (*r, *f)).collect(); + let expected_freqs: BTreeMap = + expected.iter().map(|(r, (f, _))| (*r, *f)).collect(); + assert_eq!(actual_freqs, expected_freqs, "{token} row_id -> freq"); + } + } + } + + // End-to-end: a real (multi-worker) build over rows whose row_ids arrive in + // shuffled order must produce partitions whose row_ids are strictly + // ascending in doc-id order, with no document dropped or duplicated. + #[rstest::rstest] + #[case::single_worker(1)] + #[case::multi_worker(4)] + #[tokio::test] + async fn test_build_yields_ascending_row_ids_per_partition( + #[case] num_workers: usize, + #[values(false, true)] with_position: bool, + ) -> Result<()> { + use std::collections::BTreeSet; + + let index_dir = TempDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + index_dir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + + // A deterministic shuffle of 0..NUM_DOCS (stride coprime with NUM_DOCS) + // so the worker(s) append documents out of row_id order. + const NUM_DOCS: u64 = 60; + let shuffled: Vec = { + let mut ids = Vec::with_capacity(NUM_DOCS as usize); + let mut next = 0u64; + for _ in 0..NUM_DOCS { + ids.push(next); + next = (next + 23) % NUM_DOCS; + } + ids + }; + assert_eq!( + shuffled.iter().copied().collect::>().len(), + NUM_DOCS as usize + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("doc", DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + // Several batches so more than one worker receives data. + let batches: Vec = shuffled + .chunks(5) + .map(|chunk| { + let docs: Vec = chunk + .iter() + .map(|row_id| format!("common token{row_id}")) + .collect(); + let doc_arr = + StringArray::from(docs.iter().map(|s| Some(s.as_str())).collect::>()); + let row_arr = UInt64Array::from(chunk.to_vec()); + RecordBatch::try_new(schema.clone(), vec![Arc::new(doc_arr), Arc::new(row_arr)]) + .unwrap() + }) + .collect(); + + let stream = RecordBatchStreamAdapter::new( + schema.clone(), + stream::iter(batches.into_iter().map(Ok)), + ); + let params = + InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) + .with_position(with_position) + .remove_stop_words(false) + .stem(false) + .max_token_length(None) + .num_workers(num_workers); + let mut builder = InvertedIndexBuilder::new(params); + builder + .update(Box::pin(stream), store.as_ref(), None) + .await?; + + let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; + let mut seen_row_ids = BTreeSet::new(); + for partition in &index.partitions { + let part_builder = partition.as_ref().clone().into_builder().await?; + let row_ids: Vec = part_builder.docs.iter().map(|(r, _)| *r).collect(); + assert!( + row_ids.windows(2).all(|w| w[0] < w[1]), + "partition {} row_ids must be strictly ascending, got {row_ids:?}", + partition.id(), + ); + for row_id in row_ids { + assert!( + seen_row_ids.insert(row_id), + "row_id {row_id} appeared twice" + ); + } + } + assert_eq!(seen_row_ids, (0..NUM_DOCS).collect::>()); + + Ok(()) + } + #[test] fn test_merge_from_after_remap_does_not_panic() { // `first` is the merge accumulator. Give it three tokens, then remap away the 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..9e140b07b00 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined.rs @@ -0,0 +1,1464 @@ +// 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. A term-at-a-time MAXSCORE (Lucene `CombinedFieldQuery`'s +//! constant per-term ceiling, ordinary MAXSCORE across terms) prunes candidate +//! scoring for a top-k query without changing the top-k it returns; see +//! [`combined_fields_search`]. +//! +//! On the fast path (compressed postings whose partition `row_ids` are strictly +//! ascending) the same MAXSCORE also prunes *reads*: each term's postings are a +//! lazy, block-skipping cross-column cursor, so a non-essential term's posting +//! blocks are never decoded when it is only probed at the sparse candidates the +//! essential terms discover. Anything the fast path cannot prove safe (legacy +//! layout, unsorted `row_ids`, non-compressed postings) falls back to a full +//! read whose result is bit-identical. + +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap}; +use std::sync::Arc; + +use arrow_array::Array; +use lance_core::{Error, Result}; +use lance_select::RowAddrMask; + +use super::builder::ScoredDoc; +use super::encoding::{ + MAX_POSTING_BLOCK_SIZE, decompress_posting_block, decompress_posting_remainder, +}; +use super::index::{CompressedPostingList, DocSet, InvertedIndex, PostingList, PostingTailCodec}; +use super::query::{FtsSearchParams, Operator, Tokens}; +use super::scorer::{CombinedFieldsBM25Scorer, K1}; +use crate::metrics::MetricsCollector; +use crate::prefilter::PreFilter; + +/// 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 silent mis-score. The error lists the offending columns. +pub fn validate_combined_tokenizers(columns: &[CombinedFieldColumn]) -> Result<()> { + let mut reference: Option<&CombinedFieldColumn> = None; + let mut offending = Vec::new(); + for column in columns { + let Some(index) = column.indices.first() else { + continue; + }; + match reference { + None => reference = Some(column), + Some(reference) => { + // 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 compatible = reference + .indices + .first() + .is_some_and(|r| r.params().same_tokenization(index.params())); + if !compatible { + offending.push(column.column.clone()); + } + } + } + } + if !offending.is_empty() { + let reference_column = reference.map(|c| c.column.as_str()).unwrap_or_default(); + 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(()) +} + +/// Fold the target columns' segment statistics into a single [`CombinedFieldsBM25Scorer`]. +/// +/// Generalizes [`super::build_global_bm25_scorer`] (which folds segments of one +/// column) to fold across *columns* 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`] scores. +pub async fn build_combined_bm25_scorer( + columns: &[CombinedFieldColumn], + tokens: &Tokens, +) -> 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 { + let (total_tokens, num_docs, token_docs) = index.bm25_stats_for_terms(&terms).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, + )) +} + +/// The constant per-term BM25F score ceiling `idf · (K1 + 1)` (Lucene +/// `CombinedFieldQuery::getMaxScore`: BM25 saturates the `tf` factor at +/// `K1 + 1`), clamped to 0 when the term cannot raise a score (`idf <= 0`). +#[inline] +fn term_upper_bound(idf: f32) -> f32 { + if idf > 0.0 { idf * (K1 + 1.0) } else { 0.0 } +} + +/// The MAXSCORE loop's view of one query term. Both the eager [full-read +/// fallback](MaterializedTerm) and the lazy [block-skipping fast +/// path](LazyTerm) implement this, so the term-at-a-time MAXSCORE in +/// [`combined_maxscore`] is a single algorithm whose only per-path difference +/// is *how* a `tf'` is fetched — the returned top-k is identical by +/// construction. +/// +/// A term exposes a merged cross-column cursor over the shared row-id space: +/// [`head`](Self::head) is the smallest not-yet-consumed row id, advanced by +/// the essential terms during discovery; [`probe`](Self::probe) reads `tf'` at +/// an arbitrary (non-decreasing) target for the non-essential terms. +trait TermCursor { + /// The constant per-term ceiling `idf · (K1 + 1)` (clamped, see + /// [`term_upper_bound`]). + fn upper_bound(&self) -> f32; + /// The blended IDF `idf'(t)`. + fn idf(&self) -> f32; + /// The smallest not-yet-consumed row id across the term's columns, or + /// `None` when the term is exhausted. + fn head(&self) -> Option; + /// `tf'` at [`head`](Self::head) (the weighted sum across the columns that + /// carry the head row id). Only meaningful while `head()` is `Some`. + fn head_tf(&self) -> f32; + /// Consume `row_id`: advance every column cursor currently sitting on it. + /// Called only for `row_id == head()`. + fn consume(&mut self, row_id: u64); + /// `tf'` at `target`, or 0 when the term is absent there. `target` never + /// decreases across calls to one cursor, so the lazy path can skip blocks. + fn probe(&mut self, target: u64) -> f32; +} + +/// One query term's postings, merged across every target column/partition into +/// the shared row-id space — the full-read fallback. Used whenever the fast +/// path cannot prove block skipping safe (legacy layout, unsorted `row_ids`, +/// non-compressed postings) and by the MAXSCORE unit tests. +/// +/// Entries are unique row ids sorted ascending, each carrying the blended term +/// frequency `tf'(t, d) = Σ_f w_f · freq_f(t, d)`. +struct CombinedTermPostings { + idf: f32, + upper_bound: f32, + postings: Vec<(u64, f32)>, +} + +impl CombinedTermPostings { + /// `tf'` for `row_id`, or 0 when the term does not occur in the document. + #[inline] + 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, + } + } +} + +/// [`TermCursor`] over a [`CombinedTermPostings`] with an owned scan cursor. +struct MaterializedTerm { + postings: CombinedTermPostings, + cursor: usize, +} + +impl MaterializedTerm { + fn new(postings: CombinedTermPostings) -> Self { + Self { + postings, + cursor: 0, + } + } +} + +impl TermCursor for MaterializedTerm { + #[inline] + fn upper_bound(&self) -> f32 { + self.postings.upper_bound + } + + #[inline] + fn idf(&self) -> f32 { + self.postings.idf + } + + #[inline] + fn head(&self) -> Option { + self.postings.postings.get(self.cursor).map(|(id, _)| *id) + } + + #[inline] + fn head_tf(&self) -> f32 { + self.postings + .postings + .get(self.cursor) + .map(|(_, tf)| *tf) + .unwrap_or(0.0) + } + + #[inline] + fn consume(&mut self, row_id: u64) { + if self.head() == Some(row_id) { + self.cursor += 1; + } + } + + #[inline] + fn probe(&mut self, target: u64) -> f32 { + self.postings.tf_prime(target) + } +} + +/// One `(column, index, partition)` compressed posting list feeding a term's +/// cross-field cursor on the fast path. Decodes posting blocks lazily and skips +/// (never decodes) the blocks a row-id seek jumps past. Correct only when the +/// partition's `row_ids` are strictly ascending, so doc-id order equals row-id +/// order and a block's row-id span is an interval (checked by the caller via +/// [`DocSet::row_ids_strictly_ascending`]). +struct FastPostingSource { + weight: f32, + docs: Arc, + mask: Arc, + list: CompressedPostingList, + block_size: usize, + tail_codec: PostingTailCodec, + num_blocks: usize, + remainder: usize, + /// The block currently held in `doc_ids`/`freqs`, or `usize::MAX` if none. + decoded_block: usize, + doc_ids: Vec, + freqs: Vec, + /// Index of the head posting within the decoded block. + within: usize, + buffer: Box<[u32; MAX_POSTING_BLOCK_SIZE]>, + /// Head = first selected posting at or after `within`; `None` when exhausted. + head_row_id: Option, + head_freq: u32, + // Read accounting (this source only). + blocks_decoded: u64, + postings_decoded: u64, +} + +impl FastPostingSource { + fn new( + weight: f32, + docs: Arc, + mask: Arc, + list: CompressedPostingList, + ) -> Self { + let block_size = list.block_size; + let num_blocks = list.blocks.len(); + let remainder = list.length as usize % block_size; + let mut source = Self { + weight, + docs, + mask, + tail_codec: list.posting_tail_codec, + list, + block_size, + num_blocks, + remainder, + decoded_block: usize::MAX, + doc_ids: Vec::with_capacity(block_size), + freqs: Vec::with_capacity(block_size), + within: 0, + buffer: Box::new([0; MAX_POSTING_BLOCK_SIZE]), + head_row_id: None, + head_freq: 0, + blocks_decoded: 0, + postings_decoded: 0, + }; + // Establish the initial head (block 0 is always needed: every term is + // essential until the top-k heap fills). + if source.num_blocks > 0 { + source.decode(0); + source.settle_forward(); + } + source + } + + /// Decode block `block_idx` into `doc_ids`/`freqs`, reset `within`, and + /// count the read. Doc ids come out absolute and ascending. + fn decode(&mut self, block_idx: usize) { + let block = self.list.blocks.value(block_idx); + self.doc_ids.clear(); + self.freqs.clear(); + if block_idx + 1 == self.num_blocks && self.remainder != 0 { + decompress_posting_remainder( + block, + self.remainder, + self.tail_codec, + self.block_size, + &mut self.doc_ids, + &mut self.freqs, + ); + } else { + decompress_posting_block( + block, + &mut self.buffer[..], + &mut self.doc_ids, + &mut self.freqs, + self.block_size, + ); + } + self.decoded_block = block_idx; + self.within = 0; + self.blocks_decoded += 1; + self.postings_decoded += self.doc_ids.len() as u64; + } + + /// From the current `(decoded_block, within)`, set the head to the first + /// mask-selected posting, decoding later blocks as needed. Exhausts the + /// source (head `None`) when none remains. + fn settle_forward(&mut self) { + loop { + while self.within < self.doc_ids.len() { + let row_id = self.docs.row_id(self.doc_ids[self.within]); + if self.mask.selected(row_id) { + self.head_row_id = Some(row_id); + self.head_freq = self.freqs[self.within]; + return; + } + self.within += 1; + } + if self.decoded_block + 1 >= self.num_blocks { + self.head_row_id = None; + return; + } + self.decode(self.decoded_block + 1); + } + } + + /// Advance past the current head to the next selected posting. + fn advance(&mut self) { + self.within += 1; + self.settle_forward(); + } + + /// Position the head at the first selected posting with `row_id >= target`, + /// skipping (not decoding) blocks whose entire row-id span is below + /// `target`. `target` must not decrease across calls. + fn seek(&mut self, target: u64) { + // Skip whole blocks: block `b` is entirely below `target` when the next + // block's least row id is `<= target` (row ids ascend across blocks, so + // block `b`'s max row id is strictly below block `b + 1`'s least). Start + // from the decoded block; monotone `target` never moves us backward. + let mut block_idx = if self.decoded_block == usize::MAX { + 0 + } else { + self.decoded_block + }; + while block_idx + 1 < self.num_blocks { + let next_least_doc = self.list.block_least_doc_id(block_idx + 1); + if self.docs.row_id(next_least_doc) <= target { + block_idx += 1; + } else { + break; + } + } + if self.decoded_block != block_idx { + self.decode(block_idx); + } + // Scan forward to the first selected posting at or after `target`, + // spilling into later blocks if this one ends before `target`. + loop { + while self.within < self.doc_ids.len() { + let row_id = self.docs.row_id(self.doc_ids[self.within]); + if row_id >= target && self.mask.selected(row_id) { + self.head_row_id = Some(row_id); + self.head_freq = self.freqs[self.within]; + return; + } + self.within += 1; + } + if self.decoded_block + 1 >= self.num_blocks { + self.head_row_id = None; + return; + } + self.decode(self.decoded_block + 1); + } + } + + #[inline] + fn contribution(&self) -> f32 { + self.weight * self.head_freq as f32 + } +} + +/// [`TermCursor`] merging a term's [`FastPostingSource`]s across columns in the +/// shared row-id space — the read-pruning fast path. +struct LazyTerm { + idf: f32, + upper_bound: f32, + sources: Vec, +} + +impl TermCursor for LazyTerm { + #[inline] + fn upper_bound(&self) -> f32 { + self.upper_bound + } + + #[inline] + fn idf(&self) -> f32 { + self.idf + } + + fn head(&self) -> Option { + self.sources.iter().filter_map(|s| s.head_row_id).min() + } + + fn head_tf(&self) -> f32 { + // Sum in source (column → index → partition) order so `tf'` is + // bit-identical to the eager scan's ordered accumulation. + let Some(head) = self.head() else { + return 0.0; + }; + let mut tf = 0.0f32; + for source in &self.sources { + if source.head_row_id == Some(head) { + tf += source.contribution(); + } + } + tf + } + + fn consume(&mut self, row_id: u64) { + for source in &mut self.sources { + if source.head_row_id == Some(row_id) { + source.advance(); + } + } + } + + fn probe(&mut self, target: u64) -> f32 { + let mut tf = 0.0f32; + for source in &mut self.sources { + source.seek(target); + if source.head_row_id == Some(target) { + tf += source.contribution(); + } + } + tf + } +} + +/// Candidate-work accounting for one MAXSCORE run, for observability. Counts +/// the scoring-side work the essential/non-essential split saves. +#[derive(Default, Debug, Clone, Copy)] +struct MaxscoreStats { + /// Candidates pulled from the essential terms' cursors. + discovered: u64, + /// Discovered candidates skipped by the upper-bound test (no length lookup, + /// no non-essential probe, no full score). + pruned: u64, + /// Candidates fully scored. + scored: u64, +} + +/// Read-volume accounting for one search, for observability. `*_total` is what +/// the eager full-read scan would touch; `*_read` is what the lazy fast path +/// actually decoded. Equal on the fallback path (no read pruning there). +#[derive(Default, Debug, Clone, Copy)] +struct ReadStats { + blocks_total: u64, + blocks_read: u64, + postings_total: u64, + postings_read: u64, +} + +/// A `(column, index, partition)` posting source loaded for one term, retained +/// so the fast-path/fallback decision (which needs every posting's layout) is +/// made once without re-reading. +struct LoadedSource { + weight: f32, + docs: Arc, + is_legacy: bool, + posting: PostingList, +} + +/// Exact cross-field BM25F search over the target columns. +/// +/// Loads each query term's postings across every column/partition, then runs +/// [`combined_maxscore`]: a term-at-a-time MAXSCORE that discovers candidates +/// only from the *essential* terms and probes the *non-essential* ones on +/// demand, skipping any candidate whose score ceiling cannot beat the running +/// k-th score. The returned top-k is identical to an exact merged scan — +/// MAXSCORE changes speed, not results. +/// +/// When every partition is fast-path eligible — compressed postings with +/// strictly ascending `row_ids` (so doc-id order equals row-id order and each +/// row owns one doc) — the term cursors are lazy [`LazyTerm`]s that block-skip: +/// a non-essential term's blocks are decoded only around the sparse candidates +/// the essential terms surface, so reads are pruned, not just scoring. Any +/// partition that is legacy, has unsorted `row_ids`, or serves non-compressed +/// postings forces the full-read [`MaterializedTerm`] fallback, whose result is +/// bit-identical. +/// +/// `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, deciding fast-path eligibility as we + // go. 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, Arc)> = Vec::new(); + // Fast path: compressed postings + strictly-ascending row_ids everywhere. + let mut fast_path = true; + for column in columns { + let weight = column.weight; + for index in &column.indices { + for partition in &index.partitions { + let docs = partition.docs.ensure_loaded().await?; + let is_legacy = partition.is_legacy(); + if is_legacy || !docs.row_ids_strictly_ascending() { + fast_path = false; + } + 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?; + if !matches!(posting, PostingList::Compressed(_)) { + fast_path = false; + } + loaded[term_index].push(LoadedSource { + weight, + docs: docs.clone(), + is_legacy, + posting, + }); + } + length_sources.push((weight, docs)); + } + } + } + + let dl_prime = |row_id: u64| -> f32 { + length_sources + .iter() + .map(|(weight, docs)| weight * docs.doc_length_by_row_id(row_id) as f32) + .sum() + }; + + let (top, stats, reads) = if fast_path { + let mut cursors: Vec = terms + .iter() + .zip(loaded) + .map(|(term, sources)| build_lazy_term(term, sources, &mask, scorer)) + .collect(); + if cursors.iter().all(|term| term.head().is_none()) { + return Ok((Vec::new(), Vec::new())); + } + let (top, stats) = + combined_maxscore(&mut cursors, dl_prime, limit, require_all_terms, scorer); + let reads = fast_read_stats(&cursors); + (top, stats, reads) + } else { + let mut cursors: Vec = Vec::with_capacity(terms.len()); + let mut reads = ReadStats::default(); + for (term, sources) in terms.iter().zip(loaded) { + cursors.push(build_materialized_term( + term, sources, &mask, scorer, &mut reads, + )); + } + if cursors.iter().all(|term| term.head().is_none()) { + return Ok((Vec::new(), Vec::new())); + } + let (top, stats) = + combined_maxscore(&mut cursors, dl_prime, limit, require_all_terms, scorer); + (top, stats, reads) + }; + + tracing::debug!( + target: "lance::fts::combined", + discovered = stats.discovered, + pruned = stats.pruned, + scored = stats.scored, + blocks_total = reads.blocks_total, + blocks_read = reads.blocks_read, + postings_total = reads.postings_total, + postings_read = reads.postings_read, + fast_path = fast_path, + "combined_fields MAXSCORE", + ); + + Ok(top + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id, doc.score.0)) + .unzip()) +} + +/// Build a lazy fast-path cursor for `term`. Every source is a compressed +/// posting list (verified by the caller before choosing the fast path). +fn build_lazy_term( + term: &str, + sources: Vec, + mask: &Arc, + scorer: &CombinedFieldsBM25Scorer, +) -> LazyTerm { + let idf = scorer.query_weight(term); + let sources = sources + .into_iter() + .map(|source| { + let PostingList::Compressed(list) = source.posting else { + unreachable!("fast path requires compressed postings; eligibility verified above") + }; + FastPostingSource::new(source.weight, source.docs, mask.clone(), list) + }) + .collect(); + LazyTerm { + idf, + upper_bound: term_upper_bound(idf), + sources, + } +} + +/// Build the eager full-read cursor for `term` by merging every source's +/// postings into the shared row-id space, accumulating `tf'` in the canonical +/// order and accounting the reads (the fallback reads everything). +fn build_materialized_term( + term: &str, + sources: Vec, + mask: &Arc, + scorer: &CombinedFieldsBM25Scorer, + reads: &mut ReadStats, +) -> MaterializedTerm { + let mut acc: HashMap = HashMap::new(); + for source in &sources { + let posting_len = source.posting.len() as u64; + let blocks = match &source.posting { + PostingList::Compressed(list) => list.blocks.len() as u64, + PostingList::Plain(_) => 1, + }; + reads.postings_total += posting_len; + reads.postings_read += posting_len; + reads.blocks_total += blocks; + reads.blocks_read += blocks; + for (posting_doc_id, freq, _) in source.posting.iter() { + // Compressed postings key on a partition-local doc id; legacy plain + // postings key on the row id directly. + let row_id = if source.is_legacy { + posting_doc_id + } else { + source.docs.row_id(posting_doc_id as u32) + }; + 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); + MaterializedTerm::new(CombinedTermPostings { + idf, + upper_bound: term_upper_bound(idf), + postings, + }) +} + +/// Aggregate the fast-path term cursors' per-source read counters. +fn fast_read_stats(cursors: &[LazyTerm]) -> ReadStats { + let mut reads = ReadStats::default(); + for term in cursors { + for source in &term.sources { + reads.blocks_total += source.num_blocks as u64; + reads.blocks_read += source.blocks_decoded; + reads.postings_total += source.list.length as u64; + reads.postings_read += source.postings_decoded; + } + } + reads +} + +/// Term-at-a-time MAXSCORE over the cross-field combined postings. +/// +/// Terms are ordered ascending by their constant ceiling `upper_bound`. Once the +/// heap holds `limit` docs, the leading prefix whose cumulative ceiling cannot +/// beat the k-th score (`threshold`) is *non-essential*: candidate discovery +/// walks only the *essential* terms' cursors, and a candidate is skipped whole +/// (no length lookup, no probe) when the essential terms it matches plus every +/// non-essential ceiling still cannot beat `threshold`. `threshold` only rises, +/// so a term never returns to essential once it crosses over. +/// +/// Every surviving candidate is scored exactly, summing `idf'(t) · +/// doc_weight(tf', dl')` over the terms in their original order so the result is +/// bit-identical to the exact scan (absent terms add `doc_weight(0, ·) == 0`, +/// which matches the exact scan skipping them). `dl_prime` supplies the blended +/// document length. +/// +/// Generic over [`TermCursor`] so the fast (lazy, block-skipping) and fallback +/// (eager, materialized) paths run the exact same algorithm; only `tf'` +/// retrieval differs, so both return the identical top-k. +fn combined_maxscore( + cursors: &mut [C], + dl_prime: impl Fn(u64) -> f32, + limit: usize, + require_all_terms: bool, + scorer: &CombinedFieldsBM25Scorer, +) -> (BinaryHeap>, MaxscoreStats) { + let num_terms = cursors.len(); + // Term indices ordered ascending by ceiling; ties broken by index so the + // essential/non-essential split is deterministic. + let mut order: Vec = (0..num_terms).collect(); + order.sort_by(|&a, &b| { + cursors[a] + .upper_bound() + .partial_cmp(&cursors[b].upper_bound()) + .unwrap_or(std::cmp::Ordering::Equal) + .then(a.cmp(&b)) + }); + + let mut top: BinaryHeap> = BinaryHeap::new(); + let mut stats = MaxscoreStats::default(); + // The current k-th score; pruning is armed only once the heap is full. + let mut threshold = f32::NEG_INFINITY; + // Per-term score contribution buffer, reused across candidates (no + // per-candidate allocation). Every scored candidate overwrites all slots + // before it is summed. + let mut contrib = vec![0.0f32; num_terms]; + + loop { + // Split `order` into non-essential (a leading prefix whose cumulative + // ceiling <= threshold) and essential. Recomputed each step because + // `threshold` may have risen; it never falls, so the split only grows + // and a term never returns to essential. + let mut nonessential_upper_bound = 0.0f32; + let split = if top.len() < limit { + 0 + } else { + let mut cumulative = 0.0f32; + let mut boundary = 0; + while boundary < num_terms { + let bound = cursors[order[boundary]].upper_bound(); + if cumulative + bound <= threshold { + cumulative += bound; + boundary += 1; + } else { + break; + } + } + nonessential_upper_bound = cumulative; + boundary + }; + + // Next candidate = smallest row id at any essential term's cursor. Docs + // that match only non-essential terms are never discovered: their score + // is at most `nonessential_upper_bound <= threshold`, so the score-only + // strict-`<` collector could not admit them anyway. + let mut doc = u64::MAX; + for &term in &order[split..] { + if let Some(row_id) = cursors[term].head() { + doc = doc.min(row_id); + } + } + if doc == u64::MAX { + break; // essential terms exhausted (or all terms non-essential) + } + stats.discovered += 1; + + let dl = dl_prime(doc); + // Exact contribution of the essential terms, consuming each cursor that + // sits on `doc`. `doc_weight` is < K1 + 1 for any finite tf, so a + // present non-essential term contributes strictly less than its ceiling; + // `essential_score + nonessential_upper_bound` is therefore a strict + // over-estimate of the full score and skipping a candidate that cannot + // beat the k-th score is safe under the score-only strict-`<` collector. + // + // NOTE(pagination / PR #7846): when the collector gains a + // (score DESC, row_id ASC) tiebreak, the block-skip / candidate prune + // must weaken to strict `< threshold` and equal-bound candidates must be + // scored (and their non-essential blocks read) so the row-id tiebreak + // can decide; skipping them here would drop the pagination-correct doc. + let mut essential_score = 0.0f32; + let mut missing_term = false; + for &term in &order[split..] { + let tf = if cursors[term].head() == Some(doc) { + let tf = cursors[term].head_tf(); + cursors[term].consume(doc); + tf + } else { + 0.0 + }; + missing_term |= tf <= 0.0; + contrib[term] = cursors[term].idf() * scorer.doc_weight(tf, dl); + essential_score += contrib[term]; + } + if top.len() >= limit && essential_score + nonessential_upper_bound <= threshold { + stats.pruned += 1; + continue; + } + + // Probe the non-essential terms on demand to complete the exact score. + // `doc` increases monotonically, so each probe seeks forward and skips + // the intervening blocks (fast path). + for &term in &order[..split] { + let tf = cursors[term].probe(doc); + missing_term |= tf <= 0.0; + contrib[term] = cursors[term].idf() * scorer.doc_weight(tf, dl); + } + if require_all_terms && missing_term { + continue; + } + stats.scored += 1; + + // Sum in original term order so the score is bit-identical to the exact + // scan (absent terms contribute doc_weight(0, dl) == 0, which matches the + // exact scan skipping them). + let score: f32 = contrib.iter().sum(); + + if top.len() < limit { + top.push(Reverse(ScoredDoc::new(doc, score))); + if top.len() == limit { + threshold = top.peek().expect("heap is full").0.score.0; + } + } else if top.peek().is_some_and(|worst| worst.0.score.0 < score) { + top.pop(); + top.push(Reverse(ScoredDoc::new(doc, score))); + threshold = top.peek().expect("heap is full").0.score.0; + } + } + + (top, stats) +} + +#[cfg(test)] +mod tests { + use super::super::builder::BLOCK_SIZE; + use super::super::encoding::compress_posting_list_with_tail_codec_and_block_size; + use super::*; + use arrow_array::{UInt32Array, UInt64Array}; + use lance_select::RowAddrTreeMap; + use std::collections::BTreeSet; + + fn term(idf: f32, entries: &[(u64, f32)]) -> MaterializedTerm { + let mut postings = entries.to_vec(); + postings.sort_unstable_by_key(|(row_id, _)| *row_id); + MaterializedTerm::new(CombinedTermPostings { + idf, + upper_bound: term_upper_bound(idf), + postings, + }) + } + + // A varying blended length so scores are not all equal; the exact reference + // and MAXSCORE must read it identically. + fn dl_of(row_id: u64) -> f32 { + 3.0 + (row_id % 4) as f32 + } + + /// Independent exact reference: score every candidate the merged scan would + /// (the union of the terms' postings for OR; all-terms-present for AND) and + /// return the `limit` highest scores, descending. This is what the + /// pre-MAXSCORE exact scan's top-k heap ends up holding. + fn exact_topk_scores( + cursors: &[MaterializedTerm], + limit: usize, + require_all_terms: bool, + scorer: &CombinedFieldsBM25Scorer, + ) -> Vec { + let union: BTreeSet = cursors + .iter() + .flat_map(|t| t.postings.postings.iter().map(|(row_id, _)| *row_id)) + .collect(); + let mut scores: Vec = union + .into_iter() + .filter_map(|row_id| { + let dl = dl_of(row_id); + let mut score = 0.0f32; + let mut missing = false; + for t in cursors { + let tf = t.postings.tf_prime(row_id); + if tf <= 0.0 { + missing = true; + continue; + } + score += t.postings.idf * scorer.doc_weight(tf, dl); + } + (!(require_all_terms && missing)).then_some(score) + }) + .collect(); + scores.sort_by(|a, b| b.partial_cmp(a).unwrap()); + scores.truncate(limit); + scores + } + + fn maxscore_scores( + cursors: &mut [MaterializedTerm], + limit: usize, + require_all_terms: bool, + scorer: &CombinedFieldsBM25Scorer, + ) -> (Vec, MaxscoreStats) { + let (top, stats) = combined_maxscore(cursors, dl_of, limit, require_all_terms, scorer); + let mut scores: Vec = top.into_iter().map(|Reverse(doc)| doc.score.0).collect(); + scores.sort_by(|a, b| b.partial_cmp(a).unwrap()); + (scores, stats) + } + + fn test_scorer() -> CombinedFieldsBM25Scorer { + CombinedFieldsBM25Scorer::new(1000, 5.0, HashMap::new()) + } + + #[test] + fn test_combined_maxscore_matches_exact_scan_or() { + let scorer = test_scorer(); + // Skewed idf: a rare high-value term, a medium term, a very common term. + let build = || { + [ + term(5.0, &[(1, 3.0), (2, 1.0), (3, 2.0), (17, 1.0)]), + term(1.5, &[(1, 1.0), (3, 1.0), (5, 1.0), (8, 2.0), (17, 1.0)]), + term( + 0.05, + &(0..40u64).map(|row_id| (row_id, 1.0)).collect::>(), + ), + ] + }; + // Every top-k depth must reproduce the exact scan's top-k scores, so the + // MAXSCORE pruning provably does not change the result. + for limit in [1usize, 2, 3, 5, 10, 50] { + let expected = exact_topk_scores(&build(), limit, false, &scorer); + let (actual, _stats) = maxscore_scores(&mut build(), limit, false, &scorer); + assert_eq!( + actual, expected, + "OR top-{limit} scores diverged from exact scan" + ); + } + } + + #[test] + fn test_combined_maxscore_matches_exact_scan_and() { + let scorer = test_scorer(); + let build = || { + [ + term(5.0, &[(1, 3.0), (2, 1.0), (3, 2.0), (17, 1.0)]), + term(1.5, &[(1, 1.0), (3, 1.0), (5, 1.0), (17, 1.0)]), + ] + }; + for limit in [1usize, 2, 3, 10] { + let expected = exact_topk_scores(&build(), limit, true, &scorer); + let (actual, _stats) = maxscore_scores(&mut build(), limit, true, &scorer); + assert_eq!( + actual, expected, + "AND top-{limit} scores diverged from exact scan" + ); + // AND keeps only docs that have both terms: rows 1, 3, 17. + assert!(actual.len() <= 3); + } + } + + #[test] + fn test_combined_maxscore_no_limit_is_exact_scan() { + let scorer = test_scorer(); + let build = || { + [ + term(5.0, &[(1, 3.0), (2, 1.0), (3, 2.0)]), + term(0.05, &(0..30u64).map(|r| (r, 1.0)).collect::>()), + ] + }; + // usize::MAX limit: the heap never fills, so nothing is pruned and every + // matching doc is scored, exactly like the merged scan. + let expected = exact_topk_scores(&build(), usize::MAX, false, &scorer); + let (actual, stats) = maxscore_scores(&mut build(), usize::MAX, false, &scorer); + assert_eq!(actual, expected); + assert_eq!(stats.pruned, 0, "no limit must not prune"); + // Union of the two terms is rows 0..30. + assert_eq!(stats.scored, 30); + } + + #[test] + fn test_combined_maxscore_discovery_pruning() { + let scorer = test_scorer(); + // A rare high-idf term whose docs sit early, and a common low-idf term + // over a huge posting list. With a small k the common term goes + // non-essential once the heap fills with the rare docs, so its ~1000 + // rows are never discovered — only the essential rare term drives + // discovery. + let build = || { + [ + term(6.0, &[(1, 3.0), (2, 3.0), (3, 3.0), (4, 3.0), (5, 3.0)]), + term(0.05, &(0..1000u64).map(|r| (r, 1.0)).collect::>()), + ] + }; + let union = 1000u64; // common covers every row + + let (actual, stats) = maxscore_scores(&mut build(), 3, false, &scorer); + let expected = exact_topk_scores(&build(), 3, false, &scorer); + assert_eq!( + actual, expected, + "pruned run must still match the exact top-k" + ); + + // Discovery examines only a handful of candidates, not the full union. + assert!( + stats.discovered <= 20, + "expected heavy discovery pruning, discovered {} of {union}", + stats.discovered + ); + } + + #[test] + fn test_combined_maxscore_per_candidate_pruning() { + let scorer = test_scorer(); + // Two essential terms. `a` stays essential (large ceiling) but has a + // late doc (row 100) whose tf is low, so its exact contribution plus the + // non-essential ceiling cannot beat the k-th score: that discovered + // candidate is skipped by the per-candidate upper-bound test before any + // non-essential probe. + let build = || { + [ + term(8.0, &[(1, 10.0), (2, 10.0), (100, 1.0)]), + term(0.1, &(0..50u64).map(|r| (r, 1.0)).collect::>()), + ] + }; + + let (actual, stats) = maxscore_scores(&mut build(), 2, false, &scorer); + let expected = exact_topk_scores(&build(), 2, false, &scorer); + assert_eq!( + actual, expected, + "pruned run must still match the exact top-k" + ); + assert!( + stats.pruned >= 1, + "expected the per-candidate upper-bound test to fire, stats={stats:?}" + ); + } + + // --- Fast-path (lazy, block-skipping) equivalence ------------------------ + // + // The fast path only changes how `tf'` is fetched (lazy compressed + // block-skip vs eager materialized scan); the MAXSCORE algorithm is shared. + // These tests build real compressed posting blocks + `DocSet`s and assert + // the lazy cursors return the *identical* top-k (bit-exact scores and row + // ids, in the same order) as the materialized cursors, and that both match + // the independent exact oracle — across OR/AND, every k, block skipping, + // absent columns, a masked prefilter, ties, and a non-positive-idf term. + + /// Per-column postings for one term: `columns[c]` is that column's ascending + /// `(doc_id, freq)` list (empty when the term is absent from column `c`). + struct TermSpec { + idf: f32, + columns: Vec>, + } + + 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, + ) + } + + /// Identity `DocSet`: `row_id(doc_id) == doc_id` (strictly ascending, one + /// doc per row), with the given per-doc token counts for `dl'`. + fn identity_docset(num_tokens: &[u32]) -> Arc { + let row_ids = UInt64Array::from((0..num_tokens.len() as u64).collect::>()); + let num_tokens = UInt32Array::from(num_tokens.to_vec()); + Arc::new(DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap()) + } + + fn assert_lazy_matches_materialized( + column_num_tokens: &[Vec], + weights: &[f32], + terms: &[TermSpec], + avg_doc_length: f32, + mask: Arc, + label: &str, + ) { + let docsets: Vec> = column_num_tokens + .iter() + .map(|nt| identity_docset(nt)) + .collect(); + // Every docset must be fast-path eligible for this harness to exercise + // the lazy cursors at all. + assert!( + docsets.iter().all(|d| d.row_ids_strictly_ascending()), + "{label}: identity docsets must be strictly ascending" + ); + let scorer = CombinedFieldsBM25Scorer::new(1000, avg_doc_length, HashMap::new()); + let length_sources: Vec<(f32, Arc)> = weights + .iter() + .zip(&docsets) + .map(|(w, d)| (*w, d.clone())) + .collect(); + let dl_prime = |row_id: u64| -> f32 { + length_sources + .iter() + .map(|(w, d)| w * d.doc_length_by_row_id(row_id) as f32) + .sum() + }; + + let build_lazy = || -> Vec { + terms + .iter() + .map(|spec| { + let sources = spec + .columns + .iter() + .enumerate() + .filter(|(_, postings)| !postings.is_empty()) + .map(|(c, postings)| { + FastPostingSource::new( + weights[c], + docsets[c].clone(), + mask.clone(), + compressed_list(postings), + ) + }) + .collect(); + LazyTerm { + idf: spec.idf, + upper_bound: term_upper_bound(spec.idf), + sources, + } + }) + .collect() + }; + // Reference merge, mirroring `build_materialized_term`: sum in column + // order, drop masked rows. + let build_mat = || -> Vec { + terms + .iter() + .map(|spec| { + let mut acc: HashMap = HashMap::new(); + for (c, postings) in spec.columns.iter().enumerate() { + for (doc_id, freq) in postings { + let row_id = *doc_id as u64; + if !mask.selected(row_id) { + continue; + } + *acc.entry(row_id).or_insert(0.0) += weights[c] * *freq as f32; + } + } + let mut postings: Vec<(u64, f32)> = acc.into_iter().collect(); + postings.sort_unstable_by_key(|(row_id, _)| *row_id); + MaterializedTerm::new(CombinedTermPostings { + idf: spec.idf, + upper_bound: term_upper_bound(spec.idf), + postings, + }) + }) + .collect() + }; + + for require_all in [false, true] { + for k in [1usize, 2, 3, 5, 10, 50] { + let mut lazy = build_lazy(); + let mut mat = build_mat(); + let (lazy_top, _) = combined_maxscore(&mut lazy, dl_prime, k, require_all, &scorer); + let (mat_top, _) = combined_maxscore(&mut mat, dl_prime, k, require_all, &scorer); + + // Bit-exact scores AND row ids, in identical order. + let lazy_v: Vec<(u64, u32)> = lazy_top + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id, doc.score.0.to_bits())) + .collect(); + let mat_v: Vec<(u64, u32)> = mat_top + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id, doc.score.0.to_bits())) + .collect(); + assert_eq!( + lazy_v, mat_v, + "{label}: lazy top-k != materialized top-k (and={require_all}, k={k})" + ); + + // Independent exact oracle on scores, scoring every union doc + // with `dl_prime` (`mat` postings survive the maxscore run; + // only its cursor advanced). + let union: BTreeSet = mat + .iter() + .flat_map(|t| t.postings.postings.iter().map(|(row_id, _)| *row_id)) + .collect(); + let mut expected: Vec = union + .into_iter() + .filter_map(|row_id| { + let dl = dl_prime(row_id); + let mut score = 0.0f32; + let mut missing = false; + for t in &mat { + let tf = t.postings.tf_prime(row_id); + if tf <= 0.0 { + missing = true; + continue; + } + score += t.postings.idf * scorer.doc_weight(tf, dl); + } + (!(require_all && missing)).then_some(score) + }) + .collect(); + expected.sort_by(|a, b| b.partial_cmp(a).unwrap()); + expected.truncate(k); + let mut lazy_scores: Vec = lazy_v + .iter() + .map(|(_, bits)| f32::from_bits(*bits)) + .collect(); + lazy_scores.sort_by(|a, b| b.partial_cmp(a).unwrap()); + assert_eq!( + lazy_scores, expected, + "{label}: lazy scores != exact oracle (and={require_all}, k={k})" + ); + } + } + } + + #[test] + fn test_lazy_fast_path_matches_materialized_skew() { + // 500 docs, two columns with varying lengths. A dense common term spans + // four posting blocks (so block skipping actually fires when it goes + // non-essential); rare/mid terms with far-apart docs drive discovery and + // force cross-block seeks into the common term. `body_only` is absent + // from column 0, `rare` is present in both, `zero_idf` has a clamped + // (0) ceiling but still contributes its exact score. + let col0: Vec = (0..500).map(|d| 3 + (d % 5) as u32).collect(); + let col1: Vec = (0..500).map(|d| 2 + (d % 7) as u32).collect(); + let common: Vec<(u32, u32)> = (0..500u32).map(|d| (d, 1 + d % 3)).collect(); + let terms = [ + TermSpec { + idf: 0.05, + columns: vec![common.clone(), common], + }, + TermSpec { + idf: 6.0, + columns: vec![vec![(5, 2), (250, 1), (495, 3)], vec![(5, 1)]], + }, + TermSpec { + idf: 1.5, + columns: vec![vec![(10, 1), (300, 2)], vec![(10, 1), (260, 1)]], + }, + TermSpec { + idf: 3.0, + columns: vec![vec![], vec![(3, 2), (400, 1)]], + }, + TermSpec { + idf: 0.0, + columns: vec![vec![(7, 4), (8, 1)], vec![(7, 1)]], + }, + ]; + assert_lazy_matches_materialized( + &[col0, col1], + &[2.0, 1.0], + &terms, + 15.0, + Arc::new(RowAddrMask::all_rows()), + "skew", + ); + } + + #[test] + fn test_lazy_fast_path_matches_materialized_masked() { + // Same corpus but a prefilter blocks a handful of rows, including some a + // seek would otherwise land on — the lazy cursor must skip them exactly + // as the eager merge drops them. + let col0: Vec = (0..500).map(|d| 3 + (d % 5) as u32).collect(); + let col1: Vec = (0..500).map(|d| 2 + (d % 7) as u32).collect(); + let common: Vec<(u32, u32)> = (0..500u32).map(|d| (d, 1)).collect(); + let terms = [ + TermSpec { + idf: 0.05, + columns: vec![common.clone(), common], + }, + TermSpec { + idf: 6.0, + columns: vec![vec![(5, 2), (250, 1), (495, 3)], vec![(5, 1), (250, 2)]], + }, + ]; + let blocked = RowAddrTreeMap::from_iter([5u64, 10, 128, 250, 400]); + let mask = RowAddrMask::all_rows().also_block(blocked); + assert_lazy_matches_materialized( + &[col0, col1], + &[2.0, 1.0], + &terms, + 15.0, + Arc::new(mask), + "masked", + ); + } + + #[test] + fn test_lazy_fast_path_matches_materialized_ties() { + // Constant lengths + uniform frequency make every matching doc score + // identically. The score-only collector must keep the same k docs under + // both cursors (deterministic tie handling), not an arbitrary subset. + let col0: Vec = vec![4; 40]; + let col1: Vec = vec![3; 40]; + let all: Vec<(u32, u32)> = (0..40u32).map(|d| (d, 1)).collect(); + let terms = [TermSpec { + idf: 2.0, + columns: vec![all, vec![]], + }]; + assert_lazy_matches_materialized( + &[col0, col1], + &[1.0, 1.0], + &terms, + 7.0, + Arc::new(RowAddrMask::all_rows()), + "ties", + ); + } + + #[test] + fn test_build_materialized_term_merges_legacy_and_compressed() { + // Fallback data path: a legacy (Plain, row-id-keyed, list-multiplicity) + // source and a compressed source merge into one ordered `tf'` stream, + // masked rows dropped, contributions summed in column order, and the + // read counters report the full read (no pruning on the fallback). + use 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 identity + // docset to row 20; row 42 is blocked by the mask below. + let compressed = PostingList::Compressed(compressed_list(&[(20, 5), (42, 7)])); + let docs = identity_docset(&vec![1u32; 64]); + 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 mut reads = ReadStats::default(); + let term = build_materialized_term("t", sources, &mask, &scorer, &mut reads); + + // 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.postings, + vec![(10, 2.0), (20, 15.0), (30, 2.0)] + ); + // Fallback reads everything it loaded: 4 plain + 2 compressed postings. + assert_eq!(reads.postings_total, 6); + assert_eq!(reads.postings_read, reads.postings_total); + assert_eq!(reads.blocks_read, reads.blocks_total); + } +} diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 9b1496f7140..dadad33b71a 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -6383,6 +6383,12 @@ pub struct DocSet { // partitions never set the flag and keep exact scoring. scoring_quantized: bool, norms: Arc>>, + // Lazily-cached answer to "is `row_ids` strictly ascending?" — the + // condition the combined_fields read-pruning fast path needs (doc-id order + // == row-id order, one doc per row). Computed once per loaded set and + // shared by clones. Filled on first query use; build-time appends happen + // before any query touches it. + row_ids_ascending: Arc>, } impl DeepSizeOf for DocSet { @@ -6430,6 +6436,23 @@ impl DocSet { self.row_ids[doc_id as usize] } + /// 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 legacy list multiplicity, so a term contributes at most + /// one posting per partition per row). Frag-reuse partitions carry + /// tombstoned (non-ascending) rows and are correctly rejected here. + /// + /// The answer is a property of the loaded set, so it is computed once and + /// memoized (shared across clones). Deferred-row_id sets report `false`. + pub fn row_ids_strictly_ascending(&self) -> bool { + *self + .row_ids_ascending + .get_or_init(|| self.has_row_ids() && self.row_ids.windows(2).all(|w| w[0] < w[1])) + } + /// Resolve a `row_id` to every `doc_id` it owns. /// /// Modern indexes map each row to a single document. Older list indexes @@ -6554,6 +6577,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), } } @@ -6591,6 +6615,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }); } @@ -6634,6 +6659,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }); } @@ -6655,6 +6681,7 @@ impl DocSet { total_tokens, scoring_quantized: false, norms: Arc::new(std::sync::OnceLock::new()), + row_ids_ascending: Arc::new(std::sync::OnceLock::new()), }) } @@ -6667,6 +6694,7 @@ impl DocSet { let num_tokens = std::mem::replace(&mut self.num_tokens, NumTokens::with_capacity(len)).into_owned(); 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) { @@ -6736,6 +6764,36 @@ 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). + /// + /// Mirrors [`Self::doc_ids`] so cross-field BM25F can read a candidate's + /// per-column length with a targeted lookup instead of scanning the whole + /// `DocSet`. It is exact against a full `iter()` scan filtered to the row: + /// legacy list indexes may map one row to several documents (a contiguous + /// run in the sorted `row_ids`) and every one contributes its length, while + /// non-legacy sets resolve the row through the sorted `inv` index. + #[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 { @@ -6743,6 +6801,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 } @@ -6754,6 +6813,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() @@ -7704,6 +7771,51 @@ mod tests { 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()); + } + + #[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/parser.rs b/rust/lance-index/src/scalar/inverted/parser.rs index 8068e07bfcf..0465d788952 100644 --- a/rust/lance-index/src/scalar/inverted/parser.rs +++ b/rust/lance-index/src/scalar/inverted/parser.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors 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; @@ -110,6 +111,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(); @@ -153,6 +202,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: {}", @@ -324,4 +376,46 @@ 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.boosts, vec![1.0, 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()); + + // 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 7388bc0401f..1a557a9b6e5 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -111,6 +111,7 @@ pub enum FtsQuery { // compound queries Boost(BoostQuery), MultiMatch(MultiMatchQuery), + CombinedFields(CombinedFieldsQuery), Boolean(BooleanQuery), } @@ -125,6 +126,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, @@ -153,6 +155,7 @@ impl FtsQueryNode for FtsQuery { } columns } + Self::CombinedFields(query) => query.columns(), Self::Boolean(query) => { let mut columns = HashSet::new(); for query in &query.must { @@ -174,6 +177,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.clone(), Self::Boolean(_) => { // Bool queries don't have a single query string, they are composed of multiple queries String::new() @@ -189,6 +193,7 @@ 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()), + Self::CombinedFields(query) => query.columns.is_empty(), Self::Boolean(query) => { query.must.iter().any(|q| q.is_missing_column()) || query.should.iter().any(|q| q.is_missing_column()) @@ -217,6 +222,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 @@ -267,6 +277,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) @@ -551,6 +567,146 @@ 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 `boosts` follow Lucene's `CombinedFieldQuery`: every weight must be +/// `>= 1` (fractional weights allowed) so the combined length norm stays +/// additive. +#[derive(Debug, Clone, PartialEq)] +pub struct CombinedFieldsQuery { + /// The columns combined into the virtual field. Must be non-empty. + pub columns: Vec, + /// The query string, tokenized once and matched against every target column. + pub terms: String, + /// Per-column weights, aligned with `columns`. Each weight must be `>= 1`. + pub boosts: Vec, + /// How to combine terms: `And` (all terms must match) or `Or` (default). + pub operator: Operator, +} + +impl CombinedFieldsQuery { + /// Minimum allowed per-column weight (Lucene `CombinedFieldQuery` constraint). + const MIN_BOOST: f32 = 1.0; + + /// 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 + ))); + } + } + let boosts = vec![Self::MIN_BOOST; columns.len()]; + Ok(Self { + columns, + terms, + boosts, + operator: Operator::Or, + }) + } + + /// Set per-column weights, aligned with `columns`. + /// + /// Returns an error if the number of boosts does not match the number of + /// columns, or if any weight is not a finite number `>= 1` (mirrors Lucene's + /// `CombinedFieldQuery`, which requires `weight >= 1` so the combined length + /// norm stays additive). + pub fn try_with_boosts(mut self, boosts: Vec) -> Result { + if boosts.len() != self.columns.len() { + return Err(Error::invalid_input(format!( + "The number of boosts ({}) must match the number of columns ({})", + boosts.len(), + self.columns.len() + ))); + } + for (column, &boost) in self.columns.iter().zip(&boosts) { + if !boost.is_finite() || boost < Self::MIN_BOOST { + return Err(Error::invalid_input(format!( + "combined_fields boost for column '{}' must be a finite value >= {}, got {}", + column, + Self::MIN_BOOST, + boost + ))); + } + } + self.boosts = boosts; + Ok(self) + } + + /// Set the operator used to combine terms. + pub fn with_operator(mut self, operator: Operator) -> Self { + self.operator = operator; + self + } +} + +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.columns)?; + map.serialize_entry("boost", &self.boosts)?; + 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.columns.iter().cloned().collect() + } +} + pub enum Occur { Should, Must, @@ -888,6 +1044,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 @@ -1044,4 +1205,116 @@ 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); + + 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.columns, vec!["title".to_string(), "body".to_string()]); + assert_eq!(query.boosts, vec![1.0, 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. + assert_invalid_input(query.clone().try_with_boosts(vec![1.0]), "number of boosts"); + + // Weights below 1 are rejected (Lucene CombinedFieldQuery constraint). + assert_invalid_input( + query.clone().try_with_boosts(vec![1.0, 0.5]), + "must be a finite value >= 1", + ); + assert_invalid_input( + query.clone().try_with_boosts(vec![0.0, 1.0]), + "must be a finite value >= 1", + ); + assert_invalid_input( + query.clone().try_with_boosts(vec![f32::NAN, 1.0]), + "must be a finite value >= 1", + ); + + // Fractional weights >= 1 are accepted. + assert!(query.try_with_boosts(vec![1.5, 2.0]).is_ok()); + } } diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 3a33a67ff7a..c2dc86e7eee 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -236,3 +236,102 @@ pub fn idf(token_docs: usize, num_docs: usize) -> f32 { let num_docs = num_docs as 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 locked 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`). + pub fn query_weight(&self, term: &str) -> f32 { + match self.doc_freq.get(term).copied() { + Some(token_docs) if token_docs > 0 => idf(token_docs, self.doc_count), + _ => 0.0, + } + } + + /// BM25 term contribution for a document, given the blended term frequency + /// `tf'` and blended document length `dl'`. + 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); + (K1 + 1.0) * tf_prime / (tf_prime + doc_norm) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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); + } + + #[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 5978621c462..a34bbfda2ce 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -645,6 +645,61 @@ 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 + /// deliberately 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: _, + } = 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 @@ -1059,6 +1114,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/Cargo.toml b/rust/lance/Cargo.toml index 762ca513a19..4e0186dd7f8 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -297,6 +297,11 @@ name = "mem_wal_fts_bench" path = "benches/mem_wal/fts/mem_wal_fts_bench.rs" harness = false +[[bench]] +name = "combined_fields_compare" +path = "benches/fts/combined_fields_compare.rs" +harness = false + [[bench]] name = "mem_wal_fts_read_bench" path = "benches/mem_wal/fts/mem_wal_fts_read_bench.rs" diff --git a/rust/lance/benches/fts/LuceneCombinedFieldsBench.java b/rust/lance/benches/fts/LuceneCombinedFieldsBench.java new file mode 100644 index 00000000000..e771fddc20d --- /dev/null +++ b/rust/lance/benches/fts/LuceneCombinedFieldsBench.java @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors +// +// Apache Lucene reference for Lance's combined_fields (BM25F) query — the JVM +// side of run_combined_fields_compare.sh. Reads the shared two-field corpus and +// query set produced by the `combined_fields_compare` Rust bench, builds an +// in-memory index over `title` + `body` with BM25Similarity(1.2, 0.75), and for +// each query runs a BooleanQuery of one CombinedFieldQuery per term (SHOULD/OR, +// per-field weights from weights.txt). Emits `lucene_topk.txt`: the top-k doc +// ids per query. Rankings (not absolute scores) are what the driver compares. +// +// Usage: java -cp :. LuceneCombinedFieldsBench --in-dir DIR + +import java.io.BufferedReader; +import java.io.FileReader; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.List; + +import org.apache.lucene.analysis.core.WhitespaceAnalyzer; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.FieldType; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.CombinedFieldQuery; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.similarities.BM25Similarity; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; + +public class LuceneCombinedFieldsBench { + + static String arg(String[] a, String flag, String def) { + for (int i = 0; i < a.length - 1; i++) { + if (a[i].equals(flag)) return a[i + 1]; + } + return def; + } + + static List readLines(String path) throws Exception { + List out = new ArrayList<>(); + try (BufferedReader r = new BufferedReader(new FileReader(path))) { + String line; + while ((line = r.readLine()) != null) out.add(line); + } + return out; + } + + public static void main(String[] argv) throws Exception { + String inDir = arg(argv, "--in-dir", "/tmp/combined_fields_compare"); + List titles = readLines(Paths.get(inDir, "title.txt").toString()); + List bodies = readLines(Paths.get(inDir, "body.txt").toString()); + List queries = readLines(Paths.get(inDir, "queries.txt").toString()); + String[] w = readLines(Paths.get(inDir, "weights.txt").toString()).get(0).trim().split("\\s+"); + float wTitle = Float.parseFloat(w[0]); + float wBody = Float.parseFloat(w[1]); + int k = Integer.parseInt(w[2]); + + BM25Similarity sim = new BM25Similarity(1.2f, 0.75f); + + // CombinedFieldQuery requires norms; index freqs (no positions needed). + FieldType textType = new FieldType(); + textType.setTokenized(true); + textType.setOmitNorms(false); + textType.setIndexOptions(IndexOptions.DOCS_AND_FREQS); + textType.freeze(); + + Directory dir = new ByteBuffersDirectory(); + IndexWriterConfig iwc = new IndexWriterConfig(new WhitespaceAnalyzer()); + iwc.setSimilarity(sim); + try (IndexWriter writer = new IndexWriter(dir, iwc)) { + for (int id = 0; id < titles.size(); id++) { + Document d = new Document(); + d.add(new StoredField("id", id)); + d.add(new Field("title", titles.get(id), textType)); + d.add(new Field("body", id < bodies.size() ? bodies.get(id) : "", textType)); + writer.addDocument(d); + } + writer.commit(); + } + + DirectoryReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = new IndexSearcher(reader); + searcher.setSimilarity(sim); + StoredFields storedFields = searcher.storedFields(); + + List> topk = new ArrayList<>(); + for (String line : queries) { + BooleanQuery.Builder bq = new BooleanQuery.Builder(); + for (String term : line.trim().split("\\s+")) { + if (term.isEmpty()) continue; + CombinedFieldQuery cfq = + new CombinedFieldQuery.Builder(term) + .addField("title", wTitle) + .addField("body", wBody) + .build(); + bq.add(cfq, BooleanClause.Occur.SHOULD); + } + TopDocs td = searcher.search(bq.build(), k); + List ids = new ArrayList<>(); + for (ScoreDoc sd : td.scoreDocs) { + ids.add(storedFields.document(sd.doc).getField("id").numericValue().intValue()); + } + topk.add(ids); + } + + StringBuilder out = new StringBuilder(); + for (List ids : topk) { + for (int j = 0; j < ids.size(); j++) { + if (j > 0) out.append(' '); + out.append(ids.get(j)); + } + out.append('\n'); + } + Files.writeString(Paths.get(inDir, "lucene_topk.txt"), out.toString()); + System.out.printf("lucene combined_fields: indexed %d docs, ran %d queries (k=%d)%n", + titles.size(), queries.size(), k); + + reader.close(); + dir.close(); + } +} diff --git a/rust/lance/benches/fts/combined_fields_compare.rs b/rust/lance/benches/fts/combined_fields_compare.rs new file mode 100644 index 00000000000..fcc52a6b1b9 --- /dev/null +++ b/rust/lance/benches/fts/combined_fields_compare.rs @@ -0,0 +1,826 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Self-contained `combined_fields` (BM25F) validation harness — Lance vs Lucene. +//! +//! Unlike `mem_wal_fts_bench` (FineWeb + `FtsMemIndex`), this drives the real +//! persistent `Dataset` / `InvertedIndex` scanner path — the code a +//! `CombinedFieldsQuery` actually runs through — over a small synthetic +//! two-field corpus, and cross-checks it against both an exact brute-force +//! BM25F oracle and Apache Lucene's `CombinedFieldQuery`. +//! +//! It writes into `--out-dir`: +//! `title.txt`, `body.txt` one document per line (whitespace-tokenized) +//! `queries.txt` one query per line (space-separated terms) +//! `weights.txt` ` ` +//! `lance_topk.txt` top-k doc ids per query (Lance combined_fields) +//! `truth.txt` top-k doc ids per query (exact brute-force BM25F) +//! +//! `LuceneCombinedFieldsBench.java` reads the same corpus and emits +//! `lucene_topk.txt`; `run_combined_fields_compare.sh` reports mutual top-k +//! overlap (Lance↔Lucene) and recall@k (each vs the brute-force truth). +//! +//! Tokens are lowercase whitespace-separated so Lance's `simple` tokenizer (no +//! stemming / stop words) and Lucene's `WhitespaceAnalyzer` agree, isolating the +//! BM25F scoring. Rankings are compared (not absolute scores): Lance keeps a +//! constant `(k1 + 1)` numerator Lucene lacks, and Lucene quantizes norms. +//! +//! With `--perf` it additionally times `combined_fields` against the +//! `best_fields` (MultiMatch) baseline on the same indexed dataset and prints +//! mean/p50/p95 latency plus the combined/best slowdown. `combined_fields` runs +//! an exact merged scan with WAND disabled and an O(total docs) length pass, so +//! use a large `--docs` (e.g. `--docs 200000`) to expose its cost relative to +//! the WAND-pruned per-column `best_fields` path. +//! +//! `--skew` swaps the uniform vocab for a Zipfian one: a few head tokens land in +//! a large fraction of docs (high df, low idf) while a long tail stays rare (low +//! df, high idf), and each query mixes one common head term with one or two rare +//! tail terms. That document-frequency skew is what engages `combined_fields` +//! MAXSCORE — the common term's low ceiling falls into the non-essential prefix +//! while the rare terms drive candidate discovery. With `--perf` the MAXSCORE +//! `discovered`/`pruned`/`scored` counters are aggregated over the recall pass +//! and printed, so the discovery- and per-candidate-pruning factors are visible. +//! +//! Usage: `combined_fields_compare --out-dir DIR [--docs N] [--vocab V] [--queries Q] [--k K] [--skew] [--perf [--perf-iters I]]` + +use std::collections::HashSet; +use std::io::Write; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; + +use arrow_array::{ + ArrayRef, Float32Array, Int32Array, RecordBatch, RecordBatchIterator, StringArray, +}; +use arrow_schema::{DataType, Field, Schema as ArrowSchema}; +use lance::Dataset; +use lance::dataset::WriteParams; +use lance::dataset::optimize::{CompactionOptions, compact_files}; +use lance::index::DatasetIndexExt; +use lance_core::Result; +use lance_index::IndexType; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::scalar::inverted::query::{CombinedFieldsQuery, FtsQuery, MultiMatchQuery}; +use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; +use lance_tokenizer::Language; +use tracing::field::{Field as TracingField, Visit}; +use tracing::{Event, Subscriber}; +use tracing_subscriber::filter::{LevelFilter, Targets}; +use tracing_subscriber::layer::{Context, Layer}; +use tracing_subscriber::prelude::*; + +const W_TITLE: f32 = 2.0; +const W_BODY: f32 = 1.0; +const K1: f32 = 1.2; +const B: f32 = 0.75; + +/// Tracing target the combined_fields MAXSCORE phase logs its per-query counters +/// under (see `combined_fields_search` in lance-index). +const MAXSCORE_TARGET: &str = "lance::fts::combined"; + +/// Aggregates the `discovered`/`pruned`/`scored` counters the combined_fields +/// MAXSCORE phase emits (one `tracing::debug!` per query) so the bench can report +/// how much candidate work the essential/non-essential split saves. Summed across +/// every `combined_fields_search` between a [`reset`](Self::reset) and the +/// following [`snapshot`](Self::snapshot). +#[derive(Default)] +struct MaxscoreCounters { + discovered: AtomicU64, + pruned: AtomicU64, + scored: AtomicU64, + events: AtomicU64, + // Read-volume counters: `*_total` is what an eager full-read scan would + // touch; `*_read` is what the lazy fast path actually decoded. + blocks_total: AtomicU64, + blocks_read: AtomicU64, + postings_total: AtomicU64, + postings_read: AtomicU64, + fast_path_events: AtomicU64, +} + +impl MaxscoreCounters { + fn reset(&self) { + self.discovered.store(0, Ordering::Relaxed); + self.pruned.store(0, Ordering::Relaxed); + self.scored.store(0, Ordering::Relaxed); + self.events.store(0, Ordering::Relaxed); + self.blocks_total.store(0, Ordering::Relaxed); + self.blocks_read.store(0, Ordering::Relaxed); + self.postings_total.store(0, Ordering::Relaxed); + self.postings_read.store(0, Ordering::Relaxed); + self.fast_path_events.store(0, Ordering::Relaxed); + } + + /// `(discovered, pruned, scored, events)` accumulated since the last `reset`. + fn snapshot(&self) -> (u64, u64, u64, u64) { + ( + self.discovered.load(Ordering::Relaxed), + self.pruned.load(Ordering::Relaxed), + self.scored.load(Ordering::Relaxed), + self.events.load(Ordering::Relaxed), + ) + } + + /// `(blocks_total, blocks_read, postings_total, postings_read)` since reset. + fn read_snapshot(&self) -> (u64, u64, u64, u64) { + ( + self.blocks_total.load(Ordering::Relaxed), + self.blocks_read.load(Ordering::Relaxed), + self.postings_total.load(Ordering::Relaxed), + self.postings_read.load(Ordering::Relaxed), + ) + } +} + +/// Pulls the `u64` counters out of one MAXSCORE tracing event. +#[derive(Default)] +struct MaxscoreVisitor { + discovered: u64, + pruned: u64, + scored: u64, + blocks_total: u64, + blocks_read: u64, + postings_total: u64, + postings_read: u64, + fast_path: bool, +} + +impl Visit for MaxscoreVisitor { + fn record_u64(&mut self, field: &TracingField, value: u64) { + match field.name() { + "discovered" => self.discovered = value, + "pruned" => self.pruned = value, + "scored" => self.scored = value, + "blocks_total" => self.blocks_total = value, + "blocks_read" => self.blocks_read = value, + "postings_total" => self.postings_total = value, + "postings_read" => self.postings_read = value, + _ => {} + } + } + + fn record_bool(&mut self, field: &TracingField, value: bool) { + if field.name() == "fast_path" { + self.fast_path = value; + } + } + + fn record_debug(&mut self, _field: &TracingField, _value: &dyn std::fmt::Debug) {} +} + +/// Tracing layer that sums MAXSCORE counters into a shared [`MaxscoreCounters`]. +/// A `Targets` filter (installed in `main`) enables only [`MAXSCORE_TARGET`], so +/// every other event is disabled at its callsite and adds no timing overhead. +struct MaxscoreLayer(Arc); + +impl Layer for MaxscoreLayer { + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + if event.metadata().target() != MAXSCORE_TARGET { + return; + } + let mut visitor = MaxscoreVisitor::default(); + event.record(&mut visitor); + self.0 + .discovered + .fetch_add(visitor.discovered, Ordering::Relaxed); + self.0.pruned.fetch_add(visitor.pruned, Ordering::Relaxed); + self.0.scored.fetch_add(visitor.scored, Ordering::Relaxed); + self.0.events.fetch_add(1, Ordering::Relaxed); + self.0 + .blocks_total + .fetch_add(visitor.blocks_total, Ordering::Relaxed); + self.0 + .blocks_read + .fetch_add(visitor.blocks_read, Ordering::Relaxed); + self.0 + .postings_total + .fetch_add(visitor.postings_total, Ordering::Relaxed); + self.0 + .postings_read + .fetch_add(visitor.postings_read, Ordering::Relaxed); + if visitor.fast_path { + self.0.fast_path_events.fetch_add(1, Ordering::Relaxed); + } + } +} + +/// Deterministic splitmix64-style generator so the corpus and query set are +/// reproducible across the Lance, Lucene, and brute-force sides. +struct Lcg(u64); + +impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (self.0 >> 33) ^ self.0 + } + + fn below(&mut self, n: usize) -> usize { + (self.next() % n as u64) as usize + } +} + +/// Unnormalized Zipf cumulative weights over `vocab` ranks (rank `r` has weight +/// `1/(r + 1)`), so [`sample_zipf`] can draw ranks with a `1/rank` frequency law. +fn zipf_cumulative(vocab: usize) -> Vec { + let mut cumulative = Vec::with_capacity(vocab); + let mut acc = 0.0f64; + for rank in 0..vocab { + acc += 1.0 / (rank as f64 + 1.0); + cumulative.push(acc); + } + cumulative +} + +/// Draw a token rank from the Zipf table: rank 0 (the most common token) carries +/// ~`1/H_vocab` of the mass, so a handful of head tokens dominate while the tail +/// stays rare. Deterministic given the `Lcg`. +fn sample_zipf(rng: &mut Lcg, cumulative: &[f64]) -> usize { + let total = *cumulative.last().expect("non-empty vocab"); + // 53-bit uniform in [0, 1) from the generator, scaled to [0, total). + let u = (rng.next() >> 11) as f64 / (1u64 << 53) as f64 * total; + match cumulative.binary_search_by(|w| w.partial_cmp(&u).unwrap_or(std::cmp::Ordering::Equal)) { + Ok(i) | Err(i) => i.min(cumulative.len() - 1), + } +} + +/// A synthetic two-field corpus plus the query set, all whitespace-tokenized. +struct Corpus { + titles: Vec>, + bodies: Vec>, + queries: Vec>, +} + +fn generate_corpus(docs: usize, vocab: usize, num_queries: usize, skew: bool) -> Corpus { + if skew { + return generate_corpus_skewed(docs, vocab, num_queries); + } + let mut rng = Lcg(0x1234_5678_9abc_def0); + let token = |i: usize| format!("t{i}"); + + let sample = |rng: &mut Lcg, count: usize, skew: usize| -> Vec { + // `skew` biases token selection so the same term has different document + // frequencies in `title` vs `body`, which is exactly what BM25F blends. + (0..count) + .map(|_| token((rng.below(vocab) + skew * rng.below(3)) % vocab)) + .collect() + }; + + let mut titles = Vec::with_capacity(docs); + let mut bodies = Vec::with_capacity(docs); + for _ in 0..docs { + let title_len = 1 + rng.below(3); + let body_len = 2 + rng.below(5); + titles.push(sample(&mut rng, title_len, 0)); + bodies.push(sample(&mut rng, body_len, 1)); + } + + let queries = (0..num_queries) + .map(|_| { + let terms = 2 + rng.below(2); + (0..terms).map(|_| token(rng.below(vocab))).collect() + }) + .collect(); + + Corpus { + titles, + bodies, + queries, + } +} + +/// Skewed variant of [`generate_corpus`]: tokens are drawn Zipfian, so a few head +/// tokens land in a large fraction of docs (high df, low idf) while a long tail +/// stays rare (low df, high idf). Each query mixes one common head term +/// (rank `< HEAD`) with one or two rare tail terms (rank `>= HEAD`). That df skew +/// is what engages combined_fields MAXSCORE: the common term's low ceiling falls +/// into the non-essential prefix while the rare, high-idf terms drive candidate +/// discovery. The brute-force oracle recomputes df from this corpus, so it stays +/// exact regardless of the distribution. +fn generate_corpus_skewed(docs: usize, vocab: usize, num_queries: usize) -> Corpus { + // Common-token ranks are [0, HEAD); every query pulls its one common term + // from there and its rare terms from the [HEAD, vocab) tail. + const HEAD: usize = 8; + let head = HEAD.min(vocab.max(2) / 2).max(1); + let tail = (vocab - head).max(1); + + let mut rng = Lcg(0x1234_5678_9abc_def0); + let token = |i: usize| format!("t{i}"); + let cumulative = zipf_cumulative(vocab); + let sample = |rng: &mut Lcg, count: usize| -> Vec { + (0..count) + .map(|_| token(sample_zipf(rng, &cumulative))) + .collect() + }; + + let mut titles = Vec::with_capacity(docs); + let mut bodies = Vec::with_capacity(docs); + for _ in 0..docs { + let title_len = 1 + rng.below(3); + let body_len = 2 + rng.below(5); + titles.push(sample(&mut rng, title_len)); + bodies.push(sample(&mut rng, body_len)); + } + + let queries = (0..num_queries) + .map(|_| { + let mut terms = Vec::with_capacity(3); + terms.push(token(rng.below(head))); // one common head term + for _ in 0..(1 + rng.below(2)) { + // one or two rare tail terms + terms.push(token(head + rng.below(tail))); + } + terms + }) + .collect(); + + Corpus { + titles, + bodies, + queries, + } +} + +/// Exact BM25F reference (docFreq'/docCount' = max across fields; tf'/dl'/ +/// sumTotalTermFreq' = weighted sums; keeps Lance's `(k1 + 1)` numerator). +/// Returns the top-`k` doc ids per query, ranked by score. +fn brute_force_truth(corpus: &Corpus, k: usize) -> Vec> { + let columns: [(f32, &Vec>); 2] = + [(W_TITLE, &corpus.titles), (W_BODY, &corpus.bodies)]; + let num_docs = corpus.titles.len(); + let doc_count = num_docs; // every doc has a value in both columns + + corpus + .queries + .iter() + .map(|query| { + let mut terms = query.clone(); + terms.sort(); + terms.dedup(); + + let mut sum_total_term_freq = 0f64; + let mut doc_freq = vec![0usize; terms.len()]; + for (weight, field) in &columns { + let total_tokens: usize = field.iter().map(|d| d.len()).sum(); + sum_total_term_freq += *weight as f64 * total_tokens as f64; + for (slot, term) in terms.iter().enumerate() { + let df = field.iter().filter(|d| d.contains(term)).count(); + doc_freq[slot] = doc_freq[slot].max(df); + } + } + let avgdl = (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(); + + let mut scored: Vec<(i32, f32)> = (0..num_docs) + .filter_map(|doc| { + let mut tf = vec![0f32; terms.len()]; + let mut dl = 0f32; + for (weight, field) in &columns { + dl += weight * field[doc].len() as f32; + for (slot, term) in terms.iter().enumerate() { + let count = field[doc].iter().filter(|t| *t == term).count(); + tf[slot] += weight * count as f32; + } + } + let mut score = 0.0; + for slot in 0..terms.len() { + if tf[slot] <= 0.0 { + continue; + } + let doc_norm = K1 * (1.0 - B + B * dl / avgdl); + score += idf[slot] * (K1 + 1.0) * tf[slot] / (tf[slot] + doc_norm); + } + (score > 0.0).then_some((doc as i32, score)) + }) + .collect(); + top_k_ids(&mut scored, k) + }) + .collect() +} + +/// Per-query count of docs matching at least one query term in either field — +/// the candidate set an un-pruned merged scan would score. MAXSCORE's +/// `discovered` counter is compared against this to show discovery pruning. +fn union_sizes(corpus: &Corpus) -> Vec { + corpus + .queries + .iter() + .map(|query| { + let terms: HashSet<&str> = query.iter().map(String::as_str).collect(); + (0..corpus.titles.len()) + .filter(|&doc| { + corpus.titles[doc] + .iter() + .chain(&corpus.bodies[doc]) + .any(|t| terms.contains(t.as_str())) + }) + .count() + }) + .collect() +} + +/// Sort `(id, score)` pairs by descending score (id as a stable tiebreak) and +/// return the first `k` ids. +fn top_k_ids(scored: &mut [(i32, f32)], k: usize) -> Vec { + scored.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then(a.0.cmp(&b.0)) + }); + scored.iter().take(k).map(|(id, _)| *id).collect() +} + +fn join(tokens: &[Vec]) -> String { + tokens + .iter() + .map(|d| d.join(" ")) + .collect::>() + .join("\n") +} + +fn write_topk(path: &Path, rows: &[Vec]) -> Result<()> { + let mut out = String::new(); + for ids in rows { + let line = ids + .iter() + .map(|id| id.to_string()) + .collect::>() + .join(" "); + out.push_str(&line); + out.push('\n'); + } + std::fs::write(path, out).map_err(|e| lance_core::Error::io(format!("write {path:?}: {e}"))) +} + +/// The FTS index configuration shared by every column (scoping harness). +fn fts_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) +} + +/// (Re)build the FTS indexes on `title` and `body` in place (replace=true). +async fn create_fts_indexes(dataset: &mut Dataset) -> Result<()> { + let params = fts_params(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await?; + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await?; + Ok(()) +} + +/// Build a `Dataset` with FTS indexes on `title` and `body`. +async fn build_indexed_dataset( + dir: &Path, + corpus: &Corpus, + write_params: WriteParams, +) -> Result { + let ids = Int32Array::from((0..corpus.titles.len() as i32).collect::>()); + let titles = StringArray::from( + corpus + .titles + .iter() + .map(|d| d.join(" ")) + .collect::>(), + ); + let bodies = StringArray::from( + corpus + .bodies + .iter() + .map(|d| d.join(" ")) + .collect::>(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("title", DataType::Utf8, false), + Field::new("body", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(ids) as ArrayRef, + Arc::new(titles) as ArrayRef, + Arc::new(bodies) as ArrayRef, + ], + )?; + + let ds_uri = dir.join("lance_ds"); + let ds_uri = ds_uri.to_str().unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, ds_uri, Some(write_params)).await?; + create_fts_indexes(&mut dataset).await?; + Ok(dataset) +} + +/// `combined_fields` (BM25F) query over `title`+`body` with the fixed weights. +fn combined_query(query: &[String], k: usize) -> Result { + let combined = CombinedFieldsQuery::try_new( + query.join(" "), + vec!["title".to_string(), "body".to_string()], + )? + .try_with_boosts(vec![W_TITLE, W_BODY])?; + Ok(FullTextSearchQuery::new_query(FtsQuery::CombinedFields(combined)).limit(Some(k as i64))) +} + +/// `best_fields` (MultiMatch: per-column BM25, max-fused) with the same weights. +/// This is the pre-`combined_fields` way to search multiple columns and serves +/// as the perf baseline: each column runs a WAND-pruned `MatchQueryExec`. +fn best_fields_query(query: &[String], k: usize) -> Result { + let multi = MultiMatchQuery::try_new( + query.join(" "), + vec!["title".to_string(), "body".to_string()], + )? + .try_with_boosts(vec![W_TITLE, W_BODY])?; + Ok(FullTextSearchQuery::new_query(FtsQuery::MultiMatch(multi)).limit(Some(k as i64))) +} + +/// Run one FTS query through the scanner and return `(id, score)` for every hit. +async fn run_fts(dataset: &Dataset, fts: FullTextSearchQuery) -> Result> { + let batch = dataset + .scan() + .project(&["id"])? + .full_text_search(fts)? + .try_into_batch() + .await?; + let id_col = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let score_col = batch + .column_by_name("_score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + Ok((0..batch.num_rows()) + .map(|i| (id_col.value(i), score_col.value(i))) + .collect()) +} + +/// Mean, p50, and p95 of a latency sample (ms). Sorts `samples` in place. +fn latency_stats(samples: &mut [f64]) -> (f64, f64, f64) { + let mean = samples.iter().sum::() / samples.len() as f64; + samples.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let at = |q: f64| samples[((samples.len() as f64 * q) as usize).min(samples.len() - 1)]; + (mean, at(0.50), at(0.95)) +} + +/// Time `combined_fields` against the `best_fields` baseline on the same indexed +/// dataset, interleaving the two per query so system noise hits both equally, +/// and print mean/p50/p95 latency + the combined/best slowdown ratio. +async fn bench_latency(dataset: &Dataset, corpus: &Corpus, k: usize, iters: usize) -> Result<()> { + // Warm caches (index metadata, doc lengths) for both paths before timing. + let mut sink = 0usize; + for query in &corpus.queries { + sink += run_fts(dataset, combined_query(query, k)?).await?.len(); + sink += run_fts(dataset, best_fields_query(query, k)?).await?.len(); + } + + let mut combined_ms = Vec::with_capacity(iters * corpus.queries.len()); + let mut best_ms = Vec::with_capacity(iters * corpus.queries.len()); + for _ in 0..iters { + for query in &corpus.queries { + let bfq = best_fields_query(query, k)?; + let t = Instant::now(); + sink += run_fts(dataset, bfq).await?.len(); + best_ms.push(t.elapsed().as_secs_f64() * 1e3); + + let cfq = combined_query(query, k)?; + let t = Instant::now(); + sink += run_fts(dataset, cfq).await?.len(); + combined_ms.push(t.elapsed().as_secs_f64() * 1e3); + } + } + std::hint::black_box(sink); + + let (b_mean, b_p50, b_p95) = latency_stats(&mut best_ms); + let (c_mean, c_p50, c_p95) = latency_stats(&mut combined_ms); + let b_qps = 1e3 / b_mean; + let c_qps = 1e3 / c_mean; + let slow_mean = c_mean / b_mean; + let slow_p50 = c_p50 / b_p50; + let ndocs = corpus.titles.len(); + let nq = corpus.queries.len(); + let report = format!( + "\n=== FTS multi-column latency: best_fields vs combined_fields ===\n\ + corpus {ndocs} docs, {nq} queries, k={k}, {iters} iters\n\ + {:<16} {:>10} {:>10} {:>10} {:>12}\n\ + {:<16} {b_mean:>10.3} {b_p50:>10.3} {b_p95:>10.3} {b_qps:>12.1}\n\ + {:<16} {c_mean:>10.3} {c_p50:>10.3} {c_p95:>10.3} {c_qps:>12.1}\n\ + combined/best slowdown: {slow_mean:.2}x (mean), {slow_p50:.2}x (p50)", + "mode", "mean(ms)", "p50(ms)", "p95(ms)", "QPS", "best_fields", "combined_fields", + ); + writeln!(std::io::stdout(), "{report}") + .map_err(|e| lance_core::Error::io(format!("stdout: {e}")))?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn run( + out_dir: &Path, + docs: usize, + vocab: usize, + num_queries: usize, + k: usize, + perf: bool, + perf_iters: usize, + skew: bool, + stable_row_ids: bool, + max_rows_per_file: Option, + compact: bool, + stats: &MaxscoreCounters, +) -> Result<()> { + std::fs::create_dir_all(out_dir) + .map_err(|e| lance_core::Error::io(format!("mkdir {out_dir:?}: {e}")))?; + let corpus = generate_corpus(docs, vocab, num_queries, skew); + + std::fs::write(out_dir.join("title.txt"), join(&corpus.titles)) + .map_err(|e| lance_core::Error::io(format!("write title.txt: {e}")))?; + std::fs::write(out_dir.join("body.txt"), join(&corpus.bodies)) + .map_err(|e| lance_core::Error::io(format!("write body.txt: {e}")))?; + std::fs::write( + out_dir.join("queries.txt"), + corpus + .queries + .iter() + .map(|q| q.join(" ")) + .collect::>() + .join("\n"), + ) + .map_err(|e| lance_core::Error::io(format!("write queries.txt: {e}")))?; + std::fs::write( + out_dir.join("weights.txt"), + format!("{W_TITLE} {W_BODY} {k}\n"), + ) + .map_err(|e| lance_core::Error::io(format!("write weights.txt: {e}")))?; + + let write_params = WriteParams { + enable_stable_row_ids: stable_row_ids, + max_rows_per_file: max_rows_per_file.unwrap_or(WriteParams::default().max_rows_per_file), + ..Default::default() + }; + let mut dataset = build_indexed_dataset(out_dir, &corpus, write_params).await?; + if compact { + compact_files(&mut dataset, CompactionOptions::default(), None).await?; + // Rebuild the FTS indexes so the builder re-scans the compacted physical + // layout and re-captures row ids in the new scan order. + create_fts_indexes(&mut dataset).await?; + } + let truth = brute_force_truth(&corpus, k); + // Reset MAXSCORE counters so the snapshot below covers exactly this recall + // pass (one combined_fields_search per query), not the later perf loop. + stats.reset(); + let mut lance = Vec::with_capacity(corpus.queries.len()); + for query in &corpus.queries { + let mut scored = run_fts(&dataset, combined_query(query, k)?).await?; + lance.push(top_k_ids(&mut scored, k)); + } + write_topk(&out_dir.join("truth.txt"), &truth)?; + write_topk(&out_dir.join("lance_topk.txt"), &lance)?; + + // Recall of Lance vs the exact truth, reported inline for a quick signal. + let mut hit = 0usize; + let mut total = 0usize; + for (l, t) in lance.iter().zip(&truth) { + let tset: HashSet = t.iter().copied().collect(); + hit += l.iter().filter(|id| tset.contains(id)).count(); + total += t.len(); + } + let recall = if total > 0 { + hit as f64 / total as f64 + } else { + 0.0 + }; + let mut stdout = std::io::stdout(); + writeln!( + stdout, + "lance combined_fields recall@{k} vs brute-force BM25F = {recall:.4} ({docs} docs, {num_queries} queries)" + ) + .map_err(|e| lance_core::Error::io(format!("stdout: {e}")))?; + + // MAXSCORE candidate-work stats over the recall pass. `union` is what an + // un-pruned merged scan would score, so union/discovered is the discovery- + // pruning factor and scored/discovered the per-candidate-pruning factor. + let (discovered, pruned, scored, events) = stats.snapshot(); + if events > 0 { + let union: u64 = union_sizes(&corpus).iter().map(|&u| u as u64).sum(); + let nq = events as f64; + let disc_avg = discovered as f64 / nq; + let prun_avg = pruned as f64 / nq; + let scor_avg = scored as f64 / nq; + let union_avg = union as f64 / nq; + let discovery_factor = union as f64 / discovered.max(1) as f64; + let per_candidate_kept = scored as f64 / discovered.max(1) as f64; + let mode = if skew { "skewed" } else { "uniform" }; + writeln!( + stdout, + "\n=== combined_fields MAXSCORE candidate work ({mode}, {events} queries) ===\n\ + sums: discovered={discovered} pruned={pruned} scored={scored} union={union}\n\ + avg/q: discovered={disc_avg:.1} pruned={prun_avg:.1} scored={scor_avg:.1} union={union_avg:.1}\n\ + discovery pruning: union/discovered = {discovery_factor:.2}x fewer candidates walked\n\ + per-candidate pruning: scored/discovered = {per_candidate_kept:.3}" + ) + .map_err(|e| lance_core::Error::io(format!("stdout: {e}")))?; + + // Read-volume: `*_total` is what the eager full-read path (pre-pruning) + // decoded; `*_read` is what the lazy fast path actually decoded. The + // ratio is the read-pruning factor (BEFORE vs AFTER from one build). + let (blocks_total, blocks_read, postings_total, postings_read) = stats.read_snapshot(); + let fast = stats.fast_path_events.load(Ordering::Relaxed); + let block_prune = blocks_total as f64 / blocks_read.max(1) as f64; + let posting_prune = postings_total as f64 / postings_read.max(1) as f64; + writeln!( + stdout, + "\n=== combined_fields read volume ({mode}, {events} queries, fast_path={fast}/{events}) ===\n\ + blocks: total={blocks_total} read={blocks_read} ({block_prune:.2}x fewer decoded)\n\ + postings: total={postings_total} read={postings_read} ({posting_prune:.2}x fewer decoded)\n\ + avg/q: blocks_read={:.1} postings_read={:.1} (total blocks/q={:.1} postings/q={:.1})", + blocks_read as f64 / nq, + postings_read as f64 / nq, + blocks_total as f64 / nq, + postings_total as f64 / nq, + ) + .map_err(|e| lance_core::Error::io(format!("stdout: {e}")))?; + } + + if perf { + bench_latency(&dataset, &corpus, k, perf_iters).await?; + } + Ok(()) +} + +fn main() -> Result<()> { + let argv: Vec = std::env::args() + .skip(1) + .filter(|s| s != "--bench") + .collect(); + let get = |flag: &str, def: &str| -> String { + argv.iter() + .position(|a| a == flag) + .and_then(|i| argv.get(i + 1)) + .cloned() + .unwrap_or_else(|| def.to_string()) + }; + let out_dir = get("--out-dir", "/tmp/combined_fields_compare"); + let docs: usize = get("--docs", "1000").parse().unwrap_or(1000); + let vocab: usize = get("--vocab", "40").parse().unwrap_or(40); + let num_queries: usize = get("--queries", "40").parse().unwrap_or(40); + let k: usize = get("--k", "10").parse().unwrap_or(10); + let perf = argv.iter().any(|a| a == "--perf"); + let perf_iters: usize = get("--perf-iters", "20").parse().unwrap_or(20); + let skew = argv.iter().any(|a| a == "--skew"); + let stable_row_ids = argv.iter().any(|a| a == "--stable-row-ids"); + let max_rows_per_file: Option = argv + .iter() + .position(|a| a == "--max-rows-per-file") + .and_then(|i| argv.get(i + 1)) + .and_then(|v| v.parse().ok()); + let compact = argv.iter().any(|a| a == "--compact"); + + // Aggregate the combined_fields MAXSCORE counters via a tracing layer. The + // `Targets` filter enables only MAXSCORE_TARGET, so no other event reaches a + // callsite and the timing paths stay untouched. + let stats = Arc::new(MaxscoreCounters::default()); + tracing_subscriber::registry() + .with( + MaxscoreLayer(stats.clone()) + .with_filter(Targets::new().with_target(MAXSCORE_TARGET, LevelFilter::DEBUG)), + ) + .init(); + + let rt = tokio::runtime::Runtime::new() + .map_err(|e| lance_core::Error::io(format!("runtime: {e}")))?; + rt.block_on(run( + Path::new(&out_dir), + docs, + vocab, + num_queries, + k, + perf, + perf_iters, + skew, + stable_row_ids, + max_rows_per_file, + compact, + &stats, + )) +} diff --git a/rust/lance/benches/fts/run_combined_fields_compare.sh b/rust/lance/benches/fts/run_combined_fields_compare.sh new file mode 100755 index 00000000000..ad8e0969c3a --- /dev/null +++ b/rust/lance/benches/fts/run_combined_fields_compare.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +# combined_fields (BM25F) validation — Lance persistent scanner vs Apache Lucene. +# +# Generates a shared two-field corpus, runs Lance's CombinedFieldsQuery through +# the real Dataset/InvertedIndex scanner, runs Lucene's CombinedFieldQuery over +# the same corpus, and reports: +# - recall@k of Lance vs an exact brute-force BM25F oracle, +# - recall@k of Lucene vs the same oracle, +# - Lance<->Lucene mutual top-k overlap (Jaccard). +# Rankings are compared, not absolute scores (Lance keeps a constant (k1+1) +# numerator Lucene lacks; Lucene quantizes norms). Gate: all metrics >= MIN_OK. +# +# Usage: rust/lance/benches/fts/run_combined_fields_compare.sh +# +# Env: +# DOCS, VOCAB, QUERIES, K corpus/query knobs (defaults 1000/40/40/10) +# MIN_OK pass threshold (default 0.95) +# LUCENE_DIR Lucene source checkout (default ~/repos/extern/lucene) +# LUCENE_CP prebuilt Lucene classpath (skips the gradle build) +# JAVA_HOME JDK 21+ home (falls back to `java` on PATH) + +set -uo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +REPO_ROOT="$(git -C "$SCRIPT_DIR" rev-parse --show-toplevel)" || exit 1 +[ -n "$REPO_ROOT" ] || { echo "ERROR: could not resolve repo root" >&2; exit 1; } +cd "$REPO_ROOT" || exit 1 + +DOCS="${DOCS:-1000}" +VOCAB="${VOCAB:-40}" +QUERIES="${QUERIES:-40}" +K="${K:-10}" +MIN_OK="${MIN_OK:-0.95}" +LUCENE_DIR="${LUCENE_DIR:-$HOME/repos/extern/lucene}" +WORK="${WORK:-${TMPDIR:-/tmp}/combined_fields_compare}" +rm -rf "$WORK"; mkdir -p "$WORK" + +# ---- resolve a working cargo (the ~/.cargo/bin shim can be broken) ---- +if ! cargo --version >/dev/null 2>&1; then + CARGO_BIN="$(rustup which cargo 2>/dev/null || true)" + [ -n "$CARGO_BIN" ] && export PATH="$(dirname "$CARGO_BIN"):$PATH" +fi +cargo --version >/dev/null 2>&1 || { echo "ERROR: cargo not found" >&2; exit 1; } + +# ---- locate a JDK (21+) ---- +JAVA="${JAVA_HOME:+$JAVA_HOME/bin/java}"; JAVA="${JAVA:-java}" +JAVAC="${JAVA_HOME:+$JAVA_HOME/bin/javac}"; JAVAC="${JAVAC:-javac}" +"$JAVA" -version >/dev/null 2>&1 || { echo "ERROR: java not found" >&2; exit 1; } +echo "JDK: $("$JAVA" -version 2>&1 | head -1)" + +# ---- build Lucene classpath ---- +if [ -z "${LUCENE_CP:-}" ]; then + [ -x "$LUCENE_DIR/gradlew" ] || { echo "ERROR: set LUCENE_CP or a valid LUCENE_DIR" >&2; exit 1; } + CORE_JAR="$(find "$LUCENE_DIR/lucene/core/build/libs" -name 'lucene-core-*.jar' 2>/dev/null | head -1)" + if [ -z "$CORE_JAR" ]; then + echo "=== Building Lucene jars ($LUCENE_DIR) ===" + ( cd "$LUCENE_DIR" && ./gradlew -q :lucene:core:jar :lucene:analysis:common:jar ) \ + || { echo "ERROR: Lucene build failed" >&2; exit 1; } + CORE_JAR="$(find "$LUCENE_DIR/lucene/core/build/libs" -name 'lucene-core-*.jar' | head -1)" + fi + ANALYSIS_JAR="$(find "$LUCENE_DIR/lucene/analysis/common/build/libs" -name 'lucene-analysis-common-*.jar' | head -1)" + LUCENE_CP="$CORE_JAR:$ANALYSIS_JAR" +fi +echo "Lucene classpath: $LUCENE_CP" + +# ---- build + run the Lance side ---- +echo "=== Building Lance combined_fields bench ===" +rm -f "$REPO_ROOT"/target/release/deps/combined_fields_compare-* +cargo bench -p lance --bench combined_fields_compare --no-run +LANCE_BIN="$(find "$REPO_ROOT/target/release/deps" -maxdepth 1 -type f -perm -111 \ + -name 'combined_fields_compare-*' ! -name '*.d' -exec ls -t {} + | head -1)" +echo "Lance bench: $LANCE_BIN" +echo "=== Generating corpus + running Lance ===" +"$LANCE_BIN" --out-dir "$WORK" --docs "$DOCS" --vocab "$VOCAB" --queries "$QUERIES" --k "$K" + +# ---- compile + run the Lucene side ---- +echo "=== Compiling + running Lucene ===" +"$JAVAC" -cp "$LUCENE_CP" -d "$WORK" "$SCRIPT_DIR/LuceneCombinedFieldsBench.java" \ + || { echo "ERROR: javac failed" >&2; exit 1; } +"$JAVA" -cp "$LUCENE_CP:$WORK" LuceneCombinedFieldsBench --in-dir "$WORK" + +# ---- compare ---- +echo "=== Results ===" +python3 - "$WORK" "$K" "$MIN_OK" <<'PY' +import sys +work, k, min_ok = sys.argv[1], int(sys.argv[2]), float(sys.argv[3]) + +def rows(name): + return [[t for t in line.split()] for line in open(f"{work}/{name}").read().splitlines()] + +lance, lucene, truth = rows("lance_topk.txt"), rows("lucene_topk.txt"), rows("truth.txt") +n = min(len(lance), len(lucene), len(truth)) + +def recall(pred): + tot = 0.0 + for i in range(n): + t = set(truth[i]) + if not t: + continue + tot += len(set(pred[i]) & t) / len(t) + return tot / n if n else float("nan") + +def overlap(a, b): + tot = 0.0 + for i in range(n): + sa, sb = set(a[i]), set(b[i]) + tot += len(sa & sb) / max(len(sa | sb), 1) + return tot / n if n else float("nan") + +lance_recall = recall(lance) +lucene_recall = recall(lucene) +mutual = overlap(lance, lucene) +print(f" lance recall@{k} vs brute-force BM25F = {lance_recall:.4f}") +print(f" lucene recall@{k} vs brute-force BM25F = {lucene_recall:.4f}") +print(f" lance <-> lucene mutual top-{k} overlap = {mutual:.4f}") +ok = min(lance_recall, lucene_recall, mutual) >= min_ok +print(f" gate (>= {min_ok}): {'PASS' if ok else 'FAIL'}") +sys.exit(0 if ok else 1) +PY diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 52fb6f73bca..446d74098d5 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -95,7 +95,8 @@ use crate::io::exec::filtered_read::{ FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, }; use crate::io::exec::fts::{ - BoostQueryExec, FlatMatchFilterExec, FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec, + BoostQueryExec, CombinedFieldsQueryExec, FlatMatchFilterExec, FlatMatchQueryExec, + MatchQueryExec, PhraseQueryExec, }; use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; @@ -3256,6 +3257,16 @@ 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.columns { + if !self.fragments_covered_by_fts_leaf(column, accum).await? { + return Ok(false); + } + } + Ok(true) + } FtsQuery::Boolean(bool_query) => { for query in bool_query.must.iter() { if !self @@ -3429,6 +3440,15 @@ impl Scanner { .with_fetch(self.limit.map(|l| l as usize)), ) } + // The exec runs a single unified scan that already emits the merged + // hits sorted by score and applies the top-k limit, so (like the + // fully-indexed single-Match path) it needs no union/aggregate/sort. + FtsQuery::CombinedFields(query) => Arc::new(CombinedFieldsQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + )), FtsQuery::Boolean(query) => { // TODO: rewrite the query for better performance diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index b12a7198e5a..ee14da35ab1 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -37,7 +37,7 @@ use lance_file::version::LanceFileVersion; use lance_index::optimize::OptimizeOptions; use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::inverted::{ - InvertedListFormatVersion, + InvertedListFormatVersion, Language, query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, tokenizer::InvertedIndexParams, }; @@ -51,7 +51,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_testing::datagen::generate_random_array; use rand::Rng; use rstest::rstest; @@ -863,6 +863,960 @@ async fn test_fts_on_multiple_columns() { assert_eq!(results.num_rows(), 1); } +#[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 *any* field. best_fields + // (MultiMatch) evaluates AND per field, so it only matches a single field + // that contains every term. This is the headline semantic difference. + let params = InvertedIndexParams::default(); + let id_col = Int32Array::from(vec![0, 1, 2]); + // 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 title_col = GenericStringArray::::from(vec!["john", "john smith", "john"]); + let body_col = GenericStringArray::::from(vec!["smith", "foo", "alice"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", id_col.data_type().to_owned(), false), + arrow_schema::Field::new("title", title_col.data_type().to_owned(), false), + arrow_schema::Field::new("body", body_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(title_col) as ArrayRef, + Arc::new(body_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + // Spread the rows across fragments so the merged scan exercises multi-fragment + // row-id resolution. + let mut dataset = Dataset::write( + batches, + &test_uri, + Some(WriteParams { + max_rows_per_file: 1, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let columns = vec!["title".to_string(), "body".to_string()]; + let combined = |op| { + FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new("john smith".to_string(), columns.clone()) + .unwrap() + .with_operator(op), + ) + }; + let multi = |op| { + FtsQuery::MultiMatch( + MultiMatchQuery::try_new("john smith".to_string(), columns.clone()) + .unwrap() + .with_operator(op), + ) + }; + let as_set = |ids: Vec| { + // A duplicate-row emission collapses in the set unnoticed; assert the + // raw count matches the distinct count so it fails loudly instead. + let set = ids.iter().copied().collect::>(); + assert_eq!(set.len(), ids.len(), "duplicate row ids emitted: {ids:?}"); + set + }; + + // 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!( + as_set(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!( + as_set(fts_result_ids(&dataset, multi(Operator::And)).await), + HashSet::from([1]) + ); + + // OR matches any doc containing either term: every row has "john". + assert_eq!( + as_set(fts_result_ids(&dataset, combined(Operator::Or)).await), + HashSet::from([0, 1, 2]) + ); + assert_eq!( + as_set(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); + let id_col = Int32Array::from(vec![0, 1]); + // id 0 has the term only in `title`; id 1 has it only in the shorter `body`. + let title_col = GenericStringArray::::from(vec!["lance", "other"]); + let body_col = GenericStringArray::::from(vec!["other other other", "lance"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", id_col.data_type().to_owned(), false), + arrow_schema::Field::new("title", title_col.data_type().to_owned(), false), + arrow_schema::Field::new("body", body_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(title_col) as ArrayRef, + Arc::new(body_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let columns = vec!["title".to_string(), "body".to_string()]; + let combined = |boosts: Vec| { + FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new("lance".to_string(), columns.clone()) + .unwrap() + .try_with_boosts(boosts) + .unwrap(), + ) + }; + // 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])).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])).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 id_col = Int32Array::from(vec![0, 1, 2, 3]); + let title_col = StringArray::from(vec![Some("lance"), None, None, Some("lance")]); + let body_col = StringArray::from(vec![None, Some("lance"), None, Some("lance")]); + let all_null = StringArray::from(vec![None::<&str>, None, None, None]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, true), + arrow_schema::Field::new("body", DataType::Utf8, true), + arrow_schema::Field::new("empty", DataType::Utf8, true), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(title_col) as ArrayRef, + Arc::new(body_col) as ArrayRef, + Arc::new(all_null) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + for column in ["title", "body", "empty"] { + dataset + .create_index(&[column], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + } + let as_set = |ids: Vec| { + // A duplicate-row emission collapses in the set unnoticed; assert the + // raw count matches the distinct count so it fails loudly instead. + let set = ids.iter().copied().collect::>(); + assert_eq!(set.len(), ids.len(), "duplicate row ids emitted: {ids:?}"); + set + }; + + // 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 = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "lance".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap(), + ); + assert_eq!( + as_set(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 = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "lance".to_string(), + vec!["title".to_string(), "empty".to_string()], + ) + .unwrap(), + ); + assert_eq!( + as_set(fts_result_ids(&dataset, with_empty).await), + HashSet::from([0, 3]) + ); +} + +/// Run a full-text query and return the matched `id`s in result order. +async fn fts_result_ids(dataset: &Dataset, query: FtsQuery) -> Vec { + let batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query)) + .unwrap() + .try_into_batch() + .await + .unwrap(); + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() +} + +/// Run a full-text query and return `(id, score)` pairs in result order. +async fn fts_result_id_scores(dataset: &Dataset, query: FtsQuery) -> Vec<(i32, f32)> { + 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::(); + (0..batch.num_rows()) + .map(|i| (ids.value(i), scores.value(i))) + .collect() +} + +/// Independent brute-force BM25F reference for the exact combined-fields blend +/// (docFreq'/docCount' = max across fields; tf'/dl'/sumTotalTermFreq' = weighted +/// sums; keeps Lance's `(k1 + 1)` numerator). Assumes every document has a value +/// in every column (no nulls). Returns `None` for documents that do not match. +fn brute_force_bm25f( + columns: &[(f32, Vec<&str>)], + query: &str, + require_all_terms: bool, +) -> Vec> { + const K1: f32 = 1.2; + const B: f32 = 0.75; + let tokenize = |s: &str| { + s.split_whitespace() + .map(|w| w.to_lowercase()) + .collect::>() + }; + let mut seen = HashSet::new(); + let terms: Vec = tokenize(query) + .into_iter() + .filter(|t| seen.insert(t.clone())) + .collect(); + let num_docs = columns[0].1.len(); + let tokenized: Vec>> = columns + .iter() + .map(|(_, texts)| texts.iter().map(|t| tokenize(t)).collect()) + .collect(); + + let doc_count = num_docs; // every doc has a value in every column + 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(|d| d.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 = (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(|t| *t == term) + .count(); + tf[term_index] += weight * count as f32; + } + } + let matched = if require_all_terms { + tf.iter().all(|&f| f > 0.0) + } else { + tf.iter().any(|&f| f > 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() +} + +#[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 params = InvertedIndexParams::new("simple".to_string(), Language::English) + .lower_case(true) + .stem(false) + .remove_stop_words(false) + .ascii_folding(false) + .max_token_length(None); + + let titles = vec!["aa bb", "aa", "cc dd", "aa aa bb", "cc"]; + let bodies = vec!["cc", "bb cc dd", "aa", "dd", "aa bb"]; + let id_col = Int32Array::from((0..titles.len() as i32).collect::>()); + let title_col = GenericStringArray::::from(titles.clone()); + let body_col = GenericStringArray::::from(bodies.clone()); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, false), + arrow_schema::Field::new("body", DataType::Utf8, false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(title_col) as ArrayRef, + Arc::new(body_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + &test_uri, + None, + ) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let weights = [2.0f32, 1.0f32]; + let query = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "aa bb".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap() + .try_with_boosts(weights.to_vec()) + .unwrap(), + ); + + let expected = brute_force_bm25f( + &[(weights[0], titles.clone()), (weights[1], bodies.clone())], + "aa bb", + false, + ); + let expected_ids: HashSet = (0..expected.len()) + .filter(|&i| expected[i].is_some()) + .map(|i| i as i32) + .collect(); + + let actual = fts_result_id_scores(&dataset, query).await; + let actual_ids: HashSet = actual.iter().map(|(id, _)| *id).collect(); + // The set compare would hide a duplicate-row emission; pin the row count. + assert_eq!( + actual.len(), + expected_ids.len(), + "duplicate row emitted: {actual:?}" + ); + assert_eq!(actual_ids, expected_ids); + 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}: scan={score}, brute force={want}" + ); + } +} + +#[tokio::test] +async fn test_fts_combined_fields_topk_limit_matches_brute_force() { + // A limited top-k query arms the MAXSCORE pruning path (an unlimited query + // takes the exact scan). Drive it end-to-end through the scanner on a skewed + // corpus (a rare high-idf term, a common low-idf term) so pruning fires, and + // confirm the pruned top-k equals the exact brute-force BM25F top-k. + let params = InvertedIndexParams::new("simple".to_string(), Language::English) + .lower_case(true) + .stem(false) + .remove_stop_words(false) + .ascii_folding(false) + .max_token_length(None); + + // "alpha" appears in only the first three titles (rare, high idf); "beta" + // fills every body with growing frequency (common, low idf). The top docs + // are the three that carry the rare term; the rest can be pruned. + let titles = vec![ + "alpha", "alpha", "alpha", "gamma", "gamma", "gamma", "gamma", "gamma", "gamma", "gamma", + "gamma", "gamma", + ]; + let bodies = vec![ + "beta", + "beta beta", + "beta beta beta", + "beta", + "beta beta", + "beta beta beta", + "beta", + "beta beta", + "beta beta beta", + "beta", + "beta beta", + "beta beta beta", + ]; + let id_col = Int32Array::from((0..titles.len() as i32).collect::>()); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, false), + arrow_schema::Field::new("body", DataType::Utf8, false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(GenericStringArray::::from(titles.clone())) as ArrayRef, + Arc::new(GenericStringArray::::from(bodies.clone())) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + &test_uri, + None, + ) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let weights = [2.0f32, 1.0f32]; + let limit = 3usize; + let query = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "alpha beta".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap() + .try_with_boosts(weights.to_vec()) + .unwrap(), + ); + let scan_batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query).limit(Some(limit as i64))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let actual: Vec<(i32, f32)> = { + let ids = scan_batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + let scores = scan_batch + .column_by_name("_score") + .unwrap() + .as_primitive::(); + (0..scan_batch.num_rows()) + .map(|i| (ids.value(i), scores.value(i))) + .collect() + }; + + let expected = brute_force_bm25f( + &[(weights[0], titles.clone()), (weights[1], bodies.clone())], + "alpha beta", + false, + ); + let mut expected_topk: Vec = expected.iter().filter_map(|s| *s).collect(); + expected_topk.sort_by(|a, b| b.partial_cmp(a).unwrap()); + expected_topk.truncate(limit); + + // Same number of hits, and the pruned top-k scores equal the exact top-k + // scores (a strictly-better doc could never be pruned). + assert_eq!(actual.len(), limit); + let mut actual_scores: Vec = actual.iter().map(|(_, s)| *s).collect(); + actual_scores.sort_by(|a, b| b.partial_cmp(a).unwrap()); + for (got, want) in actual_scores.iter().zip(&expected_topk) { + assert!( + (got - want).abs() < 1e-3, + "top-k score mismatch: pruned={actual_scores:?}, exact={expected_topk:?}" + ); + } + // The three rare-term docs must be the winners; every returned score must + // match that doc's exact brute-force score. + let ids: HashSet = actual.iter().map(|(id, _)| *id).collect(); + assert_eq!(ids, HashSet::from([0, 1, 2])); + for (id, score) in actual { + let want = expected[id as usize].expect("returned an unmatched doc"); + assert!( + (score - want).abs() < 1e-3, + "score mismatch for id {id}: scan={score}, brute force={want}" + ); + } +} + +#[tokio::test] +async fn test_fts_combined_fields_multi_partition_topk_matches_brute_force() { + // Drive the read-pruning fast path end-to-end on a multi-fragment (hence + // multi-partition) skewed corpus and confirm the pruned top-k equals the + // exact brute-force BM25F top-k for OR and AND across every k. The corpus + // mixes a rare high-idf term ("zeta") with a common low-idf term ("beta"), + // so MAXSCORE makes "beta" non-essential and its blocks are block-skipped; + // small `max_rows_per_file` splits the index into several partitions so the + // lazy cursors also exercise the cross-partition row-id merge. The OR query + // additionally carries an absent term ("missingterm", idf' == 0) to check + // the clamped-ceiling handling. + let params = InvertedIndexParams::new("simple".to_string(), Language::English) + .lower_case(true) + .stem(false) + .remove_stop_words(false) + .ascii_folding(false) + .max_token_length(None); + + 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 id_col = Int32Array::from((0..n as i32).collect::>()); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, false), + arrow_schema::Field::new("body", DataType::Utf8, false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(GenericStringArray::::from( + titles.iter().map(|s| s.as_str()).collect::>(), + )) as ArrayRef, + Arc::new(GenericStringArray::::from( + bodies.iter().map(|s| s.as_str()).collect::>(), + )) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + &test_uri, + Some(WriteParams { + max_rows_per_file: 7, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + 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 weights = [2.0f32, 1.0f32]; + let columns = vec!["title".to_string(), "body".to_string()]; + + 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, + ); + let mut expected_ranked: Vec = expected.iter().filter_map(|s| *s).collect(); + expected_ranked.sort_by(|a, b| b.partial_cmp(a).unwrap()); + let matches = expected_ranked.len(); + + for k in [1usize, 2, 3, 5, 10, 50] { + let query = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new(query_str.to_string(), columns.clone()) + .unwrap() + .try_with_boosts(weights.to_vec()) + .unwrap() + .with_operator(operator), + ); + let scan_batch = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query).limit(Some(k as i64))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = scan_batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + let scores = scan_batch + .column_by_name("_score") + .unwrap() + .as_primitive::(); + let actual: Vec<(i32, f32)> = (0..scan_batch.num_rows()) + .map(|i| (ids.value(i), scores.value(i))) + .collect(); + + // Right number of hits, and the pruned top-k scores equal the exact + // top-k scores (a strictly-better doc could never be pruned). + assert_eq!( + actual.len(), + matches.min(k), + "op={operator:?} k={k}: hit count mismatch" + ); + let mut actual_scores: Vec = actual.iter().map(|(_, s)| *s).collect(); + actual_scores.sort_by(|a, b| b.partial_cmp(a).unwrap()); + for (got, want) in actual_scores.iter().zip(&expected_ranked) { + assert!( + (got - want).abs() < 1e-3, + "op={operator:?} k={k}: top-k score mismatch pruned={actual_scores:?} exact={expected_ranked:?}" + ); + } + // Every returned doc's score matches its own exact brute-force score. + for (id, score) in actual { + let want = expected[id as usize].expect("scan returned an unmatched doc"); + assert!( + (score - want).abs() < 1e-3, + "op={operator:?} k={k}: score mismatch id {id}: scan={score} brute={want}" + ); + } + } + } +} + +#[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 id_col = Int32Array::from(vec![0, 1]); + let title_col = GenericStringArray::::from(vec!["aa", "bb"]); + let body_col = GenericStringArray::::from(vec!["bb", "aa"]); + let alt_col = GenericStringArray::::from(vec!["aa", "aa"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, false), + arrow_schema::Field::new("body", DataType::Utf8, false), + arrow_schema::Field::new("alt", DataType::Utf8, false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(title_col) as ArrayRef, + Arc::new(body_col) as ArrayRef, + Arc::new(alt_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + &test_uri, + None, + ) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, &with_pos, true) + .await + .unwrap(); + dataset + .create_index(&["body"], IndexType::Inverted, None, &without_pos, true) + .await + .unwrap(); + dataset + .create_index(&["alt"], IndexType::Inverted, None, &whitespace, true) + .await + .unwrap(); + + // title (with positions) and body (without) tokenize identically: accepted. + let accepted = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "aa".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap(), + ); + let as_set = |ids: Vec| { + // A duplicate-row emission collapses in the set unnoticed; assert the + // raw count matches the distinct count so it fails loudly instead. + let set = ids.iter().copied().collect::>(); + assert_eq!(set.len(), ids.len(), "duplicate row ids emitted: {ids:?}"); + set + }; + assert_eq!( + as_set(fts_result_ids(&dataset, accepted).await), + HashSet::from([0, 1]) + ); + + // title (simple) and alt (whitespace) use different tokenizers: rejected. + let rejected = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "aa".to_string(), + vec!["title".to_string(), "alt".to_string()], + ) + .unwrap(), + ); + 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 = InvertedIndexParams::new("simple".to_string(), Language::English) + .lower_case(true) + .stem(false) + .remove_stop_words(false) + .ascii_folding(false) + .max_token_length(None); + 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 id_col = Int32Array::from((0..titles.len() as i32).collect::>()); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", DataType::Int32, false), + arrow_schema::Field::new("title", DataType::Utf8, false), + arrow_schema::Field::new("body", DataType::Utf8, false), + arrow_schema::Field::new("concat", DataType::Utf8, false), + ]) + .into(), + vec![ + Arc::new(id_col) as ArrayRef, + Arc::new(GenericStringArray::::from(titles.to_vec())) as ArrayRef, + Arc::new(GenericStringArray::::from(bodies.to_vec())) as ArrayRef, + Arc::new(GenericStringArray::::from( + concat.iter().map(|s| s.as_str()).collect::>(), + )) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema), + &test_uri, + None, + ) + .await + .unwrap(); + for column in ["title", "body", "concat"] { + dataset + .create_index(&[column], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + } + + let combined = FtsQuery::CombinedFields( + CombinedFieldsQuery::try_new( + "cat dog".to_string(), + vec!["title".to_string(), "body".to_string()], + ) + .unwrap() + .try_with_boosts(vec![w_title as f32, w_body as f32]) + .unwrap(), + ); + 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) + .await + .into_iter() + .collect(); + let plain_scores: HashMap = fts_result_id_scores(&dataset, plain) + .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}" + ); + } +} + fn nested_fts_batch( ids: Vec, a_values: Vec>, diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 2c1859211f6..5b8b49a3bb8 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -45,13 +45,14 @@ 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, FtsSearchParams, MatchQuery, PhraseQuery, Tokens, collect_query_tokens, - has_query_token, + BoostQuery, CombinedFieldsQuery, FtsSearchParams, MatchQuery, PhraseQuery, Tokens, + collect_query_tokens, has_query_token, }; use lance_index::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_index::scalar::inverted::{ - FTS_SCHEMA, InvertedIndex, MemBM25Scorer, SCORE_COL, build_global_bm25_scorer, - flat_bm25_search_stream_with_metrics_and_operator, + CombinedFieldColumn, CombinedFieldsBM25Scorer, FTS_SCHEMA, InvertedIndex, MemBM25Scorer, + SCORE_COL, build_combined_bm25_scorer, build_global_bm25_scorer, combined_fields_search, + flat_bm25_search_stream_with_metrics_and_operator, validate_combined_tokenizers, }; use lance_index::{prefilter::PreFilter, scalar::inverted::query::BooleanQuery}; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; @@ -588,6 +589,294 @@ 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, + 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>, + 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, + params, + prefilter_source, + base_scorer: None, + properties, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// 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 + } + + 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(), + params: self.params.clone(), + prefilter_source, + base_scorer: self.base_scorer.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 { + let columns = self.query.columns.join(", "); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "CombinedFieldsQuery: columns=[{}], query={}", + columns, self.query.terms + ) + } + DisplayFormatType::TreeRender => { + write!( + f, + "CombinedFieldsQuery\ncolumns=[{}]\nquery={}", + columns, self.query.terms + ) + } + } + } +} + +impl ExecutionPlan for CombinedFieldsQueryExec { + fn name(&self) -> &str { + "CombinedFieldsQueryExec" + } + + fn children(&self) -> Vec<&Arc> { + match &self.prefilter_source { + PreFilterSource::None => vec![], + PreFilterSource::FilteredRowIds(src) => vec![&src], + PreFilterSource::ScalarIndexQuery(src) => vec![&src], + } + } + + fn required_input_distribution(&self) -> Vec { + self.children() + .iter() + .map(|_| Distribution::SinglePartition) + .collect() + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + let plan = match children.len() { + 0 => { + if !matches!(self.prefilter_source, PreFilterSource::None) { + return Err(DataFusionError::Internal( + "Unexpected prefilter source".to_string(), + )); + } + self.clone_with_prefilter_source(PreFilterSource::None) + } + 1 => { + let src = children.pop().unwrap(); + let prefilter_source = match &self.prefilter_source { + PreFilterSource::FilteredRowIds(_) => PreFilterSource::FilteredRowIds(src), + PreFilterSource::ScalarIndexQuery(_) => PreFilterSource::ScalarIndexQuery(src), + PreFilterSource::None => { + return Err(DataFusionError::Internal( + "Unexpected prefilter source".to_string(), + )); + } + }; + self.clone_with_prefilter_source(prefilter_source) + } + _ => { + return Err(DataFusionError::Internal( + "Unexpected number of children".to_string(), + )); + } + }; + Ok(Arc::new(plan)) + } + + #[instrument(name = "combined_fields_query_exec", level = "debug", skip_all)] + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let query = self.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 metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); + let stream = stream::once(async move { + let _timer = metrics.baseline_metrics.elapsed_compute().timer(); + + // Open every target column's segments and pair each with its boost. + let mut columns = Vec::with_capacity(query.columns.len()); + let mut all_segments = Vec::new(); + for (column, &weight) in query.columns.iter().zip(&query.boosts) { + let segments = + load_segments(&ds, column) + .await? + .ok_or(DataFusionError::Execution(format!( + "No Inverted index found for column {}", + column + )))?; + let _details = load_segment_details(&ds, column, &segments).await?; + let indices = + open_fts_segments(&ds, column, &segments, &metrics.index_metrics).await?; + all_segments.extend(segments.iter().cloned()); + columns.push(CombinedFieldColumn { + column: column.clone(), + weight, + indices, + }); + } + validate_combined_tokenizers(&columns)?; + + // Prefilter spans the union of the target columns' segments. + let mut pre_filter = build_prefilter( + context.clone(), + partition, + &prefilter_source, + ds, + &all_segments, + )?; + 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); + } + 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 + .first() + .and_then(|column| column.indices.first()) + .ok_or(DataFusionError::Execution( + "combined_fields query has no target columns".to_string(), + ))?; + let mut tokenizer = first_index.tokenizer(); + let tokens = collect_query_tokens(&query.terms, &mut tokenizer); + + let scorer = match preset_base_scorer { + Some(scorer) => scorer, + None => { + let scorer_start = std::time::Instant::now(); + let scorer = Arc::new( + build_combined_bm25_scorer(&columns, &tokens) + .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, removing rows that do not share tokens with the query #[derive(Debug)] pub struct FlatMatchFilterExec {