lance-index 13.0.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use super::super::posting_prewarm::ChunkPostingMode;
use super::*;

#[derive(Debug)]
struct ControlledPostingReads {
    started: std::sync::atomic::AtomicUsize,
    active: std::sync::atomic::AtomicUsize,
    max_active: std::sync::atomic::AtomicUsize,
    gates: Vec<tokio::sync::Semaphore>,
    completed: std::sync::Mutex<Vec<usize>>,
    changed: tokio::sync::Notify,
}

impl ControlledPostingReads {
    fn new(token_count: usize) -> Self {
        Self {
            started: std::sync::atomic::AtomicUsize::new(0),
            active: std::sync::atomic::AtomicUsize::new(0),
            max_active: std::sync::atomic::AtomicUsize::new(0),
            gates: (0..token_count)
                .map(|_| tokio::sync::Semaphore::new(0))
                .collect(),
            completed: std::sync::Mutex::new(Vec::with_capacity(token_count)),
            changed: tokio::sync::Notify::new(),
        }
    }

    async fn wait_for_started(&self, expected: usize) {
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            loop {
                let changed = self.changed.notified();
                if self.started.load(std::sync::atomic::Ordering::SeqCst) >= expected {
                    return;
                }
                changed.await;
            }
        })
        .await
        .expect("timed out waiting for controlled posting reads to start");
    }

    async fn release_and_wait(&self, token_id: usize, expected_completed: usize) {
        self.gates[token_id].add_permits(1);
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            loop {
                let changed = self.changed.notified();
                if self
                    .completed
                    .lock()
                    .expect("controlled posting completion lock poisoned")
                    .len()
                    >= expected_completed
                {
                    return;
                }
                changed.await;
            }
        })
        .await
        .expect("timed out waiting for controlled posting read to complete");
    }
}

struct ControlledPostingReader {
    inner: Arc<dyn IndexReader>,
    control: Arc<ControlledPostingReads>,
    fail_token: Option<usize>,
}

#[async_trait]
impl IndexReader for ControlledPostingReader {
    async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result<RecordBatch> {
        self.inner.read_record_batch(n, batch_size).await
    }

    async fn read_global_buffer(&self, index: u32) -> Result<bytes::Bytes> {
        self.inner.read_global_buffer(index).await
    }

    async fn read_range(
        &self,
        range: std::ops::Range<usize>,
        projection: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        let is_singleton_position_read = range.end == range.start + 1
            && projection.is_some_and(|columns| {
                columns.contains(&POSTING_COL) && columns.contains(&POSITION_COL)
            });
        if !is_singleton_position_read {
            return self.inner.read_range(range, projection).await;
        }

        let token_id = range.start;
        self.control
            .started
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let active = self
            .control
            .active
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
            + 1;
        self.control
            .max_active
            .fetch_max(active, std::sync::atomic::Ordering::SeqCst);
        self.control.changed.notify_waiters();

        let permit = self.control.gates[token_id]
            .acquire()
            .await
            .map_err(|err| Error::internal(format!("controlled posting gate closed: {err}")))?;
        permit.forget();
        let result = if self.fail_token == Some(token_id) {
            Err(Error::io(format!(
                "injected singleton posting read failure for token {token_id}"
            )))
        } else {
            self.inner.read_range(range, projection).await
        };

        self.control
            .active
            .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
        self.control
            .completed
            .lock()
            .expect("controlled posting completion lock poisoned")
            .push(token_id);
        self.control.changed.notify_waiters();
        result
    }

    async fn num_batches(&self, batch_size: u64) -> u32 {
        self.inner.num_batches(batch_size).await
    }

    fn num_rows(&self) -> usize {
        self.inner.num_rows()
    }

    fn schema(&self) -> &lance_core::datatypes::Schema {
        self.inner.schema()
    }

    fn file_size_bytes(&self) -> Option<u64> {
        self.inner.file_size_bytes()
    }
}

#[derive(Debug)]
struct ControlledMergeStore {
    inner: Arc<dyn IndexStore>,
    posting_file: String,
    io_parallelism: usize,
    control: Arc<ControlledPostingReads>,
    fail_token: Option<usize>,
}

impl DeepSizeOf for ControlledMergeStore {
    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
        self.inner.deep_size_of_children(context)
    }
}

#[async_trait]
impl IndexStore for ControlledMergeStore {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn clone_arc(&self) -> Arc<dyn IndexStore> {
        Arc::new(Self {
            inner: self.inner.clone(),
            posting_file: self.posting_file.clone(),
            io_parallelism: self.io_parallelism,
            control: self.control.clone(),
            fail_token: self.fail_token,
        })
    }

    fn io_parallelism(&self) -> usize {
        self.io_parallelism
    }

    async fn new_index_file(
        &self,
        name: &str,
        schema: Arc<Schema>,
    ) -> Result<Box<dyn crate::scalar::IndexWriter>> {
        self.inner.new_index_file(name, schema).await
    }

    async fn open_index_file(&self, name: &str) -> Result<Arc<dyn IndexReader>> {
        let reader = self.inner.open_index_file(name).await?;
        if name == self.posting_file {
            Ok(Arc::new(ControlledPostingReader {
                inner: reader,
                control: self.control.clone(),
                fail_token: self.fail_token,
            }))
        } else {
            Ok(reader)
        }
    }

    fn with_io_priority(&self, io_priority: u64) -> Arc<dyn IndexStore> {
        Arc::new(Self {
            inner: self.inner.with_io_priority(io_priority),
            posting_file: self.posting_file.clone(),
            io_parallelism: self.io_parallelism,
            control: self.control.clone(),
            fail_token: self.fail_token,
        })
    }

    async fn copy_index_file(
        &self,
        name: &str,
        dest_store: &dyn IndexStore,
    ) -> Result<crate::scalar::IndexFile> {
        self.inner.copy_index_file(name, dest_store).await
    }

    async fn rename_index_file(
        &self,
        name: &str,
        new_name: &str,
    ) -> Result<crate::scalar::IndexFile> {
        self.inner.rename_index_file(name, new_name).await
    }

    async fn delete_index_file(&self, name: &str) -> Result<()> {
        self.inner.delete_index_file(name).await
    }

    async fn list_files_with_sizes(&self) -> Result<Vec<crate::scalar::IndexFile>> {
        self.inner.list_files_with_sizes().await
    }
}

#[test]
fn address_read_concurrency_respects_payload_budget() {
    assert_eq!(address_read_concurrency(64, 0), 64);
    assert_eq!(address_read_concurrency(64, 8 * 1024 * 1024), 8);
    assert_eq!(address_read_concurrency(64, 16 * 1024 * 1024), 4);
    assert_eq!(
        address_read_concurrency(64, 2 * MAX_CONCURRENT_ADDRESS_READ_BYTES),
        1
    );
}

#[derive(Debug)]
struct MetadataAccessDeniedStore {
    inner: Arc<dyn IndexStore>,
}

impl DeepSizeOf for MetadataAccessDeniedStore {
    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
        self.inner.deep_size_of_children(context)
    }
}

#[async_trait]
impl IndexStore for MetadataAccessDeniedStore {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn clone_arc(&self) -> Arc<dyn IndexStore> {
        Arc::new(Self {
            inner: self.inner.clone(),
        })
    }

    fn io_parallelism(&self) -> usize {
        self.inner.io_parallelism()
    }

    async fn new_index_file(
        &self,
        name: &str,
        schema: Arc<Schema>,
    ) -> Result<Box<dyn crate::scalar::IndexWriter>> {
        self.inner.new_index_file(name, schema).await
    }

