From 03f29d581dbd117c33cf610538009af14e5d9946 Mon Sep 17 00:00:00 2001
From: Alex Garcia <alexsebastian.garcia@gmail.com>
Date: Wed, 13 Nov 2024 12:59:54 -0800
Subject: [PATCH] block UPDATEs on partition key values for now

---
 sqlite-vec.c                                 | 14 +++++++---
 tests/__snapshots__/test-partition-keys.ambr | 28 ++++++++++++++++++++
 tests/test-partition-keys.py                 | 21 +++++++++++++++
 3 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/sqlite-vec.c b/sqlite-vec.c
index b277f46..caa992e 100644
--- a/sqlite-vec.c
+++ b/sqlite-vec.c
@@ -6786,19 +6786,27 @@ int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) {
     rowid = sqlite3_value_int64(argv[0]);
   }
 
-  // 1. get chunk_id and chunk_offset from _rowids
+  // 1) get chunk_id and chunk_offset from _rowids
   rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
   if (rc != SQLITE_OK) {
     return rc;
   }
 
-  // 2) iterate over all new vectors, update the vectors
+  // 2) update any partition key values
   for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
     if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) {
-      // TODO handle partition key values in UPDATEs
+      continue;
+    }
+    int partition_key_idx = p->user_column_idxs[i];
+    sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
+    if(sqlite3_value_nochange(value)) {
+      continue;
     }
+    vtab_set_error(pVTab, "UPDATE on partition key columns are not supported yet. ");
+    return SQLITE_ERROR;
   }
 
+  // 3) 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-partition-keys.ambr b/tests/__snapshots__/test-partition-keys.ambr
index f20c92a..a9dca88 100644
--- a/tests/__snapshots__/test-partition-keys.ambr
+++ b/tests/__snapshots__/test-partition-keys.ambr
@@ -215,3 +215,31 @@
     }),
   })
 # ---
+# name: test_updates[1. Initial dataset]
+  OrderedDict({
+    'sql': 'select * from v',
+    'rows': list([
+      OrderedDict({
+        'rowid': 1,
+        'p': 'a',
+        'a': b'\x11\x11\x11\x11',
+      }),
+      OrderedDict({
+        'rowid': 2,
+        'p': 'a',
+        'a': b'""""',
+      }),
+      OrderedDict({
+        'rowid': 3,
+        'p': 'a',
+        'a': b'3333',
+      }),
+    ]),
+  })
+# ---
+# name: test_updates[2. update #1]
+  dict({
+    'error': 'OperationalError',
+    'message': 'UPDATE on partition key columns are not supported yet. ',
+  })
+# ---
diff --git a/tests/test-partition-keys.py b/tests/test-partition-keys.py
index cf4f5ba..41c7671 100644
--- a/tests/test-partition-keys.py
+++ b/tests/test-partition-keys.py
@@ -53,6 +53,27 @@ def test_types(db, snapshot):
     )
 
 
+def test_updates(db, snapshot):
+    db.execute(
+        "create virtual table v using vec0(p text partition key, a float[1], chunk_size=8)"
+    )
+
+    db.execute(
+        "insert into v(rowid, p, a) values (?, ?, ?)", [1, "a", b"\x11\x11\x11\x11"]
+    )
+    db.execute(
+        "insert into v(rowid, p, a) values (?, ?, ?)", [2, "a", b"\x22\x22\x22\x22"]
+    )
+    db.execute(
+        "insert into v(rowid, p, a) values (?, ?, ?)", [3, "a", b"\x33\x33\x33\x33"]
+    )
+
+    assert exec(db, "select * from v") == snapshot(name="1. Initial dataset")
+    assert exec(db, "update v set p = ? where rowid = ?", ["new", 1]) == snapshot(
+        name="2. update #1"
+    )
+
+
 class Row:
     def __init__(self):
         pass