lance-index 13.0.0

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

use super::*;

pub fn doc_index_storage_column(rank: usize) -> String {
    format!("{DOC_INDEX_STORAGE_PREFIX}{rank}")
}

pub fn document_coordinate_rank(schema: &arrow_schema::Schema) -> usize {
    (0..)
        .take_while(|rank| {
            schema
                .column_with_name(&doc_index_storage_column(*rank))
                .is_some()
        })
        .count()
}

pub fn flat_full_text_search(
    batches: &[&RecordBatch],
    doc_col: &str,
    query: &str,
    tokenizer: Option<Box<dyn LanceTokenizer>>,
) -> Result<Vec<u64>> {
    if batches.is_empty() {
        return Ok(vec![]);
    }

    let (query, phrase_slop) = match phrase_query_text(query) {
        Some(query) => (query, Some(0)),
        None => (query, None),
    };

    match batches[0][doc_col].data_type() {
        DataType::Utf8 => {
            do_flat_full_text_search::<i32>(batches, doc_col, query, tokenizer, phrase_slop)
        }
        DataType::LargeUtf8 => {
            do_flat_full_text_search::<i64>(batches, doc_col, query, tokenizer, phrase_slop)
        }
        DataType::List(_) => {
            do_flat_full_text_search_list::<i32>(batches, doc_col, query, tokenizer, phrase_slop)
        }
        DataType::LargeList(_) => {
            do_flat_full_text_search_list::<i64>(batches, doc_col, query, tokenizer, phrase_slop)
        }
        data_type => Err(Error::invalid_input(format!(
            "unsupported data type {} for inverted index",
            data_type
        ))),
    }
}

pub(super) fn do_flat_full_text_search<Offset: OffsetSizeTrait>(
    batches: &[&RecordBatch],
    doc_col: &str,
    query: &str,
    tokenizer: Option<Box<dyn LanceTokenizer>>,
    phrase_slop: Option<u32>,
) -> Result<Vec<u64>> {
    let mut results = Vec::new();
    let mut tokenizer =
        tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap());
    let query_tokens = try_collect_query_tokens(query, &mut tokenizer)?;

    for batch in batches {
        let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();
        let doc_array = batch[doc_col].as_string::<Offset>();
        for i in 0..row_id_array.len() {
            let doc = doc_array.value(i);
            if document_matches_flat_query(doc, &mut tokenizer, &query_tokens, phrase_slop)? {
                results.push(row_id_array.value(i));
            }
        }
    }

    Ok(results)
}

pub(super) fn do_flat_full_text_search_list<ListOffset: OffsetSizeTrait>(
    batches: &[&RecordBatch],
    doc_col: &str,
    query: &str,
    tokenizer: Option<Box<dyn LanceTokenizer>>,
    phrase_slop: Option<u32>,
) -> Result<Vec<u64>> {
    let mut results = Vec::new();
    let mut tokenizer =
        tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap());
    let query_tokens = try_collect_query_tokens(query, &mut tokenizer)?;

    for batch in batches {
        let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();
        let doc_array = batch[doc_col].as_list::<ListOffset>();
        match doc_array.value_type() {
            DataType::Utf8 | DataType::LargeUtf8 => {}
            data_type => {
                return Err(Error::invalid_input(format!(
                    "unsupported list item data type {} for inverted index",
                    data_type
                )));
            }
        }
        for i in 0..row_id_array.len() {
            if doc_array.is_null(i) {
                continue;
            }
            let elements = doc_array.value(i);
            let matches = if phrase_slop.is_some() {
                let document = iter_str_array(elements.as_ref())
                    .flatten()
                    .collect::<Vec<_>>()
                    .join(" ");
                document_matches_flat_query(&document, &mut tokenizer, &query_tokens, phrase_slop)?
            } else {
                iter_str_array(elements.as_ref())
                    .flatten()
                    .any(|element| has_query_token(element, &mut tokenizer, &query_tokens))
            };
            if matches {
                results.push(row_id_array.value(i));
            }
        }
    }

    Ok(results)
}

pub(super) fn document_matches_flat_query(
    document: &str,
    tokenizer: &mut Box<dyn LanceTokenizer>,
    query_tokens: &Tokens,
    phrase_slop: Option<u32>,
) -> Result<bool> {
    let Some(slop) = phrase_slop else {
        return Ok(has_query_token(document, tokenizer, query_tokens));
    };

    let mut document_positions = (0..query_tokens.len())
        .map(|_| Vec::new())
        .collect::<Vec<_>>();
    let mut stream = tokenizer.token_stream_for_doc(document);
    while let Some(token) = stream.next() {
        let position = u32::try_from(token.position).map_err(|_| {
            Error::invalid_input(format!(
                "flat FTS token position exceeds u32: {}",
                token.position
            ))
        })?;
        for (query_index, positions) in document_positions.iter_mut().enumerate() {
            if query_tokens.get_token(query_index) == token.text {
                positions.push(position);
            }
        }
    }
    Ok(phrase_matches_positions(
        query_tokens,
        &document_positions,
        slop,
    ))
}

pub(super) const FLAT_ALL_TOKENS_COL: &str = "all_tokens";
pub(super) const FLAT_QUERY_TOKEN_COUNTS_COL: &str = "query_token_counts";
pub(super) const FLAT_PHRASE_MATCH_COL: &str = "phrase_match";

