Native Arrow path support in stream transport with zero-copy transfer

[rishabhmaurya]

Context

Adds a native Arrow transport path to the Flight transport plugin, enabling zero-copy transfer of Arrow data without byte serialization. This is an alternative to #21240 that keeps all changes within the arrow-flight-rpc plugin.

Problem

The existing byte-serialized path (VectorStreamOutput.ByteSerialized) serializes Arrow data into a VarBinaryVector, which is then sent via Flight putNext(). The byte-serialized path is essential for streaming aggregation where results are produced as OpenSearch objects and need serialization into Arrow format. However, for use cases like DataFusion integration where data already originates as native Arrow vectors (potentially from C data import), a direct transfer path avoids the serialization round-trip.

Solution

Introduce ArrowBatchResponse — an abstract base class that API developers extend. When the framework detects this response type, it performs a zero-copy transferTo() of the producer's vectors into the channel's shared root, bypassing serialization entirely.

Key design decisions

Producer's allocator for the shared root (same-allocator transfer). The shared root is created from the first batch's producer allocator. This ensures same-allocator transfer, which avoids an Arrow Java bug where BufferLedger.transferOwnership() of foreign-backed buffers (from C data import via wrapForeignAllocation) doesn't properly free the ArrowArray C struct (128 bytes per batch). Same-allocator transfer sidesteps this entirely.

Long-lived allocator requirement. The allocator used for producer roots must outlive the gRPC stream. gRPC's zero-copy write path (ArrowBufRetainingCompositeByteBuf) retains ArrowBuf references beyond putNext() and even beyond completed() — they are released asynchronously by gRPC's Netty event loop. Closing the allocator while gRPC still holds these retained references causes memory accounting errors.

Producer root closed after transfer. Each batch's producer root is closed by the framework after transferTo() moves its buffers into the shared root. The producer's buffers are empty after transfer, so close is safe and immediate.

Challenges investigated

gRPC zero-copy buffer lifecycle. putNext() with setUseZeroCopy(true) creates ArrowBufRetainingCompositeByteBuf which retains ArrowBufs independently of the shared root. When transferTo() replaces the shared root's buffers on the next batch, the old ArrowBufs are released from the root's side but kept alive by gRPC's retain (refcount > 0). The byte-serialized path avoids this because it reuses the same ArrowBuf across batches (overwriting contents via setSafe()), so gRPC's retained ByteBufs always point to valid, live memory. For the native arrow path, the allocator must be long-lived so gRPC can release the retained ArrowBufs back to it at any time.

Upstream issues discovered

During this work we identified two issues in Arrow Java that are worth reporting upstream:

1. Arrow Java: BufferLedger.transferOwnership() leaks ArrowArray C struct for foreign-backed buffers. When ArrayImporter.importArray() imports data via the C Data Interface, it allocates a 128-byte ArrowArray C struct buffer from the data allocator, wrapped in ReferenceCountedArrowArray. During cross-allocator transferOwnership, the accounting for the data buffers moves correctly, but the ReferenceCountedArrowArray refcount never reaches 0 because the foreign allocation cleanup path isn't triggered. This leaks 128 bytes per imported array per cross-allocator transfer. Same-allocator transfer avoids this because transferOwnership returns the same buffer (no new allocation). Present in Arrow Java 18.1.0 and latest main.

2. Arrow Flight Java: ServerStreamListener.completed() is fire-and-forget with no flush guarantee. completed() enqueues HTTP/2 trailers to gRPC's WriteQueue but returns immediately without waiting for pending data frames to be flushed. Meanwhile, ArrowBufRetainingCompositeByteBuf holds retained references to ArrowBufs that are only released when gRPC's Netty event loop processes the write. There is no callback mechanism (ServerStreamListener doesn't expose setOnCloseHandler, which exists on ServerCallStreamObserver but is not accessible through the Flight API) to know when gRPC has fully released all buffer references. This means allocators backing the stream's data cannot be safely closed immediately after completed(). A setOnCloseHandler or similar API on ServerStreamListener would allow producers to defer cleanup until gRPC is truly done with the buffers.

Design doc

See plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit a5e7857)

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: Native Arrow transport path core implementation and tests 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/ArrowFlightChannel.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/FlightTransportChannel.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/VectorStreamOutput.java plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamInput.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowFlightChannelTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamInputTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamOutputTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md Sub-PR theme: Native Arrow stream transport example plugin Relevant files: plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataAction.java plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataRequest.java plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataResponse.java plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/TransportNativeArrowStreamDataAction.java plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/StreamTransportExamplePlugin.java plugins/examples/stream-transport-example/src/internalClusterTest/java/org/opensearch/example/stream/NativeArrowStreamTransportExampleIT.java

