diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 9d65220aa53..8f06847da8b 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -21,9 +21,9 @@ use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::inverted::{ DocumentGranularity, query::{ - BooleanQuery as FtsBooleanQuery, BoostQuery as FtsBoostQuery, FtsQuery, - MatchQuery as FtsMatchQuery, MultiMatchQuery as FtsMultiMatchQuery, Occur as FtsOccur, - PhraseQuery as FtsPhraseQuery, + BooleanQuery as FtsBooleanQuery, BoostQuery as FtsBoostQuery, + CombinedFieldsQuery as FtsCombinedFieldsQuery, FtsQuery, MatchQuery as FtsMatchQuery, + MultiMatchQuery as FtsMultiMatchQuery, Occur as FtsOccur, PhraseQuery as FtsPhraseQuery, }, }; use lance_io::ffi::to_ffi_arrow_array_stream; @@ -174,6 +174,33 @@ pub(crate) fn build_full_text_search_query<'a>( Ok(FtsQuery::MultiMatch(query)) } + "COMBINED_FIELDS" => { + let query_text = env.get_string_from_method(&java_obj, "getQueryText")?; + let columns: Vec = + import_vec_from_method(env, &java_obj, "getColumns", |env, elem| { + let jstr = JString::from(elem); + let value: String = env.get_string(&jstr)?.into(); + Ok(value) + })?; + + let boosts: Option> = + env.get_optional_from_method(&java_obj, "getBoosts", |env, list_obj| { + import_vec_to_rust(env, &list_obj, |env, elem| { + env.get_f32_from_method(&elem, "floatValue") + }) + })?; + let operator = env.get_fts_operator_from_method(&java_obj)?; + + // Column uniqueness and boost (>= 1) validation live in the Rust core; + // `?` surfaces those errors across the JNI boundary. + let mut query = FtsCombinedFieldsQuery::try_new(query_text, columns)?; + if let Some(boosts) = boosts { + query = query.try_with_boosts(boosts)?; + } + query = query.with_operator(operator); + + Ok(FtsQuery::CombinedFields(query)) + } "BOOST" => { let positive_obj = env .call_method( diff --git a/java/src/main/java/org/lance/ipc/FullTextQuery.java b/java/src/main/java/org/lance/ipc/FullTextQuery.java index 73badab9a49..193b3452ea0 100755 --- a/java/src/main/java/org/lance/ipc/FullTextQuery.java +++ b/java/src/main/java/org/lance/ipc/FullTextQuery.java @@ -36,6 +36,7 @@ public enum Type { MATCH_PHRASE, BOOST, MULTI_MATCH, + COMBINED_FIELDS, BOOLEAN } @@ -156,6 +157,15 @@ public static FullTextQuery multiMatch( return new MultiMatchQuery(queryText, columns, boosts, operator); } + public static FullTextQuery combinedFields(String queryText, List columns) { + return combinedFields(queryText, columns, null, Operator.OR); + } + + public static FullTextQuery combinedFields( + String queryText, List columns, List boosts, Operator operator) { + return new CombinedFieldsQuery(queryText, columns, boosts, operator); + } + public static FullTextQuery boost(FullTextQuery positive, FullTextQuery negative) { return boost(positive, negative, 0.5f); } @@ -431,6 +441,93 @@ public String toString() { } } + /** + * Combined-fields (BM25F) query across multiple columns. + * + *

Unlike {@link MultiMatchQuery}, which scores each column independently and keeps the best + * field, this query treats the target columns as a single virtual field so that term statistics + * are blended across fields (Lucene {@code CombinedFieldQuery}). A term that is rare in one field + * but common in another then scores consistently, and a single query term can match across fields + * (for example a first name in one column and a last name in another). + * + *

