opensearch-project · mch2 · Aug 27, 2026 · Aug 21, 2026 · Aug 21, 2026 · Aug 24, 2026
@@ -33,6 +33,7 @@ public class ShardScanExecutionContext implements CommonExecutionContext {
     private final Task task;
     private byte[] fragmentBytes;
     private BufferAllocator allocator;
+    private BufferAllocator importStagingAllocator;
     private MapperService mapperService;
     private IndexSettings indexSettings;
     private NamedWriteableRegistry namedWriteableRegistry;
@@ -89,6 +90,31 @@ public void setAllocator(BufferAllocator allocator) {
         this.allocator = allocator;
     }
 
+    /**
+     * Returns the allocator Arrow C Data Interface imports are staged on, or null if the caller did not
+     * set one. Distinct from {@link #getAllocator()} on two properties a batch import depends on:
+     * <ul>
+     *   <li><b>Unbounded, parented at the root.</b> {@code Data#importIntoVectorSchemaRoot} charges each
+     *       buffer as it walks the array; a target that fills part-way through throws, and arrow-java
+     *       (&le; 18.1.0) retains the imported array before the throwing {@code wrapForeignAllocation}
+     *       without rolling back, so the C Data release callback never fires and the whole batch leaks in
+     *       the producer's native allocator — invisible to the JVM heap and to Java Arrow accounting.</li>
+     *   <li><b>Node-scoped, never closed per request.</b> The Flight transport builds its reused stream
+     *       root on the FIRST emitted batch's vector allocator and charges that same allocator for every
+     *       later batch ({@code FlightServerChannel#transferIntoStreamRoot}: "The producer's allocator must
+     *       be long-lived (not closed per-request)"), and frees it asynchronously with its own channel.</li>
+     * </ul>
+     * Caller-owned: the engine imports onto it and must never close it.
+     */
+    public BufferAllocator getImportStagingAllocator() {
+        return importStagingAllocator;
+    }
+
+    /** Sets the node-scoped import staging allocator. The caller owns its lifecycle. */
+    public void setImportStagingAllocator(BufferAllocator importStagingAllocator) {
+        this.importStagingAllocator = importStagingAllocator;
+    }
+
     /** Returns the shard's mapper service for field type resolution. */
     public MapperService getMapperService() {
         return mapperService;

@@ -220,14 +220,17 @@ default CanMatchResult canMatchWithBounds(IndexShard shard, byte[] filterBytes,
      * @param rowIdVector Arrow BigIntVector containing global row IDs
      * @param columns column names to read
      * @param allocator Arrow buffer allocator for result import
+     * @param importStagingAllocator node-scoped allocator to stage Arrow C Data imports on; see
+     *        {@link org.opensearch.analytics.backend.ShardScanExecutionContext#getImportStagingAllocator()}
      * @return a result stream containing the requested rows
      */
     default EngineResultStream fetchByRowIds(
         Reader reader,
         BigIntVector rowIdVector,
         String[] columns,
         BufferAllocator allocator,
-        long contextId
+        long contextId,
+        BufferAllocator importStagingAllocator
     ) {
         throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]");
     }

@@ -43,12 +43,37 @@
  *   <li>{@code downstream} — sink the backend drains its reduced output
  *       into. The backend owns {@code downstream}'s lifecycle: it must
  *       feed every produced batch and close it when draining is complete.</li>
+ *   <li>{@code importStagingAllocator} — node-scoped allocator the backend stages
+ *       Arrow C Data Interface imports on. Unbounded and parented at the root so an
+ *       import cannot fail part-way through an array (which strands the whole native
+ *       batch — see {@link org.opensearch.analytics.backend.ShardScanExecutionContext#getImportStagingAllocator()}),
+ *       and long-lived because the Flight transport keeps charging it after the
+ *       importing stream closes. Caller-owned: the backend must never close it, and
+ *       must never derive a per-stream child to import onto — an un-closed child stays
+ *       registered in the root's {@code childAllocators} map until node restart.</li>
  * </ul>
  *
  * @opensearch.internal
  */
 public record ExchangeSinkContext(String queryId, int stageId, long taskId, byte[] fragmentBytes, BufferAllocator allocator, List<
-    ChildInput> childInputs, ExchangeSink downstream) implements CommonExecutionContext {
+    ChildInput> childInputs, ExchangeSink downstream, BufferAllocator importStagingAllocator) implements CommonExecutionContext {
+
+    /**
+     * Stages imports on {@code allocator} itself. For callers whose allocator is already an unbounded
+     * root — chiefly tests: that is the pre-staging behaviour, correct but without the mid-import-OOM
+     * mitigation described above. Production paths pass a dedicated staging allocator.
+     */
+    public ExchangeSinkContext(
+        String queryId,
+        int stageId,
+        long taskId,
+        byte[] fragmentBytes,
+        BufferAllocator allocator,
+        List<ChildInput> childInputs,
+        ExchangeSink downstream
+    ) {
+        this(queryId, stageId, taskId, fragmentBytes, allocator, childInputs, downstream, allocator);
+    }
 
     /**
      * Per-child input descriptor: the child stage id and the producer-side plan bytes the

@@ -168,7 +168,15 @@ public byte[] getExecutionMetrics() {
     protected final void drainOutputIntoDownstream(StreamHandle outStream) {
         BufferAllocator alloc = ctx.allocator();
         try (CDataDictionaryProvider dictProvider = new CDataDictionaryProvider()) {
-            DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(outStream, alloc, dictProvider);
+            // Imports go onto the caller-owned, node-scoped staging allocator — never a child minted here:
+            // downstream may hand the batch to the Flight transport, which keeps charging that allocator
+            // long after this drain returns.
+            DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(
+                outStream,
+                alloc,
+                ctx.importStagingAllocator(),
+                dictProvider
+            );
             while (it.hasNext()) {
                 // next() transfers ownership of the imported VSR to us. feed() takes ownership only
                 // on success; if it throws (e.g. the downstream sink was torn down on a concurrent

@@ -1079,7 +1079,8 @@ public EngineResultStream fetchByRowIds(
         BigIntVector rowIdVector,
         String[] columns,
         BufferAllocator allocator,
-        long contextId
+        long contextId,
+        BufferAllocator importStagingAllocator
     ) {
         DataFusionService dataFusionService = plugin.getDataFusionService();
         if (dataFusionService == null) {
@@ -1115,7 +1116,7 @@ public EngineResultStream fetchByRowIds(
             throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0");
         }
         StreamHandle streamHandle = new StreamHandle(streamPtr, dataFusionService.getNativeRuntime());
-        return new DatafusionResultStream(streamHandle, allocator);
+        return new DatafusionResultStream(streamHandle, allocator, importStagingAllocator);
     }
 
     @Override

@@ -27,10 +27,10 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.action.ActionListener;
 
-import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 
 import static org.apache.arrow.c.Data.importField;
@@ -48,20 +48,29 @@ public class DatafusionResultStream implements EngineResultStream, FragmentResou
 
     private final StreamHandle streamHandle;
     private final BufferAllocator allocator;
+    private final BufferAllocator stagingAllocator;
     private final CDataDictionaryProvider dictionaryProvider;
     private volatile BatchIterator iteratorInstance;
 
-    // Allocator is caller-owned; this stream imports into it but never closes it.
-    public DatafusionResultStream(StreamHandle streamHandle, BufferAllocator allocator) {
+    /**
+     * Both allocators are caller-owned; this stream imports onto them but never closes either.
+     *
+     * @param stagingAllocator node-scoped, unbounded allocator every batch of every stream is imported onto
+     *        (see {@link BatchIterator#importBatch}). MUST outlive this stream: the Flight transport builds
+     *        its reused stream root on the first batch's vector allocator and frees it asynchronously with
+     *        its own channel.
+     */
+    public DatafusionResultStream(StreamHandle streamHandle, BufferAllocator allocator, BufferAllocator stagingAllocator) {
         this.streamHandle = streamHandle;
         this.allocator = allocator;
+        this.stagingAllocator = Objects.requireNonNull(stagingAllocator, "stagingAllocator");
         this.dictionaryProvider = new CDataDictionaryProvider();
     }
 
     @Override
     public Iterator<EngineResultBatch> iterator() {
         if (iteratorInstance == null) {
-            iteratorInstance = new BatchIterator(streamHandle, allocator, dictionaryProvider);
+            iteratorInstance = new BatchIterator(streamHandle, allocator, stagingAllocator, dictionaryProvider);
         }
         return iteratorInstance;
     }
@@ -76,7 +85,6 @@ public void close() {
         try {
             if (iteratorInstance != null) {
                 iteratorInstance.closeLastBatch();
-                iteratorInstance.reclaimDrainedStaging();
             }
         } finally {
             try {
@@ -93,19 +101,34 @@ static class BatchIterator implements Iterator<EngineResultBatch> {
 
         private final StreamHandle streamHandle;
         private final BufferAllocator allocator;
+        /**
+         * Caller-owned, node-scoped, unbounded staging allocator every batch is imported onto (see
+         * {@link #importBatch}). Owned by {@code AnalyticsSearchService}, never by this stream: the Flight
+         * transport builds its reused stream root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the
+         * FIRST batch's staging allocator ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment
+         * states "The producer's allocator must be long-lived (not closed per-request)"), then charges that
+         * same allocator for every later batch via {@code BaseFixedWidthVector#transferTo} and frees the
+         * stream root asynchronously in its own {@code close()}. So this stream can neither close it at batch
+         * boundaries (the transport is still using it) nor at stream close (the transport's free may run
+         * after ours) — it does not own it at all.
+         */
+        private final BufferAllocator stagingAllocator;
         private final CDataDictionaryProvider dictionaryProvider;
         private Schema schema;
         private VectorSchemaRoot nextBatch;
         private Boolean nextAvailable;
         private boolean batchEmitted;
         private boolean nativeStreamExhausted;
-        // Per-batch staging allocators used by {@link #importBatch}. Each is reclaimed once its batch's
-        // buffers have been released by the consumer (see {@link #reclaimDrainedStaging}).
-        private final List<BufferAllocator> stagingAllocators = new ArrayList<>();
 
-        BatchIterator(StreamHandle streamHandle, BufferAllocator allocator, CDataDictionaryProvider dictionaryProvider) {
+        BatchIterator(
+            StreamHandle streamHandle,
+            BufferAllocator allocator,
+            BufferAllocator stagingAllocator,
+            CDataDictionaryProvider dictionaryProvider
+        ) {
             this.streamHandle = streamHandle;
             this.allocator = allocator;
+            this.stagingAllocator = stagingAllocator;
             this.dictionaryProvider = dictionaryProvider;
         }
 
@@ -147,48 +170,39 @@ private boolean loadNextBatch() {
         }
 
         /**
-         * Imports one native batch across the Arrow C Data Interface into a per-batch staging allocator
-         * (an unbounded child of the root) rather than directly into {@code allocator}.
+         * Imports one native batch across the Arrow C Data Interface onto the caller-supplied
+         * {@link #stagingAllocator} rather than directly into {@code allocator}.
          *
          * <p>{@link Data#importIntoVectorSchemaRoot} charges each buffer against the target allocator as it
          * walks the array. Against a bounded target that fills part-way through a wide batch the import
          * throws, and arrow-java's {@code ReferenceCountedArrowArray#unsafeAssociateAllocation} retains the
          * imported array <em>before</em> the throwing {@code wrapForeignAllocation} without rolling back, so
          * the C Data release callback never fires and the whole native batch leaks in the producer's native
-         * allocator — invisible to the JVM heap and the Java Arrow allocator (arrow-java &le; 18.1.0). An
-         * unbounded staging child can't OOM mid-array, so the release callback always fires.
+         * allocator — invisible to the JVM heap and the Java Arrow allocator (arrow-java &le; 18.1.0). The
+         * staging allocator is unbounded and parented at the root, so it can't OOM mid-array before the pool
+         * itself is exhausted and the release callback always fires.
          *
          * <p>The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close
-         * paths, which drives the C Data reference count to zero. Each staging allocator is reclaimed once
-         * drained (see {@link #reclaimDrainedStaging}); on import failure it is closed immediately.
+         * paths, which drives the C Data reference count to zero. On import failure the partially-imported
+         * root is closed by {@link #importOntoStaging}; the allocator is untouched either way.
          */
         private VectorSchemaRoot importBatch(ArrowArray arrowArray) {
-            reclaimDrainedStaging();
-            BufferAllocator staging = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE);
-            try {
-                VectorSchemaRoot root = importOntoStaging(staging, schema, arrowArray, dictionaryProvider);
-                stagingAllocators.add(staging);
-                return root;
-            } catch (RuntimeException e) {
-                staging.close();
-                throw e;
-            }
+            return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider);
+        }
+
+        /** The caller-supplied staging allocator batches are imported onto. For tests. */
+        BufferAllocator stagingAllocator() {
+            return stagingAllocator;
         }
 
         /**
-         * Closes staging allocators whose batches have been fully released (drained to zero). A batch still
-         * in flight keeps its staging allocator open so the eventual release callback frees the small C Data
-         * bookkeeping allocation against a live allocator; that allocator is a leaf child of the root and
-         * holds no batch data once drained.
+         * Drives the exact production {@link #importBatch} path with an explicitly supplied schema, so the
+         * staging-allocator lifetime regression test can import several batches without a native stream.
+         * {@code schema} is normally set by {@link #ensureSchema()} from the native handle.
          */
-        private void reclaimDrainedStaging() {
-            stagingAllocators.removeIf(a -> {
-                if (a.getAllocatedMemory() == 0) {
-                    a.close();
-                    return true;
-                }
-                return false;
-            });
+        VectorSchemaRoot importBatchForTest(Schema batchSchema, ArrowArray arrowArray) {
+            this.schema = batchSchema;
+            return importBatch(arrowArray);
         }
 
         /**
@@ -208,7 +222,14 @@ static VectorSchemaRoot importOntoStaging(
             try {
                 Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider);
             } catch (RuntimeException e) {
-                root.close();
+                // Releasing a partially-imported root can itself throw (VectorSchemaRoot#close rethrows any
+                // RuntimeException from the vectors' release). Attach it rather than let it mask the import
+                // failure that is the real diagnosis.
+                try {
+                    root.close();
+                } catch (RuntimeException releaseFailure) {
+                    e.addSuppressed(releaseFailure);
+                }
                 throw e;
             }
             return root;

@@ -48,6 +48,12 @@ public EngineResultStream execute(ShardScanExecutionContext requestContext) thro
         if (allocator == null) {
             throw new IllegalStateException("ExecutionContext.allocator must be set by the caller before execute()");
         }
+        // Node-scoped and unbounded — batches are imported onto it and the Flight transport keeps charging it
+        // after this stream closes, so it must be supplied rather than minted per stream.
+        BufferAllocator stagingAllocator = requestContext.getImportStagingAllocator();
+        if (stagingAllocator == null) {
+            throw new IllegalStateException("ExecutionContext.importStagingAllocator must be set by the caller before execute()");
+        }
 
         // Register cancellation hook so HTTP disconnect / _tasks/_cancel / timeout
         // immediately fires the Rust CancellationToken.
@@ -63,7 +69,7 @@ public EngineResultStream execute(ShardScanExecutionContext requestContext) thro
         DatafusionSearcher searcher = datafusionContext.getSearcher();
         searcher.search(datafusionContext);
         StreamHandle handle = datafusionContext.takeStreamHandle();
-        return new DatafusionResultStream(handle, allocator);
+        return new DatafusionResultStream(handle, allocator, stagingAllocator);
     }
 
     @Override