feat: Native Arrow transport path for VectorSchemaRoot in Flight plugin

[vamsimanohar]

Context

While building a query engine POC using Apache DataFusion, I found that query results are already in native Arrow format (VectorSchemaRoot) on both the sending and receiving sides. However, the current Flight transport path forces serialization for native arrow format.

Worker DataFusion → VectorSchemaRoot → writeTo() → bytes → VarBinary vector → Flight → VarBinary vector → read(StreamInput) → bytes → Coordinator DataFusion

The data starts as Arrow, gets serialized to bytes, wrapped in a VarBinary vector, sent over Flight, unwrapped, deserialized back to bytes, and then reconstructed — all redundant when both ends already speak Arrow natively.

This PR adds a zero-serialization path so Arrow data can flow directly:

Worker DataFusion → VectorSchemaRoot → Flight → VectorSchemaRoot → Coordinator DataFusion

Summary

Server side (sending)

• ArrowBatchResponse — marker interface for TransportResponse instances carrying native Arrow data. When FlightOutboundHandler detects this, it sends the VectorSchemaRoot directly via putNext(), bypassing VectorStreamOutput serialization.
• FlightServerChannel.sendArrowBatch() — sends the caller's VectorSchemaRoot directly. Tracks ownership via externalRoot flag so close() doesn't free externally-owned roots.

Client side (receiving)

• ArrowStreamHandler — interface for response handlers that can consume native VectorSchemaRoot data via readArrow(). Handlers implementing this receive typed Arrow data instead of byte streams.
• FlightTransportResponse.resolveArrowStreamHandler() — walks the decorator chain (MetricsTrackingResponseHandler → ContextRestoreResponseHandler → TraceableTransportResponseHandler → original handler) to find the ArrowStreamHandler. Result is cached after first resolution.

Handler chain walking

• Added TransportResponseHandler.getDelegate() default method to enable generic decorator chain introspection. Implemented in ContextRestoreResponseHandler, TraceableTransportResponseHandler, and MetricsTrackingResponseHandler.

Key design decisions

• Caller retains ownership of VectorSchemaRoot on the send side — Flight reads from it but does not close it.
• FlightStream reuses its root on the receive side — readArrow() consumers must deep-copy if they need to hold data across next() calls, or process inline for zero-copy.
• Fallback preserved — non-ArrowBatchResponse responses continue to use the existing byte serialization path. Non-ArrowStreamHandler handlers continue to use VectorStreamInput deserialization.

Test plan

• NativeArrowTransportIT.testSingleBatchNativeArrow — 1 batch, 3 rows, verifies typed columns (VarChar, Int) and actual data values end-to-end
• NativeArrowTransportIT.testMultipleBatchesNativeArrow — 3 batches, 2 rows each, verifies multi-batch streaming with data verification
• ArrowBatchResponseTests — unit tests for ownership, serialization bypass, and error handling
• All existing flight integration tests pass (FlightTransportIT, ClientSideChaosIT, SubAggregationIT, etc.)
• All existing flight unit tests pass

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit e5dc448)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 Multiple PR themes Sub-PR theme: Add getDelegate() to TransportResponseHandler decorator chain Relevant files: server/src/main/java/org/opensearch/transport/TransportResponseHandler.java server/src/main/java/org/opensearch/transport/TransportService.java server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java Sub-PR theme: Native Arrow transport path for VectorSchemaRoot in Flight plugin Relevant files: plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowStreamHandler.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/MetricsTrackingResponseHandler.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md

