Add sorting support (#20798)

* Adding sort support during merge operation

Signed-off-by: Dhwanil Patel <dhwanip@amazon.com>

* Add sorting support

* Fix sort flow

* Fix date field type for sort in merge

* Renamed @timestamp to timestamp

* Changing sort to EventDate

* Fix memory issue for larger merge

* Using polars for merging

* Added hybrid approach of polars and heap

* optimize merge by grouping contiguous rows

* merge with streaming using sink_batches

* streaming k way merge

* using sink_batch to stream one batch per file

* optimize - materialize chunks at flush time

* replace threaded sink_batches with sequential ParquetReader slicing to reduce memory usage

* optimize - binary search to find cut point in batch + emit entire batch by checking last row in batch

* use arrow writer - ipc

* use rayon to paralleize writes across columns

* add tokio for parallel column writes

* add rayon with thread pool for column encoding during flush

* add tokio for async write and rayon parallel decoding during reads through the shared rayon thread pool

* temp commit

* take sort column from index setting during indexing flow

* support reverse sort in merge

* store and look up settings for every index from settings store

* refactor and address comments

* add exception handling + take union of all schema to make code extensible for dynamic mapping support

* add profiler

---------

Signed-off-by: Dhwanil Patel <dhwanip@amazon.com>
Co-authored-by: Dhwanil Patel <dhwanip@amazon.com>
Co-authored-by: Mohit Godwani <mgodwan@amazon.com>
Co-authored-by: Shailesh Singh <shaikumm@amazon.com>
Co-authored-by: Shailesh Singh <shaileshkumarsingh260@gmaill.com>

Author: 5 people

modules/parquet-data-format/benchmarks/src/main/java/com/parquet/parquetdataformat/benchmark/ParquetWriterCloseBenchmark.java

@@ -60,7 +60,7 @@ public void setup() throws IOException {
         writerCreationBenchmarkData = generator.generate("simple", fieldCount, 0);
         writerWriteBenchmarkData = generator.generate("simple", fieldCount, recordCount);
         filePath = generateTempFilePath();
-        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress());
+        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress(), "sortColumn", false);
         RustBridge.write(filePath, writerWriteBenchmarkData.getArrowArray().memoryAddress(), writerWriteBenchmarkData.getArrowSchema().memoryAddress());
     }
 

modules/parquet-data-format/benchmarks/src/main/java/com/parquet/parquetdataformat/benchmark/ParquetWriterCreateBenchmark.java

@@ -81,7 +81,7 @@ public void tearDown() throws IOException {
     @Benchmark
     public void benchmarkCreate() throws IOException {
         // This is what we're benchmarking - just writer creation
-        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress());
+        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress(), null, false);
     }
 
     private String generateTempFilePath() {

modules/parquet-data-format/benchmarks/src/main/java/com/parquet/parquetdataformat/benchmark/ParquetWriterWriteBenchmark.java

@@ -39,7 +39,7 @@ public void setup() throws IOException {
         writerCreationBenchmarkData = generator.generate("simple", fieldCount, 0);
         writerWriteBenchmarkData = generator.generate("simple", fieldCount, recordCount);
         filePath = generateTempFilePath();
-        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress());
+        RustBridge.createWriter(filePath, "benchmark-index", writerCreationBenchmarkData.getArrowSchema().memoryAddress(), null, false);
     }
 
     @Benchmark

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/bridge/NativeParquetWriter.java

