Embedding vectors#

AgensGraph stores and searches embeddings through pgvector. The driver adds a value type for each of pgvector’s types, helpers that build an index and the search that uses it from the same arguments, and a way to tune a search without leaking the setting to the next user of a pooled connection.

Setting up#

pgvector is an extension, created per database:

CREATE EXTENSION vector;

Then, on each connection, tell the driver to read and write pgvector’s types:

>>> conn.has_vectors()
True
>>> conn.vector_version()
(0, 8, 1)
>>> conn.register_vectors()
('vector', 'halfvec', 'sparsevec')

register_vectors() returns the types it found, and () in a database without the extension. The types are registered on that connection only, because their ids differ between databases. A pool does it on every connection it opens when you pass configure=lambda conn: conn.register_vectors().

halfvec and sparsevec need pgvector 0.7 or later, and the iterative index scans among the search settings 0.8 or later.

Two places to keep an embedding#

In the property map, like any other property. Nothing has to be declared, and the example dataset keeps its films’ four-number embeddings this way. A property in the map is jsonb, so reading it gives a list, and searching or indexing it casts it to a vector of a stated dimension:

>>> conn.execute_query("MATCH (m:movie {title: 'Inception'}) RETURN m.embedding").records
[([0.03923657, 0.39284106, -0.20927092, -0.17770818],)]
>>> conn.execute_query(
...     "MATCH (m:movie {title: 'Inception'}) RETURN m.embedding::vector(4)"
... ).records
[(Vector(4 dimensions),)]

In a column of its own, from AgensGraph 2.18. The column has a vector type with a dimension, so the server refuses a value of the wrong length when it is written, rather than at search time, and reading it never touches the rest of the property map:

>>> from agensgraph.vector import generated_column
>>> generated_column("embedding", 4)
'embedding vector(4) generated'
>>> conn.execute(f"CREATE VLABEL doc ({generated_column('embedding', 4)})")
<agensgraph.connection.Cursor ...>
>>> conn.refresh_labels()

This is the better choice for a label whose elements all carry an embedding. ensure_labels() declares the same column with PromotedProperty("embedding", "vector(4)").

Writing a vector#

Pass a Vector:

>>> from agensgraph import Vector
>>> conn.execute_query(
...     "CREATE (:doc {title: 'a', embedding: %s})", (Vector([0.1, 0.2, 0.3, 0.4]),)
... )
Result(records=[], ...)
>>> conn.execute_query("CREATE (:doc {title: 'b', embedding: %s})", (Vector([0.1, 0.2]),))
Traceback (most recent call last):
  ...
psycopg.errors.DataException: expected 4 dimensions, not 2
>>> conn.rollback()

A Vector goes over the wire as its numbers, in binary: 1536 dimensions are about 6 KB, against about 17 KB as decimal text, and the driver’s own measurement puts one write of a 1536-dimension vector at 0.32 milliseconds as a Vector against 2.55 as a list. A Vector built from a list, a tuple or a NumPy array is single precision, as pgvector stores it.

The bulk loaders are the exception: they write property maps as JSON, and a Vector is not JSON, so load embeddings in bulk as plain lists. Written into a column of its own, a list reads back as a Vector all the same.

Reading a vector#

A vector column, or a property cast to a vector, comes back as a Vector. Ask for the composite rendering when you read many: a vector’s text is long decimal, and its composite form is the numbers themselves, which reads several times faster.

>>> (v,) = conn.execute_query(
...     "MATCH (m:movie {title: 'Inception'}) RETURN m.embedding::vector(4)", binary_=True
... ).records[0]
>>> v
Vector(4 dimensions)
>>> len(v), v[0]
(4, 0.03923657163977623)
>>> v.values
array('f', [0.03923657163977623, 0.39284107089042664, -0.2092709243297577, -0.17770817875862122])

len(v)

the dimension; for a vector read in the composite rendering, from its header, without unpacking the numbers

v[i], iterating

the numbers, as Python floats

v[i:j]

an array('f')

v.values

an array('f'), which NumPy and PyTorch take without copying

v.tolist()

a list of floats

A Vector keeps the bytes it arrived in and turns them into numbers only when something reads them. A search asks the server for the distance, so the vectors of the elements it ranked are often never read, and then never cost anything to unpack.

It is not a list: isinstance(v, list) is false, and json.dumps(v) raises. It compares equal to a sequence of the same numbers. Its own numbers are single precision and the other sequence’s are compared as they are, so Vector([1.0, 2.0]) == [1.0, 2.0], but Vector([0.1]) != [0.1], because 0.1 stored in single precision is not quite 0.1. to_builtins() and to_json() write it as a list.

halfvec reads as a Vector too, at half precision.

Indexing#

