Objectstore integration on native write and merge path

sandbox/libs/tiered-storage/src/main/rust/src/ffm.rs

@@ -58,6 +58,56 @@ unsafe fn arc_from_ptr(ptr: i64) -> Result<Arc<TieredObjectStore>, String> {
     Ok(Arc::from_raw(raw))
 }
 
+// ---------------------------------------------------------------------------
+// Generic ObjectStore handle lifecycle (write-path, shard-scoped)
+// ---------------------------------------------------------------------------
+//
+// These entry points mint a plain `Box<Arc<dyn ObjectStore>>` handle that the
+// parquet write/merge paths use as a shard-scoped sink/source. For hot indices
+// this is a `LocalFileSystem` rooted at the shard directory; the warm path can
+// later hand back the same-shaped handle backed by a `TieredObjectStore`, so
+// consumers (writer/merge/search) never change.
+//
+// Ownership model: the engine owns the master handle and is the ONLY caller of
+// `os_destroy_store` (exactly once, at shard close). Consumers (parquet
+// writer/merge) clone the `Arc` inline from the raw box pointer so the store
+// outlives any in-flight op.
+
+/// Create a local filesystem `ObjectStore` rooted at `path`, returned as a
+/// `Box<Arc<dyn ObjectStore>>` raw pointer (i64). Object paths are then plain
+/// filenames relative to `path`. Free with [`os_destroy_store`].
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn os_create_local_store(path_ptr: *const u8, path_len: i64) -> i64 {
+    let path = str_from_raw(path_ptr, path_len)
+        .map_err(|e| format!("os_create_local_store path: {}", e))?;
+    let store: Arc<dyn ObjectStore> = Arc::new(
+        object_store::local::LocalFileSystem::new_with_prefix(path).map_err(|e| {
+            format!(
+                "os_create_local_store: failed to open LocalFileSystem at '{}': {}",
+                path, e
+            )
+        })?,
+    );
+    let ptr = Box::into_raw(Box::new(store)) as i64;
+    native_bridge_common::log_info!("ffm: os_create_local_store root='{}' ok", path);
+    Ok(ptr)
+}
+
+/// Drop a `Box<Arc<dyn ObjectStore>>` created by [`os_create_local_store`],
+/// decrementing the `Arc` strong count. Must be called exactly once, by the
+/// engine, at shard close. Outstanding consumer clones keep the store alive.
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn os_destroy_store(handle: i64) -> i64 {
+    if handle == NULL_PTR {
+        return Err("os_destroy_store: null pointer (0)".to_string());
+    }
+    let _ = Box::from_raw(handle as *mut Arc<dyn ObjectStore>);
+    native_bridge_common::log_info!("ffm: os_destroy_store ok");
+    Ok(0)
+}
+
 // ---------------------------------------------------------------------------
 // Public FFM exports
 // ---------------------------------------------------------------------------

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

@@ -189,6 +189,73 @@ public void testSortedRefreshWithRandomizedData() throws Exception {
         verifyLuceneRowIdSequential();
     }
 
