Inline merge at refresh for composite engine: produce one segment per… (opensearch-project#21831)

* Merge on refresh for composite engine with inline merge

Implements merge-on-refresh: when multiple writers flush in the same
refresh cycle, consolidate all formats into a single segment.

Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>

Author: Bukhtawar

sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java

@@ -8,6 +8,8 @@
 
 package org.opensearch.be.datafusion;
 
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
+
 import org.opensearch.analytics.backend.jni.NativeHandle;
 import org.opensearch.be.datafusion.nativelib.NativeBridge;
 import org.opensearch.be.datafusion.nativelib.ReaderHandle;
@@ -28,6 +30,9 @@
  * Smoke test for the DataFusion JNI bridge.
  * Verifies native library loading, runtime creation, and reader lifecycle.
  */
+// The Tokio IO runtime thread is a process-lifetime singleton spawned by the native Rust library.
+// It persists after tests complete and cannot be interrupted (empty Java stack, RUNNABLE state).
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
 public class DataFusionNativeBridgeTests extends OpenSearchTestCase {
 
     // Note: initTokioRuntimeManager uses OnceLock and can only be initialized once per JVM.

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

@@ -252,6 +252,94 @@ public void testSortedParquetPrimaryLuceneSecondaryMerge() throws Exception {
         verifyCrossFormatConsistency(snapshot);
     }
 
+    /**
+     * Validates inline consolidation at refresh with fileMappings:
+     * When multiple writers flush in the same refresh cycle, the primary (Parquet) merges
+     * them and produces a RowIdMapping via fileMappings. The secondary (Lucene) then applies
+     * the same mapping. This test uses concurrent indexing to fill multiple writers, then
+     * a single refresh to trigger consolidation.
+     *
+     * Correctness criteria:
+     * <ol>
+     *   <li>After refresh, catalog has 1 segment (consolidated) instead of N per-writer segments</li>
+     *   <li>Both parquet and lucene formats are present</li>
+     *   <li>Lucene __row_id__ values are sequential (RowIdMapping correctly applied)</li>
+     *   <li>Cross-format consistency: parquet row data matches lucene data at same row_id</li>
+     * </ol>
+     */
+    public void testMergeOnRefresh() throws Exception {
+        client().admin()
+            .indices()
+            .prepareCreate(INDEX_NAME)
+            .setSettings(sortedParquetPrimaryLuceneSecondarySettings())
+            .setMapping("name", "type=keyword", "age", "type=integer")
+            .get();
+        ensureGreen(INDEX_NAME);
+
+        // Use concurrent threads to fill multiple writers in a single refresh cycle.
+        // With refresh_interval=-1 and no explicit refresh, all docs land in the writer pool.
+        int numThreads = 8;
+        int docsPerThread = 5;
+        int totalDocs = numThreads * docsPerThread;
+        java.util.concurrent.CountDownLatch startLatch = new java.util.concurrent.CountDownLatch(1);
+        java.util.concurrent.atomic.AtomicReference<Exception> error = new java.util.concurrent.atomic.AtomicReference<>();
+
+        Thread[] threads = new Thread[numThreads];
+        for (int t = 0; t < numThreads; t++) {
+            int threadId = t;
+            threads[t] = new Thread(() -> {
+                try {
+                    startLatch.await();
+                    for (int i = 0; i < docsPerThread; i++) {
+                        int docId = threadId * docsPerThread + i;
+                        client().prepareIndex(INDEX_NAME)
+                            .setId(String.valueOf(docId))
+                            .setSource("name", "doc_" + docId, "age", randomIntBetween(0, 100))
+                            .get();
+                    }
+                } catch (Exception e) {
+                    error.compareAndSet(null, e);
+                }
+            });
+            threads[t].start();
+        }
+        startLatch.countDown();
+        for (Thread t : threads) {
+            t.join();
+        }
+        if (error.get() != null) {
+            throw error.get();
+        }
+
+        // Single refresh flushes all writers — triggers inline consolidation via fileMappings
+        client().admin().indices().prepareRefresh(INDEX_NAME).get();
+        client().admin().indices().prepareFlush(INDEX_NAME).setForce(true).setWaitIfOngoing(true).get();
+
+        DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot();
+
+        // Merge on refresh must have consolidated: with N writers (N>1) flushing,
+        // the catalog should have exactly 1 segment (merged) instead of N.
+        // If this assertion fails, consolidation didn't trigger (all docs landed in 1 writer).
+        assertEquals("Inline merge on refresh should produce exactly 1 segment from multiple writers", 1, snapshot.getSegments().size());
+
+        // Both formats must be present in the single consolidated segment
+        Set<String> formats = snapshot.getDataFormats();
+        assertTrue("Catalog should contain 'parquet'", formats.contains("parquet"));
+        assertTrue("Catalog should contain 'lucene'", formats.contains("lucene"));
+
+        // Verify total row count
+        verifyRowCount(snapshot, totalDocs);
+
+        // Verify lucene doc count
+        verifyLuceneDocCount(totalDocs);
+
+        // RowIdMapping correctness: __row_id__ must be sequential within each segment
+        verifyLuceneRowIdSequential();
+
+        // Cross-format field value consistency
+        verifyCrossFormatConsistency(snapshot);
+    }
+
     // ══════════════════════════════════════════════════════════════════════
     // Private helpers: merge feature flag
     // ══════════════════════════════════════════════════════════════════════
@@ -908,13 +996,19 @@ public void testMergeStatsViaApi() throws Exception {
 
     private DataformatAwareCatalogSnapshot waitForMerge(int refreshCycles) throws Exception {
         flush(INDEX_NAME);
+        // With inline merge at refresh, each cycle already produces 1 consolidated segment.
+        // Background merge may further reduce, but for small segments it may not fire.
+        // Accept: segment count <= refreshCycles (inline consolidation working correctly).
         assertBusy(() -> {
             DataformatAwareCatalogSnapshot snap = getCatalogSnapshot();
             assertTrue(
-                "Expected merges to reduce segment count below " + refreshCycles + ", but got: " + snap.getSegments().size(),
-                snap.getSegments().size() < refreshCycles
+                "Expected segment count <= "
+                    + refreshCycles
+                    + " (inline consolidation or background merge), but got: "
+                    + snap.getSegments().size(),
+                snap.getSegments().size() <= refreshCycles
             );
-        });
+        }, 30, java.util.concurrent.TimeUnit.SECONDS);
         return getCatalogSnapshot();
     }