From 72b05f800343f4a1abd43cecc0f457c645c4f99f Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Fri, 10 Jul 2026 15:30:58 +0100 Subject: [PATCH 1/6] iss-986: regression tests + merge-shape probes for field-level updates Red (flip green with the partial-staging change): - scalar_indexes::update_of_unindexed_property_preserves_other_index_coverage (whole-row update merges mark every field modified, so Lance prunes the touched fragments from every index's fragment_bitmap) - writes::single_update_stages_partial_matched_only_source (a sole update should stage only key+assigned columns as a matched-only merge) Green pins (behavior that must not change): - writes::mixed_insert_update_same_table_stages_full_row_upsert - writes::chained_updates_same_table_stage_full_rows - writes::empty_match_update_stages_no_merge - validators::partial_update_completes_composite_unique_group Probe infra: MergeWriteProbes gains per-call MergeShape (source columns + unmatched-row disposition), recorded in stage_merge_insert. --- crates/omnigraph/src/instrumentation.rs | 31 ++++- crates/omnigraph/src/table_store.rs | 13 +- crates/omnigraph/tests/scalar_indexes.rs | 43 ++++++ crates/omnigraph/tests/validators.rs | 72 +++++++++- crates/omnigraph/tests/writes.rs | 159 +++++++++++++++++++++++ 5 files changed, 314 insertions(+), 4 deletions(-) diff --git a/crates/omnigraph/src/instrumentation.rs b/crates/omnigraph/src/instrumentation.rs index f4ef9ef9..80a95642 100644 --- a/crates/omnigraph/src/instrumentation.rs +++ b/crates/omnigraph/src/instrumentation.rs @@ -175,6 +175,12 @@ pub struct MergeWriteProbes { pub stage_append_rows: Arc, pub stage_merge_insert_calls: Arc, pub stage_merge_insert_rows: Arc, + /// Per-call shape of each `stage_merge_insert`: the source batch's column + /// names (schema order) and whether unmatched source rows insert + /// (`WhenNotMatched::InsertAll`) or drop (`DoNothing`). Lets a fitness test + /// assert a partial-schema matched-only update stages exactly + /// (key + assigned + completion) columns — the RFC-022 staging shape. + pub merge_shapes: Arc>>, /// Inline vector-index (IVF) builds. The fast-forward adopt path defers /// index coverage to the reconciler, so an adopt merge must do 0 of these. pub create_vector_index_calls: Arc, @@ -184,10 +190,22 @@ pub struct MergeWriteProbes { pub scan_staged_combined_calls: Arc, } +/// One `stage_merge_insert` call's observable shape (see +/// [`MergeWriteProbes::merge_shapes`]). +#[derive(Clone, Debug)] +pub struct MergeShape { + pub source_columns: Vec, + pub inserts_unmatched: bool, +} + impl MergeWriteProbes { pub fn stage_append_calls(&self) -> u64 { self.stage_append_calls.load(Ordering::Relaxed) } + /// Snapshot of the recorded per-call merge shapes, in call order. + pub fn merge_shapes(&self) -> Vec { + self.merge_shapes.lock().unwrap().clone() + } pub fn stage_append_rows(&self) -> u64 { self.stage_append_rows.load(Ordering::Relaxed) } @@ -227,12 +245,21 @@ pub(crate) fn record_stage_append(rows: u64) { }); } -/// Record one `stage_merge_insert` of `rows` rows against the active probes. +/// Record one `stage_merge_insert` of `rows` rows against the active probes, +/// with its observable shape (source columns + unmatched-row disposition). /// No-op in production (no probes installed). -pub(crate) fn record_stage_merge_insert(rows: u64) { +pub(crate) fn record_stage_merge_insert( + rows: u64, + source_columns: Vec, + inserts_unmatched: bool, +) { let _ = MERGE_WRITE_PROBES.try_with(|p| { p.stage_merge_insert_calls.fetch_add(1, Ordering::Relaxed); p.stage_merge_insert_rows.fetch_add(rows, Ordering::Relaxed); + p.merge_shapes.lock().unwrap().push(MergeShape { + source_columns, + inserts_unmatched, + }); }); } diff --git a/crates/omnigraph/src/table_store.rs b/crates/omnigraph/src/table_store.rs index d34cb899..72918fd9 100644 --- a/crates/omnigraph/src/table_store.rs +++ b/crates/omnigraph/src/table_store.rs @@ -1177,6 +1177,13 @@ impl TableStore { )); } let merged_rows = batch.num_rows() as u64; + let inserts_unmatched = matches!(when_not_matched, WhenNotMatched::InsertAll); + let source_columns: Vec = batch + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); // Precondition for the FirstSeen workaround below: every call path that // reaches stage_merge_insert (load, MutationStaging::finalize, @@ -1219,7 +1226,11 @@ impl TableStore { .map_err(|e| OmniError::Lance(e.to_string()))?; // Record only after the staging write succeeds, so a failed write does // not inflate the probe (matches `stage_append`/`stage_append_stream`). - crate::instrumentation::record_stage_merge_insert(merged_rows); + crate::instrumentation::record_stage_merge_insert( + merged_rows, + source_columns, + inserts_unmatched, + ); // Operation::Update { removed_fragment_ids, updated_fragments, new_fragments, .. } — // `new_fragments` are the freshly inserted rows; `updated_fragments` // are rewrites of existing fragments that include both retained and diff --git a/crates/omnigraph/tests/scalar_indexes.rs b/crates/omnigraph/tests/scalar_indexes.rs index 8d8a3f0a..b76b985e 100644 --- a/crates/omnigraph/tests/scalar_indexes.rs +++ b/crates/omnigraph/tests/scalar_indexes.rs @@ -72,3 +72,46 @@ async fn node_scalar_and_enum_index_columns_get_btree() { "un-annotated column should have no scalar index, got {note_cov:?}" ); } + +const TOUCH_NOTE: &str = r#" +query touch_note($slug: String, $note: String) { + update Item set { note: $note } where slug = $slug +} +"#; + +// RFC-022 acceptance cell (red → green): updating an UN-indexed property must +// not degrade index coverage on any OTHER column. The whole-row update merge +// marks every column modified, so Lance prunes the touched fragments from every +// index (`prune_updated_fields_from_indices`) and the rewritten rows land +// outside coverage; the partial-schema update (fields_modified = assigned +// columns only, in-place column patch on the same fragment) leaves the id/enum/ +// scalar BTREEs intact. This is the erosion class measured in the field-level- +// updates investigation: keyed operations degrade ~+1 read/fragment after +// whole-row updates until an optimize. +#[tokio::test] +async fn update_of_unindexed_property_preserves_other_index_coverage() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut db = Omnigraph::init(uri, SCHEMA).await.unwrap(); + load_jsonl(&mut db, DATA, LoadMode::Overwrite).await.unwrap(); + + db.mutate( + "main", + TOUCH_NOTE, + "touch_note", + ¶ms(&[("$slug", "a"), ("$note", "touched")]), + ) + .await + .unwrap(); + + let snap = snapshot_main(&db).await.unwrap(); + let ds = snap.open("node:Item").await.unwrap(); + for col in ["id", "status", "published", "rank"] { + let cov = TableStore::key_column_index_coverage(&ds, col).await.unwrap(); + assert_eq!( + cov, + IndexCoverage::Indexed, + "updating un-indexed 'note' must not degrade the index on '{col}', got {cov:?}" + ); + } +} diff --git a/crates/omnigraph/tests/validators.rs b/crates/omnigraph/tests/validators.rs index 226968f5..31da6bbb 100644 --- a/crates/omnigraph/tests/validators.rs +++ b/crates/omnigraph/tests/validators.rs @@ -8,7 +8,7 @@ mod helpers; use omnigraph::db::Omnigraph; use omnigraph::loader::{LoadMode, load_jsonl}; -use helpers::{count_rows, mutate_main, params}; +use helpers::{count_rows, mixed_params, mutate_main, params}; const ENUM_SCHEMA: &str = r#" node Person { @@ -708,3 +708,73 @@ async fn cardinality_rejected_on_jsonl_load() { err ); } + +// ─── RFC-022: composite-unique completion under partial updates ───────────── + +const COMPOSITE_UNIQUE_SCHEMA: &str = r#" +node Slot { + name: String @key + room: String? + hour: I32? + @unique(room, hour) +} +"#; + +const COMPOSITE_MUTATIONS: &str = r#" +query insert_slot($name: String, $room: String, $hour: I32) { + insert Slot { name: $name, room: $room, hour: $hour } +} +query move_slot($name: String, $room: String) { + update Slot set { room: $room } where name = $name +} +"#; + +/// An update assigning only PART of a composite `@unique` group must still +/// validate the whole group — the unassigned member's committed value completes +/// the tuple. Guards RFC-022's completion-column projection rule: a partial +/// update that assigns `room` must also carry `hour` (old value) into the +/// change-set, or the (room, hour) collision below would go undetected. +#[tokio::test] +async fn partial_update_completes_composite_unique_group() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let db = omnigraph::db::Omnigraph::init(uri, COMPOSITE_UNIQUE_SCHEMA) + .await + .unwrap(); + + for (name, room, hour) in [("s1", "roomA", 9i64), ("s2", "roomB", 9)] { + db.mutate( + "main", + COMPOSITE_MUTATIONS, + "insert_slot", + &mixed_params(&[("$name", name), ("$room", room)], &[("$hour", hour)]), + ) + .await + .unwrap(); + } + + // Moving s2 into roomA collides with s1's (roomA, 9). + let err = db + .mutate( + "main", + COMPOSITE_MUTATIONS, + "move_slot", + ¶ms(&[("$name", "s2"), ("$room", "roomA")]), + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("@unique"), + "moving s2 to (roomA, 9) must be a composite-unique violation, got: {err}" + ); + + // A non-colliding move succeeds. + db.mutate( + "main", + COMPOSITE_MUTATIONS, + "move_slot", + ¶ms(&[("$name", "s2"), ("$room", "roomC")]), + ) + .await + .unwrap(); +} diff --git a/crates/omnigraph/tests/writes.rs b/crates/omnigraph/tests/writes.rs index aac44bf7..5f7206bf 100644 --- a/crates/omnigraph/tests/writes.rs +++ b/crates/omnigraph/tests/writes.rs @@ -2065,3 +2065,162 @@ async fn filtered_read_after_append_and_delete_is_consistent() { assert_eq!(got, expected, "filtered read for {name}"); } } + +// ─── RFC-022: field-level update staging shape ────────────────────────────── + +const MIXED_INSERT_UPDATE: &str = r#" +query mix($newname: String, $age: I32, $target: String, $newage: I32) { + insert Person { name: $newname, age: $age } + update Person set { age: $newage } where name = $target +} +"#; + +const TWO_UPDATES: &str = r#" +query two($a: String, $aage: I32, $b: String, $bage: I32) { + update Person set { age: $aage } where name = $a + update Person set { age: $bage } where name = $b +} +"#; + +/// RFC-022 acceptance (red → green): a query whose ONLY op on a table is one +/// `update` stages a PARTIAL source — exactly (merge key + assigned columns) — +/// as a matched-only merge (`WhenNotMatched::DoNothing`). The whole-row path +/// stages every column with `InsertAll`. +#[tokio::test] +async fn single_update_stages_partial_matched_only_source() { + use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; + let dir = tempfile::tempdir().unwrap(); + let db = init_and_load(&dir).await; + + let probes = MergeWriteProbes::default(); + with_merge_write_probes( + probes.clone(), + db.mutate( + "main", + MUTATION_QUERIES, + "set_age", + &mixed_params(&[("$name", "Alice")], &[("$age", 41)]), + ), + ) + .await + .unwrap(); + + let shapes = probes.merge_shapes(); + assert_eq!(shapes.len(), 1, "one staged merge for the touched table"); + let mut cols = shapes[0].source_columns.clone(); + cols.sort(); + assert_eq!( + cols, + vec!["age".to_string(), "id".to_string()], + "a sole update stages only (key + assigned) columns, got {:?}", + shapes[0].source_columns + ); + assert!( + !shapes[0].inserts_unmatched, + "a sole update stages a matched-only merge (WhenNotMatched::DoNothing)" + ); +} + +/// Fallback pin: a query mixing insert + update on ONE table stages full rows +/// as an upsert — partial and full batches cannot share one merge source, and +/// one table commits at most one version per query (invariant 4). Pins the +/// RFC-022 fallback so the partial path can never split a table's commit. +#[tokio::test] +async fn mixed_insert_update_same_table_stages_full_row_upsert() { + use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; + let dir = tempfile::tempdir().unwrap(); + let db = init_and_load(&dir).await; + + let probes = MergeWriteProbes::default(); + with_merge_write_probes( + probes.clone(), + db.mutate( + "main", + MIXED_INSERT_UPDATE, + "mix", + &mixed_params( + &[("$newname", "Zed"), ("$target", "Alice")], + &[("$age", 20), ("$newage", 42)], + ), + ), + ) + .await + .unwrap(); + + let shapes = probes.merge_shapes(); + assert_eq!(shapes.len(), 1, "insert+update coalesce into one staged merge"); + let mut cols = shapes[0].source_columns.clone(); + cols.sort(); + assert!( + cols.contains(&"name".to_string()) && cols.contains(&"age".to_string()), + "mixed insert+update stages full rows, got {:?}", + shapes[0].source_columns + ); + assert!( + shapes[0].inserts_unmatched, + "mixed insert+update keeps upsert semantics (InsertAll)" + ); +} + +/// Fallback pin: two updates on one table in one query stage full rows (a +/// later update's read-your-writes scan must see the earlier update's full +/// effect; two partial batches with different assigned sets cannot share one +/// uniform-schema merge source — schema-level partial semantics make a union +/// unsound). Behavior identical to today. +#[tokio::test] +async fn chained_updates_same_table_stage_full_rows() { + use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; + let dir = tempfile::tempdir().unwrap(); + let db = init_and_load(&dir).await; + + let probes = MergeWriteProbes::default(); + with_merge_write_probes( + probes.clone(), + db.mutate( + "main", + TWO_UPDATES, + "two", + &mixed_params( + &[("$a", "Alice"), ("$b", "Bob")], + &[("$aage", 51), ("$bage", 52)], + ), + ), + ) + .await + .unwrap(); + + let shapes = probes.merge_shapes(); + assert_eq!(shapes.len(), 1, "chained updates coalesce into one staged merge"); + let mut cols = shapes[0].source_columns.clone(); + cols.sort(); + assert!( + cols.contains(&"name".to_string()), + "chained updates stage full rows (fallback), got {:?}", + shapes[0].source_columns + ); +} + +/// An update matching zero rows stages nothing and commits nothing — +/// unchanged by RFC-022 (the early return precedes staging). +#[tokio::test] +async fn empty_match_update_stages_no_merge() { + use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; + let dir = tempfile::tempdir().unwrap(); + let db = init_and_load(&dir).await; + + let probes = MergeWriteProbes::default(); + let result = with_merge_write_probes( + probes.clone(), + db.mutate( + "main", + MUTATION_QUERIES, + "set_age", + &mixed_params(&[("$name", "nobody-here")], &[("$age", 1)]), + ), + ) + .await + .unwrap(); + + assert_eq!(result.affected_nodes, 0); + assert!(probes.merge_shapes().is_empty(), "no merge staged for an empty match"); +} From 6005ba2c42f7e66762206d2774c3da5bde08923b Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Fri, 10 Jul 2026 15:43:03 +0100 Subject: [PATCH 2/6] =?UTF-8?q?feat(engine):=20field-level=20updates=20?= =?UTF-8?q?=E2=80=94=20partial-schema=20matched-only=20staging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A node table whose only op in a query is a single `update` now stages a PARTIAL merge source — (id + assigned + constraint-completion columns) — as a matched-only merge (WhenNotMatched::DoNothing). Lance's partial-schema path patches the provided columns in place on the same fragment, so: - the update never reads unassigned columns (the scan projects id + completion-minus-assigned; assignments are literal values) - write amplification drops from row-width to assigned-width - every index over unassigned columns keeps its fragment coverage (Lance prunes only fields_modified) Completion columns: every member of a @unique group intersecting the assigned set rides along, so composite-unique validation sees the whole tuple; the unique evaluator skips groups with no column present in a batch (untouched by the write) and stays loud on partially-present groups. Fallbacks (whole-row, behavior unchanged): mixed insert+update on one table, multiple updates on one table, and any non-eligible shape — partial and full batches cannot share one uniform-schema merge source, and a present column's null cell means "set NULL", so widening is unsound. Known residual (pinned as a tripwire in scalar_indexes): the merge source must carry the join key, and Lance's column patcher counts every source column as modified — so the id BTREE alone still loses patched fragments (parity with the whole-row path) until upstream excludes ON columns from column patches; every other index is preserved. --- crates/omnigraph/src/exec/mutation.rs | 142 +++++++++++++++++++++-- crates/omnigraph/src/exec/staging.rs | 61 +++++++++- crates/omnigraph/src/validate.rs | 12 ++ crates/omnigraph/tests/scalar_indexes.rs | 17 ++- 4 files changed, 220 insertions(+), 12 deletions(-) diff --git a/crates/omnigraph/src/exec/mutation.rs b/crates/omnigraph/src/exec/mutation.rs index c56897a8..1d706e58 100644 --- a/crates/omnigraph/src/exec/mutation.rs +++ b/crates/omnigraph/src/exec/mutation.rs @@ -397,6 +397,64 @@ fn predicate_to_sql( /// want to be split into separate queries. /// - If a blob column has a string-URI assignment, build the blob array /// inline. +/// RFC-022: the column plan for a partial-schema update. +struct PartialUpdatePlan { + /// Staged-batch schema: (id + assigned + completion), in catalog order. + output_schema: SchemaRef, + /// Scan schema: (id + completion-minus-assigned), in catalog order — + /// the only columns whose OLD values the update needs. + scan_schema: SchemaRef, + /// `scan_schema`'s column names (the scan projection). + scan_cols: Vec, +} + +/// Compute the partial-update column plan: the staged source carries the merge +/// key (`id`), the assigned columns (values from literals), and the +/// constraint-completion columns — every member of a `@unique` group that +/// intersects the assigned set, so the end-of-query evaluator can validate the +/// whole tuple (an update assigning only `room` of `@unique(room, hour)` must +/// still detect a (room, hour) collision). Columns outside the plan are left +/// physically untouched by the matched-only merge, and every index over them +/// keeps its coverage (Lance prunes only `fields_modified`). +fn partial_update_plan( + node_type: &omnigraph_compiler::catalog::NodeType, + full_schema: &SchemaRef, + assignments: &[IRAssignment], +) -> PartialUpdatePlan { + let assigned: HashSet<&str> = assignments.iter().map(|a| a.property.as_str()).collect(); + let mut completion: HashSet<&str> = HashSet::new(); + for group in &node_type.unique_constraints { + if group.iter().any(|col| assigned.contains(col.as_str())) { + for col in group { + completion.insert(col.as_str()); + } + } + } + + let mut output_fields = Vec::new(); + let mut scan_fields = Vec::new(); + for field in full_schema.fields() { + let name = field.name().as_str(); + let is_assigned = assigned.contains(name); + let in_completion = completion.contains(name); + if name == "id" || is_assigned || in_completion { + output_fields.push(field.clone()); + } + if name == "id" || (in_completion && !is_assigned) { + scan_fields.push(field.clone()); + } + } + let scan_cols = scan_fields + .iter() + .map(|f| f.name().clone()) + .collect::>(); + PartialUpdatePlan { + output_schema: Arc::new(Schema::new(output_fields)), + scan_schema: Arc::new(Schema::new(scan_fields)), + scan_cols, + } +} + fn apply_assignments( full_schema: &SchemaRef, batch: &RecordBatch, @@ -916,6 +974,29 @@ impl Omnigraph { staging: &mut MutationStaging, txn: Option<&crate::db::WriteTxn>, ) -> Result { + // RFC-022 eligibility: a node table whose ONLY op in this query is a + // single `update` stages a partial-schema matched-only merge (key + + // assigned + constraint-completion columns). Every other combination + // falls back to whole-row staging: partial and full batches cannot + // share one uniform-schema merge source (a present column's null cell + // means "set NULL", so widening a partial batch would null-overwrite), + // a second same-table op's read-your-writes scan must see full rows, + // and one table commits at most one version per query (invariant 4). + let mut update_op_census: std::collections::HashMap<&str, (usize, usize)> = + std::collections::HashMap::new(); + for op in &ir.ops { + let (type_name, is_update) = match op { + MutationOpIR::Update { type_name, .. } => (type_name.as_str(), true), + MutationOpIR::Insert { type_name, .. } + | MutationOpIR::Delete { type_name, .. } => (type_name.as_str(), false), + }; + let entry = update_op_census.entry(type_name).or_insert((0, 0)); + entry.0 += 1; + if is_update { + entry.1 += 1; + } + } + let mut total = MutationResult::default(); for op in &ir.ops { let result = match op { @@ -931,8 +1012,12 @@ impl Omnigraph { assignments, predicate, } => { + let partial_ok = update_op_census + .get(type_name.as_str()) + .is_some_and(|&(ops, updates)| ops == 1 && updates == 1); self.execute_update( type_name, assignments, predicate, params, branch, staging, txn, + partial_ok, ) .await? } @@ -1063,6 +1148,7 @@ impl Omnigraph { branch: Option<&str>, staging: &mut MutationStaging, txn: Option<&crate::db::WriteTxn>, + partial_ok: bool, ) -> Result { // Defense in depth: ensure this is a node type if !self.catalog().node_types.contains_key(type_name) { @@ -1086,6 +1172,20 @@ impl Omnigraph { let schema = self.catalog().node_types[type_name].arrow_schema.clone(); let blob_props = self.catalog().node_types[type_name].blob_properties.clone(); + // RFC-022: when this update is the table's only op in the query, + // stage a PARTIAL source — (id + assigned + constraint-completion) + // columns — instead of whole rows. `partial_plan` is `None` on the + // whole-row fallback path. + let partial_plan = if partial_ok { + Some(partial_update_plan( + &self.catalog().node_types[type_name], + &schema, + assignments, + )) + } else { + None + }; + let table_key = format!("node:{}", type_name); let (handle, _full_path, _table_branch) = open_table_for_mutation( self, @@ -1118,8 +1218,19 @@ impl Omnigraph { .filter(|f| !blob_props.contains(f.name())) .map(|f| f.name().as_str()) .collect(); - let projection: Option<&[&str]> = - (!blob_props.is_empty()).then_some(non_blob_cols.as_slice()); + // Partial path: scan only (id + completion-minus-assigned) — assigned + // columns' old values are never needed (`.gq` assignments are literal + // values), and the WHERE predicate evaluates by pushdown without being + // projected. Whole-row path: unchanged (full schema minus blobs). + let partial_scan_cols: Vec<&str> = partial_plan + .as_ref() + .map(|plan| plan.scan_cols.iter().map(|c| c.as_str()).collect()) + .unwrap_or_default(); + let projection: Option<&[&str]> = if partial_plan.is_some() { + Some(partial_scan_cols.as_slice()) + } else { + (!blob_props.is_empty()).then_some(non_blob_cols.as_slice()) + }; let pending_batches = staging.pending_batches(&table_key); let pending_schema = staging.pending_schema(&table_key); // Use merge semantics on the union: a committed row whose `id` @@ -1155,7 +1266,15 @@ impl Omnigraph { // diverge (typically a blob-table mid-schema-shift), the helper // surfaces a clear error directing the caller to split the // mutation. - let matched = concat_match_batches_to_schema(&schema, &blob_props, batches)?; + // Partial path: normalize/concat against the SCAN schema (id + + // completion), then apply assignments over the OUTPUT schema (id + + // assigned + completion) — assigned columns come from literals, so + // they need no scanned values. Whole-row path: unchanged. + let (concat_schema, output_schema) = match partial_plan.as_ref() { + Some(plan) => (plan.scan_schema.clone(), plan.output_schema.clone()), + None => (schema.clone(), schema.clone()), + }; + let matched = concat_match_batches_to_schema(&concat_schema, &blob_props, batches)?; let affected_count = matched.num_rows(); @@ -1163,16 +1282,19 @@ impl Omnigraph { for a in assignments { resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?); } - let updated = apply_assignments(&schema, &matched, &resolved, &blob_props)?; + let updated = apply_assignments(&output_schema, &matched, &resolved, &blob_props)?; // Validation (value/enum/unique) runs end-of-query via the evaluator. - // Accumulate the updated batch into the Merge-mode pending stream. - // The accumulator may now contain entries with the same id as a - // prior insert or update on this table; `MutationStaging::finalize` - // dedupes by id (last-occurrence wins) before issuing the single - // `stage_merge_insert` call at end-of-query. + // Accumulate the updated batch. Whole-row: the Merge-mode pending + // stream (may coalesce with prior same-table ops; finalize dedupes by + // id, last wins). Partial (RFC-022): a dedicated partial-update entry + // that stages as a matched-only merge. let updated_schema = updated.schema(); - staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?; + if partial_plan.is_some() { + staging.append_partial_update_batch(&table_key, updated_schema, updated)?; + } else { + staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?; + } Ok(MutationResult { affected_nodes: affected_count, diff --git a/crates/omnigraph/src/exec/staging.rs b/crates/omnigraph/src/exec/staging.rs index 2f865f37..08ef5085 100644 --- a/crates/omnigraph/src/exec/staging.rs +++ b/crates/omnigraph/src/exec/staging.rs @@ -52,6 +52,15 @@ pub(crate) struct PendingTable { pub(crate) schema: SchemaRef, pub(crate) mode: PendingMode, pub(crate) batches: Vec, + /// RFC-022 field-level update: the accumulated batches carry a PARTIAL + /// schema — (merge key + assigned + constraint-completion) columns only — + /// and stage as a matched-only merge (`WhenNotMatched::DoNothing`), so + /// Lance patches the assigned columns in place and leaves every other + /// column (and every index over them) untouched. Set only by + /// [`MutationStaging::append_partial_update_batch`], whose eligibility + /// rule (the table's ONLY op in this query is a single `update`) + /// guarantees no full-schema batch ever lands on the same table. + pub(crate) partial_update: bool, } impl PendingTable { @@ -60,6 +69,7 @@ impl PendingTable { schema, mode, batches: Vec::new(), + partial_update: false, } } @@ -184,6 +194,13 @@ impl MutationStaging { // caller a clearer point of failure attached to the specific // op that introduced the drift. if let Some(existing) = self.pending.get(table_key) { + if existing.partial_update { + return Err(OmniError::manifest_internal(format!( + "append_batch: table '{}' holds a partial-update batch — the \ + eligibility rule should have prevented a second op on this table", + table_key + ))); + } if existing.mode == PendingMode::Overwrite || mode == PendingMode::Overwrite { if existing.mode != mode { return Err(OmniError::manifest_internal(format!( @@ -221,6 +238,37 @@ impl MutationStaging { Ok(()) } + /// RFC-022: stage a PARTIAL-schema update batch — (key + assigned + + /// completion) columns — for a table whose only op in this query is one + /// `update`. The caller (`execute_update`) enforces that eligibility rule + /// from the lowered IR, so an existing accumulator entry here is an + /// internal invariant breach, not a user error: partial and full batches + /// cannot share one uniform-schema merge source, and silently widening + /// (or narrowing) would corrupt the staged shape. + pub(crate) fn append_partial_update_batch( + &mut self, + table_key: &str, + schema: SchemaRef, + batch: RecordBatch, + ) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + if self.pending.contains_key(table_key) { + return Err(OmniError::manifest_internal(format!( + "append_partial_update_batch: table '{}' already has pending batches — \ + partial-update staging requires the update to be the table's only op \ + (the eligibility rule in execute_named_mutation)", + table_key + ))); + } + let mut entry = PendingTable::new(schema, PendingMode::Merge); + entry.partial_update = true; + entry.batches.push(batch); + self.pending.insert(table_key.to_string(), entry); + Ok(()) + } + /// Record a delete predicate for `table_key`. The caller must have already /// called `ensure_path` (via `open_table_for_mutation`) so the table's /// path/version/op-kind are captured. D₂ guarantees a delete-touched table @@ -481,13 +529,24 @@ async fn stage_pending_table( let staged = match table.mode { PendingMode::Append => db.storage().stage_append(&ds, combined, &[]).await?, PendingMode::Merge => { + // RFC-022: a partial-update table stages matched-only. Its source + // carries a subset schema (key + assigned + completion), so Lance + // patches the provided columns in place and never inserts — the + // subset schema could not satisfy non-null target columns anyway + // (Lance rejects partial-source inserts), and the update executor + // only stages rows it matched against the committed table. + let when_not_matched = if table.partial_update { + lance::dataset::WhenNotMatched::DoNothing + } else { + lance::dataset::WhenNotMatched::InsertAll + }; db.storage() .stage_merge_insert( ds.clone(), combined, vec!["id".to_string()], lance::dataset::WhenMatched::UpdateAll, - lance::dataset::WhenNotMatched::InsertAll, + when_not_matched, ) .await? } diff --git a/crates/omnigraph/src/validate.rs b/crates/omnigraph/src/validate.rs index 5d750c8d..397fee31 100644 --- a/crates/omnigraph/src/validate.rs +++ b/crates/omnigraph/src/validate.rs @@ -671,6 +671,18 @@ async fn evaluate_unique( // became null removes the id (it no longer holds a unique key). let mut final_by_id: HashMap, Vec)> = HashMap::new(); for batch in change.value_batches() { + // RFC-022 partial-update batches carry only (id + assigned + + // completion) columns. A group with NO column in the batch is + // untouched by the write — its committed values still satisfy the + // constraint — so skip it. A PARTIALLY present group would mean the + // completion-column projection failed; keep that loud (the error + // below fires on the first missing member). + if columns + .iter() + .all(|name| batch.column_by_name(name).is_none()) + { + continue; + } let group_columns = columns .iter() .map(|name| { diff --git a/crates/omnigraph/tests/scalar_indexes.rs b/crates/omnigraph/tests/scalar_indexes.rs index b76b985e..9d643336 100644 --- a/crates/omnigraph/tests/scalar_indexes.rs +++ b/crates/omnigraph/tests/scalar_indexes.rs @@ -106,7 +106,7 @@ async fn update_of_unindexed_property_preserves_other_index_coverage() { let snap = snapshot_main(&db).await.unwrap(); let ds = snap.open("node:Item").await.unwrap(); - for col in ["id", "status", "published", "rank"] { + for col in ["status", "published", "rank"] { let cov = TableStore::key_column_index_coverage(&ds, col).await.unwrap(); assert_eq!( cov, @@ -114,4 +114,19 @@ async fn update_of_unindexed_property_preserves_other_index_coverage() { "updating un-indexed 'note' must not degrade the index on '{col}', got {cov:?}" ); } + + // Known residual (tripwire, not the goal state): the partial merge source + // must carry the join key (`id`) for matching, and Lance's column patcher + // counts every source column as modified — including the ON column whose + // values are equal by definition of the match — so the id BTREE alone + // loses the patched fragment until the reconciler folds it back. Parity + // with the whole-row path for `id`, strictly better for every other index. + // When upstream excludes ON columns from column patches this assertion + // goes red — flip it to `Indexed` and delete this comment. + let id_cov = TableStore::key_column_index_coverage(&ds, "id").await.unwrap(); + assert!( + matches!(id_cov, IndexCoverage::Degraded { .. }), + "id-BTREE currently loses patched fragments (join key rides the source); \ + if this is now Indexed, upstream fixed ON-column patching — tighten this test" + ); } From cb5125d523b51e161df0835cdab78a4a5226f7cf Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Fri, 10 Jul 2026 15:54:37 +0100 Subject: [PATCH 3/6] =?UTF-8?q?docs+guards:=20field-level=20updates=20?= =?UTF-8?q?=E2=80=94=20partial-merge=20surface=20guard,=20mutation/write?= =?UTF-8?q?=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit lance_surface_guards gains partial_schema_merge_patches_in_place_and_prunes_only_modified_fields: a partial-source matched-only merge patches columns in place (same fragment set), leaves missing columns untouched, and prunes only indexes covering a source column — with the join-key index pruning pinned as the current residual (goes red when upstream excludes ON columns from column patches; tighten scalar_indexes.rs then). docs/user/mutations: updates write only assigned columns (cost follows the assignment, not row width); @embed staleness note. docs/dev/writes: partial-schema staging shape, completion-column rule, fallback matrix, evaluator group-skip, the id-BTREE residual. --- .../omnigraph/tests/lance_surface_guards.rs | 146 ++++++++++++++++++ docs/dev/writes.md | 21 +++ docs/user/mutations/index.md | 15 ++ 3 files changed, 182 insertions(+) diff --git a/crates/omnigraph/tests/lance_surface_guards.rs b/crates/omnigraph/tests/lance_surface_guards.rs index 2abe1e23..c6c5c5b3 100644 --- a/crates/omnigraph/tests/lance_surface_guards.rs +++ b/crates/omnigraph/tests/lance_surface_guards.rs @@ -1257,3 +1257,149 @@ async fn filtered_scan_tolerates_merge_update_row_id_overlap() { assert_eq!(rows, expected, "filtered read for {slug}"); } } + +// --- Guard: partial-schema merge patches columns in place, pruning only +// --- fields_modified (the RFC-022 substrate contract) ----------------------- +// +// A merge-insert whose source carries a SUBSET of the target schema must: +// (a) patch the provided columns for matched rows IN PLACE — same fragment +// id, no new fragments (`update_mode: RewriteColumns`); +// (b) leave missing columns' values untouched; +// (c) prune ONLY indexes covering a source column from the touched +// fragment's bitmap (`prune_updated_fields_from_indices` keyed on +// `fields_modified`) — an index on an untouched column keeps coverage. +// Current residual pinned here: the join key rides the source, so the KEY +// column's index is pruned too even though matched keys are equal by +// definition. If (d) below goes red, upstream started excluding ON columns +// from column patches — tighten omnigraph's coverage test in +// scalar_indexes.rs and drop this arm. +#[tokio::test] +async fn partial_schema_merge_patches_in_place_and_prunes_only_modified_fields() { + use arrow_array::Int64Array; + use futures::TryStreamExt; + + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().join("guard_partial_merge.lance"); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["x", "y", "z"])), + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let params = WriteParams { + mode: WriteMode::Create, + enable_stable_row_ids: true, + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }; + let mut ds = Dataset::write(reader, uri.to_str().unwrap(), Some(params)) + .await + .unwrap(); + ds.create_index_builder(&["id"], IndexType::BTree, &ScalarIndexParams::default()) + .await + .unwrap(); + ds.create_index_builder(&["b"], IndexType::BTree, &ScalarIndexParams::default()) + .await + .unwrap(); + ds.checkout_latest().await.unwrap(); + let frags_before: Vec = ds.fragments().iter().map(|f| f.id).collect(); + + // Partial source: (id, a) — patch a=200 for id="y". No `b` column. + let partial_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("a", DataType::Int64, false), + ])); + let patch = RecordBatch::try_new( + partial_schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["y"])), + Arc::new(Int64Array::from(vec![200i64])), + ], + ) + .unwrap(); + let job = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap(); + let source = RecordBatchIterator::new(vec![Ok(patch)], partial_schema); + let (ds, _stats) = job.execute_reader(source).await.unwrap(); + + // (a) in place: same fragment set. + let frags_after: Vec = ds.fragments().iter().map(|f| f.id).collect(); + assert_eq!(frags_before, frags_after, "partial merge must patch in place"); + + // (b) patched + retained values. + let rows = ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let all = arrow_select::concat::concat_batches(&rows[0].schema(), &rows).unwrap(); + let ids = all + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let a = all + .column_by_name("a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let b = all + .column_by_name("b") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..all.num_rows() { + match ids.value(i) { + "y" => { + assert_eq!(a.value(i), 200, "assigned column patched"); + assert_eq!(b.value(i), 20, "missing column retained"); + } + "x" => assert_eq!((a.value(i), b.value(i)), (1, 10)), + "z" => assert_eq!((a.value(i), b.value(i)), (3, 30)), + other => panic!("unexpected id {other}"), + } + } + + // (c) index on untouched column `b` keeps its fragment coverage. + let indices = ds.load_indices().await.unwrap(); + let bitmap_of = |col_field: &str| { + let field_id = ds.schema().field(col_field).unwrap().id; + indices + .iter() + .find(|i| i.fields == vec![field_id]) + .and_then(|i| i.fragment_bitmap.as_ref()) + .map(|bm| bm.iter().collect::>()) + }; + assert_eq!( + bitmap_of("b"), + Some(frags_before.iter().map(|f| *f as u32).collect::>()), + "index on an untouched column must keep the patched fragment" + ); + // (d) residual: the join key's index IS pruned (source carries `id`). + assert_eq!( + bitmap_of("id"), + Some(vec![]), + "join-key index currently loses the patched fragment; if this went red \ + with a non-empty bitmap, upstream now excludes ON columns from patches \ + — tighten scalar_indexes.rs and update RFC-022's residual note" + ); +} diff --git a/docs/dev/writes.md b/docs/dev/writes.md index e22e85b1..8899a18b 100644 --- a/docs/dev/writes.md +++ b/docs/dev/writes.md @@ -64,6 +64,27 @@ shared by both `mutate_as` and the bulk loader: (below) still prevents inserts/updates from coexisting with deletes in one query. +- **Field-level updates (partial-schema staging).** A node table whose ONLY op + in the query is a single `update` stages a PARTIAL merge source — (id + + assigned + constraint-completion columns, where completion = every member of + a `@unique` group intersecting the assigned set) — as a **matched-only** + merge (`WhenNotMatched::DoNothing`). Lance's partial-schema path patches the + provided columns in place on the same fragment (`update_mode: + RewriteColumns`), so unassigned columns are never read or rewritten and + indexes over them keep fragment coverage (Lance prunes only + `fields_modified`). Every other shape (insert+update on one table, multiple + updates, etc.) falls back to whole-row staging: partial and full batches + cannot share one uniform-schema merge source (a present column's null cell + means "set NULL", so widening would null-overwrite), a later same-table op's + read-your-writes scan needs full rows, and one table commits at most one + version per query. The unique evaluator skips a `@unique` group with no + column present in a batch (untouched by the write) and stays loud on a + partially-present group. Known residual (pinned by + `lance_surface_guards::partial_schema_merge_patches_in_place_and_prunes_only_modified_fields` + and the `scalar_indexes` coverage cell): the join key rides the source, so + the **id BTREE** alone still loses patched fragments (parity with the + whole-row path) until upstream excludes ON columns from column patches. + This upholds the manifest-atomic mutation and read-your-writes invariants tracked in [docs/dev/invariants.md](invariants.md). diff --git a/docs/user/mutations/index.md b/docs/user/mutations/index.md index 736d3449..4eb4ac29 100644 --- a/docs/user/mutations/index.md +++ b/docs/user/mutations/index.md @@ -29,6 +29,21 @@ failure leaves the graph untouched. See [transactions](../branching/transactions for the per-query atomicity contract and [branches](../branching/index.md) for multi-query workflows. +## Updates write only the columns they assign + +When an `update` is the only statement touching its type in a query, the engine +stages just the assigned columns (plus the row key and any columns needed to +complete a `@unique` group the assignment touches) and patches them in place. +Unassigned columns — including large `Vector` embedding columns — are neither +read nor rewritten, and indexes over them keep their coverage. Semantics are +identical to a whole-row update; the difference is cost: update latency and +write volume scale with what you assign, not with row width. + +Queries that combine an update with other statements on the same type fall back +to whole-row staging (same semantics, previous cost). Note that updating a +property that feeds an `@embed` column does **not** recompute the embedding — +re-run the embedding pass after changing embed-source text. + ## Inserts/updates and deletes cannot mix in one query A single change query must be **either insert/update-only or delete-only**. From aee12416d497654cb17a9a85f31151b063b85b10 Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Fri, 10 Jul 2026 17:26:28 +0100 Subject: [PATCH 4/6] =?UTF-8?q?fix(engine):=20completion=20columns=20ride?= =?UTF-8?q?=20validation=20only=20=E2=80=94=20staged=20source=20is=20(id?= =?UTF-8?q?=20+=20assigned)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review finding: @unique-group completion columns were included in the STAGED merge source, so Lance patched them (every source column counts as modified) and pruned their indexes even though their values were unchanged — e.g. updating only `room` of @unique(room, hour) degraded the index on `hour`. Completion columns are validation inputs, never merge inputs: the pending (validation) batch keeps (id + assigned + completion) so the evaluator sees the whole unique tuple, and the staged source is projected down to (id + assigned) before stage_merge_insert. PendingTable.partial_update becomes partial_stage_cols: Option> carrying the projection. Regression test (red first with 'hour' Degraded): scalar_indexes::completion_column_index_survives_partial_update. --- crates/omnigraph/src/exec/mutation.rs | 21 ++++++-- crates/omnigraph/src/exec/staging.rs | 62 ++++++++++++++++-------- crates/omnigraph/tests/scalar_indexes.rs | 50 +++++++++++++++++++ docs/dev/writes.md | 13 +++-- 4 files changed, 120 insertions(+), 26 deletions(-) diff --git a/crates/omnigraph/src/exec/mutation.rs b/crates/omnigraph/src/exec/mutation.rs index 1d706e58..df6feff8 100644 --- a/crates/omnigraph/src/exec/mutation.rs +++ b/crates/omnigraph/src/exec/mutation.rs @@ -399,8 +399,13 @@ fn predicate_to_sql( /// inline. /// RFC-022: the column plan for a partial-schema update. struct PartialUpdatePlan { - /// Staged-batch schema: (id + assigned + completion), in catalog order. + /// Accumulated-batch schema: (id + assigned + completion), in catalog + /// order — the validation change-set's view. output_schema: SchemaRef, + /// Columns actually STAGED to the merge: (id + assigned), in catalog + /// order. Completion columns are validation-only — staging them would + /// patch unassigned columns and prune their indexes. + stage_cols: Vec, /// Scan schema: (id + completion-minus-assigned), in catalog order — /// the only columns whose OLD values the update needs. scan_schema: SchemaRef, @@ -433,6 +438,7 @@ fn partial_update_plan( let mut output_fields = Vec::new(); let mut scan_fields = Vec::new(); + let mut stage_cols = Vec::new(); for field in full_schema.fields() { let name = field.name().as_str(); let is_assigned = assigned.contains(name); @@ -440,6 +446,9 @@ fn partial_update_plan( if name == "id" || is_assigned || in_completion { output_fields.push(field.clone()); } + if name == "id" || is_assigned { + stage_cols.push(name.to_string()); + } if name == "id" || (in_completion && !is_assigned) { scan_fields.push(field.clone()); } @@ -450,6 +459,7 @@ fn partial_update_plan( .collect::>(); PartialUpdatePlan { output_schema: Arc::new(Schema::new(output_fields)), + stage_cols, scan_schema: Arc::new(Schema::new(scan_fields)), scan_cols, } @@ -1290,8 +1300,13 @@ impl Omnigraph { // id, last wins). Partial (RFC-022): a dedicated partial-update entry // that stages as a matched-only merge. let updated_schema = updated.schema(); - if partial_plan.is_some() { - staging.append_partial_update_batch(&table_key, updated_schema, updated)?; + if let Some(plan) = partial_plan.as_ref() { + staging.append_partial_update_batch( + &table_key, + updated_schema, + updated, + plan.stage_cols.clone(), + )?; } else { staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?; } diff --git a/crates/omnigraph/src/exec/staging.rs b/crates/omnigraph/src/exec/staging.rs index 08ef5085..2b8704f9 100644 --- a/crates/omnigraph/src/exec/staging.rs +++ b/crates/omnigraph/src/exec/staging.rs @@ -52,15 +52,18 @@ pub(crate) struct PendingTable { pub(crate) schema: SchemaRef, pub(crate) mode: PendingMode, pub(crate) batches: Vec, - /// RFC-022 field-level update: the accumulated batches carry a PARTIAL - /// schema — (merge key + assigned + constraint-completion) columns only — - /// and stage as a matched-only merge (`WhenNotMatched::DoNothing`), so - /// Lance patches the assigned columns in place and leaves every other - /// column (and every index over them) untouched. Set only by + /// RFC-022 field-level update: `Some(stage_cols)` marks the accumulated + /// batches as PARTIAL — they carry (merge key + assigned + constraint- + /// completion) columns for the validation change-set, but the STAGED merge + /// source is projected down to `stage_cols` = (merge key + assigned) and + /// staged matched-only (`WhenNotMatched::DoNothing`). Completion columns + /// are validation inputs, never merge inputs: Lance counts every source + /// column as modified, so staging an unassigned `@unique`-group member + /// would patch it and prune its index for no semantic reason. Set only by /// [`MutationStaging::append_partial_update_batch`], whose eligibility /// rule (the table's ONLY op in this query is a single `update`) /// guarantees no full-schema batch ever lands on the same table. - pub(crate) partial_update: bool, + pub(crate) partial_stage_cols: Option>, } impl PendingTable { @@ -69,7 +72,7 @@ impl PendingTable { schema, mode, batches: Vec::new(), - partial_update: false, + partial_stage_cols: None, } } @@ -194,7 +197,7 @@ impl MutationStaging { // caller a clearer point of failure attached to the specific // op that introduced the drift. if let Some(existing) = self.pending.get(table_key) { - if existing.partial_update { + if existing.partial_stage_cols.is_some() { return Err(OmniError::manifest_internal(format!( "append_batch: table '{}' holds a partial-update batch — the \ eligibility rule should have prevented a second op on this table", @@ -250,6 +253,7 @@ impl MutationStaging { table_key: &str, schema: SchemaRef, batch: RecordBatch, + stage_cols: Vec, ) -> Result<()> { if batch.num_rows() == 0 { return Ok(()); @@ -263,7 +267,7 @@ impl MutationStaging { ))); } let mut entry = PendingTable::new(schema, PendingMode::Merge); - entry.partial_update = true; + entry.partial_stage_cols = Some(stage_cols); entry.batches.push(batch); self.pending.insert(table_key.to_string(), entry); Ok(()) @@ -529,16 +533,36 @@ async fn stage_pending_table( let staged = match table.mode { PendingMode::Append => db.storage().stage_append(&ds, combined, &[]).await?, PendingMode::Merge => { - // RFC-022: a partial-update table stages matched-only. Its source - // carries a subset schema (key + assigned + completion), so Lance - // patches the provided columns in place and never inserts — the - // subset schema could not satisfy non-null target columns anyway - // (Lance rejects partial-source inserts), and the update executor - // only stages rows it matched against the committed table. - let when_not_matched = if table.partial_update { - lance::dataset::WhenNotMatched::DoNothing - } else { - lance::dataset::WhenNotMatched::InsertAll + // RFC-022: a partial-update table stages matched-only, and the + // staged source is projected down to (key + assigned) — completion + // columns were carried for the validation change-set only; staging + // them would patch unassigned columns and prune their indexes. + // Lance patches the provided columns in place and never inserts — + // the subset schema could not satisfy non-null target columns + // anyway (Lance rejects partial-source inserts), and the update + // executor only stages rows it matched against the committed table. + let (combined, when_not_matched) = match &table.partial_stage_cols { + Some(stage_cols) => { + let indices = stage_cols + .iter() + .map(|c| { + combined.schema().index_of(c).map_err(|e| { + OmniError::manifest_internal(format!( + "partial-update stage column '{}' missing from \ + accumulated batch: {}", + c, e + )) + }) + }) + .collect::>>()?; + ( + combined + .project(&indices) + .map_err(|e| OmniError::Lance(e.to_string()))?, + lance::dataset::WhenNotMatched::DoNothing, + ) + } + None => (combined, lance::dataset::WhenNotMatched::InsertAll), }; db.storage() .stage_merge_insert( diff --git a/crates/omnigraph/tests/scalar_indexes.rs b/crates/omnigraph/tests/scalar_indexes.rs index 9d643336..1382b708 100644 --- a/crates/omnigraph/tests/scalar_indexes.rs +++ b/crates/omnigraph/tests/scalar_indexes.rs @@ -130,3 +130,53 @@ async fn update_of_unindexed_property_preserves_other_index_coverage() { if this is now Indexed, upstream fixed ON-column patching — tighten this test" ); } + +const COMPOSITE_SCHEMA: &str = r#" +node Slot { + name: String @key + room: String? + hour: I32 @index + @unique(room, hour) +} +"#; + +const COMPOSITE_DATA: &str = r#"{"type":"Slot","data":{"name":"s1","room":"roomA","hour":9}} +{"type":"Slot","data":{"name":"s2","room":"roomB","hour":9}}"#; + +const MOVE_SLOT: &str = r#" +query move_slot($name: String, $room: String) { + update Slot set { room: $room } where name = $name +} +"#; + +// A `@unique`-group completion column is a VALIDATION input, not a merge +// input: an update assigning only `room` must not patch `hour` (its value is +// unchanged), so `hour`'s index keeps the fragment. The staged merge source is +// (id + assigned); the completion column rides only the validation change-set. +#[tokio::test] +async fn completion_column_index_survives_partial_update() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut db = Omnigraph::init(uri, COMPOSITE_SCHEMA).await.unwrap(); + load_jsonl(&mut db, COMPOSITE_DATA, LoadMode::Overwrite) + .await + .unwrap(); + + db.mutate( + "main", + MOVE_SLOT, + "move_slot", + ¶ms(&[("$name", "s2"), ("$room", "roomC")]), + ) + .await + .unwrap(); + + let snap = snapshot_main(&db).await.unwrap(); + let ds = snap.open("node:Slot").await.unwrap(); + let cov = TableStore::key_column_index_coverage(&ds, "hour").await.unwrap(); + assert_eq!( + cov, + IndexCoverage::Indexed, + "completion column 'hour' was not assigned — its index must keep coverage, got {cov:?}" + ); +} diff --git a/docs/dev/writes.md b/docs/dev/writes.md index 8899a18b..f34679d3 100644 --- a/docs/dev/writes.md +++ b/docs/dev/writes.md @@ -65,10 +65,15 @@ shared by both `mutate_as` and the bulk loader: query. - **Field-level updates (partial-schema staging).** A node table whose ONLY op - in the query is a single `update` stages a PARTIAL merge source — (id + - assigned + constraint-completion columns, where completion = every member of - a `@unique` group intersecting the assigned set) — as a **matched-only** - merge (`WhenNotMatched::DoNothing`). Lance's partial-schema path patches the + in the query is a single `update` stages a PARTIAL merge source — exactly + (id + assigned) columns — as a **matched-only** merge + (`WhenNotMatched::DoNothing`). The accumulated (pending) batch additionally + carries the constraint-completion columns (every member of a `@unique` group + intersecting the assigned set) so the end-of-query evaluator can validate the + whole tuple — but those are **validation inputs only** and are projected out + before staging: Lance counts every source column as modified, so staging an + unassigned group member would patch it and prune its index for no semantic + reason. Lance's partial-schema path patches the provided columns in place on the same fragment (`update_mode: RewriteColumns`), so unassigned columns are never read or rewritten and indexes over them keep fragment coverage (Lance prunes only From c9549d1eedf0f5a9d5a58b5ffb07c62d3eba2ec6 Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Sat, 11 Jul 2026 18:40:37 +0100 Subject: [PATCH 5/6] =?UTF-8?q?fix(engine):=20address=20review=20=E2=80=94?= =?UTF-8?q?=20transitive=20completion,=20blob-aware=20plan,=20partial=20mo?= =?UTF-8?q?de?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness (both reproduced red-first): - Constraint completion is now the FIXED-POINT CLOSURE over every @unique group AND the @key group (the evaluator registers @key as a unique constraint): overlapping groups chain, and the evaluator hard-errors on a partially present group, so a legal sole update assigning one column of @unique(room, hour) + @unique(hour, day) no longer fails with 'missing unique column'. Regression: validators::overlapping_unique_groups_allow_sole_partial_update. - The partial plan is blob-aware: the staged source is derived by SUBTRACTION (batch schema minus validation-only columns) so apply-time omissions cannot desync it — assigning null to an optional Blob is again the historical no-op instead of an internal error — and Blob columns are excluded from completion/scan projection. Regression: writes::null_blob_assignment_in_sole_update_is_a_no_op. - @key assignment rejection now checks EVERY key column (composite keys), not just the first. - Parser: @unique rejects Blob and Vector columns at compile time, mirroring the @key/@index guards (+ two parser tests). Robustness/efficiency/hygiene from the same review: - PendingMode gains a PartialUpdate variant so every match on the mode must disposition partial batches explicitly (silently widening one to full-row would null-overwrite unassigned columns). - Merge-shape probe: column names materialize inside the try_with closure (zero production allocations, the record_open pattern). - Partial eligibility is an inline count at the Update arm; the per-query census map and the discarded non_blob_cols computation are gone. - Fallback-shape assertions folded into the existing owner tests (mixed-insert-update coalesce, chained-updates) instead of duplicates. --- .../omnigraph-compiler/src/schema/parser.rs | 19 +- .../src/schema/parser_tests.rs | 34 +++ crates/omnigraph/src/exec/mutation.rs | 165 ++++++++------ crates/omnigraph/src/exec/staging.rs | 116 ++++++---- crates/omnigraph/src/instrumentation.rs | 13 +- crates/omnigraph/src/table_store.rs | 12 +- crates/omnigraph/tests/validators.rs | 69 ++++++ crates/omnigraph/tests/writes.rs | 202 +++++++++--------- 8 files changed, 403 insertions(+), 227 deletions(-) diff --git a/crates/omnigraph-compiler/src/schema/parser.rs b/crates/omnigraph-compiler/src/schema/parser.rs index 6e34e53c..78d5c91b 100644 --- a/crates/omnigraph-compiler/src/schema/parser.rs +++ b/crates/omnigraph-compiler/src/schema/parser.rs @@ -947,11 +947,28 @@ fn validate_type_constraints( if is_edge && (col == "src" || col == "dst") { continue; } - if !prop_names.contains_key(col.as_str()) { + let Some(prop) = prop_names.get(col.as_str()) else { return Err(CompilerError::Parse(format!( "@unique on {} references unknown property '{}'", type_name, col ))); + }; + // Uniqueness needs a type that reduces to a scalar key: a + // Blob cannot be compared or scanned by projection and a + // Vector has no meaningful equality — mirror the @key + // guards so the class fails at compile time, not at + // validation/scan time. + if matches!(prop.prop_type.scalar, ScalarType::Blob) { + return Err(CompilerError::Parse(format!( + "@unique is not supported on blob property {}.{}", + type_name, col + ))); + } + if matches!(prop.prop_type.scalar, ScalarType::Vector(_)) { + return Err(CompilerError::Parse(format!( + "@unique is not supported on vector property {}.{}", + type_name, col + ))); } } } diff --git a/crates/omnigraph-compiler/src/schema/parser_tests.rs b/crates/omnigraph-compiler/src/schema/parser_tests.rs index 9a2e1ba1..b18c4025 100644 --- a/crates/omnigraph-compiler/src/schema/parser_tests.rs +++ b/crates/omnigraph-compiler/src/schema/parser_tests.rs @@ -1013,3 +1013,37 @@ fn test_parse_error_diagnostic_has_span() { let err = parse_schema_diagnostic(input).unwrap_err(); assert!(err.span.is_some()); } + +#[test] +fn test_reject_unique_on_blob_property() { + let input = r#" +node Doc { +name: String @key +data: Blob? +@unique(name, data) +} +"#; + let err = parse_schema(input).unwrap_err(); + assert!( + err.to_string() + .contains("@unique is not supported on blob property Doc.data"), + "got: {err}" + ); +} + +#[test] +fn test_reject_unique_on_vector_property() { + let input = r#" +node Doc { +name: String @key +embedding: Vector(4) +@unique(embedding) +} +"#; + let err = parse_schema(input).unwrap_err(); + assert!( + err.to_string() + .contains("@unique is not supported on vector property Doc.embedding"), + "got: {err}" + ); +} diff --git a/crates/omnigraph/src/exec/mutation.rs b/crates/omnigraph/src/exec/mutation.rs index df6feff8..80f435b6 100644 --- a/crates/omnigraph/src/exec/mutation.rs +++ b/crates/omnigraph/src/exec/mutation.rs @@ -402,10 +402,14 @@ struct PartialUpdatePlan { /// Accumulated-batch schema: (id + assigned + completion), in catalog /// order — the validation change-set's view. output_schema: SchemaRef, - /// Columns actually STAGED to the merge: (id + assigned), in catalog - /// order. Completion columns are validation-only — staging them would - /// patch unassigned columns and prune their indexes. - stage_cols: Vec, + /// Completion columns that ride ONLY the validation change-set and are + /// projected OUT of the staged merge source (staging them would patch + /// unassigned columns and prune their indexes). The staged set is derived + /// at stage time as (batch schema − validation_only), so it stays + /// consistent with `apply_assignments`' own omissions (e.g. a Blob + /// assigned a non-String value is omitted from the batch and therefore + /// never listed for staging). + validation_only: Vec, /// Scan schema: (id + completion-minus-assigned), in catalog order — /// the only columns whose OLD values the update needs. scan_schema: SchemaRef, @@ -414,43 +418,78 @@ struct PartialUpdatePlan { } /// Compute the partial-update column plan: the staged source carries the merge -/// key (`id`), the assigned columns (values from literals), and the -/// constraint-completion columns — every member of a `@unique` group that -/// intersects the assigned set, so the end-of-query evaluator can validate the -/// whole tuple (an update assigning only `room` of `@unique(room, hour)` must -/// still detect a (room, hour) collision). Columns outside the plan are left -/// physically untouched by the matched-only merge, and every index over them -/// keeps its coverage (Lance prunes only `fields_modified`). +/// key (`id`) and the assigned columns; the validation change-set additionally +/// carries the constraint-completion columns so the end-of-query evaluator can +/// validate whole tuples (an update assigning only `room` of +/// `@unique(room, hour)` must still detect a (room, hour) collision). +/// +/// Completion is the FIXED-POINT CLOSURE over every `@unique` group AND the +/// `@key` group (the evaluator registers `@key` as a unique constraint too): +/// overlapping groups chain — assigning `room` pulls in `hour` via +/// `(room, hour)`, which pulls in `day` via `(hour, day)` — and the evaluator +/// hard-errors on a partially present group, so any group intersecting the +/// needed set must be completed in full. Blob columns are excluded from +/// completion/scan end-to-end (they cannot be projected through the scanner; +/// the schema layer also rejects `@unique` on Blob/Vector). +/// +/// Columns outside the plan are left physically untouched by the matched-only +/// merge, and every index over them keeps its coverage (Lance prunes only +/// `fields_modified`). fn partial_update_plan( node_type: &omnigraph_compiler::catalog::NodeType, full_schema: &SchemaRef, assignments: &[IRAssignment], ) -> PartialUpdatePlan { let assigned: HashSet<&str> = assignments.iter().map(|a| a.property.as_str()).collect(); - let mut completion: HashSet<&str> = HashSet::new(); - for group in &node_type.unique_constraints { - if group.iter().any(|col| assigned.contains(col.as_str())) { - for col in group { - completion.insert(col.as_str()); + + let mut groups: Vec<&[String]> = node_type + .unique_constraints + .iter() + .map(|g| g.as_slice()) + .collect(); + if let Some(key) = &node_type.key { + groups.push(key.as_slice()); + } + let mut needed: HashSet<&str> = assigned.clone(); + loop { + let mut grew = false; + for group in &groups { + if group.iter().any(|c| needed.contains(c.as_str())) + && !group.iter().all(|c| needed.contains(c.as_str())) + { + for c in group.iter() { + needed.insert(c.as_str()); + } + grew = true; } } + if !grew { + break; + } } let mut output_fields = Vec::new(); let mut scan_fields = Vec::new(); - let mut stage_cols = Vec::new(); + let mut validation_only = Vec::new(); for field in full_schema.fields() { let name = field.name().as_str(); let is_assigned = assigned.contains(name); - let in_completion = completion.contains(name); - if name == "id" || is_assigned || in_completion { - output_fields.push(field.clone()); + let is_completion = !is_assigned && name != "id" && needed.contains(name); + let is_blob = node_type.blob_properties.contains(name); + if is_completion && is_blob { + // Defense in depth for pre-guard schemas: a Blob can neither be + // scanned by projection nor compared by the evaluator's key + // encoding; the parser now rejects @unique on Blob outright. + continue; } - if name == "id" || is_assigned { - stage_cols.push(name.to_string()); + if name == "id" || is_assigned || is_completion { + output_fields.push(field.clone()); } - if name == "id" || (in_completion && !is_assigned) { + if is_completion { + validation_only.push(name.to_string()); scan_fields.push(field.clone()); + } else if name == "id" { + scan_fields.insert(0, field.clone()); } } let scan_cols = scan_fields @@ -459,7 +498,7 @@ fn partial_update_plan( .collect::>(); PartialUpdatePlan { output_schema: Arc::new(Schema::new(output_fields)), - stage_cols, + validation_only, scan_schema: Arc::new(Schema::new(scan_fields)), scan_cols, } @@ -984,29 +1023,6 @@ impl Omnigraph { staging: &mut MutationStaging, txn: Option<&crate::db::WriteTxn>, ) -> Result { - // RFC-022 eligibility: a node table whose ONLY op in this query is a - // single `update` stages a partial-schema matched-only merge (key + - // assigned + constraint-completion columns). Every other combination - // falls back to whole-row staging: partial and full batches cannot - // share one uniform-schema merge source (a present column's null cell - // means "set NULL", so widening a partial batch would null-overwrite), - // a second same-table op's read-your-writes scan must see full rows, - // and one table commits at most one version per query (invariant 4). - let mut update_op_census: std::collections::HashMap<&str, (usize, usize)> = - std::collections::HashMap::new(); - for op in &ir.ops { - let (type_name, is_update) = match op { - MutationOpIR::Update { type_name, .. } => (type_name.as_str(), true), - MutationOpIR::Insert { type_name, .. } - | MutationOpIR::Delete { type_name, .. } => (type_name.as_str(), false), - }; - let entry = update_op_census.entry(type_name).or_insert((0, 0)); - entry.0 += 1; - if is_update { - entry.1 += 1; - } - } - let mut total = MutationResult::default(); for op in &ir.ops { let result = match op { @@ -1022,9 +1038,31 @@ impl Omnigraph { assignments, predicate, } => { - let partial_ok = update_op_census - .get(type_name.as_str()) - .is_some_and(|&(ops, updates)| ops == 1 && updates == 1); + // RFC-022 eligibility: partial staging only when this + // update is the table's ONLY op in the query (the count + // includes this op itself, so == 1 means "no other op + // touches the type"). Every other combination falls back + // to whole-row staging: partial and full batches cannot + // share one uniform-schema merge source (a present + // column's null cell means "set NULL", so widening a + // partial batch would null-overwrite), a second same-table + // op's read-your-writes scan must see full rows, and one + // table commits at most one version per query + // (invariant 4). Op lists are tiny; the inline count + // replaces a per-query census map. + let partial_ok = ir + .ops + .iter() + .filter(|op| { + let t = match op { + MutationOpIR::Insert { type_name, .. } + | MutationOpIR::Update { type_name, .. } + | MutationOpIR::Delete { type_name, .. } => type_name, + }; + t == type_name + }) + .count() + == 1; self.execute_update( type_name, assignments, predicate, params, branch, staging, txn, partial_ok, @@ -1168,12 +1206,13 @@ impl Omnigraph { ))); } - // Reject updates to @key properties — identity is immutable - if let Some(key_prop) = self.catalog().node_types[type_name].key_property() { - if assignments.iter().any(|a| a.property == key_prop) { + // Reject updates to @key properties — identity is immutable. Checks + // EVERY key column (composite keys included), not just the first. + if let Some(key_cols) = &self.catalog().node_types[type_name].key { + if let Some(a) = assignments.iter().find(|a| key_cols.contains(&a.property)) { return Err(OmniError::manifest(format!( "cannot update @key property '{}' — delete and re-insert instead", - key_prop + a.property ))); } } @@ -1222,12 +1261,6 @@ impl Omnigraph { // from explicit assignments and omits unassigned blobs (Lance's // merge_insert leaves them untouched). Tables without blob // columns scan the full schema unprojected. - let non_blob_cols: Vec<&str> = schema - .fields() - .iter() - .filter(|f| !blob_props.contains(f.name())) - .map(|f| f.name().as_str()) - .collect(); // Partial path: scan only (id + completion-minus-assigned) — assigned // columns' old values are never needed (`.gq` assignments are literal // values), and the WHERE predicate evaluates by pushdown without being @@ -1236,6 +1269,16 @@ impl Omnigraph { .as_ref() .map(|plan| plan.scan_cols.iter().map(|c| c.as_str()).collect()) .unwrap_or_default(); + let non_blob_cols: Vec<&str> = if partial_plan.is_none() { + schema + .fields() + .iter() + .filter(|f| !blob_props.contains(f.name())) + .map(|f| f.name().as_str()) + .collect() + } else { + Vec::new() + }; let projection: Option<&[&str]> = if partial_plan.is_some() { Some(partial_scan_cols.as_slice()) } else { @@ -1305,7 +1348,7 @@ impl Omnigraph { &table_key, updated_schema, updated, - plan.stage_cols.clone(), + plan.validation_only.clone(), )?; } else { staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?; diff --git a/crates/omnigraph/src/exec/staging.rs b/crates/omnigraph/src/exec/staging.rs index 2b8704f9..2f680e50 100644 --- a/crates/omnigraph/src/exec/staging.rs +++ b/crates/omnigraph/src/exec/staging.rs @@ -41,6 +41,14 @@ use crate::error::{OmniError, Result}; pub(crate) enum PendingMode { Append, Merge, + /// RFC-022 field-level update: a Merge whose staged source is the batch + /// schema MINUS the table's `partial_validation_only` columns, staged + /// matched-only (`WhenNotMatched::DoNothing`). A distinct variant — not a + /// side-flag on `Merge` — so every `match` on the mode is forced to + /// disposition partial batches explicitly: treating one as a full-row + /// Merge (a future coalescing change, a retry path) would null-overwrite + /// unassigned columns, and the compiler now surfaces that decision. + PartialUpdate, Overwrite, } @@ -52,18 +60,22 @@ pub(crate) struct PendingTable { pub(crate) schema: SchemaRef, pub(crate) mode: PendingMode, pub(crate) batches: Vec, - /// RFC-022 field-level update: `Some(stage_cols)` marks the accumulated - /// batches as PARTIAL — they carry (merge key + assigned + constraint- - /// completion) columns for the validation change-set, but the STAGED merge - /// source is projected down to `stage_cols` = (merge key + assigned) and - /// staged matched-only (`WhenNotMatched::DoNothing`). Completion columns - /// are validation inputs, never merge inputs: Lance counts every source - /// column as modified, so staging an unassigned `@unique`-group member - /// would patch it and prune its index for no semantic reason. Set only by + /// RFC-022 field-level update: `Some(validation_only)` marks the + /// accumulated batches as PARTIAL — they carry (merge key + assigned + + /// constraint-completion) columns for the validation change-set, and the + /// STAGED merge source is the batch schema MINUS these validation-only + /// completion columns, staged matched-only (`WhenNotMatched::DoNothing`). + /// Deriving the staged set by subtraction (rather than listing it) keeps + /// it consistent with `apply_assignments`' own omissions — e.g. a Blob + /// assigned a non-String value is omitted from the batch entirely. + /// Completion columns are validation inputs, never merge inputs: Lance + /// counts every source column as modified, so staging an unassigned + /// `@unique`-group member would patch it and prune its index for no + /// semantic reason. Set only by /// [`MutationStaging::append_partial_update_batch`], whose eligibility /// rule (the table's ONLY op in this query is a single `update`) /// guarantees no full-schema batch ever lands on the same table. - pub(crate) partial_stage_cols: Option>, + pub(crate) partial_validation_only: Option>, } impl PendingTable { @@ -72,7 +84,7 @@ impl PendingTable { schema, mode, batches: Vec::new(), - partial_stage_cols: None, + partial_validation_only: None, } } @@ -197,7 +209,7 @@ impl MutationStaging { // caller a clearer point of failure attached to the specific // op that introduced the drift. if let Some(existing) = self.pending.get(table_key) { - if existing.partial_stage_cols.is_some() { + if existing.mode == PendingMode::PartialUpdate { return Err(OmniError::manifest_internal(format!( "append_batch: table '{}' holds a partial-update batch — the \ eligibility rule should have prevented a second op on this table", @@ -253,7 +265,7 @@ impl MutationStaging { table_key: &str, schema: SchemaRef, batch: RecordBatch, - stage_cols: Vec, + validation_only: Vec, ) -> Result<()> { if batch.num_rows() == 0 { return Ok(()); @@ -266,8 +278,8 @@ impl MutationStaging { table_key ))); } - let mut entry = PendingTable::new(schema, PendingMode::Merge); - entry.partial_stage_cols = Some(stage_cols); + let mut entry = PendingTable::new(schema, PendingMode::PartialUpdate); + entry.partial_validation_only = Some(validation_only); entry.batches.push(batch); self.pending.insert(table_key.to_string(), entry); Ok(()) @@ -499,7 +511,10 @@ async fn stage_pending_table( // same SchemaRewrite policy as schema apply. let stage_kind = match table.mode { PendingMode::Append => crate::db::MutationOpKind::Insert, - PendingMode::Merge => crate::db::MutationOpKind::Merge, + // PartialUpdate reopens with the same non-strict policy as Merge; + // strictness for the commit-time drift check comes from `op_kinds` + // (the update op registered MutationOpKind::Update there). + PendingMode::Merge | PendingMode::PartialUpdate => crate::db::MutationOpKind::Merge, PendingMode::Overwrite => crate::db::MutationOpKind::SchemaRewrite, }; let ds = db @@ -517,7 +532,9 @@ async fn stage_pending_table( } let combined = match table.mode { - PendingMode::Merge => dedupe_merge_batches_by_id(&table.schema, table.batches)?, + PendingMode::Merge | PendingMode::PartialUpdate => { + dedupe_merge_batches_by_id(&table.schema, table.batches)? + } PendingMode::Append | PendingMode::Overwrite => { if table.batches.len() == 1 { table.batches.into_iter().next().unwrap() @@ -533,44 +550,49 @@ async fn stage_pending_table( let staged = match table.mode { PendingMode::Append => db.storage().stage_append(&ds, combined, &[]).await?, PendingMode::Merge => { - // RFC-022: a partial-update table stages matched-only, and the - // staged source is projected down to (key + assigned) — completion - // columns were carried for the validation change-set only; staging - // them would patch unassigned columns and prune their indexes. + db.storage() + .stage_merge_insert( + ds.clone(), + combined, + vec!["id".to_string()], + lance::dataset::WhenMatched::UpdateAll, + lance::dataset::WhenNotMatched::InsertAll, + ) + .await? + } + PendingMode::PartialUpdate => { + // RFC-022: staged source = batch schema − validation-only columns + // (completion columns ride the validation change-set only; staging + // them would patch unassigned columns and prune their indexes). A + // subtraction cannot reference a column the batch lacks, so + // apply-time omissions (unassigned blobs) are immune. Matched-only: // Lance patches the provided columns in place and never inserts — - // the subset schema could not satisfy non-null target columns - // anyway (Lance rejects partial-source inserts), and the update - // executor only stages rows it matched against the committed table. - let (combined, when_not_matched) = match &table.partial_stage_cols { - Some(stage_cols) => { - let indices = stage_cols - .iter() - .map(|c| { - combined.schema().index_of(c).map_err(|e| { - OmniError::manifest_internal(format!( - "partial-update stage column '{}' missing from \ - accumulated batch: {}", - c, e - )) - }) - }) - .collect::>>()?; - ( - combined - .project(&indices) - .map_err(|e| OmniError::Lance(e.to_string()))?, - lance::dataset::WhenNotMatched::DoNothing, + // a subset schema could not satisfy non-null target columns anyway, + // and the update executor only stages rows it matched. + let validation_only = + table.partial_validation_only.as_deref().ok_or_else(|| { + OmniError::manifest_internal( + "PartialUpdate entry without validation-only columns".to_string(), ) - } - None => (combined, lance::dataset::WhenNotMatched::InsertAll), - }; + })?; + let schema = combined.schema(); + let indices: Vec = schema + .fields() + .iter() + .enumerate() + .filter(|(_, f)| !validation_only.contains(f.name())) + .map(|(i, _)| i) + .collect(); + let projected = combined + .project(&indices) + .map_err(|e| OmniError::Lance(e.to_string()))?; db.storage() .stage_merge_insert( ds.clone(), - combined, + projected, vec!["id".to_string()], lance::dataset::WhenMatched::UpdateAll, - when_not_matched, + lance::dataset::WhenNotMatched::DoNothing, ) .await? } diff --git a/crates/omnigraph/src/instrumentation.rs b/crates/omnigraph/src/instrumentation.rs index 80a95642..df2d045e 100644 --- a/crates/omnigraph/src/instrumentation.rs +++ b/crates/omnigraph/src/instrumentation.rs @@ -247,17 +247,24 @@ pub(crate) fn record_stage_append(rows: u64) { /// Record one `stage_merge_insert` of `rows` rows against the active probes, /// with its observable shape (source columns + unmatched-row disposition). -/// No-op in production (no probes installed). +/// No-op in production (no probes installed): the schema arrives as a cheap +/// `&SchemaRef` borrow and the column-name strings materialize only inside +/// `try_with` — the same classify-inside-the-closure pattern as +/// `record_open`, so production pays zero allocations. pub(crate) fn record_stage_merge_insert( rows: u64, - source_columns: Vec, + source_schema: &arrow_schema::SchemaRef, inserts_unmatched: bool, ) { let _ = MERGE_WRITE_PROBES.try_with(|p| { p.stage_merge_insert_calls.fetch_add(1, Ordering::Relaxed); p.stage_merge_insert_rows.fetch_add(rows, Ordering::Relaxed); p.merge_shapes.lock().unwrap().push(MergeShape { - source_columns, + source_columns: source_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(), inserts_unmatched, }); }); diff --git a/crates/omnigraph/src/table_store.rs b/crates/omnigraph/src/table_store.rs index 72918fd9..76e23e3e 100644 --- a/crates/omnigraph/src/table_store.rs +++ b/crates/omnigraph/src/table_store.rs @@ -856,7 +856,7 @@ impl TableStore { return self.table_state(dataset_uri, ds).await; } let schema = batch.schema(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); let params = WriteParams { mode: WriteMode::Append, allow_external_blob_outside_bases: true, @@ -1178,12 +1178,6 @@ impl TableStore { } let merged_rows = batch.num_rows() as u64; let inserts_unmatched = matches!(when_not_matched, WhenNotMatched::InsertAll); - let source_columns: Vec = batch - .schema() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); // Precondition for the FirstSeen workaround below: every call path that // reaches stage_merge_insert (load, MutationStaging::finalize, @@ -1218,7 +1212,7 @@ impl TableStore { .try_build() .map_err(|e| OmniError::Lance(e.to_string()))?; let schema = batch.schema(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); let stream = lance_datafusion::utils::reader_to_stream(Box::new(reader)); let uncommitted = job .execute_uncommitted(stream) @@ -1228,7 +1222,7 @@ impl TableStore { // not inflate the probe (matches `stage_append`/`stage_append_stream`). crate::instrumentation::record_stage_merge_insert( merged_rows, - source_columns, + &schema, inserts_unmatched, ); // Operation::Update { removed_fragment_ids, updated_fragments, new_fragments, .. } — diff --git a/crates/omnigraph/tests/validators.rs b/crates/omnigraph/tests/validators.rs index 31da6bbb..7a331318 100644 --- a/crates/omnigraph/tests/validators.rs +++ b/crates/omnigraph/tests/validators.rs @@ -778,3 +778,72 @@ async fn partial_update_completes_composite_unique_group() { .await .unwrap(); } + +const OVERLAP_UNIQUE_SCHEMA: &str = r#" +node Booking { + name: String @key + room: String? + hour: I32? + day: String? + @unique(room, hour) + @unique(hour, day) +} +"#; + +const OVERLAP_MUTATIONS: &str = r#" +query insert_booking($name: String, $room: String, $hour: I32, $day: String) { + insert Booking { name: $name, room: $room, hour: $hour, day: $day } +} +query move_room($name: String, $room: String) { + update Booking set { room: $room } where name = $name +} +"#; + +/// Overlapping `@unique` groups: a sole update assigning `room` pulls in +/// `hour` (its group partner), and `hour`'s OTHER group `(hour, day)` must +/// then be completed transitively — the evaluator hard-errors on a partially +/// present group, so without the transitive closure a perfectly legal update +/// fails with "missing unique column 'day'". Both directions pinned: the +/// legal move succeeds, and a genuine (room, hour) collision is still caught. +#[tokio::test] +async fn overlapping_unique_groups_allow_sole_partial_update() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let db = Omnigraph::init(uri, OVERLAP_UNIQUE_SCHEMA).await.unwrap(); + + for (name, room, hour, day) in [("b1", "roomA", 9i64, "mon"), ("b2", "roomB", 9, "tue")] { + db.mutate( + "main", + OVERLAP_MUTATIONS, + "insert_booking", + &mixed_params(&[("$name", name), ("$room", room), ("$day", day)], &[("$hour", hour)]), + ) + .await + .unwrap(); + } + + // Legal sole update — touches neither `day` nor collides. + db.mutate( + "main", + OVERLAP_MUTATIONS, + "move_room", + ¶ms(&[("$name", "b2"), ("$room", "roomC")]), + ) + .await + .unwrap(); + + // Genuine collision on (room, hour) still rejected. + let err = db + .mutate( + "main", + OVERLAP_MUTATIONS, + "move_room", + ¶ms(&[("$name", "b2"), ("$room", "roomA")]), + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("@unique"), + "moving b2 to (roomA, 9) must collide with b1, got: {err}" + ); +} diff --git a/crates/omnigraph/tests/writes.rs b/crates/omnigraph/tests/writes.rs index 5f7206bf..c4b56e6d 100644 --- a/crates/omnigraph/tests/writes.rs +++ b/crates/omnigraph/tests/writes.rs @@ -677,8 +677,10 @@ async fn mixed_insert_and_update_on_same_person_coalesces_to_one_merge() { let pre_version = version_main(&db).await.unwrap(); - let result = db - .mutate( + let probes = omnigraph::instrumentation::MergeWriteProbes::default(); + let result = omnigraph::instrumentation::with_merge_write_probes( + probes.clone(), + db.mutate( "main", STAGED_QUERIES, "insert_then_update_same_person", @@ -686,11 +688,28 @@ async fn mixed_insert_and_update_on_same_person_coalesces_to_one_merge() { &[("$name", "Yves")], &[("$insert_age", 10), ("$update_age", 99)], ), - ) - .await - .unwrap(); + ), + ) + .await + .unwrap(); assert_eq!(result.affected_nodes, 2, "1 insert + 1 update reported"); + // RFC-022 fallback pin: mixing insert + update on one table stages FULL + // rows as an upsert — partial and full batches cannot share one + // uniform-schema merge source, and one table commits at most one version + // per query (invariant 4). + let shapes = probes.merge_shapes(); + assert_eq!(shapes.len(), 1, "insert+update coalesce into one staged merge"); + assert!( + shapes[0].source_columns.contains(&"name".to_string()), + "mixed insert+update stages full rows, got {:?}", + shapes[0].source_columns + ); + assert!( + shapes[0].inserts_unmatched, + "mixed insert+update keeps upsert semantics (InsertAll)" + ); + // The end-state row carries the update value (last-write-wins via // dedupe in finalize), proving the staged merge_insert ran with the // correct source dedupe. Read the underlying Person table directly @@ -1083,22 +1102,38 @@ async fn chained_updates_with_overlapping_predicate_respects_intermediate_value( let pre_version = version_main(&db).await.unwrap(); - db.mutate( - "main", - STAGED_QUERIES, - "update_then_filter_by_old_value", - &mixed_params( - &[("$first_name", "Alice")], - &[ - ("$first_new_age", 99), - ("$second_threshold", 50), - ("$second_new_age", 10), - ], + let probes = omnigraph::instrumentation::MergeWriteProbes::default(); + omnigraph::instrumentation::with_merge_write_probes( + probes.clone(), + db.mutate( + "main", + STAGED_QUERIES, + "update_then_filter_by_old_value", + &mixed_params( + &[("$first_name", "Alice")], + &[ + ("$first_new_age", 99), + ("$second_threshold", 50), + ("$second_new_age", 10), + ], + ), ), ) .await .unwrap(); + // RFC-022 fallback pin: two updates on one table stage FULL rows — a + // later update's read-your-writes scan must see the earlier update's + // full effect, and two partial batches with different assigned sets + // cannot share one uniform-schema merge source. + let shapes = probes.merge_shapes(); + assert_eq!(shapes.len(), 1, "chained updates coalesce into one staged merge"); + assert!( + shapes[0].source_columns.contains(&"name".to_string()), + "chained updates stage full rows (fallback), got {:?}", + shapes[0].source_columns + ); + // After op-1: Alice = 99. After op-2 (where age > 50): Alice // matches (99 > 50) → set to 10. End state: Alice = 10. let batches = read_table(&db, "node:Person").await; @@ -2068,20 +2103,6 @@ async fn filtered_read_after_append_and_delete_is_consistent() { // ─── RFC-022: field-level update staging shape ────────────────────────────── -const MIXED_INSERT_UPDATE: &str = r#" -query mix($newname: String, $age: I32, $target: String, $newage: I32) { - insert Person { name: $newname, age: $age } - update Person set { age: $newage } where name = $target -} -"#; - -const TWO_UPDATES: &str = r#" -query two($a: String, $aage: I32, $b: String, $bage: I32) { - update Person set { age: $aage } where name = $a - update Person set { age: $bage } where name = $b -} -"#; - /// RFC-022 acceptance (red → green): a query whose ONLY op on a table is one /// `update` stages a PARTIAL source — exactly (merge key + assigned columns) — /// as a matched-only merge (`WhenNotMatched::DoNothing`). The whole-row path @@ -2121,106 +2142,75 @@ async fn single_update_stages_partial_matched_only_source() { ); } -/// Fallback pin: a query mixing insert + update on ONE table stages full rows -/// as an upsert — partial and full batches cannot share one merge source, and -/// one table commits at most one version per query (invariant 4). Pins the -/// RFC-022 fallback so the partial path can never split a table's commit. +/// An update matching zero rows stages nothing and commits nothing — +/// unchanged by RFC-022 (the early return precedes staging). #[tokio::test] -async fn mixed_insert_update_same_table_stages_full_row_upsert() { +async fn empty_match_update_stages_no_merge() { use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; let dir = tempfile::tempdir().unwrap(); let db = init_and_load(&dir).await; let probes = MergeWriteProbes::default(); - with_merge_write_probes( + let result = with_merge_write_probes( probes.clone(), db.mutate( "main", - MIXED_INSERT_UPDATE, - "mix", - &mixed_params( - &[("$newname", "Zed"), ("$target", "Alice")], - &[("$age", 20), ("$newage", 42)], - ), + MUTATION_QUERIES, + "set_age", + &mixed_params(&[("$name", "nobody-here")], &[("$age", 1)]), ), ) .await .unwrap(); - let shapes = probes.merge_shapes(); - assert_eq!(shapes.len(), 1, "insert+update coalesce into one staged merge"); - let mut cols = shapes[0].source_columns.clone(); - cols.sort(); - assert!( - cols.contains(&"name".to_string()) && cols.contains(&"age".to_string()), - "mixed insert+update stages full rows, got {:?}", - shapes[0].source_columns - ); - assert!( - shapes[0].inserts_unmatched, - "mixed insert+update keeps upsert semantics (InsertAll)" - ); + assert_eq!(result.affected_nodes, 0); + assert!(probes.merge_shapes().is_empty(), "no merge staged for an empty match"); } -/// Fallback pin: two updates on one table in one query stage full rows (a -/// later update's read-your-writes scan must see the earlier update's full -/// effect; two partial batches with different assigned sets cannot share one -/// uniform-schema merge source — schema-level partial semantics make a union -/// unsound). Behavior identical to today. -#[tokio::test] -async fn chained_updates_same_table_stage_full_rows() { - use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; - let dir = tempfile::tempdir().unwrap(); - let db = init_and_load(&dir).await; - - let probes = MergeWriteProbes::default(); - with_merge_write_probes( - probes.clone(), - db.mutate( - "main", - TWO_UPDATES, - "two", - &mixed_params( - &[("$a", "Alice"), ("$b", "Bob")], - &[("$aage", 51), ("$bage", 52)], - ), - ), - ) - .await - .unwrap(); +const BLOB_DOC_SCHEMA: &str = r#" +node Doc { + name: String @key + content: Blob? +} +"#; - let shapes = probes.merge_shapes(); - assert_eq!(shapes.len(), 1, "chained updates coalesce into one staged merge"); - let mut cols = shapes[0].source_columns.clone(); - cols.sort(); - assert!( - cols.contains(&"name".to_string()), - "chained updates stage full rows (fallback), got {:?}", - shapes[0].source_columns - ); +const BLOB_DOC_MUTATIONS: &str = r#" +query insert_doc($name: String) { + insert Doc { name: $name } +} +query set_content($name: String, $c: Blob?) { + update Doc set { content: $c } where name = $name } +"#; -/// An update matching zero rows stages nothing and commits nothing — -/// unchanged by RFC-022 (the early return precedes staging). +/// A sole update assigning `null` to an optional Blob is the historical no-op +/// (the blob column is omitted from the staged batch and the merge leaves it +/// untouched). The partial-update plan must mirror that omission: listing the +/// blob in the stage projection while `apply_assignments` omits it fails the +/// whole mutation with an internal error — and made the outcome depend on +/// whether unrelated statements shared the query (whole-row fallback +/// succeeded). Pinned: the sole-update path succeeds like the fallback does. #[tokio::test] -async fn empty_match_update_stages_no_merge() { - use omnigraph::instrumentation::{MergeWriteProbes, with_merge_write_probes}; +async fn null_blob_assignment_in_sole_update_is_a_no_op() { + use omnigraph_compiler::query::ast::Literal; let dir = tempfile::tempdir().unwrap(); - let db = init_and_load(&dir).await; + let uri = dir.path().to_str().unwrap(); + let db = Omnigraph::init(uri, BLOB_DOC_SCHEMA).await.unwrap(); - let probes = MergeWriteProbes::default(); - let result = with_merge_write_probes( - probes.clone(), - db.mutate( - "main", - MUTATION_QUERIES, - "set_age", - &mixed_params(&[("$name", "nobody-here")], &[("$age", 1)]), - ), + db.mutate( + "main", + BLOB_DOC_MUTATIONS, + "insert_doc", + ¶ms(&[("$name", "d1")]), ) .await .unwrap(); - assert_eq!(result.affected_nodes, 0); - assert!(probes.merge_shapes().is_empty(), "no merge staged for an empty match"); + let mut p = params(&[("$name", "d1")]); + p.insert("c".to_string(), Literal::Null); + let result = db + .mutate("main", BLOB_DOC_MUTATIONS, "set_content", &p) + .await + .unwrap(); + assert_eq!(result.affected_nodes, 1); } From 3d9e8ffcd9e1f5d4544015ec2f3346b44b0fba9e Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Sun, 12 Jul 2026 12:46:00 +0100 Subject: [PATCH 6/6] docs: staged partial source is (key + assigned); completion columns are read-only validation inputs The user doc, dev write-path doc, and probe doc still described the pre-review staged shape (completion columns staged; unassigned columns 'never read'). Corrected to the implemented contract: the patched Lance source is exactly (row key + assigned); @unique-completion columns may be READ as validation-only inputs and are projected out before staging, never rewritten; all other unassigned columns are neither read nor rewritten. --- crates/omnigraph/src/instrumentation.rs | 4 +++- docs/dev/writes.md | 6 ++++-- docs/user/mutations/index.md | 14 ++++++++------ 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/crates/omnigraph/src/instrumentation.rs b/crates/omnigraph/src/instrumentation.rs index f0a4bc8f..fafb38f6 100644 --- a/crates/omnigraph/src/instrumentation.rs +++ b/crates/omnigraph/src/instrumentation.rs @@ -179,7 +179,9 @@ pub struct MergeWriteProbes { /// names (schema order) and whether unmatched source rows insert /// (`WhenNotMatched::InsertAll`) or drop (`DoNothing`). Lets a fitness test /// assert a partial-schema matched-only update stages exactly - /// (key + assigned + completion) columns — the field-level-update staging shape. + /// (key + assigned) columns — the field-level-update staging shape. + /// (`@unique`-completion columns ride only the validation change-set and + /// are projected out before staging, so they never appear here.) pub merge_shapes: Arc>>, /// Inline vector-index (IVF) builds. The fast-forward adopt path defers /// index coverage to the reconciler, so an adopt merge must do 0 of these. diff --git a/docs/dev/writes.md b/docs/dev/writes.md index 586681c8..1cfa1467 100644 --- a/docs/dev/writes.md +++ b/docs/dev/writes.md @@ -240,8 +240,10 @@ shared by both `mutate_as` and the bulk loader: unassigned group member would patch it and prune its index for no semantic reason. Lance's partial-schema path patches the provided columns in place on the same fragment (`update_mode: - RewriteColumns`), so unassigned columns are never read or rewritten and - indexes over them keep fragment coverage (Lance prunes only + RewriteColumns`), so the staged source is exactly (row key + assigned): + completion columns are read as validation-only inputs but never rewritten, + every other unassigned column is neither read nor rewritten, and indexes + over all unassigned columns keep fragment coverage (Lance prunes only `fields_modified`). Every other shape (insert+update on one table, multiple updates, etc.) falls back to whole-row staging: partial and full batches cannot share one uniform-schema merge source (a present column's null cell diff --git a/docs/user/mutations/index.md b/docs/user/mutations/index.md index 3de157e1..f00011b0 100644 --- a/docs/user/mutations/index.md +++ b/docs/user/mutations/index.md @@ -38,12 +38,14 @@ multi-query workflows. ## Updates write only the columns they assign When an `update` is the only statement touching its type in a query, the engine -stages just the assigned columns (plus the row key and any columns needed to -complete a `@unique` group the assignment touches) and patches them in place. -Unassigned columns — including large `Vector` embedding columns — are neither -read nor rewritten, and indexes over them keep their coverage. Semantics are -identical to a whole-row update; the difference is cost: update latency and -write volume scale with what you assign, not with row width. Types with a +patches just the row key plus the assigned columns in place. Columns that +complete a `@unique` group the assignment touches may be **read** as +validation-only inputs (so the whole tuple is checked), but they are never +rewritten. All other unassigned columns — including large `Vector` embedding +columns — are neither read nor rewritten, and indexes over every unassigned +column keep their coverage. Semantics are identical to a whole-row update; the +difference is cost: update latency and write volume scale with what you assign, +not with row width. Types with a `Blob` property currently always use the whole-row path (including the blob materialization described above).