lance 14.0.1

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

use std::collections::{HashMap, hash_map::Entry};
use std::sync::Arc;

use futures::{FutureExt, TryStreamExt};
use lance_core::{Error, Result};
use lance_file::reader::FileReaderOptions;
use lance_file::versions::data_file_columns;
use lance_index::{
    INDEX_FILE_NAME, IndexType,
    frag_reuse::CompactFragReuseIndex,
    metrics::NoOpMetricsCollector,
    optimize::OptimizeOptions,
    progress::{IndexBuildProgress, NoopIndexBuildProgress},
    scalar::{
        BuiltinIndexType, CreatedIndex, OldIndexDataFilter, ScalarIndex, index_files_to_table,
        inverted::InvertedIndex,
        json::json_wrapped_target_type,
        lance_format::LanceIndexStore,
        seed::{FragmentSeed, SEED_META_KEY_PREFIX},
        table_files_to_index,
    },
};
use lance_io::scheduler::{ScanScheduler, SchedulerConfig};
use lance_io::utils::CachedFileSize;
use lance_select::{RowAddrTreeMap, RowSetOps};
use lance_table::format::{Fragment, IndexMetadata};
use prost::Message;
use roaring::RoaringBitmap;
use uuid::Uuid;

use super::vector::ivf::{
    SteadyStateRebalance, VectorSegmentCompatibility, index_type_for_segmented_optimize,
    optimize_vector_indices, select_steady_state_rebalance, vector_segment_compatibility,
};
use super::vector::{LogicalVectorIndex, VectorIndexParams, details, fresh_vector_segment_params};
use super::{CreateIndexBuilder, DatasetIndexExt, DatasetIndexInternalExt};
use crate::dataset::Dataset;
use crate::dataset::index::LanceIndexStoreExt;
use crate::dataset::rowids::load_row_id_sequences;
use crate::index::scalar::{
    IndexDetails, fetch_index_details, load_fts_training_data, load_training_data,
};
use crate::index::vector_index_details_default;

#[derive(Debug, Clone)]
pub struct IndexMergeResults<'a> {
    pub new_uuid: Uuid,
    pub removed_indices: Vec<&'a IndexMetadata>,
    pub new_fragment_bitmap: RoaringBitmap,
    pub new_dataset_version: u64,
    pub new_index_version: i32,
    pub new_index_details: prost_types::Any,
    /// List of files and their sizes for the merged index
    pub files: Vec<lance_table::format::IndexFile>,
}

impl IndexMergeResults<'_> {
    /// The new segment as manifest metadata, named and keyed like `template`
    /// (the segment it replaces or extends).
    pub fn into_metadata(self, template: &IndexMetadata) -> IndexMetadata {
        IndexMetadata {
            uuid: self.new_uuid,
            name: template.name.clone(),
            fields: template.fields.clone(),
            covering_fields: template.covering_fields.clone(),
            dataset_version: self.new_dataset_version,
            fragment_bitmap: Some(self.new_fragment_bitmap),
            index_details: Some(Arc::new(self.new_index_details)),
            index_version: self.new_index_version,
            created_at: Some(chrono::Utc::now()),
            base_id: None,
            files: Some(self.files),
        }
    }
}

async fn build_stable_row_id_filter(
    dataset: &Dataset,
    effective_old_frags: &RoaringBitmap,
) -> Result<RowAddrTreeMap> {
    // For stable row IDs we cannot derive fragment ownership from row_id bits.
    // Instead, we:
    // 1) keep only fragments still considered "effective" for the old index, and
    // 2) load their persisted row-id sequences from dataset metadata, then
    // 3) build one exact allow-list used to retain only still-valid old rows.
    let retained_frags = dataset
        .manifest
        .fragments
        .iter()
        .filter(|frag| effective_old_frags.contains(frag.id as u32))
        .cloned()
        .collect::<Vec<_>>();

    if retained_frags.is_empty() {
        return Ok(RowAddrTreeMap::new());
    }

    let row_id_sequences = load_row_id_sequences(dataset, &retained_frags)
        .try_collect::<Vec<_>>()
        .await?;

    let frag_by_id: std::collections::HashMap<u32, _> = dataset
        .get_fragments()
        .into_iter()
        .map(|f| (f.id() as u32, f))
        .collect();

    let mut row_id_maps = Vec::with_capacity(row_id_sequences.len());
    for (frag_id, seq) in &row_id_sequences {
        row_id_maps.push(live_row_ids(frag_by_id.get(frag_id), seq).await?);
    }
    let row_id_map_refs = row_id_maps.iter().collect::<Vec<_>>();

    // Merge all fragment-local row-id sets into one exact membership structure.
    Ok(<RowAddrTreeMap as RowSetOps>::union_all(&row_id_map_refs))
}

/// The fragment's live row ids: its persisted row-id sequence minus the rows
/// its deletion vector marks gone. A persisted sequence covers every row the
/// fragment ever held, so a row whose old copy was deleted (e.g. rewritten by an
/// update under the same stable row id) would otherwise be retained as a stale
/// old-index entry.
async fn live_row_ids(
    fragment: Option<&crate::dataset::fragment::FileFragment>,
    seq: &lance_table::rowids::RowIdSequence,
) -> Result<RowAddrTreeMap> {
    // Propagate a deletion-vector read failure rather than swallowing it: a
    // swallowed error would fall through to the "no deletions" branch below,
    // putting the deleted rows back into the allow-list as stale entries.
    let deletion_vector = match fragment {
        Some(f) if f.metadata().deletion_file.is_some() => f.get_deletion_vector().await?,
        _ => None,
    };
    Ok(match deletion_vector {
        Some(dv) => seq
            .iter()
            .enumerate()
            .filter(|(offset, _)| !dv.contains(*offset as u32))
            .map(|(_, row_id)| row_id)
            .collect(),
        None => RowAddrTreeMap::from(seq),
    })
}

/// Build the [`OldIndexDataFilter`] that must be applied to existing index
/// rows when their owning fragments have been pruned by compaction or
/// deletions.
///
/// `address_domain` must be `true` when the consuming index stores physical
/// row addresses rather than row ids (see
/// `ScalarIndex::results_are_row_addresses`). `build_stable_row_id_filter`
/// below produces row ids, which would be the wrong domain to filter such an
/// index's stored addresses against, so an address-domain index always falls
/// back to coarse `Fragments` filtering instead, regardless of the dataset's
/// row-id scheme.
///
/// That coarse filtering cannot tell a still-"effective" fragment's live rows
/// from ones a same-fragment update superseded (under stable row ids,
/// updating an indexed column deletes the row's old physical copy and writes
/// the new value elsewhere, but the old fragment keeps its id and stays
/// "effective"), so an address-domain index can merge in a stale posting this
/// way. Exact address-level filtering would avoid that; left as a future
/// optimization.
pub async fn build_old_data_filter(
    dataset: &Dataset,
    effective_old_frags: &RoaringBitmap,
    deleted_old_frags: &RoaringBitmap,
    address_domain: bool,
) -> Result<Option<OldIndexDataFilter>> {
    if !address_domain && dataset.manifest.uses_stable_row_ids() {
        let valid_old_row_ids = build_stable_row_id_filter(dataset, effective_old_frags).await?;
        Ok(Some(OldIndexDataFilter::RowIds(valid_old_row_ids)))
    } else {
        Ok(Some(OldIndexDataFilter::Fragments {
            to_keep: effective_old_frags.clone(),
            to_remove: deleted_old_frags.clone(),
        }))
    }
}

/// On a table with a tagged fragment reuse history, the fragments each segment
/// ANSWERS FOR right now: its coverage as derived by the tagged reader (live
/// destinations reached through the ledger plus directly covered live
/// fragments). `None` on every other table.
///
/// A tagged segment's stored bitmap is provenance, not coverage: it names the
/// fragments the segment was built from, which a rewrite may have retired,
/// while the segment's addresses are translated to live fragments when it is
/// opened. Maintenance that reads a segment through the translating loader
/// therefore sees rows in these derived fragments, and must filter and account
/// for old data in this domain rather than in `stored ∩ live` (which is empty
/// for a fully rewritten segment and would drop every translated row).
///
/// `staged` supplies the plan of segments the manifest does not list (a staged
/// build being merged, see `plan_staged_segments`); their coverage comes from
/// that plan. A segment that is neither staged nor registered in the manifest
/// is an error: "not registered" is not "covers nothing".
pub async fn tagged_segment_coverage(
    dataset: &Dataset,
    segments: &[&IndexMetadata],
    staged: Option<&crate::index::frag_reuse::StagedRemappingPlans>,
) -> Result<Option<HashMap<Uuid, RoaringBitmap>>> {
    use crate::index::frag_reuse::SegmentRemappingPlan;

    let stored = crate::index::load_all_indices(dataset).await?;
    if !stored
        .iter()
        .any(lance_table::system_index::frag_reuse::metadata::is_tagged)
    {
        return Ok(None);
    }
    let derived = dataset.load_indices().await?;
    let mut coverage = HashMap::with_capacity(segments.len());
    for segment in segments {
        if let Some(plan) = staged.and_then(|plans| plans.get(&segment.uuid)) {
            let bitmap = match plan {
                SegmentRemappingPlan::Identity => segment
                    .fragment_bitmap
                    .as_ref()
                    .map(|bitmap| bitmap & dataset.fragment_bitmap.as_ref())
                    .unwrap_or_default(),
                SegmentRemappingPlan::Translate { coverage, .. } => coverage.clone(),
                SegmentRemappingPlan::MissingCoverage(_) => {
                    return Err(Error::not_supported(format!(
                        "FRI query coverage is unavailable for staged segment {}",
                        segment.uuid
                    )));
                }
            };
            coverage.insert(segment.uuid, bitmap);
            continue;
        }
        if !stored.iter().any(|entry| entry.uuid == segment.uuid) {
            return Err(Error::invalid_input(format!(
                "segment {} is not registered in the manifest; a staged segment must be \
                 merged with its own translation plan",
                segment.uuid
            )));
        }
        // A registered segment the tagged reader excludes (no coverage it can
        // honor) has no old data reachable here; leave it empty rather than
        // fail, the loader reports the precise reason if it is actually opened.
        let bitmap = derived
            .iter()
            .find(|entry| entry.uuid == segment.uuid)
            .and_then(|entry| entry.fragment_bitmap.clone())
            .unwrap_or_default();
        coverage.insert(segment.uuid, bitmap);
    }
    Ok(Some(coverage))
}

/// Split the stored fragment coverage of `segments` into fragments still live in
/// `dataset` (`effective`) and fragments that compaction or deletion has already
/// retired (`deleted`).
pub fn split_segment_coverage<'a>(
    dataset: &Dataset,
    segments: impl IntoIterator<Item = &'a IndexMetadata>,
) -> (RoaringBitmap, RoaringBitmap) {
    let mut effective = RoaringBitmap::new();
    let mut deleted = RoaringBitmap::new();
    for segment in segments {
        if let Some(eff) = segment.effective_fragment_bitmap(&dataset.fragment_bitmap) {
            effective |= eff;
        }
        if let Some(del) = segment.deleted_fragment_bitmap(&dataset.fragment_bitmap) {
            deleted |= del;
        }
    }
    (effective, deleted)
}

pub fn fragment_reuse_affects_segments<'a>(
    frag_reuse_index: &CompactFragReuseIndex,
    segments: impl IntoIterator<Item = &'a IndexMetadata>,
) -> bool {
    segments.into_iter().any(|segment| {
        let Some(coverage) = segment.fragment_bitmap.as_ref() else {
            return false;
        };
        fragment_reuse_affects_segment(frag_reuse_index, coverage, segment.dataset_version)
    })
}

pub fn fragment_reuse_affects_segment(
    frag_reuse_index: &CompactFragReuseIndex,
    coverage: &RoaringBitmap,
    dataset_version: u64,
) -> bool {
    frag_reuse_index.details.versions.iter().any(|version| {
        version.groups.iter().any(|group| {
            if group.changed_row_addrs.is_empty() {
                return false;
            }
            let covers_old = group
                .old_frags
                .iter()
                .any(|fragment| coverage.contains(fragment.id as u32));
            let covers_new = group
                .new_frags
                .iter()
                .any(|fragment| coverage.contains(fragment.id as u32));
            // A segment built at or before `version.dataset_version` predates the
            // rewrite, even if its coverage already moved onto the new fragments.
            version.dataset_version >= dataset_version && (covers_old || covers_new)
        })
    })
}

/// Build one [`OldIndexDataFilter`] per segment and return their effective coverage.
///
/// `staged` carries the plans of segments the manifest does not list (see
/// [`tagged_segment_coverage`]); committed merges pass `None`.
///
/// See [`build_old_data_filter`] for the meaning of `address_domain`.
pub async fn build_per_segment_filters(
    dataset: &Dataset,
    segments: &[&IndexMetadata],
    staged: Option<&crate::index::frag_reuse::StagedRemappingPlans>,
    address_domain: bool,
) -> Result<(RoaringBitmap, Vec<Option<OldIndexDataFilter>>)> {
    if let Some(coverage) = tagged_segment_coverage(dataset, segments, staged).await? {
        // Translated addresses live in the derived coverage; retired source
        // fragments never appear in them, so there is nothing to remove.
        let mut effective_union = RoaringBitmap::new();
        let mut filters = Vec::with_capacity(segments.len());
        for segment in segments {
            let effective = coverage.get(&segment.uuid).cloned().ok_or_else(|| {
                Error::internal(format!(
                    "tagged coverage missing for segment {}",
                    segment.uuid
                ))
            })?;
            effective_union |= &effective;
            filters.push(
                build_old_data_filter(dataset, &effective, &RoaringBitmap::new(), address_domain)
                    .await?,
            );
        }
        return Ok((effective_union, filters));
    }
    if dataset.manifest.uses_stable_row_ids() && !address_domain {
        let mut effective_union = RoaringBitmap::new();
        let mut filters = Vec::with_capacity(segments.len());
        for segment in segments {
            let effective = segment
                .effective_fragment_bitmap(&dataset.fragment_bitmap)
                .ok_or_else(|| {
                    Error::invalid_input(format!(
                        "CreateIndex: segment {} is missing fragment coverage",
                        segment.uuid
                    ))
                })?;
            effective_union |= &effective;
            filters.push(
                build_old_data_filter(dataset, &effective, &RoaringBitmap::new(), address_domain)
                    .await?,
            );
        }
        return Ok((effective_union, filters));
    }

    let mut effective_union = RoaringBitmap::new();
    let mut filters = Vec::with_capacity(segments.len());
    for segment in segments {
        if segment.fragment_bitmap.is_none() {
            return Err(Error::invalid_input(format!(
                "CreateIndex: segment {} is missing fragment coverage",
                segment.uuid
            )));
        }
        let effective = segment
            .effective_fragment_bitmap(&dataset.fragment_bitmap)
            .unwrap_or_default();
        let deleted = segment
            .deleted_fragment_bitmap(&dataset.fragment_bitmap)
            .unwrap_or_default();
        effective_union |= &effective;
        filters.push(build_old_data_filter(dataset, &effective, &deleted, address_domain).await?);
    }
    Ok((effective_union, filters))
}

/// Attempt to read seed buffers for `index` from `fragments`' data files.
///
/// Returns `Some(vec)` only if every fragment has a seed entry whose field
/// identity can be verified; returns `None` if a seed is missing, its file
/// cannot be opened, or its physical column mapping cannot be verified.
/// Index-type-specific validation (e.g. `rows_per_zone` checks) is left to the
/// caller via [`FragmentSeed::metadata_value`].
async fn try_harvest_seeds(
    dataset: &Dataset,
    fragments: &[Fragment],
    index: &IndexMetadata,
) -> Result<Option<Vec<FragmentSeed>>> {
    if fragments.is_empty() {
        return Ok(Some(Vec::new()));
    }

    // SeedCollector writes seeds only for top-level keyed fields. A nested
    // leaf can share its name with a different, seeded top-level column.
    let Some(field_id) = index.keyed_field() else {
        return Ok(None);
    };
    let Some(field) = dataset.schema().fields.iter().find(|f| f.id == field_id) else {
        return Ok(None);
    };
    let mut seeds = Vec::with_capacity(fragments.len());
    let mut schedulers = HashMap::new();

    for fragment in fragments {
        // A column replacement can retain an older file for other columns.
        // Its seed still describes the old values, even after the manifest
        // stops mapping this column to that file.
        let Some(data_file) = fragment.files.iter().rev().find(|file| {
            file.schema(dataset.schema())
                .field_by_id(field.id)
                .is_some()
        }) else {
            return Ok(None);
        };

        let path = dataset
            .data_file_dir(data_file)?
            .join(data_file.path.as_str());
        let scheduler = match schedulers.entry(data_file.base_id) {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => {
                let store = dataset.object_store(data_file.base_id).await?;
                let config = SchedulerConfig::max_bandwidth(&store);
                entry.insert(ScanScheduler::new(store, config))
            }
        };
        let Ok(file_scheduler) = scheduler.open_file(&path, &CachedFileSize::unknown()).await
        else {
            return Ok(None);
        };

        let Ok(reader) = lance_file::reader::FileReader::try_open(
            file_scheduler,
            None,
            Default::default(),
            &dataset.metadata_cache.file_metadata_cache(&path),
            FileReaderOptions::default(),
        )
        .await
        else {
            return Ok(None);
        };

        // Renames preserve the manifest's field IDs, but seed keys retain
        // the names used when writing the file. Resolve the stored field
        // through its physical column, since field IDs can also be remapped.
        let Some((_, column_index)) = data_file
            .fields
            .iter()
            .zip(data_file.column_indices.iter())
            .find(|(id, column_index)| **id == field.id && **column_index >= 0)
        else {
            return Ok(None);
        };
        let metadata = reader.metadata();
        let file_schema = &metadata.file_schema;
        let (stored_ids, stored_columns) = data_file_columns(metadata.version, file_schema);
        // Blob footers can retain descriptor fields without physical columns.
        // Their inferred positions must not select another column's seed.
        if stored_columns.len() != metadata.column_metadatas.len() {
            return Ok(None);
        }
        let Some((stored_id, _)) = stored_ids
            .iter()
            .zip(&stored_columns)
            .find(|(_, stored_column)| *stored_column == column_index)
        else {
            return Ok(None);
        };
        let Some(stored_field) = file_schema
            .fields
            .iter()
            .find(|stored| stored.id == *stored_id && stored.data_type() == field.data_type())
        else {
            return Ok(None);
        };
        let meta_key = format!("{}{}", SEED_META_KEY_PREFIX, stored_field.name);
        let Some(meta_value) = file_schema.metadata.get(&meta_key).cloned() else {
            return Ok(None);
        };

        // The buf_index is always the portion before the first ':'.
        let Some(buf_index_str) = meta_value.split(':').next() else {
            return Ok(None);
        };
        let Ok(buf_index) = buf_index_str.parse::<u32>() else {
            return Ok(None);
        };

        let Ok(bytes) = reader.read_global_buffer(buf_index).await else {
            return Ok(None);
        };

        seeds.push(FragmentSeed {
            fragment_id: fragment.id,
            bytes,
            metadata_value: meta_value,
        });
    }

    Ok(Some(seeds))
}

async fn load_unindexed_training_data(
    dataset: &Dataset,
    field_path: &str,
    update_criteria: &lance_index::scalar::UpdateCriteria,
    unindexed: &[Fragment],
) -> Result<datafusion::execution::SendableRecordBatchStream> {
    let fragments = if update_criteria.requires_old_data {
        None
    } else {
        Some(unindexed.to_vec())
    };
    load_training_data(
        dataset,
        field_path,
        &update_criteria.data_criteria,
        fragments,
        true,
        None,
    )
    .await
}

/// Build a fresh, canonical (non-sharded) scalar index over `fragment_ids`,
/// reusing `reference_index`'s params and training criteria.
async fn rebuild_scalar_segment(
    dataset: &Dataset,
    reference_index: &Arc<dyn ScalarIndex>,
    field_path: &str,
    column_name: &str,
    uuid: Uuid,
    fragment_ids: Vec<u32>,
) -> Result<CreatedIndex> {
    let params = reference_index.derive_index_params()?;
    let update_criteria = reference_index.update_criteria();
    let training_data = load_training_data(
        dataset,
        field_path,
        &update_criteria.data_criteria,
        None,
        true,
        Some(fragment_ids),
    )
    .await?;
    super::scalar::build_scalar_index(
        dataset,
        column_name,
        uuid,
        &params,
        true,
        None,
        Some((training_data, update_criteria.data_criteria.clone())),
        Arc::new(NoopIndexBuildProgress),
    )
    .await
}

/// The index segments to rewrite in this optimize pass.
///
/// Normally the trailing `num_indices_to_merge` segments. Under stable row ids,
/// any *older* segment that still covers a fragment carrying deletions is added
/// too: an update deletes a row's old copy (leaving a deletion vector) and
/// rewrites it under the same row id, so its stale old-value postings survive
/// until that segment is rewritten and filtered. Only the segments that actually
/// cover a deleted-from fragment are pulled in -- clean segments in between are
/// left untouched -- so an edit to old data does not force a full reindex.
///
/// The deletion check is conservative (any current deletion vector on a covered
/// fragment), so a segment built after those deletions may be rewritten as a
/// harmless no-op; it never leaves a stale segment behind (PR #7359).
pub fn select_segments_to_merge<'a>(
    dataset: &Dataset,
    old_indices: &[&'a IndexMetadata],
    options: &OptimizeOptions,
) -> Vec<&'a IndexMetadata> {
    let num_to_merge = options
        .num_indices_to_merge
        .unwrap_or(1)
        .min(old_indices.len());
    let tail_start = old_indices.len() - num_to_merge;

    // Address-style row ids mask stale postings at search time, and append mode
    // (num_to_merge == 0) defers cleanup to a real merge; both keep the plain tail.
    if num_to_merge == 0 || !dataset.manifest.uses_stable_row_ids() {
        return old_indices[tail_start..].to_vec();
    }

    let deleted_frags: RoaringBitmap = dataset
        .get_fragments()
        .iter()
        .filter(|f| f.metadata().deletion_file.is_some())
        .map(|f| f.id() as u32)
        .collect();
    if deleted_frags.is_empty() {
        return old_indices[tail_start..].to_vec();
    }

    let mut selected = Vec::new();
    for (i, idx) in old_indices.iter().enumerate() {
        let covers_deleted = idx
            .effective_fragment_bitmap(&dataset.fragment_bitmap)
            .is_some_and(|eff| !eff.is_disjoint(&deleted_frags));
        if i >= tail_start || covers_deleted {
            selected.push(*idx);
        }
    }
    selected
}

#[allow(clippy::too_many_arguments)]
async fn merge_scalar_indices<'a>(
    dataset: Arc<Dataset>,
    old_indices: &[&'a IndexMetadata],
    unindexed: &[Fragment],
    options: &OptimizeOptions,
    index_type: IndexType,
    field_path: &str,
    column_name: &str,
    base_unindexed_bitmap: RoaringBitmap,
) -> Result<
    Option<(
        Uuid,
        Vec<&'a IndexMetadata>,
        RoaringBitmap,
        CreatedIndex,
        u64,
    )>,
