Auxiliary column support (#123)

* initial pass at PARTITION KEY support.

* Initial pass, allow auxiliary columns on vec0 virtual tables

* update TODO

* unit tests

* gha this PR branch

* fixup tests

* doc internal

* fix tests, KNN/rowids in

* define SQLITE_INDEX_CONSTRAINT_OFFSET

* whoops

* update tests, syrupy, use uv

* un ignore pyproject.toml

* dot

* tests/

* type error?

* win: .exe, update error name

* try fix macos python, paren around expr?

* win bash?

* dbg :(

* explicit error

* op

* dbg win

* win ./tests/.venv/Scripts/python.exe

* block UPDATEs on partition key values for now

* test this branch

* accidentally removved "partition key type mistmatch" block during merge

* typo ugh

* bruv

* start aux snapshots

* drop aux shadow table on destroy

* enforce column types

* block WHERE constraints on auxiliary columns in KNN queries

* support delete

* support UPDATE on auxiliary columns
This commit is contained in:
Alex Garcia 2024-11-20 00:30:23 -08:00 committed by GitHub
parent 6658624172
commit 9bfeaa7842
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1249 additions and 8 deletions

View file

@ -0,0 +1,642 @@
# serializer version: 1
# name: test_constructor_limit[max 16 auxiliary columns]
dict({
'error': 'OperationalError',
'message': 'vec0 constructor error: More than 16 auxiliary columns were provided',
})
# ---
# name: test_deletes
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_deletes.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_deletes.2
OrderedDict({
'sql': 'delete from v where rowid = 1',
'rows': list([
]),
})
# ---
# name: test_deletes.3
OrderedDict({
'sql': 'select rowid, * from v',
'rows': list([
OrderedDict({
'rowid': 2,
'vector': b'\x00\x00\x00@',
'name': 'brian',
}),
OrderedDict({
'rowid': 3,
'vector': b'\x00\x00@@',
'name': 'craig',
}),
]),
})
# ---
# name: test_deletes.4
dict({
'v_auxiliary': OrderedDict({
'sql': 'select * from v_auxiliary',
'rows': list([
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'\x06',
'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': 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_knn
OrderedDict({
'sql': 'select * 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_knn[illegal KNN w/ aux]
dict({
'error': 'OperationalError',
'message': 'An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query.',
})
# ---
# name: test_knn[legal KNN w/ aux]
OrderedDict({
'sql': "select *, distance from v where vector match '[5]' and k = 10",
'rows': list([
OrderedDict({
'rowid': 3,
'vector': b'\x00\x00@@',
'name': 'craig',
'distance': 2.0,
}),
OrderedDict({
'rowid': 2,
'vector': b'\x00\x00\x00@',
'name': 'brian',
'distance': 3.0,
}),
OrderedDict({
'rowid': 1,
'vector': b'\x00\x00\x80?',
'name': 'alex',
'distance': 4.0,
}),
]),
})
# ---
# name: test_normal.1
OrderedDict({
'sql': 'select * from v',
'rows': list([
OrderedDict({
'rowid': 1,
'a': b'\x11\x11\x11\x11',
'name': 'alex',
}),
OrderedDict({
'rowid': 2,
'a': b'""""',
'name': 'brian',
}),
OrderedDict({
'rowid': 3,
'a': b'3333',
'name': 'craig',
}),
]),
})
# ---
# name: test_normal.2
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'\x11\x11\x11\x11""""3333\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
}),
]),
}),
})
# ---
# name: test_normal[sqlite_master post drop]
OrderedDict({
'sql': 'select * from sqlite_master order by name',
'rows': list([
OrderedDict({
'type': 'table',
'name': 'sqlite_sequence',
'tbl_name': 'sqlite_sequence',
'rootpage': 3,
'sql': 'CREATE TABLE sqlite_sequence(name,seq)',
}),
]),
})
# ---
# name: test_normal[sqlite_master]
OrderedDict({
'sql': 'select * from sqlite_master order by name',
'rows': list([
OrderedDict({
'type': 'index',
'name': 'sqlite_autoindex_v_vector_chunks00_1',
'tbl_name': 'v_vector_chunks00',
'rootpage': 6,
'sql': None,
}),
OrderedDict({
'type': 'table',
'name': 'sqlite_sequence',
'tbl_name': 'sqlite_sequence',
'rootpage': 3,
'sql': 'CREATE TABLE sqlite_sequence(name,seq)',
}),
OrderedDict({
'type': 'table',
'name': 'v',
'tbl_name': 'v',
'rootpage': 0,
'sql': 'CREATE VIRTUAL TABLE v using vec0(a float[1], +name text, chunk_size=8)',
}),
OrderedDict({
'type': 'table',
'name': 'v_auxiliary',
'tbl_name': 'v_auxiliary',
'rootpage': 7,
'sql': 'CREATE TABLE "v_auxiliary"( rowid integer PRIMARY KEY , value00)',
}),
OrderedDict({
'type': 'table',
'name': 'v_chunks',
'tbl_name': 'v_chunks',
'rootpage': 2,
'sql': 'CREATE TABLE "v_chunks"(chunk_id INTEGER PRIMARY KEY AUTOINCREMENT,size INTEGER NOT NULL,validity BLOB NOT NULL,rowids BLOB NOT NULL)',
}),
OrderedDict({
'type': 'table',
'name': 'v_rowids',
'tbl_name': 'v_rowids',
'rootpage': 4,
'sql': 'CREATE TABLE "v_rowids"(rowid INTEGER PRIMARY KEY AUTOINCREMENT,id,chunk_id INTEGER,chunk_offset INTEGER)',
}),
OrderedDict({
'type': 'table',
'name': 'v_vector_chunks00',
'tbl_name': 'v_vector_chunks00',
'rootpage': 5,
'sql': 'CREATE TABLE "v_vector_chunks00"(rowid PRIMARY KEY,vectors BLOB NOT NULL)',
}),
]),
})
# ---
# name: test_types
OrderedDict({
'sql': 'select * from v',
'rows': list([
]),
})
# ---
# name: test_types.1
OrderedDict({
'sql': 'insert into v(vector, aux_int, aux_float, aux_text, aux_blob) values (?, ?, ?, ?, ?)',
'rows': list([
]),
})
# ---
# name: test_types.2
OrderedDict({
'sql': 'select * from v',
'rows': list([
OrderedDict({
'rowid': 1,
'vector': b'\x11\x11\x11\x11',
'aux_int': 1,
'aux_float': 1.22,
'aux_text': 'text',
'aux_blob': b'blob',
}),
]),
})
# ---
# name: test_types.3
dict({
'error': 'OperationalError',
'message': 'Auxiliary column type mismatch: The auxiliary column aux_int has type INTEGER, but TEXT was provided.',
})
# ---
# name: test_types.4
dict({
'error': 'OperationalError',
'message': 'Auxiliary column type mismatch: The auxiliary column aux_float has type FLOAT, but TEXT was provided.',
})
# ---
# name: test_types.5
dict({
'error': 'OperationalError',
'message': 'Auxiliary column type mismatch: The auxiliary column aux_text has type TEXT, but INTEGER was provided.',
})
# ---
# name: test_types.6
dict({
'error': 'OperationalError',
'message': 'Auxiliary column type mismatch: The auxiliary column aux_blob has type BLOB, but INTEGER was provided.',
})
# ---
# name: test_types.7
OrderedDict({
'sql': 'insert into v(vector, aux_int, aux_float, aux_text, aux_blob) values (?, ?, ?, ?, ?)',
'rows': list([
]),
})
# ---
# name: test_types.8
OrderedDict({
'sql': 'select * from v',
'rows': list([
OrderedDict({
'rowid': 1,
'vector': b'\x11\x11\x11\x11',
'aux_int': 1,
'aux_float': 1.22,
'aux_text': 'text',
'aux_blob': b'blob',
}),
OrderedDict({
'rowid': 2,
'vector': b'\x11\x11\x11\x11',
'aux_int': None,
'aux_float': None,
'aux_text': None,
'aux_blob': None,
}),
]),
})
# ---
# 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',
}),
]),
}),
})
# ---