⚡ Recommended focus areas for review Thread Safety The externalRoot flag is set to true inside sendArrowBatch() after the null check on root. If sendArrowBatch() and close() are called concurrently, close() could read externalRoot=false and attempt to close an externally-owned root before externalRoot is set to true. The assignment of externalRoot = true should happen before root is assigned. VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); externalRoot = true; if (root == null) { middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { root = arrowRoot; } Resource Leak In sendArrowBatch(), when root is not null (subsequent batches), the previous root reference is overwritten without any cleanup. If the previous root was internally owned (not external), it would be leaked. Also, callTracker.recordBatchSent() uses System.nanoTime() - batchStartTime for the duration but putNextTime (already converted to ms) is logged — the raw nanosecond value is passed to recordBatchSent which may expect nanoseconds, but this is inconsistent with the logging variable. VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); externalRoot = true; if (root == null) { middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { root = arrowRoot; } logger.debug("Sending native Arrow batch #{} for correlation ID: {}", batchNumber, correlationId); serverStreamListener.putNext(); long putNextTime = (System.nanoTime() - batchStartTime) / 1_000_000; if (callTracker != null) { long rootSize = FlightUtils.calculateVectorSchemaRootSize(root); callTracker.recordBatchSent(rootSize, System.nanoTime() - batchStartTime); logger.debug( "Native Arrow batch #{} sent for correlation ID: {}, size: {} bytes, putNext: {}ms", batchNumber, correlationId, rootSize, putNextTime ); } else { logger.debug("Native Arrow batch #{} sent for correlation ID: {}, putNext: {}ms", batchNumber, correlationId, putNextTime); } Allocator Leak In handleStreamRequest, a new RootAllocator is created per batch but never closed on the server side. The ArrowDataResponse only closes the root and allocator when close() is called on the client side, but the server-side allocator (passed to createTestBatch) is not tracked or closed if an exception occurs before the response is sent. private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); channel.sendResponseBatch(new ArrowDataResponse(root)); } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Unsafe Cast The resolveArrowStreamHandler() method uses @SuppressWarnings("unchecked") and casts current to ArrowStreamHandler<T>. If a decorator in the chain implements ArrowStreamHandler with a different generic type parameter than T, this will produce a runtime ClassCastException when readArrow() is called. There is no type-safe validation before the cast. private ArrowStreamHandler<T> resolveArrowStreamHandler() { TransportResponseHandler<T> current = handler; while (current != null) { if (current instanceof ArrowStreamHandler) { return (ArrowStreamHandler<T>) current; } current = current.getDelegate(); } return null; } Test Isolation In testMultipleBatchesNativeArrow, the test iterates over all cluster nodes and sends requests to each, but the responses list and latch are created outside the loop. If the loop runs more than once (multiple nodes), the latch will only count down once and the responses list will accumulate results from all nodes, making assertions unreliable. The same issue exists in testSingleBatchNativeArrow. public void testMultipleBatchesNativeArrow() throws Exception { for (DiscoveryNode node : getClusterState().nodes()) { StreamTransportService streamTransportService = internalCluster().getInstance(StreamTransportService.class); List<ArrowDataResponse> responses = new ArrayList<>(); CountDownLatch latch = new CountDownLatch(1); AtomicReference<Exception> failure = new AtomicReference<>(); StreamTransportResponseHandler<ArrowDataResponse> handler = createArrowHandler(responses, latch, failure); ArrowDataRequest request = new ArrowDataRequest(3, 2); // 3 batches, 2 rows each streamTransportService.sendRequest( node, ArrowDataAction.NAME, request, TransportRequestOptions.builder().withType(TransportRequestOptions.Type.STREAM).build(), handler ); assertTrue("Stream should complete within 10s", latch.await(10, TimeUnit.SECONDS)); assertNull("No exception expected", failure.get()); assertEquals("Should receive 3 batches", 3, responses.size()); for (int i = 0; i < 3; i++) { ArrowDataResponse response = responses.get(i); VectorSchemaRoot root = response.getArrowRoot(); assertEquals("Each batch should have 2 rows", 2, root.getRowCount()); assertEquals("name", root.getSchema().getFields().get(0).getName()); assertEquals("age", root.getSchema().getFields().get(1).getName()); // Verify data in each batch VarCharVector nameVector = (VarCharVector) root.getVector("name"); IntVector ageVector = (IntVector) root.getVector("age"); assertNotNull("Name vector should not be null", nameVector.get(0)); assertEquals(30, ageVector.get(0)); assertEquals(31, ageVector.get(1)); response.close(); } } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to e5dc448

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Fix allocator memory leak per batch A new RootAllocator is created for every batch but never closed on the server side. The ArrowDataResponse passed to sendResponseBatch only stores the root, and the allocator is not tracked. This causes a memory leak for each batch sent. The allocator should be closed after the batch is sent, or the response should hold and close the allocator. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [308-319] private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); - channel.sendResponseBatch(new ArrowDataResponse(root)); + ArrowDataResponse response = new ArrowDataResponse(root, allocator); + channel.sendResponseBatch(response); + // Close after send since Flight has already copied/transferred the data + root.close(); + allocator.close(); } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Suggestion importance[1-10]: 7 __ Why: A new RootAllocator is created per batch but never closed on the server side in the test code. This is a real memory leak in the test, and the ArrowDataResponse constructor that takes both root and allocator exists but is not used here. The fix is accurate and directly addresses the leak.	Medium
	Fix race condition in handler resolution There is a race condition: arrowHandlerResolved and cachedArrowHandler are two separate volatile writes, so another thread could observe arrowHandlerResolved == false and cachedArrowHandler being non-null or null inconsistently. Use a synchronized block or a single atomic reference to ensure both fields are updated atomically. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [131-134] if (!arrowHandlerResolved) { - cachedArrowHandler = resolveArrowStreamHandler(); - arrowHandlerResolved = true; + synchronized (this) { + if (!arrowHandlerResolved) { + cachedArrowHandler = resolveArrowStreamHandler(); + arrowHandlerResolved = true; + } + } } Suggestion importance[1-10]: 5 __ Why: The two separate volatile writes to cachedArrowHandler and arrowHandlerResolved could theoretically be observed inconsistently by another thread. However, nextResponse() is typically called from a single consumer thread in streaming scenarios, making this a low-risk issue in practice. The double-checked locking pattern in the improved code is correct but may be over-engineering for this use case.	Low
	Fix flag ordering to prevent race condition The externalRoot flag is set to true on the first call but never reset between batches. If sendBatch (the byte path) is called after sendArrowBatch, the close() method will incorrectly skip closing the internally-owned root. Additionally, externalRoot should be set before assigning root to avoid a race condition on the volatile field. Consider making externalRoot a per-call local concern or resetting it appropriately. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [159-167] VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); -externalRoot = true; if (root == null) { + externalRoot = true; middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { + externalRoot = true; root = arrowRoot; } Suggestion importance[1-10]: 4 __ Why: The suggestion to set externalRoot before assigning root has some merit for ordering clarity, but the actual race condition concern is limited since sendArrowBatch is called sequentially per batch. The mixed-path concern (calling sendBatch after sendArrowBatch) is a valid edge case but unlikely in practice. The improved code doesn't fundamentally change behavior.	Low
General	Prevent silent root reference leak on close When sendArrowBatch is called multiple times (multiple batches), root is reassigned to each new arrowRoot but the previous external root is never closed. The caller (via ArrowDataResponse) is responsible for closing, but the channel holds a reference to the last assigned root. If close() is called and externalRoot is true, the last root is silently leaked. Consider nulling out root after each batch in the external path, or documenting clearly that callers must close all roots. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [282-284] if (root != null && !externalRoot) { root.close(); } +// For externalRoot, the caller is responsible for closing all VectorSchemaRoot instances. +root = null; Suggestion importance[1-10]: 3 __ Why: The suggestion to null out root after close is a minor defensive improvement, but the improved_code is essentially the same as existing_code with just a comment and a root = null added after the closing brace. The actual leak concern is valid but minor since the channel itself is closed at this point.	Low

---

Previous suggestions

Suggestions up to commit 6b7b55a

Category	Suggestion	Impact
Possible issue	Prevent allocator leak per batch on server side A new RootAllocator is created per batch but never closed on the server side. The ArrowDataResponse passed to sendResponseBatch does not hold the allocator, so after putNext() completes the allocator leaks. The allocator should be passed to ArrowDataResponse so it can be closed after the batch is sent, or the channel should close it post-send. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [308-319] private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); - channel.sendResponseBatch(new ArrowDataResponse(root)); + channel.sendResponseBatch(new ArrowDataResponse(root, allocator)); } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Suggestion importance[1-10]: 7 __ Why: The server-side handleStreamRequest creates a new RootAllocator per batch but passes it to ArrowDataResponse without the allocator reference, causing a memory leak. The fix correctly passes the allocator to ArrowDataResponse(root, allocator) so it can be properly closed.	Medium
	Fix ownership flag set before root assignment The externalRoot flag is set to true on the first call but never reset between batches. If sendBatch (the byte path) is called after sendArrowBatch, the close() method will incorrectly skip closing the internally-owned root. Additionally, when switching from an external root to an internal one (or vice versa), the flag must be updated per-call rather than once. Consider tracking ownership per root assignment rather than as a single sticky flag. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [159-167] VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); -externalRoot = true; if (root == null) { middleware.setHeader(header); root = arrowRoot; + externalRoot = true; serverStreamListener.start(root); } else { root = arrowRoot; + externalRoot = true; } Suggestion importance[1-10]: 4 __ Why: The suggestion moves externalRoot = true inside the conditional branches rather than before them. While this is a minor ordering improvement, the functional behavior is identical since externalRoot is always set to true in both cases. The suggestion doesn't address the more significant concern about mixing sendBatch and sendArrowBatch calls.	Low
