Skip to content

Commit

Permalink
rm perf.sh & refactor examples.
Browse files Browse the repository at this point in the history
Signed-off-by: my-vegetable-has-exploded <[email protected]>
  • Loading branch information
my-vegetable-has-exploded committed Mar 19, 2024
1 parent 8c593e0 commit 2821faf
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 122 deletions.
5 changes: 0 additions & 5 deletions bindings/python/examples/perf.sh

This file was deleted.

57 changes: 57 additions & 0 deletions bindings/python/examples/psycopg_copy_dense.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import os

import numpy as np
import psycopg

from pgvecto_rs.psycopg import register_vector

URL = "postgresql://{username}:{password}@{host}:{port}/{db_name}".format(
port=os.getenv("DB_PORT", "5432"),
host=os.getenv("DB_HOST", "localhost"),
username=os.getenv("DB_USER", "postgres"),
password=os.getenv("DB_PASS", "mysecretpassword"),
db_name=os.getenv("DB_NAME", "postgres"),
)

# Connect to the DB and init things
with psycopg.connect(URL) as conn:
conn.execute("CREATE EXTENSION IF NOT EXISTS vectors;")
register_vector(conn)
conn.execute("DROP TABLE IF EXISTS documents;")
conn.execute(
"CREATE TABLE documents (id SERIAL PRIMARY KEY, embedding vector(3) NOT NULL);",
)
conn.commit()
try:
embeddings = [
np.array([1, 2, 3]),
np.array([1.0, 2.0, 4.0]),
np.array([1, 3, 4]),
]

with conn.cursor() as cursor, cursor.copy(
"COPY documents (embedding) FROM STDIN (FORMAT BINARY)"
) as copy:
# write row by row
for e in embeddings:
copy.write_row([e])
copy.write_row([[1, 3, 5]])
conn.commit()

# Select the rows using binary format
cur = conn.execute(
"SELECT * FROM documents;",
binary=True,
)
for row in cur.fetchall():
print(row[0], ": ", row[1])

# output will be:
# 1 : [1.0, 2.0, 3.0]
# 2 : [1.0, 2.0, 4.0]
# 3 : [1.0, 3.0, 4.0]
# 4 : [1.0, 3.0, 5.0]
finally:
# Drop the table
conn.execute("DROP TABLE IF EXISTS documents;")
conn.commit()
117 changes: 0 additions & 117 deletions bindings/python/examples/psycopg_copy_example.py

This file was deleted.

52 changes: 52 additions & 0 deletions bindings/python/examples/psycopg_copy_sparse.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
import os

import numpy as np
import psycopg

from pgvecto_rs.psycopg import register_vector
from pgvecto_rs.types import SparseVector

URL = "postgresql://{username}:{password}@{host}:{port}/{db_name}".format(
port=os.getenv("DB_PORT", "5432"),
host=os.getenv("DB_HOST", "localhost"),
username=os.getenv("DB_USER", "postgres"),
password=os.getenv("DB_PASS", "mysecretpassword"),
db_name=os.getenv("DB_NAME", "postgres"),
)


# Connect to the DB and init things
with psycopg.connect(URL) as conn:
conn.execute("CREATE EXTENSION IF NOT EXISTS vectors;")
register_vector(conn)
conn.execute("DROP TABLE IF EXISTS documents;")
conn.execute(
"CREATE TABLE documents (id SERIAL PRIMARY KEY, embedding svector NOT NULL);",
)
conn.commit()
try:
with conn.cursor() as cursor, cursor.copy(
"COPY documents (embedding) FROM STDIN (FORMAT BINARY)"
) as copy:
copy.write_row([SparseVector(3, [0, 2], [1.0, 3.0])])
copy.write_row([SparseVector(3, np.array([0, 1, 2]), [1.0, 2.0, 3.0])])
copy.write_row([SparseVector(3, np.array([1, 2]), np.array([2.0, 3.0]))])
conn.pgconn.flush()
conn.commit()

# Select the rows using binary format
cur = conn.execute(
"SELECT * FROM documents;",
binary=True,
)
for row in cur.fetchall():
print(row[0], ": ", row[1])

# output will be:
# 1 : [1.0, 0.0, 3.0]
# 2 : [1.0, 2.0, 3.0]
# 3 : [0.0, 2.0, 3.0]
finally:
# Drop the table
conn.execute("DROP TABLE IF EXISTS documents;")
conn.commit()

0 comments on commit 2821faf

Please sign in to comment.