Migrate derived source from filter to mask

Migrates derived source functionality from filter to mask based
approach. Moves old read functionality and related classes to backwards
codecs KNN9120... Removes old write as no longer necessary.

In order to support bwc, we add custom functionality in the stored
fields writer merge logic to fall back to base, non-optimized merge if
it detects older readers in the merge state. This is needed because for
these segments, we need to rebuild the source and then apply filter to
migrate to new write format.

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/KNN10010Codec.java

@@ -91,18 +91,7 @@ private StoredFieldsFormat getStoredFieldsFormat() {
             }
             return null;
 
-        }, (segmentReadState) -> {
-            if (segmentReadState.fieldInfos.hasPostings()) {
-                return postingsFormat().fieldsProducer(segmentReadState);
-            }
-            return null;
-
-        }, (segmentReadState -> {
-            if (segmentReadState.fieldInfos.hasNorms()) {
-                return normsFormat().normsProducer(segmentReadState);
-            }
-            return null;
-        }));
+        });
         return new KNN10010DerivedSourceStoredFieldsFormat(delegate.storedFieldsFormat(), derivedSourceReadersSupplier, mapperService);
     }
 }

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

@@ -9,7 +9,6 @@
 import org.apache.lucene.codecs.StoredFieldsFormat;
 import org.apache.lucene.codecs.StoredFieldsReader;
 import org.apache.lucene.codecs.StoredFieldsWriter;
