diff --git a/crates/omnigraph/src/exec/mutation.rs b/crates/omnigraph/src/exec/mutation.rs index c56897a8..df6feff8 100644 --- a/crates/omnigraph/src/exec/mutation.rs +++ b/crates/omnigraph/src/exec/mutation.rs @@ -397,6 +397,74 @@ 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 { + /// 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, + /// `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(); + let mut stage_cols = 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" || is_assigned { + stage_cols.push(name.to_string()); + } + 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)), + stage_cols, + scan_schema: Arc::new(Schema::new(scan_fields)), + scan_cols, + } +} + fn apply_assignments( full_schema: &SchemaRef, batch: &RecordBatch, @@ -916,6 +984,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 +1022,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 +1158,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 +1182,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 +1228,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 +1276,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 +1292,24 @@ 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 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)?; + } Ok(MutationResult { affected_nodes: affected_count, diff --git a/crates/omnigraph/src/exec/staging.rs b/crates/omnigraph/src/exec/staging.rs index 2f865f37..2b8704f9 100644 --- a/crates/omnigraph/src/exec/staging.rs +++ b/crates/omnigraph/src/exec/staging.rs @@ -52,6 +52,18 @@ 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 + /// [`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>, } impl PendingTable { @@ -60,6 +72,7 @@ impl PendingTable { schema, mode, batches: Vec::new(), + partial_stage_cols: None, } } @@ -184,6 +197,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_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", + table_key + ))); + } if existing.mode == PendingMode::Overwrite || mode == PendingMode::Overwrite { if existing.mode != mode { return Err(OmniError::manifest_internal(format!( @@ -221,6 +241,38 @@ 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, + stage_cols: Vec, + ) -> 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_stage_cols = Some(stage_cols); + 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 +533,44 @@ 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. + // 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( 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/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/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/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/crates/omnigraph/tests/scalar_indexes.rs b/crates/omnigraph/tests/scalar_indexes.rs index 8d8a3f0a..1382b708 100644 --- a/crates/omnigraph/tests/scalar_indexes.rs +++ b/crates/omnigraph/tests/scalar_indexes.rs @@ -72,3 +72,111 @@ 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 ["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:?}" + ); + } + + // 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" + ); +} + +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/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"); +} diff --git a/docs/dev/writes.md b/docs/dev/writes.md index e22e85b1..f34679d3 100644 --- a/docs/dev/writes.md +++ b/docs/dev/writes.md @@ -64,6 +64,32 @@ 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 — 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 + `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**.