Stream Arrow batches on the data-node fragment execution path

[bowenlan-amzn]

Description

Today's path drains DataFusion results into Object[] rows, sends one buffered response, and the coordinator converts back to Arrow. This PR replaces that with native Arrow batches over the stream transport from #21253 #21437 #21454.

When Arrow Flight stream transport is enabled, the data-node handler iterates EngineResultStream batches and sends each as an ArrowBatchResponse over the streaming channel. The coordinator receives batches directly as VectorSchemaRoot and feeds them into the stage's ExchangeSink. The row path is preserved as fallback when stream transport is not available.

Data flow

flowchart TD
    subgraph Coordinator
        HANDLER2[StreamingResponseListener]
        RESP[FragmentExecutionArrowResponse<br>extends ArrowBatchResponse]
        SFSE[ShardFragmentStageExecution]
        SINK[ExchangeSink]
    end

    subgraph Data Node
        HANDLER[StreamingFragmentHandler]
        SVC[AnalyticsSearchService]
        FR[FragmentResources<br>reader + engine + stream]
        DRS[DatafusionResultStream<br>fresh VSR per batch]
    end

    SFSE -- "dispatchFragmentStreaming<br>(STREAM type)" --> HANDLER
    HANDLER -- "executeFragmentStreaming" --> SVC
    SVC -- "startFragment" --> FR
    FR -. "iterate batches" .-> DRS
    DRS -- "FragmentExecutionArrowResponse<br>per batch" --> HANDLER
    HANDLER -- "sendResponseBatch" --> HANDLER2
    HANDLER2 -- "response.getRoot()" --> RESP
    RESP -- "VectorSchemaRoot" --> SINK

Loading

Key types

flowchart TD
    subgraph coordinator [Coordinator Node]
        direction TB
        DPE[DefaultPlanExecutor]
        QC[QueryContext<br>bufferAllocator per query]
        PW[PlanWalker]
        SFSE2[ShardFragmentStageExecution<br>extends AbstractStageExecution<br>implements DataProducer]
        LSE[LocalStageExecution<br>implements SinkProvidingStageExecution]
        DRS2[DatafusionReduceSink<br>implements ExchangeSink]
        ATXS[AnalyticsSearchTransportService<br>dispatchFragmentStreaming]
    end

    subgraph datanode [Data Node]
        direction TB
        ATXS2[AnalyticsSearchTransportService<br>handler]
        ASS[AnalyticsSearchService<br>service-level allocator]
        FR2[FragmentResources<br>reader + engine + stream]
        DFSE[DatafusionSearchExecEngine<br>implements SearchExecEngine]
        DFRS[DatafusionResultStream<br>fresh VSR per batch]
    end

    subgraph shared [Shared Root]
        AAP[ArrowAllocatorProvider<br>node-level RootAllocator]
    end

    DPE --> QC
    DPE --> PW
    PW --> SFSE2
    PW --> LSE
    SFSE2 --> ATXS
    LSE --> DRS2
    DRS2 -.-> QC

    ATXS --> ATXS2
    ATXS2 --> ASS
    ASS --> FR2
    FR2 --> DFSE
    DFSE --> DFRS

    AAP -.-> QC
    AAP -.-> ASS

Loading

Send side (data node)