    async fn open_index_file(&self, name: &str) -> Result<Arc<dyn IndexReader>> {
        if name == METADATA_FILE {
            Err(Error::io("metadata access denied"))
        } else {
            self.inner.open_index_file(name).await
        }
    }

    fn with_io_priority(&self, io_priority: u64) -> Arc<dyn IndexStore> {
        Arc::new(Self {
            inner: self.inner.with_io_priority(io_priority),
        })
    }

    async fn copy_index_file(
        &self,
        name: &str,
        dest_store: &dyn IndexStore,
    ) -> Result<crate::scalar::IndexFile> {
        self.inner.copy_index_file(name, dest_store).await
    }

    async fn rename_index_file(
        &self,
        name: &str,
        new_name: &str,
    ) -> Result<crate::scalar::IndexFile> {
        self.inner.rename_index_file(name, new_name).await
    }

    async fn delete_index_file(&self, name: &str) -> Result<()> {
        self.inner.delete_index_file(name).await
    }

    async fn list_files_with_sizes(&self) -> Result<Vec<crate::scalar::IndexFile>> {
        self.inner.list_files_with_sizes().await
    }
}

#[tokio::test]
async fn params_legacy_fallback_probes_tokens_after_metadata_access_denied() {
    let tmpdir = TempObjDir::default();
    let inner: Arc<dyn IndexStore> = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        tmpdir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let expected = InvertedIndexParams::default();
    let metadata = HashMap::from([(
        "tokenizer".to_owned(),
        serde_json::to_string(&expected).unwrap(),
    )]);
    let mut writer = inner
        .new_index_file(TOKENS_FILE, Arc::new(Schema::empty()))
        .await
        .unwrap();
    writer.finish_with_metadata(metadata).await.unwrap();
    let store = MetadataAccessDeniedStore { inner };

    let actual = InvertedIndex::load_params(&store).await.unwrap();
    assert_eq!(
        serde_json::to_value(actual).unwrap(),
        serde_json::to_value(expected).unwrap()
    );
}

#[tokio::test]
async fn params_legacy_probe_preserves_metadata_error_when_tokens_are_missing() {
    let tmpdir = TempObjDir::default();
    let inner: Arc<dyn IndexStore> = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        tmpdir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let store = MetadataAccessDeniedStore { inner };

    let error = InvertedIndex::load_params(&store).await.unwrap_err();
    assert!(matches!(error, Error::IO { .. }));
    assert!(error.to_string().contains("metadata access denied"));
}

#[tokio::test]
async fn params_metadata_ignores_unknown_fields() {
    let tmpdir = TempObjDir::default();
    let store = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        tmpdir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let expected = InvertedIndexParams::default();
    let mut params = serde_json::to_value(&expected).unwrap();
    let params = params.as_object_mut().unwrap();
    params.insert("skip_merge".to_owned(), true.into());
    params.insert(
        "future_parameter".to_owned(),
        serde_json::json!({ "enabled": true }),
    );
    let metadata = HashMap::from([("params".to_owned(), serde_json::to_string(params).unwrap())]);
    let mut writer = store
        .new_index_file(METADATA_FILE, Arc::new(Schema::empty()))
        .await
        .unwrap();
    writer.finish_with_metadata(metadata).await.unwrap();

    let actual = InvertedIndex::load_params(store.as_ref()).await.unwrap();
    assert_eq!(
        serde_json::to_value(actual).unwrap(),
        serde_json::to_value(expected).unwrap()
    );
}

#[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"));
}

#[test]
fn test_num_tokens_only_reuses_sliced_arrow_storage() {
    let docs = {
        let source = UInt32Array::from(vec![999, 7, 16, 1024, 888]);
        let sliced = source.slice(1, 3);
        let mut docs = DocSet::from_num_tokens_only(&sliced);

        let NumTokens::Shared(values) = &docs.num_tokens else {
            panic!("num-tokens-only DocSet must retain shared Arrow storage");
        };
        assert!(values.ptr_eq(sliced.values()));
        assert_eq!(values.as_ref(), &[7, 16, 1024]);
        assert_eq!(docs.total_tokens_num(), 1047);
        docs.set_quantized_scoring(true);
        assert_eq!(docs.scoring_norms().unwrap().len(), 3);
        assert_eq!(
            docs.scoring_num_tokens(0),
            dequantize_doc_length(quantize_doc_length(7))
        );
        assert_eq!(
            docs.scoring_num_tokens(2),
            dequantize_doc_length(quantize_doc_length(1024))
        );
        docs
    };

    assert_eq!(docs.len(), 3);
    assert_eq!(docs.num_tokens(0), 7);
    assert_eq!(docs.num_tokens(2), 1024);
}

#[test]
fn test_cached_num_tokens_uses_supplied_total_and_full_stays_owned() {
    const CACHED_TOTAL_MARKER: u64 = 123_456;

    let num_tokens = UInt32Array::from(vec![3, 5, 8]);
    let docs = DocSet::from_cached_num_tokens(&num_tokens, CACHED_TOTAL_MARKER);
    assert_eq!(docs.total_tokens_num(), CACHED_TOTAL_MARKER);
    assert!(matches!(&docs.num_tokens, NumTokens::Shared(_)));

    let row_ids = UInt64Array::from(vec![10, 20, 30]);
    let full = DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap();
    assert!(matches!(&full.num_tokens, NumTokens::Owned(_)));
    assert_eq!(full.total_tokens_num(), 16);
    assert_eq!(full.row_id(1), 20);
}

#[test]
fn test_combined_fields_append_invalidates_row_ids_ascending_cache() {
    let row_ids = UInt64Array::from(vec![10, 20, 30]);
    let num_tokens = UInt32Array::from(vec![1, 1, 1]);
    let mut docs = DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap();

    // Memoize the ascending answer for the current (sorted) row_ids.
    assert!(docs.row_ids_strictly_ascending());

    // Appending a smaller row_id breaks the invariant; the cached answer
    // must be dropped so the recomputed value reflects the mutation.
    docs.append(5, 1);
    assert!(!docs.row_ids_strictly_ascending());

    // The element-document append is the same mutation and owes the same
    // invalidation. A set built by `with_coordinate_rank` is shared by its
    // clones, so a stale answer would outlive the builder that produced it.
    let mut docs = DocSet::with_coordinate_rank(1);
    for (doc_id, row_id) in [10u64, 20, 30].into_iter().enumerate() {
        assert_eq!(
            docs.append_with_doc_index(row_id, 1, &[doc_id as u32])
                .unwrap(),
            doc_id as u32
        );
    }
    assert!(docs.row_ids_strictly_ascending());

    docs.append_with_doc_index(5, 1, &[0]).unwrap();
    assert!(!docs.row_ids_strictly_ascending());
}

#[test]
fn test_doc_length_by_row_id_matches_scan() {
    // Compressed layout: row ids are stored in doc-id order and resolved
    // through `inv`. The targeted lookup must equal an `iter()` scan that
    // sums the matching row's lengths, and return 0 for an absent row.
    let row_ids = UInt64Array::from(vec![30, 10, 20]);
    let num_tokens = UInt32Array::from(vec![8, 3, 5]);
    let docs = DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap();
    for target in [10u64, 20, 30, 99] {
        let expected: u64 = docs
            .iter()
            .filter(|(id, _)| **id == target)
            .map(|(_, nt)| *nt as u64)
            .sum();
        assert_eq!(docs.doc_length_by_row_id(target), expected, "row {target}");
    }
    assert_eq!(docs.doc_length_by_row_id(10), 3);
    assert_eq!(docs.doc_length_by_row_id(99), 0);

    // Legacy layout: row id == doc id, row ids are sorted, and a row indexed
    // as several list documents forms a contiguous run whose lengths sum.
    let legacy_row_ids = UInt64Array::from(vec![10, 10, 20]);
    let legacy_num_tokens = UInt32Array::from(vec![3, 4, 5]);
    let legacy = DocSet::from_columns(&legacy_row_ids, &legacy_num_tokens, true, None).unwrap();
    assert!(legacy.inv.is_empty());
    assert_eq!(legacy.doc_length_by_row_id(10), 7);
    assert_eq!(legacy.doc_length_by_row_id(20), 5);
    assert_eq!(legacy.doc_length_by_row_id(99), 0);
}

