Expose m and efConstruction params in GPU index building

Author: mayya-sharipova

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -2117,7 +2117,7 @@ public boolean updatableTo(DenseVectorIndexOptions update) {
         }
     }
 
-    static class HnswIndexOptions extends DenseVectorIndexOptions {
+    public static class HnswIndexOptions extends DenseVectorIndexOptions {
         private final int m;
         private final int efConstruction;
 
@@ -2160,6 +2160,14 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             return builder;
         }
 
+        public int m() {
+            return m;
+        }
+
+        public int efConstruction() {
+            return efConstruction;
+        }
+
         @Override
         public boolean doEquals(DenseVectorIndexOptions o) {
             if (this == o) return true;

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java

@@ -6,6 +6,8 @@
  */
 package org.elasticsearch.xpack.gpu;
 
+import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.util.hnsw.HnswGraphBuilder;
 import org.elasticsearch.common.util.FeatureFlag;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
@@ -37,12 +39,12 @@ public VectorsFormatProvider getVectorsFormatProvider() {
                             "[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node."
                         );
                     }
-                    return new GPUVectorsFormat();
+                    return getVectorsFormat(indexOptions);
                 }
                 if (gpuMode == IndexSettings.GpuMode.AUTO
                     && vectorIndexTypeSupported(indexOptions.getType())
                     && GPUSupport.isSupported(false)) {
-                    return new GPUVectorsFormat();
+                    return getVectorsFormat(indexOptions);
                 }
             }
             return null;
@@ -52,4 +54,19 @@ && vectorIndexTypeSupported(indexOptions.getType())
     private boolean vectorIndexTypeSupported(DenseVectorFieldMapper.VectorIndexType type) {
         return type == DenseVectorFieldMapper.VectorIndexType.HNSW;
     }
+
+    private static KnnVectorsFormat getVectorsFormat(DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions) {
+        if (indexOptions.getType() == DenseVectorFieldMapper.VectorIndexType.HNSW) {
+            DenseVectorFieldMapper.HnswIndexOptions hnswIndexOptions = (DenseVectorFieldMapper.HnswIndexOptions) indexOptions;
+            int efConstruction = hnswIndexOptions.efConstruction();
+            if (efConstruction == HnswGraphBuilder.DEFAULT_BEAM_WIDTH) {
+                efConstruction = GPUVectorsFormat.DEFAULT_BEAM_WIDTH; // default value for GPU graph construction is 128
+            }
+            return new GPUVectorsFormat(hnswIndexOptions.m(), efConstruction);
+        } else {
+            throw new IllegalArgumentException(
+                "GPU vector indexing is not supported on this vector type: [" + indexOptions.getType() + "]"
+            );
+        }
+    }
 }

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUToHNSWVectorsWriter.java

@@ -86,8 +86,8 @@ final class GPUToHNSWVectorsWriter extends KnnVectorsWriter {
         assert cuVSResourceManager != null : "CuVSResources must not be null";
         this.cuVSResourceManager = cuVSResourceManager;
         this.M = M;
-        this.flatVectorWriter = flatVectorWriter;
         this.beamWidth = beamWidth;
+        this.flatVectorWriter = flatVectorWriter;
         this.segmentWriteState = state;
         String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, LUCENE99_HNSW_META_EXTENSION);
         String indexDataFileName = IndexFileNames.segmentFileName(
@@ -274,10 +274,11 @@ private CagraIndex buildGPUIndex(
             case COSINE -> CagraIndexParams.CuvsDistanceType.CosineExpanded;
         };
 
-        // TODO: expose cagra index params of intermediate graph degree, graph degree, algorithm, NNDescentNumIterations
+        // TODO: expose cagra index params for algorithm, NNDescentNumIterations
         CagraIndexParams params = new CagraIndexParams.Builder().withNumWriterThreads(1) // TODO: how many CPU threads we can use?
             .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT)
-            .withGraphDegree(16)
+            .withGraphDegree(M)
+            .withIntermediateGraphDegree(beamWidth)
             .withMetric(distanceType)
             .build();
 

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormat.java

