Adding support for context aware segments

Signed-off-by: RS146BIJAY <[email protected]>

Author: RS146BIJAY

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add a toBuilder method in EngineConfig to support easy modification of configs([#19054](https://github.com/opensearch-project/OpenSearch/pull/19054))
 - Add StoreFactory plugin interface for custom Store implementations([#19091](https://github.com/opensearch-project/OpenSearch/pull/19091))
 - Add a dynamic setting to change skip_cache_factor and min_frequency for querycache ([#18351](https://github.com/opensearch-project/OpenSearch/issues/18351))
+- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098))
 
 ### Changed
 - Add CompletionStage variants to methods in the Client Interface and default to ActionListener impl ([#18998](https://github.com/opensearch-project/OpenSearch/pull/18998))

plugins/store-smb/src/main/java/org/opensearch/index/store/smbmmapfs/SmbMmapFsDirectoryFactory.java

@@ -47,7 +47,7 @@
 public final class SmbMmapFsDirectoryFactory extends FsDirectoryFactory {
 
     @Override
-    protected Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
+    public Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
         return new SmbDirectoryWrapper(
             setPreload(
                 new MMapDirectory(location, lockFactory),

plugins/store-smb/src/main/java/org/opensearch/index/store/smbniofs/SmbNIOFsDirectoryFactory.java

@@ -24,7 +24,7 @@
 public final class SmbNIOFsDirectoryFactory extends FsDirectoryFactory {
 
     @Override
-    protected Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
+    public Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
         return new SmbDirectoryWrapper(new NIOFSDirectory(location, lockFactory));
     }
 }

release-notes/opensearch.release-notes-3.0.0-alpha1.md

@@ -55,6 +55,8 @@
 - Fix Bug - Handle unsigned long in sorting order assertion of LongHashSet ([#17207](https://github.com/opensearch-project/OpenSearch/pull/17207))
 - Implemented computation of segment replication stats at shard level ([#17055](https://github.com/opensearch-project/OpenSearch/pull/17055))
 - [Rule Based Auto-tagging] Add in-memory attribute value store ([#17342](https://github.com/opensearch-project/OpenSearch/pull/17342))
+- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098)
+
 
 ### Dependencies
 - Bump Apache HttpCore5/HttpClient5 dependencies from 5.2.5/5.3.1 to 5.3.1/5.4.1 to support ExtendedSocketOption in HttpAsyncClient ([#16757](https://github.com/opensearch-project/OpenSearch/pull/16757))

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

@@ -0,0 +1,31 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+public class CriteriaBasedCompositeDirectory extends FilterDirectory {
+    /**
+     * Sole constructor, typically called from sub-classes.
+     *
+     * @param in
+     */
+    protected CriteriaBasedCompositeDirectory(Directory in) {
+        super(in);
+    }
+
+    @Override
+    public String[] listAll() throws IOException {
+        return Arrays.stream(super.listAll()).filter(fileName -> !fileName.startsWith("temp_")).distinct().toArray(String[]::new);
+    }
+}

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

@@ -0,0 +1,60 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index;
+
+import org.apache.lucene.index.MergeTrigger;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentInfos;
+import org.apache.lucene.index.TieredMergePolicy;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class CriteriaBasedMergePolicy extends TieredMergePolicy {
+    @Override
+    public MergeSpecification findMerges(
+        MergeTrigger mergeTrigger, SegmentInfos infos, MergeContext mergeContext) throws IOException {
+        final Set<SegmentCommitInfo> merging = mergeContext.getMergingSegments();
+        MergeSpecification spec = null;
+        final Map<String, List<SegmentCommitInfo>> commitInfos = new HashMap<>();
+        for (SegmentCommitInfo si : infos) {
+            if (merging.contains(si)) {
+                continue;
+            }
+
+            final String dwptGroupNumber = si.info.getAttribute("criteria");
+            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);
+                }
+
+                final MergeSpecification tieredMergePolicySpec =
+                    super.findMerges(mergeTrigger, newSIS, mergeContext);
+                if (tieredMergePolicySpec != null) {
+                    if (spec == null) {
+                        spec = new MergeSpecification();
+                    }
+
+                    spec.merges.addAll(tieredMergePolicySpec.merges);
+                }
+            }
+        }
+
+        return spec;
+    }
+}

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

@@ -38,6 +38,8 @@
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.search.similarities.BM25Similarity;
 import org.apache.lucene.search.similarities.Similarity;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.LockFactory;
 import org.apache.lucene.util.Constants;
 import org.opensearch.Version;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
@@ -76,6 +78,7 @@
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.shard.IndexingOperationListener;
 import org.opensearch.index.shard.SearchOperationListener;
+import org.opensearch.index.shard.ShardPath;
 import org.opensearch.index.similarity.SimilarityService;
 import org.opensearch.index.store.DefaultCompositeDirectoryFactory;
 import org.opensearch.index.store.FsDirectoryFactory;
@@ -97,6 +100,7 @@
 import org.opensearch.transport.client.Client;
 
 import java.io.IOException;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -837,7 +841,18 @@ private static IndexStorePlugin.DirectoryFactory getDirectoryFactory(
                 throw new IllegalArgumentException("Unknown store type [" + storeType + "]");
             }
         }
-        return factory;
+
+        return new IndexStorePlugin.DirectoryFactory() {
+            @Override
+            public Directory newDirectory(IndexSettings indexSettings, ShardPath shardPath) throws IOException {
+                return new CriteriaBasedCompositeDirectory(factory.newDirectory(indexSettings, shardPath));
+            }
+
+            @Override
+            public Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
+                return factory.newFSDirectory(location, lockFactory, indexSettings);
+            }
+        };
     }
 
     private static IndexStorePlugin.CompositeDirectoryFactory getCompositeDirectoryFactory(

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

@@ -718,7 +718,7 @@ protected void closeInternal() {
                         // Do nothing for shard lock on remote store
                     }
                 };
-                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, remoteStoreLock, Store.OnClose.EMPTY, path);
+                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, remoteStoreLock, Store.OnClose.EMPTY, path, directoryFactory);
             } else {
                 // Disallow shards with remote store based settings to be created on non-remote store enabled nodes
                 // Even though we have `RemoteStoreMigrationAllocationDecider` in place to prevent something like this from happening at the
@@ -753,7 +753,8 @@ protected void closeInternal() {
                 directory,
                 lock,
                 new StoreCloseListener(shardId, () -> eventListener.onStoreClosed(shardId)),
-                path
+                path,
+                directoryFactory
             );
             eventListener.onStoreCreated(shardId);
             indexShard = new IndexShard(

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

@@ -0,0 +1,49 @@
+/*Add commentMore actions
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.codec;
+
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.FilterCodec;
+import org.apache.lucene.codecs.SegmentInfoFormat;
+import org.apache.lucene.codecs.lucene101.Lucene101Codec;
+import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+
+import java.io.IOException;
+
+public class CriteriaBasedCodec extends FilterCodec {
+
+    private final String criteria;
+    public CriteriaBasedCodec() {
+        super("CriteriaBasedCodec", new Lucene101Codec());
+        criteria = null;
+    }
+
+    public CriteriaBasedCodec(Codec delegate, String criteria) {
+        super("CriteriaBasedCodec", delegate);
+        this.criteria = criteria;
+    }
+
+    @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 {
+                info.putAttribute("criteria", criteria);
+                delegate.segmentInfoFormat().write(directory, info, ioContext);
+            }
+        };
+    }
+}