Switch away from getMergeInstance (#2472)

Signed-off-by: John Mazanec <[email protected]>
(cherry picked from commit abae323)

src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsReader.java

@@ -107,12 +107,12 @@ public void close() throws IOException {
 
     /**
      * For merging, we need to tell the derived source stored fields reader to skip injecting the source. Otherwise,
-     * on merge we will end up just writing the source to disk
+     * on merge we will end up just writing the source to disk. We cant override
+     * {@link StoredFieldsReader#getMergeInstance()} because it is used elsewhere than just merging.
      *
      * @return Merged instance that wont inject by default
      */
-    @Override
-    public StoredFieldsReader getMergeInstance() {
+    private StoredFieldsReader cloneForMerge() {
         try {
             return new DerivedSourceStoredFieldsReader(
                 delegate.getMergeInstance(),
@@ -125,4 +125,18 @@ public StoredFieldsReader getMergeInstance() {
             throw new RuntimeException(e);
         }
     }
+
+    /**
+     * For merging, we need to tell the derived source stored fields reader to skip injecting the source. Otherwise,
+     * on merge we will end up just writing the source to disk
+     *
+     * @param storedFieldsReader stored fields reader to wrap
+     * @return wrapped stored fields reader
+     */
+    public static StoredFieldsReader wrapForMerge(StoredFieldsReader storedFieldsReader) {
+        if (storedFieldsReader instanceof DerivedSourceStoredFieldsReader) {
+            return ((DerivedSourceStoredFieldsReader) storedFieldsReader).cloneForMerge();
+        }
+        return storedFieldsReader;
+    }
 }

src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsWriter.java

@@ -65,6 +65,10 @@ public void writeField(FieldInfo info, DataInput value, int length) throws IOExc
 
     @Override
     public int merge(MergeState mergeState) throws IOException {
+        // We have to wrap these here to avoid storing the vectors during merge
+        for (int i = 0; i < mergeState.storedFieldsReaders.length; i++) {
+            mergeState.storedFieldsReaders[i] = DerivedSourceStoredFieldsReader.wrapForMerge(mergeState.storedFieldsReaders[i]);
+        }
         return delegate.merge(mergeState);
     }