Extend reflection to retrieve vector values from DefaultFlatVectorScorer$FloatScoringSupplier too (#138503)

While running gpu:yamlRestTest tests, under some circumstances the scorer is not Lucene99MemorySegmentFloatVectorScorerSupplier, but DefaultFlatVectorScorer$FloatScoringSupplier. This PR extend reflection to support the "fast path" to access vector values to this class too.

Author: ldematte

distribution/tools/server-cli/src/main/java/org/elasticsearch/server/cli/SystemJvmOptions.java

@@ -62,6 +62,7 @@ static List<String> systemJvmOptions(Settings nodeSettings, final Map<String, St
                 "-Dio.netty.recycler.maxCapacityPerThread=0",
                 // Needed to get access to raw vectors from Lucene scorers
                 "--add-opens=org.apache.lucene.core/org.apache.lucene.codecs.lucene99=org.elasticsearch.server",
+                "--add-opens=org.apache.lucene.core/org.apache.lucene.codecs.hnsw=org.elasticsearch.server",
                 "--add-opens=org.apache.lucene.core/org.apache.lucene.internal.vectorization=org.elasticsearch.server",
                 // log4j 2
                 "-Dlog4j.shutdownHookEnabled=false",

libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/ES92GpuHnswVectorsWriter.java

@@ -620,7 +620,7 @@ private void mergeByteVectorField(
                 }
             }
         } else {
-            logger.warn("Cannot get merged raw vectors from scorer.");
+            logger.warn("Cannot get merged raw vectors from scorer. Performances will be degraded.");
             var byteVectorValues = getMergedByteVectorValues(fieldInfo, mergeState);
 
             // TODO: revert to CuVSMatrix.deviceBuilder when cuvs has fixed the multiple copies problem

server/src/main/java/org/elasticsearch/index/codec/vectors/reflect/VectorsFormatReflectionUtils.java

@@ -22,11 +22,14 @@
 public class VectorsFormatReflectionUtils {
     private static final VarHandle FLOAT_SUPPLIER_HANDLE;
     private static final VarHandle BYTE_SUPPLIER_HANDLE;
-    private static final VarHandle FLOAT_VECTORS_HANDLE;
+    private static final VarHandle L99_FLOAT_VECTORS_HANDLE;
+    private static final VarHandle DEFAULT_FLOAT_VECTORS_HANDLE;
 
     private static final Class<?> FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS;
     private static final Class<?> SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS;
-    private static final Class<?> FLOAT_SCORING_SUPPLIER_CLASS;
+    private static final Class<?> L99_FLOAT_SCORING_SUPPLIER_CLASS;
+    private static final Class<?> DEFAULT_FLOAT_SCORING_SUPPLIER_CLASS;
+
     static {
         try {
             FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS = Class.forName(
@@ -39,11 +42,17 @@ public class VectorsFormatReflectionUtils {
                 RandomVectorScorerSupplier.class
             );
 
-            FLOAT_SCORING_SUPPLIER_CLASS = Class.forName(
+            L99_FLOAT_SCORING_SUPPLIER_CLASS = Class.forName(
                 "org.apache.lucene.internal.vectorization.Lucene99MemorySegmentFloatVectorScorerSupplier"
             );
-            lookup = MethodHandles.privateLookupIn(FLOAT_SCORING_SUPPLIER_CLASS, MethodHandles.lookup());
-            FLOAT_VECTORS_HANDLE = lookup.findVarHandle(FLOAT_SCORING_SUPPLIER_CLASS, "values", FloatVectorValues.class);
+            lookup = MethodHandles.privateLookupIn(L99_FLOAT_SCORING_SUPPLIER_CLASS, MethodHandles.lookup());
+            L99_FLOAT_VECTORS_HANDLE = lookup.findVarHandle(L99_FLOAT_SCORING_SUPPLIER_CLASS, "values", FloatVectorValues.class);
+
+            DEFAULT_FLOAT_SCORING_SUPPLIER_CLASS = Class.forName(
+                "org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer$FloatScoringSupplier"
+            );
+            lookup = MethodHandles.privateLookupIn(DEFAULT_FLOAT_SCORING_SUPPLIER_CLASS, MethodHandles.lookup());
+            DEFAULT_FLOAT_VECTORS_HANDLE = lookup.findVarHandle(DEFAULT_FLOAT_SCORING_SUPPLIER_CLASS, "vectors", FloatVectorValues.class);
 
             SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS = Class.forName(
                 Lucene99ScalarQuantizedVectorsWriter.class.getCanonicalName() + "$ScalarQuantizedCloseableRandomVectorScorerSupplier"
@@ -79,8 +88,13 @@ public static RandomVectorScorerSupplier getScalarQuantizedRandomVectorScorerInn
     }
 
     public static HasIndexSlice getFloatScoringSupplierVectorOrNull(RandomVectorScorerSupplier scorerSupplier) {
-        if (FLOAT_SCORING_SUPPLIER_CLASS.isAssignableFrom(scorerSupplier.getClass())) {
-            var vectorValues = FLOAT_VECTORS_HANDLE.get(scorerSupplier);
+        if (L99_FLOAT_SCORING_SUPPLIER_CLASS.isAssignableFrom(scorerSupplier.getClass())) {
+            var vectorValues = L99_FLOAT_VECTORS_HANDLE.get(scorerSupplier);
+            if (vectorValues instanceof HasIndexSlice indexSlice) {
+                return indexSlice;
+            }
+        } else if (DEFAULT_FLOAT_SCORING_SUPPLIER_CLASS.isAssignableFrom(scorerSupplier.getClass())) {
+            var vectorValues = DEFAULT_FLOAT_VECTORS_HANDLE.get(scorerSupplier);
             if (vectorValues instanceof HasIndexSlice indexSlice) {
                 return indexSlice;
             }

x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java

@@ -33,6 +33,7 @@ private static ElasticsearchCluster createCluster() {
             .setting("xpack.security.enabled", "false")
             // Needed to get access to raw vectors from Lucene scorers
             .jvmArg("--add-opens=org.apache.lucene.core/org.apache.lucene.codecs.lucene99=org.elasticsearch.server")
+            .jvmArg("--add-opens=org.apache.lucene.core/org.apache.lucene.codecs.hnsw=org.elasticsearch.server")
             .jvmArg("--add-opens=org.apache.lucene.core/org.apache.lucene.internal.vectorization=org.elasticsearch.server");
 
         var libraryPath = System.getenv("LD_LIBRARY_PATH");