Adapt VectorCommand, VectorBenchScenario, and VectorBenchConfig to support PgVectorBench

Author: yabsp

src/main/java/org/polypheny/simpleclient/cli/VectorCommand.java

@@ -47,6 +47,10 @@ public class VectorCommand implements CliRunnable {
     private List<String> args;
 
 
+    @Option(name = { "-m", "--mode" }, title = "Mode", arity = 1, description = "Execution mode: polypheny (default) or postgres.")
+    public String mode = "polypheny";
+
+
     @Option(name = { "-pdb", "--polyphenydb" }, title = "IP or Hostname", arity = 1, description = "IP or Hostname of the Polypheny-DB server (default: 127.0.0.1).")
     public static String polyphenyDbHost = "127.0.0.1";
 
@@ -78,16 +82,22 @@ public int run() throws SQLException {
 
         ExecutorFactory executorFactory;
         executorFactory = new PolyphenyDbJdbcExecutorFactory( polyphenyDbHost, false );
+        boolean usePostgres = mode.equalsIgnoreCase( "postgres" );
+        String task = args.getFirst();
 
         try {
-            if ( args.getFirst().equalsIgnoreCase( "data" ) ) {
-                VectorBenchScenario.data( executorFactory, multiplier, true );
-            } else if ( args.getFirst().equalsIgnoreCase( "workload" ) ) {
-                VectorBenchScenario.workload( executorFactory, multiplier, true, writeCsv, dumpQueryList );
+            if ( task.equalsIgnoreCase( "data" ) ) {
+                if ( usePostgres ) VectorBenchScenario.pgData( multiplier, true );
+                else VectorBenchScenario.data( executorFactory, multiplier, true );
+            } else if ( task.equalsIgnoreCase( "workload" ) ) {
+                if ( usePostgres ) VectorBenchScenario.pgWorkload( multiplier, true, writeCsv, dumpQueryList );
+                else VectorBenchScenario.workload( executorFactory, multiplier, true, writeCsv, dumpQueryList );
             } else if ( args.getFirst().equalsIgnoreCase( "schema" ) ) {
-                VectorBenchScenario.schema( executorFactory, true );
+                if ( usePostgres ) VectorBenchScenario.pgSchema( true );
+                else VectorBenchScenario.schema( executorFactory, true );
             } else if ( args.getFirst().equalsIgnoreCase( "warmup" ) ) {
-                VectorBenchScenario.warmup( executorFactory, multiplier, true, dumpQueryList );
+                if ( usePostgres ) VectorBenchScenario.pgWarmup( multiplier, true, dumpQueryList );
+                else VectorBenchScenario.warmup( executorFactory, multiplier, true, dumpQueryList );
             } else {
                 System.err.println( "Unknown task: " + args.getFirst() );
             }

src/main/java/org/polypheny/simpleclient/main/VectorBenchScenario.java

@@ -26,6 +26,8 @@
 
 import lombok.extern.slf4j.Slf4j;
 import org.polypheny.simpleclient.executor.Executor.ExecutorFactory;
+import org.polypheny.simpleclient.executor.PostgresExecutor.PostgresExecutorFactory;
+import org.polypheny.simpleclient.scenario.vectorbench.PgVectorBench;
 import org.polypheny.simpleclient.scenario.vectorbench.VectorBench;
 import org.polypheny.simpleclient.scenario.vectorbench.VectorBenchConfig;
 import java.io.File;
@@ -43,6 +45,14 @@ public static void schema( ExecutorFactory executorFactory, boolean commitAfterE
     }
 
 
+    public static void pgSchema( boolean commitAfterEveryQuery ) {
+        VectorBenchConfig config = new VectorBenchConfig( getProperties(), 1 );
+        ExecutorFactory factory = new PostgresExecutorFactory( config.postgresHost, false );
+        PgVectorBench vectorBench = new PgVectorBench( factory, config, commitAfterEveryQuery, false );
+        vectorBench.createSchema( null, true );
+    }
+
+
     public static void data( ExecutorFactory executorFactory, int multiplier, boolean commitAfterEveryQuery ) {
         VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
         VectorBench vectorBench = new VectorBench( executorFactory, config, commitAfterEveryQuery, false );
@@ -52,6 +62,15 @@ public static void data( ExecutorFactory executorFactory, int multiplier, boolea
     }
 
 
+    public static void pgData( int multiplier, boolean commitAfterEveryQuery ) {
+        VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
+        ExecutorFactory factory = new PostgresExecutorFactory( config.postgresHost, false );
+        PgVectorBench bench = new PgVectorBench( factory, config, commitAfterEveryQuery, false );
+        ProgressReporter progressReporter = new ProgressBar( config.numberOfThreads, config.progressReportBase );
+        bench.generateData( null, progressReporter );
+    }
+
+
     public static void workload( ExecutorFactory executorFactory, int multiplier, boolean commitAfterEveryQuery, boolean writeCsv, boolean dumpQueryList ) {
         VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
         VectorBench vectorBench = new VectorBench( executorFactory, config, commitAfterEveryQuery, dumpQueryList );
@@ -68,6 +87,17 @@ public static void workload( ExecutorFactory executorFactory, int multiplier, bo
     }
 
 
+    public static void pgWorkload( int multiplier, boolean commitAfterEveryQuery, boolean writeCsv, boolean dumpQueryList ) {
+        VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
+        ExecutorFactory factory = new PostgresExecutorFactory( config.postgresHost, false );
+        PgVectorBench bench = new PgVectorBench( factory, config, commitAfterEveryQuery, dumpQueryList );
+        CsvWriter csvWriter = writeCsv ? new CsvWriter( "results-pg.csv" ) : null;
+        ProgressReporter progressReporter = new ProgressBar( config.numberOfThreads, config.progressReportBase );
+        bench.execute( progressReporter, csvWriter, new File( "." ), config.numberOfThreads );
+    }
+
+
+
     public static void warmup( ExecutorFactory executorFactory, int multiplier, boolean commitAfterEveryQuery, boolean dumpQueryList ) {
         VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
         VectorBench vectorBench = new VectorBench( executorFactory, config, commitAfterEveryQuery, dumpQueryList );
@@ -77,6 +107,16 @@ public static void warmup( ExecutorFactory executorFactory, int multiplier, bool
     }
 
 
+    public static void pgWarmup( int multiplier, boolean commitAfterEveryQuery, boolean dumpQueryList ) {
+        VectorBenchConfig config = new VectorBenchConfig( getProperties(), multiplier );
+        ExecutorFactory factory = new PostgresExecutorFactory( config.postgresHost, false );
+        PgVectorBench bench = new PgVectorBench( factory, config, commitAfterEveryQuery, dumpQueryList );
+        ProgressReporter progressReporter = new ProgressBar( config.numberOfThreads, config.progressReportBase );
+        bench.warmUp( progressReporter );
+    }
+
+
+
     private static Properties getProperties() {
         Properties props = new Properties();
         try {

src/main/java/org/polypheny/simpleclient/scenario/vectorbench/VectorBenchConfig.java

@@ -38,6 +38,8 @@ public class VectorBenchConfig extends AbstractConfig {
     public String dataStoreFeature;
     public String dataStoreMetadata;
 
+    public String postgresHost;
+
     public long randomSeedInsert;
     public long randomSeedQuery;
 
@@ -80,6 +82,8 @@ public VectorBenchConfig(Properties properties, int multiplier ) {
         }
         //dataStores.add( "cottontail" );
 
+        postgresHost = getStringProperty( properties, "postgresHost" );
+
         if ( getBooleanProperty( properties, "useRandomSeeds" ) ) {
             Random tempRand = new Random();
             randomSeedInsert = tempRand.nextLong();

src/main/resources/org/polypheny/simpleclient/scenario/vectorbench/vector.properties

@@ -3,6 +3,9 @@ scenario = "vectorBench"
 dataStoreFeature = postgresql1
 dataStoreMeta = hsqldb
 
+# PostgreSQL direct connect settings
+postgresHost = 127.0.0.1
+
 numberOfThreads = 4
 progressReportBase = 100
 numberOfWarmUpIterations = 4