@@ -23,12 +23,15 @@ public class NativeParquetWriter implements Closeable {
     /**
      * Creates a new native Parquet writer.
      * @param filePath path to the Parquet file
+     * @param indexName name of the index (used for settings lookup)
      * @param schemaAddress Arrow C Data Interface schema pointer
+     * @param sortColumn column to sort by
+     * @param reverseSort whether to sort in reverse order
      * @throws IOException if writer creation fails
      */
-    public NativeParquetWriter(String filePath, String indexName, long schemaAddress) throws IOException {
+    public NativeParquetWriter(String filePath, String indexName, long schemaAddress, String sortColumn, boolean reverseSort) throws IOException {
         this.filePath = filePath;
-        RustBridge.createWriter(filePath, indexName, schemaAddress);
+        RustBridge.createWriter(filePath, indexName, schemaAddress, sortColumn, reverseSort);
     }
 
     /**

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/bridge/RustBridge.java

@@ -29,7 +29,7 @@ public class RustBridge {
     public static native void initLogger();
 
     // Enhanced native methods that handle validation and provide better error reporting
-    public static native void createWriter(String file, String indexName, long schemaAddress) throws IOException;
+    public static native void createWriter(String file, String indexName, long schemaAddress, String sortColumn, boolean reverseSort) throws IOException;
     public static native void write(String file, long arrayAddress, long schemaAddress) throws IOException;
     public static native ParquetFileMetadata closeWriter(String file) throws IOException;
     public static native void flushToDisk(String file) throws IOException;
@@ -42,5 +42,5 @@ public class RustBridge {
 
 
     // Native method declarations - these will be implemented in the JNI library
-    public static native void mergeParquetFilesInRust(List<Path> inputFiles, String outputFile, String indexName);
+    public static native void mergeParquetFilesInRust(List<Path> inputFiles, String outputFile, String indexName, String sortKey, boolean isReverse);
 }

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/engine/ParquetExecutionEngine.java

@@ -77,6 +77,8 @@ public class ParquetExecutionEngine implements IndexingExecutionEngine<ParquetDa
     private final ParquetMerger parquetMerger;
     private final ArrowBufferPool arrowBufferPool;
     private final IndexSettings indexSettings;
+    private volatile String sortColumn;
+    private volatile boolean reverseSort;
     private final boolean isPrimaryEngine;
 
     public ParquetExecutionEngine(
@@ -90,7 +92,7 @@ public ParquetExecutionEngine(
         this.shardPath = shardPath;
         this.arrowBufferPool = new ArrowBufferPool(settings);
         this.indexSettings = indexSettings;
-        this.parquetMerger = new ParquetMergeExecutor(CompactionStrategy.RECORD_BATCH, indexSettings.getIndex().getName());
+        this.parquetMerger = new ParquetMergeExecutor(CompactionStrategy.RECORD_BATCH);
         this.isPrimaryEngine = isPrimaryEngine;
         // Push current settings to Rust store once on construction, then keep in sync on updates
         pushSettingsToRust(indexSettings);
@@ -108,6 +110,16 @@ public ParquetExecutionEngine(
 //        );
     }
 
+    @Override
+    public void setSortColumn(String sortColumn) {
+        this.sortColumn = sortColumn;
+    }
+
+    @Override
+    public void setReverseSort(boolean reverseSort) {
+        this.reverseSort = reverseSort;
+    }
+
     private void pushSettingsToRust(IndexSettings indexSettings) {
         NativeSettings config = new NativeSettings();
         config.setIndexName(indexSettings.getIndex().getName());
@@ -155,7 +167,7 @@ public List<String> supportedFieldTypes(boolean isPrimaryEngine) {
     public Writer<ParquetDocumentInput> createWriter(long writerGeneration) {
         String fileName = Path.of(shardPath.getDataPath().toString(), getDataFormat().name(), FILE_NAME_PREFIX + "_" + writerGeneration + FILE_NAME_EXT).toString();
         EngineRole role = isPrimaryEngine ? EngineRole.PRIMARY : EngineRole.SECONDARY;
-        return new ParquetWriter(fileName, schema.get(), writerGeneration, arrowBufferPool, indexSettings, role);
+        return new ParquetWriter(fileName, schema.get(), writerGeneration, arrowBufferPool, indexSettings, sortColumn, reverseSort, role);
     }
 
     @Override

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

@@ -8,25 +8,21 @@
 
 package com.parquet.parquetdataformat.merge;
 
-import org.opensearch.index.engine.exec.WriterFileSet;
+import org.opensearch.index.engine.exec.merge.MergeInput;
 import org.opensearch.index.engine.exec.merge.MergeResult;
-import java.util.List;
-
 /**
  * Executes Parquet merge operations using a chosen compaction strategy.
  */
 public class ParquetMergeExecutor extends ParquetMerger {
 
     private final ParquetMergeStrategy strategy;
-    private final String indexName;
 
-    public ParquetMergeExecutor(CompactionStrategy compactionStrategy, String indexName) {
+    public ParquetMergeExecutor(CompactionStrategy compactionStrategy) {
         this.strategy = ParquetMergeStrategyFactory.getStrategy(compactionStrategy);
-        this.indexName = indexName;
     }
 
     @Override
-    public MergeResult merge(List<WriterFileSet> fileMetadataList, long writerGeneration) {
-        return strategy.mergeParquetFiles(fileMetadataList, writerGeneration, indexName);
+    public MergeResult merge(MergeInput mergeInput) {
+        return strategy.mergeParquetFiles(mergeInput);
     }
 }

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

@@ -8,9 +8,9 @@
 
 package com.parquet.parquetdataformat.merge;
 
-import org.opensearch.index.engine.exec.WriterFileSet;
+
+import org.opensearch.index.engine.exec.merge.MergeInput;
 import org.opensearch.index.engine.exec.merge.MergeResult;
-import java.util.List;
 
 /**
  * Interface defining a Parquet merge strategy.
@@ -20,6 +20,6 @@ public interface ParquetMergeStrategy {
     /**
      * Performs the actual Parquet merge.
      */
-    MergeResult mergeParquetFiles(List<WriterFileSet> files, long writerGeneration, String indexName);
+    MergeResult mergeParquetFiles(MergeInput mergeInput);
 
 }

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

@@ -17,7 +17,7 @@ public static ParquetMergeStrategy getStrategy(CompactionStrategy compactionStra
         switch (compactionStrategy) {
             case RECORD_BATCH:
             default:
-                return new RecordBatchMergeStrategy();
+                return new StreamingParquetMergeStrategy();
         }
     }
 }

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

@@ -8,18 +8,14 @@
 
 package com.parquet.parquetdataformat.merge;
 
-import org.opensearch.index.engine.exec.FileMetadata;
+import org.opensearch.index.engine.exec.merge.MergeInput;
 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(List<WriterFileSet> fileMetadataList, RowIdMapping rowIdMapping, long writerGeneration) {
+    public MergeResult merge(MergeInput mergeInput, RowIdMapping rowIdMapping) {
         throw new UnsupportedOperationException("Not supported parquet as secondary data format yet.");
     }
 }