Added HnswIndex for Django

ankane · ankane · commit 6867e8b451ba · 2023-08-30T20:14:27.000-07:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -19,7 +19,7 @@ jobs:
           dev-files: true
       - run: |
           cd /tmp
-          git clone --branch v0.4.4 https://github.com/pgvector/pgvector.git
+          git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
           cd pgvector
           make
           sudo make install
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.2.2 (unreleased)
+
+- Added `HnswIndex` for Django
+
 ## 0.2.1 (2023-07-31)
 
 - Fixed form issues with Django
diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py
@@ -73,6 +73,31 @@ def get_with_params(self):
         return with_params
 
 
+class HnswIndex(PostgresIndex):
+    suffix = 'hnsw'
+
+    def __init__(self, *expressions, m=None, ef_construction=None, **kwargs):
+        self.m = m
+        self.ef_construction = ef_construction
+        super().__init__(*expressions, **kwargs)
+
+    def deconstruct(self):
+        path, args, kwargs = super().deconstruct()
+        if self.m is not None:
+            kwargs['m'] = self.m
+        if self.ef_construction is not None:
+            kwargs['ef_construction'] = self.ef_construction
+        return path, args, kwargs
+
+    def get_with_params(self):
+        with_params = []
+        if self.m is not None:
+            with_params.append('m = %d' % self.m)
+        if self.ef_construction is not None:
+            with_params.append('ef_construction = %d' % self.ef_construction)
+        return with_params
+
+
 class DistanceBase(Func):
     output_field = FloatField()
 
diff --git a/tests/test_django.py b/tests/test_django.py
@@ -7,7 +7,7 @@
 from math import sqrt
 import numpy as np
 import pgvector.django
-from pgvector.django import VectorExtension, VectorField, IvfflatIndex, L2Distance, MaxInnerProduct, CosineDistance
+from pgvector.django import VectorExtension, VectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance
 from unittest import mock
 
 settings.configure(
@@ -28,10 +28,17 @@ class Meta:
         app_label = 'myapp'
         indexes = [
             IvfflatIndex(
-                name='my_index',
+                name='ivfflat_idx',
                 fields=['embedding'],
                 lists=100,
                 opclasses=['vector_l2_ops']
+            ),
+            HnswIndex(
+                name='hnsw_idx',
+                fields=['embedding'],
+                m=16,
+                ef_construction=100,
+                opclasses=['vector_l2_ops']
             )
         ]
 