pub(super) fn phrase_matches_positions(
    query_tokens: &Tokens,
    document_positions: &[Vec<u32>],
    slop: u32,
) -> bool {
    let Some(first_positions) = document_positions.first() else {
        return false;
    };
    if first_positions.is_empty() {
        return false;
    }

    let mut candidates = first_positions.clone();
    debug_assert_eq!(query_tokens.len(), document_positions.len());
    for (query_index, positions) in document_positions.iter().enumerate().skip(1) {
        let Some(query_delta) = query_tokens
            .position(query_index)
            .checked_sub(query_tokens.position(query_index - 1))
        else {
            return false;
        };
        let mut next_candidates = Vec::new();
        for &position in positions {
            let position = u64::from(position);
            if candidates.iter().any(|candidate| {
                let least = u64::from(*candidate) + u64::from(query_delta);
                least <= position && position <= least + u64::from(slop)
            }) {
                next_candidates.push(position as u32);
            }
        }
        if next_candidates.is_empty() {
            return false;
        }
        candidates = next_candidates;
    }
    true
}

/// If we accumulate this many bytes we warn the user they probably want to use an FTS index instead.
pub(super) const BYTES_ACCUMULATED_WARNING_THRESHOLD: u64 = 1024 * 1024 * 1024; // 1GB

