lance 14.0.1

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

#![allow(clippy::redundant_pub_crate)]

use std::{collections::BTreeMap, sync::Arc};

use arrow_array::cast::AsArray;
use arrow_array::{
    Array, ArrayRef, LargeStringArray, RecordBatch, StringArray, StringViewArray, UInt32Array,
    UInt64Array,
};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use datafusion::error::DataFusionError;
use datafusion::execution::SendableRecordBatchStream;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use futures::StreamExt;
use lance_core::{
    ROW_ID,
    datatypes::{Field, LogicalType, Schema, format_field_path, parse_field_path},
};
use lance_index::metrics::NoOpMetricsCollector;
use lance_index::pbold::InvertedIndexDetails;
use lance_index::scalar::inverted::{
    DocumentGranularity, InvertedIndex, InvertedIndexParams, doc_index_storage_column,
};
use lance_index::scalar::lance_format::LanceIndexStore;
use lance_index::scalar::registry::VALUE_COLUMN_NAME;
use lance_index::scalar::{IndexFile, IndexStore, index_files_to_table};
use lance_table::format::IndexMetadata;
use prost::Message;
use roaring::RoaringBitmap;
use uuid::Uuid;

use crate::{
    Dataset, Error, Result,
    dataset::index::LanceIndexStoreExt,
    index::{DatasetIndexExt, scalar::fetch_index_details},
};

#[derive(Debug, Clone)]
enum FtsTraversal {
    Text,
    Struct {
        child_index: usize,
        child: Box<Self>,
    },
    List {
        child: Box<Self>,
    },
}

impl FtsTraversal {
    fn list_depth(&self) -> usize {
        match self {
            Self::Text => 0,
            Self::Struct { child, .. } => child.list_depth(),
            Self::List { child } => 1 + child.list_depth(),
        }
    }
}

/// Schema-derived form of an FTS field path.
#[derive(Debug, Clone)]
pub(crate) struct ResolvedFtsField {
    pub final_field_id: i32,
    pub root_column: String,
    pub canonical_path: String,
    pub document_granularity: DocumentGranularity,
    traversal: FtsTraversal,
    list_depth: usize,
}

/// One logical document extracted from a dataset row.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct FtsDocument {
    pub row_index: usize,
    pub text: String,
    pub doc_index: Vec<u32>,
}

impl ResolvedFtsField {
    pub fn has_lists(&self) -> bool {
        self.list_depth > 0
    }

    pub fn coordinate_rank(&self) -> usize {
        if self.document_granularity.is_list_element() {
            self.list_depth
        } else {
            0
        }
    }

    pub fn documents_from_batch(&self, batch: &RecordBatch) -> Result<Vec<FtsDocument>> {
        let column = batch.column_by_name(&self.root_column).ok_or_else(|| {
            Error::invalid_input(format!(
                "FTS root column '{}' is missing from the input batch",
                self.root_column
            ))
        })?;
        self.documents_from_array(column, batch.num_rows())
    }

    fn documents_from_array(&self, column: &ArrayRef, num_rows: usize) -> Result<Vec<FtsDocument>> {
        let mut documents = Vec::new();
        match self.document_granularity {
            DocumentGranularity::Row => {
                documents.reserve(num_rows);
                for row_index in 0..num_rows {
                    let mut text = String::new();
                    append_row_text(&self.traversal, column.as_ref(), row_index, &mut text)?;
                    documents.push(FtsDocument {
                        row_index,
                        text,
                        doc_index: Vec::new(),
                    });
                }
            }
            DocumentGranularity::ListElement => {
                for row_index in 0..num_rows {
                    collect_element_documents(
                        &self.traversal,
                        column.as_ref(),
                        row_index,
                        row_index,
                        0,
                        self.list_depth,
                        &mut Vec::with_capacity(self.list_depth),
                        &mut documents,
                    )?;
                }
            }
        }
        Ok(documents)
    }
}

fn string_value(array: &dyn Array, index: usize) -> Result<Option<&str>> {
    if array.is_null(index) {
        return Ok(None);
    }
    match array.data_type() {
        DataType::Utf8 => Ok(Some(
            array
                .as_any()
                .downcast_ref::<StringArray>()
                .expect("Utf8 array type")
                .value(index),
        )),
        DataType::LargeUtf8 => Ok(Some(
            array
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .expect("LargeUtf8 array type")
                .value(index),
        )),
        DataType::Utf8View => Ok(Some(
            array
                .as_any()
                .downcast_ref::<StringViewArray>()
                .expect("Utf8View array type")
                .value(index),
        )),
        data_type => Err(Error::internal(format!(
            "FTS traversal expected a string array, got {data_type}"
        ))),
    }
}

fn append_text(output: &mut String, text: &str) {
    if !output.is_empty() {
        output.push(' ');
    }
    output.push_str(text);
}

