[Analytics Engine] Fix Arrow import staging allocator lifetime (transport race and leaked child allocators)

[LantaoJin]

Description

Arrow C Data batches are imported onto a dedicated staging allocator (unbounded, parented at the root) because arrow-java ≤ 18.1.0 doesn't roll back a mid-import failure — the C Data release callback never fires and the whole native batch leaks in the producer's native allocator. That allocator was created and closed per stream, which can't work. FlightServerChannel builds its reused stream root on the first batch's allocator ("The producer's allocator must be long-lived (not closed per-request)") and frees it asynchronously, so closing at stream close raced the transport and threw IllegalStateException: Memory was leaked by query:

java.lang.IllegalStateException: Memory was leaked by query. Memory leaked: (1024)
Allocator(datafusion-import-staging) 0/66560/3394560/9223372036854775807 (res/actual/peak/limit)
        at org.apache.arrow.memory.BaseAllocator.close(BaseAllocator.java:504)
        at org.opensearch.be.datafusion.DatafusionResultStream$BatchIterator.lambda$reclaimDrainedStaging$0(DatafusionResultStream.java:187)

Fix: one node-scoped arrow-import-staging allocator, created and closed by AnalyticsSearchService. Streams borrow it and never create or close an allocator, so there's no race and nothing accumulates. It reaches all import paths — shard scan (ShardScanExecutionContext), QTF fetch (fetchByRowIds), and coordinator reduce (ExchangeSinkContext). This also fixes a pre-existing instance of the same leak in AbstractDatafusionReduceSink#drainOutputIntoDownstream, which minted a staging child per drain and never closed it.

Related Issues

Resolves #22799

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit f278551)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Potential allocator leak on shutdown When importStagingAllocator has outstanding memory at close(), the code logs a warning and skips close(), relying on process exit to reclaim it. If the JVM is not actually shutting down (e.g. plugin/service reload, tests, or partial teardown), this leaks the allocator and its accounted memory indefinitely without a mechanism to retry closing it later. Consider documenting this constraint more strictly or scheduling a deferred close once outstanding drops to zero. 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 ); } } Breaking SPI change The fetchByRowIds default method signature adds two parameters (long contextId, BufferAllocator importStagingAllocator). Any existing external backend implementation overriding the previous signature will no longer override it and will silently fall through to the default UnsupportedOperationException. Since this is an SPI, consider whether a deprecation/overload path is warranted, or ensure all callers now route through the new signature. default EngineResultStream fetchByRowIds( Reader reader, BigIntVector rowIdVector, String[] columns, BufferAllocator allocator, long contextId, BufferAllocator importStagingAllocator ) { throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]");

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to f278551

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent shutdown failure from outstanding child Skipping importStagingAllocator.close() when there are outstanding bytes leaves this child registered in its parent's childAllocators map. If the parent (the root allocator) is subsequently closed by another plugin, it will throw IllegalStateException: Allocator[ROOT] closed with outstanding child allocators. Consider also detaching/orphaning the allocator from its parent (or ensuring the root allocator does not attempt to enforce this) so the warn path does not turn a graceful shutdown into a shutdown error. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [168-177] 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 ); + // NOTE: allocator remains registered under its parent; parent close will fail unless + // the root allocator tolerates this or the allocator is detached here. } Suggestion importance[1-10]: 6 __ Why: Valid concern: skipping close() on an outstanding child could cause the parent root allocator to throw at shutdown. However, the suggestion only notes the issue without providing a concrete fix, and the PR comment indicates this is intentional behavior deferring to process exit.	Low
General	Fail fast on null allocator allocator is dereferenced elsewhere (e.g., iterator() and BatchIterator), so a null value would produce an obscure NPE deep in the pipeline. Add an Objects.requireNonNull(allocator, "allocator") alongside the staging one to fail fast at construction, consistent with the strengthened contract. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [63-67] public DatafusionResultStream(StreamHandle streamHandle, BufferAllocator allocator, BufferAllocator stagingAllocator) { this.streamHandle = streamHandle; - this.allocator = allocator; + this.allocator = Objects.requireNonNull(allocator, "allocator"); this.stagingAllocator = Objects.requireNonNull(stagingAllocator, "stagingAllocator"); this.dictionaryProvider = new CDataDictionaryProvider(); } Suggestion importance[1-10]: 3 __ Why: Minor defensive programming improvement for consistency with the newly added Objects.requireNonNull on stagingAllocator. Low impact since a null allocator would still fail quickly at first use.	Low
	Validate staging allocator wiring at startup If searchService is null or not yet initialized here, getImportStagingAllocator() will NPE. Guard against a null searchService (or ensure ordering) and consider asserting the returned allocator is non-null, so misconfiguration surfaces here rather than later during query execution when a coordinator-reduce path tries to stage imports on null. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [230-233] +BufferAllocator stagingAlloc = Objects.requireNonNull( + searchService.getImportStagingAllocator(), + "AnalyticsSearchService must expose a non-null import staging allocator" +); coordinatorAllocatorHandle = new CoordinatorAllocatorHandle( nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).newChildAllocator("coordinator", 0, Long.MAX_VALUE), - searchService.getImportStagingAllocator() + stagingAlloc ); Suggestion importance[1-10]: 3 __ Why: Adds a fail-fast null check at wiring time. Minor benefit since a misconfiguration would surface fairly quickly anyway, and there is no evidence searchService could be null at this point.	Low