@@ -16,8 +16,6 @@
 import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
-import org.elasticsearch.logging.LogManager;
-import org.elasticsearch.logging.Logger;
 
 import java.io.IOException;
 
@@ -26,9 +24,6 @@
  * leverage GPU processing capabilities for vector search operations.
  */
 public class GPUVectorsFormat extends KnnVectorsFormat {
-
-    private static final Logger LOG = LogManager.getLogger(GPUVectorsFormat.class);
-
     public static final String NAME = "GPUVectorsFormat";
     public static final int VERSION_START = 0;
 
@@ -38,34 +33,38 @@ public class GPUVectorsFormat extends KnnVectorsFormat {
     static final String LUCENE99_HNSW_VECTOR_INDEX_EXTENSION = "vex";
     static final int LUCENE99_VERSION_CURRENT = VERSION_START;
 
-    static final int DEFAULT_MAX_CONN = 16;
-    static final int DEFAULT_BEAM_WIDTH = 100;
+    static final int DEFAULT_MAX_CONN = 16; // graph degree
+    public static final int DEFAULT_BEAM_WIDTH = 128; // intermediate graph degree
     static final int MIN_NUM_VECTORS_FOR_GPU_BUILD = 2;
 
     private static final FlatVectorsFormat flatVectorsFormat = new Lucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
+    // How many nodes each node in the graph is connected to in the final graph
+    private final int maxConn;
+    // Intermediate graph degree, the number of connections for each node before pruning
+    private final int beamWidth;
     final CuVSResourceManager cuVSResourceManager;
 
     public GPUVectorsFormat() {
-        this(CuVSResourceManager.pooling());
+        this(CuVSResourceManager.pooling(), DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH);
     }
 
-    public GPUVectorsFormat(CuVSResourceManager cuVSResourceManager) {
+    public GPUVectorsFormat(int maxConn, int beamWidth) {
+        this(CuVSResourceManager.pooling(), maxConn, beamWidth);
+    };
+
+    public GPUVectorsFormat(CuVSResourceManager cuVSResourceManager, int maxConn, int beamWidth) {
         super(NAME);
         this.cuVSResourceManager = cuVSResourceManager;
+        this.maxConn = maxConn;
+        this.beamWidth = beamWidth;
     }
 
     @Override
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        return new GPUToHNSWVectorsWriter(
-            cuVSResourceManager,
-            state,
-            DEFAULT_MAX_CONN,
-            DEFAULT_BEAM_WIDTH,
-            flatVectorsFormat.fieldsWriter(state)
-        );
+        return new GPUToHNSWVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state));
     }
 
     @Override
@@ -80,6 +79,6 @@ public int getMaxDimensions(String fieldName) {
 
     @Override
     public String toString() {
-        return NAME + "()";
+        return NAME + "(maxConn=" + maxConn + ", beamWidth=" + beamWidth + ", flatVectorFormat=" + flatVectorsFormat.getName() + ")";
     }
 }

x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormatTests.java

@@ -76,29 +76,14 @@ public void testMismatchedFields() throws Exception {
         // No bytes support
     }
 
-    @Override
-    public void testSortedIndex() throws Exception {
-        // TODO: implement sorted index
-    }
-
-    @Override
-    public void testFloatVectorScorerIteration() throws Exception {
-        // TODO: implement sorted index
-    }
-
-    @Override
-    public void testRandom() throws Exception {
-        // TODO: implement sorted index
-    }
-
     public void testToString() {
         FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) {
             @Override
             public KnnVectorsFormat knnVectorsFormat() {
                 return new GPUVectorsFormat();
             }
         };
-        String expectedPattern = "GPUVectorsFormat()";
+        String expectedPattern = "GPUVectorsFormat(maxConn=16, beamWidth=128, flatVectorFormat=Lucene99FlatVectorsFormat)";
         assertEquals(expectedPattern, customCodec.knnVectorsFormat().toString());
     }