Adding the Implementation of the RowIdMapping, this will support Parquet as Primary engine and Lucene as Secondary Engine

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

Author: Sagar Darji

sandbox/plugins/analytics-backend-lucene/src/main/java/org/apache/lucene/index/MergeIndexWriter.java

@@ -9,7 +9,6 @@
 package org.apache.lucene.index;
 
 import org.apache.lucene.store.Directory;
-import org.opensearch.common.annotation.ExperimentalApi;
 
 import java.io.IOException;
 
@@ -31,7 +30,6 @@
  *
  * @opensearch.experimental
  */
-@ExperimentalApi
 public class MergeIndexWriter extends IndexWriter {
 
     public MergeIndexWriter(Directory d, IndexWriterConfig conf) throws IOException {
@@ -68,4 +66,11 @@ public void executeMerge(MergePolicy.OneMerge oneMerge, long mergeGeneration) th
         // merge() must be called without holding the lock — mergeInit asserts !Thread.holdsLock(this)
         merge(oneMerge);
     }
+
+    @Override
+    protected void mergeSuccess(MergePolicy.OneMerge merge) {
+        // TODO update this for lucene as a primary engine
+        // https://github.com/opensearch-project/OpenSearch/issues/21505
+        super.mergeSuccess(merge);
+    }
 }

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

@@ -232,6 +232,7 @@ private IndexWriterConfig createIndexWriterConfig(EngineConfig engineConfig) {
         // 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");
 

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

@@ -66,8 +66,11 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOExcepti
         if (DocumentInput.ROW_ID_FIELD.equals(field.name)) {
             if (rowIdMapping != null) {
                 return new MappedRowIdDocValues(delegate.getSortedNumeric(field), rowIdMapping, generation);
+            } else {
+                // https://github.com/opensearch-project/OpenSearch/issues/21508
+                // TODO check how this will work for primary engine when rowIdMapping will be null.
+                throw new UnsupportedOperationException("Lucene as Primary Format is not supported yet");
             }
-            return new SequentialRowIdDocValues(maxDoc, rowIdOffset);
         }
         return delegate.getSortedNumeric(field);
     }
@@ -153,57 +156,4 @@ public long cost() {
             return delegate.cost();
         }
     }
-
-    /**
-     * Assigns sequential {@code ___row_id} = {@code rowIdOffset + docID}.
-     * Used when no RowIdMapping is provided.
-     */
-    private static class SequentialRowIdDocValues extends SortedNumericDocValues {
-
-        private final int maxDoc;
-        private final int rowIdOffset;
-        private int docID = -1;
-
-        SequentialRowIdDocValues(int maxDoc, int rowIdOffset) {
-            this.maxDoc = maxDoc;
-            this.rowIdOffset = rowIdOffset;
-        }
-
-        @Override
-        public long nextValue() {
-            return rowIdOffset + docID;
-        }
-
-        @Override
-        public int docValueCount() {
-            return 1;
-        }
-
-        @Override
-        public boolean advanceExact(int target) {
-            docID = target;
-            return true;
-        }
-
-        @Override
-        public int docID() {
-            return docID;
-        }
-
-        @Override
-        public int nextDoc() {
-            return ++docID < maxDoc ? docID : NO_MORE_DOCS;
-        }
-
-        @Override
-        public int advance(int target) {
-            docID = target;
-            return docID < maxDoc ? docID : NO_MORE_DOCS;
-        }
-
-        @Override
-        public long cost() {
-            return maxDoc;
-        }
-    }
 }

sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneMergerTests.java

@@ -42,10 +42,8 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 import static org.opensearch.be.lucene.index.LuceneWriter.WRITER_GENERATION_ATTRIBUTE;
 
@@ -117,44 +115,6 @@ public void testMergeWithNoMatchingSegments() throws IOException {
         assertTrue(result.getMergedWriterFileSet().isEmpty());
     }
 
-    /**
-     * Merge without RowIdMapping (primary merge) preserves all documents and field data.
-     */
-    public void testMergeWithoutRowIdMappingPreservesData() throws IOException {
-        writeSegment(writer, 1L, 0, 3);
-        writeSegment(writer, 2L, 3, 2);
-        writer.commit();
-
-        SegmentInfos infos = getSegmentInfos(writer);
-        assertEquals(2, infos.size());
-        assertEquals(5, writer.getDocStats().numDocs);
-
-        LuceneMerger merger = new LuceneMerger(writer, new LuceneDataFormat(), dataPath);
-        List<Segment> segments = buildSegments(infos);
-
-        MergeInput input = MergeInput.builder().segments(segments).newWriterGeneration(10L).build();
-        MergeResult result = merger.merge(input);
-        assertNotNull(result);
-
-        writer.commit();
-        // addIndexes adds a new merged segment; original segments remain.
-        // Total docs = original 5 + merged 5 = 10
-        try (DirectoryReader reader = DirectoryReader.open(writer)) {
-            assertTrue("Should have at least 5 docs after merge", reader.numDocs() >= 5);
-            Set<String> foundIds = new HashSet<>();
-            for (LeafReaderContext ctx : reader.leaves()) {
-                for (int i = 0; i < ctx.reader().maxDoc(); i++) {
-                    Document doc = ctx.reader().storedFields().document(i);
-                    foundIds.add(doc.get("id"));
-                    assertNotNull("data field should be preserved", doc.get("data"));
-                }
-            }
-            for (int i = 0; i < 5; i++) {
-                assertTrue("Missing doc id: doc_" + i, foundIds.contains("doc_" + i));
-            }
-        }
-    }
-
     /**
      * Merge with RowIdMapping remaps ___row_id doc values AND reorders documents.
      * Verifies that the merged segment has documents sorted by remapped row IDs
@@ -242,6 +202,11 @@ public void testMergeWithRowIdMappingRemapsRowIds() throws IOException {
 
     /**
      * Merge preserves keyword, numeric, and stored field data integrity.
+     *
+     * <p>Uses an identity {@link RowIdMapping} so the merge exercises the real
+     * secondary-format path; the assertions focus on field-data survival rather
+     * than on row-id remapping (which is covered by
+     * {@link #testMergeWithRowIdMappingRemapsRowIds()}).
      */
     public void testMergePreservesFieldDataIntegrity() throws IOException {
         writeSegmentWithRichFields(writer, 1L, 0, 3);
@@ -252,7 +217,11 @@ public void testMergePreservesFieldDataIntegrity() throws IOException {
         SegmentInfos infos = getSegmentInfos(writer);
         List<Segment> segments = buildSegments(infos);
 
-        MergeInput input = MergeInput.builder().segments(segments).newWriterGeneration(10L).build();
+        // Identity mapping — writeSegmentWithRichFields already writes globally-unique row IDs
+        // (0,1,2 in gen=1 and 3,4 in gen=2), so returning the original row ID is well-formed.
+        RowIdMapping identityMapping = (oldId, oldGeneration) -> oldId;
+
+        MergeInput input = MergeInput.builder().segments(segments).rowIdMapping(identityMapping).newWriterGeneration(10L).build();
         merger.merge(input);
         writer.commit();