Field-level updates: stage only assigned columns in update mutations

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
+                        )));
                     }
                 }
             }

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}"
+    );
+}

crates/omnigraph/src/exec/mutation.rs

@@ -382,6 +382,124 @@ fn predicate_to_sql(
 
 /// Replace specific columns in a RecordBatch with new literal values.
 ///
+/// Field-level update: 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,
+    /// 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<String>,
+    /// 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<String>,
+}
+
+/// Compute the partial-update column plan: the staged source carries the merge
+/// 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,
+    resolved: &HashMap<String, Literal>,
+) -> PartialUpdatePlan {
+    // A blob assigned a non-String value is a no-op on every path (the
+    // whole-row path copies the committed value through; historically the
+    // column was omitted) — exclude it from the plan so the partial batch
+    // neither requires a materializing scan nor patches the blob.
+    let assigned: HashSet<&str> = resolved
+        .iter()
+        .filter(|(name, value)| {
+            !(node_type.blob_properties.contains(name.as_str())
+                && !matches!(value, Literal::String(_)))
+        })
+        .map(|(name, _)| name.as_str())
+        .collect();
+
+    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 validation_only = Vec::new();
+    for field in full_schema.fields() {
+        let name = field.name().as_str();
+        let is_assigned = assigned.contains(name);
+        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 || is_completion {
+            output_fields.push(field.clone());
+        }
+        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
+        .iter()
+        .map(|f| f.name().clone())
+        .collect::<Vec<_>>();
+    PartialUpdatePlan {
+        output_schema: Arc::new(Schema::new(output_fields)),
+        validation_only,
+        scan_schema: Arc::new(Schema::new(scan_fields)),
+        scan_cols,
+    }
+}
+
 /// Blob-bearing updates always arrive with the full logical schema. Committed
 /// blob payloads were materialized by the caller and rebuilt as logical
 /// `Struct<data,uri>` arrays; pending batches already have that shape. An
@@ -902,6 +1020,31 @@ impl Omnigraph {
                     assignments,
                     predicate,
                 } => {
+                    // Field-level-update 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,
@@ -910,6 +1053,7 @@ impl Omnigraph {
                         branch,
                         staging,
                         txn,
+                        partial_ok,
                     )
                     .await?
                 }
@@ -1044,6 +1188,7 @@ impl Omnigraph {
         branch: Option<&str>,
         staging: &mut MutationStaging,
         txn: &crate::db::WriteTxn,
+        partial_ok: bool,
     ) -> Result<MutationResult> {
         let catalog = &txn.catalog;
         // Defense in depth: ensure this is a node type
@@ -1054,12 +1199,13 @@ impl Omnigraph {
             )));
         }
 
-        // Reject updates to @key properties — identity is immutable
-        if let Some(key_prop) = 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) = &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
                 )));
             }
         }
@@ -1068,6 +1214,34 @@ impl Omnigraph {
         let schema = catalog.node_types[type_name].arrow_schema.clone();
         let blob_props = catalog.node_types[type_name].blob_properties.clone();
 
+        let mut resolved: HashMap<String, Literal> = HashMap::new();
+        for a in assignments {
+            resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?);
+        }
+
+        // Field-level update: 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. Takes RESOLVED assignments so no-op blob
+        // assignments (non-String values) are excluded from the plan.
+        //
+        // Blob-bearing tables always take the whole-row path: Lance's
+        // partial-source column patch trips the blob writer on a blob-bearing
+        // fragment even when no blob column is in the source ("Blob struct
+        // missing `data` field", lance blob.rs) — pinned by
+        // `end_to_end::update_with_blob_mid_schema_does_not_panic`. Revisit
+        // when Lance's blob merge accepts partial sources.
+        let partial_ok = partial_ok && blob_props.is_empty();
+        let partial_plan = if partial_ok {
+            Some(partial_update_plan(
+                &catalog.node_types[type_name],
+                &schema,
+                &resolved,
+            ))
+        } else {
+            None
+        };
+
         let table_key = format!("node:{}", type_name);
         let (handle, _full_path, _table_branch) = open_table_for_mutation(
             self,
@@ -1100,7 +1274,25 @@ impl Omnigraph {
         // only those payloads into the logical blob schema before unioning
         // pending rows. This keeps correctness independent of whether an id
         // index happens to steer Lance onto its legacy partial-column plan.
-        let batches = if blob_props.is_empty() {
+        let batches = if let Some(plan) = partial_plan.as_ref() {
+            // Field-level update: scan only (id + completion) — assigned
+            // columns' old values are never needed (assignments are literal
+            // values) and the WHERE predicate evaluates by pushdown without
+            // being projected. Blobs are never in this projection (the plan
+            // excludes them), so the blob-materialization arm below is
+            // unnecessary here; pending is empty by the eligibility rule.
+            let scan_cols: Vec<&str> = plan.scan_cols.iter().map(|c| c.as_str()).collect();
+            self.storage()
+                .scan_with_pending(
+                    &ds,
+                    pending_batches,
+                    pending_schema,
+                    Some(scan_cols.as_slice()),
+                    Some(&pred_sql),
+                    Some("id"),
+                )
+                .await?
+        } else if blob_props.is_empty() {
             self.storage()
                 .scan_with_pending(
                     &ds,
@@ -1131,26 +1323,35 @@ impl Omnigraph {
         }
 
         // Concat the matched batches (committed + pending) into one. The
-        // helper binds both sides to the catalog's full logical schema. Any
-        // divergence here is an internal scan/staging contract violation.
-        let matched = concat_match_batches_to_schema(&schema, batches)?;
+        // helper binds both sides to the catalog's full logical schema (partial
+        // path: to the plan's SCAN schema — id + completion). Any divergence
+        // here is an internal scan/staging contract violation.
+        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, batches)?;
 
         let affected_count = matched.num_rows();
 
-        let mut resolved: HashMap<String, Literal> = HashMap::new();
-        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: 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.validation_only.clone(),
+            )?;
+        } else {
+            staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?;
+        }
 
         Ok(MutationResult {
             affected_nodes: affected_count,