sqlite-vec/tests/test-knn-distance-constraints.py

import sqlite3
from helpers import exec


def test_normal(db, snapshot):
    db.execute("create virtual table v using vec0(embedding float[1], is_odd boolean, chunk_size=8)")
    db.executemany(
        "insert into v(rowid, is_odd, embedding) values (?1, ?1 % 2, ?2)",
        [
            [1, "[1]"],
            [2, "[2]"],
            [3, "[3]"],
            [4, "[4]"],
            [5, "[5]"],
            [6, "[6]"],
            [7, "[7]"],
            [8, "[8]"],
            [9, "[9]"],
            [10, "[10]"],
            [11, "[11]"],
            [12, "[12]"],
            [13, "[13]"],
            [14, "[14]"],
            [15, "[15]"],
            [16, "[16]"],
            [17, "[17]"],
        ],
    )
    assert exec(db,"SELECT * FROM v") == snapshot()

    BASE_KNN = "select rowid, distance from v where embedding match ? and k = ? "
    assert exec(db, BASE_KNN, ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance > 5", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance >= 5", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance < 3", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance <= 3", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance > 7 AND distance <= 10", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND distance BETWEEN 7 AND 10", ["[1]", 5]) == snapshot()
    assert exec(db, BASE_KNN + "AND is_odd == TRUE AND distance BETWEEN 7 AND 10", ["[1]", 5]) == snapshot()


class Row:
    def __init__(self):
        pass

    def __repr__(self) -> str:
        return repr()