• Fresh VSR per batch. sendResponseBatch is async — Flight's executor transfers buffers on a separate thread while the producer advances to the next batch. Reusing one VSR would race.
• Producer does not close sent VSRs. Flight closes both the producer VSR (after transfer into the channel's staging root) and the channel VSR (after gRPC write). Closing in the producer races the flight thread.
• Service-level allocator. Because sends are async, batches outlive the query. A per-query allocator would close while Flight still holds ArrowBuf references. AnalyticsSearchService owns one allocator for its lifetime, injected into engines via ExecutionContext.

Receive side (coordinator)

• Per-query allocator from shared root. QueryContext.bufferAllocator() lazily creates a child allocator (256 MB limit) used by ExchangeSink to hold received batches during the reduce. Closed by DefaultPlanExecutor's terminal listener when the query completes or fails.
• Shared allocator root across plugins (ArrowAllocatorProvider). Arrow's associate check compares root identity (reference equality). Multi-shard streaming failed at DatafusionReduceSink.feed with "A buffer can only be associated between two allocators that share the same root" because FlightTransport and QueryContext had separate RootAllocator instances. Both now take children of a single node-level root exposed from arrow-flight-rpc.

Cross-plugin wiring

• extendedPlugins = ['arrow-flight-rpc'] on analytics-engine so cross-plugin Arrow types share a classloader. Same FQN loaded by different classloaders = different Java classes; instanceof and casts fail silently. Overlapping jars (arrow-*, jackson, guava, slf4j, flatbuffers) moved to compileOnly.

TransportService.java change

Commit a9606ffb forwards skipsDeserialization() in the sendRequestAsync anonymous wrapper. PR #21454 threaded this marker through all named wrappers but missed this one. Without it, any ArrowBatchResponseHandler dispatched under a parent task receives a byte-serialized input and throws at ArrowBatchResponse.<init>.

Compatibility

With transport.stream.enabled=false (default), behavior is unchanged. CoordinatorReduceIT and CoordinatorReduceMemtableIT validate the row path. StreamingCoordinatorReduceIT (@LockFeatureFlag(STREAM_TRANSPORT), 2-shard parquet, 20 rows) is the streaming regression gate.

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

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 Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit 1b7aed7.

Path	Line	Severity	Description
sandbox/plugins/parquet-data-format/build.gradle	23	high	New dependency added: org.apache.arrow:arrow-memory-netty. Artifact authenticity cannot be verified from diff alone; maintainers must confirm expected coordinates and SHA1.
sandbox/plugins/parquet-data-format/build.gradle	24	high	New dependency added: org.apache.arrow:arrow-memory-netty-buffer-patch. Artifact authenticity cannot be verified from diff alone.
sandbox/plugins/parquet-data-format/build.gradle	25	high	New dependency added: io.netty:netty-buffer. Introduces Netty native memory access into the plugin classpath; artifact authenticity must be verified.
sandbox/plugins/parquet-data-format/build.gradle	26	high	New dependency added: io.netty:netty-common. Paired with netty-buffer addition; artifact authenticity must be verified.
sandbox/plugins/analytics-engine/build.gradle	24	high	New build plugin configuration: extendedPlugins = ['arrow-flight-rpc']. Extends the plugin classloader chain; this is a build plugin/dependency change that must be verified by maintainers.
sandbox/plugins/analytics-engine/build.gradle	175	high	Guava version downgraded from 33.4.0-jre to 33.3.1-jre. Version downgrades can reintroduce known CVEs; maintainers must confirm this is intentional and safe.
sandbox/plugins/analytics-engine/build.gradle	176	high	failureaccess version downgraded from 1.0.2 to 1.0.1. Dependency version change; maintainers must verify no regressions or CVEs reintroduced.
sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java	115	medium	assert used for null-safety check on channelAllocator in production code path. Java assertions are disabled by default in production JVMs; a null allocator will silently pass this check and cause an NPE or memory corruption later in the streaming path.
sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java	125	low	StreamException with CANCELLED error code is silently swallowed with no logging. While labeled intentional, this suppresses all visibility into stream cancellations, which could mask unexpected terminations if the error code mapping is incorrect.

The table above displays the top 10 most important findings.

Total: 9 | Critical: 0 | High: 7 | Medium: 1 | Low: 1

---

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

---

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 089de3e)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

📝 TODO sections sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [132]: Re-wire once operatorTable() is added back to AnalyticsSearchBackendPlugin

🔀 Multiple PR themes Sub-PR theme: Add Netty and Arrow-memory-netty license files for parquet-data-format plugin Relevant files: sandbox/plugins/parquet-data-format/licenses/netty-LICENSE.txt sandbox/plugins/parquet-data-format/licenses/netty-NOTICE.txt sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-buffer-patch-NOTICE.txt sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-NOTICE.txt sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-buffer-patch-18.1.0.jar.sha1 sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-18.1.0.jar.sha1 sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-buffer-patch-LICENSE.txt sandbox/plugins/parquet-data-format/licenses/arrow-memory-netty-LICENSE.txt Sub-PR theme: Refactor allocator ownership and per-batch VSR lifecycle in DataFusion backend Relevant files: sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExecutionContext.java sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/EngineResultBatch.java server/src/main/java/org/opensearch/search/SearchExecutionContext.java sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java Sub-PR theme: Stream Arrow batches on data-node fragment execution path (coordinator + transport) Relevant files: sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionArrowResponse.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java server/src/main/java/org/opensearch/transport/TransportService.java sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/StreamingCoordinatorReduceIT.java sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java