fn append_row_text(
    traversal: &FtsTraversal,
    array: &dyn Array,
    index: usize,
    output: &mut String,
) -> Result<()> {
    match traversal {
        FtsTraversal::Text => {
            if let Some(text) = string_value(array, index)? {
                append_text(output, text);
            }
        }
        FtsTraversal::Struct { child_index, child } => {
            if !array.is_null(index) {
                let structs = array.as_struct();
                append_row_text(child, structs.column(*child_index).as_ref(), index, output)?;
            }
        }
        FtsTraversal::List { child } => match array.data_type() {
            DataType::List(_) => {
                let lists = array.as_list::<i32>();
                if !lists.is_null(index) {
                    let offsets = lists.value_offsets();
                    let start = offsets[index] as usize;
                    let end = offsets[index + 1] as usize;
                    for element_index in start..end {
                        append_row_text(child, lists.values().as_ref(), element_index, output)?;
                    }
                }
            }
            DataType::LargeList(_) => {
                let lists = array.as_list::<i64>();
                if !lists.is_null(index) {
                    let offsets = lists.value_offsets();
                    let start = offsets[index] as usize;
                    let end = offsets[index + 1] as usize;
                    for element_index in start..end {
                        append_row_text(child, lists.values().as_ref(), element_index, output)?;
                    }
                }
            }
            data_type => {
                return Err(Error::internal(format!(
                    "FTS traversal expected a list array, got {data_type}"
                )));
            }
        },
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn collect_element_documents(
    traversal: &FtsTraversal,
    array: &dyn Array,
    index: usize,
    row_index: usize,
    list_depth: usize,
    boundary_depth: usize,
    doc_index: &mut Vec<u32>,
    documents: &mut Vec<FtsDocument>,
) -> Result<()> {
    match traversal {
        FtsTraversal::Text => Err(Error::internal(
            "ListElement FTS traversal did not encounter its list boundary".to_string(),
        )),
        FtsTraversal::Struct { child_index, child } => {
            if !array.is_null(index) {
                let structs = array.as_struct();
                collect_element_documents(
                    child,
                    structs.column(*child_index).as_ref(),
                    index,
                    row_index,
                    list_depth,
                    boundary_depth,
                    doc_index,
                    documents,
                )?;
            }
            Ok(())
        }
        FtsTraversal::List { child } => {
            let mut visit = |values: &ArrayRef, start: usize, end: usize| -> Result<()> {
                for (ordinal, element_index) in (start..end).enumerate() {
                    let ordinal = u32::try_from(ordinal).map_err(|_| {
                        Error::invalid_input(format!(
                            "FTS element ordinal overflow for row index {row_index}"
                        ))
                    })?;
                    doc_index.push(ordinal);
                    if list_depth + 1 == boundary_depth {
                        let mut text = String::new();
                        append_row_text(child, values.as_ref(), element_index, &mut text)?;
                        documents.push(FtsDocument {
                            row_index,
                            text,
                            doc_index: doc_index.clone(),
                        });
                    } else {
                        collect_element_documents(
                            child,
                            values.as_ref(),
                            element_index,
                            row_index,
                            list_depth + 1,
                            boundary_depth,
                            doc_index,
                            documents,
                        )?;
                    }
                    doc_index.pop();
                }
                Ok(())
            };
            match array.data_type() {
                DataType::List(_) => {
                    let lists = array.as_list::<i32>();
                    if !lists.is_null(index) {
                        let offsets = lists.value_offsets();
                        visit(
                            lists.values(),
                            offsets[index] as usize,
                            offsets[index + 1] as usize,
                        )?;
                    }
                }
                DataType::LargeList(_) => {
                    let lists = array.as_list::<i64>();
                    if !lists.is_null(index) {
                        let offsets = lists.value_offsets();
                        visit(
                            lists.values(),
                            offsets[index] as usize,
                            offsets[index + 1] as usize,
                        )?;
                    }
                }
                data_type => {
                    return Err(Error::internal(format!(
                        "FTS traversal expected a list array, got {data_type}"
                    )));
                }
            }
            Ok(())
        }
    }
}

fn find_child_case_insensitive<'a>(field: &'a Field, name: &str) -> Option<(usize, &'a Field)> {
    field
        .children
        .iter()
        .enumerate()
        .find(|(_, child)| child.name == name)
        .or_else(|| {
            field
                .children
                .iter()
                .enumerate()
                .find(|(_, child)| child.name.eq_ignore_ascii_case(name))
        })
}

fn build_fts_traversal(
    field: &Field,
    remaining_names: &[String],
    canonical_names: &mut Vec<String>,
    final_field_id: &mut i32,
) -> Result<FtsTraversal> {
    match field.data_type() {
        DataType::List(_) | DataType::LargeList(_) => {
            let child = field.children.first().ok_or_else(|| {
                Error::invalid_input(format!(
                    "FTS list field '{}' does not have an item field",
                    field.name
                ))
            })?;
            Ok(FtsTraversal::List {
                child: Box::new(build_fts_traversal(
                    child,
                    remaining_names,
                    canonical_names,
                    final_field_id,
                )?),
            })
        }
        DataType::Struct(_) => {
            let Some((name, rest)) = remaining_names.split_first() else {
                return Err(Error::invalid_input(format!(
                    "FTS field path ends at struct '{}'; specify the final text field",
                    field.name
                )));
            };
            let (child_index, child) =
                find_child_case_insensitive(field, name).ok_or_else(|| {
                    Error::index(format!(
                        "FTS field '{}' does not contain child '{}'",
                        field.name, name
                    ))
                })?;
            canonical_names.push(child.name.clone());
            *final_field_id = child.id;
            Ok(FtsTraversal::Struct {
                child_index,
                child: Box::new(build_fts_traversal(
                    child,
                    rest,
                    canonical_names,
                    final_field_id,
                )?),
            })
        }
        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View if remaining_names.is_empty() => {
            Ok(FtsTraversal::Text)
        }
        _ if field.logical_type == LogicalType::from("json") && remaining_names.is_empty() => {
            Ok(FtsTraversal::Text)
        }
        data_type if !remaining_names.is_empty() => Err(Error::invalid_input(format!(
            "FTS field '{}' has type {data_type} and cannot contain child '{}'",
            field.name, remaining_names[0]
        ))),
        data_type => Err(Error::invalid_input(format!(
            "FTS field '{}' must resolve to Utf8, LargeUtf8, Utf8View, or JSON, got {data_type}",
            field.name
        ))),
    }
}

/// Resolve a public FTS field path and derive all list traversal from schema.
pub(crate) fn resolve_fts_field(
    schema: &Schema,
    path: &str,
    document_granularity: DocumentGranularity,
) -> Result<ResolvedFtsField> {
    let names = parse_field_path(path)?;
    let (root_name, remaining_names) = names
        .split_first()
        .ok_or_else(|| Error::invalid_input("FTS field path cannot be empty".to_string()))?;
    let root = schema
        .fields
        .iter()
        .find(|field| field.name == *root_name)
        .or_else(|| {
            schema
                .fields
                .iter()
                .find(|field| field.name.eq_ignore_ascii_case(root_name))
        })
        .ok_or_else(|| {
            Error::index(format!(
                "FTS field path '{path}' does not exist in the dataset schema"
            ))
        })?;

    let mut canonical_names = vec![root.name.clone()];
    let mut final_field_id = root.id;
    let traversal = build_fts_traversal(
        root,
        remaining_names,
        &mut canonical_names,
        &mut final_field_id,
    )?;
    let list_depth = traversal.list_depth();
    if document_granularity.is_list_element() && list_depth == 0 {
        return Err(Error::invalid_input(format!(
            "FTS field path '{}' has no List layer and cannot use ListElement document granularity",
            format_field_path(
                &canonical_names
                    .iter()
                    .map(String::as_str)
                    .collect::<Vec<_>>()
            )
        )));
    }
    if list_depth > 0 && root.logical_type == LogicalType::from("json") {
        return Err(Error::invalid_input(
            "nested List traversal is not supported for JSON FTS sources".to_string(),
        ));
    }
    let canonical_path = format_field_path(
        &canonical_names
            .iter()
            .map(String::as_str)
            .collect::<Vec<_>>(),
    );
    Ok(ResolvedFtsField {
        final_field_id,
        root_column: root.name.clone(),
        canonical_path,
        document_granularity,
        traversal,
        list_depth,
    })
}

fn find_public_path_by_id(field: &Field, field_id: i32, path: &mut Vec<String>) -> bool {
    if field.id == field_id {
        return true;
    }
    match field.data_type() {
        DataType::List(_) | DataType::LargeList(_) => field
            .children
            .first()
            .is_some_and(|child| find_physical_path_by_id(child, field_id, path)),
        DataType::Struct(_) => field.children.iter().any(|child| {
            path.push(child.name.clone());
            let found = find_public_path_by_id(child, field_id, path);
            if !found {
                path.pop();
            }
            found
        }),
        _ => false,
    }
}

fn find_physical_path_by_id(field: &Field, field_id: i32, path: &mut Vec<String>) -> bool {
    if field.id == field_id {
        return false;
    }
    match field.data_type() {
        DataType::List(_) | DataType::LargeList(_) => field
            .children
            .first()
            .is_some_and(|child| find_physical_path_by_id(child, field_id, path)),
        DataType::Struct(_) => field.children.iter().any(|child| {
            path.push(child.name.clone());
            let found = find_public_path_by_id(child, field_id, path);
            if !found {
                path.pop();
            }
            found
        }),
        _ => false,
    }
}

/// Resolve an indexed final field id against the current schema, preserving
/// routing across field renames.
pub(crate) fn resolve_fts_field_by_id(
    schema: &Schema,
    field_id: i32,
    document_granularity: DocumentGranularity,
) -> Result<ResolvedFtsField> {
    for root in &schema.fields {
        let mut path = vec![root.name.clone()];
        if find_public_path_by_id(root, field_id, &mut path) {
            let path_refs = path.iter().map(String::as_str).collect::<Vec<_>>();
            return resolve_fts_field(schema, &format_field_path(&path_refs), document_granularity);
        }
    }
    Err(Error::invalid_input(format!(
        "FTS index refers to missing or Arrow-internal field id {field_id}"
    )))
}

/// Return the persisted FTS document granularity for each logical index on a
/// public field path. Multiple physical segments with the same index name are
/// collapsed after validating that their metadata agrees.
pub(crate) async fn indexed_fts_document_granularities(
    dataset: &Dataset,
    column: &str,
) -> Result<Vec<(String, DocumentGranularity)>> {
    let resolved = resolve_fts_field(dataset.schema(), column, DocumentGranularity::Row)?;
    let mut by_name = BTreeMap::new();

    for index in dataset.load_indices().await?.iter() {
        // A covered FTS index still answers for its keyed column; only the
        // keyed prefix decides whether this index matches, not the full
        // `fields` vector including carried columns.
        if index.keyed_field() != Some(resolved.final_field_id) {
            continue;
        }
        let details_any = fetch_index_details(dataset, &resolved.canonical_path, index).await?;
        if !details_any.type_url.ends_with("InvertedIndexDetails") {
            continue;
        }
        let details =
            InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| {
                Error::corrupt_file(
                    dataset.indices_dir().join(index.uuid.to_string()),
                    format!(
                        "failed to decode InvertedIndexDetails for FTS index '{}': {error}",
                        index.name
                    ),
                )
            })?;
        let document_granularity = DocumentGranularity::try_from(details.document_granularity)?;
        if let Some(existing) = by_name.insert(index.name.clone(), document_granularity)
            && existing != document_granularity
        {
            return Err(Error::corrupt_file(
                dataset.indices_dir().join(index.uuid.to_string()),
                format!(
                    "FTS index '{}' has inconsistent document granularity across segments: \
                     {existing:?} and {document_granularity:?}",
                    index.name
                ),
            ));
        }
    }

    Ok(by_name.into_iter().collect())
}

