sqlite-vec/tests/test-insert-delete.py

486 lines
16 KiB
Python
Raw Normal View History

import sqlite3
import struct
import pytest
from helpers import _f32, _i64, _int8, exec
def test_insert_creates_chunks_and_vectors(db, snapshot):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
vecs = [
[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0],
[0.1, 0.2, 0.3, 0.4],
[10.0, 20.0, 30.0, 40.0],
[0.5, 0.5, 0.5, 0.5],
]
for i, v in enumerate(vecs, start=1):
db.execute("insert into v(rowid, emb) values (?, ?)", [i, _f32(v)])
assert exec(db, "select count(*) as cnt from v_rowids") == snapshot(
name="rowids_count"
)
assert exec(db, "select count(*) as cnt from v_vector_chunks00") == snapshot(
name="vector_chunks_count"
)
# Verify round-trip: each inserted vector comes back identical
for i, v in enumerate(vecs, start=1):
rows = db.execute("select emb from v where rowid = ?", [i]).fetchall()
assert len(rows) == 1
assert rows[0][0] == _f32(v)
def test_insert_auto_rowid(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
vecs = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]
for v in vecs:
db.execute("insert into v(emb) values (?)", [_f32(v)])
rows = db.execute("select rowid from v order by rowid").fetchall()
rowids = [r[0] for r in rows]
assert rowids == [1, 2, 3]
for i, v in enumerate(vecs, start=1):
row = db.execute("select emb from v where rowid = ?", [i]).fetchone()
assert row[0] == _f32(v)
def test_insert_text_primary_key(db, snapshot):
db.execute(
"create virtual table v using vec0(id text primary key, emb float[4], chunk_size=8)"
)
db.execute(
"insert into v(id, emb) values ('doc_a', ?)", [_f32([1.0, 2.0, 3.0, 4.0])]
)
db.execute(
"insert into v(id, emb) values ('doc_b', ?)", [_f32([5.0, 6.0, 7.0, 8.0])]
)
assert exec(db, "select rowid, id, chunk_id, chunk_offset from v_rowids order by rowid") == snapshot(
name="rowids"
)
row = db.execute("select emb from v where id = 'doc_a'").fetchone()
assert row[0] == _f32([1.0, 2.0, 3.0, 4.0])
row = db.execute("select emb from v where id = 'doc_b'").fetchone()
assert row[0] == _f32([5.0, 6.0, 7.0, 8.0])
def test_delete_clears_validity(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
for i, v in enumerate(
[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]],
start=1,
):
db.execute("insert into v(rowid, emb) values (?, ?)", [i, _f32(v)])
db.execute("delete from v where rowid = 2")
rows = db.execute("select rowid from v order by rowid").fetchall()
assert [r[0] for r in rows] == [1, 3]
rowid_rows = db.execute("select rowid from v_rowids order by rowid").fetchall()
assert 2 not in [r[0] for r in rowid_rows]
# Inserting a new vector after deletion still works
db.execute("insert into v(rowid, emb) values (4, ?)", [_f32([0.0, 0.0, 0.0, 1.0])])
row = db.execute("select emb from v where rowid = 4").fetchone()
assert row[0] == _f32([0.0, 0.0, 0.0, 1.0])
def test_insert_delete_reinsert(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
db.execute("insert into v(rowid, emb) values (1, ?)", [_f32([1.0, 1.0, 1.0, 1.0])])
db.execute("delete from v where rowid = 1")
db.execute("insert into v(rowid, emb) values (2, ?)", [_f32([2.0, 2.0, 2.0, 2.0])])
rows = db.execute("select rowid from v order by rowid").fetchall()
assert [r[0] for r in rows] == [2]
# KNN query works and returns rowid 2
knn = db.execute(
"select rowid, distance from v where emb match ? and k = 1",
[_f32([2.0, 2.0, 2.0, 2.0])],
).fetchall()
assert len(knn) == 1
assert knn[0][0] == 2
def test_insert_validates_dimensions(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
result = exec(db, "insert into v(rowid, emb) values (1, ?)", [_f32([1.0, 2.0, 3.0])])
assert result["error"] == "OperationalError"
assert "Dimension mismatch" in result["message"]
assert "Expected 4" in result["message"]
assert "3" in result["message"]
result = exec(
db, "insert into v(rowid, emb) values (1, ?)", [_f32([1.0, 2.0, 3.0, 4.0, 5.0])]
)
assert result["error"] == "OperationalError"
assert "Dimension mismatch" in result["message"]
assert "Expected 4" in result["message"]
assert "5" in result["message"]
def test_insert_validates_type(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
int8_vec = struct.pack("4b", 1, 2, 3, 4)
result = exec(
db,
"insert into v(rowid, emb) values (1, vec_int8(?))",
[int8_vec],
)
assert "error" in result
assert "float32" in result["message"]
assert "int8" in result["message"]
def test_info_table_contents(db, snapshot):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
assert exec(db, "select key, value from v_info where key not like 'CREATE_VERSION%' order by key") == snapshot()
# Smoke-check that version keys exist without pinning exact values
version_rows = exec(db, "select key from v_info where key like 'CREATE_VERSION%' order by key")
keys = [r["key"] for r in version_rows["rows"]]
assert keys == ["CREATE_VERSION", "CREATE_VERSION_MAJOR", "CREATE_VERSION_MINOR", "CREATE_VERSION_PATCH"]
def test_delete_zeroes_rowid_blob(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
for i in range(1, 4):
db.execute(
"insert into v(rowid, emb) values (?, ?)",
[i, _f32([float(i)] * 4)],
)
db.execute("delete from v where rowid = 2")
blob = db.execute("select rowids from v_chunks where rowid = 1").fetchone()[0]
rowids = struct.unpack("<8q", blob)
assert rowids[0] == 1 # slot 0 intact
assert rowids[1] == 0 # slot 1 zeroed (was rowid 2)
assert rowids[2] == 3 # slot 2 intact
def test_delete_zeroes_vector_blob(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
db.execute(
"insert into v(rowid, emb) values (1, ?)", [_f32([1.0, 2.0, 3.0, 4.0])]
)
db.execute(
"insert into v(rowid, emb) values (2, ?)", [_f32([5.0, 6.0, 7.0, 8.0])]
)
db.execute("delete from v where rowid = 1")
blob = db.execute(
"select vectors from v_vector_chunks00 where rowid = 1"
).fetchone()[0]
# First slot (4 floats = 16 bytes) should be zeroed
first_slot = struct.unpack("<4f", blob[:16])
assert first_slot == (0.0, 0.0, 0.0, 0.0)
# Second slot should be unchanged
second_slot = struct.unpack("<4f", blob[16:32])
assert second_slot == (5.0, 6.0, 7.0, 8.0)
def test_delete_all_rows_deletes_chunk(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
for i in range(1, 9):
db.execute(
"insert into v(rowid, emb) values (?, ?)",
[i, _f32([float(i)] * 4)],
)
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
assert (
db.execute("select count(*) from v_chunks").fetchone()[0] == 0
)
assert (
db.execute("select count(*) from v_vector_chunks00").fetchone()[0] == 0
)
# Inserting after full deletion still works
db.execute(
"insert into v(rowid, emb) values (100, ?)", [_f32([9.0, 9.0, 9.0, 9.0])]
)
row = db.execute("select emb from v where rowid = 100").fetchone()
assert row[0] == _f32([9.0, 9.0, 9.0, 9.0])
def test_delete_chunk_multiple_chunks(db):
db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")
for i in range(1, 17):
db.execute(
"insert into v(rowid, emb) values (?, ?)",
[i, _f32([float(i)] * 4)],
)
# Delete all rows from the first chunk (rows 1-8)
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
# Only 1 chunk should remain
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 1
# Rows 9-16 still queryable
for i in range(9, 17):
row = db.execute("select emb from v where rowid = ?", [i]).fetchone()
assert row[0] == _f32([float(i)] * 4)
def test_delete_with_metadata_columns(db):
db.execute(
"create virtual table v using vec0("
"emb float[4], "
"m_bool boolean, "
"m_int integer, "
"m_float float, "
"m_text text, "
"chunk_size=8"
")"
)
for i in range(1, 9):
db.execute(
"insert into v(rowid, emb, m_bool, m_int, m_float, m_text) "
"values (?, ?, ?, ?, ?, ?)",
[i, _f32([float(i)] * 4), i % 2 == 0, i * 10, float(i) / 2.0, f"text_{i}"],
)
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 0
assert db.execute("select count(*) from v_vector_chunks00").fetchone()[0] == 0
assert db.execute("select count(*) from v_metadatachunks00").fetchone()[0] == 0
assert db.execute("select count(*) from v_metadatachunks01").fetchone()[0] == 0
assert db.execute("select count(*) from v_metadatachunks02").fetchone()[0] == 0
assert db.execute("select count(*) from v_metadatachunks03").fetchone()[0] == 0
def test_delete_with_auxiliary_columns(db):
db.execute(
"create virtual table v using vec0("
"emb float[4], "
"+aux_text text, "
"chunk_size=8"
")"
)
for i in range(1, 9):
db.execute(
"insert into v(rowid, emb, aux_text) values (?, ?, ?)",
[i, _f32([float(i)] * 4), f"aux_{i}"],
)
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 0
assert db.execute("select count(*) from v_auxiliary").fetchone()[0] == 0
def test_delete_with_text_primary_key(db):
db.execute(
"create virtual table v using vec0("
"id text primary key, emb float[4], chunk_size=8"
")"
)
db.execute(
"insert into v(id, emb) values ('a', ?)", [_f32([1.0, 2.0, 3.0, 4.0])]
)
db.execute(
"insert into v(id, emb) values ('b', ?)", [_f32([5.0, 6.0, 7.0, 8.0])]
)
db.execute("delete from v where id = 'a'")
# Vector blob slot 0 should be zeroed
blob = db.execute(
"select vectors from v_vector_chunks00 where rowid = 1"
).fetchone()[0]
first_slot = struct.unpack("<4f", blob[:16])
assert first_slot == (0.0, 0.0, 0.0, 0.0)
# Remaining row still queryable
row = db.execute("select emb from v where id = 'b'").fetchone()
assert row[0] == _f32([5.0, 6.0, 7.0, 8.0])
def test_delete_with_partition_keys(db):
db.execute(
"create virtual table v using vec0("
"part text partition key, emb float[4], chunk_size=8"
")"
)
for i in range(1, 9):
db.execute(
"insert into v(rowid, part, emb) values (?, 'A', ?)",
[i, _f32([float(i)] * 4)],
)
for i in range(9, 17):
db.execute(
"insert into v(rowid, part, emb) values (?, 'B', ?)",
[i, _f32([float(i)] * 4)],
)
# Delete all from partition A
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
# 1 chunk should remain (partition B's)
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 1
# Partition B rows intact
for i in range(9, 17):
row = db.execute("select emb from v where rowid = ?", [i]).fetchone()
assert row[0] == _f32([float(i)] * 4)
# Re-insert into partition A works
db.execute(
"insert into v(rowid, part, emb) values (100, 'A', ?)",
[_f32([99.0, 99.0, 99.0, 99.0])],
)
row = db.execute("select emb from v where rowid = 100").fetchone()
assert row[0] == _f32([99.0, 99.0, 99.0, 99.0])
def test_delete_int8_vectors(db):
db.execute("create virtual table v using vec0(emb int8[4], chunk_size=8)")
db.execute(
"insert into v(rowid, emb) values (1, vec_int8(?))",
[_int8([1, 2, 3, 4])],
)
db.execute(
"insert into v(rowid, emb) values (2, vec_int8(?))",
[_int8([5, 6, 7, 8])],
)
db.execute("delete from v where rowid = 1")
blob = db.execute(
"select vectors from v_vector_chunks00 where rowid = 1"
).fetchone()[0]
# int8[4] = 4 bytes per slot
first_slot = struct.unpack("<4b", blob[:4])
assert first_slot == (0, 0, 0, 0)
second_slot = struct.unpack("<4b", blob[4:8])
assert second_slot == (5, 6, 7, 8)
def test_delete_bit_vectors(db):
db.execute("create virtual table v using vec0(emb bit[8], chunk_size=8)")
db.execute(
"insert into v(rowid, emb) values (1, vec_bit(?))",
[bytes([0xFF])],
)
db.execute(
"insert into v(rowid, emb) values (2, vec_bit(?))",
[bytes([0xAA])],
)
db.execute("delete from v where rowid = 1")
blob = db.execute(
"select vectors from v_vector_chunks00 where rowid = 1"
).fetchone()[0]
# bit[8] = 1 byte per slot
assert blob[0:1] == bytes([0x00])
assert blob[1:2] == bytes([0xAA])
def _file_db(tmp_path):
"""Open a file-backed DB (required for page_count to shrink after VACUUM)."""
db = sqlite3.connect(str(tmp_path / "test.db"))
db.row_factory = sqlite3.Row
db.enable_load_extension(True)
db.load_extension("dist/vec0")
db.enable_load_extension(False)
return db
def test_delete_chunk_shrinks_pages(tmp_path):
"""Use large vectors (float[256]) so each chunk blob spans multiple pages,
making the page_count difference measurable after VACUUM."""
dims = 256
db = _file_db(tmp_path)
db.execute(f"create virtual table v using vec0(emb float[{dims}], chunk_size=8)")
for i in range(1, 25): # 3 full chunks of 8
db.execute(
"insert into v(rowid, emb) values (?, ?)",
[i, _f32([float(i)] * dims)],
)
db.commit()
pages_before = db.execute("pragma page_count").fetchone()[0]
# Delete all rows
for i in range(1, 25):
db.execute("delete from v where rowid = ?", [i])
db.commit()
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 0
db.execute("vacuum")
pages_after = db.execute("pragma page_count").fetchone()[0]
assert pages_after < pages_before, (
f"page_count should shrink after deleting all chunks and vacuum: "
f"{pages_before} -> {pages_after}"
)
db.close()
def test_delete_one_chunk_of_two_shrinks_pages(tmp_path):
"""Use large vectors (float[256]) so each chunk blob spans multiple pages,
making the page_count difference measurable after VACUUM."""
dims = 256
db = _file_db(tmp_path)
db.execute(f"create virtual table v using vec0(emb float[{dims}], chunk_size=8)")
for i in range(1, 17): # 2 full chunks of 8
db.execute(
"insert into v(rowid, emb) values (?, ?)",
[i, _f32([float(i)] * dims)],
)
db.commit()
pages_before = db.execute("pragma page_count").fetchone()[0]
# Delete all rows from the first chunk (rows 1-8)
for i in range(1, 9):
db.execute("delete from v where rowid = ?", [i])
db.commit()
assert db.execute("select count(*) from v_chunks").fetchone()[0] == 1
db.execute("vacuum")
pages_after = db.execute("pragma page_count").fetchone()[0]
assert pages_after < pages_before, (
f"page_count should shrink after deleting one chunk and vacuum: "
f"{pages_before} -> {pages_after}"
)
# Remaining rows still queryable after vacuum
for i in range(9, 17):
row = db.execute("select emb from v where rowid = ?", [i]).fetchone()
assert row[0] == _f32([float(i)] * dims)
db.close()