diff --git a/Cargo.lock b/Cargo.lock index 4c943f837..0d5a5b45f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1756,7 +1756,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3034,7 +3034,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3257,7 +3257,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3455,7 +3455,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4560,7 +4560,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -4814,7 +4814,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arc-swap", "arrow", @@ -4888,7 +4888,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-buffer", @@ -4931,7 +4931,7 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrayref", "crunchy", @@ -4941,7 +4941,7 @@ dependencies = [ [[package]] name = "lance-core" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-buffer", @@ -4981,7 +4981,7 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "arrow-array", @@ -5011,7 +5011,7 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "arrow-array", @@ -5028,7 +5028,7 @@ dependencies = [ [[package]] name = "lance-derive" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "proc-macro2", "quote", @@ -5037,7 +5037,7 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-arith", "arrow-array", @@ -5071,7 +5071,7 @@ dependencies = [ [[package]] name = "lance-file" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-arith", "arrow-array", @@ -5102,7 +5102,7 @@ dependencies = [ [[package]] name = "lance-index" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arc-swap", "arrow", @@ -5169,7 +5169,7 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-schema", @@ -5191,7 +5191,7 @@ dependencies = [ [[package]] name = "lance-io" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "arrow-array", @@ -5230,7 +5230,7 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-buffer", @@ -5246,7 +5246,7 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "async-trait", @@ -5258,7 +5258,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "arrow-ipc", @@ -5288,7 +5288,6 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "time", "tokio", "tower", "tower-http 0.5.2", @@ -5312,7 +5311,7 @@ dependencies = [ [[package]] name = "lance-select" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-buffer", @@ -5327,7 +5326,7 @@ dependencies = [ [[package]] name = "lance-table" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow", "arrow-array", @@ -5337,6 +5336,7 @@ dependencies = [ "async-trait", "aws-credential-types", "aws-sdk-dynamodb", + "blake3", "byteorder", "bytes", "chrono", @@ -5366,7 +5366,7 @@ dependencies = [ [[package]] name = "lance-testing" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "arrow-array", "arrow-schema", @@ -5379,7 +5379,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" dependencies = [ "frostem", "icu_segmenter", @@ -6232,7 +6232,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -7624,8 +7624,8 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.11.0", "log", "multimap", "petgraph", @@ -7644,7 +7644,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.11.0", "proc-macro2", "quote", "syn 2.0.117", @@ -7915,7 +7915,7 @@ dependencies = [ "once_cell", "socket2 0.6.3", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -8693,7 +8693,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -8764,7 +8764,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -9320,7 +9320,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -9332,7 +9332,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -9754,7 +9754,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -10731,7 +10731,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 83f2a7da3..b3f22298c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,20 +14,20 @@ categories = ["database-implementations"] rust-version = "1.91.0" [workspace.dependencies] -lance = { "version" = "=11.0.0-beta.2", default-features = false, "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=11.0.0-beta.2", default-features = false, "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=11.0.0-beta.2", default-features = false, "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=11.0.0-beta.2", "tag" = "v11.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } ahash = "0.8" # Note that this one does not include pyarrow arrow = { version = "58.0.0", optional = false } diff --git a/java/pom.xml b/java/pom.xml index e8f030b27..e1c39c3e2 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 11.0.0-beta.2 + 11.0.0-beta.3 false 2.30.0 1.7 diff --git a/lance-artifact/Cargo.toml b/lance-artifact/Cargo.toml index a63542d7d..181f44ed0 100644 --- a/lance-artifact/Cargo.toml +++ b/lance-artifact/Cargo.toml @@ -33,7 +33,7 @@ resolver = "3" [workspace.package] -version = "11.0.0-beta.2" +version = "11.0.0-beta.3" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" @@ -58,28 +58,28 @@ rust-version = "1.91.0" [workspace.dependencies] arc-swap = "1.7" libc = "0.2.176" -lance = { version = "=11.0.0-beta.2", path = "./rust/lance", default-features = false } -lance-arrow = { version = "=11.0.0-beta.2", path = "./rust/lance-arrow" } -lance-core = { version = "=11.0.0-beta.2", path = "./rust/lance-core" } -lance-datafusion = { version = "=11.0.0-beta.2", path = "./rust/lance-datafusion" } -lance-datagen = { version = "=11.0.0-beta.2", path = "./rust/lance-datagen" } -lance-derive = { version = "=11.0.0-beta.2", path = "./rust/lance-derive" } -lance-encoding = { version = "=11.0.0-beta.2", path = "./rust/lance-encoding" } -lance-file = { version = "=11.0.0-beta.2", path = "./rust/lance-file" } -lance-geo = { version = "=11.0.0-beta.2", path = "./rust/lance-geo" } -lance-index = { version = "=11.0.0-beta.2", path = "./rust/lance-index" } -lance-index-core = { version = "=11.0.0-beta.2", path = "./rust/lance-index-core" } -lance-io = { version = "=11.0.0-beta.2", path = "./rust/lance-io", default-features = false } -lance-linalg = { version = "=11.0.0-beta.2", path = "./rust/lance-linalg" } -lance-namespace = { version = "=11.0.0-beta.2", path = "./rust/lance-namespace" } -lance-namespace-impls = { version = "=11.0.0-beta.2", path = "./rust/lance-namespace-impls" } +lance = { version = "=11.0.0-beta.3", path = "./rust/lance", default-features = false } +lance-arrow = { version = "=11.0.0-beta.3", path = "./rust/lance-arrow" } +lance-core = { version = "=11.0.0-beta.3", path = "./rust/lance-core" } +lance-datafusion = { version = "=11.0.0-beta.3", path = "./rust/lance-datafusion" } +lance-datagen = { version = "=11.0.0-beta.3", path = "./rust/lance-datagen" } +lance-derive = { version = "=11.0.0-beta.3", path = "./rust/lance-derive" } +lance-encoding = { version = "=11.0.0-beta.3", path = "./rust/lance-encoding" } +lance-file = { version = "=11.0.0-beta.3", path = "./rust/lance-file" } +lance-geo = { version = "=11.0.0-beta.3", path = "./rust/lance-geo" } +lance-index = { version = "=11.0.0-beta.3", path = "./rust/lance-index" } +lance-index-core = { version = "=11.0.0-beta.3", path = "./rust/lance-index-core" } +lance-io = { version = "=11.0.0-beta.3", path = "./rust/lance-io", default-features = false } +lance-linalg = { version = "=11.0.0-beta.3", path = "./rust/lance-linalg" } +lance-namespace = { version = "=11.0.0-beta.3", path = "./rust/lance-namespace" } +lance-namespace-impls = { version = "=11.0.0-beta.3", path = "./rust/lance-namespace-impls" } lance-namespace-datafusion = { version = "=7.0.0-beta.9", path = "./rust/lance-namespace-datafusion" } lance-namespace-reqwest-client = "0.8.6" -lance-select = { version = "=11.0.0-beta.2", path = "./rust/lance-select" } -lance-tokenizer = { version = "=11.0.0-beta.2", path = "./rust/lance-tokenizer" } -lance-table = { version = "=11.0.0-beta.2", path = "./rust/lance-table" } -lance-test-macros = { version = "=11.0.0-beta.2", path = "./rust/lance-test-macros" } -lance-testing = { version = "=11.0.0-beta.2", path = "./rust/lance-testing" } +lance-select = { version = "=11.0.0-beta.3", path = "./rust/lance-select" } +lance-tokenizer = { version = "=11.0.0-beta.3", path = "./rust/lance-tokenizer" } +lance-table = { version = "=11.0.0-beta.3", path = "./rust/lance-table" } +lance-test-macros = { version = "=11.0.0-beta.3", path = "./rust/lance-test-macros" } +lance-testing = { version = "=11.0.0-beta.3", path = "./rust/lance-testing" } approx = "0.5.1" # Note that this one does not include pyarrow arrow = { version = "58.0.0", optional = false, features = ["prettyprint"] } @@ -107,7 +107,7 @@ half = { "version" = "2.1", default-features = false, features = [ "std", "bytemuck", ] } -lance-bitpacking = { version = "=11.0.0-beta.2", path = "./rust/compression/bitpacking" } +lance-bitpacking = { version = "=11.0.0-beta.3", path = "./rust/compression/bitpacking" } bitpacking = "0.9" bitvec = "1" blake3 = "1.8.5" @@ -150,7 +150,7 @@ datafusion-substrait = { version = "54.0.0", default-features = false } dirs = "6.0.0" either = "1.0" fst = { version = "0.4.7", features = ["levenshtein"] } -fsst = { version = "=11.0.0-beta.2", path = "./rust/compression/fsst" } +fsst = { version = "=11.0.0-beta.3", path = "./rust/compression/fsst" } futures = "0.3" geoarrow-array = "0.8" geoarrow-schema = "0.8" diff --git a/lance-artifact/LANCEDB_PATCH.md b/lance-artifact/LANCEDB_PATCH.md index 6d37091f7..9185fd780 100644 --- a/lance-artifact/LANCEDB_PATCH.md +++ b/lance-artifact/LANCEDB_PATCH.md @@ -1,7 +1,7 @@ # LanceDB patch provenance -This artifact vendors the Lance 11.0.0-beta.2 Rust workspace from Lance commit -`35da5d920159b49d1b53032652f7615ab699c160`. The complete workspace keeps all mutually coupled +This artifact vendors the Lance 11.0.0-beta.3 Rust workspace from Lance commit +`f7d475539cefbd140cc46a828f3d843e68cd10f1`. The complete workspace keeps all mutually coupled Lance crates on one Cargo source identity when `lancedb` consumes the pinned artifact commit. The local patch makes AWS credential-family merging atomic before backend selection and teaches diff --git a/lance-artifact/rust/lance-arrow/src/scalar.rs b/lance-artifact/rust/lance-arrow/src/scalar.rs index 22d9d1855..c39607dc4 100644 --- a/lance-artifact/rust/lance-arrow/src/scalar.rs +++ b/lance-artifact/rust/lance-arrow/src/scalar.rs @@ -4,10 +4,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use arrow_array::{ArrayRef, make_array}; +use arrow_array::{ArrayRef, UInt64Array, make_array}; use arrow_buffer::Buffer; -use arrow_data::{ArrayDataBuilder, transform::MutableArrayData}; +use arrow_data::ArrayDataBuilder; use arrow_schema::{ArrowError, DataType}; +use arrow_select::take::take; use crate::DataTypeExt; @@ -22,10 +23,7 @@ pub fn extract_scalar_value(array: &ArrayRef, idx: usize) -> Result { )); } - let data = array.to_data(); - let mut mutable = MutableArrayData::new(vec![&data], /*use_nulls=*/ true, 1); - mutable.extend(0, idx, idx + 1); - Ok(make_array(mutable.freeze())) + take(array.as_ref(), &UInt64Array::from(vec![idx as u64]), None) } fn read_u32(buf: &[u8], offset: &mut usize) -> Result { @@ -198,7 +196,10 @@ pub fn try_inline_value(scalar: &ArrayRef) -> Option> { mod tests { use std::sync::Arc; - use arrow_array::{BooleanArray, FixedSizeBinaryArray, Int32Array, StringArray, cast::AsArray}; + use arrow_array::{ + BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int8Array, Int32Array, StringArray, + cast::AsArray, types::Int8Type, + }; use super::*; @@ -215,6 +216,24 @@ mod tests { ); } + #[test] + fn test_extract_scalar_value_from_full_dictionary() { + let values = Arc::new(StringArray::from( + (0..=i8::MAX) + .map(|value| format!("value-{value}")) + .collect::>(), + )); + let keys = Int8Array::from((0..=i8::MAX).collect::>()); + let array: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + + let scalar = extract_scalar_value(&array, i8::MAX as usize).unwrap(); + + let scalar = scalar.as_dictionary::(); + assert_eq!(scalar.len(), 1); + assert_eq!(scalar.key(0), Some(i8::MAX as usize)); + assert_eq!(scalar.values().len(), i8::MAX as usize + 1); + } + #[test] fn test_scalar_value_buffer_utf8_round_trip() { let scalar: ArrayRef = Arc::new(StringArray::from(vec!["hello"])); diff --git a/lance-artifact/rust/lance-core/src/cache/mod.rs b/lance-artifact/rust/lance-core/src/cache/mod.rs index 7ffe3d00f..5c5047feb 100644 --- a/lance-artifact/rust/lance-core/src/cache/mod.rs +++ b/lance-artifact/rust/lance-core/src/cache/mod.rs @@ -480,6 +480,45 @@ impl LanceCache { self.state.backend.insert(&key, metadata, size, None).await; } + pub async fn get_or_insert_unsized_with_key( + &self, + cache_key: K, + loader: F, + ) -> Result> + where + K: UnsizedCacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + F: FnOnce() -> Fut + Send, + Fut: Future>> + Send, + { + let key = self.unsized_key(&cache_key); + let state = self.state.clone(); + let typed_loader = Box::pin(async move { + let value = loader().await?; + let size = state.entry_size(&value); + Ok((Arc::new(value) as CacheEntry, size)) + }); + + let (entry, was_cached) = self + .state + .backend + .get_or_insert(&key, typed_loader, None) + .await?; + let entry = entry.downcast::>().map_err(|_| { + self.state.misses.fetch_add(1, Ordering::Relaxed); + Error::io(format!( + "cache backend returned a value with the wrong concrete type for unsized key type {:?}", + K::stable_type_id() + )) + })?; + if was_cached { + self.state.hits.fetch_add(1, Ordering::Relaxed); + } else { + self.state.misses.fetch_add(1, Ordering::Relaxed); + } + Ok(entry.as_ref().clone()) + } + pub async fn get_unsized_with_key(&self, cache_key: &K) -> Option> where K: UnsizedCacheKey, diff --git a/lance-artifact/rust/lance-core/src/cache/registry.rs b/lance-artifact/rust/lance-core/src/cache/registry.rs index c023df57e..aa30a01fb 100644 --- a/lance-artifact/rust/lance-core/src/cache/registry.rs +++ b/lance-artifact/rust/lance-core/src/cache/registry.rs @@ -113,10 +113,18 @@ fn registry_lock_for_test() -> MutexGuard<'static, HashMap lance_core::Result<()> { -/// lance_core::cache::register_backend("my_backend", build_my_backend) +/// ``` +/// # use std::sync::Arc; +/// # use lance_core::Result; +/// # use lance_core::cache::{BackendConfig, CacheBackend, MokaCacheBackend, register_backend}; +/// fn build_my_backend(_config: &BackendConfig) -> Result> { +/// Ok(Arc::new(MokaCacheBackend::with_capacity(1024))) /// } +/// +/// # fn main() -> Result<()> { +/// register_backend("my-backend", build_my_backend)?; +/// # Ok(()) +/// # } /// ``` pub fn register_backend(kind: &str, build: BackendBuildFn) -> Result<()> { let kind = normalize_backend_kind(kind)?; diff --git a/lance-artifact/rust/lance-core/src/utils/cpu.rs b/lance-artifact/rust/lance-core/src/utils/cpu.rs index 0107670a9..6a7c40d4a 100644 --- a/lance-artifact/rust/lance-core/src/utils/cpu.rs +++ b/lance-artifact/rust/lance-core/src/utils/cpu.rs @@ -213,6 +213,14 @@ pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { SimdSupport::None } } + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + target_arch = "loongarch64" + )))] + { + SimdSupport::None + } }); #[cfg(target_arch = "x86_64")] diff --git a/lance-artifact/rust/lance-datafusion/src/spill.rs b/lance-artifact/rust/lance-datafusion/src/spill.rs index cf19a8189..82df84560 100644 --- a/lance-artifact/rust/lance-datafusion/src/spill.rs +++ b/lance-artifact/rust/lance-datafusion/src/spill.rs @@ -392,7 +392,7 @@ impl Default for DataLocation { } } -/// A DataFusion error that be be emitted multiple times. We provide the +/// A DataFusion error that can be emitted multiple times. We provide the /// Original error first, and subsequent conversions provide a copy with a /// string representation of the original error. #[derive(Debug)] diff --git a/lance-artifact/rust/lance-file/src/reader.rs b/lance-artifact/rust/lance-file/src/reader.rs index 7fe69d5eb..343786e3f 100644 --- a/lance-artifact/rust/lance-file/src/reader.rs +++ b/lance-artifact/rust/lance-file/src/reader.rs @@ -2433,8 +2433,9 @@ mod tests { }; use arrow_array::{ - Int32Array, ListArray, RecordBatch, RecordBatchIterator, UInt32Array, - types::{Float64Type, Int32Type}, + DictionaryArray, Int8Array, Int32Array, ListArray, RecordBatch, RecordBatchIterator, + StringArray, UInt32Array, + types::{Float64Type, Int8Type, Int32Type}, }; use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; @@ -2601,6 +2602,61 @@ mod tests { assert_eq!(take, vec![batch.take(&indices).unwrap()]); } + #[tokio::test] + async fn full_int8_dictionary_v2_2_roundtrip() { + let fs = FsFixture::default(); + let values = Arc::new(StringArray::from( + (0..=i8::MAX) + .map(|value| format!("value-{value}")) + .collect::>(), + )); + let keys = Int8Array::from((0..=i8::MAX).collect::>()); + let dictionary = Arc::new(DictionaryArray::::new(keys, values)); + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "dictionary", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + )])); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![dictionary]).unwrap(); + + write_lance_file( + RecordBatchIterator::new([Ok(batch.clone())], arrow_schema), + &fs, + ConcreteFileVersion::V2_2, + FileWriterOptions::default(), + ) + .await; + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + let actual = file_reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(actual, vec![batch]); + } + async fn create_some_file(fs: &FsFixture, version: ConcreteFileVersion) -> WrittenFile { let location_type = DataType::Struct(Fields::from(vec![ Field::new("x", DataType::Float64, true), diff --git a/lance-artifact/rust/lance-index/src/scalar/inverted.rs b/lance-artifact/rust/lance-index/src/scalar/inverted.rs index d417331a3..c02c4d814 100644 --- a/lance-artifact/rust/lance-index/src/scalar/inverted.rs +++ b/lance-artifact/rust/lance-index/src/scalar/inverted.rs @@ -130,7 +130,8 @@ use crate::scalar::{ CreatedIndex, RowIdRemapper, ScalarIndex, expression::{FtsQueryParser, ScalarQueryParser}, registry::{ - BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, + BasicTrainer, ScalarIndexCacheKey, ScalarIndexLoad, ScalarIndexPlugin, TrainingCriteria, + TrainingOrdering, TrainingRequest, }, }; @@ -329,6 +330,18 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { Ok(index as Arc) } + async fn get_or_insert_in_cache( + &self, + _index_store: Arc, + _frag_reuse_index: Option>, + cache: &LanceCache, + load: ScalarIndexLoad<'_>, + ) -> Result> { + cache + .get_or_insert_unsized_with_key(ScalarIndexCacheKey, || load) + .await + } + fn details_as_json(&self, details: &prost_types::Any) -> Result { let index_details = details.to_msg::()?; let index_params = InvertedIndexParams::try_from(&index_details)?; diff --git a/lance-artifact/rust/lance-index/src/scalar/inverted/compound.rs b/lance-artifact/rust/lance-index/src/scalar/inverted/compound.rs index 14eecefec..913a1f7d9 100644 --- a/lance-artifact/rust/lance-index/src/scalar/inverted/compound.rs +++ b/lance-artifact/rust/lance-index/src/scalar/inverted/compound.rs @@ -205,6 +205,8 @@ pub(super) trait ComposableScorer: Send { Ok(true) } + /// Estimated relative cost of [`Self::matches`], stable for this scorer's + /// lifetime. `None` means no ordering hint, not that confirmation may be skipped. fn match_cost(&self) -> Option { None } @@ -1192,6 +1194,9 @@ pub(super) struct RequiredConjunctionScorer<'a> { /// already cheapest-first. `children` remains in query order so scoring and /// score-bound arithmetic stay bit-for-bit stable. approximation_order: Option>, + /// Child indices sorted by two-phase confirmation cost. Children without a + /// cost hint remain in query order after costed confirmations. + confirmation_order: Option>, current: Option, confirmed_doc: Option, confirmed: bool, @@ -1223,6 +1228,30 @@ fn align_conjunction_children( } } +fn compare_confirmation_cost( + left: &dyn ComposableScorer, + right: &dyn ComposableScorer, +) -> Ordering { + match (left.match_cost(), right.match_cost()) { + (Some(left), Some(right)) => left.total_cmp(&right), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } +} + +fn confirm_conjunction_children( + children: &mut [BoxScorer<'_>], + child_index: impl Fn(usize) -> usize, +) -> Result { + for position in 0..children.len() { + if !children[child_index(position)].matches()? { + return Ok(false); + } + } + Ok(true) +} + impl<'a> RequiredConjunctionScorer<'a> { pub(super) fn try_new(children: Vec>) -> Result { if children.is_empty() { @@ -1240,9 +1269,31 @@ impl<'a> RequiredConjunctionScorer<'a> { order.sort_by_key(|&index| (children[index].cost(), index)); Some(order) }; + for (index, child) in children.iter().enumerate() { + if let Some(match_cost) = child.match_cost() + && (!match_cost.is_finite() || match_cost < 0.0) + { + return Err(Error::internal(format!( + "FTS conjunction child {index} reported invalid two-phase match cost: {match_cost}" + ))); + } + } + let confirmation_order = if children.windows(2).all(|pair| { + compare_confirmation_cost(pair[0].as_ref(), pair[1].as_ref()) != Ordering::Greater + }) { + None + } else { + let mut order = (0..children.len()).collect::>(); + order.sort_by(|&left, &right| { + compare_confirmation_cost(children[left].as_ref(), children[right].as_ref()) + .then_with(|| left.cmp(&right)) + }); + Some(order) + }; Ok(Self { children, approximation_order, + confirmation_order, current: None, confirmed_doc: None, confirmed: false, @@ -1269,12 +1320,11 @@ impl<'a> RequiredConjunctionScorer<'a> { if self.confirmed_doc == Some(current) { return Ok(self.confirmed); } - self.confirmed = true; - for child in &mut self.children { - if !child.matches()? { - self.confirmed = false; - } - } + self.confirmed = if let Some(order) = &self.confirmation_order { + confirm_conjunction_children(&mut self.children, |position| order[position])? + } else { + confirm_conjunction_children(&mut self.children, |position| position)? + }; self.confirmed_doc = Some(current); Ok(self.confirmed) } @@ -2443,7 +2493,9 @@ mod tests { struct TwoPhaseScorer { inner: MaterializedScorer, accepted: Vec, - confirmations: usize, + match_cost: Option, + approximations: Arc, + confirmations: Arc, } impl ComposableScorer for TwoPhaseScorer { @@ -2452,11 +2504,19 @@ mod tests { } fn next(&mut self) -> Result> { - self.inner.next() + let doc = self.inner.next()?; + if doc.is_some() { + self.approximations.fetch_add(1, AtomicOrdering::Relaxed); + } + Ok(doc) } fn advance(&mut self, target: u64) -> Result> { - self.inner.advance(target) + let doc = self.inner.advance(target)?; + if doc.is_some() { + self.approximations.fetch_add(1, AtomicOrdering::Relaxed); + } + Ok(doc) } fn cost(&self) -> usize { @@ -2480,17 +2540,42 @@ mod tests { } fn matches(&mut self) -> Result { - self.confirmations += 1; + self.confirmations.fetch_add(1, AtomicOrdering::Relaxed); Ok(self .doc() .is_some_and(|doc| self.accepted.binary_search(&doc).is_ok())) } + fn match_cost(&self) -> Option { + self.match_cost + } + fn scores_non_negative(&self) -> bool { true } } + fn two_phase( + values: &[(u64, f32)], + accepted: Vec, + match_cost: Option, + ) -> ( + Box, + Arc, + Arc, + ) { + let approximations = Arc::new(AtomicUsize::new(0)); + let confirmations = Arc::new(AtomicUsize::new(0)); + let scorer = TwoPhaseScorer { + inner: MaterializedScorer::try_new(rows(values)).unwrap(), + accepted, + match_cost, + approximations: approximations.clone(), + confirmations: confirmations.clone(), + }; + (Box::new(scorer), approximations, confirmations) + } + struct CountingScorer { inner: MaterializedScorer, cost: usize, @@ -2555,17 +2640,84 @@ mod tests { #[test] fn collector_confirms_two_phase_matches_without_a_cost_hint() { - let mut scorer = TwoPhaseScorer { - inner: MaterializedScorer::try_new(rows(&[(1, 100.0), (2, 2.0), (3, 1.0)])).unwrap(), - accepted: vec![2, 3], - confirmations: 0, - }; - let results = TopKCollector::new(2).collect(&mut scorer).unwrap(); + let (mut scorer, approximations, confirmations) = + two_phase(&[(1, 100.0), (2, 2.0), (3, 1.0)], vec![2, 3], None); + let results = TopKCollector::new(2).collect(scorer.as_mut()).unwrap(); assert_eq!(results, rows(&[(2, 2.0), (3, 1.0)])); - assert_eq!(scorer.confirmations, 3); + assert_eq!(approximations.load(AtomicOrdering::Relaxed), 3); + assert_eq!(confirmations.load(AtomicOrdering::Relaxed), 3); assert_eq!(scorer.match_cost(), None); } + #[test] + fn required_conjunction_confirms_cheapest_first_and_short_circuits() { + let values = (0..100).map(|doc| (doc, 1.0)).collect::>(); + let accepted_by_cheap = (0..100).step_by(5).collect::>(); + + let (expensive, expensive_approximations, expensive_confirmations) = + two_phase(&values, (0..100).collect(), Some(10.0)); + let (cheap, cheap_approximations, cheap_confirmations) = + two_phase(&values, accepted_by_cheap.clone(), Some(1.0)); + let mut scorer = RequiredConjunctionScorer::try_new(vec![expensive, cheap]).unwrap(); + assert_eq!(scorer.confirmation_order.as_deref(), Some(&[1, 0][..])); + + let results = TopKCollector::new(100).collect(&mut scorer).unwrap(); + let expected = accepted_by_cheap + .iter() + .map(|doc| (*doc, 2.0)) + .collect::>(); + assert_eq!(results, rows(&expected)); + assert_eq!(cheap_confirmations.load(AtomicOrdering::Relaxed), 100); + assert_eq!(expensive_confirmations.load(AtomicOrdering::Relaxed), 20); + let approximations = cheap_approximations.load(AtomicOrdering::Relaxed) + + expensive_approximations.load(AtomicOrdering::Relaxed); + let confirmations = cheap_confirmations.load(AtomicOrdering::Relaxed) + + expensive_confirmations.load(AtomicOrdering::Relaxed); + assert_eq!(approximations, 200); + assert_eq!(confirmations, 120); + assert!( + confirmations * 5 <= approximations * 4, + "confirmation ordering should reduce work by at least 20%: {confirmations}/{approximations}" + ); + + let (cheap, _, _) = two_phase(&values, accepted_by_cheap, Some(1.0)); + let (expensive, _, _) = two_phase(&values, (0..100).collect(), Some(10.0)); + let scorer = RequiredConjunctionScorer::try_new(vec![cheap, expensive]).unwrap(); + assert!(scorer.confirmation_order.is_none()); + } + + #[test] + fn required_conjunction_confirms_children_without_cost_hints() { + let (unknown, _, unknown_confirmations) = two_phase(&[(0, 1.0)], Vec::new(), None); + let (costed, _, costed_confirmations) = two_phase(&[(0, 1.0)], vec![0], Some(1.0)); + let mut scorer = RequiredConjunctionScorer::try_new(vec![unknown, costed]).unwrap(); + + assert_eq!(scorer.confirmation_order.as_deref(), Some(&[1, 0][..])); + assert!( + TopKCollector::new(1) + .collect(&mut scorer) + .unwrap() + .is_empty() + ); + assert_eq!(costed_confirmations.load(AtomicOrdering::Relaxed), 1); + assert_eq!(unknown_confirmations.load(AtomicOrdering::Relaxed), 1); + } + + #[test] + fn required_conjunction_rejects_invalid_match_cost() { + let (invalid, _, _) = two_phase(&[(0, 1.0)], vec![0], Some(f32::NAN)); + + let error = RequiredConjunctionScorer::try_new(vec![invalid]) + .err() + .unwrap(); + assert!(matches!(error, Error::Internal { .. })); + assert!( + error + .to_string() + .contains("child 0 reported invalid two-phase match cost: NaN") + ); + } + #[test] fn required_conjunction_uses_all_must_scores_for_competitive_bounds() { let left = Box::new( diff --git a/lance-artifact/rust/lance-index/src/scalar/inverted/documents.rs b/lance-artifact/rust/lance-index/src/scalar/inverted/documents.rs index 9d6092833..1990d30dc 100644 --- a/lance-artifact/rust/lance-index/src/scalar/inverted/documents.rs +++ b/lance-artifact/rust/lance-index/src/scalar/inverted/documents.rs @@ -26,6 +26,7 @@ use object_store::path::Path; use roaring::RoaringBitmap; use tokio::sync::OnceCell; +use crate::FtsPrewarmDocumentStatus; use crate::scalar::{IndexReader, IndexStore, RowIdRemapper}; use super::index::{ @@ -658,6 +659,7 @@ impl PartitionDocumentStore { } } + #[cfg(test)] pub(crate) fn query_ready(&self) -> bool { match self { Self::Legacy(_) => true, @@ -665,6 +667,18 @@ impl PartitionDocumentStore { } } + pub(crate) fn prewarm_status(&self) -> FtsPrewarmDocumentStatus { + match self { + Self::Legacy(_) => FtsPrewarmDocumentStatus { + prewarm_complete: true, + scoring_ready: true, + reverse_lookup_ready: true, + projection_resident: true, + }, + Self::Modern(docs) => docs.prewarm_status(), + } + } + pub(crate) async fn load_build_docset(&self) -> Result { match self { Self::Legacy(docs) => Ok((**docs).clone()), @@ -776,17 +790,24 @@ impl PartitionDocuments { self.resident_address_projection().is_some() } + #[cfg(test)] pub(crate) fn query_ready(&self) -> bool { - self.prewarm_complete.initialized() - && self + self.prewarm_status().query_ready() + } + + pub(crate) fn prewarm_status(&self) -> FtsPrewarmDocumentStatus { + FtsPrewarmDocumentStatus { + prewarm_complete: self.prewarm_complete.initialized(), + scoring_ready: self .lengths .get() - .is_some_and(|lengths| lengths.scoring_ready()) - && self + .is_some_and(|lengths| lengths.scoring_ready()), + reverse_lookup_ready: self .projection .get() - .is_some_and(|projection| projection.doc_ids_by_address.initialized()) - && self.projection_resident() + .is_some_and(|projection| projection.doc_ids_by_address.initialized()), + projection_resident: self.projection_resident(), + } } async fn reader(&self) -> Result> { diff --git a/lance-artifact/rust/lance-index/src/scalar/inverted/index.rs b/lance-artifact/rust/lance-index/src/scalar/inverted/index.rs index e0fb83f9f..8331f45d5 100644 --- a/lance-artifact/rust/lance-index/src/scalar/inverted/index.rs +++ b/lance-artifact/rust/lance-index/src/scalar/inverted/index.rs @@ -93,14441 +93,41 @@ use crate::scalar::{ OldIndexDataFilter, RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, TokenQuery, UpdateCriteria, }; -use crate::{FtsPrewarmOptions, Index}; +use crate::{ + FtsPrewarmDiagnostics, FtsPrewarmOptions, FtsPrewarmPartitionStatus, FtsPrewarmResult, Index, +}; use crate::{prefilter::PreFilter, scalar::inverted::iter::take_fst_keys}; use std::str::FromStr; -// Version 0: Arrow TokenSetFormat (legacy) -// Version 1: Fst TokenSetFormat with per-doc compressed positions -// Version 2: Fst TokenSetFormat with shared posting-list position streams. -// Version 3: Reader capability for configurable posting blocks, analyzer -// metadata, and element-document coordinates. -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 TOKENS_FILE: &str = "tokens.lance"; -pub const INVERT_LIST_FILE: &str = "invert.lance"; -pub const DOCS_FILE: &str = "docs.lance"; -pub const METADATA_FILE: &str = "metadata.lance"; - -/// Partitions searched per CPU-pool task. Each chunk loads concurrently and -/// then scores sequentially so query concurrency does not flood the pool with -/// one small task per partition. `LANCE_FTS_SEARCH_CHUNK=1` restores the -/// per-partition task shape. -fn fts_search_chunk() -> usize { - static CHUNK: LazyLock = LazyLock::new(|| { - std::env::var("LANCE_FTS_SEARCH_CHUNK") - .ok() - .and_then(|value| value.parse().ok()) - .filter(|&value| value >= 1) - .unwrap_or(16) - }); - *CHUNK -} - -pub const TOKEN_COL: &str = "_token"; -pub const TOKEN_ID_COL: &str = "_token_id"; -pub const TOKEN_FST_BYTES_COL: &str = "_token_fst_bytes"; -pub const TOKEN_NEXT_ID_COL: &str = "_token_next_id"; -pub const TOKEN_TOTAL_LENGTH_COL: &str = "_token_total_length"; -pub const FREQUENCY_COL: &str = "_frequency"; -pub const POSITION_COL: &str = "_position"; -pub const COMPRESSED_POSITION_COL: &str = "_compressed_position"; -pub const POSITION_BLOCK_OFFSET_COL: &str = "_position_block_offset"; -pub const POSTING_COL: &str = "_posting"; -pub const IMPACT_COL: &str = "_impacts"; -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"; -pub const FTS_FORMAT_VERSION_KEY: &str = "format_version"; -pub const POSITIONS_LAYOUT_KEY: &str = "positions_layout"; -pub const POSITIONS_CODEC_KEY: &str = "positions_codec"; -pub const POSTING_BLOCK_SIZE_KEY: &str = "posting_block_size"; -pub const POSTING_TAIL_CODEC_FIXED32_V1: &str = "fixed32_v1"; -pub const POSTING_TAIL_CODEC_VARINT_DELTA_V1: &str = "varint_delta_v1"; -pub const POSITIONS_LAYOUT_SHARED_STREAM_V2: &str = "shared_stream_v2"; -pub const POSITIONS_CODEC_VARINT_DOC_DELTA_V2: &str = "varint_doc_delta_v2"; -pub const POSITIONS_CODEC_PACKED_DELTA_V1: &str = "packed_delta_v1"; -pub const DELETED_FRAGMENTS_COL: &str = "deleted_fragments"; - -// Just a heuristic when we need to pre-allocate memory for tokens -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()]))); - -pub fn resolve_fts_format_version( - value: Option<&str>, -) -> std::result::Result { - match value { - Some(value) => value.parse(), - None => Ok(default_fts_format_version()), - } -} - -pub fn default_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V2 -} - -pub fn current_fts_format_version() -> InvertedListFormatVersion { - default_fts_format_version() -} - -pub fn max_supported_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V3 -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] -pub enum InvertedListFormatVersion { - V1, - #[default] - V2, - V3, -} - -impl InvertedListFormatVersion { - pub fn from_posting_tail_codec(codec: PostingTailCodec) -> Self { - match codec { - PostingTailCodec::Fixed32 => Self::V1, - PostingTailCodec::VarintDelta => Self::V2, - } - } - - pub fn from_posting_tail_codec_and_block_size( - codec: PostingTailCodec, - block_size: usize, - ) -> Result { - validate_block_size(block_size)?; - let format_version = match (codec, block_size) { - (PostingTailCodec::Fixed32, LEGACY_BLOCK_SIZE) => Self::V1, - (PostingTailCodec::VarintDelta, LEGACY_BLOCK_SIZE) => Self::V2, - (PostingTailCodec::VarintDelta, 256) => Self::V3, - (PostingTailCodec::Fixed32, 256) => { - return Err(Error::invalid_input( - "FTS format_version=3 requires the varint-delta posting tail codec".to_string(), - )); - } - _ => unreachable!("validate_block_size limits supported block sizes"), - }; - validate_format_version_block_size(format_version, block_size)?; - Ok(format_version) - } - - pub fn index_version(self) -> u32 { - match self { - Self::V1 => INVERTED_INDEX_VERSION_V1, - Self::V2 => INVERTED_INDEX_VERSION_V2, - Self::V3 => INVERTED_INDEX_VERSION_V3, - } - } - - pub fn posting_tail_codec(self) -> PostingTailCodec { - match self { - Self::V1 => PostingTailCodec::Fixed32, - Self::V2 | Self::V3 => PostingTailCodec::VarintDelta, - } - } - - pub fn position_codec(self) -> Option { - match self { - Self::V1 => None, - Self::V2 | Self::V3 => Some(PositionStreamCodec::PackedDelta), - } - } - - pub fn uses_shared_position_stream(self) -> bool { - matches!(self, Self::V2 | Self::V3) - } -} - -impl FromStr for InvertedListFormatVersion { - type Err = Error; - - fn from_str(s: &str) -> std::result::Result { - match s.trim() { - "1" | "v1" | "V1" => Ok(Self::V1), - "2" | "v2" | "V2" => Ok(Self::V2), - "3" | "v3" | "V3" => Ok(Self::V3), - other => Err(Error::index(format!( - "unsupported FTS format version {}, expected 1, 2, or 3", - other - ))), - } - } -} - -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"), - } -} - -pub fn validate_format_version_block_size( - format_version: InvertedListFormatVersion, - block_size: usize, -) -> Result<()> { - validate_block_size(block_size)?; - match (format_version, block_size) { - (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE) - | (InvertedListFormatVersion::V3, _) => Ok(()), - (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, 256) => { - Err(Error::invalid_input(format!( - "FTS format_version={} is incompatible with block_size=256; use format_version=3", - format_version.index_version() - ))) - } - _ => unreachable!("validate_block_size limits supported block sizes"), - } -} - -#[derive(Debug)] -struct PartitionCandidates { - tokens_by_position: Vec, - grouped_expansions: Vec, - candidates: Vec>, -} - -struct ModernSearchRequest<'a> { - tokens: Arc, - params: Arc, - operator: Operator, - mask: Arc, - metrics: Arc, - scorer: &'a MemBM25Scorer, - impact_scorer: Arc, - limit: usize, -} - -/// Typed identity for one modern candidate after partition-local scoring. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct PartitionDocId { - partition_ordinal: u32, - doc_id: DocId, -} - -impl PartitionDocId { - fn try_new(partition_ordinal: usize, doc_id: DocId) -> Result { - Ok(Self { - partition_ordinal: u32::try_from(partition_ordinal).map_err(|_| { - Error::index(format!( - "FTS partition ordinal {partition_ordinal} exceeds candidate identity capacity" - )) - })?, - doc_id, - }) - } - - fn partition_ordinal(self) -> usize { - self.partition_ordinal as usize - } -} - -#[derive(Debug, Clone)] -struct ScoredPartitionDoc { - document: PartitionDocId, - score: OrderedFloat, -} - -impl ScoredPartitionDoc { - fn new(document: PartitionDocId, score: f32) -> Self { - Self { - document, - score: OrderedFloat(score), - } - } -} - -impl PartialEq for ScoredPartitionDoc { - fn eq(&self, other: &Self) -> bool { - self.score == other.score - } -} - -impl Eq for ScoredPartitionDoc {} - -impl PartialOrd for ScoredPartitionDoc { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for ScoredPartitionDoc { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.score.cmp(&other.score) - } -} - -const MAX_CONCURRENT_ADDRESS_READ_BYTES: usize = 64 * 1024 * 1024; - -fn address_read_concurrency(io_parallelism: usize, largest_read_bytes: usize) -> usize { - let io_parallelism = io_parallelism.max(1); - if largest_read_bytes == 0 { - return io_parallelism; - } - io_parallelism.min( - MAX_CONCURRENT_ADDRESS_READ_BYTES - .checked_div(largest_read_bytes) - .unwrap_or(0) - .max(1), - ) -} - -fn push_scored_key( - candidates: &mut BinaryHeap>, - limit: usize, - key: u64, - score: f32, -) { - if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(key, score))); - } else if candidates - .peek() - .is_some_and(|candidate| candidate.0.score.0 < score) - { - candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(key, score))); - } -} - -fn push_scored_partition_doc( - candidates: &mut BinaryHeap>, - limit: usize, - document: PartitionDocId, - score: f32, -) { - if candidates.len() < limit { - candidates.push(Reverse(ScoredPartitionDoc::new(document, score))); - } else if candidates - .peek() - .is_some_and(|candidate| candidate.0.score.0 < score) - { - candidates.pop(); - candidates.push(Reverse(ScoredPartitionDoc::new(document, score))); - } -} - -fn rescore_partition_candidates( - partition: PartitionCandidates, - scorer: &MemBM25Scorer, - idf_cache: &mut HashMap, -) -> Vec<(C, f32)> { - let PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates, - } = partition; - let idf_by_position = tokens_by_position - .iter() - .map(|token| { - *idf_cache - .entry(token.clone()) - .or_insert_with(|| scorer.query_weight(token)) - }) - .collect::>(); - let grouped_positions = grouped_expansions - .iter() - .map(|group| group.position) - .collect::>(); - - candidates - .into_iter() - .map( - |DocCandidate { - document, - posting_doc_id, - freqs, - doc_length, - }| { - let mut score = 0.0; - for (term_index, freq) in freqs { - if grouped_positions.contains(&term_index) { - continue; - } - debug_assert!((term_index as usize) < idf_by_position.len()); - score += - idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); - } - for group in &grouped_expansions { - for term in group.terms.iter() { - let Some(freq) = term.frequency(posting_doc_id) else { - continue; - }; - score += term.query_weight() * scorer.doc_weight(freq, doc_length); - } - } - (document, score) - }, - ) - .collect() -} - -#[derive(Debug)] -pub(super) struct LoadedPostings { - pub(super) postings: Vec, - grouped_expansions: Vec, - impact_safe: bool, - exact_scoring_required: bool, -} - -enum LoadedDocLengths { - Legacy(Arc), - Modern(Arc), -} - -impl LoadedDocLengths { - fn scoring_num_tokens(&self, doc_id: u32) -> u32 { - match self { - Self::Legacy(docs) => docs.scoring_num_tokens(doc_id), - Self::Modern(lengths) => lengths.scoring(DocId::new(doc_id)), - } - } - - fn num_tokens_by_row_id(&self, row_id: u64) -> u32 { - match self { - Self::Legacy(docs) => docs.num_tokens_by_row_id(row_id), - Self::Modern(_) => unreachable!("modern posting lists use dense DocIds"), - } - } -} - -impl LoadedPostings { - fn empty() -> Self { - Self { - postings: Vec::new(), - grouped_expansions: Vec::new(), - impact_safe: false, - exact_scoring_required: false, - } - } -} - -#[derive(Debug)] -struct GroupedExpansionTerms { - position: u32, - terms: Arc<[GroupedTermScorer]>, -} - -#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Default)] -pub enum TokenSetFormat { - Arrow, - #[default] - Fst, -} - -impl Display for TokenSetFormat { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Arrow => f.write_str("arrow"), - Self::Fst => f.write_str("fst"), - } - } -} - -impl FromStr for TokenSetFormat { - type Err = Error; - - fn from_str(s: &str) -> std::result::Result { - match s.trim() { - "" => Ok(Self::Arrow), - "arrow" => Ok(Self::Arrow), - "fst" => Ok(Self::Fst), - other => Err(Error::index(format!( - "unsupported token set format {}", - other - ))), - } - } -} - -impl DeepSizeOf for TokenSetFormat { - fn deep_size_of_children(&self, _: &mut lance_core::deepsize::Context) -> usize { - 0 - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] -pub enum PositionStreamCodec { - VarintDocDelta, - #[default] - PackedDelta, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] -pub enum PostingTailCodec { - Fixed32, - #[default] - VarintDelta, -} - -impl PostingTailCodec { - pub fn as_str(self) -> &'static str { - match self { - Self::Fixed32 => POSTING_TAIL_CODEC_FIXED32_V1, - Self::VarintDelta => POSTING_TAIL_CODEC_VARINT_DELTA_V1, - } - } - - fn from_metadata_value(value: &str) -> Result { - match value.trim() { - POSTING_TAIL_CODEC_FIXED32_V1 => Ok(Self::Fixed32), - POSTING_TAIL_CODEC_VARINT_DELTA_V1 => Ok(Self::VarintDelta), - other => Err(Error::index(format!( - "unsupported posting tail codec {}", - other - ))), - } - } -} - -pub(super) fn parse_posting_tail_codec( - metadata: &HashMap, -) -> Result { - Ok(metadata - .get(POSTING_TAIL_CODEC_KEY) - .map(|codec| PostingTailCodec::from_metadata_value(codec)) - .transpose()? - .unwrap_or(PostingTailCodec::Fixed32)) -} - -pub(super) fn parse_posting_block_size(metadata: &HashMap) -> Result { - metadata - .get(POSTING_BLOCK_SIZE_KEY) - .map(|value| { - let block_size = value.parse::().map_err(|err| { - Error::index(format!( - "invalid {POSTING_BLOCK_SIZE_KEY} metadata value {value:?}: {err}" - )) - })?; - validate_block_size(block_size) - }) - .transpose() - .map(|block_size| block_size.unwrap_or(LEGACY_BLOCK_SIZE)) -} - -impl PositionStreamCodec { - pub fn as_str(self) -> &'static str { - match self { - Self::VarintDocDelta => POSITIONS_CODEC_VARINT_DOC_DELTA_V2, - Self::PackedDelta => POSITIONS_CODEC_PACKED_DELTA_V1, - } - } - - fn from_metadata_value(value: &str) -> Result { - match value.trim() { - POSITIONS_CODEC_VARINT_DOC_DELTA_V2 => Ok(Self::VarintDocDelta), - POSITIONS_CODEC_PACKED_DELTA_V1 => Ok(Self::PackedDelta), - other => Err(Error::index(format!( - "unsupported positions codec {}", - other - ))), - } - } -} - -fn parse_shared_position_codec(metadata: &HashMap) -> Result { - if let Some(codec) = metadata.get(POSITIONS_CODEC_KEY) { - return PositionStreamCodec::from_metadata_value(codec); - } - - match metadata - .get(POSITIONS_LAYOUT_KEY) - .map(|layout| layout.as_str()) - { - Some(POSITIONS_LAYOUT_SHARED_STREAM_V2) => Ok(PositionStreamCodec::VarintDocDelta), - _ => Ok(PositionStreamCodec::VarintDocDelta), - } -} - -pub(super) fn parse_format_version_from_metadata( - metadata: &HashMap, -) -> Result { - if let Some(value) = metadata.get(FTS_FORMAT_VERSION_KEY) { - let format_version = InvertedListFormatVersion::from_str(value)?; - let block_size = parse_posting_block_size(metadata)?; - validate_format_version_block_size(format_version, block_size)?; - return Ok(format_version); - } - let block_size = parse_posting_block_size(metadata)?; - if block_size == 256 { - if metadata - .get(POSTING_TAIL_CODEC_KEY) - .map(|_| parse_posting_tail_codec(metadata)) - .transpose()? - .is_some_and(|posting_tail_codec| posting_tail_codec != PostingTailCodec::VarintDelta) - { - return Err(Error::index( - "FTS block_size=256 requires the varint-delta posting tail codec".to_string(), - )); - } - return Ok(InvertedListFormatVersion::V3); - } - if metadata.contains_key(POSITIONS_CODEC_KEY) || metadata.contains_key(POSITIONS_LAYOUT_KEY) { - return Ok(InvertedListFormatVersion::V2); - } - if parse_posting_tail_codec(metadata)? == PostingTailCodec::VarintDelta { - Ok(InvertedListFormatVersion::V2) - } else { - Ok(InvertedListFormatVersion::V1) - } -} - -#[derive(Debug, Default)] -struct InvertedPrewarmState { - query_ready: bool, - positions_ready: bool, -} - -impl InvertedPrewarmState { - fn satisfies(&self, with_position: bool) -> bool { - self.query_ready && (!with_position || self.positions_ready) - } -} - -#[derive(Clone)] -pub struct InvertedIndex { - params: InvertedIndexParams, - store: Arc, - tokenizer: Box, - token_set_format: TokenSetFormat, - format_version: InvertedListFormatVersion, - pub(crate) partitions: Vec>, - corpus_stats: Arc>, - prewarm_state: Arc>, - /// Optimistic fast-path hint. Cache eviction can make it stale; the - /// resident resolver clears it when a weak projection upgrade misses. - document_projections_resident: Arc, - // Fragments which are contained in the index, but no longer in the dataset. - // These should be pruned at search time since we don't prune them at update time. - deleted_fragments: RoaringBitmap, -} - -impl Debug for InvertedIndex { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("InvertedIndex") - .field("params", &self.params) - .field("token_set_format", &self.token_set_format) - .field("format_version", &self.format_version) - .field("partitions", &self.partitions) - .field("deleted_fragments", &self.deleted_fragments) - .finish() - } -} - -impl DeepSizeOf for InvertedIndex { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.partitions.deep_size_of_children(context) - } -} - -impl InvertedIndex { - fn format_version(&self) -> InvertedListFormatVersion { - self.format_version - } - - fn index_version(&self) -> u32 { - if self.params.get_document_granularity().is_list_element() { - return INVERTED_INDEX_VERSION_V3; - } - match (self.token_set_format, self.format_version()) { - ( - TokenSetFormat::Arrow, - InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, - ) => 0, - (_, format_version) => format_version.index_version(), - } - } - - fn posting_tail_codec(&self) -> PostingTailCodec { - self.partitions - .first() - .map(|partition| partition.inverted_list.posting_tail_codec()) - .unwrap_or_default() - } - - fn to_builder(&self) -> InvertedIndexBuilder { - self.to_builder_with_offset(None) - } - - fn to_builder_with_offset(&self, fragment_mask: Option) -> InvertedIndexBuilder { - if self.is_legacy() { - // for legacy format, we re-create the index in the new format - InvertedIndexBuilder::from_existing_index( - self.params.clone(), - None, - Vec::new(), - self.token_set_format, - fragment_mask, - self.deleted_fragments.clone(), - ) - .with_posting_tail_codec(self.posting_tail_codec()) - } else { - let partitions = match fragment_mask { - Some(fragment_mask) => self - .partitions - .iter() - // Filter partitions that belong to the specified fragment - // The mask contains fragment_id in high 32 bits, we check if partition's - // fragment_id matches by comparing the masked result with the original mask - .filter(|part| part.belongs_to_fragment(fragment_mask)) - .map(|part| part.id()) - .collect(), - None => self.partitions.iter().map(|part| part.id()).collect(), - }; - - InvertedIndexBuilder::from_existing_index( - self.params.clone(), - Some(self.store.clone()), - partitions, - self.token_set_format, - fragment_mask, - self.deleted_fragments.clone(), - ) - .with_format_version(self.format_version()) - } - } - - pub fn tokenizer(&self) -> Box { - self.tokenizer.clone() - } - - pub fn params(&self) -> &InvertedIndexParams { - &self.params - } - - /// Returns the number of partitions in this inverted index. - pub fn partition_count(&self) -> usize { - self.partitions.len() - } - /// Returns the set of fragments which are contained in the index, but no longer in the dataset. - /// - /// Most other indices remove data from deleted fragments when the index updates (copy-on-write). - /// However, this would require an expensive copy of the FTS index. Instead, we track the deleted - /// fragments and prune them at search time (merge-on-read). - pub fn deleted_fragments(&self) -> &RoaringBitmap { - &self.deleted_fragments - } - - pub async fn merge_segments( - segments: &[Arc], - new_data: SendableRecordBatchStream, - dest_store: &dyn IndexStore, - old_data_filter: Option, - progress: Arc, - ) -> Result { - let Some(first) = segments.first() else { - return Err(Error::invalid_input( - "cannot merge inverted index without at least one source segment".to_string(), - )); - }; - - for segment in segments.iter().skip(1) { - if segment.params != first.params { - return Err(Error::index( - "cannot merge inverted index segments with different parameters".to_string(), - )); - } - if segment.token_set_format != first.token_set_format { - return Err(Error::index( - "cannot merge inverted index segments with different token set formats" - .to_string(), - )); - } - if segment.format_version() != first.format_version() { - return Err(Error::index( - "cannot merge inverted index segments with different format versions" - .to_string(), - )); - } - if segment.posting_tail_codec() != first.posting_tail_codec() { - return Err(Error::index( - "cannot merge inverted index segments with different posting tail codecs" - .to_string(), - )); - } - } - - 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()); - let files = builder - .update_from_segments(new_data, dest_store, segments, old_data_filter) - .await?; - - let details = pbold::InvertedIndexDetails::try_from(&first.params)?; - - Ok(CreatedIndex { - index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: first.index_version(), - files, - }) - } - - /// Build a single-segment [`MemBM25Scorer`] whose per-term IDF table - /// covers every token that the per-partition scoring loop will look - /// up. For fuzzy queries that means the union of Levenshtein - /// expansions, not just the raw query tokens — otherwise - /// `query_weight(expanded_token)` returns 0 and the BM25 contribution - /// of every expanded match is discarded. - pub async fn bm25_base_scorer( - &self, - query_tokens: &Tokens, - params: &FtsSearchParams, - metrics: Option<&dyn MetricsCollector>, - ) -> Result { - if matches!(params.fuzziness, Some(n) if n != 0) { - let expanded = self.expand_fuzzy_tokens(query_tokens, params)?; - self.bm25_scorer_for_final_tokens(&expanded, metrics).await - } else { - self.bm25_scorer_for_final_tokens(query_tokens, metrics) - .await - } - } - - /// Scorer for a token list that needs no further fuzzy expansion: dedup - /// the terms and pull their document frequencies. `bm25_search` calls - /// this with the tokens it already expanded, so the expansion runs once - /// per query rather than once for the scorer and once per partition. - async fn bm25_scorer_for_final_tokens( - &self, - tokens: &Tokens, - metrics: Option<&dyn MetricsCollector>, - ) -> Result { - let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; - let mut terms: Vec = Vec::new(); - let mut seen = HashSet::new(); - for token in tokens { - if seen.insert(token.to_string()) { - terms.push(token.to_string()); - } - } - let mut token_docs = HashMap::with_capacity(terms.len()); - for term in &terms { - let df = self.df_for_term(term, metrics).await?; - token_docs.insert(term.clone(), df); - } - Ok(MemBM25Scorer::new(total_tokens, num_docs, token_docs)) - } - - pub async fn bm25_stats_for_terms( - &self, - terms: &[String], - metrics: Option<&dyn MetricsCollector>, - ) -> Result<(u64, usize, Vec)> { - let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; - let token_docs = - futures::future::try_join_all(terms.iter().map(|term| self.df_for_term(term, metrics))) - .await?; - Ok((total_tokens, num_docs, token_docs)) - } - - /// Aggregate immutable per-partition corpus statistics. New modern files - /// read both values from the already-opened docs footer; older partitioned - /// files scan `_num_tokens` once as a compatibility fallback. - async fn aggregate_corpus_stats(&self) -> Result<(u64, usize)> { - self.corpus_stats - .get_or_try_init(|| async { - let io_parallelism = self.store.io_parallelism(); - let futures = self - .partitions - .iter() - .map(|p| { - let part = p.clone(); - async move { part.docs.stats().await } - }) - .collect::>(); - let stats = stream::iter(futures) - .buffer_unordered(io_parallelism) - .try_collect::>() - .await?; - let mut total_tokens = 0_u64; - let mut num_docs = 0_usize; - for stat in stats { - total_tokens = total_tokens - .checked_add(stat.total_tokens) - .ok_or_else(|| Error::index("FTS corpus token count overflows u64"))?; - num_docs = num_docs - .checked_add(stat.num_docs) - .ok_or_else(|| Error::index("FTS corpus document count overflows usize"))?; - } - Ok((total_tokens, num_docs)) - }) - .await - .copied() - } - - /// Sum the posting-list length for `term` across this index's partitions - /// via single-row reads, with partition lookups bounded by the store's - /// `io_parallelism()`. - async fn df_for_term( - &self, - term: &str, - metrics: Option<&dyn MetricsCollector>, - ) -> Result { - let io_parallelism = self.store.io_parallelism(); - let futures = self - .partitions - .iter() - .map(|part| { - let part = part.clone(); - async move { - match part.tokens.get(term) { - Some(token_id) => { - part.inverted_list - .posting_len_for_token(token_id, metrics) - .await - } - None => Ok(0), - } - } - }) - .collect::>(); - let dfs: Vec = stream::iter(futures) - .buffer_unordered(io_parallelism) - .try_collect() - .await?; - Ok(dfs.into_iter().sum()) - } - - /// Expand fuzzy query tokens against all partitions in this segment. - /// - /// `params.max_expansions` caps the whole query's expansion, not any - /// single partition's: for each query token the per-partition candidates - /// (each streamed in FST key order) merge into one lexicographically - /// ordered set, and the remaining budget takes a prefix of it. The - /// selected terms are a pure function of the segment's vocabulary, so - /// splitting the same corpus into more partitions cannot change which - /// terms a fuzzy query matches. - pub fn expand_fuzzy_tokens(&self, tokens: &Tokens, params: &FtsSearchParams) -> Result { - let mut expanded_tokens = Vec::new(); - let mut expanded_positions = Vec::new(); - let mut seen = HashSet::new(); - for token_idx in 0..tokens.len() { - let remaining = params.max_expansions.saturating_sub(expanded_tokens.len()); - if remaining == 0 { - break; - } - let token = tokens.get_token(token_idx); - let position = tokens.position(token_idx); - // Each partition contributes at most its `remaining` - // lexicographically smallest candidates, so the global - // lex-smallest `remaining` selection below is unaffected by the - // per-partition truncation. - let mut candidates = BTreeSet::new(); - let base_prefix_len = tokens.token_type().prefix_len(token) as u32; - for partition in &self.partitions { - partition.collect_fuzzy_candidates( - token, - base_prefix_len, - params, - remaining, - &mut candidates, - )?; - } - for candidate in candidates { - if expanded_tokens.len() >= params.max_expansions { - break; - } - if seen.insert((candidate.clone(), position)) { - expanded_tokens.push(candidate); - expanded_positions.push(position); - } - } - } - Ok(Tokens::with_positions( - expanded_tokens, - expanded_positions, - tokens.token_type().clone(), - )) - } - - /// Search documents that match the query and return row ids sorted by BM25 score. - /// - /// When `base_scorer` is provided, search uses those corpus-level BM25 statistics - /// instead of deriving them from this segment alone. - #[instrument(level = "debug", skip_all)] - pub async fn bm25_search( - &self, - tokens: Arc, - params: Arc, - operator: Operator, - prefilter: Arc, - 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 - // corpus happens to be partitioned. - let tokens = if matches!(params.fuzziness, Some(n) if n != 0) { - let expanded = Arc::new(self.expand_fuzzy_tokens(tokens.as_ref(), params.as_ref())?); - if operator == Operator::And || params.phrase_slop.is_some() { - // AND/phrase semantics require every original token position - // to keep at least one expansion; a position that expands to - // nothing anywhere in the segment can never be matched. - let surviving = (0..expanded.len()) - .map(|idx| expanded.position(idx)) - .collect::>(); - if (0..tokens.len()).any(|idx| !surviving.contains(&tokens.position(idx))) { - return Ok(Vec::new()); - } - } - expanded - } else { - tokens - }; - - // The wand only consults `scorer.doc_weight`, which is metadata-free. - // The outer aggregation below consults `scorer.query_weight`, which - // hits per-token `posting_len`; building a `MemBM25Scorer` with - // precomputed per-term IDFs avoids the v2 bulk metadata pull. - let local_scorer; - let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer { - base_scorer - } else { - local_scorer = self - .bm25_scorer_for_final_tokens(tokens.as_ref(), Some(metrics.as_ref())) - .await?; - &local_scorer - }; - let impact_scorer = Arc::new(scorer.clone()); - - let limit = params.limit.unwrap_or(usize::MAX); - if limit == 0 { - return Ok(Vec::new()); - } - let mask = prefilter.mask(); - if self.is_legacy() { - let (row_ids, scores) = self - .bm25_search_legacy( - tokens, - params, - operator, - mask, - metrics, - scorer, - impact_scorer, - limit, - ) - .await?; - Ok(row_ids - .into_iter() - .zip(scores) - .map(|(row_id, score)| ScoredDoc::new(row_id, score)) - .collect()) - } else { - self.bm25_search_modern(ModernSearchRequest { - tokens, - params, - operator, - mask, - metrics, - scorer, - impact_scorer, - limit, - }) - .await - } - } - - #[allow(clippy::too_many_arguments)] - async fn bm25_search_legacy( - &self, - tokens: Arc, - params: Arc, - operator: Operator, - mask: Arc, - metrics: Arc, - scorer: &MemBM25Scorer, - impact_scorer: Arc, - limit: usize, - ) -> Result<(Vec, Vec)> { - let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); - let io_parallelism = self.store.io_parallelism(); - let parts = self - .partitions - .chunks(fts_search_chunk()) - .map(|chunk| { - let chunk = chunk.to_vec(); - let tokens = tokens.clone(); - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); - let impact_scorer = impact_scorer.clone(); - let impact_shared_threshold = impact_shared_threshold.clone(); - async move { - let loads = chunk.into_iter().map(|part| { - let tokens = tokens.clone(); - let params = params.clone(); - let metrics = metrics.clone(); - let impact_scorer = impact_scorer.clone(); - let impact_shared_threshold = impact_shared_threshold.clone(); - async move { - let LoadedPostings { - postings, - grouped_expansions, - impact_safe, - exact_scoring_required, - } = part - .load_posting_lists( - tokens.as_ref(), - params.as_ref(), - operator, - impact_scorer.as_ref(), - metrics.as_ref(), - false, - ) - .await?; - if postings.is_empty() { - return Result::Ok(None); - } - let max_position = postings - .iter() - .map(|posting| posting.term_index() as usize) - .max() - .unwrap_or_default(); - let mut tokens_by_position = vec![String::new(); max_position + 1]; - for posting in &postings { - tokens_by_position[posting.term_index() as usize] = - posting.token().to_owned(); - } - let docs = part.docs.legacy().cloned().ok_or_else(|| { - Error::internal("legacy index contains modern partition documents") - })?; - let use_global_scorer = impact_safe || exact_scoring_required; - let threshold = if use_global_scorer { - impact_shared_threshold - } else { - Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) - }; - let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); - Result::Ok(Some(( - part, - docs, - postings, - wand_scorer, - threshold, - tokens_by_position, - grouped_expansions, - ))) - } - }); - let loaded = stream::iter(loads) - .buffer_unordered(io_parallelism) - .try_collect::>() - .await? - .into_iter() - .flatten() - .collect::>(); - if loaded.is_empty() { - return Result::Ok(Vec::new()); - } - - let results = spawn_cpu(move || { - let mut results = Vec::with_capacity(loaded.len()); - for ( - part, - docs, - postings, - wand_scorer, - threshold, - tokens_by_position, - grouped_expansions, - ) in loaded - { - let candidates = part.bm25_search_legacy( - docs.as_ref(), - params.as_ref(), - operator, - mask.as_ref(), - postings, - wand_scorer, - metrics.as_ref(), - threshold, - )?; - results.push(PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates, - }); - } - Result::Ok(results) - }) - .await?; - Result::Ok(results) - } - }) - .collect::>(); - - let mut ranked = BinaryHeap::new(); - let mut idf_cache = HashMap::new(); - let mut parts = stream::iter(parts) - .buffer_unordered(get_num_compute_intensive_cpus().min(32)) - .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok))) - .try_flatten(); - while let Some(partition) = parts.try_next().await? { - for (row_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { - push_scored_key(&mut ranked, limit, row_id, score); - } - } - Ok(ranked - .into_sorted_vec() - .into_iter() - .map(|Reverse(doc)| (doc.row_id, doc.score.0)) - .unzip()) - } - - async fn bm25_search_modern(&self, request: ModernSearchRequest<'_>) -> Result> { - // Select a concrete completion path before candidate search. The - // fully resident future never builds deferred address-read state, while - // a cold query keeps DocIds until its final bounded I/O phase. - if self.has_resident_document_projections() { - self.bm25_search_modern_resident(request).await - } else { - self.bm25_search_modern_deferred(request).await - } - } - - fn has_resident_document_projections(&self) -> bool { - if self.document_projections_resident.load(Ordering::Acquire) { - return true; - } - let resident = self.document_projections_resident_now(); - if resident { - self.document_projections_resident - .store(true, Ordering::Release); - } - resident - } - - fn document_projections_resident_now(&self) -> bool { - self.partitions.iter().all(|partition| { - partition - .docs - .modern() - .is_some_and(|documents| documents.projection_resident()) - }) - } - - async fn bm25_search_modern_resident( - &self, - request: ModernSearchRequest<'_>, - ) -> Result> { - let ranked = self.bm25_search_modern_candidates(request).await?; - if let Some(result) = self.resolve_resident_modern_candidates(&ranked)? { - return Ok(result); - } - self.document_projections_resident - .store(false, Ordering::Release); - self.resolve_deferred_modern_candidates(ranked).await - } - - async fn bm25_search_modern_deferred( - &self, - request: ModernSearchRequest<'_>, - ) -> Result> { - // Old partitioned files without persisted stats populate their - // fallback stats before deferred candidate orchestration. A resident - // search can skip this full-index synchronization: standard prewarm - // has already initialized it, while any independently resident - // partition loads its lengths before constructing a local scorer. - if self.corpus_stats.get().is_none() { - self.aggregate_corpus_stats().await?; - } - // For new-format indexes, aggregate_corpus_stats reads corpus stats from - // persisted schema metadata (O(1)) without loading doc lengths as a side - // effect. Pre-load lengths in parallel now for partitions that contain at - // least one query token, so the scoring phase gets cache hits instead of - // issuing sequential per-partition IO. Partitions with no matching terms - // are skipped to preserve the no-load optimization for no-hit queries. - let io_parallelism = self.store.io_parallelism(); - let uncached_lengths = self - .partitions - .iter() - .filter_map(|part| { - let docs = part.docs.modern()?.clone(); - if docs.cached_lengths().is_some() { - return None; - } - let has_match = (0..request.tokens.len()) - .any(|i| part.tokens.get(request.tokens.get_token(i)).is_some()); - has_match.then_some(async move { docs.lengths().await.map(|_| ()) }) - }) - .collect::>(); - if !uncached_lengths.is_empty() { - stream::iter(uncached_lengths) - .buffer_unordered(io_parallelism) - .try_collect::>() - .await?; - } - let ranked = self.bm25_search_modern_candidates(request).await?; - self.resolve_deferred_modern_candidates(ranked).await - } - - async fn bm25_search_modern_candidates( - &self, - request: ModernSearchRequest<'_>, - ) -> Result>> { - let ModernSearchRequest { - tokens, - params, - operator, - mask, - metrics, - scorer, - impact_scorer, - limit, - } = request; - if self.partitions.len() > u32::MAX as usize { - return Err(Error::index(format!( - "FTS partition count {} exceeds candidate identity capacity", - self.partitions.len() - ))); - } - let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); - let io_parallelism = self.store.io_parallelism(); - let parts = self - .partitions - .chunks(fts_search_chunk()) - .enumerate() - .map(|(chunk_ordinal, chunk)| { - let first_partition_ordinal = chunk_ordinal * fts_search_chunk(); - let chunk = chunk - .iter() - .cloned() - .enumerate() - .map(|(offset, part)| (first_partition_ordinal + offset, part)) - .collect::>(); - let tokens = tokens.clone(); - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); - let impact_scorer = impact_scorer.clone(); - let impact_shared_threshold = impact_shared_threshold.clone(); - async move { - let loads = chunk.into_iter().map(|(partition_ordinal, part)| { - let tokens = tokens.clone(); - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); - let impact_scorer = impact_scorer.clone(); - let impact_shared_threshold = impact_shared_threshold.clone(); - async move { - let LoadedPostings { - postings, - grouped_expansions, - impact_safe, - exact_scoring_required, - } = part - .load_posting_lists( - tokens.as_ref(), - params.as_ref(), - operator, - impact_scorer.as_ref(), - metrics.as_ref(), - false, - ) - .await?; - if postings.is_empty() { - return Result::Ok(None); - } - let documents = part.docs.modern().cloned().ok_or_else(|| { - Error::internal("modern index contains legacy partition documents") - })?; - let materialize_selected = operator == Operator::Or - && mask.max_len().is_some_and(|selected| { - u128::from(selected).saturating_mul(100) - <= u128::from(*FLAT_SEARCH_PERCENT_THRESHOLD) - .saturating_mul(documents.len() as u128) - }); - let visibility = match documents - .immediate_visibility(mask.clone(), materialize_selected) - { - Some(visibility) => visibility, - None => { - documents - .visibility(mask.clone(), materialize_selected) - .await? - } - }; - if visibility.is_empty() { - return Result::Ok(None); - } - let lengths = match documents.cached_lengths() { - Some(lengths) => lengths, - None => documents.lengths().await?, - }; - let max_position = postings - .iter() - .map(|posting| posting.term_index() as usize) - .max() - .unwrap_or_default(); - let mut tokens_by_position = vec![String::new(); max_position + 1]; - for posting in &postings { - tokens_by_position[posting.term_index() as usize] = - posting.token().to_owned(); - } - let use_global_scorer = impact_safe || exact_scoring_required; - let threshold = if use_global_scorer { - impact_shared_threshold - } else { - Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) - }; - let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); - Result::Ok(Some(( - partition_ordinal, - part, - lengths, - visibility, - postings, - wand_scorer, - threshold, - tokens_by_position, - grouped_expansions, - ))) - } - }); - let loaded = stream::iter(loads) - .buffer_unordered(io_parallelism) - .try_collect::>() - .await? - .into_iter() - .flatten() - .collect::>(); - if loaded.is_empty() { - return Result::Ok(Vec::new()); - } - - let results = spawn_cpu(move || { - let mut results = Vec::with_capacity(loaded.len()); - for ( - partition_ordinal, - part, - lengths, - visibility, - postings, - wand_scorer, - threshold, - tokens_by_position, - grouped_expansions, - ) in loaded - { - let candidates = part.bm25_search_modern( - lengths.as_ref(), - &visibility, - params.as_ref(), - operator, - postings, - wand_scorer, - metrics.as_ref(), - threshold, - )?; - results.push(( - partition_ordinal, - PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates, - }, - )); - } - Result::Ok(results) - }) - .await?; - Result::Ok(results) - } - }) - .collect::>(); - - let mut ranked = BinaryHeap::new(); - let mut idf_cache = HashMap::new(); - let mut parts = stream::iter(parts) - .buffer_unordered(get_num_compute_intensive_cpus().min(32)) - .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok))) - .try_flatten(); - while let Some((partition_ordinal, partition)) = parts.try_next().await? { - for (doc_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { - push_scored_partition_doc( - &mut ranked, - limit, - PartitionDocId::try_new(partition_ordinal, doc_id)?, - score, - ); - } - } - - Ok(ranked.into_sorted_vec()) - } - - fn resolve_resident_modern_candidates( - &self, - ranked: &[Reverse], - ) -> Result>> { - if self.partitions.iter().any(|partition| { - partition - .docs - .modern() - .is_some_and(|documents| documents.coordinate_rank() > 0) - }) { - return Ok(None); - } - let mut resolved_documents = ranked - .iter() - .map(|Reverse(candidate)| ScoredDoc::new(0, candidate.score.0)) - .collect::>(); - let mut by_partition = BTreeMap::>::new(); - for (rank, Reverse(candidate)) in ranked.iter().enumerate() { - let partition_ordinal = candidate.document.partition_ordinal(); - let doc_id = candidate.document.doc_id; - by_partition - .entry(partition_ordinal) - .or_default() - .push((rank, doc_id)); - } - for (partition_ordinal, entries) in by_partition { - let documents = self - .partitions - .get(partition_ordinal) - .and_then(|partition| partition.docs.modern()) - .ok_or_else(|| { - Error::internal(format!( - "resident FTS candidates reference missing modern partition ordinal {partition_ordinal}" - )) - })?; - let doc_ids = entries - .iter() - .map(|(_, doc_id)| *doc_id) - .collect::>(); - let Some(resolved) = documents.cached_row_addresses(&doc_ids)? else { - return Ok(None); - }; - for ((rank, _), address) in entries.into_iter().zip(resolved) { - resolved_documents[rank].row_id = address; - } - } - Ok(Some(resolved_documents)) - } - - async fn resolve_deferred_modern_candidates( - &self, - ranked: Vec>, - ) -> Result> { - let mut resolved_documents = ranked - .iter() - .map(|Reverse(candidate)| ScoredDoc::new(0, candidate.score.0)) - .collect::>(); - let mut by_partition = BTreeMap::>::new(); - for (rank, Reverse(candidate)) in ranked.iter().enumerate() { - let partition_ordinal = candidate.document.partition_ordinal(); - let doc_id = candidate.document.doc_id; - by_partition - .entry(partition_ordinal) - .or_default() - .push((rank, doc_id)); - } - let mut address_reads = Vec::with_capacity(by_partition.len()); - let mut largest_read_bytes = 0; - for (partition_ordinal, entries) in by_partition { - let documents = self - .partitions - .get(partition_ordinal) - .and_then(|partition| partition.docs.modern()) - .cloned() - .ok_or_else(|| { - Error::internal(format!( - "deferred FTS candidates reference missing modern partition ordinal {partition_ordinal}" - )) - })?; - let doc_ids = entries - .iter() - .map(|(_, doc_id)| *doc_id) - .collect::>(); - largest_read_bytes = - largest_read_bytes.max(documents.estimated_address_read_bytes(&doc_ids)); - address_reads.push(async move { - let resolved = documents.resolve_document_keys(&doc_ids).await?; - Result::Ok((entries, resolved)) - }); - } - let concurrency = address_read_concurrency(self.store.io_parallelism(), largest_read_bytes); - let mut address_reads = stream::iter(address_reads).buffer_unordered(concurrency); - while let Some((entries, resolved)) = address_reads.try_next().await? { - for ((rank, _), (row_id, doc_index)) in entries.into_iter().zip(resolved) { - resolved_documents[rank].row_id = row_id; - resolved_documents[rank].doc_index = doc_index; - } - } - Ok(resolved_documents) - } - - async fn load_legacy_index( - store: Arc, - frag_reuse_index: Option>, - index_cache: &LanceCache, - ) -> Result> { - log::warn!("loading legacy FTS index"); - let tokens_fut = tokio::spawn({ - let store = store.clone(); - async move { - let token_reader = store.open_index_file(TOKENS_FILE).await?; - let tokenizer = token_reader - .schema() - .metadata - .get("tokenizer") - .map(|s| serde_json::from_str::(s)) - .transpose()? - .unwrap_or_default(); - let tokens = TokenSet::load(token_reader, TokenSetFormat::Arrow).await?; - Result::Ok((tokenizer, tokens)) - } - }); - let invert_list_fut = tokio::spawn({ - let store = store.clone(); - let index_cache_clone = index_cache.clone(); - async move { - let invert_list_reader = store.open_index_file(INVERT_LIST_FILE).await?; - let invert_list = - PostingListReader::try_new(invert_list_reader, &index_cache_clone).await?; - Result::Ok(Arc::new(invert_list)) - } - }); - let docs_fut = tokio::spawn({ - let store = store.clone(); - async move { - let docs_reader = store.open_index_file(DOCS_FILE).await?; - let docs = DocSet::load(docs_reader, true, frag_reuse_index).await?; - Result::Ok(docs) - } - }); - - let (tokenizer_config, tokens) = tokens_fut.await??; - let inverted_list = invert_list_fut.await??; - let docs = docs_fut.await??; - - let tokenizer = tokenizer_config.build()?; - - Ok(Arc::new(Self { - params: tokenizer_config, - store: store.clone(), - tokenizer, - token_set_format: TokenSetFormat::Arrow, - format_version: InvertedListFormatVersion::V1, - partitions: vec![Arc::new(InvertedPartition { - id: 0, - store, - tokens, - inverted_list, - docs: PartitionDocumentStore::Legacy(Arc::new(docs)), - token_set_format: TokenSetFormat::Arrow, - })], - corpus_stats: Arc::new(OnceCell::new()), - prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), - document_projections_resident: Arc::new(AtomicBool::new(false)), - deleted_fragments: RoaringBitmap::new(), - })) - } - - pub fn is_legacy(&self) -> bool { - self.partitions.len() == 1 && self.partitions[0].docs.legacy().is_some() - } - - /// Read only the index's [`InvertedIndexParams`], - /// Contains more complete info than manifest's lossy `InvertedIndexDetails`. - pub async fn load_params(store: &dyn IndexStore) -> Result { - match store.open_index_file(METADATA_FILE).await { - Ok(reader) => { - let params = reader - .schema() - .metadata - .get("params") - .ok_or(Error::index("params not found in metadata".to_owned()))?; - Ok(serde_json::from_str::(params)?) - } - Err(metadata_error) => { - // Legacy format: params live in the tokens file (see - // `load_legacy_index`). Some S3 configurations return 403 for - // a missing object, so the readable legacy file is the - // authoritative format probe. - let Ok(reader) = store.open_index_file(TOKENS_FILE).await else { - return Err(metadata_error); - }; - Ok(reader - .schema() - .metadata - .get("tokenizer") - .map(|s| serde_json::from_str::(s)) - .transpose()? - .unwrap_or_default()) - } - } - } - - pub async fn load( - store: Arc, - frag_reuse_index: Option>, - index_cache: &LanceCache, - ) -> Result> - where - Self: Sized, - { - // for new index format, there is a metadata file and multiple partitions, - // each partition is a separate index containing tokens, inverted list and docs. - // for old index format, there is no metadata file, and it's just like a single partition - - match store.open_index_file(METADATA_FILE).await { - Ok(reader) => { - let params = reader - .schema() - .metadata - .get("params") - .ok_or(Error::index("params not found in metadata".to_owned()))?; - let mut params = serde_json::from_str::(params)?; - let partitions = reader - .schema() - .metadata - .get("partitions") - .ok_or(Error::index("partitions not found in metadata".to_owned()))?; - let partitions: Vec = serde_json::from_str(partitions)?; - let token_set_format = reader - .schema() - .metadata - .get(TOKEN_SET_FORMAT_KEY) - .map(|name| TokenSetFormat::from_str(name)) - .transpose()? - .unwrap_or(TokenSetFormat::Arrow); - let format_version = parse_format_version_from_metadata(&reader.schema().metadata)?; - - // Load deleted_fragments if present (optional for backward compatibility) - let deleted_fragments = if reader.num_rows() > 0 { - let metadata_batch = reader.read_range(0..1, None).await?; - if let Some(col) = metadata_batch.column_by_name(DELETED_FRAGMENTS_COL) { - let arr = col.as_binary_opt::().expect_ok()?; - RoaringBitmap::deserialize_from(arr.value(0))? - } else { - RoaringBitmap::new() - } - } else { - RoaringBitmap::new() - }; - - let format = token_set_format; - let partitions = partitions.into_iter().enumerate().map(|(priority, id)| { - let store = store.with_io_priority(priority as u64); - let frag_reuse_index_clone = frag_reuse_index.clone(); - let index_cache_for_part = - index_cache.with_key_prefix(format!("part-{}", id).as_str()); - let token_set_format = format; - async move { - Result::Ok(Arc::new( - InvertedPartition::load( - store, - id, - frag_reuse_index_clone, - &index_cache_for_part, - token_set_format, - ) - .await?, - )) - } - }); - let partitions = stream::iter(partitions) - .buffer_unordered(store.io_parallelism()) - .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(), - )); - } - params.document_granularity = if coordinate_rank == 0 { - DocumentGranularity::Row - } else { - DocumentGranularity::ListElement - }; - - let tokenizer = params.build()?; - Ok(Arc::new(Self { - params, - store, - tokenizer, - token_set_format, - format_version, - partitions, - corpus_stats: Arc::new(OnceCell::new()), - prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), - document_projections_resident: Arc::new(AtomicBool::new(false)), - deleted_fragments, - })) - } - Err(_) => { - // old index format - Self::load_legacy_index(store, frag_reuse_index, index_cache).await - } - } - } -} - -#[async_trait] -impl Index for InvertedIndex { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn as_index(self: Arc) -> Arc { - self - } - - fn statistics(&self) -> Result { - let num_tokens = self - .partitions - .iter() - .map(|part| part.tokens.len()) - .sum::(); - let num_docs = self - .partitions - .iter() - .map(|part| part.docs.len()) - .sum::(); - Ok(serde_json::json!({ - "params": self.params, - "num_tokens": num_tokens, - "num_docs": num_docs, - })) - } - - async fn prewarm(&self) -> Result<()> { - self.prewarm_with_options(&FtsPrewarmOptions::default()) - .await - } - - fn index_type(&self) -> crate::IndexType { - crate::IndexType::Inverted - } - - async fn calculate_included_frags(&self) -> Result { - unimplemented!() - } -} - -/// Target on-disk size of one prewarm chunk. Keep this large enough that cloud -/// stores do not spend prewarm time on thousands of tiny range reads, but still -/// bounded so one large partition is not materialized all at once. -const PREWARM_CHUNK_TARGET_BYTES: u64 = 128 << 20; - -/// Cap on token rows per chunk, bounding the built `Vec` when posting lists are tiny. -const PREWARM_MAX_CHUNK_TOKENS: usize = 256 * 1024; - -/// Floor on token rows per chunk, so a partition always makes progress. -const PREWARM_MIN_CHUNK_TOKENS: usize = 1; - -/// Maximum number of posting lists in a runtime synthetic cache group. This is -/// deliberately token-count based so grouping works for old v2 indexes without -/// scanning posting lengths or requiring index rebuilds. -static LANCE_FTS_POSTING_GROUP_MAX_TOKENS: LazyLock = LazyLock::new(|| { - std::env::var("LANCE_FTS_POSTING_GROUP_MAX_TOKENS") - .unwrap_or_else(|_| "128".to_string()) - .parse() - .expect("failed to parse LANCE_FTS_POSTING_GROUP_MAX_TOKENS") -}); - -fn runtime_posting_group_tokens() -> usize { - (*LANCE_FTS_POSTING_GROUP_MAX_TOKENS).max(1) -} - -/// Runtime posting-list cache grouping. Non-empty v2 indexes synthesize fixed -/// groups at read time so prewarm and queries share group cache entries without -/// persisted grouping metadata or index rebuilds. -#[derive(Debug, Clone, DeepSizeOf)] -enum PostingGrouping { - /// Leaves legacy or empty partitions ungrouped. - None, - /// Uses a fixed runtime cache group size measured in token rows, not posting bytes. - SyntheticFixed { group_size: u32 }, -} - -impl PostingGrouping { - fn for_reader(is_legacy_layout: bool, token_count: usize) -> Self { - if is_legacy_layout || token_count == 0 { - return Self::None; - } - - let group_size = u32::try_from(runtime_posting_group_tokens()) - .unwrap_or(u32::MAX) - .max(1); - Self::SyntheticFixed { group_size } - } - - fn is_grouped(&self) -> bool { - !matches!(self, Self::None) - } - - fn range_for_token(&self, token_id: u32, token_count: usize) -> Option<(u32, u32)> { - match self { - Self::None => None, - Self::SyntheticFixed { group_size } => { - let token_count = u32::try_from(token_count).unwrap_or(u32::MAX); - let start = (token_id / *group_size) * *group_size; - let end = start.saturating_add(*group_size).min(token_count); - Some((start, end)) - } - } - } - - fn aligned_chunk_end(&self, token_count: usize, tok_start: usize, desired_end: usize) -> usize { - match self { - Self::None => desired_end, - Self::SyntheticFixed { group_size } => synthetic_group_aligned_chunk_end( - usize::try_from(*group_size).unwrap_or(usize::MAX).max(1), - token_count, - tok_start, - desired_end, - ), - } - } - - fn ranges_for_chunk( - &self, - tok_start: usize, - tok_end: usize, - token_count: usize, - ) -> Vec<(u32, u32)> { - match self { - Self::None => Vec::new(), - Self::SyntheticFixed { group_size } => synthetic_group_ranges_for_chunk( - usize::try_from(*group_size).unwrap_or(usize::MAX).max(1), - tok_start, - tok_end, - token_count, - ), - } - } -} - -/// Token rows per chunk: byte target / average bytes-per-token, clamped to `[MIN, MAX]`. -fn prewarm_chunk_tokens(token_count: usize, file_size_bytes: u64) -> usize { - if token_count == 0 { - return PREWARM_MIN_CHUNK_TOKENS; - } - let bytes_per_token = (file_size_bytes / token_count as u64).max(1); // >= 1: no div-by-zero - let by_bytes = (PREWARM_CHUNK_TARGET_BYTES / bytes_per_token) as usize; - by_bytes.clamp(PREWARM_MIN_CHUNK_TOKENS, PREWARM_MAX_CHUNK_TOKENS) -} - -fn synthetic_group_aligned_chunk_end( - group_size: usize, - token_count: usize, - tok_start: usize, - desired_end: usize, -) -> usize { - if desired_end >= token_count { - return token_count; - } - - let boundary = desired_end - (desired_end % group_size); - if boundary > tok_start { - boundary - } else { - tok_start.saturating_add(group_size).min(token_count) - } -} - -fn synthetic_group_ranges_for_chunk( - group_size: usize, - tok_start: usize, - tok_end: usize, - token_count: usize, -) -> Vec<(u32, u32)> { - let mut ranges = Vec::new(); - let mut start = tok_start - (tok_start % group_size); - if start < tok_start { - start = start.saturating_add(group_size).min(token_count); - } - while start < tok_end { - let end = start.saturating_add(group_size).min(token_count); - ranges.push(( - u32::try_from(start).unwrap_or(u32::MAX), - u32::try_from(end).unwrap_or(u32::MAX), - )); - start = end; - } - ranges -} - -fn prewarm_chunk_ranges( - grouping: &PostingGrouping, - token_count: usize, - chunk_tokens: usize, -) -> Vec<(usize, usize)> { - let mut ranges = Vec::new(); - let mut tok_start = 0usize; - while tok_start < token_count { - let mut tok_end = (tok_start + chunk_tokens).min(token_count); - // `tok_start` is always a group boundary; snap `tok_end` back to one too. - if grouping.is_grouped() { - tok_end = grouping.aligned_chunk_end(token_count, tok_start, tok_end); - } - ranges.push((tok_start, tok_end)); - tok_start = tok_end; - } - ranges -} - -impl InvertedIndex { - pub async fn prewarm_with_options(&self, options: &FtsPrewarmOptions) -> Result<()> { - let mut state = self.prewarm_state.lock().await; - if state.satisfies(options.with_position) - && (self.is_legacy() || self.document_projections_resident_now()) - { - return Ok(()); - } - let with_position = options.with_position || state.positions_ready; - state.query_ready = false; - state.positions_ready = false; - self.document_projections_resident - .store(false, Ordering::Release); - self.prewarm_query_state(with_position).await?; - state.query_ready = true; - state.positions_ready = with_position; - Ok(()) - } - - async fn prewarm_query_state(&self, with_position: bool) -> Result<()> { - let chunk_concurrency = self.store.io_parallelism().max(1); - let prewarm_started = Instant::now(); - info!( - partition_count = self.partitions.len(), - with_position, chunk_concurrency, "fts index prewarm started" - ); - for part in &self.partitions { - let partition_started = Instant::now(); - info!( - partition_id = part.id(), - token_count = part.tokens.len(), - with_position, - chunk_concurrency, - "fts partition prewarm started" - ); - if let Err(err) = part - .inverted_list - .prewarm_posting_lists(with_position, chunk_concurrency) - .await - { - warn!( - partition_id = part.id(), - error = %err, - elapsed_ms = partition_started.elapsed().as_millis() as u64, - "fts partition posting list prewarm failed" - ); - return Err(err); - } - info!( - partition_id = part.id(), - elapsed_ms = partition_started.elapsed().as_millis() as u64, - "fts partition posting lists prewarmed" - ); - let docs_started = Instant::now(); - if let Err(err) = part.docs.prewarm().await { - warn!( - partition_id = part.id(), - error = %err, - elapsed_ms = docs_started.elapsed().as_millis() as u64, - total_elapsed_ms = partition_started.elapsed().as_millis() as u64, - "fts partition docset prewarm failed" - ); - return Err(err); - } - info!( - partition_id = part.id(), - docset_elapsed_ms = docs_started.elapsed().as_millis() as u64, - elapsed_ms = partition_started.elapsed().as_millis() as u64, - "fts partition prewarm finished" - ); - } - self.aggregate_corpus_stats().await?; - let query_ready = self.partitions.iter().all(|partition| { - partition.docs.query_ready() - && partition.inverted_list.modern_posting_validation_ready() - }); - if !query_ready { - return Err(Error::internal( - "FTS prewarm completed without publishing a query-ready document and posting state" - .to_owned(), - )); - } - self.document_projections_resident - .store(true, Ordering::Release); - info!( - partition_count = self.partitions.len(), - query_ready, - elapsed_ms = prewarm_started.elapsed().as_millis() as u64, - "fts index prewarm finished" - ); - Ok(()) - } - /// Search docs match the input text. - async fn do_search(&self, text: &str) -> Result { - let params = FtsSearchParams::new(); - let mut tokenizer = self.tokenizer.clone(); - let tokens = collect_query_tokens(text, &mut tokenizer); - - let (doc_ids, _) = self - .bm25_search( - Arc::new(tokens), - params.into(), - Operator::And, - Arc::new(NoFilter), - Arc::new(NoOpMetricsCollector), - None, - ) - .boxed() - .await?; - - Ok(RecordBatch::try_new( - ROW_ID_SCHEMA.clone(), - vec![Arc::new(UInt64Array::from(doc_ids))], - )?) - } -} - -#[async_trait] -impl ScalarIndex for InvertedIndex { - // return the row ids of the documents that contain the query - #[instrument(level = "debug", skip_all)] - async fn search( - &self, - query: &dyn AnyQuery, - _metrics: &dyn MetricsCollector, - ) -> Result { - let query = query.as_any().downcast_ref::().unwrap(); - - match query { - TokenQuery::TokensContains(text) => { - let records = self.do_search(text).await?; - let row_ids = records - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let row_ids = row_ids.iter().flatten().collect_vec(); - Ok(SearchResult::at_most(RowAddrTreeMap::from_iter(row_ids))) - } - } - } - - fn can_remap(&self) -> bool { - true - } - - async fn remap( - &self, - mapping: &RowAddrRemap, - dest_store: &dyn IndexStore, - ) -> Result { - let files = self - .to_builder() - .remap(mapping, self.store.clone(), dest_store) - .await?; - - let details = pbold::InvertedIndexDetails::try_from(&self.params)?; - - Ok(CreatedIndex { - index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: self.index_version(), - files, - }) - } - - async fn update( - &self, - new_data: SendableRecordBatchStream, - dest_store: &dyn IndexStore, - old_data_filter: Option, - ) -> Result { - let files = self - .to_builder() - .update(new_data, dest_store, old_data_filter) - .await?; - - let details = pbold::InvertedIndexDetails::try_from(&self.params)?; - - Ok(CreatedIndex { - index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: self.index_version(), - files, - }) - } - - fn update_criteria(&self) -> UpdateCriteria { - let criteria = TrainingCriteria::new(TrainingOrdering::None).with_row_id(); - if self.is_legacy() { - UpdateCriteria::requires_old_data(criteria) - } else { - UpdateCriteria::only_new_data(criteria) - } - } - - fn derive_index_params(&self) -> Result { - let mut params = self.params.clone(); - if params.base_tokenizer.is_empty() { - // Empty tokenizer metadata only appears in legacy simple-tokenizer indexes. - params.base_tokenizer = "simple".to_string(); - } - params = params.format_version(self.format_version()); - - let params_json = params.to_training_json()?.to_string(); - - Ok(ScalarIndexParams { - index_type: BuiltinIndexType::Inverted.as_str().to_string(), - params: Some(params_json), - }) - } -} - -#[derive(Debug, Clone, DeepSizeOf)] -pub struct InvertedPartition { - // 0 for legacy format - id: u64, - store: Arc, - pub(crate) tokens: TokenSet, - pub(crate) inverted_list: Arc, - /// Legacy documents stay in their original complete `DocSet`; modern - /// documents use typed, independently-loaded lengths and addresses. - pub(super) docs: PartitionDocumentStore, - token_set_format: TokenSetFormat, -} - -impl InvertedPartition { - /// Check if this partition belongs to the specified fragment. - /// - /// This method encapsulates the bit manipulation logic for fragment filtering - /// in distributed indexing scenarios. - /// - /// # Arguments - /// * `fragment_mask` - A mask with fragment_id in high 32 bits - /// - /// # Returns - /// * `true` if the partition belongs to the fragment, `false` otherwise - pub fn belongs_to_fragment(&self, fragment_mask: u64) -> bool { - (self.id() & fragment_mask) == fragment_mask - } - - pub fn id(&self) -> u64 { - self.id - } - - pub fn store(&self) -> &dyn IndexStore { - self.store.as_ref() - } - - pub fn is_legacy(&self) -> bool { - self.inverted_list.is_legacy_layout() - } - - pub async fn load( - store: Arc, - id: u64, - frag_reuse_index: Option>, - index_cache: &LanceCache, - token_set_format: TokenSetFormat, - ) -> Result { - let token_file = store.open_index_file(&token_file_path(id)).await?; - let tokens = TokenSet::load(token_file, token_set_format).await?; - let invert_list_file = store.open_index_file(&posting_file_path(id)).await?; - let mut inverted_list = PostingListReader::try_new(invert_list_file, index_cache).await?; - let docs_path = doc_file_path(id); - let docs_reader = store.open_index_file(&docs_path).await?; - let docs = PartitionDocuments::try_new( - store.clone(), - docs_path, - id, - WeakLanceCache::from(index_cache), - docs_reader.as_ref(), - frag_reuse_index, - // 256-document blocks score with quantized document lengths. - inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, - )?; - inverted_list.modern_num_docs = Some(docs.len()); - - Ok(Self { - id, - store, - tokens, - inverted_list: Arc::new(inverted_list), - docs: PartitionDocumentStore::Modern(Arc::new(docs)), - token_set_format, - }) - } - - fn map(&self, token: &str) -> Option { - self.tokens.get(token) - } - - pub fn expand_fuzzy(&self, tokens: &Tokens, params: &FtsSearchParams) -> Result { - let mut new_tokens = Vec::with_capacity(min(tokens.len(), params.max_expansions)); - let mut new_positions = Vec::with_capacity(new_tokens.capacity()); - let mut seen = HashSet::new(); - for token_idx in 0..tokens.len() { - let remaining = params.max_expansions.saturating_sub(new_tokens.len()); - if remaining == 0 { - break; - } - let token = tokens.get_token(token_idx); - let position = tokens.position(token_idx); - let base_prefix_len = tokens.token_type().prefix_len(token) as u32; - let mut candidates = BTreeSet::new(); - self.collect_fuzzy_candidates( - token, - base_prefix_len, - params, - remaining, - &mut candidates, - )?; - for candidate in candidates { - if new_tokens.len() >= params.max_expansions { - break; - } - if seen.insert((candidate.clone(), position)) { - new_tokens.push(candidate); - new_positions.push(position); - } - } - } - Ok(Tokens::with_positions( - new_tokens, - new_positions, - tokens.token_type().clone(), - )) - } - - /// Collect up to `limit` fuzzy candidates for one query token from this - /// partition's token FST, in key (lexicographic) order. Callers merge - /// candidates across partitions and apply the query-wide - /// `max_expansions` budget; truncating each partition at `limit` is - /// lossless for that selection because any term among the merged - /// lexicographically-smallest `limit` is also among its own partition's - /// smallest `limit`. - fn collect_fuzzy_candidates( - &self, - token: &str, - base_prefix_len: u32, - params: &FtsSearchParams, - limit: usize, - candidates: &mut BTreeSet, - ) -> Result<()> { - let fuzziness = match params.fuzziness { - Some(fuzziness) => fuzziness, - None => MatchQuery::auto_fuzziness(token), - }; - let lev = fst::automaton::Levenshtein::new(token, fuzziness) - .map_err(|e| Error::index(format!("failed to construct the fuzzy query: {}", e)))?; - - if let TokenMap::Fst(ref map) = self.tokens.tokens { - let mut expanded = Vec::new(); - match base_prefix_len + params.prefix_length { - 0 => take_fst_keys(map.search(lev), &mut expanded, limit), - prefix_length => { - let prefix = &token[..min(prefix_length as usize, token.len())]; - let prefix = fst::automaton::Str::new(prefix).starts_with(); - take_fst_keys(map.search(lev.intersection(prefix)), &mut expanded, limit) - } - } - candidates.extend(expanded); - Ok(()) - } else { - Err(Error::index( - "tokens is not fst, which is not expected".to_owned(), - )) - } - } - - #[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: &LoadedDocLengths, - 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: &LoadedDocLengths, - 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() { - let entry = freqs_by_row_id.entry(row_id).or_insert(0u32); - *entry = entry.checked_add(freq).ok_or_else(|| { - Error::index(format!("posting frequency overflow for row id {}", row_id)) - })?; - } - } - 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), - Some(max_score), - None, - ))) - } - - fn union_plain_posting_lists_with_positions( - postings: Vec, - docs: &LoadedDocLengths, - 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: &LoadedDocLengths, - 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 freqs_by_doc_id = BTreeMap::new(); - for posting in postings { - for (doc_id, freq, _) 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 entry = freqs_by_doc_id.entry(doc_id).or_insert(0u32); - *entry = entry.checked_add(freq).ok_or_else(|| { - Error::index(format!("posting frequency overflow for doc id {}", doc_id)) - })?; - } - } - if freqs_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(false, block_size); - let mut doc_ids = Vec::with_capacity(freqs_by_doc_id.len()); - let mut frequencies = Vec::with_capacity(freqs_by_doc_id.len()); - for (doc_id, freq) in freqs_by_doc_id { - builder.add(doc_id, PositionRecorder::Count(freq)); - doc_ids.push(doc_id); - frequencies.push(freq); - } - 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: &LoadedDocLengths, - 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); - let length = batch[LENGTH_COL].as_primitive::().value(0); - PostingList::from_batch(&batch, Some(max_score), Some(length)) - } - - fn union_posting_lists( - postings: Vec, - docs: &LoadedDocLengths, - with_positions: bool, - query_weight: f32, - scorer: &MemBM25Scorer, - ) -> Result { - let has_plain = postings - .iter() - .any(|posting| matches!(posting, PostingList::Plain(_))); - let has_compressed = postings - .iter() - .any(|posting| matches!(posting, PostingList::Compressed(_))); - match (has_plain, has_compressed) { - (true, true) => Err(Error::index( - "cannot union mixed plain and compressed posting lists".to_owned(), - )), - (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()), - None, - None, - ))), - } - } - - // search the documents that contain the query - // return the doc info and the doc length - // ref: https://en.wikipedia.org/wiki/Okapi_BM25 - // - // `force_global_scorer` is used by compound search, where leaf scores and - // bounds must share corpus-level statistics before the global collector - // can safely propagate its threshold. Old posting formats without impacts - // fall back to a scorer-derived global upper bound in that mode. - #[instrument(level = "debug", skip_all)] - pub(super) async fn load_posting_lists( - &self, - tokens: &Tokens, - params: &FtsSearchParams, - operator: Operator, - impact_scorer: &MemBM25Scorer, - metrics: &dyn MetricsCollector, - force_global_scorer: bool, - ) -> Result { - 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(|| { - (0..tokens.len()) - .map(|index| tokens.position(index)) - .collect::>() - }); - // Fuzzy expansion already ran once at the index level (see - // `InvertedIndex::bm25_search`) under the global `max_expansions` - // 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() { - let token_id = self.map(&token); - if let Some(token_id) = token_id { - let position = token_positions[index]; - if let Some(matched_positions) = matched_positions.as_mut() { - matched_positions.insert(position); - } - token_ids.push((token_id, token, position)); - } - } - if token_ids.is_empty() { - return Ok(LoadedPostings::empty()); - } - if let Some(required_positions) = required_positions.as_ref() - && let Some(matched_positions) = matched_positions.as_ref() - && !required_positions.is_subset(matched_positions) - { - return Ok(LoadedPostings::empty()); - } - - 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) - .map(|(token_id, token, position)| async move { - let posting = self - .inverted_list - .posting_list(token_id, is_phrase_query, metrics) - .await?; - - Result::Ok((token_id, token, position, posting)) - }) - .buffered(self.store.io_parallelism()) - .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) - && !needs_union - && loaded_postings - .iter() - .any(|(_, _, _, posting)| posting.is_empty()) - { - return Ok(LoadedPostings::empty()); - } - - if !needs_union { - let impact_safe = loaded_postings - .iter() - .all(|(_, _, _, posting)| posting.has_impacts()); - return Ok(LoadedPostings { - postings: loaded_postings - .into_iter() - .map(|(token_id, token, position, posting)| { - let needs_scorer_upper_bound = (exact_scoring_required - || force_global_scorer) - && !posting.has_impacts(); - let query_weight = - if impact_safe || exact_scoring_required || force_global_scorer { - impact_scorer.query_weight(&token) - } else { - idf(posting.len(), num_docs) - }; - 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 docs_for_union = if needs_union { - Some(match &self.docs { - PartitionDocumentStore::Legacy(docs) => LoadedDocLengths::Legacy(docs.clone()), - PartitionDocumentStore::Modern(documents) => { - LoadedDocLengths::Modern(documents.lengths().await?) - } - }) - } else { - None - }; - - // WAND's AND mode treats every iterator as required, so expansions from - // one original query position must be merged before scoring. - let mut grouped_postings = Vec::new(); - let mut grouped_expansions = Vec::new(); - let mut iter = loaded_postings.into_iter().peekable(); - while let Some((token_id, token, position, posting)) = iter.next() { - let mut group = vec![(token_id, token, posting)]; - while matches!(iter.peek(), Some((_, _, next_position, _)) if *next_position == position) - { - let (token_id, token, _, posting) = iter.next().expect("peeked item must exist"); - group.push((token_id, token, posting)); - } - - let (token_id, token, posting) = if group.len() == 1 { - group.pop().expect("single-item group must exist") - } 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: terms.clone(), - }); - let postings = group - .into_iter() - .map(|(_, _, posting)| posting) - .collect::>(); - let docs = docs_for_union.as_ref().ok_or_else(|| { - Error::index("union docs were not loaded for grouped query terms".to_string()) - })?; - let posting = Self::union_posting_lists( - postings, - docs, - is_phrase_query, - query_weight, - impact_scorer, - )?; - 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() { - if is_and_query || is_phrase_query { - return Ok(LoadedPostings::empty()); - } - continue; - } - - 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, - }) - } - - #[allow(clippy::too_many_arguments)] - fn bm25_search_legacy( - &self, - docs: &DocSet, - params: &FtsSearchParams, - operator: Operator, - mask: &RowAddrMask, - postings: Vec, - impact_scorer: Option>, - metrics: &dyn MetricsCollector, - shared_threshold: Arc, - ) -> Result>> { - let documents = LegacyWandDocuments::new(docs, mask); - self.bm25_search_with_documents( - &documents, - params, - operator, - postings, - impact_scorer, - metrics, - shared_threshold, - ) - } - - #[allow(clippy::too_many_arguments)] - fn bm25_search_modern( - &self, - lengths: &DocLengths, - visibility: &DocVisibility, - params: &FtsSearchParams, - operator: Operator, - postings: Vec, - impact_scorer: Option>, - metrics: &dyn MetricsCollector, - shared_threshold: Arc, - ) -> Result>> { - if visibility.is_all() { - let documents = ModernWandDocuments::all(lengths); - self.bm25_search_with_documents( - &documents, - params, - operator, - postings, - impact_scorer, - metrics, - shared_threshold, - ) - } else { - let documents = ModernWandDocuments::filtered(lengths, visibility); - self.bm25_search_with_documents( - &documents, - params, - operator, - postings, - impact_scorer, - metrics, - shared_threshold, - ) - } - } - - #[instrument(level = "debug", skip_all)] - #[allow(clippy::too_many_arguments)] - fn bm25_search_with_documents( - &self, - documents: &D, - params: &FtsSearchParams, - operator: Operator, - postings: Vec, - impact_scorer: Option>, - metrics: &dyn MetricsCollector, - shared_threshold: Arc, - ) -> Result>> { - if postings.is_empty() { - return Ok(Vec::new()); - } - - let hits = if let Some(scorer) = impact_scorer { - let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) - .with_shared_threshold(shared_threshold); - wand.search(params, metrics)? - } else { - let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) - .with_shared_threshold(shared_threshold); - wand.search(params, metrics)? - }; - Ok(hits) - } - - pub async fn into_builder(self) -> Result { - let mut builder = InnerBuilder::new_with_posting_tail_codec_and_block_size( - self.id, - self.inverted_list.has_positions(), - self.token_set_format, - self.inverted_list.posting_tail_codec(), - self.inverted_list.block_size(), - ); - builder.tokens = self.tokens.into_mutable(); - builder.docs = self.docs.load_build_docset().await?; - - builder - .posting_lists - .reserve_exact(self.inverted_list.len()); - for posting_list in self - .inverted_list - .read_all(self.inverted_list.has_positions()) - .await? - { - let posting_list = posting_list?; - builder - .posting_lists - .push(posting_list.into_builder(&builder.docs)); - } - Ok(builder) - } -} - -// at indexing, we use HashMap because we need it to be mutable, -// at searching, we use fst::Map because it's more efficient -#[derive(Debug, Clone)] -pub enum TokenMap { - HashMap(HashMap), - Fst(fst::Map>), -} - -impl Default for TokenMap { - fn default() -> Self { - Self::HashMap(HashMap::new()) - } -} - -impl DeepSizeOf for TokenMap { - fn deep_size_of_children(&self, ctx: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::HashMap(map) => map.deep_size_of_children(ctx), - Self::Fst(map) => map.as_fst().size(), - } - } -} - -impl TokenMap { - pub fn len(&self) -> usize { - match self { - Self::HashMap(map) => map.len(), - Self::Fst(map) => map.len(), - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -// TokenSet is a mapping from tokens to token ids -#[derive(Debug, Clone, Default, DeepSizeOf)] -pub struct TokenSet { - // token -> token_id - pub(crate) tokens: TokenMap, - pub(crate) next_id: u32, - total_length: usize, -} - -impl TokenSet { - pub fn into_mut(self) -> Self { - let tokens = match self.tokens { - TokenMap::HashMap(map) => map, - TokenMap::Fst(map) => { - let mut new_map = HashMap::with_capacity(map.len()); - let mut stream = map.into_stream(); - while let Some((token, token_id)) = stream.next() { - new_map.insert(String::from_utf8_lossy(token).into_owned(), token_id as u32); - } - - new_map - } - }; - - Self { - tokens: TokenMap::HashMap(tokens), - next_id: self.next_id, - total_length: self.total_length, - } - } - - pub fn len(&self) -> usize { - self.tokens.len() - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn to_batch(self, format: TokenSetFormat) -> Result { - match format { - TokenSetFormat::Arrow => self.into_arrow_batch(), - TokenSetFormat::Fst => self.into_fst_batch(), - } - } - - fn into_arrow_batch(self) -> Result { - let mut token_builder = StringBuilder::with_capacity(self.tokens.len(), self.total_length); - let mut token_id_builder = UInt32Builder::with_capacity(self.tokens.len()); - - match self.tokens { - TokenMap::Fst(map) => { - let mut stream = map.stream(); - while let Some((token, token_id)) = stream.next() { - token_builder.append_value(String::from_utf8_lossy(token)); - token_id_builder.append_value(token_id as u32); - } - } - TokenMap::HashMap(map) => { - for (token, token_id) in map.into_iter().sorted_unstable() { - token_builder.append_value(token); - token_id_builder.append_value(token_id); - } - } - } - - let token_col = token_builder.finish(); - let token_id_col = token_id_builder.finish(); - - let schema = arrow_schema::Schema::new(vec![ - arrow_schema::Field::new(TOKEN_COL, DataType::Utf8, false), - arrow_schema::Field::new(TOKEN_ID_COL, DataType::UInt32, false), - ]); - - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(token_col) as ArrayRef, - Arc::new(token_id_col) as ArrayRef, - ], - )?; - Ok(batch) - } - - fn into_fst_batch(mut self) -> Result { - let fst_map = match std::mem::take(&mut self.tokens) { - TokenMap::Fst(map) => map, - TokenMap::HashMap(map) => Self::build_fst_from_map(map)?, - }; - let bytes = fst_map.into_fst().into_inner(); - - let mut fst_builder = LargeBinaryBuilder::with_capacity(1, bytes.len()); - fst_builder.append_value(bytes); - let fst_col = fst_builder.finish(); - - let mut next_id_builder = UInt32Builder::with_capacity(1); - next_id_builder.append_value(self.next_id); - let next_id_col = next_id_builder.finish(); - - let mut total_length_builder = UInt64Builder::with_capacity(1); - total_length_builder.append_value(self.total_length as u64); - let total_length_col = total_length_builder.finish(); - - let schema = arrow_schema::Schema::new(vec![ - arrow_schema::Field::new(TOKEN_FST_BYTES_COL, DataType::LargeBinary, false), - arrow_schema::Field::new(TOKEN_NEXT_ID_COL, DataType::UInt32, false), - arrow_schema::Field::new(TOKEN_TOTAL_LENGTH_COL, DataType::UInt64, false), - ]); - - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(fst_col) as ArrayRef, - Arc::new(next_id_col) as ArrayRef, - Arc::new(total_length_col) as ArrayRef, - ], - )?; - Ok(batch) - } - - fn build_fst_from_map(map: HashMap) -> Result>> { - let mut entries: Vec<_> = map.into_iter().collect(); - entries.sort_unstable_by(|(lhs, _), (rhs, _)| lhs.cmp(rhs)); - let mut builder = fst::MapBuilder::memory(); - for (token, token_id) in entries { - builder - .insert(&token, token_id as u64) - .map_err(|e| Error::index(format!("failed to insert token {}: {}", token, e)))?; - } - Ok(builder.into_map()) - } - - pub async fn load(reader: Arc, format: TokenSetFormat) -> Result { - match format { - TokenSetFormat::Arrow => Self::load_arrow(reader).await, - TokenSetFormat::Fst => Self::load_fst(reader).await, - } - } - - async fn load_arrow(reader: Arc) -> Result { - let batch = reader.read_range(0..reader.num_rows(), None).await?; - - let (tokens, next_id, total_length) = spawn_blocking(move || { - let mut next_id = 0; - let mut total_length = 0; - let mut tokens = fst::MapBuilder::memory(); - - let token_col = batch[TOKEN_COL].as_string::(); - let token_id_col = batch[TOKEN_ID_COL].as_primitive::(); - - for (token, &token_id) in token_col.iter().zip(token_id_col.values().iter()) { - let token = - token.ok_or(Error::index("found null token in token set".to_owned()))?; - next_id = next_id.max(token_id + 1); - total_length += token.len(); - tokens.insert(token, token_id as u64).map_err(|e| { - Error::index(format!("failed to insert token {}: {}", token, e)) - })?; - } - - Ok::<_, Error>((tokens.into_map(), next_id, total_length)) - }) - .await - .map_err(|err| Error::execution(format!("failed to spawn blocking task: {}", err)))??; - - Ok(Self { - tokens: TokenMap::Fst(tokens), - next_id, - total_length, - }) - } - - async fn load_fst(reader: Arc) -> Result { - let batch = reader.read_range(0..reader.num_rows(), None).await?; - if batch.num_rows() == 0 { - return Err(Error::index("token set batch is empty".to_owned())); - } - - let fst_col = batch[TOKEN_FST_BYTES_COL].as_binary::(); - let bytes = fst_col.value(0); - let map = fst::Map::new(bytes.to_vec()) - .map_err(|e| Error::index(format!("failed to load fst tokens: {}", e)))?; - - let total_length_col = - batch[TOKEN_TOTAL_LENGTH_COL].as_primitive::(); - - // Token ids are dense `[0, len)`, so `next_id` must equal the token count. Recompute - // it instead of trusting the persisted value, which writers before #7115 could leave - // stale. Mirrors `load_arrow`. - let next_id = map.len() as u32; - - let total_length = total_length_col - .values() - .first() - .copied() - .ok_or(Error::index( - "token total length column is empty".to_owned(), - ))?; - - Ok(Self { - tokens: TokenMap::Fst(map), - next_id, - total_length: usize::try_from(total_length).map_err(|_| { - Error::index(format!( - "token total length {} overflows usize", - total_length - )) - })?, - }) - } - - pub fn add(&mut self, token: String) -> u32 { - let next_id = self.next_id(); - let len = token.len(); - let token_id = match self.tokens { - TokenMap::HashMap(ref mut map) => *map.entry(token).or_insert(next_id), - _ => unreachable!("tokens must be HashMap while indexing"), - }; - - // add token if it doesn't exist - if token_id == next_id { - self.next_id += 1; - self.total_length += len; - } - - token_id - } - - pub(crate) fn get_or_add(&mut self, token: &str) -> u32 { - let next_id = self.next_id; - match self.tokens { - TokenMap::HashMap(ref mut map) => { - if let Some(&token_id) = map.get(token) { - return token_id; - } - - map.insert(token.to_owned(), next_id); - } - _ => unreachable!("tokens must be HashMap while indexing"), - } - - self.next_id += 1; - self.total_length += token.len(); - next_id - } - - pub(crate) fn into_mutable(self) -> Self { - let Self { - tokens, - next_id, - total_length, - } = self; - match tokens { - TokenMap::HashMap(_) => Self { - tokens, - next_id, - total_length, - }, - TokenMap::Fst(map) => { - let mut mutable = HashMap::new(); - let mut stream = map.stream(); - while let Some((token, token_id)) = stream.next() { - mutable.insert(String::from_utf8_lossy(token).into_owned(), token_id as u32); - } - Self { - tokens: TokenMap::HashMap(mutable), - next_id, - total_length, - } - } - } - } - - pub fn get(&self, token: &str) -> Option { - match self.tokens { - TokenMap::HashMap(ref map) => map.get(token).copied(), - TokenMap::Fst(ref map) => map.get(token).map(|id| id as u32), - } - } - - // the `removed_token_ids` must be sorted - pub fn remap(&mut self, removed_token_ids: &[u32]) { - if removed_token_ids.is_empty() { - return; - } - - let mut map = match std::mem::take(&mut self.tokens) { - TokenMap::HashMap(map) => map, - TokenMap::Fst(map) => { - let mut new_map = HashMap::with_capacity(map.len()); - let mut stream = map.into_stream(); - while let Some((token, token_id)) = stream.next() { - new_map.insert(String::from_utf8_lossy(token).into_owned(), token_id as u32); - } - - new_map - } - }; - - let mut retained_length = 0; - map.retain( - |token, token_id| match removed_token_ids.binary_search(token_id) { - Ok(_) => false, - Err(index) => { - *token_id -= index as u32; - retained_length += token.len(); - true - } - }, - ); - - self.tokens = TokenMap::HashMap(map); - - // The retain above compacts the surviving token ids into a dense `[0, len)` - // range, so `next_id` (handed to the next new token) must follow them down. - // `total_length` likewise must drop the removed tokens' bytes; it is persisted - // and feeds memory accounting, so a stale value drifts across remap/merge cycles. - self.next_id = self.tokens.len() as u32; - self.total_length = retained_length; - } - - pub fn next_id(&self) -> u32 { - self.next_id - } - - pub(crate) fn memory_size(&self) -> usize { - match &self.tokens { - TokenMap::HashMap(map) => { - self.total_length - + map.capacity() - * (std::mem::size_of::() - + std::mem::size_of::() - + std::mem::size_of::()) - } - TokenMap::Fst(map) => map.as_fst().size(), - } - } -} - -pub struct PostingListReader { - reader: Arc, - - /// Layout-specific metadata. V2 keeps its per-token max-score and - /// length columns lazy so opening a partition doesn't drag O(num_tokens) - /// bytes off cold storage when the caller only needs `df` for a few terms. - metadata: PostingMetadata, - - has_position: bool, - has_impacts: bool, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions_layout: PositionsLayout, - - /// Runtime posting-list cache grouping. Non-empty v2 indexes use synthetic - /// fixed groups so prewarm can improve cache density without rebuilding the - /// index or relying on persisted grouping metadata. - grouping: PostingGrouping, - - /// Modern postings contain dense DocIds into the partition document table. - /// Cache successful boundary validation per immutable token so repeated - /// queries do not decode the final posting block again. - modern_doc_id_validations: Option]>>, - /// Skips per-token readiness checks once the whole immutable table is validated. - modern_postings_validated: AtomicBool, - modern_num_docs: Option, - - index_cache: WeakLanceCache, -} - -/// Per-token metadata (max_score, length) needed by the BM25 query and stats -/// paths. The legacy and v2 formats store this metadata in different -/// places, with very different cost profiles for cold-load: the variants -/// surface that asymmetry so callers can choose a per-token or bulk access -/// pattern. -enum PostingMetadata { - /// Legacy v1: offsets and max_scores are encoded in the file's schema - /// metadata, so they are already in memory by the time `try_new` returns. - LegacyV1 { - offsets: Vec, - max_scores: Option>, - }, - /// V2: per-token `max_score` and `length` live as columns in the - /// posting file. The bulk vectors are filled lazily by - /// `ensure_metadata_loaded`, and the stats path can also fetch a single - /// token via `posting_len_for_token` without forcing the bulk load. - V2 { - metadata: OnceCell, - }, -} - -#[derive(Debug, Clone)] -struct LoadedPostingMetadata { - max_scores: Vec, - lengths: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum PositionsLayout { - None, - LegacyPerDoc, - SharedStream(PositionStreamCodec), -} - -impl std::fmt::Debug for PostingListReader { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut s = f.debug_struct("InvertedListReader"); - match &self.metadata { - PostingMetadata::LegacyV1 { - offsets, - max_scores, - } => { - s.field("layout", &"legacy_v1") - .field("offsets", offsets) - .field("max_scores", max_scores); - } - PostingMetadata::V2 { metadata } => { - s.field("layout", &"v2") - .field("metadata_loaded", &metadata.initialized()); - } - } - s.finish() - } -} - -impl DeepSizeOf for PostingListReader { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - let metadata_size = match &self.metadata { - PostingMetadata::LegacyV1 { - offsets, - max_scores, - } => offsets.deep_size_of_children(context) + max_scores.deep_size_of_children(context), - PostingMetadata::V2 { metadata } => metadata - .get() - .map(|loaded| { - loaded.max_scores.deep_size_of_children(context) - + loaded.lengths.deep_size_of_children(context) - }) - .unwrap_or(0), - }; - let validation_size = self - .modern_doc_id_validations - .as_ref() - .map(|validations| { - validations - .len() - .saturating_mul(std::mem::size_of::>()) - }) - .unwrap_or(0); - metadata_size + self.grouping.deep_size_of_children(context) + validation_size - } -} - -impl PostingListReader { - pub(crate) async fn try_new( - reader: Arc, - index_cache: &LanceCache, - ) -> Result { - let positions_layout = if reader.schema().field(COMPRESSED_POSITION_COL).is_some() { - PositionsLayout::SharedStream(parse_shared_position_codec(&reader.schema().metadata)?) - } else if reader.schema().field(POSITION_COL).is_some() { - PositionsLayout::LegacyPerDoc - } else { - PositionsLayout::None - }; - let posting_tail_codec = parse_posting_tail_codec(&reader.schema().metadata)?; - let block_size = parse_posting_block_size(&reader.schema().metadata)?; - let has_position = positions_layout != PositionsLayout::None; - let has_impacts = reader.schema().field(IMPACT_COL).is_some(); - let metadata = if reader.schema().field(POSTING_COL).is_none() { - let (offsets, max_scores) = Self::load_metadata(reader.schema())?; - PostingMetadata::LegacyV1 { - offsets, - max_scores, - } - } else { - PostingMetadata::V2 { - metadata: OnceCell::new(), - } - }; - - let is_legacy_layout = matches!(&metadata, PostingMetadata::LegacyV1 { .. }); - let grouping = PostingGrouping::for_reader(is_legacy_layout, reader.num_rows()); - let modern_doc_id_validations = (!is_legacy_layout).then(|| { - (0..reader.num_rows()) - .map(|_| OnceCell::new()) - .collect::>() - .into() - }); - - Ok(Self { - reader, - metadata, - has_position, - has_impacts, - posting_tail_codec, - block_size, - positions_layout, - grouping, - modern_doc_id_validations, - modern_postings_validated: AtomicBool::new(false), - modern_num_docs: None, - index_cache: WeakLanceCache::from(index_cache), - }) - } - - // for legacy format - // returns the offsets and max scores - fn load_metadata( - schema: &lance_core::datatypes::Schema, - ) -> Result<(Vec, Option>)> { - let offsets = schema - .metadata - .get("offsets") - .ok_or(Error::index("offsets not found in metadata".to_owned()))?; - let offsets = serde_json::from_str(offsets)?; - - let max_scores = schema - .metadata - .get("max_scores") - .map(|max_scores| serde_json::from_str(max_scores)) - .transpose()?; - Ok((offsets, max_scores)) - } - - // the number of posting lists - pub fn len(&self) -> usize { - match &self.metadata { - PostingMetadata::LegacyV1 { offsets, .. } => offsets.len(), - PostingMetadata::V2 { .. } => self.reader.num_rows(), - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub(crate) fn has_positions(&self) -> bool { - self.has_position - } - - pub(crate) fn posting_tail_codec(&self) -> PostingTailCodec { - self.posting_tail_codec - } - - pub(crate) fn block_size(&self) -> usize { - self.block_size - } - - fn is_legacy_layout(&self) -> bool { - matches!(self.metadata, PostingMetadata::LegacyV1 { .. }) - } - - /// Sync access to `posting_len`. Requires v2 metadata to already be - /// loaded via [`ensure_metadata_loaded`]; the bm25 scoring path enforces - /// that contract before kicking off wand. The stats path uses - /// [`Self::posting_len_for_token`] instead, which avoids the bulk load. - pub(crate) fn posting_len(&self, token_id: u32) -> usize { - let token_id = token_id as usize; - match &self.metadata { - PostingMetadata::LegacyV1 { offsets, .. } => { - let next_offset = offsets - .get(token_id + 1) - .copied() - .unwrap_or(self.reader.num_rows()); - next_offset - offsets[token_id] - } - PostingMetadata::V2 { metadata } => { - let metadata = metadata - .get() - .expect("v2 posting metadata must be bulk-loaded before sync posting_len; call ensure_metadata_loaded first"); - metadata.lengths[token_id] as usize - } - } - } - - /// Async access to a single token's posting list length. For v2 - /// indexes this reads one row of posting metadata if the bulk metadata has - /// not been loaded yet, and never triggers the bulk load itself. The stats - /// path uses this so a single-term `df` lookup costs O(1) bytes rather - /// than O(num_unique_tokens). - pub(crate) async fn posting_len_for_token( - &self, - token_id: u32, - metrics: Option<&dyn MetricsCollector>, - ) -> Result { - match &self.metadata { - PostingMetadata::LegacyV1 { .. } => Ok(self.posting_len(token_id)), - PostingMetadata::V2 { metadata } => { - if let Some(metadata) = metadata.get() { - return Ok(metadata.lengths[token_id as usize] as usize); - } - let (_, length) = self.posting_metadata_for_token(token_id, metrics).await?; - length - .map(|len| len as usize) - .ok_or_else(|| Error::index("posting length metadata missing".to_string())) - } - } - } - - /// Async access to a single token's `(max_score, length)` pair. Mirrors - /// [`Self::posting_len_for_token`] but covers both columns the scoring - /// path needs, in one read. For v2 indexes that have not been - /// bulk-loaded this issues one `read_range(token..token+1, [MAX_SCORE, - /// LENGTH])`; for legacy v1 the values come from in-memory schema - /// metadata. - pub(crate) async fn posting_metadata_for_token( - &self, - token_id: u32, - metrics: Option<&dyn MetricsCollector>, - ) -> Result<(Option, Option)> { - match &self.metadata { - PostingMetadata::LegacyV1 { max_scores, .. } => { - Ok((max_scores.as_ref().map(|m| m[token_id as usize]), None)) - } - PostingMetadata::V2 { metadata } => { - if let Some(loaded) = metadata.get() { - return Ok(( - Some(loaded.max_scores[token_id as usize]), - Some(loaded.lengths[token_id as usize]), - )); - } - let result = self - .index_cache - .get_or_insert_with_key_hit(PostingMetadataKey { token_id }, || async move { - let token_id = token_id as usize; - let batch = self - .reader - .read_range(token_id..token_id + 1, Some(&[MAX_SCORE_COL, LENGTH_COL])) - .await?; - let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); - let length = batch[LENGTH_COL].as_primitive::().value(0); - Ok(PostingMetadataValue { max_score, length }) - }) - .await; - if let Some(metrics) = metrics { - match &result { - Ok((_, true)) => metrics.record_index_cache_hit(), - _ => metrics.record_index_cache_miss(), - } - } - let metadata = result.map(|(value, _)| value)?; - Ok((Some(metadata.max_score), Some(metadata.length))) - } - } - } - - /// Force the v2 bulk metadata (`max_scores`, `lengths`) into - /// memory. Cheap to call repeatedly; no-op for legacy v1 indexes whose - /// metadata is already populated from schema metadata at `try_new` time. - pub(crate) async fn ensure_metadata_loaded(&self) -> Result<()> { - let PostingMetadata::V2 { metadata } = &self.metadata else { - return Ok(()); - }; - metadata - .get_or_try_init(|| async { - let batch = self - .reader - .read_range( - 0..self.reader.num_rows(), - Some(&[MAX_SCORE_COL, LENGTH_COL]), - ) - .await?; - let max_scores = batch[MAX_SCORE_COL] - .as_primitive::() - .values() - .to_vec(); - let lengths = batch[LENGTH_COL] - .as_primitive::() - .values() - .to_vec(); - Ok::(LoadedPostingMetadata { - max_scores, - lengths, - }) - }) - .await?; - Ok(()) - } - - pub(crate) async fn posting_batch( - &self, - token_id: u32, - with_position: bool, - ) -> Result { - if self.is_legacy_layout() { - self.posting_batch_legacy(token_id, with_position).await - } else { - let token_id = token_id as usize; - let mut columns = if with_position { - match self.positions_layout { - PositionsLayout::SharedStream(_) => { - vec![ - POSTING_COL, - COMPRESSED_POSITION_COL, - POSITION_BLOCK_OFFSET_COL, - ] - } - PositionsLayout::LegacyPerDoc => vec![POSTING_COL, POSITION_COL], - PositionsLayout::None => vec![POSTING_COL], - } - } else { - vec![POSTING_COL] - }; - if self.has_impacts { - columns.push(IMPACT_COL); - } - let batch = self - .reader - .read_range(token_id..token_id + 1, Some(&columns)) - .await?; - Ok(batch) - } - } - - async fn posting_batch_legacy( - &self, - token_id: u32, - with_position: bool, - ) -> Result { - let mut columns = vec![ROW_ID, FREQUENCY_COL]; - if with_position { - columns.push(POSITION_COL); - } - - let length = self.posting_len(token_id); - let PostingMetadata::LegacyV1 { offsets, .. } = &self.metadata else { - unreachable!("posting_batch_legacy is only reachable on legacy v1 layout"); - }; - let token_id = token_id as usize; - let offset = offsets[token_id]; - let batch = self - .reader - .read_range(offset..offset + length, Some(&columns)) - .await?; - Ok(batch) - } - - #[instrument(level = "debug", skip(self, metrics))] - pub(crate) async fn posting_list( - &self, - token_id: u32, - is_phrase_query: bool, - metrics: &dyn MetricsCollector, - ) -> Result { - let mut posting = match self.group_range_for_token(token_id) { - // Grouped path (issue #7040): one cache entry covers rows - // [start, end), so neighbouring rare terms share a single read. - Some((start, end)) => { - let result = self - .index_cache - .get_or_insert_with_key_hit( - posting_list_group_cache_key(start, end, self.has_impacts), - || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=start); - self.load_posting_list_group(start, end).await - }, - ) - .await; - match &result { - Ok((_, true)) => metrics.record_index_cache_hit(), - _ => metrics.record_index_cache_miss(), - } - let (group, _) = result?; - let (max_score, length) = if group.needs_external_metadata() { - self.posting_metadata_for_token(token_id, Some(metrics)) - .await? - } else { - (None, None) - }; - let slot = (token_id - start) as usize; - group - .posting_list(slot, max_score, length)? - .ok_or_else(|| { - Error::index(format!( - "token {token_id} maps to slot {slot} outside posting group [{start}, {end})" - )) - })? - } - // Fallback for layouts that cannot use row-based groups: one cache - // entry per token. - None => { - let result = self - .index_cache - .get_or_insert_with_key_hit( - posting_list_cache_key(token_id, self.has_impacts), - || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); - // Fetch the posting batch and this token's (max_score, - // length) in parallel; for cold v2 partitions this is one - // single-row metadata read plus one posting-row read, - // instead of pulling the full per-token metadata table. - let (batch, (max_score, length)) = futures::try_join!( - self.posting_batch(token_id, false), - self.posting_metadata_for_token(token_id, Some(metrics)), - )?; - self.posting_list_from_batch(&batch, max_score, length) - }, - ) - .await; - match &result { - Ok((_, true)) => metrics.record_index_cache_hit(), - _ => metrics.record_index_cache_miss(), - } - result?.0.as_ref().clone() - } - }; - - if !self.modern_posting_is_validated(token_id)? { - self.ensure_modern_posting_validated(token_id, &posting) - .await?; - } - - if is_phrase_query && !posting.has_position() { - // hit the cache and when the cache was populated, the positions column was not loaded - let positions = self.read_positions(token_id, metrics).await?; - posting.set_positions(positions); - } - - Ok(posting) - } - - async fn ensure_modern_posting_validated( - &self, - token_id: u32, - posting: &PostingList, - ) -> Result<()> { - let (Some(validations), Some(num_docs)) = - (&self.modern_doc_id_validations, self.modern_num_docs) - else { - return Ok(()); - }; - let validation = validations.get(token_id as usize).ok_or_else(|| { - Error::index(format!( - "modern FTS token id {token_id} is outside validation state [0, {})", - validations.len() - )) - })?; - validation - .get_or_try_init(|| async { - Self::validate_modern_posting(token_id, posting, num_docs) - }) - .await - .map(|_| ()) - } - - #[inline] - fn modern_posting_is_validated(&self, token_id: u32) -> Result { - if self.modern_postings_validated.load(Ordering::Acquire) { - return Ok(true); - } - let (Some(validations), Some(_)) = (&self.modern_doc_id_validations, self.modern_num_docs) - else { - return Ok(true); - }; - let validation = validations.get(token_id as usize).ok_or_else(|| { - Error::index(format!( - "modern FTS token id {token_id} is outside validation state [0, {})", - validations.len() - )) - })?; - Ok(validation.get().is_some()) - } - - fn validate_modern_posting( - token_id: u32, - posting: &PostingList, - num_docs: usize, - ) -> Result<()> { - validate_modern_posting_doc_ids(posting, &format!("token id {token_id}"), num_docs) - } - - async fn publish_modern_posting_validated(&self, token_id: u32) -> Result<()> { - let Some(validations) = &self.modern_doc_id_validations else { - return Ok(()); - }; - let validation = validations.get(token_id as usize).ok_or_else(|| { - Error::index(format!( - "modern FTS token id {token_id} is outside validation state [0, {})", - validations.len() - )) - })?; - validation - .get_or_try_init(|| async { Result::Ok(()) }) - .await - .map(|_| ()) - } - - fn modern_posting_validation_ready(&self) -> bool { - if self.modern_postings_validated.load(Ordering::Acquire) { - return true; - } - let ready = self - .modern_doc_id_validations - .as_ref() - .is_none_or(|validations| validations.iter().all(|state| state.get().is_some())); - if ready { - self.modern_postings_validated - .store(true, Ordering::Release); - } - ready - } - - /// Map a token id to its cache group's row range `[start, end)`, or `None` - /// when grouping is not available so the caller falls back to the per-token - /// path. In v2 the token id is the row offset, so the group range is also - /// the physical row range. - fn group_range_for_token(&self, token_id: u32) -> Option<(u32, u32)> { - self.grouping.range_for_token(token_id, self.len()) - } - - /// Read rows `[start, end)` into one compact Arrow-backed cache value. - /// Positions are excluded; phrase queries load them on demand via - /// [`Self::read_positions`]. - async fn load_posting_list_group(&self, start: u32, end: u32) -> Result { - let mut columns = vec![POSTING_COL, MAX_SCORE_COL, LENGTH_COL]; - if self.has_impacts { - columns.push(IMPACT_COL); - } - let batch = self - .reader - .read_range(start as usize..end as usize, Some(&columns)) - .await?; - PostingListGroup::new_packed_with_block_size( - batch.shrink_to_fit()?, - self.posting_tail_codec, - self.block_size, - ) - } - - fn posting_list_from_batch_parts( - batch: &RecordBatch, - max_score: Option, - length: Option, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions_layout: PositionsLayout, - ) -> Result { - let posting_list = PostingList::from_batch_with_tail_codec_and_positions_layout( - batch, - max_score, - length, - posting_tail_codec, - block_size, - positions_layout, - )?; - Ok(posting_list) - } - - pub(crate) fn posting_list_from_batch( - &self, - batch: &RecordBatch, - max_score: Option, - length: Option, - ) -> Result { - Self::posting_list_from_batch_parts( - batch, - max_score, - length, - self.posting_tail_codec, - self.block_size, - self.positions_layout, - ) - } - - /// Build posting lists for one chunk's token range from `chunk_batch`, rebasing - /// global offsets to chunk-local rows. Returns `(global token_id, PostingList)` - /// pairs identical to the whole-file path, only bounded to one chunk. - fn build_prewarm_posting_lists_chunk( - chunk_batch: RecordBatch, - chunk: PrewarmChunk<'_>, - ctx: &PrewarmBuildCtx<'_>, - ) -> Result> { - let mut posting_lists = Vec::with_capacity(chunk.token_count); - for local in 0..chunk.token_count { - let global = chunk.tok_start + local; - let row_batch = if let Some(chunk_offsets) = chunk.offsets { - // Legacy v1: rebase global offsets to chunk row 0; the last token - // ends at `chunk.end_row` (no trailing sentinel in chunk_offsets). - let base = chunk_offsets[0]; - let start = chunk_offsets[local] - base; - let end = if local + 1 < chunk_offsets.len() { - chunk_offsets[local + 1] - base - } else { - chunk.end_row - base - }; - chunk_batch.slice(start, end - start) - } else { - // V2: one posting row per token; row `local` within the chunk. - chunk_batch.slice(local, 1) - }; - let row_batch = row_batch.shrink_to_fit()?; - let posting_list = Self::posting_list_from_batch_parts( - &row_batch, - ctx.max_scores.map(|scores| scores[global]), - ctx.lengths.map(|lengths| lengths[global]), - ctx.posting_tail_codec, - ctx.block_size, - ctx.positions_layout, - )?; - posting_lists.push((global as u32, posting_list)); - } - - Ok(posting_lists) - } - - /// Read the posting rows for token ids `[tok_start, tok_end)` into one RecordBatch. - /// For v2 the token range is the row range; for v1 it's derived from the offsets. - async fn read_chunk_batch( - &self, - tok_start: usize, - tok_end: usize, - with_position: bool, - ) -> Result { - let columns = self.posting_columns(with_position); - let row_range = match &self.metadata { - PostingMetadata::LegacyV1 { offsets, .. } => { - let start = offsets[tok_start]; - let end = offsets - .get(tok_end) - .copied() - .unwrap_or_else(|| self.reader.num_rows()); - start..end - } - PostingMetadata::V2 { .. } => tok_start..tok_end, - }; - let batch = self.reader.read_range(row_range, Some(&columns)).await?; - Ok(batch) - } - - async fn prewarm_posting_lists( - &self, - with_position: bool, - chunk_concurrency: usize, - ) -> Result<()> { - self.prewarm_posting_lists_chunked(with_position, None, chunk_concurrency) - .await?; - Ok(()) - } - - /// Stream the partition's posting lists into the cache in bounded token-row chunks - /// (read -> build -> insert -> drop), so peak resident set is ~one chunk. Returns - /// the chunk count (tests assert it split). `chunk_tokens_override` is test-only. - async fn prewarm_posting_lists_chunked( - &self, - with_position: bool, - chunk_tokens_override: Option, - chunk_concurrency: usize, - ) -> Result { - if with_position && !self.has_positions() { - return Err(Error::invalid_input( - "cannot prewarm positions for an inverted index that was built without positions; recreate the index with with_position=true".to_owned(), - )); - } - - // Make max_scores/lengths available for query-local packed views. The - // materialized fallback also clones them into its blocking build task. - self.ensure_metadata_loaded().await?; - - // With grouping the cache stores one entry per group, so a group's - // posting lists must all be resident at once: align chunk boundaries to - // whole groups. Without grouping, chunks are plain token ranges. - let grouping = self.grouping.clone(); - let use_packed_groups = grouping.is_grouped() && !with_position; - // Packed groups reuse the reader's bulk metadata at query time, so they - // do not need the temporary full-partition metadata clones used by the - // materialized fallback. - let state = (!use_packed_groups).then(|| self.chunk_build_state()); - let token_count = self.len(); - let posting_data_size_bytes = self.posting_data_size_bytes(); - let chunk_tokens = chunk_tokens_override - .unwrap_or_else(|| prewarm_chunk_tokens(token_count, posting_data_size_bytes)) - .max(1); - let chunk_ranges = prewarm_chunk_ranges(&grouping, token_count, chunk_tokens); - let chunk_count = chunk_ranges.len(); - let chunk_concurrency = chunk_concurrency.max(1); - - let read_build_start = Instant::now(); - stream::iter(chunk_ranges) - .map(|(tok_start, tok_end)| { - let state = state.as_ref(); - let grouping = &grouping; - async move { - if use_packed_groups { - let groups = self - .build_packed_chunk_groups(tok_start, tok_end, token_count, grouping) - .await?; - for (start, end, group) in groups { - self.index_cache - .insert_with_key( - &posting_list_group_cache_key(start, end, self.has_impacts), - Arc::new(group), - ) - .await; - } - } else { - let state = state.expect( - "materialized prewarm must initialize posting-list build state", - ); - let posting_lists = self - .build_chunk_postings(tok_start, tok_end, with_position, state) - .await?; - self.publish_chunk_postings( - posting_lists, - grouping, - tok_start, - tok_end, - token_count, - with_position, - ) - .await; - } - Result::Ok(()) - } - }) - .buffer_unordered(chunk_concurrency) - .try_collect::<()>() - .await?; - let read_build_elapsed = read_build_start.elapsed(); - - info!( - legacy_layout = self.is_legacy_layout(), - with_position, - token_count, - chunk_count, - chunk_tokens, - chunk_concurrency, - posting_data_size_bytes, - read_build_ms = read_build_elapsed.as_secs_f64() * 1000.0, - "posting list prewarm timing" - ); - - Ok(chunk_count) - } - - /// Loop-invariant inputs shared by every chunk build: the metadata vecs - /// (`Arc`d so chunks share them without re-cloning) plus codec/layout. - fn chunk_build_state(&self) -> ChunkBuildState { - let (offsets, max_scores, lengths) = match &self.metadata { - PostingMetadata::LegacyV1 { - offsets, - max_scores, - } => (Some(offsets.clone()), max_scores.clone(), None), - PostingMetadata::V2 { metadata } => ( - None, - metadata.get().map(|loaded| loaded.max_scores.clone()), - metadata.get().map(|loaded| loaded.lengths.clone()), - ), - }; - ChunkBuildState { - offsets: offsets.map(Arc::new), - max_scores: max_scores.map(Arc::new), - lengths: lengths.map(Arc::new), - posting_tail_codec: self.posting_tail_codec, - block_size: self.block_size, - positions_layout: self.positions_layout, - } - } - - /// Read one token-row chunk and build its posting lists off the runtime thread. - /// The large batch is dropped inside the blocking task once built, bounding - /// resident memory to one chunk. - async fn build_chunk_postings( - &self, - tok_start: usize, - tok_end: usize, - with_position: bool, - state: &ChunkBuildState, - ) -> Result> { - let chunk_token_count = tok_end - tok_start; - let chunk_batch = self - .read_chunk_batch(tok_start, tok_end, with_position) - .await?; - - let (chunk_offsets, chunk_end_row) = match state.offsets.as_ref() { - Some(offsets) => { - let end_row = offsets - .get(tok_end) - .copied() - .unwrap_or_else(|| self.reader.num_rows()); - (Some(offsets[tok_start..tok_end].to_vec()), end_row) - } - // V2 doesn't use chunk_end_row (one row per token); pass tok_end. - None => (None, tok_end), - }; - let max_scores = state.max_scores.clone(); - let lengths = state.lengths.clone(); - let posting_tail_codec = state.posting_tail_codec; - let block_size = state.block_size; - let positions_layout = state.positions_layout; - let num_docs = self.modern_num_docs; - let posting_lists = spawn_blocking(move || { - let ctx = PrewarmBuildCtx { - max_scores: max_scores.as_deref().map(|v| v.as_slice()), - lengths: lengths.as_deref().map(|v| v.as_slice()), - posting_tail_codec, - block_size, - positions_layout, - }; - let chunk = PrewarmChunk { - tok_start, - token_count: chunk_token_count, - offsets: chunk_offsets.as_deref(), - end_row: chunk_end_row, - }; - let posting_lists = Self::build_prewarm_posting_lists_chunk(chunk_batch, chunk, &ctx)?; - if let Some(num_docs) = num_docs { - for (token_id, posting) in &posting_lists { - Self::validate_modern_posting(*token_id, posting, num_docs)?; - } - } - Result::Ok(posting_lists) - }) - .await - .map_err(|err| { - Error::internal(format!( - "Failed to build prewarm posting lists in blocking task: {err}" - )) - })??; - for (token_id, _) in &posting_lists { - self.publish_modern_posting_validated(*token_id).await?; - } - // The chunk yields its token range as contiguous ascending ids from - // `tok_start`; the group publish path relies on this to index the lists. - debug_assert_eq!(posting_lists.len(), chunk_token_count); - debug_assert!( - posting_lists - .iter() - .enumerate() - .all(|(i, (token_id, _))| *token_id as usize == tok_start + i) - ); - Ok(posting_lists) - } - - /// Build compact v2 groups directly from one posting-row chunk. Each group - /// slice is deep-copied once, so it owns only its Arrow buffers without - /// materializing a `Vec` or retaining the full chunk. - async fn build_packed_chunk_groups( - &self, - tok_start: usize, - tok_end: usize, - token_count: usize, - grouping: &PostingGrouping, - ) -> Result> { - debug_assert!(grouping.is_grouped()); - debug_assert!(!self.is_legacy_layout()); - - let chunk_batch = self.read_chunk_batch(tok_start, tok_end, false).await?; - let ranges = grouping.ranges_for_chunk(tok_start, tok_end, token_count); - let posting_tail_codec = self.posting_tail_codec; - let block_size = self.block_size; - let num_docs = self.modern_num_docs; - let (chunk_max_scores, chunk_lengths) = match &self.metadata { - PostingMetadata::V2 { metadata } => { - let loaded = metadata.get().ok_or_else(|| { - Error::internal("packed prewarm requires loaded posting metadata".to_owned()) - })?; - ( - loaded.max_scores[tok_start..tok_end].to_vec(), - loaded.lengths[tok_start..tok_end].to_vec(), - ) - } - PostingMetadata::LegacyV1 { .. } => { - return Err(Error::internal( - "packed prewarm is not supported for legacy posting metadata".to_owned(), - )); - } - }; - - let groups = spawn_blocking(move || { - let mut groups = Vec::with_capacity(ranges.len()); - for (start, end) in ranges { - let start_usize = start as usize; - let end_usize = end as usize; - let local_start = start_usize - tok_start; - let group_len = end_usize - start_usize; - let group_batch = chunk_batch.slice(local_start, group_len).shrink_to_fit()?; - let group = PostingListGroup::new_packed_with_block_size( - group_batch, - posting_tail_codec, - block_size, - )?; - if let Some(num_docs) = num_docs { - for token_id in start..end { - let chunk_slot = token_id as usize - tok_start; - let posting = group - .posting_list( - (token_id - start) as usize, - Some(chunk_max_scores[chunk_slot]), - Some(chunk_lengths[chunk_slot]), - )? - .ok_or_else(|| { - Error::index(format!( - "token {token_id} is missing from prewarm posting group [{start}, {end})" - )) - })?; - Self::validate_modern_posting(token_id, &posting, num_docs)?; - } - } - groups.push((start, end, group)); - } - Result::Ok(groups) - }) - .await - .map_err(|err| { - Error::internal(format!( - "Failed to build packed prewarm posting groups in blocking task: {err}" - )) - })??; - for (start, end, _) in &groups { - for token_id in *start..*end { - self.publish_modern_posting_validated(token_id).await?; - } - } - Ok(groups) - } - - /// Strip positions into their own per-token cache entries (the posting cache - /// holds positions-free lists), then populate the same cache keys the read - /// path uses: grouped entries when grouping is active, per-token entries - /// otherwise. Called once per chunk; the chunk's lists drop on return. - async fn publish_chunk_postings( - &self, - posting_lists: Vec<(u32, PostingList)>, - grouping: &PostingGrouping, - tok_start: usize, - tok_end: usize, - token_count: usize, - with_position: bool, - ) { - match grouping { - PostingGrouping::None => { - for (token_id, mut posting_list) in posting_lists { - self.cache_positions(&mut posting_list, token_id, with_position) - .await; - self.index_cache - .insert_with_key( - &posting_list_cache_key(token_id, self.has_impacts), - Arc::new(posting_list), - ) - .await; - } - } - PostingGrouping::SyntheticFixed { .. } => { - let mut chunk_postings = Vec::with_capacity(posting_lists.len()); - for (token_id, mut posting_list) in posting_lists { - self.cache_positions(&mut posting_list, token_id, with_position) - .await; - chunk_postings.push(posting_list); - } - // Chunk is group-aligned, so every group starting in it also ends - // in it; `chunk_postings[i]` is token `tok_start + i`. The last - // group's `end` derives from `token_count`, matching the read path - // so both produce identical `PostingListGroupKey`s. - for (start, end) in grouping.ranges_for_chunk(tok_start, tok_end, token_count) { - let start_usize = start as usize; - let lo = start_usize - tok_start; - let hi = end as usize - tok_start; - let group = PostingListGroup::new(chunk_postings[lo..hi].to_vec()); - self.index_cache - .insert_with_key( - &posting_list_group_cache_key(start, end, self.has_impacts), - Arc::new(group), - ) - .await; - } - } - } - } - - /// Move a posting list's positions (when present and requested) into the - /// dedicated per-token position cache, leaving the posting list positions-free. - async fn cache_positions( - &self, - posting_list: &mut PostingList, - token_id: u32, - with_position: bool, - ) { - if with_position && let Some(positions) = posting_list.take_positions() { - self.index_cache - .insert_with_key(&PositionKey { token_id }, Arc::new(Positions(positions))) - .await; - } - } - - /// Cheap `invert.lance` size estimate (file length from object metadata, no - /// data read), used only to size prewarm chunks. Falls back to a row-count - /// proxy when the reader can't surface the length (legacy v1). - pub(crate) fn posting_data_size_bytes(&self) -> u64 { - if let Some(size) = self.reader.file_size_bytes() { - return size; - } - // Fallback proxy for readers that don't cache their file length: just needs - // to be monotonic in partition size. - const ESTIMATED_BYTES_PER_ROW: u64 = 16; - (self.reader.num_rows() as u64).saturating_mul(ESTIMATED_BYTES_PER_ROW) - } - - pub(crate) async fn read_batch(&self, with_position: bool) -> Result { - let columns = self.posting_columns(with_position); - let batch = self - .reader - .read_range(0..self.reader.num_rows(), Some(&columns)) - .await?; - Ok(batch) - } - - pub(crate) async fn read_all( - &self, - with_position: bool, - ) -> Result> + '_> { - // read_all walks every posting list; the bulk metadata is paid for - // unconditionally, so just load it once up front and index into it - // synchronously below. - self.ensure_metadata_loaded().await?; - let batch = self.read_batch(with_position).await?; - Ok((0..self.len()).map(move |i| { - let token_id = i as u32; - let range = self.posting_list_range(token_id); - let batch = batch.slice(i, range.end - range.start); - let (max_score, length) = self.bulk_metadata_for_token(token_id); - self.posting_list_from_batch(&batch, max_score, length) - })) - } - - /// Sync lookup of `(max_score, length)` from the bulk-loaded metadata. - /// Only safe after [`Self::ensure_metadata_loaded`]; callers that hold - /// the OnceCell-loaded reference (e.g. read_all, prewarm) use this to - /// avoid the per-token IO path. - fn bulk_metadata_for_token(&self, token_id: u32) -> (Option, Option) { - match &self.metadata { - PostingMetadata::LegacyV1 { max_scores, .. } => { - (max_scores.as_ref().map(|m| m[token_id as usize]), None) - } - PostingMetadata::V2 { metadata } => { - let loaded = metadata.get().expect( - "v2 metadata must be bulk-loaded before bulk_metadata_for_token; call ensure_metadata_loaded first", - ); - ( - Some(loaded.max_scores[token_id as usize]), - Some(loaded.lengths[token_id as usize]), - ) - } - } - } - - async fn read_positions( - &self, - token_id: u32, - metrics: &dyn MetricsCollector, - ) -> Result { - let result = self.index_cache.get_or_insert_with_key_hit(PositionKey { token_id }, || async move { - let positions = match self.positions_layout { - PositionsLayout::None => { - return Err(Error::invalid_input( - "position is not found but required for phrase queries, try recreating the index with position".to_owned(), - )); - } - PositionsLayout::LegacyPerDoc => { - let batch = self - .reader - .read_range(self.posting_list_range(token_id), Some(&[POSITION_COL])) - .await - .map_err(|e| match e { - Error::Schema { .. } => Error::invalid_input("position is not found but required for phrase queries, try recreating the index with position".to_owned()), - e => e, - })?; - CompressedPositionStorage::LegacyPerDoc( - batch[POSITION_COL].as_list::().value(0).as_list::().clone(), - ) - } - PositionsLayout::SharedStream(codec) => { - let batch = self - .reader - .read_range( - self.posting_list_range(token_id), - Some(&[COMPRESSED_POSITION_COL, POSITION_BLOCK_OFFSET_COL]), - ) - .await - .map_err(|e| match e { - Error::Schema { .. } => Error::invalid_input("position is not found but required for phrase queries, try recreating the index with position".to_owned()), - e => e, - })?; - let bytes = bytes::Bytes::from( - batch[COMPRESSED_POSITION_COL] - .as_binary::() - .value(0) - .to_vec(), - ); - let block_offsets = batch[POSITION_BLOCK_OFFSET_COL] - .as_list::() - .value(0) - .as_primitive::() - .values() - .to_vec(); - CompressedPositionStorage::SharedStream(SharedPositionStream::new( - codec, - block_offsets, - bytes, - )) - } - }; - Result::Ok(Positions(positions)) - }).await; - match &result { - Ok((_, true)) => metrics.record_index_cache_hit(), - _ => metrics.record_index_cache_miss(), - } - let (positions, _) = result?; - Ok(positions.0.clone()) - } - - fn posting_list_range(&self, token_id: u32) -> Range { - match &self.metadata { - PostingMetadata::LegacyV1 { offsets, .. } => { - let offset = offsets[token_id as usize]; - let posting_len = self.posting_len(token_id); - offset..offset + posting_len - } - PostingMetadata::V2 { .. } => { - let token_id = token_id as usize; - token_id..token_id + 1 - } - } - } - - fn posting_columns(&self, with_position: bool) -> Vec<&'static str> { - let mut base_columns = if self.is_legacy_layout() { - vec![ROW_ID, FREQUENCY_COL] - } else { - vec![POSTING_COL] - }; - if with_position { - match self.positions_layout { - PositionsLayout::None => {} - PositionsLayout::LegacyPerDoc => base_columns.push(POSITION_COL), - PositionsLayout::SharedStream(_) => { - base_columns.push(COMPRESSED_POSITION_COL); - base_columns.push(POSITION_BLOCK_OFFSET_COL); - } - } - } - if self.has_impacts { - base_columns.push(IMPACT_COL); - } - base_columns - } -} - -/// Loop-invariant state for [`InvertedPartition::build_chunk_postings`]. The -/// metadata vecs are `Arc`d so each chunk's blocking build shares them cheaply. -struct ChunkBuildState { - offsets: Option>>, - max_scores: Option>>, - lengths: Option>>, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions_layout: PositionsLayout, -} - -/// Chunk-invariant inputs to [`InvertedPartition::build_prewarm_posting_lists_chunk`]: -/// the per-partition codec/layout and the (shared, whole-partition) metadata -/// slices indexed by global token id. These don't change across chunks. -struct PrewarmBuildCtx<'a> { - max_scores: Option<&'a [f32]>, - lengths: Option<&'a [u32]>, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions_layout: PositionsLayout, -} - -/// Per-chunk inputs to [`InvertedPartition::build_prewarm_posting_lists_chunk`]: -/// the token sub-range `[tok_start, tok_start + token_count)` and, for legacy -/// v1, the rebased offset slice plus the chunk's end row. -struct PrewarmChunk<'a> { - tok_start: usize, - token_count: usize, - /// Legacy v1 only: `offsets[tok_start..tok_start+token_count]` (no sentinel). - offsets: Option<&'a [usize]>, - /// Legacy v1 only: global row at which this chunk's posting rows end. - end_row: usize, -} - -/// New type just to allow Positions implement DeepSizeOf so it can be put -/// in the cache. -#[derive(Clone)] -pub struct Positions(pub(super) CompressedPositionStorage); - -/// Slice-aware cache-size charge for the Arrow array shapes stored in posting -/// caches. [`Array::get_buffer_memory_size`] reports the full capacity of shared -/// backing buffers; cached posting lists often reference only a small slice of a -/// group read. Count the referenced span for the known posting-list types and -/// fall back to Arrow's full-buffer size for anything else. -fn sliced_cache_bytes(array: &dyn Array) -> usize { - let validity_bytes = array - .nulls() - .map(|nulls| nulls.len().div_ceil(8)) - .unwrap_or(0); - match array.data_type() { - DataType::LargeBinary => { - let array = array.as_binary::(); - let data_bytes = if array.is_empty() { - 0 - } else { - let offsets = array.value_offsets(); - (offsets[array.len()] - offsets[0]) as usize - }; - data_bytes + (array.len() + 1) * std::mem::size_of::() + validity_bytes - } - DataType::List(_) => { - let array = array.as_list::(); - let (child_start, child_end) = if array.is_empty() { - (0, 0) - } else { - let offsets = array.value_offsets(); - (offsets[0] as usize, offsets[array.len()] as usize) - }; - let offset_bytes = (array.len() + 1) * std::mem::size_of::(); - let child = array.values().slice(child_start, child_end - child_start); - offset_bytes + validity_bytes + sliced_cache_bytes(child.as_ref()) - } - // Fixed-width primitives hold exactly `len * width` bytes regardless of - // buffer capacity, so this is already slice-aware. Any other type falls - // back to the full-buffer size. - other => match other.primitive_width() { - Some(width) => array.len() * width + validity_bytes, - None => array.get_buffer_memory_size(), - }, - } -} - -impl DeepSizeOf for Positions { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.0.deep_size_of_children(context) - } -} - -// Cache key implementations for type-safe cache access -#[derive(Debug, Clone)] -pub struct PostingListKey { - pub token_id: u32, -} - -impl CacheKey for PostingListKey { - type ValueType = PostingList; - - fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}", self.token_id).into() - } - - fn type_name() -> &'static str { - "PostingList" - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.scalar.inverted.posting-list-key", 1) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_u32(self.token_id); - } - - fn codec() -> Option { - Some(CacheCodec::from_impl::()) - } -} - -/// Cache key for a group of consecutive posting lists stored as a single -/// entry, covering rows `[start, end)` (issue #7040). The range, not a token -/// id, is the key so a runtime group-size change simply misses old entries -/// instead of serving a differently-shaped group. -#[derive(Debug, Clone)] -pub struct PostingListGroupKey { - pub start: u32, - pub end: u32, -} - -impl CacheKey for PostingListGroupKey { - type ValueType = PostingListGroup; - - fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}-{}", self.start, self.end).into() - } - - fn type_name() -> &'static str { - "PostingListGroup" - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.scalar.inverted.posting-list-group-key", 1) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_u32(self.start); - builder.write_u32(self.end); - } - - fn codec() -> Option { - Some(CacheCodec::from_impl::()) - } -} - -/// Internal cache-key decorator that isolates impact-bearing posting values -/// without changing the source-compatible public posting key structs. -#[derive(Debug, Clone)] -struct ImpactAwareCacheKey { - inner: K, - has_impacts: bool, -} - -impl CacheKey for ImpactAwareCacheKey { - type ValueType = K::ValueType; - - fn key(&self) -> std::borrow::Cow<'_, str> { - if self.has_impacts { - format!("{}-impacts", self.inner.key()).into() - } else { - self.inner.key() - } - } - - fn type_name() -> &'static str { - K::type_name() - } - - fn stable_type_id() -> &'static str { - K::stable_type_id() - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.scalar.inverted.impact-aware-key", 1) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - let inner_schema = K::schema(); - builder.write_str(K::stable_type_id()); - builder.write_str(inner_schema.id()); - builder.write_u32(inner_schema.version()); - builder.write_variant(if self.has_impacts { 1 } else { 0 }); - self.inner.write_key(builder); - } - - fn codec() -> Option { - K::codec() - } -} - -fn posting_list_cache_key(token_id: u32, has_impacts: bool) -> ImpactAwareCacheKey { - ImpactAwareCacheKey { - inner: PostingListKey { token_id }, - has_impacts, - } -} - -fn posting_list_group_cache_key( - start: u32, - end: u32, - has_impacts: bool, -) -> ImpactAwareCacheKey { - ImpactAwareCacheKey { - inner: PostingListGroupKey { start, end }, - has_impacts, - } -} - -#[derive(Debug, Clone, DeepSizeOf)] -struct PostingMetadataValue { - max_score: f32, - length: u32, -} - -#[derive(Debug, Clone)] -struct PostingMetadataKey { - token_id: u32, -} - -impl CacheKey for PostingMetadataKey { - type ValueType = PostingMetadataValue; - - fn key(&self) -> std::borrow::Cow<'_, str> { - format!("posting-metadata-{}", self.token_id).into() - } - - fn type_name() -> &'static str { - "PostingMetadata" - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.scalar.inverted.posting-metadata-key", 1) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_u32(self.token_id); - } -} - -#[derive(Debug, Clone)] -pub struct PositionKey { - pub token_id: u32, -} - -impl CacheKey for PositionKey { - type ValueType = Positions; - - fn key(&self) -> std::borrow::Cow<'_, str> { - format!("positions-{}", self.token_id).into() - } - - fn type_name() -> &'static str { - "Position" - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.scalar.inverted.position-key", 1) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_u32(self.token_id); - } - - fn codec() -> Option { - Some(CacheCodec::from_impl::()) - } -} - -#[derive(Debug, Clone, PartialEq)] -pub enum CompressedPositionStorage { - LegacyPerDoc(ListArray), - SharedStream(SharedPositionStream), -} - -impl DeepSizeOf for CompressedPositionStorage { - fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::LegacyPerDoc(positions) => sliced_cache_bytes(positions), - Self::SharedStream(stream) => stream.size(), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Default)] -pub struct SharedPositionStream { - codec: PositionStreamCodec, - block_offsets: Arc<[u32]>, - // Stored with shared ownership so cache hits can clone position streams - // without copying either offsets or bytes. - bytes: bytes::Bytes, -} - -impl SharedPositionStream { - pub fn new(codec: PositionStreamCodec, block_offsets: Vec, bytes: bytes::Bytes) -> Self { - Self { - codec, - block_offsets: Arc::from(block_offsets.into_boxed_slice()), - bytes, - } - } - - pub fn codec(&self) -> PositionStreamCodec { - self.codec - } - - pub fn block_count(&self) -> usize { - self.block_offsets.len() - } - - pub fn block_range(&self, index: usize) -> Range { - let start = self.block_offsets[index] as usize; - let end = self - .block_offsets - .get(index + 1) - .map(|offset| *offset as usize) - .unwrap_or(self.bytes.len()); - start..end - } - - pub fn block(&self, index: usize) -> &[u8] { - let range = self.block_range(index); - &self.bytes[range] - } - - pub fn bytes(&self) -> &[u8] { - &self.bytes - } - - pub fn block_offsets(&self) -> &[u32] { - self.block_offsets.as_ref() - } - - pub fn size(&self) -> usize { - self.block_offsets.len() * std::mem::size_of::() + self.bytes.len() - } -} - -/// A group of consecutive posting lists held in a single cache entry, in row -/// order (issue #7040). Prewarmed modern groups without positions retain only -/// the compact Arrow posting rows read from `invert.lance`; max-score/length -/// metadata stays in the reader and is injected when a query creates a -/// posting-list view. Cold-loaded groups may keep inline metadata to preserve -/// one-read query loading. Legacy and position-bearing prewarm paths use the -/// materialized fallback. -#[derive(Debug, Clone)] -pub struct PostingListGroup { - pub(super) storage: PostingListGroupStorage, -} - -#[derive(Debug, Clone)] -pub(super) enum PostingListGroupStorage { - Packed(PackedPostingListGroup), - Materialized(Vec), -} - -#[derive(Debug, Clone)] -pub(super) struct PackedPostingListGroup { - pub(super) batch: RecordBatch, - pub(super) posting_tail_codec: PostingTailCodec, - pub(super) block_size: usize, - first_docs_states: Arc<[OnceLock>]>, - first_docs_state_capacity_bytes: usize, - impact_states: Option>]>>, - impact_state_capacity_bytes: usize, -} - -impl DeepSizeOf for PostingListGroup { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - match &self.storage { - PostingListGroupStorage::Packed(group) => group - .batch - .columns() - .iter() - .map(|column| sliced_cache_bytes(column.as_ref())) - .sum::() - .saturating_add(group.first_docs_state_capacity_bytes) - .saturating_add(group.impact_state_capacity_bytes), - PostingListGroupStorage::Materialized(posting_lists) => { - posting_lists.deep_size_of_children(context) - } - } - } -} - -impl PostingListGroup { - pub(super) fn new(posting_lists: Vec) -> Self { - Self { - storage: PostingListGroupStorage::Materialized(posting_lists), - } - } - - pub(super) fn new_packed( - batch: RecordBatch, - posting_tail_codec: PostingTailCodec, - ) -> Result { - let block_size = parse_posting_block_size(batch.schema_ref().metadata())?; - Self::new_packed_with_block_size(batch, posting_tail_codec, block_size) - } - - fn new_packed_with_block_size( - batch: RecordBatch, - posting_tail_codec: PostingTailCodec, - block_size: usize, - ) -> Result { - validate_block_size(block_size)?; - if let Some(encoded_block_size) = batch.schema_ref().metadata().get(POSTING_BLOCK_SIZE_KEY) - { - let encoded_block_size = encoded_block_size.parse::().map_err(|err| { - Error::index(format!( - "invalid {POSTING_BLOCK_SIZE_KEY} metadata value {encoded_block_size:?}: {err}" - )) - })?; - if encoded_block_size != block_size { - return Err(Error::index(format!( - "packed posting group {POSTING_BLOCK_SIZE_KEY}={encoded_block_size} does not match block_size={block_size}" - ))); - } - } - - // Projected reads may drop schema metadata. Restore the reader's - // validated block size before the batch enters the packed cache so IPC - // roundtrips remain self-describing. Older packed cache entries omit - // the key and enter through new_packed with the legacy 128-doc default. - let mut schema = batch.schema().as_ref().clone(); - schema - .metadata - .insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); - let batch = batch.with_schema(Arc::new(schema))?; - let postings = batch - .column_by_name(POSTING_COL) - .and_then(|column| column.as_list_opt::()) - .ok_or_else(|| { - Error::index(format!( - "packed posting group column {POSTING_COL} must be List" - )) - })?; - if postings.values().data_type() != &DataType::LargeBinary { - return Err(Error::index(format!( - "packed posting group column {POSTING_COL} must contain LargeBinary values, got {}", - postings.values().data_type() - ))); - } - if postings.null_count() != 0 { - return Err(Error::index( - "packed posting group column must not contain nulls".to_string(), - )); - } - let total_posting_blocks = (0..batch.num_rows()) - .map(|slot| postings.value_length(slot) as usize) - .sum::(); - let first_docs_states: Arc<[OnceLock>]> = (0..batch.num_rows()) - .map(|_| OnceLock::new()) - .collect::>() - .into(); - // Reserve the compact per-slot state slab and the block-head arrays it - // can lazily retain, so warming these derived values cannot grow the - // cache beyond its admission charge. - let first_docs_state_capacity_bytes = first_docs_states - .len() - .saturating_mul(std::mem::size_of::>>()) - .saturating_add(total_posting_blocks.saturating_mul(std::mem::size_of::())); - let (impact_states, impact_state_capacity_bytes) = if let Some(impacts) = - batch.column_by_name(IMPACT_COL) - { - let impacts = impacts.as_list_opt::().ok_or_else(|| { - Error::index(format!( - "packed posting group column {IMPACT_COL} must be List" - )) - })?; - if impacts.values().data_type() != &DataType::LargeBinary { - return Err(Error::index(format!( - "packed posting group column {IMPACT_COL} must contain LargeBinary values, got {}", - impacts.values().data_type() - ))); - } - if impacts.null_count() != 0 { - return Err(Error::index(format!( - "packed posting group column {IMPACT_COL} must not contain nulls" - ))); - } - let mut derived_cache_bytes = 0usize; - for slot in 0..batch.num_rows() { - let posting_blocks = postings.value_length(slot) as usize; - let impact_entries = impacts.value_length(slot) as usize; - let expected_impact_entries = - posting_blocks.saturating_add(posting_blocks.div_ceil(IMPACT_LEVEL1_BLOCKS)); - if impact_entries != expected_impact_entries { - return Err(Error::index(format!( - "packed posting group impact slot {slot} has {impact_entries} entries, expected {expected_impact_entries} for {posting_blocks} posting blocks" - ))); - } - derived_cache_bytes = derived_cache_bytes.saturating_add( - ImpactSkipData::derived_cache_bytes_for_entries(impact_entries), - ); - } - - let states: Arc<[OnceLock>]> = (0..batch.num_rows()) - .map(|_| OnceLock::new()) - .collect::>() - .into(); - // Account up front for every allocation that the lazy states can - // eventually retain. The impact entry bytes themselves remain in - // `batch` and are already charged exactly once above. - let per_slot_bytes = std::mem::size_of::>>() - .saturating_add(std::mem::size_of::()); - let capacity_bytes = states - .len() - .saturating_mul(per_slot_bytes) - .saturating_add(derived_cache_bytes); - (Some(states), capacity_bytes) - } else { - (None, 0) - }; - match ( - batch.column_by_name(MAX_SCORE_COL), - batch.column_by_name(LENGTH_COL), - ) { - (None, None) => {} - (Some(max_scores), Some(lengths)) => { - let max_scores = max_scores - .as_primitive_opt::() - .ok_or_else(|| { - Error::index(format!( - "packed posting group column {MAX_SCORE_COL} must be Float32" - )) - })?; - let lengths = lengths.as_primitive_opt::().ok_or_else(|| { - Error::index(format!( - "packed posting group column {LENGTH_COL} must be UInt32" - )) - })?; - if max_scores.null_count() != 0 || lengths.null_count() != 0 { - return Err(Error::index( - "packed posting group metadata columns must not contain nulls".to_string(), - )); - } - } - _ => { - return Err(Error::index(format!( - "packed posting group must contain both {MAX_SCORE_COL} and {LENGTH_COL}, or neither" - ))); - } - } - - Ok(Self { - storage: PostingListGroupStorage::Packed(PackedPostingListGroup { - batch, - posting_tail_codec, - block_size, - first_docs_states, - first_docs_state_capacity_bytes, - impact_states, - impact_state_capacity_bytes, - }), - }) - } - - pub(super) fn len(&self) -> usize { - match &self.storage { - PostingListGroupStorage::Packed(group) => group.batch.num_rows(), - PostingListGroupStorage::Materialized(posting_lists) => posting_lists.len(), - } - } - - #[cfg(test)] - pub(super) fn is_packed(&self) -> bool { - matches!(&self.storage, PostingListGroupStorage::Packed(_)) - } - - fn needs_external_metadata(&self) -> bool { - match &self.storage { - PostingListGroupStorage::Packed(group) => { - group.batch.column_by_name(MAX_SCORE_COL).is_none() - } - PostingListGroupStorage::Materialized(_) => false, - } - } - - /// Build an owned posting-list view for `slot`. Packed groups clone only - /// Arrow array metadata; the compressed posting bytes remain shared with - /// the group's `List` child buffers. - pub(super) fn posting_list( - &self, - slot: usize, - max_score: Option, - length: Option, - ) -> Result> { - match &self.storage { - PostingListGroupStorage::Materialized(posting_lists) => { - Ok(posting_lists.get(slot).cloned()) - } - PostingListGroupStorage::Packed(group) => { - if slot >= group.batch.num_rows() { - return Ok(None); - } - let postings = group - .batch - .column_by_name(POSTING_COL) - .and_then(|column| column.as_list_opt::()) - .ok_or_else(|| { - Error::index(format!( - "packed posting group column {POSTING_COL} must be List" - )) - })?; - let blocks = postings.value(slot); - let blocks = blocks.as_binary_opt::().ok_or_else(|| { - Error::index(format!( - "packed posting group slot {slot} is not LargeBinary" - )) - })?; - let max_score = match group.batch.column_by_name(MAX_SCORE_COL) { - Some(column) => column - .as_primitive_opt::() - .expect("packed group metadata was validated at construction") - .value(slot), - None => max_score.ok_or_else(|| { - Error::index("packed posting group requires max-score metadata".to_string()) - })?, - }; - let length = match group.batch.column_by_name(LENGTH_COL) { - Some(column) => column - .as_primitive_opt::() - .expect("packed group metadata was validated at construction") - .value(slot), - None => length.ok_or_else(|| { - Error::index("packed posting group requires length metadata".to_string()) - })?, - }; - let impacts = match ( - group.impact_states.as_ref(), - group.batch.column_by_name(IMPACT_COL), - ) { - (Some(states), Some(column)) => { - let state = states.get(slot).ok_or_else(|| { - Error::index(format!( - "packed posting group impact state missing slot {slot}" - )) - })?; - let impact_lists = column.as_list_opt::().ok_or_else(|| { - Error::index(format!( - "packed posting group column {IMPACT_COL} must be List" - )) - })?; - let entries = impact_lists.value(slot); - let entries = entries.as_binary_opt::().ok_or_else(|| { - Error::index(format!( - "packed posting group impact slot {slot} is not LargeBinary" - )) - })?; - let impacts = - state.get_or_init(|| { - Box::new(ImpactSkipData::new(entries.clone(), blocks.len()).expect( - "packed impact entry count was validated at construction", - )) - }); - Some(impacts.as_ref().clone()) - } - (None, None) => None, - _ => { - return Err(Error::internal( - "packed posting group impact column/state mismatch".to_string(), - )); - } - }; - Ok(Some(PostingList::Compressed( - CompressedPostingList::new( - blocks.clone(), - max_score, - length, - group.posting_tail_codec, - group.block_size, - None, - impacts, - ) - .with_packed_first_docs(group.first_docs_states.clone(), slot), - ))) - } - } - } -} - -#[derive(Debug, Clone, DeepSizeOf)] -#[allow(clippy::large_enum_variant)] -pub enum PostingList { - Plain(PlainPostingList), - Compressed(CompressedPostingList), -} - -impl PostingList { - pub fn from_batch( - batch: &RecordBatch, - max_score: Option, - length: Option, - ) -> Result { - let posting_tail_codec = parse_posting_tail_codec(batch.schema_ref().metadata())?; - let block_size = parse_posting_block_size(batch.schema_ref().metadata())?; - Self::from_batch_with_tail_codec(batch, max_score, length, posting_tail_codec, block_size) - } - - pub fn from_batch_with_tail_codec( - batch: &RecordBatch, - max_score: Option, - length: Option, - posting_tail_codec: PostingTailCodec, - block_size: usize, - ) -> Result { - let positions_layout = if batch.column_by_name(COMPRESSED_POSITION_COL).is_some() { - PositionsLayout::SharedStream(parse_shared_position_codec( - batch.schema_ref().metadata(), - )?) - } else if batch.column_by_name(POSITION_COL).is_some() { - PositionsLayout::LegacyPerDoc - } else { - PositionsLayout::None - }; - Self::from_batch_with_tail_codec_and_positions_layout( - batch, - max_score, - length, - posting_tail_codec, - block_size, - positions_layout, - ) - } - - fn from_batch_with_tail_codec_and_positions_layout( - batch: &RecordBatch, - max_score: Option, - length: Option, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions_layout: PositionsLayout, - ) -> Result { - match batch.column_by_name(POSTING_COL) { - Some(_) => { - debug_assert!(max_score.is_some() && length.is_some()); - let shared_position_codec = match positions_layout { - PositionsLayout::SharedStream(codec) => Some(codec), - _ => None, - }; - let posting = CompressedPostingList::from_batch( - batch, - max_score.unwrap(), - length.unwrap(), - posting_tail_codec, - block_size, - shared_position_codec, - )?; - Ok(Self::Compressed(posting)) - } - None => { - let posting = PlainPostingList::from_batch(batch, max_score); - Ok(Self::Plain(posting)) - } - } - } - - pub fn iter(&self) -> PostingListIterator<'_> { - PostingListIterator::new(self) - } - - pub fn has_position(&self) -> bool { - match self { - Self::Plain(posting) => posting.positions.is_some(), - Self::Compressed(posting) => posting.positions.is_some(), - } - } - - pub fn has_impacts(&self) -> bool { - match self { - Self::Plain(_) => false, - Self::Compressed(posting) => posting.impacts.is_some(), - } - } - - pub fn set_positions(&mut self, positions: CompressedPositionStorage) { - match self { - Self::Plain(posting) => match positions { - CompressedPositionStorage::LegacyPerDoc(positions) => { - posting.positions = Some(positions) - } - CompressedPositionStorage::SharedStream(_) => { - unreachable!("shared position stream is not supported for plain postings") - } - }, - Self::Compressed(posting) => { - posting.positions = Some(positions); - } - } - } - - pub fn take_positions(&mut self) -> Option { - match self { - Self::Plain(posting) => posting - .positions - .take() - .map(CompressedPositionStorage::LegacyPerDoc), - Self::Compressed(posting) => posting.positions.take(), - } - } - - pub fn max_score(&self) -> Option { - match self { - Self::Plain(posting) => posting.max_score, - Self::Compressed(posting) => Some(posting.max_score), - } - } - - pub fn len(&self) -> usize { - match self { - Self::Plain(posting) => posting.len(), - Self::Compressed(posting) => posting.length as usize, - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn into_builder(self, docs: &DocSet) -> PostingListBuilder { - let posting_tail_codec = match &self { - Self::Plain(_) => PostingTailCodec::Fixed32, - Self::Compressed(posting) => posting.posting_tail_codec, - }; - let block_size = match &self { - Self::Plain(_) => LEGACY_BLOCK_SIZE, - Self::Compressed(posting) => posting.block_size, - }; - let mut builder = PostingListBuilder::new_with_posting_tail_codec_and_block_size( - self.has_position(), - posting_tail_codec, - block_size, - ); - match self { - // legacy format - Self::Plain(posting) => { - // convert the posting list to the new format: - // 1. map row ids to doc ids - // 2. sort the posting list by doc ids - struct Item { - doc_id: u32, - positions: PositionRecorder, - } - let doc_ids = docs - .row_ids - .iter() - .enumerate() - .map(|(doc_id, row_id)| (*row_id, doc_id as u32)) - .collect::>(); - let mut items = Vec::with_capacity(posting.len()); - for (row_id, freq, positions) in posting.iter() { - let freq = freq as u32; - let positions = match positions { - Some(positions) => { - PositionRecorder::Position(positions.collect::>().into()) - } - None => PositionRecorder::Count(freq), - }; - items.push(Item { - doc_id: doc_ids[&row_id], - positions, - }); - } - items.sort_unstable_by_key(|item| item.doc_id); - for item in items { - builder.add(item.doc_id, item.positions); - } - } - Self::Compressed(posting) => { - posting.iter().for_each(|(doc_id, freq, positions)| { - let positions = match positions { - Some(positions) => { - PositionRecorder::Position(positions.collect::>().into()) - } - None => PositionRecorder::Count(freq), - }; - builder.add(doc_id, positions); - }); - } - } - builder - } -} - -#[derive(Debug, PartialEq, Clone)] -pub struct PlainPostingList { - pub row_ids: ScalarBuffer, - pub frequencies: ScalarBuffer, - pub max_score: Option, - pub positions: Option, // List of Int32 -} - -impl DeepSizeOf for PlainPostingList { - fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { - self.row_ids.len() * std::mem::size_of::() - + self.frequencies.len() * std::mem::size_of::() - + self - .positions - .as_ref() - .map(|positions| sliced_cache_bytes(positions)) - .unwrap_or(0) - } -} - -impl PlainPostingList { - pub fn new( - row_ids: ScalarBuffer, - frequencies: ScalarBuffer, - max_score: Option, - positions: Option, - ) -> Self { - Self { - row_ids, - frequencies, - max_score, - positions, - } - } - - pub fn from_batch(batch: &RecordBatch, max_score: Option) -> Self { - let row_ids = batch[ROW_ID].as_primitive::().values().clone(); - let frequencies = batch[FREQUENCY_COL] - .as_primitive::() - .values() - .clone(); - let positions = batch - .column_by_name(POSITION_COL) - .map(|col| col.as_list::().clone()); - - Self::new(row_ids, frequencies, max_score, positions) - } - - pub fn len(&self) -> usize { - self.row_ids.len() - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn iter(&self) -> PlainPostingListIterator<'_> { - Box::new( - self.row_ids - .iter() - .zip(self.frequencies.iter()) - .enumerate() - .map(|(idx, (doc_id, freq))| { - ( - *doc_id, - *freq, - self.positions.as_ref().map(|p| { - let start = p.value_offsets()[idx] as usize; - let end = p.value_offsets()[idx + 1] as usize; - Box::new( - p.values().as_primitive::().values()[start..end] - .iter() - .map(|pos| *pos as u32), - ) as _ - }), - ) - }), - ) - } - - #[inline] - pub fn doc(&self, i: usize) -> LocatedDocInfo { - LocatedDocInfo::new(self.row_ids[i], self.frequencies[i]) - } - - pub fn positions(&self, index: usize) -> Option> { - self.positions - .as_ref() - .map(|positions| positions.value(index)) - } - - pub fn max_score(&self) -> Option { - self.max_score - } - - pub fn row_id(&self, i: usize) -> u64 { - self.row_ids[i] - } -} - -#[derive(Debug, Clone)] -enum FirstDocsState { - Standalone(Arc>>), - Packed { - states: Arc<[OnceLock>]>, - slot: usize, - }, -} - -impl FirstDocsState { - fn standalone() -> Self { - Self::Standalone(Arc::new(OnceLock::new())) - } - - fn state(&self) -> &OnceLock> { - match self { - Self::Standalone(state) => state, - Self::Packed { states, slot } => &states[*slot], - } - } - - fn get_or_init(&self, initialize: impl FnOnce() -> Box<[u32]>) -> &[u32] { - self.state().get_or_init(initialize) - } - - fn capacity_bytes( - &self, - block_count: usize, - context: &mut lance_core::deepsize::Context, - ) -> usize { - if context.mark_seen(self.state() as *const _ as usize) { - std::mem::size_of::>>() - .saturating_add(block_count.saturating_mul(std::mem::size_of::())) - } else { - 0 - } - } - - #[cfg(test)] - fn shares_state_with(&self, other: &Self) -> bool { - std::ptr::eq(self.state(), other.state()) - } -} - -#[derive(Debug, Clone)] -pub struct CompressedPostingList { - pub max_score: f32, - pub length: u32, - // each binary is a block of compressed data - // that contains `block_size` doc ids and then `block_size` frequencies, - // packed by the physical bitpacker matching that block size. - pub blocks: LargeBinaryArray, - pub posting_tail_codec: PostingTailCodec, - pub block_size: usize, - pub positions: Option, - pub(crate) impacts: Option, - // First doc id per block, baked lazily and shared across per-query clones - // of the cached list. See `block_first_docs`. - first_docs: FirstDocsState, -} - -impl PartialEq for CompressedPostingList { - fn eq(&self, other: &Self) -> bool { - self.max_score == other.max_score - && self.length == other.length - && self.blocks == other.blocks - && self.posting_tail_codec == other.posting_tail_codec - && self.block_size == other.block_size - && self.positions == other.positions - && self.impacts == other.impacts - } -} - -impl DeepSizeOf for CompressedPostingList { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - sliced_cache_bytes(&self.blocks) - + self - .positions - .as_ref() - .map(|positions| positions.deep_size_of_children(context)) - .unwrap_or(0) - + self - .impacts - .as_ref() - .map(|impacts| { - sliced_cache_bytes(impacts.entries()) - .saturating_add(impacts.derived_cache_bytes()) - }) - .unwrap_or(0) - + self.first_docs.capacity_bytes(self.blocks.len(), context) - } -} - -impl CompressedPostingList { - pub(crate) fn new( - blocks: LargeBinaryArray, - max_score: f32, - length: u32, - posting_tail_codec: PostingTailCodec, - block_size: usize, - positions: Option, - impacts: Option, - ) -> Self { - debug_assert!(block_size.is_power_of_two()); - Self { - max_score, - length, - blocks, - posting_tail_codec, - block_size, - positions, - impacts, - first_docs: FirstDocsState::standalone(), - } - } - - fn with_packed_first_docs(mut self, states: Arc<[OnceLock>]>, slot: usize) -> Self { - debug_assert!(slot < states.len()); - self.first_docs = FirstDocsState::Packed { states, slot }; - self - } - - /// Block sizes are validated powers of two, so per-doc hot loops derive - /// block indices with shift/mask instead of runtime division, which is - /// measurably slower in the iterator advance path. - #[inline] - pub(crate) fn block_shift(&self) -> u32 { - self.block_size.trailing_zeros() - } - - #[inline] - pub(crate) fn block_mask(&self) -> usize { - self.block_size - 1 - } - - pub fn from_batch( - batch: &RecordBatch, - max_score: f32, - length: u32, - posting_tail_codec: PostingTailCodec, - block_size: usize, - shared_position_codec: Option, - ) -> Result { - debug_assert_eq!(batch.num_rows(), 1); - let blocks = batch[POSTING_COL] - .as_list::() - .value(0) - .as_binary::() - .clone(); - let positions = if let Some(col) = batch.column_by_name(COMPRESSED_POSITION_COL) { - let bytes = bytes::Bytes::from(col.as_binary::().value(0).to_vec()); - let block_offsets = batch[POSITION_BLOCK_OFFSET_COL] - .as_list::() - .value(0) - .as_primitive::() - .values() - .to_vec(); - let codec = shared_position_codec.unwrap_or_else(|| { - parse_shared_position_codec(batch.schema_ref().metadata()) - .expect("shared position stream codec metadata should be valid") - }); - Some(CompressedPositionStorage::SharedStream( - SharedPositionStream::new(codec, block_offsets, bytes), - )) - } else { - batch.column_by_name(POSITION_COL).map(|col| { - CompressedPositionStorage::LegacyPerDoc( - col.as_list::().value(0).as_list::().clone(), - ) - }) - }; - let impacts = batch - .column_by_name(IMPACT_COL) - .map(|col| { - let entries = col.as_list::().value(0).as_binary::().clone(); - ImpactSkipData::new(entries, blocks.len()) - }) - .transpose()?; - - Ok(Self { - max_score, - length, - blocks, - posting_tail_codec, - block_size, - positions, - impacts, - first_docs: FirstDocsState::standalone(), - }) - } - - pub fn iter(&self) -> CompressedPostingListIterator { - CompressedPostingListIterator::new( - self.length as usize, - self.blocks.clone(), - self.posting_tail_codec, - self.positions.clone(), - self.block_size, - ) - } - - pub fn block_max_score(&self, block_idx: usize) -> f32 { - // 256-document blocks store no per-block max score: their impact - // skip data supplies the tight per-block bound, so callers on that - // path never reach here. Fall back to the list-level max, which is - // still a valid (looser) bound for any block. - if super::encoding::posting_block_score_prefix_len(self.block_size) == 0 { - return self.max_score; - } - let block = self.blocks.value(block_idx); - block[0..4].try_into().map(f32::from_le_bytes).unwrap() - } - - #[inline] - pub fn block_least_doc_id(&self, block_idx: usize) -> u32 { - self.block_first_docs()[block_idx] - } - - /// First doc id of every block, decoded once per cached list and shared by - /// the per-query clones. Block boundary lookups (window bounds, block - /// binary searches) are hot enough that re-reading the block headers — - /// and re-decoding the tail block — shows up in profiles. - pub(crate) fn block_first_docs(&self) -> &[u32] { - self.first_docs.get_or_init(|| { - (0..self.blocks.len()) - .map(|block_idx| { - let block = self.blocks.value(block_idx); - let remainder = self.length as usize % self.block_size; - if block_idx + 1 == self.blocks.len() && remainder > 0 { - return super::encoding::read_posting_tail_first_doc( - block, - self.posting_tail_codec, - self.block_size, - ); - } - let prefix = super::encoding::posting_block_score_prefix_len(self.block_size); - block[prefix..prefix + 4] - .try_into() - .map(u32::from_le_bytes) - .unwrap() - }) - .collect::>() - .into_boxed_slice() - }) - } - - #[cfg(test)] - fn shares_first_docs_with(&self, other: &Self) -> bool { - self.first_docs.shares_state_with(&other.first_docs) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Default)] -struct EncodedBlocks { - offsets: Vec, - bytes: Vec, -} - -impl EncodedBlocks { - fn len(&self) -> usize { - self.offsets.len() - } - - fn size(&self) -> usize { - self.offsets.capacity() * std::mem::size_of::() + self.bytes.capacity() - } - - fn push_full_block(&mut self, doc_ids: &[u32], frequencies: &[u32]) -> Result { - let start = self.bytes.len(); - self.offsets.push(start as u32); - super::encoding::encode_full_posting_block_into(doc_ids, frequencies, &mut self.bytes)?; - Ok(self.bytes.len() - start) - } - - fn block(&self, index: usize) -> &[u8] { - let (start, end) = self.block_range(index); - &self.bytes[start..end] - } - - fn block_range(&self, index: usize) -> (usize, usize) { - let start = self.offsets[index] as usize; - let end = self - .offsets - .get(index + 1) - .map(|offset| *offset as usize) - .unwrap_or(self.bytes.len()); - (start, end) - } - - fn set_block_score(&mut self, index: usize, score: f32) { - let (start, _) = self.block_range(index); - self.bytes[start..start + 4].copy_from_slice(&score.to_le_bytes()); - } - - fn append_remainder_block_with_codec( - &mut self, - doc_ids: &[u32], - frequencies: &[u32], - codec: PostingTailCodec, - block_size: usize, - ) -> Result<()> { - self.offsets.push(self.bytes.len() as u32); - super::encoding::encode_remainder_posting_block_into( - doc_ids, - frequencies, - codec, - block_size, - &mut self.bytes, - ) - } - - fn into_array(mut self) -> LargeBinaryArray { - let mut offsets = Vec::with_capacity(self.offsets.len() + 1); - offsets.extend(self.offsets.into_iter().map(i64::from)); - offsets.push(self.bytes.len() as i64); - LargeBinaryArray::new( - OffsetBuffer::new(ScalarBuffer::from(offsets)), - Buffer::from_vec(std::mem::take(&mut self.bytes)), - None, - ) - } - - fn iter(&self) -> impl Iterator { - (0..self.len()).map(|index| self.block(index)) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Default)] -struct EncodedPositionBlocks { - offsets: Vec, - bytes: Vec, -} - -impl EncodedPositionBlocks { - fn size(&self) -> usize { - self.offsets.capacity() * std::mem::size_of::() + self.bytes.capacity() - } - - fn block(&self, index: usize) -> &[u8] { - let start = self.offsets[index] as usize; - let end = self - .offsets - .get(index + 1) - .map(|offset| *offset as usize) - .unwrap_or(self.bytes.len()); - &self.bytes[start..end] - } - - fn push_encoded_block(&mut self, block: &[u8]) -> usize { - let start = self.bytes.len(); - self.offsets.push(start as u32); - self.bytes.extend_from_slice(block); - self.bytes.len() - start - } - - fn into_stream(self) -> SharedPositionStream { - SharedPositionStream::new( - PositionStreamCodec::PackedDelta, - self.offsets, - bytes::Bytes::from(self.bytes), - ) - } -} - -#[derive(Debug)] -pub struct PostingListBuilder { - with_positions: bool, - posting_tail_codec: PostingTailCodec, - encoded_blocks: Option>, - encoded_position_blocks: Option>, - tail_entries: Vec, - tail_positions: PositionBlockBuilder, - open_doc_id: Option, - open_doc_frequency: u32, - open_doc_last_position: Option, - block_size: usize, - memory_size_bytes: u32, - len: u32, -} - -pub(super) struct PostingListBatchBuilder { - schema: SchemaRef, - postings: ListBuilder, - impacts: Option>, - max_scores: Float32Builder, - lengths: UInt32Builder, - positions: BatchPositionsBuilder, - len: usize, -} - -enum BatchPositionsBuilder { - None, - Legacy(ListBuilder>), - Shared { - bytes: LargeBinaryBuilder, - block_offsets: ListBuilder, - }, -} - -struct PostingListParts<'a> { - with_positions: bool, - posting_tail_codec: PostingTailCodec, - block_size: usize, - length: usize, - encoded_blocks: EncodedBlocks, - encoded_position_blocks: EncodedPositionBlocks, - tail_entries: &'a [RawDocInfo], - tail_position_block: Option>, -} - -impl PostingListBatchBuilder { - pub fn new( - schema: SchemaRef, - with_positions: bool, - format_version: InvertedListFormatVersion, - capacity: usize, - ) -> Self { - let positions = if !with_positions { - BatchPositionsBuilder::None - } else if format_version.uses_shared_position_stream() { - BatchPositionsBuilder::Shared { - bytes: LargeBinaryBuilder::with_capacity(capacity, 0), - block_offsets: ListBuilder::with_capacity(UInt32Builder::new(), capacity), - } - } else { - BatchPositionsBuilder::Legacy(ListBuilder::with_capacity( - ListBuilder::new(LargeBinaryBuilder::new()), - capacity, - )) - }; - let impacts = schema - .field_with_name(IMPACT_COL) - .ok() - .map(|_| ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity)); - Self { - schema, - postings: ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity), - impacts, - max_scores: Float32Builder::with_capacity(capacity), - lengths: UInt32Builder::with_capacity(capacity), - positions, - len: 0, - } - } - - pub fn len(&self) -> usize { - self.len - } - - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - fn append( - &mut self, - compressed: LargeBinaryArray, - impacts: Option<&ImpactSkipData>, - max_score: f32, - length: u32, - positions: Option<&CompressedPositionStorage>, - ) -> Result<()> { - { - let values = self.postings.values(); - for index in 0..compressed.len() { - values.append_value(compressed.value(index)); - } - } - self.postings.append(true); - if let Some(impacts_builder) = &mut self.impacts { - let impacts = impacts.ok_or_else(|| { - Error::index(format!( - "impacts builder missing impact data for posting length {}", - length - )) - })?; - let values = impacts_builder.values(); - for index in 0..impacts.entries().len() { - values.append_value(impacts.entries().value(index)); - } - impacts_builder.append(true); - } - self.max_scores.append_value(max_score); - self.lengths.append_value(length); - - match &mut self.positions { - BatchPositionsBuilder::None => {} - BatchPositionsBuilder::Shared { - bytes, - block_offsets, - } => { - let positions = positions.ok_or_else(|| { - Error::index(format!( - "positions builder missing position data for posting length {}", - length - )) - })?; - let CompressedPositionStorage::SharedStream(positions) = positions else { - return Err(Error::index( - "shared positions builder received legacy positions".to_owned(), - )); - }; - bytes.append_value(positions.bytes()); - let offsets_builder = block_offsets.values(); - for &offset in positions.block_offsets() { - offsets_builder.append_value(offset); - } - block_offsets.append(true); - } - BatchPositionsBuilder::Legacy(position_lists) => { - let positions = positions.ok_or_else(|| { - Error::index(format!( - "positions builder missing position data for posting length {}", - length - )) - })?; - let CompressedPositionStorage::LegacyPerDoc(positions) = positions else { - return Err(Error::index( - "legacy positions builder received shared position stream".to_owned(), - )); - }; - let docs_builder = position_lists.values(); - for doc_idx in 0..positions.len() { - let doc_positions = positions.value(doc_idx); - let compressed_positions = doc_positions.as_binary::(); - for block_idx in 0..compressed_positions.len() { - docs_builder - .values() - .append_value(compressed_positions.value(block_idx)); - } - docs_builder.append(true); - } - position_lists.append(true); - } - } - - self.len += 1; - Ok(()) - } - - pub fn finish(&mut self) -> Result { - let mut columns = vec![ - Arc::new(self.postings.finish()) as ArrayRef, - Arc::new(self.max_scores.finish()) as ArrayRef, - Arc::new(self.lengths.finish()) as ArrayRef, - ]; - if let Some(impacts) = &mut self.impacts { - columns.push(Arc::new(impacts.finish()) as ArrayRef); - } - match &mut self.positions { - BatchPositionsBuilder::None => {} - BatchPositionsBuilder::Legacy(position_lists) => { - columns.push(Arc::new(position_lists.finish()) as ArrayRef); - } - BatchPositionsBuilder::Shared { - bytes, - block_offsets, - } => { - columns.push(Arc::new(bytes.finish()) as ArrayRef); - columns.push(Arc::new(block_offsets.finish()) as ArrayRef); - } - } - self.len = 0; - RecordBatch::try_new(self.schema.clone(), columns).map_err(Error::from) - } -} - -impl PostingListBuilder { - pub fn size(&self) -> u64 { - self.memory_size_bytes as u64 - } - - pub fn has_positions(&self) -> bool { - self.with_positions - } - - pub fn new(with_position: bool) -> Self { - Self::new_with_posting_tail_codec_and_block_size( - with_position, - current_fts_format_version().posting_tail_codec(), - LEGACY_BLOCK_SIZE, - ) - } - - pub fn new_with_posting_tail_codec( - with_position: bool, - posting_tail_codec: PostingTailCodec, - ) -> Self { - Self::new_with_posting_tail_codec_and_block_size( - with_position, - posting_tail_codec, - LEGACY_BLOCK_SIZE, - ) - } - - pub fn new_with_block_size(with_position: bool, block_size: usize) -> Self { - Self::new_with_posting_tail_codec_and_block_size( - with_position, - current_fts_format_version().posting_tail_codec(), - block_size, - ) - } - - pub fn new_with_posting_tail_codec_and_block_size( - with_position: bool, - posting_tail_codec: PostingTailCodec, - block_size: usize, - ) -> Self { - validate_block_size(block_size).expect("invalid posting list block size"); - Self { - with_positions: with_position, - posting_tail_codec, - encoded_blocks: None, - encoded_position_blocks: None, - tail_entries: Vec::new(), - tail_positions: PositionBlockBuilder::default(), - open_doc_id: None, - open_doc_frequency: 0, - open_doc_last_position: None, - block_size, - len: 0, - memory_size_bytes: 0, - } - } - - pub fn len(&self) -> usize { - self.len as usize - } - - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - pub fn iter(&self) -> std::vec::IntoIter<(u32, u32, Option>)> { - self.collect_entries().into_iter() - } - - pub fn for_each_entry( - &self, - mut visit: impl FnMut(u32, u32, Option>) -> std::result::Result<(), E>, - ) -> std::result::Result<(), E> { - let mut doc_ids = Vec::with_capacity(self.block_size); - let mut frequencies = Vec::with_capacity(self.block_size); - let mut decoded_positions = Vec::new(); - let mut position_block_index = 0usize; - - if let Some(encoded_blocks) = self.encoded_blocks.as_deref() { - for block in encoded_blocks.iter() { - doc_ids.clear(); - frequencies.clear(); - super::encoding::decode_full_posting_block( - block, - &mut doc_ids, - &mut frequencies, - self.block_size, - ); - decoded_positions.clear(); - if self.with_positions { - let position_blocks = self - .encoded_position_blocks - .as_deref() - .expect("positions must exist for posting list"); - super::encoding::decode_position_stream_block( - position_blocks.block(position_block_index), - &frequencies, - PositionStreamCodec::PackedDelta, - &mut decoded_positions, - ) - .expect("position stream decoding should succeed"); - position_block_index += 1; - } - let mut offset = 0usize; - for (doc_id, frequency) in doc_ids.iter().copied().zip(frequencies.iter().copied()) - { - let positions = self.with_positions.then(|| { - let end = offset + frequency as usize; - let doc_positions = decoded_positions[offset..end].to_vec(); - offset = end; - doc_positions - }); - visit(doc_id, frequency, positions)?; - } - } - } - - let mut decoded_tail_positions = Vec::new(); - if self.with_positions && !self.tail_entries.is_empty() { - let tail_frequencies = self - .tail_entries - .iter() - .map(|entry| entry.frequency) - .collect::>(); - self.tail_positions - .decode_into(tail_frequencies.as_slice(), &mut decoded_tail_positions) - .expect("tail position stream decoding should succeed"); - } - let mut tail_offset = 0usize; - for entry in &self.tail_entries { - let positions = self.with_positions.then(|| { - let end = tail_offset + entry.frequency as usize; - let doc_positions = decoded_tail_positions[tail_offset..end].to_vec(); - tail_offset = end; - doc_positions - }); - visit(entry.doc_id, entry.frequency, positions)?; - } - - Ok(()) - } - - pub fn add(&mut self, doc_id: u32, term_positions: PositionRecorder) { - debug_assert!( - self.open_doc_id.is_none(), - "cannot add closed doc while a positions doc is still open" - ); - let tail_entries_capacity_before = self.tail_entries.capacity(); - self.tail_entries - .push(RawDocInfo::new(doc_id, term_positions.len())); - let tail_entries_capacity_after = self.tail_entries.capacity(); - if tail_entries_capacity_after > tail_entries_capacity_before { - self.add_memory_bytes( - (tail_entries_capacity_after - tail_entries_capacity_before) - * std::mem::size_of::(), - ); - } - if let PositionRecorder::Position(positions_in_doc) = term_positions { - debug_assert!(self.with_positions); - let old_size = self.tail_positions.size(); - self.tail_positions - .append_doc_positions(positions_in_doc.as_slice()) - .expect("position stream encoding should succeed"); - self.adjust_tail_positions_size(old_size); - } - self.len += 1; - - if self.tail_entries.len() == self.block_size { - self.flush_tail_block() - .expect("posting list block compression should succeed"); - } - } - - pub fn add_occurrence(&mut self, doc_id: u32, position: u32) -> Result { - if !self.with_positions { - return Err(Error::index( - "cannot append streamed positions to a posting list without positions".to_owned(), - )); - } - - match self.open_doc_id { - Some(open_doc_id) if open_doc_id == doc_id => { - let old_size = self.tail_positions.size(); - self.tail_positions - .append_position(position, self.open_doc_last_position)?; - self.adjust_tail_positions_size(old_size); - self.open_doc_frequency += 1; - self.open_doc_last_position = Some(position); - Ok(false) - } - Some(open_doc_id) => Err(Error::index(format!( - "posting list received doc {} before finishing open doc {}", - doc_id, open_doc_id - ))), - None => { - let old_size = self.tail_positions.size(); - self.tail_positions.append_position(position, None)?; - self.adjust_tail_positions_size(old_size); - self.open_doc_id = Some(doc_id); - self.open_doc_frequency = 1; - self.open_doc_last_position = Some(position); - self.len += 1; - Ok(true) - } - } - } - - pub fn finish_open_doc(&mut self, doc_id: u32) -> Result<()> { - if !self.with_positions { - return Ok(()); - } - match self.open_doc_id { - Some(open_doc_id) if open_doc_id == doc_id => { - let tail_entries_capacity_before = self.tail_entries.capacity(); - self.tail_entries - .push(RawDocInfo::new(doc_id, self.open_doc_frequency)); - let tail_entries_capacity_after = self.tail_entries.capacity(); - if tail_entries_capacity_after > tail_entries_capacity_before { - self.add_memory_bytes( - (tail_entries_capacity_after - tail_entries_capacity_before) - * std::mem::size_of::(), - ); - } - self.open_doc_id = None; - self.open_doc_frequency = 0; - self.open_doc_last_position = None; - if self.tail_entries.len() == self.block_size { - self.flush_tail_block()?; - } - Ok(()) - } - Some(open_doc_id) => Err(Error::index(format!( - "attempted to finish doc {} while doc {} is still open", - doc_id, open_doc_id - ))), - None => Ok(()), - } - } - - fn collect_entries(&self) -> Vec<(u32, u32, Option>)> { - let mut entries = Vec::with_capacity(self.len()); - self.for_each_entry(|doc_id, frequency, positions| { - entries.push((doc_id, frequency, positions)); - Ok::<(), ()>(()) - }) - .expect("collecting posting list entries should not fail"); - entries - } - - fn encoded_blocks_mut(&mut self) -> &mut EncodedBlocks { - if self.encoded_blocks.is_none() { - self.encoded_blocks = Some(Box::default()); - self.add_memory_bytes(std::mem::size_of::()); - } - self.encoded_blocks - .as_deref_mut() - .expect("encoded blocks must exist") - } - - fn encoded_position_blocks_mut(&mut self) -> &mut EncodedPositionBlocks { - if self.encoded_position_blocks.is_none() { - self.encoded_position_blocks = Some(Box::default()); - self.add_memory_bytes(std::mem::size_of::()); - } - self.encoded_position_blocks - .as_deref_mut() - .expect("encoded position blocks must exist") - } - - fn flush_tail_block(&mut self) -> Result<()> { - if self.tail_entries.is_empty() { - return Ok(()); - } - debug_assert!( - self.open_doc_id.is_none(), - "cannot flush a posting block while a document is still open" - ); - debug_assert_eq!(self.tail_entries.len(), self.block_size); - let doc_ids = self - .tail_entries - .iter() - .map(|entry| entry.doc_id) - .collect::>(); - let frequencies = self - .tail_entries - .iter() - .map(|entry| entry.frequency) - .collect::>(); - let encoded_blocks_size_before = self - .encoded_blocks - .as_ref() - .map(|encoded_blocks| encoded_blocks.size()) - .unwrap_or(0usize); - self.encoded_blocks_mut() - .push_full_block(&doc_ids, &frequencies)?; - let encoded_blocks_size_after = self - .encoded_blocks - .as_ref() - .map(|encoded_blocks| encoded_blocks.size()) - .unwrap_or(0usize); - if encoded_blocks_size_after > encoded_blocks_size_before { - self.add_memory_bytes(encoded_blocks_size_after - encoded_blocks_size_before); - } - if self.with_positions { - let encoded_positions_size_before = self - .encoded_position_blocks - .as_ref() - .map(|encoded| encoded.size()) - .unwrap_or(0usize); - let released_tail_positions_bytes = self.tail_positions.size(); - let tail_position_block = std::mem::take(&mut self.tail_positions).finish(); - self.encoded_position_blocks_mut() - .push_encoded_block(tail_position_block.as_slice()); - let encoded_positions_size_after = self - .encoded_position_blocks - .as_ref() - .map(|encoded| encoded.size()) - .unwrap_or(0usize); - if released_tail_positions_bytes > 0 { - self.subtract_memory_bytes(released_tail_positions_bytes); - } - if encoded_positions_size_after > encoded_positions_size_before { - self.add_memory_bytes(encoded_positions_size_after - encoded_positions_size_before); - } - } - self.tail_entries.clear(); - Ok(()) - } - - fn adjust_tail_positions_size(&mut self, old_size: usize) { - let new_size = self.tail_positions.size(); - if new_size > old_size { - self.add_memory_bytes(new_size - old_size); - } else if old_size > new_size { - self.subtract_memory_bytes(old_size - new_size); - } - } - - fn add_memory_bytes(&mut self, bytes: usize) { - self.memory_size_bytes = self - .memory_size_bytes - .checked_add( - u32::try_from(bytes).expect("posting list memory size delta overflowed u32"), - ) - .expect("posting list memory size overflowed u32"); - } - - fn subtract_memory_bytes(&mut self, bytes: usize) { - self.memory_size_bytes = self - .memory_size_bytes - .checked_sub( - u32::try_from(bytes).expect("posting list memory size delta overflowed u32"), - ) - .expect("posting list memory size underflowed u32"); - } - - fn build_position_columns( - positions: Option, - ) -> Result> { - let Some(positions) = positions else { - return Ok(Vec::new()); - }; - match positions { - CompressedPositionStorage::LegacyPerDoc(positions) => { - Ok(vec![Arc::new(ListArray::try_new( - Arc::new(Field::new("item", positions.data_type().clone(), true)), - OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, positions.len() as i32])), - Arc::new(positions) as ArrayRef, - None, - )?) as ArrayRef]) - } - CompressedPositionStorage::SharedStream(positions) => { - let mut columns = Vec::with_capacity(2); - columns.push( - Arc::new(LargeBinaryArray::from(vec![Some(positions.bytes())])) as ArrayRef, - ); - - let mut offsets_builder = ListBuilder::new(UInt32Builder::new()); - for &offset in positions.block_offsets() { - offsets_builder.values().append_value(offset); - } - offsets_builder.append(true); - columns.push(Arc::new(offsets_builder.finish()) as ArrayRef); - Ok(columns) - } - } - } - - fn build_batch( - self, - compressed: LargeBinaryArray, - impacts: Option, - max_score: f32, - schema: SchemaRef, - positions: Option, - ) -> Result { - let length = self.len(); - let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, compressed.len() as i32])); - let mut columns = vec![ - Arc::new(ListArray::try_new( - Arc::new(Field::new("item", datatypes::DataType::LargeBinary, true)), - offsets, - Arc::new(compressed), - None, - )?) as ArrayRef, - Arc::new(Float32Array::from_iter_values(std::iter::once(max_score))) as ArrayRef, - Arc::new(UInt32Array::from_iter_values(std::iter::once( - length as u32, - ))) as ArrayRef, - ]; - if schema.field_with_name(IMPACT_COL).is_ok() { - let impacts = impacts.ok_or_else(|| { - Error::index(format!( - "impact column requested without impact data for posting length {}", - length - )) - })?; - let impact_offsets = - OffsetBuffer::new(ScalarBuffer::from(vec![0, impacts.entries().len() as i32])); - columns.push(Arc::new(ListArray::try_new( - Arc::new(Field::new("item", datatypes::DataType::LargeBinary, true)), - impact_offsets, - Arc::new(impacts.entries().clone()), - None, - )?) as ArrayRef); - } - columns.extend(Self::build_position_columns(positions)?); - - let batch = RecordBatch::try_new(schema, columns)?; - Ok(batch) - } - - fn build_legacy_positions(&self) -> Result { - let mut positions_builder = ListBuilder::new(LargeBinaryBuilder::new()); - self.for_each_entry(|_doc_id, frequency, positions| { - let positions = positions.ok_or_else(|| { - Error::index(format!( - "legacy position writer missing positions for frequency {}", - frequency - )) - })?; - let compressed = super::encoding::compress_positions(positions.as_slice())?; - for block_idx in 0..compressed.len() { - positions_builder - .values() - .append_value(compressed.value(block_idx)); - } - positions_builder.append(true); - Ok::<(), Error>(()) - })?; - Ok(positions_builder.finish()) - } - - pub(super) fn append_to_batch_with_docs( - self, - docs: &DocSet, - batch_builder: &mut PostingListBatchBuilder, - format_version: InvertedListFormatVersion, - ) -> Result<()> { - let legacy_positions = - if self.with_positions && !format_version.uses_shared_position_stream() { - Some(self.build_legacy_positions()?) - } else { - None - }; - let Self { - with_positions, - posting_tail_codec, - encoded_blocks, - encoded_position_blocks, - tail_entries, - tail_positions, - open_doc_id, - open_doc_frequency, - open_doc_last_position, - block_size, - len, - .. - } = self; - debug_assert!(open_doc_id.is_none()); - debug_assert_eq!(open_doc_frequency, 0); - debug_assert!(open_doc_last_position.is_none()); - let parts = PostingListParts { - with_positions, - posting_tail_codec, - block_size, - length: len as usize, - encoded_blocks: encoded_blocks - .map(|encoded_blocks| *encoded_blocks) - .unwrap_or_default(), - encoded_position_blocks: encoded_position_blocks - .map(|encoded_positions| *encoded_positions) - .unwrap_or_default(), - tail_entries: tail_entries.as_slice(), - tail_position_block: with_positions.then(|| tail_positions.finish()), - }; - let (compressed, shared_positions, max_score, impacts) = - Self::build_compressed_with_scores_from_parts(parts, docs)?; - let positions = match legacy_positions { - Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), - None => shared_positions.map(CompressedPositionStorage::SharedStream), - }; - batch_builder.append( - compressed, - Some(&impacts), - max_score, - len, - positions.as_ref(), - ) - } - - fn extend_tail_components( - tail_entries: &[RawDocInfo], - doc_ids: &mut Vec, - frequencies: &mut Vec, - ) { - doc_ids.clear(); - frequencies.clear(); - doc_ids.extend(tail_entries.iter().map(|entry| entry.doc_id)); - frequencies.extend(tail_entries.iter().map(|entry| entry.frequency)); - } - - fn build_compressed_with_scores_from_parts( - parts: PostingListParts<'_>, - docs: &DocSet, - ) -> Result<( - LargeBinaryArray, - Option, - f32, - ImpactSkipData, - )> { - let PostingListParts { - with_positions, - posting_tail_codec, - length, - block_size, - mut encoded_blocks, - mut encoded_position_blocks, - tail_entries, - tail_position_block, - } = parts; - let avgdl = docs.average_length(); - let idf_scale = idf(length, docs.len()) * (K1 + 1.0); - let mut max_score = f32::MIN; - let mut doc_ids = Vec::with_capacity(block_size); - let mut frequencies = Vec::with_capacity(block_size); - let mut impact_block = Vec::with_capacity(block_size); - let mut impact_builder = - ImpactSkipDataBuilder::with_capacity(length.div_ceil(block_size), block_size); - - for index in 0..encoded_blocks.len() { - let block = encoded_blocks.block(index); - doc_ids.clear(); - frequencies.clear(); - super::encoding::decode_full_posting_block( - block, - &mut doc_ids, - &mut frequencies, - block_size, - ); - let block_score = compute_block_score_and_impact_block( - docs, - avgdl, - idf_scale, - doc_ids.iter().copied(), - frequencies.iter().copied(), - &mut impact_block, - ); - impact_builder.append_block(impact_block.as_slice())?; - max_score = max_score.max(block_score); - if super::encoding::posting_block_score_prefix_len(block_size) > 0 { - encoded_blocks.set_block_score(index, block_score); - } - } - - if !tail_entries.is_empty() { - Self::extend_tail_components(tail_entries, &mut doc_ids, &mut frequencies); - let block_score = compute_block_score_and_impact_block( - docs, - avgdl, - idf_scale, - doc_ids.iter().copied(), - frequencies.iter().copied(), - &mut impact_block, - ); - impact_builder.append_block(impact_block.as_slice())?; - max_score = max_score.max(block_score); - encoded_blocks.append_remainder_block_with_codec( - doc_ids.as_slice(), - frequencies.as_slice(), - posting_tail_codec, - block_size, - )?; - if super::encoding::posting_block_score_prefix_len(block_size) > 0 { - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); - } - if with_positions { - encoded_position_blocks.push_encoded_block( - tail_position_block - .as_deref() - .expect("tail position block must exist for postings with positions"), - ); - } - } - - let impacts = impact_builder.finish()?; - Ok(( - encoded_blocks.into_array(), - with_positions.then(|| encoded_position_blocks.into_stream()), - max_score, - impacts, - )) - } - - #[allow(clippy::too_many_arguments)] - fn build_compressed_with_block_scores_from_parts( - with_positions: bool, - posting_tail_codec: PostingTailCodec, - block_size: usize, - mut encoded_blocks: EncodedBlocks, - mut encoded_position_blocks: EncodedPositionBlocks, - tail_entries: &[RawDocInfo], - tail_position_block: Option>, - mut block_max_scores: impl Iterator, - ) -> Result<(LargeBinaryArray, Option, f32)> { - let has_score_prefix = super::encoding::posting_block_score_prefix_len(block_size) > 0; - let mut max_score = f32::MIN; - let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); - let mut frequencies = Vec::with_capacity(BLOCK_SIZE); - - for index in 0..encoded_blocks.len() { - let block_score = block_max_scores - .next() - .ok_or_else(|| Error::index("missing block max score".to_owned()))?; - max_score = max_score.max(block_score); - if has_score_prefix { - encoded_blocks.set_block_score(index, block_score); - } - } - - if !tail_entries.is_empty() { - let block_score = block_max_scores - .next() - .ok_or_else(|| Error::index("missing tail block max score".to_owned()))?; - max_score = max_score.max(block_score); - Self::extend_tail_components(tail_entries, &mut doc_ids, &mut frequencies); - encoded_blocks.append_remainder_block_with_codec( - doc_ids.as_slice(), - frequencies.as_slice(), - posting_tail_codec, - block_size, - )?; - if has_score_prefix { - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); - } - if with_positions { - encoded_position_blocks.push_encoded_block( - tail_position_block - .as_deref() - .expect("tail position block must exist for postings with positions"), - ); - } - } - - Ok(( - encoded_blocks.into_array(), - with_positions.then(|| encoded_position_blocks.into_stream()), - max_score, - )) - } - - pub fn to_batch(self, block_max_scores: Vec) -> Result { - let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( - self.posting_tail_codec, - self.block_size, - )?; - let schema = inverted_list_schema_for_version_with_block_size_and_impacts( - self.has_positions(), - format_version, - self.block_size, - false, - ); - let legacy_positions = - if self.with_positions && !format_version.uses_shared_position_stream() { - Some(self.build_legacy_positions()?) - } else { - None - }; - let Self { - with_positions, - posting_tail_codec, - encoded_blocks, - encoded_position_blocks, - tail_entries, - tail_positions, - open_doc_id, - open_doc_frequency, - open_doc_last_position, - block_size, - len, - .. - } = self; - debug_assert!(open_doc_id.is_none()); - debug_assert_eq!(open_doc_frequency, 0); - debug_assert!(open_doc_last_position.is_none()); - let (compressed, shared_positions, max_score) = - Self::build_compressed_with_block_scores_from_parts( - with_positions, - posting_tail_codec, - block_size, - encoded_blocks - .map(|encoded_blocks| *encoded_blocks) - .unwrap_or_default(), - encoded_position_blocks - .map(|encoded_positions| *encoded_positions) - .unwrap_or_default(), - tail_entries.as_slice(), - with_positions.then(|| tail_positions.finish()), - block_max_scores.into_iter(), - )?; - let builder = Self { - with_positions, - posting_tail_codec, - encoded_blocks: None, - encoded_position_blocks: None, - tail_entries: Vec::new(), - tail_positions: PositionBlockBuilder::default(), - open_doc_id: None, - open_doc_frequency: 0, - open_doc_last_position: None, - block_size, - memory_size_bytes: 0, - len, - }; - let positions = match legacy_positions { - Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), - None => shared_positions.map(CompressedPositionStorage::SharedStream), - }; - builder.build_batch(compressed, None, max_score, schema, positions) - } - - pub fn to_batch_with_docs(self, docs: &DocSet, schema: SchemaRef) -> Result { - let format_version = parse_format_version_from_metadata(schema.metadata())?; - let legacy_positions = - if self.with_positions && !format_version.uses_shared_position_stream() { - Some(self.build_legacy_positions()?) - } else { - None - }; - let Self { - with_positions, - posting_tail_codec, - encoded_blocks, - encoded_position_blocks, - tail_entries, - tail_positions, - open_doc_id, - open_doc_frequency, - open_doc_last_position, - block_size, - len, - .. - } = self; - debug_assert!(open_doc_id.is_none()); - debug_assert_eq!(open_doc_frequency, 0); - debug_assert!(open_doc_last_position.is_none()); - let parts = PostingListParts { - with_positions, - posting_tail_codec, - block_size, - length: len as usize, - encoded_blocks: encoded_blocks - .map(|encoded_blocks| *encoded_blocks) - .unwrap_or_default(), - encoded_position_blocks: encoded_position_blocks - .map(|encoded_positions| *encoded_positions) - .unwrap_or_default(), - tail_entries: tail_entries.as_slice(), - tail_position_block: with_positions.then(|| tail_positions.finish()), - }; - let (compressed, shared_positions, max_score, impacts) = - Self::build_compressed_with_scores_from_parts(parts, docs)?; - let builder = Self { - with_positions, - posting_tail_codec, - encoded_blocks: None, - encoded_position_blocks: None, - tail_entries: Vec::new(), - tail_positions: PositionBlockBuilder::default(), - open_doc_id: None, - open_doc_frequency: 0, - open_doc_last_position: None, - block_size, - memory_size_bytes: 0, - len, - }; - let positions = match legacy_positions { - Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), - None => shared_positions.map(CompressedPositionStorage::SharedStream), - }; - builder.build_batch(compressed, Some(impacts), max_score, schema, positions) - } - - pub fn remap(&mut self, removed: &[u32]) { - let mut cursor = 0; - let mut new_builder = Self::new_with_posting_tail_codec_and_block_size( - self.has_positions(), - self.posting_tail_codec, - self.block_size, - ); - for (doc_id, freq, positions) in self.iter() { - while cursor < removed.len() && removed[cursor] < doc_id { - cursor += 1; - } - if cursor < removed.len() && removed[cursor] == doc_id { - continue; - } - let positions = match positions { - Some(positions) => PositionRecorder::Position(positions.into()), - None => PositionRecorder::Count(freq), - }; - new_builder.add(doc_id - cursor as u32, positions); - } - - *self = new_builder; - } -} - -fn compute_block_score_and_impact_block( - docs: &DocSet, - avgdl: f32, - idf_scale: f32, - doc_ids: impl Iterator, - frequencies: impl Iterator, - impact_block: &mut Vec<(u32, u32, u32)>, -) -> f32 { - impact_block.clear(); - let mut block_max_score = f32::MIN; - for (doc_id, freq) in doc_ids.zip(frequencies) { - let doc_len = docs.num_tokens(doc_id); - let doc_norm = K1 * (1.0 - B + B * doc_len as f32 / avgdl); - let freq_f32 = freq as f32; - let score = freq_f32 / (freq_f32 + doc_norm); - block_max_score = block_max_score.max(score); - impact_block.push((doc_id, freq, doc_len)); - } - block_max_score * idf_scale -} - -#[derive(Debug, Clone, DeepSizeOf, Copy)] -pub enum DocInfo { - Located(LocatedDocInfo), - Raw(RawDocInfo), -} - -impl DocInfo { - pub fn doc_id(&self) -> u64 { - match self { - Self::Raw(info) => info.doc_id as u64, - Self::Located(info) => info.row_id, - } - } - - pub fn frequency(&self) -> u32 { - match self { - Self::Raw(info) => info.frequency, - Self::Located(info) => info.frequency as u32, - } - } -} - -impl Eq for DocInfo {} - -impl PartialEq for DocInfo { - fn eq(&self, other: &Self) -> bool { - self.doc_id() == other.doc_id() - } -} - -impl PartialOrd for DocInfo { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for DocInfo { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.doc_id().cmp(&other.doc_id()) - } -} - -#[derive(Debug, Clone, Default, DeepSizeOf, Copy)] -pub struct LocatedDocInfo { - pub row_id: u64, - pub frequency: f32, -} - -impl LocatedDocInfo { - pub fn new(row_id: u64, frequency: f32) -> Self { - Self { row_id, frequency } - } -} - -impl Eq for LocatedDocInfo {} - -impl PartialEq for LocatedDocInfo { - fn eq(&self, other: &Self) -> bool { - self.row_id == other.row_id - } -} - -impl PartialOrd for LocatedDocInfo { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for LocatedDocInfo { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.row_id.cmp(&other.row_id) - } -} - -#[derive(Debug, Clone, Default, DeepSizeOf, Copy)] -pub struct RawDocInfo { - pub doc_id: u32, - pub frequency: u32, -} - -impl RawDocInfo { - pub fn new(doc_id: u32, frequency: u32) -> Self { - Self { doc_id, frequency } - } -} - -impl Eq for RawDocInfo {} - -impl PartialEq for RawDocInfo { - fn eq(&self, other: &Self) -> bool { - self.doc_id == other.doc_id - } -} - -impl PartialOrd for RawDocInfo { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for RawDocInfo { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.doc_id.cmp(&other.doc_id) - } -} - -/// Lucene SmallFloat-style document-length quantization for 256-document-block scoring and impact -/// norms: a 4-mantissa-bit float-like byte code. Values 0-7 are exact; larger -/// values keep their top four significand bits (relative error <= 6.25%) and -/// decode to their bucket floor. The floor only ever shortens a doc, so impact -/// bounds remain conservative for exact scoring as well as quantized scoring. -pub(super) fn quantize_doc_length(value: u32) -> u8 { - let num_bits = 32 - value.leading_zeros(); - if num_bits < 4 { - value as u8 - } else { - let shift = num_bits - 4; - (((value >> shift) as u8) & 0x07) | (((shift + 1) as u8) << 3) - } -} - -#[inline] -pub(super) fn dequantize_doc_length(code: u8) -> u32 { - DEQUANTIZED_DOC_LENGTHS[code as usize] -} - -pub(super) static DEQUANTIZED_DOC_LENGTHS: [u32; 256] = build_dequantized_doc_lengths(); - -const fn build_dequantized_doc_lengths() -> [u32; 256] { - let mut table = [0u32; 256]; - let mut code = 0usize; - while code < 256 { - let bits = (code & 0x07) as u64; - let shift = (code >> 3) as i64 - 1; - let decoded = if shift < 0 { - bits - } else { - (bits | 0x08) << shift - }; - // Codes past the largest u32 encoding are never produced; saturate so - // the table stays total. - table[code] = if decoded > u32::MAX as u64 { - u32::MAX - } else { - decoded as u32 - }; - code += 1; - } - table -} - -#[derive(Debug, Clone)] -enum NumTokens { - Owned(Vec), - Shared(ScalarBuffer), -} - -impl Default for NumTokens { - fn default() -> Self { - Self::Owned(Vec::new()) - } -} - -impl std::ops::Deref for NumTokens { - type Target = [u32]; - - fn deref(&self) -> &Self::Target { - match self { - Self::Owned(values) => values, - Self::Shared(values) => values, - } - } -} - -impl DeepSizeOf for NumTokens { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::Owned(values) => values.deep_size_of_children(context), - Self::Shared(values) => values.deep_size_of_children(context), - } - } -} - -impl NumTokens { - fn with_capacity(capacity: usize) -> Self { - Self::Owned(Vec::with_capacity(capacity)) - } - - fn into_owned(self) -> Vec { - match self { - Self::Owned(values) => values, - Self::Shared(values) => values.to_vec(), - } - } - - fn push(&mut self, value: u32) { - match self { - Self::Owned(values) => values.push(value), - Self::Shared(values) => { - let mut owned = values.to_vec(); - owned.push(value); - *self = Self::Owned(owned); - } - } - } - - fn memory_size(&self) -> usize { - match self { - Self::Owned(values) => values.capacity() * std::mem::size_of::(), - Self::Shared(values) => values.inner().capacity(), - } - } -} - -// DocSet is a mapping from row ids to the number of tokens in the document -// It's used to sort the documents by the bm25 score -#[derive(Debug, Clone, Default)] -pub struct DocSet { - row_ids: Vec, - num_tokens: NumTokens, - // 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)>, - - total_tokens: u64, - - // 256-document-block partitions score with quantized document lengths: the - // flag is set at partition load and the byte-norm slab bakes lazily on - // first scoring use (shared by clones of the loaded set). 128-block - // partitions never set the flag and keep exact scoring. - scoring_quantized: bool, - norms: Arc>>, -} - -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 - .get() - .map(|slab| std::mem::size_of_val(slab.as_ref())) - .unwrap_or(0) - } -} - -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 - // scoring path (which constructs a DocSet via - // [`Self::from_num_tokens_only`]) still reports the right doc - // count. - self.num_tokens.len() - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// True iff the per-doc `row_id` array is populated. The - /// deferred-row_id scoring path constructs DocSets with the array - /// left empty so wand can skip the load; callers that need to do - /// row_id lookups in the inner loop must check this and fall back - /// to async resolution otherwise. - #[inline] - pub fn has_row_ids(&self) -> bool { - !self.row_ids.is_empty() - } - - pub fn iter(&self) -> impl Iterator { - self.row_ids.iter().zip(self.num_tokens.iter()) - } - - pub fn row_id(&self, doc_id: u32) -> u64 { - 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. - /// - /// 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 + '_ { - if self.inv.is_empty() { - // in legacy format, the row id is doc id (one document per row) - let found = self.row_ids.binary_search(&row_id).is_ok(); - Either::Left(found.then_some(row_id).into_iter()) - } else { - // `inv` is sorted by row_id, so the entries sharing this key form a - // contiguous run; yield the doc_id of each. - let lo = self.inv.partition_point(|entry| entry.0 < row_id); - let hi = self.inv.partition_point(|entry| entry.0 <= row_id); - Either::Right(self.inv[lo..hi].iter().map(|entry| entry.1 as u64)) - } - } - pub fn total_tokens_num(&self) -> u64 { - self.total_tokens - } - - #[inline] - pub fn average_length(&self) -> f32 { - self.total_tokens as f32 / self.len() as f32 - } - - pub fn calculate_block_max_scores<'a>( - &self, - doc_ids: impl Iterator, - freqs: impl Iterator, - ) -> Vec { - self.calculate_block_max_scores_with_block_size(doc_ids, freqs, LEGACY_BLOCK_SIZE) - } - - pub fn calculate_block_max_scores_with_block_size<'a>( - &self, - doc_ids: impl Iterator, - freqs: impl Iterator, - block_size: usize, - ) -> Vec { - validate_block_size(block_size).expect("invalid posting list block size"); - let avgdl = self.average_length(); - let length = doc_ids.size_hint().0; - let num_blocks = length.div_ceil(block_size); - let mut block_max_scores = Vec::with_capacity(num_blocks); - let idf_scale = idf(length, self.len()) * (K1 + 1.0); - let mut max_score = f32::MIN; - for (i, (doc_id, freq)) in doc_ids.zip(freqs).enumerate() { - let doc_norm = K1 * (1.0 - B + B * self.num_tokens(*doc_id) as f32 / avgdl); - let freq = *freq as f32; - let score = freq / (freq + doc_norm); - if score > max_score { - max_score = score; - } - if (i + 1) % block_size == 0 { - max_score *= idf_scale; - block_max_scores.push(max_score); - max_score = f32::MIN; - } - } - if !length.is_multiple_of(block_size) { - max_score *= idf_scale; - block_max_scores.push(max_score); - } - block_max_scores - } - - pub fn to_batch(&self) -> Result { - 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 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), columns)?; - Ok(batch) - } - - pub async fn load( - reader: Arc, - is_legacy: bool, - frag_reuse_index: Option>, - ) -> Result { - 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::(); - 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; - /// `row_ids` and `inv` are left empty. Used by the deferred-row_id - /// scoring path: wand checks `has_row_ids()` to skip `row_id` / - /// `num_tokens_by_row_id` calls, and the per-partition caller - /// resolves doc_id → row_id for the surviving top-K post-wand. - pub fn from_num_tokens_only(num_tokens_col: &arrow_array::UInt32Array) -> Self { - let total_tokens = num_tokens_col.values().iter().map(|&n| n as u64).sum(); - Self::from_cached_num_tokens(num_tokens_col, total_tokens) - } - - /// Build a zero-copy num-tokens-only view from an Arrow column and its - /// already-computed total. The caller must guarantee that `total_tokens` - /// is the sum of `num_tokens_col`. - pub(crate) fn from_cached_num_tokens( - num_tokens_col: &arrow_array::UInt32Array, - total_tokens: u64, - ) -> Self { - Self { - row_ids: Vec::new(), - num_tokens: NumTokens::Shared(num_tokens_col.values().clone()), - doc_indices: Vec::new(), - inv: Vec::new(), - total_tokens, - scoring_quantized: false, - norms: Arc::new(std::sync::OnceLock::new()), - } - } - - /// Build a `DocSet` from already-loaded `row_id` and `num_tokens` - /// Arrow columns without re-reading either column. - pub fn from_columns( - row_id_col: &UInt64Array, - num_tokens_col: &arrow_array::UInt32Array, - 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 - .values() - .iter() - .filter_map(|id| { - if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - frag_reuse_index_ref.remap_row_id(*id) - } else { - Some(*id) - } - }) - .zip(num_tokens_col.values().iter()) - .sorted_unstable_by_key(|x| x.0) - .unzip(); - - let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); - return Ok(Self { - row_ids, - num_tokens: NumTokens::Owned(num_tokens), - doc_indices, - inv: Vec::new(), - total_tokens, - scoring_quantized: false, - norms: Arc::new(std::sync::OnceLock::new()), - }); - } - - // If frag reuse happened, remap the row_ids through it. Crucially we - // must NOT drop the rows the reuse index deleted, because the posting - // lists reference doc_ids *positionally* (a doc_id is an index into - // these arrays, fixed at build time). Dropping deleted rows would - // renumber every later doc_id and desync the posting lists, so wand - // would index `num_tokens`/`row_ids` out of bounds or score the wrong - // doc. Instead we tombstone deleted rows in place: their slot survives - // (so doc_ids stay aligned with the posting lists) carrying - // `RowAddress::TOMBSTONE_ROW`, which wand skips, and they are left out - // of `inv` so a row_id lookup never resolves to a deleted doc. The - // heavyweight physical remap (`DocSet::remap`) is what actually - // renumbers and compacts; this load-time path only has to stay - // consistent until then. - if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - let mut row_ids = Vec::with_capacity(row_id_col.len()); - let num_tokens = num_tokens_col.values().to_vec(); - let mut inv = Vec::with_capacity(row_id_col.len()); - for (doc_id, row_id) in row_id_col.values().iter().enumerate() { - match frag_reuse_index_ref.remap_row_id(*row_id) { - Some(new_row_id) => { - row_ids.push(new_row_id); - inv.push((new_row_id, doc_id as u32)); - } - None => { - // Deleted: keep the slot (doc_ids must not shift) but - // tombstone it and leave it out of `inv`. - row_ids.push(RowAddress::TOMBSTONE_ROW); - } - } - } - inv.sort_unstable_by_key(|entry| entry.0); - - let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); - return Ok(Self { - row_ids, - num_tokens: NumTokens::Owned(num_tokens), - doc_indices, - inv, - total_tokens, - scoring_quantized: false, - norms: Arc::new(std::sync::OnceLock::new()), - }); - } - - let row_ids = row_id_col.values().to_vec(); - let num_tokens = num_tokens_col.values().to_vec(); - let mut inv: Vec<(u64, u32)> = row_ids - .iter() - .enumerate() - .map(|(doc_id, row_id)| (*row_id, doc_id as u32)) - .collect(); - if !row_ids.is_sorted() { - inv.sort_unstable_by_key(|entry| entry.0); - } - let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); - Ok(Self { - row_ids, - num_tokens: NumTokens::Owned(num_tokens), - doc_indices, - inv, - total_tokens, - scoring_quantized: false, - norms: Arc::new(std::sync::OnceLock::new()), - }) - } - - // remap the row ids to the new row ids - // returns the removed doc ids - pub fn remap(&mut self, mapping: &RowAddrRemap) -> Vec { - let mut removed = Vec::new(); - 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, NumTokens::with_capacity(len)).into_owned(); - 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() { - match mapping.get(row_id) { - 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) => { - removed.push(doc_id as u32); - } - 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; - } - } - } - removed - } - - #[inline] - pub fn num_tokens(&self, doc_id: u32) -> u32 { - self.num_tokens[doc_id as usize] - } - - /// Enable quantized document-length scoring for 256-document-block partitions. - pub fn set_quantized_scoring(&mut self, quantized: bool) { - self.scoring_quantized = quantized; - } - - /// The quantized document-length slab when this set scores quantized, - /// baked on first use; `None` for exact-scoring sets. - pub fn scoring_norms(&self) -> Option<&[u8]> { - if !self.scoring_quantized { - return None; - } - Some( - self.norms - .get_or_init(|| { - self.num_tokens - .iter() - .map(|&n| quantize_doc_length(n)) - .collect() - }) - .as_ref(), - ) - } - - /// Document length as scoring sees it: the quantized bucket floor for - /// 256-document-block partitions, the exact value otherwise. - #[inline] - pub fn scoring_num_tokens(&self, doc_id: u32) -> u32 { - match self.scoring_norms() { - Some(norms) => dequantize_doc_length(norms[doc_id as usize]), - None => self.num_tokens[doc_id as usize], - } - } - - // this can be used only if it's a legacy format, - // which store the sorted row ids so that we can use binary search - #[inline] - pub fn num_tokens_by_row_id(&self, row_id: u64) -> u32 { - self.row_ids - .binary_search(&row_id) - .map(|idx| self.num_tokens[idx]) - .unwrap_or(0) - } - - // append a document to the doc set - // returns the doc_id (the number of documents before appending) - pub fn append(&mut self, row_id: u64, num_tokens: u32) -> u32 { - self.row_ids.push(row_id); - self.num_tokens.push(num_tokens); - self.total_tokens += num_tokens as u64; - self.invalidate_norms(); - 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) { - if self.norms.get().is_some() { - self.norms = Arc::new(std::sync::OnceLock::new()); - } - } - - pub(crate) fn memory_size(&self) -> usize { - self.row_ids.capacity() * std::mem::size_of::() - + self.num_tokens.memory_size() - + 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, - query: &str, - tokenizer: Option>, -) -> Result> { - if batches.is_empty() { - return Ok(vec![]); - } - - 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, 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, phrase_slop) - } - DataType::LargeList(_) => { - 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", - data_type - ))), - } -} - -fn do_flat_full_text_search( - batches: &[&RecordBatch], - doc_col: &str, - query: &str, - tokenizer: Option>, - phrase_slop: Option, -) -> Result> { - let mut results = Vec::new(); - let mut tokenizer = - tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap()); - let query_tokens = collect_query_tokens(query, &mut tokenizer); - - for batch in batches { - let row_id_array = batch[ROW_ID].as_primitive::(); - let doc_array = batch[doc_col].as_string::(); - for i in 0..row_id_array.len() { - let doc = doc_array.value(i); - if document_matches_flat_query(doc, &mut tokenizer, &query_tokens, phrase_slop)? { - results.push(row_id_array.value(i)); - } - } - } - - Ok(results) -} - -fn do_flat_full_text_search_list( - batches: &[&RecordBatch], - doc_col: &str, - query: &str, - tokenizer: Option>, - phrase_slop: Option, -) -> Result> { - let mut results = Vec::new(); - let mut tokenizer = - tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap()); - let query_tokens = collect_query_tokens(query, &mut tokenizer); - - for batch in batches { - let row_id_array = batch[ROW_ID].as_primitive::(); - let doc_array = batch[doc_col].as_list::(); - match doc_array.value_type() { - DataType::Utf8 | DataType::LargeUtf8 => {} - data_type => { - return Err(Error::invalid_input(format!( - "unsupported list item data type {} for inverted index", - data_type - ))); - } - } - for i in 0..row_id_array.len() { - if doc_array.is_null(i) { - continue; - } - let elements = doc_array.value(i); - 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)); - } - } - } - - Ok(results) -} - -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 - -/// Consumes a stream of record batches and produces token counts -/// -/// The resulting batch will have three columns: -/// - row_id: the row id of the document -/// - all_tokens: the total number of tokens in the document -/// - query_token_counts: a fixed size list of the count of each query token in the document -/// -/// This is an unbounded accumulation, however, for most queries, the per-row -/// growth will be fairly small. As a result we can process millions of tokens -/// with fairly modest memory usage. -/// -/// However, it is unwise to do a flat search across billions of rows. An FTS -/// index should be created instead. -async fn tokenize_and_count( - input: impl Stream> + Send, - tokenizer: Box, - query_tokens: Arc, - doc_col_idx: usize, - elapsed_compute: Option