⚡ Recommended focus areas for review Batch Leak on Error In registerStreamingFragmentHandler, when iterating batches and calling channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())), if sendResponseBatch throws an exception (other than StreamException), the current batch obtained from it.next() is never closed before channel.sendResponse(e) is called. The FragmentResources try-with-resources will close the stream, but closeLastBatch() only closes nextBatch (the pre-fetched one), not the batch already handed out via it.next(). This could cause a buffer leak on the data node. try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); } channel.completeStream(); } catch (StreamException e) { if (e.getErrorCode() != StreamErrorCode.CANCELLED) { channel.sendResponse(e); } // CANCELLED: channel already torn down — exit silently } catch (Exception e) { channel.sendResponse(e); } } Listener Not Called on Streaming Success executeFragmentStreaming starts the fragment and returns FragmentResources but never calls listener.onFragmentSuccess(...). The success listener is only called in executeFragment. This means streaming executions will not report success metrics/events, which may affect observability and monitoring. public FragmentResources executeFragmentStreaming(FragmentExecutionRequest request, IndexShard shard, AnalyticsShardTask task) { ResolvedFragment resolved = resolveFragment(request, shard); try { return startFragment(request, resolved, task); } catch (TaskCancelledException | IllegalStateException | IllegalArgumentException e) { listener.onFragmentFailure(resolved.queryId, resolved.stageId, resolved.shardIdStr, e); throw e; } catch (Exception e) { listener.onFragmentFailure(resolved.queryId, resolved.stageId, resolved.shardIdStr, e); throw new RuntimeException("Failed to start streaming fragment on " + shard.shardId(), e); } } Missing VSR Close on Codec Path In responseListener, when isDone() is true and the response is not an ArrowBatchResponse (i.e., the row codec path), releaseResponseResources does nothing because the VectorSchemaRoot is created inside the toVsr lambda (via responseCodec.decode). However, toVsr.apply(response) is never called in the early-return path, so the codec-decoded VSR is never created — this is actually safe. But if toVsr.apply(response) throws after partial allocation, there is no cleanup. Consider wrapping the toVsr.apply call in a try-catch that closes any partially-created VSR. public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { if (isDone()) { releaseResponseResources(response); return; } VectorSchemaRoot vsr = toVsr.apply(response); outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); Schema Not Reset on Close In BatchIterator, the schema field is set once in ensureSchema() and never cleared. If close() is called and then the iterator is somehow reused (or if closeLastBatch() is called), the schema reference remains. While this is likely not a practical issue given the one-shot usage, the schema field holds references to Arrow Field objects that could prevent GC. More importantly, closeLastBatch() does not null out schema, which is a minor inconsistency with the cleanup intent. void closeLastBatch() { // Only close batches that were loaded but never handed to the caller. Caller // owns any batch returned by next(); closing it here would double-close after // Flight's transferTo or after row-path reads. if (nextBatch != null) { nextBatch.close(); nextBatch = null; } } Flaky Doc Count indexDeterministicDocs indexes NUM_SHARDS * DOCS_PER_SHARD documents with sequential IDs but does not guarantee even distribution across shards. The assertion assertEquals("all docs across shards must be returned", expectedRows, response.getRows().size()) assumes exactly 20 rows, but with default routing, documents may not be evenly distributed. The test could be flaky if routing sends all docs to one shard and the query returns fewer rows due to shard-level limits or if the index has different actual shard counts. public void testBaselineScanAcrossShards() throws Exception { createParquetBackedIndex(); indexDeterministicDocs(); PPLResponse response = executePPL("source = " + INDEX); assertNotNull("PPLResponse must not be null", response); assertTrue("columns must contain 'value', got " + response.getColumns(), response.getColumns().contains("value")); int expectedRows = NUM_SHARDS * DOCS_PER_SHARD; assertEquals("all docs across shards must be returned", expectedRows, response.getRows().size()); int idx = response.getColumns().indexOf("value"); for (Object[] row : response.getRows()) { Object cell = row[idx]; assertNotNull("value cell must not be null", cell); assertEquals("every doc has value=" + VALUE, (long) VALUE, ((Number) cell).longValue()); } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 089de3e

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Close Arrow batch after streaming send to prevent memory leak Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot (caller-owned lifecycle per the new design). After sendResponseBatch serializes the batch to the wire, the VSR is never closed in the happy path, causing a memory leak. The batch should be closed after each send, ideally in a finally block or try-with-resources. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-122] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } channel.completeStream(); Suggestion importance[1-10]: 8 __ Why: The PR's new design makes each EngineResultBatch caller-owned, so after sendResponseBatch the VSR is never closed in the happy path, causing a real memory leak. The fix is accurate and directly addresses the ownership model introduced in this PR.	Medium
	Close freshly created VSR on import failure to prevent leak If Data.importIntoVectorSchemaRoot throws an exception, freshRoot is created but never closed, leaking the Arrow buffers already allocated for it. The freshRoot should be closed in a catch block before re-throwing. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-120] VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); +} catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; Suggestion importance[1-10]: 7 __ Why: If Data.importIntoVectorSchemaRoot throws, freshRoot is allocated but never closed, leaking Arrow buffers. The suggested catch block correctly closes freshRoot before re-throwing, which is a valid resource management fix.	Medium
	Prevent VSR leak when sink feed throws an exception If toVsr.apply(response) succeeds but outputSink.feed(vsr) throws an exception, the VectorSchemaRoot will be leaked because it is never closed. The VSR should be closed in a finally block if feeding fails, or the sink contract should guarantee ownership transfer only on success. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [125-127] VectorSchemaRoot vsr = toVsr.apply(response); -outputSink.feed(vsr); +try { + outputSink.feed(vsr); +} catch (Exception ex) { + vsr.close(); + throw ex; +} metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 6 __ Why: If outputSink.feed(vsr) throws, the VectorSchemaRoot is leaked. The suggested fix correctly wraps the feed call in a try-catch to close the VSR on failure, though the sink contract may implicitly take ownership on success.	Low
