LabelList index returns stale results after an in-place merge_insert

[mmatczuk]

LabelList index returns stale results after an in-place merge_insert

Description

After a merge_insert rewrites a list column in place, its LabelList index can still claim to cover that fragment. array_has then returns the stale index answer while an unindexed full scan returns the updated row. The trigger needs a scalar index on the merge key and a source batch that contains all fields in a different order from the dataset schema; the backslash label from the original finding is not material.

Steps to reproduce

The complete program below uses only the public Rust API:

1. Write (id: UInt64, v0: LargeList<LargeUtf8>) rows (1, ["a"]) and (2, ["b"]).
2. Create a Bitmap index on id and a LabelList index on v0.
3. Run merge_insert on id, with UpdateAll, from a source batch ordered (v0, id), changing row 2 to ["z"].
4. Query array_has(v0, 'z') with scalar indices enabled and disabled.

From this checkout:

cd bugs
cargo run --bin labellist-array-has-index-divergence

Create Cargo.toml:

[package]
name = "lance-labellist-repro"
version = "0.1.0"
edition = "2024"

[dependencies]
arrow = "58"
futures = "0.3"
tempfile = "3"
tokio = { version = "1.23", features = ["macros", "rt-multi-thread"] }
lance = { git = "https://github.com/lance-format/lance", tag = "v11.0.0-beta.4", default-features = false }
lance-index = { git = "https://github.com/lance-format/lance", tag = "v11.0.0-beta.4", default-features = false }

Create src/main.rs:

//! A LabelList index keeps answering from stale data after an in-place merge_insert.
//!
//! Setup: a `LargeList<LargeUtf8>` column `v0` with a LabelList index, plus a Bitmap index on
//! the merge-insert join key `id`. A merge_insert that carries every target column but lists
//! them in a different order than the dataset schema rewrites `v0` *in place* (the fragment
//! count does not change), and the commit leaves the LabelList index covering that fragment.
//! `array_has(v0, '<new label>')` then answers purely from the stale index and returns nothing,
//! while the same predicate on a full scan finds the row.
//!
//! The backslash label from the fuzz artifact is not load-bearing: an ordinary label diverges
//! the same way (both are checked below), so this is an index-staleness bug, not an
//! escaping/quoting bug.
//!
//! Run: cargo run --bin labellist-array-has-index-divergence

use std::sync::Arc;

use arrow::array::builder::{LargeListBuilder, LargeStringBuilder};
use arrow::array::{ArrayRef, RecordBatch, RecordBatchIterator, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema};
use futures::TryStreamExt as _;
use lance::Dataset;
use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched, WhenNotMatchedBySource};
use lance::index::DatasetIndexExt as _;
use lance_index::IndexType;
use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams};

fn id_field() -> Field {
    Field::new("id", DataType::UInt64, false)
}

fn list_field() -> Field {
    Field::new(
        "v0",
        DataType::LargeList(Arc::new(Field::new("item", DataType::LargeUtf8, true))),
        true,
    )
}

/// Dataset schema: `id` first.
fn target_schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![id_field(), list_field()]))
}

/// Merge-insert source schema: the same fields, `v0` first. Nothing is missing — only the
/// order differs, which is what makes the slow merge-insert path treat it as a partial schema.
fn source_schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![list_field(), id_field()]))
}

fn batch(schema: Arc<Schema>, rows: &[(u64, &str)]) -> RecordBatch {
    let ids: ArrayRef = Arc::new(UInt64Array::from(
        rows.iter().map(|(k, _)| *k).collect::<Vec<_>>(),
    ));
    let mut builder = LargeListBuilder::new(LargeStringBuilder::new())
        .with_field(Arc::new(Field::new("item", DataType::LargeUtf8, true)));
    for (_, label) in rows {
        builder.values().append_value(label);
        builder.append(true);
    }
    let lists: ArrayRef = Arc::new(builder.finish());
    let columns: Vec<ArrayRef> = if schema.field(0).name() == "id" {
        vec![ids, lists]
    } else {
        vec![lists, ids]
    };
    RecordBatch::try_new(schema, columns).expect("build batch")
}

fn reader(batch: RecordBatch) -> Box<dyn arrow::array::RecordBatchReader + Send> {
    let schema = batch.schema();
    Box::new(RecordBatchIterator::new([Ok(batch)], schema))
}

