Fix CriteriaBasedCodec to work with delegate codec (opensearch-project#20442)

Signed-off-by: RS146BIJAY <rishavsagar4b1@gmail.com>

Author: RS146BIJAY

CHANGELOG.md

@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Update EncryptedBlobContainer to adhere limits while listing blobs in specific sort order if wrapped blob container supports ([#20514](https://github.com/opensearch-project/OpenSearch/pull/20514))
 - [segment replication] Fix segment replication infinite retry due to stale metadata checkpoint ([#20551](https://github.com/opensearch-project/OpenSearch/pull/20551))
 - Changing opensearch.cgroups.hierarchy.override causes java.lang.SecurityException exception ([#20565](https://github.com/opensearch-project/OpenSearch/pull/20565))
+- Fix CriteriaBasedCodec to work with delegate codec. ([20442](https://github.com/opensearch-project/OpenSearch/pull/20442))
 
 ### Dependencies
 - Bump `ch.qos.logback:logback-core` and `ch.qos.logback:logback-classic` from 1.5.24 to 1.5.27 ([#20525](https://github.com/opensearch-project/OpenSearch/pull/20525))

server/src/main/java/org/opensearch/index/CriteriaBasedMergePolicy.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.index.MergeTrigger;
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentInfos;
+import org.opensearch.common.CheckedFunction;
 import org.opensearch.index.codec.CriteriaBasedCodec;
 
 import java.io.IOException;
@@ -36,47 +37,106 @@ public CriteriaBasedMergePolicy(MergePolicy in) {
         this.in = in;
     }
 
-    /**
-     * Merges the segments belonging to same group.
-     *
-     * @param mergeTrigger the event that triggered the merge
-     * @param infos the total set of segments in the index
-     * @param mergeContext the IndexWriter to find the merges on
-     * @return
-     * @throws IOException
-     */
-    @Override
-    public MergeSpecification findMerges(MergeTrigger mergeTrigger, SegmentInfos infos, MergeContext mergeContext) throws IOException {
+    private MergeSpecification findMergesInternal(
+        SegmentInfos segmentInfos,
+        MergeContext mergeContext,
+        CheckedFunction<SegmentInfos, MergeSpecification, IOException> mergeFinderFunction
+    ) throws IOException {
+
         final Set<SegmentCommitInfo> merging = mergeContext.getMergingSegments();
-        MergeSpecification spec = null;
         final Map<String, List<SegmentCommitInfo>> commitInfos = new HashMap<>();
-        for (SegmentCommitInfo si : infos) {
+
+        for (SegmentCommitInfo si : segmentInfos) {
             if (merging.contains(si)) {
                 continue;
             }
-
             final String dwptGroupNumber = si.info.getAttribute(CriteriaBasedCodec.BUCKET_NAME);
             commitInfos.computeIfAbsent(dwptGroupNumber, k -> new ArrayList<>()).add(si);
         }
 
-        for (String dwptGroupNumber : commitInfos.keySet()) {
-            if (commitInfos.get(dwptGroupNumber).size() > 1) {
-                final SegmentInfos newSIS = new SegmentInfos(infos.getIndexCreatedVersionMajor());
-                for (SegmentCommitInfo info : commitInfos.get(dwptGroupNumber)) {
-                    newSIS.add(info);
-                }
+        MergeSpecification spec = null;
+        for (Map.Entry<String, List<SegmentCommitInfo>> entry : commitInfos.entrySet()) {
+            List<SegmentCommitInfo> segments = entry.getValue();
+            if (segments.size() > 1) {
+                final SegmentInfos newSIS = new SegmentInfos(segmentInfos.getIndexCreatedVersionMajor());
+                segments.forEach(newSIS::add);
 
-                final MergeSpecification tieredMergePolicySpec = in.findMerges(mergeTrigger, newSIS, mergeContext);
-                if (tieredMergePolicySpec != null) {
+                final MergeSpecification delegateSpec = mergeFinderFunction.apply(newSIS);
+                if (delegateSpec != null) {
                     if (spec == null) {
                         spec = new MergeSpecification();
                     }
-
-                    spec.merges.addAll(tieredMergePolicySpec.merges);
+                    spec.merges.addAll(delegateSpec.merges);
                 }
             }
         }
 
         return spec;
     }
+
+    /**
+     * Merges the segments belonging to same group
+     *
+     * @param mergeTrigger the event that triggered the merge
+     * @param infos the total set of segments in the index
+     * @param mergeContext the IndexWriter to find the merges on
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public MergeSpecification findMerges(MergeTrigger mergeTrigger, SegmentInfos infos, MergeContext mergeContext) throws IOException {
+        return findMergesInternal(infos, mergeContext, newSIS -> in.findMerges(mergeTrigger, newSIS, mergeContext));
+    }
+
+    /**
+     * Force merges segments belonging to same group.
+     *
+     * @param segmentInfos the total set of segments in the index
+     * @param maxSegmentCount requested maximum number of segments in the index
+     * @param segmentsToMerge contains the specific SegmentInfo instances that must be merged away.
+     * @param mergeContext the MergeContext to find the merges on
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public MergeSpecification findForcedMerges(
+        SegmentInfos segmentInfos,
+        int maxSegmentCount,
+        Map<SegmentCommitInfo, Boolean> segmentsToMerge,
+        MergeContext mergeContext
+    ) throws IOException {
+        return findMergesInternal(
+            segmentInfos,
+            mergeContext,
+            newSIS -> in.findForcedMerges(newSIS, maxSegmentCount, segmentsToMerge, mergeContext)
+        );
+    }
+
+    /**
+     * Merges segment belonging to same group to expunge deletes.
+     *
+     * @param segmentInfos the total set of segments in the index
+     * @param mergeContext the MergeContext to find the merges on
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public MergeSpecification findForcedDeletesMerges(SegmentInfos segmentInfos, MergeContext mergeContext) throws IOException {
+        return findMergesInternal(segmentInfos, mergeContext, newSIS -> in.findForcedDeletesMerges(newSIS, mergeContext));
+    }
+
+    /**
+     * Identifies merges that we want to execute (synchronously) on commit
+     *
+     * @param mergeTrigger the event that triggered the merge (COMMIT or GET_READER).
+     * @param segmentInfos the total set of segments in the index (while preparing the commit)
+     * @param mergeContext the MergeContext to find the merges on.
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public MergeSpecification findFullFlushMerges(MergeTrigger mergeTrigger, SegmentInfos segmentInfos, MergeContext mergeContext)
+        throws IOException {
+        return findMergesInternal(segmentInfos, mergeContext, newSIS -> in.findFullFlushMerges(mergeTrigger, newSIS, mergeContext));
+    }
 }

server/src/main/java/org/opensearch/index/codec/CriteriaBasedCodec.java

@@ -9,15 +9,9 @@
 package org.opensearch.index.codec;
 
 import org.apache.lucene.codecs.Codec;
-import org.apache.lucene.codecs.DocValuesFormat;
 import org.apache.lucene.codecs.FilterCodec;
-import org.apache.lucene.codecs.SegmentInfoFormat;
-import org.apache.lucene.codecs.lucene103.Lucene103Codec;
-import org.apache.lucene.index.SegmentInfo;
-import org.apache.lucene.store.Directory;
-import org.apache.lucene.store.IOContext;
-
-import java.io.IOException;
+import org.apache.lucene.codecs.PostingsFormat;
+import org.apache.lucene.codecs.perfield.PerFieldPostingsFormat;
 
 /**
  * Filter codec used to attach bucket attributes to segments of child writer.
@@ -27,46 +21,29 @@ public class CriteriaBasedCodec extends FilterCodec {
 
     private final String bucket;
     public static final String BUCKET_NAME = "bucket";
-    private static final String PLACEHOLDER_BUCKET_FOR_PARENT_WRITER = "-2";
-
-    public CriteriaBasedCodec() {
-        super("CriteriaBasedCodec", new Lucene103Codec());
-        bucket = null;
-    }
+    public static final String ATTRIBUTE_BINDING_TARGET_FIELD = "_id";
 
     public CriteriaBasedCodec(Codec delegate, String bucket) {
-        super("CriteriaBasedCodec", delegate);
+        super(delegate.getName(), delegate);
         this.bucket = bucket;
     }
 
     @Override
-    public SegmentInfoFormat segmentInfoFormat() {
-        return new SegmentInfoFormat() {
-            @Override
-            public SegmentInfo read(Directory directory, String segmentName, byte[] segmentID, IOContext context) throws IOException {
-                return delegate.segmentInfoFormat().read(directory, segmentName, segmentID, context);
-            }
-
-            @Override
-            public void write(Directory directory, SegmentInfo info, IOContext ioContext) throws IOException {
-                if (bucket != null) {
-                    // We will set BUCKET_NAME attribute only for child writer where bucket will set.
-                    info.putAttribute(BUCKET_NAME, bucket);
-                } else if (info.getAttribute(BUCKET_NAME) == null) {
-                    // For segment belonging to parent writer, attributes will be set. In case write went to parent
-                    // writer (like for no ops writes or for temporary tombstone entry which is added for deletes/updates
-                    // to sync version across child and parent writers), segments corresponding to those writer does not
-                    // have
-                    info.putAttribute(BUCKET_NAME, PLACEHOLDER_BUCKET_FOR_PARENT_WRITER);
+    public PostingsFormat postingsFormat() {
+        PostingsFormat format = super.postingsFormat();
+        if (format instanceof PerFieldPostingsFormat) {
+            return new PerFieldPostingsFormat() {
+
+                @Override
+                public PostingsFormat getPostingsFormatForField(String field) {
+                    if (field.equals(ATTRIBUTE_BINDING_TARGET_FIELD)) {
+                        return new CriteriaBasedPostingsFormat(((PerFieldPostingsFormat) format).getPostingsFormatForField(field), bucket);
+                    } else {
+                        return ((PerFieldPostingsFormat) format).getPostingsFormatForField(field);
+                    }
                 }
-
-                delegate.segmentInfoFormat().write(directory, info, ioContext);
-            }
-        };
-    }
-
-    @Override
-    public DocValuesFormat docValuesFormat() {
-        return new CriteriaBasedDocValueFormat(bucket);
+            };
+        }
+        return format;
     }
 }