#[test]
fn test_posting_builder_writes_impacts_for_supported_block_sizes() {
    for block_size in [128, 256] {
        let format_version = default_fts_format_version_for_block_size(block_size).unwrap();
        let num_docs = block_size * 33 + 1;
        let mut docs = DocSet::default();
        let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size(
            false,
            format_version.posting_tail_codec(),
            block_size,
        );
        for doc_id in 0..num_docs {
            docs.append(doc_id as u64, (doc_id % 5 + 1) as u32);
            posting.add(
                doc_id as u32,
                PositionRecorder::Count((doc_id % 3 + 1) as u32),
            );
        }
        let schema =
            inverted_list_schema_for_version_with_block_size(false, format_version, block_size);
        let batch = posting.to_batch_with_docs(&docs, schema).unwrap();
        assert!(batch.column_by_name(IMPACT_COL).is_some());
        let max_score = batch[MAX_SCORE_COL].as_primitive::<Float32Type>().value(0);
        let length = batch[LENGTH_COL].as_primitive::<UInt32Type>().value(0);
        let posting = PostingList::from_batch(&batch, Some(max_score), Some(length)).unwrap();
        let PostingList::Compressed(posting) = posting else {
            panic!("expected compressed posting list");
        };
        let impacts = posting.impacts.expect("posting should include impacts");
        assert_eq!(impacts.level0_len(), posting.blocks.len());
        assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32));
        assert_eq!(
            impacts.entries().len(),
            impacts.level0_len() + impacts.level1_len()
        );
    }
}

#[test]
fn test_posting_builder_without_impact_column_roundtrips_without_impacts() {
    let mut posting = PostingListBuilder::new(false);
    for doc_id in 0..BLOCK_SIZE + 3 {
        posting.add(doc_id as u32, PositionRecorder::Count(1));
    }
    let batch = posting.to_batch(vec![1.0, 1.0]).unwrap();
    assert!(batch.column_by_name(IMPACT_COL).is_none());
    let posting =
        PostingList::from_batch(&batch, Some(1.0), Some((BLOCK_SIZE + 3) as u32)).unwrap();
    assert!(!posting.has_impacts());
}

#[tokio::test]
async fn test_build_search_uses_configured_posting_block_size() {
    let tmpdir = TempObjDir::default();
    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.get_or_add("needle");
    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 impacts = posting
        .impacts
        .as_ref()
        .expect("newly written posting list should include impacts");
    assert_eq!(impacts.level0_len(), posting.blocks.len());
    assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32));
    assert_eq!(
        impacts.entries().len(),
        impacts.level0_len() + impacts.level1_len()
    );

    let tokens = Arc::new(Tokens::new(vec!["needle".to_owned()], DocType::Text));
    let params = Arc::new(FtsSearchParams::new().with_limit(Some(10)));
    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));
}

#[rstest::rstest]
#[case::v1(InvertedListFormatVersion::V1, LEGACY_BLOCK_SIZE, 514, 7)]
#[case::v2(InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE, 8, 4)]
#[case::v3(InvertedListFormatVersion::V3, 256, 6, 4)]
#[tokio::test]
async fn test_into_builder_chunks_postings_by_list_children(
    #[case] format_version: InvertedListFormatVersion,
    #[case] block_size: usize,
    #[case] max_list_children: u64,
    #[case] expected_chunk_count: usize,
) {
    const NUM_TOKENS: usize = 7;
    const NUM_DOCS: usize = 257;
    // V1's nested per-document position block count is not stored in metadata,
    // so each token gets its own read. V3's exact metadata-derived bound admits
    // two token rows per read.

    let source_dir = TempObjDir::default();
    let source_store = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        source_dir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let mut source = InnerBuilder::new_with_format_version_and_block_size(
        0,
        true,
        TokenSetFormat::default(),
        format_version,
        block_size,
    );
    for token_id in 0..NUM_TOKENS {
        source.tokens.get_or_add(&format!("token_{token_id}"));
        let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size(
            true,
            format_version.posting_tail_codec(),
            block_size,
        );
        for doc_id in 0..NUM_DOCS {
            posting.add(
                doc_id as u32,
                PositionRecorder::Position(vec![token_id as u32, token_id as u32 + 10].into()),
            );
        }
        source.posting_lists.push(posting);
    }
    for doc_id in 0..NUM_DOCS {
        source
            .docs
            .append(1_000 + doc_id as u64, (NUM_TOKENS * 2) as u32);
    }
    source.write(source_store.as_ref()).await.unwrap();

    let source_partition = InvertedPartition::load(
        source_store,
        0,
        None,
        &LanceCache::no_cache(),
        TokenSetFormat::default(),
    )
    .await
    .unwrap();
    let (mut merged, chunk_count) = source_partition
        .into_builder_with_chunk_limits(NUM_TOKENS, max_list_children)
        .await
        .unwrap();
    assert_eq!(
        chunk_count, expected_chunk_count,
        "list-child limit must split merge reads"
    );
    assert_eq!(merged.posting_lists.len(), NUM_TOKENS);
    assert!(
        merged
            .posting_lists
            .iter()
            .all(|posting| posting.len() == NUM_DOCS && posting.has_positions())
    );

    // Rewriting the chunk-built builder verifies that positions survived
    // conversion and that impact data can be regenerated from every posting.
    let dest_dir = TempObjDir::default();
    let dest_store = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        dest_dir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    merged.write(dest_store.as_ref()).await.unwrap();
    let merged_cache = LanceCache::with_capacity(1 << 20);
    let merged_partition = InvertedPartition::load(
        dest_store,
        0,
        None,
        &merged_cache,
        TokenSetFormat::default(),
    )
    .await
    .unwrap();
    let posting = merged_partition
        .inverted_list
        .posting_list(0, true, &NoOpMetricsCollector)
        .await
        .unwrap();
    let PostingList::Compressed(posting) = posting else {
        panic!("expected compressed posting list");
    };
    assert_eq!(posting.block_size, block_size);
    assert!(posting.impacts.is_some());
    let actual = PostingList::Compressed(posting)
        .iter()
        .map(|(doc_id, frequency, positions)| {
            (doc_id, frequency, positions.unwrap().collect::<Vec<_>>())
        })
        .collect::<Vec<_>>();
    assert_eq!(actual.len(), NUM_DOCS);
    assert_eq!(actual[0], (0, 2, vec![0, 10]));
    assert_eq!(actual[NUM_DOCS - 1], (256, 2, vec![0, 10]));
}

