From 059ae90e4cb63ed820d7b0fa647c17f0b85e5603 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Thu, 25 Jun 2026 15:05:21 +0800 Subject: [PATCH 01/20] feat(fts): add configurable posting block size --- docs/src/format/index/scalar/fts.md | 3 + docs/src/quickstart/full-text-search.md | 1 + .../index/scalar/InvertedIndexParams.java | 22 ++ .../index/scalar/InvertedIndexParamsTest.java | 32 +- protos/index_old.proto | 1 + python/python/lance/dataset.py | 3 + python/python/tests/test_scalar_index.py | 13 + python/src/dataset.rs | 5 + rust/lance-index/protos-cache/cache.proto | 3 + rust/lance-index/src/scalar/inverted.rs | 35 ++- .../src/scalar/inverted/builder.rs | 142 ++++++++- .../src/scalar/inverted/cache_codec.rs | 17 +- .../src/scalar/inverted/encoding.rs | 158 ++++++++-- rust/lance-index/src/scalar/inverted/index.rs | 275 +++++++++++++++--- rust/lance-index/src/scalar/inverted/iter.rs | 9 +- .../src/scalar/inverted/tokenizer.rs | 151 +++++++++- rust/lance-index/src/scalar/inverted/wand.rs | 51 ++-- rust/lance/src/dataset/mem_wal/index/fts.rs | 16 +- 18 files changed, 833 insertions(+), 104 deletions(-) diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index adc7f94d65e..9c9550eb897 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -43,6 +43,8 @@ An FTS index may contain multiple partitions. Each partition has its own set of | `_length` | UInt32 | false | Number of documents containing the token | | `_compressed_position` | List> | true | Optional compressed position lists for phrase queries | +The posting-list file schema metadata includes `posting_block_size`, the number of documents encoded per compressed posting block. Older indexes that do not have this metadata use the legacy block size `128`. + ### Metadata File Schema The metadata file contains JSON-serialized configuration and partition information: @@ -67,6 +69,7 @@ The metadata file contains JSON-serialized configuration and partition informati | `min_gram` | UInt32 | 2 | Minimum n-gram length (only for ngram tokenizer) | | `max_gram` | UInt32 | 15 | Maximum n-gram length (only for ngram tokenizer) | | `prefix_only` | Boolean | false | Generate only prefix n-grams (only for ngram tokenizer) | +| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128, 256, or 512. Missing values from older indexes read as 128. | ## Tokenizers diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index f990b2bd589..c72a22602ea 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -98,6 +98,7 @@ ds.create_scalar_index( remove_stop_words=True, # Remove stop words (language-dependent) custom_stop_words=None, # Optional additional stop words (only used if remove_stop_words=True) ascii_folding=True, # Fold accents to ASCII when possible (e.g., "é" -> "e") + block_size=256, # Posting block size: 128, 256, or 512 ) ``` diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index f509efd922b..7e2d9699ae7 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -53,6 +53,7 @@ public static final class Builder { private Integer minNgramLength; private Integer maxNgramLength; private Boolean prefixOnly; + private Integer blockSize = 256; private Boolean skipMerge; /** @@ -225,6 +226,24 @@ public Builder prefixOnly(boolean prefixOnly) { return this; } + /** + * Configure the number of documents in each compressed posting block. + * + *