---

Previous suggestions

Suggestions up to commit 5ef81a4

Category	Suggestion	Impact
Possible issue	Ensure staging allocator is truly unbounded sharedAllocator has a bounded limit of 64 MiB, so importStagingAllocator (its child) inherits that cap through Arrow's parent-cap check. This contradicts the documented invariant that the staging allocator must be "unbounded" so a C Data import cannot fail part-way through an array and strand the native batch. Consider parenting the staging allocator at a truly unbounded root or making sharedAllocator unbounded. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java [98-102] private final BufferAllocator sharedAllocator = new RootAllocator(64 * 1024 * 1024); -private final BufferAllocator importStagingAllocator = sharedAllocator.newChildAllocator( - "datafusion-get-import-staging", - 0, - Long.MAX_VALUE -); +// Parent unbounded so an Arrow C Data import cannot OOM mid-array; see +// ShardScanExecutionContext#getImportStagingAllocator. +private final BufferAllocator importStagingAllocator = new RootAllocator(Long.MAX_VALUE) + .newChildAllocator("datafusion-get-import-staging", 0, Long.MAX_VALUE); Suggestion importance[1-10]: 7 __ Why: Correctly identifies that the staging allocator inherits the 64 MiB cap from its parent sharedAllocator, contradicting the documented "unbounded" invariant. This could cause the very mid-import-OOM issue the staging allocator was designed to prevent, though the GetService path drains inline which reduces likelihood.	Medium
General	Preserve SPI backward compatibility Adding a new required parameter to this default SPI method is a source-incompatible break for any external backend plugin overriding the previous signature — their override will silently no longer be called, and the default UnsupportedOperationException will be invoked instead. Consider retaining a bridge default that delegates from the old signature or clearly documenting the required migration. sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java [227-236] default EngineResultStream fetchByRowIds( Reader reader, BigIntVector rowIdVector, String[] columns, BufferAllocator allocator, long contextId, BufferAllocator importStagingAllocator ) { + // Backwards-compat bridge: delegate to legacy overload for backends not yet updated. + return fetchByRowIds(reader, rowIdVector, columns, allocator, contextId); +} + +default EngineResultStream fetchByRowIds( + Reader reader, + BigIntVector rowIdVector, + String[] columns, + BufferAllocator allocator, + long contextId +) { throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]"); } Suggestion importance[1-10]: 5 __ Why: Valid observation about SPI source-incompatibility for external backend plugins, though this is an internal experimental API in a sandbox module, limiting impact. The proposed bridge would help migration but may not be desired given the intentional signature change.	Low
	Avoid stranding allocator on shutdown leak Leaving importStagingAllocator unclosed when it has outstanding memory strands it registered in its parent's childAllocators map, which will then cause the parent's close() to throw IllegalStateException for having outstanding child allocators — masking the underlying leak and potentially preventing clean shutdown of the arrow-base root. Consider still attempting close() inside a try/catch to surface the leak diagnostics from Arrow while allowing shutdown to proceed, or explicitly document that the parent root must tolerate this. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [168-177] 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 ); + try { + importStagingAllocator.close(); + } catch (RuntimeException e) { + LOGGER.warn("Failed to close import staging allocator cleanly at shutdown", e); + } } Suggestion importance[1-10]: 3 __ Why: The suggestion contradicts the explicit intent documented in the PR comment, which states that closing on a non-zero balance would throw and strand bytes because close() sets isClosed before its leak check. The proposed change would reintroduce exactly the failure mode the code deliberately avoids.	Low
	Fail fast on null allocators importStagingAllocator is used by DefaultPlanExecutor unconditionally and coordinator-reduce imports depend on it being non-null; a null passed here would surface only later as an obscure NPE inside an import path. Add an explicit null-check at the boundary to fail fast. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java [45-48] public CoordinatorAllocatorHandle(BufferAllocator allocator, BufferAllocator importStagingAllocator) { - this.allocator = allocator; - this.importStagingAllocator = importStagingAllocator; + this.allocator = java.util.Objects.requireNonNull(allocator, "allocator"); + this.importStagingAllocator = java.util.Objects.requireNonNull(importStagingAllocator, "importStagingAllocator"); } Suggestion importance[1-10]: 3 __ Why: Minor defensive-programming improvement. Adding null checks at the boundary provides clearer error messages but has limited impact since a null would surface quickly at first use.	Low

