Skip to content

Commit

Permalink
address feedback
Browse files Browse the repository at this point in the history
  • Loading branch information
onurctirtir committed Sep 13, 2024
1 parent 381215d commit 2c30c59
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 6 deletions.
10 changes: 8 additions & 2 deletions neurips23/streaming/Dockerfile.BasePostgres
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,16 @@ RUN mkdir /var/lib/postgresql/test_database
RUN /usr/lib/postgresql/16/bin/initdb -D /var/lib/postgresql/test_database
USER root

# copy the script to start the database that we will use start the database during the tests
COPY neurips23/streaming/start_database.sh /home/app/
# create that we will use start the database during the tests
RUN echo "su - postgres -c \"/usr/lib/postgresql/16/bin/pg_ctl \
-D /var/lib/postgresql/test_database \
-l /var/lib/postgresql/test_database_logfile \
-o \\\"-F -p 5432\\\" start\"" > /home/app/start_database.sh
RUN chmod +x /home/app/start_database.sh

# install python deps
RUN pip3 install pgvector==0.3.3 psycopg==3.2.1

# install linux-tools-generic into docker so that devs can use perf if they want
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt install -y linux-tools-generic

Expand Down
12 changes: 11 additions & 1 deletion neurips23/streaming/base_postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,9 @@ def __init__(self, metric, index_params):
if self.n_insert_conns == None:
raise Exception('Missing parameter insert_conns')

# save it for later use in __str__()
self._index_params = index_params

# we'll initialize the connections later in set_query_arguments() per "query-arguments" set
self.conns = []

Expand Down Expand Up @@ -228,6 +231,9 @@ def batch_query(conn_idx, query_vec_start_idx, query_vec_end_idx):
self.res = np.vstack(result_id_lists, dtype=np.int32)

def set_query_arguments(self, query_args):
# save it for later use in __str__()
self._query_args = query_args

# close any existing connections
for conn in self.conns:
conn.close()
Expand Down Expand Up @@ -261,4 +267,8 @@ def set_query_arguments(self, query_args):
cursor_print_and_execute(cur, f"SET enable_seqscan TO OFF")

def __str__(self):
return self.name
build_args_str = ' '.join([f'{k}={v}' for k, v in sorted(self._index_params.items())])
query_args_str = ' '.join([f'{k}={v}' for k, v in sorted(self._query_args.items())])

return f'{self.name}({build_args_str} - {query_args_str})'

1 change: 0 additions & 1 deletion neurips23/streaming/start_database.sh

This file was deleted.

2 changes: 0 additions & 2 deletions requirements_py3.10.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,3 @@ scipy==1.10.1
scikit-learn
jinja2==3.1.2
pandas==2.0.0
psycopg==3.2.1
pgvector==0.3.3

0 comments on commit 2c30c59

Please sign in to comment.