[wkalt]

An update scanned every column of the rows it rewrites, so the old values of the columns it assigns were read and then discarded. For large values, such as binary or blob columns, that read dominated the update.

The scan now leaves out assigned columns that no SET expression reads, and apply_updates lays its results out in dataset order. This removes the placeholder blob v2 assignments used to avoid the same read.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[wkalt]

Measured against main at 77123c9 (15.0.0-beta.4), release build, memory:// datasets.

Time (1024 rows, median of 5 runs):

Update	main	this PR
assign a 16 KiB binary column, 50% of rows	4.0 ms	1.3 ms
assign a 256 KiB binary column, 50% of rows	15.7 ms	1.1 ms
assign a 256 KiB binary column, all rows	26.6 ms	0.9 ms
assign a 256 KiB blob v2 column, 50% of rows	1.7 ms	1.2 ms
control: assign id, 256 KiB column carried	31.5 ms	33.5 ms
control: payload = payload, 256 KiB	33.5 ms	30.1 ms

Both versions write the same new values, so the gap is the read of the old ones; the blob v2 case was already mostly avoided by the placeholder this removes.

Bytes read (64 rows of 16 KiB, 32 updated, counted by the object store's IO tracker):

Update	main	this PR
SET payload = 'x'	516 KiB in 6 requests	4 KiB in 4 requests
SET payload = payload	516 KiB in 6 requests	516 KiB in 6 requests

Benchmark (place at rust/lance/src/dataset/write/update_bench.rs, add #[cfg(test)] mod update_bench; to write.rs, run cargo test --release -p lance --lib -- --ignored update_bench --nocapture --test-threads=1)

// Bytes read and wall time of UpdateBuilder.
// Run: cargo test --release -p lance --lib -- --ignored update_bench --nocapture

use std::sync::Arc;
use std::time::{Duration, Instant};

use arrow_array::{
    ArrayRef, BinaryArray, Int32Array, RecordBatch, RecordBatchIterator,
};
use arrow_schema::{DataType, Field, Schema};

use crate::Dataset;
use crate::dataset::WriteParams;
use crate::dataset::write::update::UpdateBuilder;

const ROWS: usize = 1024;
const RUNS: usize = 5;

async fn dataset(dir: &str, payload_bytes: usize, blob: bool) -> Dataset {
    let ids = Arc::new(Int32Array::from_iter_values(0..ROWS as i32)) as ArrayRef;
    let (field, payload): (Field, ArrayRef) = if blob {
        let mut builder = crate::BlobArrayBuilder::new(ROWS);
        for row in 0..ROWS {
            builder.push_bytes(vec![row as u8; payload_bytes]).unwrap();
        }
        (crate::blob_field("payload", true), builder.finish().unwrap())
    } else {
        let values = (0..ROWS).map(|row| vec![row as u8; payload_bytes]);
        let values = values.collect::<Vec<_>>();
        let array = BinaryArray::from_iter_values(values.iter());
        (Field::new("payload", DataType::Binary, true), Arc::new(array))
    };
    let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false), field]));
    let batch = RecordBatch::try_new(schema.clone(), vec![ids, payload]).unwrap();
    let params = WriteParams {
        data_storage_version: Some(lance_file::version::LanceFileVersion::V2_2),
        ..Default::default()
    };
    Dataset::write(RecordBatchIterator::new([Ok(batch)], schema), dir, Some(params))
        .await
        .unwrap()
}

// An in-memory store: its tracker counts every byte the update reads, which
// local files bypass, and times exclude disk.
async fn measure(case: usize, payload_bytes: usize, blob: bool, filter: &str, column: &str, value: &str) -> (u64, Duration) {
    let mut times = Vec::new();
    let mut read = 0;
    for run in 0..RUNS {
        let uri = format!("memory://update-bench-{case}-{run}");
        let ds = Arc::new(dataset(&uri, payload_bytes, blob).await);
        ds.object_store.io_stats_incremental();
        let start = Instant::now();
        let mut builder = UpdateBuilder::new(ds.clone());
        if !filter.is_empty() {
            builder = builder.update_where(filter).unwrap();
        }
        builder.set(column, value).unwrap().build().unwrap().execute().await.unwrap();
        times.push(start.elapsed());
        read = ds.object_store.io_stats_incremental().read_bytes;
    }
    times.sort();
    (read, times[RUNS / 2])
}

