panic when using merge_insert with an index and # of rows * embedding dimension > u32::MAX

Lance version

4.0.0-beta.10 + 311ff1ad1b37e784dceed8d30873282ccacdc907

What happened?

See original issue for context. And one PR that adresses this https://github.com/lance-format/lance/pull/6148#issuecomment-4039498958

With an updated test that uses an index we still encounter a panic.

With the docs suggesting to use a scalar index to speed up merge insert, this should be fixed.

thread 'main' (19105945) panicked at /Users/valkum/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/arrow-data-57.3.0/src/data.rs:551:9:
assertion failed: (offset + length) <= self.len()
stack backtrace:
   0: __rustc::rust_begin_unwind
   1: core::panicking::panic_fmt
   2: core::panicking::panic
   3: arrow_data::data::ArrayData::slice
   4: <arrow_array::array::fixed_size_list_array::FixedSizeListArray as core::convert::From<arrow_data::data::ArrayData>>::from
   5: arrow_select::take::take_fixed_size_list
   6: arrow_select::take::take_impl
   7: arrow_select::take::take
   8: datafusion_physical_plan::joins::utils::build_batch_from_indices
   9: datafusion_physical_plan::joins::hash_join::stream::HashJoinStream::poll_next_impl
  10: <S as futures_core::stream::TryStream>::try_poll_next
  11: <futures_util::stream::try_stream::try_flatten::TryFlatten<St> as futures_core::stream::Stream>::poll_next
  12: <futures_util::stream::try_stream::try_flatten::TryFlatten<St> as futures_core::stream::Stream>::poll_next
  13: <futures_util::stream::stream::map::Map<St,F> as futures_core::stream::Stream>::poll_next
  14: lance::dataset::write::write_fragments_internal::{{closure}}::{{closure}}
  15: lance::dataset::write::write_fragments_internal::{{closure}}
  16: lance::dataset::write::merge_insert::MergeInsertJob::execute_uncommitted_impl::{{closure}}
  17: <futures_util::future::future::map::Map<Fut,F> as core::future::future::Future>::poll
  18: <core::pin::Pin<P> as core::future::future::Future>::poll
  19: <futures_util::future::future::map::Map<Fut,F> as core::future::future::Future>::poll
  20: <lancedb::table::NativeTable as lancedb::table::BaseTable>::merge_insert::{{closure}}
  21: with_index::main::{{closure}}
  22: tokio::runtime::park::CachedParkThread::block_on
  23: tokio::runtime::context::runtime::enter_runtime
  24: tokio::runtime::runtime::Runtime::block_on
  25: with_index::main
note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.

Are there known steps to reproduce?

main.rs:

//! Reproducer: `merge_insert` panics on tables with >23M rows containing a
//! `FixedSizeListArray` column (e.g. embeddings).
//!
//! Root cause: u32 overflow in `take_value_indices_from_fixed_size_list`
//! (arrow-select take.rs). `merge_insert` with `when_not_matched_by_source_delete`
//! uses a FULL OUTER JOIN whose HashJoin concatenates all build-side rows into a
//! single RecordBatch. When the FixedSizeList column has more than
//! `u32::MAX / value_length` rows, `take()` overflows internally and panics with:
//!
//!     assertion failed: (offset + length) <= self.len()
//!
//! For value_length=184, the threshold is 23,342,213 rows.
//!
//! Run:  cargo run --release          (requires ~1 GB RAM, panics)
//!       cargo run --release -- 23342213  (5k rows, passes)

use arrow_array::{
    BooleanArray, FixedSizeListArray, RecordBatch, RecordBatchIterator, UInt32Array,
};
use arrow_buffer::{BooleanBuffer, MutableBuffer};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;

const DIM: i32 = 184;
const NUM_ROWS: usize = 24_000_000;
const CHUNK: usize = 1_000_000;

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    let n: usize = std::env::args()
        .nth(1)
        .and_then(|s| s.parse().ok())
        .unwrap_or(NUM_ROWS);

    let threshold = (u32::MAX as u64) / (DIM as u64);
    eprintln!("rows={n}, dim={DIM}, overflow threshold={threshold}");

    let tmp = tempfile::tempdir().unwrap();
    let db = lancedb::connect(tmp.path().to_str().unwrap())
        .execute()
        .await
        .unwrap();

    eprintln!("creating table...");
    let table = db
        .create_table("t", batches(n).collect::<Result<Vec<_>, _>>().unwrap())
        .execute()
        .await
        .unwrap();
    eprintln!("table has {} rows", table.count_rows(None).await.unwrap());

    table
        .create_index(&["id"], lancedb::index::Index::Auto)
        .execute()
        .await
        .unwrap();

    // merge_insert with when_not_matched_by_source_delete forces a FULL OUTER JOIN.
    // DataFusion's HashJoin concat_batches all build-side rows into one RecordBatch,
    // then calls take() → u32 overflow → panic.
    eprintln!("running merge_insert...");
    let mut merge = table.merge_insert(&["id"]);
    merge
        .when_matched_update_all(None)
        .when_not_matched_by_source_delete(None);

    merge
        .execute(Box::new(RecordBatchIterator::new(batches(n), schema())))
        .await
        .unwrap();

    eprintln!("merge_insert succeeded (bug not triggered at this scale)");
}

fn schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::UInt32, false),
        Field::new(
            "vec",
            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Boolean, true)), DIM),
            true,
        ),
    ]))
}

// In Production there will be changes between the batches,
// but the amount of removed rows is small and thus not relevant to the bug.
fn batches(
    total: usize,
) -> impl Iterator<Item = Result<RecordBatch, arrow_schema::ArrowError>> + Send {
    let mut off = 0;
    std::iter::from_fn(move || {
        if off >= total {
            return None;
        }
        let n = (total - off).min(CHUNK);
        let b = make_batch(off, n);
        off += n;
        Some(Ok(b))
    })
}

fn make_batch(start: usize, count: usize) -> RecordBatch {
    let ids: Vec<u32> = (start..start + count).map(|i| i as u32).collect();
    let n = count * DIM as usize;
    let buf: arrow_buffer::Buffer = MutableBuffer::from_len_zeroed((n + 7) / 8).into();
    let bools = BooleanArray::new(BooleanBuffer::new(buf, 0, n), None);
    let fsl = FixedSizeListArray::new(
        Arc::new(Field::new("item", DataType::Boolean, true)),
        DIM,
        Arc::new(bools),
        None,
    );
    RecordBatch::try_new(
        schema(),
        vec![Arc::new(UInt32Array::from(ids)), Arc::new(fsl)],
    )
    .unwrap()
}