Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 116 additions & 13 deletions pgvector/django/functions.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,32 @@
from django.db.models import FloatField, Func, Value
from .. import Vector, HalfVector, SparseVector
from typing import Any
from typing import Any, ClassVar


class DistanceBase(Func):
"""Base class for pgvector distance functions that operate on vector columns.

Accepts either a Django ORM expression (any object with a
``resolve_expression`` method) or a raw vector value as ``list[float]``,
:class:`~pgvector.Vector`, :class:`~pgvector.HalfVector`, or
:class:`~pgvector.SparseVector`. Raw values are converted to their
PostgreSQL text representation before being passed to the query.
"""

output_field = FloatField() # type: ignore

def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None:
"""Initialize a distance expression between *expression* and *vector*.

Args:
expression: A Django ORM expression referencing a vector column.
vector: The comparison vector. May be a ``list[float]``,
:class:`~pgvector.Vector`, :class:`~pgvector.HalfVector`,
:class:`~pgvector.SparseVector`, or any Django expression
that resolves to a compatible vector type.
**extra: Additional keyword arguments forwarded to
:class:`django.db.models.Func`.
"""
if not hasattr(vector, 'resolve_expression'):
if isinstance(vector, (Vector, HalfVector, SparseVector)):
vector = Value(vector.to_text())
Expand All @@ -20,39 +40,122 @@ def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None:


class BitDistanceBase(Func):
"""Base class for pgvector distance functions that operate on bit columns.

Unlike :class:`DistanceBase`, the comparison value is used as-is (a plain
Python value that Django wraps in a :class:`~django.db.models.Value`).
"""

output_field = FloatField() # type: ignore

def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None:
"""Initialize a bit-distance expression between *expression* and *vector*.

Args:
expression: A Django ORM expression referencing a bit column.
vector: The comparison bit string or any Django expression that
resolves to a compatible bit type.
**extra: Additional keyword arguments forwarded to
:class:`django.db.models.Func`.
"""
if not hasattr(vector, 'resolve_expression'):
vector = Value(vector)
super().__init__(expression, vector, **extra)


class L2Distance(DistanceBase):
function = ''
arg_joiner = ' <-> '
"""L2 (Euclidean) distance via the pgvector ``<->`` operator.

Returns the straight-line distance between two vectors in Euclidean space.
Equivalent to ``sqrt(sum((a_i - b_i)^2))``. Lower values indicate
greater similarity.

Example::

Item.objects.annotate(dist=L2Distance('embedding', [1, 2, 3])).order_by('dist')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <-> '


class MaxInnerProduct(DistanceBase):
function = ''
arg_joiner = ' <#> '
"""Negative inner product via the pgvector ``<#>`` operator.

Returns the *negated* dot product of two vectors so that ORDER BY gives
the highest-similarity results first (pgvector stores it negated to
support ``ORDER BY ASC``). Equivalent to ``-sum(a_i * b_i)``.

Use this instead of :class:`CosineDistance` when vectors are already
unit-normalised (cosine similarity and inner product are equivalent in
that case, but inner product is faster).

Example::

Item.objects.annotate(score=MaxInnerProduct('embedding', [1, 2, 3])).order_by('score')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <#> '


class CosineDistance(DistanceBase):
function = ''
arg_joiner = ' <=> '
"""Cosine distance via the pgvector ``<=>`` operator.

Returns ``1 - cosine_similarity(a, b)``, ranging from 0 (identical
direction) to 2 (opposite direction). Magnitude-invariant, making it
the typical choice for semantic-embedding similarity.

Example::

Item.objects.annotate(dist=CosineDistance('embedding', [1, 2, 3])).order_by('dist')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <=> '


class L1Distance(DistanceBase):
function = ''
arg_joiner = ' <+> '
"""L1 (Manhattan / taxicab) distance via the pgvector ``<+>`` operator.

Returns ``sum(|a_i - b_i|)``. More robust than L2 distance in
high-dimensional spaces because it penalises large component differences
less aggressively.

Example::