/// Consumes a stream of record batches and produces token counts
///
/// The resulting batch will have three columns:
/// - row_id: the row id of the document
/// - all_tokens: the total number of tokens in the document
/// - query_token_counts: a fixed size list of the count of each query token in the document
///
/// This is an unbounded accumulation, however, for most queries, the per-row
/// growth will be fairly small.  As a result we can process millions of tokens
/// with fairly modest memory usage.
///
/// However, it is unwise to do a flat search across billions of rows.  An FTS
/// index should be created instead.
pub(super) async fn tokenize_and_count(
    input: impl Stream<Item = DataFusionResult<RecordBatch>> + Send,
    tokenizer: Box<dyn LanceTokenizer>,
    query_tokens: Arc<Tokens>,
    doc_col_idx: usize,
    elapsed_compute: Option<Time>,
    coordinate_rank: usize,
    phrase_slop: Option<u32>,
) -> DataFusionResult<RecordBatch> {
    let mut output_fields = vec![ROW_ID_FIELD.clone()];
    output_fields.extend(
        (0..coordinate_rank)
            .map(|rank| Field::new(doc_index_storage_column(rank), DataType::UInt32, false)),
    );
    output_fields.extend([
        Field::new(FLAT_ALL_TOKENS_COL, DataType::UInt64, false),
        Field::new(
            FLAT_QUERY_TOKEN_COUNTS_COL,
            DataType::FixedSizeList(
                Arc::new(Field::new("item", DataType::UInt64, true)),
                query_tokens.len() as i32,
            ),
            false,
        ),
    ]);
    if phrase_slop.is_some() {
        output_fields.push(Field::new(FLAT_PHRASE_MATCH_COL, DataType::Boolean, false));
    }
    let output_schema = Arc::new(Schema::new(output_fields));
    let output_schema_clone = output_schema.clone();
    let query_token_indices = Arc::new(query_token_indices(query_tokens.as_ref()));
    let bytes_accumulated = Arc::new(AtomicU64::new(0));
    let bytes_warning_emitted = Arc::new(AtomicBool::new(false));

    let batches = input
        .map(move |batch| {
            let mut tokenizer = tokenizer.box_clone();
            let output_schema = output_schema.clone();
            let query_tokens = query_tokens.clone();
            let query_token_indices = query_token_indices.clone();
            let bytes_accumulated = bytes_accumulated.clone();
            let bytes_warning_emitted = bytes_warning_emitted.clone();
            let elapsed_compute = elapsed_compute.clone();
            spawn_cpu(move || {
                // Time the per-batch CPU work so callers can attribute it to
                // `elapsed_compute` on a metric handle (the spawn_cpu worker
                // thread is invisible to the caller's poll timer otherwise).
                let start = std::time::Instant::now();
                let batch = batch?;
                let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();
                let input_doc_indices = (0..coordinate_rank)
                    .map(|rank| {
                        let column_name = doc_index_storage_column(rank);
                        batch
                            .column_by_name(&column_name)
                            .ok_or_else(|| {
                                datafusion_common::DataFusionError::Internal(format!(
                                    "flat ListElement FTS is missing {column_name}"
                                ))
                            })
                            .map(|column| column.as_primitive::<UInt32Type>())
                    })
                    .collect::<DataFusionResult<Vec<_>>>()?;
                let mut row_ids = UInt64Builder::with_capacity(batch.num_rows());
                let mut doc_indices = (0..coordinate_rank)
                    .map(|_| UInt32Builder::with_capacity(batch.num_rows()))
                    .collect::<Vec<_>>();
                let mut all_token_counts = UInt64Builder::with_capacity(batch.num_rows());
                let mut query_token_counts = FixedSizeListBuilder::with_capacity(
                    UInt64Builder::with_capacity(batch.num_rows() * query_tokens.len()),
                    query_tokens.len() as i32,
                    batch.num_rows(),
                );
                let mut temp_query_token_counts = Vec::with_capacity(query_tokens.len());
                let mut temp_query_positions = (0..query_tokens.len())
                    .map(|_| Vec::new())
                    .collect::<Vec<_>>();
                let mut phrase_matches = phrase_slop
                    .map(|_| BooleanBuilder::with_capacity(batch.num_rows()));
                let mut count_text = |doc: &str,
                                      temp_query_token_counts: &mut Vec<u64>|
                 -> DataFusionResult<(u64, bool)> {
                    for positions in &mut temp_query_positions {
                        positions.clear();
                    }
                    let mut stream = tokenizer.token_stream_for_doc(doc);
                    let mut all_tokens = 0;
                    while let Some(token) = stream.next() {
                        all_tokens += 1;
                        if let Some(token_indices) = query_token_indices.get(&token.text) {
                            for token_index in token_indices {
                                temp_query_token_counts[*token_index] += 1;
                                if phrase_slop.is_some() {
                                    temp_query_positions[*token_index].push(
                                        u32::try_from(token.position).map_err(|_| {
                                            datafusion_common::DataFusionError::Execution(format!(
                                                "flat FTS token position exceeds u32: {}",
                                                token.position
                                            ))
                                        })?,
                                    );
                                }
                            }
                        }
                    }
                    let matches = phrase_slop.is_none_or(|slop| {
                        phrase_matches_positions(
                            query_tokens.as_ref(),
                            &temp_query_positions,
                            slop,
                        )
                    });
                    Ok((all_tokens, matches))
                };
                let mut append_counts = |row_index: usize,
                                         row_id: u64,
                                         all_tokens: u64,
                                         temp_query_token_counts: &[u64],
                                         phrase_match: bool|
                 -> DataFusionResult<()> {
                        row_ids.append_value(row_id);
                        for (builder, input) in
                            doc_indices.iter_mut().zip(input_doc_indices.iter())
                        {
                            builder.append_value(input.value(row_index));
                        }
                        all_token_counts.append_value(all_tokens);
                        for count in temp_query_token_counts.iter().copied() {
                            query_token_counts.values().append_value(count);
                        }
                        query_token_counts.append(true);
                        if let Some(builder) = phrase_matches.as_mut() {
                            builder.append_value(phrase_match);
                        }
                        Ok(())
                    };
                match batch.column(doc_col_idx).data_type() {
                    DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
                        let doc_iter = iter_str_array(batch.column(doc_col_idx));
                        for (row_index, (doc, row_id)) in
                            doc_iter.zip(row_id_array.values().iter()).enumerate()
                        {
                            temp_query_token_counts.clear();
                            temp_query_token_counts
                                .extend(std::iter::repeat_n(0, query_tokens.len()));

                            let (all_tokens, phrase_match) = match doc {
                                Some(doc) => count_text(doc, &mut temp_query_token_counts)?,
                                None => (0, false),
                            };
                            if coordinate_rank > 0 || all_tokens > 0 {
                                append_counts(
                                    row_index,
                                    *row_id,
                                    all_tokens,
                                    &temp_query_token_counts,
                                    phrase_match,
                                )?;
                            }
                        }
                    }
                    DataType::List(_) => {
                        if coordinate_rank != 0 {
                            return DataFusionResult::Err(
                                datafusion_common::DataFusionError::Internal(
                                    "ListElement flat FTS input must be expanded to string documents"
                                        .to_string(),
                                ),
                            );
                        }
                        tokenize_and_count_list::<i32>(
                            batch.column(doc_col_idx),
                            row_id_array,
                            &mut count_text,
                            &mut append_counts,
                            &mut temp_query_token_counts,
                            query_tokens.len(),
                            phrase_slop.is_some(),
                        )?;
                    }
                    DataType::LargeList(_) => {
                        if coordinate_rank != 0 {
                            return DataFusionResult::Err(
                                datafusion_common::DataFusionError::Internal(
                                    "ListElement flat FTS input must be expanded to string documents"
                                        .to_string(),
                                ),
                            );
                        }
                        tokenize_and_count_list::<i64>(
                            batch.column(doc_col_idx),
                            row_id_array,
                            &mut count_text,
                            &mut append_counts,
                            &mut temp_query_token_counts,
                            query_tokens.len(),
                            phrase_slop.is_some(),
                        )?;
                    }
                    data_type => {
                        return DataFusionResult::Err(datafusion_common::DataFusionError::Execution(
                            format!("unsupported data type {} for flat full text search", data_type),
                        ));
                    }
                }
                let row_ids = row_ids.finish();
                let doc_indices = doc_indices
                    .into_iter()
                    .map(|mut builder| builder.finish())
                    .collect::<Vec<_>>();
                let all_token_counts = all_token_counts.finish();
                let query_token_counts = query_token_counts.finish();
                let mut columns = vec![Arc::new(row_ids) as ArrayRef];
                for doc_indices in doc_indices {
                    columns.push(Arc::new(doc_indices) as ArrayRef);
                }
                columns.extend([
                    Arc::new(all_token_counts) as ArrayRef,
                    Arc::new(query_token_counts) as ArrayRef,
                ]);
                if let Some(mut builder) = phrase_matches {
                    columns.push(Arc::new(builder.finish()) as ArrayRef);
                }
                let result_batch = RecordBatch::try_new(output_schema, columns)?;
                let bytes_accumulated = bytes_accumulated.fetch_add(result_batch.get_array_memory_size() as u64, Ordering::Relaxed);
                if bytes_accumulated > BYTES_ACCUMULATED_WARNING_THRESHOLD && !bytes_warning_emitted.swap(true, Ordering::Relaxed) {
                    tracing::warn!("Flat full text search is accumulating a large number of bytes.  Consider using an FTS index instead.");
                }

                if let Some(t) = &elapsed_compute {
                    t.add_duration(start.elapsed());
                }
                DataFusionResult::Ok(result_batch)
            })
        })
        .buffered(get_num_compute_intensive_cpus())
        .try_collect::<Vec<_>>()
        .await?;

    Ok(arrow::compute::concat_batches(
        &output_schema_clone,
        &batches,
    )?)
}

