diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index adc7f94d65e..d5c75158011 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 | 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 f990b2bd589..7f09c4325b7 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=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 9b29d7a0795..82513a89ee7 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 = 128; private Boolean skipMerge; private Integer formatVersion; @@ -226,6 +227,27 @@ public Builder prefixOnly(boolean prefixOnly) { return this; } + /** + * Configure the number of documents in each compressed posting block. + * + *

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 + * 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 + */ + public Builder blockSize(int blockSize) { + if (blockSize != 128 && blockSize != 256) { + throw new IllegalArgumentException("blockSize must be one of 128 or 256"); + } + 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 @@ -242,15 +264,16 @@ 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; @@ -258,6 +281,11 @@ 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); @@ -300,6 +328,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..0ccc429ec57 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,66 @@ */ 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(128, ((Number) json.get("block_size")).intValue()); + } + + @Test + void blockSizeIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build(); + + assertEquals("inverted", params.getIndexType()); + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + 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)); + } + + @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/protos/index_old.proto b/protos/index_old.proto index 601aa2681da..eb984d6fe29 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -39,4 +39,9 @@ message InvertedIndexDetails { uint32 min_ngram_length = 9; uint32 max_ngram_length = 10; bool prefix_only = 11; + // Number of documents per compressed posting block. An absent value means + // the index predates this field and must use the legacy block size of 128. + // A present value records the block size used by the index; 256 is only + // valid with format version 3. + optional uint32 block_size = 12; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 21402c3d276..4b349dd959c 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3366,8 +3366,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 @@ -3375,6 +3377,12 @@ 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 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. + 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/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index c3bf301eee0..4f54e9d31c2 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,9 +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, format_version=1 - ) + 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", format_version=1, **kwargs) def check_read(self): """Verify FTS index can be queried.""" @@ -351,6 +355,16 @@ def check_write(self): def skip_downgrade(self, version: str) -> bool: return version.startswith("0.") + def current_env(self, method_name: str) -> dict[str, str]: + if method_name == "create": + 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 {} + def compat_env(self, version: str, method_name: str) -> dict[str, str]: if method_name in {"create", "check_write"}: return {"LANCE_FTS_FORMAT_VERSION": "1"} diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index b988253b040..985072fe5de 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -949,6 +949,39 @@ 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=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() + assert results.num_rows > 0 + + with pytest.raises(ValueError, match="block_size"): + 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 + ) + + 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( pa.table({"ints": range(1024)}), tmp_path, max_rows_per_file=100 @@ -5188,6 +5221,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 a7e8b52fb1f..d2e9f0b37c3 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2447,6 +2447,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()?); } @@ -2462,7 +2467,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/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 f0e25e37e8c..c6aa6d75ad0 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/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 185e3dcf79c..5f777109124 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -146,6 +146,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 = @@ -211,7 +212,7 @@ impl BasicTrainer for InvertedIndexPlugin { .into())) } - let params = serde_json::from_str::(params)?; + let params = InvertedIndexParams::from_training_json(params)?; Ok(Box::new(InvertedIndexTrainingRequest::new(params))) } @@ -315,6 +316,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() { @@ -324,4 +326,36 @@ mod tests { max_supported_fts_format_version().index_version() ); } + + #[test] + fn test_new_training_request_defaults_missing_block_size_to_128() { + 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 f9567a3ea4b..d5ce5778348 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -4,6 +4,7 @@ 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; #[cfg(test)] use crate::scalar::lance_format::LanceIndexStore; @@ -38,9 +39,9 @@ use std::sync::LazyLock; 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 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. @@ -194,8 +195,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 } @@ -222,6 +226,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(); @@ -256,6 +261,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(); @@ -381,6 +387,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)); @@ -531,6 +538,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 @@ -547,6 +555,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() { @@ -591,6 +607,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)?), @@ -603,6 +620,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( @@ -739,6 +764,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, @@ -760,12 +786,48 @@ 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 { + 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, + 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_format_version_block_size(format_version, block_size) + .expect("invalid FTS format version for posting 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(), @@ -778,13 +840,34 @@ impl InnerBuilder { token_set_format: TokenSetFormat, posting_tail_codec: PostingTailCodec, ) -> 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); + 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 = 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, + token_set_format, + format_version, + block_size, + ); builder.posting_tail_codec = posting_tail_codec; builder } @@ -865,6 +948,7 @@ impl InnerBuilder { token_set_format, format_version, posting_tail_codec, + block_size, tokens, posting_lists, docs, @@ -894,6 +978,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 { @@ -919,7 +1009,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() { @@ -986,7 +1080,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); @@ -999,7 +1097,11 @@ impl InnerBuilder { ); let with_position = self.with_position; let format_version = self.format_version; - 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; @@ -1168,6 +1270,7 @@ struct IndexWorkerConfig { fragment_mask: Option, token_set_format: TokenSetFormat, worker_memory_limit_bytes: u64, + block_size: usize, } impl IndexWorker { @@ -1211,17 +1314,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(), @@ -1351,6 +1459,7 @@ impl IndexWorker { let memory_size = &mut self.memory_size; let posting_tail_codec = builder.posting_tail_codec; + let block_size = builder.block_size; let mut process_text = |text: &str| -> Result<()> { doc_length_bytes += text.len(); let mut token_stream = tokenizer.token_stream_for_doc(text); @@ -1363,9 +1472,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, + block_size, ), ); let new_posting_lists_overhead_size = (builder.posting_lists.capacity() @@ -1450,9 +1560,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(); @@ -1552,15 +1663,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(); @@ -1680,17 +1793,35 @@ 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_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), - InvertedListFormatVersion::V2 => inverted_list_schema_with_tail_codec_and_position_codec( - with_position, - PostingTailCodec::VarintDelta, - Some(PositionStreamCodec::PackedDelta), - ), + InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position, 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, + ) + } } } -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, @@ -1719,24 +1850,42 @@ 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()), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + InvertedListFormatVersion::V1.index_version().to_string(), + ), + ]), + )) } 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, ) } 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, ) -> SchemaRef { let mut fields = vec![ // we compress the posting lists (including row ids and frequencies), @@ -1773,6 +1922,11 @@ 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( POSITIONS_LAYOUT_KEY.to_owned(), @@ -3074,6 +3228,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3097,6 +3252,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3579,6 +3735,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3610,6 +3767,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3648,6 +3806,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 ec3ea6f92ff..885d7089e17 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -47,6 +47,7 @@ use super::index::{ Positions, PostingList, PostingListGroup, PostingListGroupStorage, PostingTailCodec, SharedPositionStream, }; +use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; // --------------------------------------------------------------------------- // Tags @@ -201,7 +202,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 { @@ -217,15 +218,24 @@ impl CacheCodecImpl for PostingList { } fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { - let variant = r.read_u8()?; - match variant { - POSTING_VARIANT_PLAIN => Ok(Self::Plain(deserialize_plain(r)?)), - POSTING_VARIANT_COMPRESSED => Ok(Self::Compressed(deserialize_compressed(r)?)), - other => Err(Error::io(format!("unknown PostingList variant: {other}"))), + match r.version() { + 1 | Self::CURRENT_VERSION => deserialize_posting_list_body(r), + other => Err(Error::io(format!( + "unsupported PostingList cache version: {other}" + ))), } } } +fn deserialize_posting_list_body(r: &mut CacheEntryReader<'_>) -> Result { + let variant = r.read_u8()?; + match variant { + POSTING_VARIANT_PLAIN => Ok(PostingList::Plain(deserialize_plain(r)?)), + POSTING_VARIANT_COMPRESSED => Ok(PostingList::Compressed(deserialize_compressed(r)?)), + other => Err(Error::io(format!("unknown PostingList variant: {other}"))), + } +} + fn serialize_plain(w: &mut CacheEntryWriter<'_>, plain: &PlainPostingList) -> Result<()> { // Plain postings carry only per-doc legacy positions (or none). let position_storage = if plain.positions.is_some() { @@ -314,6 +324,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)?; @@ -345,12 +356,18 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result) -> Result) -> Result<()> { let count = u32::try_from(self.len()) @@ -390,7 +409,7 @@ impl CacheCodecImpl for PostingListGroup { fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { match r.version() { 1 => return deserialize_materialized_group(r), - Self::CURRENT_VERSION => {} + 2 | Self::CURRENT_VERSION => {} other => { return Err(Error::io(format!( "unsupported PostingListGroup cache version: {other}" @@ -425,7 +444,7 @@ fn deserialize_materialized_group(r: &mut CacheEntryReader<'_>) -> Result arrow_array::ListArray { let mut builder = ListBuilder::new(Int32Builder::new()); @@ -500,10 +521,7 @@ mod tests { builder.finish() } - fn packed_group( - postings: &[Vec>], - posting_tail_codec: PostingTailCodec, - ) -> PostingListGroup { + fn packed_batch(postings: &[Vec>], block_size: Option) -> RecordBatch { let mut builder = ListBuilder::new(LargeBinaryBuilder::new()); for posting in postings { for block in posting { @@ -512,13 +530,24 @@ mod tests { builder.append(true); } let postings = builder.finish(); - let schema = Arc::new(Schema::new(vec![Field::new( - POSTING_COL, - postings.data_type().clone(), - false, - )])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(postings)]).unwrap(); - PostingListGroup::new_packed(batch, posting_tail_codec).unwrap() + let fields = vec![Field::new(POSTING_COL, postings.data_type().clone(), false)]; + let schema = Arc::new(match block_size { + Some(block_size) => Schema::new_with_metadata( + fields, + HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string())]), + ), + None => Schema::new(fields), + }); + RecordBatch::try_new(schema, vec![Arc::new(postings)]).unwrap() + } + + fn packed_group( + postings: &[Vec>], + posting_tail_codec: PostingTailCodec, + block_size: Option, + ) -> PostingListGroup { + PostingListGroup::new_packed(packed_batch(postings, block_size), posting_tail_codec) + .unwrap() } fn assert_plain_eq(a: &PlainPostingList, b: &PlainPostingList) { @@ -622,13 +651,14 @@ mod tests { Some(&[6, 7, 8, 9, 10][..]), ]); let posting = - CompressedPostingList::new(blocks, 3.5, 42, PostingTailCodec::VarintDelta, None); + CompressedPostingList::new(blocks, 3.5, 42, PostingTailCodec::VarintDelta, 256, None); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { PostingList::Compressed(restored) => { 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()); } @@ -644,6 +674,7 @@ mod tests { 1.25, 5, PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, Some(CompressedPositionStorage::LegacyPerDoc(legacy_positions( &[&[0, 4, 8]], ))), @@ -678,6 +709,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream(stream)), ); let entry = PostingList::Compressed(posting.clone()); @@ -706,6 +738,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), @@ -740,6 +773,7 @@ mod tests { 2.5, 7, PostingTailCodec::VarintDelta, + 256, None, )); @@ -772,6 +806,7 @@ mod tests { let group = packed_group( &[vec![vec![1, 2, 3], vec![4, 5]], vec![vec![7; 16 * 1024]]], PostingTailCodec::VarintDelta, + Some(256), ); let restored = from_body::(&body_bytes(&group)).unwrap(); assert!(restored.is_packed()); @@ -796,13 +831,27 @@ mod tests { assert_eq!(actual.max_score, expected.max_score); assert_eq!(actual.length, expected.length); assert_eq!(actual.posting_tail_codec, expected.posting_tail_codec); + assert_eq!(actual.block_size, 256); } + let legacy_packed = + packed_group(&[vec![vec![9, 8, 7]]], PostingTailCodec::VarintDelta, None); + let restored = from_body::(&body_bytes(&legacy_packed)).unwrap(); + let PostingList::Compressed(posting) = restored + .posting_list(0, Some(2.0), Some(3)) + .unwrap() + .unwrap() + else { + panic!("expected compressed legacy packed posting"); + }; + assert_eq!(posting.block_size, LEGACY_BLOCK_SIZE); + let legacy_member = PostingList::Compressed(CompressedPostingList::new( LargeBinaryArray::from_opt_vec(vec![Some(&[9u8, 8, 7][..])]), 2.0, 3, PostingTailCodec::VarintDelta, + LEGACY_BLOCK_SIZE, None, )); let mut legacy_body = Vec::new(); @@ -859,26 +908,130 @@ mod tests { use std::sync::Arc; use arrow_array::Array; - use lance_core::cache::CacheCodec; + use arrow_schema::DataType; + use lance_core::cache::{ + CacheCodec, CacheCodecImpl, CacheDecode, CacheEntryReader, CacheEntryWriter, + CacheMissReason, + }; + use lance_core::{Error, Result}; use prost::Message; + use super::super::{ + BLOCKS_COLUMN, GROUP_VARIANT_PACKED, POSTING_VARIANT_COMPRESSED, + posting_tail_codec_to_tag, + }; use super::*; - use crate::cache_pb::{CompressedPostingHeader, PostingTailCodec as PbPostingTailCodec}; + use crate::cache_pb::{ + CompressedPostingHeader, PostingListGroupHeader, PostingTailCodec as PbPostingTailCodec, + }; type ArcAny = Arc; + struct PostingListV1Codec(PostingList); + + impl CacheCodecImpl for PostingListV1Codec { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 1; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + self.0.serialize(w) + } + + fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { + PostingList::deserialize(r).map(Self) + } + } + + struct PostingListGroupV2Codec(PostingListGroup); + + impl CacheCodecImpl for PostingListGroupV2Codec { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 2; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + self.0.serialize(w) + } + + fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { + PostingListGroup::deserialize(r).map(Self) + } + } + + struct LegacyCompressedPostingV1 { + blocks: LargeBinaryArray, + } + + impl CacheCodecImpl for LegacyCompressedPostingV1 { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 1; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + w.write_u8(POSTING_VARIANT_COMPRESSED)?; + w.write_header(&CompressedPostingHeader { + max_score: 2.0, + length: 3, + posting_tail_codec: PbPostingTailCodec::VarintDelta as i32, + ..Default::default() + })?; + let schema = Arc::new(Schema::new(vec![Field::new( + BLOCKS_COLUMN, + DataType::LargeBinary, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(self.blocks.clone())])?; + w.write_ipc(&batch) + } + + fn deserialize(_r: &mut CacheEntryReader<'_>) -> Result { + Err(Error::io( + "LegacyCompressedPostingV1 is a writer-only test codec".to_string(), + )) + } + } + + struct LegacyPackedGroupV2 { + batch: RecordBatch, + posting_tail_codec: PostingTailCodec, + } + + impl CacheCodecImpl for LegacyPackedGroupV2 { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 2; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + w.write_u8(GROUP_VARIANT_PACKED)?; + let count = u32::try_from(self.batch.num_rows()) + .map_err(|_| Error::io("legacy packed group is too large".to_string()))?; + w.write_header(&PostingListGroupHeader { count })?; + w.write_u8(posting_tail_codec_to_tag(self.posting_tail_codec))?; + w.write_ipc(&self.batch) + } + + fn deserialize(_r: &mut CacheEntryReader<'_>) -> Result { + Err(Error::io( + "LegacyPackedGroupV2 is a writer-only test codec".to_string(), + )) + } + } + fn codec() -> CacheCodec { CacheCodec::from_impl::() } - /// Serialize an entry through the full codec (envelope + body). - fn serialize_entry(entry: PostingList) -> Vec { + fn serialize_typed_entry(entry: T) -> Vec { let any: ArcAny = Arc::new(entry); let mut buf = Vec::new(); - codec().serialize(&any, &mut buf).unwrap(); + CacheCodec::from_impl::() + .serialize(&any, &mut buf) + .unwrap(); buf } + /// Serialize an entry through the full codec (envelope + body). + fn serialize_entry(entry: PostingList) -> Vec { + serialize_typed_entry(entry) + } + /// A `Bytes` whose base address is 64-byte aligned, modelling a backend /// that reads cache entries into an aligned buffer. fn aligned_bytes(payload: &[u8]) -> Bytes { @@ -902,6 +1055,7 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream(stream)), )) } @@ -947,6 +1101,7 @@ mod tests { vec![vec![1; 48], vec![1; 48]], ], PostingTailCodec::VarintDelta, + Some(256), ); let group_codec = CacheCodec::from_impl::(); @@ -1054,6 +1209,73 @@ mod tests { assert!(codec().deserialize(&Bytes::from(buf)).hit().is_none()); } + #[test] + fn old_codecs_reject_new_v3_envelopes_as_version_too_new() { + let posting = Bytes::from(serialize_entry(compressed_with_shared_positions())); + match CacheCodec::from_impl::().deserialize(&posting) { + CacheDecode::Miss(reason) => { + assert_eq!(reason, CacheMissReason::VersionTooNew) + } + CacheDecode::Hit(_) => panic!("v1 PostingList codec accepted a v2 envelope"), + } + + let group = packed_group( + &[vec![vec![1, 2, 3]], vec![vec![4, 5, 6]]], + PostingTailCodec::VarintDelta, + Some(256), + ); + let group = Bytes::from(serialize_typed_entry(group)); + match CacheCodec::from_impl::().deserialize(&group) { + CacheDecode::Miss(reason) => { + assert_eq!(reason, CacheMissReason::VersionTooNew) + } + CacheDecode::Hit(_) => { + panic!("v2 PostingListGroup codec accepted a v3 envelope") + } + } + } + + #[test] + fn current_codecs_read_legacy_payloads_without_block_size() { + let legacy_posting = LegacyCompressedPostingV1 { + blocks: LargeBinaryArray::from_opt_vec(vec![Some(&[9u8, 8, 7][..])]), + }; + let legacy_posting = Bytes::from(serialize_typed_entry(legacy_posting)); + let restored = codec().deserialize(&legacy_posting).hit().unwrap(); + let restored = restored.downcast::().unwrap(); + let PostingList::Compressed(restored) = restored.as_ref() else { + panic!("expected a compressed legacy posting"); + }; + assert_eq!(restored.block_size, LEGACY_BLOCK_SIZE); + + let legacy_batch = packed_batch(&[vec![vec![1, 2, 3]], vec![vec![4, 5, 6]]], None); + assert!( + !legacy_batch + .schema_ref() + .metadata() + .contains_key(POSTING_BLOCK_SIZE_KEY) + ); + let legacy_group = LegacyPackedGroupV2 { + batch: legacy_batch, + posting_tail_codec: PostingTailCodec::VarintDelta, + }; + let legacy_group = Bytes::from(serialize_typed_entry(legacy_group)); + let restored = CacheCodec::from_impl::() + .deserialize(&legacy_group) + .hit() + .unwrap() + .downcast::() + .unwrap(); + let PostingList::Compressed(restored) = restored + .posting_list(0, Some(2.0), Some(3)) + .unwrap() + .unwrap() + else { + panic!("expected a compressed legacy packed posting"); + }; + assert_eq!(restored.block_size, LEGACY_BLOCK_SIZE); + } + /// A pre-stabilization blob (no magic) self-heals to a miss. #[test] fn pre_stabilization_blob_is_miss() { diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 22383c9fbe4..8fb7f8f4032 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -5,10 +5,15 @@ 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 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: @@ -41,18 +46,40 @@ 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 - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder( doc_ids.copied().collect::>().as_slice(), frequencies.copied().collect::>().as_slice(), @@ -63,27 +90,25 @@ 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)?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } + encode_posting_block_payload(&doc_id_buffer, &freq_buffer, &mut builder)?; builder.append_value(""); doc_id_buffer.clear(); freq_buffer.clear(); @@ -91,44 +116,187 @@ pub fn compress_posting_list_with_tail_codec<'a>( // we don't compress the last block if it is not full if !doc_id_buffer.is_empty() { - // 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())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder(&doc_id_buffer, &freq_buffer, tail_codec, &mut builder)?; builder.append_value(""); } Ok(builder.finish()) } +/// Byte length of the block-max-score prefix on posting blocks. 128-doc +/// blocks store a per-block max score, patched in at build time; 256-doc +/// (V3) blocks always carry impact skip data, which supersedes it, so they +/// store none. +#[inline] +pub fn posting_block_score_prefix_len(block_size: usize) -> usize { + if block_size == MAX_POSTING_BLOCK_SIZE { + 0 + } else { + 4 + } +} + pub fn encode_full_posting_block_into( doc_ids: &[u32], frequencies: &[u32], block: &mut Vec, ) -> Result<()> { - debug_assert_eq!(doc_ids.len(), BLOCK_SIZE); - debug_assert_eq!(frequencies.len(), BLOCK_SIZE); - block.extend_from_slice(&0f32.to_le_bytes()); - let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - compress_sorted_block(doc_ids, &mut buffer, block)?; - compress_block(frequencies, &mut buffer, block)?; + validate_block_size(doc_ids.len())?; + debug_assert_eq!(doc_ids.len(), frequencies.len()); + if posting_block_score_prefix_len(doc_ids.len()) > 0 { + 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()); + 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)?; + } + // 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_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], codec: PostingTailCodec, + block_size: usize, block: &mut Vec, ) -> Result<()> { debug_assert_eq!(doc_ids.len(), frequencies.len()); - block.extend_from_slice(&0f32.to_le_bytes()); + if posting_block_score_prefix_len(block_size) > 0 { + block.extend_from_slice(&0f32.to_le_bytes()); + } compress_posting_remainder(doc_ids, frequencies, codec, block)?; Ok(()) } #[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())?; @@ -139,7 +307,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])?; @@ -236,7 +414,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; @@ -346,7 +524,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() { @@ -617,23 +795,46 @@ 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 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(); - 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( compressed, remainder, tail_codec, + block_size, &mut doc_ids, &mut frequencies, ); @@ -668,24 +869,46 @@ 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, ) { - // 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); + debug_assert!(validate_block_size(block_size).is_ok()); + debug_assert!(buffer.len() >= block_size); + // skip the block max score prefix (128-doc blocks only) + let mut block = &block[posting_block_score_prefix_len(block_size)..]; + 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_pfor_block_with::(block, buffer, frequencies); + block = &block[num_bytes..]; + } + _ => unreachable!("validated posting block size should be supported"), + } + debug_assert!( + block.is_empty(), + "posting block has {} trailing bytes after decoding", + block.len() + ); } pub fn decompress_posting_remainder( block: &[u8], n: usize, codec: PostingTailCodec, + block_size: usize, doc_ids: &mut Vec, frequencies: &mut Vec, ) { - let block = &block[4..]; + let block = &block[posting_block_score_prefix_len(block_size)..]; match codec { PostingTailCodec::Fixed32 => { decompress_raw_remainder(block, n, doc_ids); @@ -722,9 +945,14 @@ pub fn decompress_posting_remainder( } } -pub fn decode_full_posting_block(block: &[u8], doc_ids: &mut Vec, frequencies: &mut Vec) { - let mut buffer = [0u32; BLOCK_SIZE]; - decompress_posting_block(block, &mut buffer, doc_ids, frequencies); +pub fn decode_full_posting_block( + block: &[u8], + doc_ids: &mut Vec, + frequencies: &mut Vec, + block_size: usize, +) { + let mut buffer = [0u32; MAX_POSTING_BLOCK_SIZE]; + decompress_posting_block(block, &mut buffer, doc_ids, frequencies, block_size); } pub fn decompress_sorted_block( @@ -732,7 +960,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); @@ -740,11 +978,18 @@ pub fn decompress_sorted_block( 5 + num_bytes } -fn decompress_block(block: &[u8], buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec) { - 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]; - 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) { @@ -754,11 +999,18 @@ pub fn decompress_raw_remainder(compressed: &[u8], n: usize, dest: &mut Vec } } -pub fn read_posting_tail_first_doc(block: &[u8], codec: PostingTailCodec) -> u32 { +pub fn read_posting_tail_first_doc( + block: &[u8], + codec: PostingTailCodec, + block_size: usize, +) -> u32 { + let prefix = posting_block_score_prefix_len(block_size); match codec { - PostingTailCodec::Fixed32 => u32::from_le_bytes(block[4..8].try_into().unwrap()), + PostingTailCodec::Fixed32 => { + u32::from_le_bytes(block[prefix..prefix + 4].try_into().unwrap()) + } PostingTailCodec::VarintDelta => { - let mut offset = 4usize; + let mut offset = prefix; decode_varint_u32(block, &mut offset) .expect("posting tail block should contain a valid first doc id") } @@ -809,6 +1061,84 @@ mod tests { Ok(()) } + #[test] + fn test_compress_posting_list_supported_block_sizes() -> Result<()> { + 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) + .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_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); + // 256-doc blocks carry no block-max-score prefix (impacts supply the + // per-block bound): [first_doc u32][doc num_bits u8][doc payload]... + let doc_num_bits = block[4]; + let doc_bytes = BitPacker8x::compressed_block_size(doc_num_bits); + let freq_header_offset = 5 + 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 + 2 + 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 26b0f0256b0..b07d1e207b2 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -3,8 +3,8 @@ use lance_core::utils::row_addr_remap::RowAddrRemap; 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, @@ -53,14 +53,15 @@ use std::sync::LazyLock; use tokio::{sync::OnceCell, task::spawn_blocking}; use tracing::{info, instrument, warn}; -use super::encoding::PositionBlockBuilder; +use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder}; 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, ScoredDoc, doc_file_path, inverted_list_schema_for_version, posting_file_path, - token_file_path, + BLOCK_SIZE, ScoredDoc, doc_file_path, inverted_list_schema_for_version_with_block_size, + posting_file_path, token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -86,8 +87,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,8 +113,10 @@ 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"; pub const POSTING_TAIL_CODEC_FIXED32_V1: &str = "fixed32_v1"; pub const POSTING_TAIL_CODEC_VARINT_DELTA_V1: &str = "varint_delta_v1"; pub const POSITIONS_LAYOUT_SHARED_STREAM_V2: &str = "shared_stream_v2"; @@ -147,7 +152,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)] @@ -155,6 +160,7 @@ pub enum InvertedListFormatVersion { V1, #[default] V2, + V3, } impl InvertedListFormatVersion { @@ -165,29 +171,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) } } @@ -198,14 +225,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, @@ -367,6 +427,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 { @@ -404,6 +479,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); } @@ -481,9 +575,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(), } } @@ -1641,6 +1738,8 @@ impl InvertedPartition { num_docs, false, frag_reuse_index, + // V3 (256-doc block) partitions score with quantized doc lengths. + inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, )); Ok(Self { @@ -1763,6 +1862,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() { @@ -1787,7 +1893,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 { @@ -1795,7 +1901,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); @@ -2027,11 +2137,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 @@ -2441,6 +2552,7 @@ pub struct PostingListReader { has_position: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, /// Runtime posting-list cache grouping. Non-empty v2 indexes use synthetic @@ -2538,6 +2650,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())?; @@ -2559,6 +2672,7 @@ impl PostingListReader { metadata, has_position, posting_tail_codec, + block_size, positions_layout, grouping, index_cache: WeakLanceCache::from(index_cache), @@ -2604,6 +2718,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 { .. }) } @@ -2861,7 +2979,11 @@ impl PostingListReader { Some(&[POSTING_COL, MAX_SCORE_COL, LENGTH_COL]), ) .await?; - PostingListGroup::new_packed(batch.shrink_to_fit()?, self.posting_tail_codec) + PostingListGroup::new_packed_with_block_size( + batch.shrink_to_fit()?, + self.posting_tail_codec, + self.block_size, + ) } fn posting_list_from_batch_parts( @@ -2869,6 +2991,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( @@ -2876,6 +2999,7 @@ impl PostingListReader { max_score, length, posting_tail_codec, + block_size, positions_layout, )?; Ok(posting_list) @@ -2892,6 +3016,7 @@ impl PostingListReader { max_score, length, self.posting_tail_codec, + self.block_size, self.positions_layout, ) } @@ -2928,6 +3053,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)); @@ -3084,6 +3210,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, } } @@ -3117,12 +3244,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 { @@ -3167,6 +3296,7 @@ impl PostingListReader { let chunk_batch = self.read_chunk_batch(tok_start, tok_end, false).await?; let ranges = grouping.ranges_for_chunk(tok_start, tok_end, token_count); let posting_tail_codec = self.posting_tail_codec; + let block_size = self.block_size; spawn_blocking(move || { let mut groups = Vec::with_capacity(ranges.len()); @@ -3179,7 +3309,11 @@ impl PostingListReader { groups.push(( start, end, - PostingListGroup::new_packed(group_batch, posting_tail_codec)?, + PostingListGroup::new_packed_with_block_size( + group_batch, + posting_tail_codec, + block_size, + )?, )); } Result::Ok(groups) @@ -3413,6 +3547,7 @@ struct ChunkBuildState { max_scores: Option>>, lengths: Option>>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3423,6 +3558,7 @@ struct PrewarmBuildCtx<'a> { max_scores: Option<&'a [f32]>, lengths: Option<&'a [u32]>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3654,8 +3790,8 @@ impl SharedPositionStream { } /// A group of consecutive posting lists held in a single cache entry, in row -/// order (issue #7040). Prewarmed v2 groups without positions retain only the -/// compact Arrow posting rows read from `invert.lance`; max-score/length +/// order (issue #7040). Prewarmed modern groups without positions retain only +/// the compact Arrow posting rows read from `invert.lance`; max-score/length /// metadata stays in the reader and is injected when a query creates a /// posting-list view. Cold-loaded groups may keep inline metadata to preserve /// one-read query loading. Legacy and position-bearing prewarm paths use the @@ -3675,6 +3811,7 @@ pub(super) enum PostingListGroupStorage { pub(super) struct PackedPostingListGroup { pub(super) batch: RecordBatch, pub(super) posting_tail_codec: PostingTailCodec, + pub(super) block_size: usize, } impl DeepSizeOf for PostingListGroup { @@ -3704,6 +3841,39 @@ impl PostingListGroup { batch: RecordBatch, posting_tail_codec: PostingTailCodec, ) -> Result { + let block_size = parse_posting_block_size(batch.schema_ref().metadata())?; + Self::new_packed_with_block_size(batch, posting_tail_codec, block_size) + } + + fn new_packed_with_block_size( + batch: RecordBatch, + posting_tail_codec: PostingTailCodec, + block_size: usize, + ) -> Result { + validate_block_size(block_size)?; + if let Some(encoded_block_size) = batch.schema_ref().metadata().get(POSTING_BLOCK_SIZE_KEY) + { + let encoded_block_size = encoded_block_size.parse::().map_err(|err| { + Error::index(format!( + "invalid {POSTING_BLOCK_SIZE_KEY} metadata value {encoded_block_size:?}: {err}" + )) + })?; + if encoded_block_size != block_size { + return Err(Error::index(format!( + "packed posting group {POSTING_BLOCK_SIZE_KEY}={encoded_block_size} does not match block_size={block_size}" + ))); + } + } + + // Projected reads may drop schema metadata. Restore the reader's + // validated block size before the batch enters the packed cache so IPC + // roundtrips remain self-describing. Older packed cache entries omit + // the key and enter through new_packed with the legacy 128-doc default. + let mut schema = batch.schema().as_ref().clone(); + schema + .metadata + .insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); + let batch = batch.with_schema(Arc::new(schema))?; let postings = batch .column_by_name(POSTING_COL) .and_then(|column| column.as_list_opt::()) @@ -3758,6 +3928,7 @@ impl PostingListGroup { storage: PostingListGroupStorage::Packed(PackedPostingListGroup { batch, posting_tail_codec, + block_size, }), }) } @@ -3838,6 +4009,7 @@ impl PostingListGroup { max_score, length, group.posting_tail_codec, + group.block_size, None, )))) } @@ -3858,7 +4030,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( @@ -3866,6 +4039,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( @@ -3881,6 +4055,7 @@ impl PostingList { max_score, length, posting_tail_codec, + block_size, positions_layout, ) } @@ -3890,6 +4065,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) { @@ -3904,6 +4080,7 @@ impl PostingList { max_score.unwrap(), length.unwrap(), posting_tail_codec, + block_size, shared_position_codec, ); Ok(Self::Compressed(posting)) @@ -3975,9 +4152,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 @@ -4130,15 +4312,31 @@ impl PlainPostingList { } } -#[derive(Debug, PartialEq, Clone)] +#[derive(Debug, Clone)] pub struct CompressedPostingList { pub max_score: f32, pub length: u32, // each binary is a block of compressed data - // that contains `BLOCK_SIZE` doc ids and then `BLOCK_SIZE` frequencies + // that contains `block_size` doc ids and then `block_size` frequencies, + // packed by the physical bitpacker matching that block size. pub blocks: LargeBinaryArray, pub posting_tail_codec: PostingTailCodec, + pub block_size: usize, pub positions: Option, + // 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 { @@ -4158,22 +4356,40 @@ impl CompressedPostingList { max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, positions: Option, ) -> Self { + debug_assert!(block_size.is_power_of_two()); Self { max_score, length, blocks, posting_tail_codec, + block_size, positions, + first_docs: Arc::new(OnceLock::new()), } } + /// Block sizes are validated powers of two, so per-doc hot loops derive + /// block indices with shift/mask instead of runtime division, which is + /// measurably slower in the iterator advance path. + #[inline] + pub(crate) fn block_shift(&self) -> u32 { + self.block_size.trailing_zeros() + } + + #[inline] + pub(crate) fn block_mask(&self) -> usize { + self.block_size - 1 + } + pub fn from_batch( batch: &RecordBatch, max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, shared_position_codec: Option, ) -> Self { debug_assert_eq!(batch.num_rows(), 1); @@ -4210,7 +4426,9 @@ impl CompressedPostingList { length, blocks, posting_tail_codec, + block_size, positions, + first_docs: Arc::new(OnceLock::new()), } } @@ -4220,23 +4438,51 @@ impl CompressedPostingList { self.blocks.clone(), self.posting_tail_codec, self.positions.clone(), + self.block_size, ) } pub fn block_max_score(&self, block_idx: usize) -> f32 { + // 256-doc (V3) blocks store no per-block max score: their impact + // skip data supplies the tight per-block bound, so callers on that + // path never reach here. Fall back to the list-level max, which is + // still a valid (looser) bound for any block. + if super::encoding::posting_block_score_prefix_len(self.block_size) == 0 { + return self.max_score; + } let block = self.blocks.value(block_idx); block[0..4].try_into().map(f32::from_le_bytes).unwrap() } 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 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, + self.block_size, + ); + } + let prefix = super::encoding::posting_block_score_prefix_len(self.block_size); + block[prefix..prefix + 4] + .try_into() + .map(u32::from_le_bytes) + .unwrap() + }) + .collect() + }) } } @@ -4287,12 +4533,14 @@ impl EncodedBlocks { doc_ids: &[u32], frequencies: &[u32], codec: PostingTailCodec, + block_size: usize, ) -> Result<()> { self.offsets.push(self.bytes.len() as u32); super::encoding::encode_remainder_posting_block_into( doc_ids, frequencies, codec, + block_size, &mut self.bytes, ) } @@ -4361,6 +4609,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, } @@ -4386,6 +4635,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, @@ -4537,9 +4787,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, ) } @@ -4547,6 +4798,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, @@ -4557,6 +4829,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, len: 0, memory_size_bytes: 0, } @@ -4578,8 +4851,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; @@ -4587,7 +4860,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 @@ -4667,7 +4945,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"); } @@ -4726,7 +5004,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(()) @@ -4777,13 +5055,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() @@ -4954,6 +5236,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4963,6 +5246,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) @@ -5001,6 +5285,7 @@ impl PostingListBuilder { with_positions, posting_tail_codec, length, + block_size, mut encoded_blocks, mut encoded_position_blocks, tail_entries, @@ -5009,14 +5294,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, @@ -5025,7 +5315,9 @@ impl PostingListBuilder { frequencies.iter().copied(), ); max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { @@ -5042,8 +5334,11 @@ impl PostingListBuilder { doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -5060,15 +5355,18 @@ impl PostingListBuilder { )) } + #[allow(clippy::too_many_arguments)] fn build_compressed_with_block_scores_from_parts( with_positions: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, mut encoded_blocks: EncodedBlocks, mut encoded_position_blocks: EncodedPositionBlocks, tail_entries: &[RawDocInfo], tail_position_block: Option>, mut block_max_scores: impl Iterator, ) -> Result<(LargeBinaryArray, Option, f32)> { + let has_score_prefix = super::encoding::posting_block_score_prefix_len(block_size) > 0; let mut max_score = f32::MIN; let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); let mut frequencies = Vec::with_capacity(BLOCK_SIZE); @@ -5078,7 +5376,9 @@ impl PostingListBuilder { .next() .ok_or_else(|| Error::index("missing block max score".to_owned()))?; max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { @@ -5091,8 +5391,11 @@ impl PostingListBuilder { doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -5110,12 +5413,15 @@ 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 schema = inverted_list_schema_for_version(self.has_positions(), format_version); + 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, + self.block_size, + ); let legacy_positions = if self.with_positions && !format_version.uses_shared_position_stream() { Some(self.build_legacy_positions()?) @@ -5132,6 +5438,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -5142,6 +5449,7 @@ impl PostingListBuilder { Self::build_compressed_with_block_scores_from_parts( with_positions, posting_tail_codec, + block_size, encoded_blocks .map(|encoded_blocks| *encoded_blocks) .unwrap_or_default(), @@ -5162,6 +5470,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -5173,13 +5482,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()?) @@ -5196,6 +5499,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -5205,6 +5509,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) @@ -5227,6 +5532,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -5239,8 +5545,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; @@ -5382,9 +5691,55 @@ impl Ord for RawDocInfo { } } +/// Lucene SmallFloat-style doc-length quantization for V3 (256-doc block) +/// scoring: a 4-mantissa-bit float-like byte code. Values 0-7 are exact; +/// larger values keep their top four significand bits (relative error +/// <= 6.25%) and decode to their bucket floor. The floor only ever shortens +/// a doc, and a shorter doc only raises its BM25 weight, so bounds baked +/// from quantized lengths stay valid upper bounds of quantized scores. +pub(super) fn quantize_doc_length(value: u32) -> u8 { + let num_bits = 32 - value.leading_zeros(); + if num_bits < 4 { + value as u8 + } else { + let shift = num_bits - 4; + (((value >> shift) as u8) & 0x07) | (((shift + 1) as u8) << 3) + } +} + +#[inline] +pub(super) fn dequantize_doc_length(code: u8) -> u32 { + DEQUANTIZED_DOC_LENGTHS[code as usize] +} + +pub(super) static DEQUANTIZED_DOC_LENGTHS: [u32; 256] = build_dequantized_doc_lengths(); + +const fn build_dequantized_doc_lengths() -> [u32; 256] { + let mut table = [0u32; 256]; + let mut code = 0usize; + while code < 256 { + let bits = (code & 0x07) as u64; + let shift = (code >> 3) as i64 - 1; + let decoded = if shift < 0 { + bits + } else { + (bits | 0x08) << shift + }; + // Codes past the largest u32 encoding are never produced; saturate so + // the table stays total. + table[code] = if decoded > u32::MAX as u64 { + u32::MAX + } else { + decoded as u32 + }; + code += 1; + } + table +} + // DocSet is a mapping from row ids to the number of tokens in the document // It's used to sort the documents by the bm25 score -#[derive(Debug, Clone, Default, DeepSizeOf)] +#[derive(Debug, Clone, Default)] pub struct DocSet { row_ids: Vec, num_tokens: Vec, @@ -5392,6 +5747,26 @@ pub struct DocSet { inv: Vec<(u64, u32)>, total_tokens: u64, + + // V3 (256-doc block) partitions score with quantized doc lengths: the + // flag is set at partition load and the byte-norm slab bakes lazily on + // first scoring use (shared by clones of the loaded set). 128-block + // partitions never set the flag and keep exact scoring. + scoring_quantized: bool, + norms: Arc>>, +} + +impl DeepSizeOf for DocSet { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.row_ids.deep_size_of_children(context) + + self.num_tokens.deep_size_of_children(context) + + self.inv.deep_size_of_children(context) + + self + .norms + .get() + .map(|slab| std::mem::size_of_val(slab.as_ref())) + .unwrap_or(0) + } } impl DocSet { @@ -5460,9 +5835,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; @@ -5473,13 +5858,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); } @@ -5529,6 +5914,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), } } @@ -5564,6 +5951,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5605,6 +5994,8 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5624,6 +6015,8 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }) } @@ -5634,6 +6027,7 @@ impl DocSet { let len = self.len(); let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, Vec::with_capacity(len)); + self.invalidate_norms(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { match mapping.get(row_id) { @@ -5660,6 +6054,39 @@ impl DocSet { self.num_tokens[doc_id as usize] } + /// Enable quantized doc-length scoring (V3 / 256-doc block partitions). + pub fn set_quantized_scoring(&mut self, quantized: bool) { + self.scoring_quantized = quantized; + } + + /// The quantized doc-length slab when this set scores quantized (V3 + /// partitions), baked on first use; `None` for exact-scoring sets. + pub fn scoring_norms(&self) -> Option<&[u8]> { + if !self.scoring_quantized { + return None; + } + Some( + self.norms + .get_or_init(|| { + self.num_tokens + .iter() + .map(|&n| quantize_doc_length(n)) + .collect() + }) + .as_ref(), + ) + } + + /// Doc length as scoring sees it: the quantized bucket floor for V3 + /// partitions, the exact value otherwise. + #[inline] + pub fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + match self.scoring_norms() { + Some(norms) => dequantize_doc_length(norms[doc_id as usize]), + None => self.num_tokens[doc_id as usize], + } + } + // this can be used only if it's a legacy format, // which store the sorted row ids so that we can use binary search #[inline] @@ -5676,9 +6103,18 @@ impl DocSet { self.row_ids.push(row_id); self.num_tokens.push(num_tokens); self.total_tokens += num_tokens as u64; + self.invalidate_norms(); self.row_ids.len() as u32 - 1 } + // Drop the baked norm slab after a mutation; it re-bakes on the next + // scoring use. + fn invalidate_norms(&mut self) { + if self.norms.get().is_some() { + self.norms = Arc::new(std::sync::OnceLock::new()); + } + } + pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.capacity() * std::mem::size_of::() @@ -6265,15 +6701,21 @@ mod tests { token: &str, row_id: u64, ) -> Result> { - let mut partition = InnerBuilder::new_with_format_version( + 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(false, PostingTailCodec::Fixed32); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); posting_list.add(0, PositionRecorder::Count(1)); partition.posting_lists.push(posting_list); partition.docs.append(row_id, 1); @@ -6289,6 +6731,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())) @@ -6309,6 +6760,85 @@ 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())]); + 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(); + 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 format_version = params.resolved_format_version(); + 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(), + format_version, + block_size, + ); + builder.tokens.add("needle".to_owned()); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + 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; @@ -6548,6 +7078,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] @@ -7191,13 +7734,16 @@ mod tests { /// Prewarming a large partition in multiple chunks must end up holding exactly the /// same per-token posting lists (doc ids and frequencies) as the whole-file path. /// Parametrized over layout: the legacy-v1 chunk path rebases global offsets to - /// chunk-local rows, which the v2 one-row-per-token path never exercises. + /// chunk-local rows, while the modern one-row-per-token path covers both + /// legacy-sized v2 and 256-doc v3 posting blocks. #[rstest::rstest] - #[case::v1(InvertedListFormatVersion::V1)] - #[case::v2(InvertedListFormatVersion::V2)] + #[case::v1(InvertedListFormatVersion::V1, LEGACY_BLOCK_SIZE)] + #[case::v2(InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE)] + #[case::v3(InvertedListFormatVersion::V3, 256)] #[tokio::test] async fn test_prewarm_streams_in_chunks_preserves_content( #[case] format_version: InvertedListFormatVersion, + #[case] block_size: usize, ) { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( @@ -7211,19 +7757,23 @@ mod tests { let num_tokens = runtime_posting_group_tokens() as u32 + 4; const DOCS_PER_TOKEN: u32 = 3; let posting_tail_codec = format_version.posting_tail_codec(); - let mut builder = InnerBuilder::new_with_format_version( + let mut builder = InnerBuilder::new_with_format_version_and_block_size( 0, false, TokenSetFormat::default(), format_version, + block_size, ); // expected[token] = [(doc_id, frequency)] in stored (doc-id) order. let mut expected: Vec> = Vec::new(); let mut doc_id = 0u64; for t in 0..num_tokens { builder.tokens.add(format!("tok_{t:03}")); - let mut posting = - PostingListBuilder::new_with_posting_tail_codec(false, posting_tail_codec); + let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + posting_tail_codec, + block_size, + ); let mut docs = Vec::new(); for _ in 0..DOCS_PER_TOKEN { posting.add(doc_id as u32, PositionRecorder::Count(1)); @@ -7236,15 +7786,15 @@ mod tests { } builder.write(store.as_ref()).await.unwrap(); + let params = InvertedIndexParams::default() + .block_size(block_size) + .unwrap(); let metadata = std::collections::HashMap::from_iter(vec![ ( "partitions".to_owned(), serde_json::to_string(&vec![0u64]).unwrap(), ), - ( - "params".to_owned(), - serde_json::to_string(&InvertedIndexParams::default()).unwrap(), - ), + ("params".to_owned(), serde_json::to_string(¶ms).unwrap()), ( TOKEN_SET_FORMAT_KEY.to_owned(), TokenSetFormat::default().to_string(), @@ -7253,6 +7803,11 @@ mod tests { POSTING_TAIL_CODEC_KEY.to_owned(), 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())) @@ -7266,6 +7821,7 @@ mod tests { .unwrap(); let inverted_list = &index.partitions[0].inverted_list; assert_eq!(inverted_list.len(), num_tokens as usize); + assert_eq!(inverted_list.block_size(), block_size); // Force a small target chunk. Since CHUNK_TOKENS is below the runtime // group size, synthetic group alignment should still split only at @@ -7284,6 +7840,23 @@ mod tests { "single partition must be streamed in more than one chunk, got {chunk_count}" ); + if format_version == InvertedListFormatVersion::V3 { + let (start, end) = inverted_list.group_range_for_token(0).unwrap(); + let group = inverted_list + .index_cache + .get_with_key(&PostingListGroupKey { start, end }) + .await + .expect("v3 prewarm should populate the packed group cache"); + assert!(group.is_packed()); + let (max_score, length) = inverted_list.bulk_metadata_for_token(0); + let PostingList::Compressed(posting) = + group.posting_list(0, max_score, length).unwrap().unwrap() + else { + panic!("expected compressed v3 posting list"); + }; + assert_eq!(posting.block_size, 256); + } + // (2) Correctness: every token's posting list round-trips with exactly // the doc ids and frequencies of the whole-file path. for token_id in 0..num_tokens { @@ -7974,6 +8547,7 @@ mod tests { 1.0, SLICE_LEN as u32, PostingTailCodec::Fixed32, + LEGACY_BLOCK_SIZE, None, ); @@ -8342,6 +8916,7 @@ mod tests { partition_ids: Vec, params: InvertedIndexParams, ) { + let format_version = params.resolved_format_version(); let metadata = HashMap::from([ ( "partitions".to_owned(), @@ -8352,6 +8927,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())) @@ -9248,6 +9835,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/iter.rs b/rust/lance-index/src/scalar/inverted/iter.rs index dc07b15769c..3755fa41d80 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, }, }; @@ -28,6 +27,7 @@ impl<'a> PostingListIterator<'a> { posting.blocks.clone(), posting.posting_tail_codec, posting.positions.clone(), + posting.block_size, ))) } } @@ -82,13 +82,14 @@ pub struct CompressedPostingListIterator { next_block_idx: usize, posting_tail_codec: PostingTailCodec, positions: Option, + block_size: usize, idx: usize, doc_ids: Vec, frequencies: Vec, doc_idx_in_block: usize, decoded_positions: Vec, position_offsets: Vec, - buffer: [u32; BLOCK_SIZE], + buffer: [u32; MAX_POSTING_BLOCK_SIZE], } impl CompressedPostingListIterator { @@ -97,10 +98,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,18 +110,19 @@ 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(), doc_idx_in_block: 0, decoded_positions: Vec::new(), position_offsets: Vec::new(), - buffer: [0; BLOCK_SIZE], + buffer: [0; MAX_POSTING_BLOCK_SIZE], } } } @@ -163,6 +166,7 @@ impl Iterator for CompressedPostingListIterator { compressed, self.remainder, self.posting_tail_codec, + self.block_size, &mut self.doc_ids, &mut self.frequencies, ); @@ -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/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 12b5c1ebd32..a8724513340 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -64,6 +64,9 @@ pub struct DeferredDocSet { docs_path: String, is_legacy: bool, frag_reuse_index: Option>, + /// V3 (256-doc block) partitions score with quantized doc lengths; the + /// flag is applied to every `DocSet` this deferred set materializes. + quantized_scoring: bool, /// Doc count cached at construction so `len()` stays sync + IO-free. num_rows: usize, /// `sum(num_tokens)` cached on first compute. @@ -126,18 +129,21 @@ impl lance_core::deepsize::DeepSizeOf for LazyDocSet { } impl LazyDocSet { + #[allow(clippy::too_many_arguments)] pub fn new( store: Arc, docs_path: String, num_rows: usize, is_legacy: bool, frag_reuse_index: Option>, + quantized_scoring: bool, ) -> Self { Self::Deferred(Box::new(DeferredDocSet { store, docs_path, is_legacy, frag_reuse_index, + quantized_scoring, num_rows, total_tokens: OnceCell::new(), num_tokens_col: OnceCell::new(), @@ -300,7 +306,7 @@ impl DeferredDocSet { .get_or_try_init(|| async { // If the stats path already pulled NUM_TOKEN_COL, // read only ROW_ID and rebuild from the two columns. - let docs = if self.num_tokens_col.get().is_some() { + let mut docs = if self.num_tokens_col.get().is_some() { let num_tokens = self.num_tokens_column().await?; let row_ids = self.row_ids_column().await?; DocSet::from_columns( @@ -317,6 +323,7 @@ impl DeferredDocSet { ) .await? }; + docs.set_quantized_scoring(self.quantized_scoring); Result::Ok(Arc::new(docs)) }) .await? @@ -333,7 +340,9 @@ impl DeferredDocSet { .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()))) + let mut docs = DocSet::from_num_tokens_only(num_tokens.as_ref()); + docs.set_quantized_scoring(self.quantized_scoring); + Result::Ok(Arc::new(docs)) }) .await? .clone(); diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 5080bfe624f..429ad965419 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")] @@ -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::{ @@ -32,6 +33,16 @@ use lance_tokenizer::{ WhitespaceTokenizer, }; +/// Posting block size for indexes whose metadata predates configurable block sizes. +/// +/// This must remain 128 so that legacy on-disk data is decoded correctly. +pub const LEGACY_BLOCK_SIZE: usize = 128; +/// Default posting block size for newly created indexes when none is configured. +/// +/// This intentionally matches [`LEGACY_BLOCK_SIZE`] today but may evolve independently. +pub const DEFAULT_BLOCK_SIZE: usize = 128; +pub const VALID_BLOCK_SIZES: [usize; 2] = [128, 256]; + /// Tokenizer configs #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct InvertedIndexParams { @@ -101,6 +112,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 currently default to 128. + #[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 @@ -127,7 +149,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, @@ -153,6 +178,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), }) } } @@ -180,6 +206,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, format_version: defaults.format_version, @@ -199,6 +229,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 or 256, 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) +} + fn deserialize_format_version<'de, D>( deserializer: D, ) -> std::result::Result, D::Error> @@ -215,17 +269,17 @@ where .map(Some) .map_err(serde::de::Error::custom), 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", - )); + let Some(format_version) = value.as_u64() else { + return Err(serde::de::Error::custom(format!( + "FTS format_version must be 1, 2, or 3, got {value}" + ))); }; - resolve_fts_format_version(Some(&value.to_string())) + resolve_fts_format_version(Some(&format_version.to_string())) .map(Some) .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}" ))), } } @@ -267,6 +321,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, format_version: None, @@ -358,6 +413,25 @@ impl InvertedIndexParams { self } + /// Set the compressed posting block size. + /// + /// 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 + } + pub fn memory_limit_mb(mut self, memory_limit_mb: u64) -> Self { self.memory_limit_mb = Some(memory_limit_mb); self @@ -374,17 +448,29 @@ 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) + self.format_version.unwrap_or_else(|| { + default_fts_format_version_for_block_size(self.block_size) + .expect("InvertedIndexParams block_size must be validated before use") + }) + } + + /// 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. @@ -414,6 +500,25 @@ 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()); + + let params: Self = serde_json::from_value(value)?; + params.validate_format_version()?; + Ok(params) + } + pub fn build(&self) -> Result> { let mut builder = self.build_base_tokenizer()?; if let Some(max_token_length) = self.max_token_length { @@ -519,8 +624,10 @@ pub fn language_model_home() -> Option { #[cfg(test)] mod tests { + use crate::pbold; + use super::{InvertedIndexParams, InvertedListFormatVersion}; - use lance_tokenizer::TokenStream; + use lance_tokenizer::{Language, TokenStream}; use rstest::rstest; #[test] @@ -588,6 +695,130 @@ 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_training_json_invalid_numeric_format_version_includes_value() { + let err = InvertedIndexParams::from_training_json(r#"{"format_version": -1}"#).unwrap_err(); + assert!(matches!(&err, lance_core::Error::Arrow { .. })); + assert!(err.to_string().contains("got -1")); + } + + #[test] + fn test_block_size_default_serializes() { + let params = InvertedIndexParams::default(); + 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(128))); + } + + #[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(256).unwrap(); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + assert_eq!(details.block_size, Some(256)); + + 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); + } + + #[rstest] + #[case::block_size_128(128)] + #[case::block_size_256(256)] + fn test_block_size_accepts_supported_values(#[case] block_size: usize) { + 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 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 or 256")); + } + #[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..1e0a87c7ad4 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, @@ -74,12 +74,12 @@ 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), - buffer: Box::new([0; BLOCK_SIZE]), + doc_ids: Vec::with_capacity(block_size), + freqs: Vec::with_capacity(block_size), + buffer: Box::new([0; MAX_POSTING_BLOCK_SIZE]), position_block_idx: None, position_values: Vec::new(), position_offsets: Vec::new(), @@ -95,21 +95,29 @@ 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, remainder, tail_codec, + block_size, &mut self.doc_ids, &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 +287,7 @@ impl PostingIterator { list.blocks.len(), list.length, list.posting_tail_codec, + list.block_size, ); } compressed as *mut CompressedState @@ -307,8 +316,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 +332,7 @@ impl PostingIterator { index: 0, block_idx: 0, approximate_upper_bound, - compressed: is_compressed.then(|| UnsafeCell::new(CompressedState::new())), + compressed, } } @@ -361,8 +374,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_shift(); + let block_offset = self.index & list.block_mask(); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; // Read from the decompressed block @@ -400,8 +413,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_shift(); + let block_offset = self.index & list.block_mask(); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; if compressed.position_block_idx != Some(block_idx) { @@ -441,33 +454,34 @@ 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 shift = list.block_shift(); + let mut block_idx = self.index >> shift; 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 << shift); 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 >> shift; + let block_offset = self.index & list.block_mask(); 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 << shift) + 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) << shift; } - self.block_idx = self.index / BLOCK_SIZE; + self.block_idx = self.index >> shift; } PostingList::Plain(ref list) => { self.index += list.row_ids[self.index..].partition_point(|&id| id < least_id); @@ -903,7 +917,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } let doc_length = match &doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; @@ -1077,7 +1091,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // score the doc let doc_length = match is_compressed { - true => self.docs.num_tokens(doc_id as u32), + true => self.docs.scoring_num_tokens(doc_id as u32), false => self.docs.num_tokens_by_row_id(row_id), }; if self.operator == Operator::Or && !self.refine_or_candidate(doc_id, doc_length) { @@ -1227,7 +1241,7 @@ impl<'a, S: Scorer> Wand<'a, S> { continue; }; let doc_length = match &first_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; let mut lead_score = 0.0; @@ -1309,7 +1323,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead_doc = self.lead.first().and_then(|posting| posting.doc())?; let doc_length = match &lead_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; if self.and_candidate_cannot_beat_threshold(doc_length) { @@ -2146,6 +2160,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.rs b/rust/lance/src/dataset/mem_wal/index.rs index d16d3105551..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. @@ -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 ))), } @@ -352,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); } } @@ -452,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 @@ -1043,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(), @@ -1052,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, @@ -1386,13 +1403,71 @@ 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}" ); } + #[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 61c797db041..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 @@ -1764,6 +1774,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 format_version = self.params.resolved_format_version(); let posting_tail_codec = format_version.posting_tail_codec(); let total_rows_u64 = total_rows as u64; @@ -1783,11 +1794,12 @@ impl FtsMemIndex { } } if all_docs.is_empty() { - return Ok(InnerBuilder::new_with_format_version( + return Ok(InnerBuilder::new_with_format_version_and_block_size( partition_id, with_position, Default::default(), format_version, + block_size, )); } @@ -1873,10 +1885,13 @@ 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_posting_tail_codec( - with_position, - posting_tail_codec, - )); + posting_lists.push( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + with_position, + posting_tail_codec, + block_size, + ), + ); let plb = &mut posting_lists[token_id]; for (doc_id, freq, pos) in docs_for_term { let recorder = if with_position { @@ -1888,11 +1903,12 @@ impl FtsMemIndex { } } - let mut builder = InnerBuilder::new_with_format_version( + let mut builder = InnerBuilder::new_with_format_version_and_block_size( partition_id, with_position, Default::default(), format_version, + block_size, ); builder.set_tokens(tokens); builder.set_docs(docs); @@ -2332,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, - } + }) } } @@ -4666,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![ 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::>(); diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 059fa4e2b36..1da614d2930 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}; @@ -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 b0b0eacded6..2ac243bff71 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 fn scan(self) -> Result> + 'static + Send> { + pub fn scan(self) -> Result>> { let batch_readahead = self.config.batch_readahead; let simplified_predicates = self.simplified_predicates()?; let ordered_output = self.config.ordered_output;