diff --git a/docs/src/guide/migration.md b/docs/src/guide/migration.md index 5efd7b26b7f..6b999558e7d 100644 --- a/docs/src/guide/migration.md +++ b/docs/src/guide/migration.md @@ -8,15 +8,15 @@ migrate. ## 9.0.0 -* Newly created FTS / inverted indexes now default to format v2 instead of v1. +* Newly created FTS / inverted indexes now default to format v4 instead of v1. The `LANCE_FTS_FORMAT_VERSION` environment variable no longer controls the - format used for newly created indexes. Users who need a specific index layout - should pass the index creation parameter `format_version` explicitly. + format used for newly created indexes. Users who need a specific older index + layout should pass the index creation parameter `format_version` explicitly. * This affects users who create FTS / inverted indexes and need those indexes to be readable by older Lance versions, or who depend on the v1 index layout. In those cases, pass `format_version=1` when creating the index. Otherwise, newly - created indexes will use v2 by default, and older Lance readers may not be able + created indexes will use v4 by default, and older Lance readers may not be able to read them. ```python 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..900d6b9b2fd 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -264,16 +264,16 @@ public Builder skipMerge(boolean skipMerge) { /** * Configure the on-disk FTS format version to write when creating a new index. * - *

If unset, Lance writes v2 for {@code blockSize = 128} and v3 for {@code blockSize = 256}. - * {@code formatVersion = 3} is experimental and is only valid with {@code blockSize = 256}. + *