An index serves a search only when the search is written exactly the way the index is: the same cast, with the same dimension, and an operator the index’s operator class serves. A search that differs in either way is not refused; it quietly sorts every element instead. vector_index() and nearest() build the index and the search from the same arguments, so they agree:

>>> from agensgraph import Distance
>>> from agensgraph.vector import nearest, vector_index
>>> vector_index("movie", "embedding", dimensions=4,
...              operator_class=Distance.COSINE.operator_class)
'create property index on movie using hnsw ((embedding::vector(4)) vector_cosine_ops)'
>>> nearest("movie", "embedding", dimensions=4, operator=Distance.COSINE, limit=3)
'match (n:movie) return n order by n.embedding::vector(4) <=> %s::vector(4) limit 3'

For a property in the map, pass dimensions: the index is over the cast. For a column of its own, leave it out: the column is indexed directly, and an index over a cast of it would never be used by a search on the column.

>>> vector_index("doc", "embedding", operator_class=Distance.L2.operator_class)
'create property index on doc using hnsw ((embedding) vector_l2_ops)'

method is "hnsw" (the default) or "ivfflat", name= names the index, and options= passes the method’s settings, such as options={"m": 16, "ef_construction": 64} or options={"lists": 100}. The result is a property index, the form a Cypher query matches; a SQL index on the same expression over the label’s table is never used by a Cypher search.

An index can also be declared and reconciled like any other.

Distances#

Distance

Operator

Operator class

L2

<->

vector_l2_ops

Euclidean distance

INNER_PRODUCT

<#>

vector_ip_ops

the negated inner product, so that smaller is nearer

COSINE

<=>

vector_cosine_ops

cosine distance

L1

<+>

vector_l1_ops

taxicab distance

HAMMING

<~>

bit_hamming_ops

for bit strings

JACCARD

<%>

bit_jaccard_ops

for bit strings

A Distance is a string whose value is the operator, so it goes straight into operator= or into a statement you write. .operator_class names the class for the vector type; for halfvec or sparsevec, name the class yourself, such as halfvec_cosine_ops.

Searching#

>>> conn.execute(vector_index("movie", "embedding", dimensions=4,
...                           operator_class=Distance.COSINE.operator_class))
<agensgraph.connection.Cursor ...>
>>> query = Vector([0.03923657, 0.39284106, -0.20927092, -0.17770818])
>>> search = nearest("movie", "embedding", dimensions=4, operator=Distance.COSINE, limit=3)
>>> rows = conn.execute_query(search, (query,)).records
>>> [movie.properties["title"] for (movie,) in rows]
['Inception', 'Minority Report', 'Total Recall']

nearest() returns the whole vertex. Write the statement yourself for anything else, such as the distance, a filter, or a traversal from the matches; the one rule is to cast the same way the index does:

>>> conn.execute_query(
...     "MATCH (m:movie) RETURN m.title, m.embedding::vector(4) <=> %s::vector(4) AS distance"
...     " ORDER BY distance LIMIT 3",
...     (query,),
... ).records
[('Inception', 0), ('Minority Report', 0.24371635018387505), ('Total Recall', 0.33612068966778275)]

The query vector is a parameter like any other: a Vector, or a list of floats.

Reading the indexes back#

>>> conn.vector_indexes()
[VectorIndex(label='movie', name='movie_embedding_idx', property='embedding', method='hnsw', operator_class='vector_cosine_ops', type='vector', dimensions=4, options={})]

Each index is a VectorIndex, read in parts: the label, the index’s name, the property, the method, the operator class (None when it is the method’s default, which the server leaves out), the vector type, the dimension, and the options it was built with, as strings.

Sparse vectors#

A SparseVector holds only its non-zero entries, which is what a sparse embedding is for: three non-zeros in a million dimensions are a few dozen bytes, not megabytes of zeros.

>>> from agensgraph import SparseVector
>>> s = SparseVector({0: 1.0, 3: 2.0}, 6)
>>> s.indices, s.values, s.dimensions
((0, 3), (1.0, 2.0), 6)
>>> s.to_dense()
[1.0, 0.0, 0.0, 2.0, 0.0, 0.0]
>>> conn.execute_query("SELECT %s::sparsevec", (s,)).records
[(SparseVector({0: 1.0, 3: 2.0}, 6),)]

SparseVector(entries, dimensions) takes a mapping or pairs of index and value. Indices count from zero in Python, everywhere, while pgvector’s text form counts from one: the conversion happens in one place, so an index you give is the index you get back. len(s) is the number of stored entries. What the server would refuse, the constructor refuses: an index outside the dimensions, a repeated index, a non-finite value, a dimension below one. An explicit zero is dropped, as the server drops it. Indices and values must be plain Python numbers; convert NumPy scalars with int() and float().

A sparse vector is sent as text, so in SQL write %s::sparsevec. Into a column of its own it goes directly. SparseVector.from_dense() builds one from a list.