Merge branch 'master' into feature/other_inputs_for_sqlalchemy_bit

Author: jackrua

.github/workflows/build.yml

@@ -19,7 +19,7 @@ jobs:
           dev-files: true
       - run: |
           cd /tmp
-          git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git
+          git clone --branch v0.8.1 https://github.com/pgvector/pgvector.git
           cd pgvector
           make
           sudo make install

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.4.2 (unreleased)
+
+- Added support for `str` objects for `bit` type with SQLAlchemy
+
 ## 0.4.1 (2025-04-26)
 
 - Fixed `SparseVector` constructor for SciPy sparse matrices

README.md

@@ -259,7 +259,6 @@ index = Index(
     'my_index',
     func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
     postgresql_using='hnsw',
-    postgresql_with={'m': 16, 'ef_construction': 64},
     postgresql_ops={'embedding': 'halfvec_l2_ops'}
 )
 ```
@@ -271,6 +270,37 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
 session.scalars(select(Item).order_by(order).limit(5))
 ```
 
+#### Binary Quantization
+
+Use expression indexing for binary quantization
+
+```python
+from pgvector.sqlalchemy import BIT
+from sqlalchemy.sql import func
+
+index = Index(
+    'my_index',
+    func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
+    postgresql_using='hnsw',
+    postgresql_ops={'embedding': 'bit_hamming_ops'}
+)
+```
+
+Get the nearest neighbors by Hamming distance
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+session.scalars(select(Item).order_by(order).limit(5))
+```
+
+Re-rank by the original vectors for better recall
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+subquery = session.query(Item).order_by(order).limit(20).subquery()
+session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
+```
+
 #### Arrays
 
 Add an array column

pgvector/sqlalchemy/bit.py

@@ -17,21 +17,16 @@ def get_col_spec(self, **kw):
         return 'BIT(%d)' % self.length
 
     def bind_processor(self, dialect):
-        def process(value):
-            value = Bit._to_db(value)
-            if value and isinstance(dialect, PGDialect_asyncpg): 
-                return asyncpg.BitString(value)
-            return value
-        return process
-    
-    def result_processor(self, dialect, coltype):
-        def process(value): 
-            if value is None: return None
-            else: 
-                if isinstance(dialect, PGDialect_asyncpg): 
-                    return value.as_string()
-            return Bit._from_db(value).to_text() 
-        return process
+        if dialect.__class__.__name__ == 'PGDialect_asyncpg':
+            import asyncpg
+
+            def process(value):
+                if isinstance(value, str):
+                    return asyncpg.BitString(value)
+                return value
+            return process
+        else:
+            return super().bind_processor(dialect)
 
     class comparator_factory(UserDefinedType.Comparator):
         def hamming_distance(self, other):

tests/test_sqlalchemy.py

@@ -103,7 +103,6 @@ class Item(Base):
     'sqlalchemy_orm_half_precision_index',
     func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
     postgresql_using='hnsw',
-    postgresql_with={'m': 16, 'ef_construction': 64},
     postgresql_ops={'embedding': 'halfvec_l2_ops'}
 )
 half_precision_index.create(setup_engine)
@@ -112,7 +111,6 @@ class Item(Base):
     'sqlalchemy_orm_binary_quantize_index',
     func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
     postgresql_using='hnsw',
-    postgresql_with={'m': 16, 'ef_construction': 64},
     postgresql_ops={'embedding': 'bit_hamming_ops'}
 )
 binary_quantize_index.create(setup_engine)
@@ -528,6 +526,22 @@ def test_binary_quantize(self, engine):
             items = session.query(Item).order_by(distance).all()
             assert [v.id for v in items] == [2, 3, 1]
 
+    def test_binary_quantize_reranking(self, engine):
+        # recreate index (could also vacuum table)
+        binary_quantize_index.drop(setup_engine)
+        binary_quantize_index.create(setup_engine)
+
+        with Session(engine) as session:
+            session.add(Item(id=1, embedding=[-1, -2, -3]))
+            session.add(Item(id=2, embedding=[1, -2, 3]))
+            session.add(Item(id=3, embedding=[1, 2, 3]))
+            session.commit()
+
+            distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+            subquery = session.query(Item).order_by(distance).limit(20).subquery()
+            items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all()
+            assert [v.id for v in items] == [2, 3, 1]
+
 
 @pytest.mark.parametrize('engine', array_engines)
 class TestSqlalchemyArray:
@@ -596,6 +610,11 @@ async def test_bit(self, engine):
                 item = await session.get(Item, 1)
                 assert item.binary_embedding == embedding
 
+                if engine == asyncpg_engine:
+                    session.add(Item(id=2, binary_embedding='101'))
+                    item = await session.get(Item, 2)
+                    assert item.binary_embedding == embedding
+
         await engine.dispose()
 
     @pytest.mark.asyncio