Addressed the comments

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

Author: Sagar Darji

sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java

@@ -19,37 +19,19 @@
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.engine.CommitStats;
-import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.dataformat.DataFormatDescriptor;
-import org.opensearch.index.engine.dataformat.DataFormatPlugin;
 import org.opensearch.index.engine.dataformat.DataFormatRegistry;
-import org.opensearch.index.engine.dataformat.DocumentInput;
-import org.opensearch.index.engine.dataformat.FieldTypeCapabilities;
-import org.opensearch.index.engine.dataformat.FileInfos;
-import org.opensearch.index.engine.dataformat.IndexingEngineConfig;
-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;
-import org.opensearch.index.engine.dataformat.WriteResult;
-import org.opensearch.index.engine.dataformat.Writer;
-import org.opensearch.index.engine.exec.Segment;
-import org.opensearch.index.engine.exec.WriterFileSet;
-import org.opensearch.index.engine.exec.commit.IndexStoreProvider;
+import org.opensearch.index.engine.dataformat.stub.MockDataFormat;
+import org.opensearch.index.engine.dataformat.stub.MockDataFormatPlugin;
 import org.opensearch.index.engine.exec.coord.DataformatAwareCatalogSnapshot;
-import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.merge.MergeStats;
-import org.opensearch.index.store.FormatChecksumStrategy;
 import org.opensearch.index.store.PrecomputedChecksumStrategy;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
@@ -67,163 +49,13 @@ public class CompositeMergeIT extends OpenSearchIntegTestCase {
 
     private static final String INDEX_NAME = "test-composite-merge";
 
-    // ── Stub DataFormat for "parquet" ──
+    // ── Mock DataFormatPlugin using test framework stubs ──
 
-    private static final DataFormat STUB_PARQUET_FORMAT = new DataFormat() {
-        @Override
-        public String name() {
-            return "parquet";
-        }
-
-        @Override
-        public long priority() {
-            return 0;
-        }
-
-        @Override
-        public Set<FieldTypeCapabilities> supportedFields() {
-            return Set.of();
-        }
-    };
-
-    // ── Stub DocumentInput ──
-
-    private static class StubDocumentInput implements DocumentInput<Object> {
-        @Override
-        public void addField(MappedFieldType fieldType, Object value) {}
-
-        @Override
-        public void setRowId(String rowIdFieldName, long rowId) {}
-
-        @Override
-        public Object getFinalInput() {
-            return null;
-        }
-
-        @Override
-        public void close() {}
-    }
-
-    // ── Stub Writer ──
-
-    private static class StubWriter implements Writer<StubDocumentInput> {
-        private final long generation;
-        private int docCount = 0;
-
-        StubWriter(long generation) {
-            this.generation = generation;
-        }
-
-        @Override
-        public WriteResult addDoc(StubDocumentInput documentInput) {
-            docCount++;
-            return new WriteResult.Success(1L, 1L, docCount);
-        }
-
-        @Override
-        public FileInfos flush() {
-            if (docCount == 0) {
-                return FileInfos.empty();
-            }
-            WriterFileSet wfs = new WriterFileSet("/tmp/stub", generation, Set.of("stub_" + generation + ".parquet"), docCount);
-            return new FileInfos(Map.of(STUB_PARQUET_FORMAT, wfs));
-        }
-
-        @Override
-        public void sync() {}
-
-        @Override
-        public long generation() {
-            return generation;
-        }
-
-        @Override
-        public void lock() {}
-
-        @Override
-        public boolean tryLock() {
-            return true;
-        }
-
-        @Override
-        public void unlock() {}
-
-        @Override
-        public void close() throws IOException {}
-    }
-
-    // ── Stub IndexingExecutionEngine ──
-
-    @SuppressWarnings({ "unchecked", "rawtypes" })
-    private static class StubParquetEngine implements IndexingExecutionEngine<DataFormat, StubDocumentInput> {
+    public static class MockParquetDataFormatPlugin extends MockDataFormatPlugin {
+        private static final MockDataFormat PARQUET_FORMAT = new MockDataFormat("parquet", 0L, Set.of());
 
-        @Override
-        public Writer<StubDocumentInput> createWriter(long writerGeneration) {
-            return new StubWriter(writerGeneration);
-        }
-
-        @Override
-        public Merger getMerger() {
-            return mergeInput -> {
-                long totalRows = mergeInput.writerFiles().stream().mapToLong(WriterFileSet::numRows).sum();
-                WriterFileSet merged = new WriterFileSet("/tmp/stub", mergeInput.newWriterGeneration(),
-                    Set.of("merged_" + mergeInput.newWriterGeneration() + ".parquet"), totalRows);
-                return new MergeResult(Map.of(STUB_PARQUET_FORMAT, merged));
-            };
-        }
-
-        @Override
-        public RefreshResult refresh(RefreshInput refreshInput) {
-            if (refreshInput == null) return new RefreshResult(List.of());
-            List<Segment> segments = new ArrayList<>();
-            segments.addAll(refreshInput.existingSegments());
-            segments.addAll(refreshInput.writerFiles());
-            return new RefreshResult(List.copyOf(segments));
-        }
-
-        @Override
-        public long getNextWriterGeneration() {
-            return 0;
-        }
-
-        @Override
-        public DataFormat getDataFormat() {
-            return STUB_PARQUET_FORMAT;
-        }
-
-        @Override
-        public long getNativeBytesUsed() {
-            return 0;
-        }
-
-        @Override
-        public void deleteFiles(Map<String, Collection<String>> filesToDelete) {}
-
-        @Override
-        public StubDocumentInput newDocumentInput() {
-            return new StubDocumentInput();
-        }
-
-        @Override
-        public IndexStoreProvider getProvider() {
-            return null;
-        }
-
-        @Override
-        public void close() {}
-    }
-
-    // ── Stub DataFormatPlugin ──
-
-    public static class MockParquetDataFormatPlugin extends Plugin implements DataFormatPlugin {
-        @Override
-        public DataFormat getDataFormat() {
-            return STUB_PARQUET_FORMAT;
-        }
-
-        @Override
-        public IndexingExecutionEngine<?, ?> indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy) {
-            return new StubParquetEngine();
+        public MockParquetDataFormatPlugin() {
+            super(PARQUET_FORMAT);
         }
 
         @Override
@@ -279,8 +111,7 @@ public void testBackgroundMergeSingleEngine() throws Exception {
             flush(INDEX_NAME);
             DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot();
             assertTrue(
-                "Expected merges to reduce segment count below " + totalSegmentsCreated
-                    + ", but got: " + snapshot.getSegments().size(),
+                "Expected merges to reduce segment count below " + totalSegmentsCreated + ", but got: " + snapshot.getSegments().size(),
                 snapshot.getSegments().size() < totalSegmentsCreated
             );
         });
@@ -305,7 +136,6 @@ private Settings singleEngineSettings() {
             .build();
     }
 
-
     private int indexDocsAcrossMultipleRefreshes(int refreshCycles, int docsPerCycle) {
         int totalDocs = 0;
         for (int cycle = 0; cycle < refreshCycles; cycle++) {

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormat.java

@@ -40,6 +40,14 @@ public CompositeDataFormat(DataFormat primaryDataFormat, List<DataFormat> dataFo
         this.dataFormats = List.copyOf(Objects.requireNonNull(dataFormats, "dataFormats must not be null"));
     }
 
+    /**
+     * Constructs an empty CompositeDataFormat with no constituent formats.
+     */
+    public CompositeDataFormat() {
+        this.primaryDataFormat = null;
+        this.dataFormats = List.of();
+    }
+
     /**
      * Returns the list of constituent data formats.
      *

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormatPlugin.java

@@ -18,7 +18,6 @@
 import org.opensearch.index.engine.dataformat.DataFormatDescriptor;
 import org.opensearch.index.engine.dataformat.DataFormatPlugin;
 import org.opensearch.index.engine.dataformat.DataFormatRegistry;
-import org.opensearch.index.engine.dataformat.FieldTypeCapabilities;
 import org.opensearch.index.engine.dataformat.IndexingEngineConfig;
 import org.opensearch.index.engine.dataformat.IndexingExecutionEngine;
 import org.opensearch.index.store.FormatChecksumStrategy;
@@ -29,7 +28,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * Sandbox plugin that provides a {@link CompositeIndexingExecutionEngine} for
@@ -90,22 +88,7 @@ public List<Setting<?>> getSettings() {
     @Override
     public DataFormat getDataFormat() {
         // TODO: Dataformat for Composite is per index, while this one talks about cluster level. Switching it off for now
-        return new CompositeDataFormat(new DataFormat() {
-            @Override
-            public String name() {
-                return "composite";
-            }
-
-            @Override
-            public long priority() {
-                return Long.MIN_VALUE;
-            }
-
-            @Override
-            public Set<FieldTypeCapabilities> supportedFields() {
-                return Set.of();
-            }
-        }, List.of());
+        return new CompositeDataFormat();
     }
 
     @Override

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java

@@ -13,6 +13,7 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.io.IOUtils;
+import org.opensearch.composite.merge.CompositeMerger;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.dataformat.DataFormatPlugin;
@@ -27,7 +28,6 @@
 import org.opensearch.index.engine.dataformat.Writer;
 import org.opensearch.index.engine.exec.EngineReaderManager;
 import org.opensearch.index.engine.exec.Segment;
-import org.opensearch.composite.merge.CompositeMerger;
 import org.opensearch.index.engine.exec.WriterFileSet;
 import org.opensearch.index.engine.exec.commit.Committer;
 import org.opensearch.index.engine.exec.commit.IndexStoreProvider;

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/merge/CompositeMergeExecutor.java

@@ -14,7 +14,7 @@
 import org.opensearch.index.engine.dataformat.MergeResult;
 import org.opensearch.index.engine.dataformat.Merger;
 import org.opensearch.index.engine.dataformat.RowIdMapping;
-import org.opensearch.index.engine.dataformat.merge.OneMerge;
+import org.opensearch.index.engine.exec.Segment;
 import org.opensearch.index.engine.exec.WriterFileSet;
 
 import java.io.IOException;
@@ -49,14 +49,12 @@ public CompositeMergeExecutor(Map<DataFormat, Merger> mergers) {
     public MergeResult execute(MergePlan plan) {
         List<FormatMergeResult> completed = new ArrayList<>();
         try {
-            FormatMergeResult primaryResult = mergeFormat(
-                plan, plan.primaryFormat(), null
-            );
+            FormatMergeResult primaryResult = mergeFormat(plan, plan.primaryFormat(), null);
             completed.add(primaryResult);
 
             RowIdMapping mapping = plan.hasSecondaries()
-                ? primaryResult.rowIdMappingOpt().orElseThrow(() -> new IllegalStateException(
-                    "Primary merge did not produce row-ID mapping required by secondaries"))
+                ? primaryResult.rowIdMappingOpt()
+                    .orElseThrow(() -> new IllegalStateException("Primary merge did not produce row-ID mapping required by secondaries"))
                 : null;
 
             for (DataFormat secondary : plan.secondaryFormats()) {
@@ -71,24 +69,18 @@ public MergeResult execute(MergePlan plan) {
         }
     }
 
-    private FormatMergeResult mergeFormat(
-        MergePlan plan, DataFormat format, RowIdMapping mapping
-    ) throws IOException {
+    private FormatMergeResult mergeFormat(MergePlan plan, DataFormat format, RowIdMapping mapping) throws IOException {
         Merger merger = mergers.get(format);
         List<WriterFileSet> files = plan.filesFor(format);
-        MergeResult result = merger.merge(
-            new MergeInput(files, mapping, plan.mergedWriterGeneration())
-        );
-        return new FormatMergeResult(
-            format,
-            result.getMergedWriterFileSetForDataformat(format),
-            result.rowIdMapping().orElse(null)
-        );
+        List<Segment> segments = new ArrayList<>();
+        for (WriterFileSet wfs : files) {
+            segments.add(Segment.builder(wfs.writerGeneration()).addSearchableFiles(format, wfs).build());
+        }
+        MergeResult result = merger.merge(new MergeInput(segments, mapping, plan.mergedWriterGeneration()));
+        return new FormatMergeResult(format, result.getMergedWriterFileSetForDataformat(format), result.rowIdMapping().orElse(null));
     }
 
-    private static MergeResult toMergeResult(
-        List<FormatMergeResult> results, RowIdMapping mapping
-    ) {
+    private static MergeResult toMergeResult(List<FormatMergeResult> results, RowIdMapping mapping) {
         Map<DataFormat, WriterFileSet> merged = new HashMap<>();
         for (FormatMergeResult r : results) {
             merged.put(r.format(), r.mergedFiles());