/// The effective [`InvertedIndexParams`] of every persisted FTS index on
/// `column`, paired with the document granularity it was built for.
///
/// The analyzer and positional settings are part of the query contract, not an
/// implementation detail: a phrase query needs positions, and a stemming or
/// n-gram tokenizer changes which terms a document produces. Anything that
/// evaluates the same query over rows an index does not cover — an unindexed
/// fragment, or a MemWAL active memtable — has to analyze them the same way or
/// the two disagree about what matches.
///
/// Keyed by granularity rather than flattened, because row and list-element
/// indexes may coexist on one column with *different* settings, and
/// `load_segments` selects between them by the query's resolved granularity.
/// Picking the first match instead would rebuild against the wrong contract —
/// a list-element phrase query could be analyzed with the row index's
/// positionless settings and silently match nothing.
pub(crate) async fn indexed_fts_index_params(
    dataset: &Dataset,
    column: &str,
) -> Result<Vec<(DocumentGranularity, InvertedIndexParams)>> {
    let resolved = resolve_fts_field(dataset.schema(), column, DocumentGranularity::Row)?;
    let mut found = Vec::new();
    for index in dataset.load_indices().await?.iter() {
        if index.keyed_field() != Some(resolved.final_field_id) {
            continue;
        }
        let details_any = fetch_index_details(dataset, &resolved.canonical_path, index).await?;
        if !details_any.type_url.ends_with("InvertedIndexDetails") {
            continue;
        }
        let details =
            InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| {
                Error::corrupt_file(
                    dataset.indices_dir().join(index.uuid.to_string()),
                    format!(
                        "failed to decode InvertedIndexDetails for FTS index '{}': {error}",
                        index.name
                    ),
                )
            })?;
        let document_granularity = DocumentGranularity::try_from(details.document_granularity)?;
        let params = InvertedIndexParams::try_from(&details)?;
        found.push((document_granularity, params));
    }
    Ok(found)
}

/// Resolve an optional query granularity against persisted FTS index metadata.
///
/// A unique indexed granularity is authoritative and also controls flat search
/// over unindexed fragments. An explicit request selects between coexisting
/// row and list-element indexes, but cannot contradict the only indexed
/// granularity. With no index, the established row default is retained.
pub(crate) async fn resolve_query_document_granularity(
    dataset: &Dataset,
    column: &str,
    requested: Option<DocumentGranularity>,
) -> Result<DocumentGranularity> {
    let indices = indexed_fts_document_granularities(dataset, column).await?;
    let mut available = indices
        .iter()
        .map(|(_, document_granularity)| *document_granularity)
        .collect::<Vec<_>>();
    available.sort_by_key(|document_granularity| match document_granularity {
        DocumentGranularity::Row => 0,
        DocumentGranularity::ListElement => 1,
    });
    available.dedup();

    let resolved = match requested {
        Some(requested) if available.is_empty() || available.contains(&requested) => requested,
        Some(requested) => {
            let indexed = indices
                .iter()
                .map(|(name, document_granularity)| format!("'{name}' ({document_granularity:?})"))
                .collect::<Vec<_>>()
                .join(", ");
            return Err(Error::invalid_input(format!(
                "FTS query for field '{column}' requested {requested:?} document granularity, \
                 but the existing FTS index uses a different granularity: {indexed}"
            )));
        }
        None if available.is_empty() => DocumentGranularity::Row,
        None if available.len() == 1 => available[0],
        None => {
            let indexed = indices
                .iter()
                .map(|(name, document_granularity)| format!("'{name}' ({document_granularity:?})"))
                .collect::<Vec<_>>()
                .join(", ");
            return Err(Error::invalid_input(format!(
                "FTS query for field '{column}' is ambiguous because Row and ListElement \
                 indexes coexist: {indexed}; specify document_granularity"
            )));
        }
    };

    resolve_fts_field(dataset.schema(), column, resolved)?;
    Ok(resolved)
}

/// Check that one `combined_fields` target column can take part in a BM25F
/// blend, and resolve its path.
///
/// Cross-field scoring joins the target columns on the row address and sums their
/// per-row term frequencies and document lengths, so it needs row documents. A
/// column indexed only at list-element granularity cannot supply them: that index
/// stores several documents per row, identified by element coordinates a
/// cross-column scan cannot pair up between columns and the combined result schema
/// cannot report. A column carrying both granularities is fine; the row index is
/// the one used.
///
/// Any field shape a row-document index accepts is allowed, list nesting included:
/// the flat sibling scan reaches such a leaf through
/// [`flatten_fts_document_column`], which uses the same traversal a single-column
/// match drives through [`FtsDocument`].
pub(crate) async fn validate_combined_fields_target_column(
    dataset: &Dataset,
    column: &str,
) -> Result<()> {
    let indices = indexed_fts_document_granularities(dataset, column).await?;
    if !indices.is_empty()
        && indices
            .iter()
            .all(|(_, granularity)| granularity.is_list_element())
    {
        let indexed = indices
            .iter()
            .map(|(name, granularity)| format!("'{name}' ({granularity:?})"))
            .collect::<Vec<_>>()
            .join(", ");
        return Err(Error::not_supported(format!(
            "combined_fields (BM25F) scores whole rows and needs a Row document granularity FTS \
             index on every target column, but '{column}' only has: {indexed}"
        )));
    }
    resolve_fts_field(dataset.schema(), column, DocumentGranularity::Row)?;
    Ok(())
}

pub(crate) fn fts_document_schema(coordinate_rank: usize) -> Arc<ArrowSchema> {
    let mut fields = vec![
        ArrowField::new(VALUE_COLUMN_NAME, DataType::Utf8, false),
        ArrowField::new(ROW_ID, DataType::UInt64, false),
    ];
    fields.extend(
        (0..coordinate_rank)
            .map(|rank| ArrowField::new(doc_index_storage_column(rank), DataType::UInt32, false)),
    );
    Arc::new(ArrowSchema::new(fields))
}