pub(super) fn tokenize_and_count_list<ListOffset: OffsetSizeTrait>(
    doc_col: &ArrayRef,
    row_id_array: &arrow_array::PrimitiveArray<UInt64Type>,
    count_text: &mut impl FnMut(&str, &mut Vec<u64>) -> DataFusionResult<(u64, bool)>,
    append_counts: &mut impl FnMut(usize, u64, u64, &[u64], bool) -> DataFusionResult<()>,
    temp_query_token_counts: &mut Vec<u64>,
    query_tokens_len: usize,
    match_phrase: bool,
) -> DataFusionResult<()> {
    let doc_array = doc_col.as_list::<ListOffset>();
    match doc_array.value_type() {
        DataType::Utf8 | DataType::LargeUtf8 => {}
        data_type => {
            return Err(datafusion_common::DataFusionError::Execution(format!(
                "unsupported list item data type {} for flat full text search",
                data_type
            )));
        }
    }

    for i in 0..row_id_array.len() {
        temp_query_token_counts.clear();
        temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens_len));
        let mut all_tokens = 0;
        let mut phrase_match = false;
        if !doc_array.is_null(i) {
            let elements = doc_array.value(i);
            if match_phrase {
                let mut document = String::new();
                for element in iter_str_array(elements.as_ref()).flatten() {
                    if !document.is_empty() {
                        document.push(' ');
                    }
                    document.push_str(element);
                }
                (all_tokens, phrase_match) = count_text(&document, temp_query_token_counts)?;
            } else {
                for element in iter_str_array(elements.as_ref()).flatten() {
                    all_tokens += count_text(element, temp_query_token_counts)?.0;
                }
            }
        }
        if all_tokens > 0 {
            append_counts(
                i,
                row_id_array.value(i),
                all_tokens,
                temp_query_token_counts,
                phrase_match,
            )?;
        }
    }

    Ok(())
}

/// One chunk's worth of blended per-row BM25F scoring inputs: everything a flat
/// `combined_fields` scan retains until the corpus statistics are complete.
/// `8 + 4·(1 + t)` bytes per row for `t` terms, independent of the column count.
pub(in super::super) struct BlendedRows {
    /// `_rowid` of every retained row.
    pub(in super::super) row_ids: Vec<u64>,
    /// `dl'(d) = Σ_f w_f · dl_f(d)`, one per retained row.
    pub(in super::super) doc_lengths: Vec<f32>,
    /// `tf'(t, d) = Σ_f w_f · tf_f(t, d)`, `num_terms`-strided by row.
    pub(in super::super) term_freqs: Vec<f32>,
}

