From 5d9b923c599ce7dbed1c8dc68fb233c7c5c2029d Mon Sep 17 00:00:00 2001
From: Alex Garcia <alexsebastian.garcia@gmail.com>
Date: Wed, 13 Nov 2024 15:49:04 -0800
Subject: [PATCH] support UPDATE on auxiliary columns

---
 TODO                                    |   1 -
 sqlite-vec.c                            |  40 +++++-
 tests/__snapshots__/test-auxiliary.ambr | 179 ++++++++++++++++++++++++
 tests/test-auxiliary.py                 |  14 +-
 4 files changed, 230 insertions(+), 4 deletions(-)

diff --git a/TODO b/TODO
index dd3b651..7db7b48 100644
--- a/TODO
+++ b/TODO
@@ -7,7 +7,6 @@
 
 # auxiliary columns
 
-- UPDATE support
 - later:
   - NOT NULL?
   - perf: INSERT stmt should be cached on vec0_vtab
diff --git a/sqlite-vec.c b/sqlite-vec.c
index 3fda9aa..1fc2241 100644
--- a/sqlite-vec.c
+++ b/sqlite-vec.c
@@ -7059,6 +7059,28 @@ int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
   return SQLITE_OK;
 }
 
+int vec0Update_UpdateAuxColumn(vec0_vtab *p, int auxiliary_column_idx, sqlite3_value * value, i64 rowid) {
+  int rc;
+  sqlite3_stmt *stmt;
+  const char * zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_AUXILIARY_NAME " SET value%02d = ? WHERE rowid = ?", p->schemaName, p->tableName, auxiliary_column_idx);
+  if(!zSql) {
+    return SQLITE_NOMEM;
+  }
+  rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
+  if(rc != SQLITE_OK) {
+    return rc;
+  }
+  sqlite3_bind_value(stmt, 1, value);
+  sqlite3_bind_int64(stmt, 2, rowid);
+  rc = sqlite3_step(stmt);
+  if(rc != SQLITE_DONE) {
+    sqlite3_finalize(stmt);
+    return SQLITE_ERROR;
+  }
+  sqlite3_finalize(stmt);
+  return SQLITE_OK;
+}
+
 int vec0Update_UpdateVectorColumn(vec0_vtab *p, i64 chunk_id, i64 chunk_offset,
                                   int i, sqlite3_value *valueVector) {
   int rc;
@@ -7183,9 +7205,23 @@ int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) {
     return SQLITE_ERROR;
   }
 
-  // TODO handle auxiliary column updates
+  // 3) handle auxiliary column updates
+  for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
+    if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) {
+      continue;
+    }
+    int auxiliary_column_idx = p->user_column_idxs[i];
+    sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
+    if(sqlite3_value_nochange(value)) {
+      continue;
+    }
+    rc = vec0Update_UpdateAuxColumn(p, auxiliary_column_idx, value, rowid);
+    if(rc != SQLITE_OK) {
+      return SQLITE_ERROR;
+    }
+  }
 
-  // 3) iterate over all new vectors, update the vectors
+  // 4) iterate over all new vectors, update the vectors
   for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
     if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
       continue;
diff --git a/tests/__snapshots__/test-auxiliary.ambr b/tests/__snapshots__/test-auxiliary.ambr
index 01ab03b..eb84f0f 100644
--- a/tests/__snapshots__/test-auxiliary.ambr
+++ b/tests/__snapshots__/test-auxiliary.ambr
@@ -461,3 +461,182 @@
     ]),
   })
 # ---