/// Add (or replace) a flat `Utf8` column holding each row's document text for a
/// list-bearing FTS field path, leaving every other column and row untouched.
///
/// A Lance projection cannot address a field that continues past a `List`, so a
/// scan that needs `docs.content` has to project the `docs` root instead. This
/// walks the traversal down to the leaf per row, joining the elements exactly the
/// way the index builder does at row granularity.
///
/// Row granularity yields one document per input row, in row order, keeping several
/// such columns aligned for a cross-field blend. Not valid for list-element
/// granularity, where a row expands to several documents.
pub(crate) fn flatten_fts_document_column(
    input: SendableRecordBatchStream,
    resolved: ResolvedFtsField,
) -> Result<SendableRecordBatchStream> {
    if resolved.document_granularity.is_list_element() {
        return Err(Error::internal(
            "flatten_fts_document_column needs row documents to keep one output row per input row"
                .to_string(),
        ));
    }
    let input_schema = input.schema();
    if input_schema
        .column_with_name(&resolved.root_column)
        .is_none()
    {
        return Err(Error::internal(format!(
            "FTS document input must contain the root source column '{}'",
            resolved.root_column
        )));
    }
    let text_field = ArrowField::new(&resolved.canonical_path, DataType::Utf8, false);
    let replaced = input_schema.column_with_name(&resolved.canonical_path);
    let mut fields = input_schema.fields().to_vec();
    match replaced {
        Some((index, _)) => fields[index] = Arc::new(text_field),
        None => fields.push(Arc::new(text_field)),
    }
    let replaced = replaced.map(|(index, _)| index);
    let output_schema = Arc::new(ArrowSchema::new(fields));
    let stream_schema = output_schema.clone();
    let stream = input.map(move |batch| {
        let batch = batch?;
        let source = batch
            .column_by_name(&resolved.root_column)
            .expect("FTS document input schema was validated");
        let documents = resolved
            .documents_from_array(source, batch.num_rows())
            .map_err(DataFusionError::from)?;
        let texts = Arc::new(StringArray::from_iter_values(
            documents.iter().map(|document| document.text.as_str()),
        )) as ArrayRef;
        let mut columns = batch.columns().to_vec();
        match replaced {
            Some(index) => columns[index] = texts,
            None => columns.push(texts),
        }
        RecordBatch::try_new(stream_schema.clone(), columns).map_err(DataFusionError::from)
    });
    Ok(Box::pin(RecordBatchStreamAdapter::new(
        output_schema,
        stream,
    )))
}

pub(crate) fn transform_fts_document_stream(
    input: SendableRecordBatchStream,
    resolved: ResolvedFtsField,
) -> Result<SendableRecordBatchStream> {
    let input_schema = input.schema();
    if (input_schema
        .column_with_name(&resolved.root_column)
        .is_none()
        && input_schema.column_with_name(VALUE_COLUMN_NAME).is_none())
        || input_schema.column_with_name(ROW_ID).is_none()
    {
        return Err(Error::internal(
            "FTS document input must contain the root source column and _rowid".to_string(),
        ));
    }
    let output_schema = fts_document_schema(resolved.coordinate_rank());
    let stream_schema = output_schema.clone();
    let stream = input.map(move |batch| {
        let batch = batch?;
        let source = batch
            .column_by_name(&resolved.root_column)
            .or_else(|| batch.column_by_name(VALUE_COLUMN_NAME))
            .expect("FTS document input schema was validated");
        let documents = resolved
            .documents_from_array(source, batch.num_rows())
            .map_err(DataFusionError::from)?;
        let input_row_ids = batch[ROW_ID].as_primitive::<arrow_array::types::UInt64Type>();
        let texts =
            StringArray::from_iter_values(documents.iter().map(|document| document.text.as_str()));
        let row_ids = UInt64Array::from_iter_values(
            documents
                .iter()
                .map(|document| input_row_ids.value(document.row_index)),
        );
        let mut columns = vec![Arc::new(texts) as ArrayRef, Arc::new(row_ids) as ArrayRef];
        for rank in 0..resolved.coordinate_rank() {
            columns.push(Arc::new(UInt32Array::from_iter_values(
                documents.iter().map(|document| document.doc_index[rank]),
            )) as ArrayRef);
        }
        RecordBatch::try_new(stream_schema.clone(), columns).map_err(DataFusionError::from)
    });
    Ok(Box::pin(RecordBatchStreamAdapter::new(
        output_schema,
        stream,
    )))
}

/// Fill legacy posting-format metadata while leaving protobuf-default row
/// document granularity untouched.
pub(crate) fn normalize_inverted_details(
    index: &IndexMetadata,
    mut details: InvertedIndexDetails,
) -> Result<InvertedIndexDetails> {
    if !matches!(index.index_version, 0..=3) {
        return Err(Error::invalid_input(format!(
            "FTS index '{}' has unsupported index_version {}; expected 0, 1, 2, or 3",
            index.name, index.index_version
        )));
    }
    if details.posting_format_version.is_none() {
        let posting_format_version = match index.index_version {
            0 | 1 => 1,
            2 => 2,
            3 => 3,
            _ => unreachable!("index version was validated above"),
        };
        details.posting_format_version = Some(posting_format_version);
    }
    Ok(details)
}

/// Build an empty update stream for the inverted merge API.
///
/// `InvertedIndex::merge_segments` is shaped as "merge old segments plus new
/// rows", so even a pure segment merge needs a stream with the document column
/// and `_rowid` fields. The stream intentionally contains no batches.
fn empty_inverted_update_stream(
    dataset: &Dataset,
    resolved: &ResolvedFtsField,
) -> Result<SendableRecordBatchStream> {
    let field = dataset
        .schema()
        .field_by_id(resolved.final_field_id)
        .ok_or_else(|| {
            Error::invalid_input(format!(
                "merge_existing_index_segments: field id {} does not exist",
                resolved.final_field_id
            ))
        })?;
    let schema = if resolved.has_lists() {
        fts_document_schema(resolved.coordinate_rank())
    } else {
        Arc::new(ArrowSchema::new(vec![
            ArrowField::new(VALUE_COLUMN_NAME, field.data_type(), true),
            ArrowField::new(ROW_ID, arrow_schema::DataType::UInt64, false),
        ]))
    };
    Ok(Box::pin(RecordBatchStreamAdapter::new(
        schema,
        futures::stream::empty(),
    )))
}

/// Finalize a staged segment in place. Returns the final index files, or
/// `None` when the segment already had its root metadata.
pub(crate) async fn finalize_segment_files_if_needed(
    dataset: &Dataset,
    segment: &IndexMetadata,
) -> Result<Option<Vec<IndexFile>>> {
    let index_dir = dataset.indices_dir().join(segment.uuid.to_string());
    let metadata_path = index_dir
        .clone()
        .join(lance_index::scalar::inverted::METADATA_FILE);
    if dataset.object_store.as_ref().exists(&metadata_path).await? {
        return Ok(None);
    }

    let store = Arc::new(LanceIndexStore::from_dataset_for_new(
        dataset,
        &segment.uuid,
    )?);
    let files = lance_index::scalar::inverted::builder::merge_index_files_with_options(
        dataset.object_store.as_ref(),
        &index_dir,
        store,
        lance_index::scalar::inverted::builder::FinalizeOptions::default(),
        lance_index::progress::noop_progress(),
    )
    .await?;
    Ok(Some(files))
}

/// A lone staged segment is the merged segment as it stands only while its
/// files hold live row addresses: no reuse mapping moved its coverage and no
/// tagged plan asks to translate it. Anything else must be rewritten.
fn single_segment_keeps_row_addresses(
    staged: Option<&crate::index::frag_reuse::StagedRemappingPlans>,
    uuid: &Uuid,
    source_coverage_moved: bool,
) -> bool {
    use crate::index::frag_reuse::SegmentRemappingPlan;
    if source_coverage_moved {
        return false;
    }
    match staged.and_then(|plans| plans.get(uuid)) {
        None | Some(SegmentRemappingPlan::Identity) => true,
        Some(SegmentRemappingPlan::Translate { .. })
        | Some(SegmentRemappingPlan::MissingCoverage(_)) => false,
    }
}