/// Count the query tokens in every target column of a flat batch stream and blend
/// them into per-row BM25F scoring inputs.
///
/// Unlike [`tokenize_and_count`], this walks all target columns in one pass,
/// since a row may be empty in one field and populated in another. The
/// per-column counts are dropped with each batch once blended; only the
/// fixed-size [`FlatFieldStats`] needs them, so it is folded here, gated by
/// `stats_masks`.
///
/// Every non-empty row is folded into the statistics, but only a row that can
/// score is retained: one `emit_mask` selects and that holds a query term (every
/// query term when `require_all_terms`). So the retained memory tracks the
/// matches, not the scanned rows.
#[allow(clippy::too_many_arguments)]
pub(in super::super) async fn tokenize_and_blend_multi(
    input: impl Stream<Item = DataFusionResult<RecordBatch>> + Send,
    tokenizer: Box<dyn LanceTokenizer>,
    query_tokens: Arc<Tokens>,
    doc_col_indices: Arc<Vec<usize>>,
    weights: Arc<Vec<f32>>,
    stats_masks: Arc<Vec<Arc<RowAddrMask>>>,
    emit_mask: Option<Arc<RowAddrMask>>,
    require_all_terms: bool,
    elapsed_compute: Option<Time>,
) -> DataFusionResult<(Vec<BlendedRows>, FlatFieldStats)> {
    let num_columns = doc_col_indices.len();
    let num_terms = query_tokens.len();
    let retained_bytes_per_row = (size_of::<u64>() + (1 + num_terms) * size_of::<f32>()) as u64;
    let query_token_indices = Arc::new(query_token_indices(query_tokens.as_ref()));
    let bytes_accumulated = Arc::new(AtomicU64::new(0));
    let bytes_warning_emitted = Arc::new(AtomicBool::new(false));

    let tasks = input
        .map(move |batch| {
            let mut tokenizer = tokenizer.box_clone();
            let query_token_indices = query_token_indices.clone();
            let doc_col_indices = doc_col_indices.clone();
            let weights = weights.clone();
            let stats_masks = stats_masks.clone();
            let emit_mask = emit_mask.clone();
            let bytes_accumulated = bytes_accumulated.clone();
            let bytes_warning_emitted = bytes_warning_emitted.clone();
            let elapsed_compute = elapsed_compute.clone();
            spawn_cpu(move || {
                let start = std::time::Instant::now();
                let batch = batch?;
                let num_rows = batch.num_rows();
                let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();

                // Dense `(row, column)`-strided scratch, freed with the batch:
                // whether a row is retained depends on all of its columns, so every
                // column is counted before any row is filtered.
                let mut all_tokens = vec![0u64; num_rows * num_columns];
                let mut counts = vec![0u64; num_rows * num_columns * num_terms];
                for (column_slot, doc_col_idx) in doc_col_indices.iter().enumerate() {
                    count_column_into(
                        batch.column(*doc_col_idx),
                        num_columns,
                        num_terms,
                        column_slot,
                        &mut tokenizer,
                        &query_token_indices,
                        &mut all_tokens,
                        &mut counts,
                    )?;
                }

                let mut blended = BlendedRows {
                    row_ids: Vec::with_capacity(num_rows),
                    doc_lengths: Vec::with_capacity(num_rows),
                    term_freqs: Vec::with_capacity(num_rows * num_terms),
                };
                let mut stats = FlatFieldStats::zeros(num_columns, num_terms);
                let counts_stride = num_columns * num_terms;
                for row in 0..num_rows {
                    let lengths = &all_tokens[row * num_columns..(row + 1) * num_columns];
                    // Empty in every target column: `dl' == 0` and `tf' == 0`, so the
                    // row cannot score. Dropping it here also keeps it out of the
                    // corpus statistics, matching the indexed side (a row absent from
                    // every field's DocSet contributes to no `docCount_f`).
                    if lengths.iter().all(|length| *length == 0) {
                        continue;
                    }
                    let row_id = row_id_array.value(row);
                    let row_counts = &counts[row * counts_stride..(row + 1) * counts_stride];
                    stats.fold_row(row_id, &stats_masks, lengths, row_counts);
                    if emit_mask.as_ref().is_some_and(|mask| !mask.selected(row_id)) {
                        continue;
                    }

                    // Accumulate in column-slot order. f32 addition is not
                    // associative, so this order is part of the contract: it keeps
                    // the scores identical to summing the same contributions
                    // downstream from a materialized breakdown.
                    let terms_start = blended.term_freqs.len();
                    blended.term_freqs.resize(terms_start + num_terms, 0.0);
                    let term_freqs = &mut blended.term_freqs[terms_start..];
                    let mut doc_length = 0f32;
                    for (column, weight) in weights.iter().enumerate() {
                        doc_length += weight * lengths[column] as f32;
                        let base = column * num_terms;
                        for (term, tf) in term_freqs.iter_mut().enumerate() {
                            *tf += weight * row_counts[base + term] as f32;
                        }
                    }
                    // `And` spans the virtual field, so it checks the blended `tf'`.
                    let matches = if require_all_terms {
                        term_freqs.iter().all(|tf| *tf > 0.0)
                    } else {
                        term_freqs.iter().any(|tf| *tf > 0.0)
                    };
                    if !matches {
                        blended.term_freqs.truncate(terms_start);
                        continue;
                    }
                    blended.row_ids.push(row_id);
                    blended.doc_lengths.push(doc_length);
                }

                let bytes_accumulated = bytes_accumulated.fetch_add(
                    blended.row_ids.len() as u64 * retained_bytes_per_row,
                    Ordering::Relaxed,
                );
                if bytes_accumulated > BYTES_ACCUMULATED_WARNING_THRESHOLD
                    && !bytes_warning_emitted.swap(true, Ordering::Relaxed)
                {
                    tracing::warn!(
                        "Flat combined_fields search is accumulating a large number of bytes.  Consider indexing every target column instead."
                    );
                }
                if let Some(t) = &elapsed_compute {
                    t.add_duration(start.elapsed());
                }
                DataFusionResult::Ok((blended, stats))
            })
        })
        .buffered(get_num_compute_intensive_cpus());

    // Folded incrementally rather than collected and concatenated: the retained
    // rows are already the whole accumulation, so a second contiguous copy would
    // double the peak for nothing.
    let mut tasks = std::pin::pin!(tasks);
    let mut chunks = Vec::new();
    // Seeded rather than taken from the first batch so that a scan with no batches
    // at all still reports statistics, as all-zero totals.
    let mut totals = FlatFieldStats::zeros(num_columns, num_terms);
    while let Some(counted) = tasks.next().await {
        let (blended, stats) = counted?;
        totals.merge(stats);
        chunks.push(blended);
    }
    Ok((chunks, totals))
}

/// Tokenize one column into the `(row, column_slot)`-strided slots of
/// `all_tokens` / `counts`. Mirrors [`tokenize_and_count`]'s per-type handling: a
/// list column sums over its elements and a null leaves the slots at zero.
#[allow(clippy::too_many_arguments)]
fn count_column_into(
    doc_col: &ArrayRef,
    num_columns: usize,
    num_terms: usize,
    column_slot: usize,
    tokenizer: &mut Box<dyn LanceTokenizer>,
    query_token_indices: &HashMap<String, Vec<usize>>,
    all_tokens: &mut [u64],
    counts: &mut [u64],
) -> DataFusionResult<()> {
    let mut record = |row: usize, doc: &str| {
        let slot = row * num_columns + column_slot;
        let base = slot * num_terms;
        let mut stream = tokenizer.token_stream_for_doc(doc);
        while let Some(token) = stream.next() {
            all_tokens[slot] += 1;
            if let Some(token_indices) = query_token_indices.get(&token.text) {
                for token_index in token_indices {
                    counts[base + *token_index] += 1;
                }
            }
        }
    };
    match doc_col.data_type() {
        DataType::Utf8 | DataType::LargeUtf8 => {
            for (row, doc) in iter_str_array(doc_col.as_ref()).enumerate() {
                if let Some(doc) = doc {
                    record(row, doc);
                }
            }
        }
        // One string per row is the whole contract: a caller whose field path
        // crosses a `List` owes a column already reduced to a row's document text
        // (the planner uses `flatten_fts_document_column`), because reducing it
        // here would count the elements separately and diverge from the single row
        // document the index builder wrote.
        data_type => {
            return Err(datafusion_common::DataFusionError::Execution(format!(
                "flat combined_fields search needs one string per row per column, got {}; \
                 reduce a list-bearing field path to its row document first",
                data_type
            )));
        }
    }
    Ok(())
}