General	Enforce non-null contract on required constructor parameter The constructor allows gatedReader to be null (it is passed as null in the error-path cleanup inside startFragment), but closeQuietly handles nulls safely. However, if gatedReader is non-null but engine and stream are both null (partial construction failure), the close order stream → engine → gatedReader is correct. The real risk is that a non-null gatedReader passed to the error-path new FragmentResources(gatedReader, null, null).close() in startFragment will close the reader, which is the intended behavior — this is fine. No change needed here, but consider adding a null-check assertion or documentation that gatedReader must not be null for a valid instance to make the contract explicit. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java [29-37] public FragmentResources( GatedCloseable<Reader> gatedReader, SearchExecEngine<ExecutionContext, EngineResultStream> engine, EngineResultStream stream ) { - this.gatedReader = gatedReader; + // gatedReader must not be null for a valid instance; engine and stream may be null during partial construction cleanup + this.gatedReader = Objects.requireNonNull(gatedReader, "gatedReader must not be null"); this.engine = engine; this.stream = stream; } Suggestion importance[1-10]: 2 __ Why: The suggestion itself acknowledges that closeQuietly handles nulls safely and the current behavior is correct. Adding a requireNonNull would actually break the error-path cleanup in startFragment which intentionally passes null for engine and stream, making this suggestion potentially harmful rather than helpful.	Low

---

Previous suggestions

Suggestions up to commit 2b5a112

Category	Suggestion	Impact
Possible issue	Close Arrow batch after sending to prevent memory leak Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot and the caller is responsible for closing it. In the streaming handler the batch is never explicitly closed after sendResponseBatch, so Arrow memory is leaked for every batch sent. The batch should be closed after the response is sent (or in a finally block). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-121] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } Suggestion importance[1-10]: 8 __ Why: The PR's new ownership model makes each EngineResultBatch caller-responsible for closing its VectorSchemaRoot. In the streaming handler, batch.getArrowRoot() is passed to sendResponseBatch but never closed, leaking off-heap Arrow memory for every batch sent. The fix is accurate and addresses a real resource leak.	Medium
	Release Arrow batch when stage is already done When isDone() returns true the method returns early without calling toVsr, so the response (and its underlying VectorSchemaRoot) is never closed. For the Arrow streaming path this leaks the batch's off-heap memory. The response should be closed even when the stage is already done. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [118-125] public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { - if (isDone()) return; + if (isDone()) { + response.close(); + return; + } VectorSchemaRoot vsr = toVsr.apply(response); outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 7 __ Why: When isDone() is true, the method returns early without closing the response, leaking the underlying VectorSchemaRoot off-heap memory in the Arrow streaming path. The suggestion correctly identifies and fixes this resource leak, though ActionResponse doesn't implement AutoCloseable so the response.close() call may not compile without a cast or interface check.	Medium
	Close freshly allocated VSR on import failure If Data.importIntoVectorSchemaRoot throws, freshRoot has already been allocated but is never closed, leaking Arrow memory. The freshRoot should be closed in a catch block on failure before re-throwing. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-121] private boolean loadNextBatch() { - ... + long arrayAddr = callNativeFn( + listener -> NativeBridge.streamNext(streamHandle.getRuntimeHandle().get(), streamHandle.getPointer(), listener) + ); if (arrayAddr == 0) return false; VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + } catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; return true; } Suggestion importance[1-10]: 7 __ Why: If Data.importIntoVectorSchemaRoot throws, freshRoot is allocated but never closed, leaking Arrow off-heap memory. The suggested catch block correctly closes freshRoot before re-throwing, preventing the leak.	Medium
	Guard against null backend causing resource leak When resolved.plan.getBackendId() does not match any registered backend, backends.get(...) returns null, causing a NullPointerException that bypasses the cleanup path's suppressed-exception handling and leaks gatedReader. Add an explicit null-check for the backend before using it. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [108-126] private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, Task task) throws IOException { GatedCloseable<Reader> gatedReader = resolved.readerProvider.acquireReader(); SearchExecEngine<ExecutionContext, EngineResultStream> engine = null; EngineResultStream stream = null; try { ExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, task); AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); + if (backend == null) { + throw new IllegalStateException("No backend registered for id: " + resolved.plan.getBackendId()); + } engine = backend.getSearchExecEngineProvider().createSearchExecEngine(ctx); stream = engine.execute(ctx); return new FragmentResources(gatedReader, engine, stream); } catch (Exception e) { try { new FragmentResources(gatedReader, engine, stream).close(); } catch (Exception suppressed) { e.addSuppressed(suppressed); } throw e; } } Suggestion importance[1-10]: 6 __ Why: If backends.get(resolved.plan.getBackendId()) returns null, a NullPointerException is thrown before the cleanup path can run, potentially leaking gatedReader. The suggestion is valid and improves robustness, though this is a defensive check for a configuration error rather than a critical runtime bug.	Low