/// Promote a lone staged segment under its own uuid: finalize its files in
/// place and return its metadata with the final file list, with no second
/// copy of the index.
async fn promote_finalized_segment(
    dataset: &Dataset,
    segment: &IndexMetadata,
    details: InvertedIndexDetails,
) -> Result<IndexMetadata> {
    let files = match finalize_segment_files_if_needed(dataset, segment).await? {
        Some(files) => files,
        None => {
            let store = LanceIndexStore::from_dataset_for_new(dataset, &segment.uuid)?;
            store
                .list_files_with_sizes()
                .await?
                .into_iter()
                .filter(|file| !lance_index::scalar::inverted::builder::is_staged_path(&file.path))
                .collect()
        }
    };
    let details_any = prost_types::Any::from_msg(&details).map_err(|error| {
        Error::io(format!(
            "failed to encode InvertedIndexDetails payload: {error}"
        ))
    })?;
    Ok(IndexMetadata {
        dataset_version: dataset.manifest.version,
        index_details: Some(Arc::new(details_any)),
        files: Some(index_files_to_table(files)),
        ..segment.clone()
    })
}

/// Merge one caller-defined group of source FTS segments into a single segment.
pub(crate) async fn merge_segments(
    dataset: &Dataset,
    segments: Vec<IndexMetadata>,
    staged: Option<&crate::index::frag_reuse::StagedRemappingPlans>,
    source_coverage_moved: bool,
) -> Result<IndexMetadata> {
    if segments.is_empty() {
        return Err(Error::index("No segment metadata was provided".to_string()));
    }

    let field_id = *segments[0].fields.first().ok_or_else(|| {
        Error::invalid_input(format!(
            "CreateIndex: segment {} is missing field ids",
            segments[0].uuid
        ))
    })?;
    let details = match segments[0].index_details.as_ref() {
        Some(details_any) => {
            let details =
                InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| {
                    Error::io(format!(
                        "failed to decode InvertedIndexDetails payload: {error}"
                    ))
                })?;
            normalize_inverted_details(&segments[0], details)?
        }
        None => normalize_inverted_details(&segments[0], InvertedIndexDetails::default())?,
    };
    let document_granularity = DocumentGranularity::try_from(details.document_granularity)?;
    let resolved = resolve_fts_field_by_id(dataset.schema(), field_id, document_granularity)?;
    load_segment_details(dataset, &resolved.canonical_path, &segments).await?;

    if let [segment] = segments.as_slice()
        && single_segment_keeps_row_addresses(staged, &segment.uuid, source_coverage_moved)
    {
        return promote_finalized_segment(dataset, segment, details).await;
    }

    let mut source_indices = Vec::with_capacity(segments.len());
    let mut fragment_bitmap = RoaringBitmap::new();
    for segment in &segments {
        finalize_segment_files_if_needed(dataset, segment).await?;
        fragment_bitmap |= segment.fragment_bitmap.as_ref().cloned().ok_or_else(|| {
            Error::invalid_input(format!(
                "CreateIndex: segment {} is missing fragment coverage",
                segment.uuid
            ))
        })?;
        if segment.fields != segments[0].fields {
            return Err(Error::invalid_input(format!(
                "FTS index {} has inconsistent fields across segments",
                segments[0].name
            )));
        }
        let scalar_index = super::open_scalar_index_with_plan(
            dataset,
            &resolved.canonical_path,
            segment,
            staged.and_then(|plans| plans.get(&segment.uuid)),
            crate::index::frag_reuse::OpenPurpose::Maintenance,
            &NoOpMetricsCollector,
        )
        .await?;
        let inverted_index = scalar_index
            .as_any()
            .downcast_ref::<InvertedIndex>()
            .ok_or_else(|| {
                Error::index(format!(
                    "merge_existing_index_segments: expected inverted segment {}, got {:?}",
                    segment.uuid,
                    scalar_index.index_type()
                ))
            })?;
        source_indices.push(Arc::new(inverted_index.clone()));
    }

    let new_uuid = Uuid::new_v4();
    let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?;
    let created_index = InvertedIndex::merge_segments(
        &source_indices,
        empty_inverted_update_stream(dataset, &resolved)?,
        &new_store,
        None,
        lance_index::progress::noop_progress(),
    )
    .await?;

    Ok(IndexMetadata {
        uuid: new_uuid,
        dataset_version: dataset.manifest.version,
        fragment_bitmap: Some(fragment_bitmap),
        index_details: Some(Arc::new(created_index.index_details)),
        index_version: created_index.index_version as i32,
        created_at: Some(chrono::Utc::now()),
        base_id: None,
        files: Some(index_files_to_table(created_index.files)),
        ..segments[0].clone()
    })
}

/// Load all committed inverted-index segments that belong to the same named
/// FTS index on `column`.
///
/// Returns `Ok(None)` if no FTS index exists on the column. When an index
/// exists, the returned vector contains every committed segment's
/// [`IndexMetadata`] (UUID, fragment coverage, index details). All segments
/// must share the same indexed fields; mismatched fields return an error.
pub async fn load_segments(
    dataset: &Dataset,
    column: &str,
    document_granularity: DocumentGranularity,
) -> Result<Option<Vec<IndexMetadata>>> {
    let Some(index_meta) = dataset
        .load_scalar_index(
            lance_index::IndexCriteria::default()
                .for_column(column)
                .supports_fts()
                .with_fts_document_granularity(document_granularity),
        )
        .await?
    else {
        return Ok(None);
    };

    let indices = dataset.load_indices_by_name(&index_meta.name).await?;
    if indices.is_empty() {
        return Ok(None);
    }

    let expected_fields = indices[0].fields.clone();
    for meta in &indices {
        if meta.fields != expected_fields {
            return Err(Error::invalid_input(format!(
                "FTS index {} has inconsistent fields across segments",
                index_meta.name
            )));
        }
    }

    Ok(Some(indices))
}

/// Return the fragment coverage for the logical FTS index selected by `column`
/// and `document_granularity`.
///
/// Returns `None` when no FTS index is available or when any physical segment
/// has unknown coverage. In either case, callers must not use the result for
/// pruning.
pub(crate) async fn fts_index_fragment_bitmap(
    dataset: &Dataset,
    column: &str,
    document_granularity: DocumentGranularity,
) -> Result<Option<RoaringBitmap>> {
    let Some(segments) = load_segments(dataset, column, document_granularity).await? else {
        return Ok(None);
    };

    let fragment_bitmap =
        segments
            .iter()
            .try_fold(RoaringBitmap::new(), |mut coverage, segment| {
                coverage |= segment.fragment_bitmap.as_ref()?.clone();
                Some(coverage)
            });

    Ok(fragment_bitmap)
}