#[tokio::test]
async fn test_v1_position_merge_reads_are_bounded_and_ordered() {
    const NUM_TOKENS: usize = 10;
    const NUM_DOCS: usize = 3;
    const EXPECTED_CONCURRENCY: usize = 8;

    let source_dir = TempObjDir::default();
    let source_store: Arc<dyn IndexStore> = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        source_dir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let mut source = InnerBuilder::new_with_format_version_and_block_size(
        0,
        true,
        TokenSetFormat::default(),
        InvertedListFormatVersion::V1,
        LEGACY_BLOCK_SIZE,
    );
    for token_id in 0..NUM_TOKENS {
        source.tokens.get_or_add(&format!("token_{token_id}"));
        let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size(
            true,
            InvertedListFormatVersion::V1.posting_tail_codec(),
            LEGACY_BLOCK_SIZE,
        );
        for doc_id in 0..NUM_DOCS {
            posting.add(
                doc_id as u32,
                PositionRecorder::Position(vec![token_id as u32, token_id as u32 + 100].into()),
            );
        }
        source.posting_lists.push(posting);
    }
    for doc_id in 0..NUM_DOCS {
        source
            .docs
            .append(1_000 + doc_id as u64, (NUM_TOKENS * 2) as u32);
    }
    source.write(source_store.as_ref()).await.unwrap();

    let control = Arc::new(ControlledPostingReads::new(NUM_TOKENS));
    let controlled_store: Arc<dyn IndexStore> = Arc::new(ControlledMergeStore {
        inner: source_store.clone(),
        posting_file: posting_file_path(0),
        io_parallelism: 64,
        control: control.clone(),
        fail_token: None,
    });
    let source_partition = InvertedPartition::load(
        controlled_store,
        0,
        None,
        &LanceCache::no_cache(),
        TokenSetFormat::default(),
    )
    .await
    .unwrap();

    let merge_task = tokio::spawn(async move {
        source_partition
            .into_builder_with_chunk_limits(NUM_TOKENS, u64::MAX)
            .await
    });
    control.wait_for_started(EXPECTED_CONCURRENCY).await;
    assert_eq!(
        control.active.load(std::sync::atomic::Ordering::SeqCst),
        EXPECTED_CONCURRENCY
    );
    assert_eq!(
        control.max_active.load(std::sync::atomic::Ordering::SeqCst),
        EXPECTED_CONCURRENCY,
        "store I/O parallelism must be capped"
    );

    let mut completed_count = 0;
    for token_id in (1..EXPECTED_CONCURRENCY).rev() {
        completed_count += 1;
        control.release_and_wait(token_id, completed_count).await;
    }
    assert_eq!(
        control.started.load(std::sync::atomic::Ordering::SeqCst),
        EXPECTED_CONCURRENCY,
        "completed chunks behind token 0 must remain inside the bounded ordered window"
    );
    completed_count += 1;
    control.release_and_wait(0, completed_count).await;
    control.wait_for_started(NUM_TOKENS).await;
    for token_id in (EXPECTED_CONCURRENCY..NUM_TOKENS).rev() {
        completed_count += 1;
        control.release_and_wait(token_id, completed_count).await;
    }

    let (merged, chunk_count) = tokio::time::timeout(std::time::Duration::from_secs(5), merge_task)
        .await
        .expect("timed out waiting for controlled V1 position merge")
        .unwrap()
        .unwrap();
    assert_eq!(chunk_count, NUM_TOKENS);
    assert!(
        control.max_active.load(std::sync::atomic::Ordering::SeqCst) <= EXPECTED_CONCURRENCY,
        "singleton posting reads exceeded the bounded concurrency cap"
    );
    assert_eq!(
        *control
            .completed
            .lock()
            .expect("controlled posting completion lock poisoned"),
        vec![7, 6, 5, 4, 3, 2, 1, 0, 9, 8],
        "test reads must complete out of token order"
    );
    assert_eq!(merged.posting_lists.len(), NUM_TOKENS);
    for (token_id, posting) in merged.posting_lists.iter().enumerate() {
        let entries = posting.iter().collect::<Vec<_>>();
        assert_eq!(entries.len(), NUM_DOCS);
        for (expected_doc_id, (doc_id, frequency, positions)) in entries.into_iter().enumerate() {
            assert_eq!(frequency, 2);
            assert_eq!(doc_id as usize, expected_doc_id);
            assert_eq!(
                positions.unwrap(),
                vec![token_id as u32, token_id as u32 + 100],
                "posting positions must remain aligned with token order"
            );
        }
    }

    let error_control = Arc::new(ControlledPostingReads::new(NUM_TOKENS));
    let error_store: Arc<dyn IndexStore> = Arc::new(ControlledMergeStore {
        inner: source_store,
        posting_file: posting_file_path(0),
        io_parallelism: 64,
        control: error_control.clone(),
        fail_token: Some(1),
    });
    let error_partition = InvertedPartition::load(
        error_store,
        0,
        None,
        &LanceCache::no_cache(),
        TokenSetFormat::default(),
    )
    .await
    .unwrap();
    let error_task = tokio::spawn(async move {
        error_partition
            .into_builder_with_chunk_limits(NUM_TOKENS, u64::MAX)
            .await
    });
    error_control.wait_for_started(EXPECTED_CONCURRENCY).await;
    for (completed_count, token_id) in (0..EXPECTED_CONCURRENCY).rev().enumerate() {
        error_control
            .release_and_wait(token_id, completed_count + 1)
            .await;
    }
    let error = tokio::time::timeout(std::time::Duration::from_secs(5), error_task)
        .await
        .expect("timed out waiting for injected V1 position merge failure")
        .unwrap()
        .unwrap_err();
    assert!(matches!(error, Error::IO { .. }));
    assert!(
        error
            .to_string()
            .contains("injected singleton posting read failure for token 1")
    );
}

#[tokio::test]
async fn test_chunk_posting_mode_controls_buffer_sharing() {
    const NUM_TOKENS: usize = 2;
    const NUM_DOCS: usize = 257;

    let source_dir = TempObjDir::default();
    let source_store = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        source_dir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));
    let mut source = InnerBuilder::new_with_format_version_and_block_size(
        0,
        false,
        TokenSetFormat::default(),
        InvertedListFormatVersion::V3,
        MAX_POSTING_BLOCK_SIZE,
    );
    for token_id in 0..NUM_TOKENS {
        source.tokens.get_or_add(&format!("token_{token_id}"));
        let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size(
            false,
            PostingTailCodec::VarintDelta,
            MAX_POSTING_BLOCK_SIZE,
        );
        for doc_id in 0..NUM_DOCS {
            posting.add(doc_id as u32, PositionRecorder::Count(1));
        }
        source.posting_lists.push(posting);
    }
    for doc_id in 0..NUM_DOCS {
        source.docs.append(1_000 + doc_id as u64, NUM_TOKENS as u32);
    }
    source.write(source_store.as_ref()).await.unwrap();

    let source_partition = InvertedPartition::load(
        source_store,
        0,
        None,
        &LanceCache::no_cache(),
        TokenSetFormat::default(),
    )
    .await
    .unwrap();
    let merge_postings = source_partition
        .inverted_list
        .build_chunk_postings_for_test(0, NUM_TOKENS, false, ChunkPostingMode::Merge)
        .await
        .unwrap();
    let prewarm_postings = source_partition
        .inverted_list
        .build_chunk_postings_for_test(0, NUM_TOKENS, false, ChunkPostingMode::Prewarm)
        .await
        .unwrap();

    let [
        PostingList::Compressed(merge_first),
        PostingList::Compressed(merge_second),
    ] = merge_postings.as_slice()
    else {
        panic!("expected two compressed merge posting lists");
    };
    assert!(
        merge_first
            .blocks
            .values()
            .ptr_eq(merge_second.blocks.values()),
        "merge posting lists from one chunk should share backing storage"
    );

    let [
        PostingList::Compressed(prewarm_first),
        PostingList::Compressed(prewarm_second),
    ] = prewarm_postings.as_slice()
    else {
        panic!("expected two compressed prewarm posting lists");
    };
    assert!(
        !prewarm_first
            .blocks
            .values()
            .ptr_eq(prewarm_second.blocks.values()),
        "prewarm posting lists should own compact backing storage"
    );
}