Suggestions up to commit 53e284b

Category	Suggestion	Impact
Possible issue	Close Arrow batch after streaming send Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot and the caller is responsible for closing it. In the streaming handler, after channel.sendResponseBatch(...) the batch's VSR is never closed, leaking Arrow memory. The batch should be closed after the send (or in a try-finally block) to release its buffers. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-121] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } Suggestion importance[1-10]: 8 __ Why: The PR comment in DatafusionResultStream.BatchIterator.next() states "Caller owns the returned VSR's lifecycle. Streaming handler transfers it to Flight (Flight closes after wire write)". However, if channel.sendResponseBatch does not guarantee closing the VSR, the batch's VectorSchemaRoot would leak. This is a valid memory management concern that could cause Arrow buffer leaks in production.	Medium
	Close VSR on import failure to prevent leak If Data.importIntoVectorSchemaRoot(...) throws an exception, freshRoot is created but never closed, leaking Arrow memory. The freshRoot should be closed in a catch/finally block if the import fails before it is assigned to nextBatch. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-121] private boolean loadNextBatch() { - ... + long arrayAddr = callNativeFn( + listener -> NativeBridge.streamNext(streamHandle.getRuntimeHandle().get(), streamHandle.getPointer(), listener) + ); if (arrayAddr == 0) return false; VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + } catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; return true; } Suggestion importance[1-10]: 7 __ Why: If Data.importIntoVectorSchemaRoot throws, freshRoot is created but never closed, causing an Arrow memory leak. Adding a catch block to close freshRoot on failure is a correct and important fix for resource safety.	Medium
	Prevent VSR leak on sink feed failure When isDone() returns true early, the VectorSchemaRoot obtained from toVsr.apply(response) is never created, but the response itself (which may own Arrow buffers via FragmentExecutionArrowResponse) is never released. Additionally, if toVsr.apply(response) succeeds but outputSink.feed(vsr) throws, the VSR leaks. The response/VSR should be closed in a finally block or the early-exit path should also release the response. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [118-125] public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { if (isDone()) return; VectorSchemaRoot vsr = toVsr.apply(response); - outputSink.feed(vsr); + try { + outputSink.feed(vsr); + } catch (Exception e) { + vsr.close(); + throw e; + } metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 6 __ Why: If outputSink.feed(vsr) throws an exception, the VectorSchemaRoot would leak. The suggestion to wrap in try-catch is valid, though the impact depends on how often feed can throw and whether the sink takes ownership of the VSR.	Low
General	Enforce memory limit on service allocator The allocator limit is set to Long.MAX_VALUE, which means there is no effective memory cap for the analytics search service. This could allow unbounded memory allocation. A more reasonable limit (e.g., a configurable value or a sensible default) should be used to prevent out-of-memory conditions. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [63-66] public AnalyticsSearchService(Map<String, AnalyticsSearchBackendPlugin> backends, List<AnalyticsOperationListener> listeners) { this.backends = backends; this.listener = new AnalyticsOperationListener.CompositeListener(listeners); - this.allocator = ArrowAllocatorProvider.newChildAllocator("analytics-search-service", Long.MAX_VALUE); + this.allocator = ArrowAllocatorProvider.newChildAllocator("analytics-search-service", DEFAULT_SERVICE_MEMORY_LIMIT); } Suggestion importance[1-10]: 4 __ Why: Using Long.MAX_VALUE as the allocator limit is a common pattern in Arrow for "no effective cap at this level" when the root allocator already enforces a global limit. The suggestion to use a configurable limit is reasonable but may be a design choice rather than a bug.	Low

Suggestions up to commit e6ee337

Category	Suggestion	Impact
Possible issue	Close Arrow batch after sending to prevent memory leak Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot and the caller is responsible for closing it. In the streaming handler, after channel.sendResponseBatch(...) the batch's VSR is never closed, leaking Arrow memory. The batch should be closed after it has been sent (or transferred to the channel). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-121] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } Suggestion importance[1-10]: 7 __ Why: The PR comment in DatafusionResultStream.BatchIterator.next() explicitly states "Caller owns the returned VSR's lifecycle. Streaming handler transfers it to Flight (Flight closes after wire write)". If channel.sendResponseBatch transfers ownership to Flight which closes it, this suggestion may be incorrect. However, if Flight does NOT close it, this is a real memory leak. The score reflects the uncertainty about whether Flight's sendResponseBatch takes ownership.	Medium
	Close freshly created VSR on import failure If Data.importIntoVectorSchemaRoot throws an exception, freshRoot has already been created but is never closed, leaking the Arrow buffers allocated for it. The freshRoot should be closed in a catch block if the import fails. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-121] private boolean loadNextBatch() { ... if (arrayAddr == 0) return false; VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + } catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; return true; } Suggestion importance[1-10]: 7 __ Why: This is a legitimate resource leak: if Data.importIntoVectorSchemaRoot throws, freshRoot is allocated but never closed. The fix correctly adds a catch block to close freshRoot before re-throwing, preventing Arrow buffer leaks on import failures.	Medium
	Release Arrow resources when stage is already done When isDone() is true, the method returns early without calling toVsr.apply(response), but the response object (which may own an Arrow VectorSchemaRoot in the Arrow streaming path) is never closed. This leaks Arrow memory for every batch received after the stage is done (e.g., after cancellation). The response should be closed when the stage is already done. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [118-125] public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { - if (isDone()) return; + if (isDone()) { + if (response instanceof AutoCloseable ac) { + try { ac.close(); } catch (Exception ignored) {} + } + return; + } VectorSchemaRoot vsr = toVsr.apply(response); outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 6 __ Why: This is a valid memory leak concern — when isDone() returns true after cancellation, Arrow VectorSchemaRoot batches in the streaming path are never closed. The instanceof AutoCloseable pattern is a reasonable approach, though it could be more explicit about the FragmentExecutionArrowResponse type.	Low
General	Guard against null backend before acquiring resources The buildContext method sets allocator from the service-level shared allocator, but createSearchExecEngine(ctx) is called with the same ctx that is then passed again to engine.execute(ctx). If createSearchExecEngine internally calls prepare(ctx) (which it does per DataFusionAnalyticsBackendPlugin), the allocator is already set. However, if backends.get(resolved.plan.getBackendId()) returns null (unknown backend ID), a NullPointerException will be thrown without the gatedReader being closed via the catch block. The catch block does handle this, but it's worth verifying the null backend case is explicitly handled with a clear error message before resource acquisition. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [108-126] private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, Task task) throws IOException { + AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); + if (backend == null) { + throw new IllegalArgumentException("Unknown backend: " + resolved.plan.getBackendId()); + } GatedCloseable<Reader> gatedReader = resolved.readerProvider.acquireReader(); SearchExecEngine<ExecutionContext, EngineResultStream> engine = null; EngineResultStream stream = null; try { ExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, task); - AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); engine = backend.getSearchExecEngineProvider().createSearchExecEngine(ctx); stream = engine.execute(ctx); return new FragmentResources(gatedReader, engine, stream); } catch (Exception e) { try { new FragmentResources(gatedReader, engine, stream).close(); } catch (Exception suppressed) { e.addSuppressed(suppressed); } throw e; } } Suggestion importance[1-10]: 4 __ Why: While the null backend check is a valid defensive improvement, the existing catch block already handles the NullPointerException by closing gatedReader via FragmentResources. The suggestion improves error clarity but doesn't fix a critical resource leak.	Low

Suggestions up to commit e231cb1

Category	Suggestion	Impact
Possible issue	Close Arrow batch after streaming send Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot and the caller is responsible for closing it. In the streaming handler, after sendResponseBatch the batch is never closed, leaking Arrow memory. Each batch should be closed after it has been sent (or transferred to the transport layer). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-121] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } Suggestion importance[1-10]: 7 __ Why: The PR comment on ArrowResultBatch states "Caller owns the returned VSR's lifecycle. Streaming handler transfers it to Flight (Flight closes after wire write)". If sendResponseBatch transfers ownership to Flight which closes it, there's no leak. However, if Flight does NOT close the VSR after writing, this is a real memory leak. The suggestion is valid as a defensive measure but depends on sendResponseBatch semantics.	Medium
	Close VSR on import failure to prevent leak If Data.importIntoVectorSchemaRoot throws an exception, freshRoot is created but never closed, leaking Arrow memory. The freshRoot should be closed in a catch block if the import fails. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-121] private boolean loadNextBatch() { ... if (arrayAddr == 0) return false; VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + } catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; return true; } Suggestion importance[1-10]: 7 __ Why: If Data.importIntoVectorSchemaRoot throws, freshRoot is created but never closed, causing an Arrow memory leak. Adding a catch block to close freshRoot on failure is a correct and important fix for resource safety.	Medium
	Release response resources on early exit When isDone() returns true early, the VectorSchemaRoot obtained from toVsr.apply(response) is never created, but the response itself (which may own an Arrow VSR in the Arrow streaming path) is never closed. This leaks memory for every response received after the stage is done. The response should be closed in the early-exit branch. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [118-124] public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { - if (isDone()) return; + if (isDone()) { + if (response instanceof AutoCloseable ac) { + try { ac.close(); } catch (Exception ignored) {} + } + return; + } VectorSchemaRoot vsr = toVsr.apply(response); outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 6 __ Why: When isDone() returns true early, the response (which may own an Arrow VectorSchemaRoot in the streaming path) is never closed, potentially leaking Arrow memory. This is a valid concern for the Arrow streaming path where FragmentExecutionArrowResponse holds a VSR.	Low

Suggestions up to commit 2d2d2c1

Category	Suggestion	Impact
Possible issue	Close Arrow batch after streaming send Each EngineResultBatch returned by it.next() now owns its VectorSchemaRoot and the caller is responsible for closing it. In the streaming handler, batch.getArrowRoot() is passed to FragmentExecutionArrowResponse but the batch itself is never closed, leaking the Arrow buffers. Each batch should be closed after it has been sent (or after the send fails). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [116-122] try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { EngineResultBatch batch = it.next(); - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } finally { + batch.getArrowRoot().close(); + } } channel.completeStream(); Suggestion importance[1-10]: 8 __ Why: The PR introduces per-batch VectorSchemaRoot ownership where each batch returned by it.next() must be closed by the caller. In the streaming handler, batch.getArrowRoot() is passed to FragmentExecutionArrowResponse but the batch is never closed, causing Arrow buffer leaks. This is a real resource management bug introduced by the PR's ownership model change.	Medium
	Prevent Arrow buffer leak on import failure If Data.importIntoVectorSchemaRoot throws an exception, freshRoot is created but never closed, leaking Arrow memory. The freshRoot should be closed in a catch block (or try-with-resources wrapping the import) before re-throwing. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [116-121] private boolean loadNextBatch() { - ... + long arrayAddr = callNativeFn( + listener -> NativeBridge.streamNext(streamHandle.getRuntimeHandle().get(), streamHandle.getPointer(), listener) + ); if (arrayAddr == 0) return false; VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + } catch (Exception e) { + freshRoot.close(); + throw e; } nextBatch = freshRoot; return true; } Suggestion importance[1-10]: 7 __ Why: If Data.importIntoVectorSchemaRoot throws, freshRoot is created but never closed, leaking Arrow memory. The PR changed to a fresh-VSR-per-batch model, making this leak path newly relevant and worth fixing with a catch block.	Medium
	Release Arrow resources on early stream exit When isDone() returns true early, the response (which may carry an Arrow VectorSchemaRoot for the streaming path) is silently dropped without being closed, leaking Arrow memory. The early-exit branch should close the response's underlying resources before returning. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java [118-125] public void onStreamResponse(T response, boolean isLast) { config.searchExecutor().execute(() -> { - if (isDone()) return; + if (isDone()) { + // Release Arrow resources even when we short-circuit. + VectorSchemaRoot vsr = toVsr.apply(response); + if (vsr != null) vsr.close(); + return; + } VectorSchemaRoot vsr = toVsr.apply(response); outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); Suggestion importance[1-10]: 6 __ Why: When isDone() returns true, the streaming response carrying an Arrow VectorSchemaRoot is dropped without being closed, leaking memory. This is a valid concern given the PR's new per-batch ownership model, though the toVsr.apply(response) call itself may have side effects that complicate the fix.	Low