⚡ Recommended focus areas for review Null Dereference Risk When creating the shared root for the native Arrow path, the code accesses arrowResponse.getRoot().getFieldVectors().get(0).getAllocator(). If the producer root has no field vectors (empty schema), this will throw an IndexOutOfBoundsException. There is no guard for an empty vector list before calling .get(0). sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() ); Shared Root Leak The shared root created in the native Arrow path (when sharedRoot == null) is not stored back into flightChannel. After transferTo() and sendBatch(), the newly created shared root is passed to VectorStreamOutput.forNativeArrow(sharedRoot) and then closed via try (out). On the next batch, flightChannel.getRoot() will again return null, causing a new shared root to be created each time instead of reusing it. This defeats the purpose of the shared root pattern and may cause resource management issues. VectorSchemaRoot sharedRoot = flightChannel.getRoot(); if (sharedRoot == null) { // Create shared root using the producer's allocator for same-allocator transfer. // This avoids an Arrow bug where cross-allocator transferOwnership of foreign-backed // buffers (from C data import) doesn't properly free the ArrowArray C struct. // The producer's allocator must be long-lived (not closed per-request). sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() ); } arrowResponse.transferTo(sharedRoot); arrowResponse.getRoot().close(); // release producer's buffers — safe, they've been moved out = VectorStreamOutput.forNativeArrow(sharedRoot); } else { out = VectorStreamOutput.create(flightChannel.getAllocator(), flightChannel.getRoot()); task.response().writeTo(out); } try (out) { flightChannel.sendBatch(getHeaderBuffer(task.requestId(), task.nodeVersion(), task.features()), out); messageListener.onResponseSent(task.requestId(), task.action(), task.response()); Root Closed on Close In ByteSerialized.close(), the vector is closed but the root field (if created) is not closed. This may leave the VectorSchemaRoot wrapper alive without its underlying vector, potentially causing resource tracking issues depending on how Arrow accounts for the root vs. the vector. public void close() throws IOException { row = 0; vector.close(); } AwaitsFix Tests Both integration tests in this file are annotated with @AwaitsFix(bugUrl = "") with an empty bug URL. These tests will be skipped in CI. The empty bugUrl is also non-standard — it should reference a tracking issue. Consider either fixing the tests or providing a valid bug URL. @AwaitsFix(bugUrl = "") @LockFeatureFlag(STREAM_TRANSPORT)

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to a5e7857

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Guard against empty field vector list If the producer root has no field vectors (empty schema), calling getFieldVectors().get(0) will throw an IndexOutOfBoundsException. Add a guard to fall back to the channel allocator when the vector list is empty. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [164-167] +List<FieldVector> fieldVectors = arrowResponse.getRoot().getFieldVectors(); +BufferAllocator rootAllocator = fieldVectors.isEmpty() + ? flightChannel.getAllocator() + : fieldVectors.get(0).getAllocator(); sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), - arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() + rootAllocator ); Suggestion importance[1-10]: 7 __ Why: Calling getFieldVectors().get(0) on an empty schema will throw an IndexOutOfBoundsException. This is a valid edge case that should be handled, though schemas with no fields are uncommon in practice.	Medium
	Ensure executor shutdown on timeout or interruption producers.shutdown() is called after producersDone.await(), but if the await times out or is interrupted, the executor is never shut down, causing a thread leak. Call producers.shutdownNow() in a finally block to ensure cleanup regardless of outcome. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [326-327] -producersDone.await(30, TimeUnit.SECONDS); -producers.shutdown(); +try { + producersDone.await(30, TimeUnit.SECONDS); +} finally { + producers.shutdownNow(); +} Suggestion importance[1-10]: 5 __ Why: The producers executor is not shut down if producersDone.await() times out or is interrupted, causing a thread leak. However, this is test code, so the impact is lower than in production code.	Low
	Prevent Arrow memory leak on enqueue failure If createBatch succeeds but queue.put() throws (e.g., InterruptedException), the created VectorSchemaRoot is never closed, leaking Arrow memory. The root should be closed in the catch block if it was not successfully enqueued. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [305-314] producers.submit(() -> { try { VectorSchemaRoot root = createBatch(allocator, batchIndex, request.rowsPerBatch); - queue.put(new TestArrowResponse(root)); + try { + queue.put(new TestArrowResponse(root)); + } catch (Exception e) { + root.close(); + throw e; + } } catch (Exception e) { throw new RuntimeException(e); } finally { producersDone.countDown(); } }); Suggestion importance[1-10]: 5 __ Why: If queue.put() throws after createBatch succeeds, the VectorSchemaRoot is leaked. While LinkedBlockingQueue.put() rarely throws in practice (only on interruption), closing the root on failure is correct resource management even in test code.	Low
