diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 50bf36ee1f0..1892558bdcc 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -49,7 +49,7 @@ pub type RowPosition = u64; // Re-export public types used externally pub use btree::{BTreeIndexConfig, BTreeMemIndex}; -pub use fts::{FtsIndexConfig, FtsMemIndex, FtsQueryExpr, SearchOptions}; +pub use fts::{FtsIndexConfig, FtsMemIndex, FtsQueryExpr, SearchOptions, search_cross_column}; pub(crate) use fts::{QueryLocalFtsIndex, QueryLocalFtsStats}; pub use hnsw::{HnswIndexConfig, HnswMemIndex}; pub use pk_key::encode_pk_tuple; diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index a2d84bcb0ba..5c1caa94583 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -122,6 +122,9 @@ fn public_doc_index(coordinates: &[u32]) -> Option> { pub enum FtsQueryExpr { /// Simple term match query. Match { + /// Column this leaf searches. `None` binds to the query's single column; + /// a tree whose leaves name different columns carries `Some` on each. + column: Option, /// The search query string. query: String, /// The operator used to combine tokenized query terms. @@ -131,6 +134,9 @@ pub enum FtsQueryExpr { }, /// Phrase query with optional slop. Phrase { + /// Column this leaf searches. `None` binds to the query's single column; + /// a tree whose leaves name different columns carries `Some` on each. + column: Option, /// The phrase to search for. query: String, /// Maximum allowed distance between consecutive tokens. @@ -140,6 +146,9 @@ pub enum FtsQueryExpr { }, /// Fuzzy match query with typo tolerance. Fuzzy { + /// Column this leaf searches. `None` binds to the query's single column; + /// a tree whose leaves name different columns carries `Some` on each. + column: Option, /// The search query string. query: String, /// Maximum edit distance (Levenshtein distance). @@ -252,6 +261,7 @@ impl FtsQueryExpr { pub fn match_query_with_operator(query: impl Into, operator: Operator) -> Self { Self::Match { + column: None, query: query.into(), operator, boost: 1.0, @@ -260,6 +270,7 @@ impl FtsQueryExpr { pub fn phrase(query: impl Into) -> Self { Self::Phrase { + column: None, query: query.into(), slop: 0, boost: 1.0, @@ -268,6 +279,7 @@ impl FtsQueryExpr { pub fn phrase_with_slop(query: impl Into, slop: u32) -> Self { Self::Phrase { + column: None, query: query.into(), slop, boost: 1.0, @@ -276,6 +288,7 @@ impl FtsQueryExpr { pub fn fuzzy(query: impl Into) -> Self { Self::Fuzzy { + column: None, query: query.into(), fuzziness: None, prefix_length: 0, @@ -286,6 +299,7 @@ impl FtsQueryExpr { pub fn fuzzy_with_distance(query: impl Into, fuzziness: u32) -> Self { Self::Fuzzy { + column: None, query: query.into(), fuzziness: Some(fuzziness), prefix_length: 0, @@ -301,6 +315,7 @@ impl FtsQueryExpr { max_expansions: usize, ) -> Self { Self::Fuzzy { + column: None, query: query.into(), fuzziness, prefix_length, @@ -332,20 +347,36 @@ impl FtsQueryExpr { pub fn with_boost(self, boost: f32) -> Self { match self { Self::Match { - query, operator, .. + column, + query, + operator, + .. } => Self::Match { + column, query, operator, boost, }, - Self::Phrase { query, slop, .. } => Self::Phrase { query, slop, boost }, + Self::Phrase { + column, + query, + slop, + .. + } => Self::Phrase { + column, + query, + slop, + boost, + }, Self::Fuzzy { + column, query, fuzziness, prefix_length, max_expansions, .. } => Self::Fuzzy { + column, query, fuzziness, prefix_length, @@ -357,6 +388,153 @@ impl FtsQueryExpr { other @ (Self::Boolean { .. } | Self::Boost { .. }) => other, } } + + /// Bind this leaf to `column`. Compound nodes carry no binding of their + /// own; bind their leaves instead. + pub fn with_column(self, column: impl Into) -> Self { + let bound = Some(column.into()); + match self { + Self::Match { + query, + operator, + boost, + .. + } => Self::Match { + column: bound, + query, + operator, + boost, + }, + Self::Phrase { + query, slop, boost, .. + } => Self::Phrase { + column: bound, + query, + slop, + boost, + }, + Self::Fuzzy { + query, + fuzziness, + prefix_length, + max_expansions, + boost, + .. + } => Self::Fuzzy { + column: bound, + query, + fuzziness, + prefix_length, + max_expansions, + boost, + }, + other @ (Self::Boolean { .. } | Self::Boost { .. }) => other, + } + } + + /// The column this leaf searches, or `None` for a compound node or an + /// unbound leaf. + pub fn column(&self) -> Option<&str> { + match self { + Self::Match { column, .. } + | Self::Phrase { column, .. } + | Self::Fuzzy { column, .. } => column.as_deref(), + Self::Boolean { .. } | Self::Boost { .. } => None, + } + } + + /// Bind every leaf that names no column to `column`, leaving bound leaves + /// alone. Lets a tree built from bare terms be attached to one field. + pub fn bind_unbound_leaves(self, column: &str) -> Self { + match self { + Self::Boolean { + must, + should, + must_not, + } => Self::Boolean { + must: bind_all(must, column), + should: bind_all(should, column), + must_not: bind_all(must_not, column), + }, + Self::Boost { + positive, + negative, + negative_boost, + } => Self::Boost { + positive: Box::new(positive.bind_unbound_leaves(column)), + negative: negative.map(|n| Box::new(n.bind_unbound_leaves(column))), + negative_boost, + }, + leaf if leaf.column().is_some() => leaf, + leaf => leaf.with_column(column), + } + } + + /// Whether any leaf of this tree names no column. Cross-column routing has + /// no single column to fall back to, so such a tree cannot be evaluated. + pub fn has_unbound_leaf(&self) -> bool { + match self { + Self::Boolean { + must, + should, + must_not, + } => must + .iter() + .chain(should) + .chain(must_not) + .any(Self::has_unbound_leaf), + Self::Boost { + positive, negative, .. + } => { + positive.has_unbound_leaf() + || negative.as_ref().is_some_and(|n| n.has_unbound_leaf()) + } + leaf => leaf.column().is_none(), + } + } + + /// Distinct columns this tree's leaves name, in tree order. An unbound leaf + /// contributes nothing, so an all-unbound tree yields an empty vec. + pub fn columns(&self) -> Vec<&str> { + fn visit<'a>(expr: &'a FtsQueryExpr, out: &mut Vec<&'a str>) { + match expr { + FtsQueryExpr::Boolean { + must, + should, + must_not, + } => { + for child in must.iter().chain(should).chain(must_not) { + visit(child, out); + } + } + FtsQueryExpr::Boost { + positive, negative, .. + } => { + visit(positive, out); + if let Some(negative) = negative { + visit(negative, out); + } + } + leaf => { + if let Some(column) = leaf.column() + && !out.contains(&column) + { + out.push(column); + } + } + } + } + let mut out = Vec::new(); + visit(self, &mut out); + out + } +} + +fn bind_all(exprs: Vec, column: &str) -> Vec { + exprs + .into_iter() + .map(|expr| expr.bind_unbound_leaves(column)) + .collect() } /// Auto-fuzziness based on token length: @@ -2316,10 +2494,35 @@ impl FtsMemIndex { tail_skip: bool, ) -> Vec { match query { + FtsQueryExpr::Boolean { .. } | FtsQueryExpr::Boost { .. } => { + // Every leaf of this subtree searches this index, so the leaf + // evaluator ignores the binding and keeps the one snapshot. + combine_compound(query, &|leaf| { + self.search_leaf_with_state(leaf, st, None, include_tail, true) + }) + } + leaf => self.search_leaf_with_state(leaf, st, limit, include_tail, tail_skip), + } + } + + /// Score one `Match` / `Phrase` / `Fuzzy` leaf against this index. + /// + /// The leaf's own column binding is not consulted: routing a leaf to the + /// index holding its column is the caller's job ([`search_cross_column`]). + fn search_leaf_with_state( + &self, + leaf: &FtsQueryExpr, + st: &IndexState, + limit: Option, + include_tail: bool, + tail_skip: bool, + ) -> Vec { + match leaf { FtsQueryExpr::Match { query, operator, boost, + .. } => { let tokens = self.analyze_for_search(query); let mut results = @@ -2327,7 +2530,9 @@ impl FtsMemIndex { apply_boost(&mut results, *boost); results } - FtsQueryExpr::Phrase { query, slop, boost } => { + FtsQueryExpr::Phrase { + query, slop, boost, .. + } => { let tokens = self.analyze_for_search(query); let mut results = self.search_phrase_tokens(st, &tokens, *slop, include_tail); apply_boost(&mut results, *boost); @@ -2339,6 +2544,7 @@ impl FtsMemIndex { prefix_length, max_expansions, boost, + .. } => { let tokens = self.tokenize_for_search(query); let mut results = self.search_fuzzy_tokens( @@ -2352,25 +2558,34 @@ impl FtsMemIndex { apply_boost(&mut results, *boost); results } - FtsQueryExpr::Boolean { - must, - should, - must_not, - } => self.search_boolean(must, should, must_not, st, include_tail), - FtsQueryExpr::Boost { - positive, - negative, - negative_boost, - } => self.search_boost( - positive, - negative.as_deref(), - *negative_boost, - st, - include_tail, - ), + // `combine_compound` routes compound nodes itself and only ever + // hands a leaf here. + FtsQueryExpr::Boolean { .. } | FtsQueryExpr::Boost { .. } => Vec::new(), } } + /// Score one leaf, dropping hits past `max_row_position`. + /// + /// The bound is what makes leaves from *different* indexes safe to combine: + /// each index snapshots its own `{partitions, tail}` view, and those views + /// can disagree about how far the memtable has advanced. Clamping every + /// leaf to one row-position ceiling before the clauses meet keeps a MUST + /// across columns from dropping a row both columns actually contain. + pub fn search_leaf_bounded( + &self, + leaf: &FtsQueryExpr, + include_tail: bool, + max_row_position: Option, + ) -> Vec { + let st = self.state.load_full(); + // No limit: a clause needs its full result set before the combine. + let mut results = self.search_leaf_with_state(leaf, &st, None, include_tail, true); + if let Some(max) = max_row_position { + results.retain(|entry| entry.row_position <= max); + } + results + } + /// Execute a query with options (sort + WAND prune + limit). pub fn search_with_options( &self, @@ -2412,103 +2627,6 @@ impl FtsMemIndex { results } - fn search_boost( - &self, - positive: &FtsQueryExpr, - negative: Option<&FtsQueryExpr>, - negative_boost: f32, - st: &IndexState, - include_tail: bool, - ) -> Vec { - let mut results = self.search_query_with_state(positive, st, None, include_tail, true); - let Some(neg) = negative else { - return results; - }; - let negative_results = self.search_query_with_state(neg, st, None, include_tail, true); - // Subtractive, matching the compound scorer's contract exactly: - // `positive - negative_boost * negative_score` for a document the - // negative clause also matches, `positive` otherwise - // (`BoostScorer::score`). The negative *score* is needed, not just - // membership, which is why this keeps a map rather than a set. - // - // Scores may go negative; only non-finite values are an error upstream - // (`checked_score`), so nothing is clamped here. - let negative_scores: HashMap = negative_results - .iter() - .map(|entry| (entry.key(), entry.score)) - .collect(); - for entry in &mut results { - if let Some(negative) = negative_scores.get(&entry.key()) { - entry.score -= negative_boost * negative; - } - } - results - } - - fn search_boolean( - &self, - must: &[FtsQueryExpr], - should: &[FtsQueryExpr], - must_not: &[FtsQueryExpr], - st: &IndexState, - include_tail: bool, - ) -> Vec { - let excluded: HashSet = must_not - .iter() - .flat_map(|q| self.search_query_with_state(q, st, None, include_tail, true)) - .map(|entry| entry.key()) - .collect(); - - let mut result_map: HashMap = if must.is_empty() { - let mut map: HashMap = HashMap::new(); - for q in should { - for entry in self.search_query_with_state(q, st, None, include_tail, true) { - *map.entry(entry.key()).or_default() += entry.score; - } - } - map - } else { - let first_results = - self.search_query_with_state(&must[0], st, None, include_tail, true); - let mut map: HashMap = first_results - .into_iter() - .map(|entry| (entry.key(), entry.score)) - .collect(); - for q in must.iter().skip(1) { - let results = self.search_query_with_state(q, st, None, include_tail, true); - let result_set: HashMap = results - .into_iter() - .map(|entry| (entry.key(), entry.score)) - .collect(); - map = map - .into_iter() - .filter_map(|(pos, score)| result_set.get(&pos).map(|s| (pos, score + s))) - .collect(); - } - for q in should { - for entry in self.search_query_with_state(q, st, None, include_tail, true) { - if let Some(score) = map.get_mut(&entry.key()) { - *score += entry.score; - } - } - } - map - }; - - for pos in &excluded { - result_map.remove(pos); - } - - result_map - .into_iter() - .map(|(key, score)| FtsEntry { - row_position: key.row_position, - doc_index: public_doc_index(&key.doc_index), - score, - }) - .collect() - } - fn tokenize_for_search(&self, text: &str) -> Vec { query_tokens_to_vec(&self.analyze_for_search(text)) } @@ -3272,6 +3390,173 @@ fn relaxed_score_threshold(anchor: f32, factor: f32) -> f32 { anchor - (1.0 - factor) * anchor.abs() } +/// Evaluate a query tree, combining its clauses over leaves scored by +/// `eval_leaf`. +/// +/// The clause algebra is the contract the committed compound scorer applies: +/// MUST intersects and sums (`RequiredConjunctionScorer`), SHOULD sums into the +/// surviving set (`DisjunctionScore::Sum`), MUST_NOT excludes, and a boost +/// subtracts `negative_boost * negative_score`. It is the same whether the +/// leaves all come from one index or from one index per column — only +/// `eval_leaf` differs. +fn combine_compound(expr: &FtsQueryExpr, eval_leaf: &F) -> Vec +where + F: Fn(&FtsQueryExpr) -> Vec, +{ + match expr { + FtsQueryExpr::Boolean { + must, + should, + must_not, + } => combine_boolean(must, should, must_not, eval_leaf), + FtsQueryExpr::Boost { + positive, + negative, + negative_boost, + } => combine_boost(positive, negative.as_deref(), *negative_boost, eval_leaf), + leaf => eval_leaf(leaf), + } +} + +fn combine_boost( + positive: &FtsQueryExpr, + negative: Option<&FtsQueryExpr>, + negative_boost: f32, + eval_leaf: &F, +) -> Vec +where + F: Fn(&FtsQueryExpr) -> Vec, +{ + let mut results = combine_compound(positive, eval_leaf); + let Some(neg) = negative else { + return results; + }; + let negative_results = combine_compound(neg, eval_leaf); + // Subtractive, matching the compound scorer's contract exactly: + // `positive - negative_boost * negative_score` for a document the + // negative clause also matches, `positive` otherwise + // (`BoostScorer::score`). The negative *score* is needed, not just + // membership, which is why this keeps a map rather than a set. + // + // Scores may go negative; only non-finite values are an error upstream + // (`checked_score`), so nothing is clamped here. + let negative_scores: HashMap = negative_results + .iter() + .map(|entry| (entry.key(), entry.score)) + .collect(); + for entry in &mut results { + if let Some(negative) = negative_scores.get(&entry.key()) { + entry.score -= negative_boost * negative; + } + } + results +} + +fn combine_boolean( + must: &[FtsQueryExpr], + should: &[FtsQueryExpr], + must_not: &[FtsQueryExpr], + eval_leaf: &F, +) -> Vec +where + F: Fn(&FtsQueryExpr) -> Vec, +{ + let excluded: HashSet = must_not + .iter() + .flat_map(|q| combine_compound(q, eval_leaf)) + .map(|entry| entry.key()) + .collect(); + + let mut result_map: HashMap = if must.is_empty() { + let mut map: HashMap = HashMap::new(); + for q in should { + for entry in combine_compound(q, eval_leaf) { + *map.entry(entry.key()).or_default() += entry.score; + } + } + map + } else { + let first_results = combine_compound(&must[0], eval_leaf); + let mut map: HashMap = first_results + .into_iter() + .map(|entry| (entry.key(), entry.score)) + .collect(); + for q in must.iter().skip(1) { + let results = combine_compound(q, eval_leaf); + let result_set: HashMap = results + .into_iter() + .map(|entry| (entry.key(), entry.score)) + .collect(); + map = map + .into_iter() + .filter_map(|(pos, score)| result_set.get(&pos).map(|s| (pos, score + s))) + .collect(); + } + for q in should { + for entry in combine_compound(q, eval_leaf) { + if let Some(score) = map.get_mut(&entry.key()) { + *score += entry.score; + } + } + } + map + }; + + for pos in &excluded { + result_map.remove(pos); + } + + result_map + .into_iter() + .map(|(key, score)| FtsEntry { + row_position: key.row_position, + doc_index: public_doc_index(&key.doc_index), + score, + }) + .collect() +} + +/// Evaluate a tree whose leaves may name different columns, routing each leaf +/// to the index holding its column. +/// +/// `max_row_position` bounds every leaf to one visibility cut; see +/// [`FtsMemIndex::search_leaf_bounded`] for why that matters here and not on +/// the single-index path. +/// +/// The routing is resolved before any leaf is scored, so a tree naming a column +/// with no index fails outright instead of contributing a silently short arm. +pub fn search_cross_column( + expr: &FtsQueryExpr, + indexes: &HashMap<&str, &FtsMemIndex>, + include_tail: bool, + max_row_position: Option, +) -> Result> { + if expr.has_unbound_leaf() { + return Err(Error::invalid_input( + "cross-column full-text search needs every leaf bound to a column; \ + there is no single index to fall back to" + .to_string(), + )); + } + if let Some(missing) = expr + .columns() + .into_iter() + .find(|column| !indexes.contains_key(column)) + { + return Err(Error::invalid_input(format!( + "cross-column full-text search has no in-memory FTS index for column '{missing}'" + ))); + } + Ok(combine_compound(expr, &|leaf| { + // Resolved above, so the tree and the index set cannot disagree here. + let Some(index) = leaf.column().and_then(|column| indexes.get(column)) else { + debug_assert!(false, "cross-column FTS leaf routing was validated"); + return Vec::new(); + }; + index.search_leaf_bounded(leaf, include_tail, max_row_position) + })) +} + fn apply_boost(results: &mut [FtsEntry], boost: f32) { if boost == 1.0 { return; diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index 50f79260583..cb4f9b8dc79 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -17,7 +17,7 @@ use lance_core::{Error, ROW_ID, Result}; use lance_datafusion::expr::safe_coerce_scalar; use lance_datafusion::planner::Planner; use lance_index::scalar::FullTextSearchQuery; -use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, FtsQueryNode, Operator}; +use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, Operator}; use lance_index::scalar::inverted::{DOC_INDEX_FIELD, DocumentGranularity}; use lance_linalg::distance::DistanceType; @@ -62,13 +62,14 @@ pub struct VectorQuery { /// inverted index evaluates — so every shape the index supports (nested /// boolean, boost with a negative clause, phrase, fuzzy) reaches it without a /// lossy intermediate form. Everything outside the tree here is execution -/// policy rather than the query: which column, which document unit, and the -/// recall/latency knobs. +/// policy rather than the query: which document unit and the recall/latency +/// knobs. The columns searched live on the tree's leaves — [`Self::columns`] +/// reads them back out — so there is no second copy to fall out of step with it. #[derive(Debug, Clone)] pub struct FtsQuery { - /// Column name to search. - pub column: String, /// The query tree, evaluated as given by the memtable's inverted index. + /// Every leaf names the column it searches; the memtable routes each leaf + /// to that column's index. pub expr: FtsQueryExpr, /// Logical document unit. Defaults to one document per dataset row. pub document_granularity: DocumentGranularity, @@ -88,11 +89,11 @@ pub struct FtsQuery { pub const DEFAULT_WAND_FACTOR: f32 = 1.0; impl FtsQuery { - /// Wrap an already-built query tree. + /// Wrap an already-built query tree, binding every unbound leaf to + /// `column`. Leaves that already name a column keep it. pub fn new(column: impl Into, expr: FtsQueryExpr) -> Self { Self { - column: column.into(), - expr, + expr: expr.bind_unbound_leaves(&column.into()), document_granularity: DocumentGranularity::Row, wand_factor: DEFAULT_WAND_FACTOR, limit: None, @@ -100,6 +101,38 @@ impl FtsQuery { } } + /// Wrap a tree whose leaves name several columns. Every leaf must already + /// carry its binding — there is no single column to fall back to. + pub fn cross_column(expr: FtsQueryExpr) -> Result { + let num_columns = expr.columns().len(); + if num_columns < 2 { + return Err(Error::invalid_input(format!( + "cross-column MemTable full-text search needs at least two bound columns, \ + got {num_columns}" + ))); + } + if expr.has_unbound_leaf() { + return Err(Error::invalid_input( + "cross-column MemTable full-text search needs every leaf bound to a column; \ + there is no single index to fall back to" + .to_string(), + )); + } + Ok(Self { + expr, + document_granularity: DocumentGranularity::Row, + wand_factor: DEFAULT_WAND_FACTOR, + limit: None, + include_tail: true, + }) + } + + /// Distinct columns this query's leaves name, in tree order. A tree with + /// no leaves at all (an empty boolean) names none. + pub fn columns(&self) -> Vec<&str> { + self.expr.columns() + } + /// Create a simple term match query. pub fn match_query(column: impl Into, query: impl Into) -> Self { Self::new(column, FtsQueryExpr::match_query(query)) @@ -208,9 +241,10 @@ impl FtsQuery { /// entry type. /// /// Every shape the in-memory inverted index can evaluate is carried across: -/// match (exact and fuzzy), phrase, boolean and boost, nested to any depth. -/// Multi-match is the exception — it spans columns, and the memtable's indexes -/// are per-column, so it has no single tree to evaluate and is refused. +/// match (exact and fuzzy), phrase, boolean and boost, nested to any depth, and +/// over one column or several. Multi-match is the exception — its fields are +/// scored independently and fused by max, which the LSM planner decomposes +/// above this layer rather than expressing as one tree. fn resolve_memtable_document_granularity( column: &str, requested: Option, @@ -257,33 +291,43 @@ fn local_fts_query(query: FullTextSearchQuery, indexes: Option<&IndexStore>) -> ) }) }; - let column = require_column(single_query_column(&query.query)?)?; - let document_granularity = resolve_memtable_document_granularity( - &column, - requested_document_granularity(&query.query)?, - indexes, - )?; + let requested = requested_document_granularity(&query.query)?; let expr = to_local_expr(&query.query)?; - Ok(FtsQuery::new(column, expr) - .with_document_granularity(document_granularity) - .with_wand_factor(wand_factor) - .with_limit(limit)) -} - -/// The single column the query targets, or `None` when it binds none. A query -/// spanning several columns is refused before this — the memtable holds one -/// inverted index per column, so there is no single index to evaluate against. -fn single_query_column(query: &IndexFtsQuery) -> Result> { - let mut columns = query.columns().into_iter(); - let first = columns.next(); - if columns.next().is_some() { - return Err(Error::not_supported( - "MemTable full-text search is single-column; a query spanning several \ - columns has no single in-memory index to evaluate against" - .to_string(), - )); - } - Ok(first) + // Taken from the mapped tree rather than the index-level query: it is what + // the arm evaluates, and its order is the tree's rather than a hash set's. + let columns = expr + .columns() + .into_iter() + .map(str::to_string) + .collect::>(); + let local = if columns.len() > 1 { + // Each column resolves its granularity against its own index, and they + // have to agree: the arm emits one schema, and `_doc_index` is present + // or absent for the whole batch. + let mut resolved = None; + for column in &columns { + let granularity = resolve_memtable_document_granularity(column, requested, indexes)?; + match resolved { + None => resolved = Some(granularity), + Some(previous) if previous != granularity => { + return Err(Error::invalid_input(format!( + "cross-column full-text search resolved {previous:?} document \ + granularity for an earlier column and {granularity:?} for \ + '{column}'; they must agree" + ))); + } + Some(_) => {} + } + } + let document_granularity = resolved.unwrap_or(DocumentGranularity::Row); + FtsQuery::cross_column(expr)?.with_document_granularity(document_granularity) + } else { + let column = require_column(columns.into_iter().next())?; + let document_granularity = + resolve_memtable_document_granularity(&column, requested, indexes)?; + FtsQuery::new(column, expr).with_document_granularity(document_granularity) + }; + Ok(local.with_wand_factor(wand_factor).with_limit(limit)) } /// The document granularity the query asks for, erroring if its leaves disagree @@ -328,30 +372,44 @@ fn requested_document_granularity(query: &IndexFtsQuery) -> Result Result { + fn bind(expr: FtsQueryExpr, column: Option<&String>) -> FtsQueryExpr { + match column { + Some(column) => expr.with_column(column.clone()), + None => expr, + } + } Ok(match query { - IndexFtsQuery::Match(m) => match m.fuzziness { - // `Some(0)` is an exact match in the index model. - Some(0) => FtsQueryExpr::match_query_with_operator(m.terms.clone(), m.operator) + IndexFtsQuery::Match(m) => bind( + match m.fuzziness { + // `Some(0)` is an exact match in the index model. + Some(0) => FtsQueryExpr::match_query_with_operator(m.terms.clone(), m.operator) + .with_boost(m.boost), + // The fuzzy path expands each term independently and unions the + // expansions, so it cannot also require every term to match. + _ if m.operator != Operator::Or => { + return Err(Error::not_supported( + "MemTable fuzzy full-text search only supports OR match operators" + .to_string(), + )); + } + fuzziness => FtsQueryExpr::fuzzy_with_options( + m.terms.clone(), + fuzziness, + m.prefix_length, + m.max_expansions, + ) .with_boost(m.boost), - // The fuzzy path expands each term independently and unions the - // expansions, so it cannot also require every term to match. - _ if m.operator != Operator::Or => { - return Err(Error::not_supported( - "MemTable fuzzy full-text search only supports OR match operators".to_string(), - )); - } - fuzziness => FtsQueryExpr::fuzzy_with_options( - m.terms.clone(), - fuzziness, - m.prefix_length, - m.max_expansions, - ) - .with_boost(m.boost), - }, - IndexFtsQuery::Phrase(p) => FtsQueryExpr::phrase_with_slop(p.terms.clone(), p.slop), + }, + m.column.as_ref(), + ), + IndexFtsQuery::Phrase(p) => bind( + FtsQueryExpr::phrase_with_slop(p.terms.clone(), p.slop), + p.column.as_ref(), + ), IndexFtsQuery::Boost(b) => FtsQueryExpr::boosting_with_negative( to_local_expr(&b.positive)?, to_local_expr(&b.negative)?, @@ -1199,7 +1257,13 @@ impl MemTableScanner { /// Uses the effective visibility (min of max_readable and max_indexed) to ensure /// queries only see indexed data. async fn plan_fts_search(&self, query: &FtsQuery) -> Result> { - if !self.has_fts_index(&query.column, query.document_granularity) { + // Every queried column needs an index: a cross-column predicate is one + // predicate, so a missing arm is a missing answer, not a smaller one. + if !query + .columns() + .into_iter() + .all(|column| self.has_fts_index(column, query.document_granularity)) + { return self.empty_fts_plan(query.document_granularity); } @@ -1867,7 +1931,7 @@ mod tests { .with_column("text".to_string()) .unwrap(); let local = local_fts_query(q, None).unwrap(); - assert_eq!(local.column, "text"); + assert_eq!(local.columns(), ["text"]); assert!( matches!(local.expr, FtsQueryExpr::Match { query, operator, .. } if query == "hello" && operator == Operator::Or) @@ -1928,7 +1992,7 @@ mod tests { )); let local = local_fts_query(exact_and, None).unwrap(); assert!( - matches!(local.expr, FtsQueryExpr::Match { query, operator, boost } + matches!(local.expr, FtsQueryExpr::Match { query, operator, boost, .. } if query == "hello world" && operator == Operator::And && boost == 3.0) ); diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs index 73d6710d7ed..e8dbfb21a59 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs @@ -3,6 +3,7 @@ //! FtsIndexExec - Full-text search with MVCC visibility. +use std::collections::HashMap; use std::fmt::{Debug, Formatter}; use std::sync::Arc; @@ -27,7 +28,7 @@ use lance_index::scalar::inverted::DOC_INDEX_FIELD; use super::super::builder::FtsQuery; use super::newest_pk_positions; -use crate::dataset::mem_wal::index::SearchOptions; +use crate::dataset::mem_wal::index::{SearchOptions, search_cross_column}; use crate::dataset::mem_wal::scanner::exec::resolve_pk_indices; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; @@ -42,6 +43,10 @@ struct BatchRange { batch_id: usize, } +/// One scored hit: row position, the element ordinal for a list-element +/// document, and the BM25 score. +type FtsHit = (u64, Option>, f32); + type MaterializedFtsRows = ( Vec>, Vec, @@ -79,7 +84,7 @@ pub struct FtsIndexExec { impl Debug for FtsIndexExec { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("FtsIndexExec") - .field("column", &self.query.column) + .field("columns", &self.query.columns()) .field("expr", &self.query.expr) .field("readable_count", &self.readable_count) .field("with_row_id", &self.with_row_id) @@ -108,16 +113,19 @@ impl FtsIndexExec { base_schema: SchemaRef, with_row_id: bool, ) -> Result { - // Verify the index exists for this column - let column = &query.column; - let Some(_index) = - indexes.get_fts_by_column_and_granularity(column, query.document_granularity) - else { - return Err(Error::invalid_input(format!( - "No FTS index found for column '{}'", - column - ))); - }; + // Every queried column must resolve an index. A cross-column predicate + // is one predicate: a column with no arm is a missing answer rather + // than a narrower one. + for column in query.columns() { + if indexes + .get_fts_by_column_and_granularity(column, query.document_granularity) + .is_none() + { + return Err(Error::invalid_input(format!( + "No FTS index found for column '{column}'" + ))); + } + } let with_doc_index = query.document_granularity.is_list_element(); // Build output schema: base fields + optional _doc_index + _score + optional _rowid @@ -215,12 +223,21 @@ impl FtsIndexExec { } /// Query the index and return matching rows with BM25 scores. - fn query_index(&self) -> Vec<(u64, Option>, f32)> { + fn query_index(&self) -> Result> { + let columns = self.query.columns(); + if columns.len() > 1 { + return self.query_across_columns(&columns); + } + let Some(&column) = columns.first() else { + return Err(Error::invalid_input( + "full-text search names no column to search".to_string(), + )); + }; let Some(index) = self .indexes - .get_fts_by_column_and_granularity(&self.query.column, self.query.document_granularity) + .get_fts_by_column_and_granularity(column, self.query.document_granularity) else { - return vec![]; + return Ok(vec![]); }; // The scanner carries the tree the index evaluates, so there is nothing @@ -248,17 +265,45 @@ impl FtsIndexExec { let entries = index.search_with_options(&query_expr, options); // Convert to (row_position, element ordinal, score) tuples. - entries + Ok(entries .into_iter() .map(|entry| (entry.row_position, entry.doc_index, entry.score)) - .collect() + .collect()) + } + + /// Route each leaf of a cross-column tree to the index holding its column. + /// + /// No WAND pruning and no query limit: the clauses combine over full result + /// sets, the same way a single-index compound query already does. The + /// visibility ceiling goes *in* rather than being applied after, so leaves + /// read from indexes whose tails have advanced differently still meet over + /// one cut. + fn query_across_columns(&self, columns: &[&str]) -> Result> { + let mut indexes = HashMap::with_capacity(columns.len()); + for &column in columns { + let Some(index) = self + .indexes + .get_fts_by_column_and_granularity(column, self.query.document_granularity) + else { + return Err(Error::invalid_input(format!( + "No FTS index found for column '{column}'" + ))); + }; + indexes.insert(column, index); + } + Ok(search_cross_column( + &self.query.expr, + &indexes, + self.query.include_tail, + self.max_readable_row, + )? + .into_iter() + .map(|entry| (entry.row_position, entry.doc_index, entry.score)) + .collect()) } /// Filter results by MVCC visibility using max_row_position. O(n). - fn filter_by_visibility( - &self, - results: Vec<(u64, Option>, f32)>, - ) -> Vec<(u64, Option>, f32)> { + fn filter_by_visibility(&self, results: Vec) -> Vec { let Some(max_readable) = self.max_readable_row else { return vec![]; }; @@ -272,10 +317,7 @@ impl FtsIndexExec { /// /// This method processes results one at a time to preserve the score-sorted order, /// then combines them into a single batch. - fn materialize_rows_sorted( - &self, - results: &[(u64, Option>, f32)], - ) -> DataFusionResult> { + fn materialize_rows_sorted(&self, results: &[FtsHit]) -> DataFusionResult> { if results.is_empty() { return Ok(vec![]); } @@ -578,15 +620,19 @@ impl DisplayAs for FtsIndexExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!( f, - "FtsIndexExec: column={}, query_type={:?}, with_row_id={}", - self.query.column, self.query.expr, self.with_row_id + "FtsIndexExec: columns={:?}, query_type={:?}, with_row_id={}", + self.query.columns(), + self.query.expr, + self.with_row_id ) } DisplayFormatType::TreeRender => { write!( f, - "FtsIndexExec\ncolumn={}\nquery_type={:?}\nwith_row_id={}", - self.query.column, self.query.expr, self.with_row_id + "FtsIndexExec\ncolumns={:?}\nquery_type={:?}\nwith_row_id={}", + self.query.columns(), + self.query.expr, + self.with_row_id ) } } @@ -624,7 +670,7 @@ impl ExecutionPlan for FtsIndexExec { _context: Arc, ) -> DataFusionResult { // Query the index - let results = self.query_index(); + let results = self.query_index()?; // Filter by visibility let mut visible_results = self.filter_by_visibility(results); diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 5d8885d01cd..8092d1b0329 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -35,6 +35,7 @@ //! base/SSTable Lance datasets, `MemTableScanner` for the active //! memtable) and requires no changes to `lance-index`. +use std::collections::HashMap; use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; @@ -49,7 +50,7 @@ use datafusion::prelude::Expr; use lance_core::{Error, Result, is_system_column}; use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::InvertedIndexParams; -use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, FtsQueryNode, Operator}; +use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, Operator}; use lance_index::scalar::inverted::{DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity}; use tracing::instrument; @@ -292,21 +293,29 @@ fn transient_fts_index_store( batch_store: &Arc, source: &IndexStore, schema: &SchemaRef, - column: &str, + columns: &[String], document_granularity: DocumentGranularity, pk_columns: &[String], - index_params: Option<&InvertedIndexParams>, + index_params: &HashMap<&str, InvertedIndexParams>, ) -> Result>> { let visible_batches = source.visible_count(); if visible_batches == 0 { return Ok(None); } - let field_id = schema.index_of(column).map_err(|_| { - Error::invalid_input(format!( - "FTS query column '{column}' is not in the MemWAL schema" - )) - })? as i32; + let field_ids = columns + .iter() + .map(|column| { + schema + .index_of(column) + .map(|field_id| (column.as_str(), field_id as i32)) + .map_err(|_| { + Error::invalid_input(format!( + "FTS query column '{column}' is not in the MemWAL schema" + )) + }) + }) + .collect::>>()?; // PK columns first: `enable_pk_index` refuses to run once a search index // holds rows, and the FTS exec needs the PK index to drop superseded @@ -335,16 +344,19 @@ fn transient_fts_index_store( // rows about what matches — silently, by returning fewer rows. Defaults are // right only when no persisted index covers the column, where there is no // contract to match. - let params = index_params - .cloned() - .unwrap_or_default() - .document_granularity(document_granularity); - store.add_fts_with_params( - format!("__transient_fts_{column}"), - field_id, - column.to_string(), - params, - )?; + for (column, field_id) in field_ids { + let params = index_params + .get(column) + .cloned() + .unwrap_or_default() + .document_granularity(document_granularity); + store.add_fts_with_params( + format!("__transient_fts_{column}"), + field_id, + column.to_string(), + params, + )?; + } // Exactly the prefix the real store publishes. A bare `IndexStore` carries // no durability cursors, so its own `visible_count` is its indexed prefix — @@ -362,18 +374,28 @@ fn transient_fts_index_store( Ok(Some(Arc::new(store))) } -/// Split a query targeting several columns into one sub-query per column, or -/// `None` when it targets at most one and takes the single-column path unchanged. +/// How the planner routes a query across the columns it names. +enum FtsPlanShape { + /// One predicate, evaluated whole by every source against these columns. + /// Usually one; several when the tree's leaves name different fields. + Bound(Vec), + /// A top-level multi-match: independent per-column searches, unioned and + /// collapsed to the best field per row. + PerColumn(Vec<(String, IndexFtsQuery)>), +} + +/// Decide how `query` reaches the columns it names. /// -/// Only a top-level multi-match decomposes. Its leaves are independent -/// per-column matches, which is exactly what the base-table path scores -/// separately before taking the best per row. A boolean or boost spanning -/// columns is one predicate over several fields — `must: [a in title, b in -/// body]` is a conjunction, not a union of per-column results — so unioning -/// per-column plans would answer a different query, and it is refused instead. -fn cross_column_targets(query: &IndexFtsQuery) -> Result>> { - if query.columns().len() <= 1 { - return Ok(None); +/// A top-level multi-match decomposes: its leaves are independent per-column +/// matches, which is exactly what the base-table path scores separately before +/// taking the best per row. Every other shape spanning columns is *one* +/// predicate over several fields — `must: [a in title, b in body]` is a +/// conjunction, not a union of per-column results — so it stays whole and each +/// source evaluates it across all of them. +fn fts_plan_shape(query: &IndexFtsQuery) -> Result { + let columns = collect_query_columns(query); + if columns.len() <= 1 { + return Ok(FtsPlanShape::Bound(columns)); } // The on-disk cross-column path accepts Row documents only // (`validate_row_leaf_granularities`); match that contract rather than @@ -387,10 +409,7 @@ fn cross_column_targets(query: &IndexFtsQuery) -> Result Result>>() - .map(Some) + .map(FtsPlanShape::PerColumn) +} + +/// The columns `query` names, in tree order and deduplicated. +/// +/// `FtsQueryNode::columns` returns a `HashSet`, and the order decides which +/// column a single-column plan binds to and the order arms are built in, so it +/// is derived from the tree here instead. +fn collect_query_columns(query: &IndexFtsQuery) -> Vec { + fn visit(query: &IndexFtsQuery, out: &mut Vec) { + let mut push = |column: &Option| { + if let Some(column) = column + && !out.contains(column) + { + out.push(column.clone()); + } + }; + match query { + IndexFtsQuery::Match(query) => push(&query.column), + IndexFtsQuery::Phrase(query) => push(&query.column), + IndexFtsQuery::MultiMatch(query) => { + for leaf in &query.match_queries { + push(&leaf.column); + } + } + IndexFtsQuery::Boost(query) => { + visit(&query.positive, out); + visit(&query.negative, out); + } + IndexFtsQuery::Boolean(query) => { + for child in query + .must + .iter() + .chain(&query.should) + .chain(&query.must_not) + { + visit(child, out); + } + } + } + } + let mut out = Vec::new(); + visit(query, &mut out); + out } /// Plans local-scoring FTS queries over LSM data. @@ -525,15 +587,19 @@ impl LsmFtsSearchPlanner { )); } - let Some(per_column) = cross_column_targets(&query.query)? else { - // `cross_column_targets` returns `None` only for a query naming at - // most one column, and the guard above rules out naming none. - let column = query.columns().into_iter().next().ok_or_else(|| { - Error::internal("full-text query names no column after the bound check".to_string()) - })?; - return self - .plan_single_column_search(&column, query, limit, projection) - .await; + let per_column = match fts_plan_shape(&query.query)? { + FtsPlanShape::Bound(columns) => { + // The bound check above rules out a query naming no column. + if columns.is_empty() { + return Err(Error::internal( + "full-text query names no column after the bound check".to_string(), + )); + } + return self + .plan_bound_search(&columns, query, limit, projection) + .await; + } + FtsPlanShape::PerColumn(per_column) => per_column, }; // Collapsing field hits needs a row identity to collapse *by*. Without @@ -568,8 +634,8 @@ impl LsmFtsSearchPlanner { bounded.limit = query.limit; bounded.wand_factor = query.wand_factor; per_column_plans.push( - Box::pin(self.plan_single_column_search( - &field, + Box::pin(self.plan_bound_search( + std::slice::from_ref(&field), bounded, candidate_limit, projection, @@ -634,87 +700,80 @@ impl LsmFtsSearchPlanner { }) } + /// Plan one query evaluated whole against `columns` on every source. + /// + /// One column is the ordinary case. Several means the tree's leaves name + /// different fields and the predicate spans them — a MUST across columns is + /// an intersection, so it cannot be split into per-column arms the way a + /// multi-match can. Each source evaluates the whole tree instead: the base + /// and SSTable arms through the dataset scanner's own cross-column path, + /// the memtable arm by routing each leaf to that column's in-memory index. #[instrument( - name = "lsm_fts_search_column", + name = "lsm_fts_search_columns", level = "info", skip_all, - fields(column = %column, limit) + fields(columns = ?columns, limit) )] - async fn plan_single_column_search( + async fn plan_bound_search( &self, - column: &str, + columns: &[String], mut query: FullTextSearchQuery, limit: Option, projection: Option<&[String]>, ) -> Result> { let sources = self.collector.collect()?; let requested = requested_query_document_granularity(&query.query)?; - let mut available = Vec::new(); - let mut source_granularities = Vec::with_capacity(sources.len()); - // The analyzer/positional contract each granularity's persisted index - // was built with. Which one applies is only known once the query's - // granularity is resolved below — row and list-element indexes may - // coexist on a column with different settings, so choosing early would - // rebuild the transient arm against the wrong one. - let mut index_params: Vec<(DocumentGranularity, InvertedIndexParams)> = Vec::new(); - for source in &sources { - let granularities = match source { - LsmDataSource::BaseTable { dataset } => { - if index_params.is_empty() { - index_params = indexed_fts_index_params(dataset, column).await?; - } - indexed_fts_document_granularities(dataset, column) - .await? - .into_iter() - .map(|(_, document_granularity)| document_granularity) - .collect::>() - } - LsmDataSource::SsTable { path, .. } => { - let dataset = open_sstable( - path, - self.session.as_ref(), - self.store_params.as_ref(), - self.sstable_cache.as_ref(), - self.warmer.as_ref(), - ) - .await?; - if index_params.is_empty() { - index_params = indexed_fts_index_params(&dataset, column).await?; - } - indexed_fts_document_granularities(&dataset, column) - .await? - .into_iter() - .map(|(_, document_granularity)| document_granularity) - .collect::>() - } - LsmDataSource::ActiveMemTable { index_store, .. } => { - index_store.fts_document_granularities_by_column(column) + + // Resolve each column independently, then require the results to agree. + // The arm emits one schema per source, and `_doc_index` is present for + // the whole batch or not at all, so columns that resolved to different + // document units could not be merged. + let mut document_granularity: Option = None; + let mut index_params: HashMap<&str, InvertedIndexParams> = + HashMap::with_capacity(columns.len()); + for column in columns { + let (granularity, params) = self + .resolve_column_index_contract(&sources, column, requested) + .await?; + match document_granularity { + None => document_granularity = Some(granularity), + Some(previous) if previous != granularity => { + return Err(Error::not_supported(format!( + "cross-column full-text search resolved {previous:?} document \ + granularity for an earlier column and {granularity:?} for \ + '{column}'; they must agree" + ))); } - }; - available.extend(granularities.iter().copied()); - source_granularities.push(granularities); + Some(_) => {} + } + if let Some(params) = params { + index_params.insert(column.as_str(), params); + } } - let document_granularity = - resolve_document_granularity_from_candidates(column, requested, available)?; - validate_source_document_granularities( - column, - document_granularity, - &source_granularities, - )?; - // Now that the granularity is settled, take the settings of the index - // the query actually selects — `load_segments` picks the persisted index - // the same way. - let index_params = index_params - .into_iter() - .find(|(granularity, _)| *granularity == document_granularity) - .map(|(_, params)| params); + let document_granularity = document_granularity.ok_or_else(|| { + Error::internal("full-text query names no column after the bound check".to_string()) + })?; + // Same contract the committed cross-column scorer enforces + // (`validate_row_leaf_granularities`): an element hit carries a + // per-element coordinate, and the clauses meet on row identity. + if columns.len() > 1 && document_granularity.is_list_element() { + return Err(Error::not_supported( + "cross-column full-text search supports row documents only: element hits \ + carry a per-element coordinate the clauses cannot be joined on" + .to_string(), + )); + } + let schema = lance_core::datatypes::Schema::try_from(self.base_schema.as_ref())?; - resolve_fts_field(&schema, column, document_granularity)?; + for column in columns { + resolve_fts_field(&schema, column, document_granularity)?; + } set_query_document_granularity(&mut query.query, document_granularity); - if sources - .iter() - .any(|source| active_source_can_execute_fts(source, column, document_granularity)) - { + if sources.iter().any(|source| { + columns + .iter() + .any(|column| active_source_can_execute_fts(source, column, document_granularity)) + }) { validate_lsm_fts_query(&query)?; } let allowed_system_columns: &[&str] = if document_granularity.is_list_element() { @@ -777,11 +836,11 @@ impl LsmFtsSearchPlanner { futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, fetch_limit)| { Box::pin(self.build_source_plan( source, - column, + columns, &query, *fetch_limit, projection, - index_params.as_ref(), + &index_params, )) })) .await?; @@ -828,6 +887,76 @@ impl LsmFtsSearchPlanner { self.sort_by_score(merged, limit) } + /// Resolve one column's document granularity across every source, plus the + /// analyzer/positional settings its persisted index was built with. + /// + /// The settings are part of the query contract — a phrase needs positions, + /// and a stemming or n-gram tokenizer changes which terms a document + /// produces — so the transient memtable arm rebuilds against them rather + /// than against defaults. Which one applies is only known once the + /// granularity is settled: row and list-element indexes may coexist on a + /// column with different settings. + async fn resolve_column_index_contract( + &self, + sources: &[LsmDataSource], + column: &str, + requested: Option, + ) -> Result<(DocumentGranularity, Option)> { + let mut available = Vec::new(); + let mut source_granularities = Vec::with_capacity(sources.len()); + let mut index_params: Vec<(DocumentGranularity, InvertedIndexParams)> = Vec::new(); + for source in sources { + let granularities = match source { + LsmDataSource::BaseTable { dataset } => { + if index_params.is_empty() { + index_params = indexed_fts_index_params(dataset, column).await?; + } + indexed_fts_document_granularities(dataset, column) + .await? + .into_iter() + .map(|(_, document_granularity)| document_granularity) + .collect::>() + } + LsmDataSource::SsTable { path, .. } => { + let dataset = open_sstable( + path, + self.session.as_ref(), + self.store_params.as_ref(), + self.sstable_cache.as_ref(), + self.warmer.as_ref(), + ) + .await?; + if index_params.is_empty() { + index_params = indexed_fts_index_params(&dataset, column).await?; + } + indexed_fts_document_granularities(&dataset, column) + .await? + .into_iter() + .map(|(_, document_granularity)| document_granularity) + .collect::>() + } + LsmDataSource::ActiveMemTable { index_store, .. } => { + index_store.fts_document_granularities_by_column(column) + } + }; + available.extend(granularities.iter().copied()); + source_granularities.push(granularities); + } + let document_granularity = + resolve_document_granularity_from_candidates(column, requested, available)?; + validate_source_document_granularities( + column, + document_granularity, + &source_granularities, + )?; + // `load_segments` picks the persisted index the same way. + let params = index_params + .into_iter() + .find(|(granularity, _)| *granularity == document_granularity) + .map(|(_, params)| params); + Ok((document_granularity, params)) + } + /// Order a merged FTS result by `_score` descending, capped at `limit`. /// /// Per-partition sort with `fetch=k` so each upstream partition can @@ -869,12 +998,30 @@ impl LsmFtsSearchPlanner { async fn build_source_plan( &self, source: &LsmDataSource, - column: &str, + columns: &[String], query: &FullTextSearchQuery, limit: Option, projection: Option<&[String]>, - index_params: Option<&InvertedIndexParams>, + index_params: &HashMap<&str, InvertedIndexParams>, ) -> Result> { + // One column: bind every leaf to it, which is what lets a tree built + // from bare terms reach the right field. Several: the leaves already + // carry their bindings and rebinding would collapse the query onto one + // column, so the dataset scanner gets the tree untouched. + let bind_column = match columns { + [column] => Some(column.as_str()), + _ => None, + }; + let bind = |query: &FullTextSearchQuery| -> Result { + let bound = match bind_column { + Some(column) => query.clone().with_column(column.to_string())?, + None => query.clone(), + }; + Ok(match limit { + Some(limit) => bound.limit(Some(limit as i64)), + None => bound.limit(None), + }) + }; match source { LsmDataSource::BaseTable { dataset } => { let mut scanner = dataset.scan(); @@ -887,13 +1034,7 @@ impl LsmFtsSearchPlanner { scanner.filter_expr(filter.clone()); scanner.prefilter(true); } - let mut bound_query = query.clone().with_column(column.to_string())?; - if let Some(limit) = limit { - bound_query = bound_query.limit(Some(limit as i64)); - } else { - bound_query = bound_query.limit(None); - } - scanner.full_text_search(bound_query)?; + scanner.full_text_search(bind(query)?)?; scanner.create_plan().await } LsmDataSource::SsTable { path, .. } => { @@ -914,13 +1055,7 @@ impl LsmFtsSearchPlanner { scanner.filter_expr(filter.clone()); scanner.prefilter(true); } - let mut bound_query = query.clone().with_column(column.to_string())?; - if let Some(limit) = limit { - bound_query = bound_query.limit(Some(limit as i64)); - } else { - bound_query = bound_query.limit(None); - } - scanner.full_text_search(bound_query)?; + scanner.full_text_search(bind(query)?)?; scanner.create_plan().await } LsmDataSource::ActiveMemTable { @@ -936,27 +1071,34 @@ impl LsmFtsSearchPlanner { // this query rather than contributing nothing: an empty arm is a // silently short answer, since those rows are present and do // match. `None` means there is nothing to index. - let index_store = - if active_source_can_execute_fts(source, column, document_granularity) { - index_store.clone() - } else { - match transient_fts_index_store( - batch_store, - index_store, - schema, - column, - document_granularity, - &self.pk_columns, - index_params, - )? { - Some(store) => store, - None => { - return self.empty_plan( - &self.canonical_fts_schema(projection, document_granularity), - ); - } + // + // One missing column sends every queried column through the + // transient store: the arm routes leaves to indexes from a + // single store, and re-indexing a maintained column costs one + // tokenize pass over a memtable that is already paying for the + // missing one. + let index_store = if columns.iter().all(|column| { + active_source_can_execute_fts(source, column, document_granularity) + }) { + index_store.clone() + } else { + match transient_fts_index_store( + batch_store, + index_store, + schema, + columns, + document_granularity, + &self.pk_columns, + index_params, + )? { + Some(store) => store, + None => { + return self.empty_plan( + &self.canonical_fts_schema(projection, document_granularity), + ); } - }; + } + }; validate_lsm_fts_query(query)?; let mut scanner = MemTableScanner::new(batch_store.clone(), index_store, schema.clone()); @@ -973,16 +1115,7 @@ impl LsmFtsSearchPlanner { if !self.pk_columns.is_empty() { scanner.with_pk_columns(self.pk_columns.clone()); } - // `MemTableScanner::full_text_search` now takes a structured - // `FullTextSearchQuery` (match/phrase); it rejects compound - // shapes the MemTable path can't model. - let mut bound_query = query.clone().with_column(column.to_string())?; - if let Some(limit) = limit { - bound_query = bound_query.limit(Some(limit as i64)); - } else { - bound_query = bound_query.limit(None); - } - scanner.full_text_search(bound_query)?; + scanner.full_text_search(bind(query)?)?; scanner.create_plan().await } } @@ -3157,36 +3290,261 @@ mod tests { ); } - /// A boolean spanning columns is one predicate over several fields, not a - /// union of per-column results, so it is refused rather than decomposed - /// into an answer to a different question. + /// Rows and their `_score`s from a cross-column plan over the standard + /// two-column fixture, in plan order. + async fn run_cross_column( + rows: &[(i32, &str, &str)], + indexed_columns: &[&str], + query: IndexFtsQuery, + limit: usize, + ) -> Vec<(i32, f32)> { + let schema = two_column_fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + for column in indexed_columns { + let field_id = match *column { + "title" => 1, + "body" => 2, + other => panic!("unknown fixture column '{other}'"), + }; + indexes.add_fts(format!("{column}_fts"), field_id, column.to_string()); + } + let active_batch = make_two_column_batch(&schema, rows); + let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, row_offset, Some(batch_position)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let plan = planner + .plan_search( + FullTextSearchQuery::new_query(query), + Some(limit), + Some(&["id".to_string()]), + ) + .await + .expect("a cross-column predicate must plan"); + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + batches + .iter() + .flat_map(|batch| { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = batch + .column_by_name(SCORE_COLUMN) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|row| (ids.value(row), scores.value(row))) + .collect::>() + }) + .collect() + } + + fn match_leaf(terms: &str, column: &str) -> IndexFtsQuery { + use lance_index::scalar::inverted::query::MatchQuery; + IndexFtsQuery::Match( + MatchQuery::new(terms.to_string()).with_column(Some(column.to_string())), + ) + } + + /// A MUST spanning columns is an intersection: the row has to match in + /// *both* fields. Decomposing it into per-column arms and unioning them — + /// the way a multi-match decomposes — would return rows matching in only + /// one, which is a different query. #[tokio::test] - async fn cross_column_boolean_is_refused() { + async fn cross_column_boolean_must_intersects_the_columns() { + use lance_index::scalar::inverted::query::{BooleanQuery, Occur}; + + let results = run_cross_column( + &[ + (1, "alpha title", "beta body"), + (2, "alpha title", "unrelated body"), + (3, "unrelated title", "beta body"), + (4, "unrelated title", "unrelated body"), + ], + &["title", "body"], + IndexFtsQuery::Boolean(BooleanQuery::new(vec![ + (Occur::Must, match_leaf("alpha", "title")), + (Occur::Must, match_leaf("beta", "body")), + ])), + 10, + ) + .await; + let ids: Vec = results.iter().map(|(id, _)| *id).collect(); + assert_eq!( + ids, + vec![1], + "only the row matching in both columns satisfies a cross-column MUST" + ); + } + + /// The MUST_NOT clause reaches across columns too: a row matching the + /// excluded term in the *other* field is dropped. This is the hazard the + /// sophon-side fragment descent used to invert — arming only the MUST_NOT + /// clause returns exactly the rows the caller asked to leave out. + #[tokio::test] + async fn cross_column_boolean_must_not_excludes_across_columns() { + use lance_index::scalar::inverted::query::{BooleanQuery, Occur}; + + let results = run_cross_column( + &[ + (1, "alpha title", "beta body"), + (2, "alpha title", "unrelated body"), + ], + &["title", "body"], + IndexFtsQuery::Boolean(BooleanQuery::new(vec![ + (Occur::Must, match_leaf("alpha", "title")), + (Occur::MustNot, match_leaf("beta", "body")), + ])), + 10, + ) + .await; + let ids: Vec = results.iter().map(|(id, _)| *id).collect(); + assert_eq!(ids, vec![2], "id=1 carries the excluded term in `body`"); + } + + /// A boost demotes rather than drops: the row matching the negative clause + /// in the other column still comes back, ranked below the one that does + /// not. Dropping it would be a boolean MUST_NOT, a different query. + #[tokio::test] + async fn cross_column_boost_demotes_rather_than_drops() { + use lance_index::scalar::inverted::query::BoostQuery; + + let results = run_cross_column( + &[ + (1, "alpha title", "beta body"), + (2, "alpha title", "unrelated body"), + ], + &["title", "body"], + IndexFtsQuery::Boost(BoostQuery::new( + match_leaf("alpha", "title"), + match_leaf("beta", "body"), + Some(1.0), + )), + 10, + ) + .await; + let ids: Vec = results.iter().map(|(id, _)| *id).collect(); + assert_eq!( + ids, + vec![2, 1], + "both rows come back, and the demoted one ranks last" + ); + let demoted = results.iter().find(|(id, _)| *id == 1).unwrap().1; + let kept = results.iter().find(|(id, _)| *id == 2).unwrap().1; + assert!( + demoted < kept, + "the negative clause must subtract from the score: {demoted} vs {kept}" + ); + } + + /// A column outside the maintained set still contributes: the transient + /// store covers every queried column, so the clause on `body` is evaluated + /// rather than silently matching nothing — which on a MUST would empty the + /// whole result. + #[tokio::test] + async fn cross_column_builds_transient_indexes_for_unmaintained_columns() { + use lance_index::scalar::inverted::query::{BooleanQuery, Occur}; + + let results = run_cross_column( + &[ + (1, "alpha title", "beta body"), + (2, "alpha title", "unrelated body"), + ], + // `body` has no maintained in-memory index. + &["title"], + IndexFtsQuery::Boolean(BooleanQuery::new(vec![ + (Occur::Must, match_leaf("alpha", "title")), + (Occur::Must, match_leaf("beta", "body")), + ])), + 10, + ) + .await; + let ids: Vec = results.iter().map(|(id, _)| *id).collect(); + assert_eq!( + ids, + vec![1], + "the unmaintained column must still be searched" + ); + } + + /// A leaf naming no column cannot be routed: once the tree spans fields + /// there is no single index to fall back to, and binding it to one of the + /// others would silently answer a different query. Refused at planning. + #[tokio::test] + async fn cross_column_unbound_leaf_is_refused() { use lance_index::scalar::inverted::query::{BooleanQuery, MatchQuery, Occur}; let schema = two_column_fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("title_fts".to_string(), 1, "title".to_string()); + indexes.add_fts("body_fts".to_string(), 2, "body".to_string()); + let active_batch = make_two_column_batch(&schema, &[(1, "alpha title", "beta body")]); + let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, row_offset, Some(batch_position)) + .unwrap(); + let indexes = Arc::new(indexes); + let tmp = tempfile::tempdir().unwrap(); let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); - let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); - - let leaf = |terms: &str, column: &str| { - IndexFtsQuery::Match( - MatchQuery::new(terms.to_string()).with_column(Some(column.to_string())), - ) - }; let query = FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new(vec![ - (Occur::Must, leaf("alpha", "title")), - (Occur::Must, leaf("beta", "body")), + (Occur::Must, match_leaf("alpha", "title")), + ( + Occur::Must, + IndexFtsQuery::Match(MatchQuery::new("beta".to_string())), + ), ]))); let err = planner .plan_search(query, Some(10), Some(&["id".to_string()])) .await .unwrap_err(); assert!( - err.to_string().contains("across several columns"), - "unexpected error for a cross-column boolean: {err}" + err.to_string().contains("name the columns to search"), + "unexpected error for an unbound leaf: {err}" ); } diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index f53fea5a461..a8dec9ff753 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -790,6 +790,32 @@ impl CompoundQueryExec { self.base_scorer.as_ref() } + /// Re-cut this scorer at `limit`, keeping its segment selection, prepared + /// scorers and masks. + /// + /// The FTS top-k lives in `FtsSearchParams`, not in an enclosing fetch + /// node, so a caller that needs more candidates than the user asked for — + /// over-fetching to survive a later dedup, say — has no other way to raise + /// it. Everything that decides *which* rows are eligible is carried over + /// untouched, so this widens the cut without widening the domain. + pub fn with_limit(&self, limit: usize) -> Self { + let mut params = self.params.clone(); + params.limit = Some(limit); + Self { + dataset: self.dataset.clone(), + query: self.query.clone(), + tokenized_query: self.tokenized_query.clone(), + params, + prefilter_source: self.prefilter_source.clone(), + base_scorer: self.base_scorer.clone(), + prepared_match: self.prepared_match.clone(), + segment_selection: self.segment_selection.clone(), + external_mask: self.external_mask.clone(), + properties: self.properties.clone(), + metrics: ExecutionPlanMetricsSet::new(), + } + } + /// See [`MatchQueryExec::explicit_segment_uuids`]. pub fn explicit_segment_uuids(&self) -> Option> { self.segment_selection.explicit_segment_uuids() @@ -1828,6 +1854,30 @@ impl CrossColumnCompoundQueryExec { pub fn prefilter_source(&self) -> &PreFilterSource { &self.prefilter_source } + + /// Re-cut this scorer at `limit`, keeping its segment selection, prepared + /// scorers and masks. + /// + /// The FTS top-k lives in `FtsSearchParams`, not in an enclosing fetch + /// node, so a caller that needs more candidates than the user asked for — + /// over-fetching to survive a later dedup, say — has no other way to raise + /// it. Everything that decides *which* rows are eligible is carried over + /// untouched, so this widens the cut without widening the domain. + pub fn with_limit(&self, limit: usize) -> Self { + let mut params = self.params.clone(); + params.limit = Some(limit); + Self { + dataset: self.dataset.clone(), + query: self.query.clone(), + tokenized_query: self.tokenized_query.clone(), + params, + prefilter_source: self.prefilter_source.clone(), + columns: self.columns.clone(), + external_mask: self.external_mask.clone(), + properties: self.properties.clone(), + metrics: ExecutionPlanMetricsSet::new(), + } + } } impl DisplayAs for CrossColumnCompoundQueryExec {