Add write buffer parameter; add log for write time vs doc add time

Author: ldematte

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java

@@ -53,7 +53,8 @@ record CmdLineArgs(
     VectorEncoding vectorEncoding,
     int dimensions,
     boolean earlyTermination,
-    KnnIndexTester.MergePolicyType mergePolicy
+    KnnIndexTester.MergePolicyType mergePolicy,
+    double writerBufferSizeInMb
 ) implements ToXContentObject {
 
     static final ParseField DOC_VECTORS_FIELD = new ParseField("doc_vectors");
@@ -82,6 +83,9 @@ record CmdLineArgs(
     static final ParseField FILTER_SELECTIVITY_FIELD = new ParseField("filter_selectivity");
     static final ParseField SEED_FIELD = new ParseField("seed");
     static final ParseField MERGE_POLICY_FIELD = new ParseField("merge_policy");
+    static final ParseField WRITER_BUFFER_FIELD = new ParseField("writer_buffer_mb");
+
+    static final double DEFAULT_WRITER_BUFFER_MB = 128;
 
     static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         Builder builder = PARSER.apply(parser, null);
@@ -117,6 +121,7 @@ static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
         PARSER.declareFloat(Builder::setFilterSelectivity, FILTER_SELECTIVITY_FIELD);
         PARSER.declareLong(Builder::setSeed, SEED_FIELD);
         PARSER.declareString(Builder::setMergePolicy, MERGE_POLICY_FIELD);
+        PARSER.declareDouble(Builder::setWriterBufferMb, WRITER_BUFFER_FIELD);
     }
 
     @Override
@@ -186,6 +191,7 @@ static class Builder {
         private float filterSelectivity = 1f;
         private long seed = 1751900822751L;
         private KnnIndexTester.MergePolicyType mergePolicy = null;
+        private double writerBufferSizeInMb = DEFAULT_WRITER_BUFFER_MB;
 
         public Builder setDocVectors(List<String> docVectors) {
             if (docVectors == null || docVectors.isEmpty()) {
@@ -316,6 +322,11 @@ public Builder setMergePolicy(String mergePolicy) {
             return this;
         }
 
+        public Builder setWriterBufferMb(double writerBufferSizeInMb) {
+            this.writerBufferSizeInMb = writerBufferSizeInMb;
+            return this;
+        }
+
         public CmdLineArgs build() {
             if (docVectors == null) {
                 throw new IllegalArgumentException("Document vectors path must be provided");
@@ -350,7 +361,8 @@ public CmdLineArgs build() {
                 vectorEncoding,
                 dimensions,
                 earlyTermination,
-                mergePolicy
+                mergePolicy,
+                writerBufferSizeInMb
             );
         }
     }

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -240,7 +240,8 @@ public static void main(String[] args) throws Exception {
                     cmdLineArgs.dimensions(),
                     cmdLineArgs.vectorSpace(),
                     cmdLineArgs.numDocs(),
-                    mergePolicy
+                    mergePolicy,
+                    cmdLineArgs.writerBufferSizeInMb()
                 );
                 if (cmdLineArgs.reindex() == false && Files.exists(indexPath) == false) {
                     throw new IllegalArgumentException("Index path does not exist: " + indexPath);

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java

@@ -29,6 +29,7 @@
 import org.apache.lucene.index.ConcurrentMergeScheduler;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.IndexableField;
 import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -65,7 +66,6 @@
 import static org.elasticsearch.test.knn.KnnIndexTester.logger;
 
 class KnnIndexer {
-    private static final double WRITER_BUFFER_MB = 128;
     static final String ID_FIELD = "id";
     static final String VECTOR_FIELD = "vector";
 
@@ -78,6 +78,7 @@ class KnnIndexer {
     private final int numDocs;
     private final int numIndexThreads;
     private final MergePolicy mergePolicy;
+    private final double writerBufferSizeInMb;
 
     KnnIndexer(
         List<Path> docsPath,
@@ -88,7 +89,8 @@ class KnnIndexer {
         int dim,
         VectorSimilarityFunction similarityFunction,
         int numDocs,
-        MergePolicy mergePolicy
+        MergePolicy mergePolicy,
+        double writerBufferSizeInMb
     ) {
         this.docsPath = docsPath;
         this.indexPath = indexPath;
@@ -99,12 +101,14 @@ class KnnIndexer {
         this.similarityFunction = similarityFunction;
         this.numDocs = numDocs;
         this.mergePolicy = mergePolicy;
+        this.writerBufferSizeInMb = writerBufferSizeInMb;
     }
 
     void createIndex(KnnIndexTester.Results result) throws IOException, InterruptedException, ExecutionException {
         IndexWriterConfig iwc = new IndexWriterConfig().setOpenMode(IndexWriterConfig.OpenMode.CREATE);
         iwc.setCodec(codec);
-        iwc.setRAMBufferSizeMB(WRITER_BUFFER_MB);
+        iwc.setMaxBufferedDocs(IndexWriterConfig.DISABLE_AUTO_FLUSH);
+        iwc.setRAMBufferSizeMB(writerBufferSizeInMb);
         iwc.setUseCompoundFile(false);
         if (mergePolicy != null) {
             iwc.setMergePolicy(mergePolicy);
@@ -248,6 +252,9 @@ static class IndexerThread extends Thread {
         private final float[] floatVectorBuffer;
         private final VectorReader in;
 
+        long readTime;
+        long docAddTime;
+
         private IndexerThread(
             IndexWriter iw,
             VectorReader in,
@@ -283,6 +290,7 @@ public void run() {
             } catch (IOException ioe) {
                 throw new UncheckedIOException(ioe);
             }
+            logger.info("Index thread times: [{}] read, [{}] add doc", readTime, docAddTime);
         }
 
         private void _run() throws IOException {
@@ -294,23 +302,32 @@ private void _run() throws IOException {
                     continue;
                 }
 
-                Document doc = new Document();
+                var startRead = System.nanoTime();
+                final IndexableField field;
                 switch (vectorEncoding) {
                     case BYTE -> {
                         in.next(byteVectorBuffer);
-                        doc.add(new KnnByteVectorField(VECTOR_FIELD, byteVectorBuffer, fieldType));
+                        field = new KnnByteVectorField(VECTOR_FIELD, byteVectorBuffer, fieldType);
                     }
                     case FLOAT32 -> {
                         in.next(floatVectorBuffer);
-                        doc.add(new KnnFloatVectorField(VECTOR_FIELD, floatVectorBuffer, fieldType));
+                        field = new KnnFloatVectorField(VECTOR_FIELD, floatVectorBuffer, fieldType);
                     }
+                    default -> throw new UnsupportedOperationException();
                 }
+                long endRead = System.nanoTime();
+                readTime += (endRead - startRead);
+
+                Document doc = new Document();
+                doc.add(field);
 
                 if ((id + 1) % 25000 == 0) {
                     logger.debug("Done indexing " + (id + 1) + " documents.");
                 }
                 doc.add(new StoredField(ID_FIELD, id));
                 iw.addDocument(doc);
+
+                docAddTime += (System.nanoTime() - endRead);
             }
         }
     }