[Arrow Flight RPC] Add metadata channel on ArrowBatchResponse

[rishabhmaurya]

Summary

Adds an opaque byte[] metadata channel on ArrowBatchResponse so a stream-transport action can attach bytes to a batch (or to its terminal batch) and read them back on the consumer. Wire path is Arrow Flight's putNext(ArrowBuf) / getLatestMetadata(). No changes to the public stream-transport SPI in server/.

This is a self-contained prototype. It's published as a draft so the use case in #21972 (DataFusion operator metrics → coordinator profile output) can borrow from it if helpful.

Design rationale

The stream-transport tenets in #21253 and #21465 are: the transport moves TransportResponse instances; the action defines the response shape. Anything an action needs to send between nodes belongs on a request or response object. The transport stays narrow and is shared across actions.

Putting the metadata bytes on ArrowBatchResponse keeps the abstraction faithful to that:

• The metadata is a property of the response, where the action already lives. The action encodes; the consumer decodes.
• server/transport/stream/StreamTransportResponse doesn't grow new methods — every transport implementation that exists or will exist remains decoupled from this feature.
• FlightTransportChannel doesn't grow new methods either — plugins talk to it through the same sendResponseBatch they already use.
• The wire path is Arrow Flight's intended per-frame metadata mechanism (putNext(ArrowBuf)), not a custom convention layered on top.

The receive-side flow is symmetrical: FlightTransportResponse copies getLatestMetadata() into a byte[] on each frame (Flight retains buffer ownership otherwise), and the response constructor pulls it via ArrowStreamInput.getMetadata().

The mechanism covers both shapes naturally: per-batch metadata (row offsets, watermarks, batch-level stats) and stream-terminal metadata (attach to the last batch — profile counters, summary stats).