Item.objects.annotate(dist=L1Distance('embedding', [1, 2, 3])).order_by('dist')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <+> '


class HammingDistance(BitDistanceBase):
function = ''
arg_joiner = ' <~> '
"""Hamming distance between two bit strings via the pgvector ``<~>`` operator.

Counts the number of bit positions at which the two values differ.
Operates on ``BIT(n)`` columns (use :class:`~pgvector.django.bit.BitField`).

Example::

Item.objects.annotate(dist=HammingDistance('signature', b'\\x01')).order_by('dist')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <~> '


class JaccardDistance(BitDistanceBase):
function = ''
arg_joiner = ' <%%> '
"""Jaccard distance between two bit strings via the pgvector ``<%%>`` operator.

Returns ``1 - |A ∩ B| / |A ∪ B|``, i.e. 0 for identical bit strings and
1 for disjoint ones. Operates on ``BIT(n)`` columns.

Example::

Item.objects.annotate(dist=JaccardDistance('signature', b'\\x01')).order_by('dist')
"""

function: ClassVar[str] = ''
arg_joiner: ClassVar[str] = ' <%%> '
155 changes: 154 additions & 1 deletion pgvector/django/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,34 +3,181 @@


class IvfflatIndex(PostgresIndex):
"""Django index class for the pgvector IVFFlat (Inverted File Flat) index type.

IVFFlat partitions vectors into *lists* (Voronoi cells) at build time and
restricts each query to the *probes* nearest lists. It offers fast approximate
nearest-neighbour search with a controllable recall/speed trade-off.

Build-time cost is proportional to ``lists``; query-time cost is proportional to
``probes`` (a session-level GUC, ``ivfflat.probes``, not an index parameter).
The index **requires** the table to contain data before it is created — an empty
table produces degenerate clusters.

Supported distance operators (set via ``opclasses``):

