Skip to content

Commit

Permalink
Added test for COPY TO with Psycopg 3
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Jul 1, 2024
1 parent 25074e0 commit a55d2c2
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions tests/test_psycopg.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,25 +118,34 @@ def test_sparsevec_text_format(self):
assert res.to_list() == [1.5, 0, 2, 0, 3, 0]
assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))

def test_text_copy(self):
def test_text_copy_from(self):
embedding = np.array([1.5, 2, 3])
cur = conn.cursor()
with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy:
copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)])

def test_binary_copy(self):
def test_binary_copy_from(self):
embedding = np.array([1.5, 2, 3])
cur = conn.cursor()
with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)])

def test_binary_copy_set_types(self):
def test_binary_copy_from_set_types(self):
embedding = np.array([1.5, 2, 3])
cur = conn.cursor()
with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec'])
copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)])

def test_binary_copy_to_set_types(self):
embedding = np.array([1.5, 2, 3])
conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,))
cur = conn.cursor()
with cur.copy("COPY psycopg_items (embedding) TO STDOUT WITH (FORMAT BINARY)") as copy:
copy.set_types(['vector'])
for row in copy.rows():
assert np.array_equal(row[0], embedding)

@pytest.mark.asyncio
async def test_async(self):
conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True)
Expand Down

0 comments on commit a55d2c2

Please sign in to comment.