If unset, Lance writes v4 for either supported block size. {@code formatVersion = 3} is + * experimental and is only valid with {@code blockSize = 256}. * - * @param formatVersion FTS format version, must be 1, 2, or 3 + * @param formatVersion FTS format version, must be 1, 2, 3, or 4 * @return this builder * @throws IllegalArgumentException */ public Builder formatVersion(int formatVersion) { - if (formatVersion != 1 && formatVersion != 2 && formatVersion != 3) { - throw new IllegalArgumentException("formatVersion must be 1, 2, or 3"); + if (formatVersion != 1 && formatVersion != 2 && formatVersion != 3 && formatVersion != 4) { + throw new IllegalArgumentException("formatVersion must be 1, 2, 3, or 4"); } this.formatVersion = formatVersion; return this; @@ -283,8 +283,10 @@ public Builder formatVersion(int formatVersion) { public ScalarIndexParams build() { if (formatVersion != null) { Preconditions.checkArgument( - (blockSize == 256 && formatVersion == 3) || (blockSize == 128 && formatVersion != 3), - "formatVersion 3 requires blockSize 256, and blockSize 256 requires formatVersion 3"); + formatVersion == 4 + || (blockSize == 256 && formatVersion == 3) + || (blockSize == 128 && formatVersion != 3), + "formatVersion 3 requires blockSize 256, and legacy formats require blockSize 128"); } Map params = new HashMap<>(); if (baseTokenizer != null) { 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..8618d3ce091 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -75,4 +75,15 @@ void formatVersionThreeRequiresBlockSize256() { IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(256).formatVersion(2).build()); } + + @Test + void formatVersionFourSupportsBothBlockSizes() { + for (int blockSize : new int[] {128, 256}) { + ScalarIndexParams params = + InvertedIndexParams.builder().blockSize(blockSize).formatVersion(4).build(); + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(blockSize, ((Number) json.get("block_size")).intValue()); + assertEquals(4, ((Number) json.get("format_version")).intValue()); + } + } } diff --git a/protos/index_old.proto b/protos/index_old.proto index eb984d6fe29..10b4dfa602d 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -26,6 +26,27 @@ message LabelListIndexDetails {} message NGramIndexDetails {} message ZoneMapIndexDetails {} message InvertedIndexDetails { + message CodeTokenizerConfig { + // Split one lexical identifier into subwords, e.g. getUserName -> + // get/user/name. + bool split_identifiers = 1; + // Split identifier subwords across letter/number boundaries, e.g. + // HTML2JSON -> html/2/json. An absent value uses the code tokenizer default; + // a present value records the explicit index-time choice. + optional bool split_on_numerics = 2; + // Keep the complete lexical identifier in addition to subwords, e.g. + // user_name plus user/name. An absent value uses the code tokenizer default; + // a present value records the explicit index-time choice. + optional bool preserve_original = 3; + // Index operator tokens such as "::", "->", and "!=". Operators are not + // indexed by default because they are often high-frequency noise. + bool index_operators = 4; + } + + // Lexical tokenizer used after document-level text extraction. This is an + // implementation component such as "simple", "icu", "ngram", or "code". + // Input-time analyzer profiles are expanded into this field and the concrete + // options below before these details are persisted. // 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; @@ -41,7 +62,11 @@ message InvertedIndexDetails { bool prefix_only = 11; // Number of documents per compressed posting block. An absent value means // the index predates this field and must use the legacy block size of 128. - // A present value records the block size used by the index; 256 is only - // valid with format version 3. + // A present value records the block size used by the index; 256 is valid + // with format versions 3 and 4. optional uint32 block_size = 12; + // Options for base_tokenizer = "code". Presence records the code tokenizer + // configuration used to build the index; absence means there is no + // code-specific configuration to apply. + CodeTokenizerConfig code_config = 13; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..f4d5ccd8380 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3370,8 +3370,8 @@ def create_scalar_index( format_version: int or str, optional This is for the ``INVERTED`` / ``FTS`` index. Explicit on-disk FTS format version to write when creating a new index. Accepts ``1``, - ``2``, ``3``, ``"v1"``, ``"v2"``, or ``"v3"``. If unset, Lance - writes v2 for ``block_size=128`` and v3 for ``block_size=256``. + ``2``, ``3``, ``4``, ``"v1"``, ``"v2"``, ``"v3"``, or ``"v4"``. + If unset, Lance writes v4 for either supported block size. ``format_version=3`` is experimental and is only valid with ``block_size=256``. diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 985072fe5de..0c5b41592ae 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, + FullTextOperator, MatchQuery, MultiMatchQuery, Occur, @@ -806,6 +807,321 @@ def test_full_text_search(dataset, with_position, base_tokenizer): ) +def test_code_analyzer_does_not_split_identifiers_by_default(tmp_path): + table = pa.table({"code": ["GetUserName", "GetUserEmail", "user"]}) + ds = lance.write_dataset(table, tmp_path) + ds.create_scalar_index("code", index_type="INVERTED", analyzer="code") + + results = ds.to_table( + columns=["code"], + full_text_query=MatchQuery("user", "code"), + ) + assert results["code"].to_pylist() == ["user"] + + stats = ds.stats.index_stats("code_idx")["indices"][0] + params = stats["params"] + assert "analyzer" not in params + assert params["base_tokenizer"] == "code" + assert params["split_identifiers"] is False + + +def test_code_analyzer_full_text_search_with_identifier_splitting(tmp_path): + table = pa.table( + { + "code": [ + "getUserName", + "set_user_name", + "user-name", + "username", + "other", + ] + } + ) + ds = lance.write_dataset(table, tmp_path) + ds.create_scalar_index( + "code", + index_type="INVERTED", + analyzer="code", + split_identifiers=True, + ) + + results = ds.to_table( + columns=["code"], + full_text_query=MatchQuery("user", "code"), + ) + assert set(results["code"].to_pylist()) == { + "getUserName", + "set_user_name", + "user-name", + } + + stats = ds.stats.index_stats("code_idx")["indices"][0] + params = stats["params"] + assert "analyzer" not in params + assert params["base_tokenizer"] == "code" + assert params["split_identifiers"] is True + assert params["split_on_numerics"] is True + assert params["preserve_original"] is True + assert params["stem"] is False + assert params["remove_stop_words"] is False + + +def test_code_analyzer_operator_search_matches_rust_turbofish(tmp_path): + table = pa.table( + { + "path": ["turbofish.rs", "comparison.rs"], + "code": ["value.parse::()", "value.parse()"], + } + ) + ds = lance.write_dataset(table, tmp_path) + ds.create_scalar_index( + "code", + index_type="INVERTED", + analyzer="code", + index_operators=True, + ) + + results = ds.to_table( + columns=["path"], + full_text_query=MatchQuery("::", "code", operator=FullTextOperator.OR), + ) + assert results["path"].to_pylist() == ["turbofish.rs"] + + +def test_code_analyzer_exact_identifier_survives_grouped_top_k(tmp_path): + table = pa.table( + { + "path": ["split_0.rs", "split_1.rs", "split_2.rs", "exact.rs"], + "code": ["get user name", "get user name", "get user name", "getUserName"], + } + ) + ds = lance.write_dataset(table, tmp_path) + ds.create_scalar_index( + "code", + index_type="INVERTED", + analyzer="code", + split_identifiers=True, + ) + + results = ds.scanner( + columns=["path", "_score"], + full_text_query=MatchQuery( + "getUserName", "code", operator=FullTextOperator.AND + ), + limit=1, + ).to_table() + assert results["path"].to_pylist() == ["exact.rs"] + + +def test_code_analyzer_flags_require_code_analyzer(tmp_path): + table = pa.table({"text": ["getUserName"]}) + ds = lance.write_dataset(table, tmp_path) + + with pytest.raises(ValueError, match="code analyzer flags require analyzer='code'"): + ds.create_scalar_index( + "text", + index_type="INVERTED", + split_identifiers=True, + ) + + +def test_code_analyzer_complex_code_constructs(tmp_path): + table = pa.table( + { + "path": [ + "edge/trait.rs", + "edge/impl.rs", + "edge/fn_pointer.rs", + "edge/unit_result.rs", + "edge/hrtb.rs", + "edge/associated.rs", + "edge/operators.rs", + ], + "code": [ + """ +pub trait EdgeAsyncRepository<'a, T: Send + Sync> +where + T: TryFrom<&'a str, Error = EdgeParseError>, +{ + type Output<'b>: Iterator> + where + Self: 'b; + + async fn fetch_by_key( + &'a self, + key: [u8; N], + ) -> Result, EdgeRepoError>; +} +""", + """ +impl<'a, T, S> EdgeAsyncRepository<'a, T> for EdgeStore +where + T: TryFrom<&'a str, Error = EdgeParseError> + Clone + Send + Sync, + S: EdgeBackend + ?Sized, +{ + type Output<'b> = std::vec::IntoIter> where Self: 'b; + + async fn fetch_by_key( + &'a self, + key: [u8; N], + ) -> Result, EdgeRepoError> { + self.backend.fetch::(key).await + } +} +""", + """ +pub fn build_edge_handler( + factory: F, +) -> impl Fn() -> Result, EdgeError> +where + F: FnOnce() -> Result + Send + 'static, + T: Default + Send + Sync + 'static, +{ + move || factory().map(EdgeHandler::new) +} +""", + """ +pub fn edge_unit_result_callback() -> Result<()> { + Ok(()) +} +""", + """ +pub fn edge_higher_ranked<'a, T>( + visitor: impl for<'b> Fn(&'b T) -> Result<&'b str, EdgeVisitError>, + value: &'a T, +) -> Result<&'a str, EdgeVisitError> { + visitor(value) +} +""", + """ +pub fn edge_collect_stream(items: I) -> Result, E::Error> +where + I: IntoIterator, + E: EdgeExtract, +{ + items.into_iter().map(E::extract).collect() +} +""", + """ +pub fn edge_operator_arrow() -> Result { + let variant = EdgeModule::EdgeVariant; + if variant != EdgeModule::Default && EdgeMask::enabled() { + return Ok(EdgeArrow::new(variant)); + } + Err(EdgeError::empty()) +} +""", + ], + } + ) + table = table.append_column("code_ops", table["code"]) + ds = lance.write_dataset(table, tmp_path) + ds.create_scalar_index("code", index_type="INVERTED", analyzer="code") + ds.create_scalar_index( + "code_ops", + index_type="INVERTED", + analyzer="code", + index_operators=True, + ) + + ds.insert( + pa.table( + { + "path": ["edge/flat_unindexed.rs"], + "code": [ + """ +pub async fn edge_flat_generic_return() -> Result +where + T: TryFrom + Send, + E: Into, +{ + T::try_from(String::new()).map_err(Into::into) +} +""" + ], + "code_ops": [ + """ +pub async fn edge_flat_operator() -> Result { + EdgeFlat::try_new() -> Result +} +""" + ], + } + ) + ) + ds = lance.dataset(tmp_path) + + def assert_search(column, query, expected_path, operator=FullTextOperator.AND): + result = ds.scanner( + columns=["path", "_score"], + full_text_query=MatchQuery(query, column, operator=operator), + limit=50, + ).to_table() + assert expected_path in result["path"].to_pylist() + + assert_search( + "code", + "EdgeAsyncRepository fetch_by_key TryFrom EdgeRepoError", + "edge/trait.rs", + ) + assert_search( + "code", + "EdgeStore fetch_by_key const usize where Result", + "edge/impl.rs", + ) + assert_search( + "code", + "build_edge_handler FnOnce Result EdgeHandler", + "edge/fn_pointer.rs", + ) + assert_search( + "code", + "edge_unit_result_callback fn () -> Result", + "edge/unit_result.rs", + ) + assert_search( + "code", + "edge_higher_ranked for Fn EdgeVisitError Result", + "edge/hrtb.rs", + ) + assert_search( + "code", + "edge_collect_stream IntoIterator Item Error Result", + "edge/associated.rs", + ) + assert_search( + "code", + "edge_flat_generic_return TryFrom EdgeFlatError Result", + "edge/flat_unindexed.rs", + ) + assert_search( + "code_ops", + "edge_operator_arrow -> Result", + "edge/operators.rs", + ) + assert_search( + "code_ops", + "EdgeModule :: EdgeVariant !=", + "edge/operators.rs", + ) + assert_search( + "code_ops", + "edge_flat_operator -> Result EdgeFlatError", + "edge/flat_unindexed.rs", + ) + + default_operator_results = ds.scanner( + columns=["path", "_score"], + full_text_query=MatchQuery("->", "code", operator=FullTextOperator.OR), + ).to_table() + operator_results = ds.scanner( + columns=["path", "_score"], + full_text_query=MatchQuery("->", "code_ops", operator=FullTextOperator.OR), + ).to_table() + assert default_operator_results.num_rows == 0 + assert operator_results.num_rows > 0 + + def test_unindexed_full_text_search_on_empty_index(tmp_path): # Create fts index on empty table. schema = pa.schema({"text": pa.string()}) @@ -955,7 +1271,7 @@ def test_create_scalar_index_fts_block_size(dataset): ) indices = dataset.describe_indices() doc_index = next(index for index in indices if index.name == "doc_idx") - assert doc_index.segments[0].index_version == 3 + assert doc_index.segments[0].index_version == 4 row = dataset.take(indices=[0], columns=["doc"]) query = row.column(0)[0].as_py().split(" ")[0] @@ -5074,7 +5390,7 @@ def test_json_inverted_match_query(tmp_path): @pytest.mark.parametrize( ("format_version", "expected_format_version"), - [(1, 1), (2, 2), ("v1", 1), ("v2", 2)], + [(1, 1), (2, 2), (4, 4), ("v1", 1), ("v2", 2), ("v4", 4)], ) def test_describe_indices(tmp_path, format_version, expected_format_version): data = pa.table( @@ -5124,7 +5440,6 @@ def test_describe_indices(tmp_path, format_version, expected_format_version): assert details["lower_case"] assert details["stem"] assert details["remove_stop_words"] - assert details["custom_stop_words"] is None assert details["ascii_folding"] assert details["min_ngram_length"] == 3 assert details["max_ngram_length"] == 3 @@ -5205,7 +5520,7 @@ def test_describe_indices(tmp_path, format_version, expected_format_version): assert index.num_rows_indexed == 50 -def test_create_inverted_index_defaults_to_v2_and_ignores_env(tmp_path, monkeypatch): +def test_create_inverted_index_defaults_to_v4_and_ignores_env(tmp_path, monkeypatch): monkeypatch.setenv("LANCE_FTS_FORMAT_VERSION", "1") data = pa.table({"text": ["document about lance database"]}) ds = lance.write_dataset(data, tmp_path) @@ -5213,7 +5528,7 @@ def test_create_inverted_index_defaults_to_v2_and_ignores_env(tmp_path, monkeypa ds.create_scalar_index("text", index_type="INVERTED") indices = ds.describe_indices() - assert indices[0].segments[0].index_version == 2 + assert indices[0].segments[0].index_version == 4 def test_create_inverted_index_rejects_invalid_format_version(tmp_path): @@ -5221,7 +5536,7 @@ def test_create_inverted_index_rejects_invalid_format_version(tmp_path): ds = lance.write_dataset(data, tmp_path) with pytest.raises(ValueError, match="unsupported FTS format version"): - ds.create_scalar_index("text", index_type="INVERTED", format_version="v4") + ds.create_scalar_index("text", index_type="INVERTED", format_version="v5") with pytest.raises(ValueError, match="format_version=3"): ds.create_scalar_index("text", index_type="INVERTED", format_version="v3") diff --git a/python/src/dataset.rs b/python/src/dataset.rs index a361e0b63a8..eddf555fd35 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -70,6 +70,7 @@ use lance_core::datatypes::BlobHandling; use lance_datafusion::utils::reader_to_stream; use lance_encoding::decoder::DecoderConfig; use lance_file::reader::FileReaderOptions; +use lance_index::scalar::inverted::InvertedListFormatVersion; use lance_index::scalar::inverted::query::Occur; use lance_index::scalar::inverted::query::{ BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Operator, PhraseQuery, @@ -78,7 +79,6 @@ use lance_index::{ FtsPrewarmOptions, IndexParams, IndexType, PrewarmOptions, optimize::OptimizeOptions, progress::{IndexBuildProgress, NoopIndexBuildProgress}, - scalar::inverted::InvertedListFormatVersion, scalar::{FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams}, vector::{ ApproxMode, DEFAULT_QUERY_PARALLELISM, Query as VectorQuery, @@ -2421,19 +2421,106 @@ impl Dataset { "INVERTED" | "FTS" => { let mut params = InvertedIndexParams::default(); if let Some(kwargs) = kwargs { + let allowed_kwargs = [ + "analyzer", + "with_position", + "base_tokenizer", + "language", + "max_token_length", + "lower_case", + "stem", + "remove_stop_words", + "custom_stop_words", + "ascii_folding", + "min_ngram_length", + "max_ngram_length", + "prefix_only", + "block_size", + "split_identifiers", + "split_on_numerics", + "preserve_original", + "index_operators", + "memory_limit", + "num_workers", + "format_version", + "fragment_ids", + "index_uuid", + "progress_callback", + ]; + for (key, _) in kwargs.iter() { + let key: String = key.extract()?; + if !allowed_kwargs.contains(&key.as_str()) { + return Err(PyValueError::new_err(format!( + "unknown FTS index parameter '{}'", + key + ))); + } + } + + let analyzer: Option = kwargs + .get_item("analyzer")? + .map(|value| value.extract()) + .transpose()?; + let base_tokenizer: Option = kwargs + .get_item("base_tokenizer")? + .map(|value| value.extract()) + .transpose()?; + + match (analyzer.as_deref(), base_tokenizer.as_deref()) { + (Some("text"), Some("code")) => { + return Err(PyValueError::new_err( + "base_tokenizer='code' requires analyzer='code'", + )); + } + (Some("code"), Some(base_tokenizer)) if base_tokenizer != "code" => { + return Err(PyValueError::new_err(format!( + "analyzer='code' requires base_tokenizer='code', got '{}'", + base_tokenizer + ))); + } + _ => {} + } + + let uses_code_analyzer = match analyzer.as_deref() { + Some("code") => true, + Some("text") | None => base_tokenizer.as_deref() == Some("code"), + Some(_) => true, + }; + if !uses_code_analyzer { + for flag in [ + "split_identifiers", + "split_on_numerics", + "preserve_original", + "index_operators", + ] { + if let Some(value) = kwargs.get_item(flag)? + && value.extract::()? + { + return Err(PyValueError::new_err( + "code analyzer flags require analyzer='code'", + )); + } + } + } + + if let Some(analyzer) = analyzer { + params = params + .analyzer(&analyzer) + .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()?); } - if let Some(base_tokenizer) = kwargs.get_item("base_tokenizer")? { - params = params.base_tokenizer(base_tokenizer.extract()?); + if let Some(base_tokenizer) = base_tokenizer { + params = params.base_tokenizer(base_tokenizer); } if let Some(language) = kwargs.get_item("language")? { let language: PyBackedStr = language.cast::()?.clone().try_into()?; - params = params.language(&language).map_err(|e| { + params = params.language(&language).map_err(|err| { PyValueError::new_err(format!( - "can't set tokenizer language to {}: {:?}", - language, e + "can't set tokenizer language to {}: {}", + language, err )) })?; } @@ -2449,8 +2536,8 @@ impl Dataset { if let Some(remove_stop_words) = kwargs.get_item("remove_stop_words")? { params = params.remove_stop_words(remove_stop_words.extract()?); } - if let Some(stop_words_file) = kwargs.get_item("custom_stop_words")? { - params = params.custom_stop_words(stop_words_file.extract()?); + if let Some(custom_stop_words) = kwargs.get_item("custom_stop_words")? { + params = params.custom_stop_words(custom_stop_words.extract()?); } if let Some(ascii_folding) = kwargs.get_item("ascii_folding")? { params = params.ascii_folding(ascii_folding.extract()?); @@ -2469,6 +2556,18 @@ impl Dataset { .block_size(block_size.extract()?) .map_err(|e| PyValueError::new_err(e.to_string()))?; } + if let Some(split_identifiers) = kwargs.get_item("split_identifiers")? { + params = params.split_identifiers(split_identifiers.extract()?); + } + if let Some(split_on_numerics) = kwargs.get_item("split_on_numerics")? { + params = params.split_on_numerics(split_on_numerics.extract()?); + } + if let Some(preserve_original) = kwargs.get_item("preserve_original")? { + params = params.preserve_original(preserve_original.extract()?); + } + if let Some(index_operators) = kwargs.get_item("index_operators")? { + params = params.index_operators(index_operators.extract()?); + } if let Some(memory_limit) = kwargs.get_item("memory_limit")? { params = params.memory_limit_mb(memory_limit.extract()?); } @@ -2484,7 +2583,7 @@ impl Dataset { value.to_string() } else { return Err(PyValueError::new_err( - "format_version must be 1, 2, 3, 'v1', 'v2', or 'v3'", + "format_version must be 1, 2, 3, 4, 'v1', 'v2', 'v3', or 'v4'", )); }; let format_version = value diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 926bfad6a69..ca27367febc 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -268,7 +268,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { } fn version(&self) -> u32 { - max_supported_fts_format_version().index_version() + INVERTED_INDEX_VERSION_V4 } fn new_query_parser( @@ -320,12 +320,9 @@ mod tests { use crate::scalar::{BuiltinIndexType, ScalarIndexParams}; #[test] - fn test_plugin_version_tracks_max_supported_format() { + fn test_plugin_version_tracks_current_index_version() { let plugin = InvertedIndexPlugin; - assert_eq!( - plugin.version(), - max_supported_fts_format_version().index_version() - ); + assert_eq!(plugin.version(), INVERTED_INDEX_VERSION_V4); } #[test] diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 41804e94fd1..060e85a4d67 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1825,16 +1825,16 @@ pub(crate) fn inverted_list_schema_for_version_with_block_size_and_impacts( InvertedListFormatVersion::V1 => { inverted_list_schema_v1(with_position, block_size, with_impacts) } - InvertedListFormatVersion::V2 | InvertedListFormatVersion::V3 => { - inverted_list_schema_with_tail_codec_and_position_codec( - with_position, - format_version, - PostingTailCodec::VarintDelta, - Some(PositionStreamCodec::PackedDelta), - block_size, - with_impacts, - ) - } + InvertedListFormatVersion::V2 + | InvertedListFormatVersion::V3 + | InvertedListFormatVersion::V4 => inverted_list_schema_with_tail_codec_and_position_codec( + with_position, + format_version, + PostingTailCodec::VarintDelta, + Some(PositionStreamCodec::PackedDelta), + block_size, + with_impacts, + ), } } @@ -2105,7 +2105,6 @@ async fn merge_metadata_files( let mut params = None; let mut token_set_format = None; let mut format_version = None; - let mut posting_tail_codec = None; let mut deleted_fragments = RoaringBitmap::new(); progress .stage_start( @@ -2147,9 +2146,6 @@ async fn merge_metadata_files( if format_version.is_none() { format_version = Some(parse_format_version_from_metadata(metadata)?); } - if posting_tail_codec.is_none() { - posting_tail_codec = Some(parse_posting_tail_codec(metadata)?); - } if reader.num_rows() > 0 { let metadata_batch = reader.read_range(0..1, None).await?; @@ -2215,8 +2211,7 @@ async fn merge_metadata_files( None, deleted_fragments, ) - .with_format_version(format_version.unwrap_or(InvertedListFormatVersion::V1)) - .with_posting_tail_codec(posting_tail_codec.unwrap_or(PostingTailCodec::Fixed32)); + .with_format_version(format_version.unwrap_or(InvertedListFormatVersion::V1)); progress .stage_start("write_merged_metadata", Some(1), "files") .await?; @@ -2922,6 +2917,10 @@ mod tests { expected_partitions.dedup(); let remapped_partitions = (0..expected_partitions.len() as u64).collect::>(); assert_eq!(written_partitions, remapped_partitions); + assert_eq!( + parse_format_version_from_metadata(metadata)?, + InvertedListFormatVersion::V4 + ); for (new_id, old_id) in expected_partitions.iter().enumerate() { assert_partition_file_markers(base_store.as_ref(), new_id as u64, *old_id).await?; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index c1740545226..0ec3248be53 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::{FixedSizeListBuilder, Float32Builder, Int32Builder}; use arrow::datatypes::{self, Float32Type, Int32Type, UInt64Type}; use arrow::{ array::{ @@ -90,9 +90,11 @@ use std::str::FromStr; // Version 1: Fst TokenSetFormat with per-doc compressed positions // Version 2: Fst TokenSetFormat with shared posting-list position streams. // Version 3: Version 2 layout with 256-document physical posting blocks. +// Version 4: Code analyzer support with configurable posting block sizes. pub const INVERTED_INDEX_VERSION_V1: u32 = 1; pub const INVERTED_INDEX_VERSION_V2: u32 = 2; pub const INVERTED_INDEX_VERSION_V3: u32 = 3; +pub const INVERTED_INDEX_VERSION_V4: u32 = 4; pub const TOKENS_FILE: &str = "tokens.lance"; pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; @@ -147,7 +149,7 @@ pub fn resolve_fts_format_version( } pub fn default_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V2 + InvertedListFormatVersion::V4 } pub fn current_fts_format_version() -> InvertedListFormatVersion { @@ -155,15 +157,16 @@ pub fn current_fts_format_version() -> InvertedListFormatVersion { } pub fn max_supported_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V3 + InvertedListFormatVersion::V4 } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] pub enum InvertedListFormatVersion { V1, - #[default] V2, V3, + #[default] + V4, } impl InvertedListFormatVersion { @@ -199,25 +202,26 @@ impl InvertedListFormatVersion { Self::V1 => INVERTED_INDEX_VERSION_V1, Self::V2 => INVERTED_INDEX_VERSION_V2, Self::V3 => INVERTED_INDEX_VERSION_V3, + Self::V4 => INVERTED_INDEX_VERSION_V4, } } pub fn posting_tail_codec(self) -> PostingTailCodec { match self { Self::V1 => PostingTailCodec::Fixed32, - Self::V2 | Self::V3 => PostingTailCodec::VarintDelta, + Self::V2 | Self::V3 | Self::V4 => PostingTailCodec::VarintDelta, } } pub fn position_codec(self) -> Option { match self { Self::V1 => None, - Self::V2 | Self::V3 => Some(PositionStreamCodec::PackedDelta), + Self::V2 | Self::V3 | Self::V4 => Some(PositionStreamCodec::PackedDelta), } } pub fn uses_shared_position_stream(self) -> bool { - matches!(self, Self::V2 | Self::V3) + matches!(self, Self::V2 | Self::V3 | Self::V4) } } @@ -229,8 +233,9 @@ impl FromStr for InvertedListFormatVersion { "1" | "v1" | "V1" => Ok(Self::V1), "2" | "v2" | "V2" => Ok(Self::V2), "3" | "v3" | "V3" => Ok(Self::V3), + "4" | "v4" | "V4" => Ok(Self::V4), other => Err(Error::index(format!( - "unsupported FTS format version {}, expected 1, 2, or 3", + "unsupported FTS format version {}, expected 1, 2, 3, or 4", other ))), } @@ -241,11 +246,7 @@ pub fn default_fts_format_version_for_block_size( block_size: usize, ) -> Result { validate_block_size(block_size)?; - match block_size { - LEGACY_BLOCK_SIZE => Ok(InvertedListFormatVersion::V2), - 256 => Ok(InvertedListFormatVersion::V3), - _ => unreachable!("validate_block_size limits supported block sizes"), - } + Ok(InvertedListFormatVersion::V4) } pub fn validate_format_version_block_size( @@ -255,7 +256,8 @@ pub fn validate_format_version_block_size( validate_block_size(block_size)?; match (format_version, block_size) { (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE) - | (InvertedListFormatVersion::V3, 256) => Ok(()), + | (InvertedListFormatVersion::V3, 256) + | (InvertedListFormatVersion::V4, _) => Ok(()), (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, 256) => { Err(Error::invalid_input(format!( "FTS format_version={} is incompatible with block_size=256; use format_version=3", @@ -291,6 +293,7 @@ struct LoadedPostings { postings: Vec, grouped_expansions: Vec, impact_safe: bool, + exact_scoring_required: bool, } impl LoadedPostings { @@ -299,6 +302,7 @@ impl LoadedPostings { postings: Vec::new(), grouped_expansions: Vec::new(), impact_safe: false, + exact_scoring_required: false, } } } @@ -306,48 +310,7 @@ impl LoadedPostings { #[derive(Debug)] struct GroupedExpansionTerms { position: u32, - terms: Vec, -} - -fn grouped_rescore_wand_limit( - limit: Option, - grouped_expansions: &[GroupedExpansionTerms], -) -> Option { - let limit = limit?; - // Grouped fuzzy AND rescoring needs a small candidate cushion because WAND - // ranks by the unioned group posting first and the exact expansion IDF later. - let expansion_terms = grouped_expansions - .iter() - .map(|group| group.terms.len()) - .sum::() - .max(1); - Some(limit.saturating_mul(expansion_terms)) -} - -#[derive(Debug)] -struct ExpansionTermFreqs { - token: String, - freqs_by_posting_doc_id: Vec<(u64, u32)>, -} - -impl ExpansionTermFreqs { - fn new(token: String, posting: &PostingList) -> Self { - let freqs_by_posting_doc_id = posting - .iter() - .map(|(posting_doc_id, freq, _)| (posting_doc_id, freq)) - .collect(); - Self { - token, - freqs_by_posting_doc_id, - } - } - - fn frequency(&self, posting_doc_id: u64) -> Option { - self.freqs_by_posting_doc_id - .binary_search_by_key(&posting_doc_id, |(doc_id, _)| *doc_id) - .ok() - .map(|idx| self.freqs_by_posting_doc_id[idx].1) - } + terms: Arc<[GroupedTermScorer]>, } #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Default)] @@ -701,8 +664,7 @@ impl InvertedIndex { let mut builder = InvertedIndexBuilder::new(first.params.clone()).with_progress(progress); builder = builder .with_token_set_format(first.token_set_format) - .with_format_version(first.format_version()) - .with_posting_tail_codec(first.posting_tail_codec()); + .with_format_version(first.format_version()); let files = builder .update_from_segments(new_data, dest_store, segments, old_data_filter) .await?; @@ -983,6 +945,7 @@ impl InvertedIndex { postings, grouped_expansions, impact_safe, + exact_scoring_required, } = loaded_postings; if postings.is_empty() { // No hits in this partition; its DocSet stays @@ -1005,28 +968,17 @@ impl InvertedIndex { let mask = mask.clone(); let metrics = metrics.clone(); let part_for_wand = part.clone(); - let has_grouped_expansions = !grouped_expansions.is_empty(); - let use_impact_path = impact_safe && !has_grouped_expansions; - let wand_params = if has_grouped_expansions { - let mut rescoring_params = params.as_ref().clone(); - rescoring_params.limit = - grouped_rescore_wand_limit(params.limit, &grouped_expansions); - Arc::new(rescoring_params) - } else { - params.clone() - }; - let partition_threshold = if has_grouped_expansions { - Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) - } else if use_impact_path { + let use_global_scorer = impact_safe || exact_scoring_required; + let partition_threshold = if use_global_scorer { impact_shared_threshold } else { legacy_shared_threshold }; - let wand_scorer = use_impact_path.then(|| impact_scorer.clone()); + let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); let candidates = spawn_cpu(move || { let candidates = part_for_wand.bm25_search( docs_for_wand.as_ref(), - wand_params.as_ref(), + params.as_ref(), operator, mask, postings, @@ -1110,19 +1062,11 @@ impl InvertedIndex { * scorer.doc_weight(freq, doc_length); } for group in &grouped_expansions { - for term in &group.terms { + for term in group.terms.iter() { let Some(freq) = term.frequency(posting_doc_id) else { continue; }; - let idf_weight = match idf_cache.get(&term.token) { - Some(weight) => *weight, - None => { - let weight = scorer.query_weight(&term.token); - idf_cache.insert(term.token.clone(), weight); - weight - } - }; - score += idf_weight * scorer.doc_weight(freq, doc_length); + score += term.query_weight() * scorer.doc_weight(freq, doc_length); } } push_scored_candidate(&mut candidates, limit, addr, score)?; @@ -1877,7 +1821,53 @@ impl InvertedPartition { } } - fn union_plain_posting_lists(postings: Vec) -> Result { + #[inline] + fn grouped_score_upper_bound( + query_weight: f32, + union_freq: u32, + doc_length: u32, + scorer: &MemBM25Scorer, + ) -> f32 { + // BM25's document weight is monotonic in frequency and every IDF is + // non-negative. Scoring the summed frequency with the summed IDF is + // therefore an upper bound on the sum of the individual term scores. + query_weight * scorer.doc_weight(union_freq, doc_length) + } + + fn grouped_block_max_scores( + doc_ids: &[u32], + frequencies: &[u32], + block_size: usize, + docs: &DocSet, + query_weight: f32, + scorer: &MemBM25Scorer, + ) -> Vec { + doc_ids + .chunks(block_size) + .zip(frequencies.chunks(block_size)) + .map(|(doc_ids, frequencies)| { + doc_ids + .iter() + .zip(frequencies) + .map(|(doc_id, freq)| { + Self::grouped_score_upper_bound( + query_weight, + *freq, + docs.scoring_num_tokens(*doc_id), + scorer, + ) + }) + .fold(0.0, f32::max) + }) + .collect() + } + + fn union_plain_posting_lists( + postings: Vec, + docs: &DocSet, + query_weight: f32, + scorer: &MemBM25Scorer, + ) -> Result { let mut freqs_by_row_id = BTreeMap::new(); for posting in postings { for (row_id, freq, _) in posting.iter() { @@ -1889,21 +1879,86 @@ impl InvertedPartition { } let mut row_ids = Vec::with_capacity(freqs_by_row_id.len()); let mut frequencies = Vec::with_capacity(freqs_by_row_id.len()); + let mut max_score = 0.0_f32; for (row_id, freq) in freqs_by_row_id { + max_score = max_score.max(Self::grouped_score_upper_bound( + query_weight, + freq, + docs.num_tokens_by_row_id(row_id), + scorer, + )); row_ids.push(row_id); frequencies.push(freq as f32); } Ok(PostingList::Plain(PlainPostingList::new( ScalarBuffer::from(row_ids), ScalarBuffer::from(frequencies), - None, + Some(max_score), None, ))) } + fn union_plain_posting_lists_with_positions( + postings: Vec, + docs: &DocSet, + query_weight: f32, + scorer: &MemBM25Scorer, + ) -> Result { + let mut positions_by_row_id = BTreeMap::>::new(); + for posting in postings { + for (row_id, _, positions) in posting.iter() { + let positions = positions.ok_or_else(|| { + Error::index("cannot union grouped phrase terms without positions".to_string()) + })?; + positions_by_row_id + .entry(row_id) + .or_default() + .extend(positions); + } + } + if positions_by_row_id.is_empty() { + return Ok(PostingList::Plain(PlainPostingList::new( + ScalarBuffer::from(Vec::::new()), + ScalarBuffer::from(Vec::::new()), + None, + None, + ))); + } + + let mut row_ids = Vec::with_capacity(positions_by_row_id.len()); + let mut frequencies = Vec::with_capacity(positions_by_row_id.len()); + let mut positions_builder = ListBuilder::new(Int32Builder::new()); + let mut max_score = 0.0_f32; + for (row_id, mut positions) in positions_by_row_id { + positions.sort_unstable(); + let frequency = positions.len() as u32; + max_score = max_score.max(Self::grouped_score_upper_bound( + query_weight, + frequency, + docs.num_tokens_by_row_id(row_id), + scorer, + )); + row_ids.push(row_id); + frequencies.push(frequency as f32); + for position in positions { + positions_builder.values().append_value(position as i32); + } + positions_builder.append(true); + } + + Ok(PostingList::Plain(PlainPostingList::new( + ScalarBuffer::from(row_ids), + ScalarBuffer::from(frequencies), + Some(max_score), + Some(positions_builder.finish()), + ))) + } + fn union_compressed_posting_lists( postings: Vec, docs: &DocSet, + query_weight: f32, + scorer: &MemBM25Scorer, ) -> Result { let block_size = postings .iter() @@ -1944,10 +1999,77 @@ impl InvertedPartition { doc_ids.push(doc_id); frequencies.push(freq); } - let block_max_scores = docs.calculate_block_max_scores_with_block_size( - doc_ids.iter(), - frequencies.iter(), + let block_max_scores = Self::grouped_block_max_scores( + &doc_ids, + &frequencies, + block_size, + docs, + query_weight, + scorer, + ); + let batch = builder.to_batch(block_max_scores)?; + let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); + let length = batch[LENGTH_COL].as_primitive::().value(0); + PostingList::from_batch(&batch, Some(max_score), Some(length)) + } + + fn union_compressed_posting_lists_with_positions( + postings: Vec, + docs: &DocSet, + query_weight: f32, + scorer: &MemBM25Scorer, + ) -> Result { + let block_size = postings + .iter() + .find_map(|posting| match posting { + PostingList::Compressed(posting) => Some(posting.block_size), + PostingList::Plain(_) => None, + }) + .unwrap_or(LEGACY_BLOCK_SIZE); + let mut positions_by_doc_id = BTreeMap::>::new(); + for posting in postings { + for (doc_id, _, positions) in posting.iter() { + let doc_id = u32::try_from(doc_id).map_err(|_| { + Error::index(format!( + "compressed posting doc id {} exceeds u32::MAX", + doc_id + )) + })?; + let positions = positions.ok_or_else(|| { + Error::index("cannot union grouped phrase terms without positions".to_string()) + })?; + positions_by_doc_id + .entry(doc_id) + .or_default() + .extend(positions); + } + } + if positions_by_doc_id.is_empty() { + return Ok(PostingList::Plain(PlainPostingList::new( + ScalarBuffer::from(Vec::::new()), + ScalarBuffer::from(Vec::::new()), + None, + None, + ))); + } + + let mut builder = PostingListBuilder::new_with_block_size(true, block_size); + let mut doc_ids = Vec::with_capacity(positions_by_doc_id.len()); + let mut frequencies = Vec::with_capacity(positions_by_doc_id.len()); + for (doc_id, mut positions) in positions_by_doc_id { + positions.sort_unstable(); + let frequency = positions.len() as u32; + builder.add(doc_id, PositionRecorder::Position(positions.into())); + doc_ids.push(doc_id); + frequencies.push(frequency); + } + let block_max_scores = Self::grouped_block_max_scores( + &doc_ids, + &frequencies, block_size, + docs, + query_weight, + scorer, ); let batch = builder.to_batch(block_max_scores)?; let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); @@ -1955,7 +2077,13 @@ impl InvertedPartition { PostingList::from_batch(&batch, Some(max_score), Some(length)) } - fn union_posting_lists(postings: Vec, docs: &DocSet) -> Result { + fn union_posting_lists( + postings: Vec, + docs: &DocSet, + with_positions: bool, + query_weight: f32, + scorer: &MemBM25Scorer, + ) -> Result { let has_plain = postings .iter() .any(|posting| matches!(posting, PostingList::Plain(_))); @@ -1966,8 +2094,19 @@ impl InvertedPartition { (true, true) => Err(Error::index( "cannot union mixed plain and compressed posting lists".to_owned(), )), - (true, false) => Self::union_plain_posting_lists(postings), - (false, true) => Self::union_compressed_posting_lists(postings, docs), + (true, false) if with_positions => { + Self::union_plain_posting_lists_with_positions(postings, docs, query_weight, scorer) + } + (true, false) => Self::union_plain_posting_lists(postings, docs, query_weight, scorer), + (false, true) if with_positions => Self::union_compressed_posting_lists_with_positions( + postings, + docs, + query_weight, + scorer, + ), + (false, true) => { + Self::union_compressed_posting_lists(postings, docs, query_weight, scorer) + } (false, false) => Ok(PostingList::Plain(PlainPostingList::new( ScalarBuffer::from(Vec::::new()), ScalarBuffer::from(Vec::::new()), @@ -1989,7 +2128,6 @@ impl InvertedPartition { impact_scorer: &MemBM25Scorer, metrics: &dyn MetricsCollector, ) -> Result { - let is_fuzzy = matches!(params.fuzziness, Some(n) if n != 0); let is_phrase_query = params.phrase_slop.is_some(); let is_and_query = operator == Operator::And; let required_positions = (is_and_query || is_phrase_query).then(|| { @@ -1999,12 +2137,16 @@ impl InvertedPartition { }); // Fuzzy expansion already ran once at the index level (see // `InvertedIndex::bm25_search`) under the global `max_expansions` - // budget; the incoming tokens are final and `is_fuzzy` only drives - // the grouped dedup/scoring semantics below. + // budget. Positions identify alternatives that must share one posting + // iterator, including code identifier subwords and fuzzy expansions. let tokens = tokens.clone(); let token_positions = (0..tokens.len()) .map(|index| tokens.position(index)) .collect::>(); + let mut seen_positions = HashSet::with_capacity(token_positions.len()); + let exact_scoring_required = token_positions + .iter() + .any(|position| !seen_positions.insert(*position)); let mut token_ids = Vec::with_capacity(tokens.len()); let mut matched_positions = required_positions.as_ref().map(|_| HashSet::new()); for (index, token) in tokens.into_iter().enumerate() { @@ -2015,9 +2157,6 @@ impl InvertedPartition { matched_positions.insert(position); } token_ids.push((token_id, token, position)); - } else if is_phrase_query || is_and_query { - // if the token is not found, we can't do phrase or AND query - return Ok(LoadedPostings::empty()); } } if token_ids.is_empty() { @@ -2030,16 +2169,8 @@ impl InvertedPartition { return Ok(LoadedPostings::empty()); } - let is_fuzzy_and_query = is_fuzzy && is_and_query && !is_phrase_query; - if !is_phrase_query { - if is_fuzzy_and_query { - token_ids.sort_unstable_by_key(|(token_id, _, position)| (*position, *token_id)); - token_ids.dedup_by(|lhs, rhs| lhs.0 == rhs.0 && lhs.2 == rhs.2); - } else { - token_ids.sort_unstable_by_key(|(token_id, _, _)| *token_id); - token_ids.dedup_by_key(|(token_id, _, _)| *token_id); - } - } + token_ids.sort_unstable_by_key(|(token_id, _, position)| (*position, *token_id)); + token_ids.dedup_by(|lhs, rhs| lhs.0 == rhs.0 && lhs.2 == rhs.2); let num_docs = self.docs.len(); let loaded_postings = stream::iter(token_ids) @@ -2055,8 +2186,11 @@ impl InvertedPartition { .try_collect::>() .await?; + let needs_union = loaded_postings + .windows(2) + .any(|window| window[0].2 == window[1].2); if (is_and_query || is_phrase_query) - && !is_fuzzy_and_query + && !needs_union && loaded_postings .iter() .any(|(_, _, _, posting)| posting.is_empty()) @@ -2064,7 +2198,7 @@ impl InvertedPartition { return Ok(LoadedPostings::empty()); } - if !is_fuzzy_and_query { + if !needs_union { let impact_safe = loaded_postings .iter() .all(|(_, _, _, posting)| posting.has_impacts()); @@ -2072,29 +2206,34 @@ impl InvertedPartition { postings: loaded_postings .into_iter() .map(|(token_id, token, position, posting)| { - let query_weight = if impact_safe { + let needs_scorer_upper_bound = + exact_scoring_required && !posting.has_impacts(); + let query_weight = if impact_safe || exact_scoring_required { impact_scorer.query_weight(&token) } else { idf(posting.len(), num_docs) }; - PostingIterator::with_query_weight( + let posting = PostingIterator::with_query_weight( token, token_id, position, query_weight, posting, num_docs, - ) + ); + if needs_scorer_upper_bound { + posting.with_scorer_upper_bound() + } else { + posting + } }) .collect(), grouped_expansions: Vec::new(), impact_safe, + exact_scoring_required, }); } - let needs_union = loaded_postings - .windows(2) - .any(|window| window[0].2 == window[1].2); let docs_for_union = if needs_union { Some(self.docs.ensure_num_tokens_loaded().await?) } else { @@ -2119,44 +2258,77 @@ impl InvertedPartition { } else { let token_id = group[0].0; let token = group[0].1.clone(); + let terms = group + .iter() + .map(|(_, token, posting)| { + GroupedTermScorer::new(impact_scorer.query_weight(token), posting) + }) + .collect::>(); + let terms = Arc::<[GroupedTermScorer]>::from(terms); + let query_weight = terms.iter().map(GroupedTermScorer::query_weight).sum(); grouped_expansions.push(GroupedExpansionTerms { position, - terms: group - .iter() - .map(|(_, token, posting)| ExpansionTermFreqs::new(token.clone(), posting)) - .collect(), + terms: terms.clone(), }); let postings = group .into_iter() .map(|(_, _, posting)| posting) .collect::>(); + let docs = docs_for_union.as_deref().ok_or_else(|| { + Error::index("union docs were not loaded for grouped query terms".to_string()) + })?; let posting = Self::union_posting_lists( postings, - docs_for_union - .as_deref() - .expect("union docs must be loaded for grouped fuzzy AND"), + docs, + is_phrase_query, + query_weight, + impact_scorer, )?; - (token_id, token, posting) + if posting.is_empty() && (is_and_query || is_phrase_query) { + return Ok(LoadedPostings::empty()); + } + grouped_postings.push( + PostingIterator::with_query_weight( + token, + token_id, + position, + query_weight, + posting, + num_docs, + ) + .with_grouped_terms(terms), + ); + continue; }; if posting.is_empty() { - return Ok(LoadedPostings::empty()); + if is_and_query || is_phrase_query { + return Ok(LoadedPostings::empty()); + } + continue; } - let query_weight = idf(posting.len(), num_docs); - grouped_postings.push(PostingIterator::with_query_weight( + let query_weight = impact_scorer.query_weight(&token); + let needs_scorer_upper_bound = !posting.has_impacts(); + let posting = PostingIterator::with_query_weight( token, token_id, position, query_weight, posting, num_docs, - )); + ); + grouped_postings.push(if needs_scorer_upper_bound { + posting.with_scorer_upper_bound() + } else { + posting + }); } Ok(LoadedPostings { postings: grouped_postings, grouped_expansions, impact_safe: false, + exact_scoring_required: true, }) } @@ -6729,6 +6901,7 @@ async fn tokenize_and_count( ), ])); let output_schema_clone = output_schema.clone(); + let query_token_indices = Arc::new(query_token_indices(query_tokens.as_ref())); let bytes_accumulated = Arc::new(AtomicU64::new(0)); let bytes_warning_emitted = Arc::new(AtomicBool::new(false)); @@ -6737,6 +6910,7 @@ async fn tokenize_and_count( let mut tokenizer = tokenizer.box_clone(); let output_schema = output_schema.clone(); let query_tokens = query_tokens.clone(); + let query_token_indices = query_token_indices.clone(); let bytes_accumulated = bytes_accumulated.clone(); let bytes_warning_emitted = bytes_warning_emitted.clone(); let elapsed_compute = elapsed_compute.clone(); @@ -6760,8 +6934,10 @@ async fn tokenize_and_count( let mut all_tokens = 0; while let Some(token) = stream.next() { all_tokens += 1; - if let Some(token_index) = query_tokens.token_index(&token.text) { - temp_query_token_counts[token_index] += 1; + if let Some(token_indices) = query_token_indices.get(&token.text) { + for token_index in token_indices { + temp_query_token_counts[*token_index] += 1; + } } } all_tokens @@ -6891,6 +7067,17 @@ fn tokenize_and_count_list( Ok(()) } +fn query_token_indices(query_tokens: &Tokens) -> HashMap> { + let mut indices = HashMap::new(); + for idx in 0..query_tokens.len() { + indices + .entry(query_tokens.get_token(idx).to_string()) + .or_insert_with(Vec::new) + .push(idx); + } + indices +} + /// Initialize the BM25 scorer /// /// In order to calculate BM25 scores we need to know token counts for the entire corpus. We extract these from the @@ -6954,9 +7141,11 @@ fn flat_bm25_score( query_tokens: &Tokens, counted_input: &RecordBatch, scorer: &MemBM25Scorer, + operator: Operator, ) -> Result { let mut row_ids_builder = UInt64Builder::with_capacity(counted_input.num_rows()); let mut scores_builder = Float32Builder::with_capacity(counted_input.num_rows()); + let query_groups = query_position_groups(query_tokens); let mut row_ids_iter = counted_input .column(FLAT_ROW_ID_COL_IDX) @@ -6981,16 +7170,24 @@ fn flat_bm25_score( for _ in 0..counted_input.num_rows() { let num_tokens_in_doc = all_token_counts_iter.next().expect_ok()?; let row_id = row_ids_iter.next().expect_ok()?; + let mut query_token_counts = Vec::with_capacity(query_tokens.len()); + for _ in query_tokens { + query_token_counts.push(query_token_counts_iter.next().expect_ok()?); + } if num_tokens_in_doc == 0 { - for _ in query_tokens { - query_token_counts_iter.next().expect_ok()?; - } + continue; + } + if operator == Operator::And + && !query_groups + .iter() + .all(|group| group.iter().any(|idx| query_token_counts[*idx] > 0)) + { continue; } let doc_norm = K1 * (1.0 - B + B * num_tokens_in_doc as f32 / scorer.avg_doc_length()); let mut score = 0.0; - for token in query_tokens { - let freq = query_token_counts_iter.next().expect_ok()? as f32; + for (token, freq) in query_tokens.into_iter().zip(query_token_counts) { + let freq = freq as f32; let idf = idf(scorer.num_docs_containing_token(token), scorer.num_docs()); score += idf * (freq * (K1 + 1.0) / (freq + doc_norm)); } @@ -7009,6 +7206,23 @@ fn flat_bm25_score( Ok(batch) } +fn query_position_groups(query_tokens: &Tokens) -> Vec> { + let mut groups = Vec::new(); + let mut current_position = None; + for idx in 0..query_tokens.len() { + let position = query_tokens.position(idx); + if current_position != Some(position) { + current_position = Some(position); + groups.push(Vec::new()); + } + groups + .last_mut() + .expect("a group should exist after pushing for position") + .push(idx); + } + groups +} + #[deprecated( note = "use `flat_bm25_search_stream_with_metrics` to record CPU compute \ time on a metric handle; pass `None` for the old behavior" @@ -7046,6 +7260,58 @@ pub async fn flat_bm25_search_stream_with_metrics( base_scorer: Option, target_batch_size: usize, elapsed_compute: Option