Skip to content

Commit 3668191

Browse files
erikbernerikbern
authored
Merge pull request #539 from C0rWin/pgvectorscale
Add pgvectorscale
2 parents d9274ba + 264cb02 commit 3668191

File tree

5 files changed

+215
-0
lines changed

5 files changed

+215
-0
lines changed

‎.github/workflows/benchmarks.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ jobs:
5757
- parlayann
5858
- pg_embedding
5959
- pgvector
60+
- pgvectorscale
6061
- pgvecto_rs
6162
- pynndescent
6263
- redisearch

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ Evaluated
5151
* [Descartes(01AI)](https://github.com/xiaoming-01ai/descartes)
5252
* [kgn](https://github.com/Henry-yan/kgn)
5353
* [vsag](https://github.com/antgroup/vsag)
54+
* [PGVectorScale](https://github.com/timescale/pgvectorscale/tree/main)
5455

5556
Data sets
5657
=========

‎ann_benchmarks/algorithms/pgvectorscale/Dockerfile‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
FROM ann-benchmarks
2+
3+
RUN git clone https://github.com/pgvector/pgvector /tmp/pgvector
4+
5+
RUN DEBIAN_FRONTEND=noninteractive apt-get -y install tzdata
6+
RUN apt-get update && apt-get install -y --no-install-recommends build-essential postgresql-common
7+
RUN /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y
8+
RUN apt-get install -y --no-install-recommends postgresql-16 postgresql-server-dev-16
9+
RUN sh -c 'echo "local all all trust" > /etc/postgresql/16/main/pg_hba.conf'
10+
11+
# Dynamically set OPTFLAGS based on the architecture
12+
RUN ARCH=$(uname -m) && \
13+
if [ "$ARCH" = "aarch64" ]; then \
14+
OPTFLAGS="-march=native -msve-vector-bits=512"; \
15+
elif [ "$ARCH" = "x86_64" ]; then \
16+
OPTFLAGS="-march=native -mprefer-vector-width=512"; \
17+
else \
18+
OPTFLAGS="-march=native"; \
19+
fi && \
20+
cd /tmp/pgvector && \
21+
make clean && \
22+
make OPTFLAGS="$OPTFLAGS" && \
23+
make install
24+
25+
# Install necessary dependencies
26+
RUN apt-get update && apt-get install -y \
27+
curl \
28+
build-essential \
29+
libssl-dev \
30+
pkg-config \
31+
&& rm -rf /var/lib/apt/lists/*
32+
33+
# Install Rust and Cargo
34+
RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
35+
36+
# Ensure the cargo bin directory is in the PATH
37+
ENV PATH="/root/.cargo/bin:${PATH}"
38+
39+
# Install pgrx using cargo
40+
RUN cargo install --locked [email protected]
41+
# Initialize pgrx with the PostgreSQL version
42+
RUN cargo pgrx init --pg16 pg_config
43+
RUN git clone https://github.com/timescale/pgvectorscale /tmp/pgvectorscale
44+
RUN cd /tmp/pgvectorscale/pgvectorscale && \
45+
cargo pgrx install --release
46+
47+
USER postgres
48+
RUN service postgresql start && \
49+
psql -c "CREATE USER ann WITH ENCRYPTED PASSWORD 'ann'" && \
50+
psql -c "CREATE DATABASE ann" && \
51+
psql -c "GRANT ALL PRIVILEGES ON DATABASE ann TO ann" && \
52+
psql -d ann -c "GRANT ALL ON SCHEMA public TO ann" && \
53+
psql -d ann -c "CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE" && \
54+
psql -c "ALTER USER ann SET maintenance_work_mem = '4GB'" && \
55+
psql -c "ALTER USER ann SET max_parallel_maintenance_workers = 0" && \
56+
psql -c "ALTER SYSTEM SET shared_buffers = '4GB'"
57+
USER root
58+
59+
RUN pip install psycopg[binary] pgvector

‎ann_benchmarks/algorithms/pgvectorscale/config.yml‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
float:
2+
any:
3+
- base_args: ['@metric', 1]
4+
constructor: PGDiskANN
5+
disabled: false
6+
docker_tag: ann-benchmarks-pgvectorscale
7+
module: ann_benchmarks.algorithms.pgvectorscale
8+
name: pgvectorscale1
9+
run_groups:
10+
N-50:
11+
arg_groups: [{num_neighbors: 50, search_list_size: 100, max_alpha: 1.2}]
12+
args: {}
13+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
14+
N-100:
15+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
16+
args: {}
17+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
18+
ALPHA-1.00:
19+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1}]
20+
args: {}
21+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
22+
ALPHA-1.20:
23+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
24+
args: {}
25+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
26+
ALPHA-1.50:
27+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.5}]
28+
args: {}
29+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
30+
ALPHA-2.00:
31+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 2}]
32+
args: {}
33+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
34+
- base_args: ['@metric', 2]
35+
constructor: PGDiskANN
36+
disabled: false
37+
docker_tag: ann-benchmarks-pgvectorscale
38+
module: ann_benchmarks.algorithms.pgvectorscale
39+
name: pgvectorscale2
40+
run_groups:
41+
N-50:
42+
arg_groups: [{num_neighbors: 50, search_list_size: 100, max_alpha: 1.2}]
43+
args: {}
44+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
45+
N-100:
46+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
47+
args: {}
48+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
49+
ALPHA-1.00:
50+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1}]
51+
args: {}
52+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
53+
ALPHA-1.20:
54+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
55+
args: {}
56+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
57+
ALPHA-1.50:
58+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.5}]
59+
args: {}
60+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
61+
ALPHA-2.00:
62+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 2}]
63+
args: {}
64+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
65+
- base_args: ['@metric', 4]
66+
constructor: PGDiskANN
67+
disabled: false
68+
docker_tag: ann-benchmarks-pgvectorscale
69+
module: ann_benchmarks.algorithms.pgvectorscale
70+
name: pgvectorscale4
71+
run_groups:
72+
N-50:
73+
arg_groups: [{num_neighbors: 50, search_list_size: 100, max_alpha: 1.2}]
74+
args: {}
75+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
76+
N-100:
77+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
78+
args: {}
79+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
80+
ALPHA-1.00:
81+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1}]
82+
args: {}
83+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
84+
ALPHA-1.20:
85+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.2}]
86+
args: {}
87+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
88+
ALPHA-1.50:
89+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 1.5}]
90+
args: {}
91+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
92+
ALPHA-2.00:
93+
arg_groups: [{num_neighbors: 100, search_list_size: 100, max_alpha: 2}]
94+
args: {}
95+
query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]

‎ann_benchmarks/algorithms/pgvectorscale/module.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import subprocess
2+
import sys
3+
4+
import pgvector.psycopg
5+
import psycopg
6+
7+
from ..base.module import BaseANN
8+
9+
10+
class PGDiskANN(BaseANN):
11+
def __init__(self, metric, nbits, method_param):
12+
print(f"running constructor")
13+
self._metric = metric
14+
self._nbits = nbits
15+
self._cur = None
16+
self._query = "SELECT id FROM items ORDER BY embedding <=> %s LIMIT %s"
17+
self._num_neighbors = method_param["num_neighbors"]
18+
self._search_list_size = method_param["search_list_size"]
19+
self._max_alpha = method_param["max_alpha"]
20+
print(f"running only {self._metric} and {self._query}")
21+
22+
def fit(self, X):
23+
print("running fit")
24+
subprocess.run("service postgresql start", shell=True, check=True, stdout=sys.stdout, stderr=sys.stderr)
25+
conn = psycopg.connect(user="ann", password="ann", dbname="ann", autocommit=True)
26+
pgvector.psycopg.register_vector(conn)
27+
cur = conn.cursor()
28+
cur.execute("DROP TABLE IF EXISTS items")
29+
cur.execute("CREATE TABLE items (id int, embedding vector(%d))" % X.shape[1])
30+
cur.execute("ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN")
31+
print("copying data...")
32+
with cur.copy("COPY items (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
33+
copy.set_types(["int4", "vector"])
34+
for i, embedding in enumerate(X):
35+
copy.write_row((i, embedding))
36+
print("creating index...")
37+
cur.execute(
38+
"CREATE INDEX ON items USING diskann(embedding) WITH (num_neighbors = %d, search_list_size = %d, max_alpha = %d, num_bits_per_dimension = %d)"
39+
% (self._num_neighbors, self._search_list_size, self._max_alpha, self._nbits)
40+
)
41+
print("done!")
42+
self._cur = cur
43+
44+
def query(self, v, n):
45+
self._cur.execute(self._query, (v, n), binary=True, prepare=True)
46+
return [id for id, in self._cur.fetchall()]
47+
48+
def set_query_arguments(self, list_size):
49+
self._list_size = list_size
50+
self._cur.execute("SET diskann.query_search_list_size = %d" % list_size)
51+
52+
def get_memory_usage(self):
53+
if self._cur is None:
54+
return 0
55+
self._cur.execute("SELECT pg_relation_size('items_embedding_idx')")
56+
return self._cur.fetchone()[0] / 1024
57+
58+
def __str__(self):
59+
return f"PGDiskANN(metric={self._metric}), num_neighbors={self._num_neighbors}, search_list_size={self._search_list_size}, max_alpha={self._max_alpha}), nbits={self._nbits})"

0 commit comments

Comments
 (0)