Refactoring merge code to act at composite level instead of Parquet (#20119)

Signed-off-by: Dhwanil Patel <[email protected]>

Author: dhwanilpatel

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/merge/ParquetMergeExecutor.java

@@ -9,8 +9,10 @@
 package com.parquet.parquetdataformat.merge;
 
 import org.opensearch.index.engine.exec.FileMetadata;
+import org.opensearch.index.engine.exec.WriterFileSet;
 import org.opensearch.index.engine.exec.merge.MergeResult;
 import java.util.Collection;
+import java.util.List;
 
 /**
  * Executes Parquet merge operations using a chosen compaction strategy.
@@ -24,7 +26,7 @@ public ParquetMergeExecutor(CompactionStrategy compactionStrategy) {
     }
 
     @Override
-    public MergeResult merge(Collection<FileMetadata> fileMetadataList, long writerGeneration) {
+    public MergeResult merge(List<WriterFileSet> fileMetadataList, long writerGeneration) {
         MergeResult result = strategy.mergeParquetFiles(fileMetadataList, writerGeneration);
         strategy.postMerge();
         return result;

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/merge/ParquetMergeStrategy.java

@@ -10,9 +10,12 @@
 
 
 import org.opensearch.index.engine.exec.FileMetadata;
+import org.opensearch.index.engine.exec.WriterFileSet;
 import org.opensearch.index.engine.exec.merge.MergeResult;
 
 import java.util.Collection;
+import java.util.List;
+
 /**
  * Interface defining a Parquet merge strategy.
  */
@@ -21,7 +24,7 @@ public interface ParquetMergeStrategy {
     /**
      * Performs the actual Parquet merge.
      */
-    MergeResult mergeParquetFiles(Collection<FileMetadata> files, long writerGeneration);
+    MergeResult mergeParquetFiles(List<WriterFileSet> files, long writerGeneration);
 
     /**
      * Optional post-merge hook.

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/merge/ParquetMerger.java

@@ -10,14 +10,16 @@
 
 import org.opensearch.index.engine.exec.FileMetadata;
 import org.opensearch.index.engine.exec.Merger;
+import org.opensearch.index.engine.exec.WriterFileSet;
 import org.opensearch.index.engine.exec.merge.MergeResult;
 import org.opensearch.index.engine.exec.merge.RowIdMapping;
 
 import java.util.Collection;
+import java.util.List;
 
 public abstract class ParquetMerger implements Merger {
     @Override
-    public MergeResult merge(Collection<FileMetadata> fileMetadataList, RowIdMapping rowIdMapping, long writerGeneration) {
+    public MergeResult merge(List<WriterFileSet> fileMetadataList, RowIdMapping rowIdMapping, long writerGeneration) {
         throw new UnsupportedOperationException("Not supported parquet as secondary data format yet.");
     }
 }

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/merge/RecordBatchMergeStrategy.java

@@ -18,6 +18,7 @@
 import org.opensearch.index.engine.exec.merge.RowIdMapping;
 
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -34,17 +35,17 @@
 public class RecordBatchMergeStrategy implements ParquetMergeStrategy {
 
     @Override
-    public MergeResult mergeParquetFiles(Collection<FileMetadata> files, long writerGeneration) {
+    public MergeResult mergeParquetFiles(List<WriterFileSet> files, long writerGeneration) {
 
         if (files.isEmpty()) {
             throw new IllegalArgumentException("No files to merge");
         }
 
-        List<Path> filePaths = files.stream()
-            .map(fm -> Path.of(fm.directory(), fm.file()))
-            .collect(Collectors.toList());
+        List<Path> filePaths = new ArrayList<>();
+        files.forEach(writerFileSet ->  writerFileSet.getFiles().forEach(
+            file -> filePaths.add(Path.of(writerFileSet.getDirectory(), file))));
 
-        String outputDirectory = files.iterator().next().directory();
+        String outputDirectory = files.iterator().next().getDirectory();
         String mergedFilePath = getMergedFilePath(writerGeneration, outputDirectory);
         String mergedFileName = getMergedFileName(writerGeneration);
 

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

@@ -39,10 +39,8 @@
 import org.apache.lucene.index.TieredMergePolicy;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
-import org.opensearch.index.engine.exec.merge.ParquetTieredMergePolicy;
 
 /**
  * A shard in opensearch is a Lucene index, and a Lucene index is broken

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

@@ -11,21 +11,21 @@
 import org.opensearch.index.engine.exec.merge.MergeResult;
 import org.opensearch.index.engine.exec.merge.RowIdMapping;
 
-import java.util.Collection;
+import java.util.List;
 
 public interface Merger {
     /**
      *
      * @param fileMetadataList List of FileMetadata to merge
      * @return MergeResult - having RowIdMapping and mergedFileMetadata
      */
-    MergeResult merge(Collection<FileMetadata> fileMetadataList, long writerGeneration);
+    MergeResult merge(List<WriterFileSet> fileMetadataList, long writerGeneration);
 
     /**
      *
      * @param fileMetadataList List of FileMetadata to merge
      * @param rowIdMapping Mapping of old segment + old rowId to new rowId
      * @return MergeResult - having mergedFileMetadata
      */
-    MergeResult merge(Collection<FileMetadata> fileMetadataList, RowIdMapping rowIdMapping, long writerGeneration);
+    MergeResult merge(List<WriterFileSet> fileMetadataList, RowIdMapping rowIdMapping, long writerGeneration);
 }

server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java

@@ -77,16 +77,14 @@ public synchronized void applyMergeResults(MergeResult mergeResult, OneMerge one
         List<CatalogSnapshot.Segment> segmentList = latestCatalogSnapshot.getSegments();
 
         CatalogSnapshot.Segment segmentToAdd = getSegment(mergeResult.getMergedWriterFileSet());
-
-        Set<FileMetadata> filesToRemove = new HashSet<>();
-        oneMerge.getFilesToMerge().forEach(file -> filesToRemove.add(file));
+        Set<CatalogSnapshot.Segment> segmentsToRemove = new HashSet<>(oneMerge.getSegmentsToMerge());
 
         boolean inserted = false;
         int newSegIdx = 0;
         for (int segIdx = 0, cnt = segmentList.size(); segIdx < cnt; segIdx++) {
             assert segIdx >= newSegIdx;
             CatalogSnapshot.Segment currSegment = segmentList.get(segIdx);
-            if(filesToRemove.containsAll(currSegment.getSearchableFiles(oneMerge.getDataFormat().name()))) {
+            if(segmentsToRemove.contains(currSegment)) {
                 if (!inserted) {
                     segmentList.set(segIdx, segmentToAdd);
                     inserted = true;

server/src/main/java/org/opensearch/index/engine/exec/coord/CompositeEngine.java

@@ -38,7 +38,7 @@
 import org.opensearch.index.engine.exec.merge.MergeResult;
 import org.opensearch.index.engine.exec.merge.MergeScheduler;
 import org.opensearch.index.engine.exec.merge.OneMerge;
-import org.opensearch.index.engine.exec.merge.ParquetMergeHandler;
+import org.opensearch.index.engine.exec.merge.CompositeMergeHandler;
 import org.opensearch.index.mapper.IdFieldMapper;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.SeqNoFieldMapper;
@@ -268,7 +268,7 @@ public void onFailure(String reason, Exception ex) {
             this.throttle = new IndexThrottle();
             this.historyUUID = loadHistoryUUID(userData);
             this.mergeHandler =
-                new ParquetMergeHandler(this, this.engine, this.engine.getDataFormat(), indexSettings);
+                new CompositeMergeHandler(this, this.engine, this.engine.getDataFormat(), indexSettings);
             this.mergeScheduler = new MergeScheduler(this.mergeHandler, this);
 
             // Refresh here so that catalog snapshot gets initialized

server/src/main/java/org/opensearch/index/engine/exec/merge/CompositeMergeHandler.java

@@ -0,0 +1,95 @@
+/*
+ * 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.exec.merge;
+
+import org.opensearch.index.IndexSettings;
+import org.opensearch.index.engine.exec.composite.CompositeIndexingExecutionEngine;
+import org.opensearch.index.engine.exec.coord.Any;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+public class CompositeMergeHandler extends MergeHandler {
+
+    private final CompositeMergePolicy mergePolicy;
+    private final CompositeEngine compositeEngine;
+    private final CompositeIndexingExecutionEngine compositeIndexingExecutionEngine;
+
+    public CompositeMergeHandler(
+        CompositeEngine compositeEngine,
+        CompositeIndexingExecutionEngine compositeIndexingExecutionEngine,
+        Any dataFormats,
+        IndexSettings indexSettings
+    ) {
+        super(compositeEngine, compositeIndexingExecutionEngine, dataFormats);
+        this.compositeEngine = compositeEngine;
+        this.compositeIndexingExecutionEngine = compositeIndexingExecutionEngine;
+
+        mergePolicy = new CompositeMergePolicy(indexSettings.getMergePolicy(true));
+    }
+
+    @Override
+    public Collection<OneMerge> findForceMerges(int maxSegmentCount) {
+        List<OneMerge> oneMerges = new ArrayList<>();
+        try (CompositeEngine.ReleasableRef<CatalogSnapshot> catalogSnapshotReleasableRef = compositeEngine.acquireSnapshot()) {
+            CatalogSnapshot catalogSnapshot = catalogSnapshotReleasableRef.getRef();
+
+            List<CatalogSnapshot.Segment> segmentList = catalogSnapshot.getSegments();
+            List<List<CatalogSnapshot.Segment>> mergeCandidates =
+                mergePolicy.findForceMergeCandidates(segmentList, maxSegmentCount);
+
+            // Process merge candidates
+            for (List<CatalogSnapshot.Segment> mergeGroup : mergeCandidates) {
+                oneMerges.add(new OneMerge(mergeGroup));
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        return oneMerges;
+    }
+
+    @Override
+    public Collection<OneMerge> findMerges() {
+        List<OneMerge> oneMerges = new ArrayList<>();
+        try (CompositeEngine.ReleasableRef<CatalogSnapshot> catalogSnapshotReleasableRef = compositeEngine.acquireSnapshot()) {
+            CatalogSnapshot catalogSnapshot = catalogSnapshotReleasableRef.getRef();
+
+            List<CatalogSnapshot.Segment> segmentList = catalogSnapshot.getSegments();
+            List<List<CatalogSnapshot.Segment>> mergeCandidates =
+                mergePolicy.findMergeCandidates(segmentList);
+
+            // Process merge candidates
+            for (List<CatalogSnapshot.Segment> mergeGroup : mergeCandidates) {
+                oneMerges.add(new OneMerge(mergeGroup));
+            }
+        } catch (Exception e) {
+            e.printStackTrace();
+            throw new RuntimeException(e);
+        }
+        return oneMerges;
+    }
+
+    public synchronized void registerMerge(OneMerge oneMerge) {
+        super.registerMerge(oneMerge);
+        mergePolicy.addMergingSegment(oneMerge.getSegmentsToMerge());
+    }
+
+    public synchronized void onMergeFinished(OneMerge oneMerge) {
+        super.onMergeFinished(oneMerge);
+        mergePolicy.removeMergingSegment(oneMerge.getSegmentsToMerge());
+    }
+
+    public synchronized void onMergeFailure(OneMerge oneMerge) {
+        super.onMergeFailure(oneMerge);
+        mergePolicy.removeMergingSegment(oneMerge.getSegmentsToMerge());
+    }
+}