Integrating the analytics-engine profile flow (#21972)

The Rust FFM additions (df_stream_get_metrics, df_free_metrics_buf), the profile flag on QueryContext / FragmentExecutionRequest, and the TaskProfile.dataNodeMetrics parsing land as-is. Only the transport hop changes.

Producer (data node, AnalyticsSearchService.executeFragmentStreamingAsync): attach the metrics blob to the last batch's response.

boolean isLast = (batchIdx == lastBatchIdx);
byte[] metrics = (isLast && request.profile()) ? exec.resources().getExecutionMetrics() : null;
channel.sendResponseBatch(metrics != null
    ? new FragmentExecutionArrowResponse(root, metrics)
    : new FragmentExecutionArrowResponse(root));

This drops the onCompleteWithMetrics callback, the reflection-based sendStreamMetadata call in AnalyticsSearchTransportService, and the pendingStreamMetadata field on FlightTransportChannel.

Consumer (coordinator, AnalyticsSearchTransportService.handleStreamResponse): read inside the existing batch loop.

while ((next = stream.nextResponse()) != null) {
    /* existing batch handling */
    byte[] meta = next.getMetadata();
    if (meta != null) listener.onMetrics(meta);
}

This drops the getTrailingMetadata() call after the loop and the corresponding addition to StreamTransportResponse.

Empty-shard case: when no batches are produced, send a single FragmentExecutionArrowResponse(emptyRoot, metricsBytes) — the metadata rides on a real (zero-row but legitimately-shaped) batch, no sentinel convention needed.

Tests

• StreamMetadataIT — profile=true: metadata observed once on the last batch, bytes intact end-to-end. profile=false: zero observations, data path identical.
• FlightOutboundHandlerTests updated for the new sendBatch signature.
• Existing unit tests + NativeArrowTransportIT pass.

Test plan

• ./gradlew :plugins:arrow-base:check :plugins:arrow-flight-rpc:check
• ./gradlew :plugins:arrow-flight-rpc:internalClusterTest --tests "*StreamMetadataIT" --tests "*NativeArrowTransportIT"

Related

• Native Arrow transport: Native Arrow path support in stream transport with zero-copy transfer #21253
• arrow-base extraction: Extract arrow-base plugin #21465
• Motivating use case: [Analytics Engine] Add DF metrics to profile API response #21972

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit b070add)

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 Resource Leak The metadata ArrowBuf is allocated but never closed if an exception occurs between allocation (line 181) and putNext (line 183). If serverStreamListener.putNext throws, the buffer leaks. Wrap the allocation and putNext in a try-catch that closes metadataBuf on failure. if (metadata != null) { // Flight takes ownership of metadataBuf via putNext(ArrowBuf). ArrowBuf metadataBuf = allocator.buffer(metadata.length); metadataBuf.writeBytes(metadata); serverStreamListener.putNext(metadataBuf); } else { serverStreamListener.putNext(); } Mutable Array Exposure getMetadata() returns the internal byte[] directly. Callers can mutate the array, affecting all subsequent readers of the same response instance. Return a defensive copy or document that callers must not modify the array. public byte[] getMetadata() { return metadata; }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to b070add

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent metadata buffer leak on allocation failure If allocator.buffer() throws an exception (e.g., out of memory), metadataBuf will not be closed, causing a memory leak. Wrap the buffer allocation and write operations in a try-catch block to ensure the buffer is released on failure. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [179-186] if (metadata != null) { - // Flight takes ownership of metadataBuf via putNext(ArrowBuf). - ArrowBuf metadataBuf = allocator.buffer(metadata.length); - metadataBuf.writeBytes(metadata); - serverStreamListener.putNext(metadataBuf); + ArrowBuf metadataBuf = null; + try { + metadataBuf = allocator.buffer(metadata.length); + metadataBuf.writeBytes(metadata); + serverStreamListener.putNext(metadataBuf); + } catch (Exception e) { + if (metadataBuf != null) { + metadataBuf.close(); + } + throw e; + } } else { serverStreamListener.putNext(); } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential resource leak if allocator.buffer() throws an exception. However, since putNext(ArrowBuf) takes ownership of the buffer, the leak scenario is limited to the narrow window between allocation and the putNext call. The fix is valid and improves robustness.	Medium
	Validate metadata buffer size before casting Casting buf.readableBytes() (a long) to int can silently truncate large metadata buffers exceeding Integer.MAX_VALUE, leading to incomplete data copies. Add a bounds check to fail fast if the buffer size exceeds the maximum array length. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [167-173] static byte[] copyMetadata(ArrowBuf buf) { if (buf == null || buf.readableBytes() == 0) return null; - int len = (int) buf.readableBytes(); + long readableBytes = buf.readableBytes(); + if (readableBytes > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Metadata buffer too large: " + readableBytes + " bytes"); + } + int len = (int) readableBytes; byte[] copy = new byte[len]; buf.getBytes(0, copy); return copy; } Suggestion importance[1-10]: 6 __ Why: The suggestion identifies a theoretical issue with casting long to int for very large buffers. While the scenario is unlikely in practice (metadata buffers are typically small), adding a bounds check improves defensive programming and prevents silent data truncation.	Low
General	Store defensive copy of metadata array The metadata array is stored directly without defensive copying, allowing external modifications to affect the response's internal state. Consider storing a defensive copy to prevent unintended mutations by the caller after construction. plugins/arrow-base/src/main/java/org/opensearch/arrow/transport/ArrowBatchResponse.java [102-105] protected ArrowBatchResponse(VectorSchemaRoot batchRoot, byte[] metadata) { this.batchRoot = batchRoot; - this.metadata = metadata; + this.metadata = metadata != null ? metadata.clone() : null; } Suggestion importance[1-10]: 5 __ Why: The suggestion recommends defensive copying of the metadata array to prevent external modifications. While this is a valid defensive programming practice, the impact is moderate since the metadata is typically constructed and passed immediately without retention by the caller. The improvement enhances encapsulation but is not critical.	Low

---

Previous suggestions

Suggestions up to commit a50a09e

Category	Suggestion	Impact
General	Ensure stream cleanup in finally block If an exception occurs during the loop, stream.close() is never called before stream.cancel() in the catch block. This can leave the stream in an inconsistent state. Move stream.close() to the finally block to ensure it's always invoked, or remove the redundant stream.close() call since cancel() should handle cleanup. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/StreamMetadataIT.java [189-212] @Override public void handleStreamResponse(StreamTransportResponse<FragmentResponse> stream) { try { FragmentResponse response; while ((response = stream.nextResponse()) != null) { int idx = sink.batchesSeen++; VectorSchemaRoot root = response.getRoot(); sink.dataRowTotal += root.getRowCount(); if (response.getMetadata() != null) { sink.metadataObservations++; sink.metadataBatchIndex = idx; sink.metrics = response.getMetadata(); } retainedRoots.add(root); } - stream.close(); } catch (Exception e) { failure.set(e); stream.cancel("test error", e); } finally { + try { + stream.close(); + } catch (Exception ignored) {} for (VectorSchemaRoot r : retainedRoots) r.close(); latch.countDown(); } } Suggestion importance[1-10]: 8 __ Why: The suggestion correctly identifies that stream.close() is only called in the try block and not when an exception occurs. Moving it to the finally block ensures proper cleanup regardless of success or failure, which is important for resource management in streaming scenarios.	Medium
Possible issue	Handle metadata buffer allocation failures The metadataBuf allocation can fail if the allocator runs out of memory, but the code doesn't handle this scenario. If allocator.buffer() throws an exception after incrementing batchNumber, the batch count becomes inconsistent. Wrap the metadata buffer allocation in a try-catch block to ensure proper error handling and resource cleanup. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [179-186] if (metadata != null) { - // Flight takes ownership of metadataBuf via putNext(ArrowBuf). - ArrowBuf metadataBuf = allocator.buffer(metadata.length); - metadataBuf.writeBytes(metadata); - serverStreamListener.putNext(metadataBuf); + ArrowBuf metadataBuf = null; + try { + metadataBuf = allocator.buffer(metadata.length); + metadataBuf.writeBytes(metadata); + serverStreamListener.putNext(metadataBuf); + } catch (Exception e) { + if (metadataBuf != null) { + metadataBuf.close(); + } + throw e; + } } else { serverStreamListener.putNext(); } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that allocator.buffer() can throw an exception if memory allocation fails. However, Arrow's ArrowBuf allocation typically throws OutOfMemoryException which would propagate naturally. The suggested try-catch adds defensive cleanup, but the impact is moderate since the allocator should handle cleanup internally.	Medium
	Validate metadata size before casting Casting buf.readableBytes() to int can silently truncate metadata larger than 2GB, leading to incomplete data copies and potential buffer overruns. Validate that the buffer size fits within int range before casting, or throw an exception if the metadata exceeds the maximum supported size. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [159-166] private byte[] readMetadata() { ArrowBuf buf = flightStream.getLatestMetadata(); if (buf == null || buf.readableBytes() == 0) return null; - int len = (int) buf.readableBytes(); + long readableBytes = buf.readableBytes(); + if (readableBytes > Integer.MAX_VALUE) { + throw new IllegalStateException("Metadata size exceeds maximum supported: " + readableBytes); + } + int len = (int) readableBytes; byte[] copy = new byte[len]; buf.getBytes(0, copy); return copy; } Suggestion importance[1-10]: 6 __ Why: The suggestion identifies a theoretical issue with casting long to int for metadata larger than 2GB. While technically correct, metadata of that size is highly impractical for per-batch application metadata. The validation adds safety but addresses an edge case unlikely to occur in practice.	Low

Suggestions up to commit b10e5ba

Category	Suggestion	Impact
Possible issue	Add buffer leak protection The ArrowBuf allocation for appMetadata lacks error handling. If allocator.buffer() fails or writeBytes() throws an exception, the buffer will leak. Wrap the allocation and write operations in a try-catch block and ensure the buffer is released on failure before rethrowing. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [179-186] if (appMetadata != null) { - // Flight takes ownership of metadataBuf via putNext(ArrowBuf). - ArrowBuf metadataBuf = allocator.buffer(appMetadata.length); - metadataBuf.writeBytes(appMetadata); - serverStreamListener.putNext(metadataBuf); + ArrowBuf metadataBuf = null; + try { + metadataBuf = allocator.buffer(appMetadata.length); + metadataBuf.writeBytes(appMetadata); + serverStreamListener.putNext(metadataBuf); + } catch (Exception e) { + if (metadataBuf != null) { + metadataBuf.close(); + } + throw e; + } } else { serverStreamListener.putNext(); } Suggestion importance[1-10]: 8 __ Why: The suggestion correctly identifies a resource leak risk. If writeBytes() or putNext() throws an exception after allocator.buffer() succeeds, the ArrowBuf will leak. The improved code properly handles cleanup on failure.	Medium
	Validate metadata size before casting Casting buf.readableBytes() to int can silently truncate metadata larger than 2GB, leading to incomplete data or buffer overruns. Validate that readableBytes() fits within Integer.MAX_VALUE before casting, and throw an exception if it exceeds this limit. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [159-166] private byte[] readAppMetadata() { ArrowBuf buf = flightStream.getLatestMetadata(); if (buf == null || buf.readableBytes() == 0) return null; - int len = (int) buf.readableBytes(); + long readableBytes = buf.readableBytes(); + if (readableBytes > Integer.MAX_VALUE) { + throw new IllegalStateException("appMetadata exceeds maximum size: " + readableBytes); + } + int len = (int) readableBytes; byte[] copy = new byte[len]; buf.getBytes(0, copy); return copy; } Suggestion importance[1-10]: 7 __ Why: The suggestion identifies a potential issue with casting long to int that could truncate large metadata. While metadata exceeding 2GB is unlikely in practice, adding validation prevents silent data corruption and provides clear error messaging.	Medium

Suggestions up to commit 444d754

Category	Suggestion	Impact
Possible issue	Prevent buffer leak on exception If serverStreamListener.putNext(metadataBuf) throws an exception, the allocated metadataBuf will leak because it's not released. Wrap the allocation and write operations in a try-catch block to ensure the buffer is released on failure. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [179-186] if (appMetadata != null) { ArrowBuf metadataBuf = allocator.buffer(appMetadata.length); - metadataBuf.writeBytes(appMetadata); - serverStreamListener.putNext(metadataBuf); + try { + metadataBuf.writeBytes(appMetadata); + serverStreamListener.putNext(metadataBuf); + } catch (Exception e) { + metadataBuf.close(); + throw e; + } } Suggestion importance[1-10]: 8 __ Why: Critical resource leak issue. If writeBytes() or putNext() throws, the allocated ArrowBuf won't be released, causing native memory leak. The comment "Flight takes ownership of metadataBuf via putNext(ArrowBuf)" suggests ownership transfer only happens on success, so cleanup on failure is necessary.	Medium
Security	Add defensive copy for appMetadata The appMetadata byte array is stored directly without defensive copying, creating a potential security risk. If the caller retains a reference and modifies the array after construction, it will affect the internal state. Consider creating a defensive copy to prevent external mutation. plugins/arrow-base/src/main/java/org/opensearch/arrow/transport/ArrowBatchResponse.java [95-98] protected ArrowBatchResponse(VectorSchemaRoot batchRoot, byte[] appMetadata) { this.batchRoot = batchRoot; - this.appMetadata = appMetadata; + this.appMetadata = appMetadata != null ? appMetadata.clone() : null; } Suggestion importance[1-10]: 7 __ Why: Valid security concern about storing a mutable array reference. However, the impact is moderate since appMetadata is typically used in controlled contexts (transport layer), and the PR already documents ownership semantics. Defensive copying would prevent external mutation.	Medium
	Return defensive copy of appMetadata Returning the internal appMetadata array directly exposes it to external modification. Callers can mutate the array, breaking encapsulation and potentially causing unexpected behavior. Return a defensive copy to protect the internal state. plugins/arrow-base/src/main/java/org/opensearch/arrow/transport/ArrowBatchResponse.java [123-125] public byte[] getAppMetadata() { - return appMetadata; + return appMetadata != null ? appMetadata.clone() : null; } Suggestion importance[1-10]: 7 __ Why: Valid encapsulation concern about exposing internal mutable state. The getter returns the raw byte[] which callers could modify. However, the impact is moderate as the metadata is typically consumed once and the PR documents that "bytes are owned by the consumer."	Medium

[github-actions]

❌ Gradle check result for 444d754: 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 b10e5ba

[github-actions]

Persistent review updated to latest commit a50a09e

[github-actions]

✅ Gradle check result for a50a09e: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 96.77419% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.46%. Comparing base (db84aa5) to head (b070add).
⚠️ Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
...rrow/flight/transport/FlightTransportResponse.java	88.88%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22003      +/-   ##
============================================
- Coverage     73.47%   73.46%   -0.01%     
+ Complexity    75588    75567      -21     
============================================
  Files          6037     6038       +1     
  Lines        342786   342811      +25     
  Branches      49311    49312       +1     
============================================
- Hits         251859   251857       -2     
- Misses        70924    70929       +5     
- Partials      20003    20025      +22     

☔ 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 b070add

[github-actions]

❌ Gradle check result for b070add: 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 b070add: 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 b070add: 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 b070add: SUCCESS