[github-actions]

Persistent review updated to latest commit f5bde4e

[github-actions]

Persistent review updated to latest commit 5347011

[github-actions]

Persistent review updated to latest commit f748228

[github-actions]

❌ Gradle check result for f748228: 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 96923e3

[github-actions]

❌ Gradle check result for 96923e3: 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 0f8c8b5

[github-actions]

✅ Gradle check result for 0f8c8b5: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.51%. Comparing base (fbfcabe) to head (089de3e).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
...ava/org/opensearch/transport/TransportService.java	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21418      +/-   ##
============================================
- Coverage     73.55%   73.51%   -0.05%     
+ Complexity    74490    74480      -10     
============================================
  Files          5970     5970              
  Lines        338262   338261       -1     
  Branches      48758    48752       -6     
============================================
- Hits         248824   248656     -168     
- Misses        69581    69785     +204     
+ Partials      19857    19820      -37     

☔ View full report in Codecov by Sentry.
📢 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 83c0660

[github-actions]

❌ Gradle check result for 83c0660: 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 f450ad2

[github-actions]

Persistent review updated to latest commit 6ab8fff

[github-actions]

Persistent review updated to latest commit 53e284b

[github-actions]

✅ Gradle check result for 53e284b: SUCCESS

[github-actions]

Persistent review updated to latest commit 2b5a112

[github-actions]

Persistent review updated to latest commit 089de3e

[bowenlan-amzn]

Breaking change check failure is because of #21143 (comment)
Not related to this PR

[github-actions]

✅ Gradle check result for 089de3e: SUCCESS