Updated on 2024-06-03 GMT+08:00

Example: Vector Scenario

import psycopg2
import os

# Obtain a username and a password from environment variables.
user = os.getenv('user')
password = os.getenv('password')

# Create a connection object.
conn=psycopg2.connect(database="db_test",user=user,password=password,host="localhost",port=port)
cur=conn.cursor() # Create a pointer object.

# Create a connection object (using SSL).
#conn = psycopg2.connect(dbname="database", user=user, password=password, host="localhost", port=port,
#         sslmode="verify-ca", sslcert="client.crt",sslkey="client.key",sslrootcert="cacert.pem")
# Note: If sslcert, sslkey, and sslrootcert are not set, the following files in the .postgresql directory of the current user are used by default: client.crt, client.key, and root.crt.

# Create a table.
cur.execute("CREATE TABLE IF NOT EXISTS diskann_table (id int, repr FloatVector (128)) with (storage_type=astore);")

# Import local data in copy mode.
cur.execute("copy diskann_table from '/data/dataSet/sift1b_10w_128.csv' WITH(format 'text', delimiter E'\t', ignore_extra_data 'true', noescaping 'true');")

# Create an index. Specify the index parameters as required.
cur.execute("set maintenance_work_mem=1024000; CREATE Index diskannidx on diskann_table using gsdiskann (repr COSINE) with (pq_nseg = 128, pq_nclus = 16, queue_size = 100, enable_pq = false);")

# Insert a data record.
cur.execute("INSERT INTO diskann_table(id,repr) VALUES(%s,%s)",(0,'[0,16,35,5,32,31,14,10,11,78,55,10,45,83,11,6,14,57,102,75,20,8,3,5,67,17,19,26,5,0,1,22,60,26,7,1,18,22,84,53,85,119,119,4,24,18,7,7,1,81,106,102,72,30,6,0,9,1,9,119,72,1,4,33,119,29,6,1,0,1,14,52,119,30,3,0,0,55,92,111,2,5,4,9,22,89,96,14,1,0,1,82,59,16,20,5,25,14,11,4,0,0,1,26,47,23,4,0,0,4,38,83,30,14,9,4,9,17,23,41,0,0,2,8,19,25,23,1]'))

# Insert data in batches.
data = ((900000,'[78,9,0,0,0,0,0,0,163,43,1,5,15,1,0,0,14,4,1,42,75,5,0,0,23,2,7,30,13,0,0,25,83,6,0,0,0,1,0,3,163,22,0,0,14,9,1,7,38,2,1,21,163,35,3,4,61,43,3,18,29,3,6,20,76,0,0,0,0,1,0,10,163,6,0,0,17,3,0,10,62,1,0,2,163,43,11,7,32,7,0,1,18,13,43,96,60,0,0,0,0,0,0,18,163,1,0,0,12,0,0,33,56,1,0,1,141,6,0,5,97,10,0,0,11,0,2,27]'),(900001,'[7,7,0,1,9,71,39,6,62,1,0,0,0,19,26,54,119,4,0,0,16,6,7,136,15,0,0,1,136,23,4,14,22,120,11,3,5,39,20,3,78,23,5,5,21,44,33,37,136,73,0,0,34,10,8,75,11,5,0,2,136,31,3,5,88,68,6,1,0,1,8,61,80,31,6,2,8,49,40,50,136,46,2,7,29,5,6,41,19,10,2,32,136,1,0,1,62,31,0,0,0,0,12,51,41,34,3,0,1,5,18,25,83,34,0,5,20,2,2,2,12,8,0,22,136,1,0,0]'))
cur.executemany("INSERT INTO diskann_table(id,repr) VALUES(%s,%s)",data)

# Perform a commit.
conn.commit()

# Query the result.
cur.execute("SELECT id, repr<+>'[14,35,19,20,3,1,13,11,16,119,85,5,0,5,24,26,0,27,119,13,3,9,19,0,0,11,73,9,10,3,5,0,92,38,17,39,32,7,15,47,119,111,53,27,8,0,0,52,5,7,63,51,84,43,0,1,12,8,20,25,33,30,2,5,59,23,25,105,25,23,5,18,119,15,7,13,14,19,95,119,5,0,0,14,119,103,93,39,11,4,1,4,13,43,62,18,2,0,0,8,44,65,7,1,3,0,0,1,19,45,94,95,13,7,0,0,3,52,119,52,15,2,0,0,0,11,21,33]' as dist from diskann_table order by dist limit 1;")
results=cur.fetchall()
print (results)

# Close the connection.
cur.close()
conn.close()