fix `:update` on relations with default val cols

cozo-core/src/query/stored.rs

@@ -223,7 +223,8 @@ impl<'a> SessionTx<'a> {
         force_collect: &str,
         span: SourceSpan,
     ) -> Result<()> {
-        let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == &relation_store.name;
+        let is_callback_target = callback_targets.contains(&relation_store.name)
+            || force_collect == &relation_store.name;
 
         if relation_store.access_level < AccessLevel::Protected {
             bail!(InsufficientAccessLevel(
@@ -240,7 +241,8 @@ impl<'a> SessionTx<'a> {
             headers,
         )?;
 
-        let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
+        let need_to_collect = !force_collect.is_empty()
+            || (!relation_store.is_temp
                 && (is_callback_target
                     || (propagate_triggers && !relation_store.put_triggers.is_empty())));
         let has_indices = !relation_store.indices.is_empty();
@@ -464,11 +466,9 @@ impl<'a> SessionTx<'a> {
     ) -> Result<BTreeMap<SmartString<LazyCompact>, (Arc<TextAnalyzer>, Vec<Bytecode>)>> {
         let mut processors = BTreeMap::new();
         for (name, (_, manifest)) in relation_store.fts_indices.iter() {
-            let tokenizer = self.tokenizers.get(
-                &name,
-                &manifest.tokenizer,
-                &manifest.filters,
-            )?;
+            let tokenizer = self
+                .tokenizers
+                .get(&name, &manifest.tokenizer, &manifest.filters)?;
 
             let parsed = CozoScriptParser::parse(Rule::expr, &manifest.extractor)
                 .into_diagnostic()?
@@ -481,11 +481,9 @@ impl<'a> SessionTx<'a> {
             processors.insert(name.clone(), (tokenizer, extractor));
         }
         for (name, (_, _, manifest)) in relation_store.lsh_indices.iter() {
-            let tokenizer = self.tokenizers.get(
-                &name,
-                &manifest.tokenizer,
-                &manifest.filters,
-            )?;
+            let tokenizer = self
+                .tokenizers
+                .get(&name, &manifest.tokenizer, &manifest.filters)?;
 
             let parsed = CozoScriptParser::parse(Rule::expr, &manifest.extractor)
                 .into_diagnostic()?
@@ -535,7 +533,8 @@ impl<'a> SessionTx<'a> {
         force_collect: &str,
         span: SourceSpan,
     ) -> Result<()> {
-        let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == &relation_store.name;
+        let is_callback_target = callback_targets.contains(&relation_store.name)
+            || force_collect == &relation_store.name;
 
         if relation_store.access_level < AccessLevel::Protected {
             bail!(InsufficientAccessLevel(
@@ -552,7 +551,8 @@ impl<'a> SessionTx<'a> {
             headers,
         )?;
 
-        let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
+        let need_to_collect = !force_collect.is_empty()
+            || (!relation_store.is_temp
                 && (is_callback_target
                     || (propagate_triggers && !relation_store.put_triggers.is_empty())));
         let has_indices = !relation_store.indices.is_empty();
@@ -567,7 +567,7 @@ impl<'a> SessionTx<'a> {
             &metadata.keys,
             key_bindings,
             headers,
-        );
+        )?;
 
         let mut stack = vec![];
         let hnsw_filters = Self::make_hnsw_filters(relation_store)?;
@@ -928,7 +928,8 @@ impl<'a> SessionTx<'a> {
         force_collect: &str,
         span: SourceSpan,
     ) -> Result<()> {
-        let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == relation_store.name;
+        let is_callback_target =
+            callback_targets.contains(&relation_store.name) || force_collect == relation_store.name;
 
         if relation_store.access_level < AccessLevel::Protected {
             bail!(InsufficientAccessLevel(
@@ -944,7 +945,8 @@ impl<'a> SessionTx<'a> {
             headers,
         )?;
 
-        let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
+        let need_to_collect = !force_collect.is_empty()
+            || (!relation_store.is_temp
                 && (is_callback_target
                     || (propagate_triggers && !relation_store.rm_triggers.is_empty())));
         let has_indices = !relation_store.indices.is_empty();
@@ -1149,11 +1151,17 @@ fn make_update_extractors(
     input: &[ColumnDef],
     bindings: &[Symbol],
     tuple_headers: &[Symbol],
-) -> Vec<Option<DataExtractor>> {
-    stored
-        .iter()
-        .map(|s| make_extractor(s, input, bindings, tuple_headers).ok())
-        .collect_vec()
+) -> Result<Vec<Option<DataExtractor>>> {
+    let input_keys: BTreeSet<_> = input.iter().map(|b| &b.name).collect();
+    let mut extractors = Vec::with_capacity(stored.len());
+    for col in stored.iter() {
+        if input_keys.contains(&col.name) {
+            extractors.push(Some(make_extractor(col, input, bindings, tuple_headers)?));
+        } else {
+            extractors.push(None);
+        }
+    }
+    Ok(extractors)
 }
 
 fn make_extractor(

cozo-core/src/runtime/tests.rs

@@ -1203,3 +1203,26 @@ fn as_store_in_imperative_script() {
         println!("{}", row);
     }
 }
+
+#[test]
+fn update_shall_not_destroy_values() {
+    let db = DbInstance::default();
+    db.run_default(r"?[x, y] <- [[1, 2]] :create z {x => y default 0}").unwrap();
+    let r = db.run_default(r"?[x, y] := *z {x, y}").unwrap();
+    assert_eq!(r.into_json()["rows"], json!([[1, 2]]));
+    db.run_default(r"?[x] <- [[1]] :update z {x}").unwrap();
+    let r = db.run_default(r"?[x, y] := *z {x, y}").unwrap();
+    assert_eq!(r.into_json()["rows"], json!([[1, 2]]));
+}
+
+
+#[test]
+fn update_shall_work() {
+    let db = DbInstance::default();
+    db.run_default(r"?[x, y, z] <- [[1, 2, 3]] :create z {x => y, z}").unwrap();
+    let r = db.run_default(r"?[x, y, z] := *z {x, y, z}").unwrap();
+    assert_eq!(r.into_json()["rows"], json!([[1, 2, 3]]));
+    db.run_default(r"?[x, y] <- [[1, 4]] :update z {x, y}").unwrap();
+    let r = db.run_default(r"?[x, y, z] := *z {x, y, z}").unwrap();
+    assert_eq!(r.into_json()["rows"], json!([[1, 4, 3]]));
+}