Target columns must be unique and share the same tokenizer/index configuration. When {@code + * boosts} is given, its length must equal the number of columns and each per-column weight must + * be finite and {@code >= 1} (fractional weights allowed); when {@code null}, every column + * defaults to {@code 1.0}. A {@code null} operator defaults to {@link Operator#OR}. These + * constraints are validated in the Rust core and surface as an exception when the query runs. + */ + public static final class CombinedFieldsQuery extends FullTextQuery { + private final String queryText; + private final List columns; + private final Optional> boosts; + private final Operator operator; + + CombinedFieldsQuery( + String queryText, List columns, List boosts, Operator operator) { + Preconditions.checkArgument( + queryText != null && !queryText.isEmpty(), "queryText must not be null or empty"); + Preconditions.checkArgument( + columns != null && !columns.isEmpty(), "columns must not be null or empty"); + + this.queryText = queryText; + this.columns = + Collections.unmodifiableList(new java.util.ArrayList<>(Objects.requireNonNull(columns))); + this.boosts = + boosts == null + ? Optional.empty() + : Optional.of(Collections.unmodifiableList(new java.util.ArrayList<>(boosts))); + this.operator = operator == null ? Operator.OR : operator; + } + + @Override + public Type getType() { + return Type.COMBINED_FIELDS; + } + + public String getQueryText() { + return queryText; + } + + public List getColumns() { + return columns; + } + + public Optional> getBoosts() { + return boosts; + } + + public Operator getOperator() { + return operator; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof CombinedFieldsQuery)) return false; + CombinedFieldsQuery other = (CombinedFieldsQuery) o; + return operator == other.operator + && Objects.equals(queryText, other.queryText) + && Objects.equals(columns, other.columns) + && Objects.equals(boosts, other.boosts); + } + + @Override + public int hashCode() { + return Objects.hash(queryText, columns, operator, boosts); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("type", getType()) + .add("queryText", queryText) + .add("columns", columns) + .add("boosts", boosts) + .add("operator", operator) + .toString(); + } + } + /** Boost query combining positive and negative queries. */ public static final class BoostQuery extends FullTextQuery { private final FullTextQuery positive; diff --git a/java/src/test/java/org/lance/ipc/FullTextQueryTest.java b/java/src/test/java/org/lance/ipc/FullTextQueryTest.java index b84a19db4c8..7b6fe54784d 100755 --- a/java/src/test/java/org/lance/ipc/FullTextQueryTest.java +++ b/java/src/test/java/org/lance/ipc/FullTextQueryTest.java @@ -331,6 +331,101 @@ void testMultiMatchQueryInequalityDifferentBoosts() { assertFalse(a.equals(b), "MultiMatchQuery instances with different boosts must not be equal"); } + @Test + void testCombinedFieldsWithoutBoosts() { + FullTextQuery.CombinedFieldsQuery q = + (FullTextQuery.CombinedFieldsQuery) + FullTextQuery.combinedFields("hello", Arrays.asList("title", "body")); + + assertEquals(FullTextQuery.Type.COMBINED_FIELDS, q.getType()); + assertEquals("hello", q.getQueryText()); + assertEquals(Arrays.asList("title", "body"), q.getColumns()); + assertFalse(q.getBoosts().isPresent()); + assertEquals(FullTextQuery.Operator.OR, q.getOperator()); + } + + @Test + void testCombinedFieldsWithBoosts() { + FullTextQuery.CombinedFieldsQuery q = + (FullTextQuery.CombinedFieldsQuery) + FullTextQuery.combinedFields( + "hello", + Arrays.asList("title", "body"), + Arrays.asList(2.0f, 1.0f), + FullTextQuery.Operator.AND); + + assertEquals(FullTextQuery.Type.COMBINED_FIELDS, q.getType()); + assertTrue(q.getBoosts().isPresent()); + assertEquals(2, q.getBoosts().get().size()); + assertEquals(2.0f, q.getBoosts().get().get(0)); + assertEquals(1.0f, q.getBoosts().get().get(1)); + assertEquals(FullTextQuery.Operator.AND, q.getOperator()); + assertNotNull(q.toString()); + } + + @Test + void testCombinedFieldsQueryEquality() { + FullTextQuery a = FullTextQuery.combinedFields("hello", Arrays.asList("title", "body")); + FullTextQuery b = FullTextQuery.combinedFields("hello", Arrays.asList("title", "body")); + assertEquals(a, b, "Two CombinedFieldsQuery instances with the same fields must be equal"); + assertEquals(a.hashCode(), b.hashCode()); + } + + @Test + void testCombinedFieldsQueryInequalityDifferentOperator() { + FullTextQuery a = + FullTextQuery.combinedFields( + "hello", Arrays.asList("title", "body"), null, FullTextQuery.Operator.AND); + FullTextQuery b = + FullTextQuery.combinedFields( + "hello", Arrays.asList("title", "body"), null, FullTextQuery.Operator.OR); + assertFalse( + a.equals(b), "CombinedFieldsQuery instances with different operator must not be equal"); + } + + @Test + void testCombinedFieldsQueryInequalityDifferentBoosts() { + FullTextQuery a = + FullTextQuery.combinedFields( + "hello", + Arrays.asList("title", "body"), + Arrays.asList(2.0f, 1.0f), + FullTextQuery.Operator.OR); + FullTextQuery b = + FullTextQuery.combinedFields( + "hello", + Arrays.asList("title", "body"), + Arrays.asList(1.0f, 1.0f), + FullTextQuery.Operator.OR); + assertFalse( + a.equals(b), "CombinedFieldsQuery instances with different boosts must not be equal"); + } + + @Test + void testCombinedFieldsNotEqualMultiMatch() { + FullTextQuery combined = FullTextQuery.combinedFields("hello", Arrays.asList("title", "body")); + FullTextQuery multi = FullTextQuery.multiMatch("hello", Arrays.asList("title", "body")); + assertFalse( + combined.equals(multi), + "CombinedFieldsQuery and MultiMatchQuery must not be equal even with same fields"); + } + + @Test + void testCombinedFieldsBoostsDefensivelyCopied() { + java.util.List boosts = new java.util.ArrayList<>(Arrays.asList(2.0f, 1.0f)); + FullTextQuery.CombinedFieldsQuery q = + (FullTextQuery.CombinedFieldsQuery) + FullTextQuery.combinedFields( + "hello", Arrays.asList("title", "body"), boosts, FullTextQuery.Operator.OR); + + // Mutating the caller-provided list must not change the stored query. + boosts.set(0, 9.0f); + boosts.clear(); + + assertTrue(q.getBoosts().isPresent()); + assertEquals(Arrays.asList(2.0f, 1.0f), q.getBoosts().get()); + } + @Test void testDifferentSubtypesNotEqual() { FullTextQuery match = FullTextQuery.match("hello", "body"); diff --git a/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java b/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java index 06438d30706..e79b76a331f 100755 --- a/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java +++ b/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java @@ -49,6 +49,15 @@ class LanceScannerFullTextSearchTest { + private static final List DEFAULT_DOCUMENTS = + Arrays.asList("hello world", "hello lance", "other text"); + + private static ScalarIndexParams defaultIndexParams() { + return ScalarIndexParams.create( + "inverted", + "{\"base_tokenizer\":\"simple\",\"language\":\"English\",\"with_position\":true}"); + } + @Test void testMatchQuery() throws Exception { runFtsQuery( @@ -120,6 +129,55 @@ void testMultiMatch() throws Exception { runFtsQuery("memory://fts_java_multimatch", multiMatch, 3); } + @Test + void testCombinedFields() throws Exception { + // "hello" appears in doc or title of every row, so all 3 match. + FullTextQuery combined = FullTextQuery.combinedFields("hello", Arrays.asList("doc", "title")); + runFtsQuery("memory://fts_java_combined", combined, 3); + } + + @Test + void testCombinedFieldsWithBoosts() throws Exception { + FullTextQuery combined = + FullTextQuery.combinedFields( + "hello", + Arrays.asList("doc", "title"), + Arrays.asList(2.0f, 1.0f), + FullTextQuery.Operator.OR); + runFtsQuery("memory://fts_java_combined_boosts", combined, 3); + } + + @Test + void testCombinedFieldsInvalidBoostPropagates() throws Exception { + // Per-column weights must be >= 1. The check lives in the Rust core and must + // surface across the JNI boundary when the query runs. + FullTextQuery combined = + FullTextQuery.combinedFields( + "hello", + Arrays.asList("doc", "title"), + Arrays.asList(0.5f, 1.0f), + FullTextQuery.Operator.OR); + withIndexedDataset( + "memory://fts_java_combined_bad_boost", + dataset -> { + ScanOptions scanOptions = new ScanOptions.Builder().fullTextQuery(combined).build(); + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> { + try (LanceScanner scanner = dataset.newScan(scanOptions); + ArrowReader arrowReader = scanner.scanBatches()) { + while (arrowReader.loadNextBatch()) { + // Drain batches to force query execution. + } + } + }); + assertTrue( + ex.getMessage().contains("combined_fields boost for column 'doc'"), + "expected invalid-boost validation error, got: " + ex.getMessage()); + }); + } + @Test void testBooleanQuery() throws Exception { FullTextQuery.MatchQuery shouldMatch = @@ -139,16 +197,7 @@ void testBooleanQuery() throws Exception { } private void runFtsQuery(String uri, FullTextQuery query, long expectedTotal) throws Exception { - ScalarIndexParams indexParams = - ScalarIndexParams.create( - "inverted", - "{\"base_tokenizer\":\"simple\",\"language\":\"English\",\"with_position\":true}"); - runFtsQuery( - uri, - query, - expectedTotal, - Arrays.asList("hello world", "hello lance", "other text"), - indexParams); + runFtsQuery(uri, query, expectedTotal, DEFAULT_DOCUMENTS, defaultIndexParams()); } private void runFtsQuery( @@ -158,7 +207,41 @@ private void runFtsQuery( List documents, ScalarIndexParams scalarParams) throws Exception { + withIndexedDataset( + uri, + documents, + scalarParams, + dataset -> { + ScanOptions scanOptions = new ScanOptions.Builder().fullTextQuery(query).build(); + + try (LanceScanner scanner = dataset.newScan(scanOptions)) { + long total = 0L; + try (ArrowReader arrowReader = scanner.scanBatches()) { + while (arrowReader.loadNextBatch()) { + total += arrowReader.getVectorSchemaRoot().getRowCount(); + } + } + assertEquals(expectedTotal, total); + } + }); + } + + /** Same dataset as the three-argument {@link #runFtsQuery}, without running a query on it. */ + private void withIndexedDataset(String uri, IndexedDatasetConsumer consumer) throws Exception { + withIndexedDataset(uri, DEFAULT_DOCUMENTS, defaultIndexParams(), consumer); + } + /** + * Create an in-memory dataset with two text columns ({@code doc}, {@code title}), each backed by + * an inverted index, and hand it to {@code consumer}. Resources are released when the consumer + * returns. + */ + private void withIndexedDataset( + String uri, + List documents, + ScalarIndexParams scalarParams, + IndexedDatasetConsumer consumer) + throws Exception { Schema schema = new Schema( Arrays.asList( @@ -212,20 +295,15 @@ private void runFtsQuery( .withIndexName("title_idx") .build()); - ScanOptions scanOptions = new ScanOptions.Builder().fullTextQuery(query).build(); - - try (LanceScanner scanner = dataset.newScan(scanOptions)) { - long total = 0L; - try (ArrowReader arrowReader = scanner.scanBatches()) { - while (arrowReader.loadNextBatch()) { - total += arrowReader.getVectorSchemaRoot().getRowCount(); - } - } - assertEquals(expectedTotal, total); - } + consumer.accept(dataset); } } } } } + + @FunctionalInterface + private interface IndexedDatasetConsumer { + void accept(Dataset dataset) throws Exception; + } } diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 11a2b6e513f..a5c7dd05f8b 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -30,8 +30,8 @@ use arrow_schema::{DataType, Field}; use async_trait::async_trait; pub use builder::InvertedIndexBuilder; pub use combined::{ - CombinedFieldColumn, build_combined_bm25_scorer, combined_fields_search, - validate_combined_tokenizers, + CombinedCorpusStats, CombinedFieldColumn, FlatFieldStats, build_combined_bm25_scorer, + combined_fields_search, flat_combined_fields_search_stream, validate_combined_tokenizers, }; pub use compound::{ compound_search, compound_search_prepared_match, diff --git a/rust/lance-index/src/scalar/inverted/combined.rs b/rust/lance-index/src/scalar/inverted/combined.rs index 9de77185ce7..e48f245b8a0 100644 --- a/rust/lance-index/src/scalar/inverted/combined.rs +++ b/rust/lance-index/src/scalar/inverted/combined.rs @@ -23,6 +23,7 @@ //! scored; see [`combined_fields_search`]. mod cursor; +mod flat; mod search; mod stats; #[cfg(test)] @@ -31,9 +32,11 @@ mod testing; use std::sync::Arc; use lance_core::{Error, Result}; +use lance_select::RowAddrTreeMap; +pub use flat::flat_combined_fields_search_stream; pub use search::combined_fields_search; -pub use stats::build_combined_bm25_scorer; +pub use stats::{CombinedCorpusStats, FlatFieldStats, build_combined_bm25_scorer}; use super::index::InvertedIndex; use super::query::Tokens; @@ -47,6 +50,12 @@ pub struct CombinedFieldColumn { pub weight: f32, /// Opened inverted-index segments for this column. pub indices: Vec>, + /// Rows an overlay made stale, in the row-id domain index results use. + /// + /// The flat scan folds their current values into the corpus, so their old + /// values are subtracted from `indices` to avoid counting them twice. + /// `None` when no overlay touched this column. + pub stale_rows: Option>, } /// Deduplicate the query tokens into the unique terms of the virtual field, diff --git a/rust/lance-index/src/scalar/inverted/combined/flat.rs b/rust/lance-index/src/scalar/inverted/combined/flat.rs new file mode 100644 index 00000000000..6391db557ec --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/combined/flat.rs @@ -0,0 +1,388 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The unindexed `combined_fields` plan: blend the scanned column values into +//! `dl'`/`tf'` and score the rows no target column's index covers. + +use std::sync::Arc; + +use arrow::array::{Float32Builder, UInt64Builder}; +use arrow_array::{ArrayRef, RecordBatch}; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::metrics::Time; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{FutureExt, stream}; +use lance_core::error::DataFusionResult; +use lance_core::utils::tokio::spawn_cpu; +use lance_select::RowAddrMask; + +use super::super::index::{BlendedRows, FTS_SCHEMA, slice_into_batches, tokenize_and_blend_multi}; +use super::super::query::{Operator, Tokens}; +use super::super::scorer::CombinedFieldsBM25Scorer; +use super::super::tokenizer::document_tokenizer::LanceTokenizer; +use super::stats::{CombinedCorpusStats, FlatFieldStats, build_combined_bm25_scorer}; +use super::{CombinedFieldColumn, unique_terms}; +use crate::metrics::MetricsCollector; + +/// Exact cross-field BM25F search over rows that no index fully covers, scored +/// straight from their column values. Together with the indexed +/// [`combined_fields_search`](super::combined_fields_search) this covers the +/// whole dataset, with the same term deduplication and `operator` semantics. +/// +/// - `input` carries `_rowid` plus every target column; `doc_col_indices` locates +/// them, in `columns` order. +/// - `stats_masks[column]` selects the rows folded into that column's corpus +/// statistics, so a row indexed for some columns is not counted twice. +/// - `emit_mask`, when set, selects the rows to emit from an unfiltered `input`. +/// - `flat_covers_whole_corpus` selects [`CombinedCorpusStats::FlatOnly`]. +/// - `metrics` receives only the scorer build's index reads. +/// +/// Returns the scorer alongside the stream. The whole input is consumed before +/// the first output batch, so an indexed sibling can score against the same +/// statistics. +#[allow(clippy::too_many_arguments)] +pub async fn flat_combined_fields_search_stream( + input: SendableRecordBatchStream, + columns: &[CombinedFieldColumn], + doc_col_indices: Vec, + stats_masks: &[Arc], + emit_mask: Option>, + flat_covers_whole_corpus: bool, + tokens: &Tokens, + tokenizer: Box, + operator: Operator, + target_batch_size: usize, + elapsed_compute: Option