-import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.index.SegmentReadState;
@@ -19,14 +18,15 @@
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.knn.index.KNNSettings;
-import org.opensearch.knn.index.codec.KNN9120Codec.KNN9120DerivedSourceStoredFieldsReader;
+import org.opensearch.knn.index.codec.derivedsource.DerivedFieldInfo;
 import org.opensearch.knn.index.codec.derivedsource.DerivedSourceReadersSupplier;
 import org.opensearch.knn.index.codec.derivedsource.DerivedSourceSegmentAttributeParser;
 import org.opensearch.knn.index.mapper.KNNVectorFieldType;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Stream;
 
 @AllArgsConstructor
 public class KNN10010DerivedSourceStoredFieldsFormat extends StoredFieldsFormat {
@@ -40,16 +40,22 @@ public class KNN10010DerivedSourceStoredFieldsFormat extends StoredFieldsFormat
     @Override
     public StoredFieldsReader fieldsReader(Directory directory, SegmentInfo segmentInfo, FieldInfos fieldInfos, IOContext ioContext)
         throws IOException {
-        List<FieldInfo> derivedVectorFields = DerivedSourceSegmentAttributeParser.parseDerivedVectorFields(segmentInfo)
-            .stream()
-            .filter(field -> fieldInfos.fieldInfo(field) != null)
-            .map(fieldInfos::fieldInfo)
-            .toList();
+        List<DerivedFieldInfo> derivedVectorFields = Stream.concat(
+            DerivedSourceSegmentAttributeParser.parseDerivedVectorFields(segmentInfo, false)
+                .stream()
+                .filter(field -> fieldInfos.fieldInfo(field) != null)
+                .map(field -> new DerivedFieldInfo(fieldInfos.fieldInfo(field), false)),
+            DerivedSourceSegmentAttributeParser.parseDerivedVectorFields(segmentInfo, true)
+                .stream()
+                .filter(field -> fieldInfos.fieldInfo(field) != null)
+                .map(field -> new DerivedFieldInfo(fieldInfos.fieldInfo(field), true))
+        ).toList();
+
         // If no fields have it enabled, we can just short-circuit and return the delegate's fieldReader
-        if (derivedVectorFields == null || derivedVectorFields.isEmpty()) {
+        if (derivedVectorFields.isEmpty()) {
             return delegate.fieldsReader(directory, segmentInfo, fieldInfos, ioContext);
         }
-        return new KNN9120DerivedSourceStoredFieldsReader(
+        return new KNN10010DerivedSourceStoredFieldsReader(
             delegate.fieldsReader(directory, segmentInfo, fieldInfos, ioContext),
             derivedVectorFields,
             derivedSourceReadersSupplier,
@@ -60,18 +66,34 @@ public StoredFieldsReader fieldsReader(Directory directory, SegmentInfo segmentI
     @Override
     public StoredFieldsWriter fieldsWriter(Directory directory, SegmentInfo segmentInfo, IOContext ioContext) throws IOException {
         StoredFieldsWriter delegateWriter = delegate.fieldsWriter(directory, segmentInfo, ioContext);
-        if (mapperService != null && KNNSettings.isKNNDerivedSourceEnabled(mapperService.getIndexSettings().getSettings())) {
-            List<String> vectorFieldTypes = new ArrayList<>();
-            for (MappedFieldType fieldType : mapperService.fieldTypes()) {
-                if (fieldType instanceof KNNVectorFieldType) {
+        if (mapperService == null || KNNSettings.isKNNDerivedSourceEnabled(mapperService.getIndexSettings().getSettings()) == false) {
+            return delegateWriter;
+        }
+
+        List<String> vectorFieldTypes = new ArrayList<>();
+        List<String> nestedVectorFieldTypes = new ArrayList<>();
+        for (MappedFieldType fieldType : mapperService.fieldTypes()) {
+            if (fieldType instanceof KNNVectorFieldType) {
+                boolean isNested = mapperService.documentMapper().mappers().getNestedScope(fieldType.name()) != null;
+                if (isNested) {
+                    nestedVectorFieldTypes.add(fieldType.name());
+                } else {
                     vectorFieldTypes.add(fieldType.name());
                 }
             }
-            if (vectorFieldTypes.isEmpty() == false) {
-                DerivedSourceSegmentAttributeParser.addDerivedVectorFieldsSegmentInfoAttribute(segmentInfo, vectorFieldTypes);
-                return new KNN10010DerivedSourceStoredFieldsWriter(delegateWriter, vectorFieldTypes);
-            }
         }
-        return delegateWriter;
+        if (vectorFieldTypes.isEmpty() && nestedVectorFieldTypes.isEmpty()) {
+            return delegateWriter;
+        }
+
+        // Store nested fields separately from non-nested for easy handling on read
+        if (vectorFieldTypes.isEmpty() == false) {
+            DerivedSourceSegmentAttributeParser.addDerivedVectorFieldsSegmentInfoAttribute(segmentInfo, vectorFieldTypes, false);
+        }
+        if (nestedVectorFieldTypes.isEmpty() == false) {
+            vectorFieldTypes.addAll(nestedVectorFieldTypes);
+            DerivedSourceSegmentAttributeParser.addDerivedVectorFieldsSegmentInfoAttribute(segmentInfo, nestedVectorFieldTypes, true);
+        }
+        return new KNN10010DerivedSourceStoredFieldsWriter(delegateWriter, vectorFieldTypes);
     }
 }

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

@@ -3,29 +3,29 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.knn.index.codec.KNN9120Codec;
+package org.opensearch.knn.index.codec.KNN10010Codec;
 
 import org.apache.lucene.codecs.StoredFieldsReader;
-import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.StoredFieldVisitor;
 import org.apache.lucene.util.IOUtils;
 import org.opensearch.index.fieldvisitor.FieldsVisitor;
+import org.opensearch.knn.index.codec.derivedsource.DerivedFieldInfo;
 import org.opensearch.knn.index.codec.derivedsource.DerivedSourceReadersSupplier;
 import org.opensearch.knn.index.codec.derivedsource.DerivedSourceStoredFieldVisitor;
-import org.opensearch.knn.index.codec.derivedsource.DerivedSourceVectorInjector;
+import org.opensearch.knn.index.codec.derivedsource.DerivedSourceVectorTransformer;
 
 import java.io.IOException;
 import java.util.List;
 
-public class KNN9120DerivedSourceStoredFieldsReader extends StoredFieldsReader {
+public class KNN10010DerivedSourceStoredFieldsReader extends StoredFieldsReader {
     private final StoredFieldsReader delegate;
-    private final List<FieldInfo> derivedVectorFields;
+    private final List<DerivedFieldInfo> derivedVectorFields;
     private final DerivedSourceReadersSupplier derivedSourceReadersSupplier;
     private final SegmentReadState segmentReadState;
     private final boolean shouldInject;
 
-    private final DerivedSourceVectorInjector derivedSourceVectorInjector;
+    private final DerivedSourceVectorTransformer derivedSourceVectorTransformer;
 
     /**
      *
@@ -35,18 +35,18 @@ public class KNN9120DerivedSourceStoredFieldsReader extends StoredFieldsReader {
      * @param segmentReadState SegmentReadState for the segment
      * @throws IOException in case of I/O error
      */
-    public KNN9120DerivedSourceStoredFieldsReader(
+    public KNN10010DerivedSourceStoredFieldsReader(
         StoredFieldsReader delegate,
-        List<FieldInfo> derivedVectorFields,
+        List<DerivedFieldInfo> derivedVectorFields,
         DerivedSourceReadersSupplier derivedSourceReadersSupplier,
         SegmentReadState segmentReadState
     ) throws IOException {
         this(delegate, derivedVectorFields, derivedSourceReadersSupplier, segmentReadState, true);
     }
 
-    private KNN9120DerivedSourceStoredFieldsReader(
+    private KNN10010DerivedSourceStoredFieldsReader(
         StoredFieldsReader delegate,
-        List<FieldInfo> derivedVectorFields,
+        List<DerivedFieldInfo> derivedVectorFields,
         DerivedSourceReadersSupplier derivedSourceReadersSupplier,
         SegmentReadState segmentReadState,
         boolean shouldInject
@@ -56,25 +56,25 @@ private KNN9120DerivedSourceStoredFieldsReader(
         this.derivedSourceReadersSupplier = derivedSourceReadersSupplier;
         this.segmentReadState = segmentReadState;
         this.shouldInject = shouldInject;
-        this.derivedSourceVectorInjector = createDerivedSourceVectorInjector();
+        this.derivedSourceVectorTransformer = createDerivedSourceVectorTransformer();
     }
 
-    private DerivedSourceVectorInjector createDerivedSourceVectorInjector() throws IOException {
-        return new DerivedSourceVectorInjector(derivedSourceReadersSupplier, segmentReadState, derivedVectorFields);
+    private DerivedSourceVectorTransformer createDerivedSourceVectorTransformer() throws IOException {
+        return new DerivedSourceVectorTransformer(derivedSourceReadersSupplier, segmentReadState, derivedVectorFields);
     }
 
     @Override
     public void document(int docId, StoredFieldVisitor storedFieldVisitor) throws IOException {
         // If the visitor has explicitly indicated it does not need the fields, we should not inject them
         boolean isVisitorNeedFields = true;
         if (storedFieldVisitor instanceof FieldsVisitor) {
-            isVisitorNeedFields = derivedSourceVectorInjector.shouldInject(
+            isVisitorNeedFields = derivedSourceVectorTransformer.shouldInject(
                 ((FieldsVisitor) storedFieldVisitor).includes(),
                 ((FieldsVisitor) storedFieldVisitor).excludes()
             );
         }
         if (shouldInject && isVisitorNeedFields) {
-            delegate.document(docId, new DerivedSourceStoredFieldVisitor(storedFieldVisitor, docId, derivedSourceVectorInjector));
+            delegate.document(docId, new DerivedSourceStoredFieldVisitor(storedFieldVisitor, docId, derivedSourceVectorTransformer));
             return;
         }
         delegate.document(docId, storedFieldVisitor);
@@ -83,7 +83,7 @@ public void document(int docId, StoredFieldVisitor storedFieldVisitor) throws IO
     @Override
     public StoredFieldsReader clone() {
         try {
-            return new KNN9120DerivedSourceStoredFieldsReader(
+            return new KNN10010DerivedSourceStoredFieldsReader(
                 delegate.clone(),
                 derivedVectorFields,
                 derivedSourceReadersSupplier,
@@ -102,7 +102,7 @@ public void checkIntegrity() throws IOException {
 
     @Override
     public void close() throws IOException {
-        IOUtils.close(delegate, derivedSourceVectorInjector);
+        IOUtils.close(delegate, derivedSourceVectorTransformer);
     }
 
     /**
@@ -114,7 +114,7 @@ public void close() throws IOException {
      */
     private StoredFieldsReader cloneForMerge() {
         try {
-            return new KNN9120DerivedSourceStoredFieldsReader(
+            return new KNN10010DerivedSourceStoredFieldsReader(
                 delegate.getMergeInstance(),
                 derivedVectorFields,
                 derivedSourceReadersSupplier,
@@ -134,8 +134,8 @@ private StoredFieldsReader cloneForMerge() {
      * @return wrapped stored fields reader
      */
     public static StoredFieldsReader wrapForMerge(StoredFieldsReader storedFieldsReader) {
-        if (storedFieldsReader instanceof KNN9120DerivedSourceStoredFieldsReader) {
-            return ((KNN9120DerivedSourceStoredFieldsReader) storedFieldsReader).cloneForMerge();
+        if (storedFieldsReader instanceof KNN10010DerivedSourceStoredFieldsReader) {
+            return ((KNN10010DerivedSourceStoredFieldsReader) storedFieldsReader).cloneForMerge();
         }
         return storedFieldsReader;
     }

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

@@ -5,7 +5,6 @@
 
 package org.opensearch.knn.index.codec.KNN10010Codec;
 
-import lombok.RequiredArgsConstructor;
 import org.apache.lucene.codecs.StoredFieldsWriter;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.MergeState;
@@ -14,25 +13,46 @@
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.xcontent.XContentHelper;
-import org.opensearch.common.xcontent.support.XContentMapValues;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaType;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.SourceFieldMapper;
-import org.opensearch.knn.index.codec.KNN9120Codec.KNN9120DerivedSourceStoredFieldsReader;
+import org.opensearch.knn.index.codec.backward_codecs.KNN9120Codec.KNN9120DerivedSourceStoredFieldsReader;
+import org.opensearch.knn.index.codec.derivedsource.DerivedMapHelper;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
-@RequiredArgsConstructor
 public class KNN10010DerivedSourceStoredFieldsWriter extends StoredFieldsWriter {
 
     private final StoredFieldsWriter delegate;
-    private final List<String> vectorFieldTypes;
+    private final Function<Map<String, Object>, Map<String, Object>> vectorMask;
+
+    // Keeping the mask as small as possible.
+    private final static char MASK = 'z';
+
+    /**
+     *
+     * @param delegate StoredFieldsWriter to wrap
+     * @param vectorFieldTypesArg List of vector field types to mask. If empty, no masking will be done
+     */
+    public KNN10010DerivedSourceStoredFieldsWriter(StoredFieldsWriter delegate, List<String> vectorFieldTypesArg) {
+        this.delegate = delegate;
+        List<String> vectorFieldTypes = vectorFieldTypesArg.stream().map(String::toLowerCase).toList();
+        if (vectorFieldTypes.isEmpty() == false) {
+            this.vectorMask = DerivedMapHelper.transform(
+                vectorFieldTypes.stream().collect(Collectors.toMap(k -> k, k -> (Object o) -> o == null ? o : MASK))
+            );
+        } else {
+            this.vectorMask = null;
+        }
+    }
 
     @Override
     public void startDocument() throws IOException {
@@ -66,26 +86,33 @@ 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
+        // 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);
+        }
+
+        // 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);
     }
 
     @Override
     public void writeField(FieldInfo fieldInfo, BytesRef bytesRef) throws IOException {
         // Parse out the vectors from the source
-        if (Objects.equals(fieldInfo.name, SourceFieldMapper.NAME) && !vectorFieldTypes.isEmpty()) {
+        if (Objects.equals(fieldInfo.name, SourceFieldMapper.NAME) && vectorMask != null) {
             // Reference:
             // https://github.com/opensearch-project/OpenSearch/blob/2.18.0/server/src/main/java/org/opensearch/index/mapper/SourceFieldMapper.java#L322
             Tuple<? extends MediaType, Map<String, Object>> mapTuple = XContentHelper.convertToMap(
                 BytesReference.fromByteBuffer(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length)),
                 true,
                 MediaTypeRegistry.JSON
             );
-            Map<String, Object> filteredSource = XContentMapValues.filter(null, vectorFieldTypes.toArray(new String[0]))
-                .apply(mapTuple.v2());
+            Map<String, Object> filteredSource = vectorMask.apply(mapTuple.v2());
             BytesStreamOutput bStream = new BytesStreamOutput();
             MediaType actualContentType = mapTuple.v1();
             XContentBuilder builder = MediaTypeRegistry.contentBuilder(actualContentType, bStream).map(filteredSource);

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

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.knn.index.codec.derivedsource;
+package org.opensearch.knn.index.codec.backward_codecs.KNN9120Codec;
 
 import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.index.FieldInfo;