Skip to content

Commit b2dafb6

Browse files
committed
Improved readme and tests
1 parent 69e1e92 commit b2dafb6

File tree

3 files changed

+51
-3
lines changed

3 files changed

+51
-3
lines changed

README.md

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,13 +182,17 @@ Also supports `sum`
182182
Add an approximate index
183183

184184
```python
185-
index = Index('my_index', Item.embedding,
185+
index = Index(
186+
'my_index',
187+
Item.embedding,
186188
postgresql_using='hnsw',
187189
postgresql_with={'m': 16, 'ef_construction': 64},
188190
postgresql_ops={'embedding': 'vector_l2_ops'}
189191
)
190192
# or
191-
index = Index('my_index', Item.embedding,
193+
index = Index(
194+
'my_index',
195+
Item.embedding,
192196
postgresql_using='ivfflat',
193197
postgresql_with={'lists': 100},
194198
postgresql_ops={'embedding': 'vector_l2_ops'}
@@ -255,6 +259,32 @@ session.exec(select(func.avg(Item.embedding))).first()
255259

256260
Also supports `sum`
257261

262+
Add an approximate index
263+
264+
```python
265+
from sqlalchemy import Index
266+
267+
index = Index(
268+
'my_index',
269+
Item.embedding,
270+
postgresql_using='hnsw',
271+
postgresql_with={'m': 16, 'ef_construction': 64},
272+
postgresql_ops={'embedding': 'vector_l2_ops'}
273+
)
274+
# or
275+
index = Index(
276+
'my_index',
277+
Item.embedding,
278+
postgresql_using='ivfflat',
279+
postgresql_with={'lists': 100},
280+
postgresql_ops={'embedding': 'vector_l2_ops'}
281+
)
282+
283+
index.create(engine)
284+
```
285+
286+
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
287+
258288
## Psycopg 3
259289

260290
Enable the extension

tests/test_sqlalchemy.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,15 @@ class Item(Base):
2525
Base.metadata.drop_all(engine)
2626
Base.metadata.create_all(engine)
2727

28+
index = Index(
29+
'orm_index',
30+
Item.embedding,
31+
postgresql_using='hnsw',
32+
postgresql_with={'m': 16, 'ef_construction': 64},
33+
postgresql_ops={'embedding': 'vector_l2_ops'}
34+
)
35+
index.create(engine)
36+
2837

2938
def create_items():
3039
vectors = [

tests/test_sqlmodel.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import numpy as np
22
from pgvector.sqlalchemy import Vector
33
import pytest
4-
from sqlalchemy import Column
4+
from sqlalchemy import Column, Index
55
from sqlalchemy.exc import StatementError
66
from sqlalchemy.sql import func
77
from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text
@@ -22,6 +22,15 @@ class Item(SQLModel, table=True):
2222
SQLModel.metadata.drop_all(engine)
2323
SQLModel.metadata.create_all(engine)
2424

25+
index = Index(
26+
'sqlmodel_index',
27+
Item.embedding,
28+
postgresql_using='hnsw',
29+
postgresql_with={'m': 16, 'ef_construction': 64},
30+
postgresql_ops={'embedding': 'vector_l2_ops'}
31+
)
32+
index.create(engine)
33+
2534

2635
def create_items():
2736
vectors = [

0 commit comments

Comments
 (0)