- ``vector_l2_ops`` — L2 / Euclidean distance (``<->``)
- ``vector_ip_ops`` — inner product (``<#>``)
- ``vector_cosine_ops`` — cosine distance (``<=>``)
- ``halfvec_l2_ops`` / ``halfvec_ip_ops`` / ``halfvec_cosine_ops`` — half-precision variants
- ``sparsevec_l2_ops`` / ``sparsevec_ip_ops`` / ``sparsevec_cosine_ops`` — sparse-vector variants

Example::

from pgvector.django import IvfflatIndex

class Meta:
indexes = [
IvfflatIndex(
name='item_embedding_ivfflat',
fields=['embedding'],
lists=100,
opclasses=['vector_cosine_ops'],
)
]

See also: https://github.com/pgvector/pgvector#ivfflat
"""

suffix = 'ivfflat'

def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) -> None:
"""Create an IVFFlat index.

Args:
*expressions: Column expressions to index (forwarded to
:class:`~django.contrib.postgres.indexes.PostgresIndex`).
lists: Number of lists (Voronoi cells) to partition the vector
space into during index construction. Controls the
accuracy/speed trade-off at *build* time.

- Higher values improve recall but increase build time and
memory consumption.
- As a rule of thumb pgvector recommends ``rows / 1000`` for
up to 1 million rows, and ``sqrt(rows)`` beyond that.
- Valid range: 1 – 32768. pgvector default: ``100``.
- ``None`` omits the parameter and lets the server use its
compiled-in default.
**kwargs: Additional keyword arguments forwarded to
:class:`~django.contrib.postgres.indexes.PostgresIndex`
(e.g. ``name``, ``fields``, ``opclasses``, ``condition``).
"""
self.lists = lists
super().__init__(*expressions, **kwargs)

def deconstruct(self) -> tuple[Any, Any, Any]:
"""Return enough information to recreate the index via migrations.

Returns:
A ``(path, args, kwargs)`` tuple as required by Django's
deconstruction protocol. ``lists`` is included in *kwargs* only
when it was explicitly set.
"""
path, args, kwargs = super().deconstruct()
if self.lists is not None:
kwargs['lists'] = self.lists
return path, args, kwargs

def get_with_params(self) -> list[str]:
"""Return the ``WITH (...)`` storage parameters for the ``CREATE INDEX`` DDL.

Returns:
A list of ``'key = value'`` strings passed to the PostgreSQL
``WITH`` clause. Empty list when no parameters were specified.
"""
with_params = []
if self.lists is not None:
with_params.append('lists = %d' % self.lists)
return with_params


class HnswIndex(PostgresIndex):
"""Django index class for the pgvector HNSW (Hierarchical Navigable Small World) index type.

HNSW builds a multi-layer proximity graph over the vectors. It generally
provides higher recall and faster queries than :class:`IvfflatIndex` at the
cost of more memory and a longer build time. Unlike IVFFlat it can be built
on an empty table and updated incrementally.

Query-time recall is tuned with ``hnsw.ef_search`` (a session-level GUC),
not an index parameter.

Supported distance operators (set via ``opclasses``):

- ``vector_l2_ops`` — L2 / Euclidean distance (``<->``)
- ``vector_ip_ops`` — inner product (``<#>``)
- ``vector_cosine_ops`` — cosine distance (``<=>``)
- ``vector_l1_ops`` — L1 / Manhattan distance (``<+>``)
- ``halfvec_*`` / ``sparsevec_*`` variants are also available

Example::

from pgvector.django import HnswIndex

class Meta:
indexes = [
HnswIndex(
name='item_embedding_hnsw',
fields=['embedding'],
m=16,
ef_construction=64,
opclasses=['vector_cosine_ops'],
)
]

See also: https://github.com/pgvector/pgvector#hnsw
"""

suffix = 'hnsw'

def __init__(self, *expressions: Any, m: int | None = None, ef_construction: int | None = None, **kwargs: Any) -> None:
def __init__(
self,
*expressions: Any,
m: int | None = None,
ef_construction: int | None = None,
**kwargs: Any,
) -> None:
"""Create an HNSW index.

Args:
*expressions: Column expressions to index (forwarded to
:class:`~django.contrib.postgres.indexes.PostgresIndex`).
m: Maximum number of bidirectional links (edges) created for each
new node during graph construction. Controls the graph
connectivity at *build* time.

- Higher values yield better recall and faster queries but
consume more memory (roughly ``m * 8`` bytes per vector per
layer beyond the base layer).
- Valid range: 2 – 100. pgvector default: ``16``.
- ``None`` omits the parameter and lets the server use its
compiled-in default.
ef_construction: Size of the dynamic candidate list maintained
while inserting a new vector into the graph. Controls the
quality of the *build*-time graph structure.

- Higher values produce a denser, better-connected graph with
improved recall, at the cost of a slower build.
- Must satisfy ``ef_construction >= 2 * m``.
- Valid range: 4 – 1000. pgvector default: ``64``.
- ``None`` omits the parameter and lets the server use its
compiled-in default.
**kwargs: Additional keyword arguments forwarded to
:class:`~django.contrib.postgres.indexes.PostgresIndex`
(e.g. ``name``, ``fields``, ``opclasses``, ``condition``).
"""
self.m = m
self.ef_construction = ef_construction
super().__init__(*expressions, **kwargs)

def deconstruct(self) -> tuple[Any, Any, Any]:
"""Return enough information to recreate the index via migrations.

Returns:
A ``(path, args, kwargs)`` tuple as required by Django's
deconstruction protocol. ``m`` and ``ef_construction`` are
included in *kwargs* only when they were explicitly set.
"""
path, args, kwargs = super().deconstruct()
if self.m is not None:
kwargs['m'] = self.m
Expand All @@ -39,6 +186,12 @@ def deconstruct(self) -> tuple[Any, Any, Any]:
return path, args, kwargs

def get_with_params(self) -> list[str]:
"""Return the ``WITH (...)`` storage parameters for the ``CREATE INDEX`` DDL.

Returns:
A list of ``'key = value'`` strings passed to the PostgreSQL
``WITH`` clause. Empty list when no parameters were specified.
"""
with_params = []
if self.m is not None:
with_params.append('m = %d' % self.m)
Expand Down