+# name: test_updates
+  OrderedDict({
+    'sql': 'select rowid, * from v',
+    'rows': list([
+      OrderedDict({
+        'rowid': 1,
+        'vector': b'\x00\x00\x80?',
+        'name': 'alex',
+      }),
+      OrderedDict({
+        'rowid': 2,
+        'vector': b'\x00\x00\x00@',
+        'name': 'brian',
+      }),
+      OrderedDict({
+        'rowid': 3,
+        'vector': b'\x00\x00@@',
+        'name': 'craig',
+      }),
+    ]),
+  })
+# ---
+# name: test_updates.1
+  dict({
+    'v_auxiliary': OrderedDict({
+      'sql': 'select * from v_auxiliary',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'value00': 'alex',
+        }),
+        OrderedDict({
+          'rowid': 2,
+          'value00': 'brian',
+        }),
+        OrderedDict({
+          'rowid': 3,
+          'value00': 'craig',
+        }),
+      ]),
+    }),
+    'v_chunks': OrderedDict({
+      'sql': 'select * from v_chunks',
+      'rows': list([
+        OrderedDict({
+          'chunk_id': 1,
+          'size': 8,
+          'validity': b'\x07',
+          'rowids': b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
+        }),
+      ]),
+    }),
+    'v_rowids': OrderedDict({
+      'sql': 'select * from v_rowids',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 0,
+        }),
+        OrderedDict({
+          'rowid': 2,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 1,
+        }),
+        OrderedDict({
+          'rowid': 3,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 2,
+        }),
+      ]),
+    }),
+    'v_vector_chunks00': OrderedDict({
+      'sql': 'select * from v_vector_chunks00',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'vectors': b'\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
+        }),
+      ]),
+    }),
+  })
+# ---
+# name: test_updates.2
+  OrderedDict({
+    'sql': "update v set name = 'ALEX' where rowid = 1",
+    'rows': list([
+    ]),
+  })
+# ---
+# name: test_updates.3
+  OrderedDict({
+    'sql': 'select rowid, * from v',
+    'rows': list([
+      OrderedDict({
+        'rowid': 1,
+        'vector': b'\x00\x00\x80?',
+        'name': 'ALEX',
+      }),
+      OrderedDict({
+        'rowid': 2,
+        'vector': b'\x00\x00\x00@',
+        'name': 'brian',
+      }),
+      OrderedDict({
+        'rowid': 3,
+        'vector': b'\x00\x00@@',
+        'name': 'craig',
+      }),
+    ]),
+  })
+# ---
+# name: test_updates.4
+  dict({
+    'v_auxiliary': OrderedDict({
+      'sql': 'select * from v_auxiliary',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'value00': 'ALEX',
+        }),
+        OrderedDict({
+          'rowid': 2,
+          'value00': 'brian',
+        }),
+        OrderedDict({
+          'rowid': 3,
+          'value00': 'craig',
+        }),
+      ]),
+    }),
+    'v_chunks': OrderedDict({
+      'sql': 'select * from v_chunks',
+      'rows': list([
+        OrderedDict({
+          'chunk_id': 1,
+          'size': 8,
+          'validity': b'\x07',
+          'rowids': b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
+        }),
+      ]),
+    }),
+    'v_rowids': OrderedDict({
+      'sql': 'select * from v_rowids',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 0,
+        }),
+        OrderedDict({
+          'rowid': 2,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 1,
+        }),
+        OrderedDict({
+          'rowid': 3,
+          'id': None,
+          'chunk_id': 1,
+          'chunk_offset': 2,
+        }),
+      ]),
+    }),
+    'v_vector_chunks00': OrderedDict({
+      'sql': 'select * from v_vector_chunks00',
+      'rows': list([
+        OrderedDict({
+          'rowid': 1,
+          'vectors': b'\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
+        }),
+      ]),
+    }),
+  })
+# ---
diff --git a/tests/test-auxiliary.py b/tests/test-auxiliary.py
index e8a51b7..9f03cbd 100644
--- a/tests/test-auxiliary.py
+++ b/tests/test-auxiliary.py
@@ -73,7 +73,19 @@ def test_types(db, snapshot):
 
 
 def test_updates(db, snapshot):
-    pass
+    db.execute(
+        "create virtual table v using vec0(vector float[1], +name text, chunk_size=8)"
+    )
+    db.executemany(
+        "insert into v(vector, name) values (?, ?)",
+        [("[1]", "alex"), ("[2]", "brian"), ("[3]", "craig")],
+    )
+    assert exec(db, "select rowid, * from v") == snapshot()
+    assert vec0_shadow_table_contents(db, "v") == snapshot()
+
+    assert exec(db, "update v set name = 'ALEX' where rowid = 1") == snapshot()
+    assert exec(db, "select rowid, * from v") == snapshot()
+    assert vec0_shadow_table_contents(db, "v") == snapshot()
 
 
 def test_deletes(db, snapshot):