+    /**
+     * Forces many small sorted chunks via a tiny {@code index.parquet.sort_in_memory_threshold},
+     * so that finalize runs the store-backed k-way merge over N>1 chunks. This is the only test
+     * that exercises the multi-chunk merge path: sorted chunks written to / read from / deleted in
+     * the ObjectStore plus the merge output streamed back through the store. (With a 16kb threshold
+     * and 10k docs this flushes dozens of chunks — see the "merging N pre-sorted chunks" native
+     * log.) Correctness of the merged output is asserted via full row count and global sort order.
+     */
+    public void testSortedRefreshMultiChunkThroughStore() throws Exception {
+        int numFields = 5;
+        String[] fieldNames = new String[numFields];
+        String[] sortFields = new String[numFields];
+        String[] sortOrders = new String[numFields];
+        String[] sortMissing = new String[numFields];
+        for (int f = 0; f < numFields; f++) {
+            fieldNames[f] = "field_" + f;
+            sortFields[f] = fieldNames[f];
+            sortOrders[f] = (f % 2 == 0) ? "asc" : "desc";
+            sortMissing[f] = sortOrders[f].equals("asc") ? "_last" : "_first";
+        }
+
+        Settings settings = Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            .put("index.refresh_interval", "-1")
+            .put("index.pluggable.dataformat.enabled", true)
+            .put("index.pluggable.dataformat", "composite")
+            .put("index.composite.primary_data_format", "parquet")
+            .putList("index.composite.secondary_data_formats", "lucene")
+            // Tiny threshold -> many sorted chunks -> multi-chunk k-way merge through the store.
+            .put("index.parquet.sort_in_memory_threshold", "16kb")
+            .putList("index.sort.field", sortFields)
+            .putList("index.sort.order", sortOrders)
+            .putList("index.sort.missing", sortMissing)
+            .build();
+
+        String[] mappingArgs = new String[numFields * 2];
+        for (int f = 0; f < numFields; f++) {
+            mappingArgs[f * 2] = fieldNames[f];
+            mappingArgs[f * 2 + 1] = "type=integer";
+        }
+        client().admin().indices().prepareCreate(INDEX_NAME).setSettings(settings).setMapping(mappingArgs).get();
+        ensureGreen(INDEX_NAME);
+
+        int totalDocs = 10000;
+        for (int i = 0; i < totalDocs; i++) {
+            Map<String, Object> source = new HashMap<>();
+            for (int f = 0; f < numFields; f++) {
+                // Occasionally emit null values to exercise null handling across chunk boundaries
+                if (randomIntBetween(0, 9) == 0) {
+                    continue; // skip field → null
+                }
+                source.put(fieldNames[f], randomIntBetween(0, 1000));
+            }
+            IndexResponse response = client().prepareIndex().setIndex(INDEX_NAME).setSource(source).get();
+            assertEquals(RestStatus.CREATED, response.status());
+        }
+
+        flushAndRefresh();
+
+        DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot();
+        verifyParquetRowCount(snapshot, totalDocs);
+        verifyParquetSortOrderMultiField(snapshot, fieldNames, sortOrders, sortMissing);
+        verifyLuceneDocCount(totalDocs);
+        verifyLuceneRowIdSequential();
+    }
+
     /**
      * Parquet primary + Lucene secondary without sort. Without a sort
      * permutation, Parquet does not produce a {@code RowIdMapping} and the

sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java

@@ -14,6 +14,7 @@
 import org.opensearch.plugin.stats.StatsRecorder;
 
 import java.io.IOException;
+import java.lang.ref.Cleaner;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -37,8 +38,28 @@ public class NativeParquetWriter {
     private final SetOnce<ParquetFileMetadata> metadata = new SetOnce<>();
     private final SetOnce<RowIdMapping> rowIdMapping = new SetOnce<>();
     private final ParquetShardStatsTracker stats;
+    /**
+     * Shard-scoped native ObjectStore handle ({@code Box<Arc<dyn ObjectStore>>} pointer) minted by
+     * {@code os_create_local_store} and owned by the engine, or 0 for the legacy local-file path.
+     * Passed to the native writer at {@link #initialize} time so finalize publishes through the store.
+     */
+    private final long storePtr;
     private volatile boolean initialized = false;
 
+    /** Reclaims leaked native writers if this object is GC'd without an explicit close/flush. */
+    private static final Cleaner CLEANER = Cleaner.create();
+
+    /**
+     * Opaque native handle (a {@code Box<WriterState>} pointer) minted by {@link #initialize};
+     * 0 until initialized. Every native call that dereferences it runs inside a {@code synchronized}
+     * method, so they are serialized (write/finalize/free/memory never race on the same handle).
+     */
+    private volatile long handle = 0L;
+    /** Set once the handle has been consumed (by finalize) or freed (by close/Cleaner). */
+    private final AtomicBoolean released = new AtomicBoolean(false);
+    /** Cleaner registration; {@code clean()} on close deregisters the GC backstop. */
+    private Cleaner.Cleanable cleanable;
+
     /**
      * Creates a new NativeParquetWriter handle. Does not create the native writer —
      * call {@link #initialize(String, long, ParquetSortConfig, long)} before the first write.
@@ -47,8 +68,21 @@ public class NativeParquetWriter {
      * @param stats shard-level stats tracker
      */
     public NativeParquetWriter(String filePath, ParquetShardStatsTracker stats) {
+        this(filePath, stats, 0L);
+    }
+
+    /**
+     * Creates a new NativeParquetWriter handle bound to a shard-scoped ObjectStore.
+     *
+     * @param filePath the path to the Parquet file to write
+     * @param stats shard-level stats tracker
+     * @param storePtr shard ObjectStore handle ({@code os_create_local_store} pointer), or 0 for
+     *                 the legacy local-file path
+     */
+    public NativeParquetWriter(String filePath, ParquetShardStatsTracker stats, long storePtr) {
         this.filePath = filePath;
         this.stats = stats;
+        this.storePtr = storePtr;
     }
 
     /**
@@ -71,12 +105,16 @@ public NativeParquetWriter(String filePath) {
      * @throws IOException if the native writer creation fails
      * @throws IllegalStateException if already initialized
      */
-    public void initialize(String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) throws IOException {
+    public synchronized void initialize(String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration)
+        throws IOException {
         if (initialized) {
             throw new IllegalStateException("Writer already initialized: " + filePath);
         }
-        RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration);
+        handle = RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration, storePtr);
         initialized = true;
+        // GC backstop: if this writer is dropped without close()/flush(), free the native handle.
+        // HandleCleanup holds only the primitive handle + shared released flag (never `this`).
+        cleanable = CLEANER.register(this, new HandleCleanup(handle, released));
     }
 
     /**
@@ -96,15 +134,15 @@ public boolean isInitialized() {
      * @throws IOException if the write fails or the writer is flushed
      * @throws IllegalStateException if the writer has not been initialized
      */