Supported values are {@code 128}, {@code 256}, and {@code 512}. New indexes default to + * {@code 256} when this is not set. + * + * @param blockSize posting block size + * @return this builder + * @throws IllegalArgumentException if {@code blockSize} is unsupported + */ + public Builder blockSize(int blockSize) { + if (blockSize != 128 && blockSize != 256 && blockSize != 512) { + throw new IllegalArgumentException("blockSize must be one of 128, 256, or 512"); + } + this.blockSize = blockSize; + return this; + } + /** * Configure whether to skip the partition merge stage after indexing. If true, skip the * partition merge stage after indexing. This can be useful for distributed indexing where merge @@ -282,6 +301,9 @@ public ScalarIndexParams build() { if (prefixOnly != null) { params.put("prefix_only", prefixOnly); } + if (blockSize != null) { + params.put("block_size", blockSize); + } if (skipMerge != null) { params.put("skip_merge", skipMerge); } diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index e5024a95c2a..757d831c54a 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -13,19 +13,47 @@ */ package org.lance.index.scalar; +import org.lance.util.JsonUtils; + import org.junit.jupiter.api.Test; +import java.util.Map; + import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -public class InvertedIndexParamsTest { +class InvertedIndexParamsTest { @Test - public void testIcuSplitTokenizerVariant() { + void testIcuSplitTokenizerVariant() { ScalarIndexParams params = InvertedIndexParams.builder().baseTokenizer("icu/split").build(); assertEquals("inverted", params.getIndexType()); String jsonParams = params.getJsonParams().orElseThrow(AssertionError::new); assertTrue(jsonParams.contains("\"base_tokenizer\":\"icu/split\"")); } + + @Test + void defaultBlockSizeIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(256, ((Number) json.get("block_size")).intValue()); + } + + @Test + void blockSizeIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().blockSize(512).build(); + + assertEquals("inverted", params.getIndexType()); + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(512, ((Number) json.get("block_size")).intValue()); + } + + @Test + void invalidBlockSizeIsRejected() { + assertThrows( + IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(129)); + } } diff --git a/protos/index_old.proto b/protos/index_old.proto index 601aa2681da..94429e8d964 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -39,4 +39,5 @@ message InvertedIndexDetails { uint32 min_ngram_length = 9; uint32 max_ngram_length = 10; bool prefix_only = 11; + optional uint32 block_size = 12; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 833d0f1321d..f61fcc6918c 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3315,6 +3315,9 @@ def create_scalar_index( query. This will significantly increase the index size. It won't impact the performance of non-phrase queries even if it is set to True. + block_size: int, default 256 + This is for the ``INVERTED`` index. Number of documents per compressed + posting block. Must be one of ``128``, ``256``, or ``512``. memory_limit: int, optional This is for the ``INVERTED`` index. Total build-time memory limit in MiB. If set, Lance divides this budget evenly across the workers. If unset, diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 4450f6821cc..682304fcd23 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -949,6 +949,19 @@ def test_create_scalar_index_fts_alias(dataset): assert any(idx.index_type == "Inverted" for idx in dataset.describe_indices()) +def test_create_scalar_index_fts_block_size(dataset): + dataset.create_scalar_index( + "doc", index_type="INVERTED", with_position=False, block_size=512 + ) + row = dataset.take(indices=[0], columns=["doc"]) + query = row.column(0)[0].as_py().split(" ")[0] + results = dataset.scanner(columns=["doc"], full_text_query=query).to_table() + assert results.num_rows > 0 + + with pytest.raises(ValueError, match="block_size"): + dataset.create_scalar_index("doc2", index_type="INVERTED", block_size=129) + + def test_multi_index_create(tmp_path): dataset = lance.write_dataset( pa.table({"ints": range(1024)}), tmp_path, max_rows_per_file=100 diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 7c525b0aefe..c2c936911d5 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2429,6 +2429,11 @@ impl Dataset { if let Some(prefix_only) = kwargs.get_item("prefix_only")? { params = params.ngram_prefix_only(prefix_only.extract()?); } + if let Some(block_size) = kwargs.get_item("block_size")? { + params = params + .block_size(block_size.extract()?) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + } if let Some(memory_limit) = kwargs.get_item("memory_limit")? { params = params.memory_limit_mb(memory_limit.extract()?); } diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index b24a27055d7..e92da05815f 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -28,6 +28,9 @@ message CompressedPostingHeader { PositionStorage position_storage = 4; // Only meaningful when position_storage == POSITION_STORAGE_SHARED. PositionStreamCodec position_stream_codec = 5; + // Number of documents in each compressed posting block. Older cache entries + // omit this field and decode as the legacy 128-doc block size. + uint32 block_size = 6; } // Header for a serialized `PlainPostingList` cache entry. Followed by an Arrow diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index d0bb0e40d3a..8300dc1043c 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -215,7 +215,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { .into())) } - let params = serde_json::from_str::(params)?; + let params = InvertedIndexParams::from_training_json(params)?; Ok(Box::new(InvertedIndexTrainingRequest::new(params))) } @@ -308,6 +308,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { #[cfg(test)] mod tests { use super::*; + use crate::scalar::{BuiltinIndexType, ScalarIndexParams}; #[test] fn test_plugin_version_tracks_max_supported_format() { @@ -317,4 +318,36 @@ mod tests { max_supported_fts_format_version().index_version() ); } + + #[test] + fn test_new_training_request_defaults_missing_block_size_to_256() { + let plugin = InvertedIndexPlugin; + let field = Field::new("text", DataType::Utf8, true); + + let cases = [ + ( + ScalarIndexParams::for_builtin(BuiltinIndexType::Inverted), + false, + ), + (ScalarIndexParams::new("inverted".to_string()), false), + ( + ScalarIndexParams::new("inverted".to_string()) + .with_params(&serde_json::json!({ "with_position": true })), + true, + ), + ]; + + for (params, expected_with_position) in cases { + let request = plugin + .new_training_request(params.params.as_deref().unwrap_or("{}"), &field) + .unwrap(); + let request = request + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(request.parameters.posting_block_size(), DEFAULT_BLOCK_SIZE); + assert_eq!(request.parameters.has_positions(), expected_with_position); + } + } } diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index a53b6ddd7dc..02ac612ce37 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -6,6 +6,7 @@ use super::{InvertedIndexParams, index::*}; use crate::scalar::inverted::document_tokenizer::DocType; use crate::scalar::inverted::json::JsonTextStream; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; +use crate::scalar::inverted::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; #[cfg(test)] use crate::scalar::lance_format::LanceIndexStore; use crate::scalar::{IndexFile, IndexStore, OldIndexDataFilter}; @@ -41,9 +42,8 @@ use std::task::{Context, Poll}; use std::{fmt::Debug, sync::atomic::AtomicU64}; use tracing::instrument; -// the number of elements in each block -// each block contains 128 row ids and 128 frequencies -// WARNING: changing this value will break the compatibility with existing indexes +// The physical bitpacking block size. Logical FTS posting blocks may contain +// 1, 2, or 4 of these chunks, configured by InvertedIndexParams::block_size. pub const BLOCK_SIZE: usize = BitPacker4x::BLOCK_LEN; // The default number of workers to use for FTS builds. @@ -448,6 +448,7 @@ impl InvertedIndexBuilder { fragment_mask: self.fragment_mask, token_set_format: self.token_set_format, worker_memory_limit_bytes, + block_size: self.params.block_size, }; let next_id = self.next_partition_id(); let id_alloc = Arc::new(AtomicU64::new(next_id)); @@ -794,6 +795,7 @@ pub struct InnerBuilder { token_set_format: TokenSetFormat, format_version: InvertedListFormatVersion, posting_tail_codec: PostingTailCodec, + block_size: usize, pub(crate) tokens: TokenSet, pub(crate) posting_lists: Vec, pub(crate) docs: DocSet, @@ -816,12 +818,45 @@ impl InnerBuilder { token_set_format: TokenSetFormat, format_version: InvertedListFormatVersion, ) -> Self { + Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + format_version, + LEGACY_BLOCK_SIZE, + ) + } + + pub fn new_with_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + block_size: usize, + ) -> Self { + Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + current_fts_format_version(), + block_size, + ) + } + + pub fn new_with_format_version_and_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + format_version: InvertedListFormatVersion, + block_size: usize, + ) -> Self { + validate_block_size(block_size).expect("invalid posting list block size"); Self { id, with_position, token_set_format, format_version, posting_tail_codec: format_version.posting_tail_codec(), + block_size, tokens: TokenSet::default(), posting_lists: Vec::new(), docs: DocSet::default(), @@ -834,14 +869,35 @@ impl InnerBuilder { with_position: bool, token_set_format: TokenSetFormat, posting_tail_codec: PostingTailCodec, + ) -> Self { + Self::new_with_posting_tail_codec_and_block_size( + id, + with_position, + token_set_format, + posting_tail_codec, + LEGACY_BLOCK_SIZE, + ) + } + + pub fn new_with_posting_tail_codec_and_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + posting_tail_codec: PostingTailCodec, + block_size: usize, ) -> Self { let format_version = if posting_tail_codec == PostingTailCodec::Fixed32 { InvertedListFormatVersion::V1 } else { InvertedListFormatVersion::V2 }; - let mut builder = - Self::new_with_format_version(id, with_position, token_set_format, format_version); + let mut builder = Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + format_version, + block_size, + ); builder.posting_tail_codec = posting_tail_codec; builder } @@ -922,6 +978,7 @@ impl InnerBuilder { token_set_format, format_version, posting_tail_codec, + block_size, tokens, posting_lists, docs, @@ -952,6 +1009,12 @@ impl InnerBuilder { self.posting_tail_codec, posting_tail_codec ))); } + if self.block_size != block_size { + return Err(Error::index(format!( + "cannot merge partitions with mismatched FTS block sizes: {} vs {}", + self.block_size, block_size + ))); + } let mut token_id_map = vec![u32::MAX; posting_lists.len()]; match tokens.tokens { @@ -977,7 +1040,11 @@ impl InnerBuilder { self.docs.append(*row_id, *num_tokens); } self.posting_lists.resize_with(self.tokens.len(), || { - PostingListBuilder::new_with_posting_tail_codec(with_position, self.posting_tail_codec) + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + with_position, + self.posting_tail_codec, + self.block_size, + ) }); for (token_id, posting_list) in posting_lists.into_iter().enumerate() { @@ -1044,7 +1111,11 @@ impl InnerBuilder { let mut writer = store .new_index_file( path, - inverted_list_schema_for_version(self.with_position, self.format_version), + inverted_list_schema_for_version_with_block_size( + self.with_position, + self.format_version, + self.block_size, + ), ) .await?; let posting_lists = std::mem::take(&mut self.posting_lists); @@ -1058,7 +1129,11 @@ impl InnerBuilder { let with_position = self.with_position; let format_version = self.format_version; let group_config = self.group_config; - let schema = inverted_list_schema_for_version(self.with_position, self.format_version); + let schema = inverted_list_schema_for_version_with_block_size( + self.with_position, + self.format_version, + self.block_size, + ); let docs_for_batches = docs.clone(); let schema_for_batches = schema.clone(); let batch_rows = *LANCE_FTS_POSTING_BATCH_ROWS; @@ -1214,6 +1289,7 @@ struct IndexWorkerConfig { fragment_mask: Option, token_set_format: TokenSetFormat, worker_memory_limit_bytes: u64, + block_size: usize, } impl IndexWorker { @@ -1257,17 +1333,22 @@ impl IndexWorker { id_alloc: Arc, config: IndexWorkerConfig, ) -> Result { - let schema = inverted_list_schema_for_version(config.with_position, config.format_version); + let schema = inverted_list_schema_for_version_with_block_size( + config.with_position, + config.format_version, + config.block_size, + ); Ok(Self { tokenizer, dest_store, - builder: InnerBuilder::new_with_format_version( + builder: InnerBuilder::new_with_format_version_and_block_size( id_alloc.fetch_add(1, std::sync::atomic::Ordering::Relaxed) | config.fragment_mask.unwrap_or(0), config.with_position, config.token_set_format, config.format_version, + config.block_size, ), partitions: Vec::new(), files: Vec::new(), @@ -1326,9 +1407,10 @@ impl IndexWorker { * std::mem::size_of::()) as u64; builder.posting_lists.push( - PostingListBuilder::new_with_posting_tail_codec( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( true, posting_tail_codec, + builder.block_size, ), ); let new_posting_lists_overhead_size = (builder.posting_lists.capacity() @@ -1374,9 +1456,10 @@ impl IndexWorker { self.builder .posting_lists .resize_with(self.builder.tokens.len(), || { - PostingListBuilder::new_with_posting_tail_codec( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( false, self.builder.posting_tail_codec, + self.builder.block_size, ) }); let new_posting_lists_overhead_size = self.posting_lists_overhead_size(); @@ -1481,15 +1564,17 @@ impl IndexWorker { self.memory_size = self.temporary_memory_size(); let with_position = self.has_position(); let format_version = self.builder.format_version; + let block_size = self.builder.block_size; let builder = std::mem::replace( &mut self.builder, - InnerBuilder::new_with_format_version( + InnerBuilder::new_with_format_version_and_block_size( self.id_alloc .fetch_add(1, std::sync::atomic::Ordering::Relaxed) | self.fragment_mask.unwrap_or(0), with_position, self.token_set_format, format_version, + block_size, ), ); let written_partition_id = builder.id(); @@ -1609,17 +1694,31 @@ pub fn inverted_list_schema_for_version( with_position: bool, format_version: InvertedListFormatVersion, ) -> SchemaRef { + inverted_list_schema_for_version_with_block_size( + with_position, + format_version, + LEGACY_BLOCK_SIZE, + ) +} + +pub fn inverted_list_schema_for_version_with_block_size( + with_position: bool, + format_version: InvertedListFormatVersion, + block_size: usize, +) -> SchemaRef { + validate_block_size(block_size).expect("invalid posting list block size"); match format_version { - InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position), + InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position, block_size), InvertedListFormatVersion::V2 => inverted_list_schema_with_tail_codec_and_position_codec( with_position, PostingTailCodec::VarintDelta, Some(PositionStreamCodec::PackedDelta), + block_size, ), } } -fn inverted_list_schema_v1(with_position: bool) -> SchemaRef { +fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef { let mut fields = vec![ arrow_schema::Field::new( POSTING_COL, @@ -1648,7 +1747,10 @@ fn inverted_list_schema_v1(with_position: bool) -> SchemaRef { false, )); } - Arc::new(arrow_schema::Schema::new(fields)) + Arc::new(arrow_schema::Schema::new_with_metadata( + fields, + HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string())]), + )) } pub fn inverted_list_schema_with_tail_codec( @@ -1659,6 +1761,7 @@ pub fn inverted_list_schema_with_tail_codec( with_position, posting_tail_codec, Some(PositionStreamCodec::PackedDelta), + LEGACY_BLOCK_SIZE, ) } @@ -1666,6 +1769,7 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( with_position: bool, posting_tail_codec: PostingTailCodec, position_codec: Option, + block_size: usize, ) -> SchemaRef { let mut fields = vec![ // we compress the posting lists (including row ids and frequencies), @@ -1702,6 +1806,7 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( POSTING_TAIL_CODEC_KEY.to_owned(), posting_tail_codec.as_str().to_owned(), )]); + metadata.insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); if let Some(position_codec) = position_codec.filter(|_| with_position) { metadata.insert( POSITIONS_LAYOUT_KEY.to_owned(), @@ -3184,6 +3289,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3207,6 +3313,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3680,6 +3787,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3711,6 +3819,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3749,6 +3858,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; diff --git a/rust/lance-index/src/scalar/inverted/cache_codec.rs b/rust/lance-index/src/scalar/inverted/cache_codec.rs index a676455d5c9..afc2f00cffe 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -43,6 +43,7 @@ use super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, }; +use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; // --------------------------------------------------------------------------- // Tags @@ -291,6 +292,7 @@ fn serialize_compressed( posting_tail_codec: posting_tail_codec_to_proto(posting.posting_tail_codec) as i32, position_storage: position_storage as i32, position_stream_codec: position_stream_codec as i32, + block_size: posting.block_size as u32, }; w.write_header(&header)?; @@ -322,12 +324,18 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result { assert_eq!(restored.max_score, posting.max_score); assert_eq!(restored.length, posting.length); assert_eq!(restored.posting_tail_codec, posting.posting_tail_codec); + assert_eq!(restored.block_size, posting.block_size); assert_eq!(restored.blocks, posting.blocks); assert!(restored.positions.is_none()); } @@ -555,6 +564,7 @@ mod tests { 1.25, 5, PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, Some(CompressedPositionStorage::LegacyPerDoc(legacy_positions( &[&[0, 4, 8]], ))), @@ -589,6 +599,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 512, Some(CompressedPositionStorage::SharedStream(stream)), ); let entry = PostingList::Compressed(posting.clone()); @@ -617,6 +628,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 512, Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), @@ -651,6 +663,7 @@ mod tests { 2.5, 7, PostingTailCodec::VarintDelta, + 256, None, )); @@ -760,6 +773,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream(stream)), )) } @@ -810,6 +824,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, None, )) }; diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index fabc4ac9261..67567ec8cee 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -5,6 +5,9 @@ use std::io::Write; use super::builder::BLOCK_SIZE; use super::index::{PositionStreamCodec, PostingTailCodec}; +#[cfg(test)] +use super::tokenizer::LEGACY_BLOCK_SIZE; +use super::tokenizer::validate_block_size; use arrow::array::LargeBinaryBuilder; use bitpacking::{BitPacker, BitPacker4x}; use lance_core::{Error, Result}; @@ -41,13 +44,33 @@ pub fn compress_posting_list<'a>( #[cfg(test)] pub fn compress_posting_list_with_tail_codec<'a>( + length: usize, + doc_ids: impl Iterator, + frequencies: impl Iterator, + block_max_scores: impl Iterator, + tail_codec: PostingTailCodec, +) -> Result { + compress_posting_list_with_tail_codec_and_block_size( + length, + doc_ids, + frequencies, + block_max_scores, + tail_codec, + LEGACY_BLOCK_SIZE, + ) +} + +#[cfg(test)] +pub fn compress_posting_list_with_tail_codec_and_block_size<'a>( length: usize, doc_ids: impl Iterator, frequencies: impl Iterator, mut block_max_scores: impl Iterator, tail_codec: PostingTailCodec, + block_size: usize, ) -> Result { - if length < BLOCK_SIZE { + let block_size = validate_block_size(block_size)?; + if length < block_size { // directly do remainder compression to avoid overhead of creating buffer let mut builder = LargeBinaryBuilder::with_capacity(1, length * 4 * 2 + 1); // write the max score of the block @@ -63,27 +86,23 @@ pub fn compress_posting_list_with_tail_codec<'a>( return Ok(builder.finish()); } - let mut builder = LargeBinaryBuilder::with_capacity(length.div_ceil(BLOCK_SIZE), length * 3); - let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - let mut doc_id_buffer = Vec::with_capacity(BLOCK_SIZE); - let mut freq_buffer = Vec::with_capacity(BLOCK_SIZE); + let mut builder = LargeBinaryBuilder::with_capacity(length.div_ceil(block_size), length * 3); + let mut doc_id_buffer = Vec::with_capacity(block_size); + let mut freq_buffer = Vec::with_capacity(block_size); for (doc_id, freq) in std::iter::zip(doc_ids, frequencies) { doc_id_buffer.push(*doc_id); freq_buffer.push(*freq); - if doc_id_buffer.len() < BLOCK_SIZE { + if doc_id_buffer.len() < block_size { continue; } - assert_eq!(doc_id_buffer.len(), BLOCK_SIZE); + assert_eq!(doc_id_buffer.len(), block_size); // write the max score of the block let max_score = block_max_scores.next().unwrap(); let _ = builder.write(max_score.to_le_bytes().as_ref())?; - // delta encoding + bitpacking for doc ids - compress_sorted_block(&doc_id_buffer, &mut buffer, &mut builder)?; - // bitpacking for frequencies - compress_block(&freq_buffer, &mut buffer, &mut builder)?; + encode_posting_block_payload(&doc_id_buffer, &freq_buffer, &mut builder)?; builder.append_value(""); doc_id_buffer.clear(); freq_buffer.clear(); @@ -105,12 +124,28 @@ pub fn encode_full_posting_block_into( frequencies: &[u32], block: &mut Vec, ) -> Result<()> { - debug_assert_eq!(doc_ids.len(), BLOCK_SIZE); - debug_assert_eq!(frequencies.len(), BLOCK_SIZE); + validate_block_size(doc_ids.len())?; + debug_assert_eq!(doc_ids.len(), frequencies.len()); block.extend_from_slice(&0f32.to_le_bytes()); + encode_posting_block_payload(doc_ids, frequencies, block)?; + Ok(()) +} + +fn encode_posting_block_payload( + doc_ids: &[u32], + frequencies: &[u32], + block: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(doc_ids.len(), frequencies.len()); + debug_assert!(doc_ids.len().is_multiple_of(BLOCK_SIZE)); let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - compress_sorted_block(doc_ids, &mut buffer, block)?; - compress_block(frequencies, &mut buffer, block)?; + for (doc_ids, frequencies) in doc_ids + .chunks_exact(BLOCK_SIZE) + .zip(frequencies.chunks_exact(BLOCK_SIZE)) + { + compress_sorted_block(doc_ids, &mut buffer, block)?; + compress_block(frequencies, &mut buffer, block)?; + } Ok(()) } @@ -650,17 +685,39 @@ pub fn decompress_posting_list_with_tail_codec( posting_list: &arrow::array::LargeBinaryArray, tail_codec: PostingTailCodec, ) -> Result<(Vec, Vec)> { + decompress_posting_list_with_tail_codec_and_block_size( + num_docs, + posting_list, + tail_codec, + LEGACY_BLOCK_SIZE, + ) +} + +#[cfg(test)] +pub fn decompress_posting_list_with_tail_codec_and_block_size( + num_docs: u32, + posting_list: &arrow::array::LargeBinaryArray, + tail_codec: PostingTailCodec, + block_size: usize, +) -> Result<(Vec, Vec)> { + let block_size = validate_block_size(block_size)?; let mut doc_ids: Vec = Vec::with_capacity(num_docs as usize); let mut frequencies: Vec = Vec::with_capacity(num_docs as usize); let mut buffer = [0u32; BLOCK_SIZE]; - let bitpacking_blocks = num_docs as usize / BLOCK_SIZE; + let bitpacking_blocks = num_docs as usize / block_size; for compressed in posting_list.iter().take(bitpacking_blocks) { let compressed = compressed.unwrap(); - decompress_posting_block(compressed, &mut buffer, &mut doc_ids, &mut frequencies); + decompress_posting_block( + compressed, + &mut buffer, + &mut doc_ids, + &mut frequencies, + block_size, + ); } - let remainder = num_docs as usize % BLOCK_SIZE; + let remainder = num_docs as usize % block_size; if remainder > 0 { let compressed = posting_list.value(bitpacking_blocks); decompress_posting_remainder( @@ -704,11 +761,22 @@ pub fn decompress_posting_block( buffer: &mut [u32; BLOCK_SIZE], doc_ids: &mut Vec, frequencies: &mut Vec, + block_size: usize, ) { + debug_assert!(validate_block_size(block_size).is_ok()); // skip the first 4 bytes for the max block score - let block = &block[4..]; - let num_bytes = decompress_sorted_block(block, buffer, doc_ids); - decompress_block(&block[num_bytes..], buffer, frequencies); + let mut block = &block[4..]; + for _ in 0..(block_size / BLOCK_SIZE) { + let num_bytes = decompress_sorted_block(block, buffer, doc_ids); + block = &block[num_bytes..]; + let num_bytes = decompress_block(block, buffer, frequencies); + block = &block[num_bytes..]; + } + debug_assert!( + block.is_empty(), + "posting block has {} trailing bytes after decoding", + block.len() + ); } pub fn decompress_posting_remainder( @@ -755,9 +823,14 @@ pub fn decompress_posting_remainder( } } -pub fn decode_full_posting_block(block: &[u8], doc_ids: &mut Vec, frequencies: &mut Vec) { +pub fn decode_full_posting_block( + block: &[u8], + doc_ids: &mut Vec, + frequencies: &mut Vec, + block_size: usize, +) { let mut buffer = [0u32; BLOCK_SIZE]; - decompress_posting_block(block, &mut buffer, doc_ids, frequencies); + decompress_posting_block(block, &mut buffer, doc_ids, frequencies, block_size); } pub fn decompress_sorted_block( @@ -773,11 +846,12 @@ pub fn decompress_sorted_block( 5 + num_bytes } -fn decompress_block(block: &[u8], buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec) { +fn decompress_block(block: &[u8], buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec) -> usize { let compressor = BitPacker4x::new(); let num_bits = block[0]; - compressor.decompress(&block[1..], buffer, num_bits); + let num_bytes = compressor.decompress(&block[1..], buffer, num_bits); res.extend_from_slice(&buffer[..]); + 1 + num_bytes } pub fn decompress_raw_remainder(compressed: &[u8], n: usize, dest: &mut Vec) { @@ -867,6 +941,40 @@ mod tests { Ok(()) } + #[test] + fn test_compress_posting_list_supported_block_sizes() -> Result<()> { + for block_size in [128, 256, 512] { + let num_rows: usize = block_size * 2 + 7; + let doc_ids = (0..num_rows as u32).collect::>(); + let frequencies = (0..num_rows as u32) + .map(|value| value % 7 + 1) + .collect::>(); + let block_max_scores = + (0..num_rows.div_ceil(block_size)).map(|value| value as f32 + 1.0); + + let posting_list = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + block_max_scores, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(posting_list.len(), num_rows.div_ceil(block_size)); + + let (decoded_doc_ids, decoded_frequencies) = + decompress_posting_list_with_tail_codec_and_block_size( + num_rows as u32, + &posting_list, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(decoded_doc_ids, doc_ids); + assert_eq!(decoded_frequencies, frequencies); + } + Ok(()) + } + #[test] fn test_compress_posting_list_fixed32_tail_still_roundtrips() -> Result<()> { let doc_ids = vec![3_u32, 10_u32, 24_u32]; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index fc0b02ec209..b7b730930a8 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -55,11 +55,12 @@ use tracing::{info, instrument}; use super::encoding::{PositionBlockBuilder, decode_group_starts}; use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; +use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ BLOCK_SIZE, PostingGroupAccumulator, PostingGroupConfig, ScoredDoc, doc_file_path, - inverted_list_schema_for_version, posting_file_path, token_file_path, + inverted_list_schema_for_version_with_block_size, posting_file_path, token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -111,6 +112,7 @@ pub const TOKEN_SET_FORMAT_KEY: &str = "token_set_format"; pub const POSTING_TAIL_CODEC_KEY: &str = "posting_tail_codec"; 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"; /// Schema-metadata key holding the 1-indexed global-buffer id of the /// varint-delta-encoded posting-list cache-group boundaries (issue #7040). /// Absent on indexes written before grouping was introduced, which fall back @@ -365,6 +367,21 @@ pub(super) fn parse_posting_tail_codec( .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 { @@ -1556,6 +1573,13 @@ impl InvertedPartition { postings: Vec, docs: &DocSet, ) -> 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() { @@ -1580,7 +1604,7 @@ impl InvertedPartition { ))); } - let mut builder = PostingListBuilder::new(false); + 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 { @@ -1588,7 +1612,11 @@ impl InvertedPartition { doc_ids.push(doc_id); frequencies.push(freq); } - let block_max_scores = docs.calculate_block_max_scores(doc_ids.iter(), frequencies.iter()); + let block_max_scores = docs.calculate_block_max_scores_with_block_size( + doc_ids.iter(), + frequencies.iter(), + block_size, + ); 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); @@ -1819,11 +1847,12 @@ impl InvertedPartition { } pub async fn into_builder(self) -> Result { - let mut builder = InnerBuilder::new_with_posting_tail_codec( + 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(); // into_builder rewrites every doc, so materialize the full @@ -2233,6 +2262,7 @@ pub struct PostingListReader { has_position: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, /// First row of each posting-list cache group, decoded at open from the @@ -2331,6 +2361,7 @@ impl PostingListReader { 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 metadata = if reader.schema().field(POSTING_COL).is_none() { let (offsets, max_scores) = Self::load_metadata(reader.schema())?; @@ -2351,6 +2382,7 @@ impl PostingListReader { metadata, has_position, posting_tail_codec, + block_size, positions_layout, group_starts, index_cache: WeakLanceCache::from(index_cache), @@ -2413,6 +2445,10 @@ impl PostingListReader { 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 { .. }) } @@ -2704,6 +2740,7 @@ impl PostingListReader { 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( @@ -2711,6 +2748,7 @@ impl PostingListReader { max_score, length, posting_tail_codec, + block_size, positions_layout, )?; Ok(posting_list) @@ -2727,6 +2765,7 @@ impl PostingListReader { max_score, length, self.posting_tail_codec, + self.block_size, self.positions_layout, ) } @@ -2763,6 +2802,7 @@ impl PostingListReader { 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)); @@ -2899,6 +2939,7 @@ impl PostingListReader { 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, } } @@ -2932,12 +2973,14 @@ impl PostingListReader { 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 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 { @@ -3188,6 +3231,7 @@ struct ChunkBuildState { max_scores: Option>>, lengths: Option>>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3198,6 +3242,7 @@ struct PrewarmBuildCtx<'a> { max_scores: Option<&'a [f32]>, lengths: Option<&'a [u32]>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3461,7 +3506,8 @@ impl PostingList { length: Option, ) -> Result { let posting_tail_codec = parse_posting_tail_codec(batch.schema_ref().metadata())?; - Self::from_batch_with_tail_codec(batch, max_score, length, posting_tail_codec) + 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( @@ -3469,6 +3515,7 @@ impl PostingList { 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( @@ -3484,6 +3531,7 @@ impl PostingList { max_score, length, posting_tail_codec, + block_size, positions_layout, ) } @@ -3493,6 +3541,7 @@ impl PostingList { max_score: Option, length: Option, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, ) -> Result { match batch.column_by_name(POSTING_COL) { @@ -3507,6 +3556,7 @@ impl PostingList { max_score.unwrap(), length.unwrap(), posting_tail_codec, + block_size, shared_position_codec, ); Ok(Self::Compressed(posting)) @@ -3578,9 +3628,14 @@ impl PostingList { Self::Plain(_) => PostingTailCodec::Fixed32, Self::Compressed(posting) => posting.posting_tail_codec, }; - let mut builder = PostingListBuilder::new_with_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 @@ -3738,9 +3793,11 @@ 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 + // that contains `block_size` doc ids and then `block_size` frequencies, + // packed internally in 128-value bitpacking chunks. pub blocks: LargeBinaryArray, pub posting_tail_codec: PostingTailCodec, + pub block_size: usize, pub positions: Option, } @@ -3761,6 +3818,7 @@ impl CompressedPostingList { max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, positions: Option, ) -> Self { Self { @@ -3768,6 +3826,7 @@ impl CompressedPostingList { length, blocks, posting_tail_codec, + block_size, positions, } } @@ -3777,6 +3836,7 @@ impl CompressedPostingList { max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, shared_position_codec: Option, ) -> Self { debug_assert_eq!(batch.num_rows(), 1); @@ -3813,6 +3873,7 @@ impl CompressedPostingList { length, blocks, posting_tail_codec, + block_size, positions, } } @@ -3823,6 +3884,7 @@ impl CompressedPostingList { self.blocks.clone(), self.posting_tail_codec, self.positions.clone(), + self.block_size, ) } @@ -3833,7 +3895,7 @@ impl CompressedPostingList { pub fn block_least_doc_id(&self, block_idx: usize) -> u32 { let block = self.blocks.value(block_idx); - let remainder = self.length as usize % BLOCK_SIZE; + let remainder = self.length as usize % self.block_size; let is_remainder_block = remainder > 0 && block_idx + 1 == self.blocks.len(); if is_remainder_block { super::encoding::read_posting_tail_first_doc(block, self.posting_tail_codec) @@ -3964,6 +4026,7 @@ pub struct PostingListBuilder { open_doc_id: Option, open_doc_frequency: u32, open_doc_last_position: Option, + block_size: usize, memory_size_bytes: u32, len: u32, } @@ -3993,6 +4056,7 @@ enum BatchPositionsBuilder { struct PostingListParts<'a> { with_positions: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, length: usize, encoded_blocks: EncodedBlocks, encoded_position_blocks: EncodedPositionBlocks, @@ -4155,9 +4219,10 @@ impl PostingListBuilder { } pub fn new(with_position: bool) -> Self { - Self::new_with_posting_tail_codec( + Self::new_with_posting_tail_codec_and_block_size( with_position, current_fts_format_version().posting_tail_codec(), + LEGACY_BLOCK_SIZE, ) } @@ -4165,6 +4230,27 @@ impl PostingListBuilder { 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, @@ -4175,6 +4261,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, len: 0, memory_size_bytes: 0, } @@ -4196,8 +4283,8 @@ impl PostingListBuilder { &self, mut visit: impl FnMut(u32, u32, Option>) -> std::result::Result<(), E>, ) -> std::result::Result<(), E> { - let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); - let mut frequencies = Vec::with_capacity(BLOCK_SIZE); + 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; @@ -4205,7 +4292,12 @@ impl PostingListBuilder { for block in encoded_blocks.iter() { doc_ids.clear(); frequencies.clear(); - super::encoding::decode_full_posting_block(block, &mut doc_ids, &mut frequencies); + 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 @@ -4285,7 +4377,7 @@ impl PostingListBuilder { } self.len += 1; - if self.tail_entries.len() == BLOCK_SIZE { + if self.tail_entries.len() == self.block_size { self.flush_tail_block() .expect("posting list block compression should succeed"); } @@ -4344,7 +4436,7 @@ impl PostingListBuilder { self.open_doc_id = None; self.open_doc_frequency = 0; self.open_doc_last_position = None; - if self.tail_entries.len() == BLOCK_SIZE { + if self.tail_entries.len() == self.block_size { self.flush_tail_block()?; } Ok(()) @@ -4395,13 +4487,17 @@ impl PostingListBuilder { self.open_doc_id.is_none(), "cannot flush a posting block while a document is still open" ); - debug_assert_eq!(self.tail_entries.len(), BLOCK_SIZE); - let mut doc_ids = [0u32; BLOCK_SIZE]; - let mut frequencies = [0u32; BLOCK_SIZE]; - for (index, entry) in self.tail_entries.iter().enumerate() { - doc_ids[index] = entry.doc_id; - frequencies[index] = entry.frequency; - } + 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() @@ -4572,6 +4668,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4581,6 +4678,7 @@ impl PostingListBuilder { let parts = PostingListParts { with_positions, posting_tail_codec, + block_size, length: len as usize, encoded_blocks: encoded_blocks .map(|encoded_blocks| *encoded_blocks) @@ -4619,6 +4717,7 @@ impl PostingListBuilder { with_positions, posting_tail_codec, length, + block_size, mut encoded_blocks, mut encoded_position_blocks, tail_entries, @@ -4627,14 +4726,19 @@ impl PostingListBuilder { 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 doc_ids = Vec::with_capacity(block_size); + let mut frequencies = Vec::with_capacity(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); + super::encoding::decode_full_posting_block( + block, + &mut doc_ids, + &mut frequencies, + block_size, + ); let block_score = compute_block_score( docs, avgdl, @@ -4733,7 +4837,11 @@ impl PostingListBuilder { } else { InvertedListFormatVersion::V2 }; - let schema = inverted_list_schema_for_version(self.has_positions(), format_version); + let schema = inverted_list_schema_for_version_with_block_size( + self.has_positions(), + format_version, + self.block_size, + ); let legacy_positions = if self.with_positions && !format_version.uses_shared_position_stream() { Some(self.build_legacy_positions()?) @@ -4750,6 +4858,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4780,6 +4889,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -4814,6 +4924,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4823,6 +4934,7 @@ impl PostingListBuilder { let parts = PostingListParts { with_positions, posting_tail_codec, + block_size, length: len as usize, encoded_blocks: encoded_blocks .map(|encoded_blocks| *encoded_blocks) @@ -4845,6 +4957,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -4857,8 +4970,11 @@ impl PostingListBuilder { pub fn remap(&mut self, removed: &[u32]) { let mut cursor = 0; - let mut new_builder = - Self::new_with_posting_tail_codec(self.has_positions(), self.posting_tail_codec); + 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; @@ -5079,9 +5195,19 @@ impl DocSet { 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 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; @@ -5092,13 +5218,13 @@ impl DocSet { if score > max_score { max_score = score; } - if (i + 1) % BLOCK_SIZE == 0 { + 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) { + if !length.is_multiple_of(block_size) { max_score *= idf_scale; block_max_scores.push(max_score); } @@ -5756,15 +5882,20 @@ mod tests { token: &str, row_id: u64, ) -> Result> { - let mut partition = InnerBuilder::new_with_format_version( + let block_size = params.posting_block_size(); + let mut partition = InnerBuilder::new_with_format_version_and_block_size( 0, false, token_set_format, InvertedListFormatVersion::V1, + block_size, ); partition.tokens.add(token.to_owned()); - let mut posting_list = - PostingListBuilder::new_with_posting_tail_codec(false, PostingTailCodec::Fixed32); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + PostingTailCodec::Fixed32, + block_size, + ); posting_list.add(0, PositionRecorder::Count(1)); partition.posting_lists.push(posting_list); partition.docs.append(row_id, 1); @@ -5800,6 +5931,83 @@ mod tests { )) } + #[test] + fn test_posting_block_size_schema_metadata() { + assert_eq!(parse_posting_block_size(&HashMap::new()).unwrap(), 128); + + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "512".to_owned())]); + assert_eq!(parse_posting_block_size(&metadata).unwrap(), 512); + + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "129".to_owned())]); + let err = parse_posting_block_size(&metadata).unwrap_err(); + assert!(err.to_string().contains("block_size")); + } + + #[tokio::test] + async fn test_build_search_uses_configured_posting_block_size() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let block_size = params.posting_block_size(); + let num_docs = block_size + 7; + + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + block_size, + ); + builder.tokens.add("needle".to_owned()); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + PostingTailCodec::Fixed32, + block_size, + ); + for doc_id in 0..num_docs { + posting_list.add(doc_id as u32, PositionRecorder::Count(1)); + builder.docs.append(1_000 + doc_id as u64, 1); + } + builder.posting_lists.push(posting_list); + builder.write(store.as_ref()).await.unwrap(); + write_test_metadata(&store, vec![0], params).await; + + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + assert_eq!(index.partitions[0].inverted_list.block_size(), block_size); + + let posting = index.partitions[0] + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed posting list"); + }; + assert_eq!(posting.block_size, block_size); + assert_eq!(posting.blocks.len(), num_docs.div_ceil(block_size)); + + let tokens = Arc::new(Tokens::new(vec!["needle".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); + let prefilter = Arc::new(NoFilter); + let metrics = Arc::new(NoOpMetricsCollector); + let (row_ids, scores) = index + .bm25_search(tokens, params, Operator::Or, prefilter, metrics, None) + .await + .unwrap(); + + assert_eq!(row_ids.len(), 10); + assert_eq!(scores.len(), 10); + assert!(row_ids.iter().all(|row_id| *row_id >= 1_000)); + } + #[tokio::test] async fn test_posting_builder_remap() { let posting_tail_codec = PostingTailCodec::Fixed32; @@ -7397,6 +7605,7 @@ mod tests { 1.0, SLICE_LEN as u32, PostingTailCodec::Fixed32, + LEGACY_BLOCK_SIZE, None, ); diff --git a/rust/lance-index/src/scalar/inverted/iter.rs b/rust/lance-index/src/scalar/inverted/iter.rs index dc07b15769c..52c09fa8dab 100644 --- a/rust/lance-index/src/scalar/inverted/iter.rs +++ b/rust/lance-index/src/scalar/inverted/iter.rs @@ -28,6 +28,7 @@ impl<'a> PostingListIterator<'a> { posting.blocks.clone(), posting.posting_tail_codec, posting.positions.clone(), + posting.block_size, ))) } } @@ -82,6 +83,7 @@ pub struct CompressedPostingListIterator { next_block_idx: usize, posting_tail_codec: PostingTailCodec, positions: Option, + block_size: usize, idx: usize, doc_ids: Vec, frequencies: Vec, @@ -97,10 +99,11 @@ impl CompressedPostingListIterator { blocks: LargeBinaryArray, posting_tail_codec: PostingTailCodec, positions: Option, + block_size: usize, ) -> Self { debug_assert!(length > 0, "length: {}", length); debug_assert_eq!( - length.div_ceil(BLOCK_SIZE), + length.div_ceil(block_size), blocks.len(), "length: {}, num_blocks: {}", length, @@ -108,11 +111,12 @@ impl CompressedPostingListIterator { ); Self { - remainder: length % BLOCK_SIZE, + remainder: length % block_size, blocks, next_block_idx: 0, posting_tail_codec, positions, + block_size, idx: 0, doc_ids: Vec::new(), frequencies: Vec::new(), @@ -172,6 +176,7 @@ impl Iterator for CompressedPostingListIterator { &mut self.buffer, &mut self.doc_ids, &mut self.frequencies, + self.block_size, ); } self.doc_idx_in_block = 0; diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 1785a23fedd..c4181a65847 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use lance_core::{Error, Result}; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; use std::{env, path::PathBuf}; #[cfg(feature = "tokenizer-jieba")] @@ -29,6 +29,10 @@ use lance_tokenizer::{ WhitespaceTokenizer, }; +pub const LEGACY_BLOCK_SIZE: usize = 128; +pub const DEFAULT_BLOCK_SIZE: usize = 256; +pub const VALID_BLOCK_SIZES: [usize; 3] = [128, 256, 512]; + /// Tokenizer configs #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct InvertedIndexParams { @@ -98,6 +102,17 @@ pub struct InvertedIndexParams { #[serde(default)] pub(crate) prefix_only: bool, + /// Number of documents in each compressed posting block. + /// + /// Missing serialized values come from indexes written before this + /// parameter existed and must read as 128 for backwards compatibility. New + /// indexes default to 256. + #[serde( + default = "legacy_block_size", + deserialize_with = "deserialize_block_size" + )] + pub(crate) block_size: usize, + /// Total memory limit in MiB for the build stage. /// /// This is split evenly across FTS workers at build time. By default Lance @@ -138,6 +153,7 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { min_ngram_length: params.min_ngram_length, max_ngram_length: params.max_ngram_length, prefix_only: params.prefix_only, + block_size: Some(params.block_size as u32), }) } } @@ -165,6 +181,10 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { min_ngram_length: details.min_ngram_length, max_ngram_length: details.max_ngram_length, prefix_only: details.prefix_only, + block_size: match details.block_size { + Some(block_size) => validate_block_size(block_size as usize)?, + None => LEGACY_BLOCK_SIZE, + }, memory_limit_mb: defaults.memory_limit_mb, num_workers: defaults.num_workers, }) @@ -183,6 +203,30 @@ fn default_max_ngram_length() -> u32 { 3 } +fn legacy_block_size() -> usize { + LEGACY_BLOCK_SIZE +} + +fn invalid_block_size_message(block_size: usize) -> String { + format!("FTS inverted index block_size must be one of 128, 256, or 512, got {block_size}") +} + +pub fn validate_block_size(block_size: usize) -> Result { + if VALID_BLOCK_SIZES.contains(&block_size) { + Ok(block_size) + } else { + Err(Error::invalid_input(invalid_block_size_message(block_size))) + } +} + +fn deserialize_block_size<'de, D>(deserializer: D) -> std::result::Result +where + D: Deserializer<'de>, +{ + let block_size = usize::deserialize(deserializer)?; + validate_block_size(block_size).map_err(serde::de::Error::custom) +} + impl Default for InvertedIndexParams { fn default() -> Self { Self::new("simple".to_owned(), Language::English) @@ -220,6 +264,7 @@ impl InvertedIndexParams { min_ngram_length: default_min_ngram_length(), max_ngram_length: default_max_ngram_length(), prefix_only: false, + block_size: DEFAULT_BLOCK_SIZE, memory_limit_mb: None, num_workers: None, } @@ -310,6 +355,21 @@ impl InvertedIndexParams { self } + /// Set the compressed posting block size. + /// + /// Supported values are 128, 256, and 512. Larger values reduce block-max + /// metadata and WAND skip granularity; smaller values preserve the legacy + /// layout. + pub fn block_size(mut self, block_size: usize) -> Result { + self.block_size = validate_block_size(block_size)?; + Ok(self) + } + + /// Get the compressed posting block size. + pub fn posting_block_size(&self) -> usize { + self.block_size + } + pub fn memory_limit_mb(mut self, memory_limit_mb: u64) -> Self { self.memory_limit_mb = Some(memory_limit_mb); self @@ -345,6 +405,23 @@ impl InvertedIndexParams { Ok(value) } + /// Deserialize params for new index training, using current creation defaults + /// for omitted fields. + pub(crate) fn from_training_json(params: &str) -> Result { + let supplied = serde_json::from_str::(params)?; + let mut value = serde_json::to_value(Self::default())?; + + let supplied = supplied.as_object().ok_or_else(|| { + Error::invalid_input("FTS inverted index params must be a JSON object".to_string()) + })?; + let object = value + .as_object_mut() + .expect("inverted index params should serialize to a JSON object"); + object.extend(supplied.clone()); + + Ok(serde_json::from_value(value)?) + } + pub fn build(&self) -> Result> { let mut builder = self.build_base_tokenizer()?; if let Some(max_token_length) = self.max_token_length { @@ -450,8 +527,10 @@ pub fn language_model_home() -> Option { #[cfg(test)] mod tests { + use crate::pbold; + use super::InvertedIndexParams; - use lance_tokenizer::TokenStream; + use lance_tokenizer::{Language, TokenStream}; use rstest::rstest; #[test] @@ -502,6 +581,74 @@ mod tests { assert_eq!(json.get("num_workers"), Some(&serde_json::Value::from(3))); } + #[test] + fn test_block_size_new_default_serializes() { + let params = InvertedIndexParams::default(); + assert_eq!(params.block_size, 256); + let json = serde_json::to_value(¶ms).unwrap(); + assert_eq!(json.get("block_size"), Some(&serde_json::Value::from(256))); + } + + #[test] + fn test_block_size_missing_metadata_falls_back_to_128() { + let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + json.as_object_mut().unwrap().remove("block_size"); + + let params: InvertedIndexParams = serde_json::from_value(json).unwrap(); + assert_eq!(params.block_size, 128); + } + + #[test] + fn test_block_size_details_conversion() { + let params = InvertedIndexParams::default().block_size(512).unwrap(); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + assert_eq!(details.block_size, Some(512)); + + let old_details = pbold::InvertedIndexDetails { + base_tokenizer: Some("simple".to_string()), + language: serde_json::to_string(&Language::English).unwrap(), + with_position: false, + max_token_length: Some(40), + lower_case: true, + stem: true, + remove_stop_words: true, + ascii_folding: true, + min_ngram_length: 3, + max_ngram_length: 3, + prefix_only: false, + block_size: None, + }; + let params = InvertedIndexParams::try_from(&old_details).unwrap(); + assert_eq!(params.block_size, 128); + } + + #[test] + fn test_block_size_accepts_supported_values() { + for block_size in [128, 256, 512] { + let params = InvertedIndexParams::default() + .block_size(block_size) + .unwrap(); + assert_eq!(params.block_size, block_size); + + let roundtrip: InvertedIndexParams = + serde_json::from_value(serde_json::to_value(¶ms).unwrap()).unwrap(); + assert_eq!(roundtrip.block_size, block_size); + } + } + + #[test] + fn test_block_size_rejects_invalid_values() { + let err = InvertedIndexParams::default().block_size(129).unwrap_err(); + assert!(err.to_string().contains("block_size")); + + let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + json.as_object_mut() + .unwrap() + .insert("block_size".to_string(), serde_json::Value::from(1024)); + let err = serde_json::from_value::(json).unwrap_err(); + assert!(err.to_string().contains("128, 256, or 512")); + } + #[test] fn test_build_icu_tokenizer() { let mut tokenizer = InvertedIndexParams::default() diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 0fc8b95cddb..1f51afa0e48 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -74,11 +74,11 @@ struct CompressedState { } impl CompressedState { - fn new() -> Self { + fn new(block_size: usize) -> Self { Self { block_idx: 0, - doc_ids: Vec::with_capacity(BLOCK_SIZE), - freqs: Vec::with_capacity(BLOCK_SIZE), + doc_ids: Vec::with_capacity(block_size), + freqs: Vec::with_capacity(block_size), buffer: Box::new([0; BLOCK_SIZE]), position_block_idx: None, position_values: Vec::new(), @@ -95,11 +95,12 @@ impl CompressedState { num_blocks: usize, length: u32, tail_codec: super::PostingTailCodec, + block_size: usize, ) { self.doc_ids.clear(); self.freqs.clear(); - let remainder = length as usize % BLOCK_SIZE; + let remainder = length as usize % block_size; if block_idx + 1 == num_blocks && remainder != 0 { decompress_posting_remainder( block, @@ -109,7 +110,13 @@ impl CompressedState { &mut self.freqs, ); } else { - decompress_posting_block(block, &mut self.buffer, &mut self.doc_ids, &mut self.freqs); + decompress_posting_block( + block, + &mut self.buffer, + &mut self.doc_ids, + &mut self.freqs, + block_size, + ); } self.block_idx = block_idx; self.position_block_idx = None; @@ -279,6 +286,7 @@ impl PostingIterator { list.blocks.len(), list.length, list.posting_tail_codec, + list.block_size, ); } compressed as *mut CompressedState @@ -307,8 +315,12 @@ impl PostingIterator { Some(max_score) => max_score, None => idf(list.len(), num_doc) * (K1 + 1.0), }; - - let is_compressed = matches!(list, PostingList::Compressed(_)); + let compressed = match &list { + PostingList::Compressed(list) => { + Some(UnsafeCell::new(CompressedState::new(list.block_size))) + } + PostingList::Plain(_) => None, + }; Self { token, @@ -319,7 +331,7 @@ impl PostingIterator { index: 0, block_idx: 0, approximate_upper_bound, - compressed: is_compressed.then(|| UnsafeCell::new(CompressedState::new())), + compressed, } } @@ -361,8 +373,8 @@ impl PostingIterator { match self.list { PostingList::Compressed(ref list) => { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index / list.block_size; + let block_offset = self.index % list.block_size; let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; // Read from the decompressed block @@ -400,8 +412,8 @@ impl PostingIterator { )) } CompressedPositionStorage::SharedStream(stream) => { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index / list.block_size; + let block_offset = self.index % list.block_size; let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; if compressed.position_block_idx != Some(block_idx) { @@ -441,33 +453,33 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - let mut block_idx = self.index / BLOCK_SIZE; + let mut block_idx = self.index / list.block_size; while block_idx + 1 < list.blocks.len() && list.block_least_doc_id(block_idx + 1) <= least_id { block_idx += 1; } - self.index = self.index.max(block_idx * BLOCK_SIZE); + self.index = self.index.max(block_idx * list.block_size); let length = list.length as usize; while self.index < length { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index / list.block_size; + let block_offset = self.index % list.block_size; let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; let in_block = &compressed.doc_ids[block_offset..]; let offset_in_block = in_block.partition_point(|&doc_id| doc_id < least_id); let new_offset = block_offset + offset_in_block; if new_offset < compressed.doc_ids.len() { - self.index = block_idx * BLOCK_SIZE + new_offset; + self.index = block_idx * list.block_size + new_offset; break; } if block_idx + 1 >= list.blocks.len() { self.index = length; break; } - self.index = (block_idx + 1) * BLOCK_SIZE; + self.index = (block_idx + 1) * list.block_size; } - self.block_idx = self.index / BLOCK_SIZE; + self.block_idx = self.index / list.block_size; } PostingList::Plain(ref list) => { self.index += list.row_ids[self.index..].partition_point(|&id| id < least_id); @@ -2146,6 +2158,7 @@ mod tests { max_score, doc_ids.len() as u32, crate::scalar::inverted::PostingTailCodec::VarintDelta, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, None, )) } else { diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 8c6c9d85de7..54abaea2de4 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1699,6 +1699,7 @@ impl FtsMemIndex { let st = self.state.load_full(); let with_position = self.params.has_positions(); + let block_size = self.params.posting_block_size(); let total_rows_u64 = total_rows as u64; // Step 1: collect (original_pos, num_tokens) for every doc across all @@ -1716,10 +1717,11 @@ impl FtsMemIndex { } } if all_docs.is_empty() { - return Ok(InnerBuilder::new( + return Ok(InnerBuilder::new_with_block_size( partition_id, with_position, Default::default(), + block_size, )); } @@ -1805,7 +1807,10 @@ impl FtsMemIndex { docs_for_term.sort_by_key(|(doc_id, _, _)| *doc_id); let token_id = tokens.add(token) as usize; debug_assert_eq!(token_id, posting_lists.len()); - posting_lists.push(PostingListBuilder::new(with_position)); + posting_lists.push(PostingListBuilder::new_with_block_size( + with_position, + block_size, + )); let plb = &mut posting_lists[token_id]; for (doc_id, freq, pos) in docs_for_term { let recorder = if with_position { @@ -1817,7 +1822,12 @@ impl FtsMemIndex { } } - let mut builder = InnerBuilder::new(partition_id, with_position, Default::default()); + let mut builder = InnerBuilder::new_with_block_size( + partition_id, + with_position, + Default::default(), + block_size, + ); builder.set_tokens(tokens); builder.set_docs(docs); builder.set_posting_lists(posting_lists); From 45f58e4fa4f4b318f614515912f26bb2ebf3c7dd Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Thu, 25 Jun 2026 18:36:29 +0800 Subject: [PATCH 02/20] fix(fts): reject unsupported posting block size --- docs/src/format/index/scalar/fts.md | 2 +- docs/src/quickstart/full-text-search.md | 2 +- .../index/scalar/InvertedIndexParams.java | 8 ++++---- .../index/scalar/InvertedIndexParamsTest.java | 6 ++++-- python/python/lance/dataset.py | 2 +- python/python/tests/test_scalar_index.py | 11 ++++++++-- .../src/scalar/inverted/cache_codec.rs | 4 ++-- .../src/scalar/inverted/encoding.rs | 2 +- rust/lance-index/src/scalar/inverted/index.rs | 3 ++- .../src/scalar/inverted/tokenizer.rs | 20 ++++++++++--------- 10 files changed, 36 insertions(+), 24 deletions(-) diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index 9c9550eb897..fc1e97f678c 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -69,7 +69,7 @@ The metadata file contains JSON-serialized configuration and partition informati | `min_gram` | UInt32 | 2 | Minimum n-gram length (only for ngram tokenizer) | | `max_gram` | UInt32 | 15 | Maximum n-gram length (only for ngram tokenizer) | | `prefix_only` | Boolean | false | Generate only prefix n-grams (only for ngram tokenizer) | -| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128, 256, or 512. Missing values from older indexes read as 128. | +| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. | ## Tokenizers diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index c72a22602ea..777a8310ad2 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -98,7 +98,7 @@ ds.create_scalar_index( remove_stop_words=True, # Remove stop words (language-dependent) custom_stop_words=None, # Optional additional stop words (only used if remove_stop_words=True) ascii_folding=True, # Fold accents to ASCII when possible (e.g., "é" -> "e") - block_size=256, # Posting block size: 128, 256, or 512 + block_size=256, # Posting block size: 128 or 256 ) ``` diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index 7e2d9699ae7..dcf0c6747a6 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -229,16 +229,16 @@ public Builder prefixOnly(boolean prefixOnly) { /** * Configure the number of documents in each compressed posting block. * - *

Supported values are {@code 128}, {@code 256}, and {@code 512}. New indexes default to - * {@code 256} when this is not set. + *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 256} + * when this is not set. * * @param blockSize posting block size * @return this builder * @throws IllegalArgumentException if {@code blockSize} is unsupported */ public Builder blockSize(int blockSize) { - if (blockSize != 128 && blockSize != 256 && blockSize != 512) { - throw new IllegalArgumentException("blockSize must be one of 128, 256, or 512"); + if (blockSize != 128 && blockSize != 256) { + throw new IllegalArgumentException("blockSize must be one of 128 or 256"); } this.blockSize = blockSize; return this; diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index 757d831c54a..20001f04e4e 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -44,16 +44,18 @@ void defaultBlockSizeIsSerialized() { @Test void blockSizeIsSerialized() { - ScalarIndexParams params = InvertedIndexParams.builder().blockSize(512).build(); + ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build(); assertEquals("inverted", params.getIndexType()); Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); - assertEquals(512, ((Number) json.get("block_size")).intValue()); + assertEquals(128, ((Number) json.get("block_size")).intValue()); } @Test void invalidBlockSizeIsRejected() { assertThrows( IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(129)); + assertThrows( + IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(512)); } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index f61fcc6918c..03ec1c52d04 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3317,7 +3317,7 @@ def create_scalar_index( True. block_size: int, default 256 This is for the ``INVERTED`` index. Number of documents per compressed - posting block. Must be one of ``128``, ``256``, or ``512``. + posting block. Must be one of ``128`` or ``256``. memory_limit: int, optional This is for the ``INVERTED`` index. Total build-time memory limit in MiB. If set, Lance divides this budget evenly across the workers. If unset, diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 682304fcd23..ec7c4985c3b 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -951,7 +951,7 @@ def test_create_scalar_index_fts_alias(dataset): def test_create_scalar_index_fts_block_size(dataset): dataset.create_scalar_index( - "doc", index_type="INVERTED", with_position=False, block_size=512 + "doc", index_type="INVERTED", with_position=False, block_size=256 ) row = dataset.take(indices=[0], columns=["doc"]) query = row.column(0)[0].as_py().split(" ")[0] @@ -959,7 +959,14 @@ def test_create_scalar_index_fts_block_size(dataset): assert results.num_rows > 0 with pytest.raises(ValueError, match="block_size"): - dataset.create_scalar_index("doc2", index_type="INVERTED", block_size=129) + dataset.create_scalar_index( + "doc", index_type="INVERTED", name="doc_invalid_129", block_size=129 + ) + + with pytest.raises(ValueError, match="block_size"): + dataset.create_scalar_index( + "doc", index_type="INVERTED", name="doc_invalid_512", block_size=512 + ) def test_multi_index_create(tmp_path): diff --git a/rust/lance-index/src/scalar/inverted/cache_codec.rs b/rust/lance-index/src/scalar/inverted/cache_codec.rs index afc2f00cffe..fffb080bda1 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -599,7 +599,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, - 512, + 256, Some(CompressedPositionStorage::SharedStream(stream)), ); let entry = PostingList::Compressed(posting.clone()); @@ -628,7 +628,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, - 512, + 256, Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 67567ec8cee..40eabe00169 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -943,7 +943,7 @@ mod tests { #[test] fn test_compress_posting_list_supported_block_sizes() -> Result<()> { - for block_size in [128, 256, 512] { + for block_size in [128, 256] { let num_rows: usize = block_size * 2 + 7; let doc_ids = (0..num_rows as u32).collect::>(); let frequencies = (0..num_rows as u32) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index b7b730930a8..2cc0c10b0f0 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -5936,7 +5936,8 @@ mod tests { assert_eq!(parse_posting_block_size(&HashMap::new()).unwrap(), 128); let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "512".to_owned())]); - assert_eq!(parse_posting_block_size(&metadata).unwrap(), 512); + let err = parse_posting_block_size(&metadata).unwrap_err(); + assert!(err.to_string().contains("block_size")); let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "129".to_owned())]); let err = parse_posting_block_size(&metadata).unwrap_err(); diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index c4181a65847..3cfce1ae279 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -31,7 +31,7 @@ use lance_tokenizer::{ pub const LEGACY_BLOCK_SIZE: usize = 128; pub const DEFAULT_BLOCK_SIZE: usize = 256; -pub const VALID_BLOCK_SIZES: [usize; 3] = [128, 256, 512]; +pub const VALID_BLOCK_SIZES: [usize; 2] = [128, 256]; /// Tokenizer configs #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -208,7 +208,7 @@ fn legacy_block_size() -> usize { } fn invalid_block_size_message(block_size: usize) -> String { - format!("FTS inverted index block_size must be one of 128, 256, or 512, got {block_size}") + format!("FTS inverted index block_size must be one of 128 or 256, got {block_size}") } pub fn validate_block_size(block_size: usize) -> Result { @@ -357,9 +357,8 @@ impl InvertedIndexParams { /// Set the compressed posting block size. /// - /// Supported values are 128, 256, and 512. Larger values reduce block-max - /// metadata and WAND skip granularity; smaller values preserve the legacy - /// layout. + /// Supported values are 128 and 256. Larger values reduce block-max metadata + /// and WAND skip granularity; smaller values preserve the legacy layout. pub fn block_size(mut self, block_size: usize) -> Result { self.block_size = validate_block_size(block_size)?; Ok(self) @@ -600,9 +599,9 @@ mod tests { #[test] fn test_block_size_details_conversion() { - let params = InvertedIndexParams::default().block_size(512).unwrap(); + let params = InvertedIndexParams::default().block_size(256).unwrap(); let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); - assert_eq!(details.block_size, Some(512)); + assert_eq!(details.block_size, Some(256)); let old_details = pbold::InvertedIndexDetails { base_tokenizer: Some("simple".to_string()), @@ -624,7 +623,7 @@ mod tests { #[test] fn test_block_size_accepts_supported_values() { - for block_size in [128, 256, 512] { + for block_size in [128, 256] { let params = InvertedIndexParams::default() .block_size(block_size) .unwrap(); @@ -641,12 +640,15 @@ mod tests { let err = InvertedIndexParams::default().block_size(129).unwrap_err(); assert!(err.to_string().contains("block_size")); + let err = InvertedIndexParams::default().block_size(512).unwrap_err(); + assert!(err.to_string().contains("128 or 256")); + let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); json.as_object_mut() .unwrap() .insert("block_size".to_string(), serde_json::Value::from(1024)); let err = serde_json::from_value::(json).unwrap_err(); - assert!(err.to_string().contains("128, 256, or 512")); + assert!(err.to_string().contains("128 or 256")); } #[test] From db2bcf17396c46d169852b0c705edee48667f861 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Thu, 25 Jun 2026 18:44:06 +0800 Subject: [PATCH 03/20] style(java): format inverted index javadocs --- .../main/java/org/lance/index/scalar/InvertedIndexParams.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index dcf0c6747a6..97d75bfc20e 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -229,8 +229,8 @@ public Builder prefixOnly(boolean prefixOnly) { /** * Configure the number of documents in each compressed posting block. * - *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 256} - * when this is not set. + *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 256} when + * this is not set. * * @param blockSize posting block size * @return this builder From 120ab7cf4aa8e93f0fdc287de70db709de17a4a5 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Sat, 27 Jun 2026 00:07:43 +0800 Subject: [PATCH 04/20] feat(fts): use physical 256 posting blocks --- .../src/bitpacker_internal/bitpacker8x.rs | 7 +- .../bitpacking/src/bitpacker_internal/mod.rs | 1 + rust/compression/bitpacking/src/lib.rs | 2 +- .../src/scalar/inverted/builder.rs | 5 +- .../src/scalar/inverted/encoding.rs | 134 +++++++++++++++--- rust/lance-index/src/scalar/inverted/index.rs | 2 +- rust/lance-index/src/scalar/inverted/iter.rs | 9 +- rust/lance-index/src/scalar/inverted/wand.rs | 10 +- 8 files changed, 130 insertions(+), 40 deletions(-) diff --git a/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs index 188a1f4ce2a..b17edacabb2 100644 --- a/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs +++ b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs @@ -420,12 +420,11 @@ enum InstructionSet { Scalar, } -/// Internal 8-wide bitpacker implementation. +/// 8-wide bitpacker implementation. /// -/// One block contains 256 integers. This stays private to avoid exposing a new -/// block-size choice through the public Lance bitpacking API. +/// One block contains 256 integers. #[derive(Clone, Copy)] -pub(crate) struct BitPacker8x(InstructionSet); +pub struct BitPacker8x(InstructionSet); impl BitPacker8x { #[cfg(target_arch = "x86_64")] diff --git a/rust/compression/bitpacking/src/bitpacker_internal/mod.rs b/rust/compression/bitpacking/src/bitpacker_internal/mod.rs index c287a29da0b..80803e50ec8 100644 --- a/rust/compression/bitpacking/src/bitpacker_internal/mod.rs +++ b/rust/compression/bitpacking/src/bitpacker_internal/mod.rs @@ -20,6 +20,7 @@ mod bitpacker4x; mod bitpacker8x; pub use bitpacker4x::BitPacker4x; +pub use bitpacker8x::BitPacker8x; pub(crate) trait Available { fn available() -> bool; diff --git a/rust/compression/bitpacking/src/lib.rs b/rust/compression/bitpacking/src/lib.rs index df2313bed5d..c4163a5324d 100644 --- a/rust/compression/bitpacking/src/lib.rs +++ b/rust/compression/bitpacking/src/lib.rs @@ -18,7 +18,7 @@ use core::mem::size_of; mod bitpacker_internal; -pub use bitpacker_internal::{BitPacker, BitPacker4x}; +pub use bitpacker_internal::{BitPacker, BitPacker4x, BitPacker8x}; pub const FL_ORDER: [usize; 8] = [0, 4, 2, 6, 1, 5, 3, 7]; diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 9324b706fc6..504f1ab1b8a 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -42,8 +42,9 @@ use std::task::{Context, Poll}; use std::{fmt::Debug, sync::atomic::AtomicU64}; use tracing::instrument; -// The physical bitpacking block size. Logical FTS posting blocks may contain -// 1, 2, or 4 of these chunks, configured by InvertedIndexParams::block_size. +// The legacy bitpacking block size. Position streams still use this block size; +// FTS posting blocks choose their physical bitpacker from the configured +// InvertedIndexParams::block_size. pub const BLOCK_SIZE: usize = BitPacker4x::BLOCK_LEN; // The default number of workers to use for FTS builds. diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 8a9174cfd75..28ece2a0d32 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -9,9 +9,11 @@ use super::index::{PositionStreamCodec, PostingTailCodec}; use super::tokenizer::LEGACY_BLOCK_SIZE; use super::tokenizer::validate_block_size; use arrow::array::LargeBinaryBuilder; -use lance_bitpacking::{BitPacker, BitPacker4x}; +use lance_bitpacking::{BitPacker, BitPacker4x, BitPacker8x}; use lance_core::{Error, Result}; +pub const MAX_POSTING_BLOCK_SIZE: usize = BitPacker8x::BLOCK_LEN; + // we compress the posting list to multiple blocks of fixed number of elements (BLOCK_SIZE), // returns a LargeBinaryArray, where each binary is a compressed block (128 row ids + 128 frequencies) // each block is: @@ -137,14 +139,18 @@ fn encode_posting_block_payload( block: &mut impl Write, ) -> Result<()> { debug_assert_eq!(doc_ids.len(), frequencies.len()); - debug_assert!(doc_ids.len().is_multiple_of(BLOCK_SIZE)); - let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - for (doc_ids, frequencies) in doc_ids - .chunks_exact(BLOCK_SIZE) - .zip(frequencies.chunks_exact(BLOCK_SIZE)) - { - compress_sorted_block(doc_ids, &mut buffer, block)?; - compress_block(frequencies, &mut buffer, block)?; + validate_block_size(doc_ids.len())?; + let mut buffer = [0u8; MAX_POSTING_BLOCK_SIZE * 4 + 5]; + match doc_ids.len() { + BitPacker4x::BLOCK_LEN => { + compress_sorted_block_with::(doc_ids, &mut buffer, block)?; + compress_block_with::(frequencies, &mut buffer, block)?; + } + BitPacker8x::BLOCK_LEN => { + compress_sorted_block_with::(doc_ids, &mut buffer, block)?; + compress_block_with::(frequencies, &mut buffer, block)?; + } + _ => unreachable!("validated posting block size should be supported"), } Ok(()) } @@ -163,7 +169,17 @@ pub fn encode_remainder_posting_block_into( #[inline] fn compress_sorted_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Write) -> Result<()> { - let compressor = BitPacker4x::new(); + compress_sorted_block_with::(data, buffer, builder) +} + +#[inline] +fn compress_sorted_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let compressor = P::new(); let num_bits = compressor.num_bits_sorted(data[0], data); let num_bytes = compressor.compress_sorted(data[0], data, buffer, num_bits); let _ = builder.write(data[0].to_le_bytes().as_ref())?; @@ -174,7 +190,17 @@ fn compress_sorted_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Wri #[inline] fn compress_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Write) -> Result<()> { - let compressor = BitPacker4x::new(); + compress_block_with::(data, buffer, builder) +} + +#[inline] +fn compress_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let compressor = P::new(); let num_bits = compressor.num_bits(data); let num_bytes = compressor.compress(data, buffer, num_bits); let _ = builder.write(&[num_bits])?; @@ -704,7 +730,7 @@ pub fn decompress_posting_list_with_tail_codec_and_block_size( let mut doc_ids: Vec = Vec::with_capacity(num_docs as usize); let mut frequencies: Vec = Vec::with_capacity(num_docs as usize); - let mut buffer = [0u32; BLOCK_SIZE]; + let mut buffer = [0u32; MAX_POSTING_BLOCK_SIZE]; let bitpacking_blocks = num_docs as usize / block_size; for compressed in posting_list.iter().take(bitpacking_blocks) { let compressed = compressed.unwrap(); @@ -758,19 +784,29 @@ pub fn read_num_positions(compressed: &arrow::array::LargeBinaryArray) -> u32 { pub fn decompress_posting_block( block: &[u8], - buffer: &mut [u32; BLOCK_SIZE], + buffer: &mut [u32], doc_ids: &mut Vec, frequencies: &mut Vec, block_size: usize, ) { debug_assert!(validate_block_size(block_size).is_ok()); + debug_assert!(buffer.len() >= block_size); // skip the first 4 bytes for the max block score let mut block = &block[4..]; - for _ in 0..(block_size / BLOCK_SIZE) { - let num_bytes = decompress_sorted_block(block, buffer, doc_ids); - block = &block[num_bytes..]; - let num_bytes = decompress_block(block, buffer, frequencies); - block = &block[num_bytes..]; + match block_size { + BitPacker4x::BLOCK_LEN => { + let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); + block = &block[num_bytes..]; + let num_bytes = decompress_block_with::(block, buffer, frequencies); + block = &block[num_bytes..]; + } + BitPacker8x::BLOCK_LEN => { + let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); + block = &block[num_bytes..]; + let num_bytes = decompress_block_with::(block, buffer, frequencies); + block = &block[num_bytes..]; + } + _ => unreachable!("validated posting block size should be supported"), } debug_assert!( block.is_empty(), @@ -829,7 +865,7 @@ pub fn decode_full_posting_block( frequencies: &mut Vec, block_size: usize, ) { - let mut buffer = [0u32; BLOCK_SIZE]; + let mut buffer = [0u32; MAX_POSTING_BLOCK_SIZE]; decompress_posting_block(block, &mut buffer, doc_ids, frequencies, block_size); } @@ -838,7 +874,17 @@ pub fn decompress_sorted_block( buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec, ) -> usize { - let compressor = BitPacker4x::new(); + decompress_sorted_block_with::(block, buffer, res) +} + +fn decompress_sorted_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let compressor = P::new(); let initial = u32::from_le_bytes(block[0..4].try_into().unwrap()); let num_bits = block[4]; let num_bytes = compressor.decompress_sorted(initial, &block[5..], buffer, num_bits); @@ -846,8 +892,14 @@ pub fn decompress_sorted_block( 5 + num_bytes } -fn decompress_block(block: &[u8], buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec) -> usize { - let compressor = BitPacker4x::new(); +fn decompress_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let compressor = P::new(); let num_bits = block[0]; let num_bytes = compressor.decompress(&block[1..], buffer, num_bits); res.extend_from_slice(&buffer[..]); @@ -975,6 +1027,44 @@ mod tests { Ok(()) } + #[test] + fn test_256_posting_block_uses_single_physical_bitpack_chunk() -> Result<()> { + let block_size = BitPacker8x::BLOCK_LEN; + let doc_ids = (0..block_size as u32).collect::>(); + let frequencies = (0..block_size as u32) + .map(|value| value % 13 + 1) + .collect::>(); + + let posting_list = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + std::iter::once(1.0), + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(posting_list.len(), 1); + + let block = posting_list.value(0); + let doc_num_bits = block[8]; + let doc_bytes = BitPacker8x::compressed_block_size(doc_num_bits); + let freq_header_offset = 9 + doc_bytes; + let freq_num_bits = block[freq_header_offset]; + let freq_bytes = BitPacker8x::compressed_block_size(freq_num_bits); + assert_eq!(block.len(), freq_header_offset + 1 + freq_bytes); + + let (decoded_doc_ids, decoded_frequencies) = + decompress_posting_list_with_tail_codec_and_block_size( + doc_ids.len() as u32, + &posting_list, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(decoded_doc_ids, doc_ids); + assert_eq!(decoded_frequencies, frequencies); + Ok(()) + } + #[test] fn test_compress_posting_list_fixed32_tail_still_roundtrips() -> Result<()> { let doc_ids = vec![3_u32, 10_u32, 24_u32]; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 2cc0c10b0f0..48695d624e2 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -3794,7 +3794,7 @@ pub struct CompressedPostingList { pub length: u32, // each binary is a block of compressed data // that contains `block_size` doc ids and then `block_size` frequencies, - // packed internally in 128-value bitpacking chunks. + // packed by the physical bitpacker matching that block size. pub blocks: LargeBinaryArray, pub posting_tail_codec: PostingTailCodec, pub block_size: usize, diff --git a/rust/lance-index/src/scalar/inverted/iter.rs b/rust/lance-index/src/scalar/inverted/iter.rs index 52c09fa8dab..51970c46c33 100644 --- a/rust/lance-index/src/scalar/inverted/iter.rs +++ b/rust/lance-index/src/scalar/inverted/iter.rs @@ -6,10 +6,9 @@ use arrow_array::{Array, LargeBinaryArray}; use super::{ CompressedPositionStorage, PostingList, PostingTailCodec, - builder::BLOCK_SIZE, encoding::{ - decode_position_stream_block, decompress_positions, decompress_posting_block, - decompress_posting_remainder, + MAX_POSTING_BLOCK_SIZE, decode_position_stream_block, decompress_positions, + decompress_posting_block, decompress_posting_remainder, }, }; @@ -90,7 +89,7 @@ pub struct CompressedPostingListIterator { doc_idx_in_block: usize, decoded_positions: Vec, position_offsets: Vec, - buffer: [u32; BLOCK_SIZE], + buffer: [u32; MAX_POSTING_BLOCK_SIZE], } impl CompressedPostingListIterator { @@ -123,7 +122,7 @@ impl CompressedPostingListIterator { doc_idx_in_block: 0, decoded_positions: Vec::new(), position_offsets: Vec::new(), - buffer: [0; BLOCK_SIZE], + buffer: [0; MAX_POSTING_BLOCK_SIZE], } } } diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 1f51afa0e48..b94ac61aa0a 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -29,8 +29,8 @@ use super::{ CompressedPostingList, DocSet, PostingList, RawDocInfo, builder::ScoredDoc, encoding::{ - decode_position_stream_block, decompress_positions, decompress_posting_block, - decompress_posting_remainder, + MAX_POSTING_BLOCK_SIZE, decode_position_stream_block, decompress_positions, + decompress_posting_block, decompress_posting_remainder, }, query::FtsSearchParams, scorer::Scorer, @@ -66,7 +66,7 @@ struct CompressedState { block_idx: usize, doc_ids: Vec, freqs: Vec, - buffer: Box<[u32; BLOCK_SIZE]>, + buffer: Box<[u32; MAX_POSTING_BLOCK_SIZE]>, position_block_idx: Option, position_values: Vec, position_offsets: Vec, @@ -79,7 +79,7 @@ impl CompressedState { block_idx: 0, doc_ids: Vec::with_capacity(block_size), freqs: Vec::with_capacity(block_size), - buffer: Box::new([0; BLOCK_SIZE]), + buffer: Box::new([0; MAX_POSTING_BLOCK_SIZE]), position_block_idx: None, position_values: Vec::new(), position_offsets: Vec::new(), @@ -112,7 +112,7 @@ impl CompressedState { } else { decompress_posting_block( block, - &mut self.buffer, + &mut self.buffer[..], &mut self.doc_ids, &mut self.freqs, block_size, From 000c857885c2fc46a0101caaa7ad6774bf53aa49 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Sat, 27 Jun 2026 00:55:03 +0800 Subject: [PATCH 05/20] test(fts): keep compat indexes on legacy block size --- python/python/tests/compat/test_scalar_indices.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index 35022df3b12..67f07299e6e 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -9,6 +9,7 @@ and written by other versions. """ +import os import shutil from pathlib import Path @@ -320,7 +321,12 @@ def create(self): max_rows_per_file=100, data_storage_version=safe_data_storage_version(self.compat_version), ) - dataset.create_scalar_index("text", "INVERTED", with_position=True) + kwargs = {"with_position": True} + # Downgrade reads use older wheels, so current-created FTS indexes must + # stay on the legacy posting block layout. + if os.environ.get("LANCE_COMPAT_FTS_LEGACY_BLOCK_SIZE") == "1": + kwargs["block_size"] = 128 + dataset.create_scalar_index("text", "INVERTED", **kwargs) def check_read(self): """Verify FTS index can be queried.""" @@ -351,7 +357,10 @@ def skip_downgrade(self, version: str) -> bool: def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": - return {"LANCE_FTS_FORMAT_VERSION": "1"} + return { + "LANCE_COMPAT_FTS_LEGACY_BLOCK_SIZE": "1", + "LANCE_FTS_FORMAT_VERSION": "1", + } if method_name == "check_write": return {"LANCE_FTS_FORMAT_VERSION": "2"} return {} From 9d8d77e67f9900d3bc6d304d89271b571fb72db7 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Mon, 29 Jun 2026 17:39:26 +0800 Subject: [PATCH 06/20] docs(fts): mark 256 block size experimental --- docs/src/format/index/scalar/fts.md | 2 +- docs/src/quickstart/full-text-search.md | 2 +- .../java/org/lance/index/scalar/InvertedIndexParams.java | 3 +++ python/python/lance/dataset.py | 3 +++ rust/lance-index/src/scalar/inverted/tokenizer.rs | 5 +++++ 5 files changed, 13 insertions(+), 2 deletions(-) diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index fc1e97f678c..40ffbf6b450 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -69,7 +69,7 @@ The metadata file contains JSON-serialized configuration and partition informati | `min_gram` | UInt32 | 2 | Minimum n-gram length (only for ngram tokenizer) | | `max_gram` | UInt32 | 15 | Maximum n-gram length (only for ngram tokenizer) | | `prefix_only` | Boolean | false | Generate only prefix n-grams (only for ngram tokenizer) | -| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. | +| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. `256` is experimental and may introduce breaking changes. | ## Tokenizers diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index 777a8310ad2..a06366c557d 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -98,7 +98,7 @@ ds.create_scalar_index( remove_stop_words=True, # Remove stop words (language-dependent) custom_stop_words=None, # Optional additional stop words (only used if remove_stop_words=True) ascii_folding=True, # Fold accents to ASCII when possible (e.g., "é" -> "e") - block_size=256, # Posting block size: 128 or 256 + block_size=256, # Posting block size: 128 or 256; 256 is experimental ) ``` diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index 97d75bfc20e..f8d674a7d8b 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -232,6 +232,9 @@ public Builder prefixOnly(boolean prefixOnly) { *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 256} when * this is not set. * + *

{@code blockSize = 256} is experimental and may introduce breaking changes. Use {@code + * 128} when stable compatibility with the legacy posting layout is required. + * * @param blockSize posting block size * @return this builder * @throws IllegalArgumentException if {@code blockSize} is unsupported diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 03ec1c52d04..8e10cdda5f6 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3318,6 +3318,9 @@ def create_scalar_index( block_size: int, default 256 This is for the ``INVERTED`` index. Number of documents per compressed posting block. Must be one of ``128`` or ``256``. + ``block_size=256`` is experimental and may introduce breaking changes. + Use ``128`` when stable compatibility with the legacy posting layout is + required. memory_limit: int, optional This is for the ``INVERTED`` index. Total build-time memory limit in MiB. If set, Lance divides this budget evenly across the workers. If unset, diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 3cfce1ae279..3bdfedfe717 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -359,12 +359,17 @@ impl InvertedIndexParams { /// /// Supported values are 128 and 256. Larger values reduce block-max metadata /// and WAND skip granularity; smaller values preserve the legacy layout. + /// + /// `block_size = 256` is experimental and may introduce breaking changes. + /// Use `128` when stable compatibility with the legacy posting layout is required. pub fn block_size(mut self, block_size: usize) -> Result { self.block_size = validate_block_size(block_size)?; Ok(self) } /// Get the compressed posting block size. + /// + /// `256` is experimental and may introduce breaking changes. pub fn posting_block_size(&self) -> usize { self.block_size } From f06f2c4ba9b0d6fdac99a68ce8b49190fd27a6f2 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Mon, 29 Jun 2026 20:44:31 +0800 Subject: [PATCH 07/20] fix(fts): keep posting block size default at 128 --- docs/src/format/index/scalar/fts.md | 2 +- docs/src/quickstart/full-text-search.md | 2 +- .../org/lance/index/scalar/InvertedIndexParams.java | 4 ++-- .../lance/index/scalar/InvertedIndexParamsTest.java | 2 +- python/python/lance/dataset.py | 2 +- rust/lance-index/src/scalar/inverted.rs | 2 +- rust/lance-index/src/scalar/inverted/tokenizer.rs | 10 +++++----- 7 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index 40ffbf6b450..d5c75158011 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -69,7 +69,7 @@ The metadata file contains JSON-serialized configuration and partition informati | `min_gram` | UInt32 | 2 | Minimum n-gram length (only for ngram tokenizer) | | `max_gram` | UInt32 | 15 | Maximum n-gram length (only for ngram tokenizer) | | `prefix_only` | Boolean | false | Generate only prefix n-grams (only for ngram tokenizer) | -| `block_size` | UInt32 | 256 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. `256` is experimental and may introduce breaking changes. | +| `block_size` | UInt32 | 128 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. `256` is experimental and may introduce breaking changes. | ## Tokenizers diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index a06366c557d..7f09c4325b7 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -98,7 +98,7 @@ ds.create_scalar_index( remove_stop_words=True, # Remove stop words (language-dependent) custom_stop_words=None, # Optional additional stop words (only used if remove_stop_words=True) ascii_folding=True, # Fold accents to ASCII when possible (e.g., "é" -> "e") - block_size=256, # Posting block size: 128 or 256; 256 is experimental + block_size=128, # Posting block size: 128 or 256; 256 is experimental ) ``` diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index f8d674a7d8b..023eb1813e1 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -53,7 +53,7 @@ public static final class Builder { private Integer minNgramLength; private Integer maxNgramLength; private Boolean prefixOnly; - private Integer blockSize = 256; + private Integer blockSize = 128; private Boolean skipMerge; /** @@ -229,7 +229,7 @@ public Builder prefixOnly(boolean prefixOnly) { /** * Configure the number of documents in each compressed posting block. * - *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 256} when + *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 128} when * this is not set. * *

{@code blockSize = 256} is experimental and may introduce breaking changes. Use {@code diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index 20001f04e4e..f067579dfcd 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -39,7 +39,7 @@ void defaultBlockSizeIsSerialized() { ScalarIndexParams params = InvertedIndexParams.builder().build(); Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); - assertEquals(256, ((Number) json.get("block_size")).intValue()); + assertEquals(128, ((Number) json.get("block_size")).intValue()); } @Test diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 8e10cdda5f6..701bbd05aa8 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3315,7 +3315,7 @@ def create_scalar_index( query. This will significantly increase the index size. It won't impact the performance of non-phrase queries even if it is set to True. - block_size: int, default 256 + block_size: int, default 128 This is for the ``INVERTED`` index. Number of documents per compressed posting block. Must be one of ``128`` or ``256``. ``block_size=256`` is experimental and may introduce breaking changes. diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 8300dc1043c..72a48b32004 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -320,7 +320,7 @@ mod tests { } #[test] - fn test_new_training_request_defaults_missing_block_size_to_256() { + fn test_new_training_request_defaults_missing_block_size_to_128() { let plugin = InvertedIndexPlugin; let field = Field::new("text", DataType::Utf8, true); diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 3bdfedfe717..bac9dcaa064 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -30,7 +30,7 @@ use lance_tokenizer::{ }; pub const LEGACY_BLOCK_SIZE: usize = 128; -pub const DEFAULT_BLOCK_SIZE: usize = 256; +pub const DEFAULT_BLOCK_SIZE: usize = 128; pub const VALID_BLOCK_SIZES: [usize; 2] = [128, 256]; /// Tokenizer configs @@ -106,7 +106,7 @@ pub struct InvertedIndexParams { /// /// Missing serialized values come from indexes written before this /// parameter existed and must read as 128 for backwards compatibility. New - /// indexes default to 256. + /// indexes currently default to 128. #[serde( default = "legacy_block_size", deserialize_with = "deserialize_block_size" @@ -586,11 +586,11 @@ mod tests { } #[test] - fn test_block_size_new_default_serializes() { + fn test_block_size_default_serializes() { let params = InvertedIndexParams::default(); - assert_eq!(params.block_size, 256); + assert_eq!(params.block_size, 128); let json = serde_json::to_value(¶ms).unwrap(); - assert_eq!(json.get("block_size"), Some(&serde_json::Value::from(256))); + assert_eq!(json.get("block_size"), Some(&serde_json::Value::from(128))); } #[test] From dbea49ed687b850e2519a994bf3525f0a34b6e3a Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 30 Jun 2026 20:18:08 +0800 Subject: [PATCH 08/20] fix(fts): gate 256 block size behind v3 --- .../index/scalar/InvertedIndexParams.java | 16 +- .../index/scalar/InvertedIndexParamsTest.java | 17 ++ python/python/lance/dataset.py | 6 +- python/python/tests/test_scalar_index.py | 16 ++ python/src/dataset.rs | 2 +- rust/lance-index/src/scalar/inverted.rs | 1 + .../src/scalar/inverted/builder.rs | 83 +++++-- rust/lance-index/src/scalar/inverted/index.rs | 222 ++++++++++++++++-- .../src/scalar/inverted/tokenizer.rs | 80 ++++++- rust/lance/src/dataset/mem_wal/index.rs | 8 +- .../src/dataset/mem_wal/memtable/flush.rs | 13 +- 11 files changed, 401 insertions(+), 63 deletions(-) diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index eccfe33082a..a688626d18d 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -264,15 +264,17 @@ public Builder skipMerge(boolean skipMerge) { /** * Configure the on-disk FTS format version to write when creating a new index. * - *

If unset, Lance chooses the current default format. + *

If unset, Lance writes v2 for {@code blockSize = 128} and v3 for {@code blockSize = + * 256}. {@code formatVersion = 3} is experimental and is only valid with {@code blockSize = + * 256}. * - * @param formatVersion FTS format version, must be 1 or 2 + * @param formatVersion FTS format version, must be 1, 2, or 3 * @return this builder * @throws IllegalArgumentException */ public Builder formatVersion(int formatVersion) { - if (formatVersion != 1 && formatVersion != 2) { - throw new IllegalArgumentException("formatVersion must be 1 or 2"); + if (formatVersion != 1 && formatVersion != 2 && formatVersion != 3) { + throw new IllegalArgumentException("formatVersion must be 1, 2, or 3"); } this.formatVersion = formatVersion; return this; @@ -280,6 +282,12 @@ public Builder formatVersion(int formatVersion) { /** Build a {@link ScalarIndexParams} instance for an inverted index. */ public ScalarIndexParams build() { + if (formatVersion != null) { + Preconditions.checkArgument( + (blockSize == 256 && formatVersion == 3) + || (blockSize == 128 && formatVersion != 3), + "formatVersion 3 requires blockSize 256, and blockSize 256 requires formatVersion 3"); + } Map params = new HashMap<>(); if (baseTokenizer != null) { params.put("base_tokenizer", baseTokenizer); diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index f067579dfcd..0ccc429ec57 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -58,4 +58,21 @@ void invalidBlockSizeIsRejected() { assertThrows( IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(512)); } + + @Test + void formatVersionThreeRequiresBlockSize256() { + ScalarIndexParams params = + InvertedIndexParams.builder().blockSize(256).formatVersion(3).build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(256, ((Number) json.get("block_size")).intValue()); + assertEquals(3, ((Number) json.get("format_version")).intValue()); + + assertThrows( + IllegalArgumentException.class, + () -> InvertedIndexParams.builder().formatVersion(3).build()); + assertThrows( + IllegalArgumentException.class, + () -> InvertedIndexParams.builder().blockSize(256).formatVersion(2).build()); + } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index be55a3a2eb0..0985ca62a19 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3312,8 +3312,10 @@ def create_scalar_index( format_version: int or str, optional This is for the ``INVERTED`` / ``FTS`` index. Explicit on-disk FTS format version to write when creating a new index. Accepts ``1``, - ``2``, ``"v1"``, or ``"v2"``. If unset, Lance chooses the current - default format. + ``2``, ``3``, ``"v1"``, ``"v2"``, or ``"v3"``. If unset, Lance + writes v2 for ``block_size=128`` and v3 for ``block_size=256``. + ``format_version=3`` is experimental and is only valid with + ``block_size=256``. with_position: bool, default False This is for the ``INVERTED`` index. If True, the index will store the diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 1504ded0e0a..4866f39c3d2 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -953,6 +953,10 @@ def test_create_scalar_index_fts_block_size(dataset): dataset.create_scalar_index( "doc", index_type="INVERTED", with_position=False, block_size=256 ) + indices = dataset.describe_indices() + doc_index = next(index for index in indices if index.name == "doc_idx") + assert doc_index.segments[0].index_version == 3 + row = dataset.take(indices=[0], columns=["doc"]) query = row.column(0)[0].as_py().split(" ")[0] results = dataset.scanner(columns=["doc"], full_text_query=query).to_table() @@ -968,6 +972,15 @@ def test_create_scalar_index_fts_block_size(dataset): "doc", index_type="INVERTED", name="doc_invalid_512", block_size=512 ) + with pytest.raises(ValueError, match="block_size=256"): + dataset.create_scalar_index( + "doc", + index_type="INVERTED", + name="doc_invalid_v2_256", + block_size=256, + format_version=2, + ) + def test_multi_index_create(tmp_path): dataset = lance.write_dataset( @@ -5020,6 +5033,9 @@ def test_create_inverted_index_rejects_invalid_format_version(tmp_path): ds = lance.write_dataset(data, tmp_path) with pytest.raises(ValueError, match="unsupported FTS format version"): + ds.create_scalar_index("text", index_type="INVERTED", format_version="v4") + + with pytest.raises(ValueError, match="format_version=3"): ds.create_scalar_index("text", index_type="INVERTED", format_version="v3") diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 3534c9ac1de..eaa9d0014e6 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2450,7 +2450,7 @@ impl Dataset { value.to_string() } else { return Err(PyValueError::new_err( - "format_version must be 1, 2, 'v1', or 'v2'", + "format_version must be 1, 2, 3, 'v1', 'v2', or 'v3'", )); }; let format_version = value diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index fb763789224..68b2b662785 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -147,6 +147,7 @@ impl InvertedIndexPlugin { } }); + params.validate_format_version()?; let format_version = params.resolved_format_version(); let details = pbold::InvertedIndexDetails::try_from(¶ms)?; let mut inverted_index = diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index f1d62e9f423..3301c1b1277 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -5,8 +5,8 @@ use super::encoding::encode_group_starts; use super::{InvertedIndexParams, index::*}; use crate::scalar::inverted::document_tokenizer::DocType; use crate::scalar::inverted::json::JsonTextStream; +use crate::scalar::inverted::tokenizer::LEGACY_BLOCK_SIZE; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; -use crate::scalar::inverted::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; #[cfg(test)] use crate::scalar::lance_format::LanceIndexStore; use crate::scalar::{IndexFile, IndexStore, OldIndexDataFilter}; @@ -263,8 +263,11 @@ impl InvertedIndexBuilder { } pub fn with_posting_tail_codec(mut self, posting_tail_codec: PostingTailCodec) -> Self { - self.format_version = - InvertedListFormatVersion::from_posting_tail_codec(posting_tail_codec); + self.format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + self.params.block_size, + ) + .expect("invalid posting tail codec for posting block size"); self.posting_tail_codec = posting_tail_codec; self } @@ -291,6 +294,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, old_data_filter: Option, ) -> Result> { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let schema = new_data.schema(); let doc_col = schema.field(0).name(); @@ -325,6 +329,7 @@ impl InvertedIndexBuilder { old_segments: &[Arc], old_data_filter: Option, ) -> Result> { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let schema = new_data.schema(); let doc_col = schema.field(0).name(); @@ -589,6 +594,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, partitions: &[u64], ) -> Result { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let mut serialized_deleted_fragments = Vec::with_capacity(self.deleted_fragments.serialized_size()); self.deleted_fragments @@ -605,6 +611,14 @@ impl InvertedIndexBuilder { POSTING_TAIL_CODEC_KEY.to_owned(), self.posting_tail_codec.as_str().to_owned(), ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + self.format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + self.params.block_size.to_string(), + ), ]); if self.params.with_position && self.format_version.uses_shared_position_stream() { @@ -649,6 +663,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, partition: u64, // Modify parameter type ) -> Result { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let partitions = vec![partition]; let mut metadata = HashMap::from_iter(vec![ ("partitions".to_owned(), serde_json::to_string(&partitions)?), @@ -661,6 +676,14 @@ impl InvertedIndexBuilder { POSTING_TAIL_CODEC_KEY.to_owned(), self.posting_tail_codec.as_str().to_owned(), ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + self.format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + self.params.block_size.to_string(), + ), ]); if self.params.with_position && self.format_version.uses_shared_position_stream() { metadata.insert( @@ -835,11 +858,13 @@ impl InnerBuilder { token_set_format: TokenSetFormat, block_size: usize, ) -> Self { + let format_version = default_fts_format_version_for_block_size(block_size) + .expect("invalid posting list block size"); Self::new_with_format_version_and_block_size( id, with_position, token_set_format, - current_fts_format_version(), + format_version, block_size, ) } @@ -851,7 +876,8 @@ impl InnerBuilder { format_version: InvertedListFormatVersion, block_size: usize, ) -> Self { - validate_block_size(block_size).expect("invalid posting list block size"); + validate_format_version_block_size(format_version, block_size) + .expect("invalid FTS format version for posting block size"); Self { id, with_position, @@ -888,11 +914,11 @@ impl InnerBuilder { posting_tail_codec: PostingTailCodec, block_size: usize, ) -> Self { - let format_version = if posting_tail_codec == PostingTailCodec::Fixed32 { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; + let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + block_size, + ) + .expect("invalid posting tail codec for posting block size"); let mut builder = Self::new_with_format_version_and_block_size( id, with_position, @@ -1708,15 +1734,19 @@ pub fn inverted_list_schema_for_version_with_block_size( format_version: InvertedListFormatVersion, block_size: usize, ) -> SchemaRef { - validate_block_size(block_size).expect("invalid posting list block size"); + validate_format_version_block_size(format_version, block_size) + .expect("invalid FTS format version for posting block size"); match format_version { InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position, block_size), - InvertedListFormatVersion::V2 => inverted_list_schema_with_tail_codec_and_position_codec( - with_position, - PostingTailCodec::VarintDelta, - Some(PositionStreamCodec::PackedDelta), - block_size, - ), + InvertedListFormatVersion::V2 | InvertedListFormatVersion::V3 => { + inverted_list_schema_with_tail_codec_and_position_codec( + with_position, + format_version, + PostingTailCodec::VarintDelta, + Some(PositionStreamCodec::PackedDelta), + block_size, + ) + } } } @@ -1751,7 +1781,13 @@ fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef } Arc::new(arrow_schema::Schema::new_with_metadata( fields, - HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string())]), + HashMap::from([ + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + InvertedListFormatVersion::V1.index_version().to_string(), + ), + ]), )) } @@ -1759,8 +1795,14 @@ pub fn inverted_list_schema_with_tail_codec( with_position: bool, posting_tail_codec: PostingTailCodec, ) -> SchemaRef { + let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + LEGACY_BLOCK_SIZE, + ) + .expect("invalid posting tail codec for posting block size"); inverted_list_schema_with_tail_codec_and_position_codec( with_position, + format_version, posting_tail_codec, Some(PositionStreamCodec::PackedDelta), LEGACY_BLOCK_SIZE, @@ -1769,6 +1811,7 @@ pub fn inverted_list_schema_with_tail_codec( fn inverted_list_schema_with_tail_codec_and_position_codec( with_position: bool, + format_version: InvertedListFormatVersion, posting_tail_codec: PostingTailCodec, position_codec: Option, block_size: usize, @@ -1808,6 +1851,10 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( POSTING_TAIL_CODEC_KEY.to_owned(), posting_tail_codec.as_str().to_owned(), )]); + metadata.insert( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ); metadata.insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); if let Some(position_codec) = position_codec.filter(|_| with_position) { metadata.insert( diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index b3be8fb1a3e..db39f9e8bdc 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -86,8 +86,10 @@ 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: Version 2 layout with 256-document physical posting blocks. 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"; @@ -110,6 +112,7 @@ pub const NUM_TOKEN_COL: &str = "_num_tokens"; 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"; @@ -153,7 +156,7 @@ pub fn current_fts_format_version() -> InvertedListFormatVersion { } pub fn max_supported_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V2 + InvertedListFormatVersion::V3 } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] @@ -161,6 +164,7 @@ pub enum InvertedListFormatVersion { V1, #[default] V2, + V3, } impl InvertedListFormatVersion { @@ -171,29 +175,50 @@ impl InvertedListFormatVersion { } } + 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 => PostingTailCodec::VarintDelta, + Self::V2 | Self::V3 => PostingTailCodec::VarintDelta, } } pub fn position_codec(self) -> Option { match self { Self::V1 => None, - Self::V2 => Some(PositionStreamCodec::PackedDelta), + Self::V2 | Self::V3 => Some(PositionStreamCodec::PackedDelta), } } pub fn uses_shared_position_stream(self) -> bool { - matches!(self, Self::V2) + matches!(self, Self::V2 | Self::V3) } } @@ -204,14 +229,47 @@ impl FromStr for InvertedListFormatVersion { 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 or 2", + "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, 256) => 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() + ))) + } + (InvertedListFormatVersion::V3, other) => Err(Error::invalid_input(format!( + "FTS format_version=3 requires block_size=256, got {other}" + ))), + _ => unreachable!("validate_block_size limits supported block sizes"), + } +} + #[derive(Debug)] struct PartitionCandidates { tokens_by_position: Vec, @@ -425,6 +483,25 @@ fn parse_shared_position_codec(metadata: &HashMap) -> Result, ) -> Result { + if let Some(value) = metadata.get(FTS_FORMAT_VERSION_KEY) { + let format_version = InvertedListFormatVersion::from_str(value)?; + validate_format_version_block_size(format_version, parse_posting_block_size(metadata)?)?; + 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); } @@ -502,9 +579,12 @@ impl InvertedIndex { } fn index_version(&self) -> u32 { - match self.token_set_format { - TokenSetFormat::Arrow => 0, - TokenSetFormat::Fst => self.format_version().index_version(), + match (self.token_set_format, self.format_version()) { + ( + TokenSetFormat::Arrow, + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, + ) => 0, + (_, format_version) => format_version.index_version(), } } @@ -4837,11 +4917,10 @@ impl PostingListBuilder { } pub fn to_batch(self, block_max_scores: Vec) -> Result { - let format_version = if self.posting_tail_codec == PostingTailCodec::Fixed32 { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; + 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( self.has_positions(), format_version, @@ -4906,13 +4985,7 @@ impl PostingListBuilder { } pub fn to_batch_with_docs(self, docs: &DocSet, schema: SchemaRef) -> Result { - let format_version = if schema.column_with_name(POSITION_COL).is_some() - && schema.column_with_name(COMPRESSED_POSITION_COL).is_none() - { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; + 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()?) @@ -5888,17 +5961,18 @@ mod tests { row_id: u64, ) -> Result> { let block_size = params.posting_block_size(); + let format_version = params.resolved_format_version(); let mut partition = InnerBuilder::new_with_format_version_and_block_size( 0, false, token_set_format, - InvertedListFormatVersion::V1, + format_version, block_size, ); partition.tokens.add(token.to_owned()); let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( false, - PostingTailCodec::Fixed32, + format_version.posting_tail_codec(), block_size, ); posting_list.add(0, PositionRecorder::Count(1)); @@ -5916,6 +5990,15 @@ mod tests { TOKEN_SET_FORMAT_KEY.to_owned(), token_set_format.to_string(), ), + ( + POSTING_TAIL_CODEC_KEY.to_owned(), + format_version.posting_tail_codec().as_str().to_owned(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), ]); let mut writer = store .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) @@ -5959,6 +6042,7 @@ mod tests { )); let params = InvertedIndexParams::default().block_size(256).unwrap(); + let format_version = params.resolved_format_version(); let block_size = params.posting_block_size(); let num_docs = block_size + 7; @@ -5966,13 +6050,13 @@ mod tests { 0, false, TokenSetFormat::default(), - InvertedListFormatVersion::V1, + format_version, block_size, ); builder.tokens.add("needle".to_owned()); let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( false, - PostingTailCodec::Fixed32, + format_version.posting_tail_codec(), block_size, ); for doc_id in 0..num_docs { @@ -6253,6 +6337,19 @@ mod tests { resolve_fts_format_version(Some("2")).unwrap(), InvertedListFormatVersion::V2 ); + assert_eq!( + resolve_fts_format_version(Some("3")).unwrap(), + InvertedListFormatVersion::V3 + ); + } + + #[test] + fn test_block_size_256_metadata_resolves_to_v3() { + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "256".to_owned())]); + assert_eq!( + parse_format_version_from_metadata(&metadata).unwrap(), + InvertedListFormatVersion::V3 + ); } #[test] @@ -7972,6 +8069,7 @@ mod tests { partition_ids: Vec, params: InvertedIndexParams, ) { + let format_version = params.resolved_format_version(); let metadata = HashMap::from([ ( "partitions".to_owned(), @@ -7982,6 +8080,18 @@ mod tests { TOKEN_SET_FORMAT_KEY.to_owned(), TokenSetFormat::default().to_string(), ), + ( + POSTING_TAIL_CODEC_KEY.to_owned(), + format_version.posting_tail_codec().as_str().to_owned(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + params.posting_block_size().to_string(), + ), ]); let mut writer = store .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) @@ -8673,6 +8783,70 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_block_size_256_writes_v3_metadata_and_index_version() -> Result<()> { + let src_dir = TempObjDir::default(); + let dest_dir = TempObjDir::default(); + let src_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + src_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let dest_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + dest_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let params = InvertedIndexParams::default().block_size(256)?; + let format_version = params.resolved_format_version(); + assert_eq!(format_version, InvertedListFormatVersion::V3); + + let mut partition = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + format_version, + params.posting_block_size(), + ); + partition.tokens.add("hello".to_owned()); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + params.posting_block_size(), + ); + posting_list.add(0, PositionRecorder::Count(1)); + partition.posting_lists.push(posting_list); + partition.docs.append(100, 1); + partition.write(src_store.as_ref()).await?; + + write_test_metadata(&src_store, vec![0], params).await; + + let index = InvertedIndex::load(src_store, None, &LanceCache::no_cache()).await?; + assert_eq!(index.format_version(), InvertedListFormatVersion::V3); + assert_eq!( + index.index_version(), + InvertedListFormatVersion::V3.index_version() + ); + + let created = index + .update(empty_doc_stream(), dest_store.as_ref(), None) + .await?; + assert_eq!( + created.index_version, + InvertedListFormatVersion::V3.index_version() + ); + + let updated = InvertedIndex::load(dest_store, None, &LanceCache::no_cache()).await?; + assert_eq!(updated.format_version(), InvertedListFormatVersion::V3); + assert_eq!( + updated.index_version(), + InvertedListFormatVersion::V3.index_version() + ); + + Ok(()) + } + #[tokio::test] async fn test_merge_segments_preserves_arrow_token_set_format() -> Result<()> { let src_dir = TempObjDir::default(); diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 2f31850834a..09f92804413 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -23,7 +23,8 @@ use crate::scalar::inverted::tokenizer::document_tokenizer::{ JsonTokenizer, LanceTokenizer, TextTokenizer, }; use crate::scalar::inverted::{ - InvertedListFormatVersion, default_fts_format_version, resolve_fts_format_version, + InvertedListFormatVersion, default_fts_format_version_for_block_size, + resolve_fts_format_version, validate_format_version_block_size, }; pub use lance_tokenizer::Language; use lance_tokenizer::{ @@ -142,7 +143,10 @@ pub struct InvertedIndexParams { /// On-disk FTS format version to write when creating a new index. /// /// This is a build-time only parameter and is not persisted with the index. - /// If unset, Lance writes the current default FTS format. + /// If unset, Lance writes v2 for `block_size = 128` and v3 for + /// `block_size = 256`. + /// `format_version = 3` is experimental and is only valid with + /// `block_size = 256`. #[serde( rename = "format_version", skip_serializing, @@ -261,7 +265,7 @@ where serde_json::Value::Number(value) => { let Some(value) = value.as_u64() else { return Err(serde::de::Error::custom( - "FTS format_version must be 1 or 2", + "FTS format_version must be 1, 2, or 3", )); }; resolve_fts_format_version(Some(&value.to_string())) @@ -269,7 +273,7 @@ where .map_err(serde::de::Error::custom) } other => Err(serde::de::Error::custom(format!( - "FTS format_version must be 1 or 2, got {other}" + "FTS format_version must be 1, 2, or 3, got {other}" ))), } } @@ -438,17 +442,27 @@ impl InvertedIndexParams { /// Set the on-disk FTS format version to use when creating a new index. /// - /// If unset, Lance writes the current default FTS format. Existing indexes - /// keep their own on-disk format during update and optimize operations. + /// If unset, Lance writes v2 for `block_size = 128` and v3 for + /// `block_size = 256`. Existing indexes keep their own on-disk format + /// during update and optimize operations. + /// `format_version = 3` is experimental and is only valid with + /// `block_size = 256`. pub fn format_version(mut self, format_version: InvertedListFormatVersion) -> Self { self.format_version = Some(format_version); self } - /// Resolve the requested FTS format version, falling back to Lance's default. + /// Resolve the requested FTS format version, falling back to the default for + /// the configured block size. pub fn resolved_format_version(&self) -> InvertedListFormatVersion { self.format_version - .unwrap_or_else(default_fts_format_version) + .unwrap_or_else(|| default_fts_format_version_for_block_size(self.block_size).unwrap()) + } + + /// Validate that the requested FTS format version can safely encode the + /// configured posting block size. + pub fn validate_format_version(&self) -> Result<()> { + validate_format_version_block_size(self.resolved_format_version(), self.block_size) } /// Serialize params for the build/training path, including build-only fields. @@ -492,7 +506,9 @@ impl InvertedIndexParams { .expect("inverted index params should serialize to a JSON object"); object.extend(supplied.clone()); - Ok(serde_json::from_value(value)?) + let params: Self = serde_json::from_value(value)?; + params.validate_format_version()?; + Ok(params) } pub fn build(&self) -> Result> { @@ -671,6 +687,52 @@ mod tests { ); } + #[test] + fn test_block_size_256_defaults_to_v3() { + assert_eq!( + InvertedIndexParams::default() + .block_size(256) + .unwrap() + .resolved_format_version(), + InvertedListFormatVersion::V3 + ); + } + + #[test] + fn test_format_version_must_match_block_size() { + InvertedIndexParams::default() + .format_version(InvertedListFormatVersion::V2) + .validate_format_version() + .unwrap(); + InvertedIndexParams::default() + .block_size(256) + .unwrap() + .validate_format_version() + .unwrap(); + + let err = InvertedIndexParams::default() + .block_size(256) + .unwrap() + .format_version(InvertedListFormatVersion::V2) + .validate_format_version() + .unwrap_err(); + assert!(err.to_string().contains("block_size=256")); + + let err = InvertedIndexParams::default() + .format_version(InvertedListFormatVersion::V3) + .validate_format_version() + .unwrap_err(); + assert!(err.to_string().contains("format_version=3")); + } + + #[test] + fn test_training_json_rejects_incompatible_format_version_and_block_size() { + let err = + InvertedIndexParams::from_training_json(r#"{"block_size": 256, "format_version": 2}"#) + .unwrap_err(); + assert!(err.to_string().contains("block_size=256")); + } + #[test] fn test_block_size_default_serializes() { let params = InvertedIndexParams::default(); diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index e5a38275292..f7b6971c06d 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -205,8 +205,9 @@ impl MemIndexConfig { // the maintained-index path can only write the modern format. 0 | 1 => Ok(InvertedListFormatVersion::V1), 2 => Ok(InvertedListFormatVersion::V2), + 3 => Ok(InvertedListFormatVersion::V3), version => Err(Error::invalid_input(format!( - "FTS index '{}' has unsupported index_version {}; expected 0, 1, or 2", + "FTS index '{}' has unsupported index_version {}; expected 0, 1, 2, or 3", index_meta.name, version ))), } @@ -1061,6 +1062,7 @@ mod tests { (0, InvertedListFormatVersion::V1), (1, InvertedListFormatVersion::V1), (2, InvertedListFormatVersion::V2), + (3, InvertedListFormatVersion::V3), ] { let config = MemIndexConfig::fts_from_metadata(&fts_index_metadata(index_version), &schema) @@ -1083,9 +1085,9 @@ mod tests { let arrow_schema = create_test_schema(); let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); - let err = MemIndexConfig::fts_from_metadata(&fts_index_metadata(3), &schema).unwrap_err(); + let err = MemIndexConfig::fts_from_metadata(&fts_index_metadata(4), &schema).unwrap_err(); assert!( - err.to_string().contains("unsupported index_version 3"), + err.to_string().contains("unsupported index_version 4"), "{err}" ); } diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs index 410823c31db..50c21a3eaf5 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/flush.rs @@ -749,8 +749,9 @@ impl MemTableFlusher { use std::sync::Arc; use lance_index::scalar::inverted::{ - POSITIONS_CODEC_KEY, POSITIONS_CODEC_PACKED_DELTA_V1, POSITIONS_LAYOUT_KEY, - POSITIONS_LAYOUT_SHARED_STREAM_V2, POSTING_TAIL_CODEC_KEY, TokenSetFormat, + FTS_FORMAT_VERSION_KEY, POSITIONS_CODEC_KEY, POSITIONS_CODEC_PACKED_DELTA_V1, + POSITIONS_LAYOUT_KEY, POSITIONS_LAYOUT_SHARED_STREAM_V2, POSTING_BLOCK_SIZE_KEY, + POSTING_TAIL_CODEC_KEY, TokenSetFormat, }; // Create metadata with params and partitions in schema metadata (this is what InvertedIndex expects) @@ -766,6 +767,14 @@ impl MemTableFlusher { POSTING_TAIL_CODEC_KEY.to_string(), format_version.posting_tail_codec().as_str().to_string(), ), + ( + FTS_FORMAT_VERSION_KEY.to_string(), + format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_string(), + config.params.posting_block_size().to_string(), + ), ] .into_iter() .collect::>(); From 82689bbbfc49f3bfa6a60c50cdefe25cb207cfc9 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 30 Jun 2026 20:45:25 +0800 Subject: [PATCH 09/20] style(java): format fts index params --- .../java/org/lance/index/scalar/InvertedIndexParams.java | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index a688626d18d..82513a89ee7 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -264,9 +264,8 @@ public Builder skipMerge(boolean skipMerge) { /** * Configure the on-disk FTS format version to write when creating a new index. * - *

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

If unset, Lance writes v2 for {@code blockSize = 128} and v3 for {@code blockSize = 256}. + * {@code formatVersion = 3} is experimental and is only valid with {@code blockSize = 256}. * * @param formatVersion FTS format version, must be 1, 2, or 3 * @return this builder @@ -284,8 +283,7 @@ public Builder formatVersion(int formatVersion) { public ScalarIndexParams build() { if (formatVersion != null) { Preconditions.checkArgument( - (blockSize == 256 && formatVersion == 3) - || (blockSize == 128 && formatVersion != 3), + (blockSize == 256 && formatVersion == 3) || (blockSize == 128 && formatVersion != 3), "formatVersion 3 requires blockSize 256, and blockSize 256 requires formatVersion 3"); } Map params = new HashMap<>(); From aed44091b995d700703e4982e197882f0252c8d2 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 30 Jun 2026 20:57:59 +0800 Subject: [PATCH 10/20] fix(index): avoid nightly coverage stream ice --- rust/lance-index/src/scalar/ngram.rs | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/rust/lance-index/src/scalar/ngram.rs b/rust/lance-index/src/scalar/ngram.rs index 582c7aab157..9dd62579d00 100644 --- a/rust/lance-index/src/scalar/ngram.rs +++ b/rust/lance-index/src/scalar/ngram.rs @@ -61,6 +61,8 @@ const NGRAM_INDEX_VERSION: u32 = 0; use std::sync::LazyLock; +type NGramSpillStream = futures::stream::BoxStream<'static, Result>; + pub static TOKENS_FIELD: LazyLock = LazyLock::new(|| Field::new(TOKENS_COL, DataType::UInt32, true)); pub static POSTINGS_FIELD: LazyLock = @@ -967,12 +969,10 @@ impl NGramIndexBuilder { Ok(()) } - async fn stream_spill_reader( - reader: Arc, - ) -> Result>> { + fn stream_spill_reader(reader: Arc) -> NGramSpillStream { let num_rows = reader.num_rows(); - Ok(stream::try_unfold(0, move |offset| { + stream::try_unfold(0, move |offset| { let reader = reader.clone(); async move { // These are small batches but, in the worst case scenario, each row could @@ -987,17 +987,15 @@ impl NGramIndexBuilder { Ok(Some((state, new_offset))) } .boxed() - })) + }) + .boxed() } - async fn stream_spill( - spill_store: Arc, - id: usize, - ) -> Result>> { + async fn stream_spill(spill_store: Arc, id: usize) -> Result { let reader = spill_store .open_index_file(&Self::spill_filename(id)) .await?; - Self::stream_spill_reader(reader).await + Ok(Self::stream_spill_reader(reader)) } fn merge_spill_states( @@ -1203,7 +1201,7 @@ impl NGramIndexBuilder { let left_stream = Self::stream_spill(self.spill_store.clone(), new_data_num).await?; let old_reader = old_index.open_index_file(POSTINGS_FILENAME).await?; - let right_stream = Self::stream_spill_reader(old_reader).await?; + let right_stream = Self::stream_spill_reader(old_reader); Self::merge_spill_streams(left_stream, right_stream, writer.as_mut()).await?; From 01d5cd595c58e63fae79904475c278bb319fee33 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 30 Jun 2026 21:11:58 +0800 Subject: [PATCH 11/20] fix(io): avoid coverage stream ice --- rust/lance/src/io/exec/filtered_read.rs | 14 ++++++-------- rust/lance/src/io/exec/pushdown_scan.rs | 5 +++-- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index b50cba1f9ce..78bd4fabdf7 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::any::Any; use std::collections::{BTreeMap, HashMap}; -use std::pin::Pin; use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::{ops::Range, sync::Arc}; @@ -1079,7 +1078,7 @@ impl FilteredReadStream { mut fragment_read_task: ScopedFragmentRead, global_metrics: Arc, fragment_soft_limit: Option, - ) -> Result>> { + ) -> Result>> { let output_schema = Arc::new(fragment_read_task.projection.to_arrow_schema()); if let Some(filter) = &fragment_read_task.filter { @@ -1159,12 +1158,11 @@ impl FilteredReadStream { ))) .map(|(batch_fut, args)| Self::wrap_with_filter(batch_fut, args.0, args.1)); - let result: Pin> + Send>> = - if let Some(limit) = fragment_soft_limit { - Box::pin(Self::apply_soft_limit(fragment_stream, limit)) - } else { - Box::pin(fragment_stream) - }; + let result = if let Some(limit) = fragment_soft_limit { + Self::apply_soft_limit(fragment_stream, limit).boxed() + } else { + fragment_stream.boxed() + }; Ok(result) } diff --git a/rust/lance/src/io/exec/pushdown_scan.rs b/rust/lance/src/io/exec/pushdown_scan.rs index 0d49232c094..79073e4c15d 100644 --- a/rust/lance/src/io/exec/pushdown_scan.rs +++ b/rust/lance/src/io/exec/pushdown_scan.rs @@ -25,7 +25,8 @@ use datafusion::{ }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_physical_expr::EquivalenceProperties; -use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; +use futures::stream::BoxStream; +use futures::{FutureExt, StreamExt, TryStreamExt}; use lance_arrow::{RecordBatchExt, SchemaExt}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID_FIELD}; @@ -325,7 +326,7 @@ impl FragmentScanner { }) } - pub async fn scan(self) -> Result> + 'static + Send> { + pub async fn scan(self) -> Result>> { let batch_readahead = self.config.batch_readahead; let simplified_predicates = self.simplified_predicates()?; let ordered_output = self.config.ordered_output; From 96d9697184636765b0539bb82193563a7bcf926a Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Wed, 1 Jul 2026 16:40:25 +0800 Subject: [PATCH 12/20] fix(fts): validate memwal block size version --- rust/lance/src/dataset/mem_wal/index.rs | 101 +++++++++++++++++--- rust/lance/src/dataset/mem_wal/index/fts.rs | 45 +++++++-- 2 files changed, 126 insertions(+), 20 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 0ad159a1907..8077f06149e 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -149,12 +149,12 @@ impl MemIndexConfig { }; let params = params.format_version(Self::fts_format_version_from_metadata(index_meta)?); - Ok(Self::Fts(FtsIndexConfig::with_params( + Ok(Self::Fts(FtsIndexConfig::try_with_params( index_meta.name.clone(), field_id, column, params, - ))) + )?)) } /// Create an HNSW vector index config. @@ -353,8 +353,11 @@ impl IndexStore { registry.hnsw_indexes.insert(c.name.clone(), index); } MemIndexConfig::Fts(c) => { - let index = - FtsMemIndex::with_params(c.field_id, c.column.clone(), c.params.clone()); + let index = FtsMemIndex::try_with_params( + c.field_id, + c.column.clone(), + c.params.clone(), + )?; registry.fts_indexes.insert(c.name.clone(), index); } } @@ -453,13 +456,16 @@ impl IndexStore { field_id: i32, column: String, params: InvertedIndexParams, - ) { + ) -> Result<()> { assert!( self.pk_index.is_none() || self.pk_is_empty(), "FTS indexes must be configured before inserting rows into a PK memtable" ); - self.fts_indexes - .insert(name, FtsMemIndex::with_params(field_id, column, params)); + self.fts_indexes.insert( + name, + FtsMemIndex::try_with_params(field_id, column, params)?, + ); + Ok(()) } /// Maintain a primary-key index so the memtable can answer "newest visible @@ -1044,8 +1050,21 @@ mod tests { fn fts_index_metadata(index_version: i32) -> IndexMetadata { let details = pbold::InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(); - let mut value = Vec::new(); - details.encode(&mut value).unwrap(); + fts_index_metadata_with_details(index_version, Some(details)) + } + + fn fts_index_metadata_with_details( + index_version: i32, + details: Option, + ) -> IndexMetadata { + let index_details = details.map(|details| { + let mut value = Vec::new(); + details.encode(&mut value).unwrap(); + Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.InvertedIndexDetails".to_string(), + value, + }) + }); IndexMetadata { uuid: Uuid::new_v4(), @@ -1053,10 +1072,7 @@ mod tests { name: "desc_idx".to_string(), dataset_version: 1, fragment_bitmap: None, - index_details: Some(Arc::new(prost_types::Any { - type_url: "type.googleapis.com/lance.index.InvertedIndexDetails".to_string(), - value, - })), + index_details, index_version, created_at: None, base_id: None, @@ -1365,7 +1381,6 @@ mod tests { (0, InvertedListFormatVersion::V1), (1, InvertedListFormatVersion::V1), (2, InvertedListFormatVersion::V2), - (3, InvertedListFormatVersion::V3), ] { let config = MemIndexConfig::fts_from_metadata(&fts_index_metadata(index_version), &schema) @@ -1395,6 +1410,64 @@ mod tests { ); } + #[test] + fn fts_from_metadata_rejects_v3_without_256_block_size() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let missing_details = + MemIndexConfig::fts_from_metadata(&fts_index_metadata_with_details(3, None), &schema) + .unwrap_err(); + assert!( + missing_details + .to_string() + .contains("requires block_size=256"), + "{missing_details}" + ); + assert!( + missing_details.to_string().contains("got 128"), + "{missing_details}" + ); + + let default_details = + MemIndexConfig::fts_from_metadata(&fts_index_metadata(3), &schema).unwrap_err(); + assert!( + default_details + .to_string() + .contains("requires block_size=256"), + "{default_details}" + ); + assert!( + default_details.to_string().contains("got 128"), + "{default_details}" + ); + } + + #[test] + fn fts_from_metadata_accepts_v3_with_256_block_size() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + + let config = MemIndexConfig::fts_from_metadata( + &fts_index_metadata_with_details(3, Some(details)), + &schema, + ) + .unwrap(); + + match config { + MemIndexConfig::Fts(config) => { + assert_eq!( + config.params.resolved_format_version(), + InvertedListFormatVersion::V3 + ); + assert_eq!(config.params.posting_block_size(), 256); + } + _ => unreachable!("fts metadata should create an FTS config"), + } + } + #[test] fn test_from_configs() { let configs = vec![ diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index a81b136cd2b..6e0d12a276e 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -960,10 +960,20 @@ impl FtsMemIndex { /// Create a new FTS index with custom tokenizer parameters. pub fn with_params(field_id: i32, column_name: String, params: InvertedIndexParams) -> Self { - let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP) - .expect("Failed to build tokenizer"); + Self::try_with_params(field_id, column_name, params) + .expect("invalid MemWAL FTS index parameters") + } + + /// Try to create a new FTS index with custom tokenizer parameters. + pub fn try_with_params( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + ) -> Result { + params.validate_format_version()?; + let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP)?; let writer_tokenizer = pool.template.box_clone(); - Self { + Ok(Self { field_id, column_name, params, @@ -972,7 +982,7 @@ impl FtsMemIndex { state: ArcSwap::from(IndexState::empty()), freeze_threshold_rows: Self::DEFAULT_FREEZE_THRESHOLD_ROWS, merge: Arc::new(Mutex::new(None)), - } + }) } /// Override the tail freeze threshold (docs) — the analogue of Lucene's @@ -2338,12 +2348,23 @@ impl FtsIndexConfig { column: String, params: InvertedIndexParams, ) -> Self { - Self { + Self::try_with_params(name, field_id, column, params) + .expect("invalid MemWAL FTS index config parameters") + } + + pub fn try_with_params( + name: String, + field_id: i32, + column: String, + params: InvertedIndexParams, + ) -> Result { + params.validate_format_version()?; + Ok(Self { name, field_id, column, params, - } + }) } } @@ -4672,6 +4693,18 @@ mod tests { assert!(builder.id() > 0 || builder.id() == 42); } + #[test] + fn test_to_index_builder_supports_block_size_256() { + let schema = create_test_schema(); + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let index = FtsMemIndex::try_with_params(1, "description".to_string(), params).unwrap(); + let batch = create_test_batch(&schema); + index.insert(&batch, 0).unwrap(); + + let builder = index.to_index_builder(42, 3).unwrap(); + assert_eq!(builder.id(), 42); + } + #[test] fn test_unsupported_column_type_errors() { let schema = Arc::new(ArrowSchema::new(vec![ From 14655b9d70936e9cd0905a154f7eb52679565a51 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 2 Jul 2026 23:22:13 +0800 Subject: [PATCH 13/20] fix(index): split fts tail-partition merge by the worker memory budget merge_all_tail_partitions folded every leftover worker builder into a single partition unconditionally; on a build whose workers never hit the flush threshold (large LANCE_FTS_PARTITION_SIZE), the entire index collapsed into one partition and queries lost all intra-partition parallelism. Fold with the same memory/doc-count checks as the partition-merge path instead, so the final partition count converges to total builder memory / budget. --- .../src/scalar/inverted/builder.rs | 132 ++++++++++++------ 1 file changed, 92 insertions(+), 40 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index c739c5d5d22..da0c4798f25 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -181,25 +181,44 @@ fn resolve_worker_memory_limit_bytes(params: &InvertedIndexParams, num_workers: .unwrap_or(default_worker_memory_limit_bytes) } -fn merge_all_tail_partitions(tails: Vec) -> Result> { - if tails.is_empty() { - return Ok(None); +/// Merge the workers' leftover tail builders into as few partitions as the +/// memory budget allows. Folding unconditionally would collapse every build +/// whose workers never hit the flush threshold into a single partition, which +/// destroys intra-query parallelism; splitting by the same per-partition +/// budget as the flush path makes the final partition count converge to +/// roughly total_builder_memory / memory_limit_bytes regardless of worker +/// layout. +fn merge_all_tail_partitions( + tails: Vec, + memory_limit_bytes: u64, +) -> Result> { + let mut merged_builders: Vec = Vec::new(); + let mut merged: Option = None; + for tail in tails { + let builder = tail.builder; + if builder.is_empty() { + continue; + } + match &mut merged { + Some(current) => { + let would_exceed_memory = + current.memory_size().saturating_add(builder.memory_size()) + >= memory_limit_bytes; + let would_exceed_doc_ids = + current.docs.len().saturating_add(builder.docs.len()) > u32::MAX as usize; + if would_exceed_memory || would_exceed_doc_ids { + merged_builders.push(std::mem::replace(current, builder)); + } else { + current.merge_from(builder)?; + } + } + None => merged = Some(builder), + } } - merge_tail_partition_group(tails).map(Some) -} - -fn merge_tail_partition_group(group: Vec) -> Result { - let mut group = group.into_iter(); - let mut merged = group - .next() - .ok_or_else(|| { - Error::invalid_input("cannot merge an empty tail partition group".to_owned()) - })? - .builder; - for tail in group { - merged.merge_from(tail.builder)?; - } - Ok(merged) + if let Some(builder) = merged { + merged_builders.push(builder); + } + Ok(merged_builders) } #[derive(Debug)] @@ -529,11 +548,12 @@ impl InvertedIndexBuilder { tail_partitions.push(tail_partition); } } - let merged_tail_partitions = - spawn_cpu(move || merge_all_tail_partitions(tail_partitions)).await?; - if let Some(builder) = merged_tail_partitions { + let merged_tail_partitions = spawn_cpu(move || { + merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes) + }) + .await?; + for mut builder in merged_tail_partitions { self.new_partitions.push(builder.id()); - let mut builder = builder; files.extend( builder .write_to(dest_store.as_ref(), self.partition_write_target()) @@ -3807,27 +3827,54 @@ mod tests { assert_eq!(resolve_worker_memory_limit_bytes(¶ms, 16), 256 << 20); } + fn tail_with_docs(id: u64, num_docs: u64) -> TailPartition { + let mut builder = InnerBuilder::new(id, false, TokenSetFormat::default()); + let token = builder.tokens.add(format!("token{}", id)); + builder + .posting_lists + .resize_with(builder.tokens.len(), || PostingListBuilder::new(false)); + for row in 0..num_docs { + let doc = builder.docs.append(row, 1); + builder.posting_lists[token as usize].add(doc, PositionRecorder::Count(1)); + } + TailPartition { builder } + } + #[test] - fn test_merge_all_tail_partitions_combines_everything() -> Result<()> { - let merged = merge_all_tail_partitions(vec![ - TailPartition { - builder: InnerBuilder::new(0, false, TokenSetFormat::default()), - }, - TailPartition { - builder: InnerBuilder::new(1, false, TokenSetFormat::default()), - }, - TailPartition { - builder: InnerBuilder::new(2, false, TokenSetFormat::default()), - }, - ])?; + fn test_merge_all_tail_partitions_combines_under_budget() -> Result<()> { + let merged = merge_all_tail_partitions( + vec![ + tail_with_docs(0, 4), + tail_with_docs(1, 4), + tail_with_docs(2, 4), + ], + u64::MAX, + )?; + assert_eq!(merged.len(), 1); + assert_eq!(merged[0].id(), 0); + assert_eq!(merged[0].docs.len(), 12); + Ok(()) + } - assert_eq!(merged.expect("merged builder should exist").id(), 0); + #[test] + fn test_merge_all_tail_partitions_splits_on_memory_budget() -> Result<()> { + let tails = vec![ + tail_with_docs(0, 64), + tail_with_docs(1, 64), + tail_with_docs(2, 64), + tail_with_docs(3, 64), + ]; + let single = tails[0].builder.memory_size(); + // A budget below two builders' footprint must keep them separate. + let merged = merge_all_tail_partitions(tails, single + 1)?; + assert_eq!(merged.len(), 4); + assert!(merged.iter().all(|builder| builder.docs.len() == 64)); Ok(()) } #[test] fn test_merge_all_tail_partitions_returns_none_for_empty_input() -> Result<()> { - assert!(merge_all_tail_partitions(Vec::new())?.is_none()); + assert!(merge_all_tail_partitions(Vec::new(), u64::MAX)?.is_empty()); Ok(()) } @@ -3849,10 +3896,15 @@ mod tests { let second_doc = second.docs.append(20, 2); second.posting_lists[world as usize].add(second_doc, PositionRecorder::Count(2)); - let merged = merge_tail_partition_group(vec![ - TailPartition { builder: first }, - TailPartition { builder: second }, - ])?; + let merged = merge_all_tail_partitions( + vec![ + TailPartition { builder: first }, + TailPartition { builder: second }, + ], + u64::MAX, + )?; + assert_eq!(merged.len(), 1); + let merged = &merged[0]; assert_eq!(merged.id(), 0); assert_eq!(merged.docs.len(), 2); From 7872bef34d5369402de672acfc45d01b99c4b011 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 02:47:06 +0800 Subject: [PATCH 14/20] perf(index): cache the num_tokens-only DocSet on LazyDocSet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ensure_num_tokens_loaded rebuilt the DocSet from the cached arrow column on every call, copying the whole num_tokens array (tens of MB per partition) once per query per partition — profiled at ~50% of single-term query time and a large slice of every warm query. Materialize it once in a OnceCell like the full DocSet, and account for it in deep_size_of. --- .../src/scalar/inverted/lazy_docset.rs | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index f103c2cefa4..12b5c1ebd32 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -74,6 +74,11 @@ pub struct DeferredDocSet { row_ids_col: OnceCell>, /// Full DocSet, materialized on first `ensure_loaded`. full: OnceCell>, + /// num_tokens-only DocSet, materialized on first + /// `ensure_num_tokens_loaded`. Cached because wand scoring calls this + /// once per query per partition; rebuilding it copied the whole + /// num_tokens column (tens of MB per partition) on every query. + tokens_only: OnceCell>, } impl std::fmt::Debug for LazyDocSet { @@ -103,6 +108,10 @@ impl lance_core::deepsize::DeepSizeOf for LazyDocSet { .get() .map(|d| d.deep_size_of_children(ctx)) .unwrap_or(0) + + d.tokens_only + .get() + .map(|d| d.deep_size_of_children(ctx)) + .unwrap_or(0) + d.num_tokens_col .get() .map(|arr| arr.len() * std::mem::size_of::()) @@ -134,6 +143,7 @@ impl LazyDocSet { num_tokens_col: OnceCell::new(), row_ids_col: OnceCell::new(), full: OnceCell::new(), + tokens_only: OnceCell::new(), })) } @@ -319,8 +329,14 @@ impl DeferredDocSet { if let Some(full) = self.full.get() { return Ok(full.clone()); } - let num_tokens = self.num_tokens_column().await?; - let docs = Arc::new(DocSet::from_num_tokens_only(num_tokens.as_ref())); + let docs = self + .tokens_only + .get_or_try_init(|| async { + let num_tokens = self.num_tokens_column().await?; + Result::Ok(Arc::new(DocSet::from_num_tokens_only(num_tokens.as_ref()))) + }) + .await? + .clone(); let _ = self.total_tokens.set(docs.total_tokens_num()); Ok(docs) } From 2adb6778b5aead031df0a42c0a07c00dcdb8103f Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 01:32:08 +0800 Subject: [PATCH 15/20] perf(index): compress and write fts tail partitions concurrently Tail partitions hold most of the index when workers rarely hit the flush threshold; writing them sequentially serialized the posting-list compression of nearly the whole index behind one producer thread at the end of the build. Write them through buffer_unordered instead. --- .../src/scalar/inverted/builder.rs | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index da0c4798f25..744eaa4cff1 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -552,13 +552,25 @@ impl InvertedIndexBuilder { merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes) }) .await?; - for mut builder in merged_tail_partitions { - self.new_partitions.push(builder.id()); - files.extend( - builder - .write_to(dest_store.as_ref(), self.partition_write_target()) - .await?, - ); + // Tail partitions hold most of the data when workers rarely hit the + // flush threshold; writing them one at a time serializes the + // posting-list compression of nearly the whole index behind a + // single producer thread. Compress and write them concurrently. + let write_target = self.partition_write_target(); + let mut tail_writes = futures::stream::iter( + merged_tail_partitions.into_iter().map(|mut builder| { + let dest_store = dest_store.clone(); + async move { + let partition_id = builder.id(); + let files = builder.write_to(dest_store.as_ref(), write_target).await?; + Result::Ok((partition_id, files)) + } + }), + ) + .buffer_unordered(get_num_compute_intensive_cpus().clamp(1, 16)); + while let Some((partition_id, partition_files)) = tail_writes.try_next().await? { + self.new_partitions.push(partition_id); + files.extend(partition_files); } log::info!("wait workers indexing elapsed: {:?}", start.elapsed()); Result::Ok(files) From d5aa24caaccf9ff120e22b1f00124cc398c34a90 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:02:09 +0800 Subject: [PATCH 16/20] feat(fts): patched-FOR frequency encoding for 256-doc posting blocks 256-doc blocks pack frequencies with Lucene PForUtil-style patched FOR: the body uses the bit width that minimizes total bytes and up to 31 outliers are appended as (index u8, high-bits varint) exceptions. Plain FOR let a single large tf widen the whole block; measured on a 200M-doc corpus (avg tf 4) this was 4.7 bits/posting for frequencies, patched FOR brings it to ~2.5. 128-doc blocks are unchanged. --- .../src/scalar/inverted/encoding.rs | 107 +++++++++++++++++- 1 file changed, 102 insertions(+), 5 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 28ece2a0d32..c8731dc1929 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -146,15 +146,108 @@ fn encode_posting_block_payload( compress_sorted_block_with::(doc_ids, &mut buffer, block)?; compress_block_with::(frequencies, &mut buffer, block)?; } + // 256-doc blocks (format V3) store frequencies with patched FOR: + // outliers no longer widen the whole block, which matters because one + // large tf per block otherwise doubles the frequency payload. BitPacker8x::BLOCK_LEN => { compress_sorted_block_with::(doc_ids, &mut buffer, block)?; - compress_block_with::(frequencies, &mut buffer, block)?; + compress_pfor_block_with::(frequencies, &mut buffer, block)?; } _ => unreachable!("validated posting block size should be supported"), } Ok(()) } +/// Patched FOR (Lucene PForUtil style): pick the body bit width that +/// minimizes total bytes, pack all values masked to that width, and append +/// up to [`PFOR_MAX_EXCEPTIONS`] exceptions as (index u8, high-bits varint). +const PFOR_MAX_EXCEPTIONS: usize = 31; + +#[inline] +fn u32_bits(value: u32) -> usize { + (32 - value.leading_zeros()) as usize +} + +#[inline] +fn varint_u32_len(value: u32) -> usize { + u32_bits(value).max(1).div_ceil(7) +} + +fn compress_pfor_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let max_bits = data.iter().map(|&v| u32_bits(v)).max().unwrap_or(0); + let mut best_width = max_bits; + let mut best_cost = P::BLOCK_LEN * max_bits / 8; + for width in (0..max_bits).rev() { + let mut exceptions = 0usize; + let mut exception_bytes = 0usize; + for &value in data { + if u32_bits(value) > width { + exceptions += 1; + exception_bytes += 1 + varint_u32_len(value >> width); + } + } + if exceptions > PFOR_MAX_EXCEPTIONS { + break; + } + let cost = P::BLOCK_LEN * width / 8 + exception_bytes; + if cost < best_cost { + best_cost = cost; + best_width = width; + } + } + + let mask = if best_width >= 32 { + u32::MAX + } else { + (1u32 << best_width) - 1 + }; + let mut body = [0u32; MAX_POSTING_BLOCK_SIZE]; + let mut exception_buf = Vec::new(); + let mut exception_count = 0u8; + for (index, &value) in data.iter().enumerate() { + body[index] = value & mask; + if u32_bits(value) > best_width { + exception_buf.push(index as u8); + encode_varint_u32(&mut exception_buf, value >> best_width); + exception_count += 1; + } + } + let compressor = P::new(); + let num_bytes = compressor.compress(&body[..P::BLOCK_LEN], buffer, best_width as u8); + let _ = builder.write(&[best_width as u8, exception_count])?; + let _ = builder.write(&buffer[..num_bytes])?; + let _ = builder.write(&exception_buf)?; + Ok(()) +} + +fn decompress_pfor_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let width = block[0]; + let exception_count = block[1] as usize; + let compressor = P::new(); + let num_bytes = compressor.decompress(&block[2..], buffer, width); + let mut offset = 2 + num_bytes; + for _ in 0..exception_count { + let index = block[offset] as usize; + offset += 1; + let high = decode_varint_u32(block, &mut offset) + .expect("pfor exception high bits should be a valid varint"); + buffer[index] |= high << width; + } + res.extend_from_slice(buffer); + offset +} + pub fn encode_remainder_posting_block_into( doc_ids: &[u32], frequencies: &[u32], @@ -297,7 +390,7 @@ pub fn compress_positions(positions: &[u32]) -> Result, mut value: u32) { +pub fn encode_varint_u32(dst: &mut Vec, mut value: u32) { while value >= 0x80 { dst.push((value as u8) | 0x80); value >>= 7; @@ -440,7 +533,7 @@ impl PositionBlockBuilder { } #[inline] -fn decode_varint_u32(src: &[u8], offset: &mut usize) -> Result { +pub fn decode_varint_u32(src: &[u8], offset: &mut usize) -> Result { let mut value = 0u32; let mut shift = 0u32; while *offset < src.len() { @@ -803,7 +896,7 @@ pub fn decompress_posting_block( BitPacker8x::BLOCK_LEN => { let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); block = &block[num_bytes..]; - let num_bytes = decompress_block_with::(block, buffer, frequencies); + let num_bytes = decompress_pfor_block_with::(block, buffer, frequencies); block = &block[num_bytes..]; } _ => unreachable!("validated posting block size should be supported"), @@ -1049,9 +1142,13 @@ mod tests { let doc_num_bits = block[8]; let doc_bytes = BitPacker8x::compressed_block_size(doc_num_bits); let freq_header_offset = 9 + doc_bytes; + // 256-doc blocks use patched FOR for frequencies: + // [width u8][exception_count u8][body][exceptions...] let freq_num_bits = block[freq_header_offset]; + let exception_count = block[freq_header_offset + 1] as usize; + assert_eq!(exception_count, 0, "uniform freqs need no exceptions"); let freq_bytes = BitPacker8x::compressed_block_size(freq_num_bits); - assert_eq!(block.len(), freq_header_offset + 1 + freq_bytes); + assert_eq!(block.len(), freq_header_offset + 2 + freq_bytes); let (decoded_doc_ids, decoded_frequencies) = decompress_posting_list_with_tail_codec_and_block_size( From 91095c120cf08828e7d86971ee1889018241262b Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:02:09 +0800 Subject: [PATCH 17/20] perf(fts): bake a shared first-doc-per-block slab on CompressedPostingList Block boundary lookups re-read block headers and re-decoded the tail block on every call; bake first doc ids once per cached list and share the slab across per-query clones. --- rust/lance-index/src/scalar/inverted/index.rs | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index fa232131117..f1bfffa9fb4 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::fmt::{Debug, Display}; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::{Arc, OnceLock}; use std::{ cmp::{Reverse, min}, collections::BinaryHeap, @@ -3873,7 +3873,7 @@ impl PlainPostingList { } } -#[derive(Debug, PartialEq, Clone)] +#[derive(Debug, Clone)] pub struct CompressedPostingList { pub max_score: f32, pub length: u32, @@ -3884,6 +3884,20 @@ pub struct CompressedPostingList { pub posting_tail_codec: PostingTailCodec, pub block_size: usize, pub positions: Option, + // First doc id per block, baked lazily and shared across per-query clones + // of the cached list. See `block_first_docs`. + first_docs: Arc>>, +} + +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 + } } impl DeepSizeOf for CompressedPostingList { @@ -3913,6 +3927,7 @@ impl CompressedPostingList { posting_tail_codec, block_size, positions, + first_docs: Arc::new(OnceLock::new()), } } @@ -3960,6 +3975,7 @@ impl CompressedPostingList { posting_tail_codec, block_size, positions, + first_docs: Arc::new(OnceLock::new()), } } @@ -3979,14 +3995,29 @@ impl CompressedPostingList { } pub fn block_least_doc_id(&self, block_idx: usize) -> u32 { - let block = self.blocks.value(block_idx); - let remainder = self.length as usize % self.block_size; - let is_remainder_block = remainder > 0 && block_idx + 1 == self.blocks.len(); - if is_remainder_block { - super::encoding::read_posting_tail_first_doc(block, self.posting_tail_codec) - } else { - block[4..8].try_into().map(u32::from_le_bytes).unwrap() - } + 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, + ); + } + block[4..8].try_into().map(u32::from_le_bytes).unwrap() + }) + .collect() + }) } } From 0ec1da3100af344b572ecd14061bce34927656a6 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:10:54 +0800 Subject: [PATCH 18/20] feat(fts): impact skip data for posting lists Store per-block (freq, doc_len) impact frontiers alongside 256-doc posting blocks (varint-encoded: 2-3 bytes per pair) plus one level1 entry per 32 blocks, and drive block-max WAND pruning from them instead of build-time scores that go stale as index stats drift: - Bounds bake once per cached list into an Arc-shared slab (max doc weight per entry plus the list-wide max); per-query clones reuse it, so query time pays one multiply per bound instead of frontier rescans. - Entry doc_up_tos decode once at construction. - Lagging iterators park in the WAND tail under the data-driven global bound (query_weight x baked list max) instead of INFINITY. 128-doc-block indexes keep fixed-width u32 impact entries. --- rust/lance-index/protos-cache/cache.proto | 2 + rust/lance-index/src/scalar/inverted.rs | 1 + .../src/scalar/inverted/builder.rs | 49 +- .../src/scalar/inverted/cache_codec.rs | 203 ++++- .../lance-index/src/scalar/inverted/impact.rs | 749 ++++++++++++++++++ rust/lance-index/src/scalar/inverted/index.rs | 568 +++++++++++-- .../lance-index/src/scalar/inverted/scorer.rs | 11 + rust/lance-index/src/scalar/inverted/wand.rs | 508 ++++++++++-- 8 files changed, 1961 insertions(+), 130 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/impact.rs diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index e92da05815f..a861e2ef652 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -31,6 +31,8 @@ message CompressedPostingHeader { // Number of documents in each compressed posting block. Older cache entries // omit this field and decode as the legacy 128-doc block size. uint32 block_size = 6; + // Whether an impact IPC section follows the posting/position sections. + bool has_impacts = 7; } // Header for a serialized `PlainPostingList` cache entry. Followed by an Arrow diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 55e2b2abce5..0ba7e52e530 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors pub mod builder; +mod impact; mod cache_codec; mod encoding; mod index; diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 3301c1b1277..ddd2ed9ef82 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1733,11 +1733,27 @@ pub fn inverted_list_schema_for_version_with_block_size( with_position: bool, format_version: InvertedListFormatVersion, block_size: usize, +) -> SchemaRef { + inverted_list_schema_for_version_with_block_size_and_impacts( + with_position, + format_version, + block_size, + true, + ) +} + +pub(crate) fn inverted_list_schema_for_version_with_block_size_and_impacts( + with_position: bool, + format_version: InvertedListFormatVersion, + block_size: usize, + with_impacts: bool, ) -> SchemaRef { validate_format_version_block_size(format_version, block_size) .expect("invalid FTS format version for posting block size"); match format_version { - InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position, block_size), + InvertedListFormatVersion::V1 => { + inverted_list_schema_v1(with_position, block_size, with_impacts) + } InvertedListFormatVersion::V2 | InvertedListFormatVersion::V3 => { inverted_list_schema_with_tail_codec_and_position_codec( with_position, @@ -1745,12 +1761,17 @@ pub fn inverted_list_schema_for_version_with_block_size( PostingTailCodec::VarintDelta, Some(PositionStreamCodec::PackedDelta), block_size, + with_impacts, ) } } } -fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef { +fn inverted_list_schema_v1( + with_position: bool, + block_size: usize, + with_impacts: bool, +) -> SchemaRef { let mut fields = vec![ arrow_schema::Field::new( POSTING_COL, @@ -1764,6 +1785,17 @@ fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( POSITION_COL, @@ -1806,6 +1838,7 @@ pub fn inverted_list_schema_with_tail_codec( posting_tail_codec, Some(PositionStreamCodec::PackedDelta), LEGACY_BLOCK_SIZE, + false, ) } @@ -1815,6 +1848,7 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( posting_tail_codec: PostingTailCodec, position_codec: Option, block_size: usize, + with_impacts: bool, ) -> SchemaRef { let mut fields = vec![ // we compress the posting lists (including row ids and frequencies), @@ -1831,6 +1865,17 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( COMPRESSED_POSITION_COL, diff --git a/rust/lance-index/src/scalar/inverted/cache_codec.rs b/rust/lance-index/src/scalar/inverted/cache_codec.rs index fffb080bda1..bf36209d449 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -39,6 +39,7 @@ use crate::cache_pb::{ PostingTailCodec as PbPostingTailCodec, }; +use super::impact::ImpactSkipData; use super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, @@ -96,6 +97,7 @@ const BLOCK_OFFSETS_COLUMN: &str = "block_offsets"; const ROW_IDS_COLUMN: &str = "row_ids"; const FREQUENCIES_COLUMN: &str = "frequencies"; const BLOCKS_COLUMN: &str = "blocks"; +const IMPACTS_COLUMN: &str = "impacts"; fn legacy_positions_batch(list: &ListArray) -> Result { let schema = Arc::new(Schema::new(vec![Field::new( @@ -179,7 +181,7 @@ fn read_position_sections( impl CacheCodecImpl for PostingList { const TYPE_ID: &'static str = "lance.fts.PostingList"; - const CURRENT_VERSION: u32 = 1; + const CURRENT_VERSION: u32 = 2; fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { match self { @@ -293,6 +295,7 @@ fn serialize_compressed( position_storage: position_storage as i32, position_stream_codec: position_stream_codec as i32, block_size: posting.block_size as u32, + has_impacts: posting.impacts.is_some(), }; w.write_header(&header)?; @@ -307,6 +310,15 @@ fn serialize_compressed( if let Some(storage) = &posting.positions { write_position_sections(w, storage)?; } + if let Some(impacts) = &posting.impacts { + let schema = Arc::new(Schema::new(vec![Field::new( + IMPACTS_COLUMN, + DataType::LargeBinary, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(impacts.entries().clone())])?; + w.write_ipc(&batch)?; + } Ok(()) } @@ -329,6 +341,22 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result= 2 && header.has_impacts { + let batch = r.read_ipc()?; + let entries = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::io("impacts column is not a LargeBinaryArray".to_string()))? + .clone(); + Some(ImpactSkipData::new( + entries, + blocks.len(), + super::impact::ImpactFormat::for_block_size(block_size), + )?) + } else { + None + }; Ok(CompressedPostingList::new( blocks, @@ -337,6 +365,7 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result) -> Result) -> Result<()> { let count = u32::try_from(self.posting_lists.len()) @@ -418,12 +447,13 @@ impl CacheCodecImpl for Positions { #[cfg(test)] mod tests { use arrow::buffer::ScalarBuffer; - use arrow_array::LargeBinaryArray; - use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::builder::{Int32Builder, LargeBinaryBuilder, ListBuilder}; + use arrow_array::{Array, LargeBinaryArray}; use bytes::Bytes; use lance_core::Result; use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter}; + use super::super::impact::ImpactSkipData; use super::super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, @@ -469,6 +499,44 @@ mod tests { } } + fn impact_entry(doc_up_to: u32, pairs: &[(u32, u32)]) -> Vec { + let mut bytes = Vec::with_capacity(8 + pairs.len() * 8); + bytes.extend_from_slice(&doc_up_to.to_le_bytes()); + bytes.extend_from_slice(&(pairs.len() as u32).to_le_bytes()); + for (freq, doc_len) in pairs { + bytes.extend_from_slice(&freq.to_le_bytes()); + bytes.extend_from_slice(&doc_len.to_le_bytes()); + } + bytes + } + + fn impact_skip_data(level0_len: usize) -> ImpactSkipData { + let mut entries = Vec::with_capacity(level0_len + level0_len.div_ceil(32)); + for block_idx in 0..level0_len { + entries.push(impact_entry( + (block_idx as u32 + 1) * 10 - 1, + &[(block_idx as u32 + 1, 10)], + )); + } + let num_level1 = level0_len.div_ceil(32); + for group_idx in 0..num_level1 { + let group_end = ((group_idx + 1) * 32).min(level0_len); + let doc_up_to = group_end as u32 * 10 - 1; + entries.push(impact_entry(doc_up_to, &[(1, 10), (2, 8)])); + } + let mut builder = LargeBinaryBuilder::with_capacity(entries.len(), 0); + for entry in entries { + builder.append_value(entry); + } + let array = builder.finish(); + ImpactSkipData::new( + array, + level0_len, + super::super::impact::ImpactFormat::FixedU32, + ) + .unwrap() + } + /// Serialize a codec body (no envelope) into a standalone buffer. fn body_bytes(entry: &T) -> Bytes { let mut buf = Vec::new(); @@ -483,6 +551,11 @@ mod tests { T::deserialize(&mut r) } + fn from_body_version(data: &Bytes, version: u32) -> Result { + let mut r = CacheEntryReader::new(data, 0, version); + T::deserialize(&mut r) + } + fn roundtrip_posting_list(entry: &PostingList) -> PostingList { from_body::(&body_bytes(entry)).unwrap() } @@ -540,8 +613,15 @@ mod tests { Some(&[1u8, 2, 3, 4, 5][..]), Some(&[6, 7, 8, 9, 10][..]), ]); - let posting = - CompressedPostingList::new(blocks, 3.5, 42, PostingTailCodec::VarintDelta, 256, None); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + None, + ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { PostingList::Compressed(restored) => { @@ -556,6 +636,53 @@ mod tests { } } + #[test] + fn compressed_posting_list_impacts_roundtrip() { + let blocks = LargeBinaryArray::from_opt_vec(vec![ + Some(&[1u8, 2, 3, 4, 5][..]), + Some(&[6, 7, 8, 9, 10][..]), + ]); + let impacts = impact_skip_data(blocks.len()); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + Some(impacts.clone()), + ); + let entry = PostingList::Compressed(posting); + match roundtrip_posting_list(&entry) { + PostingList::Compressed(restored) => { + let restored = restored.impacts.expect("impacts should roundtrip"); + assert_eq!(restored.level0_len(), impacts.level0_len()); + assert_eq!(restored.level1_len(), impacts.level1_len()); + assert_eq!(restored.entries(), impacts.entries()); + } + PostingList::Plain(_) => panic!("expected Compressed variant"), + } + } + + #[test] + fn compressed_posting_list_v1_cache_without_impacts_decodes() { + let posting = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]), + 1.25, + 5, + PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, + None, + ); + let data = body_bytes(&PostingList::Compressed(posting)); + let restored = from_body_version::(&data, 1).unwrap(); + let PostingList::Compressed(restored) = restored else { + panic!("expected Compressed variant"); + }; + assert!(restored.impacts.is_none()); + } + #[test] fn compressed_posting_list_legacy_positions_roundtrip() { let blocks = LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]); @@ -568,6 +695,7 @@ mod tests { Some(CompressedPositionStorage::LegacyPerDoc(legacy_positions( &[&[0, 4, 8]], ))), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -601,6 +729,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -632,6 +761,7 @@ mod tests { Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), + None, ); let serialized = body_bytes(&PostingList::Compressed(posting)); @@ -665,6 +795,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, None, + None, )); for members in [ @@ -689,6 +820,64 @@ mod tests { } } + #[test] + fn posting_list_group_impacted_compressed_members_roundtrip() { + let first = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..]), Some(&[4u8, 5, 6][..])]), + 3.0, + 256, + PostingTailCodec::VarintDelta, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, + Some(impact_skip_data(2)), + ); + let second = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[7u8, 8, 9][..])]), + 5.0, + 128, + PostingTailCodec::Fixed32, + 256, + Some(CompressedPositionStorage::SharedStream( + SharedPositionStream::new( + PositionStreamCodec::PackedDelta, + vec![0u32, 12], + Bytes::from(vec![0xABu8; 32]), + ), + )), + Some(impact_skip_data(1)), + ); + let members = vec![ + PostingList::Compressed(first.clone()), + PostingList::Compressed(second.clone()), + ]; + let group = PostingListGroup::new(members); + let restored = from_body::(&body_bytes(&group)).unwrap(); + assert_eq!(restored.posting_lists.len(), 2); + + let expected = [&first, &second]; + for (expected, restored) in expected.iter().zip(restored.posting_lists.iter()) { + let PostingList::Compressed(restored) = restored else { + panic!("expected compressed member"); + }; + assert_eq!(restored.blocks, expected.blocks); + assert_eq!(restored.length, expected.length); + assert_eq!(restored.max_score, expected.max_score); + assert_eq!(restored.posting_tail_codec, expected.posting_tail_codec); + assert_eq!(restored.block_size, expected.block_size); + assert_eq!( + restored.impacts.as_ref().unwrap().entries(), + expected.impacts.as_ref().unwrap().entries() + ); + match (&expected.positions, &restored.positions) { + (Some(expected), Some(restored)) => { + assert_position_storage_eq(expected, restored); + } + (None, None) => {} + _ => panic!("position storage mismatch"), + } + } + } + #[test] fn positions_legacy_roundtrip() { let positions = Positions(CompressedPositionStorage::LegacyPerDoc(legacy_positions( @@ -775,6 +964,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, )) } @@ -826,6 +1016,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, None, + None, )) }; let group = PostingListGroup::new(vec![make_member(9), make_member(1)]); diff --git a/rust/lance-index/src/scalar/inverted/impact.rs b/rust/lance-index/src/scalar/inverted/impact.rs new file mode 100644 index 00000000000..8b16a4f57f8 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/impact.rs @@ -0,0 +1,749 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::{Arc, OnceLock}; + +use arrow_array::builder::LargeBinaryBuilder; +use arrow_array::{Array, LargeBinaryArray}; +use lance_core::{Error, Result}; + +use super::scorer::Scorer; + +pub const IMPACT_LEVEL1_BLOCKS: usize = 32; +const SMALL_FRONTIER_FREQ_LIMIT: usize = 256; + +/// On-disk encoding of one impact entry. +/// +/// `FixedU32` (128-doc blocks): [doc_up_to u32][pair_count u32][(freq u32, +/// doc_len u32)...]. `Varint` (256-doc blocks, format V3): [doc_up_to varint] +/// [pair_count varint][(freq delta varint, doc_len varint)...] with the +/// frontier's freqs delta-encoded in ascending order — pairs shrink from 8 +/// bytes to ~2-3. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ImpactFormat { + FixedU32, + Varint, +} + +impl ImpactFormat { + pub fn for_block_size(block_size: usize) -> Self { + if block_size == 256 { + Self::Varint + } else { + Self::FixedU32 + } + } +} + +#[derive(Debug, Clone)] +pub struct ImpactSkipData { + entries: LargeBinaryArray, + level0_len: usize, + format: ImpactFormat, + // Last doc id covered by each entry (level0 entries then level1 entries), + // decoded once at construction; u32::MAX marks malformed entries. + entry_doc_up_tos: Arc<[u32]>, + // Per-entry max doc weight plus the list-wide max, baked on first use and + // shared across the per-query clones of a cached list. Doc weights depend + // only on (freq, doc_len) and index-wide stats (e.g. BM25 avgdl), not the + // query. Malformed entries bake to INFINITY so pruning stays safe. + doc_weight_bounds: Arc, f32)>>, +} + +impl PartialEq for ImpactSkipData { + fn eq(&self, other: &Self) -> bool { + self.entries == other.entries && self.level0_len == other.level0_len + } +} + +#[derive(Debug, Clone, Copy)] +pub struct ImpactScore { + pub score: f32, + pub entries_scanned: usize, +} + +#[derive(Debug, Default, Clone)] +pub struct ImpactScoreCache {} + +impl ImpactScoreCache { + fn entry_score( + &mut self, + impacts: &ImpactSkipData, + entry_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if query_weight <= 0.0 { + return 0.0; + } + query_weight * impacts.doc_weight_bounds(scorer)[entry_idx] + } +} + +#[derive(Debug, Clone, Copy)] +struct ImpactEntryHeader { + doc_up_to: u32, + pair_count: usize, +} + +impl ImpactSkipData { + pub fn new(entries: LargeBinaryArray, level0_len: usize, format: ImpactFormat) -> Result { + let expected_len = level0_len + level1_len(level0_len); + if entries.len() != expected_len { + return Err(Error::index(format!( + "impact entry count mismatch: got {}, expected {} for {} level0 blocks", + entries.len(), + expected_len, + level0_len + ))); + } + let entry_doc_up_tos = (0..entries.len()) + .map(|entry_idx| { + if entries.is_null(entry_idx) { + return u32::MAX; + } + decode_entry_doc_up_to(entries.value(entry_idx), format).unwrap_or(u32::MAX) + }) + .collect::>(); + Ok(Self { + entries, + level0_len, + format, + entry_doc_up_tos, + doc_weight_bounds: Arc::new(OnceLock::new()), + }) + } + + fn baked_bounds(&self, scorer: &S) -> &(Box<[f32]>, f32) { + self.doc_weight_bounds.get_or_init(|| { + let per_entry = (0..self.entries.len()) + .map(|entry_idx| { + let bytes = self.entries.value(entry_idx); + let mut max_doc_weight = 0.0_f32; + match for_each_entry_pair(bytes, self.format, |freq, doc_len| { + max_doc_weight = max_doc_weight.max(scorer.doc_weight(freq, doc_len)); + }) { + Ok(()) => max_doc_weight, + Err(_) => f32::INFINITY, + } + }) + .collect::>(); + // The level1 entries cover every block, so their max is the + // list-wide max doc weight; fall back to level0 entries for lists + // too short to have a level1 entry. + let global = if per_entry.len() > self.level0_len { + per_entry[self.level0_len..] + .iter() + .copied() + .fold(0.0_f32, f32::max) + } else { + per_entry.iter().copied().fold(0.0_f32, f32::max) + }; + (per_entry, global) + }) + } + + fn doc_weight_bounds(&self, scorer: &S) -> &[f32] { + &self.baked_bounds(scorer).0 + } + + /// List-wide max doc weight, from the baked bounds. The tightest valid + /// global score bound for this list is `query_weight * this`, matching + /// what the non-impact format stores as `max_score` at build time — but + /// computed against the current index stats. + pub fn global_max_doc_weight(&self, scorer: &S) -> f32 { + self.baked_bounds(scorer).1 + } + + pub fn entries(&self) -> &LargeBinaryArray { + &self.entries + } + + #[cfg(test)] + pub fn level0_len(&self) -> usize { + self.level0_len + } + + #[cfg(test)] + pub fn level1_len(&self) -> usize { + level1_len(self.level0_len) + } + + pub(crate) fn level1_doc_up_to(&self, group_idx: usize) -> Option { + if group_idx >= level1_len(self.level0_len) { + return None; + } + match self.entry_doc_up_tos[self.level0_len + group_idx] { + u32::MAX => None, + doc_up_to => Some(doc_up_to), + } + } + + /// Max score of the docs covered by the level0 entry of `block_idx`, + /// answered from the baked bounds slab. + pub fn level0_score( + &self, + block_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if block_idx >= self.level0_len || query_weight <= 0.0 { + return 0.0; + } + query_weight * self.doc_weight_bounds(scorer)[block_idx] + } + + pub fn level0_score_cached( + &self, + block_idx: usize, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> f32 { + if block_idx >= self.level0_len { + return 0.0; + } + cache.entry_score(self, block_idx, query_weight, scorer) + } + + #[cfg(test)] + pub fn max_score_up_to( + &self, + start_block_idx: usize, + up_to: u64, + _block_least_doc_id: F, + query_weight: f32, + scorer: &S, + ) -> ImpactScore + where + S: Scorer + ?Sized, + F: FnMut(usize) -> u32, + { + self.max_score_up_to_with(start_block_idx, up_to, |impacts, entry_idx| { + impacts.entry_score(entry_idx, query_weight, scorer) + }) + } + + pub fn max_score_up_to_cached( + &self, + start_block_idx: usize, + up_to: u64, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> ImpactScore + where + S: Scorer + ?Sized, + { + self.max_score_up_to_with(start_block_idx, up_to, |impacts, entry_idx| { + cache.entry_score(impacts, entry_idx, query_weight, scorer) + }) + } + + fn max_score_up_to_with( + &self, + start_block_idx: usize, + up_to: u64, + mut entry_score: E, + ) -> ImpactScore + where + E: FnMut(&Self, usize) -> f32, + { + let mut block_idx = start_block_idx; + let mut max_score = 0.0_f32; + let mut entries_scanned = 0usize; + + while block_idx < self.level0_len { + let group_idx = block_idx / IMPACT_LEVEL1_BLOCKS; + let group_start = group_idx * IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(self.level0_len); + if block_idx == group_start { + let level1_entry_idx = self.level0_len + group_idx; + match self.entry_doc_up_tos[level1_entry_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned: entries_scanned + 1, + }; + } + doc_up_to if u64::from(doc_up_to) <= up_to => { + max_score = max_score.max(entry_score(self, level1_entry_idx)); + entries_scanned += 1; + block_idx = group_end; + continue; + } + _ => {} + } + } + + max_score = max_score.max(entry_score(self, block_idx)); + entries_scanned += 1; + match self.entry_doc_up_tos[block_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned, + }; + } + doc_up_to if u64::from(doc_up_to) >= up_to => break, + _ => {} + } + block_idx += 1; + } + + ImpactScore { + score: max_score, + entries_scanned, + } + } + + #[cfg(test)] + fn entry_score( + &self, + entry_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if query_weight <= 0.0 { + return 0.0; + } + let bytes = self.entries.value(entry_idx); + let mut max_doc_weight = 0.0_f32; + if for_each_entry_pair(bytes, self.format, |freq, doc_len| { + max_doc_weight = max_doc_weight.max(scorer.doc_weight(freq, doc_len)); + }) + .is_err() + { + return f32::INFINITY; + } + query_weight * max_doc_weight + } +} + +pub struct ImpactSkipDataBuilder { + entries: LargeBinaryBuilder, + level0_len: usize, + level1_entries: Vec>, + level1_docs: Vec<(u32, u32, u32)>, + format: ImpactFormat, +} + +impl ImpactSkipDataBuilder { + pub fn with_capacity(level0_blocks: usize, block_size: usize) -> Self { + Self { + entries: LargeBinaryBuilder::with_capacity( + level0_blocks + level1_len(level0_blocks), + 0, + ), + level0_len: 0, + level1_entries: Vec::with_capacity(level1_len(level0_blocks)), + level1_docs: Vec::with_capacity(IMPACT_LEVEL1_BLOCKS * block_size), + format: ImpactFormat::for_block_size(block_size), + } + } + + pub fn append_block(&mut self, docs: &[(u32, u32, u32)]) -> Result<()> { + let bytes = encode_impact_entry(docs, self.format)?; + self.entries.append_value(bytes.as_slice()); + self.level0_len += 1; + self.level1_docs.extend_from_slice(docs); + if self.level0_len.is_multiple_of(IMPACT_LEVEL1_BLOCKS) { + self.flush_level1()?; + } + Ok(()) + } + + pub fn finish(mut self) -> Result { + if !self.level1_docs.is_empty() { + self.flush_level1()?; + } + for entry in self.level1_entries { + self.entries.append_value(entry.as_slice()); + } + ImpactSkipData::new(self.entries.finish(), self.level0_len, self.format) + } + + fn flush_level1(&mut self) -> Result<()> { + let bytes = encode_impact_entry(self.level1_docs.as_slice(), self.format)?; + self.level1_entries.push(bytes); + self.level1_docs.clear(); + Ok(()) + } +} + +#[cfg(test)] +pub fn build_impact_skip_data(blocks: &[Vec<(u32, u32, u32)>]) -> Result { + let block_size = blocks.iter().map(Vec::len).max().unwrap_or(0).max(1); + let mut builder = ImpactSkipDataBuilder::with_capacity(blocks.len(), block_size); + for block in blocks { + builder.append_block(block)?; + } + builder.finish() +} + +fn encode_impact_entry(docs: &[(u32, u32, u32)], format: ImpactFormat) -> Result> { + let doc_up_to = docs + .last() + .map(|(doc_id, _, _)| *doc_id) + .unwrap_or_default(); + let frontier = impact_frontier(docs); + let pair_count = u32::try_from(frontier.len()).map_err(|_| { + Error::index("impact frontier too large to encode as u32 pair count".to_string()) + })?; + let mut bytes = Vec::with_capacity(8 + frontier.len() * 8); + match format { + ImpactFormat::FixedU32 => { + bytes.extend_from_slice(&doc_up_to.to_le_bytes()); + bytes.extend_from_slice(&pair_count.to_le_bytes()); + for (freq, doc_len) in frontier { + bytes.extend_from_slice(&freq.to_le_bytes()); + bytes.extend_from_slice(&doc_len.to_le_bytes()); + } + } + ImpactFormat::Varint => { + super::encoding::encode_varint_u32(&mut bytes, doc_up_to); + super::encoding::encode_varint_u32(&mut bytes, pair_count); + let mut previous_freq = 0u32; + for (freq, doc_len) in frontier { + super::encoding::encode_varint_u32(&mut bytes, freq - previous_freq); + super::encoding::encode_varint_u32(&mut bytes, doc_len); + previous_freq = freq; + } + } + } + Ok(bytes) +} + +fn decode_entry_doc_up_to(bytes: &[u8], format: ImpactFormat) -> Result { + match format { + ImpactFormat::FixedU32 => decode_header(bytes).map(|header| header.doc_up_to), + ImpactFormat::Varint => { + let mut offset = 0usize; + super::encoding::decode_varint_u32(bytes, &mut offset) + } + } +} + +/// Walk an entry's (freq, doc_len) frontier pairs, validating the layout. +fn for_each_entry_pair( + bytes: &[u8], + format: ImpactFormat, + mut visit: impl FnMut(u32, u32), +) -> Result<()> { + match format { + ImpactFormat::FixedU32 => { + let header = decode_header(bytes)?; + let mut offset = 8usize; + for _ in 0..header.pair_count { + visit(read_u32_le(bytes, offset), read_u32_le(bytes, offset + 4)); + offset += 8; + } + } + ImpactFormat::Varint => { + let mut offset = 0usize; + let _doc_up_to = super::encoding::decode_varint_u32(bytes, &mut offset)?; + let pair_count = super::encoding::decode_varint_u32(bytes, &mut offset)?; + let mut freq = 0u32; + for _ in 0..pair_count { + freq = freq + .checked_add(super::encoding::decode_varint_u32(bytes, &mut offset)?) + .ok_or_else(|| Error::index("impact freq delta overflow".to_owned()))?; + let doc_len = super::encoding::decode_varint_u32(bytes, &mut offset)?; + visit(freq, doc_len); + } + if offset != bytes.len() { + return Err(Error::index(format!( + "impact varint entry has {} trailing bytes", + bytes.len() - offset + ))); + } + } + } + Ok(()) +} + +fn impact_frontier(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let max_freq = docs.iter().map(|(_, freq, _)| *freq).max().unwrap_or(0) as usize; + if max_freq <= SMALL_FRONTIER_FREQ_LIMIT { + return impact_frontier_small_freq(docs, max_freq); + } + + impact_frontier_sparse_freq(docs) +} + +fn impact_frontier_small_freq(docs: &[(u32, u32, u32)], max_freq: usize) -> Vec<(u32, u32)> { + let mut min_doc_len_by_freq = [u32::MAX; SMALL_FRONTIER_FREQ_LIMIT + 1]; + for (_, freq, doc_len) in docs { + min_doc_len_by_freq[*freq as usize] = min_doc_len_by_freq[*freq as usize].min(*doc_len); + } + + let min_doc_lens = min_doc_len_by_freq[..=max_freq] + .iter() + .enumerate() + .filter_map(|(freq, doc_len)| (*doc_len != u32::MAX).then_some((freq as u32, *doc_len))) + .collect::>(); + frontier_from_min_doc_lens(min_doc_lens) +} + +fn impact_frontier_sparse_freq(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let mut pairs = docs + .iter() + .map(|(_, freq, doc_len)| (*freq, *doc_len)) + .collect::>(); + pairs.sort_unstable_by_key(|(freq, _)| *freq); + + let mut min_doc_lens: Vec<(u32, u32)> = Vec::with_capacity(pairs.len()); + for (freq, doc_len) in pairs { + match min_doc_lens.last_mut() { + Some((last_freq, last_doc_len)) if *last_freq == freq => { + *last_doc_len = (*last_doc_len).min(doc_len); + } + _ => min_doc_lens.push((freq, doc_len)), + } + } + + frontier_from_min_doc_lens(min_doc_lens) +} + +fn frontier_from_min_doc_lens(min_doc_lens: Vec<(u32, u32)>) -> Vec<(u32, u32)> { + let mut best_doc_len = u32::MAX; + let mut frontier = Vec::with_capacity(min_doc_lens.len()); + for (freq, doc_len) in min_doc_lens.into_iter().rev() { + if doc_len < best_doc_len { + frontier.push((freq, doc_len)); + best_doc_len = doc_len; + } + } + frontier.reverse(); + frontier +} + +fn decode_header(bytes: &[u8]) -> Result { + if bytes.len() < 8 { + return Err(Error::index(format!( + "impact entry too short: {} bytes", + bytes.len() + ))); + } + let pair_count = read_u32_le(bytes, 4) as usize; + let expected_len = 8 + pair_count * 8; + if bytes.len() != expected_len { + return Err(Error::index(format!( + "impact entry length mismatch: got {} bytes, expected {} for {} pairs", + bytes.len(), + expected_len, + pair_count + ))); + } + Ok(ImpactEntryHeader { + doc_up_to: read_u32_le(bytes, 0), + pair_count, + }) +} + +#[inline] +fn read_u32_le(bytes: &[u8], offset: usize) -> u32 { + let mut value = [0u8; 4]; + value.copy_from_slice(&bytes[offset..offset + 4]); + u32::from_le_bytes(value) +} + +fn level1_len(level0_len: usize) -> usize { + level0_len.div_ceil(IMPACT_LEVEL1_BLOCKS) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::scalar::inverted::scorer::{MemBM25Scorer, Scorer}; + + #[test] + fn impact_entry_frontier_drops_dominated_pairs() { + let docs = vec![(0, 1, 10), (1, 1, 8), (2, 2, 9), (3, 3, 20)]; + assert_eq!(impact_frontier(&docs), vec![(1, 8), (2, 9), (3, 20)]); + } + + #[test] + fn impact_entry_frontier_handles_sparse_large_frequencies() { + let docs = vec![ + (0, 1, 100), + (1, 1, 80), + (2, 512, 90), + (3, 1_000, 120), + (4, 1_000, 110), + ]; + assert_eq!( + impact_frontier(&docs), + vec![(1, 80), (512, 90), (1_000, 110)] + ); + } + + #[test] + fn impact_max_score_can_use_level1_entry() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level0_len(), 40); + assert_eq!(impacts.level1_len(), 2); + let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); + let score = impacts.max_score_up_to(0, 31, |idx| idx as u32, 1.0, &scorer); + assert!(score.entries_scanned < IMPACT_LEVEL1_BLOCKS); + assert!(score.score > 0.0); + } + + #[test] + fn impact_level1_doc_up_to_reports_full_and_partial_groups() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + + assert_eq!( + impacts.level1_doc_up_to(0), + Some((IMPACT_LEVEL1_BLOCKS - 1) as u32) + ); + assert_eq!(impacts.level1_doc_up_to(1), Some(39)); + assert_eq!(impacts.level1_doc_up_to(2), None); + } + + #[test] + fn impact_level1_doc_up_to_returns_none_for_malformed_entry() { + let level0 = encode_impact_entry(&[(0, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let malformed_level1 = vec![1, 2, 3]; + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0.as_slice()), + Some(malformed_level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 1, ImpactFormat::FixedU32).unwrap(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + } + + #[test] + fn impact_score_cache_matches_uncached_scores() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); + let mut cache = ImpactScoreCache::default(); + + let uncached_level0 = impacts.level0_score(3, 1.0, &scorer); + let cached_level0 = impacts.level0_score_cached(3, 1.0, &scorer, &mut cache); + assert_eq!(cached_level0, uncached_level0); + + let uncached = impacts.max_score_up_to(0, 31, |idx| idx as u32, 1.0, &scorer); + let cached = impacts.max_score_up_to_cached(0, 31, 1.0, &scorer, &mut cache); + assert_eq!(cached.score, uncached.score); + assert_eq!(cached.entries_scanned, uncached.entries_scanned); + } + + #[test] + fn impact_entries_are_decoded_lazily() { + let level0_0 = encode_impact_entry(&[(0, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let malformed_level0_1 = vec![1, 2, 3]; + let level1 = + encode_impact_entry(&[(0, 1, 10), (1, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0_0.as_slice()), + Some(malformed_level0_1.as_slice()), + Some(level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 2, ImpactFormat::FixedU32).unwrap(); + let scorer = MemBM25Scorer::new(10, 10, HashMap::from([(String::from("token"), 2usize)])); + + let score = impacts.max_score_up_to(0, 0, |_| 0, 1.0, &scorer); + assert!(score.score.is_finite()); + assert_eq!(score.entries_scanned, 1); + + let mut cache = ImpactScoreCache::default(); + assert_eq!( + impacts.level0_score_cached(1, 1.0, &scorer, &mut cache), + f32::INFINITY + ); + } + + #[test] + fn impact_varint_entries_roundtrip_and_match_fixed() { + let docs = vec![(3, 1, 100), (9, 2, 40), (200, 7, 80), (4095, 130, 900)]; + let fixed = encode_impact_entry(&docs, ImpactFormat::FixedU32).unwrap(); + let varint = encode_impact_entry(&docs, ImpactFormat::Varint).unwrap(); + assert!(varint.len() < fixed.len()); + assert_eq!( + decode_entry_doc_up_to(&fixed, ImpactFormat::FixedU32).unwrap(), + decode_entry_doc_up_to(&varint, ImpactFormat::Varint).unwrap() + ); + let mut fixed_pairs = Vec::new(); + for_each_entry_pair(&fixed, ImpactFormat::FixedU32, |f, l| { + fixed_pairs.push((f, l)) + }) + .unwrap(); + let mut varint_pairs = Vec::new(); + for_each_entry_pair(&varint, ImpactFormat::Varint, |f, l| { + varint_pairs.push((f, l)) + }) + .unwrap(); + assert_eq!(fixed_pairs, varint_pairs); + assert!(!fixed_pairs.is_empty()); + + // a 256-doc-block skip data goes through the varint path end to end + let blocks: Vec> = (0..3) + .map(|b| (0..256).map(|i| (b * 256 + i, 1 + i % 5, 10)).collect()) + .collect(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level1_doc_up_to(0), Some(767)); + let scorer = MemBM25Scorer::new(400, 768, HashMap::from([(String::from("t"), 768usize)])); + assert!(impacts.level0_score(0, 1.0, &scorer).is_finite()); + let level1 = impacts.max_score_up_to(0, 767, |idx| (idx * 256) as u32, 1.0, &scorer); + assert!(level1.score.is_finite() && level1.score > 0.0); + } + + #[test] + fn impact_upper_bound_covers_real_scores() { + let blocks = vec![ + vec![(0, 1, 100), (3, 2, 40), (7, 4, 80)], + vec![(9, 3, 15), (10, 1, 5), (12, 5, 30)], + vec![(16, 2, 10), (18, 6, 70), (21, 3, 12)], + vec![(24, 1, 4), (28, 7, 100), (30, 2, 8)], + ]; + let impacts = build_impact_skip_data(&blocks).unwrap(); + let scorer = MemBM25Scorer::new(474, 31, HashMap::from([(String::from("token"), 4usize)])); + let query_weight = scorer.query_weight("token"); + + for start_block_idx in 0..blocks.len() { + let up_to = blocks + .iter() + .skip(start_block_idx) + .take(2) + .flatten() + .map(|(doc_id, _, _)| *doc_id) + .max() + .unwrap(); + let upper_bound = impacts.max_score_up_to( + start_block_idx, + u64::from(up_to), + |idx| blocks[idx][0].0, + query_weight, + &scorer, + ); + let exact_max = blocks + .iter() + .skip(start_block_idx) + .flatten() + .take_while(|(doc_id, _, _)| *doc_id <= up_to) + .map(|(_, freq, doc_len)| query_weight * scorer.doc_weight(*freq, *doc_len)) + .fold(0.0_f32, f32::max); + assert!( + upper_bound.score + 1e-6 >= exact_max, + "upper bound {} should cover exact max {} from block {} up to doc {}", + upper_bound.score, + exact_max, + start_block_idx, + up_to + ); + } + } +} diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index f1bfffa9fb4..269d7ff8dac 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -53,6 +53,7 @@ use tokio::{sync::OnceCell, task::spawn_blocking}; use tracing::{info, instrument}; use super::encoding::{PositionBlockBuilder, decode_group_starts}; +use super::impact::{ImpactSkipData, ImpactSkipDataBuilder}; use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; @@ -60,7 +61,8 @@ use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ BLOCK_SIZE, PostingGroupAccumulator, PostingGroupConfig, ScoredDoc, doc_file_path, - inverted_list_schema_for_version_with_block_size, posting_file_path, token_file_path, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -105,6 +107,7 @@ 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"; @@ -291,6 +294,7 @@ impl PartitionCandidates { struct LoadedPostings { postings: Vec, grouped_expansions: Vec, + impact_safe: bool, } impl LoadedPostings { @@ -298,6 +302,7 @@ impl LoadedPostings { Self { postings: Vec::new(), grouped_expansions: Vec::new(), + impact_safe: false, } } } @@ -854,7 +859,7 @@ impl InvertedIndex { // hits per-token `posting_len`; building a `MemBM25Scorer` with // precomputed per-term IDFs avoids the v2 bulk metadata pull. let local_scorer; - let scorer: &dyn Scorer = if let Some(base_scorer) = base_scorer { + let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer { base_scorer } else { local_scorer = self @@ -862,6 +867,7 @@ impl InvertedIndex { .await?; &local_scorer }; + let impact_scorer = Arc::new(scorer.clone()); let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { @@ -900,8 +906,9 @@ impl InvertedIndex { // Shared top-k floor across this query's partitions. Seeded to -inf so // the first real score wins; each partition publishes its local k-th // and prunes against the running global k-th (a lower bound on the true - // global k-th — see `Wand::shared_threshold`). - let shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + // global k-th - see `Wand::shared_threshold`). + let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self .partitions .iter() @@ -911,19 +918,23 @@ impl InvertedIndex { let params = params.clone(); let mask = mask.clone(); let metrics = metrics.clone(); - let shared_threshold = shared_threshold.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { let loaded_postings = part .load_posting_lists( tokens.as_ref(), params.as_ref(), operator, + Some(impact_scorer.as_ref()), metrics.as_ref(), ) .await?; let LoadedPostings { postings, grouped_expansions, + impact_safe, } = loaded_postings; if postings.is_empty() { // No hits in this partition; its DocSet stays @@ -947,6 +958,7 @@ impl InvertedIndex { let metrics = metrics.clone(); let part_for_wand = part.clone(); let has_grouped_expansions = !grouped_expansions.is_empty(); + let use_impact_path = impact_safe && !has_grouped_expansions; let wand_params = if has_grouped_expansions { let mut rescoring_params = params.as_ref().clone(); rescoring_params.limit = @@ -957,9 +969,12 @@ impl InvertedIndex { }; let partition_threshold = if has_grouped_expansions { Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) + } else if use_impact_path { + impact_shared_threshold } else { - shared_threshold + legacy_shared_threshold }; + let wand_scorer = use_impact_path.then(|| impact_scorer.clone()); let candidates = spawn_cpu(move || { let candidates = part_for_wand.bm25_search( docs_for_wand.as_ref(), @@ -967,6 +982,7 @@ impl InvertedIndex { operator, mask, postings, + wand_scorer, metrics.as_ref(), partition_threshold, )?; @@ -1739,6 +1755,7 @@ impl InvertedPartition { tokens: &Tokens, params: &FtsSearchParams, operator: Operator, + impact_scorer: Option<&MemBM25Scorer>, metrics: &dyn MetricsCollector, ) -> Result { let is_fuzzy = matches!(params.fuzziness, Some(n) if n != 0); @@ -1816,11 +1833,18 @@ impl InvertedPartition { } if !is_fuzzy_and_query { + let impact_safe = impact_scorer.is_some() + && loaded_postings + .iter() + .all(|(_, _, _, posting)| posting.has_impacts()); return Ok(LoadedPostings { postings: loaded_postings .into_iter() .map(|(token_id, token, position, posting)| { - let query_weight = idf(posting.len(), num_docs); + let query_weight = impact_scorer + .filter(|_| impact_safe) + .map(|scorer| scorer.query_weight(&token)) + .unwrap_or_else(|| idf(posting.len(), num_docs)); PostingIterator::with_query_weight( token, token_id, @@ -1832,6 +1856,7 @@ impl InvertedPartition { }) .collect(), grouped_expansions: Vec::new(), + impact_safe, }); } @@ -1899,6 +1924,7 @@ impl InvertedPartition { Ok(LoadedPostings { postings: grouped_postings, grouped_expansions, + impact_safe: false, }) } @@ -1914,6 +1940,7 @@ impl InvertedPartition { operator: Operator, mask: Arc, postings: Vec, + impact_scorer: Option>, metrics: &dyn MetricsCollector, shared_threshold: Arc, ) -> Result> { @@ -1924,10 +1951,16 @@ impl InvertedPartition { // Caller selects the DocSet shape via `LazyDocSet::docs_for_wand` // and passes it in here; wand uses `docs.has_row_ids()` to // handle the num_tokens-only case. - let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) - .with_shared_threshold(shared_threshold); - let hits = wand.search(params, mask, metrics)?; + let hits = if let Some(scorer) = impact_scorer { + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + } else { + let scorer = IndexBM25Scorer::new(std::iter::once(self)); + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + }; Ok(hits) } @@ -2346,6 +2379,7 @@ pub struct PostingListReader { metadata: PostingMetadata, has_position: bool, + has_impacts: bool, posting_tail_codec: PostingTailCodec, block_size: usize, positions_layout: PositionsLayout, @@ -2448,6 +2482,7 @@ impl PostingListReader { 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 { @@ -2466,6 +2501,7 @@ impl PostingListReader { reader, metadata, has_position, + has_impacts, posting_tail_codec, block_size, positions_layout, @@ -2662,7 +2698,7 @@ impl PostingListReader { self.posting_batch_legacy(token_id, with_position).await } else { let token_id = token_id as usize; - let columns = if with_position { + let mut columns = if with_position { match self.positions_layout { PositionsLayout::SharedStream(_) => { vec![ @@ -2677,6 +2713,9 @@ impl PostingListReader { } 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)) @@ -2721,11 +2760,18 @@ impl PostingListReader { Some((start, end)) => { let group = self .index_cache - .get_or_insert_with_key(PostingListGroupKey { start, end }, || 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 - }) + .get_or_insert_with_key( + PostingListGroupKey { + start, + end, + has_impacts: 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?; let slot = (token_id - start) as usize; group @@ -2741,19 +2787,25 @@ impl PostingListReader { // per token. None => self .index_cache - .get_or_insert_with_key(PostingListKey { token_id }, || 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), - )?; - self.posting_list_from_batch(&batch, max_score, length) - }) + .get_or_insert_with_key( + PostingListKey { + token_id, + has_impacts: 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), + )?; + self.posting_list_from_batch(&batch, max_score, length) + }, + ) .await? .as_ref() .clone(), @@ -2798,12 +2850,13 @@ impl PostingListReader { /// [`PostingListGroup`] cache value (issue #7040). 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(&[POSTING_COL, MAX_SCORE_COL, LENGTH_COL]), - ) + .read_range(start as usize..end as usize, Some(&columns)) .await?; let max_scores = batch[MAX_SCORE_COL].as_primitive::(); let lengths = batch[LENGTH_COL].as_primitive::(); @@ -3126,7 +3179,14 @@ impl PostingListReader { let hi = end as usize - tok_start; let group = PostingListGroup::new(chunk_postings[lo..hi].to_vec()); self.index_cache - .insert_with_key(&PostingListGroupKey { start, end }, Arc::new(group)) + .insert_with_key( + &PostingListGroupKey { + start, + end, + has_impacts: self.has_impacts, + }, + Arc::new(group), + ) .await; } } @@ -3135,7 +3195,13 @@ impl PostingListReader { self.cache_positions(&mut posting_list, token_id, with_position) .await; self.index_cache - .insert_with_key(&PostingListKey { token_id }, Arc::new(posting_list)) + .insert_with_key( + &PostingListKey { + token_id, + has_impacts: self.has_impacts, + }, + Arc::new(posting_list), + ) .await; } } @@ -3305,6 +3371,9 @@ impl PostingListReader { } } } + if self.has_impacts { + base_columns.push(IMPACT_COL); + } base_columns } } @@ -3401,13 +3470,18 @@ impl DeepSizeOf for Positions { #[derive(Debug, Clone)] pub struct PostingListKey { pub token_id: u32, + pub has_impacts: bool, } impl CacheKey for PostingListKey { type ValueType = PostingList; fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}", self.token_id).into() + if self.has_impacts { + format!("postings-{}-impacts", self.token_id).into() + } else { + format!("postings-{}", self.token_id).into() + } } fn type_name() -> &'static str { @@ -3427,13 +3501,18 @@ impl CacheKey for PostingListKey { pub struct PostingListGroupKey { pub start: u32, pub end: u32, + pub has_impacts: bool, } impl CacheKey for PostingListGroupKey { type ValueType = PostingListGroup; fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}-{}", self.start, self.end).into() + if self.has_impacts { + format!("postings-{}-{}-impacts", self.start, self.end).into() + } else { + format!("postings-{}-{}", self.start, self.end).into() + } } fn type_name() -> &'static str { @@ -3579,6 +3658,7 @@ impl PostingListGroup { } #[derive(Debug, Clone, DeepSizeOf)] +#[allow(clippy::large_enum_variant)] pub enum PostingList { Plain(PlainPostingList), Compressed(CompressedPostingList), @@ -3643,7 +3723,7 @@ impl PostingList { posting_tail_codec, block_size, shared_position_codec, - ); + )?; Ok(Self::Compressed(posting)) } None => { @@ -3664,6 +3744,13 @@ impl PostingList { } } + 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 { @@ -3884,6 +3971,7 @@ pub struct CompressedPostingList { 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: Arc>>, @@ -3897,6 +3985,7 @@ impl PartialEq for CompressedPostingList { && self.posting_tail_codec == other.posting_tail_codec && self.block_size == other.block_size && self.positions == other.positions + && self.impacts == other.impacts } } @@ -3908,17 +3997,23 @@ impl DeepSizeOf for CompressedPostingList { .as_ref() .map(|positions| positions.deep_size_of_children(context)) .unwrap_or(0) + + self + .impacts + .as_ref() + .map(|impacts| sliced_cache_bytes(impacts.entries())) + .unwrap_or(0) } } impl CompressedPostingList { - pub fn new( + pub(crate) fn new( blocks: LargeBinaryArray, max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, block_size: usize, positions: Option, + impacts: Option, ) -> Self { Self { max_score, @@ -3927,6 +4022,7 @@ impl CompressedPostingList { posting_tail_codec, block_size, positions, + impacts, first_docs: Arc::new(OnceLock::new()), } } @@ -3938,7 +4034,7 @@ impl CompressedPostingList { posting_tail_codec: PostingTailCodec, block_size: usize, shared_position_codec: Option, - ) -> Self { + ) -> Result { debug_assert_eq!(batch.num_rows(), 1); let blocks = batch[POSTING_COL] .as_list::() @@ -3967,16 +4063,28 @@ impl CompressedPostingList { ) }) }; + 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(), + super::impact::ImpactFormat::for_block_size(block_size), + ) + }) + .transpose()?; - Self { + Ok(Self { max_score, length, blocks, posting_tail_codec, block_size, positions, + impacts, first_docs: Arc::new(OnceLock::new()), - } + }) } pub fn iter(&self) -> CompressedPostingListIterator { @@ -4150,6 +4258,7 @@ pub struct PostingListBuilder { pub(super) struct PostingListBatchBuilder { schema: SchemaRef, postings: ListBuilder, + impacts: Option>, max_scores: Float32Builder, lengths: UInt32Builder, positions: BatchPositionsBuilder, @@ -4201,9 +4310,14 @@ impl PostingListBatchBuilder { 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, @@ -4223,11 +4337,19 @@ impl PostingListBatchBuilder { fn append( &mut self, compressed: LargeBinaryArray, + impacts: Option<&ImpactSkipData>, max_score: f32, length: u32, positions: Option<&CompressedPositionStorage>, ) -> Result<()> { - let posting_bytes = compressed.value_data().len(); + let impact_bytes = if self.impacts.is_some() { + impacts + .map(|impacts| impacts.entries().value_data().len()) + .unwrap_or(0) + } else { + 0 + }; + let posting_bytes = compressed.value_data().len() + impact_bytes; { let values = self.postings.values(); for index in 0..compressed.len() { @@ -4235,6 +4357,19 @@ impl PostingListBatchBuilder { } } 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.group_accumulator.push(posting_bytes); self.max_scores.append_value(max_score); self.lengths.append_value(length); @@ -4300,6 +4435,9 @@ impl PostingListBatchBuilder { 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) => { @@ -4717,6 +4855,7 @@ impl PostingListBuilder { fn build_batch( self, compressed: LargeBinaryArray, + impacts: Option, max_score: f32, schema: SchemaRef, positions: Option, @@ -4735,6 +4874,22 @@ impl PostingListBuilder { 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)?; @@ -4805,13 +4960,19 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + 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, max_score, len, positions.as_ref()) + batch_builder.append( + compressed, + Some(&impacts), + max_score, + len, + positions.as_ref(), + ) } fn extend_tail_components( @@ -4828,7 +4989,12 @@ impl PostingListBuilder { fn build_compressed_with_scores_from_parts( parts: PostingListParts<'_>, docs: &DocSet, - ) -> Result<(LargeBinaryArray, Option, f32)> { + ) -> Result<( + LargeBinaryArray, + Option, + f32, + ImpactSkipData, + )> { let PostingListParts { with_positions, posting_tail_codec, @@ -4844,6 +5010,9 @@ impl PostingListBuilder { 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); @@ -4855,26 +5024,30 @@ impl PostingListBuilder { &mut frequencies, block_size, ); - let block_score = compute_block_score( + 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.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( + 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(), @@ -4891,10 +5064,12 @@ impl PostingListBuilder { } } + let impacts = impact_builder.finish()?; Ok(( encoded_blocks.into_array(), with_positions.then(|| encoded_position_blocks.into_stream()), max_score, + impacts, )) } @@ -4952,10 +5127,11 @@ impl PostingListBuilder { self.posting_tail_codec, self.block_size, )?; - let schema = inverted_list_schema_for_version_with_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() { @@ -5012,7 +5188,7 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, None, max_score, schema, positions) } pub fn to_batch_with_docs(self, docs: &DocSet, schema: SchemaRef) -> Result { @@ -5054,7 +5230,7 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + let (compressed, shared_positions, max_score, impacts) = Self::build_compressed_with_scores_from_parts(parts, docs)?; let builder = Self { with_positions, @@ -5074,7 +5250,7 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, Some(impacts), max_score, schema, positions) } pub fn remap(&mut self, removed: &[u32]) { @@ -5102,19 +5278,23 @@ impl PostingListBuilder { } } -fn compute_block_score( +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_norm = K1 * (1.0 - B + B * docs.num_tokens(doc_id) as f32 / avgdl); - let freq = freq as f32; - let score = freq / (freq + doc_norm); + 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 } @@ -5963,7 +6143,10 @@ mod tests { use crate::prefilter::NoFilter; use crate::scalar::ScalarIndex; use crate::scalar::inverted::builder::{ - InnerBuilder, InvertedIndexBuilder, PositionRecorder, inverted_list_schema, + InnerBuilder, InvertedIndexBuilder, PositionRecorder, doc_file_path, inverted_list_schema, + inverted_list_schema_for_version_with_block_size, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }; use crate::scalar::inverted::encoding::{ compress_positions, compress_posting_list_with_tail_codec, @@ -6063,6 +6246,57 @@ mod tests { assert!(err.to_string().contains("block_size")); } + #[test] + fn test_posting_builder_writes_impacts_for_supported_block_sizes() { + for block_size in [128, 256] { + let format_version = default_fts_format_version_for_block_size(block_size).unwrap(); + let num_docs = block_size * 33 + 1; + let mut docs = DocSet::default(); + let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); + for doc_id in 0..num_docs { + docs.append(doc_id as u64, (doc_id % 5 + 1) as u32); + posting.add( + doc_id as u32, + PositionRecorder::Count((doc_id % 3 + 1) as u32), + ); + } + let schema = + inverted_list_schema_for_version_with_block_size(false, format_version, block_size); + let batch = posting.to_batch_with_docs(&docs, schema).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_some()); + let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); + let length = batch[LENGTH_COL].as_primitive::().value(0); + let posting = PostingList::from_batch(&batch, Some(max_score), Some(length)).unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed posting list"); + }; + let impacts = posting.impacts.expect("posting should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); + } + } + + #[test] + fn test_posting_builder_without_impact_column_roundtrips_without_impacts() { + let mut posting = PostingListBuilder::new(false); + for doc_id in 0..BLOCK_SIZE + 3 { + posting.add(doc_id as u32, PositionRecorder::Count(1)); + } + let batch = posting.to_batch(vec![1.0, 1.0]).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_none()); + let posting = + PostingList::from_batch(&batch, Some(1.0), Some((BLOCK_SIZE + 3) as u32)).unwrap(); + assert!(!posting.has_impacts()); + } + #[tokio::test] async fn test_build_search_uses_configured_posting_block_size() { let tmpdir = TempObjDir::default(); @@ -6114,6 +6348,16 @@ mod tests { }; assert_eq!(posting.block_size, block_size); assert_eq!(posting.blocks.len(), num_docs.div_ceil(block_size)); + let impacts = posting + .impacts + .as_ref() + .expect("newly written posting list should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); let tokens = Arc::new(Tokens::new(vec!["needle".to_owned()], DocType::Text)); let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); @@ -6865,7 +7109,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); @@ -7177,7 +7425,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(token_id).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); let slot = (token_id - start) as usize; @@ -7652,7 +7904,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); @@ -7741,6 +7997,7 @@ mod tests { PostingTailCodec::Fixed32, LEGACY_BLOCK_SIZE, None, + None, ); let full_backing = full.get_buffer_memory_size(); @@ -7887,7 +8144,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); assert!( @@ -8131,6 +8392,178 @@ mod tests { writer.finish_with_metadata(metadata).await.unwrap(); } + async fn write_test_partition_with_optional_impacts( + store: &Arc, + partition_id: u64, + mut builder: InnerBuilder, + token_set_format: TokenSetFormat, + with_impacts: bool, + ) { + let format_version = InvertedListFormatVersion::V1; + let block_size = LEGACY_BLOCK_SIZE; + let docs = std::mem::take(&mut builder.docs); + let schema = inverted_list_schema_for_version_with_block_size_and_impacts( + false, + format_version, + block_size, + with_impacts, + ); + + let mut posting_writer = store + .new_index_file(&posting_file_path(partition_id), schema.clone()) + .await + .unwrap(); + for posting_list in std::mem::take(&mut builder.posting_lists) { + let batch = posting_list + .to_batch_with_docs(&docs, schema.clone()) + .unwrap(); + posting_writer.write_record_batch(batch).await.unwrap(); + } + posting_writer.finish().await.unwrap(); + + let token_batch = std::mem::take(&mut builder.tokens) + .to_batch(token_set_format) + .unwrap(); + let mut token_writer = store + .new_index_file(&token_file_path(partition_id), token_batch.schema()) + .await + .unwrap(); + token_writer.write_record_batch(token_batch).await.unwrap(); + token_writer.finish().await.unwrap(); + + let doc_batch = docs.to_batch().unwrap(); + let mut doc_writer = store + .new_index_file(&doc_file_path(partition_id), doc_batch.schema()) + .await + .unwrap(); + doc_writer.write_record_batch(doc_batch).await.unwrap(); + doc_writer.finish().await.unwrap(); + } + + #[tokio::test] + async fn test_mixed_impact_and_legacy_partitions_use_global_final_scores() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let mut impact_builder = InnerBuilder::new_with_format_version( + 0, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + ); + impact_builder.tokens.add("alpha".to_owned()); + impact_builder + .posting_lists + .push(PostingListBuilder::new_with_posting_tail_codec( + false, + InvertedListFormatVersion::V1.posting_tail_codec(), + )); + impact_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + impact_builder.docs.append(100, 5_000); + for row_id in 101..111 { + impact_builder.docs.append(row_id, 5_000); + } + write_test_partition_with_optional_impacts( + &store, + 0, + impact_builder, + TokenSetFormat::default(), + true, + ) + .await; + + let mut legacy_builder = InnerBuilder::new_with_format_version( + 1, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + ); + legacy_builder.tokens.add("alpha".to_owned()); + legacy_builder + .posting_lists + .push(PostingListBuilder::new_with_posting_tail_codec( + false, + InvertedListFormatVersion::V1.posting_tail_codec(), + )); + legacy_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + legacy_builder.docs.append(200, 1_000); + for row_id in 201..301 { + legacy_builder.docs.append(row_id, 1); + } + write_test_partition_with_optional_impacts( + &store, + 1, + legacy_builder, + TokenSetFormat::default(), + false, + ) + .await; + + write_test_metadata(&store, vec![0, 1], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + let impact_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 0) + .unwrap(); + let legacy_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 1) + .unwrap(); + + let impact_posting = impact_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(impact_posting.has_impacts()); + + let legacy_posting = legacy_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(!legacy_posting.has_impacts()); + + let tokens = Arc::new(Tokens::new(vec!["alpha".to_string()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); + let (row_ids, scores) = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + + assert_eq!(row_ids, vec![200]); + assert_eq!(row_ids.len(), scores.len()); + + let scorer = index + .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .await + .unwrap(); + let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); + assert!( + (scores[0] - expected_score).abs() < 1e-6, + "score: {}, expected: {}", + scores[0], + expected_score + ); + } + #[tokio::test] async fn test_and_query_returns_empty_when_exact_term_missing() { let tmpdir = TempObjDir::default(); @@ -9316,7 +9749,11 @@ mod tests { assert!( posting_reader .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: posting_reader.has_impacts, + }) .await .is_some(), "prewarm did not populate group [{start}, {end}) that the read \ @@ -9463,7 +9900,10 @@ mod tests { assert!( posting_reader .index_cache - .get_with_key(&PostingListKey { token_id }) + .get_with_key(&PostingListKey { + token_id, + has_impacts: posting_reader.has_impacts, + }) .await .is_some(), "fallback prewarm should populate per-token entry {token_id}", diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index eb7d78ff397..b6b48ce4149 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -3,6 +3,7 @@ use super::InvertedPartition; use std::collections::HashMap; +use std::sync::Arc; // the Scorer trait is used to calculate the score of a token in a document // in general, the score is calculated as: @@ -12,6 +13,16 @@ pub trait Scorer: Send + Sync { fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32; } +impl Scorer for Arc { + fn query_weight(&self, token: &str) -> f32 { + self.as_ref().query_weight(token) + } + + fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { + self.as_ref().doc_weight(freq, doc_tokens) + } +} + // BM25 parameters pub const K1: f32 = 1.2; pub const B: f32 = 0.75; diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index b94ac61aa0a..8ed1c62eba6 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -22,6 +22,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, + impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, query::Operator, scorer::{K1, idf}, }; @@ -38,6 +39,7 @@ use super::{ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; +const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") .unwrap_or_else(|_| "10".to_string()) @@ -45,6 +47,33 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); +#[inline] +fn posting_block_idx(index: usize, block_size: usize) -> usize { + match block_size { + 128 => index >> 7, + 256 => index >> 8, + _ => index / block_size, + } +} + +#[inline] +fn posting_block_offset(index: usize, block_size: usize) -> usize { + match block_size { + 128 => index & 127, + 256 => index & 255, + _ => index % block_size, + } +} + +#[inline] +fn posting_block_start(block_idx: usize, block_size: usize) -> usize { + match block_size { + 128 => block_idx << 7, + 256 => block_idx << 8, + _ => block_idx * block_size, + } +} + pub struct PostingIterator { token: String, token_id: u32, @@ -55,6 +84,7 @@ pub struct PostingIterator { index: usize, // the index of current block, this can be changed by `next() and shallow_next()` block_idx: usize, + current_doc: Option, approximate_upper_bound: f32, // for compressed posting list @@ -71,6 +101,8 @@ struct CompressedState { position_values: Vec, position_offsets: Vec, block_max_window: BlockMaxWindow, + current_block_max_score: Option<(usize, f32)>, + block_least_doc_ids: Vec>, } impl CompressedState { @@ -84,6 +116,8 @@ impl CompressedState { position_values: Vec::new(), position_offsets: Vec::new(), block_max_window: BlockMaxWindow::new(), + current_block_max_score: None, + block_least_doc_ids: Vec::new(), } } @@ -100,7 +134,7 @@ impl CompressedState { self.doc_ids.clear(); self.freqs.clear(); - let remainder = length as usize % block_size; + let remainder = posting_block_offset(length as usize, block_size); if block_idx + 1 == num_blocks && remainder != 0 { decompress_posting_remainder( block, @@ -132,6 +166,7 @@ struct BlockMaxWindow { start_block_idx: usize, next_block_idx: usize, max_scores: VecDeque<(usize, f32)>, + impact_score_cache: ImpactScoreCache, } struct BlockMaxScore { @@ -145,6 +180,7 @@ impl BlockMaxWindow { start_block_idx: 0, next_block_idx: 0, max_scores: VecDeque::new(), + impact_score_cache: ImpactScoreCache::default(), } } @@ -154,12 +190,28 @@ impl BlockMaxWindow { self.max_scores.clear(); } - fn max_score_up_to( + fn max_score_up_to( &mut self, list: &CompressedPostingList, start_block_idx: usize, up_to: u64, + query_weight: f32, + scorer: &S, ) -> BlockMaxScore { + if let Some(impacts) = &list.impacts { + let score = impacts.max_score_up_to_cached( + start_block_idx, + up_to, + query_weight, + scorer, + &mut self.impact_score_cache, + ); + return BlockMaxScore { + score: score.score, + blocks_scanned: score.entries_scanned, + }; + } + if start_block_idx >= list.blocks.len() { self.reset(start_block_idx); return BlockMaxScore { @@ -263,6 +315,96 @@ impl Ord for PostingIterator { } impl PostingIterator { + #[inline] + fn block_least_doc_id(&self, list: &CompressedPostingList, block_idx: usize) -> u32 { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if compressed.block_least_doc_ids.len() < list.blocks.len() { + compressed + .block_least_doc_ids + .resize(list.blocks.len(), None); + } + if let Some(doc_id) = compressed.block_least_doc_ids[block_idx] { + return doc_id; + } + let doc_id = list.block_least_doc_id(block_idx); + compressed.block_least_doc_ids[block_idx] = Some(doc_id); + doc_id + } + + fn block_idx_for_doc( + &self, + list: &CompressedPostingList, + mut block_idx: usize, + least_id: u32, + ) -> usize { + let mut linear_skips = 0; + while block_idx + 1 < list.blocks.len() && linear_skips < LINEAR_BLOCK_SKIP_LIMIT { + if self.block_least_doc_id(list, block_idx + 1) > least_id { + return block_idx; + } + block_idx += 1; + linear_skips += 1; + } + + if block_idx + 1 >= list.blocks.len() { + return block_idx; + } + + if let Some(impacts) = list.impacts.as_ref() + && let Some(block_idx) = + self.block_idx_for_doc_with_impacts(list, impacts, block_idx, least_id) + { + return block_idx; + } + + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, list.blocks.len()) + } + + fn block_idx_for_doc_with_impacts( + &self, + list: &CompressedPostingList, + impacts: &ImpactSkipData, + mut block_idx: usize, + least_id: u32, + ) -> Option { + while block_idx + 1 < list.blocks.len() { + let group_idx = (block_idx + 1) / IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(list.blocks.len()); + let group_doc_up_to = impacts.level1_doc_up_to(group_idx)?; + if group_doc_up_to < least_id { + block_idx = group_end - 1; + continue; + } + if group_doc_up_to == least_id { + return Some(group_end - 1); + } + return Some( + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, group_end), + ); + } + Some(block_idx) + } + + fn block_idx_for_doc_by_least_doc_id( + &self, + list: &CompressedPostingList, + block_idx: usize, + least_id: u32, + right: usize, + ) -> usize { + let mut left = block_idx + 1; + let mut right = right; + while left < right { + let mid = left + (right - left) / 2; + if self.block_least_doc_id(list, mid) <= least_id { + left = mid + 1; + } else { + right = mid; + } + } + left - 1 + } + #[inline] fn compressed_state_ptr(&self) -> *mut CompressedState { debug_assert!(self.compressed.is_some()); @@ -311,9 +453,12 @@ impl PostingIterator { list: PostingList, num_doc: usize, ) -> Self { - let approximate_upper_bound = match list.max_score() { - Some(max_score) => max_score, - None => idf(list.len(), num_doc) * (K1 + 1.0), + let approximate_upper_bound = match &list { + PostingList::Compressed(posting) if posting.impacts.is_some() => f32::INFINITY, + _ => match list.max_score() { + Some(max_score) => max_score, + None => idf(list.len(), num_doc) * (K1 + 1.0), + }, }; let compressed = match &list { PostingList::Compressed(list) => { @@ -322,7 +467,7 @@ impl PostingIterator { PostingList::Plain(_) => None, }; - Self { + let mut posting = Self { token, token_id, position, @@ -330,9 +475,12 @@ impl PostingIterator { list, index: 0, block_idx: 0, + current_doc: None, approximate_upper_bound, compressed, - } + }; + posting.refresh_current_doc(); + posting } #[inline] @@ -350,8 +498,24 @@ impl PostingIterator { self.approximate_upper_bound } + /// Tightest known list-wide score bound. Impact lists answer from the + /// baked doc-weight slab (the data-driven equivalent of the max_score the + /// non-impact format bakes at build time); everything else falls back to + /// `approximate_upper_bound`. A finite, tight global bound is what lets + /// lagging iterators park in the WAND tail instead of being force-advanced + /// on every candidate. + #[inline] + fn global_upper_bound(&self, scorer: &S) -> f32 { + if let PostingList::Compressed(ref list) = self.list + && let Some(impacts) = list.impacts.as_ref() + { + return self.query_weight * impacts.global_max_doc_weight(scorer); + } + self.approximate_upper_bound + } + #[inline] - fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { + fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { self.query_weight * scorer.doc_weight(freq, doc_length) } @@ -367,14 +531,19 @@ impl PostingIterator { #[inline] fn doc(&self) -> Option { + self.current_doc + } + + fn refresh_current_doc(&mut self) { if self.empty() { - return None; + self.current_doc = None; + return; } - match self.list { + let current_doc = match self.list { PostingList::Compressed(ref list) => { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; // Read from the decompressed block @@ -384,7 +553,8 @@ impl PostingIterator { Some(doc) } PostingList::Plain(ref list) => Some(DocInfo::Located(list.doc(self.index))), - } + }; + self.current_doc = current_doc; } fn position_cursor(&self) -> Option> { @@ -412,8 +582,8 @@ impl PostingIterator { )) } CompressedPositionStorage::SharedStream(stream) => { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; if compressed.position_block_idx != Some(block_idx) { @@ -453,36 +623,46 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - let mut block_idx = self.index / list.block_size; - while block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(block_idx + 1) <= least_id - { - block_idx += 1; - } - self.index = self.index.max(block_idx * list.block_size); + let block_idx = self.block_idx_for_doc( + list, + posting_block_idx(self.index, list.block_size), + least_id, + ); + self.index = self + .index + .max(posting_block_start(block_idx, list.block_size)); let length = list.length as usize; while self.index < length { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; let in_block = &compressed.doc_ids[block_offset..]; let offset_in_block = in_block.partition_point(|&doc_id| doc_id < least_id); let new_offset = block_offset + offset_in_block; if new_offset < compressed.doc_ids.len() { - self.index = block_idx * list.block_size + new_offset; - break; + self.index = posting_block_start(block_idx, list.block_size) + new_offset; + self.block_idx = block_idx; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[new_offset], + frequency: compressed.freqs[new_offset], + })); + return; } if block_idx + 1 >= list.blocks.len() { self.index = length; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; break; } - self.index = (block_idx + 1) * list.block_size; + self.index = posting_block_start(block_idx + 1, list.block_size); } - self.block_idx = self.index / list.block_size; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; } PostingList::Plain(ref list) => { self.index += list.row_ids[self.index..].partition_point(|&id| id < least_id); + self.current_doc = (!self.empty()).then(|| DocInfo::Located(list.doc(self.index))); } } } @@ -492,11 +672,7 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - while self.block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(self.block_idx + 1) <= least_id - { - self.block_idx += 1; - } + self.block_idx = self.block_idx_for_doc(list, self.block_idx, least_id); } PostingList::Plain(_) => { // we don't have block max score for legacy index, @@ -506,21 +682,48 @@ impl PostingIterator { } #[inline] - fn block_max_score(&self) -> f32 { + fn block_max_score(&self, scorer: &S) -> f32 { match self.list { - PostingList::Compressed(ref list) => list.block_max_score(self.block_idx), + PostingList::Compressed(ref list) => { + if let Some(impacts) = list.impacts.as_ref() { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((block_idx, score)) = compressed.current_block_max_score + && block_idx == self.block_idx + { + return score; + } + + let score = impacts.level0_score_cached( + self.block_idx, + self.query_weight, + scorer, + &mut compressed.block_max_window.impact_score_cache, + ); + compressed.current_block_max_score = Some((self.block_idx, score)); + return score; + } + list.block_max_score(self.block_idx) + } PostingList::Plain(_) => self.approximate_upper_bound, } } #[inline] - fn block_max_score_up_to_with_stats(&mut self, up_to: u64) -> BlockMaxScore { + fn block_max_score_up_to_with_stats( + &mut self, + up_to: u64, + scorer: &S, + ) -> BlockMaxScore { match self.list { PostingList::Compressed(ref list) => { let compressed = unsafe { &mut *self.compressed_state_ptr() }; - compressed - .block_max_window - .max_score_up_to(list, self.block_idx, up_to) + compressed.block_max_window.max_score_up_to( + list, + self.block_idx, + up_to, + self.query_weight, + scorer, + ) } PostingList::Plain(_) => BlockMaxScore { score: self.approximate_upper_bound, @@ -545,7 +748,7 @@ impl PostingIterator { fn block_first_doc(&self) -> Option { match self.list { PostingList::Compressed(ref list) => { - Some(list.block_least_doc_id(self.block_idx) as u64) + Some(self.block_least_doc_id(list, self.block_idx) as u64) } PostingList::Plain(ref plain) => plain.row_ids.get(self.index).cloned(), } @@ -558,7 +761,7 @@ impl PostingIterator { if self.block_idx + 1 >= list.blocks.len() { return None; } - Some(list.block_least_doc_id(self.block_idx + 1) as u64) + Some(self.block_least_doc_id(list, self.block_idx + 1) as u64) } PostingList::Plain(ref plain) => plain.row_ids.get(self.index + 1).cloned(), } @@ -1186,7 +1389,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let remaining_upper_bound = remaining .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); first.score(&self.scorer, doc.frequency(), doc_length) + remaining_upper_bound <= self.threshold @@ -1373,7 +1576,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let narrow_max_score = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); if narrow_max_score >= self.threshold { @@ -1396,7 +1599,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut wide_max_score = 0.0; let mut range_blocks_scanned = 0; for posting in &mut self.lead { - let block_max = posting.block_max_score_up_to_with_stats(lead_up_to); + let block_max = posting.block_max_score_up_to_with_stats(lead_up_to, &self.scorer); wide_max_score += block_max.score; range_blocks_scanned += block_max.blocks_scanned; } @@ -1460,7 +1663,7 @@ impl<'a, S: Scorer> Wand<'a, S> { fn move_head_before_target_to_tail(&mut self, target: u64) { while matches!(self.head_doc(), Some(doc_id) if doc_id < target) { if let Some(posting) = self.head.pop() { - let upper_bound = posting.posting.approximate_upper_bound(); + let upper_bound = posting.posting.global_upper_bound(&self.scorer); if let Some(mut evicted) = self.insert_tail_with_overflow(posting.posting, upper_bound) { @@ -1479,12 +1682,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead: f32 = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum(); let head: f32 = self .head .iter() - .map(|posting| posting.posting.block_max_score()) + .map(|posting| posting.posting.block_max_score(&self.scorer)) .sum(); lead + head + self.tail_max_score } @@ -1497,12 +1700,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut sum = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); let mut possible_matches = self.lead.len(); for posting in &self.tail { if matches!(posting.posting.block_first_doc(), Some(block_doc) if block_doc <= target) { - sum += posting.posting.block_max_score(); + sum += posting.posting.block_max_score(&self.scorer); possible_matches += 1; } } @@ -1568,7 +1771,7 @@ impl<'a, S: Scorer> Wand<'a, S> { Some(block_doc) if block_doc <= target => { tail_posting .posting - .block_max_score_up_to_with_stats(up_to) + .block_max_score_up_to_with_stats(up_to, &self.scorer) .score } _ => 0.0, @@ -1705,9 +1908,9 @@ impl<'a, S: Scorer> Wand<'a, S> { && posting.is_compressed() && self.up_to.is_some_and(|up_to| target <= up_to) { - posting.block_max_score() + posting.block_max_score(&self.scorer) } else { - posting.approximate_upper_bound() + posting.global_upper_bound(&self.scorer) } } @@ -1989,13 +2192,17 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use super::super::impact::build_impact_skip_data; use super::*; use crate::scalar::inverted::scorer::IndexBM25Scorer; use crate::{ metrics::{MetricsCollector, NoOpMetricsCollector}, scalar::inverted::{ - CompressedPostingList, PlainPostingList, PostingListBuilder, builder::PositionRecorder, - encoding::compress_posting_list, + CompressedPostingList, PlainPostingList, PostingListBuilder, + builder::PositionRecorder, + encoding::{ + compress_posting_list, compress_posting_list_with_tail_codec_and_block_size, + }, }, }; @@ -2160,6 +2367,7 @@ mod tests { crate::scalar::inverted::PostingTailCodec::VarintDelta, crate::scalar::inverted::LEGACY_BLOCK_SIZE, None, + None, )) } else { PostingList::Plain(PlainPostingList::new( @@ -2171,6 +2379,75 @@ mod tests { } } + fn generate_impact_posting_list_with_freqs( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + doc_ids, + freqs, + doc_lengths, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + ) + } + + fn generate_impact_posting_list_with_freqs_and_block_size( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + block_size: usize, + ) -> PostingList { + assert_eq!(doc_ids.len(), freqs.len()); + assert_eq!(doc_ids.len(), doc_lengths.len()); + let block_max_scores = vec![0.0; doc_ids.len().div_ceil(block_size)]; + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + freqs.iter(), + block_max_scores.into_iter(), + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + ) + .unwrap(); + let impact_blocks = doc_ids + .chunks(block_size) + .zip(freqs.chunks(block_size)) + .zip(doc_lengths.chunks(block_size)) + .map(|((doc_ids, freqs), doc_lengths)| { + doc_ids + .iter() + .copied() + .zip(freqs.iter().copied()) + .zip(doc_lengths.iter().copied()) + .map(|((doc_id, freq), doc_length)| (doc_id, freq, doc_length)) + .collect::>() + }) + .collect::>(); + let impacts = build_impact_skip_data(impact_blocks.as_slice()).unwrap(); + PostingList::Compressed(CompressedPostingList::new( + blocks, + 0.0, + doc_ids.len() as u32, + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + None, + Some(impacts), + )) + } + + fn generate_contiguous_impact_posting_list_with_block_size( + total: usize, + block_size: usize, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + (0..total as u32).collect(), + vec![1; total], + vec![1; total], + block_size, + ) + } + fn generate_posting_list_with_positions( doc_ids: Vec, positions_by_doc: Vec>, @@ -2995,7 +3272,7 @@ mod tests { posting.shallow_next(0); assert_eq!( posting - .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 4.0 ); @@ -3003,7 +3280,7 @@ mod tests { posting.shallow_next((2 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 5.0 ); @@ -3011,12 +3288,127 @@ mod tests { posting.shallow_next((4 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 3.0 ); } + #[test] + fn test_impact_level1_skip_keeps_boundary_equality_in_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 1) * block_size; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (IMPACT_LEVEL1_BLOCKS * block_size - 1) as u64; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + + posting.next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_handles_partial_final_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 3) * block_size + 17; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (total - 1) as u64; + let expected_block = total.div_ceil(block_size) - 1; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, expected_block); + + posting.next(target); + assert_eq!(posting.block_idx, expected_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_reaches_far_target_doc() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS * 3 + 5) * block_size; + let target_block = IMPACT_LEVEL1_BLOCKS * 2 + 2; + let target = (target_block * block_size + 17) as u64; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + + posting.shallow_next(target); + assert_eq!(posting.block_idx, target_block); + + posting.next(target); + assert_eq!(posting.block_idx, target_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_compressed_impact_block_max_score_memoizes_current_block() { + let total = 2 * BLOCK_SIZE as u32; + let doc_ids = (0..total).collect::>(); + let freqs = doc_ids + .iter() + .map(|doc_id| if *doc_id < BLOCK_SIZE as u32 { 1 } else { 2 }) + .collect::>(); + let doc_lengths = vec![1; total as usize]; + let posting_list = generate_impact_posting_list_with_freqs(doc_ids, freqs, doc_lengths); + let mut posting = + PostingIterator::new(String::from("term"), 0, 0, posting_list, total as usize); + let scored = Arc::new(AtomicUsize::new(0)); + let scorer = CountingScorer { + scored: scored.clone(), + }; + + let first_score = posting.block_max_score(&scorer); + assert_eq!(first_score, 1.0); + // Baking the shared doc-weight bounds visits every frontier pair once + // (two level0 entries plus one level1 entry for this list). + let baked = scored.load(Ordering::Relaxed); + assert!(baked >= 2); + { + let compressed = unsafe { &mut *posting.compressed_state_ptr() }; + assert_eq!(compressed.current_block_max_score, Some((0, first_score))); + compressed.block_max_window.impact_score_cache = ImpactScoreCache::default(); + } + + let second_score = posting.block_max_score(&scorer); + assert_eq!(second_score, first_score); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "repeated block max scores must not recompute doc weights" + ); + + posting.shallow_next(BLOCK_SIZE as u64); + let next_block_score = posting.block_max_score(&scorer); + assert_eq!(next_block_score, 2.0); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "other blocks answer from the baked bounds without rescoring" + ); + } + #[test] fn test_and_candidate_prune_scores_first_term_before_full_score() { let total_docs = 2 * BLOCK_SIZE as u32 + 1; @@ -3353,7 +3745,7 @@ mod tests { let posting = PostingIterator::new(String::from("test"), 0, 0, posting_list, 1); - let actual = posting.block_max_score(); + let actual = posting.block_max_score(&UnitScorer); assert!( (actual - expected).abs() < 1e-6, "block max score should match stored value" From 0781681173ed5864af101878f0da9474bc2cfb35 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:13:21 +0800 Subject: [PATCH 19/20] perf(fts): bulk MAXSCORE search path for top-k disjunctions Port of Lucene's MaxScoreBulkScorer, opt-in via LANCE_FTS_MAXSCORE=1: per outer window (bounded by the essential clauses' blocks with adaptive growth), clauses split into a non-essential prefix and essential rest by window max score vs the running threshold. Essential clauses bulk-stream decompressed blocks (single-essential windows stream with no accumulator); non-essential clauses are only probed for candidates that can still beat the threshold. Dead ranges with one live clause skip by scanning the baked per-block bound slab. Candidate emission matches the classic path (must beat the running threshold), so results are score-identical. Measured on a 200M-doc warm benchmark at 24 partitions: 3-word OR match k10 0.137s -> 0.035s, k100 0.250s -> 0.064s; hot single-term 250ms -> 3ms. --- .../lance-index/src/scalar/inverted/impact.rs | 32 + rust/lance-index/src/scalar/inverted/wand.rs | 1030 +++++++++++++++-- 2 files changed, 955 insertions(+), 107 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/impact.rs b/rust/lance-index/src/scalar/inverted/impact.rs index 8b16a4f57f8..687142b9941 100644 --- a/rust/lance-index/src/scalar/inverted/impact.rs +++ b/rust/lance-index/src/scalar/inverted/impact.rs @@ -147,6 +147,12 @@ impl ImpactSkipData { &self.baked_bounds(scorer).0 } + /// Baked per-block max doc weights (level0 entries only), for bulk skip + /// scans over dead ranges without per-block window bookkeeping. + pub(crate) fn level0_doc_weight_bounds(&self, scorer: &S) -> &[f32] { + &self.baked_bounds(scorer).0[..self.level0_len] + } + /// List-wide max doc weight, from the baked bounds. The tightest valid /// global score bound for this list is `query_weight * this`, matching /// what the non-impact format stores as `max_score` at build time — but @@ -179,6 +185,18 @@ impl ImpactSkipData { } } + /// Last doc id covered by the level0 entry of `block_idx`, or `None` when + /// the entry is missing or malformed. + pub(crate) fn level0_doc_up_to(&self, block_idx: usize) -> Option { + if block_idx >= self.level0_len { + return None; + } + match self.entry_doc_up_tos[block_idx] { + u32::MAX => None, + doc_up_to => Some(doc_up_to), + } + } + /// Max score of the docs covered by the level0 entry of `block_idx`, /// answered from the baked bounds slab. pub fn level0_score( @@ -193,6 +211,20 @@ impl ImpactSkipData { query_weight * self.doc_weight_bounds(scorer)[block_idx] } + /// Max score of the docs covered by the level1 entry of `group_idx`, + /// answered from the baked bounds slab. + pub fn level1_score( + &self, + group_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if group_idx >= level1_len(self.level0_len) || query_weight <= 0.0 { + return 0.0; + } + query_weight * self.doc_weight_bounds(scorer)[self.level0_len + group_idx] + } + pub fn level0_score_cached( &self, block_idx: usize, diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 8ed1c62eba6..6df1a2122e6 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -22,7 +22,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, - impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, + impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData}, query::Operator, scorer::{K1, idf}, }; @@ -39,6 +39,9 @@ use super::{ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; + +/// Top-k heap entry: (scored doc, (term, freq) pairs, doc length, posting doc id). +type TopKHeap = BinaryHeap, u32, u64)>>; const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") @@ -46,6 +49,11 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .parse::() .unwrap_or(10) }); +// Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer +// style). Off by default while it bakes; LANCE_FTS_MAXSCORE=1 opts in. Its +// results are score-identical to the classic WAND loop. +static USE_MAXSCORE_SEARCH: LazyLock = + LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() == Ok("1")); #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize { @@ -101,8 +109,11 @@ struct CompressedState { position_values: Vec, position_offsets: Vec, block_max_window: BlockMaxWindow, - current_block_max_score: Option<(usize, f32)>, - block_least_doc_ids: Vec>, + // Lucene-style anchored impact score caches: one slot per level, keyed by + // the entry the block cursor currently sits in. Each holds + // (entry_idx, doc_up_to, max_score). See `impact_level0`/`impact_level1`. + level0_cache: Option<(usize, u32, f32)>, + level1_cache: Option<(usize, u32, f32)>, } impl CompressedState { @@ -116,8 +127,8 @@ impl CompressedState { position_values: Vec::new(), position_offsets: Vec::new(), block_max_window: BlockMaxWindow::new(), - current_block_max_score: None, - block_least_doc_ids: Vec::new(), + level0_cache: None, + level1_cache: None, } } @@ -163,10 +174,11 @@ impl CompressedState { struct BlockMaxWindow { // Sliding block range used for Lucene-style getMaxScore(upTo). The deque is // monotonic by score and covers blocks in [start_block_idx, next_block_idx). + // Only used for compressed lists without impact skip data; impact lists + // answer window max scores from the anchored level caches instead. start_block_idx: usize, next_block_idx: usize, max_scores: VecDeque<(usize, f32)>, - impact_score_cache: ImpactScoreCache, } struct BlockMaxScore { @@ -180,7 +192,6 @@ impl BlockMaxWindow { start_block_idx: 0, next_block_idx: 0, max_scores: VecDeque::new(), - impact_score_cache: ImpactScoreCache::default(), } } @@ -198,20 +209,6 @@ impl BlockMaxWindow { query_weight: f32, scorer: &S, ) -> BlockMaxScore { - if let Some(impacts) = &list.impacts { - let score = impacts.max_score_up_to_cached( - start_block_idx, - up_to, - query_weight, - scorer, - &mut self.impact_score_cache, - ); - return BlockMaxScore { - score: score.score, - blocks_scanned: score.entries_scanned, - }; - } - if start_block_idx >= list.blocks.len() { self.reset(start_block_idx); return BlockMaxScore { @@ -241,7 +238,10 @@ impl BlockMaxWindow { while self.next_block_idx < list.blocks.len() && list.block_least_doc_id(self.next_block_idx) as u64 <= up_to { - let score = list.block_max_score(self.next_block_idx); + let score = match list.impacts.as_ref() { + Some(impacts) => impacts.level0_score(self.next_block_idx, query_weight, scorer), + None => list.block_max_score(self.next_block_idx), + }; while matches!(self.max_scores.back(), Some((_, old_score)) if *old_score <= score) { self.max_scores.pop_back(); } @@ -317,18 +317,7 @@ impl Ord for PostingIterator { impl PostingIterator { #[inline] fn block_least_doc_id(&self, list: &CompressedPostingList, block_idx: usize) -> u32 { - let compressed = unsafe { &mut *self.compressed_state_ptr() }; - if compressed.block_least_doc_ids.len() < list.blocks.len() { - compressed - .block_least_doc_ids - .resize(list.blocks.len(), None); - } - if let Some(doc_id) = compressed.block_least_doc_ids[block_idx] { - return doc_id; - } - let doc_id = list.block_least_doc_id(block_idx); - compressed.block_least_doc_ids[block_idx] = Some(doc_id); - doc_id + list.block_least_doc_id(block_idx) } fn block_idx_for_doc( @@ -405,6 +394,30 @@ impl PostingIterator { left - 1 } + #[inline] + fn block_end_doc(&self) -> u64 { + self.next_block_first_doc() + .map(|doc| doc.saturating_sub(1)) + .unwrap_or(TERMINATED_DOC_ID) + } + + /// Level1 bound of the group holding the current block, for group-wide + /// skipping: (group doc_up_to, group max score). `None` when the list has + /// no impact skip data or the group entry is missing/malformed. + fn impact_group_bound(&self, scorer: &S) -> Option<(u64, f32)> { + match self.list { + PostingList::Compressed(ref list) => { + let impacts = list.impacts.as_ref()?; + let (doc_up_to, score) = self.impact_level1(impacts, scorer); + if doc_up_to == u32::MAX { + return None; + } + Some((u64::from(doc_up_to), score)) + } + PostingList::Plain(_) => None, + } + } + #[inline] fn compressed_state_ptr(&self) -> *mut CompressedState { debug_assert!(self.compressed.is_some()); @@ -453,8 +466,15 @@ impl PostingIterator { list: PostingList, num_doc: usize, ) -> Self { + // BM25's doc weight is bounded by K1 + 1 for any freq and doc length, + // so query_weight * (K1 + 1) is a valid global bound even when index + // stats drift after appends. Keeping it finite matters: an INFINITY + // bound can never park the iterator in the WAND tail, forcing a deep + // advance on every candidate. let approximate_upper_bound = match &list { - PostingList::Compressed(posting) if posting.impacts.is_some() => f32::INFINITY, + PostingList::Compressed(posting) if posting.impacts.is_some() => { + query_weight * (K1 + 1.0) + } _ => match list.max_score() { Some(max_score) => max_score, None => idf(list.len(), num_doc) * (K1 + 1.0), @@ -501,9 +521,8 @@ impl PostingIterator { /// Tightest known list-wide score bound. Impact lists answer from the /// baked doc-weight slab (the data-driven equivalent of the max_score the /// non-impact format bakes at build time); everything else falls back to - /// `approximate_upper_bound`. A finite, tight global bound is what lets - /// lagging iterators park in the WAND tail instead of being force-advanced - /// on every candidate. + /// `approximate_upper_bound`. A finite, tight global bound lets lagging + /// iterators park in the WAND tail instead of being force-advanced. #[inline] fn global_upper_bound(&self, scorer: &S) -> f32 { if let PostingList::Compressed(ref list) = self.list @@ -681,26 +700,55 @@ impl PostingIterator { } } + /// Anchored level0 impact bound of the current block: (doc_up_to, max_score), + /// memoized until the block cursor moves. Malformed entries degrade to + /// (u32::MAX, INFINITY), which keeps pruning safe by making the block look + /// unskippable. + #[inline] + fn impact_level0( + &self, + impacts: &ImpactSkipData, + scorer: &S, + ) -> (u32, f32) { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((block_idx, doc_up_to, score)) = compressed.level0_cache + && block_idx == self.block_idx + { + return (doc_up_to, score); + } + let doc_up_to = impacts.level0_doc_up_to(self.block_idx).unwrap_or(u32::MAX); + let score = impacts.level0_score(self.block_idx, self.query_weight, scorer); + compressed.level0_cache = Some((self.block_idx, doc_up_to, score)); + (doc_up_to, score) + } + + /// Anchored level1 impact bound of the group holding the current block, + /// memoized until the cursor crosses a group boundary. + #[inline] + fn impact_level1( + &self, + impacts: &ImpactSkipData, + scorer: &S, + ) -> (u32, f32) { + let group_idx = self.block_idx / IMPACT_LEVEL1_BLOCKS; + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((cached_group_idx, doc_up_to, score)) = compressed.level1_cache + && cached_group_idx == group_idx + { + return (doc_up_to, score); + } + let doc_up_to = impacts.level1_doc_up_to(group_idx).unwrap_or(u32::MAX); + let score = impacts.level1_score(group_idx, self.query_weight, scorer); + compressed.level1_cache = Some((group_idx, doc_up_to, score)); + (doc_up_to, score) + } + #[inline] fn block_max_score(&self, scorer: &S) -> f32 { match self.list { PostingList::Compressed(ref list) => { if let Some(impacts) = list.impacts.as_ref() { - let compressed = unsafe { &mut *self.compressed_state_ptr() }; - if let Some((block_idx, score)) = compressed.current_block_max_score - && block_idx == self.block_idx - { - return score; - } - - let score = impacts.level0_score_cached( - self.block_idx, - self.query_weight, - scorer, - &mut compressed.block_max_window.impact_score_cache, - ); - compressed.current_block_max_score = Some((self.block_idx, score)); - return score; + return self.impact_level0(impacts, scorer).1; } list.block_max_score(self.block_idx) } @@ -708,14 +756,27 @@ impl PostingIterator { } } + /// Tight max-score bound over `[current block, up_to]`. The common case — + /// a window ending inside the current block — answers from the anchored + /// level0 memo; wider windows fall back to the sliding block-max deque, + /// which scores each block once as it slides forward. #[inline] fn block_max_score_up_to_with_stats( - &mut self, + &self, up_to: u64, scorer: &S, ) -> BlockMaxScore { match self.list { PostingList::Compressed(ref list) => { + if let Some(impacts) = list.impacts.as_ref() { + let (level0_up_to, level0_score) = self.impact_level0(impacts, scorer); + if up_to <= u64::from(level0_up_to) { + return BlockMaxScore { + score: level0_score, + blocks_scanned: 0, + }; + } + } let compressed = unsafe { &mut *self.compressed_state_ptr() }; compressed.block_max_window.max_score_up_to( list, @@ -732,6 +793,16 @@ impl PostingIterator { } } + fn window_max_score(&self, up_to: Option, scorer: &S) -> f32 { + if let Some(up_to) = up_to + && let PostingList::Compressed(ref list) = self.list + && list.impacts.is_some() + { + return self.block_max_score_up_to_with_stats(up_to, scorer).score; + } + self.block_max_score(scorer) + } + #[inline] fn is_compressed(&self) -> bool { matches!(self.list, PostingList::Compressed(_)) @@ -754,6 +825,87 @@ impl PostingIterator { } } + /// Bulk-score every posting in `[current doc, up_to]` into the window + /// accumulator (slot = doc - window_min) and leave the iterator on the + /// first doc beyond `up_to`. This is the Lucene `nextDocsAndScores` + /// equivalent: it walks the decompressed block arrays directly, with no + /// per-doc heap traffic. + fn collect_window_scores( + &mut self, + window_min: u64, + up_to: u64, + clause_idx: usize, + docs: &DocSet, + scorer: &S, + acc: &mut WindowAccumulator, + ) { + if self.doc().is_some_and(|doc| doc.doc_id() < window_min) { + self.next(window_min); + } + match self.list { + PostingList::Compressed(ref list) => { + 'blocks: while let Some(doc) = self.current_doc { + if doc.doc_id() > up_to { + break; + } + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); + let compressed = + unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; + for offset in block_offset..compressed.doc_ids.len() { + let doc_id = compressed.doc_ids[offset]; + if u64::from(doc_id) > up_to { + self.index = posting_block_start(block_idx, list.block_size) + offset; + self.block_idx = block_idx; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id, + frequency: compressed.freqs[offset], + })); + break 'blocks; + } + let freq = compressed.freqs[offset]; + let doc_length = docs.num_tokens(doc_id); + let score = self.query_weight * scorer.doc_weight(freq, doc_length); + let slot = (u64::from(doc_id) - window_min) as usize; + acc.add(clause_idx, slot, score, freq); + } + // Block exhausted: step into the next block (or finish). + let next_start = posting_block_start(block_idx + 1, list.block_size); + if next_start >= list.length as usize { + self.index = list.length as usize; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; + break; + } + self.index = next_start; + self.block_idx = block_idx + 1; + let compressed = + unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx + 1) }; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[0], + frequency: compressed.freqs[0], + })); + } + } + PostingList::Plain(_) => { + while let Some(doc) = self.doc() { + let doc_id = doc.doc_id(); + if doc_id > up_to { + break; + } + let doc_length = match &doc { + DocInfo::Raw(raw) => docs.num_tokens(raw.doc_id), + DocInfo::Located(located) => docs.num_tokens_by_row_id(located.row_id), + }; + let score = self.score(scorer, doc.frequency(), doc_length); + let slot = (doc_id - window_min) as usize; + acc.add(clause_idx, slot, score, doc.frequency()); + self.next(doc_id + 1); + } + } + } + } + #[inline] fn next_block_first_doc(&self) -> Option { match self.list { @@ -768,6 +920,50 @@ impl PostingIterator { } } +/// Inner window span (in doc ids) of the bulk MAXSCORE path. Same as Lucene's +/// `MaxScoreBulkScorer.INNER_WINDOW_SIZE`. +const MAXSCORE_INNER_WINDOW: usize = 1 << 12; + +/// Per-window score/frequency accumulator for the bulk MAXSCORE path. Slot i +/// covers doc id `window_min + i`; `freqs` is laid out clause-major so accept +/// paths can recover (term, freq) pairs of the essential clauses. +struct WindowAccumulator { + scores: Vec, + freqs: Vec, + words: Vec, + num_clauses: usize, +} + +impl WindowAccumulator { + fn new(num_clauses: usize) -> Self { + Self { + scores: vec![0.0; MAXSCORE_INNER_WINDOW], + freqs: vec![0; num_clauses * MAXSCORE_INNER_WINDOW], + words: vec![0; MAXSCORE_INNER_WINDOW / 64], + num_clauses, + } + } + + #[inline] + fn add(&mut self, clause_idx: usize, slot: usize, score: f32, freq: u32) { + self.scores[slot] += score; + // Doc-major layout: one slot's clause frequencies share a cache line. + self.freqs[slot * self.num_clauses + clause_idx] = freq; + self.words[slot >> 6] |= 1u64 << (slot & 63); + } + + #[inline] + fn clause_freq(&self, clause_idx: usize, slot: usize) -> u32 { + self.freqs[slot * self.num_clauses + clause_idx] + } + + #[inline] + fn clear_slot(&mut self, slot: usize) { + self.scores[slot] = 0.0; + self.freqs[slot * self.num_clauses..(slot + 1) * self.num_clauses].fill(0); + } +} + /// How wand identified a candidate: either it already had the real /// row_id (DocSet carried row_ids), or only the partition-local /// doc_id (deferred-row_id path; the caller must resolve via @@ -1065,6 +1261,22 @@ impl<'a, S: Scorer> Wand<'a, S> { _ => {} } + // Top-k disjunctions over compressed lists can opt into the bulk + // MAXSCORE path (Lucene MaxScoreBulkScorer style): it streams whole + // blocks of the essential clauses into a window accumulator instead of + // advancing doc-at-a-time through a heap. + if *USE_MAXSCORE_SEARCH + && self.operator == Operator::Or + && params.phrase_slop.is_none() + && !self.head.is_empty() + && self + .head + .iter() + .all(|posting| posting.posting.is_compressed()) + { + return self.maxscore_search(params, mask, metrics); + } + // Deferred-row_id path: when the DocSet was built without // row_ids, wand emits candidates carrying just the // partition-local doc_id; the outer caller resolves them to @@ -1351,6 +1563,484 @@ impl<'a, S: Scorer> Wand<'a, S> { .collect()) } + /// Bulk MAXSCORE top-k disjunction, mirroring Lucene's MaxScoreBulkScorer. + /// + /// Per outer window (bounded by the essential clauses' block boundaries): + /// clauses are partitioned into a non-essential prefix — sorted by window + /// max score, as many as fit under the threshold — and the essential rest. + /// Essential clauses stream their postings into a window accumulator in + /// bulk; only accumulated candidates that could still beat the threshold + /// probe the non-essential clauses. No per-doc heap maintenance happens + /// anywhere on this path. + fn maxscore_search( + &mut self, + params: &FtsSearchParams, + mask: Arc, + metrics: &dyn MetricsCollector, + ) -> Result> { + struct MaxScoreClause { + posting: Box, + bound: f32, + prefix_bound: f32, + } + + #[inline] + fn float_sum_upper(sum: f32, num_terms: usize) -> f32 { + // Mirror of Lucene MathUtil.sumUpperBound: widen the float sum so + // it stays a true upper bound of the exact sum. + sum + sum.abs() * (num_terms as f32) * f32::EPSILON + } + + let limit = params.limit.unwrap_or(usize::MAX); + let docs_has_row_ids = self.docs.has_row_ids(); + let mut clauses = std::mem::take(&mut self.head) + .into_vec() + .into_iter() + .map(|head| MaxScoreClause { + posting: head.posting, + bound: 0.0, + prefix_bound: 0.0, + }) + .collect::>(); + + let mut acc = WindowAccumulator::new(clauses.len()); + let mut candidates: TopKHeap = + BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); + let mut num_comparisons = 0usize; + // Adaptive minimum window size (Lucene): grow windows when they yield + // too few candidates to amortize the per-window bound computations. + let mut min_window_size = 1u64; + let mut num_windows = 0u64; + let mut prev_first_essential = 0usize; + + let mut window_min = clauses + .iter() + .filter_map(|clause| clause.posting.doc().map(|doc| doc.doc_id())) + .min() + .unwrap_or(TERMINATED_DOC_ID); + + while window_min < TERMINATED_DOC_ID { + clauses.retain(|clause| clause.posting.doc().is_some()); + if clauses.is_empty() { + break; + } + self.raise_to_shared_floor(params.wand_factor); + + // Window boundary from the previous window's essential clauses + // only: dense non-essential clauses must not fragment the window. + let first_window_lead = prev_first_essential.min(clauses.len() - 1); + let mut window_max = TERMINATED_DOC_ID; + for clause in &mut clauses { + let doc = clause + .posting + .doc() + .map(|doc| doc.doc_id()) + .expect("exhausted clauses were retained out"); + clause.posting.shallow_next(doc.max(window_min)); + } + for clause in &clauses[first_window_lead..] { + window_max = window_max.min(clause.posting.block_end_doc()); + } + if clauses.len() > 1 { + // Target at least 32 candidates per clause per window on + // average before shrinking windows back to block granularity. + if (num_comparisons as u64) < num_windows * 32 * clauses.len() as u64 { + min_window_size = (min_window_size * 2).min(MAXSCORE_INNER_WINDOW as u64); + } else { + min_window_size = 1; + } + window_max = window_max.max(window_min.saturating_add(min_window_size - 1)); + } + + for clause in &mut clauses { + let doc = clause + .posting + .doc() + .map(|doc| doc.doc_id()) + .expect("exhausted clauses were retained out"); + clause.bound = if doc > window_max { + 0.0 + } else { + clause + .posting + .block_max_score_up_to_with_stats(window_max, &self.scorer) + .score + }; + } + clauses.sort_unstable_by(|a, b| a.bound.total_cmp(&b.bound)); + let mut first_essential = 0; + let mut prefix = 0.0_f32; + if self.threshold > 0.0 { + for (i, clause) in clauses.iter_mut().enumerate() { + let widened = float_sum_upper(prefix + clause.bound, i + 1); + if widened > self.threshold { + break; + } + prefix += clause.bound; + clause.prefix_bound = prefix; + first_essential = i + 1; + } + } + prev_first_essential = first_essential; + num_windows += 1; + + if first_essential == clauses.len() { + // No clause combination inside this window can beat the + // threshold: skip it wholesale. + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + // Single live clause: instead of re-running the window + // machinery once per block, scan the baked per-block bounds + // for the next block that can beat the threshold. This is the + // slab form of Lucene's getSkipUpTo and turns dead stretches + // of a dominant term into a tight load-mul-compare loop. + if clauses.len() == 1 && window_min != TERMINATED_DOC_ID && self.threshold > 0.0 { + let posting = &clauses[0].posting; + if let PostingList::Compressed(ref list) = posting.list + && let Some(impacts) = list.impacts.as_ref() + { + let bounds = impacts.level0_doc_weight_bounds(&self.scorer); + let query_weight = posting.query_weight; + // Position by binary search on the first-doc slab: the + // deep cursor lags arbitrarily far behind during long + // skip runs, and walking from it re-scans the same + // blocks on every dead window. + let first_docs = list.block_first_docs(); + let mut block_idx = first_docs + .partition_point(|&first| u64::from(first) <= window_min) + .saturating_sub(1); + while block_idx < bounds.len() + && query_weight * bounds[block_idx] <= self.threshold + { + block_idx += 1; + } + window_min = if block_idx < bounds.len() { + window_min.max(u64::from(list.block_least_doc_id(block_idx))) + } else { + TERMINATED_DOC_ID + }; + } + } + continue; + } + + let total_non_essential_bound = if first_essential > 0 { + clauses[first_essential - 1].prefix_bound + } else { + 0.0 + }; + + // Single essential clause (the common case once the threshold is + // competitive): stream it directly against the non-essential + // prefix, skipping the accumulator entirely. + if first_essential + 1 == clauses.len() { + let (non_essential, essential) = clauses.split_at_mut(first_essential); + let posting = &mut essential[0].posting; + if posting.doc().is_some_and(|doc| doc.doc_id() < window_min) { + posting.next(window_min); + } + let essential_term = posting.term_index(); + let essential_weight = posting.query_weight; + + macro_rules! consider_candidate { + ($doc:expr, $freq:expr) => {{ + let doc = $doc; + let freq = $freq; + num_comparisons += 1; + let doc_length = self.docs.num_tokens(doc as u32); + let score = essential_weight * self.scorer.doc_weight(freq, doc_length); + if !(self.threshold > 0.0 + && score + total_non_essential_bound <= self.threshold) + { + let row_id = if docs_has_row_ids { + self.docs.row_id(doc as u32) + } else { + doc + }; + let masked_out = docs_has_row_ids + && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)); + if !masked_out { + let mut total = score; + let mut rejected = false; + for i in (0..non_essential.len()).rev() { + if self.threshold > 0.0 + && total + non_essential[i].prefix_bound <= self.threshold + { + rejected = true; + break; + } + let probe = &mut non_essential[i].posting; + if probe.doc().is_some_and(|d| d.doc_id() < doc) { + probe.next(doc); + } + if let Some(d) = probe.doc() + && d.doc_id() == doc + { + total += + probe.score(&self.scorer, d.frequency(), doc_length); + } + } + + // Match the classic path's emission rule: a + // candidate must beat the running threshold, + // which drops zero-score matches (e.g. terms + // with idf 0) exactly like Wand::next does. + if !rejected && total > self.threshold { + let full = candidates.len() >= limit; + let beats_kth = + !full || total > candidates.peek().unwrap().0.0.score.0; + if beats_kth { + let mut freqs = Vec::with_capacity(non_essential.len() + 1); + freqs.push((essential_term, freq)); + for clause in non_essential.iter() { + if let Some(d) = clause.posting.doc() + && d.doc_id() == doc + { + freqs.push(( + clause.posting.term_index(), + d.frequency(), + )); + } + } + if full { + candidates.pop(); + } + candidates.push(Reverse(( + ScoredDoc::new(row_id, total), + freqs, + doc_length, + doc, + ))); + if candidates.len() == limit { + let kth = candidates.peek().unwrap().0.0.score.0; + self.update_threshold(kth, params.wand_factor); + } + } + } + } + } + }}; + } + + match posting.list { + PostingList::Compressed(ref list) => { + 'stream: while let Some(cur) = posting.current_doc { + if cur.doc_id() > window_max { + break; + } + let block_idx = posting_block_idx(posting.index, list.block_size); + let block_offset = posting_block_offset(posting.index, list.block_size); + let compressed = unsafe { + &mut *posting.ensure_compressed_block_ptr(list, block_idx) + }; + for offset in block_offset..compressed.doc_ids.len() { + let doc_id = compressed.doc_ids[offset]; + if u64::from(doc_id) > window_max { + posting.index = + posting_block_start(block_idx, list.block_size) + offset; + posting.block_idx = block_idx; + posting.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id, + frequency: compressed.freqs[offset], + })); + break 'stream; + } + consider_candidate!(u64::from(doc_id), compressed.freqs[offset]); + } + let next_start = posting_block_start(block_idx + 1, list.block_size); + if next_start >= list.length as usize { + posting.index = list.length as usize; + posting.block_idx = + posting_block_idx(posting.index, list.block_size); + posting.current_doc = None; + break; + } + posting.index = next_start; + posting.block_idx = block_idx + 1; + let compressed = unsafe { + &mut *posting.ensure_compressed_block_ptr(list, block_idx + 1) + }; + posting.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[0], + frequency: compressed.freqs[0], + })); + } + } + PostingList::Plain(_) => { + while let Some(cur) = posting.doc() { + let doc = cur.doc_id(); + if doc > window_max { + break; + } + consider_candidate!(doc, cur.frequency()); + posting.next(doc + 1); + } + } + } + + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + continue; + } + + // Stream the essential clauses through inner windows. + let mut inner_min = window_min; + loop { + let mut next_essential_doc = TERMINATED_DOC_ID; + for clause in &clauses[first_essential..] { + if let Some(doc) = clause.posting.doc() { + next_essential_doc = next_essential_doc.min(doc.doc_id()); + } + } + inner_min = inner_min.max(next_essential_doc); + if inner_min == TERMINATED_DOC_ID || inner_min > window_max { + break; + } + let inner_max = + window_max.min(inner_min.saturating_add(MAXSCORE_INNER_WINDOW as u64 - 1)); + + for (clause_idx, clause) in clauses.iter_mut().enumerate().skip(first_essential) { + clause.posting.collect_window_scores( + inner_min, + inner_max, + clause_idx, + self.docs, + &self.scorer, + &mut acc, + ); + } + + // Drain candidates in doc order, completing them with the + // non-essential clauses ordered by descending bound. + for word_idx in 0..acc.words.len() { + let mut word = acc.words[word_idx]; + if word == 0 { + continue; + } + acc.words[word_idx] = 0; + while word != 0 { + let bit = word.trailing_zeros() as usize; + word &= word - 1; + let slot = (word_idx << 6) | bit; + let doc = inner_min + slot as u64; + let mut score = acc.scores[slot]; + num_comparisons += 1; + + if self.threshold > 0.0 + && score + total_non_essential_bound <= self.threshold + { + acc.clear_slot(slot); + continue; + } + + let row_id = if docs_has_row_ids { + self.docs.row_id(doc as u32) + } else { + doc + }; + if docs_has_row_ids + && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) + { + acc.clear_slot(slot); + continue; + } + + let doc_length = self.docs.num_tokens(doc as u32); + let mut rejected = false; + for i in (0..first_essential).rev() { + if self.threshold > 0.0 + && score + clauses[i].prefix_bound <= self.threshold + { + rejected = true; + break; + } + let posting = &mut clauses[i].posting; + if posting.doc().is_some_and(|d| d.doc_id() < doc) { + posting.next(doc); + } + if let Some(d) = posting.doc() + && d.doc_id() == doc + { + score += posting.score(&self.scorer, d.frequency(), doc_length); + } + } + + if !rejected && score > self.threshold { + let full = candidates.len() >= limit; + let beats_kth = !full || score > candidates.peek().unwrap().0.0.score.0; + if beats_kth { + let freqs = clauses + .iter() + .enumerate() + .filter_map(|(i, clause)| { + if i >= first_essential { + let freq = acc.clause_freq(i, slot); + (freq > 0).then(|| (clause.posting.term_index(), freq)) + } else { + clause.posting.doc().and_then(|d| { + (d.doc_id() == doc).then(|| { + (clause.posting.term_index(), d.frequency()) + }) + }) + } + }) + .collect::>(); + if full { + candidates.pop(); + } + candidates.push(Reverse(( + ScoredDoc::new(row_id, score), + freqs, + doc_length, + doc, + ))); + if candidates.len() == limit { + let kth = candidates.peek().unwrap().0.0.score.0; + self.update_threshold(kth, params.wand_factor); + } + } + } + acc.clear_slot(slot); + } + } + if inner_max >= window_max { + break; + } + inner_min = inner_max + 1; + } + + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + } + + metrics.record_comparisons(num_comparisons); + + let to_addr = |row_id_slot: u64| { + if docs_has_row_ids { + CandidateAddr::RowId(row_id_slot) + } else { + CandidateAddr::Pending(row_id_slot as u32) + } + }; + Ok(candidates + .into_iter() + .map( + |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + addr: to_addr(doc.row_id), + posting_doc_id, + freqs, + doc_length, + }, + ) + .collect()) + } + // calculate the score of the current document fn score(&self, doc_length: u32) -> f32 { let mut score = 0.0; @@ -1429,10 +2119,16 @@ impl<'a, S: Scorer> Wand<'a, S> { if self.threshold > 0.0 && self.or_block_window_max() <= self.threshold { // On the final block `up_to` is the `u64::MAX` sentinel; step once // there to avoid seeking past the valid doc id range. - let skip_to = match self.up_to { + let mut skip_to = match self.up_to { Some(up_to) if up_to < u32::MAX as u64 => up_to + 1, _ => target + 1, }; + // The narrow window is dead; if the whole level1 group is dead + // too, hop over it in one advance. + let group_skip = self.or_group_skip_to(); + if let Some(group_skip_to) = group_skip { + skip_to = skip_to.max(group_skip_to); + } self.push_back_leads(skip_to); continue; } @@ -1682,12 +2378,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead: f32 = self .lead .iter() - .map(|posting| posting.block_max_score(&self.scorer)) + .map(|posting| posting.window_max_score(self.up_to, &self.scorer)) .sum(); let head: f32 = self .head .iter() - .map(|posting| posting.posting.block_max_score(&self.scorer)) + .map(|posting| posting.posting.window_max_score(self.up_to, &self.scorer)) .sum(); lead + head + self.tail_max_score } @@ -1700,12 +2396,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut sum = self .lead .iter() - .map(|posting| posting.block_max_score(&self.scorer)) + .map(|posting| posting.window_max_score(self.up_to, &self.scorer)) .sum::(); let mut possible_matches = self.lead.len(); for posting in &self.tail { if matches!(posting.posting.block_first_doc(), Some(block_doc) if block_doc <= target) { - sum += posting.posting.block_max_score(&self.scorer); + sum += posting.posting.window_max_score(self.up_to, &self.scorer); possible_matches += 1; } } @@ -1719,72 +2415,121 @@ impl<'a, S: Scorer> Wand<'a, S> { fn update_max_scores(&mut self, target: u64) { // Refresh the block-max window for the current target. The resulting // `up_to` is the furthest doc id for which this block-max view remains - // valid. + // valid. Like Lucene's WANDScorer, the boundary comes from the cheap + // clauses only, and the refresh avoids allocating: heaps are recycled + // through their backing vectors (shallow_next never changes the doc a + // head entry is ordered by, so heapify restores the same shape). let lead_cost = self .lead .iter() .map(|posting| posting.cost()) .min() .unwrap_or(usize::MAX); - let mut up_to = TERMINATED_DOC_ID; + let mut narrow_up_to = TERMINATED_DOC_ID; for posting in &mut self.lead { posting.shallow_next(target); - let block_end = posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end); - } - let head = std::mem::take(&mut self.head); - let mut rebuilt_head = BinaryHeap::with_capacity(head.len()); - for mut posting in head.into_vec() { - if posting.posting.cost() <= lead_cost { - posting.posting.shallow_next(posting.doc_id()); - let block_end = posting - .posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end); - } - rebuilt_head.push(posting); + narrow_up_to = narrow_up_to.min(posting.block_end_doc()); + } + + let mut head_postings = std::mem::take(&mut self.head).into_vec(); + for posting in &mut head_postings { + // Unlike Lucene, every head clause participates in the boundary: + // the refresh is allocation-free and answers from the anchored + // level caches, so frequent refreshes are cheap, while keeping the + // window inside every clause's current block keeps all the bounds + // at tight level0 values. + let doc_id = posting.doc_id(); + posting.posting.shallow_next(doc_id); + narrow_up_to = narrow_up_to.min(posting.posting.block_end_doc()); } - self.head = rebuilt_head; - if up_to == TERMINATED_DOC_ID - && let Some(top) = self.tail.peek() - && top.cost <= lead_cost + + let mut tail_postings = std::mem::take(&mut self.tail).into_vec(); + for tail_posting in &mut tail_postings { + tail_posting.posting.shallow_next(target); + } + + if narrow_up_to == TERMINATED_DOC_ID + && let Some(top) = tail_postings + .iter() + .min_by_key(|posting| posting.posting.cost()) + && top.posting.cost() <= lead_cost { - let block_end = top - .posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end.max(target)); + narrow_up_to = narrow_up_to.min(top.posting.block_end_doc().max(target)); } - self.up_to = Some(up_to); - let tail = std::mem::take(&mut self.tail); + self.up_to = Some(narrow_up_to); + self.head = BinaryHeap::from(head_postings); + self.tail_max_score = 0.0; - for mut tail_posting in tail.into_vec() { - tail_posting.posting.shallow_next(target); - let upper_bound = match tail_posting.posting.block_first_doc() { + for tail_posting in tail_postings { + let posting = tail_posting.posting; + let upper_bound = match posting.block_first_doc() { Some(block_doc) if block_doc <= target => { - tail_posting - .posting - .block_max_score_up_to_with_stats(up_to, &self.scorer) - .score + posting.window_max_score(self.up_to, &self.scorer) } _ => 0.0, }; - if let Some(mut evicted) = - self.insert_tail_with_overflow(tail_posting.posting, upper_bound) - { + if let Some(mut evicted) = self.insert_tail_with_overflow(posting, upper_bound) { evicted.next(target); self.push_head(evicted); } } } + /// After the narrow window proved skippable, try widening the skip to the + /// level1 group boundary, in the spirit of Lucene's `getSkipUpTo`. All + /// bounds come from the anchored level caches, so a failed attempt costs a + /// few loads and float adds — unlike an eager wide-window probe. + /// + /// The group boundary is the minimum current-group end over every live + /// iterator, so each iterator's level1 score is a valid bound over + /// `[target, group_up_to]`. + fn or_group_skip_to(&self) -> Option { + let mut group_up_to = TERMINATED_DOC_ID; + for posting in &self.lead { + let (doc_up_to, _) = posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + for posting in self.head.iter() { + let (doc_up_to, _) = posting.posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + for tail_posting in self.tail.iter() { + let (doc_up_to, _) = tail_posting.posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + if self.up_to.is_some_and(|up_to| group_up_to <= up_to) { + // No gain over the narrow window skip. + return None; + } + + // Second pass over the memoized bounds: sum only the iterators that + // can produce a doc inside the skipped range. + let mut bounds_sum = 0.0_f32; + for posting in &self.lead { + let (_, score) = posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + for posting in self.head.iter() { + if posting.doc_id() > group_up_to { + continue; + } + let (_, score) = posting.posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + for tail_posting in self.tail.iter() { + if !matches!( + tail_posting.posting.block_first_doc(), + Some(block_doc) if block_doc <= group_up_to + ) { + continue; + } + let (_, score) = tail_posting.posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + (bounds_sum <= self.threshold).then_some(group_up_to.saturating_add(1)) + } + fn refine_or_candidate(&mut self, target: u64, doc_length: u32) -> bool { if self.threshold <= 0.0 { return true; @@ -1908,7 +2653,7 @@ impl<'a, S: Scorer> Wand<'a, S> { && posting.is_compressed() && self.up_to.is_some_and(|up_to| target <= up_to) { - posting.block_max_score(&self.scorer) + posting.window_max_score(self.up_to, &self.scorer) } else { posting.global_upper_bound(&self.scorer) } @@ -3362,6 +4107,75 @@ mod tests { } } + #[test] + fn test_or_impact_level1_window_skips_low_group_with_single_score() { + let total = (IMPACT_LEVEL1_BLOCKS + 1) * BLOCK_SIZE; + let target = (IMPACT_LEVEL1_BLOCKS * BLOCK_SIZE) as u64; + let mut docs = DocSet::default(); + for doc_id in 0..total as u64 { + docs.append(doc_id, 1); + } + + let doc_ids = (0..total as u32).collect::>(); + let freqs = doc_ids + .iter() + .map(|doc_id| if u64::from(*doc_id) < target { 1 } else { 10 }) + .collect::>(); + let posting_list = generate_impact_posting_list_with_freqs(doc_ids, freqs, vec![1; total]); + let mut probe = PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + 1.0, + posting_list.clone(), + docs.len(), + ); + probe.shallow_next(0); + let counting_scorer = CountingScorer { + scored: Arc::new(AtomicUsize::new(0)), + }; + let (group_up_to, group_score) = probe.impact_group_bound(&counting_scorer).unwrap(); + assert_eq!(group_up_to, target - 1); + assert_eq!(group_score, 1.0); + // A window ending past the current block must answer from level1. + assert_eq!( + probe.window_max_score(Some(target - 1), &counting_scorer), + 1.0 + ); + + let posting = PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + 1.0, + posting_list, + docs.len(), + ); + let scored = Arc::new(AtomicUsize::new(0)); + let mut wand = Wand::new( + Operator::Or, + std::iter::once(posting), + &docs, + CountingScorer { + scored: scored.clone(), + }, + ); + wand.threshold = 2.0; + + let (candidate, score) = wand.next().unwrap().unwrap(); + assert_eq!(candidate.doc_id(), target); + assert_eq!(score, 10.0); + // The doc-weight bounds bake exactly once (one doc_weight call per + // frontier pair across all entries); beyond that only the returned + // candidate is scored. + let total_entries = (IMPACT_LEVEL1_BLOCKS + 1) + 2; + assert!( + scored.load(Ordering::Relaxed) <= total_entries + 8, + "bounds should be baked once instead of recomputed per window; scored={}", + scored.load(Ordering::Relaxed) + ); + } + #[test] fn test_compressed_impact_block_max_score_memoizes_current_block() { let total = 2 * BLOCK_SIZE as u32; @@ -3381,14 +4195,16 @@ mod tests { let first_score = posting.block_max_score(&scorer); assert_eq!(first_score, 1.0); - // Baking the shared doc-weight bounds visits every frontier pair once - // (two level0 entries plus one level1 entry for this list). + // Baking the doc-weight bounds visits every frontier pair once: two + // level0 entries plus one level1 entry for this list. let baked = scored.load(Ordering::Relaxed); assert!(baked >= 2); { let compressed = unsafe { &mut *posting.compressed_state_ptr() }; - assert_eq!(compressed.current_block_max_score, Some((0, first_score))); - compressed.block_max_window.impact_score_cache = ImpactScoreCache::default(); + assert_eq!( + compressed.level0_cache, + Some((0, BLOCK_SIZE as u32 - 1, first_score)) + ); } let second_score = posting.block_max_score(&scorer); From 4732787bcc8269eecfc256a9c356fc2fb989a36a Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:13:21 +0800 Subject: [PATCH 20/20] perf(fts)!: enable the bulk MAXSCORE path by default With right-sized partitions the bulk path reaches Lucene-parity latency (k10 0.034-0.037s/213-235qps vs Lucene 10.4's 0.037s/216qps on the same 200M-doc corpus, queries, and warm protocol) and its results are score-identical to the classic WAND loop. LANCE_FTS_MAXSCORE=0 opts back into the classic loop. --- rust/lance-index/src/scalar/inverted/wand.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 6df1a2122e6..56365724211 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -50,10 +50,11 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); // Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer -// style). Off by default while it bakes; LANCE_FTS_MAXSCORE=1 opts in. Its -// results are score-identical to the classic WAND loop. +// style). Default on: with right-sized partitions it wins by a wide margin +// (Lucene-parity latency) and its results are score-identical to the classic +// WAND loop. LANCE_FTS_MAXSCORE=0 opts back into the classic loop. static USE_MAXSCORE_SEARCH: LazyLock = - LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() == Ok("1")); + LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() != Ok("0")); #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize {