Remove duplicated vector storage and rely on inline vectors for RAVV

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

@@ -36,16 +36,14 @@ class KNN80DocValuesConsumer extends DocValuesConsumer {
     private final Logger logger = LogManager.getLogger(KNN80DocValuesConsumer.class);
 
     private final DocValuesConsumer delegatee;
-    private final SegmentWriteState state;
 
     KNN80DocValuesConsumer(DocValuesConsumer delegatee, SegmentWriteState state) {
         this.delegatee = delegatee;
-        this.state = state;
     }
 
     @Override
     public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException {
-        delegatee.addBinaryField(field, valuesProducer);
+        if (!(field.hasVectorValues() && extractKNNEngine(field) == KNNEngine.JVECTOR)) delegatee.addBinaryField(field, valuesProducer);
         if (isKNNBinaryFieldRequired(field)) {
             StopWatch stopWatch = new StopWatch();
             stopWatch.start();
@@ -77,9 +75,38 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer,
     @Override
     public void merge(MergeState mergeState) {
         try {
-            delegatee.merge(mergeState);
             assert mergeState != null;
             assert mergeState.mergeFieldInfos != null;
+
+            for (DocValuesProducer docValuesProducer : mergeState.docValuesProducers) {
+                if (docValuesProducer != null) {
+                    docValuesProducer.checkIntegrity();
+                }
+            }
+
+            for (FieldInfo mergeFieldInfo : mergeState.mergeFieldInfos) {
+                if (mergeFieldInfo.hasVectorValues() && extractKNNEngine(mergeFieldInfo) == KNNEngine.JVECTOR) {
+                    continue;
+                }
+
+                DocValuesType type = mergeFieldInfo.getDocValuesType();
+                if (type != DocValuesType.NONE) {
+                    if (type == DocValuesType.NUMERIC) {
+                        delegatee.mergeNumericField(mergeFieldInfo, mergeState);
+                    } else if (type == DocValuesType.BINARY) {
+                        delegatee.mergeBinaryField(mergeFieldInfo, mergeState);
+                    } else if (type == DocValuesType.SORTED) {
+                        delegatee.mergeSortedField(mergeFieldInfo, mergeState);
+                    } else if (type == DocValuesType.SORTED_SET) {
+                        delegatee.mergeSortedSetField(mergeFieldInfo, mergeState);
+                    } else if (type == DocValuesType.SORTED_NUMERIC) {
+                        delegatee.mergeSortedNumericField(mergeFieldInfo, mergeState);
+                    } else {
+                        throw new AssertionError("type=" + type);
+                    }
+                }
+            }
+
             for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
                 DocValuesType type = fieldInfo.getDocValuesType();
                 if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java

@@ -14,19 +14,33 @@
 import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.codecs.DocValuesProducer;
 import org.apache.lucene.index.*;
+import org.opensearch.knn.index.codec.jvector.JVectorFloatVectorValues;
+import org.opensearch.knn.index.codec.jvector.JVectorReader;
+import org.opensearch.knn.index.engine.KNNEngine;
 
 import java.io.IOException;
 
+import static org.opensearch.knn.common.FieldInfoExtractor.extractKNNEngine;
+
 @Log4j2
 public class KNN80DocValuesProducer extends DocValuesProducer {
     private final DocValuesProducer delegate;
+    private final SegmentReadState state;
+    private volatile JVectorReader openReader;
 
     public KNN80DocValuesProducer(DocValuesProducer delegate, SegmentReadState state) {
         this.delegate = delegate;
+        this.state = state;
+        this.openReader = null;
     }
 
     @Override
     public BinaryDocValues getBinary(FieldInfo field) throws IOException {
+        if (field.hasVectorValues() && extractKNNEngine(field) == KNNEngine.JVECTOR) {
+            if (openReader == null) openReader = new JVectorReader(state);
+            return ((JVectorFloatVectorValues) openReader.getFloatVectorValues(field.name)).asBinaryDocValues();
+        }
+
         return delegate.getBinary(field);
     }
 
@@ -69,6 +83,12 @@ public void checkIntegrity() throws IOException {
 
     @Override
     public void close() throws IOException {
+
+        if (openReader != null) {
+            openReader.close();
+            openReader = null;
+        }
+
         delegate.close();
     }
 }

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFloatVectorValues.java

@@ -11,10 +11,15 @@
 import io.github.jbellis.jvector.vector.VectorizationProvider;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.util.BytesRef;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 
 public class JVectorFloatVectorValues extends FloatVectorValues {
     private static final VectorTypeSupport VECTOR_TYPE_SUPPORT = VectorizationProvider.getInstance().getVectorTypeSupport();
@@ -26,7 +31,7 @@ public class JVectorFloatVectorValues extends FloatVectorValues {
 
     public JVectorFloatVectorValues(OnDiskGraphIndex onDiskGraphIndex, VectorSimilarityFunction similarityFunction) throws IOException {
         this.dimension = onDiskGraphIndex.getDimension();
-        this.size = onDiskGraphIndex.size();
+        this.size = onDiskGraphIndex.getIdUpperBound();
         this.view = onDiskGraphIndex.getView();
         this.similarityFunction = similarityFunction;
     }
@@ -42,7 +47,9 @@ public int size() {
     }
 
     public VectorFloat<?> vectorFloatValue(int ord) {
-        return view.getVector(ord);
+        VectorFloat<?> value = VECTOR_TYPE_SUPPORT.createFloatVector(dimension);
+        view.getVectorInto(ord, value, 0);
+        return value;
     }
 
     public DocIndexIterator iterator() {
@@ -107,4 +114,47 @@ public VectorScorer scorer(float[] query) throws IOException {
         return new JVectorVectorScorer(this, VECTOR_TYPE_SUPPORT.createFloatVector(query), similarityFunction);
     }
 
+    public BinaryDocValues asBinaryDocValues() {
+
+        final DocIdSetIterator it = iterator();
+        final BytesRef bytes = new BytesRef(dimension * Float.BYTES);
+
+        return new BinaryDocValues() {
+            @Override
+            public BytesRef binaryValue() throws IOException {
+                float[] f = vectorValue(docID());
+                ByteBuffer bb = ByteBuffer.wrap(bytes.bytes).order(ByteOrder.LITTLE_ENDIAN);
+                for (int i = 0; i < f.length; i++)
+                    bb.putFloat(f[i]);
+
+                return bytes;
+            }
+
+            @Override
+            public boolean advanceExact(int target) throws IOException {
+                return it.advance(target) == target;
+            }
+
+            @Override
+            public int docID() {
+                return it.docID();
+            }
+
+            @Override
+            public int nextDoc() throws IOException {
+                return it.nextDoc();
+            }
+
+            @Override
+            public int advance(int target) throws IOException {
+                return it.advance(target);
+            }
+
+            @Override
+            public long cost() {
+                return it.cost();
+            }
+        };
+    }
+
 }

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java

@@ -127,7 +127,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
 
     @Override
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new JVectorReader(state, mergeOnDisk);
+        return new JVectorReader(state);
     }
 
     @Override

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorRandomAccessReader.java

@@ -52,7 +52,7 @@ public float readFloat() throws IOException {
     }
 
     // TODO: bring back to override when upgrading jVector again
-    // @Override
+    @Override
     public long readLong() throws IOException {
         return indexInputDelegate.readLong();
     }

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorReader.java

@@ -22,7 +22,6 @@
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
-import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.*;
 import org.apache.lucene.search.KnnCollector;
@@ -54,13 +53,9 @@ public class JVectorReader extends KnnVectorsReader {
     private final Map<String, FieldEntry> fieldEntryMap = new HashMap<>(1);
     private final Directory directory;
     private final SegmentReadState state;
-    private final FlatVectorsReader flatVectorsReader;
-    private final boolean mergeOnDisk;
 
-    public JVectorReader(SegmentReadState state, boolean mergeOnDisk) throws IOException {
+    public JVectorReader(SegmentReadState state) throws IOException {
         this.state = state;
-        this.mergeOnDisk = mergeOnDisk;
-        this.flatVectorsReader = FLAT_VECTORS_FORMAT.fieldsReader(state);
         this.fieldInfos = state.fieldInfos;
         this.baseDataFileName = state.segmentInfo.name + "_" + state.segmentSuffix;
         final String metaFileName = IndexFileNames.segmentFileName(
@@ -92,7 +87,6 @@ public JVectorReader(SegmentReadState state, boolean mergeOnDisk) throws IOExcep
 
     @Override
     public void checkIntegrity() throws IOException {
-        flatVectorsReader.checkIntegrity();
         for (FieldEntry fieldEntry : fieldEntryMap.values()) {
             try (var indexInput = state.directory.openInput(fieldEntry.vectorIndexFieldDataFileName, IOContext.READONCE)) {
                 CodecUtil.checksumEntireFile(indexInput);
@@ -102,9 +96,6 @@ public void checkIntegrity() throws IOException {
 
     @Override
     public FloatVectorValues getFloatVectorValues(String field) throws IOException {
-        if (mergeOnDisk) {
-            return flatVectorsReader.getFloatVectorValues(field);
-        }
         final FieldEntry fieldEntry = fieldEntryMap.get(field);
         return new JVectorFloatVectorValues(fieldEntry.index, fieldEntry.similarityFunction);
     }
@@ -212,7 +203,6 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
 
     @Override
     public void close() throws IOException {
-        IOUtils.close(flatVectorsReader);
         for (FieldEntry fieldEntry : fieldEntryMap.values()) {
             IOUtils.close(fieldEntry);
         }