> {
    if old_indices.is_empty() {
        return Err(Error::index(
            "merge_scalar_indices: no previous index found".to_string(),
        ));
    }

    let mut selected_old_indices = select_segments_to_merge(dataset.as_ref(), old_indices, options);

    // No new data + ≤1 old selected = rewriting one segment to itself.
    if unindexed.is_empty() && selected_old_indices.len() <= 1 {
        return Ok(None);
    }

    // `select_segments_to_merge` only looks at the trailing `num_indices_to_merge`
    // segments (plus any deletion-affected one); with `num_indices_to_merge: Some(0)`
    // (append mode) it selects none at all. A row-id-domain BTree segment left
    // outside that selection would otherwise survive untouched beside the
    // freshly built address-domain segment this call is about to create.
    // `LogicalScalarIndex` takes the whole named index's result domain from its
    // first segment alone, so that mix makes every match the new segment finds
    // look like a row id needing resolution, and matching rows silently vanish
    // when that resolution misses. Fold every such segment into this rebuild,
    // regardless of which ones the trailing-window heuristic already picked,
    // so the named index never ends up straddling both domains.
    if index_type == IndexType::BTree && dataset.manifest.uses_stable_row_ids() {
        for idx in old_indices {
            if !idx.results_are_row_addrs()
                && !selected_old_indices
                    .iter()
                    .any(|selected| selected.uuid == idx.uuid)
            {
                selected_old_indices.push(idx);
            }
        }
    }

    // For the delta case (`selected` empty) the reference is purely
    // for reading params; fall back to the last old index then.
    let reference_idx = selected_old_indices
        .first()
        .copied()
        .unwrap_or(old_indices[old_indices.len() - 1]);
    let reference_index = dataset
        .open_scalar_index_for_maintenance(field_path, &reference_idx.uuid, &NoOpMetricsCollector)
        .await?;

    // A JSON-wrapped BTREE segment cannot yet carry the row-address-domain
    // migration safely on a stable-row-id dataset -- see `build_scalar_index`'s
    // matching guard, which this mirrors so maintenance is rejected too, not
    // just creation. Checked here, before deciding between a merge or a
    // rebuild, rather than relying on that guard to be reached only if the
    // rebuild branch happens to be picked.
    if dataset.manifest.uses_stable_row_ids() {
        let reference_params = reference_index.derive_index_params()?;
        if reference_params.index_type.eq_ignore_ascii_case("json")
            && reference_params
                .params
                .as_deref()
                .and_then(json_wrapped_target_type)
                .is_some_and(|target| target.eq_ignore_ascii_case(BuiltinIndexType::BTree.as_str()))
        {
            return Err(Error::not_supported(
                "a JSON-wrapped BTREE index is not yet supported on a dataset with stable row \
                 ids enabled -- create a plain BTREE index on this column instead, or disable \
                 stable row ids for this dataset"
                    .to_string(),
            ));
        }
    }

    let update_criteria = reference_index.update_criteria();

    // Effective = bitmap ∩ live fragments; deleted = bitmap \ live fragments.
    // On a tagged table the selected segments' old data is instead the coverage
    // the tagged reader derives for them (translated, live), and retired source
    // fragments are not something to remove from translated addresses.
    let tagged_coverage =
        tagged_segment_coverage(dataset.as_ref(), &selected_old_indices, None).await?;
    let (effective_old_frags, deleted_old_frags) = match &tagged_coverage {
        Some(coverage) => (
            selected_old_indices
                .iter()
                .filter_map(|index| coverage.get(&index.uuid))
                .fold(RoaringBitmap::new(), |acc, bitmap| acc | bitmap),
            RoaringBitmap::new(),
        ),
        None => split_segment_coverage(dataset.as_ref(), selected_old_indices.iter().copied()),
    };

    // The fragments a from-scratch rebuild must scan: everything the merged
    // segment will answer for, all of it live.
    let mut rebuild_frags = base_unindexed_bitmap.clone();
    rebuild_frags |= &effective_old_frags;
    let new_uuid = Uuid::new_v4();

    // Scalar Index that expos an N:1 segment-merge primitive reachable without
    // rescanning the dataset
    let has_segment_merge_primitive = matches!(
        index_type,
        IndexType::BTree | IndexType::Bitmap | IndexType::NGram
    );
    let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?;
    let ngram_requires_rebuild = index_type == IndexType::NGram
        && frag_reuse_index.as_ref().is_some_and(|frag_reuse_index| {
            fragment_reuse_affects_segments(frag_reuse_index, selected_old_indices.iter().copied())
        });

    // A BTree segment persisted before address-domain support stores row ids
    // directly. Merging it with newly scanned address-domain data on a
    // stable-row-id dataset would silently combine two different domains in
    // the same `ids` column, so such a segment must be rebuilt from scratch
    // (a full rescan, never touching its stale page data) rather than merged.
    // Harmless elsewhere: without stable row ids the two domains coincide.
    let btree_legacy_domain_mismatch = index_type == IndexType::BTree
        && dataset.manifest.uses_stable_row_ids()
        && selected_old_indices
            .iter()
            .any(|segment| !segment.results_are_row_addrs());

    // Merge new data into the existing segment(s) without rebuilding from
    // scratch, when all hold:
    //   - `effective_old_frags`: the selected segments' coverage intersected
    //     with live fragments is non-empty, i.e. there is old data worth keeping.
    //   - `update_criteria` only requires the newly appended data. Indexes that
    //     need old data must rebuild over `frag_bitmap` so the scanned rows
    //     exactly match the segment coverage being committed.
    //   - `has_segment_merge_primitive` (Indices supports N:1 segments merge) OR
    //     `selected_old_indices.len() == 1` (any scalar type can `update` one).
    // Otherwise (e.g. ≥2 selected segments of a type without an N:1 merge
    // primitive) the index is rebuilt from scratch over `frag_bitmap`.
    let can_merge_segments = !effective_old_frags.is_empty()
        && !update_criteria.requires_old_data
        && !ngram_requires_rebuild
        && !btree_legacy_domain_mismatch
        && (has_segment_merge_primitive || selected_old_indices.len() == 1);

    // The bitmap the merged segment commits, decided by how its content is
    // produced:
    //   - rebuilt from a scan: exactly the live fragments scanned (direct
    //     coverage of every row it holds);
    //   - merged from the old segments' pages: on a tagged table the union of
    //     the selected segments' STORED bitmaps, i.e. provenance, retired
    //     sources included. The pages were translated to live addresses when
    //     the segments were opened, and the tagged reader stops translating
    //     at a live fragment, so provenance keeps deriving coverage for this
    //     segment and for any unselected sibling that shares a destination,
    //     and keeps the mapping they need retained until they are rebuilt;
    //     elsewhere the effective (live) coverage, as before.
    let mut frag_bitmap = rebuild_frags.clone();
    let (created_index, new_dataset_version) = if !can_merge_segments {
        (
            rebuild_scalar_segment(
                dataset.as_ref(),
                &reference_index,
                field_path,
                column_name,
                new_uuid,
                rebuild_frags.iter().collect(),
            )
            .await?,
            dataset.manifest.version,
        )
    } else {
        if tagged_coverage.is_some() {
            frag_bitmap = base_unindexed_bitmap.clone();
            for index in &selected_old_indices {
                if let Some(stored) = &index.fragment_bitmap {
                    frag_bitmap |= stored;
                }
            }
        }
        let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?;

        // Try a seed-based update before falling back to a full column scan.
        // Seeds are only available if every unindexed fragment had a seed buffer
        // written during its data file write, and the plugin validates that the
        // seeds are compatible with the current index configuration.
        let index_details =
            fetch_index_details(dataset.as_ref(), field_path, reference_idx).await?;
        let details = IndexDetails(index_details.clone());
        let plugin = details.get_plugin()?;
        // Only open data files looking for seeds when the plugin confirms this
        // index type and configuration can actually produce them.
        let maybe_created = if plugin.might_use_seeds(&index_details) {
            if let Some(seeds) =
                try_harvest_seeds(dataset.as_ref(), unindexed, reference_idx).await?
            {
                plugin
                    .update_from_seeds(seeds, reference_index.clone(), &index_details, &new_store)
                    .await?
            } else {
                None
            }
        } else {
            None
        };

        let created_index = if let Some(created) = maybe_created {
            created
        } else {
            let new_data_stream = load_unindexed_training_data(
                dataset.as_ref(),
                field_path,
                &update_criteria,
                unindexed,
            )
            .await?;

            match index_type {
                IndexType::BTree => {
                    let (_, old_data_filters) = build_per_segment_filters(
                        dataset.as_ref(),
                        &selected_old_indices,
                        None,
                        true,
                    )
                    .await?;
                    crate::index::scalar::btree::open_and_merge_segments(
                        dataset.as_ref(),
                        field_path,
                        &selected_old_indices,
                        new_data_stream,
                        &new_store,
                        &old_data_filters,
                        None,
                    )
                    .await?
                }
                IndexType::Bitmap => {
                    let (_, old_data_filters) = build_per_segment_filters(
                        dataset.as_ref(),
                        &selected_old_indices,
                        None,
                        false,
                    )
                    .await?;
                    crate::index::scalar::bitmap::open_and_merge_segments(
                        dataset.as_ref(),
                        field_path,
                        &selected_old_indices,
                        new_data_stream,
                        &new_store,
                        &old_data_filters,
                        None,
                    )
                    .await?
                }
                IndexType::NGram => {
                    let (_, old_data_filters) = build_per_segment_filters(
                        dataset.as_ref(),
                        &selected_old_indices,
                        None,
                        false,
                    )
                    .await?;
                    crate::index::scalar::ngram::open_and_merge_segments(
                        dataset.as_ref(),
                        &selected_old_indices,
                        Some(new_data_stream),
                        &new_store,
                        &old_data_filters,
                        None,
                    )
                    .await?
                }
                _ => {
                    let old_data_filter = build_old_data_filter(
                        dataset.as_ref(),
                        &effective_old_frags,
                        &deleted_old_frags,
                        false,
                    )
                    .await?;
                    reference_index
                        .update(new_data_stream, &new_store, old_data_filter)
                        .await?
                }
            }
        };
        let source_dataset_version = selected_old_indices
            .iter()
            .map(|index| index.dataset_version)
            .min()
            .unwrap_or(dataset.manifest.version);
        (created_index, source_dataset_version)
    };

    Ok(Some((
        new_uuid,
        selected_old_indices.to_vec(),
        frag_bitmap,
        created_index,
        new_dataset_version,
    )))
}

async fn metadata_is_vector_index(dataset: &Dataset, index: &IndexMetadata) -> Result<bool> {
    // Declared type first: an index awaiting training declares itself a vector
    // index and has no file to find it by.
    if index.index_details.is_some() {
        return Ok(crate::index::segment_has_vector_details(index));
    }

    if let Some(files) = &index.files {
        return Ok(files.iter().any(|file| file.path == INDEX_FILE_NAME));
    }

    let index_dir = dataset.indice_files_dir(index)?;
    let index_file = index_dir
        .clone()
        .join(index.uuid.to_string())
        .join(INDEX_FILE_NAME);
    let object_store = dataset.object_store_for_index(index).await?;
    object_store.exists(&index_file).await
}

/// Whether a rebuild of `params` should expect the definition alone, because
/// the column does not hold the vectors its quantizer needs.
pub async fn expects_definition_only(
    dataset: &Dataset,
    field_path: &str,
    params: &VectorIndexParams,
) -> Result<bool> {
    let Some(minimum) = super::vector::vector_quantizer_minimum_rows(&params.stages) else {
        return Ok(false);
    };
    Ok(!super::vector::has_vectors_to_train(dataset, field_path, minimum).await?)
}

/// Retrain a whole segment over `fragment_bitmap` from index parameters alone.
async fn rebuild_vector_segment(
    dataset: &Dataset,
    name: &str,
    params: &VectorIndexParams,
    field_path: &str,
    fragment_bitmap: &RoaringBitmap,
    progress: Arc<dyn IndexBuildProgress>,
) -> Result<IndexMetadata> {
    let mut build_dataset = dataset.clone();
    CreateIndexBuilder::new(&mut build_dataset, &[field_path], IndexType::Vector, params)
        .name(name.to_string())
        .replace(true)
        .fragments(fragment_bitmap.iter().collect())
        .progress(progress)
        .execute_uncommitted()
        .await
}

/// True when a segment carries only its definition: nothing to open, nothing
/// to append to, so training is what it is waiting for.
///
/// Happens when an index is created against an empty table, or one without
/// enough rows to justify training it.
pub fn is_definition_only_segment(metadata: &IndexMetadata) -> bool {
    let no_files_recorded = metadata.files.as_ref().is_none_or(|files| files.is_empty());
    let covers_nothing = metadata
        .fragment_bitmap
        .as_ref()
        .is_some_and(RoaringBitmap::is_empty);
    no_files_recorded && covers_nothing
}

async fn scan_vector_fragments(
    dataset: &Dataset,
    field_path: &str,
    column_nullable: bool,
    fragments: &[Fragment],
) -> Result<crate::dataset::scanner::DatasetRecordBatchStream> {
    let mut scanner = dataset.scan();
    scanner
        .with_fragments(fragments.to_vec())
        .with_row_id()
        .project(&[field_path])?;
    if column_nullable {
        let column_expr = lance_datafusion::logical_expr::field_path_to_expr(field_path)?;
        scanner.filter_expr(column_expr.is_not_null());
    }
    scanner.try_into_stream().await
}

/// Train a column whose index is still only a definition, from the parameters
/// that definition carries.
///
/// The result supersedes every segment the index had: they hold no data, and
/// what this writes covers the whole column.
async fn train_from_definition<'a>(
    dataset: &Arc<Dataset>,
    old_indices: &[&'a IndexMetadata],
    field_path: &str,
    options: &OptimizeOptions,
) -> Result<IndexMergeResults<'a>> {
    let params = old_indices
        .last()
        .and_then(|metadata| metadata.index_details.as_deref())
        .and_then(details::vector_params_from_details)
        .ok_or_else(|| {
            Error::index(format!(
                "Optimize vector index: index '{}' awaits training but carries no parameters",
                old_indices[0].name
            ))
        })?;
    let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone();
    let segment = rebuild_vector_segment(
        dataset.as_ref(),
        &old_indices[0].name,
        &params,
        field_path,
        &fragment_bitmap,
        options.progress.clone(),
    )
    .await?;
    // Only a result that looks degraded is worth confirming; a rebuild that
    // covered rows needs no count to justify it.
    let definition_expected = is_definition_only_segment(&segment)
        && expects_definition_only(dataset.as_ref(), field_path, &params).await?;
    fresh_vector_segment_result(
        segment,
        &fragment_bitmap,
        old_indices.to_vec(),
        definition_expected,
    )
}

fn fresh_vector_segment_result<'a>(
    segment: IndexMetadata,
    expected_fragment_bitmap: &RoaringBitmap,
    removed_indices: Vec<&'a IndexMetadata>,
    definition_expected: bool,
) -> Result<IndexMergeResults<'a>> {
    let fragment_bitmap = segment.fragment_bitmap.ok_or_else(|| {
        Error::index(
            "Optimize vector index: newly built segment has no fragment bitmap".to_string(),
        )
    })?;
    // A column with too few vectors to train the quantizer yields the
    // definition alone: its rows return to unindexed, which is the state a
    // table below the threshold belongs in. Accepted only when the caller
    // established that beforehand and the segment wrote no files to match, so
    // coverage lost for any other reason still fails below.
    let wrote_no_files = segment.files.as_ref().is_some_and(|files| files.is_empty());
    if definition_expected && wrote_no_files && fragment_bitmap.is_empty() {
        return Ok(IndexMergeResults {
            new_uuid: segment.uuid,
            removed_indices,
            new_fragment_bitmap: fragment_bitmap,
            new_dataset_version: segment.dataset_version,
            new_index_version: segment.index_version,
            new_index_details: segment
                .index_details
                .map(|details| details.as_ref().clone())
                .ok_or_else(|| {
                    Error::index(
                        "Optimize vector index: rebuilt definition has no index details"
                            .to_string(),
                    )
                })?,
            files: Vec::new(),
        });
    }
    if fragment_bitmap != *expected_fragment_bitmap {
        return Err(Error::index(format!(
            "Optimize vector index: newly built segment covers fragments {:?}, expected {:?}",
            fragment_bitmap, expected_fragment_bitmap
        )));
    }
    let index_details = segment.index_details.ok_or_else(|| {
        Error::index("Optimize vector index: newly built segment has no index details".to_string())
    })?;
    let files = segment.files.ok_or_else(|| {
        Error::index("Optimize vector index: newly built segment has no file metadata".to_string())
    })?;

    Ok(IndexMergeResults {
        new_uuid: segment.uuid,
        removed_indices,
        new_fragment_bitmap: fragment_bitmap,
        new_dataset_version: segment.dataset_version,
        new_index_version: segment.index_version,
        new_index_details: index_details.as_ref().clone(),
        files,
    })
}

/// The column path an index is maintained under: the keyed field's path, or
/// for a full-text index the canonical path its document granularity resolves
/// to (returned alongside, for the inverted merge).
pub async fn index_field_path(
    dataset: &Dataset,
    metadata: &IndexMetadata,
) -> Result<(
    String,
    Option<crate::index::scalar::inverted::ResolvedFtsField>,
)> {
    let field_id = *metadata.fields.first().ok_or_else(|| {
        Error::index(format!(
            "Append index: segment {} is missing field ids",
            metadata.uuid
        ))
    })?;
    let raw_field_path = dataset.schema().field_path(field_id)?;
    let details = super::scalar::fetch_index_details(dataset, &raw_field_path, metadata).await?;
    let resolved_fts = if details.type_url.ends_with("InvertedIndexDetails") {
        let details = lance_index::pbold::InvertedIndexDetails::decode(details.value.as_slice())
            .map_err(|error| {
                Error::io(format!(
                    "failed to decode InvertedIndexDetails payload: {error}"
                ))
            })?;
        let granularity = lance_index::scalar::inverted::DocumentGranularity::try_from(
            details.document_granularity,
        )?;
        Some(crate::index::scalar::inverted::resolve_fts_field_by_id(
            dataset.schema(),
            field_id,
            granularity,
        )?)
    } else {
        None
    };
    let field_path = resolved_fts
        .as_ref()
        .map(|resolved| resolved.canonical_path.clone())
        .unwrap_or(raw_field_path);
    Ok((field_path, resolved_fts))
}

/// Build one new vector segment over `fragments`, encoded against the model
/// (IVF centroids and quantizer) of `reference`, without touching any existing
/// segment. The reference is opened by the caller, so it need not be in the
/// manifest.
pub async fn append_vector_segment<'a>(
    dataset: &Dataset,
    field_path: &str,
    column_nullable: bool,
    reference_metadata: &IndexMetadata,
    reference_index: &Arc<dyn lance_index::vector::VectorIndex>,
    fragments: &[Fragment],
    options: &OptimizeOptions,
) -> Result<IndexMergeResults<'a>> {
    let reference_logical_index = LogicalVectorIndex::try_new(
        reference_metadata.name.clone(),
        field_path.to_string(),
        vec![(reference_metadata.clone(), reference_index.clone())],
    )?;
    let reference_ivf_view = reference_logical_index.as_ivf()?;
    let new_data_stream =
        scan_vector_fragments(dataset, field_path, column_nullable, fragments).await?;
    let mut append_options = options.clone();
    append_options.num_indices_to_merge = Some(0);
    append_options.enable_partition_rebalancing = false;
    append_options.retrain = false;
    let (new_uuid, indices_merged, files) = optimize_vector_indices(
        dataset.clone(),
        Some(new_data_stream),
        field_path,
        &reference_ivf_view,
        &append_options,
    )
    .boxed()
    .await?;
    if indices_merged != 0 {
        return Err(Error::index(format!(
            "Optimize vector index append unexpectedly merged {indices_merged} existing segments"
        )));
    }
    Ok(IndexMergeResults {
        new_uuid,
        removed_indices: Vec::new(),
        new_fragment_bitmap: fragments.iter().map(|frag| frag.id as u32).collect(),
        new_dataset_version: dataset.manifest.version,
        new_index_version: index_type_for_segmented_optimize(reference_index.as_ref())?.version(),
        new_index_details: reference_metadata
            .index_details
            .as_deref()
            .cloned()
            .unwrap_or_else(vector_index_details_default),
        files,
    })
}

/// Split vector segments into the live ones and the dormant ones.
///
/// Segments with stored rows (raw coverage) but no live coverage (e.g. a
/// definition retained through a full rewrite) are dormant. With stable row
/// ids their stored postings still share ids with live rows, so merging them
/// would resurrect stale vectors; they may only be replaced. A born-empty
/// segment (deferred build) has no stored rows and stays mergeable.
///
/// Under a tagged history an empty bitmap does not mean empty pages: an
/// in-place rewrite withdraws a segment's whole coverage and leaves its files.
/// Such a segment is dormant too, rebuilt from the live fragments; only a
/// definition without files (the deferred build) stays mergeable there.
/// Elsewhere an index initialized on an empty table has files and no rows and
/// keeps merging as before.
///
/// On a tagged table a segment's stored bitmap is provenance: its rows
/// translate to the coverage the tagged reader derives for it, so THAT is its
/// live coverage. A segment the reader excludes (derived coverage empty) is
/// the dormant one.
///
/// `tagged_coverage` is [`tagged_segment_coverage`] of `segments`.
pub fn partition_dormant_vector_segments<'a>(
    dataset: &Dataset,
    segments: &[&'a IndexMetadata],
    tagged_coverage: Option<&HashMap<Uuid, RoaringBitmap>>,
) -> (Vec<&'a IndexMetadata>, Vec<&'a IndexMetadata>) {
    segments.iter().copied().partition(|idx| {
        let has_stored_rows = match tagged_coverage {
            Some(_) => !is_definition_only_segment(idx),
            None => idx
                .fragment_bitmap
                .as_ref()
                .is_some_and(|bitmap| !bitmap.is_empty()),
        };
        let has_live_coverage = match tagged_coverage {
            Some(coverage) => coverage
                .get(&idx.uuid)
                .is_some_and(|bitmap| !bitmap.is_empty()),
            None => idx
                .effective_fragment_bitmap(&dataset.fragment_bitmap)
                .is_none_or(|bitmap| !bitmap.is_empty()),
        };
        !has_stored_rows || has_live_coverage
    })
}