Suggestions up to commit 1deaf21

Category	Suggestion	Impact
General	Ensure parent allocator closes on child failure If importStagingAllocator.close() throws (e.g. outstanding memory from a still-in-flight release), sharedAllocator.close() is skipped, leaking the root and its children. Wrap the staging close in a try/finally so the parent close always runs. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java [109-112] private final BufferAllocator importStagingAllocator = sharedAllocator.newChildAllocator( "datafusion-get-import-staging", 0, Long.MAX_VALUE ); NativeBridgeExecutor(DataFusionPlugin dfPlugin) { this.dfPlugin = dfPlugin; } @Override public void close() { - importStagingAllocator.close(); - sharedAllocator.close(); + try { + importStagingAllocator.close(); + } finally { + sharedAllocator.close(); + } } Suggestion importance[1-10]: 6 __ Why: Correct observation: if the staging child throws on close, the parent sharedAllocator is leaked. Using try/finally is a reasonable robustness improvement for resource cleanup.	Low
	Guard allocator close against exceptions during shutdown allocator.close() may throw if children are outstanding, which would propagate and skip subsequent shutdown work of callers even though the finally handles the staging allocator locally. Consider catching/logging exceptions from allocator.close() so shutdown proceeds deterministically, mirroring the tolerant handling applied to the staging allocator below. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [158-162] @Override public void close() { // finally, so a leak report from the service allocator cannot skip the staging allocator's release. try { - allocator.close(); + try { + allocator.close(); + } catch (RuntimeException e) { + LOGGER.warn("Analytics search service allocator [{}] failed to close cleanly", allocator.getName(), e); + } } finally { Suggestion importance[1-10]: 4 __ Why: Adding try/catch around allocator.close() improves shutdown robustness, but throwing on close typically indicates a leak that should surface; the change is minor and stylistic.	Low
	Validate non-null import staging allocator The new record component importStagingAllocator has no null-check, but its Javadoc says "Caller-owned: the backend must never close it". A null value would silently propagate to backends and NPE at first import. Add a compact constructor that requires it non-null (the convenience constructor already forwards allocator, which is validated by usage). sandbox/plugins/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java [58-59] public record ExchangeSinkContext(String queryId, int stageId, long taskId, byte[] fragmentBytes, BufferAllocator allocator, List< ChildInput> childInputs, ExchangeSink downstream, BufferAllocator importStagingAllocator) implements CommonExecutionContext { + public ExchangeSinkContext { + java.util.Objects.requireNonNull(importStagingAllocator, "importStagingAllocator"); + } + Suggestion importance[1-10]: 4 __ Why: A compact constructor null-check would fail fast rather than NPE later, but this is a defensive validation suggestion with limited impact given callers control the input.	Low

Suggestions up to commit 6aa42f5

Category	Suggestion	Impact
General	Avoid leaking allocator on first-import failure If importOntoStaging throws on the very first import, the newly created stagingAllocator field is left non-null and still open, but no batch was ever produced; a subsequent import will still reuse it (fine), yet if the stream is closed after only a failed first import the allocator holds only bookkeeping and closes cleanly — however, if newChildAllocator succeeded but a later import fails and the field was just created for this call, tracking it is still correct. The real concern: on repeated failed imports the allocator persists correctly, but consider ensuring stagingAllocator creation is only "committed" after at least one successful import to avoid holding a child allocator across a failure path where the stream may never emit a batch. Alternatively, document that this is intentional. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [183-188] private VectorSchemaRoot importBatch(ArrowArray arrowArray) { if (stagingAllocator == null) { - stagingAllocator = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); + BufferAllocator candidate = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); + try { + VectorSchemaRoot root = importOntoStaging(candidate, schema, arrowArray, dictionaryProvider); + stagingAllocator = candidate; + return root; + } catch (RuntimeException e) { + candidate.close(); + throw e; + } } return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider); } Suggestion importance[1-10]: 3 __ Why: The suggestion is a minor defensive improvement. Even in the current code, a failed first import leaves an empty allocator that will either be reused by a subsequent import or cleanly closed on stream close (since importOntoStaging closes the partially-imported root, draining the allocator). The impact is marginal.	Low

