diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN10010Codec/KNN10010Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN10010Codec/KNN10010Codec.java index 73371bb61..acec0fa7f 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN10010Codec/KNN10010Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN10010Codec/KNN10010Codec.java @@ -104,6 +104,11 @@ private StoredFieldsFormat getStoredFieldsFormat() { } return null; })); - return new DerivedSourceStoredFieldsFormat(delegate.storedFieldsFormat(), derivedSourceReadersSupplier, mapperService); + return new DerivedSourceStoredFieldsFormat( + delegate.storedFieldsFormat(), + derivedSourceReadersSupplier, + mapperService, + delegate.getName() + ); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsFormat.java index e60b82b2e..109431d08 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/DerivedSourceStoredFieldsFormat.java @@ -6,6 +6,7 @@ package org.opensearch.knn.index.codec.KNN9120Codec; import lombok.AllArgsConstructor; +import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.StoredFieldsFormat; import org.apache.lucene.codecs.StoredFieldsReader; import org.apache.lucene.codecs.StoredFieldsWriter; @@ -32,15 +33,26 @@ @AllArgsConstructor public class DerivedSourceStoredFieldsFormat extends StoredFieldsFormat { + private static final String DELEGATE_CODEC_KEY = "knn_delegate_codec"; + private final StoredFieldsFormat delegate; private final DerivedSourceReadersSupplier derivedSourceReadersSupplier; // IMPORTANT Do not rely on this for the reader, it will be null if SPI is used @Nullable private final MapperService mapperService; + // IMPORTANT Do not rely on this for the reader, it will be null if SPI is used + @Nullable + private final String delegateCodecName; @Override public StoredFieldsReader fieldsReader(Directory directory, SegmentInfo segmentInfo, FieldInfos fieldInfos, IOContext ioContext) throws IOException { + StoredFieldsFormat delegateFromWriting = delegate; + if (segmentInfo.getAttribute(DELEGATE_CODEC_KEY) != null) { + String delegateCodecName = segmentInfo.getAttribute(DELEGATE_CODEC_KEY); + delegateFromWriting = Codec.forName(delegateCodecName).storedFieldsFormat(); + } + List derivedVectorFields = null; for (FieldInfo fieldInfo : fieldInfos) { if (DERIVED_VECTOR_FIELD_ATTRIBUTE_TRUE_VALUE.equals(fieldInfo.attributes().get(DERIVED_VECTOR_FIELD_ATTRIBUTE_KEY))) { @@ -53,10 +65,10 @@ public StoredFieldsReader fieldsReader(Directory directory, SegmentInfo segmentI } // If no fields have it enabled, we can just short-circuit and return the delegate's fieldReader if (derivedVectorFields == null || derivedVectorFields.isEmpty()) { - return delegate.fieldsReader(directory, segmentInfo, fieldInfos, ioContext); + return delegateFromWriting.fieldsReader(directory, segmentInfo, fieldInfos, ioContext); } return new DerivedSourceStoredFieldsReader( - delegate.fieldsReader(directory, segmentInfo, fieldInfos, ioContext), + delegateFromWriting.fieldsReader(directory, segmentInfo, fieldInfos, ioContext), derivedVectorFields, derivedSourceReadersSupplier, new SegmentReadState(directory, segmentInfo, fieldInfos, ioContext) @@ -65,6 +77,24 @@ public StoredFieldsReader fieldsReader(Directory directory, SegmentInfo segmentI @Override public StoredFieldsWriter fieldsWriter(Directory directory, SegmentInfo segmentInfo, IOContext ioContext) throws IOException { + // We write the delegate codec name into the segmentInfo attributes so that we can read it when loading the + // codec from SPI. + // This is similar to whats done in + // https://github.com/opensearch-project/custom-codecs/blob/2.19.0.0/src/main/java/org/opensearch/index/codec/customcodecs/Lucene912CustomStoredFieldsFormat.java#L95-L100 + String previous = segmentInfo.putAttribute(DELEGATE_CODEC_KEY, delegateCodecName); + if (previous != null && previous.equals(delegateCodecName) == false) { + throw new IllegalStateException( + "found existing value for " + + DELEGATE_CODEC_KEY + + " for segment: " + + segmentInfo.name + + " old = " + + previous + + ", new = " + + delegateCodecName + ); + } + StoredFieldsWriter delegateWriter = delegate.fieldsWriter(directory, segmentInfo, ioContext); if (mapperService != null && KNNSettings.isKNNDerivedSourceEnabled(mapperService.getIndexSettings().getSettings())) { List vectorFieldTypes = new ArrayList<>(); diff --git a/src/main/java/org/opensearch/knn/index/codec/backward_codecs/KNN9120Codec/KNN9120Codec.java b/src/main/java/org/opensearch/knn/index/codec/backward_codecs/KNN9120Codec/KNN9120Codec.java index 375b9fe48..ecd02705b 100644 --- a/src/main/java/org/opensearch/knn/index/codec/backward_codecs/KNN9120Codec/KNN9120Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/backward_codecs/KNN9120Codec/KNN9120Codec.java @@ -100,6 +100,11 @@ private StoredFieldsFormat getStoredFieldsFormat() { } return null; })); - return new DerivedSourceStoredFieldsFormat(delegate.storedFieldsFormat(), derivedSourceReadersSupplier, mapperService); + return new DerivedSourceStoredFieldsFormat( + delegate.storedFieldsFormat(), + derivedSourceReadersSupplier, + mapperService, + delegate.getName() + ); } }