General	Add safe cast with descriptive error message The cast to FlightServerChannel will throw a ClassCastException if the underlying TcpChannel is not a FlightServerChannel. Add a type check and throw a descriptive IllegalStateException to make failures easier to diagnose. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [154-156] +@Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel tcpChannel = getChannel(); + if (!(tcpChannel instanceof FlightServerChannel)) { + throw new IllegalStateException("Expected FlightServerChannel but got: " + tcpChannel.getClass().getName()); + } + return ((FlightServerChannel) tcpChannel).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The unchecked cast to FlightServerChannel could throw a ClassCastException with a cryptic message. Adding a type check improves diagnostics, though in practice FlightTransportChannel should always wrap a FlightServerChannel.	Low

---

Previous suggestions

Suggestions up to commit 2dce758

Category	Suggestion	Impact
Possible issue	Guard against empty field vectors when creating shared root If the producer root has no field vectors (empty schema), calling getFieldVectors().get(0) will throw an IndexOutOfBoundsException. Add a guard to fall back to the channel allocator when the vector list is empty. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [164-167] +List<FieldVector> fieldVectors = arrowResponse.getRoot().getFieldVectors(); +BufferAllocator rootAllocator = fieldVectors.isEmpty() + ? flightChannel.getAllocator() + : fieldVectors.get(0).getAllocator(); sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), - arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() + rootAllocator ); Suggestion importance[1-10]: 7 __ Why: If a producer root has an empty schema (no field vectors), getFieldVectors().get(0) will throw an IndexOutOfBoundsException. This is a real edge case that could cause a runtime crash, and the fix is straightforward.	Medium
General	Close Arrow root on enqueue failure to prevent memory leak If createBatch succeeds but queue.put throws (e.g., InterruptedException), the created VectorSchemaRoot is never closed, leaking Arrow memory. Close the root in the catch block if it was created but not enqueued. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [305-314] producers.submit(() -> { try { VectorSchemaRoot root = createBatch(allocator, batchIndex, request.rowsPerBatch); - queue.put(new TestArrowResponse(root)); + try { + queue.put(new TestArrowResponse(root)); + } catch (Exception e) { + root.close(); + throw e; + } } catch (Exception e) { throw new RuntimeException(e); } finally { producersDone.countDown(); } }); Suggestion importance[1-10]: 5 __ Why: If queue.put throws (e.g., InterruptedException), the created VectorSchemaRoot would leak Arrow memory. The fix properly closes the root in the failure path, preventing resource leaks in the test's parallel producer code.	Low
	Shut down executor forcefully on producer timeout producers.shutdown() is called after await, but if producers are still running (timeout expired), tasks may continue executing after the method returns. Call producers.shutdownNow() when the await times out to interrupt lingering threads and prevent resource leaks. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [326-327] -producersDone.await(30, TimeUnit.SECONDS); +boolean completed = producersDone.await(30, TimeUnit.SECONDS); producers.shutdown(); +if (!completed) { + producers.shutdownNow(); + throw new IOException("Timed out waiting for all producers to finish"); +} Suggestion importance[1-10]: 4 __ Why: The suggestion is about test code in NativeArrowTransportIT.java, not FlightOutboundHandler.java as stated. The producersDone.await and producers.shutdown() logic is in the IT test's handleStreamRequest. While the improvement is valid for robustness, this is test code and the impact is limited.	Low
	Add type safety check before casting channel The cast to FlightServerChannel will throw a ClassCastException if the underlying TcpChannel is not a FlightServerChannel. Add a type check and throw a descriptive IllegalStateException to make the failure easier to diagnose. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [153-156] @Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel tcpChannel = getChannel(); + if (!(tcpChannel instanceof FlightServerChannel)) { + throw new IllegalStateException("Expected FlightServerChannel but got: " + tcpChannel.getClass().getName()); + } + return ((FlightServerChannel) tcpChannel).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The cast to FlightServerChannel could throw a ClassCastException if the underlying channel is not a FlightServerChannel. Adding a type check improves error diagnostics, though in practice FlightTransportChannel is always constructed with a FlightServerChannel.	Low

