diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index caf837b371a..879240ed04d 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -3402,6 +3402,7 @@ fn inner_describe_indices<'local>( for_column: for_column.as_deref(), has_name: has_name.as_deref(), must_support_fts, + fts_document_granularity: None, must_support_exact_equality, }) })?; @@ -3579,6 +3580,7 @@ fn inner_get_zonemap_stats<'local>( for_column: Some(&column_name), has_name: None, must_support_fts: false, + fts_document_granularity: None, must_support_exact_equality: false, })) .await diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 335cb2a4fa3..ac3f9384113 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -18,10 +18,13 @@ use lance::dataset::scanner::{ ExecutionSummaryCounts, Scanner, }; use lance_index::scalar::FullTextSearchQuery; -use lance_index::scalar::inverted::query::{ - BooleanQuery as FtsBooleanQuery, BoostQuery as FtsBoostQuery, FtsQuery, - MatchQuery as FtsMatchQuery, MultiMatchQuery as FtsMultiMatchQuery, Occur as FtsOccur, - PhraseQuery as FtsPhraseQuery, +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, + }, }; use lance_io::ffi::to_ffi_arrow_array_stream; use lance_linalg::distance::DistanceType; @@ -114,6 +117,7 @@ pub(crate) fn build_full_text_search_query<'a>( let max_expansions = env.get_int_as_usize_from_method(&java_obj, "getMaxExpansions")?; let operator = env.get_fts_operator_from_method(&java_obj)?; let prefix_length = env.get_u32_from_method(&java_obj, "getPrefixLength")?; + let document_granularity = get_document_granularity(env, &java_obj)?; let mut query = FtsMatchQuery::new(query_text); query = query.with_column(Some(column)); @@ -123,6 +127,9 @@ pub(crate) fn build_full_text_search_query<'a>( .with_max_expansions(max_expansions) .with_operator(operator) .with_prefix_length(prefix_length); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity(document_granularity); + } Ok(FtsQuery::Match(query)) } @@ -130,10 +137,14 @@ pub(crate) fn build_full_text_search_query<'a>( let query_text = env.get_string_from_method(&java_obj, "getQueryText")?; let column = env.get_string_from_method(&java_obj, "getColumn")?; let slop = env.get_u32_from_method(&java_obj, "getSlop")?; + let document_granularity = get_document_granularity(env, &java_obj)?; let mut query = FtsPhraseQuery::new(query_text); query = query.with_column(Some(column)); query = query.with_slop(slop); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity(document_granularity); + } Ok(FtsQuery::Phrase(query)) } @@ -229,6 +240,20 @@ pub(crate) fn build_full_text_search_query<'a>( } } +fn get_document_granularity( + env: &mut JNIEnv<'_>, + java_obj: &JObject, +) -> Result> { + env.get_optional_from_method( + java_obj, + "getDocumentGranularity", + |env, granularity_obj| { + let value = env.get_string_from_method(&granularity_obj, "toRustString")?; + DocumentGranularity::try_from(value.as_str()).map_err(Error::from) + }, + ) +} + /// Scanner options passed from JNI - shared between blocking and async scanners pub(crate) struct ScannerOptions<'a> { pub fragment_ids_obj: JObject<'a>, diff --git a/java/src/main/java/org/lance/DocumentGranularity.java b/java/src/main/java/org/lance/DocumentGranularity.java new file mode 100644 index 00000000000..c540cc1ff9c --- /dev/null +++ b/java/src/main/java/org/lance/DocumentGranularity.java @@ -0,0 +1,34 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance; + +/** The unit treated as one full-text-search document. */ +public enum DocumentGranularity { + /** All text selected from one dataset row belongs to one document. */ + ROW("row"), + + /** Each element of the deepest list on the field path is one document. */ + LIST_ELEMENT("list_element"); + + private final String rustString; + + DocumentGranularity(String rustString) { + this.rustString = rustString; + } + + /** Return the stable value understood by the Rust API and serialized index parameters. */ + public String toRustString() { + return rustString; + } +} diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index 82513a89ee7..a5d82a57812 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -13,6 +13,7 @@ */ package org.lance.index.scalar; +import org.lance.DocumentGranularity; import org.lance.util.JsonUtils; import com.google.common.base.Preconditions; @@ -56,6 +57,7 @@ public static final class Builder { private Integer blockSize = 128; private Boolean skipMerge; private Integer formatVersion; + private DocumentGranularity documentGranularity = DocumentGranularity.ROW; /** * Configure the base tokenizer. @@ -279,6 +281,21 @@ public Builder formatVersion(int formatVersion) { return this; } + /** + * Configure the unit treated as one FTS document. + * + *

{@link DocumentGranularity#LIST_ELEMENT} uses each element of the deepest list on the + * indexed field path as one document. The default is {@link DocumentGranularity#ROW}. + * + * @param documentGranularity document boundary semantics + * @return this builder + */ + public Builder documentGranularity(DocumentGranularity documentGranularity) { + this.documentGranularity = + Objects.requireNonNull(documentGranularity, "documentGranularity must not be null"); + return this; + } + /** Build a {@link ScalarIndexParams} instance for an inverted index. */ public ScalarIndexParams build() { if (formatVersion != null) { @@ -337,6 +354,7 @@ public ScalarIndexParams build() { if (formatVersion != null) { params.put("format_version", formatVersion); } + params.put("document_granularity", documentGranularity.toRustString()); String json = JsonUtils.toJson(params); return ScalarIndexParams.create(INDEX_TYPE, json); diff --git a/java/src/main/java/org/lance/ipc/FullTextQuery.java b/java/src/main/java/org/lance/ipc/FullTextQuery.java index 0163def2f76..4fdb4741fd5 100755 --- a/java/src/main/java/org/lance/ipc/FullTextQuery.java +++ b/java/src/main/java/org/lance/ipc/FullTextQuery.java @@ -13,6 +13,8 @@ */ package org.lance.ipc; +import org.lance.DocumentGranularity; + import com.google.common.base.MoreObjects; import org.apache.arrow.util.Preconditions; @@ -21,7 +23,13 @@ import java.util.Objects; import java.util.Optional; -/** Base type for full text search queries used by Lance scanner. */ +/** + * Base type for full text search queries used by Lance scanner. + * + *

Match and phrase overloads without a {@link DocumentGranularity} infer the unique indexed + * granularity for the field. They use row documents when no index exists and fail as ambiguous when + * row and list-element indexes coexist. + */ public abstract class FullTextQuery { public enum Type { MATCH, @@ -63,7 +71,14 @@ public FullTextQuery getQuery() { public abstract Type getType(); public static FullTextQuery match(String queryText, String column) { - return match(queryText, column, 1.0f, Optional.empty(), 50, Operator.OR, 0); + return new MatchQuery( + queryText, column, 1.0f, Optional.empty(), 50, Operator.OR, 0, Optional.empty()); + } + + public static FullTextQuery match( + String queryText, String column, DocumentGranularity documentGranularity) { + return match( + queryText, column, 1.0f, Optional.empty(), 50, Operator.OR, 0, documentGranularity); } public static FullTextQuery match( @@ -75,15 +90,58 @@ public static FullTextQuery match( Operator operator, int prefixLength) { return new MatchQuery( - queryText, column, boost, fuzziness, maxExpansions, operator, prefixLength); + queryText, + column, + boost, + fuzziness, + maxExpansions, + operator, + prefixLength, + Optional.empty()); + } + + public static FullTextQuery match( + String queryText, + String column, + float boost, + Optional fuzziness, + int maxExpansions, + Operator operator, + int prefixLength, + DocumentGranularity documentGranularity) { + return new MatchQuery( + queryText, + column, + boost, + fuzziness, + maxExpansions, + operator, + prefixLength, + Optional.of( + Objects.requireNonNull(documentGranularity, "documentGranularity must not be null"))); } public static FullTextQuery phrase(String queryText, String column) { - return phrase(queryText, column, 0); + return new PhraseQuery(queryText, column, 0, Optional.empty()); + } + + public static FullTextQuery phrase( + String queryText, String column, DocumentGranularity documentGranularity) { + return phrase(queryText, column, 0, documentGranularity); } public static FullTextQuery phrase(String queryText, String column, int slop) { - return new PhraseQuery(queryText, column, slop); + return new PhraseQuery(queryText, column, slop, Optional.empty()); + } + + public static FullTextQuery phrase( + String queryText, String column, int slop, DocumentGranularity documentGranularity) { + return new PhraseQuery( + queryText, + column, + slop, + Optional.of( + Objects.requireNonNull(documentGranularity, "documentGranularity must not be null"))); } public static FullTextQuery multiMatch(String queryText, List columns) { @@ -117,6 +175,7 @@ public static final class MatchQuery extends FullTextQuery { private final int maxExpansions; private final Operator operator; private final int prefixLength; + private final Optional documentGranularity; MatchQuery( String queryText, @@ -125,7 +184,8 @@ public static final class MatchQuery extends FullTextQuery { Optional fuzziness, int maxExpansions, Operator operator, - int prefixLength) { + int prefixLength, + Optional documentGranularity) { Preconditions.checkArgument( queryText != null && !queryText.isEmpty(), "queryText must not be null or empty"); Preconditions.checkArgument( @@ -140,6 +200,7 @@ public static final class MatchQuery extends FullTextQuery { this.maxExpansions = maxExpansions; this.operator = operator == null ? Operator.OR : operator; this.prefixLength = prefixLength; + this.documentGranularity = Objects.requireNonNull(documentGranularity); } @Override @@ -175,6 +236,11 @@ public int getPrefixLength() { return prefixLength; } + /** Returns the explicit granularity, or empty when query planning should infer it. */ + public Optional getDocumentGranularity() { + return documentGranularity; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -184,6 +250,7 @@ public boolean equals(Object o) { && maxExpansions == other.maxExpansions && prefixLength == other.prefixLength && operator == other.operator + && Objects.equals(documentGranularity, other.documentGranularity) && Objects.equals(queryText, other.queryText) && Objects.equals(column, other.column) && Objects.equals(fuzziness, other.fuzziness); @@ -192,7 +259,14 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash( - queryText, column, boost, fuzziness, maxExpansions, operator, prefixLength); + queryText, + column, + boost, + fuzziness, + maxExpansions, + operator, + prefixLength, + documentGranularity); } @Override @@ -206,6 +280,7 @@ public String toString() { .add("maxExpansions", maxExpansions) .add("operator", operator) .add("prefixLength", prefixLength) + .add("documentGranularity", documentGranularity) .toString(); } } @@ -215,8 +290,13 @@ public static final class PhraseQuery extends FullTextQuery { private final String queryText; private final String column; private final int slop; + private final Optional documentGranularity; - PhraseQuery(String queryText, String column, int slop) { + PhraseQuery( + String queryText, + String column, + int slop, + Optional documentGranularity) { Preconditions.checkArgument( queryText != null && !queryText.isEmpty(), "queryText must not be null or empty"); Preconditions.checkArgument( @@ -226,6 +306,7 @@ public static final class PhraseQuery extends FullTextQuery { this.queryText = queryText; this.column = column; this.slop = slop; + this.documentGranularity = Objects.requireNonNull(documentGranularity); } @Override @@ -245,19 +326,25 @@ public int getSlop() { return slop; } + /** Returns the explicit granularity, or empty when query planning should infer it. */ + public Optional getDocumentGranularity() { + return documentGranularity; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (!(o instanceof PhraseQuery)) return false; PhraseQuery other = (PhraseQuery) o; return slop == other.slop + && Objects.equals(documentGranularity, other.documentGranularity) && Objects.equals(queryText, other.queryText) && Objects.equals(column, other.column); } @Override public int hashCode() { - return Objects.hash(queryText, column, slop); + return Objects.hash(queryText, column, slop, documentGranularity); } @Override @@ -267,6 +354,7 @@ public String toString() { .add("queryText", queryText) .add("column", column) .add("slop", slop) + .add("documentGranularity", documentGranularity) .toString(); } } diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index 0ccc429ec57..4ea1c0f9ca1 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -13,6 +13,7 @@ */ package org.lance.index.scalar; +import org.lance.DocumentGranularity; import org.lance.util.JsonUtils; import org.junit.jupiter.api.Test; @@ -40,6 +41,16 @@ void defaultBlockSizeIsSerialized() { Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); assertEquals(128, ((Number) json.get("block_size")).intValue()); + assertEquals("row", json.get("document_granularity")); + } + + @Test + void documentGranularityIsSerialized() { + ScalarIndexParams params = + InvertedIndexParams.builder().documentGranularity(DocumentGranularity.LIST_ELEMENT).build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals("list_element", json.get("document_granularity")); } @Test diff --git a/java/src/test/java/org/lance/ipc/FullTextQueryTest.java b/java/src/test/java/org/lance/ipc/FullTextQueryTest.java index 3dcf4276277..b84a19db4c8 100755 --- a/java/src/test/java/org/lance/ipc/FullTextQueryTest.java +++ b/java/src/test/java/org/lance/ipc/FullTextQueryTest.java @@ -13,6 +13,8 @@ */ package org.lance.ipc; +import org.lance.DocumentGranularity; + import org.junit.jupiter.api.Test; import java.util.Arrays; @@ -39,6 +41,7 @@ void testMatchQueryDefaults() { assertEquals(50, q.getMaxExpansions()); assertEquals(FullTextQuery.Operator.OR, q.getOperator()); assertEquals(0, q.getPrefixLength()); + assertEquals(Optional.empty(), q.getDocumentGranularity()); } @Test @@ -46,7 +49,14 @@ void testMatchQueryCustomParameters() { FullTextQuery.MatchQuery q = (FullTextQuery.MatchQuery) FullTextQuery.match( - "hello", "title", 2.0f, Optional.of(1), 10, FullTextQuery.Operator.AND, 3); + "hello", + "title", + 2.0f, + Optional.of(1), + 10, + FullTextQuery.Operator.AND, + 3, + DocumentGranularity.LIST_ELEMENT); assertEquals(FullTextQuery.Type.MATCH, q.getType()); assertEquals("hello", q.getQueryText()); @@ -56,6 +66,7 @@ void testMatchQueryCustomParameters() { assertEquals(10, q.getMaxExpansions()); assertEquals(FullTextQuery.Operator.AND, q.getOperator()); assertEquals(3, q.getPrefixLength()); + assertEquals(Optional.of(DocumentGranularity.LIST_ELEMENT), q.getDocumentGranularity()); } @Test @@ -67,17 +78,20 @@ void testPhraseQueryDefaults() { assertEquals("exact match", q.getQueryText()); assertEquals("content", q.getColumn()); assertEquals(0, q.getSlop()); + assertEquals(Optional.empty(), q.getDocumentGranularity()); } @Test void testPhraseQueryCustomSlop() { FullTextQuery.PhraseQuery q = - (FullTextQuery.PhraseQuery) FullTextQuery.phrase("ordered terms", "content", 2); + (FullTextQuery.PhraseQuery) + FullTextQuery.phrase("ordered terms", "content", 2, DocumentGranularity.LIST_ELEMENT); assertEquals(FullTextQuery.Type.MATCH_PHRASE, q.getType()); assertEquals("ordered terms", q.getQueryText()); assertEquals("content", q.getColumn()); assertEquals(2, q.getSlop()); + assertEquals(Optional.of(DocumentGranularity.LIST_ELEMENT), q.getDocumentGranularity()); } @Test diff --git a/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java b/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java index 1c46b399195..9098084ac9f 100755 --- a/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java +++ b/java/src/test/java/org/lance/ipc/LanceScannerFullTextSearchTest.java @@ -14,6 +14,7 @@ package org.lance.ipc; import org.lance.Dataset; +import org.lance.DocumentGranularity; import org.lance.WriteParams; import org.lance.index.IndexOptions; import org.lance.index.IndexParams; @@ -41,12 +42,31 @@ import java.util.Collections; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; class LanceScannerFullTextSearchTest { @Test void testMatchQuery() throws Exception { - runFtsQuery("memory://fts_java_match", FullTextQuery.match("hello", "doc"), 2L); + runFtsQuery( + "memory://fts_java_match", + FullTextQuery.match("hello", "doc", DocumentGranularity.ROW), + 2L); + } + + @Test + void testExplicitListElementGranularityReachesRustRouting() { + RuntimeException error = + assertThrows( + RuntimeException.class, + () -> + runFtsQuery( + "memory://fts_java_list_element_validation", + FullTextQuery.match("hello", "doc", DocumentGranularity.LIST_ELEMENT), + 0L)); + assertTrue(error.getMessage().contains("requested ListElement"), error.getMessage()); + assertTrue(error.getMessage().contains("'doc_idx' (Row)"), error.getMessage()); } @Test diff --git a/protos/index_old.proto b/protos/index_old.proto index eb984d6fe29..c422292083d 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -26,6 +26,11 @@ message LabelListIndexDetails {} message NGramIndexDetails {} message ZoneMapIndexDetails {} message InvertedIndexDetails { + enum DocumentGranularity { + ROW = 0; + LIST_ELEMENT = 1; + } + // Marking this field as optional as old versions of the index store blank details and we // need to make sure we have a proper optional field to detect this. optional string base_tokenizer = 1; @@ -44,4 +49,10 @@ message InvertedIndexDetails { // A present value records the block size used by the index; 256 is only // valid with format version 3. optional uint32 block_size = 12; + // The logical FTS document boundary. The protobuf default preserves the + // legacy row-document behavior when this field is absent. + DocumentGranularity document_granularity = 13; + // The posting-list payload format. This is separate from index_version, + // which is an element-document compatibility marker for version 4. + optional uint32 posting_format_version = 14; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..71e4b5fc992 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -70,7 +70,7 @@ ) from .lance import __version__ as __version__ from .lance import _Session as Session -from .query import FullTextQuery +from .query import DocumentGranularity, FullTextQuery from .types import _coerce_reader from .udf import BatchUDF, normalize_transform from .udf import BatchUDFCheckpoint as BatchUDFCheckpoint @@ -3124,8 +3124,6 @@ def _prepare_scalar_index_request( column = column[0] lance_field = self._ds.lance_schema.field_case_insensitive(column) - if lance_field is None: - raise KeyError(f"{column} not found in schema") if isinstance(index_type, str): index_type = index_type.upper() @@ -3148,6 +3146,12 @@ def _prepare_scalar_index_request( ) ) + if lance_field is None: + if index_type in ["INVERTED", "FTS"]: + # Rust resolves public FTS paths through intervening list layers. + return column, index_type, index_type + raise KeyError(f"{column} not found in schema") + field = lance_field.to_arrow() field_type = field.type @@ -3201,7 +3205,12 @@ def _prepare_scalar_index_request( return column, index_type, index_type elif isinstance(index_type, IndexConfig): logical_index_type = index_type.index_type.upper() - config = json.dumps(index_type.parameters) + if lance_field is None and logical_index_type not in ["INVERTED", "FTS"]: + raise KeyError(f"{column} not found in schema") + parameters = dict(index_type.parameters) + if logical_index_type in ["INVERTED", "FTS"]: + parameters["document_granularity"] = kwargs["document_granularity"] + config = json.dumps(parameters) kwargs["config"] = indices.IndexConfig(index_type.index_type, config) return column, "scalar", logical_index_type else: @@ -3262,6 +3271,7 @@ def create_scalar_index( index_uuid: Optional[str] = None, progress_callback: Optional[Callable[[IndexProgress], None]] = None, format_version: Optional[Union[int, str]] = None, + document_granularity: DocumentGranularity = DocumentGranularity.ROW, **kwargs, ): """Create a scalar index on a column. @@ -3375,6 +3385,11 @@ def create_scalar_index( ``format_version=3`` is experimental and is only valid with ``block_size=256``. + document_granularity: DocumentGranularity, default ROW + This is for the ``INVERTED`` / ``FTS`` index. ``ROW`` treats all + selected text in one dataset row as one document. ``LIST_ELEMENT`` + treats each element of the deepest list on ``column`` as one document. + with_position: bool, default False This is for the ``INVERTED`` index. If True, the index will store the positions of the words in the document, so that you can conduct phrase @@ -3468,6 +3483,11 @@ def create_scalar_index( ``MaterializeIndex`` operator. """ + if not isinstance(document_granularity, DocumentGranularity): + raise TypeError( + "document_granularity must be a lance.query.DocumentGranularity" + ) + kwargs["document_granularity"] = document_granularity.value column, index_type, logical_index_type = self._prepare_scalar_index_request( column, index_type, kwargs ) @@ -3489,7 +3509,6 @@ def create_scalar_index( kwargs["progress_callback"] = progress_callback if format_version is not None: kwargs["format_version"] = format_version - self._ds.create_index([column], index_type, name, replace, train, None, kwargs) def _create_index_impl( diff --git a/python/python/lance/query.py b/python/python/lance/query.py index 85bbf121e6e..a6098cc0096 100644 --- a/python/python/lance/query.py +++ b/python/python/lance/query.py @@ -22,6 +22,13 @@ class FullTextOperator(Enum): OR = "OR" +class DocumentGranularity(str, Enum): + """The unit treated as one full-text-search document.""" + + ROW = "row" + LIST_ELEMENT = "list_element" + + class Occur(Enum): SHOULD = "SHOULD" MUST = "MUST" @@ -98,6 +105,7 @@ def __init__( max_expansions: int = 50, operator: FullTextOperator = FullTextOperator.OR, prefix_length: int = 0, + document_granularity: Optional[DocumentGranularity] = None, ): """ Match query for full-text search. @@ -129,6 +137,10 @@ def __init__( prefix_length : int, default 0 The number of beginning characters being unchanged for fuzzy matching. This is useful to achieve prefix matching. + document_granularity : DocumentGranularity, optional + Explicitly select row or deepest-list-element documents. If omitted, + the indexed granularity is inferred. When both granularities are indexed + for the field, this must be specified. With no index, ``ROW`` is used. """ self._inner = PyFullTextQuery.match_query( query, @@ -138,6 +150,9 @@ def __init__( max_expansions=max_expansions, operator=operator.value, prefix_length=prefix_length, + document_granularity=( + document_granularity.value if document_granularity is not None else None + ), ) def query_type(self) -> FullTextQueryType: @@ -145,7 +160,14 @@ def query_type(self) -> FullTextQueryType: class PhraseQuery(FullTextQuery): - def __init__(self, query: str, column: str, *, slop: int = 0): + def __init__( + self, + query: str, + column: str, + *, + slop: int = 0, + document_granularity: Optional[DocumentGranularity] = None, + ): """ Phrase query for full-text search. @@ -155,8 +177,21 @@ def __init__(self, query: str, column: str, *, slop: int = 0): The query string to match against. column : str The name of the column to match against. + slop : int, default 0 + The maximum number of intervening positions permitted in the phrase. + document_granularity : DocumentGranularity, optional + Explicitly select row or deepest-list-element documents. If omitted, + the indexed granularity is inferred. When both granularities are indexed + for the field, this must be specified. With no index, ``ROW`` is used. """ - self._inner = PyFullTextQuery.phrase_query(query, column, slop) + self._inner = PyFullTextQuery.phrase_query( + query, + column, + slop, + document_granularity=( + document_granularity.value if document_granularity is not None else None + ), + ) def query_type(self) -> FullTextQueryType: return FullTextQueryType.MATCH_PHRASE diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 985072fe5de..22e754b37b5 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -21,6 +21,7 @@ from lance.query import ( BooleanQuery, BoostQuery, + DocumentGranularity, MatchQuery, MultiMatchQuery, Occur, @@ -1351,6 +1352,136 @@ def test_fts_on_list(tmp_path): assert results.num_rows == 6 +@pytest.mark.parametrize( + "list_type", + [ + pa.list_(pa.string()), + pa.list_(pa.large_string()), + pa.large_list(pa.string()), + pa.large_list(pa.large_string()), + ], +) +def test_fts_on_list_elements(tmp_path, list_type): + data = pa.table( + { + "id": pa.array([0, 1, 2]), + "tags": pa.array( + [ + ["alpha beta", "gamma alpha", None, "", "delta"], + ["beta", "gamma"], + None, + ], + type=list_type, + ), + } + ) + ds = lance.write_dataset(data, tmp_path) + + def hits(table): + return sorted(zip(table["id"].to_pylist(), table["_doc_index"].to_pylist())) + + list_element = DocumentGranularity.LIST_ELEMENT + query = MatchQuery("alpha", "tags", document_granularity=list_element) + flat = ds.to_table(full_text_query=query) + assert hits(flat) == [(0, [0]), (0, [1])] + assert pa.types.is_list(flat.schema.field("_doc_index").type) + assert flat.schema.field("_doc_index").type.value_type == pa.uint32() + assert hits( + ds.to_table( + full_text_query=MatchQuery( + "delta", "tags", document_granularity=list_element + ) + ) + ) == [(0, [4])] + + ds.create_scalar_index( + "tags", + "INVERTED", + with_position=True, + document_granularity=list_element, + ) + ds.create_scalar_index( + "tags", + IndexConfig(index_type="inverted", parameters={"with_position": True}), + document_granularity=list_element, + ) + inferred = ds.to_table(full_text_query=MatchQuery("alpha", "tags")) + assert hits(inferred) == [(0, [0]), (0, [1])] + with pytest.raises(ValueError, match=r"requested Row.*ListElement"): + ds.to_table( + full_text_query=MatchQuery( + "alpha", "tags", document_granularity=DocumentGranularity.ROW + ) + ) + + ds.create_scalar_index("tags", "INVERTED", with_position=True) + index_names = {index.name for index in ds.describe_indices()} + assert {"tags_idx", "tags_list_element_idx"}.issubset(index_names) + row_auto = ds.to_table(full_text_query="alpha") + assert row_auto["id"].to_pylist() == [0] + assert "_doc_index" not in row_auto.column_names + with pytest.raises(ValueError, match=r"ambiguous.*document_granularity"): + ds.to_table(full_text_query=MatchQuery("alpha", "tags")) + indexed = ds.to_table(full_text_query=query) + assert hits(indexed) == [(0, [0]), (0, [1])] + assert pa.types.is_list(indexed.schema.field("_doc_index").type) + assert indexed.schema.field("_doc_index").type.value_type == pa.uint32() + filtered = ds.to_table(full_text_query=query, filter="id = 0", prefilter=True) + assert hits(filtered) == [(0, [0]), (0, [1])] + assert hits( + ds.to_table( + full_text_query=MatchQuery( + "delta", "tags", document_granularity=list_element + ) + ) + ) == [(0, [4])] + + phrase = ds.to_table( + full_text_query=PhraseQuery( + "beta gamma", "tags", document_granularity=list_element + ) + ) + assert phrase.num_rows == 0 + assert hits( + ds.to_table( + full_text_query=PhraseQuery( + "alpha beta", "tags", document_granularity=list_element + ) + ) + ) == [(0, [0])] + row_phrase = ds.to_table( + full_text_query=PhraseQuery( + "beta gamma", + "tags", + document_granularity=DocumentGranularity.ROW, + ) + ) + assert sorted(row_phrase["id"].to_pylist()) == [0, 1] + assert "_doc_index" not in row_phrase.column_names + + ds.insert( + pa.table( + { + "id": pa.array([3]), + "tags": pa.array([["alpha", "alpha again"]], type=list_type), + } + ) + ) + assert hits(ds.to_table(full_text_query=query)) == [ + (0, [0]), + (0, [1]), + (3, [0]), + (3, [1]), + ] + + with pytest.raises(RuntimeError, match=r"tags\[\*\]"): + ds.create_scalar_index( + "tags[*]", + "INVERTED", + document_granularity=list_element, + ) + + def test_fts_fuzzy_query(tmp_path): data = pa.table( { diff --git a/python/src/dataset.rs b/python/src/dataset.rs index a361e0b63a8..ad2f9e25245 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -78,7 +78,7 @@ use lance_index::{ FtsPrewarmOptions, IndexParams, IndexType, PrewarmOptions, optimize::OptimizeOptions, progress::{IndexBuildProgress, NoopIndexBuildProgress}, - scalar::inverted::InvertedListFormatVersion, + scalar::inverted::{DocumentGranularity, InvertedListFormatVersion}, scalar::{FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams}, vector::{ ApproxMode, DEFAULT_QUERY_PARALLELISM, Query as VectorQuery, @@ -1211,6 +1211,13 @@ impl Dataset { let is_phrase = query.len() >= 2 && query.starts_with('"') && query.ends_with('"'); let is_multi_match = columns.as_ref().map(|cols| cols.len() > 1).unwrap_or(false); + let document_granularity = match full_text_query.get_item("document_granularity")? { + Some(value) if !value.is_none() => Some( + DocumentGranularity::try_from(value.extract::()?.as_str()) + .map_err(|err| PyValueError::new_err(err.to_string()))?, + ), + _ => None, + }; if is_phrase { // Remove the surrounding quotes for phrase queries @@ -1218,8 +1225,20 @@ impl Dataset { } let query: FtsQuery = match (is_phrase, is_multi_match) { - (false, _) => MatchQuery::new(query).into(), - (true, false) => PhraseQuery::new(query).into(), + (false, _) => { + let mut query = MatchQuery::new(query); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity(document_granularity); + } + query.into() + } + (true, false) => { + let mut query = PhraseQuery::new(query); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity(document_granularity); + } + query.into() + } (true, true) => { return Err(PyValueError::new_err( "Phrase queries cannot be used with multiple columns.", @@ -2421,6 +2440,13 @@ impl Dataset { "INVERTED" | "FTS" => { let mut params = InvertedIndexParams::default(); if let Some(kwargs) = kwargs { + if let Some(document_granularity) = kwargs.get_item("document_granularity")? { + let document_granularity: String = document_granularity.extract()?; + params = params.document_granularity( + DocumentGranularity::try_from(document_granularity.as_str()) + .map_err(|err| PyValueError::new_err(err.to_string()))?, + ); + } if let Some(with_position) = kwargs.get_item("with_position")? { params = params.with_position(with_position.extract()?); } @@ -5081,7 +5107,8 @@ pub struct PyFullTextQuery { #[pymethods] impl PyFullTextQuery { #[staticmethod] - #[pyo3(signature = (query, column, boost=1.0, fuzziness=Some(0), max_expansions=50, operator="OR", prefix_length=0))] + #[pyo3(signature = (query, column, boost=1.0, fuzziness=Some(0), max_expansions=50, operator="OR", prefix_length=0, document_granularity=None))] + #[allow(clippy::too_many_arguments)] fn match_query( query: String, column: String, @@ -5090,30 +5117,48 @@ impl PyFullTextQuery { max_expansions: usize, operator: &str, prefix_length: u32, + document_granularity: Option<&str>, ) -> PyResult { + let mut query = MatchQuery::new(query) + .with_column(Some(column)) + .with_boost(boost) + .with_fuzziness(fuzziness) + .with_max_expansions(max_expansions) + .with_operator( + Operator::try_from(operator) + .map_err(|e| PyValueError::new_err(format!("Invalid operator: {}", e)))?, + ) + .with_prefix_length(prefix_length); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity( + DocumentGranularity::try_from(document_granularity) + .map_err(|err| PyValueError::new_err(err.to_string()))?, + ); + } Ok(Self { - inner: MatchQuery::new(query) - .with_column(Some(column)) - .with_boost(boost) - .with_fuzziness(fuzziness) - .with_max_expansions(max_expansions) - .with_operator( - Operator::try_from(operator) - .map_err(|e| PyValueError::new_err(format!("Invalid operator: {}", e)))?, - ) - .with_prefix_length(prefix_length) - .into(), + inner: query.into(), }) } #[staticmethod] - #[pyo3(signature = (query, column, slop))] - fn phrase_query(query: String, column: String, slop: u32) -> PyResult { + #[pyo3(signature = (query, column, slop, document_granularity=None))] + fn phrase_query( + query: String, + column: String, + slop: u32, + document_granularity: Option<&str>, + ) -> PyResult { + let mut query = PhraseQuery::new(query) + .with_column(Some(column)) + .with_slop(slop); + if let Some(document_granularity) = document_granularity { + query = query.with_document_granularity( + DocumentGranularity::try_from(document_granularity) + .map_err(|err| PyValueError::new_err(err.to_string()))?, + ); + } Ok(Self { - inner: PhraseQuery::new(query) - .with_column(Some(column)) - .with_slop(slop) - .into(), + inner: query.into(), }) } diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 6f9fc61b334..d71274b6545 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -1771,6 +1771,11 @@ mod tests { vec!["simple".to_string()] ); + assert_eq!( + parse_field_path("tags[*]").unwrap(), + vec!["tags[*]".to_string()] + ); + // Quoted field at the end assert_eq!( parse_field_path("parent.`field.with.dot`").unwrap(), diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 926bfad6a69..d0c0560a01d 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -149,6 +149,7 @@ impl InvertedIndexPlugin { params.validate_format_version()?; let format_version = params.resolved_format_version(); + let is_element_document = params.get_document_granularity().is_list_element(); let details = pbold::InvertedIndexDetails::try_from(¶ms)?; let mut inverted_index = InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask) @@ -156,7 +157,11 @@ impl InvertedIndexPlugin { let files = inverted_index.update(data, index_store, None).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: format_version.index_version(), + index_version: if is_element_document { + INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT + } else { + format_version.index_version() + }, files, }) } @@ -268,7 +273,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { } fn version(&self) -> u32 { - max_supported_fts_format_version().index_version() + INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT } fn new_query_parser( @@ -301,10 +306,8 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { - Ok( - InvertedIndex::load(index_store, frag_reuse_index, cache).await? - as Arc, - ) + let index = InvertedIndex::load(index_store, frag_reuse_index, cache).await?; + Ok(index as Arc) } fn details_as_json(&self, details: &prost_types::Any) -> Result { @@ -322,10 +325,7 @@ mod tests { #[test] fn test_plugin_version_tracks_max_supported_format() { let plugin = InvertedIndexPlugin; - assert_eq!( - plugin.version(), - max_supported_fts_format_version().index_version() - ); + assert_eq!(plugin.version(), INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT); } #[test] diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 41804e94fd1..1d498549d9b 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -108,9 +108,13 @@ fn merge_all_tail_partitions( ) -> Result> { let mut merged_builders: Vec = Vec::new(); let mut merged: Option = None; + let mut empty_coordinate_builder: Option = None; for tail in tails { let builder = tail.builder; if builder.is_empty() { + if builder.docs.coordinate_rank() > 0 && empty_coordinate_builder.is_none() { + empty_coordinate_builder = Some(builder); + } continue; } match &mut merged { @@ -132,6 +136,11 @@ fn merge_all_tail_partitions( if let Some(builder) = merged { merged_builders.push(builder); } + if merged_builders.is_empty() + && let Some(builder) = empty_coordinate_builder + { + merged_builders.push(builder); + } Ok(merged_builders) } @@ -388,6 +397,7 @@ impl InvertedIndexBuilder { token_set_format: self.token_set_format, worker_memory_limit_bytes, block_size: self.params.block_size, + coordinate_rank: document_coordinate_rank(&stream.schema()), }; let next_id = self.next_partition_id(); let id_alloc = Arc::new(AtomicU64::new(next_id)); @@ -1005,8 +1015,16 @@ impl InnerBuilder { } let doc_id_offset = self.docs.len() as u32; - for (row_id, num_tokens) in docs.iter() { - self.docs.append(*row_id, *num_tokens); + for doc_id in 0..docs.len() as u32 { + let row_id = docs.row_id(doc_id); + let num_tokens = docs.num_tokens(doc_id); + let doc_index = docs.doc_index(doc_id); + if doc_index.is_empty() { + self.docs.append(row_id, num_tokens); + } else { + self.docs + .append_with_doc_index(row_id, num_tokens, &doc_index)?; + } } self.posting_lists.resize_with(self.tokens.len(), || { PostingListBuilder::new_with_posting_tail_codec_and_block_size( @@ -1246,6 +1264,7 @@ struct IndexWorker { token_set_format: TokenSetFormat, token_ids: Vec, last_token_count: usize, + coordinate_rank: usize, } struct TailPartition { @@ -1271,6 +1290,7 @@ struct IndexWorkerConfig { token_set_format: TokenSetFormat, worker_memory_limit_bytes: u64, block_size: usize, + coordinate_rank: usize, } impl IndexWorker { @@ -1320,17 +1340,20 @@ impl IndexWorker { config.block_size, ); + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + id_alloc.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + | config.fragment_mask.unwrap_or(0), + config.with_position, + config.token_set_format, + config.format_version, + config.block_size, + ); + builder.docs = DocSet::with_coordinate_rank(config.coordinate_rank); + Ok(Self { tokenizer, dest_store, - builder: InnerBuilder::new_with_format_version_and_block_size( - id_alloc.fetch_add(1, std::sync::atomic::Ordering::Relaxed) - | config.fragment_mask.unwrap_or(0), - config.with_position, - config.token_set_format, - config.format_version, - config.block_size, - ), + builder, partitions: Vec::new(), files: Vec::new(), id_alloc, @@ -1342,6 +1365,7 @@ impl IndexWorker { token_set_format: config.token_set_format, token_ids: Vec::new(), last_token_count: 0, + coordinate_rank: config.coordinate_rank, }) } @@ -1355,22 +1379,55 @@ impl IndexWorker { async fn process_batch(&mut self, batch: RecordBatch) -> Result<()> { let doc_col = batch.column(0); let row_id_col = batch[ROW_ID].as_primitive::(); + let doc_index_columns = (0..self.coordinate_rank) + .map(|rank| { + let column_name = doc_index_storage_column(rank); + batch + .column_by_name(&column_name) + .ok_or_else(|| { + Error::index(format!( + "FTS document input is missing coordinate column {column_name}" + )) + }) + .map(|column| column.as_primitive::()) + }) + .collect::>>()?; match doc_col.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => { - let docs = iter_str_array(doc_col.as_ref()) + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + for (row_index, (doc, row_id)) in iter_str_array(doc_col.as_ref()) .zip(row_id_col.values().iter()) - .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); - - for (doc, row_id) in docs { - self.process_document(row_id, DocumentSource::Text(doc), false) - .await?; + .enumerate() + { + let doc_index = doc_index_columns + .iter() + .map(|column| column.value(row_index)) + .collect::>(); + self.process_document( + *row_id, + DocumentSource::Text(doc.unwrap_or_default()), + &doc_index, + false, + ) + .await?; } } DataType::List(_) => { + if self.coordinate_rank > 0 { + return Err(Error::index( + "ListElement FTS input must be expanded to string documents before indexing" + .to_string(), + )); + } self.process_string_list_batch::(doc_col, row_id_col) .await?; } DataType::LargeList(_) => { + if self.coordinate_rank > 0 { + return Err(Error::index( + "ListElement FTS input must be expanded to string documents before indexing" + .to_string(), + )); + } self.process_string_list_batch::(doc_col, row_id_col) .await?; } @@ -1402,12 +1459,21 @@ impl IndexWorker { } for (doc, row_id) in docs.iter().zip(row_id_col.values().iter()) { - let Some(doc) = doc else { - continue; - }; - - self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), true) - .await?; + match doc { + Some(doc) => { + self.process_document( + *row_id, + DocumentSource::StringList(doc.as_ref()), + &[], + false, + ) + .await?; + } + None => { + self.process_document(*row_id, DocumentSource::Text(""), &[], false) + .await?; + } + } } Ok(()) @@ -1436,6 +1502,7 @@ impl IndexWorker { &mut self, row_id: u64, document: DocumentSource<'_>, + doc_index: &[u32], skip_empty_document: bool, ) -> Result<()> { let with_position = self.has_position(); @@ -1575,7 +1642,13 @@ impl IndexWorker { } let old_doc_memory_size = self.builder.docs.memory_size() as u64; - let appended_doc_id = self.builder.docs.append(row_id, token_num); + let appended_doc_id = if doc_index.is_empty() { + self.builder.docs.append(row_id, token_num) + } else { + self.builder + .docs + .append_with_doc_index(row_id, token_num, doc_index)? + }; debug_assert_eq!(appended_doc_id, doc_id); self.adjust_tracked_memory_size( old_doc_memory_size, @@ -1652,7 +1725,7 @@ impl IndexWorker { #[instrument(level = "debug", skip_all)] async fn flush(&mut self) -> Result<()> { - if self.builder.tokens.is_empty() { + if self.builder.docs.is_empty() { return Ok(()); } @@ -1664,18 +1737,17 @@ impl IndexWorker { let with_position = self.has_position(); let format_version = self.builder.format_version; let block_size = self.builder.block_size; - let builder = std::mem::replace( - &mut self.builder, - InnerBuilder::new_with_format_version_and_block_size( - self.id_alloc - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) - | self.fragment_mask.unwrap_or(0), - with_position, - self.token_set_format, - format_version, - block_size, - ), + let mut replacement = InnerBuilder::new_with_format_version_and_block_size( + self.id_alloc + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + | self.fragment_mask.unwrap_or(0), + with_position, + self.token_set_format, + format_version, + block_size, ); + replacement.docs = DocSet::with_coordinate_rank(self.coordinate_rank); + let builder = std::mem::replace(&mut self.builder, replacement); let written_partition_id = builder.id(); let mut builder = builder; let target = if self.fragment_mask.is_some() { @@ -1698,7 +1770,7 @@ impl IndexWorker { } async fn finish(self) -> Result { - let tail_partition = if self.builder.tokens.is_empty() { + let tail_partition = if self.builder.docs.is_empty() && self.coordinate_rank == 0 { None } else { Some(TailPartition { @@ -1742,6 +1814,7 @@ impl PositionRecorder { #[derive(Debug, Eq, PartialEq, Clone, DeepSizeOf)] pub struct ScoredDoc { pub row_id: u64, + pub doc_index: Vec, pub score: OrderedFloat, } @@ -1749,6 +1822,15 @@ impl ScoredDoc { pub fn new(row_id: u64, score: f32) -> Self { Self { row_id, + doc_index: Vec::new(), + score: OrderedFloat(score), + } + } + + pub fn with_doc_index(row_id: u64, doc_index: Vec, score: f32) -> Self { + Self { + row_id, + doc_index, score: OrderedFloat(score), } } @@ -2254,7 +2336,7 @@ pub fn document_input( let schema = input.schema(); let field = schema.column_with_name(column).expect_ok()?.1; match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => Ok(input), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(input), DataType::List(field) | DataType::LargeList(field) if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) => { @@ -3274,6 +3356,7 @@ mod tests { token_set_format, worker_memory_limit_bytes: u64::MAX, block_size: params.block_size, + coordinate_rank: 0, }, ) .await?; @@ -3298,6 +3381,7 @@ mod tests { token_set_format, worker_memory_limit_bytes: u64::MAX, block_size: params.block_size, + coordinate_rank: 0, }, ) .await?; @@ -3781,6 +3865,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; @@ -3813,6 +3898,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; @@ -3852,6 +3938,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 07795209331..bfeafb5d6ac 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -18,7 +18,7 @@ use std::{ use crate::metrics::NoOpMetricsCollector; use crate::prefilter::NoFilter; use crate::scalar::registry::{TrainingCriteria, TrainingOrdering}; -use arrow::array::{FixedSizeListBuilder, Float32Builder}; +use arrow::array::{BooleanBuilder, FixedSizeListBuilder, Float32Builder}; use arrow::datatypes::{self, Float32Type, Int32Type, UInt64Type}; use arrow::{ array::{ @@ -28,8 +28,8 @@ use arrow::{ }; use arrow::{buffer::ScalarBuffer, datatypes::UInt32Type}; use arrow_array::{ - Array, ArrayRef, Float32Array, LargeBinaryArray, ListArray, OffsetSizeTrait, RecordBatch, - UInt32Array, UInt64Array, + Array, ArrayRef, BooleanArray, Float32Array, LargeBinaryArray, ListArray, OffsetSizeTrait, + RecordBatch, UInt32Array, UInt64Array, }; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use async_trait::async_trait; @@ -58,7 +58,7 @@ use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder} use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; -use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; +use super::{DocumentGranularity, InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ BLOCK_SIZE, ScoredDoc, doc_file_path, @@ -93,6 +93,8 @@ use std::str::FromStr; pub const INVERTED_INDEX_VERSION_V1: u32 = 1; pub const INVERTED_INDEX_VERSION_V2: u32 = 2; pub const INVERTED_INDEX_VERSION_V3: u32 = 3; +/// Element-document indexes add persisted document coordinates. +pub const INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT: u32 = 4; pub const TOKENS_FILE: &str = "tokens.lance"; pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; @@ -113,6 +115,8 @@ pub const MAX_SCORE_COL: &str = "_max_score"; pub const LENGTH_COL: &str = "_length"; pub const BLOCK_MAX_SCORE_COL: &str = "_block_max_score"; pub const NUM_TOKEN_COL: &str = "_num_tokens"; +pub const DOC_INDEX_COL: &str = "_doc_index"; +pub const DOC_INDEX_STORAGE_PREFIX: &str = "_doc_index_"; pub const SCORE_COL: &str = "_score"; pub const TOKEN_SET_FORMAT_KEY: &str = "token_set_format"; pub const POSTING_TAIL_CODEC_KEY: &str = "posting_tail_codec"; @@ -132,8 +136,30 @@ pub const ESTIMATED_MAX_TOKENS_PER_ROW: usize = 4 * 1024; pub static SCORE_FIELD: LazyLock = LazyLock::new(|| Field::new(SCORE_COL, DataType::Float32, true)); +pub static DOC_INDEX_FIELD: LazyLock = LazyLock::new(|| { + Field::new( + DOC_INDEX_COL, + DataType::List(Arc::new(Field::new("item", DataType::UInt32, false))), + false, + ) +}); pub static FTS_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(Schema::new(vec![ROW_ID_FIELD.clone(), SCORE_FIELD.clone()]))); +pub static ELEMENT_FTS_SCHEMA: LazyLock = LazyLock::new(|| { + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + DOC_INDEX_FIELD.clone(), + SCORE_FIELD.clone(), + ])) +}); + +pub fn fts_schema(document_granularity: DocumentGranularity) -> SchemaRef { + if document_granularity.is_list_element() { + ELEMENT_FTS_SCHEMA.clone() + } else { + FTS_SCHEMA.clone() + } +} static ROW_ID_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(Schema::new(vec![ROW_ID_FIELD.clone()]))); @@ -551,7 +577,33 @@ impl DeepSizeOf for InvertedIndex { async fn resolve_deferred_candidates( docs: &LazyDocSet, candidates: &mut [DocCandidate], + coordinate_rank: usize, ) -> Result<()> { + if coordinate_rank > 0 { + let doc_ids = candidates + .iter() + .map(|candidate| { + u32::try_from(candidate.posting_doc_id).map_err(|_| { + Error::internal(format!( + "FTS partition-local doc id {} exceeds u32", + candidate.posting_doc_id + )) + }) + }) + .collect::>>()?; + let keys = docs.resolve_document_keys(&doc_ids).await?; + if keys.len() != candidates.len() { + return Err(Error::internal( + "resolve_document_keys returned an unexpected number of items".to_string(), + )); + } + for (candidate, (row_id, doc_index)) in candidates.iter_mut().zip(keys) { + candidate.addr = CandidateAddr::RowId(row_id); + candidate.doc_index = doc_index; + } + return Ok(()); + } + let pending: Vec = candidates .iter() .filter_map(|c| match c.addr { @@ -580,6 +632,9 @@ impl InvertedIndex { } fn index_version(&self) -> u32 { + if self.params.get_document_granularity().is_list_element() { + return INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT; + } match (self.token_set_format, self.format_version()) { ( TokenSetFormat::Arrow, @@ -882,6 +937,26 @@ impl InvertedIndex { metrics: Arc, base_scorer: Option<&MemBM25Scorer>, ) -> Result<(Vec, Vec)> { + let documents = self + .bm25_search_documents(tokens, params, operator, prefilter, metrics, base_scorer) + .await?; + Ok(documents + .into_iter() + .map(|document| (document.row_id, document.score.0)) + .unzip()) + } + + /// Search logical FTS documents, retaining element coordinates when present. + #[instrument(level = "debug", skip_all)] + pub async fn bm25_search_documents( + &self, + tokens: Arc, + params: Arc, + operator: Operator, + prefilter: Arc, + metrics: Arc, + base_scorer: Option<&MemBM25Scorer>, + ) -> Result> { // Fuzzy expansion runs once here, with the global `max_expansions` // budget, instead of once per partition: partitions receive the // final token list, so the matched terms cannot depend on how the @@ -896,7 +971,7 @@ impl InvertedIndex { .map(|idx| expanded.position(idx)) .collect::>(); if (0..tokens.len()).any(|idx| !surviving.contains(&tokens.position(idx))) { - return Ok((Vec::new(), Vec::new())); + return Ok(Vec::new()); } } expanded @@ -919,13 +994,14 @@ impl InvertedIndex { let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { - return Ok((Vec::new(), Vec::new())); + return Ok(Vec::new()); } fn push_scored_candidate( candidates: &mut BinaryHeap>, limit: usize, addr: CandidateAddr, + doc_index: Vec, score: f32, ) -> Result<()> { // resolve_deferred_candidates ran upstream, so every candidate @@ -940,10 +1016,10 @@ impl InvertedIndex { }; if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(ScoredDoc::with_doc_index(row_id, doc_index, score))); } else if candidates.peek().unwrap().0.score.0 < score { candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(ScoredDoc::with_doc_index(row_id, doc_index, score))); } Ok(()) } @@ -957,6 +1033,11 @@ impl InvertedIndex { // global k-th - see `Wand::shared_threshold`). let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let coordinate_rank = self + .partitions + .first() + .map(|partition| partition.docs.coordinate_rank()) + .unwrap_or(0); let parts = self .partitions .iter() @@ -1042,8 +1123,12 @@ impl InvertedIndex { grouped_expansions, candidates, }; - resolve_deferred_candidates(&part.docs, &mut partition_result.candidates) - .await?; + resolve_deferred_candidates( + &part.docs, + &mut partition_result.candidates, + coordinate_rank, + ) + .await?; Result::Ok(partition_result) } }) @@ -1075,6 +1160,7 @@ impl InvertedIndex { if grouped_expansions.is_empty() { for DocCandidate { addr, + doc_index, freqs, doc_length, .. @@ -1086,7 +1172,7 @@ impl InvertedIndex { score += idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, addr, doc_index, score)?; } } else { let grouped_positions = grouped_expansions @@ -1095,6 +1181,7 @@ impl InvertedIndex { .collect::>(); for DocCandidate { addr, + doc_index, posting_doc_id, freqs, doc_length, @@ -1125,7 +1212,7 @@ impl InvertedIndex { score += idf_weight * scorer.doc_weight(freq, doc_length); } } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, addr, doc_index, score)?; } } } @@ -1133,8 +1220,8 @@ impl InvertedIndex { Ok(candidates .into_sorted_vec() .into_iter() - .map(|Reverse(doc)| (doc.row_id, doc.score.0)) - .unzip()) + .map(|Reverse(document)| document) + .collect()) } async fn load_legacy_index( @@ -1306,6 +1393,34 @@ impl InvertedIndex { .try_collect::>() .await?; + let coordinate_rank = partitions + .first() + .map(|partition| partition.docs.coordinate_rank()) + .unwrap_or(0); + if partitions + .iter() + .any(|partition| partition.docs.coordinate_rank() != coordinate_rank) + { + return Err(Error::index( + "FTS partitions have inconsistent document coordinate ranks".to_string(), + )); + } + match params.get_document_granularity() { + DocumentGranularity::Row if coordinate_rank != 0 => { + return Err(Error::index( + "row-document FTS metadata has persisted document coordinates" + .to_string(), + )); + } + DocumentGranularity::ListElement if coordinate_rank == 0 => { + return Err(Error::index( + "ListElement FTS metadata is missing persisted document coordinates" + .to_string(), + )); + } + _ => {} + } + let tokenizer = params.build()?; Ok(Arc::new(Self { params, @@ -1774,7 +1889,10 @@ impl InvertedPartition { // the store + path instead of an open reader keeps a cached partition // from pinning a docs-file handle for its whole lifetime. let docs_path = doc_file_path(id); - let num_docs = store.open_index_file(&docs_path).await?.num_rows(); + let docs_reader = store.open_index_file(&docs_path).await?; + let num_docs = docs_reader.num_rows(); + let docs_schema = arrow_schema::Schema::from(docs_reader.schema()); + let coordinate_rank = document_coordinate_rank(&docs_schema); let docs = Arc::new(LazyDocSet::new( store.clone(), docs_path, @@ -1783,6 +1901,7 @@ impl InvertedPartition { frag_reuse_index, // V3 (256-doc block) partitions score with quantized doc lengths. inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, + coordinate_rank, )); Ok(Self { @@ -6137,6 +6256,9 @@ const fn build_dequantized_doc_lengths() -> [u32; 256] { pub struct DocSet { row_ids: Vec, num_tokens: Vec, + // One flat u32 column per list boundary. This avoids a Vec allocation per + // document while preserving the full logical document coordinate. + doc_indices: Vec>, // (row_id, doc_id) pairs sorted by row_id inv: Vec<(u64, u32)>, @@ -6154,6 +6276,7 @@ impl DeepSizeOf for DocSet { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { self.row_ids.deep_size_of_children(context) + self.num_tokens.deep_size_of_children(context) + + self.doc_indices.deep_size_of_children(context) + self.inv.deep_size_of_children(context) + self .norms @@ -6164,6 +6287,13 @@ impl DeepSizeOf for DocSet { } impl DocSet { + pub(crate) fn with_coordinate_rank(coordinate_rank: usize) -> Self { + Self { + doc_indices: (0..coordinate_rank).map(|_| Vec::new()).collect(), + ..Default::default() + } + } + #[inline] pub fn len(&self) -> usize { // Use num_tokens instead of row_ids so the deferred-row_ids @@ -6195,11 +6325,26 @@ impl DocSet { self.row_ids[doc_id as usize] } + pub fn doc_index(&self, doc_id: u32) -> Vec { + self.doc_indices + .iter() + .map(|coordinates| coordinates[doc_id as usize]) + .collect() + } + + pub fn coordinate(&self, doc_id: u32, rank: usize) -> u32 { + self.doc_indices[rank][doc_id as usize] + } + + pub fn coordinate_rank(&self) -> usize { + self.doc_indices.len() + } + /// Resolve a `row_id` to every `doc_id` it owns. /// - /// Modern indexes map each row to a single document. Older list indexes - /// may have indexed each list element as its own document, so a single - /// `row_id` can still own several `doc_id`s sharing that key in `inv`. + /// Row-document indexes map each row to a single document. Element-document + /// indexes (and older list indexes) can map one row to several documents, + /// so a single `row_id` may own multiple `doc_id`s sharing that key in `inv`. /// The prefilter path (`flat_search`) walks an allow-list of row_ids and /// must evaluate all legacy documents for that row. pub fn doc_ids(&self, row_id: u64) -> impl Iterator + '_ { @@ -6269,18 +6414,27 @@ impl DocSet { let row_id_col = UInt64Array::from_iter_values(self.row_ids.iter().cloned()); let num_tokens_col = UInt32Array::from_iter_values(self.num_tokens.iter().cloned()); - let schema = arrow_schema::Schema::new(vec![ + let mut fields = vec![ arrow_schema::Field::new(ROW_ID, DataType::UInt64, false), arrow_schema::Field::new(NUM_TOKEN_COL, DataType::UInt32, false), - ]); + ]; + let mut columns = vec![ + Arc::new(row_id_col) as ArrayRef, + Arc::new(num_tokens_col) as ArrayRef, + ]; + for (rank, coordinates) in self.doc_indices.iter().enumerate() { + fields.push(arrow_schema::Field::new( + doc_index_storage_column(rank), + DataType::UInt32, + false, + )); + columns.push( + Arc::new(UInt32Array::from_iter_values(coordinates.iter().copied())) as ArrayRef, + ); + } + let schema = arrow_schema::Schema::new(fields); - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(row_id_col) as ArrayRef, - Arc::new(num_tokens_col) as ArrayRef, - ], - )?; + let batch = RecordBatch::try_new(Arc::new(schema), columns)?; Ok(batch) } @@ -6292,7 +6446,21 @@ impl DocSet { let batch = reader.read_range(0..reader.num_rows(), None).await?; let row_id_col = batch[ROW_ID].as_primitive::(); let num_tokens_col = batch[NUM_TOKEN_COL].as_primitive::(); - Self::from_columns(row_id_col, num_tokens_col, is_legacy, frag_reuse_index) + let mut doc_indices = Vec::new(); + for rank in 0.. { + let column_name = doc_index_storage_column(rank); + let Some(column) = batch.column_by_name(&column_name) else { + break; + }; + doc_indices.push(column.as_primitive::()); + } + Self::from_columns_with_doc_indices( + row_id_col, + num_tokens_col, + &doc_indices, + is_legacy, + frag_reuse_index, + ) } /// Build a `DocSet` carrying only the per-doc `num_tokens` array; @@ -6306,6 +6474,7 @@ impl DocSet { Self { row_ids: Vec::new(), num_tokens, + doc_indices: Vec::new(), inv: Vec::new(), total_tokens, scoring_quantized: false, @@ -6323,6 +6492,34 @@ impl DocSet { is_legacy: bool, frag_reuse_index: Option>, ) -> Result { + Self::from_columns_with_doc_indices( + row_id_col, + num_tokens_col, + &[], + is_legacy, + frag_reuse_index, + ) + } + + pub fn from_columns_with_doc_indices( + row_id_col: &UInt64Array, + num_tokens_col: &arrow_array::UInt32Array, + doc_index_cols: &[&arrow_array::UInt32Array], + is_legacy: bool, + frag_reuse_index: Option>, + ) -> Result { + if doc_index_cols + .iter() + .any(|column| column.len() != row_id_col.len()) + { + return Err(Error::index( + "FTS document coordinate columns must have the same length as row ids".to_string(), + )); + } + let doc_indices = doc_index_cols + .iter() + .map(|column| column.values().to_vec()) + .collect::>(); // for legacy format, the row id is doc id; sorting keeps binary search viable if is_legacy { let (row_ids, num_tokens): (Vec<_>, Vec<_>) = row_id_col @@ -6343,6 +6540,7 @@ impl DocSet { return Ok(Self { row_ids, num_tokens, + doc_indices, inv: Vec::new(), total_tokens, scoring_quantized: false, @@ -6386,6 +6584,7 @@ impl DocSet { return Ok(Self { row_ids, num_tokens, + doc_indices, inv, total_tokens, scoring_quantized: false, @@ -6407,6 +6606,7 @@ impl DocSet { Ok(Self { row_ids, num_tokens, + doc_indices, inv, total_tokens, scoring_quantized: false, @@ -6421,6 +6621,11 @@ impl DocSet { let len = self.len(); let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, Vec::with_capacity(len)); + let doc_indices = std::mem::take(&mut self.doc_indices); + self.doc_indices = doc_indices + .iter() + .map(|_| Vec::with_capacity(len)) + .collect(); self.invalidate_norms(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { @@ -6428,6 +6633,11 @@ impl DocSet { Some(Some(new_row_id)) => { self.row_ids.push(new_row_id); self.num_tokens.push(num_token); + for (new_coordinates, old_coordinates) in + self.doc_indices.iter_mut().zip(&doc_indices) + { + new_coordinates.push(old_coordinates[doc_id]); + } self.total_tokens += num_token as u64; } Some(None) => { @@ -6436,6 +6646,11 @@ impl DocSet { None => { self.row_ids.push(row_id); self.num_tokens.push(num_token); + for (new_coordinates, old_coordinates) in + self.doc_indices.iter_mut().zip(&doc_indices) + { + new_coordinates.push(old_coordinates[doc_id]); + } self.total_tokens += num_token as u64; } } @@ -6501,6 +6716,32 @@ impl DocSet { self.row_ids.len() as u32 - 1 } + pub fn append_with_doc_index( + &mut self, + row_id: u64, + num_tokens: u32, + doc_index: &[u32], + ) -> Result { + if self.row_ids.is_empty() && self.doc_indices.is_empty() { + self.doc_indices = (0..doc_index.len()).map(|_| Vec::new()).collect(); + } + if self.doc_indices.len() != doc_index.len() { + return Err(Error::index(format!( + "all documents in an FTS partition must have the same coordinate rank: expected {}, got {}", + self.doc_indices.len(), + doc_index.len() + ))); + } + self.row_ids.push(row_id); + self.num_tokens.push(num_tokens); + for (coordinates, value) in self.doc_indices.iter_mut().zip(doc_index) { + coordinates.push(*value); + } + self.total_tokens += num_tokens as u64; + self.invalidate_norms(); + Ok(self.row_ids.len() as u32 - 1) + } + // Drop the baked norm slab after a mutation; it re-bakes on the next // scoring use. fn invalidate_norms(&mut self) { @@ -6512,10 +6753,29 @@ impl DocSet { pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.capacity() * std::mem::size_of::() + + self + .doc_indices + .iter() + .map(|coordinates| coordinates.capacity() * std::mem::size_of::()) + .sum::() + self.inv.capacity() * std::mem::size_of::<(u64, u32)>() } } +pub fn doc_index_storage_column(rank: usize) -> String { + format!("{DOC_INDEX_STORAGE_PREFIX}{rank}") +} + +pub fn document_coordinate_rank(schema: &arrow_schema::Schema) -> usize { + (0..) + .take_while(|rank| { + schema + .column_with_name(&doc_index_storage_column(*rank)) + .is_some() + }) + .count() +} + pub fn flat_full_text_search( batches: &[&RecordBatch], doc_col: &str, @@ -6526,20 +6786,23 @@ pub fn flat_full_text_search( return Ok(vec![]); } - if is_phrase_query(query) { - return Err(Error::invalid_input( - "phrase query is not supported for flat full text search, try using FTS index", - )); - } + let (query, phrase_slop) = match phrase_query_text(query) { + Some(query) => (query, Some(0)), + None => (query, None), + }; match batches[0][doc_col].data_type() { - DataType::Utf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), - DataType::LargeUtf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), + DataType::Utf8 => { + do_flat_full_text_search::(batches, doc_col, query, tokenizer, phrase_slop) + } + DataType::LargeUtf8 => { + do_flat_full_text_search::(batches, doc_col, query, tokenizer, phrase_slop) + } DataType::List(_) => { - do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer, phrase_slop) } DataType::LargeList(_) => { - do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer, phrase_slop) } data_type => Err(Error::invalid_input(format!( "unsupported data type {} for inverted index", @@ -6553,6 +6816,7 @@ fn do_flat_full_text_search( doc_col: &str, query: &str, tokenizer: Option>, + phrase_slop: Option, ) -> Result> { let mut results = Vec::new(); let mut tokenizer = @@ -6564,11 +6828,8 @@ fn do_flat_full_text_search( let doc_array = batch[doc_col].as_string::(); for i in 0..row_id_array.len() { let doc = doc_array.value(i); - if has_query_token(doc, &mut tokenizer, &query_tokens) { + if document_matches_flat_query(doc, &mut tokenizer, &query_tokens, phrase_slop)? { results.push(row_id_array.value(i)); - // What is this assertion for? Why would doc contain query? Don't we reach - // here only if they share at least one token? Why is it not debug_assert? - assert!(doc.contains(query)); } } } @@ -6581,6 +6842,7 @@ fn do_flat_full_text_search_list( doc_col: &str, query: &str, tokenizer: Option>, + phrase_slop: Option, ) -> Result> { let mut results = Vec::new(); let mut tokenizer = @@ -6604,10 +6866,18 @@ fn do_flat_full_text_search_list( continue; } let elements = doc_array.value(i); - if iter_str_array(elements.as_ref()) - .flatten() - .any(|element| has_query_token(element, &mut tokenizer, &query_tokens)) - { + let matches = if phrase_slop.is_some() { + let document = iter_str_array(elements.as_ref()) + .flatten() + .collect::>() + .join(" "); + document_matches_flat_query(&document, &mut tokenizer, &query_tokens, phrase_slop)? + } else { + iter_str_array(elements.as_ref()) + .flatten() + .any(|element| has_query_token(element, &mut tokenizer, &query_tokens)) + }; + if matches { results.push(row_id_array.value(i)); } } @@ -6616,9 +6886,82 @@ fn do_flat_full_text_search_list( Ok(results) } -const FLAT_ROW_ID_COL_IDX: usize = 0; -const FLAT_ALL_TOKENS_COL_IDX: usize = 1; -const FLAT_QUERY_TOKEN_COUNTS_COL_IDX: usize = 2; +fn document_matches_flat_query( + document: &str, + tokenizer: &mut Box, + query_tokens: &Tokens, + phrase_slop: Option, +) -> Result { + let Some(slop) = phrase_slop else { + return Ok(has_query_token(document, tokenizer, query_tokens)); + }; + + let mut document_positions = (0..query_tokens.len()) + .map(|_| Vec::new()) + .collect::>(); + let mut stream = tokenizer.token_stream_for_doc(document); + while let Some(token) = stream.next() { + let position = u32::try_from(token.position).map_err(|_| { + Error::invalid_input(format!( + "flat FTS token position exceeds u32: {}", + token.position + )) + })?; + for (query_index, positions) in document_positions.iter_mut().enumerate() { + if query_tokens.get_token(query_index) == token.text { + positions.push(position); + } + } + } + Ok(phrase_matches_positions( + query_tokens, + &document_positions, + slop, + )) +} + +const FLAT_ALL_TOKENS_COL: &str = "all_tokens"; +const FLAT_QUERY_TOKEN_COUNTS_COL: &str = "query_token_counts"; +const FLAT_PHRASE_MATCH_COL: &str = "phrase_match"; + +fn phrase_matches_positions( + query_tokens: &Tokens, + document_positions: &[Vec], + slop: u32, +) -> bool { + let Some(first_positions) = document_positions.first() else { + return false; + }; + if first_positions.is_empty() { + return false; + } + + let mut candidates = first_positions.clone(); + debug_assert_eq!(query_tokens.len(), document_positions.len()); + for (query_index, positions) in document_positions.iter().enumerate().skip(1) { + let Some(query_delta) = query_tokens + .position(query_index) + .checked_sub(query_tokens.position(query_index - 1)) + else { + return false; + }; + let mut next_candidates = Vec::new(); + for &position in positions { + let position = u64::from(position); + if candidates.iter().any(|candidate| { + let least = u64::from(*candidate) + u64::from(query_delta); + least <= position && position <= least + u64::from(slop) + }) { + next_candidates.push(position as u32); + } + } + if next_candidates.is_empty() { + return false; + } + candidates = next_candidates; + } + true +} /// If we accumulate this many bytes we warn the user they probably want to use an FTS index instead. const BYTES_ACCUMULATED_WARNING_THRESHOLD: u64 = 1024 * 1024 * 1024; // 1GB @@ -6642,19 +6985,29 @@ async fn tokenize_and_count( query_tokens: Arc, doc_col_idx: usize, elapsed_compute: Option