/// Ids matching `filter`, with the scalar index enabled or disabled.
async fn ids(dataset: &Dataset, filter: &str, use_scalar_index: bool) -> Vec<u64> {
    let mut scan = dataset.scan();
    scan.project(&["id"]).expect("project");
    scan.filter(filter).expect("filter");
    scan.use_scalar_index(use_scalar_index);
    scan.scan_in_order(true);
    let batches = scan
        .try_into_stream()
        .await
        .expect("plan scan")
        .try_collect::<Vec<_>>()
        .await
        .expect("execute scan");
    let mut out = Vec::new();
    for batch in &batches {
        let column = batch
            .column_by_name("id")
            .expect("id column")
            .as_any()
            .downcast_ref::<UInt64Array>()
            .expect("id is u64");
        out.extend((0..column.len()).map(|i| column.value(i)));
    }
    out.sort_unstable();
    out
}

/// Indexed and full-scan answers to `array_has(v0, '<label>')` after row 2's label has been
/// replaced by an in-place merge_insert. Both should be `[2]`.
async fn indexed_vs_full_scan(label: &str) -> (Vec<u64>, Vec<u64>) {
    let dir = tempfile::tempdir().expect("tempdir");
    let uri = dir.path().to_str().expect("utf8 path");

    let initial = batch(target_schema(), &[(1, "a"), (2, "b")]);
    let mut dataset = Dataset::write(reader(initial), uri, None)
        .await
        .expect("write dataset");

    // A scalar index on the join key takes merge_insert off its fast path. Only the slow path
    // has the order-sensitive full-schema check that routes this write to the in-place update.
    dataset
        .create_index(
            &["id"],
            IndexType::Bitmap,
            Some("id_idx".to_owned()),
            &ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
            true,
        )
        .await
        .expect("create bitmap index on id");
    dataset
        .create_index(
            &["v0"],
            IndexType::LabelList,
            Some("v0_idx".to_owned()),
            &ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
            true,
        )
        .await
        .expect("create labellist index on v0");

    // Replace row 2's label with `label`, through a source that carries both columns in the
    // opposite order.
    let mut builder =
        MergeInsertBuilder::try_new(Arc::new(dataset.clone()), vec!["id".into()]).expect("merge");
    builder
        .when_matched(WhenMatched::UpdateAll)
        .when_not_matched(WhenNotMatched::InsertAll)
        .when_not_matched_by_source(WhenNotMatchedBySource::Keep);
    let job = builder.try_build().expect("build merge job");
    let source = batch(source_schema(), &[(2, label)]);
    let (updated, _stats) = job
        .execute_reader(reader(source))
        .await
        .expect("execute merge_insert");
    let dataset = (*updated).clone();

    // The row data really did change (this is the harness's full-scan oracle check, in miniature).
    let stored = ids(&dataset, &format!("array_has(v0, '{label}')"), false).await;
    assert_eq!(stored, vec![2], "merge_insert did not store the new label");
    // One fragment: the column was rewritten in place rather than into a new fragment.
    assert_eq!(
        dataset.get_fragments().len(),
        1,
        "expected the in-place column-rewrite path"
    );

    let filter = format!("array_has(v0, '{label}')");
    let with_index = ids(&dataset, &filter, true).await;
    let full_scan = ids(&dataset, &filter, false).await;
    (with_index, full_scan)
}

#[tokio::main(flavor = "multi_thread")]
async fn main() {
    // The artifact's label is a single backslash; an ordinary label behaves identically.
    for label in ["z", "\\"] {
        let (with_index, full_scan) = indexed_vs_full_scan(label).await;
        println!(
            "array_has(v0, '{label}'): scalar_index=true -> {with_index:?}, \
             scalar_index=false -> {full_scan:?} (expected [2])"
        );
    }

    let (with_index, full_scan) = indexed_vs_full_scan("\\").await;
    assert_eq!(
        with_index, full_scan,
        "`array_has(v0, '\\')` disagrees with the full scan: the LabelList index still covers \
         the fragment its column was rewritten in"
    );
}

Run:

cargo run

Expected behavior

The LabelList index must be invalidated or updated when its indexed column is rewritten. Both scans must return row 2.

Lance version

v11.0.0-beta.4

Language binding

Rust

Environment

Linux x86_64, local filesystem

Logs / traceback

array_has(v0, 'z'): scalar_index=true -> [], scalar_index=false -> [2] (expected [2])
array_has(v0, '\\'): scalar_index=true -> [], scalar_index=false -> [2] (expected [2])

thread 'main' panicked at labellist-array-has-index-divergence/repro.rs:181:5:
assertion `left == right` failed: `array_has(v0, '\\')` disagrees with the full scan: the LabelList index still covers the fragment its column was rewritten in
  left: []
 right: [2]