ZoneMap index drops live rows after deferred-remap compaction

[mmatczuk]

ZoneMap index drops live rows after deferred-remap compaction

Description

When compaction runs with defer_index_remap: true, a ZoneMap index still answers in the pre-compaction fragment-address domain. A scalar-index scan therefore yields no rows although a full scan finds them. For IS NULL, the ZoneMap answer is marked exact, so the scan has no refinement filter and silently returns the false negative. This is distinct from #8076: that issue requires stable row IDs and fails while translating stale address-domain index results; this trigger uses enable_stable_row_ids = false and returns an incorrect empty result after deferred remapping.

Steps to reproduce

The complete program below uses only the public Rust API:

1. Write 12 rows in three fragments with nullable v0: Int64; rows 1, 5, and 9 are null.
2. Create a ZoneMap index on v0, and confirm v0 IS NULL returns {1, 5, 9} with and without the index.
3. Compact to one fragment with defer_index_remap: true.
4. Run v0 IS NULL with scalar indices enabled and disabled.

From this checkout:

cd bugs
cargo run --bin zonemap-is-null-false-negative

Create Cargo.toml:

[package]
name = "lance-zonemap-remap-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:

//! ZoneMap index: `v0 IS NULL` silently returns no rows after a compaction that defers the
//! index remap.
//!
//! Minimized from `crashes/dataset_ops/index-index-query-idxn-zonemap-on-vn-vn-null-26433b81.bin`
//! (harness assertion at `src/dataset_ops/index.rs:411`).
//!
//! Shape: one nullable value column, a handful of rows spread over several small fragments, a
//! ZoneMap index on that column, then `compact_files` with `defer_index_remap: true`. Compaction
//! rewrites the fragments under new ids and records the old -> new mapping in the fragment reuse
//! index instead of rewriting the scalar index. Every other scalar index applies that mapping to
//! its search output; `ZoneMapIndex` loads the remapper into `self.fri` and never uses it, so the
//! row addresses it returns still name the pre-compaction fragments and match nothing live.
//!
//! For `IS NULL` the miss is total and silent: `ZoneMapIndex::search` answers that predicate from
//! an exact null-row-address bitmap (`SearchResult::exact`), so the scan trusts the empty
//! intersection instead of rechecking, and the query returns zero rows.
//!
//! Run with: cargo run --bin zonemap-is-null-false-negative

use std::collections::BTreeSet;
use std::sync::Arc;

use arrow::array::{ArrayRef, Int64Array, RecordBatch, RecordBatchIterator, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema};
use futures::TryStreamExt as _;
use lance::Dataset;
use lance::dataset::optimize::{CompactionOptions, compact_files};
use lance::dataset::{WriteMode, WriteParams};
use lance::index::DatasetIndexExt as _;
use lance_index::IndexType;
use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams};

/// Rows written, as ids 0..ROWS.
const ROWS: u64 = 12;
/// Ids whose `v0` is null.
const NULL_IDS: [u64; 3] = [1, 5, 9];
/// Small enough that the 12 rows land in several fragments, which is what gives compaction
/// something to merge (and therefore new fragment ids to remap onto).
const ROWS_PER_FILE: usize = 4;

fn schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::UInt64, false),
        Field::new("v0", DataType::Int64, true),
    ]))
}

fn rows() -> RecordBatch {
    let ids: ArrayRef = Arc::new(UInt64Array::from_iter_values(0..ROWS));
    let values: ArrayRef = Arc::new(Int64Array::from_iter((0..ROWS).map(|id| {
        if NULL_IDS.contains(&id) {
            None
        } else {
            Some(id as i64 * 10)
        }
    })));
    RecordBatch::try_new(schema(), vec![ids, values]).expect("build batch")
}

/// Ids matching `filter`, with scalar index use either enabled or disabled.
async fn ids_where(dataset: &Dataset, filter: &str, use_scalar_index: bool) -> BTreeSet<u64> {
    let mut scan = dataset.scan();
    scan.project(&["id"]).expect("project id");
    scan.filter(filter).expect("set filter");
    scan.use_scalar_index(use_scalar_index);
    let batches: Vec<RecordBatch> = scan
        .try_into_stream()
        .await
        .expect("plan scan")
        .try_collect()
        .await
        .expect("run scan");
    batches
        .iter()
        .flat_map(|batch| {
            let column = batch
                .column(0)
                .as_any()
                .downcast_ref::<UInt64Array>()
                .expect("id is UInt64")
                .clone();
            (0..column.len()).map(move |i| column.value(i))
        })
        .collect()
}

