Revert "Attempt to use Java 22 DatasetImpl with MemorySegment"

This reverts commit ec3330e.

Author: mayya-sharipova

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

@@ -42,13 +42,8 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.lang.foreign.Arena;
-import java.lang.foreign.MemorySegment;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.nio.channels.FileChannel;
-import java.nio.file.Path;
-import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -465,56 +460,36 @@ public NodesIterator getNodesOnLevel(int level) {
         };
     }
 
+    // TODO check with deleted documents
     @Override
     @SuppressForbidden(reason = "require usage of Lucene's IOUtils#deleteFilesIgnoringExceptions(...)")
     public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
-        int dims = fieldInfo.getVectorDimension();
         flatVectorWriter.mergeOneField(fieldInfo, mergeState);
-        FloatVectorValues mergeFloatVectorValues = MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
-
-        if (mergeFloatVectorValues.size() < MIN_NUM_VECTORS_FOR_GPU_BUILD) {
-            // TODO: check how deleted documents affect size value
-            KnnVectorValues.DocIndexIterator iter = mergeFloatVectorValues.iterator();
-            float[] vector = new float[dims];
-            List<float[]> vectorsList = new ArrayList<>();
-            for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) {
-                System.arraycopy(mergeFloatVectorValues.vectorValue(iter.index()), 0, vector, 0, dims);
-                vectorsList.add(vector);
-            }
-            float[][] vectors = vectorsList.toArray(new float[0][]);
-            DatasetOrVectors datasetOrVectors = new DatasetOrVectors(vectors);
-            writeFieldInternal(fieldInfo, datasetOrVectors);
-            return;
-        }
-
-
+        FloatVectorValues vectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
+        // save merged vector values to a temp file
         final int numVectors;
         String tempRawVectorsFileName = null;
         boolean success = false;
-        // save merged vectors to a temporary file
         try (IndexOutput out = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "vec_", IOContext.DEFAULT)) {
             tempRawVectorsFileName = out.getName();
-            numVectors = writeFloatVectorValues(fieldInfo, out, mergeFloatVectorValues);
+            numVectors = writeFloatVectorValues(fieldInfo, out, MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState));
             CodecUtil.writeFooter(out);
             success = true;
         } finally {
             if (success == false && tempRawVectorsFileName != null) {
                 org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName);
             }
         }
-        // Use MemorySegment to map the temp file and pass it as a dataset for building the GPU index
-        try {
-            final Path path = ((org.apache.lucene.store.FSDirectory) mergeState.segmentInfo.dir).getDirectory().resolve(tempRawVectorsFileName);
-            Arena arena = Arena.ofShared();
-            FileChannel fileChannel = FileChannel.open(path, StandardOpenOption.READ);
-            final MemorySegment memorySegment = fileChannel.map(
-                FileChannel.MapMode.READ_ONLY,
-                0,
-                fileChannel.size() - CodecUtil.footerLength(),
-                arena
-            );
-            Dataset dataset = new DatasetImpl(arena, memorySegment, numVectors, fieldInfo.getVectorDimension());
-            DatasetOrVectors datasetOrVectors = new DatasetOrVectors(dataset, null);
+        try (IndexInput in = mergeState.segmentInfo.dir.openInput(tempRawVectorsFileName, IOContext.DEFAULT)) {
+            // TODO: Improve this (not acceptable): pass tempRawVectorsFileName for the gpuIndex construction through MemorySegment
+            final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, in, numVectors);
+            float[][] vectors = new float[numVectors][fieldInfo.getVectorDimension()];
+            float[] vector;
+            for (int i = 0; i < numVectors; i++) {
+                vector = floatVectorValues.vectorValue(i);
+                System.arraycopy(vector, 0, vectors[i], 0, vector.length);
+            }
+            DatasetOrVectors datasetOrVectors = new DatasetOrVectors(vectors);
             writeFieldInternal(fieldInfo, datasetOrVectors);
         } finally {
             org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName);
@@ -536,6 +511,47 @@ private static int writeFloatVectorValues(FieldInfo fieldInfo, IndexOutput out,
         return numVectors;
     }
 
+    private static FloatVectorValues getFloatVectorValues(FieldInfo fieldInfo, IndexInput randomAccessInput, int numVectors) {
+        if (numVectors == 0) {
+            return FloatVectorValues.fromFloats(List.of(), fieldInfo.getVectorDimension());
+        }
+        final long length = (long) Float.BYTES * fieldInfo.getVectorDimension() + Integer.BYTES;
+        final float[] vector = new float[fieldInfo.getVectorDimension()];
+        return new FloatVectorValues() {
+            @Override
+            public float[] vectorValue(int ord) throws IOException {
+                randomAccessInput.seek(ord * length + Integer.BYTES);
+                randomAccessInput.readFloats(vector, 0, vector.length);
+                return vector;
+            }
+
+            @Override
+            public FloatVectorValues copy() {
+                return this;
+            }
+
+            @Override
+            public int dimension() {
+                return fieldInfo.getVectorDimension();
+            }
+
+            @Override
+            public int size() {
+                return numVectors;
+            }
+
+            @Override
+            public int ordToDoc(int ord) {
+                try {
+                    randomAccessInput.seek(ord * length);
+                    return randomAccessInput.readInt();
+                } catch (IOException e) {
+                    throw new UncheckedIOException(e);
+                }
+            }
+        };
+    }
+
     private void writeMeta(
         FieldInfo field,
         long vectorIndexOffset,