L1 distance support for sqlalchemy and sqlmodel (#69)

domenicocinque · web-flow · commit 62c69fca1b63 · 2024-05-16T11:54:41.000-04:00
diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py
@@ -39,6 +39,9 @@ class comparator_factory(UserDefinedType.Comparator):
         def l2_distance(self, other):
             return self.op('<->', return_type=Float)(other)
 
+        def l1_distance(self, other):
+            return self.op('<+>', return_type=Float)(other)
+
         def max_inner_product(self, other):
             return self.op('<#>', return_type=Float)(other)
 
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
@@ -119,6 +119,18 @@ def test_l2_distance_orm(self):
             items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1])))
             assert [v.id for v in items] == [1, 3, 2]
 
+    def test_l1_distance(self):
+        create_items()
+        with Session(engine) as session:
+            items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 2])).all()
+            assert [v.id for v in items] == [3, 1, 2]
+
+    def test_l1_distance_orm(self):
+        create_items()
+        with Session(engine) as session:
+            items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 2])))
+            assert [v.id for v in items] == [3, 1, 2]
+
     def test_max_inner_product(self):
         create_items()
         with Session(engine) as session:
diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py
@@ -79,6 +79,12 @@ def test_l2_distance(self):
             items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1])))
             assert [v.id for v in items] == [1, 3, 2]
 
+    def test_l1_distance(self):
+        create_items()
+        with Session(engine) as session:
+            items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1])))
+            assert [v.id for v in items] == [1, 3, 2]
+
     def test_max_inner_product(self):
         create_items()
         with Session(engine) as session:
