Add interface for the Multi format merge flow

server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java

@@ -0,0 +1,110 @@
+/*
+ * 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.engine.dataformat.merge;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.engine.dataformat.MergeResult;
+
+import java.util.ArrayDeque;
+import java.util.Collection;
+import java.util.Deque;
+
+/**
+ * Abstract handler responsible for managing segment merge operations.
+ * <p>
+ * Subclasses define the merge policy by implementing {@link #findMerges()} and
+ * {@link #findForceMerges(int)}, while this base class manages the pending merge
+ * queue and lifecycle callbacks.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public abstract class MergeHandler {
+
+    private final Deque<OneMerge> mergingSegments = new ArrayDeque<>();
+
+    /**
+     * Finds merges that should be executed based on the current segment state.
+     *
+     * @return a collection of merges to execute, or an empty collection if none are needed
+     */
+    public abstract Collection<OneMerge> findMerges();
+
+    /**
+     * Finds merges required to reduce the number of segments to at most {@code maxSegmentCount}.
+     *
+     * @param maxSegmentCount the maximum number of segments allowed after merging
+     * @return a collection of merges to execute
+     */
+    public abstract Collection<OneMerge> findForceMerges(int maxSegmentCount);
+
+    /**
+     * Updates the set of pending merges. Called to refresh the merge queue
+     * when the segment state changes.
+     */
+    public synchronized void updatePendingMerges() {
+
+    }
+
+    /**
+     * Registers a merge to be executed.
+     *
+     * @param merge the merge to register
+     */
+    public synchronized void registerMerge(OneMerge merge) {
+
+    }
+
+    /**
+     * Returns whether there are any pending merges in the queue.
+     *
+     * @return {@code true} if there are pending merges
+     */
+    public boolean hasPendingMerges() {
+        return !mergingSegments.isEmpty();
+    }
+
+    /**
+     * Retrieves and removes the next pending merge from the queue.
+     *
+     * @return the next merge to execute, or {@code null} if the queue is empty
+     */
+    public synchronized OneMerge getNextMerge() {
+        if(mergingSegments.isEmpty()) {
+            return null;
+        }
+        return mergingSegments.removeFirst();
+    }
+
+    /**
+     * Callback invoked when a merge completes successfully.
+     *
+     * @param oneMerge the merge that finished
+     */
+    public synchronized void onMergeFinished(OneMerge oneMerge) {
+    }
+
+    /**
+     * Callback invoked when a merge fails.
+     *
+     * @param oneMerge the merge that failed
+     */
+    public synchronized void onMergeFailure(OneMerge oneMerge) {
+    }
+
+    /**
+     * Executes the given merge operation.
+     *
+     * @param oneMerge the merge to execute
+     * @return the result of the merge
+     */
+    public MergeResult doMerge(OneMerge oneMerge) {
+        return null;
+    }
+}

server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeScheduler.java

@@ -0,0 +1,129 @@
+/*
+ * 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.engine.dataformat.merge;
+
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.index.IndexSettings;
+import org.opensearch.index.MergeSchedulerConfig;
+import org.opensearch.index.merge.MergeStats;
+
+/**
+ * Schedules and coordinates segment merge operations for a shard.
+ * <p>
+ * This scheduler delegates merge selection to a {@link MergeHandler} and controls
+ * concurrency via configurable thread and merge count limits sourced from
+ * {@link MergeSchedulerConfig}.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class MergeScheduler {
+
+    private static Logger logger;
+    private volatile int maxConcurrentMerges;
+    private volatile int maxMergeCount;
+    private final MergeSchedulerConfig mergeSchedulerConfig;
+
+    /** true if we should rate-limit writes for each merge */
+    private boolean doAutoIOThrottle = false;
+
+    /** Initial value for IO write rate limit when doAutoIOThrottle is true */
+    private static final double START_MB_PER_SEC = 20.0;
+
+    /** Current IO writes throttle rate */
+    protected double targetMBPerSec = START_MB_PER_SEC;
+
+    /**
+     * Creates a new merge scheduler.
+     *
+     * @param mergeHandler   the handler that selects and executes merges
+     * @param shardId        the shard this scheduler is associated with
+     * @param indexSettings  the index settings providing merge scheduler configuration
+     */
+    public MergeScheduler(
+        MergeHandler mergeHandler,
+        ShardId shardId,
+        IndexSettings indexSettings
+    ) {
+        this.mergeSchedulerConfig = indexSettings.getMergeSchedulerConfig();
+        refreshConfig();
+    }
+
+    /**
+     * Refreshes the max concurrent merge thread count and max merge count from
+     * the current {@link MergeSchedulerConfig}. No-op if the values have not changed.
+     */
+    public synchronized void refreshConfig() {
+        int newMaxThreadCount = mergeSchedulerConfig.getMaxThreadCount();
+        int newMaxMergeCount = mergeSchedulerConfig.getMaxMergeCount();
+
+        if (newMaxThreadCount == this.maxConcurrentMerges && newMaxMergeCount == this.maxMergeCount) {
+            return;
+        }
+
+        logger.info(() -> new ParameterizedMessage("Updating from merge scheduler config: maxThreadCount {} -> {}, " +
+            "maxMergeCount {} -> {}", this.maxConcurrentMerges, newMaxThreadCount, this.maxMergeCount, newMaxMergeCount));
+
+        this.maxConcurrentMerges = newMaxThreadCount;
+        this.maxMergeCount = newMaxMergeCount;
+    }
+
+    /**
+     * Triggers pending merge operations. Merges are selected by the
+     * underlying {@link MergeHandler} and executed up to the configured
+     * concurrency limits.
+     */
+    public void triggerMerges() {
+
+    }
+
+    /**
+     * Forces a merge down to at most {@code maxNumSegment} segments.
+     *
+     * @param maxNumSegment the maximum number of segments after the force merge
+     */
+    public void forceMerge(int maxNumSegment) {
+
+    }
+
+    /**
+     * Turn on dynamic IO throttling, to adaptively rate limit writes bytes/sec to the minimal rate
+     * necessary so merges do not fall behind. By default, this is disabled and writes are not
+     * rate-limited.
+     */
+    public synchronized void enableAutoIOThrottle() {
+        doAutoIOThrottle = true;
+        targetMBPerSec = START_MB_PER_SEC;
+    }
+
+    /**
+     * Returns the currently set per-merge IO writes rate limit, if {@link #enableAutoIOThrottle} was
+     * called, else {@code Double.POSITIVE_INFINITY}.
+     */
+    public synchronized double getIORateLimitMBPerSec() {
+        if (doAutoIOThrottle) {
+            return targetMBPerSec;
+        }
+
+        return Double.POSITIVE_INFINITY;
+    }
+
+    /**
+     * Returns the current merge statistics for this scheduler.
+     *
+     * @return the merge stats
+     */
+    public MergeStats stats() {
+        return new MergeStats();
+    }
+}
+

server/src/main/java/org/opensearch/index/engine/dataformat/merge/OneMerge.java

@@ -0,0 +1,88 @@
+/*
+ * 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.engine.dataformat.merge;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.engine.exec.WriterFileSet;
+import org.opensearch.index.engine.exec.Segment;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Represents a single merge operation over a set of {@link Segment}s.
+ * <p>
+ * Precomputes and caches the total size in bytes and total document count
+ * across all segments selected for merging.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class OneMerge {
+    private final List<Segment> segmentsToMerge;
+    private final long totalSize;
+    private final long totalNumDocs;
+
+    /**
+     * Creates a new merge operation for the given segments.
+     *
+     * @param segmentsToMerge the segments to be merged
+     */
+    public OneMerge(List<Segment> segmentsToMerge) {
+        this.segmentsToMerge = Collections.unmodifiableList(segmentsToMerge);
+        this.totalSize = calculateTotalSizeInBytes();
+        this.totalNumDocs = calculateTotalNumDocs();
+    }
+
+    /**
+     * Returns the unmodifiable list of segments participating in this merge.
+     *
+     * @return the segments to merge
+     */
+    public List<Segment> getSegmentsToMerge() {
+        return segmentsToMerge;
+    }
+
+    /**
+     * Returns the total size in bytes across all segments in this merge.
+     *
+     * @return total size in bytes
+     */
+    public long getTotalSizeInBytes() {
+        return totalSize;
+    }
+
+    /**
+     * Returns the total number of documents across all segments in this merge.
+     *
+     * @return total document count
+     */
+    public long getTotalNumDocs() {
+        return totalNumDocs;
+    }
+
+    private long calculateTotalSizeInBytes() {
+        return segmentsToMerge.stream()
+            .flatMap(segment -> segment.dfGroupedSearchableFiles().values().stream())
+            .mapToLong(WriterFileSet::getTotalSize)
+            .sum();
+    }
+
+    private long calculateTotalNumDocs() {
+        return segmentsToMerge.stream()
+            .flatMap(segment -> segment.dfGroupedSearchableFiles().values().stream())
+            .mapToLong(WriterFileSet::getNumRows)
+            .sum();
+    }
+
+    @Override
+    public String toString() {
+        return "Merge [SegmentsToMerge=" + segmentsToMerge + "] ";
+    }
+}

server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java

@@ -36,6 +36,10 @@ public long getTotalSize() {
         }).sum();
     }
 
+    public long getNumRows() {
+        return numRows;
+    }
+
     @Override
     public String toString() {
         return "WriterFileSet{" + "directory=" + directory + ", writerGeneration=" + writerGeneration + ", files=" + files + '}';

server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java

@@ -34,6 +34,8 @@
 
 import static org.mockito.Mockito.mock;
 
+import org.opensearch.index.engine.dataformat.DataFormatTestUtils.MockDataFormat;
+
 /**
  * Tests demonstrating the data format plugin API lifecycle:
  * <ol>
@@ -147,7 +149,7 @@ public void testFullDataFormatLifecycle() throws IOException {
      * Tests DataFormat equality semantics and field capabilities.
      */
     public void testDataFormatCapabilities() {
-        MockDataFormat format = new MockDataFormat();
+        DataFormat format = new MockDataFormat();
         Set<FieldTypeCapabilities> fields = format.supportedFields();
         assertEquals(1, fields.size());
 
@@ -222,28 +224,6 @@ public void testRefreshInput() {
         assertEquals(1, input.existingSegments().size());
     }
 
-    static class MockDataFormat implements DataFormat {
-        @Override
-        public String name() {
-            return "mock-columnar";
-        }
-
-        @Override
-        public long priority() {
-            return 100L;
-        }
-
-        @Override
-        public Set<FieldTypeCapabilities> supportedFields() {
-            return Set.of(
-                new FieldTypeCapabilities(
-                    "integer",
-                    Set.of(FieldTypeCapabilities.Capability.COLUMNAR_STORAGE, FieldTypeCapabilities.Capability.STORED_FIELDS)
-                )
-            );
-        }
-    }
-
     static class MockDocumentInput implements DocumentInput<Map<String, Object>> {
         private final Map<String, Object> fields = new HashMap<>();