#[tokio::test]
async fn test_posting_builder_remap() {
    let posting_tail_codec = PostingTailCodec::Fixed32;
    let mut builder = PostingListBuilder::new_with_posting_tail_codec(false, posting_tail_codec);
    let n = BLOCK_SIZE + 3;
    for i in 0..n {
        builder.add(i as u32, PositionRecorder::Count(1));
    }
    let removed = vec![5, 7];
    builder.remap(&removed);

    let mut expected = PostingListBuilder::new_with_posting_tail_codec(false, posting_tail_codec);
    for i in 0..n - removed.len() {
        expected.add(i as u32, PositionRecorder::Count(1));
    }
    let expected_entries = expected.iter().collect::<Vec<_>>();
    let actual_entries = builder.iter().collect::<Vec<_>>();
    assert_eq!(actual_entries, expected_entries);

    // BLOCK_SIZE + 3 elements should be reduced to BLOCK_SIZE + 1,
    // there are still 2 blocks.
    let batch = builder.to_batch(vec![1.0, 2.0]).unwrap();
    let (doc_ids, freqs) = decompress_posting_list_with_tail_codec(
        (n - removed.len()) as u32,
        batch[POSTING_COL]
            .as_list::<i32>()
            .value(0)
            .as_binary::<i64>(),
        posting_tail_codec,
    )
    .unwrap();
    assert!(
        doc_ids
            .iter()
            .zip(expected_entries.iter().map(|(doc_id, _, _)| doc_id))
            .all(|(a, b)| a == b)
    );
    assert!(
        freqs
            .iter()
            .zip(expected_entries.iter().map(|(_, freq, _)| freq))
            .all(|(a, b)| a == b)
    );
}

/// Doc id gaps and frequencies that exercise one- to five-byte varints in the
/// builder's tail encoding.
fn varied_postings(len: usize) -> (Vec<u32>, Vec<u32>) {
    let gaps = [1_u32, 7, 130, 20_000, 3_000_000];
    let frequencies = [1_u32, 2, 127, 128, 70_000];
    let mut doc_id = 5u32;
    let mut doc_ids = Vec::with_capacity(len);
    for index in 0..len {
        doc_ids.push(doc_id);
        doc_id += if index == 1 {
            300_000_000
        } else {
            gaps[index % gaps.len()]
        };
    }
    let freqs = (0..len)
        .map(|index| frequencies[index % frequencies.len()])
        .collect();
    (doc_ids, freqs)
}

#[rstest::rstest]
#[case::legacy_fixed32_tail(PostingTailCodec::Fixed32, LEGACY_BLOCK_SIZE)]
#[case::varint_tail_128(PostingTailCodec::VarintDelta, LEGACY_BLOCK_SIZE)]
#[case::varint_tail_256(PostingTailCodec::VarintDelta, 256)]
fn test_posting_builder_matches_reference_encoding(
    #[case] posting_tail_codec: PostingTailCodec,
    #[case] block_size: usize,
) {
    for len in [
        1,
        3,
        block_size - 1,
        block_size,
        block_size + 1,
        3 * block_size + 5,
    ] {
        let (doc_ids, frequencies) = varied_postings(len);
        let mut builder = PostingListBuilder::new_with_posting_tail_codec_and_block_size(
            false,
            posting_tail_codec,
            block_size,
        );
        for (&doc_id, &frequency) in doc_ids.iter().zip(&frequencies) {
            builder.add(doc_id, PositionRecorder::Count(frequency));
        }
        assert_eq!(builder.len(), doc_ids.len(), "len {len}");
        let expected_entries = doc_ids
            .iter()
            .zip(&frequencies)
            .map(|(&doc_id, &frequency)| (doc_id, frequency, None))
            .collect::<Vec<_>>();
        assert_eq!(builder.iter().collect::<Vec<_>>(), expected_entries);

        let block_max_scores = (0..doc_ids.len().div_ceil(block_size))
            .map(|block| block as f32 + 0.5)
            .collect::<Vec<_>>();
        let expected = compress_posting_list_with_tail_codec_and_block_size(
            doc_ids.len(),
            doc_ids.iter(),
            frequencies.iter(),
            block_max_scores.iter().copied(),
            posting_tail_codec,
            block_size,
        )
        .unwrap();
        let batch = builder.to_batch(block_max_scores).unwrap();
        let actual = batch[POSTING_COL].as_list::<i32>().value(0);
        assert_eq!(actual.as_binary::<i64>(), &expected, "len {len}");
    }
}

#[test]
fn test_posting_builder_short_lists_stay_inline() {
    let mut builder = PostingListBuilder::new_with_block_size(false, 256);
    // Most tokens of a corpus occur in only a few documents; they must not
    // pay for a heap allocation.
    for doc_id in [9_000, 90_000, 900_000] {
        builder.add(doc_id, PositionRecorder::Count(1));
    }
    assert_eq!(builder.size(), 0);

    for doc_id in 900_001..900_100 {
        builder.add(doc_id, PositionRecorder::Count(3));
    }
    assert!(builder.size() as usize >= builder.tail.len());
    assert_eq!(builder.len(), 102);
}

#[rstest::rstest]
#[case::added(false)]
#[case::streamed(true)]
fn test_posting_builder_short_positional_lists_skip_block_storage(#[case] streamed: bool) {
    let mut builder = PostingListBuilder::new_with_block_size(true, 256);
    if streamed {
        builder.add_occurrence(9_000, 4).unwrap();
        builder.finish_open_doc(9_000).unwrap();
    } else {
        builder.add(9_000, PositionRecorder::Position(smallvec::smallvec![4]));
    }
    // Every token of a positional build needs position state from its first
    // posting, but only lists that fill a block need block storage. One
    // posting costs the position state box plus its first position buffer.
    assert!(builder.size() <= 96, "{}", builder.size());
    assert_eq!(
        builder.iter().collect::<Vec<_>>(),
        vec![(9_000_u32, 1_u32, Some(vec![4_u32]))]
    );
}

#[test]
fn test_posting_builder_flush_releases_tail_position_capacity() {
    let mut builder = PostingListBuilder::new(true);
    let positions = smallvec::SmallVec::<[u32; 2]>::from_vec((0..1024).collect());
    for doc_id in 0..BLOCK_SIZE as u32 {
        builder.add(doc_id, PositionRecorder::Position(positions.clone()));
    }

    let overflow = builder.overflow.as_deref().unwrap();
    assert_eq!(overflow.tail_positions.size(), 0);
    assert!(builder.tail.is_empty());
}

#[test]
fn test_posting_builder_streamed_positions_roundtrip() {
    let mut builder = PostingListBuilder::new(true);
    assert!(builder.add_occurrence(0, 1).unwrap());
    assert!(!builder.add_occurrence(0, 4).unwrap());
    assert!(!builder.add_occurrence(0, 9).unwrap());
    builder.finish_open_doc(0).unwrap();

    assert!(builder.add_occurrence(2, 3).unwrap());
    builder.finish_open_doc(2).unwrap();

    let entries = builder.iter().collect::<Vec<_>>();
    assert_eq!(
        entries,
        vec![
            (0_u32, 3_u32, Some(vec![1_u32, 4_u32, 9_u32])),
            (2_u32, 1_u32, Some(vec![3_u32])),
        ]
    );

    // Streaming across block boundaries must build the same posting list as
    // adding whole documents.
    let mut streamed = PostingListBuilder::new(true);
    let mut added = PostingListBuilder::new(true);
    let mut expected = Vec::new();
    for doc_id in (0..(2 * BLOCK_SIZE + 3) as u32).map(|doc| doc * 3 + 1) {
        let positions = (0..doc_id % 4 + 1)
            .map(|index| index * 5 + doc_id % 7)
            .collect::<Vec<_>>();
        for &position in &positions {
            streamed.add_occurrence(doc_id, position).unwrap();
        }
        streamed.finish_open_doc(doc_id).unwrap();
        added.add(doc_id, PositionRecorder::Position(positions.clone().into()));
        expected.push((doc_id, positions.len() as u32, Some(positions)));
    }
    assert_eq!(streamed.len(), expected.len());
    assert_eq!(streamed.iter().collect::<Vec<_>>(), expected);
    let block_max_scores = vec![1.0; expected.len().div_ceil(BLOCK_SIZE)];
    assert_eq!(
        streamed.to_batch(block_max_scores.clone()).unwrap(),
        added.to_batch(block_max_scores).unwrap()
    );
}