/// Load and validate the shared [`InvertedIndexDetails`] across committed
/// segments returned by [`load_segments`].
///
/// All segments are required to agree on their semantic `InvertedIndexDetails`
/// payload (tokenizer, position settings, etc.); inconsistent
/// segments return an error. Details are canonicalized before comparison so
/// legacy segments that omit default fields remain compatible with newly
/// written text FTS segments. `posting_format_version` is a physical
/// per-segment property and may differ when a legacy FTS v1 segment is
/// combined with a newly written one. Returns the first segment's
/// canonicalized details for tokenizer construction and query planning.
pub async fn load_segment_details(
    dataset: &Dataset,
    column: &str,
    segments: &[IndexMetadata],
) -> Result<InvertedIndexDetails> {
    let mut expected_details: Option<InvertedIndexDetails> = None;
    for meta in segments {
        let details_any = fetch_index_details(dataset, column, meta).await?;
        let details =
            InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|err| {
                Error::io(format!(
                    "failed to decode InvertedIndexDetails payload: {err}"
                ))
            })?;
        let details = canonicalize_inverted_index_details(details)?;
        match &expected_details {
            Some(expected) if !inverted_index_details_semantically_equal(expected, &details) => {
                return Err(Error::invalid_input(format!(
                    "FTS index {} has inconsistent inverted index details across segments",
                    meta.name
                )));
            }
            Some(_) => {}
            None => expected_details = Some(details),
        }
    }
    expected_details.ok_or_else(|| {
        Error::invalid_input(format!(
            "FTS index for column {} requires at least one segment",
            column
        ))
    })
}

fn canonicalize_inverted_index_details(
    details: InvertedIndexDetails,
) -> Result<InvertedIndexDetails> {
    let params = InvertedIndexParams::try_from(&details)?;
    InvertedIndexDetails::try_from(&params)
}

/// Compare canonicalized inverted-index details for shared semantic configuration.
///
/// `posting_format_version` records how a single segment physically stores
/// postings, so mixed-version FTS segments may disagree on it without being
/// incompatible. Every other field remains part of the equality check.
fn inverted_index_details_semantically_equal(
    left: &InvertedIndexDetails,
    right: &InvertedIndexDetails,
) -> bool {
    let mut left = left.clone();
    let mut right = right.clone();
    left.posting_format_version = None;
    right.posting_format_version = None;
    left == right
}