pub(super) fn query_token_indices(query_tokens: &Tokens) -> HashMap<String, Vec<usize>> {
    let mut indices = HashMap::new();
    for idx in 0..query_tokens.len() {
        indices
            .entry(query_tokens.get_token(idx).to_string())
            .or_insert_with(Vec::new)
            .push(idx);
    }
    indices
}

/// Initialize the BM25 scorer
///
/// In order to calculate BM25 scores we need to know token counts for the entire corpus.  We extract these from the
/// counted input of the flat search combined with any counts recorded for the indexed portion.
pub(super) fn initialize_scorer(
    base_scorer: Option<&MemBM25Scorer>,
    query_tokens: &Tokens,
    counted_input: &RecordBatch,
) -> MemBM25Scorer {
    let mut total_tokens = 0;
    let mut num_docs = 0;
    let mut all_token_counts = vec![0; query_tokens.len()];

    if let Some(base_scorer) = base_scorer {
        total_tokens += base_scorer.total_tokens;
        num_docs += base_scorer.num_docs;
        for (token_index, token) in query_tokens.into_iter().enumerate() {
            all_token_counts[token_index] = base_scorer.num_docs_containing_token(token) as u64;
        }
    }

    num_docs += counted_input.num_rows();
    total_tokens +=
        arrow::compute::sum(counted_input[FLAT_ALL_TOKENS_COL].as_primitive::<UInt64Type>())
            .unwrap_or_default();

    let mut input_token_counters = counted_input[FLAT_QUERY_TOKEN_COUNTS_COL]
        .as_fixed_size_list()
        .values()
        .as_primitive::<UInt64Type>()
        .values()
        .iter()
        .copied();

    for _ in 0..counted_input.num_rows() {
        for token_count in all_token_counts.iter_mut() {
            if input_token_counters.next().unwrap_or_default() > 0 {
                *token_count += 1;
            }
        }
    }

    let token_counts_map = all_token_counts
        .into_iter()
        .enumerate()
        .map(|(token_index, count)| {
            (
                query_tokens.get_token(token_index).to_string(),
                count as usize,
            )
        })
        .collect::<HashMap<String, usize>>();
    MemBM25Scorer::new(total_tokens, num_docs, token_counts_map)
}

pub(super) fn flat_bm25_score(
    query_tokens: &Tokens,
    counted_input: &RecordBatch,
    scorer: &MemBM25Scorer,
    document_granularity: DocumentGranularity,
    operator: Operator,
    boost: f32,
    phrase_slop: Option<u32>,
) -> Result<RecordBatch> {
    let mut row_ids_builder = UInt64Builder::with_capacity(counted_input.num_rows());
    let mut scores_builder = Float32Builder::with_capacity(counted_input.num_rows());
    let coordinate_rank = document_coordinate_rank(counted_input.schema().as_ref());
    let input_doc_indices = (0..coordinate_rank)
        .map(|rank| {
            let coordinate_column = doc_index_storage_column(rank);
            counted_input
                .column_by_name(&coordinate_column)
                .ok_or_else(|| {
                    Error::internal(format!(
                        "flat ListElement FTS is missing {coordinate_column}"
                    ))
                })
                .map(|column| column.as_primitive::<UInt32Type>())
        })
        .collect::<Result<Vec<_>>>()?;
    if document_granularity.is_list_element() != (coordinate_rank > 0) {
        return Err(Error::internal(
            "flat FTS document granularity does not match its coordinate columns".to_string(),
        ));
    }
    let mut doc_indices_builder = document_granularity.is_list_element().then(|| {
        ListBuilder::new(UInt32Builder::new()).with_field(Field::new(
            "item",
            DataType::UInt32,
            false,
        ))
    });
    let query_groups = query_position_groups(query_tokens);

    let mut row_ids_iter = counted_input[ROW_ID]
        .as_primitive::<UInt64Type>()
        .values()
        .iter()
        .copied();
    let mut all_token_counts_iter = counted_input[FLAT_ALL_TOKENS_COL]
        .as_primitive::<UInt64Type>()
        .values()
        .iter()
        .copied();
    let mut query_token_counts_iter = counted_input[FLAT_QUERY_TOKEN_COUNTS_COL]
        .as_fixed_size_list()
        .values()
        .as_primitive::<UInt64Type>()
        .values()
        .iter()
        .copied();
    let phrase_matches = phrase_slop
        .map(|_| {
            counted_input
                .column_by_name(FLAT_PHRASE_MATCH_COL)
                .ok_or_else(|| Error::internal("flat phrase FTS is missing phrase matches"))
                .and_then(|column| {
                    column
                        .as_any()
                        .downcast_ref::<BooleanArray>()
                        .ok_or_else(|| {
                            Error::internal("flat phrase matches must be a Boolean column")
                        })
                })
        })
        .transpose()?;
    for input_index in 0..counted_input.num_rows() {
        let num_tokens_in_doc = all_token_counts_iter.next().expect_ok()?;
        let row_id = row_ids_iter.next().expect_ok()?;
        let mut query_token_counts = Vec::with_capacity(query_tokens.len());
        for _ in query_tokens {
            query_token_counts.push(query_token_counts_iter.next().expect_ok()?);
        }
        if num_tokens_in_doc == 0 {
            continue;
        }
        if operator == Operator::And
            && !query_groups
                .iter()
                .all(|group| group.iter().any(|idx| query_token_counts[*idx] > 0))
        {
            continue;
        }
        if phrase_matches.is_some_and(|matches| !matches.value(input_index)) {
            continue;
        }
        let doc_norm = K1 * (1.0 - B + B * num_tokens_in_doc as f32 / scorer.avg_doc_length());
        let mut score = 0.0;
        for (token, freq) in query_tokens.into_iter().zip(query_token_counts) {
            let freq = freq as f32;
            let idf = idf(scorer.num_docs_containing_token(token), scorer.num_docs());
            score += idf * (freq * (K1 + 1.0) / (freq + doc_norm));
        }
        if score > 0.0 {
            row_ids_builder.append_value(row_id);
            if let Some(builder) = doc_indices_builder.as_mut() {
                for input_doc_index in &input_doc_indices {
                    builder
                        .values()
                        .append_value(input_doc_index.value(input_index));
                }
                builder.append(true);
            }
            scores_builder.append_value(score * boost);
        }
    }

    let row_ids = row_ids_builder.finish();
    let scores = scores_builder.finish();
    let mut columns = vec![Arc::new(row_ids) as ArrayRef];
    if let Some(mut builder) = doc_indices_builder {
        columns.push(Arc::new(builder.finish()) as ArrayRef);
    }
    columns.push(Arc::new(scores) as ArrayRef);
    let batch = RecordBatch::try_new(fts_schema(document_granularity), columns)?;
    Ok(batch)
}

