Skip to content

vector

Code specific to the pgvector extension.

The base vector column is still defined in embar.column.pg

CosineDistance

Bases: ClauseBase

Get the Cosine Distance using pgvector.

Assumes pgvector extension is installed and activated.

Creates a query like col_a <=> '[1,2,3]' or col_a <=> col_b.

Source code in src/embar/query/vector.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class CosineDistance(ClauseBase):
    """
    Get the Cosine Distance using pgvector.

    Assumes pgvector extension is installed and activated.

    Creates a query like col_a <=> '[1,2,3]' or col_a <=> col_b.
    """

    left: ColumnInfo
    right: list[float] | ColumnInfo

    def __init__(self, left: Column[list[float]], right: list[float] | Column[list[float]]):
        self.left = left.info
        self.right = right.info if isinstance(right, Column) else right

    @override
    def sql(self, get_count: GetCount) -> QuerySingle:
        count = get_count()
        name = f"l2distance_{self.left.name}_{count}"
        if isinstance(self.right, ColumnInfo):
            return QuerySingle(sql=f"{self.left.fqn()} <=> {self.right.fqn()}")

        # pgvector expects an argument of the form '[1,2,3]'
        stringified = str(self.right).replace(" ", "")

        return QuerySingle(sql=f"{self.left.fqn()} <=> %({name})s", params={name: stringified})

L2Distance

Bases: ClauseBase

Get the L2 Distance using pgvector.

Assumes pgvector extension is installed and activated.

Creates a query like col_a <-> '[1,2,3]' or col_a <-> col_b.

Source code in src/embar/query/vector.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class L2Distance(ClauseBase):
    """
    Get the L2 Distance using pgvector.

    Assumes pgvector extension is installed and activated.

    Creates a query like col_a <-> '[1,2,3]' or col_a <-> col_b.
    """

    left: ColumnInfo
    right: list[float] | ColumnInfo

    def __init__(self, left: Column[list[float]], right: list[float] | Column[list[float]]):
        self.left = left.info
        self.right = right.info if isinstance(right, Column) else right

    @override
    def sql(self, get_count: GetCount) -> QuerySingle:
        count = get_count()
        name = f"l2distance_{self.left.name}_{count}"
        if isinstance(self.right, ColumnInfo):
            return QuerySingle(sql=f"{self.left.fqn()} <-> {self.right.fqn()}")

        # pgvector expects an argument of the form '[1,2,3]'
        stringified = str(self.right).replace(" ", "")

        return QuerySingle(sql=f"{self.left.fqn()} <-> %({name})s", params={name: stringified})