Suggestions up to commit 70f32a2

Category	Suggestion	Impact
General	Reclaim empty staging allocator on import failure If importOntoStaging throws on the very first batch, stagingAllocator is left non-null but the caller may never retry, and on closeStagingAllocator() a zero-balance close is fine — but if a subsequent import succeeds, the failure-path partial buffers were already released by importOntoStaging's catch. However, if the first import fails and the stream is closed without further imports, this is fine. The real risk: on import failure, stagingAllocator remains set so retries reuse it — ensure that's intentional (it is, per javadoc), but consider nulling it out and closing if empty to avoid a lingering empty allocator when no retry occurs. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [183-188] private VectorSchemaRoot importBatch(ArrowArray arrowArray) { if (stagingAllocator == null) { stagingAllocator = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); } - return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider); + try { + return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider); + } catch (RuntimeException e) { + if (stagingAllocator.getAllocatedMemory() == 0) { + stagingAllocator.close(); + stagingAllocator = null; + } + throw e; + } } Suggestion importance[1-10]: 3 __ Why: The suggestion is a minor cleanup: closing an empty staging allocator on import failure is not strictly necessary since closeStagingAllocator() already handles it at stream close, and the PR javadoc explicitly documents the intended behavior of keeping the allocator open across failures.	Low

Suggestions up to commit 83ec94b

Category	Suggestion	Impact
General	Preserve original exception on cleanup failure If root.close() itself throws in the catch block (e.g. because the partial import left the root in an inconsistent reference-count state), the original import exception will be suppressed and lost. Use addSuppressed to preserve the original failure cause for diagnostics. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java [173-187] private VectorSchemaRoot importBatch() { if (stagingAllocator == null) { stagingAllocator = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); } VectorSchemaRoot root = VectorSchemaRoot.create(schema, stagingAllocator); try { Data.importIntoVectorSchemaRoot(stagingAllocator, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { - // Close the partially-imported root (fires the native release) but KEEP the allocator: it is - // stream-scoped and a later batch may still use it. - root.close(); + try { + root.close(); + } catch (RuntimeException closeEx) { + e.addSuppressed(closeEx); + } throw e; } return root; } Suggestion importance[1-10]: 5 __ Why: Using addSuppressed to preserve the original exception if root.close() throws is a valid defensive improvement for diagnostics, though the scenario is uncommon.	Low
	Log when deferring allocator close When the allocator has a non-zero balance, it is silently leaked to the root without any signal. Log a warning in that branch so operators can diagnose transport-side leaks; the current implementation makes leaks completely invisible. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [198-206] void closeStagingAllocator() { if (stagingAllocator == null) { return; } if (stagingAllocator.getAllocatedMemory() == 0) { stagingAllocator.close(); stagingAllocator = null; + } else { + logger.warn("Staging allocator has {} bytes outstanding at stream close; deferring to root allocator", + stagingAllocator.getAllocatedMemory()); } } Suggestion importance[1-10]: 4 __ Why: Adding a warning log for the non-zero balance branch is a reasonable observability improvement, but it's minor and assumes a logger exists in the class which isn't shown in the diff.	Low

[github-actions]

✅ Gradle check result for 70f32a2: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.58%. Comparing base (e8c8c09) to head (f278551).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22800      +/-   ##
============================================
- Coverage     71.66%   71.58%   -0.08%     
+ Complexity    77365    77330      -35     
============================================
  Files          6170     6170              
  Lines        359698   359710      +12     
  Branches      52458    52460       +2     
============================================
- Hits         257782   257510     -272     
- Misses        81509    81753     +244     
- Partials      20407    20447      +40     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 6aa42f5

[github-actions]

❌ Gradle check result for 6aa42f5: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

✅ Gradle check result for 6aa42f5: SUCCESS

[github-actions]

Persistent review updated to latest commit 1deaf21

[github-actions]

✅ Gradle check result for 1deaf21: SUCCESS

[github-actions]

Persistent review updated to latest commit 5ef81a4

[github-actions]

❌ Gradle check result for 5ef81a4: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit f278551

[github-actions]

✅ Gradle check result for f278551: SUCCESS

[mch2]

thanks @LantaoJin this lgtm, once apache/arrow-java#1240 is deployed in the next arrow-java version we can likely revert to a per-query allocator, and avoid leaks on import ooms.