/// Read one segment's [`InvertedIndexParams`]
pub async fn load_segment_params(
    dataset: &Dataset,
    segment: &IndexMetadata,
) -> Result<InvertedIndexParams> {
    let store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?;
    InvertedIndex::load_params(&store).await
}

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

    /// A covered FTS index (`fields=[tags, id]`, `id` carried) must still be
    /// recognized by `indexed_fts_document_granularities`: matching on
    /// `idx.fields` as a whole (rather than its keyed prefix) would leave
    /// `available` empty and silently fall back to the `Row` default even
    /// though a `ListElement` index exists.
    #[tokio::test]
    async fn test_resolve_query_document_granularity_recognizes_covered_index() {
        use arrow_array::builder::{ListBuilder, StringBuilder};
        use arrow_array::{Int32Array, RecordBatchIterator};
        use lance_core::utils::tempfile::TempStrDir;
        use lance_index::IndexType;

        let mut tags_builder = ListBuilder::new(StringBuilder::new());
        for tag in ["alpha", "beta", "gamma", "delta"] {
            tags_builder.values().append_value(tag);
            tags_builder.append(true);
        }
        let tags: ArrayRef = Arc::new(tags_builder.finish());
        let ids: ArrayRef = Arc::new(Int32Array::from_iter_values(0..4));
        let batch = RecordBatch::try_from_iter(vec![("id", ids), ("tags", tags)]).unwrap();
        let schema = batch.schema();

        let test_dir = TempStrDir::default();
        let mut dataset = Dataset::write(
            RecordBatchIterator::new(vec![Ok(batch)], schema),
            test_dir.as_str(),
            None,
        )
        .await
        .unwrap();

        let tags_field_id = dataset.schema().field("tags").unwrap().id;
        let id_field_id = dataset.schema().field("id").unwrap().id;

        let params =
            InvertedIndexParams::default().document_granularity(DocumentGranularity::ListElement);
        let segment = dataset
            .create_index_builder(&["tags"], IndexType::Inverted, &params)
            .name("tags_idx".to_string())
            .execute_uncommitted()
            .await
            .unwrap();

        // Nothing writes a covering declaration yet, and read plumbing never
        // touches covered-column storage, so declaring the uninvolved `id`
        // column as carried is a faithful stand-in for a genuinely covered
        // FTS segment.
        let covered = IndexMetadata {
            fields: vec![tags_field_id, id_field_id],
            covering_fields: vec![id_field_id],
            ..segment
        };
        dataset
            .commit_existing_index_segments("tags_idx", "tags", vec![covered])
            .await
            .unwrap();

        let resolved = resolve_query_document_granularity(&dataset, "tags", None)
            .await
            .unwrap();
        assert_eq!(
            resolved,
            DocumentGranularity::ListElement,
            "a covered ListElement FTS index must still be recognized, not silently \
             skipped in favor of the Row default"
        );
    }

    fn fts_test_schema() -> Schema {
        let schema = ArrowSchema::new(vec![
            ArrowField::new("text", DataType::Utf8, true),
            ArrowField::new(
                "tags",
                DataType::List(Arc::new(ArrowField::new("item", DataType::Utf8, true))),
                true,
            ),
        ]);
        Schema::try_from(&schema).unwrap()
    }

    #[test]
    fn decode_legacy_inverted_details_type_url() {
        let mut details_any = prost_types::Any::from_msg(&InvertedIndexDetails::default()).unwrap();
        details_any.type_url = "/lance.index.pb.InvertedIndexDetails".to_string();

        let decoded = InvertedIndexDetails::decode(details_any.value.as_slice()).unwrap();
        assert_eq!(decoded, InvertedIndexDetails::default());
    }

    #[test]
    fn resolve_element_document_field() {
        let schema = fts_test_schema();
        let tags = schema.field("tags").unwrap();
        let resolved =
            resolve_fts_field(&schema, "tags", DocumentGranularity::ListElement).unwrap();
        assert_eq!(resolved.final_field_id, tags.id);
        assert_eq!(resolved.root_column, "tags");
        assert_eq!(resolved.canonical_path, "tags");
        assert_eq!(resolved.coordinate_rank(), 1);

        let err = resolve_fts_field(&schema, "text", DocumentGranularity::ListElement).unwrap_err();
        assert!(err.to_string().contains("has no List layer"), "{err}");

        let err =
            resolve_fts_field(&schema, "tags[*]", DocumentGranularity::ListElement).unwrap_err();
        assert!(err.to_string().contains("does not exist"), "{err}");
    }

    #[test]
    fn normalize_legacy_metadata_as_row_document() {
        let metadata = IndexMetadata {
            uuid: Uuid::new_v4(),
            fields: vec![1],
            covering_fields: vec![],
            name: "tags_idx".to_string(),
            dataset_version: 1,
            fragment_bitmap: None,
            index_details: None,
            index_version: 3,
            created_at: None,
            base_id: None,
            files: None,
        };
        let details =
            normalize_inverted_details(&metadata, InvertedIndexDetails::default()).unwrap();
        assert_eq!(
            DocumentGranularity::try_from(details.document_granularity).unwrap(),
            DocumentGranularity::Row
        );
        assert_eq!(details.posting_format_version, Some(3));
    }

    #[test]
    fn normalize_rejects_unreleased_v4_metadata() {
        let metadata = IndexMetadata {
            uuid: Uuid::new_v4(),
            fields: vec![1],
            covering_fields: vec![],
            name: "tags_idx".to_string(),
            dataset_version: 1,
            fragment_bitmap: None,
            index_details: None,
            index_version: 4,
            created_at: None,
            base_id: None,
            files: None,
        };
        let details = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap();

        let err = normalize_inverted_details(&metadata, details).unwrap_err();

        assert!(err.to_string().contains("unsupported index_version 4"));
    }

    #[test]
    fn canonicalize_inverted_details_accepts_legacy_empty_details() {
        let legacy = InvertedIndexDetails::default();
        let current = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap();

        assert_ne!(legacy, current);
        assert_eq!(
            canonicalize_inverted_index_details(legacy).unwrap(),
            canonicalize_inverted_index_details(current).unwrap()
        );
    }

    #[test]
    fn inverted_details_equal_when_only_posting_format_version_differs() {
        let left = canonicalize_inverted_index_details(
            InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(),
        )
        .unwrap();
        let mut right = left.clone();
        right.posting_format_version = Some(1);

        assert_ne!(left, right);
        assert!(inverted_index_details_semantically_equal(&left, &right));
    }

    #[test]
    fn inverted_details_reject_with_position_mismatch() {
        let left = canonicalize_inverted_index_details(
            InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(),
        )
        .unwrap();
        let mut right = left.clone();
        right.with_position = !left.with_position;

        assert!(!inverted_index_details_semantically_equal(&left, &right));
    }

    #[test]
    fn test_single_segment_keeps_row_addresses_only_without_translation() {
        use crate::index::frag_reuse::{MissingCoverageReason, SegmentRemappingPlan};
        let uuid = Uuid::new_v4();
        let other = Uuid::new_v4();
        assert!(single_segment_keeps_row_addresses(None, &uuid, false));
        assert!(!single_segment_keeps_row_addresses(None, &uuid, true));

        let mut plans = crate::index::frag_reuse::StagedRemappingPlans::new();
        plans.insert(other, SegmentRemappingPlan::Identity);
        assert!(
            single_segment_keeps_row_addresses(Some(&plans), &uuid, false),
            "a plan for another segment says nothing about this one"
        );
        plans.insert(uuid, SegmentRemappingPlan::Identity);
        assert!(single_segment_keeps_row_addresses(
            Some(&plans),
            &uuid,
            false
        ));
        plans.insert(
            uuid,
            SegmentRemappingPlan::Translate {
                coverage: RoaringBitmap::new(),
                excluded_fragments: RoaringBitmap::new(),
                fingerprint: [0; 32],
            },
        );
        assert!(!single_segment_keeps_row_addresses(
            Some(&plans),
            &uuid,
            false
        ));
        plans.insert(
            uuid,
            SegmentRemappingPlan::MissingCoverage(MissingCoverageReason::Withdrawn),
        );
        assert!(!single_segment_keeps_row_addresses(
            Some(&plans),
            &uuid,
            false
        ));
    }

    async fn write_text_dataset(uri: &str, rows_per_fragment: usize, fragments: usize) -> Dataset {
        use lance_datagen::{BatchCount, RowCount, array, gen_batch};
        let reader = gen_batch()
            // Named `i` so the fragment-reuse test helpers can label its rows.
            .col("i", array::step::<arrow_array::types::Int32Type>())
            .col("text", array::random_sentence(4, 12, false))
            .into_reader_rows(
                RowCount::from(rows_per_fragment as u64),
                BatchCount::from(fragments as u32),
            );
        Dataset::write(
            reader,
            uri,
            Some(crate::dataset::WriteParams {
                max_rows_per_file: rows_per_fragment,
                ..Default::default()
            }),
        )
        .await
        .unwrap()
    }

    async fn fts_hits(dataset: &Dataset, column: &str, terms: &str) -> Vec<(u64, u32)> {
        use lance_index::scalar::FullTextSearchQuery;
        let mut scan = dataset.scan();
        scan.full_text_search(
            FullTextSearchQuery::new(terms.to_string())
                .with_column(column.to_string())
                .unwrap()
                .limit(Some(1_000_000)),
        )
        .unwrap();
        scan.with_row_id();
        let batch = scan.try_into_batch().await.unwrap();
        let row_ids = batch
            .column_by_name(ROW_ID)
            .unwrap()
            .as_primitive::<arrow_array::types::UInt64Type>();
        let scores = batch
            .column_by_name(lance_index::scalar::inverted::SCORE_COL)
            .unwrap()
            .as_primitive::<arrow_array::types::Float32Type>();
        let mut hits = row_ids
            .values()
            .iter()
            .zip(scores.values())
            .map(|(row_id, score)| (*row_id, score.to_bits()))
            .collect::<Vec<_>>();
        // Score descending, then row id: ties are ordered deterministically
        // and scores compare bit for bit.
        hits.sort_by(|a, b| {
            f32::from_bits(b.1)
                .total_cmp(&f32::from_bits(a.1))
                .then(a.0.cmp(&b.0))
        });
        hits
    }

    /// A single staged intermediate covering every fragment is finalized in
    /// place and committed under its own uuid (ENT-1419 kept: the root
    /// metadata exists), with no second copy of the index.
    #[tokio::test]
    async fn test_single_intermediate_is_promoted_in_place() {
        use futures::TryStreamExt;
        use lance_core::utils::tempfile::TempStrDir;
        use lance_index::IndexType;
        use lance_index::scalar::inverted::METADATA_FILE;

        let test_dir = TempStrDir::default();
        let mut dataset = write_text_dataset(&test_dir, 200, 4).await;
        let fragment_ids = dataset
            .get_fragments()
            .iter()
            .map(|fragment| fragment.id() as u32)
            .collect::<Vec<_>>();
        // A tiny worker budget leaves several small staged partitions.
        let params = InvertedIndexParams::default()
            .num_workers(2)
            .memory_limit_mb(1);
        let segment = dataset
            .create_index_builder(&["text"], IndexType::Inverted, &params)
            .name("text_idx".to_string())
            .fragments(fragment_ids)
            .execute_uncommitted()
            .await
            .unwrap();
        let index_dir = dataset.indices_dir().join(segment.uuid.to_string());
        let staged_paths = dataset
            .object_store
            .read_dir_all(&index_dir, None)
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        assert!(
            staged_paths
                .iter()
                .any(|meta| meta.location.as_ref().contains("/staging/")),
            "an uncommitted fragment-scoped build must be staged"
        );
        assert!(
            !dataset
                .object_store
                .exists(&index_dir.clone().join(METADATA_FILE))
                .await
                .unwrap()
        );

        let merged = dataset
            .merge_existing_index_segments(vec![segment.clone()])
            .await
            .unwrap();
        assert_eq!(merged.uuid, segment.uuid, "the lone segment keeps its uuid");
        assert_eq!(merged.fragment_bitmap, segment.fragment_bitmap);
        assert!(
            dataset
                .object_store
                .exists(&index_dir.clone().join(METADATA_FILE))
                .await
                .unwrap(),
            "finalize publishes the root metadata"
        );
        let final_paths = dataset
            .object_store
            .read_dir_all(&index_dir, None)
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        assert!(
            final_paths
                .iter()
                .all(|meta| !meta.location.as_ref().contains("/staging/")),
            "staging is cleared after the root metadata is published"
        );
        let listed = final_paths
            .iter()
            .map(|meta| meta.location.filename().unwrap().to_string())
            .collect::<std::collections::BTreeSet<_>>();
        let reported = merged
            .files
            .as_ref()
            .unwrap()
            .iter()
            .map(|file| file.path.clone())
            .collect::<std::collections::BTreeSet<_>>();
        assert_eq!(listed, reported, "metadata lists exactly the final files");
        let index_dirs = dataset
            .object_store
            .read_dir(dataset.indices_dir())
            .await
            .unwrap();
        assert_eq!(index_dirs.len(), 1, "no second index directory is written");

        dataset
            .commit_existing_index_segments("text_idx", "text", vec![merged])
            .await
            .unwrap();
        let committed = dataset.load_indices().await.unwrap();
        assert_eq!(committed.len(), 1);
        assert_eq!(committed[0].uuid, segment.uuid);

        // Reference: the same column indexed in one committed build.
        let reference_dir = TempStrDir::default();
        let mut reference = Dataset::write(
            dataset
                .scan()
                .try_into_batch()
                .await
                .map(|batch| {
                    arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema())
                })
                .unwrap(),
            &reference_dir,
            Some(crate::dataset::WriteParams {
                max_rows_per_file: 200,
                ..Default::default()
            }),
        )
        .await
        .unwrap();
        reference
            .create_index(
                &["text"],
                IndexType::Inverted,
                Some("text_idx".into()),
                &InvertedIndexParams::default(),
                true,
            )
            .await
            .unwrap();
        let probe = dataset
            .scan()
            .limit(Some(1), None)
            .unwrap()
            .try_into_batch()
            .await
            .unwrap();
        let first_text = probe
            .column_by_name("text")
            .unwrap()
            .as_string::<i32>()
            .value(0)
            .to_string();
        let first_word = first_text.split_whitespace().next().unwrap().to_string();
        let expected = fts_hits(&reference, "text", &first_word).await;
        let actual = fts_hits(&dataset, "text", &first_word).await;
        assert!(!expected.is_empty());
        assert_eq!(expected, actual);
    }

    /// Under a tagged fragment-reuse history a staged segment whose fragments
    /// no rewrite touched keeps its row addresses, so merging it alone still
    /// promotes it under its own uuid; a segment over rewritten fragments is
    /// translated into a new one.
    #[tokio::test]
    async fn test_single_intermediate_promotion_under_tagged_history() {
        use crate::index::frag_reuse_reader::tests::{install, prepare_partition};
        use lance_core::utils::tempfile::TempStrDir;
        use lance_index::IndexType;
        use lance_table::format::pb::fragment_reuse_index_details::InlineContent;
        use prost::Message;

        for rewrite_indexed_fragments in [false, true] {
            let test_dir = TempStrDir::default();
            let mut dataset = write_text_dataset(&test_dir, 50, 4).await;
            let fragment_ids = dataset
                .get_fragments()
                .iter()
                .map(|fragment| fragment.id() as u32)
                .collect::<Vec<_>>();
            let (indexed, rewritten): (Vec<u32>, Vec<u32>) = if rewrite_indexed_fragments {
                (fragment_ids[2..].to_vec(), fragment_ids[2..].to_vec())
            } else {
                (fragment_ids[2..].to_vec(), fragment_ids[..2].to_vec())
            };
            let segment = dataset
                .create_index_builder(
                    &["text"],
                    IndexType::Inverted,
                    &InvertedIndexParams::default(),
                )
                .name("text_idx".to_string())
                .fragments(indexed.clone())
                .execute_uncommitted()
                .await
                .unwrap();

            // Reference: the same rows in the same fragments, indexed over the
            // same fragments through the same staged path, on an untagged
            // table. Partial coverage is part of the scenario, so a full
            // fresh index would not score the same.
            let rows = dataset
                .scan()
                .project(&["i", "text"])
                .unwrap()
                .try_into_batch()
                .await
                .unwrap();
            let reference_dir = TempStrDir::default();
            let mut reference = Dataset::write(
                arrow_array::RecordBatchIterator::new(vec![Ok(rows.clone())], rows.schema()),
                &reference_dir,
                Some(crate::dataset::WriteParams {
                    max_rows_per_file: 50,
                    ..Default::default()
                }),
            )
            .await
            .unwrap();
            let reference_segment = reference
                .create_index_builder(
                    &["text"],
                    IndexType::Inverted,
                    &InvertedIndexParams::default(),
                )
                .name("text_idx".to_string())
                .fragments(indexed.clone())
                .execute_uncommitted()
                .await
                .unwrap();
            let reference_merged = reference
                .merge_existing_index_segments(vec![reference_segment])
                .await
                .unwrap();
            reference
                .commit_existing_index_segments("text_idx", "text", vec![reference_merged])
                .await
                .unwrap();

            // The rewrite replaces the manifest's fragment list with the
            // destinations, so the fragments it did not touch are carried
            // over and stay live.
            let rewritten_ids = rewritten.iter().map(|id| *id as u64).collect::<Vec<_>>();
            let (transition, mut destinations) =
                prepare_partition(&dataset, &rewritten_ids, 10).await;
            destinations.extend(
                dataset
                    .fragments()
                    .iter()
                    .filter(|fragment| !rewritten_ids.contains(&fragment.id))
                    .cloned(),
            );
            let content = InlineContent {
                legacy_versions: vec![],
                transitions: vec![transition],
            }
            .encode_to_vec();
            let reuse_index = install(&mut dataset, content, destinations, false).await;

            let merged = dataset
                .merge_existing_index_segments(vec![segment.clone()])
                .await
                .unwrap();
            let index_dirs = dataset
                .object_store
                .read_dir(dataset.indices_dir())
                .await
                .unwrap()
                .into_iter()
                .filter(|dir| dir != &reuse_index.uuid.to_string())
                .count();
            if rewrite_indexed_fragments {
                assert_ne!(
                    merged.uuid, segment.uuid,
                    "translated rows need a new segment"
                );
                assert_eq!(index_dirs, 2, "the staged segment and its rewrite");
            } else {
                assert_eq!(
                    merged.uuid, segment.uuid,
                    "untouched coverage keeps the uuid"
                );
                assert_eq!(index_dirs, 1, "no second index directory is written");
            }

            // Either way the committed index answers like the reference,
            // compared by the `i` column since the rewrite changed row addresses.
            dataset
                .commit_existing_index_segments("text_idx", "text", vec![merged])
                .await
                .unwrap();
            // A long word: short ones are mostly stop words and match nothing.
            let texts = rows.column_by_name("text").unwrap().as_string::<i32>();
            let probe = texts
                .iter()
                .flatten()
                .flat_map(|text| text.split_whitespace())
                .find(|word| word.len() >= 6)
                .unwrap()
                .to_string();
            let expected = fts_hits_by_id(&reference, "text", &probe).await;
            let actual = fts_hits_by_id(&dataset, "text", &probe).await;
            assert!(!expected.is_empty());
            assert_eq!(
                expected, actual,
                "rewrite_indexed_fragments={rewrite_indexed_fragments}"
            );
        }
    }

    /// FTS hits keyed by the `i` column rather than the row address, for
    /// datasets whose rows were rewritten; ordered by score then id, scores
    /// compared bit for bit.
    async fn fts_hits_by_id(dataset: &Dataset, column: &str, terms: &str) -> Vec<(i32, u32)> {
        use lance_index::scalar::FullTextSearchQuery;
        let mut scan = dataset.scan();
        scan.full_text_search(
            FullTextSearchQuery::new(terms.to_string())
                .with_column(column.to_string())
                .unwrap()
                .limit(Some(1_000_000)),
        )
        .unwrap();
        scan.project(&["i"]).unwrap();
        let batch = scan.try_into_batch().await.unwrap();
        let ids = batch
            .column_by_name("i")
            .unwrap()
            .as_primitive::<arrow_array::types::Int32Type>();
        let scores = batch
            .column_by_name(lance_index::scalar::inverted::SCORE_COL)
            .unwrap()
            .as_primitive::<arrow_array::types::Float32Type>();
        let mut hits = ids
            .values()
            .iter()
            .zip(scores.values())
            .map(|(id, score)| (*id, score.to_bits()))
            .collect::<Vec<_>>();
        hits.sort_by(|a, b| {
            f32::from_bits(b.1)
                .total_cmp(&f32::from_bits(a.1))
                .then(a.0.cmp(&b.0))
        });
        hits
    }
}