-    public void write(long arrayAddress, long schemaAddress) throws IOException {
+    public synchronized void write(long arrayAddress, long schemaAddress) throws IOException {
         if (writerFlushed.get()) {
             throw new IOException("Cannot write to flushed Parquet writer: " + filePath);
         }
         if (initialized == false) {
             throw new IllegalStateException("Writer not initialized: " + filePath);
         }
         StatsRecorder.recordOutcome(
-            () -> RustBridge.write(filePath, arrayAddress, schemaAddress),
+            () -> RustBridge.write(handle, arrayAddress, schemaAddress),
             stats::addNativeWriteTimeMillis,
             stats::incNativeWriteTotal,
             stats::incNativeWriteFailures
@@ -119,18 +157,24 @@ public void write(long arrayAddress, long schemaAddress) throws IOException {
      * @return the file metadata, or null if the writer was never initialized
      * @throws IOException if the finalization fails
      */
-    public ParquetFileMetadata flush() throws IOException {
+    public synchronized ParquetFileMetadata flush() throws IOException {
         if (writerFlushed.compareAndSet(false, true)) {
             if (initialized) {
-                StatsRecorder.recordOutcome(() -> {
-                    RustBridge.WriterFinalizeResult result = RustBridge.finalizeWriter(filePath);
-                    if (result != null) {
-                        metadata.set(result.metadata());
-                        if (result.rowIdMapping() != null) {
-                            rowIdMapping.set(result.rowIdMapping());
+                try {
+                    StatsRecorder.recordOutcome(() -> {
+                        RustBridge.WriterFinalizeResult result = RustBridge.finalizeWriter(handle);
+                        if (result != null) {
+                            metadata.set(result.metadata());
+                            if (result.rowIdMapping() != null) {
+                                rowIdMapping.set(result.rowIdMapping());
+                            }
                         }
-                    }
-                }, stats::addNativeFinalizeTimeMillis, stats::incNativeFinalizeTotal, stats::incNativeFinalizeFailures);
+                    }, stats::addNativeFinalizeTimeMillis, stats::incNativeFinalizeTotal, stats::incNativeFinalizeFailures);
+                } finally {
+                    // finalize_writer reclaims the native Box regardless of outcome, so the
+                    // handle is consumed — mark released so close()/Cleaner never double-free.
+                    released.set(true);
+                }
             }
         }
         return metadata.get();
@@ -153,4 +197,54 @@ public RowIdMapping getRowIdMapping() {
         return rowIdMapping.get();
     }
 
+    /**
+     * Returns the native memory currently reserved by this writer, or 0 if it was never
+     * initialized or has already been finalized/freed. Blocks if a native write/finalize is in
+     * flight (stats tolerate waiting; the write path is unaffected functionally).
+     */
+    public synchronized long getNativeBytesUsed() {
+        if (initialized == false || released.get()) {
+            return 0L;
+        }
+        return RustBridge.getWriterMemoryUsage(handle);
+    }
+
+    /**
+     * Releases the native writer if it was not already consumed by {@link #flush()}. Idempotent —
+     * safe to call multiple times and after flush. This is the guaranteed cleanup path invoked on
+     * writer teardown; the {@link Cleaner} is only a backstop for the case where {@code close()} is
+     * never called (e.g. the writer is dropped after a failure).
+     */
+    public synchronized void close() {
+        if (handle != 0L && released.compareAndSet(false, true)) {
+            RustBridge.freeWriter(handle);
+        }
+        if (cleanable != null) {
+            cleanable.clean();
+        }
+    }
+
+    /**
+     * Cleaner action reclaiming the native writer if the {@link NativeParquetWriter} becomes
+     * unreachable without an explicit {@link #close()}/{@link #flush()}. It captures only the
+     * primitive handle and the shared {@code released} flag — never the enclosing writer — so it
+     * cannot keep the writer reachable, and the CAS guarantees free happens at most once.
+     */
+    private static final class HandleCleanup implements Runnable {
+        private final long handle;
+        private final AtomicBoolean released;
+
+        HandleCleanup(long handle, AtomicBoolean released) {
+            this.handle = handle;
+            this.released = released;
+        }
+
+        @Override
+        public void run() {
+            if (handle != 0L && released.compareAndSet(false, true)) {
+                RustBridge.freeWriter(handle);
+            }
+        }
+    }
+
 }

sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java

@@ -11,7 +11,7 @@
 /**
  * Metadata extracted from a Parquet file after the native writer is closed.
  *
- * <p>Returned by {@link RustBridge#finalizeWriter(String)} and {@link RustBridge#getFileMetadata(String)}.
+ * <p>Returned by {@link RustBridge#finalizeWriter(long)} and {@link RustBridge#getFileMetadata(String)}.
  * Contains the Parquet format version, total row count, the creator identifier string
  * embedded in the file footer, and the whole-file CRC32 checksum computed during write.
  *