address feedback

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
 

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})'
+

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