pub(super) fn query_position_groups(query_tokens: &Tokens) -> Vec<Vec<usize>> {
    let mut groups = Vec::new();
    let mut current_position = None;
    for idx in 0..query_tokens.len() {
        let position = query_tokens.position(idx);
        if current_position != Some(position) {
            current_position = Some(position);
            groups.push(Vec::new());
        }
        groups
            .last_mut()
            .expect("a group should exist after pushing for position")
            .push(idx);
    }
    groups
}

#[deprecated(
    note = "use `flat_bm25_search_stream_with_metrics` to record CPU compute \
            time on a metric handle; pass `None` for the old behavior"
)]
pub async fn flat_bm25_search_stream(
    input: SendableRecordBatchStream,
    doc_col: String,
    query: String,
    tokenizer: Box<dyn LanceTokenizer>,
    base_scorer: Option<MemBM25Scorer>,
    target_batch_size: usize,
) -> DataFusionResult<SendableRecordBatchStream> {
    flat_bm25_search_stream_with_metrics(
        input,
        doc_col,
        query,
        tokenizer,
        base_scorer,
        target_batch_size,
        None,
    )
    .await
}

/// Same as [`flat_bm25_search_stream`] but accepts an optional `Time` handle
/// that, if provided, will receive the CPU time spent in (a) per-batch
/// tokenization on the `spawn_cpu` worker threads and (b) the synchronous
/// scoring phase. This lets a calling `ExecutionPlan` report accurate
/// `elapsed_compute` without double-counting upstream poll time.
pub async fn flat_bm25_search_stream_with_metrics(
    input: SendableRecordBatchStream,
    doc_col: String,
    query: String,
    tokenizer: Box<dyn LanceTokenizer>,
    base_scorer: Option<MemBM25Scorer>,
    target_batch_size: usize,
    elapsed_compute: Option<Time>,
) -> DataFusionResult<SendableRecordBatchStream> {
    flat_bm25_search_stream_with_metrics_and_operator(
        input,
        doc_col,
        query,
        tokenizer,
        base_scorer,
        target_batch_size,
        Operator::Or,
        elapsed_compute,
    )
    .await
}

/// Same as [`flat_bm25_search_stream_with_metrics`] but applies the provided
/// match operator when deciding whether a flat-scanned row is a hit.
///
/// # Examples
///
/// ```no_run
/// # async fn example(
/// #     input: datafusion::execution::SendableRecordBatchStream,
/// # ) -> Result<(), Box<dyn std::error::Error>> {
/// use lance_index::scalar::inverted::{
///     flat_bm25_search_stream_with_metrics_and_operator, query::Operator, InvertedIndexParams,
/// };
///
/// let tokenizer = InvertedIndexParams::code().build()?;
/// let _stream = flat_bm25_search_stream_with_metrics_and_operator(
///     input,
///     "code".to_string(),
///     "Result".to_string(),
///     tokenizer,
///     None,
///     1024,
///     Operator::And,
///     None,
/// )
/// .await?;
/// # Ok(())
/// # }
/// ```
#[allow(clippy::too_many_arguments)]
pub async fn flat_bm25_search_stream_with_metrics_and_operator(
    input: SendableRecordBatchStream,
    doc_col: String,
    query: String,
    tokenizer: Box<dyn LanceTokenizer>,
    base_scorer: Option<MemBM25Scorer>,
    target_batch_size: usize,
    operator: Operator,
    elapsed_compute: Option<Time>,
) -> DataFusionResult<SendableRecordBatchStream> {
    Ok(flat_bm25_search_stream_with_options_and_scorer(
        input,
        doc_col,
        query,
        tokenizer,
        base_scorer,
        FlatBm25SearchOptions {
            target_batch_size,
            elapsed_compute,
            document_granularity: DocumentGranularity::Row,
            operator,
            boost: 1.0,
            phrase_slop: None,
        },
    )
    .await?
    .0)
}

