@@ -29,6 +29,8 @@ def connect(dbapi_connection, connection_record):
29
29
register_vector (dbapi_connection , globally = False , arrays = True )
30
30
31
31
32
+ psycopg3_engine = create_engine ('postgresql+psycopg://localhost/pgvector_python_test' )
33
+
32
34
Base = declarative_base ()
33
35
34
36
@@ -493,6 +495,34 @@ def test_binary_quantize(self):
493
495
items = session .query (Item ).order_by (distance ).all ()
494
496
assert [v .id for v in items ] == [2 , 3 , 1 ]
495
497
498
+ def test_psycopg_vector (self ):
499
+ with Session (psycopg3_engine ) as session :
500
+ session .add (Item (id = 1 , embedding = [1 , 2 , 3 ]))
501
+ session .commit ()
502
+ item = session .get (Item , 1 )
503
+ assert item .embedding .tolist () == [1 , 2 , 3 ]
504
+
505
+ def test_psycopg_halfvec (self ):
506
+ with Session (psycopg3_engine ) as session :
507
+ session .add (Item (id = 1 , half_embedding = [1 , 2 , 3 ]))
508
+ session .commit ()
509
+ item = session .get (Item , 1 )
510
+ assert item .half_embedding .to_list () == [1 , 2 , 3 ]
511
+
512
+ def test_psycopg_bit (self ):
513
+ with Session (psycopg3_engine ) as session :
514
+ session .add (Item (id = 1 , binary_embedding = '101' ))
515
+ session .commit ()
516
+ item = session .get (Item , 1 )
517
+ assert item .binary_embedding == '101'
518
+
519
+ def test_psycopg_sparsevec (self ):
520
+ with Session (psycopg3_engine ) as session :
521
+ session .add (Item (id = 1 , sparse_embedding = [1 , 2 , 3 ]))
522
+ session .commit ()
523
+ item = session .get (Item , 1 )
524
+ assert item .sparse_embedding .to_list () == [1 , 2 , 3 ]
525
+
496
526
@pytest .mark .asyncio
497
527
@pytest .mark .skipif (sqlalchemy_version == 1 , reason = 'Requires SQLAlchemy 2+' )
498
528
async def test_psycopg_async_avg (self ):
0 commit comments