async fn compare(dataset: &Dataset, filter: &str, when: &str) -> (BTreeSet<u64>, BTreeSet<u64>) {
    let indexed = ids_where(dataset, filter, true).await;
    let full_scan = ids_where(dataset, filter, false).await;
    println!("{when}: `{filter}` indexed={indexed:?} full_scan={full_scan:?}");
    (indexed, full_scan)
}

#[tokio::main(flavor = "multi_thread")]
async fn main() {
    let dir = tempfile::tempdir().expect("temp dir");
    let uri = dir
        .path()
        .join("zonemap")
        .to_str()
        .expect("utf8")
        .to_owned();

    let reader = RecordBatchIterator::new([Ok(rows())], schema());
    let mut dataset = Dataset::write(
        reader,
        &uri,
        Some(WriteParams {
            mode: WriteMode::Overwrite,
            max_rows_per_file: ROWS_PER_FILE,
            // `defer_index_remap` is rejected on datasets with stable row ids: the deferred
            // remap this bug lives in only exists for address-style row ids.
            enable_stable_row_ids: false,
            ..Default::default()
        }),
    )
    .await
    .expect("write dataset");
    println!(
        "wrote {ROWS} rows in {} fragments",
        dataset.get_fragments().len()
    );

    dataset
        .create_index(
            &["v0"],
            IndexType::ZoneMap,
            Some("idx0".to_owned()),
            &ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
            true,
        )
        .await
        .expect("create zonemap index");

    // Sanity check: the freshly built index answers the predicate correctly, so what follows is
    // about compaction and not about the index build.
    let (indexed, full_scan) = compare(&dataset, "v0 IS NULL", "after create_index").await;
    assert_eq!(
        indexed, full_scan,
        "precondition failed: the fresh ZoneMap index already disagrees with the full scan"
    );
    assert_eq!(
        full_scan,
        BTreeSet::from(NULL_IDS),
        "precondition failed: the full scan does not see the null rows we wrote"
    );

    // The one load-bearing operation: merge the small fragments and leave the index remap to the
    // fragment reuse index.
    compact_files(
        &mut dataset,
        CompactionOptions {
            target_rows_per_fragment: 512,
            defer_index_remap: true,
            ..Default::default()
        },
        None,
    )
    .await
    .expect("compact");
    println!("compacted into {} fragments", dataset.get_fragments().len());

    let mut plan_scan = dataset.scan();
    plan_scan.project(&["id"]).expect("project id");
    plan_scan.filter("v0 IS NULL").expect("set filter");
    plan_scan.use_scalar_index(true);
    println!(
        "indexed plan after compaction:\n{}",
        plan_scan.explain_plan(true).await.expect("explain plan")
    );

    // The same staleness hits every ZoneMap predicate, printed for context: equality and range
    // probes go through the inexact zone path and come back empty too. `IS NULL` is the worst
    // case only because its answer is reported as exact, so nothing rechecks it.
    let _ = compare(&dataset, "v0 = 20", "after compaction (equality)").await;
    let _ = compare(&dataset, "v0 > 20", "after compaction (range)").await;

    let (indexed, full_scan) = compare(&dataset, "v0 IS NULL", "after compaction").await;
    assert_eq!(
        indexed, full_scan,
        "BUG: `v0 IS NULL` through the ZoneMap index dropped live null rows \
         (indexed scan {indexed:?} vs full scan {full_scan:?})"
    );

    println!("no divergence: the bug did not reproduce");
}

Run:

cargo run

Expected behavior

The ZoneMap result must be remapped to live fragment addresses, or the index must not be used. Both scans must return {1, 5, 9}.

Lance version

v11.0.0-beta.4

Language binding

Rust

Environment

Linux x86_64, local filesystem

Logs / traceback

after create_index: `v0 IS NULL` indexed={1, 5, 9} full_scan={1, 5, 9}
compacted into 1 fragments
indexed plan after compaction:
LanceRead: ... full_filter=v0 IS NULL, refine_filter=--
  ScalarIndexQuery: query=[v0 IS NULL]@idx0(ZoneMap)

after compaction: `v0 IS NULL` indexed={} full_scan={1, 5, 9}

thread 'main' panicked: BUG: `v0 IS NULL` through the ZoneMap index dropped live null rows
  left: {}
 right: {1, 5, 9}