General	Close Arrow responses after sending to prevent leaks When sendArrowBatch is called multiple times, root is updated to point to the latest external root but the previous external roots are never closed. Since the caller (server action) is responsible for closing them, this is correct only if the caller always closes each ArrowDataResponse after sendResponseBatch returns. However, the current server-side code in handleStreamRequest does not close the response after sending, so those roots and their allocators will leak if the channel is closed before the caller cleans up. Add a note or enforce cleanup contract. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [277-286] -public void close() { - if (!open.get()) { - return; - } - open.set(false); - if (root != null && !externalRoot) { - root.close(); - } - notifyCloseListeners(); +// In handleStreamRequest, close each response after sending: +for (int batch = 0; batch < request.getBatchCount(); batch++) { + BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); + ArrowDataResponse response = new ArrowDataResponse(root, allocator); + channel.sendResponseBatch(response); + response.close(); // close root and allocator after putNext() completes } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that server-side ArrowDataResponse objects (with their allocators) are not closed after sendResponseBatch, but the improved_code targets the test's handleStreamRequest rather than the close() method in FlightServerChannel.java that is cited as the relevant_file. The fix is valid but mismatched to the stated relevant file.	Low
	Fix non-atomic double-checked resolution of Arrow handler This resolution is not thread-safe. arrowHandlerResolved and cachedArrowHandler are two separate volatile fields, so a thread can read arrowHandlerResolved == false, another thread sets both fields, and the first thread overwrites cachedArrowHandler with a redundant resolution. While the result is the same, a more robust approach is to use a single atomic check-and-set or synchronize the resolution block to avoid redundant work and ensure visibility of both fields together. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [131-134] if (!arrowHandlerResolved) { - cachedArrowHandler = resolveArrowStreamHandler(); - arrowHandlerResolved = true; + synchronized (this) { + if (!arrowHandlerResolved) { + cachedArrowHandler = resolveArrowStreamHandler(); + arrowHandlerResolved = true; + } + } } Suggestion importance[1-10]: 4 __ Why: The double-checked locking concern is valid in theory, but since resolveArrowStreamHandler() is idempotent and the result is the same regardless of which thread wins, the practical impact is minimal. The volatile fields provide visibility guarantees, making this a low-priority concern.	Low

Suggestions up to commit ae33c2b

Category	Suggestion	Impact
Possible issue	Prevent allocator memory leak per batch A new RootAllocator is created for every batch but is never closed on the server side. The ArrowDataResponse passed to sendResponseBatch only stores the root, not the allocator, so the allocator will leak if the channel does not close it. The allocator should be passed to ArrowDataResponse so it can be closed after the batch is sent, or the channel must guarantee cleanup. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [308-319] private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); - channel.sendResponseBatch(new ArrowDataResponse(root)); + channel.sendResponseBatch(new ArrowDataResponse(root, allocator)); } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Suggestion importance[1-10]: 6 __ Why: The server-side RootAllocator created per batch is never closed since ArrowDataResponse(root) (without allocator) is used. Passing the allocator to ArrowDataResponse(root, allocator) ensures proper cleanup, which is a real memory leak in the test code.	Low
	Fix non-atomic double-checked initialization race This double-checked pattern is not thread-safe without synchronization. Since arrowHandlerResolved and cachedArrowHandler are volatile but the check-then-act is not atomic, two threads could both see arrowHandlerResolved == false and both call resolveArrowStreamHandler(). While the result is idempotent, it could cause a race. Use a synchronized block or AtomicBoolean to ensure the resolution happens exactly once. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [131-134] if (!arrowHandlerResolved) { - cachedArrowHandler = resolveArrowStreamHandler(); - arrowHandlerResolved = true; + synchronized (this) { + if (!arrowHandlerResolved) { + cachedArrowHandler = resolveArrowStreamHandler(); + arrowHandlerResolved = true; + } + } } Suggestion importance[1-10]: 4 __ Why: While the race condition is real, nextResponse() is typically called sequentially from a single consumer thread in stream processing, making this a low-risk issue in practice. The volatile fields do provide visibility guarantees, and the resolution is idempotent, so the practical impact is minimal.	Low
	Fix flag set order for external root tracking The externalRoot flag is set to true on the first call but is never reset between batches. If sendBatch (the byte path) is called after sendArrowBatch, the close() method will incorrectly skip closing the internally-owned root. Additionally, externalRoot should be set before any conditional logic to ensure consistency. Consider using a per-call local variable or resetting the flag appropriately for each batch. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [159-167] VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); -externalRoot = true; if (root == null) { + externalRoot = true; middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { + externalRoot = true; root = arrowRoot; } Suggestion importance[1-10]: 3 __ Why: The suggestion moves externalRoot = true inside the conditional branches, but the functional behavior is identical since externalRoot is set to true in both cases regardless. The concern about mixing sendBatch and sendArrowBatch is valid but the proposed fix doesn't actually address that scenario.	Low
General	Reset external root flag on byte-path batch send When sendArrowBatch is called multiple times, root is updated to point to the latest external VectorSchemaRoot but the previous external roots are not tracked. If the channel is closed after multiple sendArrowBatch calls, none of the externally-owned roots will be closed here (which is correct), but the root field still holds a stale reference to the last batch's root. This is acceptable only if the caller guarantees cleanup; however, if sendBatch (byte path) is ever called after sendArrowBatch, externalRoot remains true and the internally-created root will also be leaked. Consider resetting externalRoot to false when sendBatch is called. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [277-286] -public void close() { - if (!open.get()) { - return; - } - open.set(false); - if (root != null && !externalRoot) { - root.close(); - } - notifyCloseListeners(); +public void sendBatch(ByteBuffer header, VectorStreamOutput output) { + // ... existing code ... + externalRoot = false; + // ... rest of existing sendBatch logic ... } Suggestion importance[1-10]: 4 __ Why: The concern about mixing sendBatch and sendArrowBatch causing a leak is valid, but the improved_code shows a skeleton with // ... existing code ... placeholders rather than the actual implementation, making it an incomplete suggestion. The underlying issue is real but the fix is not properly demonstrated.	Low

Suggestions up to commit 57d5693

Category	Suggestion	Impact
Possible issue	Prevent allocator memory leak per batch A new RootAllocator is created per batch but never closed on the server side. The ArrowDataResponse sent via sendResponseBatch only stores the root (not the allocator), so the allocator leaks memory after each batch is sent. The allocator should be passed to ArrowDataResponse or closed after the batch is confirmed sent. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [308-319] private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); - channel.sendResponseBatch(new ArrowDataResponse(root)); + channel.sendResponseBatch(new ArrowDataResponse(root, allocator)); } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Suggestion importance[1-10]: 7 __ Why: A new RootAllocator is created per batch but never closed on the server side since ArrowDataResponse(root) doesn't store the allocator. Passing the allocator to ArrowDataResponse(root, allocator) ensures proper cleanup, preventing a real memory leak in the test.	Medium
	Fix shared ownership flag set before conditional logic The externalRoot flag is set to true on the first call but is never reset between batches. If sendBatch (the byte path) is called after sendArrowBatch, the close() method will incorrectly skip closing the internally-owned root. Additionally, the externalRoot flag should be set before any conditional logic to ensure it's always applied. Consider using a per-call local approach or tracking ownership per root rather than a single shared flag. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [159-167] VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); -externalRoot = true; if (root == null) { + externalRoot = true; middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { + externalRoot = true; root = arrowRoot; } Suggestion importance[1-10]: 4 __ Why: The suggestion points out that externalRoot is set before the conditional but the improved code still sets it in both branches, which is functionally equivalent to the original. The real concern about mixing byte/arrow paths is valid but the fix doesn't actually address the core issue described.	Low