#[test]
fn test_shared_position_stream_clone_shares_block_offsets() {
    let stream = SharedPositionStream::new(
        PositionStreamCodec::PackedDelta,
        vec![0_u32, 4, 11],
        bytes::Bytes::from_static(b"shared position bytes"),
    );
    let original_offsets = stream.block_offsets().as_ptr();

    let cloned = stream.clone();

    assert_eq!(cloned.block_offsets(), stream.block_offsets());
    assert_eq!(cloned.block_offsets().as_ptr(), original_offsets);
}

#[test]
fn test_posting_builder_roundtrip_shared_positions() {
    let entries = vec![
        (0_u32, vec![1_u32, 5]),
        (2, vec![0, 4, 9]),
        (4, vec![7]),
        (8, vec![3, 10]),
        (13, vec![2, 11, 30]),
    ];
    let mut builder =
        PostingListBuilder::new_with_posting_tail_codec(true, PostingTailCodec::VarintDelta);
    for (doc_id, positions) in &entries {
        builder.add(
            *doc_id,
            PositionRecorder::Position(positions.clone().into()),
        );
    }

    let batch = builder.to_batch(vec![1.0]).unwrap();
    assert!(batch.column_by_name(COMPRESSED_POSITION_COL).is_some());
    assert!(batch.column_by_name(POSITION_COL).is_none());
    assert_eq!(
        batch.schema_ref().metadata().get(POSTING_TAIL_CODEC_KEY),
        Some(&PostingTailCodec::VarintDelta.as_str().to_owned())
    );
    assert_eq!(
        batch.schema_ref().metadata().get(POSITIONS_LAYOUT_KEY),
        Some(&POSITIONS_LAYOUT_SHARED_STREAM_V2.to_owned())
    );
    assert_eq!(
        batch.schema_ref().metadata().get(POSITIONS_CODEC_KEY),
        Some(&PositionStreamCodec::PackedDelta.as_str().to_owned())
    );

    let posting = PostingList::from_batch(&batch, Some(1.0), Some(entries.len() as u32)).unwrap();
    let actual = posting
        .iter()
        .map(|(doc_id, freq, positions)| {
            (doc_id as u32, freq, positions.unwrap().collect::<Vec<_>>())
        })
        .collect::<Vec<_>>();
    let expected = entries
        .iter()
        .map(|(doc_id, positions)| (*doc_id, positions.len() as u32, positions.clone()))
        .collect::<Vec<_>>();
    assert_eq!(actual, expected);
}

#[test]
fn test_posting_builder_roundtrip_legacy_positions() {
    let entries = vec![(0_u32, vec![1_u32, 5]), (2, vec![0, 4, 9]), (4, vec![7])];
    let mut builder =
        PostingListBuilder::new_with_posting_tail_codec(true, PostingTailCodec::Fixed32);
    for (doc_id, positions) in &entries {
        builder.add(
            *doc_id,
            PositionRecorder::Position(positions.clone().into()),
        );
    }

    let batch = builder.to_batch(vec![1.0]).unwrap();
    assert!(batch.column_by_name(POSITION_COL).is_some());
    assert!(batch.column_by_name(COMPRESSED_POSITION_COL).is_none());
    assert_eq!(
        batch.schema_ref().metadata().get(POSTING_TAIL_CODEC_KEY),
        None
    );
    assert_eq!(
        batch.schema_ref().metadata().get(POSITIONS_LAYOUT_KEY),
        None
    );
    assert_eq!(batch.schema_ref().metadata().get(POSITIONS_CODEC_KEY), None);

    let posting = PostingList::from_batch(&batch, Some(1.0), Some(entries.len() as u32)).unwrap();
    let actual = posting
        .iter()
        .map(|(doc_id, freq, positions)| {
            (doc_id as u32, freq, positions.unwrap().collect::<Vec<_>>())
        })
        .collect::<Vec<_>>();
    let expected = entries
        .iter()
        .map(|(doc_id, positions)| (*doc_id, positions.len() as u32, positions.clone()))
        .collect::<Vec<_>>();
    assert_eq!(actual, expected);
}

#[test]
fn test_resolve_fts_format_version_defaults_to_v2() {
    assert_eq!(
        resolve_fts_format_version(None).unwrap(),
        InvertedListFormatVersion::V2
    );
    assert_eq!(
        resolve_fts_format_version(Some("2")).unwrap(),
        InvertedListFormatVersion::V2
    );
    assert_eq!(
        resolve_fts_format_version(Some("3")).unwrap(),
        InvertedListFormatVersion::V3
    );
    assert!(resolve_fts_format_version(Some("4")).is_err());
}

#[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]
fn test_legacy_compressed_positions_still_readable() {
    let doc_ids = [1_u32, 3_u32];
    let frequencies = [2_u32, 3_u32];
    let posting = compress_posting_list_with_tail_codec(
        doc_ids.len(),
        doc_ids.iter(),
        frequencies.iter(),
        std::iter::once(1.0_f32),
        PostingTailCodec::Fixed32,
    )
    .unwrap();

    let mut posting_builder = ListBuilder::new(LargeBinaryBuilder::new());
    for idx in 0..posting.len() {
        posting_builder.values().append_value(posting.value(idx));
    }
    posting_builder.append(true);

    let mut positions_builder = ListBuilder::new(ListBuilder::new(LargeBinaryBuilder::new()));
    for positions in [vec![1_u32, 5_u32], vec![0_u32, 4_u32, 9_u32]] {
        let compressed = compress_positions(&positions).unwrap();
        let doc_builder = positions_builder.values();
        for idx in 0..compressed.len() {
            doc_builder.values().append_value(compressed.value(idx));
        }
        doc_builder.append(true);
    }
    positions_builder.append(true);

    let schema = Arc::new(Schema::new(vec![
        Field::new(
            POSTING_COL,
            DataType::List(Arc::new(Field::new("item", DataType::LargeBinary, true))),
            false,
        ),
        Field::new(MAX_SCORE_COL, DataType::Float32, false),
        Field::new(LENGTH_COL, DataType::UInt32, false),
        Field::new(
            POSITION_COL,
            DataType::List(Arc::new(Field::new(
                "item",
                DataType::List(Arc::new(Field::new("item", DataType::LargeBinary, true))),
                true,
            ))),
            false,
        ),
    ]));
    let batch = RecordBatch::try_new(
        schema,
        vec![
            Arc::new(posting_builder.finish()) as ArrayRef,
            Arc::new(Float32Array::from(vec![1.0])) as ArrayRef,
            Arc::new(UInt32Array::from(vec![doc_ids.len() as u32])) as ArrayRef,
            Arc::new(positions_builder.finish()) as ArrayRef,
        ],
    )
    .unwrap();

    let posting = PostingList::from_batch(&batch, Some(1.0), Some(doc_ids.len() as u32)).unwrap();
    let actual = posting
        .iter()
        .map(|(doc_id, freq, positions)| {
            (doc_id as u32, freq, positions.unwrap().collect::<Vec<_>>())
        })
        .collect::<Vec<_>>();
    assert_eq!(actual, vec![(1, 2, vec![1, 5]), (3, 3, vec![0, 4, 9]),]);
}

