diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java index 6681cd65a0cbc..714295b9342d5 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java @@ -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: + * + * 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; diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index 071d1578561f1..1fc98dc7e7da7 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -220,6 +220,8 @@ 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( @@ -227,7 +229,8 @@ default EngineResultStream fetchByRowIds( BigIntVector rowIdVector, String[] columns, BufferAllocator allocator, - long contextId + long contextId, + BufferAllocator importStagingAllocator ) { throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]"); } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java index f374c689aa722..ae14fbde36538 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java @@ -43,12 +43,37 @@ *
  • {@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.
    • + *
  • {@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.
    • * * * @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 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 diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java index eb2c3855f498f..5cbb68415954e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java @@ -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 diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 1e3f75b5d9800..d4e24eb555176 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -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 diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index c07574b76abed..be6b68ac08306 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -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 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 { 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 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}. * *

    {@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 before 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 ≤ 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 ≤ 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. * *

    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; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java index cd0721d212552..72b37f0c95ec7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java @@ -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 diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java index 52a7b82f40de1..bdc3bc1455707 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java @@ -89,6 +89,17 @@ static final class NativeBridgeExecutor implements DocumentRowReader, Closeable private final DataFusionPlugin dfPlugin; private final BufferAllocator sharedAllocator = new RootAllocator(64 * 1024 * 1024); + /** + * Unbounded staging child every batch of every stream this executor opens is imported onto — never + * one per stream, see {@code DatafusionResultStream.BatchIterator#stagingAllocator}. These streams are + * drained inline (not handed to the Flight transport), so it is always drained by the time + * {@link #close()} runs. + */ + private final BufferAllocator importStagingAllocator = sharedAllocator.newChildAllocator( + "datafusion-get-import-staging", + 0, + Long.MAX_VALUE + ); NativeBridgeExecutor(DataFusionPlugin dfPlugin) { this.dfPlugin = dfPlugin; @@ -96,6 +107,7 @@ static final class NativeBridgeExecutor implements DocumentRowReader, Closeable @Override public void close() { + importStagingAllocator.close(); sharedAllocator.close(); } @@ -165,7 +177,7 @@ private List> readAllRows(long streamPtr) { List> results = new ArrayList<>(); try ( StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); - DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator, importStagingAllocator) ) { var iter = stream.iterator(); while (iter.hasNext()) { @@ -188,7 +200,7 @@ private List> readAllRows(long streamPtr) { private Map readSingleRow(long streamPtr) { try ( StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); - DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator, importStagingAllocator) ) { var iter = stream.iterator(); if (!iter.hasNext()) return null; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java index 1136afa93b98b..8138100a2c93f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java @@ -40,6 +40,8 @@ public class DatafusionResultStreamTests extends OpenSearchTestCase { private long queryConfigPtr; private long tieredStorePtr; private NativeStoreHandle storeHandle; + /** Stands in for the node-scoped staging allocator AnalyticsSearchService hands to every stream. */ + private BufferAllocator stagingAllocator; private final java.util.List allocatorsToClose = new java.util.ArrayList<>(); @Override @@ -50,6 +52,8 @@ public void setUp() throws Exception { long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024); runtimeHandle = new NativeRuntimeHandle(ptr); testRootAllocator = new RootAllocator(Long.MAX_VALUE); + stagingAllocator = testRootAllocator.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); + allocatorsToClose.add(stagingAllocator); // Create a real TieredObjectStore (local-only) and wrap in NativeStoreHandle tieredStorePtr = NativeStoreTestHelper.createTieredObjectStore(0L, 0L); @@ -272,7 +276,8 @@ public void onFailure(Exception e) { allocatorsToClose.add(failureAlloc); DatafusionResultStream stream = new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, tempRuntime), - failureAlloc + failureAlloc, + stagingAllocator ); // Close runtime — streamNext should now fail with IllegalStateException from NativeRuntimeHandle.get() @@ -323,7 +328,8 @@ private DatafusionResultStream createStreamWithLimit(String sql, long limitBytes allocatorsToClose.add(childAllocator); return new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, runtimeHandle), - childAllocator + childAllocator, + stagingAllocator ); } @@ -358,7 +364,8 @@ private DatafusionResultStream createStream(String sql) { allocatorsToClose.add(childAllocator); return new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, runtimeHandle), - childAllocator + childAllocator, + stagingAllocator ); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java index 48b380ea44056..ebee77c1022ba 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java @@ -8,6 +8,7 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; @@ -87,10 +88,12 @@ public void testEngineExecuteSelectAll() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); @@ -116,10 +119,12 @@ public void testEngineExecuteAggregation() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); @@ -142,10 +147,12 @@ public void testEngineExecuteFilter() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java new file mode 100644 index 0000000000000..fad807a96ac67 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java @@ -0,0 +1,198 @@ +/* + * 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.be.datafusion; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.CDataDictionaryProvider; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * Regression tests for the staging-allocator OWNERSHIP contract: the allocator batches are imported onto is + * supplied by the caller ({@code AnalyticsSearchService} in production, via + * {@code ShardScanExecutionContext#getImportStagingAllocator}), node-scoped, and never created or closed by + * the stream. Two separate bugs are pinned here, both rooted in a stream owning its staging allocator: + * + *

      + *
    • Use-after-close. {@code FlightServerChannel#transferIntoStreamRoot} builds its long-lived + * stream root on the FIRST emitted batch's vector allocator — its comment states "The producer's + * allocator must be long-lived (not closed per-request)" — and charges that same allocator for every + * later batch. A per-batch staging allocator closed as soon as its own buffers drained was therefore + * closed while the transport was still using it, throwing {@code IllegalStateException: Memory was + * leaked by query} from {@code BaseAllocator#close}. The window is a couple of statements wide with no + * injection point, so what these tests assert is the deterministic invariant behind it: every batch of a + * stream imports onto the SAME allocator, which is still OPEN afterwards.
      • + *
    • Stranded child allocators. Keeping a per-stream allocator open instead (to dodge the race) + * leaks it until node reboot: {@code BaseAllocator#newChildAllocator} registers every child in the + * parent's {@code childAllocators} map unconditionally, and only the child's own {@code close()} + * deregisters it — so an un-closed staging child stays strongly referenced by the node-lifetime root, + * one per query. {@link #testStreamsDoNotMintOrStrandChildAllocators} pins that the stream mints no + * child allocator at all.
      • + *
    + */ +public class DatafusionStagingAllocatorLifecycleTests extends OpenSearchTestCase { + + private static final int ROWS = 8192; // DataFusion's default batch size + private RootAllocator root; + /** Stands in for the node-scoped allocator AnalyticsSearchService owns and hands to every stream. */ + private BufferAllocator staging; + /** Stands in for the native (Rust) allocator owning the exported buffers; long-lived like the real one. */ + private BufferAllocator producer; + // Resources the production DatafusionResultStream#close would own; the test must release them itself. + private final List providers = new ArrayList<>(); + + @Override + public void setUp() throws Exception { + super.setUp(); + root = new RootAllocator(Long.MAX_VALUE); + staging = root.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); + producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); + } + + @Override + public void tearDown() throws Exception { + // No stream closes the staging allocator — the caller does, exactly as AnalyticsSearchService#close + // does in production. + for (CDataDictionaryProvider dp : providers) { + dp.close(); + } + staging.close(); + producer.close(); + root.close(); + super.tearDown(); + } + + /** + * Imports three batches through the production path and asserts all three land on the SAME staging + * allocator — the caller-supplied one — which is still open at the end. + */ + public void testAllBatchesOfAStreamShareOneOpenStagingAllocator() throws Exception { + DatafusionResultStream.BatchIterator it = newIterator(); + + for (int i = 0; i < 3; i++) { + VectorSchemaRoot imported = importOneBatch(it); + assertSame( + "every batch of a stream must import onto the caller-supplied staging allocator — the Flight " + + "stream root is built on the first batch's allocator and reused for all later batches", + staging, + imported.getFieldVectors().getFirst().getAllocator() + ); + // The consumer releases this batch. Under the old per-batch code this drained the allocator to + // zero and made it eligible for the sweep on the NEXT import. + imported.close(); + } + + assertSame("the iterator must not substitute an allocator of its own", staging, it.stagingAllocator()); + // The decisive assertion: still usable by the transport after later imports + earlier releases. + assertEquals("staging allocator must NOT have been closed per-batch", 0L, staging.getAllocatedMemory()); + staging.assertOpen(); + } + + /** + * Pins the exact sequence that threw in production: import batch 1, release it (allocator drains to + * zero), then import batch 2 — under the old code the sweep at the head of {@code importBatch} closed + * batch 1's allocator, which is the very allocator the Flight stream root was built on. + */ + public void testReleasingABatchDoesNotCloseTheAllocatorTheTransportHolds() throws Exception { + DatafusionResultStream.BatchIterator it = newIterator(); + + VectorSchemaRoot batch1 = importOneBatch(it); + BufferAllocator transportAllocator = batch1.getFieldVectors().getFirst().getAllocator(); + batch1.close(); // consumer drains it to zero + assertEquals( + "precondition: allocator is drained, i.e. sweep-eligible under the old design", + 0L, + transportAllocator.getAllocatedMemory() + ); + + VectorSchemaRoot batch2 = importOneBatch(it); // old code swept batch1's allocator here + try { + // Simulates FlightServerChannel#transferIntoStreamRoot charging the long-lived allocator for a + // LATER batch. On a closed allocator this throws (IllegalStateException under -ea, which tests + // run with) — that is the production failure, reproduced deterministically. + try (VectorSchemaRoot streamRoot = VectorSchemaRoot.create(batch2.getSchema(), transportAllocator)) { + streamRoot.allocateNew(); + streamRoot.setRowCount(1); + } + } finally { + batch2.close(); + } + } + + /** + * The leak regression: streaming N results must not add a single child allocator to the root. A stream + * that mints its own staging child either has to close it (racing the transport — see the tests above) or + * leave it open, in which case the root's {@code childAllocators} map retains it for the node's lifetime. + * Neither is acceptable, so the stream mints nothing. + */ + public void testStreamsDoNotMintOrStrandChildAllocators() throws Exception { + Set before = new HashSet<>(root.getChildAllocators()); + + for (int stream = 0; stream < 3; stream++) { + DatafusionResultStream.BatchIterator it = newIterator(); + VectorSchemaRoot imported = importOneBatch(it); + imported.close(); // consumer releases the batch + it.closeLastBatch(); // what DatafusionResultStream#close does to the iterator + } + + assertEquals( + "importing on a caller-supplied staging allocator must add no child allocator to the root — an " + + "un-closed child stays registered in the parent's childAllocators map until node reboot", + before, + new HashSet<>(root.getChildAllocators()) + ); + staging.assertOpen(); + } + + // ── helpers ──────────────────────────────────────────────────────────────────────────────────── + + private DatafusionResultStream.BatchIterator newIterator() { + // streamHandle is never touched by importBatch, so a null handle is safe and keeps the test free of + // the native runtime (no .so needed). + CDataDictionaryProvider dp = new CDataDictionaryProvider(); + providers.add(dp); + return new DatafusionResultStream.BatchIterator(null, root, staging, dp); + } + + /** Exports a fresh batch across the C Data Interface and imports it via the production path. */ + private VectorSchemaRoot importOneBatch(DatafusionResultStream.BatchIterator it) { + try (ArrowArray array = ArrowArray.allocateNew(producer); ArrowSchema cSchema = ArrowSchema.allocateNew(producer)) { + try (VectorSchemaRoot source = VectorSchemaRoot.create(intSchema(producer), producer)) { + IntVector v = (IntVector) source.getVector(0); + v.allocateNew(ROWS); + for (int i = 0; i < ROWS; i++) { + v.set(i, i); + } + source.setRowCount(ROWS); + Data.exportVectorSchemaRoot(producer, source, null, array, cSchema); + } + // Consume the exported C schema (not source.getSchema()) so its C-side release callback fires; + // an unconsumed ArrowSchema strands its exported children in the producer allocator. + Schema schema = Data.importSchema(root, cSchema, providers.getLast()); + return it.importBatchForTest(schema, array); + } + } + + private static Schema intSchema(BufferAllocator alloc) { + try (IntVector probe = new IntVector("n", alloc)) { + return new Schema(List.of(probe.getField())); + } + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java index 357b40fcaea76..5495d8aac3fa6 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java @@ -23,10 +23,10 @@ import org.opensearch.analytics.exec.ArrowValues; import org.opensearch.common.annotation.ExperimentalApi; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; +import java.util.Objects; import static org.apache.arrow.c.Data.importField; @@ -53,24 +53,31 @@ public class LuceneResultStream implements EngineResultStream { /** C-Data schema describing {@link #arrowArray}. */ private final ArrowSchema arrowSchema; private final BufferAllocator allocator; + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private volatile BatchIterator iteratorInstance; /** * Caller hands over ownership of {@code arrowArray} and {@code arrowSchema}; this stream - * closes them in {@link #close}. + * closes them in {@link #close}. Both allocators stay caller-owned. + * + * @param stagingAllocator node-scoped, unbounded allocator the batch 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 LuceneResultStream(ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator) { + public LuceneResultStream(ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator, BufferAllocator stagingAllocator) { this.arrowArray = arrowArray; this.arrowSchema = arrowSchema; this.allocator = allocator; + this.stagingAllocator = Objects.requireNonNull(stagingAllocator, "stagingAllocator"); this.dictionaryProvider = new CDataDictionaryProvider(); } @Override public Iterator iterator() { if (iteratorInstance == null) { - iteratorInstance = new BatchIterator(arrowArray, arrowSchema, allocator, dictionaryProvider); + iteratorInstance = new BatchIterator(arrowArray, arrowSchema, allocator, stagingAllocator, dictionaryProvider); } return iteratorInstance; } @@ -80,7 +87,6 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); } } finally { try { @@ -106,25 +112,34 @@ static class BatchIterator implements Iterator { private final ArrowArray arrowArray; private final ArrowSchema arrowSchema; private final BufferAllocator allocator; + /** + * Caller-owned, node-scoped, unbounded staging allocator the 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 and frees the stream root asynchronously in its own + * {@code close()} — which may run after ours. + */ + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private Schema schema; private VectorSchemaRoot nextBatch; private Boolean nextAvailable; private boolean batchEmitted; private boolean exhausted; - // 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 stagingAllocators = new ArrayList<>(); BatchIterator( ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator, + BufferAllocator stagingAllocator, CDataDictionaryProvider dictionaryProvider ) { this.arrowArray = arrowArray; this.arrowSchema = arrowSchema; this.allocator = allocator; + this.stagingAllocator = stagingAllocator; this.dictionaryProvider = dictionaryProvider; } @@ -147,52 +162,41 @@ private boolean loadNextBatch() { } /** - * Imports the 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 the batch across the Arrow C Data Interface onto the caller-supplied + * {@link #stagingAllocator} rather than directly into {@code allocator}. * *

    {@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 before the throwing {@code wrapForeignAllocation} without rolling back, so * the C Data release callback never fires and the whole native batch leaks — invisible to the JVM - * heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). An unbounded staging child can't OOM - * mid-array, so the release callback always fires. + * heap and the Java Arrow allocator (arrow-java ≤ 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. * *

    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. The staging allocator is caller-owned and + * outlives this stream (see {@link #stagingAllocator}), so nothing here closes it. */ private VectorSchemaRoot importBatch() { - reclaimDrainedStaging(); - BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, stagingAllocator); try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + Data.importIntoVectorSchemaRoot(stagingAllocator, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { - root.close(); - staging.close(); + // Close the partially-imported root (fires the native release); the allocator is caller-owned + // and untouched. The release can itself throw (VectorSchemaRoot#close rethrows any + // RuntimeException from the vectors); 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; } - stagingAllocators.add(staging); return root; } - /** - * 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. - */ - private void reclaimDrainedStaging() { - stagingAllocators.removeIf(a -> { - if (a.getAllocatedMemory() == 0) { - a.close(); - return true; - } - return false; - }); - } - @Override public boolean hasNext() { if (nextAvailable == null) { diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java index 701840299c8b1..bd8e4dd6c28e1 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java @@ -77,13 +77,19 @@ public EngineResultStream execute(ShardScanExecutionContext context) throws IOEx state.outputColumnNames() ); BufferAllocator allocator = context.getAllocator(); + // Node-scoped and unbounded — the batch is 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 = context.getImportStagingAllocator(); + if (stagingAllocator == null) { + throw new IllegalStateException("ExecutionContext.importStagingAllocator must be set by the caller before execute()"); + } Schema schema = buildSchema(state.outputColumnNames()); ArrowArray array = ArrowArray.allocateNew(allocator); ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); boolean transferred = false; try { populateBatchToCData(allocator, schema, state.outputColumnNames(), count, array, arrowSchema); - LuceneResultStream stream = new LuceneResultStream(array, arrowSchema, allocator); + LuceneResultStream stream = new LuceneResultStream(array, arrowSchema, allocator, stagingAllocator); transferred = true; return stream; } finally { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java index 39f6d54c871e2..2c30ea396038a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java @@ -225,8 +225,11 @@ public Collection createComponents( // plugin's lifecycle owns its lifetime. The Guice-bound DefaultPlanExecutor consumes // it via the handle without taking on close responsibility — mirroring how // AnalyticsSearchService's allocator is owned and closed by this plugin. + // The import staging allocator is created and closed by the search service; the handle only + // carries it so the coordinator-reduce path stages its imports on the same node-scoped allocator. coordinatorAllocatorHandle = new CoordinatorAllocatorHandle( - nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).newChildAllocator("coordinator", 0, Long.MAX_VALUE) + nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).newChildAllocator("coordinator", 0, Long.MAX_VALUE), + searchService.getImportStagingAllocator() ); // Returned as components so Guice can inject them into DefaultPlanExecutor diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index 1046cbaa873f1..520647468289f 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -97,6 +97,19 @@ public class AnalyticsSearchService implements AutoCloseable { private org.opensearch.threadpool.ThreadPool threadPool; private org.opensearch.cluster.service.ClusterService clusterService; private final BufferAllocator allocator; + /** + * ONE node-scoped allocator that every backend stages its Arrow C Data Interface imports on, handed to + * the backends via {@link ShardScanExecutionContext#setImportStagingAllocator} and + * {@link AnalyticsSearchBackendPlugin#fetchByRowIds}. Deliberately not per-stream: 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)"), then frees it asynchronously with its own + * channel. A per-stream staging child therefore cannot be closed at stream close without racing the + * transport, and one that is left open instead stays registered in the root's {@code childAllocators} + * map for the node's lifetime ({@code BaseAllocator#newChildAllocator} registers unconditionally; only + * the child's own {@code close()} deregisters) — one stranded allocator per query. + */ + private final BufferAllocator importStagingAllocator; private final ArrowNativeAllocator nativeAllocator; public AnalyticsSearchService(Map backends, ArrowNativeAllocator nativeAllocator) { @@ -131,13 +144,47 @@ public AnalyticsSearchService( // effect immediately via Arrow's parent-cap check at allocateBytes — no listener needed. BufferAllocator queryPool = nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); this.allocator = queryPool.newChildAllocator("analytics-search-service", 0, Long.MAX_VALUE); + // Parented at the ROOT rather than the query pool, and unbounded: a C Data import that fails + // part-way through an array strands the whole native batch (see + // ShardScanExecutionContext#getImportStagingAllocator), so the staging target must not be able to + // fill up before the pool itself is exhausted. + this.importStagingAllocator = allocator.getRoot().newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); this.namedWriteableRegistry = namedWriteableRegistry; this.readerContextStore = readerContextStore; } @Override public void close() { - allocator.close(); + // finally, so a leak report from the service allocator cannot skip the staging allocator's release. + try { + allocator.close(); + } finally { + // A stream still in flight at shutdown keeps its batches charged to the staging allocator until + // the Flight channel frees its stream root — and that free is posted to the flight executor, + // which may already refuse tasks (FlightServerChannel#close logs "root reclaimed at process + // exit"). Closing on a non-zero balance would throw from BaseAllocator#close AND strand the + // bytes, because close() sets isClosed before its leak check. Report and let process exit + // reclaim them instead. + long outstanding = importStagingAllocator.getAllocatedMemory(); + if (outstanding == 0) { + importStagingAllocator.close(); + } else { + LOGGER.warn( + "Arrow import staging allocator [{}] still holds {} bytes at shutdown; reclaimed at process exit", + importStagingAllocator.getName(), + outstanding + ); + } + } + } + + /** + * The node-scoped allocator Arrow C Data imports are staged on. Owned (and closed) here; the + * coordinator-reduce path borrows the same one via {@code CoordinatorAllocatorHandle} so a node has + * exactly one, whatever the execution path. + */ + public BufferAllocator getImportStagingAllocator() { + return importStagingAllocator; } public void setTaskResourceTrackingService(TaskResourceTrackingService service) { @@ -488,6 +535,7 @@ public void executeWorkerFragmentStreamingAsync( ShardScanExecutionContext ctx = new ShardScanExecutionContext(/* tableName */ "", task, /* reader */ null); ctx.setFragmentBytes(plan.getFragmentBytes()); ctx.setAllocator(allocator); + ctx.setImportStagingAllocator(importStagingAllocator); ctx.setNamedWriteableRegistry(namedWriteableRegistry); ctx.setShuffleBufferRegistry(shuffleBufferRegistry); @@ -707,7 +755,8 @@ private void drainFetchByRowIds( rowIdVector, columns, allocator, - task.getNativeTaskId() + task.getNativeTaskId(), + importStagingAllocator ); // FragmentResources keeps the rowIdVector alive until the stream drains — closing // it earlier would pull off-heap memory out from under the native FFM call. @@ -914,7 +963,8 @@ private ExchangeSink buildPartitionedSink( new byte[0], ctx.getAllocator(), List.of(), - /* downstream */ null + /* downstream */ null, + importStagingAllocator ); return backend.getExchangeSinkProvider() .createPartitionedSink( @@ -1061,6 +1111,7 @@ private ShardScanExecutionContext buildContext( ShardScanExecutionContext ctx = new ShardScanExecutionContext(tableName, task, reader); ctx.setFragmentBytes(plan.getFragmentBytes()); ctx.setAllocator(allocator); + ctx.setImportStagingAllocator(importStagingAllocator); ctx.setMapperService(shard.mapperService()); ctx.setIndexSettings(shard.indexSettings()); ctx.setNamedWriteableRegistry(namedWriteableRegistry); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java index ab6f4a6d786b4..b53122422f07f 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java @@ -24,18 +24,27 @@ * *

    The plugin owns the lifetime; consumers ({@code DefaultPlanExecutor}) only * read the underlying allocator via {@link #getAllocator()}. + * + *

    Also carries the node-scoped Arrow import staging allocator so the coordinator-reduce path can + * reach it. That one is borrowed, not owned: {@code AnalyticsSearchService} creates and closes + * it, so a node has exactly one whatever the execution path. */ public final class CoordinatorAllocatorHandle implements AutoCloseable { private final BufferAllocator allocator; + private final BufferAllocator importStagingAllocator; /** * Wraps an already-constructed coordinator allocator (typically a child of * {@code POOL_QUERY}). The caller (the plugin) retains responsibility for * having created the allocator under the right parent. + * + * @param importStagingAllocator the node-scoped import staging allocator owned by + * {@code AnalyticsSearchService}; borrowed here, never closed by this handle */ - public CoordinatorAllocatorHandle(BufferAllocator allocator) { + public CoordinatorAllocatorHandle(BufferAllocator allocator, BufferAllocator importStagingAllocator) { this.allocator = allocator; + this.importStagingAllocator = importStagingAllocator; } /** Returns the wrapped allocator. Consumers must not close it directly. */ @@ -43,6 +52,16 @@ public BufferAllocator getAllocator() { return allocator; } + /** + * Returns the node-scoped allocator Arrow C Data imports are staged on — unbounded and parented at + * the root so an import cannot fail part-way through an array, and long-lived because the Flight + * transport keeps charging it after the importing stream closes. Owned by + * {@code AnalyticsSearchService}; consumers must not close it. + */ + public BufferAllocator getImportStagingAllocator() { + return importStagingAllocator; + } + @Override public void close() { allocator.close(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java index ad4f94be09dcc..6861772cd55d4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java @@ -125,6 +125,8 @@ public class DefaultPlanExecutor extends HandledTransportAction perQueryBufferLimit = v); @@ -530,6 +533,9 @@ private void executeInternal( ownsAllocator, profile ); + // Node-scoped and unbounded: coordinator-reduce batches are imported onto it and the Flight + // transport keeps charging it after the importing stream closes, so it must not be per-query. + context.setImportStagingAllocator(importStagingAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); throw e; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java index 113e1e855f7a3..2a52fd813a52d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java @@ -53,6 +53,8 @@ public class QueryContext { private final List operationListeners; private final BufferAllocator allocator; private final boolean ownsAllocator; + /** Caller-owned; see {@link #importStagingAllocator()}. Never closed by this context. */ + private BufferAllocator importStagingAllocator; /** Whether profiling is enabled for this query (data nodes should collect and return metrics). */ private final boolean profile; /** @@ -311,6 +313,25 @@ public BufferAllocator bufferAllocator() { return allocator; } + /** + * Returns the node-scoped allocator coordinator-side Arrow C Data imports are staged on — unbounded + * and parented at the root so an import cannot fail part-way through an array (which strands the whole + * native batch, see {@code ShardScanExecutionContext#getImportStagingAllocator}), and long-lived + * because the Flight transport keeps charging it after the importing stream closes. + * + *

    Falls back to {@link #bufferAllocator()} when unset — that is the pre-staging behaviour, correct + * but without the mid-import-OOM mitigation, and it keeps test contexts (whose allocators are unbounded + * roots anyway) working without wiring one. + */ + public BufferAllocator importStagingAllocator() { + return importStagingAllocator != null ? importStagingAllocator : allocator; + } + + /** Set once by {@code DefaultPlanExecutor} after construction. The caller owns the allocator. */ + public void setImportStagingAllocator(BufferAllocator importStagingAllocator) { + this.importStagingAllocator = importStagingAllocator; + } + /** Lazy per-query virtual-thread executor for LOCAL tasks. Shared across phased contexts. */ public ExecutorService localTaskExecutor() { ExecutorService exec = sharedState.localTaskExecutor; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java index 60a1803750064..292fd258967b0 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java @@ -50,7 +50,8 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte chosenBytes(stage), config.bufferAllocator(), List.of(), - sink + sink, + config.importStagingAllocator() ); ExchangeSink backendSink = provider.createSink(context, null); if (backendSink instanceof ReducingExchangeSink reducing) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java index 0e5aad1cb61d1..494086f93bfa2 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java @@ -51,7 +51,8 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte chosenBytes(stage), config.bufferAllocator(), buildChildInputs(stage), - sink + sink, + config.importStagingAllocator() ); // Apply instruction handlers for the reduce stage.