#[tokio::test(flavor = "multi_thread")]
#[ignore]
async fn update_bench() {
    let cases: &[(&str, usize, bool, &str, &str, &str)] = &[
        ("a 16KiB binary, assign payload, 50%", 16 << 10, false, "id < 512", "payload", "X'00'"),
        ("b 256KiB binary, assign payload, 50%", 256 << 10, false, "id < 512", "payload", "X'00'"),
        ("c 256KiB binary, assign payload, 100%", 256 << 10, false, "", "payload", "X'00'"),
        ("d 256KiB binary, assign id only, 50%", 256 << 10, false, "id < 512", "id", "id + 10000"),
        ("e 256KiB binary, payload = payload, 50%", 256 << 10, false, "id < 512", "payload", "payload"),
        ("f 256KiB blob v2, assign payload, 50%", 256 << 10, true, "id < 512", "payload", "X'00'"),
    ];
    println!("BENCH case | read MiB | median ms");
    for (case, (name, bytes, blob, filter, column, value)) in cases.iter().enumerate() {
        let (read, time) = measure(case, *bytes, *blob, filter, column, value).await;
        println!(
            "BENCH {name} | {:.1} | {:.1}",
            read as f64 / (1 << 20) as f64,
            time.as_secs_f64() * 1000.0
        );
    }
}

/// Bytes read, counted with the setup the update tests count reliably:
/// 64 rows of 16 KiB in a `memory://` dataset, half of them updated.
#[tokio::test(flavor = "multi_thread")]
#[ignore]
async fn update_bench_bytes() {
    use arrow_array::types::Int32Type;
    for (name, value) in [("literal", "'x'"), ("payload = payload", "payload")] {
        let batch = lance_datagen::gen_batch()
            .col("id", lance_datagen::array::step::<Int32Type>())
            .col(
                "payload",
                lance_datagen::array::rand_fixedbin(lance_datagen::ByteCount::from(16 * 1024), true),
            )
            .into_batch_rows(lance_datagen::RowCount::from(64))
            .unwrap();
        let ds = Dataset::write(
            RecordBatchIterator::new([Ok(batch.clone())], batch.schema()),
            "memory://",
            None,
        )
        .await
        .unwrap();
        ds.object_store.io_stats_incremental();
        let ds = Arc::new(ds);
        let result = UpdateBuilder::new(ds.clone())
            .update_where("id < 32")
            .unwrap()
            .set("payload", value)
            .unwrap()
            .build()
            .unwrap()
            .execute()
            .await
            .unwrap();
        let stats = ds.object_store.io_stats_incremental();
        println!(
            "BYTES {name} | rows {} | read {} KiB in {} requests (matched payloads: 512 KiB)",
            result.rows_updated,
            stats.read_bytes / 1024,
            stats.read_iops
        );
    }
}

[zhangyue19921010]

Could we rebuild the scan projection when retrying against a newer dataset? If another writer adds column c after the job is built, the retry uses the new schema but retains the old scan_columns, causing the update to fail with update scan is missing 'c'.

[zhangyue19921010]

Is that make sense that use project_with_schema() here to preserve literal field names?

[lance-gatekeeper]

✅ Gate recommendation: approve.

Per-attempt planning resolves the schema-change retry failure, and the blob-only update fix remains intact. No blocking issues remain.

[wkalt]

@zhangyue19921010 thanks, updated

[LuQQiu]

@claude review

[LuQQiu]

See if claude find other small issues.
Looks good.
Guess it also depends on the table shape, if the table has other even heavier columns, then reading out other even heavier columns may become a high cost.
But very heavy column can also go through blob V2.
Not matter what, this PR is an improvement

[claude]

🔴 A caller who passes a backtick-quoted column name to set now gets a silent no-op that still reports the rows as updated. On the base branch the update fails. apply_updates keys results by the raw set argument but fills columns by field.name(), so an unmatched result is dropped and the old value is copied through. Fix: store the resolved field.name as the key at update.rs:223 rather than column.as_ref(), so every name-keyed lookup sees the real name. Also return an internal error when assigned still holds entries after the output columns are assembled.

Why this was flagged

UpdateBuilder::set (update.rs:155) looks the column up with dataset.schema().field(column). That goes through parse_field_path, which strips backticks, so set("a", "1") finds field a. update.rs:223 still stores the key as the raw string "a". In UpdatePlan::try_new (update.rs:376), updates.contains_key(&field.name) is false for a, so the old column is scanned. In apply_updates, the new values go into assigned under "a". The output loop at update.rs:795 calls assigned.remove("a"), finds nothing, and copies the scanned old column instead. The rows are rewritten unchanged, and rows_updated reports them as updated. On the base branch, batch.replace_column_by_name("a", ...) failed with "Field a does not exist", so the caller got an error. Blob v2 columns are affected the same way: the blob_columns key is also the quoted string. Nothing checks that assigned is empty after assembly. The Python binding (python/src/dataset.rs:2095) passes user column strings straight to set.

Verification: Trigger: a caller passes a backtick-quoted column name to UpdateBuilder::set, e.g. set("a", "1"). update.rs:223 still stores the raw key. The output loop (update.rs:792-800) calls assigned.remove(field.name()) with "a" and falls back to the old values. The leftover "a" entry is never checked, so no error is raised. On the base branch the update fails.

[wkalt]

addressed