Adding Lucene Merge supports for pluggable dataformat support (opensearch-project#21422)

Adding Lucene Merge support for pluggable dataformat support

Signed-off-by: Sagar Darji <darsaga@amazon.com>
Co-authored-by: Sagar Darji <darsaga@amazon.com>

Author: darjisagar7

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java

@@ -9,9 +9,12 @@
 package org.opensearch.be.lucene;
 
 import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentReader;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.exec.EngineReaderManager;
+import org.opensearch.index.engine.exec.Segment;
 import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.io.IOException;
@@ -20,6 +23,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.opensearch.be.lucene.index.LuceneWriter.WRITER_GENERATION_ATTRIBUTE;
+
 /**
  * Lucene implementation of {@link EngineReaderManager}.
  * <p>
@@ -72,11 +77,60 @@ public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) th
         }
         DirectoryReader refreshed = DirectoryReader.openIfChanged(currentReader);
         if (refreshed != null) {
+            // Guard against refresh/merge-apply races: a prior IT regression surfaced when
+            // overlapping threads produced a refreshed reader whose leaves disagreed with the
+            // catalog snapshot being registered, effectively pairing the snapshot with a stale
+            // reader. This assert catches that drift in test builds before the mismatched pair
+            // is published to readers.
+            assert readersAreSame(catalogSnapshot, refreshed);
             currentReader = refreshed;
         }
         readers.put(catalogSnapshot, currentReader);
     }
 
+    /**
+     * Consistency check: verifies that the refreshed {@link DirectoryReader} reflects exactly
+     * the set of segments the given {@link CatalogSnapshot} references. Compares the sorted
+     * list of writer generations drawn from the snapshot's {@link Segment Segments} against
+     * the sorted list of writer generations read off each leaf of the reader (via the
+     * {@link org.opensearch.be.lucene.index.LuceneWriter#WRITER_GENERATION_ATTRIBUTE} stamped
+     * onto every Lucene segment at write time).
+     *
+     * <p>Used only in an {@code assert} to catch refresh/catalog drift in test builds — if
+     * this ever returns {@code false} in production, it means a Lucene reader has been paired
+     * with the wrong catalog snapshot.
+     *
+     * @param catalogSnapshot catalog snapshot whose referenced generations are the expected set
+     * @param readers         DirectoryReader whose leaves' generations are the actual set
+     * @return {@code true} iff both lists contain the same generations in the same (sorted) order
+     */
+    private boolean readersAreSame(CatalogSnapshot catalogSnapshot, DirectoryReader readers) {
+        Collection<Long> generationsReferenced = catalogSnapshot.getSegments().stream().map(Segment::generation).sorted().toList();
+        return generationsReferenced.equals(collectReferencedGenerations(readers));
+    }
+
+    /**
+     * Extracts the writer generation from each leaf of the given {@link DirectoryReader} and
+     * returns them as a sorted list. Each leaf's {@link SegmentReader} carries a
+     * {@link SegmentCommitInfo} whose {@code SegmentInfo} is stamped with the
+     * {@link org.opensearch.be.lucene.index.LuceneWriter#WRITER_GENERATION_ATTRIBUTE} when the
+     * segment is written; parsing that attribute yields the generation that produced the leaf.
+     *
+     * @param reader the DirectoryReader to inspect
+     * @return generations of all leaves, sorted ascending
+     * @throws NumberFormatException if a leaf is missing the writer-generation attribute or
+     *                               its value is not parseable as a long (indicates a segment
+     *                               not produced by {@link org.opensearch.be.lucene.index.LuceneWriter})
+     * @throws ClassCastException    if any leaf reader is not a {@link SegmentReader}
+     */
+    private Collection<Long> collectReferencedGenerations(DirectoryReader reader) {
+        return reader.leaves().stream().map(lrc -> {
+            SegmentReader segmentReader = (SegmentReader) lrc.reader();
+            SegmentCommitInfo sci = segmentReader.getSegmentInfo();
+            return Long.parseLong(sci.info.getAttribute(WRITER_GENERATION_ATTRIBUTE));
+        }).sorted().toList();
+    }
+
     @Override
     public void onDeleted(CatalogSnapshot catalogSnapshot) throws IOException {
         DirectoryReader reader = readers.remove(catalogSnapshot);

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneCommitter.java

@@ -15,12 +15,18 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.MergeIndexWriter;
 import org.apache.lucene.index.NoMergePolicy;
 import org.apache.lucene.index.SegmentInfos;
+import org.apache.lucene.index.SerialMergeScheduler;
+import org.apache.lucene.search.Sort;
+import org.apache.lucene.search.SortField;
+import org.apache.lucene.search.SortedNumericSortField;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.CommitStats;
 import org.opensearch.index.engine.EngineConfig;
 import org.opensearch.index.engine.SafeCommitInfo;
+import org.opensearch.index.engine.dataformat.DocumentInput;
 import org.opensearch.index.engine.exec.CombinedCatalogSnapshotDeletionPolicy;
 import org.opensearch.index.engine.exec.commit.Committer;
 import org.opensearch.index.engine.exec.commit.CommitterConfig;
@@ -59,6 +65,19 @@
  * The store reference is incremented on construction and decremented on {@link #close()}.
  * Closing the committer also closes the underlying IndexWriter.
  *
+ * <h2>Refresh-lock coordination</h2>
+ *
+ * <p>The engine passes a {@code preMergeCommitHook} via {@link CommitterConfig}. We wire it
+ * into Lucene as a {@code MergedSegmentWarmer} on the {@link IndexWriterConfig}. The warmer
+ * runs between {@code mergeMiddle} and {@code commitMerge} while the {@link IndexWriter}
+ * monitor is <em>not</em> held, so invoking the hook there establishes the ordering
+ * {@code refreshLock → IW monitor} on the merge thread — matching the refresh path and
+ * avoiding the lock inversion that would occur if coordination happened inside
+ * {@code commitMerge}. Ownership of whatever the hook acquires (currently the engine's
+ * refresh lock) is transferred to the engine's {@code applyMergeChanges} callback, which
+ * releases it after the catalog is updated. This committer never touches the refresh lock
+ * directly.
+ *
  * @opensearch.experimental
  */
 @ExperimentalApi
@@ -67,7 +86,7 @@ public class LuceneCommitter extends SafeBootstrapCommitter {
     private static final Logger logger = LogManager.getLogger(LuceneCommitter.class);
 
     private final Store store;
-    private final IndexWriter indexWriter;
+    private final MergeIndexWriter indexWriter;
     private final LuceneCommitDeletionPolicy deletionPolicy;
     private final AtomicBoolean isClosed = new AtomicBoolean();
 
@@ -84,8 +103,8 @@ public LuceneCommitter(CommitterConfig committerConfig) throws IOException {
         this.store.incRef();
         try {
             this.deletionPolicy = new LuceneCommitDeletionPolicy();
-            IndexWriterConfig iwc = createIndexWriterConfig(committerConfig.engineConfig());
-            this.indexWriter = new IndexWriter(store.directory(), iwc);
+            IndexWriterConfig iwc = createIndexWriterConfig(committerConfig);
+            this.indexWriter = new MergeIndexWriter(store.directory(), iwc);
         } catch (Exception e) {
             store.decRef();
             throw e;
@@ -197,18 +216,20 @@ public boolean isCommitManagedFile(String fileName) {
      *
      * @return the index writer, or null if closed
      */
-    IndexWriter getIndexWriter() {
+    MergeIndexWriter getIndexWriter() {
         ensureOpen();
         return indexWriter;
     }
 
     // --- Internal ---
 
-    private IndexWriterConfig createIndexWriterConfig(EngineConfig engineConfig) {
+    private IndexWriterConfig createIndexWriterConfig(CommitterConfig committerConfig) {
+        EngineConfig engineConfig = committerConfig.engineConfig();
         if (engineConfig == null) {
             IndexWriterConfig iwc = new IndexWriterConfig();
             iwc.setIndexDeletionPolicy(deletionPolicy);
             iwc.setMergePolicy(NoMergePolicy.INSTANCE);
+            iwc.setMergeScheduler(new SerialMergeScheduler());
             return iwc;
         }
         // TODO:: Merge Config needs to be wired in
@@ -219,13 +240,34 @@ private IndexWriterConfig createIndexWriterConfig(EngineConfig engineConfig) {
         }
         iwc.setRAMBufferSizeMB(engineConfig.getIndexingBufferSize().getMbFrac());
         iwc.setUseCompoundFile(engineConfig.useCompoundFile());
-        if (engineConfig.getIndexSort() != null) {
+        // Refresh-lock hand-off: the MergedSegmentWarmer fires on the merge thread between
+        // mergeMiddle and commitMerge, while the IndexWriter monitor is NOT held. Invoking
+        // the engine-provided preMergeCommitHook here gives the merge path the ordering
+        // refreshLock → IW monitor, which matches the refresh path (DataFormatAwareEngine#refresh
+        // takes refreshLock before calling IndexWriter#addIndexes). Ownership of whatever the
+        // hook acquires is transferred to applyMergeChanges, which releases it after the
+        // catalog is updated. See the class Javadoc.
+        iwc.setMergedSegmentWarmer(_ -> committerConfig.preMergeCommitHook().run());
+
+        // Determine if Lucene is a secondary format in a composite setup.
+        // When secondary, use a SortedNumericSortField on the row ID so MultiSorter can reorder
+        // documents by remapped row ID during merge. When primary (or standalone), use the
+        // engine config's IndexSort (which may be user-configured).
+        // TODO Check what is the right way to get this information as the below one is leaky
+        // https://github.com/opensearch-project/OpenSearch/issues/21506
+        List<String> secondaryFormats = engineConfig.getIndexSettings().getSettings().getAsList("index.composite.secondary_data_formats");
+        boolean isSecondary = secondaryFormats.contains("lucene");
+
+        if (isSecondary) {
+            iwc.setIndexSort(new Sort(new SortedNumericSortField(DocumentInput.ROW_ID_FIELD, SortField.Type.LONG)));
+        } else if (engineConfig.getIndexSort() != null) {
             iwc.setIndexSort(engineConfig.getIndexSort());
         }
         iwc.setCommitOnClose(false);
         iwc.setOpenMode(IndexWriterConfig.OpenMode.CREATE_OR_APPEND);
         iwc.setIndexDeletionPolicy(deletionPolicy);
         iwc.setMergePolicy(NoMergePolicy.INSTANCE);
+        iwc.setMergeScheduler(new SerialMergeScheduler());
         return iwc;
     }
 

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneDocumentInput.java

@@ -9,7 +9,7 @@
 package org.opensearch.be.lucene.index;
 
 import org.apache.lucene.document.Document;
-import org.apache.lucene.document.NumericDocValuesField;
+import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.opensearch.be.lucene.LuceneFieldFactory;
 import org.opensearch.be.lucene.LuceneFieldFactoryRegistry;
 import org.opensearch.common.annotation.ExperimentalApi;
@@ -26,8 +26,9 @@
  * Only field types registered in the registry are accepted. Attempting to add a field
  * of an unregistered type throws {@link IllegalArgumentException}.
  *
- * The row ID field is stored as a {@link NumericDocValuesField} for efficient doc-value
- * access, maintaining 1:1 correspondence between Lucene doc IDs and Parquet row offsets.
+ * The row ID field is stored as a {@link SortedNumericDocValuesField} for efficient doc-value
+ * access and compatibility with the {@code SortedNumericSortField}-based IndexSort,
+ * maintaining 1:1 correspondence between Lucene doc IDs and Parquet row offsets.
  *
  * @opensearch.experimental
  */
@@ -95,15 +96,15 @@ public void addField(MappedFieldType fieldType, Object value) {
     }
 
     /**
-     * Stores the row ID as a {@link NumericDocValuesField} to maintain 1:1 correspondence
+     * Stores the row ID as a {@link SortedNumericDocValuesField} to maintain 1:1 correspondence
      * between Lucene doc IDs and Parquet row offsets.
      *
      * @param rowIdFieldName the name of the row ID field
      * @param rowId          the row ID value (0-based sequential within the writer)
      */
     @Override
     public void setRowId(String rowIdFieldName, long rowId) {
-        document.add(new NumericDocValuesField(rowIdFieldName, rowId));
+        document.add(new SortedNumericDocValuesField(rowIdFieldName, rowId));
     }
 
     /** No-op — this document input holds no closeable resources. */

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java

@@ -15,6 +15,7 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MergeIndexWriter;
 import org.apache.lucene.index.NoMergePolicy;
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentReader;
@@ -23,10 +24,10 @@
 import org.apache.lucene.store.MMapDirectory;
 import org.opensearch.be.lucene.LuceneDataFormat;
 import org.opensearch.be.lucene.LuceneFieldFactoryRegistry;
+import org.opensearch.be.lucene.merge.LuceneMerger;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.dataformat.IndexingExecutionEngine;
-import org.opensearch.index.engine.dataformat.MergeResult;
 import org.opensearch.index.engine.dataformat.Merger;
 import org.opensearch.index.engine.dataformat.RefreshInput;
 import org.opensearch.index.engine.dataformat.RefreshResult;
@@ -74,11 +75,12 @@ public class LuceneIndexingExecutionEngine implements IndexingExecutionEngine<Lu
     private static final Logger logger = LogManager.getLogger(LuceneIndexingExecutionEngine.class);
 
     private final LuceneDataFormat dataFormat;
-    private final IndexWriter sharedWriter;
+    private final MergeIndexWriter sharedWriter;
     private final Store store;
     private final Path baseDirectory;
     private final Analyzer analyzer;
     private final Codec codec;
+    private final LuceneMerger luceneMerger;
     private final LuceneFieldFactoryRegistry fieldFactoryRegistry;
 
     /**
@@ -105,6 +107,8 @@ public LuceneIndexingExecutionEngine(
         this.codec = sharedWriter.getConfig().getCodec();
         this.fieldFactoryRegistry = new LuceneFieldFactoryRegistry();
 
+        this.luceneMerger = new LuceneMerger(sharedWriter, dataFormat, store.shardPath().resolveIndex());
+
         // Create the lucene subdirectory if it doesn't exist
         try {
             Files.createDirectories(baseDirectory);
@@ -120,7 +124,7 @@ public LuceneIndexingExecutionEngine(
      *
      * @return the index writer
      */
-    public IndexWriter getWriter() {
+    public MergeIndexWriter getWriter() {
         return sharedWriter;
     }
 
@@ -153,7 +157,7 @@ public FormatStore getStore(DataFormat dataFormat) {
     public Writer<LuceneDocumentInput> createWriter(long writerGeneration) {
         assert sharedWriter.isOpen() : "Cannot create writer — shared IndexWriter is closed";
         try {
-            return new LuceneWriter(writerGeneration, dataFormat, baseDirectory, analyzer, codec);
+            return new LuceneWriter(writerGeneration, dataFormat, baseDirectory, analyzer, codec, sharedWriter.getConfig().getIndexSort());
         } catch (IOException e) {
             throw new RuntimeException("Failed to create LuceneWriter for generation " + writerGeneration, e);
         }
@@ -279,8 +283,7 @@ public RefreshResult refresh(RefreshInput refreshInput) throws IOException {
     /** Returns {@code null} — merge scheduling is not yet implemented for the Lucene format. */
     @Override
     public Merger getMerger() {
-        // TODO: Implement merge support as ParquetMerger
-        return mergeInput -> new MergeResult(Map.of());
+        return this.luceneMerger;
     }
 
     /**

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneWriter.java

@@ -17,6 +17,7 @@
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentInfos;
+import org.apache.lucene.search.Sort;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.MMapDirectory;
 import org.opensearch.be.lucene.LuceneDataFormat;
@@ -80,10 +81,18 @@ public class LuceneWriter implements Writer<LuceneDocumentInput> {
      * @param dataFormat       the Lucene data format descriptor
      * @param baseDirectory    the base directory under which to create the temp directory
      * @param analyzer         the analyzer to use for tokenized fields, or null for default
+     * @param codec            the codec to use, or null for default
+     * @param indexSort        the index sort to apply to segments, or null for no sort
      * @throws IOException if directory creation or IndexWriter opening fails
      */
-    public LuceneWriter(long writerGeneration, LuceneDataFormat dataFormat, Path baseDirectory, Analyzer analyzer, Codec codec)
-        throws IOException {
+    public LuceneWriter(
+        long writerGeneration,
+        LuceneDataFormat dataFormat,
+        Path baseDirectory,
+        Analyzer analyzer,
+        Codec codec,
+        Sort indexSort
+    ) throws IOException {
         this.writerGeneration = writerGeneration;
         this.dataFormat = dataFormat;
         this.docCount = 0;
@@ -97,6 +106,9 @@ public LuceneWriter(long writerGeneration, LuceneDataFormat dataFormat, Path bas
         IndexWriterConfig iwc = analyzer != null ? new IndexWriterConfig(analyzer) : new IndexWriterConfig();
         iwc.setOpenMode(IndexWriterConfig.OpenMode.CREATE);
         iwc.setRAMBufferSizeMB(RAM_BUFFER_SIZE_MB);
+        if (indexSort != null) {
+            iwc.setIndexSort(indexSort);
+        }
 
         iwc.setCodec(new LuceneWriterCodec(codec, writerGeneration));
         this.indexWriter = new IndexWriter(directory, iwc);