Suggestions up to commit af5cc9c

Category	Suggestion	Impact
Possible issue	Ensure producer root is closed on transfer failure If transferTo throws an exception, arrowResponse.getRoot() will never be closed, leaking the producer's Arrow buffers. The close should be in a finally block or use a try-with-resources pattern to ensure cleanup on failure. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [169-171] -arrowResponse.transferTo(sharedRoot); -arrowResponse.getRoot().close(); // release producer's buffers — safe, they've been moved +try { + arrowResponse.transferTo(sharedRoot); +} finally { + arrowResponse.getRoot().close(); +} out = VectorStreamOutput.forNativeArrow(sharedRoot); Suggestion importance[1-10]: 8 __ Why: If transferTo throws an exception, arrowResponse.getRoot().close() is never called, leaking Arrow buffers. Wrapping the transfer in a try-finally ensures the producer root is always closed, preventing memory leaks.	Medium
	Guard against empty schema when accessing allocator If the producer root has no field vectors (empty schema), calling getFieldVectors().get(0) will throw an IndexOutOfBoundsException. The allocator should be retrieved from the root itself or from the channel, not from the first field vector. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [164-167] sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), - arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() + arrowResponse.getRoot().getAllocator() ); Suggestion importance[1-10]: 7 __ Why: If the producer root has no field vectors (empty schema), getFieldVectors().get(0) will throw an IndexOutOfBoundsException. Using arrowResponse.getRoot().getAllocator() is safer and more direct, though this edge case may be rare in practice.	Medium
	Prevent vector root leak on enqueue failure If createBatch succeeds but queue.put throws (e.g., InterruptedException), the created VectorSchemaRoot will be leaked because it is never closed. The root should be closed in the catch block if it was not successfully enqueued. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [305-314] producers.submit(() -> { + VectorSchemaRoot root = null; try { - VectorSchemaRoot root = createBatch(allocator, batchIndex, request.rowsPerBatch); + root = createBatch(allocator, batchIndex, request.rowsPerBatch); queue.put(new TestArrowResponse(root)); + root = null; // ownership transferred to queue } catch (Exception e) { + if (root != null) root.close(); throw new RuntimeException(e); } finally { producersDone.countDown(); } }); Suggestion importance[1-10]: 5 __ Why: If queue.put throws (e.g., InterruptedException), the created VectorSchemaRoot is leaked. While this is test code and LinkedBlockingQueue.put rarely throws in practice, the fix is still a good defensive practice.	Low
General	Add safe cast guard for allocator retrieval The cast to FlightServerChannel is unchecked and will throw a ClassCastException if the underlying TcpChannel is not a FlightServerChannel. A guard or a more descriptive error message should be added to make failures easier to diagnose. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [153-156] @Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel tcpChannel = getChannel(); + if (!(tcpChannel instanceof FlightServerChannel)) { + throw new IllegalStateException("Expected FlightServerChannel but got: " + tcpChannel.getClass().getName()); + } + return ((FlightServerChannel) tcpChannel).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The unchecked cast to FlightServerChannel could throw a ClassCastException with an unhelpful message. Adding a guard with a descriptive error improves debuggability, though by design FlightTransportChannel should always wrap a FlightServerChannel.	Low

Suggestions up to commit 7f0c5f0

Category	Suggestion	Impact
Possible issue	Persist newly created shared root back to channel When sharedRoot is newly created (first batch), it is never stored back into flightChannel, so subsequent batches will always enter the sharedRoot == null branch and create a new root each time, leaking the previously created one. The newly created sharedRoot must be set on flightChannel so it is reused across batches. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [156-164] if (task.response() instanceof ArrowBatchResponse arrowResponse) { // Native Arrow path: zero-copy transfer producer's vectors into shared root VectorSchemaRoot sharedRoot = flightChannel.getRoot(); if (sharedRoot == null) { sharedRoot = VectorSchemaRoot.create(arrowResponse.getRoot().getSchema(), flightChannel.getAllocator()); + flightChannel.setRoot(sharedRoot); } arrowResponse.transferTo(sharedRoot); arrowResponse.getRoot().close(); // release producer's buffers — safe, they've been moved out = VectorStreamOutput.forNativeArrow(sharedRoot); } Suggestion importance[1-10]: 7 __ Why: The suggestion identifies a potential issue where a newly created sharedRoot is not stored back into flightChannel, which could cause repeated root creation on subsequent batches. However, looking at the code more carefully, flightChannel.getRoot() may be a getter that returns a field managed elsewhere, and the PR may rely on FlightServerChannel managing this state internally. The suggestion's improved_code calls flightChannel.setRoot(sharedRoot) but there's no evidence this method exists in the codebase shown, making the fix potentially incorrect as written.	Medium
General	Validate schema compatibility before vector transfer There is no validation that sourceVectors and targetVectors have the same size before iterating. If the schemas differ (e.g., a producer root with a different schema is accidentally passed), the code will silently transfer only the first N vectors or throw an IndexOutOfBoundsException with no context. An explicit size check with a descriptive error would catch schema mismatches early. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java [97-104] void transferTo(VectorSchemaRoot target) { List<FieldVector> sourceVectors = producerRoot.getFieldVectors(); List<FieldVector> targetVectors = target.getFieldVectors(); + if (sourceVectors.size() != targetVectors.size()) { + throw new IllegalArgumentException( + "Schema mismatch: source has " + sourceVectors.size() + " vectors, target has " + targetVectors.size() + ); + } for (int i = 0; i < sourceVectors.size(); i++) { TransferPair transfer = sourceVectors.get(i).makeTransferPair(targetVectors.get(i)); transfer.transfer(); } target.setRowCount(producerRoot.getRowCount()); } Suggestion importance[1-10]: 5 __ Why: Adding a size check before the transfer loop is a reasonable defensive measure that would catch schema mismatches with a clear error message instead of an IndexOutOfBoundsException. The improved_code accurately reflects the suggested change and the fix is straightforward.	Low
	Ensure executor is shut down on exception The ExecutorService producers is never shut down if an exception is thrown before producers.shutdown() is reached (e.g., during the drain loop or producersDone.await()). This leaks threads. The executor should be shut down in a finally block to guarantee cleanup. plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java [305-314] producers.submit(() -> { try { VectorSchemaRoot root = createBatch(allocator, batchIndex, request.rowsPerBatch); queue.put(new TestArrowResponse(root)); } catch (Exception e) { throw new RuntimeException(e); } finally { producersDone.countDown(); } }); +// ... (after the loop) +try { + int sent = 0; + while (sent < request.batchCount) { + TestArrowResponse response = queue.poll(10, TimeUnit.SECONDS); + if (response == null) throw new IOException("Timed out waiting for producer"); + channel.sendResponseBatch(response); + sent++; + } + producersDone.await(30, TimeUnit.SECONDS); +} finally { + producers.shutdown(); +} Suggestion importance[1-10]: 4 __ Why: The suggestion correctly identifies that producers executor may not be shut down if an exception occurs before producers.shutdown(). However, the improved_code snippet is inconsistent with the existing_code — it mixes the lambda body with surrounding drain loop code, making it confusing. The actual fix should wrap the drain loop and producersDone.await() in a try-finally, not modify the lambda itself.	Low
	Guard against unexpected channel type in allocator lookup The cast to FlightServerChannel will throw a ClassCastException if getChannel() returns a different TcpChannel implementation. A guard or a more descriptive error message would prevent a confusing runtime failure. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [153-156] @Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel ch = getChannel(); + if (!(ch instanceof FlightServerChannel)) { + throw new IllegalStateException("Expected FlightServerChannel but got: " + ch.getClass().getName()); + } + return ((FlightServerChannel) ch).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The suggestion improves error clarity by replacing a raw ClassCastException with an IllegalStateException containing a descriptive message. This is a minor defensive improvement, but in practice FlightTransportChannel is always constructed with a FlightServerChannel, so the risk is low.	Low

Suggestions up to commit 416abe6

Category	Suggestion	Impact
Possible issue	Guard against empty field vectors when creating shared root When the producer root has no field vectors (e.g., an empty schema), arrowResponse.getRoot().getFieldVectors().get(0) will throw an IndexOutOfBoundsException. The allocator should be retrieved from the producer root itself via arrowResponse.getRoot().getVector(0).getAllocator() only when vectors exist, or fall back to the channel allocator. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [158-174] -if (task.response() instanceof ArrowBatchResponse arrowResponse) { - // Native Arrow path: zero-copy transfer producer's vectors into shared root - VectorSchemaRoot sharedRoot = flightChannel.getRoot(); - if (sharedRoot == null) { - sharedRoot = VectorSchemaRoot.create( - arrowResponse.getRoot().getSchema(), - arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() - ); - } - arrowResponse.transferTo(sharedRoot); - arrowResponse.getRoot().close(); // release producer's buffers — safe, they've been moved - out = VectorStreamOutput.forNativeArrow(sharedRoot); +if (sharedRoot == null) { + List<FieldVector> fieldVectors = arrowResponse.getRoot().getFieldVectors(); + BufferAllocator rootAllocator = fieldVectors.isEmpty() + ? flightChannel.getAllocator() + : fieldVectors.get(0).getAllocator(); + sharedRoot = VectorSchemaRoot.create( + arrowResponse.getRoot().getSchema(), + rootAllocator + ); +} Suggestion importance[1-10]: 6 __ Why: This is a valid defensive fix — getFieldVectors().get(0) will throw IndexOutOfBoundsException for schemas with no fields. The improved code correctly falls back to flightChannel.getAllocator() when no field vectors exist, preventing a potential runtime crash.	Low
	Ensure executor shutdown in finally block In the parallel production path in NativeArrowTransportIT, producers.shutdown() is called after producersDone.await(), but the executor is never awaited for termination after shutdown. If a producer thread throws an exception after producersDone.countDown(), it may be silently swallowed. Additionally, producers.shutdownNow() should be called in a finally block to prevent thread leaks if an exception occurs before producers.shutdown(). plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [341-367] -if (request.parallelism <= 1) { - // Serial production +ExecutorService producers = Executors.newFixedThreadPool(request.parallelism); +try { for (int batch = 0; batch < request.batchCount; batch++) { - channel.sendResponseBatch(new TestArrowResponse(createBatch(allocator, batch, request.rowsPerBatch))); + final int batchIndex = batch; + producers.submit(() -> { + try { + VectorSchemaRoot root = createBatch(allocator, batchIndex, request.rowsPerBatch); + queue.put(new TestArrowResponse(root)); + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + producersDone.countDown(); + } + }); } -} else { - ... + + int sent = 0; + while (sent < request.batchCount) { + TestArrowResponse response = queue.poll(10, TimeUnit.SECONDS); + if (response == null) throw new IOException("Timed out waiting for producer"); + channel.sendResponseBatch(response); + sent++; + } + producersDone.await(30, TimeUnit.SECONDS); - producers.shutdown(); +} finally { + producers.shutdownNow(); } Suggestion importance[1-10]: 4 __ Why: The suggestion addresses a real resource leak issue where producers executor may not be shut down if an exception occurs before producers.shutdown(). However, the existing_code references FlightOutboundHandler.java incorrectly — the parallel production code is in NativeArrowTransportIT.java. The improvement is valid but applies to test code, limiting its impact.	Low
General	Validate schema compatibility before vector transfer There is no validation that sourceVectors.size() equals targetVectors.size(). If the producer root and target root have different schemas (e.g., due to a bug or schema mismatch), the loop will silently transfer only the minimum number of vectors or throw an IndexOutOfBoundsException on the target side. A schema compatibility check should be added. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java [89-97] void transferTo(VectorSchemaRoot target) { List<FieldVector> sourceVectors = producerRoot.getFieldVectors(); List<FieldVector> targetVectors = target.getFieldVectors(); + if (sourceVectors.size() != targetVectors.size()) { + throw new IllegalArgumentException( + "Schema mismatch: source has " + sourceVectors.size() + " vectors, target has " + targetVectors.size() + ); + } for (int i = 0; i < sourceVectors.size(); i++) { TransferPair transfer = sourceVectors.get(i).makeTransferPair(targetVectors.get(i)); transfer.transfer(); } target.setRowCount(producerRoot.getRowCount()); } Suggestion importance[1-10]: 5 __ Why: Adding a size check before the transfer loop is a reasonable defensive measure that would produce a clearer error message on schema mismatch. However, in practice the shared root is always created from the producer's schema, making this scenario unlikely in normal usage.	Low
	Add safe cast guard for allocator retrieval The cast to FlightServerChannel is unchecked and will throw a ClassCastException at runtime if the underlying TcpChannel is not a FlightServerChannel. This can happen in test or non-flight contexts. A guard with a meaningful error message should be added. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [167-170] @Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel tcpChannel = getChannel(); + if (!(tcpChannel instanceof FlightServerChannel)) { + throw new IllegalStateException( + "Expected FlightServerChannel but got: " + (tcpChannel == null ? "null" : tcpChannel.getClass().getName()) + ); + } + return ((FlightServerChannel) tcpChannel).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The unchecked cast to FlightServerChannel could throw a ClassCastException in non-flight contexts. Adding a guard with a meaningful error message improves debuggability, though in practice FlightTransportChannel is only used with FlightServerChannel as the underlying channel.	Low

Suggestions up to commit d65e0f6

Category	Suggestion	Impact
Possible issue	Persist newly created shared root to channel The newly created sharedRoot is never stored back to flightChannel, so on subsequent batches flightChannel.getRoot() will still return null, causing a new VectorSchemaRoot to be created for every batch and leaking all but the last one. The created shared root must be set on flightChannel before use. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [161-174] if (sharedRoot == null) { sharedRoot = VectorSchemaRoot.create( arrowResponse.getRoot().getSchema(), arrowResponse.getRoot().getFieldVectors().get(0).getAllocator() ); + flightChannel.setRoot(sharedRoot); } arrowResponse.transferTo(sharedRoot); arrowResponse.getRoot().close(); Suggestion importance[1-10]: 8 __ Why: If flightChannel.setRoot() doesn't exist or the root isn't persisted, every batch would create a new VectorSchemaRoot causing memory leaks. However, the PR code may rely on flightChannel tracking the root internally via sendBatch. Without seeing FlightServerChannel.getRoot()/setRoot() implementation, this is a potentially critical bug worth flagging.	Medium
	Ensure executor shutdown on timeout exception The producers executor is never shut down if queue.poll times out and throws IOException, leaking threads. Additionally, producers.shutdown() is called after producersDone.await() but the executor should be shut down (and awaited) even in the exception path. Wrap the executor lifecycle in a try-finally block. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java [358-367] int sent = 0; -while (sent < request.batchCount) { - TestArrowResponse response = queue.poll(10, TimeUnit.SECONDS); - if (response == null) throw new IOException("Timed out waiting for producer"); - channel.sendResponseBatch(response); - sent++; +try { + while (sent < request.batchCount) { + TestArrowResponse response = queue.poll(10, TimeUnit.SECONDS); + if (response == null) { + producers.shutdownNow(); + throw new IOException("Timed out waiting for producer"); + } + channel.sendResponseBatch(response); + sent++; + } + producersDone.await(30, TimeUnit.SECONDS); +} finally { + producers.shutdown(); } -producersDone.await(30, TimeUnit.SECONDS); -producers.shutdown(); - Suggestion importance[1-10]: 3 __ Why: The code in question is in the test file NativeArrowTransportIT.java, not in FlightOutboundHandler.java. The suggestion's relevant_file is wrong. Additionally, the existing_code snippet is from the IT test class, not the handler. The suggestion has merit for the test code but is misattributed.	Low
General	Validate schema compatibility before vector transfer There is no validation that sourceVectors and targetVectors have the same size before indexing into targetVectors. If the schemas differ (e.g., due to a schema mismatch between producer and shared root), this will throw an IndexOutOfBoundsException with no useful error message. A guard check would make failures easier to diagnose. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java [89-97] void transferTo(VectorSchemaRoot target) { List<FieldVector> sourceVectors = producerRoot.getFieldVectors(); List<FieldVector> targetVectors = target.getFieldVectors(); + if (sourceVectors.size() != targetVectors.size()) { + throw new IllegalArgumentException( + "Schema mismatch: source has " + sourceVectors.size() + " vectors, target has " + targetVectors.size() + ); + } for (int i = 0; i < sourceVectors.size(); i++) { TransferPair transfer = sourceVectors.get(i).makeTransferPair(targetVectors.get(i)); transfer.transfer(); } target.setRowCount(producerRoot.getRowCount()); } Suggestion importance[1-10]: 4 __ Why: Adding a size check before indexing into targetVectors is a reasonable defensive measure that improves error diagnostics, though schema mismatches would be a programming error caught early in testing.	Low
	Guard unsafe cast with descriptive error This cast to FlightServerChannel will throw a ClassCastException at runtime if the underlying TcpChannel is not a FlightServerChannel (e.g., in tests or non-Flight deployments). A defensive check with a clear error message would prevent a confusing failure. plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java [167-170] @Override public BufferAllocator getAllocator() { - return ((FlightServerChannel) getChannel()).getAllocator(); + TcpChannel tcpChannel = getChannel(); + if (!(tcpChannel instanceof FlightServerChannel)) { + throw new IllegalStateException("Expected FlightServerChannel but got: " + tcpChannel.getClass().getName()); + } + return ((FlightServerChannel) tcpChannel).getAllocator(); } Suggestion importance[1-10]: 4 __ Why: The cast to FlightServerChannel is reasonable given the class invariant, but adding a guard with a descriptive error message improves debuggability. The existing test testGetAllocator already uses a FlightServerChannel mock, suggesting this path is expected to always hold.	Low

[github-actions]

Persistent review updated to latest commit 7d506e8

[github-actions]

Persistent review updated to latest commit ac83ec6

[github-actions]

❌ Gradle check result for ac83ec6: 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 e3d1a5e

[github-actions]

Persistent review updated to latest commit ac213cc

[github-actions]

Persistent review updated to latest commit 9280e70

[github-actions]

Persistent review updated to latest commit 3118626

[github-actions]

Persistent review updated to latest commit 73d7245

[github-actions]

Persistent review updated to latest commit d76afed

[github-actions]

❌ Gradle check result for d76afed: 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 52d8ab3

[github-actions]

Persistent review updated to latest commit c9e78c5

[github-actions]

❌ Gradle check result for c9e78c5: null

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 c037252

[github-actions]

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