/// Options for flat BM25 search execution.
pub struct FlatBm25SearchOptions {
    /// Maximum output record batch size.
    pub target_batch_size: usize,
    /// Optional DataFusion metric tracking compute time.
    pub elapsed_compute: Option<Time>,
    /// Logical FTS document boundary.
    pub document_granularity: DocumentGranularity,
    /// Match operator applied within each logical document.
    pub operator: Operator,
    /// Score multiplier for the match query.
    pub boost: f32,
    /// Phrase slop. When set, documents must contain the query tokens in order.
    pub phrase_slop: Option<u32>,
}

/// Run a flat BM25 search and return the scorer initialized from both the
/// optional indexed corpus and the flat input corpus.
pub async fn flat_bm25_search_stream_with_options_and_scorer(
    input: SendableRecordBatchStream,
    doc_col: String,
    query: String,
    tokenizer: Box<dyn LanceTokenizer>,
    base_scorer: Option<MemBM25Scorer>,
    options: FlatBm25SearchOptions,
) -> DataFusionResult<(SendableRecordBatchStream, MemBM25Scorer)> {
    let FlatBm25SearchOptions {
        target_batch_size,
        elapsed_compute,
        document_granularity,
        operator,
        boost,
        phrase_slop,
    } = options;
    let mut tokenizer = tokenizer;
    let output_schema = fts_schema(document_granularity);

    // Pre-await synchronous work: query tokenization + chunk-stream setup.
    let pre_await_start = std::time::Instant::now();
    let query_tokens = Arc::new(try_collect_query_tokens(&query, &mut tokenizer)?);

    // A query that tokenizes to no terms (e.g. only stop words) has no
    // searchable content and matches nothing. Return early rather than
    // proceeding. This mirrors the indexed search path, which already
    // short-circuits on empty query tokens.
    if query_tokens.is_empty() {
        let scorer = base_scorer.unwrap_or_else(|| MemBM25Scorer::new(0, 0, HashMap::new()));
        return Ok((
            Box::pin(RecordBatchStreamAdapter::new(
                output_schema,
                stream::empty::<DataFusionResult<RecordBatch>>(),
            )),
            scorer,
        ));
    }

    let input_schema = input.schema();
    let doc_col_idx = input_schema.index_of(&doc_col)?;
    let coordinate_rank = document_coordinate_rank(input_schema.as_ref());
    if document_granularity.is_list_element() != (coordinate_rank > 0) {
        return Err(datafusion_common::DataFusionError::Execution(
            "flat FTS document granularity does not match its input coordinate columns".to_string(),
        ));
    }

    // Accumulate small batches until this threshold before dispatching a task.
    const ACCUMULATE_BYTES: usize = 256 * 1024;
    // Slice oversized batches down to roughly this size.
    const SLICE_BYTES: usize = 512 * 1024;

    // Phase 1 - rechunk the input stream into appropriately sized chunks.  Tokenization is
    // fairly CPU-intensive, and we don't need too much data to justify a new thread task.
    let chunked = lance_arrow::stream::rechunk_stream_by_size(
        input,
        input_schema,
        ACCUMULATE_BYTES,
        SLICE_BYTES,
    );
    if let Some(t) = &elapsed_compute {
        t.add_duration(pre_await_start.elapsed());
    }

    // Phase 2 - For each row we need to know the total number of tokens and the count of each
    // of the query tokens.  For example, if the query is "book" and the row is "the book shop"
    // and we are tokenizing with a whitespace tokenizer, we need to know that there are 3 tokens
    // and the token book appears once.
    let counted_input = tokenize_and_count(
        chunked,
        tokenizer,
        query_tokens.clone(),
        doc_col_idx,
        elapsed_compute.clone(),
        coordinate_rank,
        phrase_slop,
    )
    .await?;

    // Phase 3 - Calculate final scores (this is fairly cheap, probably don't need to parallelize).
    // All post-await work is synchronous; time the scorer + score + slicing loop together.
    let post_await_start = std::time::Instant::now();
    let scorer = initialize_scorer(base_scorer.as_ref(), query_tokens.as_ref(), &counted_input);
    let scores = flat_bm25_score(
        query_tokens.as_ref(),
        &counted_input,
        &scorer,
        document_granularity,
        operator,
        boost,
        phrase_slop,
    )?;

    // Finally we emit batches according to the target batch size
    let batches = slice_into_batches(scores, target_batch_size);
    if let Some(t) = &elapsed_compute {
        t.add_duration(post_await_start.elapsed());
    }
    Ok((
        Box::pin(RecordBatchStreamAdapter::new(
            output_schema,
            stream::iter(batches),
        )),
        scorer,
    ))
}

/// Cut one scored batch into the `target_batch_size` chunks a flat FTS scan
/// emits. Slicing is zero-copy, so the whole scan's scores stay in one
/// allocation.
pub(in super::super) fn slice_into_batches(
    batch: RecordBatch,
    target_batch_size: usize,
) -> Vec<DataFusionResult<RecordBatch>> {
    let num_out_batches = batch.num_rows().div_ceil(target_batch_size);
    let mut batches = Vec::with_capacity(num_out_batches);
    for i in 0..num_out_batches {
        let start = i * target_batch_size;
        let len = (batch.num_rows() - start).min(target_batch_size);
        batches.push(Ok(batch.slice(start, len)));
    }
    batches
}

pub fn is_phrase_query(query: &str) -> bool {
    phrase_query_text(query).is_some()
}

pub(super) fn phrase_query_text(query: &str) -> Option<&str> {
    query.strip_prefix('\"')?.strip_suffix('\"')
}