#[test]
fn test_shared_stream_v2_without_codec_still_readable() {
    let doc_ids = [1_u32, 3_u32];
    let frequencies = [2_u32, 3_u32];
    let posting = compress_posting_list_with_tail_codec(
        doc_ids.len(),
        doc_ids.iter(),
        frequencies.iter(),
        std::iter::once(1.0_f32),
        PostingTailCodec::Fixed32,
    )
    .unwrap();

    let mut posting_builder = ListBuilder::new(LargeBinaryBuilder::new());
    for idx in 0..posting.len() {
        posting_builder.values().append_value(posting.value(idx));
    }
    posting_builder.append(true);

    let positions = vec![1_u32, 5_u32, 0_u32, 4_u32, 9_u32];
    let mut encoded_positions = Vec::new();
    encode_position_stream_block_into(
        &positions,
        &frequencies,
        PositionStreamCodec::VarintDocDelta,
        &mut encoded_positions,
    )
    .unwrap();

    let mut position_offsets = ListBuilder::new(UInt32Builder::new());
    position_offsets.values().append_value(0);
    position_offsets.append(true);

    let schema = Arc::new(Schema::new_with_metadata(
        vec![
            Field::new(
                POSTING_COL,
                DataType::List(Arc::new(Field::new("item", DataType::LargeBinary, true))),
                false,
            ),
            Field::new(MAX_SCORE_COL, DataType::Float32, false),
            Field::new(LENGTH_COL, DataType::UInt32, false),
            Field::new(COMPRESSED_POSITION_COL, DataType::LargeBinary, false),
            Field::new(
                POSITION_BLOCK_OFFSET_COL,
                DataType::List(Arc::new(Field::new("item", DataType::UInt32, true))),
                false,
            ),
        ],
        HashMap::from([(
            POSITIONS_LAYOUT_KEY.to_owned(),
            POSITIONS_LAYOUT_SHARED_STREAM_V2.to_owned(),
        )]),
    ));
    let batch = RecordBatch::try_new(
        schema,
        vec![
            Arc::new(posting_builder.finish()) as ArrayRef,
            Arc::new(Float32Array::from(vec![1.0])) as ArrayRef,
            Arc::new(UInt32Array::from(vec![doc_ids.len() as u32])) as ArrayRef,
            Arc::new(arrow_array::LargeBinaryArray::from(vec![Some(
                encoded_positions.as_slice(),
            )])) as ArrayRef,
            Arc::new(position_offsets.finish()) as ArrayRef,
        ],
    )
    .unwrap();

    let posting = PostingList::from_batch(&batch, Some(1.0), Some(doc_ids.len() as u32)).unwrap();
    let actual = posting
        .iter()
        .map(|(doc_id, freq, positions)| {
            (doc_id as u32, freq, positions.unwrap().collect::<Vec<_>>())
        })
        .collect::<Vec<_>>();
    assert_eq!(actual, vec![(1, 2, vec![1, 5]), (3, 3, vec![0, 4, 9]),]);
}

#[test]
fn test_shared_position_stream_is_smaller_for_sparse_positions() {
    let mut builder =
        PostingListBuilder::new_with_posting_tail_codec(true, PostingTailCodec::VarintDelta);
    let mut legacy_positions = Vec::with_capacity(BLOCK_SIZE * 4);
    for doc_id in 0..(BLOCK_SIZE * 4) as u32 {
        let mut positions = vec![doc_id * 3 + 1];
        if doc_id % 8 == 0 {
            positions.push(doc_id * 3 + 2);
        }
        builder.add(doc_id, PositionRecorder::Position(positions.clone().into()));
        legacy_positions.push(positions);
    }

    let batch = builder.to_batch(vec![1.0; 4]).unwrap();
    let shared_positions_size = batch[COMPRESSED_POSITION_COL].get_buffer_memory_size()
        + batch[POSITION_BLOCK_OFFSET_COL].get_buffer_memory_size();

    let mut positions_builder = ListBuilder::new(ListBuilder::new(LargeBinaryBuilder::new()));
    for positions in legacy_positions {
        let compressed = compress_positions(&positions).unwrap();
        let doc_builder = positions_builder.values();
        for idx in 0..compressed.len() {
            doc_builder.values().append_value(compressed.value(idx));
        }
        doc_builder.append(true);
    }
    positions_builder.append(true);
    let legacy_positions_size = positions_builder.finish().get_buffer_memory_size();

    assert!(
        shared_positions_size < legacy_positions_size,
        "expected shared position stream to be smaller than legacy per-doc storage, shared={shared_positions_size}, legacy={legacy_positions_size}",
    );
}

#[test]
fn test_posting_list_batch_matches_docset_scoring() {
    let mut docs = DocSet::default();
    let num_docs = BLOCK_SIZE + 3;
    for doc_id in 0..num_docs as u32 {
        docs.append(doc_id as u64, doc_id % 7 + 1);
    }

    let doc_ids = (0..num_docs as u32).collect::<Vec<_>>();
    let freqs = doc_ids
        .iter()
        .map(|doc_id| doc_id % 5 + 1)
        .collect::<Vec<_>>();

    let mut builder_scores = PostingListBuilder::new(false);
    let mut builder_docs = PostingListBuilder::new(false);
    for (&doc_id, &freq) in doc_ids.iter().zip(freqs.iter()) {
        builder_scores.add(doc_id, PositionRecorder::Count(freq));
        builder_docs.add(doc_id, PositionRecorder::Count(freq));
    }

    let block_max_scores = docs.calculate_block_max_scores(doc_ids.iter(), freqs.iter());
    let batch_scores = builder_scores.to_batch(block_max_scores).unwrap();
    let batch_docs = builder_docs
        .to_batch_with_docs(&docs, inverted_list_schema(false))
        .unwrap();

    let scores_posting = batch_scores[POSTING_COL].as_list::<i32>().value(0);
    let scores_posting = scores_posting.as_binary::<i64>();
    let docs_posting = batch_docs[POSTING_COL].as_list::<i32>().value(0);
    let docs_posting = docs_posting.as_binary::<i64>();
    assert_eq!(scores_posting, docs_posting);

    let score_left = batch_scores[MAX_SCORE_COL]
        .as_primitive::<Float32Type>()
        .value(0);
    let score_right = batch_docs[MAX_SCORE_COL]
        .as_primitive::<Float32Type>()
        .value(0);
    assert!((score_left - score_right).abs() < 1e-6);

    let len_left = batch_scores[LENGTH_COL]
        .as_primitive::<UInt32Type>()
        .value(0);
    let len_right = batch_docs[LENGTH_COL].as_primitive::<UInt32Type>().value(0);
    assert_eq!(len_left, len_right);
}