General	Cache resolved handler to avoid repeated chain traversal The resolveArrowStreamHandler() is called on every call to nextResponse(), walking the decorator chain each time. This is inefficient for streams with many batches. The resolved handler (or a null sentinel) should be cached after the first resolution to avoid repeated traversal on every batch. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [129-139] -ArrowStreamHandler<T> arrowHandler = resolveArrowStreamHandler(); -if (arrowHandler != null) { +if (cachedArrowHandler == null && !arrowHandlerResolved) { + cachedArrowHandler = resolveArrowStreamHandler(); + arrowHandlerResolved = true; +} +if (cachedArrowHandler != null) { // Native Arrow path: hand VectorSchemaRoot directly to the handler - return arrowHandler.readArrow(root); + return cachedArrowHandler.readArrow(root); } // Existing byte path: deserialize via VectorStreamInput try (VectorStreamInput input = new VectorStreamInput(root, namedWriteableRegistry)) { input.setVersion(initialHeader.getVersion()); return handler.read(input); } Suggestion importance[1-10]: 4 __ Why: Caching the resolved ArrowStreamHandler is a valid optimization for multi-batch streams, but the decorator chain traversal is typically very short (1-2 levels), making the performance impact minimal. The suggestion also requires adding new fields (cachedArrowHandler, arrowHandlerResolved) not shown in the diff.	Low
	Clarify and fix root ownership on channel close When multiple sendArrowBatch calls are made, root is reassigned to each new arrowRoot but the previous external root is not tracked. If close() is called and externalRoot is true, none of the roots will be closed — but if externalRoot is false (e.g., after a mixed byte/arrow path), only the last root is considered. The ownership model needs to be clarified to ensure all externally-owned roots are properly managed by their callers and the flag accurately reflects the current root reference. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [277-286] public void close() { if (!open.get()) { return; } open.set(false); + // Only close root if it was internally allocated (not owned by the caller) if (root != null && !externalRoot) { root.close(); } + // If externalRoot == true, the caller (ArrowDataResponse) is responsible for closing + root = null; notifyCloseListeners(); } Suggestion importance[1-10]: 3 __ Why: The suggestion to null out root after close is a minor defensive improvement, but the improved_code is nearly identical to the existing_code with only a comment and root = null added. The ownership concern is valid but the fix is marginal.	Low

Suggestions up to commit c530fec

Category	Suggestion	Impact
Possible issue	Fix allocator leak in server-side batch handler A new RootAllocator is created for every batch but is never closed on the server side. The ArrowDataResponse passed to sendResponseBatch only stores the root, not the allocator, so the allocator leaks after the batch is sent. The allocator should be tracked and closed after the batch is sent, or the ArrowDataResponse constructor that accepts both root and allocator should be used so the caller can close it. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [308-319] private void handleStreamRequest(ArrowDataRequest request, TransportChannel channel, Task task) throws IOException { try { for (int batch = 0; batch < request.getBatchCount(); batch++) { BufferAllocator allocator = new RootAllocator(); VectorSchemaRoot root = createTestBatch(allocator, request.getRowsPerBatch(), batch); - channel.sendResponseBatch(new ArrowDataResponse(root)); + ArrowDataResponse response = new ArrowDataResponse(root, allocator); + channel.sendResponseBatch(response); + response.close(); // close root and allocator after batch is sent } channel.completeStream(); } catch (Exception e) { channel.sendResponse(e); } } Suggestion importance[1-10]: 7 __ Why: A new RootAllocator is created per batch but never closed on the server side, causing a real memory leak in the test. The improved code correctly closes the response (and thus the allocator) after sending, though care must be taken that the channel has finished reading the root before closing.	Medium
	Fix shared ownership flag set before conditional logic The externalRoot flag is set to true on the first call but is never reset between batches. If sendBatch (the byte path) is called after sendArrowBatch, the close() method will incorrectly skip closing the internally-owned root. Additionally, externalRoot should be set before any conditional logic to ensure it's always applied. Consider using a per-call local approach or tracking ownership per root rather than a single shared flag. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java [159-167] VectorSchemaRoot arrowRoot = arrowResponse.getArrowRoot(); -externalRoot = true; if (root == null) { + externalRoot = true; middleware.setHeader(header); root = arrowRoot; serverStreamListener.start(root); } else { + externalRoot = true; root = arrowRoot; } Suggestion importance[1-10]: 5 __ Why: The externalRoot flag is set once and never reset, which could cause issues if sendBatch is called after sendArrowBatch. However, the improved code still sets externalRoot = true in both branches without addressing the core concern of mixed-path usage, making the fix incomplete. The suggestion identifies a real but edge-case concern.	Low
General	Document root reuse contract to prevent data corruption The root obtained from flightStream.getRoot() is reused by the Flight stream on subsequent next() calls. If readArrow returns without copying the data (e.g., a handler that wraps the root directly), the data will be overwritten on the next batch. The comment in the integration test acknowledges this, but the core transport path has no guard. Consider documenting this contract clearly in ArrowStreamHandler.readArrow() Javadoc, or asserting/logging a warning if the returned response still references the original root. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java [129-133] ArrowStreamHandler<T> arrowHandler = resolveArrowStreamHandler(); if (arrowHandler != null) { - // Native Arrow path: hand VectorSchemaRoot directly to the handler + // Native Arrow path: hand VectorSchemaRoot directly to the handler. + // CONTRACT: the handler MUST copy any data it needs to retain, as the + // Flight stream reuses this root on the next call to flightStream.next(). return arrowHandler.readArrow(root); } Suggestion importance[1-10]: 2 __ Why: The suggestion only adds a comment to document an existing contract without changing any logic. The existing_code and improved_code are functionally identical, making this a documentation-only change with minimal impact.	Low
	Improve unsupported fallback error message clarity The ArrowDataAction constructor registers ArrowDataResponse::new as the reader, which maps to the StreamInput constructor. If this constructor is ever invoked (e.g., during action registration or response deserialization in non-Flight paths), it will throw UnsupportedOperationException rather than a more informative error. While acceptable for a test, the ActionType super constructor call with this reader means it could be triggered unexpectedly during cluster state operations. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [202-206] public ArrowDataResponse(StreamInput in) throws IOException { super(in); // Fallback deserialization for Netty4 — not expected in native Arrow path - throw new UnsupportedOperationException("Netty4 fallback not implemented in this test"); + throw new UnsupportedOperationException( + "ArrowDataResponse does not support Netty4 byte deserialization. " + + "This response is only valid over the native Arrow Flight transport path." + ); } Suggestion importance[1-10]: 1 __ Why: This is a test-only class and the suggestion only improves the error message text in an UnsupportedOperationException, which has minimal impact on correctness or functionality.	Low

[github-actions]

Persistent review updated to latest commit 57d5693

[github-actions]

❌ Gradle check result for 57d5693: 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 ae33c2b

[vamsimanohar]

Responses to review suggestions

1. Allocator leak in test (handleStreamRequest)

This is test-only code. The core issue is that sendResponseBatch() is async — it dispatches to an executor that calls putNext() later. The server has no callback to know when putNext() completes, so it can't safely close the allocator/root immediately after sending. Closing prematurely causes the executor to read freed/zeroed memory (confirmed by test failure when we tried response.close() after send).

In a real production consumer (e.g., the lakehouse distributed query engine), DataFusion owns the VectorSchemaRoot and manages its own memory lifecycle. The ownership model for production use will be addressed when we integrate the first real consumer. Deferring this for now.

2. Cache resolveArrowStreamHandler() (per-batch decorator chain walk)

Fixed in ae33c2b. The resolved handler is now cached after the first call to nextResponse(), avoiding repeated decorator chain traversal on every batch.

3. Thread safety of externalRoot

The class Javadoc states: "This implementation is not thread safe; consumer must ensure to invoke sendBatch serially and call completeStream() at the end." sendArrowBatch() and close() are not called concurrently — close() runs after completeStream() or on cancellation, both of which happen after all batches are sent.

4. Mixed byte/arrow path (externalRoot never reset)

A single stream is either all native Arrow or all byte path. Mixing sendArrowBatch() and sendBatch() on the same channel is not a supported use case and would not make semantic sense. The externalRoot flag correctly reflects the ownership for the stream's lifetime.

Additional detail on why we're not fixing this:

• Can't happen in practice — a FlightServerChannel is created per stream, and a stream is either all native Arrow batches or all byte batches. The server handler picks one path based on the response type and uses it consistently. There's no scenario where sendArrowBatch() and sendBatch() would be mixed on the same channel.
• Same pattern as existing code — sendBatch() doesn't validate this either. Neither path checks what the previous batch type was. The channel is a fire-once-per-stream construct, not a general-purpose reusable channel.
• Class Javadoc already establishes the contract — FlightServerChannel Javadoc states: "This implementation is not thread safe; consumer must ensure to invoke sendBatch serially and call completeStream() at the end." The caller controls what goes through the channel and is responsible for using it correctly.
• Adding a guard for mixed paths would add complexity for a scenario that can't arise and would be inconsistent with how the byte path already works.

5. Schema mismatch in subsequent batches

This is the caller's responsibility, same as the existing byte path (sendBatch) which also doesn't validate schema consistency between batches. Adding validation only to the native Arrow path would be inconsistent.

6. Document root reuse contract in ArrowStreamHandler.readArrow() Javadoc

The ArrowStreamHandler interface Javadoc already documents that the root is borrowed from the Flight stream. Will add an explicit note about reuse in a follow-up if needed.

7. Gradle check failure

The failure is FlightOutboundHandlerTests.testSendResponseBatchPropagatesContextToExecutorThread — a pre-existing flaky test. Confirmed by running it on the base commit without any of our changes — same failure. Unrelated to this PR.

[github-actions]

❌ Gradle check result for ae33c2b: 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?

[vamsimanohar]

@rishabhmaurya @reta @finnegancarroll — requesting review.

[github-actions]

Persistent review updated to latest commit 6b7b55a

[github-actions]

❌ Gradle check result for 6b7b55a: 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?

[rishabhmaurya]

thanks @vamsimanohar for working on it.

However, the current Flight transport path forces unnecessary serialization:

I won't call it unnecessary as the very first use case we onboarded was streaming aggregations, where we have to do ser/de to/from StreamOutput/Input. Maybe oneday we will start making use of Arrow format within opensearch aggregation directly and call it unnecessary :)

The main question here is - who is doing buffer management? Java or DF?
I'm of the opinion that Java should be doing all buffer management and pass the allocators to DF where it needs.

I think the changes can be simplified a lot as it was always designed to support this primary use case. I will add details in my next comment in a while (running the approach with kiro to implement and check).

[mch2]

nitpick - this is an alternate to Writeable's handler.read(StreamInput) - lets call it something like ArrowStreamReader to not confuse with transport response handlers.

[mch2]

The main question here is - who is doing buffer management? Java or DF?

Ideally, the buffer management should be java side with all allocation a child of the same root. Right now DataFusionService is creating its own root allocator java side.

Not an expert here, but I'm assuming your concern @rishabhmaurya is the externalRoot = true piece where FlightServerChannel skips closing the root and doesn't claim ownership? Regardless of whether Java or a backend engine allocated the VSR, should flight close it or make it the caller's responsibility?

[vamsimanohar]

thanks @vamsimanohar for working on it.

However, the current Flight transport path forces unnecessary serialization:

I won't call it unnecessary as the very first use case we onboarded was streaming aggregations, where we have to do ser/de to/from StreamOutput/Input. Maybe oneday we will start making use of Arrow format within opensearch aggregation directly and call it unnecessary :)

My bad. It slipped into description from the context of my POC and the discussion with agent. Anyways edited it. 👍 I know we can't avoid it for existing usecases.

[vamsimanohar]

The main question here is - who is doing buffer management? Java or DF? I'm of the opinion that Java should be doing all buffer management and pass the allocators to DF where it needs.

I think the changes can be simplified a lot as it was always designed to support this primary use case. I will add details in my next comment in a while (running the approach with kiro to implement and check).

Thanks for the feedback @rishabhmaurya

I could see few problems with the current implementation particularly on sending side.

Server-side(Sending Arrow Batch):
sendResponseBatch() is async, so the caller can't immediately close the root after calling sendResponseBatch and there is no callback as well. externalRoot=true was added so the channel doesn't close caller-owned roots, but then nobody closes them.

Either we provide a completion callback so the caller can close it after consumption, or FlightServerChannel takes ownership and closes it. I think FlightServerChannel closing it is fine, unless the caller is using the batch for further computation. Depends on the use case. I think once we send a batch from one node to another there won't be further computation. I should maybe revert the externalRoot change. Let me know if your idea is to introduce new interfaces altogether.

==============================================================================

Client-side:
FlightStream reuses one VectorSchemaRoot across next() calls, so consumers must either process each batch inline before calling next(), or copy it to hold multiple batches. The consumer can use any BufferAllocator for the copy, so memory tracks under whichever pool they choose. The question is whether to copy or read inline, which depends on the use case. [Not a big concern]

==============================================================================

Multiple allocator pools:
There are three independent memory pools today — Rust GreedyMemoryPool (DataFusion execution), DataFusionService.RootAllocator (Java-side Arrow buffers), and Flight's RootAllocator (transport buffers). Today, DataFusion creates RecordBatches in Rust and transfers them to a child allocator of DataFusionService.RootAllocator via the C Data Interface import(zero copy). But during execution, these pools are blind to each other. There's no unified budget. Do we need coordinated limits...what do you think?

Will give it a thought on these problems.

[rishabhmaurya]

@vamsimanohar Yes - I like the way you're going about it, however to avoid back n forth, I created some rough changes on how I think this should addressed, please take a look as a reference while you work on it - #21253 ; it have few tests and passes the checks but there might be issues.

Concerns with the current approach (which you're mostly aware of, still listing them here)

1. Transport-level branching

FlightOutboundHandler.processBatchTask() and FlightTransportResponse.nextResponse() both have instanceof checks to branch between byte and native Arrow paths. This couples the transport layer to specific response types. The transport should be agnostic — the response type should control its own serialization behavior.

2. Core server changes for handler introspection

The PR adds getDelegate() to TransportResponseHandler (core interface), TransportService.ContextRestoreResponseHandler, and TraceableTransportResponseHandler. These changes exist solely to walk the decorator chain and discover if the underlying handler supports native Arrow. This is invasive — a plugin-level feature should not require changes to core transport interfaces.

3. Buffer ownership confusion

FlightServerChannel introduces an externalRoot boolean to track whether the current root is owned by the caller or the channel. This flag is channel-wide but ownership varies per-batch. If a channel ever mixes byte and native Arrow batches (the API allows it), or if close() is called at the wrong time, roots can leak or be double-freed.

4. No pipelining support

The PR passes the VectorSchemaRoot directly from the producer to the transport. Since Flight binds to one root via start(), and putNext() reads from that root, the producer cannot queue batches ahead. If the producer modifies the root for the next batch before the executor calls putNext(), the data is corrupted. This is a race condition for any non-trivial streaming use case.

Proposed alternative

Key tenet

• Response controls ser/de: The response type decides how to serialize/deserialize — not the transport.
• Java owns buffer management: The channel's allocator is the parent. Producers create roots from it. The framework manages the shared root bound to Flight.
• Zero-copy via transfer: TransferPair.transfer() moves buffer pointers from producer vectors to the shared root — no memcpy.
• Pipelining safe: Each batch has independent buffers. The executor transfers and sends them serially.
• No core changes: Everything within the arrow-flight-rpc plugin.

How it works

Send side: API developer extends ArrowBatchResponse. Its writeTo() is a final no-op. FlightOutboundHandler detects this type, creates a shared root on the first batch, and calls transferTo() to zero-copy the producer's buffers into the shared root before putNext(). All transfer happens on the executor — safe for pipelined production.

Receive side: ArrowBatchResponse(StreamInput in) calls ((VectorStreamInput) in).getRoot(). The handler's read(StreamInput) constructs the response via this path. FlightTransportResponse always creates the same VectorStreamInput — no branching, no factory selection. The handler decides whether to call getRoot() (native) or byte methods (existing path).

Allocator access: ArrowFlightChannel interface with getAllocator(). ArrowFlightChannel.from(channel) unwraps the TaskTransportChannel/BaseTcpTransportChannel wrapper chain to find the underlying FlightServerChannel.

[github-actions]

Persistent review updated to latest commit e5dc448

[github-actions]

❌ Gradle check result for e5dc448: 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?

[vamsimanohar]

Closing this in favour of #21253