/// Merge a list of provided unindexed data, with a specific number of previous indices
/// into a new index, to improve the query performance.
pub async fn merge_indices_with_unindexed_frags<'a>(
    dataset: Arc<Dataset>,
    old_indices: &[&'a IndexMetadata],
    unindexed: &[Fragment],
    options: &OptimizeOptions,
) -> Result<Option<IndexMergeResults<'a>>> {
    if old_indices.is_empty() {
        return Err(Error::index(
            "Append index: no previous index found".to_string(),
        ));
    };

    let column = dataset
        .schema()
        .field_by_id(old_indices[0].fields[0])
        .ok_or(Error::index(format!(
            "Append index: column {} does not exist",
            old_indices[0].fields[0]
        )))?;

    let (field_path, resolved_fts) = index_field_path(dataset.as_ref(), old_indices[0]).await?;
    let first_is_vector_index = metadata_is_vector_index(dataset.as_ref(), old_indices[0]).await?;
    for idx in old_indices.iter().skip(1) {
        let is_vector_index = metadata_is_vector_index(dataset.as_ref(), idx).await?;
        if is_vector_index != first_is_vector_index {
            return Err(Error::index(format!(
                "Append index: invalid mixed index deltas: {:?}",
                old_indices
            )));
        }
    }

    let mut base_unindexed_bitmap = RoaringBitmap::new();
    unindexed.iter().for_each(|frag| {
        base_unindexed_bitmap.insert(frag.id as u32);
    });

    let (new_uuid, removed_indices, new_fragment_bitmap, created_index, new_dataset_version) =
        if first_is_vector_index {
            let tagged_coverage =
                tagged_segment_coverage(dataset.as_ref(), old_indices, None).await?;
            let (live_segments, dormant_segments) = partition_dormant_vector_segments(
                dataset.as_ref(),
                old_indices,
                tagged_coverage.as_ref(),
            );
            if !dormant_segments.is_empty() && !live_segments.is_empty() && !options.retrain {
                // Optimize the live segments as usual; the dormant segments
                // are superseded by whatever that produces.
                let mut merged = Box::pin(merge_indices_with_unindexed_frags(
                    dataset.clone(),
                    &live_segments,
                    unindexed,
                    options,
                ))
                .await?;
                if let Some(results) = merged.as_mut() {
                    results.removed_indices.extend(dormant_segments);
                }
                return Ok(merged);
            }
            let rebuild_dormant = live_segments.is_empty() && !dormant_segments.is_empty();
            // A deferred build leaves a segment with no file to open, so the
            // column is trained from its definition rather than appended to.
            // Opening the index by name opens every segment, so any one of
            // them settles it.
            let awaits_training = old_indices
                .iter()
                .any(|idx| is_definition_only_segment(idx));

            // Rebuilding and first training both read the whole column, so
            // neither has anything to do until a fragment sits outside the
            // index.
            if unindexed.is_empty() && (rebuild_dormant || awaits_training) {
                return Ok(None);
            }

            if awaits_training {
                return train_from_definition(&dataset, old_indices, &field_path, options)
                    .await
                    .map(Some);
            }

            // The maintenance opener: a segment the tagged reader excludes
            // from the listing (its coverage was withdrawn) is still here,
            // opened as contributing nothing, so it can be replaced.
            let full_logical_index = dataset
                .open_logical_vector_index_for_maintenance(&field_path, &old_indices[0].name)
                .await?;
            let mut opened_indices_by_uuid = full_logical_index
                .iter()
                .map(|(metadata, index)| (metadata.uuid, (metadata.clone(), index.clone())))
                .collect::<std::collections::HashMap<_, _>>();
            let mut selected_metadatas = Vec::with_capacity(old_indices.len());
            let mut selected_indices = Vec::with_capacity(old_indices.len());
            for metadata in old_indices {
                let (selected_metadata, selected_index) = opened_indices_by_uuid.remove(&metadata.uuid).ok_or_else(|| {
                Error::index(format!(
                    "Append index: logical vector index '{}' does not contain requested segment {}",
                    old_indices[0].name, metadata.uuid
                ))
            })?;
                selected_metadatas.push(selected_metadata);
                selected_indices.push(selected_index);
            }
            let logical_index = LogicalVectorIndex::try_new(
                old_indices[0].name.clone(),
                field_path.clone(),
                selected_metadatas
                    .into_iter()
                    .zip(selected_indices)
                    .collect(),
            )?;

            if options.retrain || rebuild_dormant {
                let (reference_metadata, reference_index) =
                    logical_index.iter().last().ok_or_else(|| {
                        Error::index(format!(
                            "Optimize vector index: logical index '{}' has no physical segments",
                            logical_index.name()
                        ))
                    })?;
                let params =
                    fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?;
                let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone();
                let segment = rebuild_vector_segment(
                    dataset.as_ref(),
                    logical_index.name(),
                    &params,
                    &field_path,
                    &fragment_bitmap,
                    options.progress.clone(),
                )
                .await?;
                // Only a result that looks degraded is worth confirming; a
                // rebuild that covered rows needs no count to justify it.
                let definition_expected = is_definition_only_segment(&segment)
                    && expects_definition_only(dataset.as_ref(), &field_path, &params).await?;
                return fresh_vector_segment_result(
                    segment,
                    &fragment_bitmap,
                    old_indices.to_vec(),
                    definition_expected,
                )
                .map(Some);
            }

            let ivf_view = logical_index.as_ivf()?;
            let compatibility =
                vector_segment_compatibility(&ivf_view, "optimizing logical vector index")?;
            let explicit_append = options.num_indices_to_merge == Some(0);
            let default_append_for_heterogeneous_models = options.num_indices_to_merge.is_none()
                && compatibility == VectorSegmentCompatibility::QueryCompatibleModelsDiffer;

            if !unindexed.is_empty() && (explicit_append || default_append_for_heterogeneous_models)
            {
                // CreateIndex commits append new segments at the end of the manifest's
                // stable index order. Reusing the current suffix model keeps consecutive
                // appends directly mergeable without requiring unrelated older segments
                // to share that model.
                let (reference_metadata, reference_index) =
                    logical_index.iter().last().ok_or_else(|| {
                        Error::index(format!(
                            "Optimize vector index: logical index '{}' has no physical segments",
                            logical_index.name()
                        ))
                    })?;
                return append_vector_segment(
                    dataset.as_ref(),
                    &field_path,
                    column.nullable,
                    reference_metadata,
                    reference_index,
                    unindexed,
                    options,
                )
                .await
                .map(Some);
            }

            // Append is a steady-state no-op when every fragment is already
            // indexed. Default optimize may still rebalance one segment.
            if unindexed.is_empty() && options.num_indices_to_merge == Some(0) {
                return Ok(None);
            }
            let steady_state_rebalance =
                if unindexed.is_empty() && options.num_indices_to_merge.is_none() {
                    let Some(rebalance) = select_steady_state_rebalance(&ivf_view, compatibility)?
                    else {
                        return Ok(None);
                    };
                    Some(rebalance)
                } else {
                    None
                };

            if let Some(SteadyStateRebalance::Segment(selected_segment_id)) = steady_state_rebalance
                && logical_index.num_segments() > 1
            {
                let removed_segment = old_indices
                .iter()
                .copied()
                .find(|metadata| metadata.uuid == selected_segment_id)
                .ok_or_else(|| {
                    Error::index(format!(
                        "Append index: logical vector index '{}' does not contain selected segment {}",
                        old_indices[0].name, selected_segment_id
                    ))
                })?;
                let (selected_metadata, selected_index) = logical_index
                .iter()
                .find(|(metadata, _)| metadata.uuid == selected_segment_id)
                .map(|(metadata, index)| (metadata.clone(), index.clone()))
                .ok_or_else(|| {
                    Error::index(format!(
                        "Append index: failed to materialize selected segment {} from logical vector index '{}'",
                        selected_segment_id, old_indices[0].name
                    ))
                })?;
                let selected_logical_index = LogicalVectorIndex::try_new(
                    old_indices[0].name.clone(),
                    field_path.clone(),
                    vec![(selected_metadata, selected_index)],
                )?;
                let selected_ivf_view = selected_logical_index.as_ivf()?;
                let (new_uuid, indices_merged, files) = Box::pin(optimize_vector_indices(
                    dataset.as_ref().clone(),
                    Option::<
                        lance_io::stream::RecordBatchStreamAdapter<
                            futures::stream::Empty<lance_core::Result<arrow_array::RecordBatch>>,
                        >,
                    >::None,
                    &field_path,
                    &selected_ivf_view,
                    options,
                ))
                .await?;
                if indices_merged == 0 {
                    return Ok(None);
                }

                // The rebalanced file was read through the translating
                // loader, so on a tagged table it holds live addresses under
                // the segment's stored provenance: publish that provenance,
                // as a merge does, not stored ∩ live, which is empty for
                // retired sources and would turn indexed rows into scans.
                let new_fragment_bitmap = if tagged_coverage.is_some() {
                    removed_segment.fragment_bitmap.clone().unwrap_or_default()
                } else {
                    removed_segment
                        .effective_fragment_bitmap(&dataset.fragment_bitmap)
                        .or_else(|| removed_segment.fragment_bitmap.clone())
                        .unwrap_or_default()
                };

                Ok((
                    new_uuid,
                    vec![removed_segment],
                    new_fragment_bitmap,
                    CreatedIndex {
                        index_details: removed_segment
                            .index_details
                            .as_deref()
                            .cloned()
                            .unwrap_or_else(vector_index_details_default),
                        index_version: removed_segment.index_version as u32,
                        files: table_files_to_index(files),
                    },
                    removed_segment.dataset_version,
                ))
            } else {
                let mut frag_bitmap = base_unindexed_bitmap.clone();
                let num_segments_to_merge = options
                    .num_indices_to_merge
                    .unwrap_or(logical_index.num_segments())
                    .min(logical_index.num_segments());
                let merge_start = logical_index
                    .num_segments()
                    .saturating_sub(num_segments_to_merge);
                let merge_logical_index = LogicalVectorIndex::try_new(
                    old_indices[0].name.clone(),
                    field_path.clone(),
                    logical_index
                        .iter()
                        .skip(merge_start)
                        .map(|(metadata, index)| (metadata.clone(), index.clone()))
                        .collect(),
                )?;
                let merge_ivf_view = merge_logical_index.as_ivf()?;

                let new_data_stream = if unindexed.is_empty() {
                    None
                } else {
                    Some(
                        scan_vector_fragments(
                            dataset.as_ref(),
                            &field_path,
                            column.nullable,
                            unindexed,
                        )
                        .await?,
                    )
                };

                let (new_uuid, indices_merged, files) = optimize_vector_indices(
                    dataset.as_ref().clone(),
                    new_data_stream,
                    &field_path,
                    &merge_ivf_view,
                    options,
                )
                .boxed()
                .await?;

                let removed_indices = old_indices[old_indices.len() - indices_merged..].to_vec();
                let new_dataset_version = removed_indices
                    .iter()
                    .map(|index| index.dataset_version)
                    .min()
                    .unwrap_or(dataset.manifest.version);
                removed_indices.iter().for_each(|idx| {
                    frag_bitmap.extend(idx.fragment_bitmap.as_ref().unwrap().iter());
                });
                for removed in removed_indices.iter() {
                    if let Some(effective) =
                        removed.effective_fragment_bitmap(&dataset.fragment_bitmap)
                    {
                        frag_bitmap |= &effective;
                    }
                }

                // Metadata must come from an actual merge source. Older incompatible
                // segments outside the selected suffix may describe a different model.
                let (reference_metadata, reference_index) =
                    merge_logical_index.iter().last().ok_or_else(|| {
                        Error::index(
                            "Optimize vector index merge did not select a reference segment"
                                .to_string(),
                        )
                    })?;
                let index_details = removed_indices
                    .iter()
                    .rev()
                    .filter_map(|idx| idx.index_details.as_ref())
                    .find(|d| !d.value.is_empty())
                    .map(|d| d.as_ref().clone())
                    .or_else(|| {
                        reference_metadata
                            .index_details
                            .as_deref()
                            .filter(|details| !details.value.is_empty())
                            .cloned()
                    })
                    .unwrap_or_else(vector_index_details_default);
                let index_version = if let Some(metadata) = removed_indices.first() {
                    metadata.index_version as u32
                } else {
                    index_type_for_segmented_optimize(reference_index.as_ref())?.version() as u32
                };

                Ok((
                    new_uuid,
                    removed_indices,
                    frag_bitmap,
                    CreatedIndex {
                        index_details,
                        index_version,
                        files: table_files_to_index(files),
                    },
                    new_dataset_version,
                ))
            }
        } else {
            let mut indices = Vec::with_capacity(old_indices.len());
            for idx in old_indices {
                match dataset
                    .open_generic_index_for_maintenance(
                        &field_path,
                        &idx.uuid,
                        &NoOpMetricsCollector,
                    )
                    .await
                {
                    Ok(index) => indices.push(index),
                    Err(e) => {
                        log::warn!(
                            "Cannot open index on column '{}': {}. \
                         Skipping index merge for this column.",
                            field_path,
                            e
                        );
                        return Ok(None);
                    }
                }
            }

            if indices
                .windows(2)
                .any(|w| w[0].index_type() != w[1].index_type())
            {
                return Err(Error::index(format!(
                    "Append index: invalid index deltas: {:?}",
                    old_indices
                )));
            }

            let index_type = indices[0].index_type();
            match index_type {
                IndexType::Inverted => {
                    let selected_old_indices =
                        select_segments_to_merge(dataset.as_ref(), old_indices, options);
                    if unindexed.is_empty() && selected_old_indices.len() <= 1 {
                        return Ok(None);
                    }
                    let reference_idx = selected_old_indices
                        .first()
                        .copied()
                        .unwrap_or(old_indices[old_indices.len() - 1]);
                    let reference_index = dataset
                        .open_scalar_index_for_maintenance(
                            &field_path,
                            &reference_idx.uuid,
                            &NoOpMetricsCollector,
                        )
                        .await?;
                    let update_criteria = reference_index.update_criteria();
                    if update_criteria.requires_old_data {
                        let params = reference_index.derive_index_params()?;
                        let resolved = resolved_fts.as_ref().ok_or_else(|| {
                            Error::internal(
                                "Inverted index metadata did not resolve an FTS field".to_string(),
                            )
                        })?;
                        let new_data_stream = load_fts_training_data(
                            dataset.as_ref(),
                            resolved,
                            &update_criteria.data_criteria,
                            None,
                            true,
                            None,
                        )
                        .await?;
                        let new_uuid = Uuid::new_v4();
                        let created_index = super::scalar::build_scalar_index(
                            dataset.as_ref(),
                            &resolved.canonical_path,
                            new_uuid,
                            &params,
                            true,
                            None,
                            Some((new_data_stream, update_criteria.data_criteria.clone())),
                            Arc::new(NoopIndexBuildProgress),
                        )
                        .await?;
                        return Ok(Some(IndexMergeResults {
                            new_uuid,
                            removed_indices: old_indices.to_vec(),
                            new_fragment_bitmap: dataset.fragment_bitmap.as_ref().clone(),
                            new_dataset_version: dataset.manifest.version,
                            new_index_version: created_index.index_version as i32,
                            new_index_details: created_index.index_details,
                            files: index_files_to_table(created_index.files),
                        }));
                    }

                    let fragments = Some(unindexed.to_vec());
                    let resolved = resolved_fts.as_ref().ok_or_else(|| {
                        Error::internal(
                            "Inverted index metadata did not resolve an FTS field".to_string(),
                        )
                    })?;
                    let new_data_stream = load_fts_training_data(
                        dataset.as_ref(),
                        resolved,
                        &update_criteria.data_criteria,
                        fragments,
                        true,
                        None,
                    )
                    .await?;

                    let mut frag_bitmap = base_unindexed_bitmap;
                    let mut effective_old_frags = RoaringBitmap::new();
                    let mut selected_indices = Vec::with_capacity(selected_old_indices.len());
                    // On a tagged table a selected segment's old data is the
                    // coverage the reader derives for it (translated, live)
                    // and the merged segment keeps the stored provenance,
                    // exactly as the scalar merge does; stored ∩ live would
                    // be empty for retired sources and drop indexed rows.
                    let tagged_coverage =
                        tagged_segment_coverage(dataset.as_ref(), &selected_old_indices, None)
                            .await?;
                    for idx in &selected_old_indices {
                        match &tagged_coverage {
                            Some(coverage) => {
                                if let Some(derived) = coverage.get(&idx.uuid) {
                                    effective_old_frags |= derived;
                                }
                                if let Some(stored) = &idx.fragment_bitmap {
                                    frag_bitmap |= stored;
                                }
                            }
                            None => {
                                if let Some(effective) =
                                    idx.effective_fragment_bitmap(&dataset.fragment_bitmap)
                                {
                                    frag_bitmap |= &effective;
                                    effective_old_frags |= &effective;
                                }
                            }
                        }
                        let scalar_index = dataset
                            .open_scalar_index_for_maintenance(
                                &field_path,
                                &idx.uuid,
                                &NoOpMetricsCollector,
                            )
                            .await?;
                        let inverted_index = scalar_index
                            .as_any()
                            .downcast_ref::<InvertedIndex>()
                            .ok_or_else(|| {
                                Error::index(format!(
                                    "Append index: expected inverted index segment {}, got {:?}",
                                    idx.uuid,
                                    scalar_index.index_type()
                                ))
                            })?;
                        selected_indices.push(Arc::new(inverted_index.clone()));
                    }

                    let old_data_filter = if selected_indices.is_empty() {
                        None
                    } else if dataset.manifest.uses_stable_row_ids() {
                        let valid_old_row_ids =
                            build_stable_row_id_filter(dataset.as_ref(), &effective_old_frags)
                                .await?;
                        Some(OldIndexDataFilter::RowIds(valid_old_row_ids))
                    } else {
                        Some(OldIndexDataFilter::Fragments {
                            to_keep: effective_old_frags,
                            to_remove: RoaringBitmap::new(),
                        })
                    };

                    let new_uuid = Uuid::new_v4();
                    let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?;
                    let (created_index, new_dataset_version) = if selected_indices.is_empty() {
                        (
                            super::scalar::build_scalar_index(
                                dataset.as_ref(),
                                &resolved.canonical_path,
                                new_uuid,
                                &reference_index.derive_index_params()?,
                                true,
                                None,
                                Some((new_data_stream, update_criteria.data_criteria.clone())),
                                Arc::new(NoopIndexBuildProgress),
                            )
                            .await?,
                            dataset.manifest.version,
                        )
                    } else {
                        (
                            InvertedIndex::merge_segments(
                                &selected_indices,
                                new_data_stream,
                                &new_store,
                                old_data_filter,
                                options.progress.clone(),
                            )
                            .await?,
                            selected_old_indices
                                .iter()
                                .map(|index| index.dataset_version)
                                .min()
                                .unwrap_or(dataset.manifest.version),
                        )
                    };

                    Ok((
                        new_uuid,
                        selected_old_indices.to_vec(),
                        frag_bitmap,
                        created_index,
                        new_dataset_version,
                    ))
                }
                it if it.is_scalar() => {
                    let Some(result) = merge_scalar_indices(
                        dataset.clone(),
                        old_indices,
                        unindexed,
                        options,
                        it,
                        &field_path,
                        column.name.as_str(),
                        base_unindexed_bitmap,
                    )
                    .await?
                    else {
                        return Ok(None);
                    };
                    Ok(result)
                }
                _ => Err(Error::index(format!(
                    "Append index: invalid index type: {:?}",
                    indices[0].index_type()
                ))),
            }
        }?;

    Ok(Some(IndexMergeResults {
        new_uuid,
        removed_indices,
        new_fragment_bitmap,
        new_dataset_version,
        new_index_version: created_index.index_version as i32,
        new_index_details: created_index.index_details,
        files: index_files_to_table(created_index.files),
    }))
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::index::DatasetIndexExt;
    use crate::index::DatasetIndexInternalExt;
    use crate::index::optimize::{
        OptimizeIndicesPlan, commit_optimize_indices, plan_optimize_indices,
    };
    use arrow::datatypes::{Float32Type, UInt32Type};
    use arrow_array::cast::AsArray;
    use arrow_array::{
        Array, ArrayRef, FixedSizeListArray, Float32Array, Int32Array, ListArray, RecordBatch,
        RecordBatchIterator, StringArray, StructArray, UInt32Array, UInt64Array,
    };
    use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer};
    use arrow_schema::{DataType, Field, Schema};
    use futures::TryStreamExt;
    use lance_arrow::FixedSizeListArrayExt;
    use lance_core::utils::tempfile::TempStrDir;
    use lance_datafusion::utils::reader_to_stream;
    use lance_datagen::{Dimension, RowCount, array};
    use lance_file::version::LanceFileVersion;
    use lance_index::vector::hnsw::builder::HnswBuildParams;
    use lance_index::vector::sq::builder::SQBuildParams;
    use lance_index::{
        IndexType,
        scalar::{
            BuiltinIndexType, InvertedIndexParams, ScalarIndexParams, SearchResult, TextQuery,
        },
        vector::{ivf::IvfBuildParams, pq::PQBuildParams},
    };
    use lance_linalg::distance::MetricType;
    use lance_testing::datagen::{generate_random_array, generate_random_array_with_seed};
    use rstest::rstest;

    use crate::dataset::builder::DatasetBuilder;
    use crate::dataset::optimize::{CompactionOptions, compact_files};
    use crate::dataset::transaction::{Operation, Transaction};
    use crate::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched, WriteMode, WriteParams};
    use crate::index::CreateIndexBuilder;
    use crate::index::vector::VectorIndexParams;
    use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount, copy_test_data_to_tmp};

    #[test]
    fn test_fragment_reuse_at_source_version_affects_segment() {
        use lance_index::frag_reuse::{
            FragDigest, FragReuseGroup, FragReuseIndexDetails, FragReuseVersion,
        };

        let segment = IndexMetadata {
            uuid: Uuid::new_v4(),
            name: "text_ngram".to_string(),
            fields: vec![0],
            covering_fields: vec![],
            dataset_version: 5,
            fragment_bitmap: Some(RoaringBitmap::from_iter([1u32])),
            index_details: None,
            index_version: 0,
            created_at: None,
            base_id: None,
            files: None,
        };
        let frag_reuse_index = CompactFragReuseIndex::from_row_id_maps(
            Uuid::new_v4(),
            vec![],
            FragReuseIndexDetails {
                versions: vec![FragReuseVersion {
                    dataset_version: 5,
                    groups: vec![FragReuseGroup {
                        changed_row_addrs: vec![1],
                        old_frags: vec![FragDigest {
                            id: 1,
                            physical_rows: 1,
                            num_deleted_rows: 0,
                        }],
                        new_frags: vec![FragDigest {
                            id: 2,
                            physical_rows: 1,
                            num_deleted_rows: 0,
                        }],
                    }],
                }],
            },
        );

        assert!(fragment_reuse_affects_segments(
            &frag_reuse_index,
            [&segment]
        ));

        let rebuilt_segment = IndexMetadata {
            dataset_version: 6,
            fragment_bitmap: Some(RoaringBitmap::from_iter([2u32])),
            ..segment
        };
        assert!(!fragment_reuse_affects_segments(
            &frag_reuse_index,
            [&rebuilt_segment]
        ));

        for stale_version in [4, 5] {
            let stale_remapped_segment = IndexMetadata {
                dataset_version: stale_version,
                ..rebuilt_segment.clone()
            };
            assert!(
                fragment_reuse_affects_segments(&frag_reuse_index, [&stale_remapped_segment]),
                "{stale_version}"
            );
        }
    }

    fn clustered_vector_batch(
        schema: Arc<Schema>,
        start_id: i32,
        rows: usize,
        dimension: usize,
        center: f32,
    ) -> (RecordBatch, Arc<FixedSizeListArray>) {
        let values = (0..rows)
            .flat_map(|row| {
                (0..dimension)
                    .map(move |column| center + row as f32 * 0.01 + column as f32 * 0.0001)
            })
            .collect::<Vec<_>>();
        let vectors = Arc::new(
            FixedSizeListArray::try_new_from_values(
                arrow_array::Float32Array::from(values),
                dimension as i32,
            )
            .unwrap(),
        );
        let ids = Arc::new(Int32Array::from_iter_values(
            start_id..start_id + rows as i32,
        ));
        let batch = RecordBatch::try_new(schema, vec![ids, vectors.clone() as ArrayRef]).unwrap();
        (batch, vectors)
    }

    async fn nearest_id(
        dataset: &Dataset,
        query: &arrow_array::Float32Array,
        num_probes: usize,
    ) -> i32 {
        let result = dataset
            .scan()
            .project(&["id"])
            .unwrap()
            .nearest("vector", query, 1)
            .unwrap()
            .minimum_nprobes(num_probes)
            .maximum_nprobes(num_probes)
            .refine(1)
            .try_into_batch()
            .await
            .unwrap();
        assert_eq!(result.num_rows(), 1);
        result["id"]
            .as_primitive::<arrow::datatypes::Int32Type>()
            .value(0)
    }

    #[rstest]
    #[case::append(OptimizeOptions::append().with_partition_rebalancing(), "num_indices_to_merge", "Some(0)")]
    #[case::task_bound(OptimizeOptions::merge(1).with_partition_rebalancing().max_rows_per_task(Some(64)), "max_rows_per_task", "Some(64)")]
    #[case::segment_bound(OptimizeOptions::merge(1).with_partition_rebalancing().max_rows_per_segment(Some(64)), "max_rows_per_segment", "Some(64)")]
    #[tokio::test]
    async fn test_partition_rebalancing_rejects_incompatible_options(
        #[case] options: OptimizeOptions,
        #[case] parameter: &str,
        #[case] value: &str,
    ) {
        let mut dataset = lance_datagen::gen_batch()
            .col("id", array::step::<UInt32Type>())
            .into_dataset("memory://", FragmentCount(1), FragmentRowCount(1))
            .await
            .unwrap();
        let error = dataset.optimize_indices(&options).await.unwrap_err();
        assert!(matches!(error, Error::InvalidInput { .. }));
        assert!(error.to_string().contains(parameter));
        assert!(error.to_string().contains(value));
    }

    #[rstest]
    #[case::unchanged(1, 128, std::cmp::Ordering::Equal)]
    #[case::split(1, 8, std::cmp::Ordering::Greater)]
    #[case::join(8, 64, std::cmp::Ordering::Less)]
    #[tokio::test]
    async fn test_merge_latest_segment_with_partition_rebalancing(
        #[case] num_partitions: usize,
        #[case] target_partition_size: usize,
        #[case] expected_partition_change: std::cmp::Ordering,
        #[values(false, true)] enable_stable_row_ids: bool,
        #[values(false, true)] distributed: bool,
    ) {
        let test_dir = TempStrDir::default();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
                false,
            ),
        ]));
        let (first, first_vectors) = clustered_vector_batch(schema.clone(), 0, 64, 2, 0.0);
        let (last, last_vectors) = clustered_vector_batch(schema.clone(), 64, 64, 2, 100.0);
        let mut dataset = Dataset::write(
            RecordBatchIterator::new(vec![Ok(first), Ok(last)], schema.clone()),
            test_dir.as_str(),
            Some(WriteParams {
                max_rows_per_file: 64,
                enable_stable_row_ids,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        let params = VectorIndexParams::with_ivf_flat_params(
            MetricType::L2,
            IvfBuildParams {
                num_partitions: Some(num_partitions),
                target_partition_size: Some(target_partition_size),
                max_iters: 2,
                sample_rate: 2,
                ..Default::default()
            },
        );
        let mut segments = Vec::with_capacity(2);
        for fragment in dataset.get_fragments() {
            segments.push(
                CreateIndexBuilder::new(&mut dataset, &["vector"], IndexType::Vector, &params)
                    .name("vector_idx".to_string())
                    .fragments(vec![fragment.id() as u32])
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        let first_uuid = segments[0].uuid;
        dataset
            .commit_existing_index_segments("vector_idx", "vector", segments)
            .await
            .unwrap();

        for round in 0..3 {
            let (batch, vectors) = clustered_vector_batch(
                schema.clone(),
                128 + round * 16,
                16,
                2,
                101.0 + round as f32,
            );
            dataset
                .append(
                    RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
                    None,
                )
                .await
                .unwrap();
            let options = OptimizeOptions::merge(1).with_partition_rebalancing();
            if distributed {
                let plan = plan_optimize_indices(&dataset, &options).await.unwrap();
                let plan: OptimizeIndicesPlan =
                    serde_json::from_str(&serde_json::to_string(&plan).unwrap()).unwrap();
                assert_eq!(plan.tasks.len(), 1);
                assert!(plan.tasks[0].enable_partition_rebalancing);
                let mut built = Vec::new();
                for task in &plan.tasks {
                    built.extend(task.execute(&dataset).await.unwrap());
                }
                let mut merged = Vec::new();
                for merge in &plan.merges {
                    merged.extend(merge.execute(&dataset, &built).await.unwrap());
                }
                commit_optimize_indices(&mut dataset, merged, &options)
                    .await
                    .unwrap();
            } else {
                dataset.optimize_indices(&options).await.unwrap();
            }
            let segments = dataset.load_indices_by_name("vector_idx").await.unwrap();
            assert_eq!(segments.len(), 2, "each append must replace only the tail");
            assert_eq!(
                segments[0].uuid, first_uuid,
                "earlier segments must not change"
            );
            assert!(
                dataset
                    .unindexed_fragments("vector_idx")
                    .await
                    .unwrap()
                    .is_empty()
            );
            let logical = dataset
                .open_logical_vector_index("vector", "vector_idx")
                .await
                .unwrap();
            let partitions = logical.as_ivf().unwrap().num_partitions_per_segment();
            assert_eq!(partitions[0].1, num_partitions);
            assert_eq!(
                partitions[1].1.cmp(&num_partitions),
                expected_partition_change
            );
            let nprobes = partitions.iter().map(|(_, count)| *count).max().unwrap();
            for (expected, query) in [
                (0, first_vectors.value(0)),
                (64, last_vectors.value(0)),
                (128 + round * 16, vectors.value(0)),
            ] {
                assert_eq!(
                    nearest_id(&dataset, query.as_primitive::<Float32Type>(), nprobes).await,
                    expected
                );
            }
        }
    }

    #[rstest]
    #[case::append(OptimizeOptions::append(), false)]
    #[case::default_with_rebalancing(OptimizeOptions::new().with_partition_rebalancing(), false)]
    #[case::default_with_stable_row_ids_and_delete(OptimizeOptions::default(), true)]
    #[tokio::test]
    async fn test_vector_append_is_segment_set_native_with_distinct_models(
        #[case] options: OptimizeOptions,
        #[case] use_stable_row_ids: bool,
    ) {
        const DIMENSION: usize = 8;
        const ROWS_PER_FRAGMENT: usize = 64;
        const INDEX_NAME: &str = "vector_idx";

        let test_dir = TempStrDir::default();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIMENSION as i32,
                ),
                false,
            ),
        ]));
        let (first_batch, first_vectors) =
            clustered_vector_batch(schema.clone(), 0, ROWS_PER_FRAGMENT, DIMENSION, 0.0);
        let (second_batch, _) = clustered_vector_batch(
            schema.clone(),
            ROWS_PER_FRAGMENT as i32,
            ROWS_PER_FRAGMENT,
            DIMENSION,
            100.0,
        );
        let reader =
            RecordBatchIterator::new(vec![Ok(first_batch), Ok(second_batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_dir.as_str(),
            Some(WriteParams {
                enable_stable_row_ids: use_stable_row_ids,
                max_rows_per_file: ROWS_PER_FRAGMENT,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let initial_fragments = dataset.get_fragments();
        assert_eq!(initial_fragments.len(), 2);
        // Two partitions per segment: summed over the segments they are under
        // the join threshold, which must not tempt a steady-state merge of
        // segments whose centroids differ.
        let params = VectorIndexParams::ivf_flat(2, MetricType::L2);
        let mut initial_segments = Vec::with_capacity(initial_fragments.len());
        for fragment in &initial_fragments {
            initial_segments.push(
                CreateIndexBuilder::new(&mut dataset, &["vector"], IndexType::Vector, &params)
                    .name(INDEX_NAME.to_string())
                    .fragments(vec![fragment.id() as u32])
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        dataset
            .commit_existing_index_segments(INDEX_NAME, "vector", initial_segments)
            .await
            .unwrap();
        dataset = DatasetBuilder::from_uri(test_dir.as_str())
            .load()
            .await
            .unwrap();

        let old_segments = dataset.load_indices_by_name(INDEX_NAME).await.unwrap();
        assert_eq!(old_segments.len(), 2);
        let logical_index = dataset
            .open_logical_vector_index("vector", INDEX_NAME)
            .await
            .unwrap();
        let centroids = logical_index
            .as_ivf()
            .unwrap()
            .segments()
            .map(|(_, index)| index.ivf_model().centroids_array().unwrap().to_data())
            .collect::<Vec<_>>();
        assert_ne!(
            centroids[0], centroids[1],
            "test setup must use distinct IVF centroid models"
        );
        let version_before_default_optimize = dataset.version().version;
        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        assert_eq!(
            dataset.version().version,
            version_before_default_optimize,
            "default optimize must leave incompatible steady-state segments unmerged"
        );
        assert_eq!(
            dataset.load_indices_by_name(INDEX_NAME).await.unwrap(),
            old_segments,
            "default optimize changed incompatible steady-state segments"
        );

        let appended_start_id = (2 * ROWS_PER_FRAGMENT) as i32;
        let (appended_batch, appended_vectors) = clustered_vector_batch(
            schema.clone(),
            appended_start_id,
            ROWS_PER_FRAGMENT,
            DIMENSION,
            200.0,
        );
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(appended_batch)], schema),
                None,
            )
            .await
            .unwrap();
        let appended_fragment_id = dataset.get_fragments().last().unwrap().id() as u32;
        if use_stable_row_ids {
            dataset
                .delete(&format!("id = {appended_start_id}"))
                .await
                .unwrap();
        }

        dataset.optimize_indices(&options).await.unwrap();

        let mut dataset = DatasetBuilder::from_uri(test_dir.as_str())
            .load()
            .await
            .unwrap();
        let committed = dataset.load_indices_by_name(INDEX_NAME).await.unwrap();
        assert_eq!(
            committed.len(),
            3,
            "append must add one segment without rewriting existing segments"
        );
        for old in &old_segments {
            let retained = committed
                .iter()
                .find(|segment| segment.uuid == old.uuid)
                .expect("old physical segment UUID must remain committed");
            assert_eq!(retained, old, "old physical segment metadata changed");
        }

        let new_segment = committed
            .iter()
            .find(|segment| {
                old_segments
                    .iter()
                    .all(|old_segment| old_segment.uuid != segment.uuid)
            })
            .expect("one new physical segment must be committed");
        assert_eq!(
            new_segment.fragment_bitmap.as_ref().unwrap(),
            &RoaringBitmap::from_iter([appended_fragment_id]),
            "the appended segment must cover only the new fragment"
        );

        let logical_index = dataset
            .open_logical_vector_index("vector", INDEX_NAME)
            .await
            .unwrap();
        let latest_old_uuid = old_segments.last().unwrap().uuid;
        let latest_old_centroids = logical_index
            .iter()
            .find(|(metadata, _)| metadata.uuid == latest_old_uuid)
            .unwrap()
            .1
            .ivf_model()
            .centroids_array()
            .unwrap()
            .to_data();
        let new_centroids = logical_index
            .iter()
            .find(|(metadata, _)| metadata.uuid == new_segment.uuid)
            .unwrap()
            .1
            .ivf_model()
            .centroids_array()
            .unwrap()
            .to_data();
        assert_eq!(
            new_centroids, latest_old_centroids,
            "append must reuse the latest segment's complete IVF model"
        );

        let mut covered = RoaringBitmap::new();
        for segment in &committed {
            let segment_coverage = segment.fragment_bitmap.as_ref().unwrap();
            assert!(
                covered.is_disjoint(segment_coverage),
                "physical vector segment coverage must be pairwise disjoint"
            );
            covered |= segment_coverage;
        }
        assert_eq!(
            covered,
            dataset.fragment_bitmap.as_ref().clone(),
            "physical vector segments must exactly cover indexed fragments"
        );

        let old_query = first_vectors.value(1);
        assert_eq!(
            nearest_id(
                &dataset,
                old_query.as_primitive::<Float32Type>(),
                committed.len()
            )
            .await,
            1,
            "query must return a representative row from an old segment"
        );
        let appended_query_offset = usize::from(use_stable_row_ids);
        let appended_query = appended_vectors.value(appended_query_offset);
        assert_eq!(
            nearest_id(
                &dataset,
                appended_query.as_primitive::<Float32Type>(),
                committed.len()
            )
            .await,
            appended_start_id + appended_query_offset as i32,
            "query must return a representative row from the appended segment"
        );
        if use_stable_row_ids {
            assert_eq!(
                dataset
                    .scan()
                    .filter(&format!("id = {appended_start_id}"))
                    .unwrap()
                    .count_rows()
                    .await
                    .unwrap(),
                0,
                "deleted stable-row-id data must not reappear through the appended segment"
            );
        }

        dataset
            .optimize_indices(&OptimizeOptions::merge(2))
            .await
            .unwrap();
        let merged_suffix = dataset.load_indices_by_name(INDEX_NAME).await.unwrap();
        assert_eq!(
            merged_suffix.len(),
            2,
            "the reference-compatible suffix should merge without rewriting the incompatible base"
        );
        assert!(
            merged_suffix
                .iter()
                .any(|segment| segment.uuid == old_segments[0].uuid),
            "the incompatible base segment must remain unchanged"
        );
        assert!(
            merged_suffix
                .iter()
                .all(|segment| segment.uuid != latest_old_uuid && segment.uuid != new_segment.uuid),
            "the compatible suffix segments must be replaced"
        );

        dataset
            .optimize_indices(&OptimizeOptions::retrain())
            .await
            .unwrap();
        let retrained = dataset.load_indices_by_name(INDEX_NAME).await.unwrap();
        assert_eq!(retrained.len(), 1);
        assert_eq!(
            retrained[0].fragment_bitmap.as_ref().unwrap(),
            dataset.fragment_bitmap.as_ref(),
            "explicit retrain must source-rebuild one segment over all current fragments"
        );
        assert_eq!(
            nearest_id(
                &dataset,
                old_query.as_primitive::<Float32Type>(),
                retrained.len()
            )
            .await,
            1
        );
        assert_eq!(
            nearest_id(
                &dataset,
                appended_query.as_primitive::<Float32Type>(),
                retrained.len()
            )
            .await,
            appended_start_id + appended_query_offset as i32
        );
        if use_stable_row_ids {
            assert_eq!(
                dataset
                    .scan()
                    .filter(&format!("id = {appended_start_id}"))
                    .unwrap()
                    .count_rows()
                    .await
                    .unwrap(),
                0,
                "explicit retrain must not restore a deleted stable-row-id row"
            );
        }
    }

    #[rstest]
    #[case::metric(
        VectorIndexParams::ivf_flat(1, MetricType::L2),
        VectorIndexParams::ivf_flat(1, MetricType::Cosine),
        "has metric"
    )]
    #[case::index_family(
        VectorIndexParams::ivf_flat(1, MetricType::L2),
        VectorIndexParams::ivf_hnsw(
            MetricType::L2,
            IvfBuildParams::new(1),
            HnswBuildParams::default()
        ),
        "has type"
    )]
    #[tokio::test]
    async fn test_vector_append_validates_logical_query_compatibility(
        #[case] first_params: VectorIndexParams,
        #[case] second_params: VectorIndexParams,
        #[case] expected_error: &str,
    ) {
        const DIMENSION: usize = 8;
        const ROWS_PER_FRAGMENT: usize = 64;
        const INDEX_NAME: &str = "vector_idx";

        let test_dir = TempStrDir::default();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIMENSION as i32,
                ),
                false,
            ),
        ]));
        let (first_batch, _) =
            clustered_vector_batch(schema.clone(), 0, ROWS_PER_FRAGMENT, DIMENSION, 1.0);
        let (second_batch, _) = clustered_vector_batch(
            schema.clone(),
            ROWS_PER_FRAGMENT as i32,
            ROWS_PER_FRAGMENT,
            DIMENSION,
            100.0,
        );
        let reader =
            RecordBatchIterator::new(vec![Ok(first_batch), Ok(second_batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_dir.as_str(),
            Some(WriteParams {
                max_rows_per_file: ROWS_PER_FRAGMENT,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let initial_fragments = dataset.get_fragments();
        let mut segments = Vec::with_capacity(initial_fragments.len());
        for (fragment, params) in initial_fragments.iter().zip([first_params, second_params]) {
            segments.push(
                CreateIndexBuilder::new(&mut dataset, &["vector"], IndexType::Vector, &params)
                    .name(INDEX_NAME.to_string())
                    .fragments(vec![fragment.id() as u32])
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        let version_before = dataset.version().version;
        let error = dataset
            .commit_existing_index_segments(INDEX_NAME, "vector", segments.clone())
            .await
            .unwrap_err();
        assert!(matches!(error, Error::InvalidInput { .. }));
        assert!(error.to_string().contains(expected_error), "{error}");
        assert_eq!(dataset.version().version, version_before);
        assert!(
            dataset
                .load_indices_by_name(INDEX_NAME)
                .await
                .unwrap()
                .is_empty()
        );

        // Bypass commit-time validation only to exercise append's independent guard
        // against an incompatible segment set already present in the manifest.
        let transaction = Transaction::new(
            dataset.manifest.version,
            Operation::CreateIndex {
                new_indices: segments,
                removed_indices: vec![],
            },
            None,
        );
        dataset
            .apply_commit(transaction, &Default::default(), &Default::default())
            .await
            .unwrap();

        let (appended_batch, _) = clustered_vector_batch(
            schema.clone(),
            (2 * ROWS_PER_FRAGMENT) as i32,
            ROWS_PER_FRAGMENT,
            DIMENSION,
            200.0,
        );
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(appended_batch)], schema),
                None,
            )
            .await
            .unwrap();

        let version_before = dataset.version().version;
        let segments_before = dataset.load_indices_by_name(INDEX_NAME).await.unwrap();
        let object_store = dataset.object_store.clone();
        let directories_before = object_store
            .read_dir(dataset.indices_dir())
            .await
            .unwrap()
            .into_iter()
            .collect::<std::collections::HashSet<_>>();

        let error = dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap_err();
        assert!(
            error.to_string().contains(expected_error),
            "expected logical query compatibility error containing '{expected_error}', got {error}"
        );

        let latest = DatasetBuilder::from_uri(test_dir.as_str())
            .load()
            .await
            .unwrap();
        assert_eq!(
            latest.version().version,
            version_before,
            "incompatible logical segments must fail before committing"
        );
        let mut segments_after = latest.load_indices_by_name(INDEX_NAME).await.unwrap();
        let mut segments_before = segments_before;
        for segment in segments_after.iter_mut().chain(segments_before.iter_mut()) {
            segment.created_at = None;
        }
        assert_eq!(
            segments_after, segments_before,
            "incompatible logical segments must remain unchanged"
        );
        let directories_after = object_store
            .read_dir(latest.indices_dir())
            .await
            .unwrap()
            .into_iter()
            .collect::<std::collections::HashSet<_>>();
        assert_eq!(
            directories_after, directories_before,
            "query compatibility validation must run before staging a new segment"
        );
    }

    #[tokio::test]
    async fn test_vector_segment_native_append_preserves_commit_conflicts() {
        const DIMENSION: usize = 8;
        const ROWS_PER_FRAGMENT: usize = 64;
        const INDEX_NAME: &str = "vector_idx";

        let test_dir = TempStrDir::default();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIMENSION as i32,
                ),
                false,
            ),
        ]));
        let (initial_batch, _) =
            clustered_vector_batch(schema.clone(), 0, ROWS_PER_FRAGMENT, DIMENSION, 0.0);
        let mut dataset = Dataset::write(
            RecordBatchIterator::new(vec![Ok(initial_batch)], schema.clone()),
            test_dir.as_str(),
            Some(WriteParams {
                max_rows_per_file: ROWS_PER_FRAGMENT,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .create_index(
                &["vector"],
                IndexType::Vector,
                Some(INDEX_NAME.to_string()),
                &VectorIndexParams::ivf_flat(1, MetricType::L2),
                true,
            )
            .await
            .unwrap();
        let old_uuid = dataset.load_indices_by_name(INDEX_NAME).await.unwrap()[0].uuid;

        let (first_append, _) = clustered_vector_batch(
            schema.clone(),
            ROWS_PER_FRAGMENT as i32,
            ROWS_PER_FRAGMENT,
            DIMENSION,
            100.0,
        );
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(first_append)], schema.clone()),
                None,
            )
            .await
            .unwrap();
        let stale_version = dataset.version().version;
        let mut first_optimizer = dataset.checkout_version(stale_version).await.unwrap();
        let mut stale_optimizer = dataset.checkout_version(stale_version).await.unwrap();
        first_optimizer
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();

        let error = stale_optimizer
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap_err();
        assert!(
            matches!(error, Error::RetryableCommitConflict { .. }),
            "stale vector append optimize must retain retryable conflict semantics: {error}"
        );

        let latest = DatasetBuilder::from_uri(test_dir.as_str())
            .load()
            .await
            .unwrap();
        let committed = latest.load_indices_by_name(INDEX_NAME).await.unwrap();
        assert_eq!(committed.len(), 2);
        assert!(
            committed.iter().any(|segment| segment.uuid == old_uuid),
            "the successful concurrent optimize must retain the old segment"
        );
        assert_eq!(
            latest.unindexed_fragments(INDEX_NAME).await.unwrap().len(),
            0,
            "a failed stale commit must not publish the staged vector segment"
        );
    }

    #[tokio::test]
    async fn test_append_index() {
        const DIM: usize = 64;
        const IVF_PARTITIONS: usize = 2;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let vectors = generate_random_array(1000 * DIM);

        let schema = Arc::new(Schema::new(vec![Field::new(
            "vector",
            DataType::FixedSizeList(
                Arc::new(Field::new("item", DataType::Float32, true)),
                DIM as i32,
            ),
            true,
        )]));
        let array = Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap());
        let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();

        let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
        let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap();

        let ivf_params = IvfBuildParams::new(IVF_PARTITIONS);
        let pq_params = PQBuildParams {
            num_sub_vectors: 2,
            ..Default::default()
        };
        let params = VectorIndexParams::with_ivf_pq_params(MetricType::L2, ivf_params, pq_params);

        dataset
            .create_index(&["vector"], IndexType::Vector, None, &params, true)
            .await
            .unwrap();

        let vectors = generate_random_array(1000 * DIM);
        let array = Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap());
        let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();

        let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
        dataset.append(batches, None).await.unwrap();

        let index = &dataset.load_indices().await.unwrap()[0];
        assert!(
            !dataset
                .unindexed_fragments(&index.name)
                .await
                .unwrap()
                .is_empty()
        );

        let q = array.value(5);
        let mut scanner = dataset.scan();
        scanner
            .nearest("vector", q.as_primitive::<Float32Type>(), 10)
            .unwrap();
        let results = scanner
            .try_into_stream()
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        assert_eq!(results[0].num_rows(), 10); // Flat search.

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let indices = dataset.load_indices().await.unwrap();

        assert!(
            dataset
                .unindexed_fragments(&index.name)
                .await
                .unwrap()
                .is_empty()
        );

        // There should be two indices directories existed.
        let object_store = dataset.object_store.as_ref();
        let index_dirs = object_store.read_dir(dataset.indices_dir()).await.unwrap();
        assert_eq!(index_dirs.len(), 2);

        let mut scanner = dataset.scan();
        scanner
            .nearest("vector", q.as_primitive::<Float32Type>(), 10)
            .unwrap();
        let results = scanner
            .try_into_stream()
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        let vectors = &results[0]["vector"];
        // Second batch of vectors should be in the index.
        let contained = vectors.as_fixed_size_list().iter().any(|v| {
            let vec = v.as_ref().unwrap();
            array.iter().any(|a| a.as_ref().unwrap() == vec)
        });
        assert!(contained);

        // Check that the index has all 2000 rows.
        let mut num_rows = 0;
        for index in indices.iter() {
            let index = dataset
                .open_vector_index("vector", &index.uuid, &NoOpMetricsCollector)
                .await
                .unwrap();
            num_rows += index.num_rows();
        }
        assert_eq!(num_rows, 2000);
    }

    #[tokio::test]
    async fn test_optimize_append_preserves_case_sensitive_nullable_vector_column() {
        const DIM: usize = 64;
        const ROWS: usize = 1000;

        fn make_vectors(rows: usize, dim: usize, include_null: bool) -> FixedSizeListArray {
            if include_null {
                let mut nulls_builder = BooleanBufferBuilder::new(rows);
                for row_idx in 0..rows {
                    nulls_builder.append(row_idx != 0);
                }
                let nulls = NullBuffer::new(nulls_builder.finish());
                FixedSizeListArray::try_new(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    dim as i32,
                    Arc::new(generate_random_array(rows * dim)),
                    Some(nulls),
                )
                .unwrap()
            } else {
                FixedSizeListArray::try_new_from_values(
                    generate_random_array(rows * dim),
                    dim as i32,
                )
                .unwrap()
            }
        }

        fn make_batch(
            schema: Arc<Schema>,
            start_id: u32,
            vectors: Arc<FixedSizeListArray>,
        ) -> RecordBatch {
            let columns: Vec<ArrayRef> = vec![
                Arc::new(UInt32Array::from_iter_values(
                    start_id..start_id + ROWS as u32,
                )) as ArrayRef,
                vectors as ArrayRef,
            ];
            RecordBatch::try_new(schema, columns).unwrap()
        }

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let vector_type = DataType::FixedSizeList(
            Arc::new(Field::new("item", DataType::Float32, true)),
            DIM as i32,
        );
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("VECTOR", vector_type, true),
        ]));

        let initial_vectors = Arc::new(make_vectors(ROWS, DIM, false));
        let initial_batch = make_batch(schema.clone(), 0, initial_vectors);
        let batches = RecordBatchIterator::new(std::iter::once(Ok(initial_batch)), schema.clone());
        let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap();

        let params = VectorIndexParams::with_ivf_pq_params(
            MetricType::L2,
            IvfBuildParams::new(2),
            PQBuildParams {
                num_sub_vectors: 2,
                ..Default::default()
            },
        );
        dataset
            .create_index(&["VECTOR"], IndexType::Vector, None, &params, true)
            .await
            .unwrap();

        let appended_vectors = Arc::new(make_vectors(ROWS, DIM, true));
        let query = appended_vectors.value(5);
        let appended_batch = make_batch(schema.clone(), ROWS as u32, appended_vectors);
        let batches = RecordBatchIterator::new(std::iter::once(Ok(appended_batch)), schema);
        dataset.append(batches, None).await.unwrap();

        let index_name = dataset.load_indices().await.unwrap()[0].name.clone();
        assert!(
            !dataset
                .unindexed_fragments(&index_name)
                .await
                .unwrap()
                .is_empty()
        );

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert!(
            dataset
                .unindexed_fragments(&index_name)
                .await
                .unwrap()
                .is_empty()
        );

        let mut scanner = dataset.scan();
        scanner
            .nearest("VECTOR", query.as_primitive::<Float32Type>(), 10)
            .unwrap();
        let results = scanner.try_into_batch().await.unwrap();
        assert_eq!(
            results.num_rows(),
            10,
            "expected the requested k=10 nearest-neighbor results"
        );
    }

    /// Regression: a second `OptimizeOptions::append()` call on a steady-state
    /// vector index used to fall through to `optimize_vector_indices` and write
    /// a new UUID directory + manifest even though nothing had changed. The
    /// no-op gate inside `merge_indices_with_unindexed_frags` should bail out
    /// once `select_steady_state_rebalance` returns `None`.
    #[tokio::test]
    async fn test_optimize_indices_append_is_noop_on_steady_state() {
        const DIM: usize = 64;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new(
            "vector",
            DataType::FixedSizeList(
                Arc::new(Field::new("item", DataType::Float32, true)),
                DIM as i32,
            ),
            true,
        )]));
        let make_batch = || {
            let arr = Arc::new(
                FixedSizeListArray::try_new_from_values(
                    generate_random_array(1000 * DIM),
                    DIM as i32,
                )
                .unwrap(),
            );
            RecordBatch::try_new(schema.clone(), vec![arr]).unwrap()
        };

        let batches =
            RecordBatchIterator::new(vec![make_batch()].into_iter().map(Ok), schema.clone());
        let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap();

        // num_partitions = 1 so the auto-rebalance heuristic has no join/split
        // candidate after the initial build — keeps this test focused on the
        // append-no-op behavior rather than the rebalance path.
        let params = VectorIndexParams::with_ivf_pq_params(
            MetricType::L2,
            IvfBuildParams::new(1),
            PQBuildParams {
                num_sub_vectors: 2,
                ..Default::default()
            },
        );
        dataset
            .create_index(&["vector"], IndexType::Vector, None, &params, true)
            .await
            .unwrap();

        let batches =
            RecordBatchIterator::new(vec![make_batch()].into_iter().map(Ok), schema.clone());
        dataset.append(batches, None).await.unwrap();

        // First append: writes the new fragment into a reference-compatible delta segment.
        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let version_before = dataset.version().version;
        let object_store = dataset.object_store.as_ref();
        let dirs_before = object_store
            .read_dir(dataset.indices_dir())
            .await
            .unwrap()
            .into_iter()
            .collect::<std::collections::HashSet<_>>();

        // Second append: nothing changed since the previous call. Must not
        // bump the manifest version or add a new UUID directory.
        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let dirs_after = object_store
            .read_dir(dataset.indices_dir())
            .await
            .unwrap()
            .into_iter()
            .collect::<std::collections::HashSet<_>>();

        assert_eq!(
            dataset.version().version,
            version_before,
            "second optimize_indices(append()) bumped the dataset version"
        );
        assert_eq!(
            dirs_after, dirs_before,
            "second optimize_indices(append()) created a new index directory"
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_query_delta_indices(
        #[values(
            VectorIndexParams::ivf_pq(2, 4, 2, MetricType::L2, 2),
            VectorIndexParams::ivf_rq(2, 1, MetricType::L2),
            VectorIndexParams::ivf_hnsw(
                MetricType::L2,
                IvfBuildParams::new(2),
                HnswBuildParams {
                    max_level: 2,
                    m: 4,
                    ef_construction: 16,
                    prefetch_distance: Some(1),
                }
            ),
            VectorIndexParams::with_ivf_hnsw_pq_params(
                MetricType::L2,
                IvfBuildParams {
                    num_partitions: Some(2),
                    max_iters: 2,
                    sample_rate: 2,
                    ..Default::default()
                },
                HnswBuildParams {
                    max_level: 2,
                    m: 4,
                    ef_construction: 16,
                    prefetch_distance: Some(1),
                },
                PQBuildParams {
                    num_sub_vectors: 2,
                    num_bits: 4,
                    max_iters: 2,
                    sample_rate: 2,
                    ..Default::default()
                }
            ),
            VectorIndexParams::with_ivf_hnsw_sq_params(
                MetricType::L2,
                IvfBuildParams::new(2),
                HnswBuildParams {
                    max_level: 2,
                    m: 4,
                    ef_construction: 16,
                    prefetch_distance: Some(1),
                },
                SQBuildParams::default()
            )
        )]
        index_params: VectorIndexParams,
    ) {
        const DIM: usize = 8;
        const INITIAL_ROWS: usize = 64;
        const APPENDED_ROWS: usize = 16;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let vectors = generate_random_array_with_seed::<Float32Type>(INITIAL_ROWS * DIM, [42; 32]);

        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIM as i32,
                ),
                true,
            ),
            Field::new("id", DataType::UInt32, false),
        ]));
        let array = Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap());
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                array.clone(),
                Arc::new(UInt32Array::from_iter_values(0..INITIAL_ROWS as u32)),
            ],
        )
        .unwrap();

        let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
        let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap();
        dataset
            .create_index(&["vector"], IndexType::Vector, None, &index_params, true)
            .await
            .unwrap();
        let stats: serde_json::Value =
            serde_json::from_str(&dataset.index_statistics("vector_idx").await.unwrap()).unwrap();
        assert_eq!(stats["num_indices"], 1);
        assert_eq!(stats["num_indexed_fragments"], 1);
        assert_eq!(stats["num_unindexed_fragments"], 0);

        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(array.slice(0, APPENDED_ROWS)),
                Arc::new(UInt32Array::from_iter_values(
                    INITIAL_ROWS as u32..(INITIAL_ROWS + APPENDED_ROWS) as u32,
                )),
            ],
        )
        .unwrap();

        let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
        dataset.append(batches, None).await.unwrap();
        let appended_fragment_id = dataset.get_fragments().last().unwrap().id() as u32;
        let stats: serde_json::Value =
            serde_json::from_str(&dataset.index_statistics("vector_idx").await.unwrap()).unwrap();
        assert_eq!(stats["num_indices"], 1);
        assert_eq!(stats["num_indexed_fragments"], 1);
        assert_eq!(stats["num_unindexed_fragments"], 1);

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let stats: serde_json::Value =
            serde_json::from_str(&dataset.index_statistics("vector_idx").await.unwrap()).unwrap();
        assert_eq!(stats["num_indices"], 2);
        assert_eq!(stats["num_indexed_fragments"], 2);
        assert_eq!(stats["num_unindexed_fragments"], 0);
        let appended_segments = dataset.load_indices_by_name("vector_idx").await.unwrap();
        assert!(
            appended_segments
                .iter()
                .all(|segment| segment.index_version == index_params.index_type().version()),
            "append must preserve the storage type's index version"
        );
        let logical_index = dataset
            .open_logical_vector_index("vector", "vector_idx")
            .await
            .unwrap();
        assert_eq!(logical_index.num_segments(), 2);
        assert_eq!(
            logical_index
                .num_rows_per_segment()
                .into_iter()
                .map(|(_, num_rows)| num_rows)
                .sum::<u64>(),
            (INITIAL_ROWS + APPENDED_ROWS) as u64
        );
        if matches!(
            index_params.index_type(),
            IndexType::IvfHnswFlat | IndexType::IvfHnswPq | IndexType::IvfHnswSq
        ) {
            let hnsw_params = logical_index
                .iter()
                .map(|(_, index)| index.statistics().unwrap()["sub_index"]["params"].clone())
                .collect::<Vec<_>>();
            assert!(
                hnsw_params.iter().all(|params| params == &hnsw_params[0]),
                "append must preserve the reference segment's HNSW build parameters: {hnsw_params:?}"
            );
        }

        let mut fanout_scanner = dataset.scan();
        fanout_scanner
            .project(&["id"])
            .unwrap()
            .nearest("vector", array.value(0).as_primitive::<Float32Type>(), 2)
            .unwrap()
            .minimum_nprobes(2)
            .maximum_nprobes(2)
            .refine(1);
        let fanout_plan = fanout_scanner.explain_plan(true).await.unwrap();
        assert!(
            fanout_plan.contains("ANNSubIndex: name=vector_idx, k=2, deltas=2"),
            "logical vector query must fan out across both physical segments, plan was:\n{fanout_plan}"
        );
        let results = fanout_scanner.try_into_batch().await.unwrap();
        assert_eq!(results.num_rows(), 2);

        for segment in &appended_segments {
            let fragment_bitmap = segment
                .fragment_bitmap
                .as_ref()
                .expect("vector segment must record fragment coverage");
            let expected_appended_row = fragment_bitmap.contains(appended_fragment_id);
            let mut segment_scanner = dataset.scan();
            segment_scanner
                .project(&["id"])
                .unwrap()
                .nearest("vector", array.value(0).as_primitive::<Float32Type>(), 1)
                .unwrap()
                .minimum_nprobes(2)
                .maximum_nprobes(2)
                .refine(1)
                .with_index_segments(vec![segment.uuid])
                .unwrap();
            let segment_result = segment_scanner.try_into_batch().await.unwrap();
            assert_eq!(segment_result.num_rows(), 1);
            let id = segment_result["id"].as_primitive::<UInt32Type>().value(0);
            assert_eq!(
                id >= INITIAL_ROWS as u32,
                expected_appended_row,
                "segment {} returned row {id} outside its fragment coverage {:?}",
                segment.uuid,
                fragment_bitmap
            );
        }

        dataset
            .optimize_indices(&OptimizeOptions::merge(2))
            .await
            .unwrap();
        let merged = dataset.load_indices_by_name("vector_idx").await.unwrap();
        assert_eq!(
            merged.len(),
            1,
            "the reference-compatible append segment must merge with its source"
        );
        assert_eq!(
            merged[0].index_version,
            index_params.index_type().version(),
            "merge must preserve the storage type's index version"
        );
        assert_eq!(
            merged[0].fragment_bitmap.as_ref().unwrap(),
            dataset.fragment_bitmap.as_ref(),
            "the compatible merge must preserve exact fragment coverage"
        );
    }

    #[tokio::test]
    async fn test_merge_indices_after_merge_insert() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        // Create initial dataset using lance_datagen
        let mut dataset = lance_datagen::gen_batch()
            .col("id", array::step::<UInt32Type>())
            .col("value", array::cycle_utf8_literals(&["a", "b", "c"]))
            .col(
                "vector",
                array::rand_vec::<Float32Type>(Dimension::from(64)),
            )
            .into_dataset_with_params(
                test_uri,
                FragmentCount(1),
                FragmentRowCount(1000),
                Some(WriteParams {
                    max_rows_per_file: 1000,
                    ..Default::default()
                }),
            )
            .await
            .unwrap();

        // Create initial index
        let ivf_params = IvfBuildParams::new(2);
        let pq_params = PQBuildParams {
            num_sub_vectors: 2,
            ..Default::default()
        };
        let params = VectorIndexParams::with_ivf_pq_params(MetricType::L2, ivf_params, pq_params);

        dataset
            .create_index(&["vector"], IndexType::Vector, None, &params, true)
            .await
            .unwrap();

        // Load initial index metadata
        let initial_indices = dataset.load_indices().await.unwrap();
        assert_eq!(initial_indices.len(), 1);
        let index_name = initial_indices[0].name.clone();

        // Prepare new data for merge insert (updates to existing rows)
        let new_batch = lance_datagen::gen_batch()
            .col("id", array::step_custom::<UInt32Type>(500, 1)) // IDs 500-999
            .col("value", array::cycle_utf8_literals(&["d", "e", "f"])) // Different values
            .col(
                "vector",
                array::rand_vec::<Float32Type>(Dimension::from(64)),
            )
            .into_batch_rows(RowCount::from(500))
            .unwrap();

        // Record the maximum fragment ID before merge insert
        let max_fragment_id_before = dataset.manifest.max_fragment_id().unwrap_or(0);

        // Execute merge insert operation
        let merge_job =
            MergeInsertBuilder::try_new(Arc::new(dataset.clone()), vec!["id".to_string()])
                .unwrap()
                .when_matched(WhenMatched::UpdateAll)
                .when_not_matched(WhenNotMatched::InsertAll)
                .try_build()
                .unwrap();

        let schema = new_batch.schema();
        let new_reader = Box::new(RecordBatchIterator::new([Ok(new_batch)], schema.clone()));
        let new_stream = reader_to_stream(new_reader);
        let (updated_dataset, merge_stats) = merge_job.execute(new_stream).await.unwrap();

        // Check merge stats
        assert_eq!(merge_stats.num_updated_rows, 500); // Updates for rows 500-999
        assert_eq!(merge_stats.num_inserted_rows, 0); // No new inserts in this case

        // Get the newly added fragments by comparing fragment IDs
        let unindexed_fragments: Vec<Fragment> = updated_dataset
            .get_fragments()
            .into_iter()
            .filter(|f| f.id() as u64 > max_fragment_id_before)
            .map(|f| f.metadata().clone())
            .collect();

        // Now run merge with known unindexed fragments
        let old_indices = updated_dataset
            .load_indices_by_name(&index_name)
            .await
            .unwrap();
        let old_indices_refs: Vec<&IndexMetadata> = old_indices.iter().collect();

        let merge_result = merge_indices_with_unindexed_frags(
            updated_dataset.clone(),
            &old_indices_refs,
            &unindexed_fragments,
            &OptimizeOptions::merge(old_indices.len()),
        )
        .await
        .unwrap();

        assert!(merge_result.is_some());
        let merge_result = merge_result.unwrap();

        // Verify that the new index covers all fragments
        let new_fragment_bitmap = &merge_result.new_fragment_bitmap;

        // Check that unindexed fragments are now included
        for fragment in &unindexed_fragments {
            assert!(new_fragment_bitmap.contains(fragment.id as u32));
        }

        // Check that old indexed fragments are still included
        // All fragments with ID <= max_fragment_id_before should be included
        for frag_id in 0..=max_fragment_id_before as u32 {
            assert!(new_fragment_bitmap.contains(frag_id));
        }

        // Verify the index can be used for search
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let indices = dataset.load_indices().await.unwrap();

        // There should still be indices (old one might be kept plus new one)
        assert!(!indices.is_empty());

        // Test that search works by querying for nearest neighbors
        let query_batch = lance_datagen::gen_batch()
            .col("query", array::rand_vec::<Float32Type>(Dimension::from(64)))
            .into_batch_rows(RowCount::from(1))
            .unwrap();

        let q = query_batch.column(0).as_fixed_size_list();
        let mut scanner = dataset.scan();
        scanner
            .nearest("vector", q.value(0).as_primitive::<Float32Type>(), 10)
            .unwrap();
        let results = scanner
            .try_into_stream()
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        assert_eq!(results[0].num_rows(), 10);
    }

    #[tokio::test]
    async fn test_vector_merge_filters_stable_row_id_replacements() {
        const DIMENSION: usize = 4;

        let test_dir = TempStrDir::default();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIMENSION as i32,
                ),
                false,
            ),
        ]));
        let initial_values = (0..40)
            .flat_map(|row| [if row < 20 { 1.0 } else { 0.0 }; DIMENSION])
            .collect::<Vec<_>>();
        let initial = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(0..40)),
                Arc::new(
                    FixedSizeListArray::try_new_from_values(
                        arrow_array::Float32Array::from(initial_values),
                        DIMENSION as i32,
                    )
                    .unwrap(),
                ),
            ],
        )
        .unwrap();
        let mut dataset = Dataset::write(
            RecordBatchIterator::new([Ok(initial)], schema.clone()),
            test_dir.as_str(),
            Some(WriteParams {
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .create_index(
                &["vector"],
                IndexType::Vector,
                Some("vector_idx".to_string()),
                &VectorIndexParams::ivf_flat(1, MetricType::L2),
                true,
            )
            .await
            .unwrap();

        let replacements = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(20..40)),
                Arc::new(
                    FixedSizeListArray::try_new_from_values(
                        arrow_array::Float32Array::from(vec![10.0; 20 * DIMENSION]),
                        DIMENSION as i32,
                    )
                    .unwrap(),
                ),
            ],
        )
        .unwrap();
        let merge_job = MergeInsertBuilder::try_new(Arc::new(dataset), vec!["id".to_string()])
            .unwrap()
            .when_matched(WhenMatched::UpdateAll)
            .try_build()
            .unwrap();
        let (dataset, stats) = merge_job
            .execute(reader_to_stream(Box::new(RecordBatchIterator::new(
                [Ok(replacements)],
                schema,
            ))))
            .await
            .unwrap();
        assert_eq!(stats.num_updated_rows, 20);

        let mut dataset = dataset.as_ref().clone();
        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        assert_eq!(
            dataset
                .load_indices_by_name("vector_idx")
                .await
                .unwrap()
                .len(),
            2
        );
        dataset
            .optimize_indices(&OptimizeOptions::merge(2))
            .await
            .unwrap();

        let logical_index = dataset
            .open_logical_vector_index("vector", "vector_idx")
            .await
            .unwrap();
        assert_eq!(logical_index.num_segments(), 1);
        assert_eq!(
            logical_index
                .num_rows_per_segment()
                .into_iter()
                .map(|(_, rows)| rows)
                .sum::<u64>(),
            40,
            "the merged index must contain one current copy of every stable row id"
        );

        let query = arrow_array::Float32Array::from(vec![0.0; DIMENSION]);
        let result = dataset
            .scan()
            .project(&["id"])
            .unwrap()
            .nearest("vector", &query, 5)
            .unwrap()
            .minimum_nprobes(1)
            .maximum_nprobes(1)
            .try_into_batch()
            .await
            .unwrap();
        let ids = result["id"].as_primitive::<arrow::datatypes::Int32Type>();
        assert!(
            ids.values().iter().all(|id| *id < 20),
            "stale pre-update vectors must not survive the optimize merge: {ids:?}"
        );
    }

    #[tokio::test]
    async fn test_merge_indices_with_unindexed_frags_vector_subset() {
        const DIM: usize = 64;
        const TOTAL: usize = 1000;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let vectors = generate_random_array(TOTAL * DIM);
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIM as i32,
                ),
                true,
            ),
            Field::new("id", DataType::UInt32, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()),
                Arc::new(UInt32Array::from_iter_values(0..TOTAL as u32)),
            ],
        )
        .unwrap();
        let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
        let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap();
        let index_params = VectorIndexParams::ivf_pq(2, 8, 4, MetricType::L2, 2);
        dataset
            .create_index(&["vector"], IndexType::Vector, None, &index_params, true)
            .await
            .unwrap();

        let next_batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(
                    FixedSizeListArray::try_new_from_values(
                        generate_random_array(TOTAL * DIM),
                        DIM as i32,
                    )
                    .unwrap(),
                ),
                Arc::new(UInt32Array::from_iter_values(
                    TOTAL as u32..(TOTAL * 2) as u32,
                )),
            ],
        )
        .unwrap();
        let batches = RecordBatchIterator::new(vec![next_batch].into_iter().map(Ok), schema);
        dataset.append(batches, None).await.unwrap();
        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();

        let indices = dataset.load_indices_by_name("vector_idx").await.unwrap();
        assert_eq!(indices.len(), 2);
        let subset = vec![&indices[1]];
        let merge_result = merge_indices_with_unindexed_frags(
            Arc::new(dataset),
            &subset,
            &[],
            &OptimizeOptions::merge(1),
        )
        .await
        .unwrap();

        assert!(
            merge_result.is_some(),
            "subset merges should respect the caller-provided indices"
        );
    }

    #[tokio::test]
    async fn test_optimize_btree_multi_segment_optimize_default() {
        async fn query_id_count(dataset: &Dataset, id: &str) -> usize {
            dataset
                .scan()
                .filter(&format!("id = '{}'", id))
                .unwrap()
                .project(&["id"])
                .unwrap()
                .try_into_batch()
                .await
                .unwrap()
                .num_rows()
        }

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
        let make_batch = |start: i32, end: i32| {
            let ids = StringArray::from_iter_values((start..end).map(|i| format!("song-{i}")));
            RecordBatch::try_new(schema.clone(), vec![Arc::new(ids)]).unwrap()
        };

        // Three fragments of 64 rows each; each commits as its own BTree
        // segment so optimize sees a multi-segment scalar logical index.
        let reader = RecordBatchIterator::new(
            vec![
                Ok(make_batch(0, 64)),
                Ok(make_batch(64, 128)),
                Ok(make_batch(128, 192)),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::BTree);
        let fragments = dataset.get_fragments();
        assert_eq!(fragments.len(), 3);

        let mut staged_segments = Vec::new();
        for fragment in &fragments {
            let segment = crate::index::create::CreateIndexBuilder::new(
                &mut dataset,
                &["id"],
                IndexType::BTree,
                &params,
            )
            .name("id_idx".into())
            .fragments(vec![fragment.id() as u32])
            .execute_uncommitted()
            .await
            .unwrap();
            staged_segments.push(segment);
        }
        dataset
            .commit_existing_index_segments("id_idx", "id", staged_segments)
            .await
            .unwrap();
        assert_eq!(
            dataset.load_indices_by_name("id_idx").await.unwrap().len(),
            3
        );

        let appended = RecordBatchIterator::new(vec![Ok(make_batch(192, 256))], schema.clone());
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        assert_eq!(dataset.get_fragments().len(), 4);

        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();

        // Reload from disk to ensure we're reading committed manifest state.
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        // Each of these IDs lives in a distinct old segment / fragment.
        // song-10 lives in fragment 0, song-80 in fragment 1, song-160 in
        // fragment 2, and song-200 in the appended fragment. After optimize
        // every row must still be reachable through the logical index,
        // regardless of which segment absorbed the new data.
        for id in ["song-10", "song-80", "song-160", "song-200"] {
            assert_eq!(
                query_id_count(&dataset, id).await,
                1,
                "expected exactly one row for {id} after multi-segment optimize"
            );
        }

        // `OptimizeOptions::default()` (= num_indices_to_merge: None) merges
        // the newest segment with the unindexed fragment, like the
        // inverted/vector default. The three old segments minus the merged one
        // plus the new delta means three segments remain, and together they
        // must still cover every dataset fragment without overlap.
        let segments_after = dataset.load_indices_by_name("id_idx").await.unwrap();
        assert_eq!(
            segments_after.len(),
            3,
            "default optimize must merge one delta, not all segments, got {segments_after:?}"
        );
        let mut covered = RoaringBitmap::new();
        for segment in &segments_after {
            let bitmap = segment
                .fragment_bitmap
                .as_ref()
                .expect("each segment should carry fragment coverage");
            assert!(
                covered.is_disjoint(bitmap),
                "post-optimize segments must not overlap, got {segments_after:?}"
            );
            covered |= bitmap;
        }
        let mut expected = RoaringBitmap::new();
        for frag in dataset.get_fragments() {
            expected.insert(frag.id() as u32);
        }
        assert_eq!(
            covered, expected,
            "post-optimize segments should cover every dataset fragment"
        );
    }

    #[tokio::test]
    async fn test_optimize_fmindex_default_rebuilds_old_and_new_rows() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)]));
        let make_batch = |values: &[&str]| {
            RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(StringArray::from_iter_values(
                    values.iter().copied(),
                ))],
            )
            .unwrap()
        };

        let reader = RecordBatchIterator::new(
            vec![Ok(make_batch(&["old alpha needle", "old beta"]))],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                max_rows_per_file: 2,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Fm);
        dataset
            .create_index(
                &["text"],
                IndexType::Fm,
                Some("text_fmindex".to_string()),
                &params,
                true,
            )
            .await
            .unwrap();

        let appended = RecordBatchIterator::new(
            vec![Ok(make_batch(&["new gamma needle", "new delta"]))],
            schema.clone(),
        );
        dataset.append(appended, None).await.unwrap();

        assert!(
            !dataset
                .unindexed_fragments("text_fmindex")
                .await
                .unwrap()
                .is_empty()
        );

        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert!(
            dataset
                .unindexed_fragments("text_fmindex")
                .await
                .unwrap()
                .is_empty()
        );

        let committed = dataset.load_indices_by_name("text_fmindex").await.unwrap();
        assert_eq!(committed.len(), 1);
        assert_eq!(
            committed[0]
                .fragment_bitmap
                .as_ref()
                .expect("FMIndex segment should carry fragment coverage")
                .len(),
            2
        );

        let logical = crate::index::scalar_logical::open_named_scalar_index(
            &dataset,
            "text",
            "text_fmindex",
            &NoOpMetricsCollector,
        )
        .await
        .unwrap();

        for (pattern, expected) in [("old alpha", 1), ("new gamma", 1), ("needle", 2)] {
            let query = TextQuery::StringContains(pattern.to_string());
            let result = logical.search(&query, &NoOpMetricsCollector).await.unwrap();
            let row_addrs = match result {
                SearchResult::Exact(row_addrs) => row_addrs,
                other => panic!("expected exact result for {pattern}, got {other:?}"),
            };
            let count = row_addrs.true_rows().row_addrs().unwrap().count();
            assert_eq!(
                count, expected,
                "expected {expected} matches for {pattern}, got {count}"
            );
        }
    }

    #[tokio::test]
    async fn test_optimize_ngram_merge_remaps_deferred_compaction() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, true)]));
        let make_batch = |values: &[&str]| {
            RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(StringArray::from_iter_values(
                    values.iter().copied(),
                ))],
            )
            .unwrap()
        };
        let reader = RecordBatchIterator::new(
            vec![
                Ok(make_batch(&["alpha needle", "beta needle"])),
                Ok(make_batch(&["gamma needle", "delta needle"])),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 2,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .create_index(
                &["text"],
                IndexType::NGram,
                Some("text_ngram".into()),
                &ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
                false,
            )
            .await
            .unwrap();

        let metrics = compact_files(
            &mut dataset,
            CompactionOptions {
                target_rows_per_fragment: 10,
                defer_index_remap: true,
                ..Default::default()
            },
            None,
        )
        .await
        .unwrap();
        assert!(metrics.fragments_removed > 0 && metrics.fragments_added > 0);

        let appended =
            RecordBatchIterator::new(vec![Ok(make_batch(&["epsilon needle"]))], schema.clone());
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 2,
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .optimize_indices(&OptimizeOptions::merge(1))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let logical = crate::index::scalar_logical::open_named_scalar_index(
            &dataset,
            "text",
            "text_ngram",
            &NoOpMetricsCollector,
        )
        .await
        .unwrap();
        let result = logical
            .search(
                &TextQuery::StringContains("needle".to_string()),
                &NoOpMetricsCollector,
            )
            .await
            .unwrap();
        let row_addrs = match result {
            SearchResult::AtMost(row_addrs) => row_addrs,
            other => panic!("expected AtMost result from ngram, got {other:?}"),
        };
        assert_eq!(row_addrs.true_rows().row_addrs().unwrap().count(), 5);
    }

    #[tokio::test]
    async fn test_optimize_btree_optimize_append() {
        async fn query_id_count(dataset: &Dataset, id: &str) -> usize {
            dataset
                .scan()
                .filter(&format!("id = '{}'", id))
                .unwrap()
                .project(&["id"])
                .unwrap()
                .try_into_batch()
                .await
                .unwrap()
                .num_rows()
        }

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
        let make_batch = |start: i32, end: i32| {
            let ids = StringArray::from_iter_values((start..end).map(|i| format!("song-{i}")));
            RecordBatch::try_new(schema.clone(), vec![Arc::new(ids)]).unwrap()
        };

        // Start with two fragments + two committed BTree segments.
        let reader = RecordBatchIterator::new(
            vec![Ok(make_batch(0, 64)), Ok(make_batch(64, 128))],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::BTree);
        let original_segment_uuids: Vec<_> = {
            let mut staged = Vec::new();
            for fragment in dataset.get_fragments() {
                let segment = crate::index::create::CreateIndexBuilder::new(
                    &mut dataset,
                    &["id"],
                    IndexType::BTree,
                    &params,
                )
                .name("id_idx".into())
                .fragments(vec![fragment.id() as u32])
                .execute_uncommitted()
                .await
                .unwrap();
                staged.push(segment);
            }
            let uuids = staged.iter().map(|s| s.uuid).collect::<Vec<_>>();
            dataset
                .commit_existing_index_segments("id_idx", "id", staged)
                .await
                .unwrap();
            uuids
        };
        assert_eq!(original_segment_uuids.len(), 2);

        // Append a third fragment, leave it unindexed, then run append-mode optimize.
        let appended = RecordBatchIterator::new(vec![Ok(make_batch(128, 192))], schema.clone());
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();

        // Read fresh from disk to make sure we're inspecting committed state.
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        // append() must preserve every original old segment unchanged and add
        // exactly one new segment covering only the newly appended fragments.
        let committed = dataset.load_indices_by_name("id_idx").await.unwrap();
        let committed_uuids: std::collections::HashSet<_> =
            committed.iter().map(|idx| idx.uuid).collect();
        for original in &original_segment_uuids {
            assert!(
                committed_uuids.contains(original),
                "append() must not remove pre-existing segment {original}, \
                 but the committed UUIDs are {committed_uuids:?}"
            );
        }
        assert_eq!(
            committed.len(),
            original_segment_uuids.len() + 1,
            "append() should add exactly one new delta segment, got {committed:?}"
        );
        let new_segment = committed
            .iter()
            .find(|idx| !original_segment_uuids.contains(&idx.uuid))
            .expect("append() must add a new delta segment");
        let new_segment_frags: Vec<_> = new_segment
            .fragment_bitmap
            .as_ref()
            .unwrap()
            .iter()
            .collect();
        // The appended fragment should be the only one covered by the new delta;
        // old segments retain their own coverage.
        assert_eq!(new_segment_frags.len(), 1);

        // Sanity check: queries across all fragments still return their rows.
        for id in ["song-10", "song-100", "song-160"] {
            assert_eq!(query_id_count(&dataset, id).await, 1, "missing row {id}");
        }
    }

    /// A stray row-id-domain segment left beside a freshly built address-domain
    /// one is not just imprecise, it is wrong: `LogicalScalarIndex` takes the
    /// whole named index's result domain from its first segment alone, so the
    /// new segment's real addresses get treated as row ids wherever they land
    /// outside the old segment's row-id numbering, and those rows vanish from
    /// every query. `select_segments_to_merge` selects nothing in append mode
    /// (`num_indices_to_merge: Some(0)`), so without folding every remaining
    /// legacy segment into the rebuild regardless of what it selected, this is
    /// exactly what append-mode optimize would do under stable row ids.
    #[tokio::test]
    async fn test_optimize_btree_append_rebuilds_stray_legacy_segment_with_stable_row_ids() {
        async fn query_id_count(dataset: &Dataset, id: &str) -> usize {
            dataset
                .scan()
                .filter(&format!("id = '{}'", id))
                .unwrap()
                .project(&["id"])
                .unwrap()
                .try_into_batch()
                .await
                .unwrap()
                .num_rows()
        }

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
        let make_batch = |start: i32, end: i32| {
            let ids = StringArray::from_iter_values((start..end).map(|i| format!("song-{i}")));
            RecordBatch::try_new(schema.clone(), vec![Arc::new(ids)]).unwrap()
        };

        // One fragment, one BTree segment, built under stable row ids so row
        // ids and row addresses are already two independent numbering schemes.
        let reader = RecordBatchIterator::new(vec![Ok(make_batch(0, 64))], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::BTree);
        let original_fragment_id = dataset.get_fragments()[0].id() as u32;
        let segment = crate::index::create::CreateIndexBuilder::new(
            &mut dataset,
            &["id"],
            IndexType::BTree,
            &params,
        )
        .name("id_idx".into())
        .fragments(vec![original_fragment_id])
        .execute_uncommitted()
        .await
        .unwrap();
        assert!(
            segment.results_are_row_addrs(),
            "a segment built by the current code is always address-domain"
        );

        // Relabel it as a pre-migration (row-id-domain) segment: index_version
        // is exactly how `results_are_row_addrs` tells the two apart, and real
        // legacy segments on disk differ from a fresh one only in that label
        // (`BTreeIndexPlugin::load_index` derives domain purely from the
        // version its caller passes in, never from the file itself).
        let legacy_segment = IndexMetadata {
            index_version: 0,
            ..segment
        };
        dataset
            .commit_existing_index_segments("id_idx", "id", vec![legacy_segment])
            .await
            .unwrap();

        // Append a second fragment and leave it unindexed, then optimize in
        // append mode -- the scenario that selects zero old segments to merge.
        let appended = RecordBatchIterator::new(vec![Ok(make_batch(64, 128))], schema.clone());
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        // The stray legacy segment must not survive append-mode optimize: it
        // has to be folded into one consolidated, address-domain segment
        // rather than left beside the new delta.
        let committed = dataset.load_indices_by_name("id_idx").await.unwrap();
        assert_eq!(
            committed.len(),
            1,
            "the legacy segment must be rebuilt together with the new data, \
             not left beside it: {committed:?}"
        );
        assert!(
            committed[0].results_are_row_addrs(),
            "the rebuilt segment must be address-domain"
        );

        // Rows from both the original and the newly appended fragment must be
        // findable -- in particular rows only the new segment's data covers,
        // which a stray legacy segment reporting the wrong domain would drop.
        for id in ["song-10", "song-100", "song-127"] {
            assert_eq!(query_id_count(&dataset, id).await, 1, "missing row {id}");
        }
    }

    /// A trailing merge window must also rebuild legacy segments outside it.
    /// Use a real pre-migration index: committing a mixture of row domains is
    /// rejected before optimization can run.
    #[tokio::test]
    async fn test_optimize_btree_merge_rebuilds_legacy_segment_outside_merge_window() {
        let dir = copy_test_data_to_tmp("v14.0.0-beta.15/stable_row_id_btree.lance").unwrap();
        let uri = dir.path_str();
        let mut dataset = Dataset::open(&uri).await.unwrap();
        assert!(dataset.manifest.uses_stable_row_ids());
        let legacy = dataset.load_indices_by_name("value_idx").await.unwrap();
        assert_eq!(legacy.len(), 2);
        assert!(legacy.iter().all(|index| !index.results_are_row_addrs()));

        let batch = arrow_array::record_batch!(("value", Int32, [Some(20), None])).unwrap();
        let schema = batch.schema();
        dataset
            .append(RecordBatchIterator::new([Ok(batch)], schema), None)
            .await
            .unwrap();
        assert_eq!(dataset.get_fragments().len(), 3);

        // One of the two legacy segments is outside this trailing window.
        // Both must join the rebuild before its address-domain result commits.
        dataset
            .optimize_indices(&OptimizeOptions::merge(1))
            .await
            .unwrap();

        let dataset = Dataset::open(&uri).await.unwrap();
        let committed = dataset.load_indices_by_name("value_idx").await.unwrap();
        assert_eq!(
            committed.len(),
            1,
            "the legacy segment outside the merge window must join the rebuild: {committed:?}"
        );
        assert!(committed[0].results_are_row_addrs());
        assert_eq!(committed[0].fragment_bitmap.as_ref().unwrap().len(), 3);

        // Check both original fragments and the appended rows, including NULLs.
        for (filter, expected) in [
            ("value = 1", 1),
            ("value = 9", 1),
            ("value = 20", 1),
            ("value IS NULL", 3),
        ] {
            for indexed in [false, true] {
                let mut scan = dataset.scan();
                scan.use_scalar_index(indexed).filter(filter).unwrap();
                assert_eq!(
                    scan.try_into_batch().await.unwrap().num_rows(),
                    expected,
                    "{filter}, indexed={indexed}"
                );
            }
        }
    }

    fn json_btree_schema_and_batch(ids: std::ops::Range<i32>) -> (Arc<Schema>, RecordBatch) {
        // A plain Utf8 field tagged with the `arrow.json` extension is read as
        // JSON text (`JsonEncoding::of_field`) -- simpler to construct by hand
        // than genuine JSONB bytes, and exercised the same way by
        // `test_json_text_input_extracts_like_jsonb` in lance-index.
        let doc_field = Field::new("doc", DataType::Utf8, false).with_metadata(
            std::collections::HashMap::from([(
                lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
                lance_arrow::json::ARROW_JSON_EXT_NAME.to_string(),
            )]),
        );
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            doc_field,
        ]));
        let doc_values: Vec<String> = ids.clone().map(|i| format!(r#"{{"x":{i}}}"#)).collect();
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(ids)),
                Arc::new(StringArray::from(doc_values)),
            ],
        )
        .unwrap();
        (schema, batch)
    }

    fn json_btree_params() -> ScalarIndexParams {
        ScalarIndexParams {
            index_type: "json".to_string(),
            params: Some(r#"{"target_index_type":"btree","path":"x"}"#.to_string()),
        }
    }

    /// A JSON-wrapped BTREE index cannot yet carry the row-address-domain
    /// migration through its own metadata (see the guard in
    /// `super::super::scalar::build_scalar_index`), so creating one on a
    /// stable-row-id dataset must be rejected rather than silently produce an
    /// address-domain segment none of the domain-safety checks elsewhere can
    /// see through the wrapper.
    #[tokio::test]
    async fn test_json_wrapped_btree_create_rejected_with_stable_row_ids() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let (schema, batch) = json_btree_schema_and_batch(0..10);
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = json_btree_params();
        let err = dataset
            .create_index_builder(&["doc"], IndexType::Scalar, &params)
            .name("doc_idx".to_string())
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("JSON-wrapped BTREE"),
            "unexpected error: {err}"
        );
        assert!(
            dataset
                .load_indices_by_name("doc_idx")
                .await
                .unwrap()
                .is_empty(),
            "the rejected index must not be committed"
        );
    }

    /// The same combination is unaffected -- and must keep working exactly as
    /// before this guard was added -- on a dataset that does not use stable
    /// row ids, where a row id and a row address are the same value.
    #[tokio::test]
    async fn test_json_wrapped_btree_create_and_optimize_without_stable_row_ids() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let (schema, batch) = json_btree_schema_and_batch(0..10);
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(reader, test_uri, Some(WriteParams::default()))
            .await
            .unwrap();

        let params = json_btree_params();
        dataset
            .create_index_builder(&["doc"], IndexType::Scalar, &params)
            .name("doc_idx".to_string())
            .await
            .unwrap();
        assert_eq!(
            dataset.load_indices_by_name("doc_idx").await.unwrap().len(),
            1
        );

        // Maintenance (insert + optimize) on the same, still-unaffected
        // combination must keep working too.
        let (_, more) = json_btree_schema_and_batch(10..20);
        let appended = RecordBatchIterator::new(vec![Ok(more)], schema);
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        assert_eq!(
            dataset.load_indices_by_name("doc_idx").await.unwrap().len(),
            1,
            "optimize should merge into a single segment, not fail or leave extras"
        );
    }

    #[tokio::test]
    async fn test_optimize_bitmap_index_append() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new(
            "category",
            DataType::Utf8,
            false,
        )]));
        let make_batch = |labels: &[&str]| {
            let arr = StringArray::from_iter_values(labels.iter().copied());
            RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
        };

        // One fragment + one Bitmap segment.
        let reader =
            RecordBatchIterator::new(vec![Ok(make_batch(&["a", "b", "a", "c"]))], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 4,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::Bitmap);
        dataset
            .create_index(
                &["category"],
                IndexType::Bitmap,
                Some("cat_idx".into()),
                &params,
                true,
            )
            .await
            .unwrap();
        let original_uuid = {
            let committed = dataset.load_indices_by_name("cat_idx").await.unwrap();
            assert_eq!(committed.len(), 1);
            committed[0].uuid
        };

        // Append a second fragment, leave it unindexed, then optimize with
        // `append()` (= num_indices_to_merge: Some(0)).
        let appended =
            RecordBatchIterator::new(vec![Ok(make_batch(&["b", "d", "d", "a"]))], schema.clone());
        let mut dataset = Dataset::write(
            appended,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 4,
                mode: WriteMode::Append,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        // append() (= num_indices_to_merge: Some(0)) is now honored uniformly:
        // Bitmap, like BTree, must keep the original segment untouched and add
        // exactly one delta segment covering only the appended fragment.
        let committed = dataset.load_indices_by_name("cat_idx").await.unwrap();
        assert_eq!(
            committed.len(),
            2,
            "Bitmap optimize append() must add a delta segment, not merge, got {committed:?}"
        );
        assert!(
            committed.iter().any(|idx| idx.uuid == original_uuid),
            "append() must preserve the pre-existing segment {original_uuid}, got {committed:?}"
        );
        let new_segment = committed
            .iter()
            .find(|idx| idx.uuid != original_uuid)
            .expect("append() must add a new delta segment");
        let new_segment_frags: std::collections::BTreeSet<u32> = new_segment
            .fragment_bitmap
            .as_ref()
            .expect("delta Bitmap should carry fragment coverage")
            .iter()
            .collect();
        assert_eq!(
            new_segment_frags,
            [1u32].into_iter().collect(),
            "the delta segment must cover only the appended fragment"
        );

        // Data correctness: a value that lives only in the appended fragment
        // is queryable through the (now multi-segment) index.
        let rows = dataset
            .scan()
            .filter("category = 'd'")
            .unwrap()
            .project(&["category"])
            .unwrap()
            .try_into_batch()
            .await
            .unwrap()
            .num_rows();
        assert_eq!(rows, 2, "value 'd' lives in appended fragment");
    }

    #[tokio::test]
    async fn test_optimize_append_with_backpressured_multichunk_read() {
        const WIDE_DIMENSION: usize = 140_000;
        const NARROW_DIMENSION: usize = 4_096;
        const SHORT_ROWS: usize = 68;
        const LONG_ROWS: usize = 128;

        fn make_batch(schema: Arc<Schema>, rows: usize, base: usize) -> RecordBatch {
            let docs_field = schema.field(1);
            let DataType::List(item_field) = docs_field.data_type() else {
                unreachable!("docs must be a list");
            };
            let DataType::Struct(doc_fields) = item_field.data_type() else {
                unreachable!("docs items must be structs");
            };
            let wide_values = Float32Array::from_iter_values(
                (0..rows * WIDE_DIMENSION).map(|index| ((index + base) % 1009) as f32),
            );
            let narrow_values = Float32Array::from_iter_values(
                (0..rows * NARROW_DIMENSION).map(|index| ((index + base) % 251) as f32),
            );
            let docs = StructArray::new(
                doc_fields.clone(),
                vec![
                    Arc::new(
                        FixedSizeListArray::try_new_from_values(wide_values, WIDE_DIMENSION as i32)
                            .unwrap(),
                    ),
                    Arc::new(
                        FixedSizeListArray::try_new_from_values(
                            narrow_values,
                            NARROW_DIMENSION as i32,
                        )
                        .unwrap(),
                    ),
                    Arc::new(StringArray::from_iter_values(
                        (0..rows).map(|index| format!("document-{}", index + base)),
                    )),
                ],
                None,
            );
            let docs = ListArray::new(
                item_field.clone(),
                OffsetBuffer::from_lengths(std::iter::repeat_n(1, rows)),
                Arc::new(docs),
                None,
            );
            let ids = UInt64Array::from_iter_values(base as u64..(base + rows) as u64);
            RecordBatch::try_new(schema, vec![Arc::new(ids), Arc::new(docs)]).unwrap()
        }

        let vector_item = Arc::new(Field::new("item", DataType::Float32, true));
        let doc_fields = vec![
            Field::new(
                "wide",
                DataType::FixedSizeList(vector_item.clone(), WIDE_DIMENSION as i32),
                true,
            ),
            Field::new(
                "narrow",
                DataType::FixedSizeList(vector_item, NARROW_DIMENSION as i32),
                true,
            ),
            Field::new("text", DataType::Utf8, true),
        ]
        .into();
        let docs_item = Arc::new(Field::new("item", DataType::Struct(doc_fields), true));
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt64, false),
            Field::new("docs", DataType::List(docs_item), true),
        ]));

        let initial_batch = make_batch(schema.clone(), 1, 1_000);
        let initial_reader = RecordBatchIterator::new(vec![Ok(initial_batch)], schema.clone());
        let mut dataset = Dataset::write(
            initial_reader,
            "memory://",
            Some(WriteParams {
                data_storage_version: Some(LanceFileVersion::V2_1),
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        dataset
            .create_index(
                &["docs.text"],
                IndexType::Inverted,
                Some("docs_text".to_string()),
                &InvertedIndexParams::default(),
                true,
            )
            .await
            .unwrap();

        let appended_batches = vec![
            Ok(make_batch(schema.clone(), SHORT_ROWS, 0)),
            Ok(make_batch(schema.clone(), LONG_ROWS, SHORT_ROWS)),
        ];
        let appended_reader = RecordBatchIterator::new(appended_batches, schema);
        dataset
            .append(
                appended_reader,
                Some(WriteParams {
                    data_storage_version: Some(LanceFileVersion::V2_1),
                    ..Default::default()
                }),
            )
            .await
            .unwrap();
        // The deleted first batch makes the live delta start exactly at the
        // second write-batch boundary, matching a merge-insert style update.
        dataset.delete("id < 68").await.unwrap();

        let optimize_options = OptimizeOptions::append();
        // Nested FTS reads the entire `docs` root. Under this budget the wide
        // sibling is split into same-priority chunks, then the narrow sibling's
        // higher-priority I/O consumes the remaining budget while the wide read
        // is awaited. Admitted chunks must continue despite that backpressure.
        let optimize = crate::index::scalar::TEST_TRAINING_IO_BUFFER_SIZE.scope(
            70 * 1024 * 1024,
            dataset.optimize_indices(&optimize_options),
        );
        tokio::time::timeout(std::time::Duration::from_secs(20), optimize)
            .await
            .expect("incremental nested FTS optimization timed out")
            .unwrap();

        let segments = dataset
            .load_indices()
            .await
            .unwrap()
            .iter()
            .filter(|index| index.name == "docs_text")
            .count();
        assert_eq!(segments, 2, "optimization should add one FTS delta segment");
    }

    /// BTree stores row addresses (`BTREE_ROW_ADDR_DOMAIN_VERSION`), so a plain
    /// eager compaction on a stable-row-id dataset needs explicit repair, unlike
    /// a row-id-domain index whose ids survive the rewrite for free. An eager
    /// commit remaps the btree in place -- the same way it would on a
    /// non-stable-row-id table -- and writes no fragment reuse index.
    /// `defer_index_remap` instead leaves the btree untouched and writes an FRI
    /// that lets it follow its data to the new fragment at read time; draining
    /// that FRI via `optimize_indices` (a segment merge, not a rescan) must not
    /// change any answer either.
    #[rstest]
    #[case::eager(false)]
    #[case::deferred(true)]
    #[tokio::test]
    async fn test_optimize_btree_keeps_rows_with_stable_row_ids_after_compaction(
        #[case] defer_index_remap: bool,
    ) {
        async fn query_id_count(dataset: &Dataset, id: &str) -> usize {
            dataset
                .scan()
                .filter(&format!("id = '{}'", id))
                .unwrap()
                .project(&["id"])
                .unwrap()
                .try_into_batch()
                .await
                .unwrap()
                .num_rows()
        }

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
        let ids = StringArray::from_iter_values((0..256).map(|i| format!("song-{i}")));
        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids)]).unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 64,
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .create_index(
                &["id"],
                IndexType::BTree,
                Some("id_idx".into()),
                &ScalarIndexParams::default(),
                true,
            )
            .await
            .unwrap();

        assert_eq!(query_id_count(&dataset, "song-42").await, 1);

        compact_files(
            &mut dataset,
            crate::dataset::optimize::CompactionOptions {
                target_rows_per_fragment: 512,
                defer_index_remap,
                ..Default::default()
            },
            None,
        )
        .await
        .unwrap();

        let has_fri = dataset
            .load_index_by_name(lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME)
            .await
            .unwrap()
            .is_some();
        assert_eq!(
            has_fri, defer_index_remap,
            "only a deferred compaction over an address-domain index writes an FRI; an \
             eager one remaps the btree in place instead"
        );

        let frags = dataset.get_fragments();
        assert!(!frags.is_empty());
        assert!(frags.iter().all(|frag| frag.id() > 0));
        assert!(
            dataset
                .unindexed_fragments("id_idx")
                .await
                .unwrap()
                .is_empty(),
            "the btree must cover the rewritten fragment, either because the FRI lets it \
             follow its data (deferred) or because the eager commit remapped it in place"
        );
        assert_eq!(query_id_count(&dataset, "song-42").await, 1);

        // Draining the FRI into the index (a segment merge, not a rescan), or
        // simply reopening an eagerly-remapped index, must not change any answer.
        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert_eq!(query_id_count(&dataset, "song-42").await, 1);
        for i in [0, 42, 128, 255] {
            assert_eq!(query_id_count(&dataset, &format!("song-{i}")).await, 1);
        }
    }

    /// Updating an indexed vector column in place (`update_columns` +
    /// `Operation::Update`) keeps the fragment id and the row address, so after the
    /// fragment is pruned from the old segment's bitmap that segment still physically
    /// holds the pre-update vector. Merging the segment must drop those rows: the merge
    /// concatenates its physical rows with the freshly scanned ones and commits the
    /// union of both coverages, so a surviving stale row is authorized by the same
    /// bitmap as its fresh copy and no read-time filter can separate them.
    #[rstest]
    // IVF_FLAT keeps the vector verbatim, so the query distance tells which of the two
    // copies survived. IVF_PQ re-uses the codebook trained before the update, which
    // cannot represent the new value, so it only pins the row count and covers the
    // transposed-code path in `take_partition_batches`.
    #[case::ivf_flat(VectorIndexParams::ivf_flat(1, MetricType::L2), true)]
    #[case::ivf_pq(VectorIndexParams::ivf_pq(1, 8, 16, MetricType::L2, 2), false)]
    #[tokio::test]
    async fn test_optimize_vector_index_drops_stale_rows_on_merge(
        #[case] index_params: VectorIndexParams,
        #[case] distance_is_exact: bool,
        // The two row-id schemes take different filter branches: an address carries its
        // fragment, a stable row id needs the fragments' persisted row-id sequences.
        #[values(false, true)] enable_stable_row_ids: bool,
    ) {
        use crate::dataset::transaction::{Operation, UpdateMode, UpdatedFragmentOffsets};
        use arrow::datatypes::UInt64Type;
        use arrow_array::{Float32Array, RecordBatchReader, UInt64Array};
        use lance_core::ROW_ID;
        use lance_index::vector::DIST_COL;
        use std::collections::HashMap;

        // Must match the sub-vector count in the IVF_PQ case above.
        const DIM: usize = 16;
        const BULK_ROWS: usize = 1000;
        const UPDATED_ID: u32 = 10_000;
        const STALE_VALUE: f32 = 2.0;
        const FRESH_VALUE: f32 = 10.8;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let vector_type = DataType::FixedSizeList(
            Arc::new(Field::new("item", DataType::Float32, true)),
            DIM as i32,
        );
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("vector", vector_type.clone(), true),
        ]));
        let constant_vector = |value: f32| {
            FixedSizeListArray::try_new_from_values(
                Float32Array::from_iter_values(std::iter::repeat_n(value, DIM)),
                DIM as i32,
            )
            .unwrap()
        };

        // Fragment 0: bulk vectors in [0, 1). None of them is near the query, so a
        // surviving stale copy of the updated row ranks well inside top-k.
        let bulk = generate_random_array_with_seed::<Float32Type>(BULK_ROWS * DIM, [42; 32]);
        let bulk_batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt32Array::from_iter_values(0..BULK_ROWS as u32)),
                Arc::new(FixedSizeListArray::try_new_from_values(bulk, DIM as i32).unwrap()),
            ],
        )
        .unwrap();
        let mut dataset = Dataset::write(
            RecordBatchIterator::new(vec![Ok(bulk_batch)], schema.clone()),
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        // Fragment 1: the single row that gets updated in place.
        let stale_batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt32Array::from_iter_values([UPDATED_ID])),
                Arc::new(constant_vector(STALE_VALUE)),
            ],
        )
        .unwrap();
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(stale_batch)], schema.clone()),
                None,
            )
            .await
            .unwrap();
        assert_eq!(dataset.get_fragments().len(), 2);

        // One segment covering both fragments.
        dataset
            .create_index(&["vector"], IndexType::Vector, None, &index_params, true)
            .await
            .unwrap();

        let mut fragment = dataset.get_fragment(1).unwrap();
        let mut fragment_scan = fragment.scan();
        fragment_scan.with_row_id();
        let updated_row_id = fragment_scan.try_into_batch().await.unwrap()[ROW_ID]
            .as_primitive::<UInt64Type>()
            .value(0);

        let update_schema = Arc::new(Schema::new(vec![
            Field::new(ROW_ID, DataType::UInt64, false),
            Field::new("vector", vector_type, true),
        ]));
        let update_batch = RecordBatch::try_new(
            update_schema.clone(),
            vec![
                Arc::new(UInt64Array::from_iter_values([updated_row_id])),
                Arc::new(constant_vector(FRESH_VALUE)),
            ],
        )
        .unwrap();
        let right_stream: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
            vec![Ok(update_batch)],
            update_schema,
        ));
        let updated = fragment
            .update_columns_with_offsets(right_stream, ROW_ID, ROW_ID)
            .await
            .unwrap();
        let updated_fragment_id = updated.fragment.id;
        let dataset = Dataset::commit(
            test_uri,
            Operation::Update {
                removed_fragment_ids: vec![],
                updated_fragments: vec![updated.fragment],
                new_fragments: vec![],
                fields_modified: updated.fields_modified,
                compacted_sstables: Vec::new(),
                fields_for_preserving_frag_bitmap: vec![],
                update_mode: Some(UpdateMode::RewriteColumns),
                inserted_rows_filter: None,
                updated_fragment_offsets: Some(UpdatedFragmentOffsets(HashMap::from([(
                    updated_fragment_id,
                    updated.matched_offsets,
                )]))),
            },
            Some(dataset.version().version),
            None,
            None,
            Default::default(),
            true,
        )
        .await
        .unwrap();

        let mut dataset = dataset;
        dataset
            .optimize_indices(&OptimizeOptions::merge(1))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let segments = dataset.load_indices_by_name("vector_idx").await.unwrap();
        assert_eq!(segments.len(), 1, "merge must leave a single segment");

        let mut scanner = dataset.scan();
        scanner
            .nearest("vector", &constant_vector(FRESH_VALUE).value(0), 10)
            .unwrap()
            .with_row_id()
            .project(&["id"])
            .unwrap();
        let results = scanner.try_into_batch().await.unwrap();
        let hits = results["id"]
            .as_primitive::<UInt32Type>()
            .values()
            .iter()
            .zip(results[DIST_COL].as_primitive::<Float32Type>().values())
            .filter(|(id, _)| **id == UPDATED_ID)
            .map(|(_, distance)| *distance)
            .collect::<Vec<_>>();
        assert_eq!(
            hits.len(),
            1,
            "updated row returned {} times after merge; row ids = {:?}",
            hits.len(),
            results[ROW_ID].as_primitive::<UInt64Type>().values()
        );
        if distance_is_exact {
            // The pre-update vector would sit at (FRESH - STALE)^2 * DIM from the query,
            // so anything but ~0 means the stale copy is the one that survived.
            assert!(
                hits[0] < 1.0,
                "surviving copy is the pre-update one: distance {} to the updated vector",
                hits[0]
            );
        }
    }

    /// The ordinary update path (`UpdateBuilder`, `UpdateMode::RewriteRows`) commits
    /// `fields_modified: vec![]`, so the old fragment keeps its place in the segment's
    /// bitmap while its physical row is only deletion-marked. Under stable row ids the
    /// rewritten copy reuses the same row id, so a merge that keeps every covered row
    /// emits that id twice and the read-time filter admits both: the id is live again
    /// at its new address. Coverage alone is therefore not a sufficient merge filter.
    #[rstest]
    // The two optimize modes read existing rows through different code, and both must
    // filter: an explicit merge through `take_partition_batches`, and the default
    // options over under-sized partitions through `partition_row_ids`, which feeds the
    // join that rebuilds them.
    #[case::merge(1, OptimizeOptions::merge(1))]
    #[case::join(4, OptimizeOptions::new())]
    #[tokio::test]
    async fn test_optimize_vector_index_drops_rewritten_rows_on_merge(
        #[case] num_partitions: usize,
        #[case] optimize_options: OptimizeOptions,
        // Only the stable-row-id scheme can duplicate: an address-domain rewrite lands
        // at a new address, so the stale posting is masked at query time.
        #[values(false, true)] enable_stable_row_ids: bool,
    ) {
        use crate::dataset::UpdateBuilder;
        use arrow::datatypes::UInt64Type;
        use arrow_array::Float32Array;
        use lance_core::ROW_ID;
        use lance_index::vector::DIST_COL;

        const DIM: usize = 4;
        const BULK_ROWS: usize = 1000;
        const UPDATED_ID: u32 = 10_000;
        const STALE_VALUE: f32 = 2.0;
        const FRESH_VALUE: f32 = 10.8;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    DIM as i32,
                ),
                true,
            ),
        ]));
        let constant_vector = |value: f32| {
            FixedSizeListArray::try_new_from_values(
                Float32Array::from_iter_values(std::iter::repeat_n(value, DIM)),
                DIM as i32,
            )
            .unwrap()
        };

        // The updated row shares its fragment with the bulk rows, so the update leaves
        // that fragment covered by the segment with one row deletion-marked instead of
        // retiring it. The bulk vectors sit in [0, 1), far from the query, so a
        // surviving stale copy of the updated row ranks well inside top-k.
        let bulk = generate_random_array_with_seed::<Float32Type>(BULK_ROWS * DIM, [42; 32]);
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt32Array::from_iter_values(
                    (0..BULK_ROWS as u32).chain(std::iter::once(UPDATED_ID)),
                )),
                Arc::new(
                    FixedSizeListArray::try_new_from_values(
                        Float32Array::from_iter_values(
                            bulk.values()
                                .iter()
                                .copied()
                                .chain(std::iter::repeat_n(STALE_VALUE, DIM)),
                        ),
                        DIM as i32,
                    )
                    .unwrap(),
                ),
            ],
        )
        .unwrap();
        let mut dataset = Dataset::write(
            RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .create_index(
                &["vector"],
                IndexType::Vector,
                None,
                &VectorIndexParams::ivf_flat(num_partitions, MetricType::L2),
                true,
            )
            .await
            .unwrap();

        let fresh_literal = format!(
            "array[{}]",
            std::iter::repeat_n(FRESH_VALUE.to_string(), DIM)
                .collect::<Vec<_>>()
                .join(", ")
        );
        let mut dataset = UpdateBuilder::new(Arc::new(dataset))
            .update_where(&format!("id = {UPDATED_ID}"))
            .unwrap()
            .set("vector", &fresh_literal)
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap()
            .new_dataset
            .as_ref()
            .clone();
        dataset.optimize_indices(&optimize_options).await.unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let segments = dataset.load_indices_by_name("vector_idx").await.unwrap();
        assert_eq!(segments.len(), 1, "merge must leave a single segment");
        let partitions_after = dataset
            .open_vector_index("vector", &segments[0].uuid, &NoOpMetricsCollector)
            .await
            .unwrap()
            .ivf_model()
            .num_partitions();
        // A join is what routes existing rows through `partition_row_ids`; without it
        // this case would silently degrade into a second copy of the merge case.
        assert!(
            num_partitions == 1 || partitions_after < num_partitions,
            "expected a join, but the index still has {partitions_after} partitions"
        );

        let mut scanner = dataset.scan();
        scanner
            .nearest("vector", &constant_vector(FRESH_VALUE).value(0), 10)
            .unwrap()
            .with_row_id()
            .project(&["id"])
            .unwrap();
        let results = scanner.try_into_batch().await.unwrap();
        let hits = results["id"]
            .as_primitive::<UInt32Type>()
            .values()
            .iter()
            .zip(results[DIST_COL].as_primitive::<Float32Type>().values())
            .filter(|(id, _)| **id == UPDATED_ID)
            .map(|(_, distance)| *distance)
            .collect::<Vec<_>>();
        assert_eq!(
            hits.len(),
            1,
            "rewritten row returned {} times after optimize; row ids = {:?}",
            hits.len(),
            results[ROW_ID].as_primitive::<UInt64Type>().values()
        );
        // The pre-update vector would sit at (FRESH - STALE)^2 * DIM from the query, so
        // anything but ~0 means the stale copy is the one that survived.
        assert!(
            hits[0] < 1.0,
            "surviving copy is the pre-update one: distance {} to the updated vector",
            hits[0]
        );
    }

    /// Under stable row ids, updating an indexed column and then calling
    /// `optimize_indices` must not leave stale entries (old value -> updated row)
    /// in the scalar index. An update deletes the old copy of each row and
    /// rewrites it under the same stable row id, so the old index entry is stale
    /// and must be dropped on merge. Covers BTree, Bitmap, and Inverted (FTS),
    /// which take three different merge paths.
    #[tokio::test]
    async fn test_optimize_scalar_index_drops_stale_rows_after_update() {
        use crate::dataset::UpdateBuilder;
        use arrow_array::Int32Array;
        use lance_index::scalar::FullTextSearchQuery;
        use lance_index::scalar::inverted::InvertedIndexParams;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        // 100 rows: num == id; cat = "A" for id<50 else "B"; body = "alpha" for
        // id<50 else "beta".
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
            Field::new("cat", DataType::Utf8, false),
            Field::new("body", DataType::Utf8, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(0..100)),
                Arc::new(Int32Array::from_iter_values(0..100)),
                Arc::new(StringArray::from_iter_values(
                    (0..100).map(|i| if i < 50 { "A" } else { "B" }),
                )),
                Arc::new(StringArray::from_iter_values(
                    (0..100).map(|i| if i < 50 { "alpha" } else { "beta" }),
                )),
            ],
        )
        .unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .create_index(
                &["num"],
                IndexType::BTree,
                None,
                &ScalarIndexParams::default(),
                true,
            )
            .await
            .unwrap();
        dataset
            .create_index(
                &["cat"],
                IndexType::Bitmap,
                None,
                &ScalarIndexParams::default(),
                true,
            )
            .await
            .unwrap();
        dataset
            .create_index(
                &["body"],
                IndexType::Inverted,
                None,
                &InvertedIndexParams::default(),
                true,
            )
            .await
            .unwrap();

        // Update the first 25 rows (id < 25): num -> -1, cat -> 'B', body -> 'beta'.
        let res = UpdateBuilder::new(Arc::new(dataset.clone()))
            .update_where("id < 25")
            .unwrap()
            .set("num", "-1")
            .unwrap()
            .set("cat", "'B'")
            .unwrap()
            .set("body", "'beta'")
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();
        dataset = res.new_dataset.as_ref().clone();

        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        // BTree: `num >= 0` matches ids 25..99 (75 rows); the 25 updated rows
        // hold num = -1 and must not appear.
        let btree_count = dataset
            .scan()
            .filter("num >= 0")
            .unwrap()
            .count_rows()
            .await
            .unwrap();
        assert_eq!(btree_count, 75, "btree returned stale/incorrect rows");

        // Bitmap: only the 25 rows (ids 25..49) that still carry cat = 'A' match;
        // the 25 rows updated to 'B' must not.
        let bitmap_count = dataset
            .scan()
            .filter("cat = 'A'")
            .unwrap()
            .count_rows()
            .await
            .unwrap();
        assert_eq!(bitmap_count, 25, "bitmap returned stale rows");

        // FTS: only the 25 rows (ids 25..49) whose body still reads "alpha" match;
        // the 25 rows updated to "beta" must not.
        let mut scan = dataset.scan();
        scan.full_text_search(FullTextSearchQuery::new("alpha".to_owned()))
            .unwrap();
        let fts_count = scan.count_rows().await.unwrap();
        assert_eq!(fts_count, 25, "FTS index returned stale rows");
    }

    /// Multi-segment variant (Jack Ye's repro, PR #7359): with one BTree segment
    /// per fragment, default optimize merges only the tail segment. A stable-row-id
    /// update to a row in an older segment's fragment must still drop that
    /// segment's stale postings -- the merge has to reach back to cover it.
    #[tokio::test]
    async fn test_optimize_btree_drops_stale_rows_across_segments_after_update() {
        use crate::dataset::UpdateBuilder;
        use crate::index::CreateIndexBuilder;
        use arrow_array::Int32Array;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(0..100)),
                Arc::new(Int32Array::from_iter_values(0..100)),
            ],
        )
        .unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        // Two fragments (0..49, 50..99) -> one BTree segment each.
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                max_rows_per_file: 50,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree);
        let fragments = dataset.get_fragments();
        let mut segments = Vec::new();
        for fragment in &fragments {
            segments.push(
                CreateIndexBuilder::new(&mut dataset, &["num"], IndexType::BTree, &params)
                    .name("num_idx".to_string())
                    .fragments(vec![fragment.id() as u32])
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        dataset
            .commit_existing_index_segments("num_idx", "num", segments)
            .await
            .unwrap();

        // Update the first 25 rows (in the first/older segment's fragment).
        let res = UpdateBuilder::new(Arc::new(dataset.clone()))
            .update_where("id < 25")
            .unwrap()
            .set("num", "-1")
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();
        dataset = res.new_dataset.as_ref().clone();

        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        assert_eq!(
            dataset
                .scan()
                .filter("num = 0")
                .unwrap()
                .count_rows()
                .await
                .unwrap(),
            0,
            "stale entry leaked from the older, unmerged segment"
        );
        assert_eq!(
            dataset
                .scan()
                .filter("num >= 0")
                .unwrap()
                .count_rows()
                .await
                .unwrap(),
            75
        );
    }

    /// Same multi-segment gap for FTS, which takes the separate Inverted dispatch
    /// path. One Inverted segment per fragment; an update to the older segment's
    /// fragment must not leave its old-token postings searchable.
    #[tokio::test]
    async fn test_optimize_fts_drops_stale_rows_across_segments_after_update() {
        use crate::dataset::UpdateBuilder;
        use crate::index::CreateIndexBuilder;
        use arrow_array::Int32Array;
        use lance_index::scalar::FullTextSearchQuery;
        use lance_index::scalar::inverted::InvertedIndexParams;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("body", DataType::Utf8, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(0..100)),
                Arc::new(StringArray::from_iter_values(
                    (0..100).map(|i| if i < 50 { "alpha" } else { "beta" }),
                )),
            ],
        )
        .unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                max_rows_per_file: 50,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        let params = InvertedIndexParams::default();
        let fragments = dataset.get_fragments();
        let mut segments = Vec::new();
        for fragment in &fragments {
            segments.push(
                CreateIndexBuilder::new(&mut dataset, &["body"], IndexType::Inverted, &params)
                    .name("body_idx".to_string())
                    .fragments(vec![fragment.id() as u32])
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        dataset
            .commit_existing_index_segments("body_idx", "body", segments)
            .await
            .unwrap();

        // Update the first 25 rows (older segment's fragment): body -> "beta".
        let res = UpdateBuilder::new(Arc::new(dataset.clone()))
            .update_where("id < 25")
            .unwrap()
            .set("body", "'beta'")
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();
        dataset = res.new_dataset.as_ref().clone();

        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();

        let mut scan = dataset.scan();
        scan.full_text_search(FullTextSearchQuery::new("alpha".to_owned()))
            .unwrap();
        assert_eq!(
            scan.count_rows().await.unwrap(),
            25,
            "FTS stale rows leaked from the older, unmerged segment"
        );
    }

    /// `optimize_indices` builds the stable-row-id allow-list by subtracting each
    /// fragment's deletion vector. If a deletion vector cannot be read, the merge
    /// must fail loudly: swallowing the error (treating the load as "no
    /// deletions") would put every deleted row back into the allow-list and
    /// silently reintroduce the stale entries this fix removes. Simulate an
    /// unreadable deletion vector by deleting the file the manifest still
    /// references, then assert optimize errors instead of succeeding.
    #[tokio::test]
    async fn test_optimize_errors_when_deletion_vector_unreadable() {
        use crate::dataset::UpdateBuilder;
        use arrow_array::Int32Array;
        use lance_index::scalar::BuiltinIndexType;
        use lance_table::io::deletion::deletion_file_path;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(0..100)),
                Arc::new(Int32Array::from_iter_values(0..100)),
            ],
        )
        .unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();

        dataset
            .create_index(
                &["num"],
                IndexType::Bitmap,
                None,
                &ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
                true,
            )
            .await
            .unwrap();

        // Update rewrites the first 25 rows under the same stable row ids,
        // leaving a deletion vector on the original fragment.
        UpdateBuilder::new(Arc::new(dataset.clone()))
            .update_where("id < 25")
            .unwrap()
            .set("num", "-1")
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();

        // Reload cold (nothing has cached the deletion vector), then remove the
        // deletion file the manifest still references so the next read fails.
        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let mut removed = 0;
        for fragment in dataset.get_fragments() {
            if let Some(deletion_file) = fragment.metadata().deletion_file.clone() {
                let path =
                    deletion_file_path(&dataset.base, fragment.metadata().id, &deletion_file);
                dataset.object_store.delete(&path).await.unwrap();
                removed += 1;
            }
        }
        assert_eq!(
            removed, 1,
            "update should have left exactly one deletion file"
        );

        let result = dataset.optimize_indices(&OptimizeOptions::default()).await;
        assert!(
            result.is_err(),
            "optimize must fail when a deletion vector cannot be read, not \
             silently keep the deleted rows in the index"
        );
    }

    #[tokio::test]
    async fn test_optimize_scalar_no_unindexed_fragments() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
        let ids = StringArray::from_iter_values((0..32).map(|i| format!("song-{i}")));
        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids)]).unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap();

        dataset
            .create_index(
                &["id"],
                IndexType::BTree,
                Some("id_idx".into()),
                &ScalarIndexParams::default(),
                true,
            )
            .await
            .unwrap();

        let before = dataset.load_indices_by_name("id_idx").await.unwrap();
        assert_eq!(before.len(), 1);
        let original_uuid = before[0].uuid;
        let original_version = dataset.manifest.version;

        // `merge(1)` would historically rebuild the single existing segment
        // (steady state, nothing unindexed) and replace its UUID; with the
        // short-circuit it must skip work entirely.
        dataset
            .optimize_indices(&OptimizeOptions::merge(1))
            .await
            .unwrap();

        let after = dataset.load_indices_by_name("id_idx").await.unwrap();
        assert_eq!(after.len(), 1, "no new segment should be produced");
        assert_eq!(
            after[0].uuid, original_uuid,
            "no-op optimize must not churn the index UUID"
        );
        assert_eq!(
            dataset.manifest.version, original_version,
            "no-op optimize must not advance the dataset version"
        );

        // The default options also short-circuit (num_to_merge defaults to 1
        // when there is a single old segment).
        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();
        let after_default = dataset.load_indices_by_name("id_idx").await.unwrap();
        assert_eq!(after_default[0].uuid, original_uuid);
        assert_eq!(dataset.manifest.version, original_version);
    }

    #[rstest]
    #[case::address_row_ids(false)]
    #[case::stable_row_ids(true)]
    #[tokio::test]
    async fn test_optimize_btree_no_duplicate_row_addr(#[case] use_stable_row_ids: bool) {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("payload", DataType::Int32, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from(vec![1])),
                Arc::new(Int32Array::from(vec![10])),
            ],
        )
        .unwrap();
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let write_params = WriteParams {
            enable_stable_row_ids: use_stable_row_ids,
            ..Default::default()
        };
        let mut dataset = Dataset::write(reader, test_uri, Some(write_params))
            .await
            .unwrap();

        let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree);
        dataset
            .create_index(
                &["id"],
                IndexType::BTree,
                Some("id_idx".into()),
                &params,
                true,
            )
            .await
            .unwrap();

        // Reordered source columns (payload, id) force the partial-schema
        // RewriteColumns path instead of a full row rewrite.
        let source_schema = Arc::new(Schema::new(vec![
            Field::new("payload", DataType::Int32, false),
            Field::new("id", DataType::Int32, false),
        ]));
        let source_batch = RecordBatch::try_new(
            source_schema.clone(),
            vec![
                Arc::new(Int32Array::from(vec![100])),
                Arc::new(Int32Array::from(vec![1])),
            ],
        )
        .unwrap();
        let merge_job =
            MergeInsertBuilder::try_new(Arc::new(dataset.clone()), vec!["id".to_string()])
                .unwrap()
                .when_matched(WhenMatched::UpdateAll)
                .try_build()
                .unwrap();
        let source_reader = Box::new(RecordBatchIterator::new(
            [Ok(source_batch)],
            source_schema.clone(),
        ));
        merge_job
            .execute(reader_to_stream(source_reader))
            .await
            .unwrap();

        // Build a delta BTree segment over the now-unindexed fragment.
        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        dataset
            .optimize_indices(&OptimizeOptions::append())
            .await
            .unwrap();
        assert_eq!(
            dataset.load_indices_by_name("id_idx").await.unwrap().len(),
            2,
            "append must create a delta segment over the rewritten fragment"
        );

        // Force the old segment + delta segment to merge.
        dataset
            .optimize_indices(&OptimizeOptions::merge(2))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        let rows = dataset
            .scan()
            .filter("id = 1")
            .unwrap()
            .project(&["id"])
            .unwrap()
            .try_into_batch()
            .await
            .unwrap()
            .num_rows();
        assert_eq!(rows, 1, "id = 1 must return exactly one row after merge");
    }

    #[tokio::test]
    async fn test_optimize_btree_merge_remaps_deferred_compaction() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();

        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
        let make = |range: std::ops::Range<i32>| {
            RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(Int32Array::from_iter_values(range))],
            )
            .unwrap()
        };

        // Two fragments: [0, 50) and [50, 100).
        let reader =
            RecordBatchIterator::new(vec![Ok(make(0..50)), Ok(make(50..100))], schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 50,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        assert_eq!(dataset.get_fragments().len(), 2);

        let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree);
        dataset
            .create_index(
                &["id"],
                IndexType::BTree,
                Some("id_idx".into()),
                &params,
                true,
            )
            .await
            .unwrap();

        // Deferred-remap compaction fuses the two fragments into one and leaves a
        // pending FragReuseIndex; the index segment is not eagerly remapped.
        compact_files(
            &mut dataset,
            CompactionOptions {
                defer_index_remap: true,
                ..Default::default()
            },
            None,
        )
        .await
        .unwrap();

        // Append a third fragment, left unindexed.
        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(make(100..150))], schema.clone()),
                None,
            )
            .await
            .unwrap();

        // Merge the deferred-remapped old segment with the new delta.
        dataset
            .optimize_indices(&OptimizeOptions::merge(2))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        // A value from the compacted fragments must still be found via the index.
        let hit = dataset
            .scan()
            .filter("id = 25")
            .unwrap()
            .project(&["id"])
            .unwrap()
            .try_into_batch()
            .await
            .unwrap()
            .num_rows();
        assert_eq!(
            hit, 1,
            "compacted-then-merged row must remain queryable via the index"
        );
        let total = dataset
            .scan()
            .filter("id >= 0")
            .unwrap()
            .project(&["id"])
            .unwrap()
            .try_into_batch()
            .await
            .unwrap()
            .num_rows();
        assert_eq!(total, 150, "no rows may be lost across compaction + merge");
    }

    /// Build one Bitmap segment per fragment of `dataset` under `index_name`.
    async fn commit_bitmap_segment_per_fragment(dataset: &mut Dataset, index_name: &str) {
        let groups = dataset
            .get_fragments()
            .iter()
            .map(|frag| vec![frag.id() as u32])
            .collect::<Vec<_>>();
        commit_bitmap_segments(dataset, index_name, &groups).await;
    }

    /// One committed Bitmap segment per fragment group.
    async fn commit_bitmap_segments(dataset: &mut Dataset, index_name: &str, groups: &[Vec<u32>]) {
        let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap);
        let mut staged = Vec::with_capacity(groups.len());
        for group in groups {
            staged.push(
                CreateIndexBuilder::new(dataset, &["cat"], IndexType::Bitmap, &params)
                    .name(index_name.into())
                    .fragments(group.clone())
                    .execute_uncommitted()
                    .await
                    .unwrap(),
            );
        }
        dataset
            .commit_existing_index_segments(index_name, "cat", staged)
            .await
            .unwrap();
    }

    /// The oldest `dataset_version` across an index's segments, which is what a
    /// segment merge carries over to the segment it produces. Checked here to
    /// predate the current manifest, since otherwise it could not tell a merge
    /// from a rebuild.
    async fn oldest_segment_version(dataset: &Dataset, index_name: &str) -> u64 {
        let oldest = dataset
            .load_indices_by_name(index_name)
            .await
            .unwrap()
            .iter()
            .map(|segment| segment.dataset_version)
            .min()
            .expect("the index must have at least one segment");
        assert!(
            oldest < dataset.manifest.version,
            "the source segments must predate the manifest, or their version \
             cannot distinguish a merge from a rebuild"
        );
        oldest
    }

    /// Pins that `optimize_indices` took the segment-merge path. A full rebuild
    /// produces the same rows, so the row counts these tests check cannot tell
    /// the two apart; the stamped `dataset_version` can, because the merge keeps
    /// the oldest source segment's version while a rebuild stamps the manifest
    /// version current at the time of the optimize.
    async fn assert_merged_not_rebuilt(dataset: &Dataset, index_name: &str, oldest_source: u64) {
        let segments = dataset.load_indices_by_name(index_name).await.unwrap();
        assert_eq!(
            segments.len(),
            1,
            "the optimize must consolidate every bitmap segment, got {segments:?}"
        );
        assert_eq!(
            segments[0].dataset_version, oldest_source,
            "the optimize rebuilt the index instead of merging the segments: a \
             merged segment carries the oldest source segment's dataset version"
        );
    }

    fn id_cat_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("cat", DataType::Utf8, false),
        ]))
    }

    /// `id` ascending, `cat` cycling through A/B/C.
    fn id_cat_batch(schema: &Arc<Schema>, range: std::ops::Range<i32>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(range.clone())),
                Arc::new(StringArray::from_iter_values(
                    range.map(|i| ["A", "B", "C"][(i % 3) as usize]),
                )),
            ],
        )
        .unwrap()
    }

    async fn count_cat(dataset: &Dataset, cat: &str) -> usize {
        count_rows_where(dataset, &format!("cat = '{cat}'")).await
    }

    async fn count_rows_where(dataset: &Dataset, predicate: &str) -> usize {
        dataset
            .scan()
            .filter(predicate)
            .unwrap()
            .count_rows()
            .await
            .unwrap() as usize
    }

    /// `id` ascending, `cat` cycling through A/B/C/NULL.
    fn id_cat_nullable_batch(schema: &Arc<Schema>, range: std::ops::Range<i32>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int32Array::from_iter_values(range.clone())),
                Arc::new(StringArray::from_iter(
                    range.map(|i| ["A", "B", "C"].get((i % 4) as usize).copied()),
                )),
            ],
        )
        .unwrap()
    }

    /// The K-way merge over a nullable column: nulls must survive consolidation of
    /// several segments plus an unindexed tail, and `IS NULL` is served from
    /// `null_map`, which is filled separately from the value row sets.
    #[tokio::test]
    async fn test_optimize_bitmap_multi_segment_merge_keeps_nulls() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("cat", DataType::Utf8, true),
        ]));

        // 36 rows over three 12-row fragments; every fourth row is null.
        let reader = RecordBatchIterator::new(
            vec![
                Ok(id_cat_nullable_batch(&schema, 0..12)),
                Ok(id_cat_nullable_batch(&schema, 12..24)),
                Ok(id_cat_nullable_batch(&schema, 24..36)),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 12,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        assert_eq!(dataset.get_fragments().len(), 3);

        commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await;
        dataset
            .append(
                RecordBatchIterator::new(
                    vec![Ok(id_cat_nullable_batch(&schema, 36..48))],
                    schema.clone(),
                ),
                None,
            )
            .await
            .unwrap();

        let oldest_source = oldest_segment_version(&dataset, "cat_idx").await;
        dataset
            .optimize_indices(&OptimizeOptions::merge(200))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert_merged_not_rebuilt(&dataset, "cat_idx", oldest_source).await;

        // 48 rows cycling A/B/C/NULL: 12 of each.
        for cat in ["A", "B", "C"] {
            assert_eq!(
                count_cat(&dataset, cat).await,
                12,
                "wrong row count for cat = {cat} after merging a nullable column"
            );
        }
        assert_eq!(
            count_rows_where(&dataset, "cat IS NULL").await,
            12,
            "nulls lost across the bitmap merge"
        );
        assert_eq!(
            count_rows_where(&dataset, "cat IS NOT NULL").await,
            36,
            "IS NOT NULL disagrees with the value row sets"
        );
        assert_eq!(dataset.scan().count_rows().await.unwrap(), 48);
    }

    /// A 200-way merge over three Bitmap segments plus an unindexed fragment must
    /// consolidate into one segment (via the N:1 segment-merge primitive) and keep
    /// every posting.
    #[tokio::test]
    async fn test_optimize_bitmap_multi_segment_merge_consolidates() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = id_cat_schema();

        let reader = RecordBatchIterator::new(
            vec![
                Ok(id_cat_batch(&schema, 0..50)),
                Ok(id_cat_batch(&schema, 50..100)),
                Ok(id_cat_batch(&schema, 100..150)),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 50,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        assert_eq!(dataset.get_fragments().len(), 3);

        commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await;
        assert_eq!(
            dataset.load_indices_by_name("cat_idx").await.unwrap().len(),
            3
        );

        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(id_cat_batch(&schema, 150..200))], schema.clone()),
                None,
            )
            .await
            .unwrap();

        let oldest_source = oldest_segment_version(&dataset, "cat_idx").await;
        dataset
            .optimize_indices(&OptimizeOptions::merge(200))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert_merged_not_rebuilt(&dataset, "cat_idx", oldest_source).await;
        let segments = dataset.load_indices_by_name("cat_idx").await.unwrap();
        let expected_coverage = dataset
            .get_fragments()
            .iter()
            .map(|frag| frag.id() as u32)
            .collect::<RoaringBitmap>();
        assert_eq!(
            segments[0].fragment_bitmap.as_ref(),
            Some(&expected_coverage),
            "merged bitmap segment must cover every dataset fragment"
        );

        for (cat, expected) in [("A", 67), ("B", 67), ("C", 66)] {
            assert_eq!(
                count_cat(&dataset, cat).await,
                expected,
                "wrong row count for cat = {cat} after multi-segment bitmap merge"
            );
        }
    }

    /// A wide consolidation: 26 Bitmap delta segments merged into one by a single
    /// `optimize_indices` call, with no unindexed data on top.
    #[tokio::test]
    async fn test_optimize_bitmap_wide_consolidation() {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = id_cat_schema();

        let num_segments = 26;
        let rows_per_fragment = 12;
        let total_rows = num_segments * rows_per_fragment;
        let batches = (0..num_segments)
            .map(|i| {
                Ok(id_cat_batch(
                    &schema,
                    i * rows_per_fragment..(i + 1) * rows_per_fragment,
                ))
            })
            .collect::<Vec<_>>();
        let reader = RecordBatchIterator::new(batches, schema.clone());
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: rows_per_fragment as usize,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        assert_eq!(dataset.get_fragments().len(), num_segments as usize);

        commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await;
        assert_eq!(
            dataset.load_indices_by_name("cat_idx").await.unwrap().len(),
            num_segments as usize
        );

        let oldest_source = oldest_segment_version(&dataset, "cat_idx").await;
        dataset
            .optimize_indices(&OptimizeOptions::merge(num_segments as usize))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert_merged_not_rebuilt(&dataset, "cat_idx", oldest_source).await;
        let segments = dataset.load_indices_by_name("cat_idx").await.unwrap();
        let expected_coverage = dataset
            .get_fragments()
            .iter()
            .map(|frag| frag.id() as u32)
            .collect::<RoaringBitmap>();
        assert_eq!(
            segments[0].fragment_bitmap.as_ref(),
            Some(&expected_coverage)
        );

        let per_cat = (total_rows / 3) as usize;
        for cat in ["A", "B", "C"] {
            assert_eq!(
                count_cat(&dataset, cat).await,
                per_cat,
                "wrong row count for cat = {cat} after wide bitmap consolidation"
            );
        }
    }

    /// Deferred-remap compaction leaves the bitmap segments pointing at retired
    /// fragment ids; a K-way segment merge must remap them through the
    /// FragReuseIndex instead of dropping or misattributing the rows.
    ///
    /// Null rows ride in `null_map`, which the merge fills separately from the
    /// value row sets, so they need remapping of their own and the nullable
    /// case is not covered by the other one.
    #[rstest]
    #[case::non_null(false)]
    #[case::nullable(true)]
    #[tokio::test]
    async fn test_optimize_bitmap_merge_remaps_deferred_compaction(#[case] nullable: bool) {
        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("cat", DataType::Utf8, nullable),
        ]));
        // A/B/C cycling, with every fourth row null in the nullable case.
        let batch = |range: std::ops::Range<i32>| {
            if nullable {
                id_cat_nullable_batch(&schema, range)
            } else {
                id_cat_batch(&schema, range)
            }
        };

        let reader = RecordBatchIterator::new(
            vec![
                Ok(batch(0..50)),
                Ok(batch(50..100)),
                Ok(batch(100..150)),
                Ok(batch(150..200)),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 50,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        // Two fragments per segment: compaction only ever bins fragments that
        // carry the same set of index segments, so one segment per fragment
        // would leave every bin a single fragment and plan no work at all.
        commit_bitmap_segments(&mut dataset, "cat_idx", &[vec![0, 1], vec![2, 3]]).await;

        compact_files(
            &mut dataset,
            CompactionOptions {
                defer_index_remap: true,
                ..Default::default()
            },
            None,
        )
        .await
        .unwrap();

        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        // Every original fragment was retired, so each segment's postings now
        // need the FragReuseIndex to reach a live row.
        assert!(
            dataset.get_fragments().iter().all(|frag| frag.id() >= 4),
            "compaction must have retired every indexed fragment"
        );
        dataset
            .append(
                RecordBatchIterator::new(vec![Ok(batch(200..240))], schema.clone()),
                None,
            )
            .await
            .unwrap();

        let oldest_source = oldest_segment_version(&dataset, "cat_idx").await;
        dataset
            .optimize_indices(&OptimizeOptions::merge(200))
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        assert_merged_not_rebuilt(&dataset, "cat_idx", oldest_source).await;
        // 240 rows: A/B/C/NULL evenly when nullable, A/B/C otherwise.
        let (expected_per_cat, expected_nulls) = if nullable {
            ([60, 60, 60], 60)
        } else {
            ([80, 80, 80], 0)
        };
        for (cat, expected) in ["A", "B", "C"].into_iter().zip(expected_per_cat) {
            assert_eq!(
                count_cat(&dataset, cat).await,
                expected,
                "cat = {cat} lost or gained rows across deferred compaction + bitmap merge"
            );
        }
        assert_eq!(
            count_rows_where(&dataset, "cat IS NULL").await,
            expected_nulls,
            "null rows were not remapped across deferred compaction + bitmap merge"
        );
        assert_eq!(
            count_rows_where(&dataset, "cat IS NOT NULL").await,
            240 - expected_nulls,
            "IS NOT NULL disagrees with the value row sets after the remap"
        );
        assert_eq!(
            dataset.scan().count_rows().await.unwrap(),
            240,
            "no rows may be lost across compaction + bitmap merge"
        );
    }

    /// Stable-row-id update against an older Bitmap segment's fragment: the K-way
    /// merge must drop that segment's stale postings, not just the tail segment's.
    #[tokio::test]
    async fn test_optimize_bitmap_drops_stale_rows_across_segments_after_update() {
        use crate::dataset::UpdateBuilder;

        let test_dir = TempStrDir::default();
        let test_uri = test_dir.as_str();
        let schema = id_cat_schema();

        let reader = RecordBatchIterator::new(
            vec![
                Ok(id_cat_batch(&schema, 0..50)),
                Ok(id_cat_batch(&schema, 50..100)),
            ],
            schema.clone(),
        );
        let mut dataset = Dataset::write(
            reader,
            test_uri,
            Some(WriteParams {
                max_rows_per_file: 50,
                enable_stable_row_ids: true,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        commit_bitmap_segment_per_fragment(&mut dataset, "cat_idx").await;

        // Rows 0..25 live in the *older* segment's fragment; rewrite their cat.
        let res = UpdateBuilder::new(Arc::new(dataset.clone()))
            .update_where("id < 25")
            .unwrap()
            .set("cat", "'Z'")
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();
        let mut dataset = res.new_dataset.as_ref().clone();

        let oldest_source = oldest_segment_version(&dataset, "cat_idx").await;
        dataset
            .optimize_indices(&OptimizeOptions::default())
            .await
            .unwrap();

        let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
        // Two segments went in, so this exercised the K-way path rather than the
        // single-segment `update` this change does not touch.
        assert_merged_not_rebuilt(&dataset, "cat_idx", oldest_source).await;
        // ids 0..25 are now 'Z'; the remaining A/B/C counts come from ids 25..100.
        assert_eq!(count_cat(&dataset, "Z").await, 25, "updated rows missing");
        for (cat, expected) in [("A", 25), ("B", 25), ("C", 25)] {
            assert_eq!(
                count_cat(&dataset, cat).await,
                expected,
                "bitmap merge returned stale rows for cat = {cat}"
            );
        }
    }
}