155
tests/test-auxiliary.py Normal file
View file

@ -0,0 +1,155 @@
import sqlite3
from collections import OrderedDict
def test_constructor_limit(db, snapshot):
assert exec(
db,
f"""
create virtual table v using vec0(
{",".join([f"+aux{x} integer" for x in range(17)])}
v float[1]
)
""",
) == snapshot(name="max 16 auxiliary columns")
def test_normal(db, snapshot):
db.execute(
"create virtual table v using vec0(a float[1], +name text, chunk_size=8)"
)
assert exec(db, "select * from sqlite_master order by name") == snapshot(
name="sqlite_master"
)
db.execute("insert into v(a, name) values (?, ?)", [b"\x11\x11\x11\x11", "alex"])
db.execute("insert into v(a, name) values (?, ?)", [b"\x22\x22\x22\x22", "brian"])
db.execute("insert into v(a, name) values (?, ?)", [b"\x33\x33\x33\x33", "craig"])
assert exec(db, "select * from v") == snapshot()
assert vec0_shadow_table_contents(db, "v") == snapshot()
db.execute("drop table v;")
assert exec(db, "select * from sqlite_master order by name") == snapshot(
name="sqlite_master post drop"
)
def test_types(db, snapshot):
db.execute(
"""
create virtual table v using vec0(
vector float[1],
+aux_int integer,
+aux_float float,
+aux_text text,
+aux_blob blob
)
"""
)
assert exec(db, "select * from v") == snapshot()
INSERT = "insert into v(vector, aux_int, aux_float, aux_text, aux_blob) values (?, ?, ?, ?, ?)"
assert (
exec(db, INSERT, [b"\x11\x11\x11\x11", 1, 1.22, "text", b"blob"]) == snapshot()
)
assert exec(db, "select * from v") == snapshot()
# bad types
assert (
exec(db, INSERT, [b"\x11\x11\x11\x11", "not int", 1.2, "text", b"blob"])
== snapshot()
)
assert (
exec(db, INSERT, [b"\x11\x11\x11\x11", 1, "not float", "text", b"blob"])
== snapshot()
)
assert exec(db, INSERT, [b"\x11\x11\x11\x11", 1, 1.2, 1, b"blob"]) == snapshot()
assert exec(db, INSERT, [b"\x11\x11\x11\x11", 1, 1.2, "text", 1]) == snapshot()
# NULLs are totally chill
assert exec(db, INSERT, [b"\x11\x11\x11\x11", None, None, None, None]) == snapshot()
assert exec(db, "select * from v") == snapshot()
def test_updates(db, snapshot):
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):
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, "delete from v where rowid = 1") == snapshot()
assert exec(db, "select rowid, * from v") == snapshot()
assert vec0_shadow_table_contents(db, "v") == snapshot()
def test_knn(db, snapshot):
db.execute("create virtual table v using vec0(vector float[1], +name text)")
db.executemany(
"insert into v(vector, name) values (?, ?)",
[("[1]", "alex"), ("[2]", "brian"), ("[3]", "craig")],
)
assert exec(db, "select * from v") == snapshot()
assert exec(
db, "select *, distance from v where vector match '[5]' and k = 10"
) == snapshot(name="legal KNN w/ aux")
# EVIDENCE-OF: V25623_09693 No aux constraint allowed on KNN queries
assert exec(
db,
"select *, distance from v where vector match '[5]' and k = 10 and name = 'alex'",
) == snapshot(name="illegal KNN w/ aux")
def exec(db, sql, parameters=[]):
try:
rows = db.execute(sql, parameters).fetchall()
except (sqlite3.OperationalError, sqlite3.DatabaseError) as e:
return {
"error": e.__class__.__name__,
"message": str(e),
}
a = []
for row in rows:
o = OrderedDict()
for k in row.keys():
o[k] = row[k]
a.append(o)
result = OrderedDict()
result["sql"] = sql
result["rows"] = a
return result
def vec0_shadow_table_contents(db, v):
shadow_tables = [
row[0]
for row in db.execute(
"select name from sqlite_master where name like ? order by 1", [f"{v}_%"]
).fetchall()
]
o = {}
for shadow_table in shadow_tables:
o[shadow_table] = exec(db, f"select * from {shadow_table}")
return o