#[tokio::test]
async fn test_remap_to_empty_posting_list() {
    let tmpdir = TempObjDir::default();
    let store = Arc::new(LanceIndexStore::new(
        ObjectStore::local().into(),
        tmpdir.clone(),
        Arc::new(LanceCache::no_cache()),
    ));

    let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default());

    // index of docs:
    // 0: lance
    // 1: lake lake
    // 2: lake lake lake
    builder.tokens.get_or_add("lance");
    builder.tokens.get_or_add("lake");
    builder.posting_lists.push(PostingListBuilder::new(false));
    builder.posting_lists.push(PostingListBuilder::new(false));
    builder.posting_lists[0].add(0, PositionRecorder::Count(1));
    builder.posting_lists[1].add(1, PositionRecorder::Count(2));
    builder.posting_lists[1].add(2, PositionRecorder::Count(3));
    builder.docs.append(0, 1);
    builder.docs.append(1, 1);
    builder.docs.append(2, 1);
    builder.write(store.as_ref()).await.unwrap();

    let index = InvertedPartition::load(
        store.clone(),
        0,
        None,
        &LanceCache::no_cache(),
        TokenSetFormat::default(),
    )
    .await
    .unwrap();
    let mut builder = index.into_builder().await.unwrap();

    let mapping = HashMap::from([(0, None), (2, Some(3))]);
    builder.remap(&RowAddrRemap::direct(mapping)).await.unwrap();

    // after remap, the doc 0 is removed, and the doc 2 is updated to 3
    assert_eq!(builder.tokens.len(), 1);
    assert_eq!(builder.tokens.get("lake"), Some(0));
    assert_eq!(builder.posting_lists.len(), 1);
    assert_eq!(builder.posting_lists[0].len(), 2);
    assert_eq!(builder.docs.len(), 2);
    assert_eq!(builder.docs.row_id(0), 1);
    assert_eq!(builder.docs.row_id(1), 3);

    builder.write(store.as_ref()).await.unwrap();

    // remap to delete all docs
    let mapping = HashMap::from([(1, None), (3, None)]);
    builder.remap(&RowAddrRemap::direct(mapping)).await.unwrap();

    assert_eq!(builder.tokens.len(), 0);
    assert_eq!(builder.posting_lists.len(), 0);
    assert_eq!(builder.docs.len(), 0);

    builder.write(store.as_ref()).await.unwrap();
}

#[test]
fn test_docset_remap_preserves_element_coordinates() {
    let mut docs = DocSet::default();
    docs.append_with_doc_index(10, 2, &[0]).unwrap();
    docs.append_with_doc_index(10, 3, &[3]).unwrap();
    docs.append_with_doc_index(11, 1, &[1]).unwrap();

    let removed = docs.remap(&RowAddrRemap::direct(HashMap::from([
        (10, Some(20)),
        (11, None),
    ])));

    assert_eq!(removed, vec![2]);
    assert_eq!(docs.len(), 2);
    assert_eq!(docs.row_id(0), 20);
    assert_eq!(docs.row_id(1), 20);
    assert_eq!(docs.doc_index(0), vec![0]);
    assert_eq!(docs.doc_index(1), vec![3]);
}

/// `contains_each` has to agree with a full decode for probes in every block,
/// before the first doc and past the last, for both tail codecs and both block
/// sizes, and for a list that is only a tail or has no tail at all.
#[rstest::rstest]
#[case::fixed32_128(PostingTailCodec::Fixed32, 128, 2 * 128 + 128 / 3)]
#[case::varint_256(PostingTailCodec::VarintDelta, 256, 2 * 256 + 256 / 3)]
#[case::tail_only(PostingTailCodec::VarintDelta, 128, 128 / 3)]
#[case::no_tail(PostingTailCodec::Fixed32, 128, 2 * 128)]
fn test_compressed_posting_contains_each_matches_full_decode(
    #[case] tail_codec: PostingTailCodec,
    #[case] block_size: usize,
    #[case] length: usize,
) {
    // Every third doc id from 5 on.
    let doc_ids: Vec<u32> = (0..length as u32).map(|i| 5 + 3 * i).collect();
    let freqs = vec![1u32; length];
    let blocks =
        crate::scalar::inverted::encoding::compress_posting_list_with_tail_codec_and_block_size(
            length,
            doc_ids.iter(),
            freqs.iter(),
            std::iter::repeat(0.0f32),
            tail_codec,
            block_size,
        )
        .unwrap();
    let posting = CompressedPostingList::new(
        blocks,
        1.0,
        length as u32,
        tail_codec,
        block_size,
        None,
        None,
    );

    let last = *doc_ids.last().unwrap();
    let mut probes: Vec<u32> = vec![0, 4, 5, 6, 8, last, last + 1, last + 3];
    for block_start in (0..length).step_by(block_size) {
        let first = doc_ids[block_start];
        probes.extend([first - 1, first, first + 3]);
    }
    probes.sort_unstable();
    probes.dedup();
    let expected: Vec<bool> = probes
        .iter()
        .map(|probe| doc_ids.binary_search(probe).is_ok())
        .collect();
    assert!(expected.iter().any(|found| *found) && expected.iter().any(|found| !*found));
    assert_eq!(posting.contains_each(&probes), expected);
    assert!(posting.contains_each(&[]).is_empty());
}

fn dictionary_test_tokens() -> Vec<String> {
    let mut tokens = ["b", "a", "ab", "abc", "a b", "z", "é", "日本語"]
        .into_iter()
        .map(str::to_owned)
        .collect::<Vec<_>>();
    tokens.extend((0..1000).map(|index| format!("tok{index:04}")));
    tokens
}

#[rstest::rstest]
#[case::arrow(TokenSetFormat::Arrow)]
#[case::fst(TokenSetFormat::Fst)]
fn test_token_dictionary_writes_same_batch_as_token_set(#[case] format: TokenSetFormat) {
    let mut expected = TokenSet::default();
    let mut dictionary = TokenDictionary::default();
    for token in dictionary_test_tokens() {
        let token_id = dictionary.get_or_add(&token);
        assert_eq!(dictionary.get_or_add(&token), token_id);
        assert_eq!(expected.add(token), token_id);
    }
    assert_eq!(dictionary.len(), expected.len());
    assert_eq!(
        dictionary.to_batch(format).unwrap(),
        expected.to_batch(format).unwrap()
    );
}

#[rstest::rstest]
#[case::arrow(TokenSetFormat::Arrow)]
#[case::fst(TokenSetFormat::Fst)]
#[tokio::test]
async fn test_token_dictionary_keeps_loaded_token_ids(#[case] format: TokenSetFormat) {
    let tmpdir = TempObjDir::default();
    let store = LanceIndexStore::new(
        ObjectStore::local().into(),
        tmpdir.clone(),
        Arc::new(LanceCache::no_cache()),
    );
    let mut written = TokenSet::default();
    for token in dictionary_test_tokens() {
        written.add(token);
    }
    let batch = written.clone().to_batch(format).unwrap();
    let mut writer = store
        .new_index_file("tokens.lance", batch.schema())
        .await
        .unwrap();
    writer.write_record_batch(batch).await.unwrap();
    writer.finish().await.unwrap();
    let reader = store.open_index_file("tokens.lance").await.unwrap();
    let loaded = TokenSet::load(reader, format).await.unwrap();

    let mut dictionary = TokenDictionary::try_from_token_set(loaded).unwrap();
    assert_eq!(dictionary.len(), written.len());
    for token in dictionary_test_tokens() {
        assert_eq!(dictionary.get(&token), written.get(&token), "{token}");
    }
    assert_eq!(dictionary.get_or_add("new token"), written.len() as u32);
}

#[test]
fn test_token_dictionary_rejects_sparse_token_ids() {
    let mut tokens = TokenSet::default();
    tokens.tokens = TokenMap::HashMap(HashMap::from([("a".to_owned(), 0), ("b".to_owned(), 5)]));
    let err = TokenDictionary::try_from_token_set(tokens).unwrap_err();
    assert!(err.to_string().contains("distinct ids below 2"), "{err}");
}