bench_code_vector_storage.py

"""One-shot shipped Q12 10,000-vector Turso calibration; the temporary database is deleted."""

from __future__ import annotations

import json
import tempfile
import time
from pathlib import Path

import turso

from tests._fixture_isolation import inherited_test_root


SAMPLE_VECTORS = 10_000
DIMENSIONS = 4096
BATCH_SIZE = 16
WORK_ROOT = inherited_test_root()


def _bytes(paths: list[Path]) -> int:
    return sum(path.stat().st_size for path in paths if path.is_file())


def main() -> int:
    vector = json.dumps([1.0, *([0.0] * (DIMENSIONS - 1))], separators=(",", ":"))
    with tempfile.TemporaryDirectory(prefix="q12-f32-calibration-", dir=WORK_ROOT) as raw:
        root = Path(raw)
        database = root / "fixture.db"
        connection = turso.connect(str(database))
        try:
            connection.execute(
                "CREATE TABLE file_vectors (profile_key INTEGER NOT NULL,file_id INTEGER NOT NULL,"
                "vector F32_BLOB(4096) NOT NULL,PRIMARY KEY(profile_key,file_id))"
            )
            connection.execute(
                "CREATE TABLE line_vectors (profile_key INTEGER NOT NULL,line_id INTEGER NOT NULL,"
                "vector F32_BLOB(4096) NOT NULL,PRIMARY KEY(profile_key,line_id))"
            )
            connection.execute("CREATE INDEX file_reverse ON file_vectors(file_id,profile_key)")
            connection.execute("CREATE INDEX line_reverse ON line_vectors(line_id,profile_key)")
            connection.commit()
            started = time.perf_counter()
            peak_companion = 0
            for start in range(0, SAMPLE_VECTORS, BATCH_SIZE):
                stop = min(SAMPLE_VECTORS, start + BATCH_SIZE)
                for identity in range(start + 1, stop + 1):
                    connection.execute(
                        "INSERT INTO line_vectors(profile_key,line_id,vector) VALUES (1,?,vector32(?))",
                        (identity, vector),
                    )
                connection.commit()
                peak_companion = max(
                    peak_companion,
                    _bytes([path for path in root.iterdir() if path != database]),
                )
            page_size = int(connection.execute("PRAGMA page_size").fetchone()[0])
            page_count = int(connection.execute("PRAGMA page_count").fetchone()[0])
            freelist_count = int(connection.execute("PRAGMA freelist_count").fetchone()[0])
            result = {
                "sample_vector_count": SAMPLE_VECTORS,
                "sample_line_vector_count": SAMPLE_VECTORS,
                "sample_file_vector_count": 0,
                "dimensions": DIMENSIONS,
                "raw_bytes_per_vector": DIMENSIONS * 4,
                "allocated_page_bytes": (page_count - freelist_count) * page_size,
                "database_file_bytes": database.stat().st_size,
                "peak_companion_bytes": peak_companion,
                "page_size": page_size,
                "page_count": page_count,
                "freelist_count": freelist_count,
                "elapsed_seconds": time.perf_counter() - started,
            }
            print(json.dumps(result, indent=2))
        finally:
            connection.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())