diff --git a/neurips23/streaming/Dockerfile.BasePostgres b/neurips23/streaming/Dockerfile.BasePostgres index 2fdd2d49..564f0762 100644 --- a/neurips23/streaming/Dockerfile.BasePostgres +++ b/neurips23/streaming/Dockerfile.BasePostgres @@ -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 diff --git a/neurips23/streaming/base_postgres.py b/neurips23/streaming/base_postgres.py index 59b10868..0f8b3b38 100644 --- a/neurips23/streaming/base_postgres.py +++ b/neurips23/streaming/base_postgres.py @@ -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 = [] @@ -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() @@ -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})' + diff --git a/neurips23/streaming/start_database.sh b/neurips23/streaming/start_database.sh deleted file mode 100644 index 80c6978d..00000000 --- a/neurips23/streaming/start_database.sh +++ /dev/null @@ -1 +0,0 @@ -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" diff --git a/requirements_py3.10.txt b/requirements_py3.10.txt index 290184e8..3c5b5490 100644 --- a/requirements_py3.10.txt +++ b/requirements_py3.10.txt @@ -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