Break bwc to test tests

Signed-off-by: John Mazanec <[email protected]>

Author: jmazanec15

qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/DerivedSourceBWCRestartIT.java

@@ -52,7 +52,9 @@ private void testIndexAndForceMergeOnOld_injectOnNew(List<IndexConfigContext> in
 
             // Delete
             testDelete(indexConfigContexts);
+            assertDocsMatch(indexConfigContexts);
         } else {
+            assertDocsMatch(indexConfigContexts);
             // Search
             testSearch(indexConfigContexts);
 
@@ -81,7 +83,9 @@ private void testIndexOnOld_forceMergeAndInjectOnNew(List<IndexConfigContext> in
         if (isRunningAgainstOldCluster()) {
             prepareOriginalIndices(indexConfigContexts);
         } else {
+            assertDocsMatch(indexConfigContexts);
             testMerging(indexConfigContexts);
+            assertDocsMatch(indexConfigContexts);
             // Update. Skipping update tests for nested docs for now. Will add in the future.
             if (indexConfigContexts.get(0).isNested() == false) {
                 testUpdate(indexConfigContexts);

src/main/java/org/opensearch/knn/index/codec/KNN10010Codec/KNN10010DerivedSourceStoredFieldsWriter.java

@@ -81,13 +81,14 @@ public int merge(MergeState mergeState) throws IOException {
         // In case of backwards compatibility, with old segments, we need to perform a non-optimal merge. Basically, it
         // will repopulate each source and then inject the vector and then remove it. This allows us to migrate
         // segments from filter approach to mask approach
-        if (KNN9120DerivedSourceStoredFieldsReader.doesMergeContainLegacySegments(mergeState)) {
-            super.merge(mergeState);
-        }
+        // if (KNN9120DerivedSourceStoredFieldsReader.doesMergeContainLegacySegments(mergeState)) {
+        // super.merge(mergeState);
+        // }
 
         // We wrap the segments to avoid injecting back vectors and then removing. If this is not done, then we will
         // inject and then just write to disk potentially.
         for (int i = 0; i < mergeState.storedFieldsReaders.length; i++) {
+            mergeState.storedFieldsReaders[i] = KNN9120DerivedSourceStoredFieldsReader.wrapForMerge(mergeState.storedFieldsReaders[i]);
             mergeState.storedFieldsReaders[i] = KNN10010DerivedSourceStoredFieldsReader.wrapForMerge(mergeState.storedFieldsReaders[i]);
         }
         return delegate.merge(mergeState);

src/main/java/org/opensearch/knn/index/codec/backward_codecs/KNN9120Codec/KNN9120DerivedSourceStoredFieldsReader.java

@@ -132,4 +132,32 @@ private static boolean doesSegmentContainLegacyFields(FieldInfos fieldInfos) {
         }
         return false;
     }
+
+    private StoredFieldsReader cloneForMerge() {
+        try {
+            return new KNN9120DerivedSourceStoredFieldsReader(
+                delegate.getMergeInstance(),
+                derivedVectorFields,
+                derivedSourceReadersSupplier,
+                segmentReadState,
+                false
+            );
+        } catch (IOException e) {
+            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 KNN9120DerivedSourceStoredFieldsReader) {
+            return ((KNN9120DerivedSourceStoredFieldsReader) storedFieldsReader).cloneForMerge();
+        }
+        return storedFieldsReader;
+    }
 }

src/testFixtures/java/org/opensearch/knn/DerivedSourceTestCase.java

@@ -31,8 +31,8 @@
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
 
 public class DerivedSourceTestCase extends KNNRestTestCase {
-    protected final int TEST_DIMENSION = 128;
-    protected final int DOCS = 50;
+    protected final int TEST_DIMENSION = 16;
+    protected final int DOCS = 500;
     protected final static String NESTED_NAME = "test_nested";
     protected final static String FIELD_NAME = "test_vector";
 
@@ -78,6 +78,8 @@ protected void prepareOriginalIndices(List<IndexConfigContext> indexConfigContex
             derivedSourceDisabledContext.indexName,
             derivedSourceEnabledContext.indexName
         );
+        flush(derivedSourceEnabledContext.indexName, true);
+        flush(derivedSourceDisabledContext.indexName, true);
     }
 
     @SneakyThrows
@@ -86,8 +88,9 @@ protected void testMerging(List<IndexConfigContext> indexConfigContexts) {
         IndexConfigContext derivedSourceDisabledContext = indexConfigContexts.get(1);
         String originalIndexNameDerivedSourceEnabled = derivedSourceEnabledContext.indexName;
         String originalIndexNameDerivedSourceDisabled = derivedSourceDisabledContext.indexName;
-        forceMergeKnnIndex(originalIndexNameDerivedSourceEnabled, 10);
-        forceMergeKnnIndex(originalIndexNameDerivedSourceDisabled, 10);
+        forceMergeKnnIndex(originalIndexNameDerivedSourceEnabled, 1);
+        forceMergeKnnIndex(originalIndexNameDerivedSourceDisabled, 1);
+
         refreshAllIndices();
         assertIndexBigger(originalIndexNameDerivedSourceDisabled, originalIndexNameDerivedSourceEnabled);
         assertDocsMatch(
@@ -105,6 +108,20 @@ protected void testMerging(List<IndexConfigContext> indexConfigContexts) {
             originalIndexNameDerivedSourceDisabled,
             originalIndexNameDerivedSourceEnabled
         );
+        flush(derivedSourceEnabledContext.indexName, true);
+        flush(derivedSourceDisabledContext.indexName, true);
+    }
+
+    public void assertDocsMatch(List<IndexConfigContext> indexConfigContexts) {
+        IndexConfigContext derivedSourceEnabledContext = indexConfigContexts.get(0);
+        IndexConfigContext derivedSourceDisabledContext = indexConfigContexts.get(1);
+        String originalIndexNameDerivedSourceEnabled = derivedSourceEnabledContext.indexName;
+        String originalIndexNameDerivedSourceDisabled = derivedSourceDisabledContext.indexName;
+        assertDocsMatch(
+            derivedSourceDisabledContext.docCount,
+            originalIndexNameDerivedSourceDisabled,
+            originalIndexNameDerivedSourceEnabled
+        );
     }
 
     @SneakyThrows

src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java

@@ -772,7 +772,7 @@ protected void addKnnDoc(String index, String docId, List<String> fieldNames, Li
      * Adds a doc where document is represented as a string.
      */
     protected void addKnnDoc(final String index, final String docId, final String document) throws IOException {
-        Request request = new Request("POST", "/" + index + "/_doc/" + docId);
+        Request request = new Request("POST", "/" + index + "/_doc/" + docId + "?refresh=true");
         request.setJsonEntity(document);
         client().performRequest(request);
     }