diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index c2a446d48dd..428ac2fc415 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -48,7 +48,7 @@ #![allow(clippy::type_complexity)] use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; @@ -66,7 +66,7 @@ use lance_index::scalar::inverted::query::{FtsQuery, Operator, Tokens}; use lance_index::scalar::inverted::tokenizer::document_tokenizer::{DocType, LanceTokenizer}; use lance_index::scalar::inverted::{DocSet, MemBM25Scorer, Scorer, TokenSet}; use lance_table::format::IndexMetadata; -use lance_tokenizer::TokenStream; +use lance_tokenizer::{SimpleTokenizer, TextAnalyzer, TokenStream}; use rayon::prelude::*; use rustc_hash::{FxHashMap, FxHashSet}; @@ -171,6 +171,9 @@ pub enum FtsQueryExpr { prefix_length: u32, /// Maximum number of terms to expand to. max_expansions: usize, + /// `Or` unions every expansion; `And` requires a match from each query + /// term's expansions. + operator: Operator, /// Boost factor applied to the score (default 1.0). boost: f32, }, @@ -308,6 +311,7 @@ impl FtsQueryExpr { } } + /// Fuzzy leaf with AUTO distance, OR semantics, and default options. pub fn fuzzy(query: impl Into) -> Self { Self::Fuzzy { column: None, @@ -315,10 +319,12 @@ impl FtsQueryExpr { fuzziness: None, prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, + operator: Operator::Or, boost: 1.0, } } + /// Fuzzy leaf with a fixed edit distance and OR semantics. pub fn fuzzy_with_distance(query: impl Into, fuzziness: u32) -> Self { Self::Fuzzy { column: None, @@ -326,10 +332,12 @@ impl FtsQueryExpr { fuzziness: Some(fuzziness), prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, + operator: Operator::Or, boost: 1.0, } } + /// Fuzzy leaf with OR semantics; `fuzziness: None` selects AUTO distance. pub fn fuzzy_with_options( query: impl Into, fuzziness: Option, @@ -342,6 +350,7 @@ impl FtsQueryExpr { fuzziness, prefix_length, max_expansions, + operator: Operator::Or, boost: 1.0, } } @@ -366,6 +375,43 @@ impl FtsQueryExpr { } } + /// Set how a `Match` or `Fuzzy` leaf combines its terms; compound shapes + /// carry no operator and are returned unchanged. + pub fn with_operator(self, operator: Operator) -> Self { + match self { + Self::Match { + column, + query, + boost, + .. + } => Self::Match { + column, + query, + operator, + boost, + }, + Self::Fuzzy { + column, + query, + fuzziness, + prefix_length, + max_expansions, + boost, + .. + } => Self::Fuzzy { + column, + query, + fuzziness, + prefix_length, + max_expansions, + operator, + boost, + }, + other => other, + } + } + + /// Set a leaf's boost; compound shapes are returned unchanged. pub fn with_boost(self, boost: f32) -> Self { match self { Self::Match { @@ -396,6 +442,7 @@ impl FtsQueryExpr { fuzziness, prefix_length, max_expansions, + operator, .. } => Self::Fuzzy { column, @@ -403,6 +450,7 @@ impl FtsQueryExpr { fuzziness, prefix_length, max_expansions, + operator, boost, }, // Compound nodes don't carry a top-level boost field today. @@ -440,6 +488,7 @@ impl FtsQueryExpr { fuzziness, prefix_length, max_expansions, + operator, boost, .. } => Self::Fuzzy { @@ -448,6 +497,7 @@ impl FtsQueryExpr { fuzziness, prefix_length, max_expansions, + operator, boost, }, other @ (Self::Boolean { .. } | Self::Boost { .. } | Self::MultiMatch { .. }) => other, @@ -2154,9 +2204,8 @@ impl FtsMemIndex { } } - /// Expand a term to fuzzy matches within the specified edit distance. - /// Returns `(matched_term, distance)` pairs sorted by distance, capped - /// at `max_expansions`. + /// The `max_expansions` lexicographically smallest terms within + /// `max_distance` of `term`, as `(term, distance)` pairs in term order. pub fn expand_fuzzy( &self, term: &str, @@ -2175,8 +2224,16 @@ impl FtsMemIndex { max_expansions: usize, ) -> Vec { let st = self.state.load_full(); - let tokens = self.tokenize_for_search(query); - self.search_fuzzy_tokens(&st, &tokens, fuzziness, 0, max_expansions, true) + let tokens = self.tokenize_fuzzy_query(query, fuzziness); + self.search_fuzzy_tokens( + &st, + &tokens, + Operator::Or, + fuzziness, + 0, + max_expansions, + true, + ) } /// BM25 OR-search over the query tokens, scored with one corpus-wide @@ -2342,6 +2399,7 @@ impl FtsMemIndex { results } + /// Conjunction over the query's position groups; see [`Self::search_groups_and`]. fn search_grouped_and( &self, st: &IndexState, @@ -2349,9 +2407,28 @@ impl FtsMemIndex { limit: Option, include_tail: bool, tail_skip: bool, + ) -> Vec { + self.search_groups_and( + st, + query_position_groups(query_tokens), + limit, + include_tail, + tail_skip, + ) + } + + /// Conjunction over term groups: a document must match at least one term of + /// every group, and its score is the sum of the per-group OR scores. + fn search_groups_and( + &self, + st: &IndexState, + groups: Vec>, + limit: Option, + include_tail: bool, + tail_skip: bool, ) -> Vec { let mut result_map: Option> = None; - for group in query_position_groups(query_tokens) { + for group in groups { let group_results = self.search_match_strings(st, &group, Operator::Or, None, include_tail, tail_skip); let group_map = group_results @@ -2515,10 +2592,15 @@ impl FtsMemIndex { results } + /// Expand each token within its edit distance, then combine per `operator`: + /// `Or` is one BM25 OR over all expansions; `And` requires a term from each + /// token's expansions and scores like [`Self::search_groups_and`]. + #[allow(clippy::too_many_arguments)] fn search_fuzzy_tokens( &self, st: &IndexState, tokens: &[String], + operator: Operator, fuzziness: Option, prefix_length: u32, max_expansions: usize, @@ -2527,31 +2609,45 @@ impl FtsMemIndex { if tokens.is_empty() { return Vec::new(); } - let mut expanded: Vec = Vec::new(); - let mut seen: HashSet = HashSet::new(); + // One leaf-wide budget spent in token order, matching the committed + // index's `expand_fuzzy_tokens`. + let mut remaining = max_expansions; + let mut groups: Vec> = Vec::with_capacity(tokens.len()); for tok in tokens { let max_dist = fuzziness.unwrap_or_else(|| auto_fuzziness(tok)); - for (matched, _) in self.expand_fuzzy_term( - st, - tok, - max_dist, - prefix_length, - max_expansions, - include_tail, - ) { - if seen.insert(matched.clone()) { - expanded.push(matched); - } + let group: Vec = self + .expand_fuzzy_term(st, tok, max_dist, prefix_length, remaining, include_tail) + .into_iter() + .map(|(matched, _)| matched) + .collect(); + remaining -= group.len(); + // A token with no expansions (none in range or budget spent) makes + // the conjunction unsatisfiable. + if group.is_empty() && operator == Operator::And { + return Vec::new(); } + groups.push(group); } - if expanded.is_empty() { - return Vec::new(); + match operator { + Operator::And => self.search_groups_and(st, groups, None, include_tail, true), + Operator::Or => { + let mut seen: HashSet = HashSet::new(); + let expanded: Vec = groups + .into_iter() + .flatten() + .filter(|term| seen.insert(term.clone())) + .collect(); + if expanded.is_empty() { + return Vec::new(); + } + self.search_match_strings(st, &expanded, Operator::Or, None, include_tail, true) + } } - self.search_match_strings(st, &expanded, Operator::Or, None, include_tail, true) } - /// Expand `term` against the term dictionaries of every partition (and the - /// visible tail, when `include_tail`). + /// Expand `term` against every partition's dictionary (and the visible tail, + /// when `include_tail`), keeping the `max_expansions` lexicographically + /// smallest matches, as the committed index does. fn expand_fuzzy_term( &self, st: &IndexState, @@ -2561,6 +2657,9 @@ impl FtsMemIndex { max_expansions: usize, include_tail: bool, ) -> Vec<(String, u32)> { + if max_expansions == 0 { + return Vec::new(); + } let tail_snap = st.tail.snapshot(); if max_distance == 0 { let in_tail = include_tail @@ -2577,8 +2676,7 @@ impl FtsMemIndex { Vec::new() }; } - let mut matches: Vec<(String, u32)> = Vec::new(); - let mut seen: HashSet = HashSet::new(); + let mut matches: BTreeMap = BTreeMap::new(); let prefix = char_prefix(term, prefix_length); for entry in st.tail.terms.iter() { if !include_tail || !has_visible_chunk(&entry.value().load(), tail_snap.visible_count) { @@ -2589,8 +2687,8 @@ impl FtsMemIndex { continue; } let dist = levenshtein_distance(term, key); - if dist <= max_distance && seen.insert(key.to_string()) { - matches.push((key.to_string(), dist)); + if dist <= max_distance { + matches.insert(key.to_string(), dist); } } for p in st.partitions.iter() { @@ -2599,14 +2697,12 @@ impl FtsMemIndex { continue; } let dist = levenshtein_distance(term, key.as_ref()); - if dist <= max_distance && seen.insert(key.to_string()) { - matches.push((key.to_string(), dist)); + if dist <= max_distance { + matches.insert(key.to_string(), dist); } } } - matches.sort_by_key(|(_, d)| *d); - matches.truncate(max_expansions); - matches + matches.into_iter().take(max_expansions).collect() } /// Execute a query expression and return matching documents with scores. @@ -2686,13 +2782,15 @@ impl FtsMemIndex { fuzziness, prefix_length, max_expansions, + operator, boost, .. } => { - let tokens = self.tokenize_for_search(query); + let tokens = self.tokenize_fuzzy_query(query, *fuzziness); let mut results = self.search_fuzzy_tokens( st, &tokens, + *operator, *fuzziness, *prefix_length, *max_expansions, @@ -2754,6 +2852,22 @@ impl FtsMemIndex { query_tokens_to_vec(&self.analyze_for_search(text)) } + /// Tokenize like the committed index's `tokenizer_for_match_query`: an + /// explicit positive distance uses a bare `SimpleTokenizer` (no stop words + /// or stemming); otherwise the index analyzer. + fn tokenize_fuzzy_query(&self, text: &str, fuzziness: Option) -> Vec { + if !matches!(fuzziness, Some(distance) if distance > 0) { + return self.tokenize_for_search(text); + } + let mut analyzer = TextAnalyzer::from(SimpleTokenizer::default()); + let mut stream = analyzer.token_stream(text); + let mut tokens = Vec::new(); + while let Some(token) = stream.next() { + tokens.push(token.text.clone()); + } + tokens + } + fn analyze_for_search(&self, text: &str) -> Tokens { let mut tok = PooledTokenizer::new(&self.tokenizer_pool); let mut stream = tok.get_mut().token_stream_for_search(text); @@ -6465,6 +6579,63 @@ mod tests { assert!(!entries.is_empty()); } + /// `alpho beta` at distance 1 expands to {alpha, alpho} and {beta, zeta}; + /// AND keeps only rows with a term from each group. + #[test] + fn test_search_fuzzy_and_requires_every_term() { + let schema = create_test_schema(); + let index = FtsMemIndex::new(1, "description".to_string()); + index.insert(&create_fuzzy_test_batch(&schema), 0).unwrap(); + let st = index.state.load_full(); + let tokens = ["alpho".to_string(), "beta".to_string()]; + let rows = |operator: Operator| { + let mut rows: Vec = index + .search_fuzzy_tokens(&st, &tokens, operator, Some(1), 0, 50, true) + .into_iter() + .map(|entry| entry.row_position) + .collect(); + rows.sort_unstable(); + rows + }; + assert_eq!(rows(Operator::And), vec![0, 1]); + assert_eq!(rows(Operator::Or), vec![0, 1, 2, 3]); + + // `alpho` spends the whole two-term budget, leaving `beta` none. + assert!( + index + .search_fuzzy_tokens(&st, &tokens, Operator::And, Some(1), 0, 2, true) + .is_empty() + ); + + // A term that expands to nothing empties the conjunction. + let unmatched = ["alpho".to_string(), "xyz".to_string()]; + assert!( + index + .search_fuzzy_tokens(&st, &unmatched, Operator::And, Some(1), 0, 50, true) + .is_empty() + ); + } + + /// AUTO leaves `xy` exact, but it still needs budget: `alpha` spends the + /// only slot, so the conjunction fails. + #[test] + fn test_search_fuzzy_and_zero_distance_term_respects_exhausted_budget() { + let schema = create_test_schema(); + let index = FtsMemIndex::new(1, "description".to_string()); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![0])), + Arc::new(StringArray::from(vec!["alpha xy"])), + ], + ) + .unwrap(); + index.insert(&batch, 0).unwrap(); + let query = + FtsQueryExpr::fuzzy_with_options("alpha xy", None, 0, 1).with_operator(Operator::And); + assert!(index.search_query(&query).is_empty()); + } + #[test] fn test_search_fuzzy_no_match() { let schema = create_test_schema(); @@ -6520,10 +6691,17 @@ mod tests { ); } + /// `prefix_length` counts characters, not bytes. #[test] fn test_search_query_fuzzy_prefix_length_uses_char_boundaries() { let schema = create_test_schema(); - let index = FtsMemIndex::new(1, "description".to_string()); + // Explicit fuzziness skips the index analyzer, so the stored term must + // keep its `é` for the query's raw `é` prefix to reach it. + let index = FtsMemIndex::with_params( + 1, + "description".to_string(), + InvertedIndexParams::default().ascii_folding(false), + ); let batch = RecordBatch::try_new( schema, 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 514d5fadfe1..1abfb795e74 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -454,20 +454,13 @@ fn to_local_expr(query: &IndexFtsQuery) -> Result { // `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_operator(m.operator) .with_boost(m.boost), }, m.column.as_ref(), @@ -1976,6 +1969,7 @@ mod tests { ); } + /// Each index query shape maps to its local form; a leaf without a column is rejected. #[test] fn local_fts_query_maps_leaf_shapes_and_rejects_the_rest() { use lance_index::scalar::inverted::query::{ @@ -2073,9 +2067,17 @@ mod tests { .with_fuzziness(Some(1)) .with_column(Some("text".to_string())), )); + let local = local_fts_query(fuzzy_and, None).unwrap(); assert!( - local_fts_query(fuzzy_and, None).is_err(), - "fuzzy AND cannot be represented by the local memtable query" + matches!( + local.expr, + FtsQueryExpr::Fuzzy { + operator: Operator::And, + fuzziness: Some(1), + .. + } + ), + "fuzzy AND maps to a conjunctive Fuzzy leaf" ); // Phrase -> local Phrase. 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 d22d64acfac..2663cdac7dd 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -30,6 +30,10 @@ //! copy lives in an SSTable; cross-tier deduplication is //! left to the caller in local mode. //! +//! Fuzzy expansion is also per source: each spends `max_expansions` over its +//! own vocabulary, not the merged one, so under a binding cap a source can +//! match rows a merged index would not, and results shift as rows change tiers. +//! //! Everything here is contained in the `mem_wal` module — it reuses the //! existing per-source FTS read paths (`scanner.full_text_search` for //! base/SSTable Lance datasets, `MemTableScanner` for the active @@ -50,7 +54,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, Operator}; +use lance_index::scalar::inverted::query::FtsQuery as IndexFtsQuery; use lance_index::scalar::inverted::{DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity}; use tracing::instrument; @@ -66,6 +70,7 @@ use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; use crate::dataset::mem_wal::index::FtsMemQuery; use crate::dataset::mem_wal::memtable::scanner::{MemTableScanner, local_fts_query}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; +use crate::dataset::scanner::apply_dataset_planner_auto_fuzziness_compatibility_gate; use crate::index::scalar::inverted::{ indexed_fts_document_granularities, indexed_fts_index_params, resolve_fts_field, }; @@ -224,22 +229,12 @@ fn validate_source_document_granularities( } /// Reject the query shapes the active memtable arm cannot evaluate. -/// -/// Only one remains: a fuzzy Match cannot also require every term, because the -/// fuzzy path expands each term independently and unions the expansions. A -/// multi-match is checked leaf by leaf under the same rule. fn validate_lsm_fts_query(query: &FullTextSearchQuery) -> Result<()> { fn visit(query: &IndexFtsQuery) -> Result<()> { match query { - IndexFtsQuery::Match(m) => { - if m.fuzziness != Some(0) && m.operator != Operator::Or { - return Err(Error::not_supported( - "LSM fuzzy full-text search only supports OR match operators".to_string(), - )); - } + IndexFtsQuery::Match(_) | IndexFtsQuery::Phrase(_) | IndexFtsQuery::MultiMatch(_) => { Ok(()) } - IndexFtsQuery::Phrase(_) => Ok(()), IndexFtsQuery::Boost(b) => { visit(&b.positive)?; visit(&b.negative) @@ -250,12 +245,6 @@ fn validate_lsm_fts_query(query: &FullTextSearchQuery) -> Result<()> { } Ok(()) } - IndexFtsQuery::MultiMatch(m) => { - for leaf in &m.match_queries { - visit(&IndexFtsQuery::Match(leaf.clone()))?; - } - Ok(()) - } // BM25F needs corpus-wide field statistics that the memtable's // in-memory index does not maintain. IndexFtsQuery::CombinedFields(_) => Err(Error::not_supported( @@ -876,6 +865,9 @@ impl LsmFtsSearchPlanner { resolve_fts_field(&schema, column, document_granularity)?; } set_query_document_granularity(&mut query.query, document_granularity); + // The dataset scanner runs AUTO as exact on the base and SSTable arms; + // do the same for the memtable arm. + apply_dataset_planner_auto_fuzziness_compatibility_gate(&mut query.query); if sources.iter().any(|source| { columns .iter() @@ -2775,8 +2767,10 @@ mod tests { ); } + /// `lance memwel` at distance 1 with AND keeps the row with both words and + /// drops the one with only `lance`. #[tokio::test] - async fn fuzzy_and_query_is_rejected_when_active_memtable_is_present() { + async fn fuzzy_and_query_is_served_by_the_active_memtable() { use lance_index::scalar::inverted::query::{ FtsQuery as IndexFtsQuery, MatchQuery, Operator, }; @@ -2785,7 +2779,7 @@ mod tests { let batch_store = Arc::new(BatchStore::with_capacity(16)); let mut indexes = IndexStore::new(); indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); - let active_batch = make_batch(&schema, &[1], &["lance memwal"]); + let active_batch = make_batch(&schema, &[1, 2], &["lance memwal", "lance other"]); let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap(); indexes .insert_with_batch_position(&active_batch, row_offset, Some(batch_position)) @@ -2809,25 +2803,135 @@ mod tests { ); let planner = LsmFtsSearchPlanner::new(collector, vec![], schema); let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match( - MatchQuery::new("lance memwal".to_string()) + MatchQuery::new("lance memwel".to_string()) .with_operator(Operator::And) .with_fuzziness(Some(1)), )); - let err = planner + let plan = planner .plan_search( query.with_column("text".to_string()).unwrap(), Some(10), None, ) .await - .expect_err("fuzzy AND should be rejected consistently"); - assert!( - err.to_string().contains("fuzzy full-text search"), - "unexpected error for fuzzy AND query: {err}" + .expect("fuzzy AND is served by the active memtable"); + let stream = plan + .execute(0, datafusion::prelude::SessionContext::new().task_ctx()) + .unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + ids, + vec![1], + "only the row matching every fuzzy term survives" ); } + /// Fuzzy AND treats an indexed base row and an identical active row alike. + #[rstest::rstest] + #[case::fuzzy_and_matches_both_tiers("alphb beta", Some(1), 50, vec![1, 2])] + #[case::auto_fuzziness_runs_exact("alphb beta", None, 50, vec![])] + #[case::explicit_fuzzy_keeps_stop_words("alpha the", Some(1), 50, vec![])] + #[case::expansion_budget_spans_the_leaf("alpha beta", Some(1), 1, vec![])] + #[tokio::test] + async fn fuzzy_and_matches_identically_on_base_and_active( + #[case] terms: &str, + #[case] fuzziness: Option, + #[case] max_expansions: usize, + #[case] expected_ids: Vec, + ) { + use crate::index::DatasetIndexExt; + use lance_index::IndexType; + use lance_index::scalar::inverted::query::Operator; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = + write_dataset(&base_uri, vec![make_batch(&schema, &[1], &["alpha beta"])]).await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + + 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("text_fts".to_string(), 1, "text".to_string()); + let active_batch = make_batch(&schema, &[2], &["alpha beta"]); + batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, 0, Some(0)) + .unwrap(); + let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: Arc::new(indexes), + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new(terms.to_string()) + .with_fuzziness(fuzziness) + .with_operator(Operator::And) + .with_max_expansions(max_expansions), + )); + let plan = planner + .plan_search( + query.with_column("text".to_string()).unwrap(), + Some(10), + None, + ) + .await + .unwrap(); + let batches: Vec = plan + .execute(0, datafusion::prelude::SessionContext::new().task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|b| { + b["id"] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + ids.sort_unstable(); + assert_eq!(ids, expected_ids); + } + #[tokio::test] async fn base_only_boolean_query_uses_dataset_scanner_support() { use crate::index::DatasetIndexExt; diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b72574da974..e058b76f1f3 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -466,7 +466,7 @@ fn normalize_fts_zero_boosts(query: &mut FtsQuery) { /// unindexed rows. AUTO activation is deferred until OSS-2105 lands that /// current-row preparation atomically. Until then, recursively rewrite AUTO to /// exact while preserving explicit positive fuzziness. -fn apply_dataset_planner_auto_fuzziness_compatibility_gate(query: &mut FtsQuery) { +pub(crate) fn apply_dataset_planner_auto_fuzziness_compatibility_gate(query: &mut FtsQuery) { match query { FtsQuery::Match(query) => { query.fuzziness.get_or_insert(0); diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index 537a5e6025d..1bbcb089663 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -34,7 +34,7 @@ pub mod utils; pub use filter::LanceFilterExec; pub use knn::{ ANNIvfBatchExec, ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec, - QUERY_INDEX_COL, + MultivectorScoringExec, QUERY_INDEX_COL, }; pub use lance_datafusion::planner::Planner; pub use lance_index::scalar::expression::FilterPlan;