From 2dce75854c562207ef506dd6ea62d67fbca4727c Mon Sep 17 00:00:00 2001 From: Rishabh Maurya Date: Tue, 21 Apr 2026 14:54:32 -0700 Subject: [PATCH] Native Arrow transport path with zero-copy transfer Signed-off-by: Rishabh Maurya --- .../docs/native-arrow-transport-design.md | 117 +++++ .../arrow/flight/NativeArrowTransportIT.java | 409 ++++++++++++++++++ .../flight/transport/ArrowBatchResponse.java | 112 +++++ .../flight/transport/ArrowFlightChannel.java | 64 +++ .../transport/FlightOutboundHandler.java | 23 +- .../flight/transport/FlightServerChannel.java | 2 +- .../transport/FlightTransportChannel.java | 8 +- .../flight/transport/VectorStreamInput.java | 23 +- .../flight/transport/VectorStreamOutput.java | 181 +++++--- .../transport/ArrowBatchResponseTests.java | 129 ++++++ .../transport/ArrowFlightChannelTests.java | 56 +++ .../ArrowStreamSerializationTests.java | 2 +- .../transport/FlightOutboundHandlerTests.java | 177 ++++++++ .../FlightTransportChannelTests.java | 22 + .../transport/VectorStreamInputTests.java | 70 +++ .../transport/VectorStreamOutputTests.java | 243 +++++++++++ .../stream-transport-example/build.gradle | 2 + .../NativeArrowStreamTransportExampleIT.java | 180 ++++++++ .../stream/NativeArrowStreamDataAction.java | 20 + .../stream/NativeArrowStreamDataRequest.java | 52 +++ .../stream/NativeArrowStreamDataResponse.java | 37 ++ .../stream/StreamTransportExamplePlugin.java | 12 +- .../TransportNativeArrowStreamDataAction.java | 110 +++++ 23 files changed, 1979 insertions(+), 72 deletions(-) create mode 100644 plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md create mode 100644 plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java create mode 100644 plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java create mode 100644 plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowFlightChannel.java create mode 100644 plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java create mode 100644 plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowFlightChannelTests.java create mode 100644 plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamInputTests.java create mode 100644 plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamOutputTests.java create mode 100644 plugins/examples/stream-transport-example/src/internalClusterTest/java/org/opensearch/example/stream/NativeArrowStreamTransportExampleIT.java create mode 100644 plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataAction.java create mode 100644 plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataRequest.java create mode 100644 plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataResponse.java create mode 100644 plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/TransportNativeArrowStreamDataAction.java diff --git a/plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md b/plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md new file mode 100644 index 0000000000000..a91717e6b7322 --- /dev/null +++ b/plugins/arrow-flight-rpc/docs/native-arrow-transport-design.md @@ -0,0 +1,117 @@ +# Native Arrow Transport Path + +## Overview + +The Arrow Flight transport supports a native Arrow path where typed `VectorSchemaRoot` data +flows directly over Flight without byte serialization. This is for APIs that produce +Arrow-columnar data natively (e.g., query engines like DataFusion). + +The existing byte-serialized path (`writeTo`/`read` via `StreamOutput`/`StreamInput`) is unchanged. + +## Quick Start + +### 1. Define your response + +Extend `ArrowBatchResponse`. No `writeTo`/`read` override needed — the framework handles it. + +```java +public class MyQueryResponse extends ArrowBatchResponse { + public MyQueryResponse(VectorSchemaRoot root) { super(root); } + public MyQueryResponse(StreamInput in) throws IOException { super(in); } +} +``` + +### 2. Server-side handler — produce Arrow data + +```java +void handleRequest(MyRequest request, TransportChannel channel, Task task) throws IOException { + // Get the channel's allocator. Use this directly for producer roots + // (not a child allocator) to avoid Arrow's cross-allocator transfer bug + // with foreign-backed buffers from C data import. + BufferAllocator allocator = ArrowFlightChannel.from(channel).getAllocator(); + + Schema schema = new Schema(List.of( + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("score", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), null) + )); + + try { + for (int i = 0; i < batchCount; i++) { + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + // populate vectors... + channel.sendResponseBatch(new MyQueryResponse(root)); + // root is now owned by the framework — don't reuse or close it + } + // Cleanup callback runs on executor after all batches are flushed. + channel.completeStream(); + } catch (Exception e) { + channel.sendResponse(e); + } +} +``` + +### 3. Client-side handler — consume Arrow data + +```java +class MyQueryHandler implements StreamTransportResponseHandler { + + public MyQueryResponse read(StreamInput in) throws IOException { + return new MyQueryResponse(in); + } + + public void handleStreamResponse(StreamTransportResponse stream) { + MyQueryResponse response; + while ((response = stream.nextResponse()) != null) { + VectorSchemaRoot root = response.getRoot(); + VarCharVector names = (VarCharVector) root.getVector("name"); + Float8Vector scores = (Float8Vector) root.getVector("score"); + // process typed vectors... + } + stream.close(); + } + + public void handleException(TransportException exp) { /* handle error */ } + public String executor() { return ThreadPool.Names.GENERIC; } +} +``` + +## Allocator Management + +The allocator used for producer roots must be **long-lived** — it must outlive the gRPC +stream. This is because gRPC's zero-copy write path (`ArrowBufRetainingCompositeByteBuf`) +retains ArrowBuf references beyond `putNext()` and `completed()`, releasing them +asynchronously on the Netty event loop. Closing the allocator while gRPC still holds +these references causes memory accounting errors. + +Use `ArrowFlightChannel.from(channel).getAllocator()` to get the channel's allocator, +or use your own long-lived application allocator. Do not create and close a child +allocator per request. + +### Important: C Data Import and allocator choice + +If your producer imports data via Arrow's C Data Interface (`Data.importIntoVector`, +`Data.importIntoVectorSchemaRoot`), the imported buffers are foreign-backed. Arrow Java +has a bug where cross-allocator `transferOwnership` of foreign-backed buffers doesn't +properly release the internal `ArrowArray` C struct buffer (128 bytes per import call), +causing a memory leak in the source allocator. + +The framework creates the shared Flight root from the **producer's allocator** (the +allocator of the first batch's vectors), ensuring same-allocator transfer which avoids +this bug. All subsequent batches should use the same allocator. + +## Ownership Contract + +| Resource | Created by | Closed by | +|----------|-----------|-----------| +| Channel allocator | Framework | Framework (on channel close) | +| Producer root (per batch) | Producer | Framework (after zero-copy transfer on executor) | +| Shared Flight root | Framework | Framework (on channel close) | + +After calling `sendResponseBatch(response)`, the framework owns the response's root. +Do not reuse or close it — the framework transfers its buffers and closes it on the executor. + +## Pipelining + +Batches can be produced in parallel. Each batch must have its own `VectorSchemaRoot` +(created from the channel's allocator). The framework serializes the transfer and send +on the executor thread. The producer can queue batches without waiting for each to flush. diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java new file mode 100644 index 0000000000000..849c63a594e3b --- /dev/null +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeArrowTransportIT.java @@ -0,0 +1,409 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.ActionType; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.TransportAction; +import org.opensearch.arrow.flight.transport.ArrowBatchResponse; +import org.opensearch.arrow.flight.transport.ArrowFlightChannel; +import org.opensearch.arrow.flight.transport.FlightStreamPlugin; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.plugins.ActionPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.StreamTransportResponseHandler; +import org.opensearch.transport.StreamTransportService; +import org.opensearch.transport.TransportChannel; +import org.opensearch.transport.TransportException; +import org.opensearch.transport.TransportRequestOptions; +import org.opensearch.transport.stream.StreamErrorCode; +import org.opensearch.transport.stream.StreamException; +import org.opensearch.transport.stream.StreamTransportResponse; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.opensearch.common.util.FeatureFlags.STREAM_TRANSPORT; + +/** + * Integration test for the native Arrow transport path. + * Tests serial and parallel batch production with zero-copy transfer, + * verifying typed data integrity end-to-end. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, minNumDataNodes = 2, maxNumDataNodes = 2) +public class NativeArrowTransportIT extends OpenSearchIntegTestCase { + + private static final Schema TEST_SCHEMA = new Schema( + List.of( + new Field("batch_id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("value", FieldType.nullable(new ArrowType.Int(32, true)), null) + ) + ); + + @Override + public void setUp() throws Exception { + super.setUp(); + internalCluster().ensureAtLeastNumDataNodes(2); + } + + @Override + protected Collection> nodePlugins() { + return List.of(NativeArrowTestPlugin.class, FlightStreamPlugin.class); + } + + // ── Tests ── + + @LockFeatureFlag(STREAM_TRANSPORT) + public void testSingleBatchNativeArrow() throws Exception { + { + DiscoveryNode node = getClusterState().nodes().iterator().next(); + List batches = sendAndReceive(node, 1, 3, 1); + + assertEquals("Should receive 1 batch", 1, batches.size()); + ReceivedBatch batch = batches.get(0); + assertEquals(3, batch.rowCount); + assertEquals(0, batch.batchId); + assertBatchIntegrity(batch); + } + } + + @LockFeatureFlag(STREAM_TRANSPORT) + public void testMultipleBatchesSerialNativeArrow() throws Exception { + { + DiscoveryNode node = getClusterState().nodes().iterator().next(); + List batches = sendAndReceive(node, 5, 4, 1); + + assertEquals("Should receive 5 batches", 5, batches.size()); + Set batchIds = new HashSet<>(); + for (ReceivedBatch batch : batches) { + assertEquals(4, batch.rowCount); + assertBatchIntegrity(batch); + batchIds.add(batch.batchId); + } + assertEquals("All batch IDs should be unique", 5, batchIds.size()); + } + } + + @LockFeatureFlag(STREAM_TRANSPORT) + public void testParallelBatchProduction() throws Exception { + // 100 batches, 10 rows each, produced by 5 parallel threads. + // Each batch has a unique batch_id. Verifies: + // - All 100 batches arrive + // - No batch is lost or duplicated + // - Each batch's data is internally consistent (batch_id matches across all rows) + { + DiscoveryNode node = getClusterState().nodes().iterator().next(); + List batches = sendAndReceive(node, 100, 10, 5); + + assertEquals("Should receive 100 batches", 100, batches.size()); + + Set batchIds = new HashSet<>(); + for (ReceivedBatch batch : batches) { + assertEquals("Each batch should have 10 rows", 10, batch.rowCount); + assertBatchIntegrity(batch); + assertTrue("Batch ID should be in range [0, 100)", batch.batchId >= 0 && batch.batchId < 100); + batchIds.add(batch.batchId); + } + assertEquals("All 100 batch IDs must be present (no lost/duplicated batches)", 100, batchIds.size()); + } + } + + // ── Helpers ── + + private List sendAndReceive(DiscoveryNode node, int batchCount, int rowsPerBatch, int parallelism) throws Exception { + StreamTransportService sts = internalCluster().getInstance(StreamTransportService.class); + List batches = new ArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + + sts.sendRequest( + node, + TestArrowAction.NAME, + new TestArrowRequest(batchCount, rowsPerBatch, parallelism), + TransportRequestOptions.builder().withType(TransportRequestOptions.Type.STREAM).build(), + new TestArrowResponseHandler(batches, latch, failure) + ); + + assertTrue("Stream should complete within 30s", latch.await(30, TimeUnit.SECONDS)); + assertNull("No exception expected: " + failure.get(), failure.get()); + return batches; + } + + /** Verifies that all rows in a batch have the same batch_id and consistent name/value. */ + private void assertBatchIntegrity(ReceivedBatch batch) { + for (int i = 0; i < batch.rowCount; i++) { + assertEquals("batch_id must be consistent across all rows", batch.batchId, batch.batchIds.get(i).intValue()); + // name = "row-{batchId}-{rowIndex}" + String expectedName = "row-" + batch.batchId + "-" + i; + assertEquals("Name must match expected pattern", expectedName, batch.names.get(i)); + // value = batchId * 1000 + rowIndex + assertEquals("Value must match expected pattern", batch.batchId * 1000 + i, batch.values.get(i).intValue()); + } + } + + /** Deep-copies data from a VectorSchemaRoot. */ + static class ReceivedBatch { + final int rowCount; + final int batchId; + final List batchIds; + final List names; + final List values; + + ReceivedBatch(VectorSchemaRoot root) { + this.rowCount = root.getRowCount(); + IntVector batchIdVector = (IntVector) root.getVector("batch_id"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + IntVector valueVector = (IntVector) root.getVector("value"); + this.batchIds = new ArrayList<>(); + this.names = new ArrayList<>(); + this.values = new ArrayList<>(); + for (int i = 0; i < rowCount; i++) { + batchIds.add(batchIdVector.get(i)); + names.add(new String(nameVector.get(i), StandardCharsets.UTF_8)); + values.add(valueVector.get(i)); + } + this.batchId = rowCount > 0 ? batchIds.get(0) : -1; + } + } + + // ── Inner classes: Action, Request, Response, Handler, Plugin ── + + public static class TestArrowResponse extends ArrowBatchResponse { + public TestArrowResponse(VectorSchemaRoot root) { + super(root); + } + + public TestArrowResponse(StreamInput in) throws IOException { + super(in); + } + } + + public static class TestArrowRequest extends ActionRequest { + private final int batchCount; + private final int rowsPerBatch; + private final int parallelism; + + public TestArrowRequest(int batchCount, int rowsPerBatch, int parallelism) { + this.batchCount = batchCount; + this.rowsPerBatch = rowsPerBatch; + this.parallelism = parallelism; + } + + public TestArrowRequest(StreamInput in) throws IOException { + super(in); + this.batchCount = in.readInt(); + this.rowsPerBatch = in.readInt(); + this.parallelism = in.readInt(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeInt(batchCount); + out.writeInt(rowsPerBatch); + out.writeInt(parallelism); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + } + + public static class TestArrowAction extends ActionType { + public static final TestArrowAction INSTANCE = new TestArrowAction(); + public static final String NAME = "cluster:internal/test/native_arrow"; + + private TestArrowAction() { + super(NAME, TestArrowResponse::new); + } + } + + /** + * Server-side handler. Produces batches using a thread pool. + * Each producer thread creates a batch with its own child allocator, + * puts it on a queue. The main thread drains the queue and sends + * batches via sendResponseBatch(). The framework does zero-copy transfer + * on the executor thread. + */ + public static class TransportTestArrowAction extends TransportAction { + + @Inject + public TransportTestArrowAction(StreamTransportService streamTransportService, ActionFilters actionFilters) { + super(TestArrowAction.NAME, actionFilters, streamTransportService.getTaskManager()); + streamTransportService.registerRequestHandler( + TestArrowAction.NAME, + ThreadPool.Names.GENERIC, + TestArrowRequest::new, + this::handleStreamRequest + ); + } + + @Override + protected void doExecute(Task task, TestArrowRequest request, ActionListener listener) { + listener.onFailure(new UnsupportedOperationException("Use StreamTransportService")); + } + + private void handleStreamRequest(TestArrowRequest request, TransportChannel channel, Task task) throws IOException { + BufferAllocator allocator = ArrowFlightChannel.from(channel).getAllocator(); + + try { + if (request.parallelism <= 1) { + // Serial production + for (int batch = 0; batch < request.batchCount; batch++) { + channel.sendResponseBatch(new TestArrowResponse(createBatch(allocator, batch, request.rowsPerBatch))); + } + } else { + // Parallel production: N threads produce batches into a queue, + // main thread drains and sends serially. + BlockingQueue queue = new LinkedBlockingQueue<>(); + CountDownLatch producersDone = new CountDownLatch(request.batchCount); + ExecutorService producers = Executors.newFixedThreadPool(request.parallelism); + + for (int batch = 0; batch < request.batchCount; batch++) { + 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(); + } + }); + } + + // Drain: send batches as they become available + 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(); + } + channel.completeStream(); + } catch (StreamException e) { + if (e.getErrorCode() != StreamErrorCode.CANCELLED) channel.sendResponse(e); + } catch (Exception e) { + channel.sendResponse(e); + } + } + + private VectorSchemaRoot createBatch(BufferAllocator allocator, int batchIndex, int rowCount) { + VectorSchemaRoot root = VectorSchemaRoot.create(TEST_SCHEMA, allocator); + + IntVector batchIdVector = (IntVector) root.getVector("batch_id"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + IntVector valueVector = (IntVector) root.getVector("value"); + batchIdVector.allocateNew(); + nameVector.allocateNew(); + valueVector.allocateNew(); + + for (int i = 0; i < rowCount; i++) { + batchIdVector.setSafe(i, batchIndex); + nameVector.setSafe(i, ("row-" + batchIndex + "-" + i).getBytes(StandardCharsets.UTF_8)); + valueVector.setSafe(i, batchIndex * 1000 + i); + } + root.setRowCount(rowCount); + return root; + } + } + + static class TestArrowResponseHandler implements StreamTransportResponseHandler { + private final List batches; + private final CountDownLatch latch; + private final AtomicReference failure; + + TestArrowResponseHandler(List batches, CountDownLatch latch, AtomicReference failure) { + this.batches = batches; + this.latch = latch; + this.failure = failure; + } + + @Override + public void handleStreamResponse(StreamTransportResponse streamResponse) { + try { + TestArrowResponse response; + while ((response = streamResponse.nextResponse()) != null) { + batches.add(new ReceivedBatch(response.getRoot())); + } + streamResponse.close(); + latch.countDown(); + } catch (Exception e) { + failure.set(e); + streamResponse.cancel("Test error", e); + latch.countDown(); + } + } + + @Override + public void handleException(TransportException exp) { + failure.set(exp); + latch.countDown(); + } + + @Override + public String executor() { + return ThreadPool.Names.GENERIC; + } + + @Override + public TestArrowResponse read(StreamInput in) throws IOException { + return new TestArrowResponse(in); + } + } + + public static class NativeArrowTestPlugin extends Plugin implements ActionPlugin { + public NativeArrowTestPlugin() {} + + @Override + public List> getActions() { + return List.of(new ActionHandler<>(TestArrowAction.INSTANCE, TransportTestArrowAction.class)); + } + } +} diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java new file mode 100644 index 0000000000000..2e3c0939f0467 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowBatchResponse.java @@ -0,0 +1,112 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.util.TransferPair; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; +import java.util.List; + +/** + * Base class for transport responses carrying native Arrow data. + * + *

The producer creates vectors using the channel's allocator and populates them freely + * on any thread. When the executor processes this batch, it does a zero-copy transfer + * of the producer's buffers into the channel's shared root — no memcpy, no serialization. + * After transfer, the framework closes the producer's root, releasing its buffers back + * to the allocator. + * + *

Allocator guidelines: The allocator used for producer roots must outlive the + * gRPC stream — do not create and close a child allocator per request. gRPC's zero-copy + * write path retains buffer references beyond stream completion, and closing the allocator + * while gRPC still holds these references causes memory accounting errors. Use either the + * channel allocator (via {@code ArrowFlightChannel.from(channel).getAllocator()}) or a + * long-lived application allocator. The framework creates the shared root from the + * producer's allocator to ensure same-allocator transfer, which avoids an Arrow bug with + * cross-allocator transfer of foreign-backed buffers from C data import. + * + *

Usage (send side): + * 

{@code
+ * BufferAllocator allocator = ArrowFlightChannel.from(channel).getAllocator();
+ * VectorSchemaRoot producerRoot = VectorSchemaRoot.create(schema, allocator);
+ * // populate producerRoot on any thread...
+ * channel.sendResponseBatch(new MyResponse(producerRoot));
+ * // producerRoot is now owned by the framework — don't reuse or close it
+ * }
+ * + *

Usage (receive side): + * 

{@code
+ * public class MyResponse extends ArrowBatchResponse {
+ *     public MyResponse(VectorSchemaRoot root) { super(root); }
+ *     public MyResponse(StreamInput in) throws IOException { super(in); }
+ * }
+ * }
+ * + * @opensearch.experimental + */ +@ExperimentalApi +public abstract class ArrowBatchResponse extends ActionResponse { + + private final VectorSchemaRoot producerRoot; + + /** + * Creates a response with the given producer root (send side). + * @param producerRoot the root populated by the producer + */ + protected ArrowBatchResponse(VectorSchemaRoot producerRoot) { + this.producerRoot = producerRoot; + } + + /** + * Deserializes a response from a StreamInput (receive side). + * @param in the stream input containing the Arrow root + * @throws IOException if deserialization fails + */ + protected ArrowBatchResponse(StreamInput in) throws IOException { + super(in); + this.producerRoot = ((VectorStreamInput) in).getRoot(); + } + + /** + * Returns the producer's root. On the send side, this is the root populated + * by the producer. On the receive side, this is the root from the Flight stream. + */ + public VectorSchemaRoot getRoot() { + return producerRoot; + } + + /** + * Zero-copy transfers the producer's vectors into the target root. + * Called by the framework on the executor thread before {@code putNext()}. + * After transfer, the producer's buffers are moved to the target — the producer + * root becomes empty. + * + * @param target the channel's shared root (bound to the Flight stream via start()) + */ + void transferTo(VectorSchemaRoot target) { + List sourceVectors = producerRoot.getFieldVectors(); + List targetVectors = target.getFieldVectors(); + for (int i = 0; i < sourceVectors.size(); i++) { + TransferPair transfer = sourceVectors.get(i).makeTransferPair(targetVectors.get(i)); + transfer.transfer(); + } + target.setRowCount(producerRoot.getRowCount()); + } + + @Override + public final void writeTo(StreamOutput out) throws IOException { + // no-op: the framework handles transfer via transferTo() + } +} diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowFlightChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowFlightChannel.java new file mode 100644 index 0000000000000..681a373094e22 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowFlightChannel.java @@ -0,0 +1,64 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.transport.BaseTcpTransportChannel; +import org.opensearch.transport.TaskTransportChannel; +import org.opensearch.transport.TcpChannel; +import org.opensearch.transport.TransportChannel; + +/** + * Provides access to the Arrow {@link BufferAllocator} for request handlers + * that produce native Arrow responses. + * + *

Use {@link #from(TransportChannel)} to obtain an instance from any + * {@code TransportChannel}, regardless of wrapper layers. + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface ArrowFlightChannel { + + /** + * Returns the Arrow allocator for this channel. + */ + BufferAllocator getAllocator(); + + /** + * Unwraps the given {@link TransportChannel} to find the {@link ArrowFlightChannel}. + * Walks through {@link TaskTransportChannel} and {@link BaseTcpTransportChannel} + * wrapper layers to find the underlying channel. + * + * @param channel the transport channel (may be wrapped) + * @return the ArrowFlightChannel + * @throws IllegalArgumentException if the channel is not backed by an ArrowFlightChannel + */ + static ArrowFlightChannel from(TransportChannel channel) { + TransportChannel current = channel; + while (current != null) { + if (current instanceof ArrowFlightChannel afc) { + return afc; + } + if (current instanceof TaskTransportChannel ttc) { + current = ttc.getChannel(); + } else if (current instanceof BaseTcpTransportChannel btc) { + TcpChannel tcpChannel = btc.getChannel(); + if (tcpChannel instanceof ArrowFlightChannel afc) { + return afc; + } + break; + } else { + break; + } + } + throw new IllegalArgumentException("Channel is not backed by an ArrowFlightChannel: " + channel.getClass().getName()); + } +} diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java index 03da4ff3a14b5..eb0f90b83c675 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java @@ -17,6 +17,7 @@ package org.opensearch.arrow.flight.transport; import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.vector.VectorSchemaRoot; import org.opensearch.Version; import org.opensearch.cluster.node.DiscoveryNode; import org.opensearch.common.io.stream.BytesStreamOutput; @@ -151,8 +152,28 @@ private void processBatchTask(BatchTask task) { } try { - try (VectorStreamOutput out = new VectorStreamOutput(flightChannel.getAllocator(), flightChannel.getRoot())) { + VectorStreamOutput out; + 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) { + // 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()); } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java index a923c7843d987..637badab4b9fa 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java @@ -39,7 +39,7 @@ * TcpChannel implementation for Arrow Flight. It is created per call in ArrowFlightProducer. * This implementation is not thread safe; consumer must ensure to invoke sendBatch serially and call completeStream() at the end */ -class FlightServerChannel implements TcpChannel { +class FlightServerChannel implements TcpChannel, ArrowFlightChannel { private static final String PROFILE_NAME = "flight"; private final Logger logger = LogManager.getLogger(FlightServerChannel.class); diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java index cba53819ff3c1..57f595c4135f9 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportChannel.java @@ -8,6 +8,7 @@ package org.opensearch.arrow.flight.transport; +import org.apache.arrow.memory.BufferAllocator; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.Version; @@ -30,7 +31,7 @@ * The underlying TcpChannel is closed when release is called. * @opensearch.internal */ -class FlightTransportChannel extends TcpTransportChannel { +class FlightTransportChannel extends TcpTransportChannel implements ArrowFlightChannel { private static final Logger logger = LogManager.getLogger(FlightTransportChannel.class); private final AtomicBoolean streamOpen = new AtomicBoolean(true); @@ -148,4 +149,9 @@ public String getChannelType() { public void releaseChannel(boolean isExceptionResponse) { release(isExceptionResponse); } + + @Override + public BufferAllocator getAllocator() { + return ((FlightServerChannel) getChannel()).getAllocator(); + } } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamInput.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamInput.java index 272ead2abaf15..6951805560572 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamInput.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamInput.java @@ -19,18 +19,37 @@ import java.io.IOException; import java.nio.ByteBuffer; +/** + * A {@link StreamInput} backed by a {@link VectorSchemaRoot} from the Flight transport. + * + * @opensearch.internal + */ class VectorStreamInput extends StreamInput { private final VarBinaryVector vector; + private final VectorSchemaRoot root; private final NamedWriteableRegistry registry; private int row = 0; private ByteBuffer buffer = null; + /** + * Creates a new VectorStreamInput. + * @param root the Arrow root containing the data + * @param registry the named writeable registry + */ public VectorStreamInput(VectorSchemaRoot root, NamedWriteableRegistry registry) { + this.root = root; vector = (VarBinaryVector) root.getVector("0"); this.registry = registry; } + /** + * Returns the underlying {@link VectorSchemaRoot}. + */ + public VectorSchemaRoot getRoot() { + return root; + } + @Override public byte readByte() throws IOException { // Check if buffer has remaining bytes @@ -112,7 +131,9 @@ public NamedWriteableRegistry namedWriteableRegistry() { @Override public void close() throws IOException { - vector.close(); + if (vector != null) { + vector.close(); + } } @Override diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java index 9755c4c77dbda..cffb35026c1be 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java @@ -19,82 +19,145 @@ import java.io.IOException; import java.util.List; -class VectorStreamOutput extends StreamOutput { +/** + * A {@link StreamOutput} that produces a {@link VectorSchemaRoot} for the Flight transport. + * + *

Two factory methods create the appropriate instance: + *

    + *
  • {@link #create(BufferAllocator, VectorSchemaRoot)} — byte serialization. + * {@code writeTo()} writes bytes into a VarBinary vector.
    • + *
  • {@link #forNativeArrow(VectorSchemaRoot)} — native Arrow. + * The root is already populated; writes are no-ops.
    • + *
+ * + *

The framework selects the right factory based on whether the response + * is an {@link ArrowBatchResponse}. + * + * @opensearch.internal + */ +abstract class VectorStreamOutput extends StreamOutput { - private int row = 0; - private final VarBinaryVector vector; - private VectorSchemaRoot root; - private final byte[] tempBuffer = new byte[8192]; - private int tempBufferPos = 0; + /** Creates a VectorStreamOutput. */ + protected VectorStreamOutput() {} - public VectorStreamOutput(BufferAllocator allocator, VectorSchemaRoot root) { - if (root != null) { - vector = (VarBinaryVector) root.getVector(0); - this.root = root; - } else { - Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); - vector = (VarBinaryVector) field.createVector(allocator); - // Pre-allocate with reasonable capacity to avoid repeated allocations - vector.setInitialCapacity(16); - vector.allocateNew(); - } + /** + * Creates a byte-serialization output. + */ + static VectorStreamOutput create(BufferAllocator allocator, VectorSchemaRoot existingRoot) { + return new ByteSerialized(allocator, existingRoot); } - @Override - public void writeByte(byte b) { - // Buffer small writes to reduce vector operations - if (tempBufferPos >= tempBuffer.length) { - flushTempBuffer(); - } - tempBuffer[tempBufferPos++] = b; + /** + * Creates a native Arrow output. The root is already populated. + */ + static VectorStreamOutput forNativeArrow(VectorSchemaRoot root) { + return new NativeArrow(root); } - @Override - public void writeBytes(byte[] b, int offset, int length) { - if (length == 0) { - return; + /** + * Returns the {@link VectorSchemaRoot} to send over Flight. + */ + public abstract VectorSchemaRoot getRoot(); + + // ── Byte serialization ── + + static final class ByteSerialized extends VectorStreamOutput { + private int row = 0; + private final VarBinaryVector vector; + private VectorSchemaRoot root; + private final byte[] tempBuffer = new byte[8192]; + private int tempBufferPos = 0; + + ByteSerialized(BufferAllocator allocator, VectorSchemaRoot existingRoot) { + if (existingRoot != null) { + vector = (VarBinaryVector) existingRoot.getVector(0); + this.root = existingRoot; + } else { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + vector = (VarBinaryVector) field.createVector(allocator); + vector.setInitialCapacity(16); + vector.allocateNew(); + } } - if (b.length < (offset + length)) { - throw new IllegalArgumentException("Illegal offset " + offset + "/length " + length + " for byte[] of length " + b.length); + + @Override + public void writeByte(byte b) { + if (tempBufferPos >= tempBuffer.length) flushTempBuffer(); + tempBuffer[tempBufferPos++] = b; } - if (tempBufferPos > 0) { - flushTempBuffer(); + + @Override + public void writeBytes(byte[] b, int offset, int length) { + if (length == 0) return; + if (b.length < (offset + length)) { + throw new IllegalArgumentException("Illegal offset " + offset + "/length " + length + " for byte[] of length " + b.length); + } + if (tempBufferPos > 0) flushTempBuffer(); + vector.setSafe(row++, b, offset, length); } - vector.setSafe(row++, b, offset, length); - } - private void flushTempBuffer() { - if (tempBufferPos > 0) { - vector.setSafe(row++, tempBuffer, 0, tempBufferPos); - tempBufferPos = 0; + private void flushTempBuffer() { + if (tempBufferPos > 0) { + vector.setSafe(row++, tempBuffer, 0, tempBufferPos); + tempBufferPos = 0; + } } - } - @Override - public void flush() { + @Override + public void flush() {} - } + @Override + public void close() throws IOException { + row = 0; + vector.close(); + } - @Override - public void close() throws IOException { - row = 0; - vector.close(); - } + @Override + public void reset() { + row = 0; + tempBufferPos = 0; + vector.clear(); + } - @Override - public void reset() { - row = 0; - tempBufferPos = 0; - vector.clear(); + @Override + public VectorSchemaRoot getRoot() { + flushTempBuffer(); + vector.setValueCount(row); + if (root == null) { + root = new VectorSchemaRoot(List.of(vector)); + } + root.setRowCount(row); + return root; + } } - public VectorSchemaRoot getRoot() { - flushTempBuffer(); - vector.setValueCount(row); - if (root == null) { - root = new VectorSchemaRoot(List.of(vector)); + // ── Native Arrow ── + + static final class NativeArrow extends VectorStreamOutput { + private final VectorSchemaRoot root; + + NativeArrow(VectorSchemaRoot root) { + this.root = root; + } + + @Override + public VectorSchemaRoot getRoot() { + return root; } - root.setRowCount(row); - return root; + + @Override + public void writeByte(byte b) {} + + @Override + public void writeBytes(byte[] b, int offset, int length) {} + + @Override + public void flush() {} + + @Override + public void close() {} + + @Override + public void reset() {} } } diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java new file mode 100644 index 0000000000000..fcc2947cce2b0 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowBatchResponseTests.java @@ -0,0 +1,129 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.List; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; + +public class ArrowBatchResponseTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private Schema schema; + + static class TestResponse extends ArrowBatchResponse { + TestResponse(VectorSchemaRoot root) { + super(root); + } + + TestResponse(StreamInput in) throws IOException { + super(in); + } + } + + @Before + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); + } + + @After + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testGetRootReturnsProducerRoot() { + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + TestResponse response = new TestResponse(root); + assertSame(root, response.getRoot()); + root.close(); + } + + public void testWriteToIsNoOp() throws IOException { + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + TestResponse response = new TestResponse(root); + StreamOutput mockOut = mock(StreamOutput.class); + response.writeTo(mockOut); + verifyNoInteractions(mockOut); + root.close(); + } + + public void testTransferToMovesBuffers() { + VectorSchemaRoot src = VectorSchemaRoot.create(schema, allocator); + IntVector srcVec = (IntVector) src.getVector("val"); + srcVec.allocateNew(); + srcVec.setSafe(0, 42); + srcVec.setSafe(1, 99); + srcVec.setValueCount(2); + src.setRowCount(2); + + VectorSchemaRoot dst = VectorSchemaRoot.create(schema, allocator); + TestResponse response = new TestResponse(src); + response.transferTo(dst); + + assertEquals(2, dst.getRowCount()); + IntVector dstVec = (IntVector) dst.getVector("val"); + assertEquals(42, dstVec.get(0)); + assertEquals(99, dstVec.get(1)); + + // Source should be empty after transfer + assertEquals(0, srcVec.getValueCount()); + + src.close(); + dst.close(); + } + + public void testTransferToWithMultipleVectors() { + Schema multiSchema = new Schema( + List.of( + new Field("a", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("b", FieldType.nullable(new ArrowType.Int(32, true)), null) + ) + ); + + VectorSchemaRoot src = VectorSchemaRoot.create(multiSchema, allocator); + ((IntVector) src.getVector("a")).allocateNew(); + ((IntVector) src.getVector("a")).setSafe(0, 1); + ((IntVector) src.getVector("a")).setValueCount(1); + ((IntVector) src.getVector("b")).allocateNew(); + ((IntVector) src.getVector("b")).setSafe(0, 2); + ((IntVector) src.getVector("b")).setValueCount(1); + src.setRowCount(1); + + VectorSchemaRoot dst = VectorSchemaRoot.create(multiSchema, allocator); + new TestResponse(src).transferTo(dst); + + assertEquals(1, dst.getRowCount()); + assertEquals(1, ((IntVector) dst.getVector("a")).get(0)); + assertEquals(2, ((IntVector) dst.getVector("b")).get(0)); + + src.close(); + dst.close(); + } +} diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowFlightChannelTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowFlightChannelTests.java new file mode 100644 index 0000000000000..13229bee401ef --- /dev/null +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowFlightChannelTests.java @@ -0,0 +1,56 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.transport.BaseTcpTransportChannel; +import org.opensearch.transport.TaskTransportChannel; +import org.opensearch.transport.TcpChannel; +import org.opensearch.transport.TransportChannel; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ArrowFlightChannelTests extends OpenSearchTestCase { + + public void testFromWithTaskChannelWrappingBaseTcpChannel() { + // Simulate: TaskTransportChannel -> BaseTcpTransportChannel -> FlightServerChannel + FlightServerChannel serverChannel = mock(FlightServerChannel.class); + BufferAllocator mockAllocator = mock(BufferAllocator.class); + when(serverChannel.getAllocator()).thenReturn(mockAllocator); + + BaseTcpTransportChannel baseTcpChannel = mock(BaseTcpTransportChannel.class); + when(baseTcpChannel.getChannel()).thenReturn(serverChannel); + + TaskTransportChannel taskChannel = mock(TaskTransportChannel.class); + when(taskChannel.getChannel()).thenReturn(baseTcpChannel); + + ArrowFlightChannel result = ArrowFlightChannel.from(taskChannel); + assertSame(mockAllocator, result.getAllocator()); + } + + public void testFromWithNonFlightChannelThrows() { + TransportChannel plainChannel = mock(TransportChannel.class); + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> ArrowFlightChannel.from(plainChannel)); + assertTrue(ex.getMessage().contains("not backed by an ArrowFlightChannel")); + } + + public void testFromWithTaskChannelWrappingNonFlightTcpChannel() { + TcpChannel regularTcpChannel = mock(TcpChannel.class); + BaseTcpTransportChannel baseTcpChannel = mock(BaseTcpTransportChannel.class); + when(baseTcpChannel.getChannel()).thenReturn(regularTcpChannel); + + TaskTransportChannel taskChannel = mock(TaskTransportChannel.class); + when(taskChannel.getChannel()).thenReturn(baseTcpChannel); + + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> ArrowFlightChannel.from(taskChannel)); + assertTrue(ex.getMessage().contains("not backed by an ArrowFlightChannel")); + } +} diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java index 16024496216da..e85225bae0c42 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java @@ -50,7 +50,7 @@ public void tearDown() throws Exception { public void testInternalAggregationSerializationDeserialization() throws IOException { StringTerms original = createTestStringTerms(); - try (VectorStreamOutput output = new VectorStreamOutput(allocator, null)) { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { output.writeNamedWriteable(original); VectorSchemaRoot unifiedRoot = output.getRoot(); diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java index 265dab56d4fa7..f4eed5d2a36f2 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java @@ -9,9 +9,16 @@ package org.opensearch.arrow.flight.transport; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.Version; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.transport.TransportResponse; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; @@ -21,7 +28,9 @@ import org.junit.After; import org.junit.Before; +import java.io.IOException; import java.util.Collections; +import java.util.List; import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -33,7 +42,9 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; public class FlightOutboundHandlerTests extends OpenSearchTestCase { @@ -200,4 +211,170 @@ public void testMultipleBatchesMaintainCallerContext() throws Exception { assertEquals("Caller's thread context should be preserved after completeStream", HEADER_VALUE, threadContext.getHeader(HEADER_KEY)); } + + // --- Native Arrow branch in processBatchTask --- + + public void testProcessBatchTaskNativeArrowFirstBatch() throws Exception { + try (RootAllocator allocator = new RootAllocator()) { + Schema schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); + VectorSchemaRoot producerRoot = VectorSchemaRoot.create(schema, allocator); + IntVector vec = (IntVector) producerRoot.getVector("val"); + vec.allocateNew(); + vec.setSafe(0, 42); + vec.setValueCount(1); + producerRoot.setRowCount(1); + + // First batch: sharedRoot is null, so it should be created + when(mockFlightChannel.getRoot()).thenReturn(null); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference error = new AtomicReference<>(); + + doAnswer(invocation -> { + latch.countDown(); + return null; + }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); + + doAnswer(invocation -> { + // Verify the output has a root with transferred data + VectorStreamOutput out = invocation.getArgument(1); + VectorSchemaRoot sentRoot = out.getRoot(); + assertNotNull(sentRoot); + assertEquals(1, sentRoot.getRowCount()); + assertEquals(42, ((IntVector) sentRoot.getVector("val")).get(0)); + // Clean up the shared root created by the handler + sentRoot.close(); + return null; + }).when(mockFlightChannel).sendBatch(any(), any(VectorStreamOutput.class)); + + TestArrowResponse response = new TestArrowResponse(producerRoot); + handler.sendResponseBatch( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action", + response, + false, + false + ); + + assertTrue("Task should complete", latch.await(5, TimeUnit.SECONDS)); + assertNull("No error expected", error.get()); + } + } + + public void testProcessBatchTaskNativeArrowWithExistingSharedRoot() throws Exception { + try (RootAllocator allocator = new RootAllocator()) { + Schema schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); + + // Simulate existing shared root (second batch scenario) + VectorSchemaRoot sharedRoot = VectorSchemaRoot.create(schema, allocator); + when(mockFlightChannel.getRoot()).thenReturn(sharedRoot); + + VectorSchemaRoot producerRoot = VectorSchemaRoot.create(schema, allocator); + IntVector vec = (IntVector) producerRoot.getVector("val"); + vec.allocateNew(); + vec.setSafe(0, 99); + vec.setValueCount(1); + producerRoot.setRowCount(1); + + CountDownLatch latch = new CountDownLatch(1); + + doAnswer(invocation -> { + VectorStreamOutput out = invocation.getArgument(1); + VectorSchemaRoot sentRoot = out.getRoot(); + // Should reuse the existing shared root + assertSame(sharedRoot, sentRoot); + assertEquals(1, sentRoot.getRowCount()); + assertEquals(99, ((IntVector) sentRoot.getVector("val")).get(0)); + return null; + }).when(mockFlightChannel).sendBatch(any(), any(VectorStreamOutput.class)); + + doAnswer(invocation -> { + latch.countDown(); + return null; + }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); + + TestArrowResponse response = new TestArrowResponse(producerRoot); + handler.sendResponseBatch( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action", + response, + false, + false + ); + + assertTrue("Task should complete", latch.await(5, TimeUnit.SECONDS)); + sharedRoot.close(); + } + } + + // --- processCompleteTask error path --- + + public void testProcessCompleteTaskErrorPath() throws Exception { + RuntimeException completeError = new RuntimeException("complete failed"); + doThrow(completeError).when(mockFlightChannel).completeStream(any()); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference capturedError = new AtomicReference<>(); + + doAnswer(invocation -> { + capturedError.set(invocation.getArgument(2)); + latch.countDown(); + return null; + }).when(mockListener).onResponseSent(anyLong(), anyString(), any(Exception.class)); + + handler.completeStream( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action" + ); + + assertTrue("Task should complete", latch.await(5, TimeUnit.SECONDS)); + assertSame("Error should be passed to listener", completeError, capturedError.get()); + } + + public void testBatchTaskCloseWithIsErrorCallsReleaseChannelWithTrue() { + FlightTransportChannel mockTransportChannel = mock(FlightTransportChannel.class); + + FlightOutboundHandler.BatchTask task = new FlightOutboundHandler.BatchTask( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mockTransportChannel, + 1L, + "test-action", + null, + false, + false, + false, // isComplete + true, // isError + new RuntimeException("error") + ); + + task.close(); + + verify(mockTransportChannel).releaseChannel(true); + } + + // --- Test helper --- + + static class TestArrowResponse extends ArrowBatchResponse { + TestArrowResponse(VectorSchemaRoot root) { + super(root); + } + + TestArrowResponse(StreamInput in) throws IOException { + super(in); + } + } } diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java index 7c01f0c9afe1e..4f0157c1b461e 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java @@ -7,6 +7,7 @@ package org.opensearch.arrow.flight.transport; +import org.apache.arrow.memory.BufferAllocator; import org.opensearch.Version; import org.opensearch.arrow.flight.stats.FlightStatsCollector; import org.opensearch.common.lease.Releasable; @@ -31,6 +32,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; public class FlightTransportChannelTests extends OpenSearchTestCase { @@ -208,4 +210,24 @@ public void testMultipleSendResponseBatchAfterComplete() { assertEquals(StreamErrorCode.UNAVAILABLE, exception1.getErrorCode()); assertEquals(StreamErrorCode.UNAVAILABLE, exception2.getErrorCode()); } + + public void testGetAllocator() { + BufferAllocator mockAllocator = mock(BufferAllocator.class); + FlightServerChannel mockServerChannel = mock(FlightServerChannel.class); + when(mockServerChannel.getAllocator()).thenReturn(mockAllocator); + + FlightTransportChannel ch = new FlightTransportChannel( + mockOutboundHandler, + mockServerChannel, + "test-action", + 1L, + Version.CURRENT, + Collections.emptySet(), + false, + false, + mockReleasable + ); + + assertSame(mockAllocator, ch.getAllocator()); + } } diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamInputTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamInputTests.java new file mode 100644 index 0000000000000..37b470dc29bdc --- /dev/null +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamInputTests.java @@ -0,0 +1,70 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.After; +import org.junit.Before; + +import java.util.Collections; +import java.util.List; + +public class VectorStreamInputTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private NamedWriteableRegistry registry; + + @Before + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + registry = new NamedWriteableRegistry(Collections.emptyList()); + } + + @After + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testGetRootReturnsRoot() { + Schema schema = new Schema(List.of(new Field("0", FieldType.nullable(new ArrowType.Binary()), null))); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + VarBinaryVector vec = (VarBinaryVector) root.getVector("0"); + vec.allocateNew(); + vec.setValueCount(0); + root.setRowCount(0); + + VectorStreamInput input = new VectorStreamInput(root, registry); + assertSame(root, input.getRoot()); + root.close(); + } + + public void testCloseWithNullVector() throws Exception { + // Create a root with no vector named "0" so vector field is null + Schema schema = new Schema(List.of(new Field("other", FieldType.nullable(new ArrowType.Utf8()), null))); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + + VectorStreamInput input = new VectorStreamInput(root, registry); + // close() should not throw even though vector is null + input.close(); + root.close(); + } +} diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamOutputTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamOutputTests.java new file mode 100644 index 0000000000000..5ddbd6da4f41d --- /dev/null +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/VectorStreamOutputTests.java @@ -0,0 +1,243 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight.transport; + +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.List; + +public class VectorStreamOutputTests extends OpenSearchTestCase { + private RootAllocator allocator; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + allocator.close(); + } + + // ── ByteSerialized tests ── + + public void testByteSerializedConstructorWithoutExistingRoot() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + output.writeByte((byte) 42); + VectorSchemaRoot root = output.getRoot(); + assertNotNull(root); + assertEquals(1, root.getRowCount()); + } + } + + public void testByteSerializedConstructorWithExistingRoot() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.setInitialCapacity(16); + vector.allocateNew(); + + VectorSchemaRoot existingRoot = new VectorSchemaRoot(List.of(vector)); + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, existingRoot)) { + output.writeByte((byte) 7); + VectorSchemaRoot root = output.getRoot(); + assertSame(existingRoot, root); + assertEquals(1, root.getRowCount()); + byte[] value = ((VarBinaryVector) root.getVector(0)).get(0); + assertEquals(1, value.length); + assertEquals((byte) 7, value[0]); + } + } + + public void testWriteByteTriggersFlushTempBuffer() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + // Write exactly 8192 bytes to fill the temp buffer, then one more to trigger flush + for (int i = 0; i < 8193; i++) { + output.writeByte((byte) (i & 0xFF)); + } + VectorSchemaRoot root = output.getRoot(); + // First 8192 bytes flushed as one row, remaining 1 byte flushed as second row + assertEquals(2, root.getRowCount()); + VarBinaryVector vector = (VarBinaryVector) root.getVector(0); + assertEquals(8192, vector.get(0).length); + assertEquals(1, vector.get(1).length); + } + } + + public void testWriteBytesWithInvalidOffsetThrows() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + byte[] data = new byte[5]; + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> output.writeBytes(data, 3, 5)); + assertTrue(e.getMessage().contains("Illegal offset")); + } + } + + public void testWriteBytesWithInvalidLengthThrows() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + byte[] data = new byte[5]; + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> output.writeBytes(data, 0, 10)); + assertTrue(e.getMessage().contains("Illegal offset")); + } + } + + public void testWriteBytesWithZeroLengthIsNoOp() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + output.writeBytes(new byte[5], 0, 0); + VectorSchemaRoot root = output.getRoot(); + assertEquals(0, root.getRowCount()); + } + } + + public void testWriteBytesFlushesPendingTempBuffer() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + // Write some bytes via writeByte to fill tempBuffer partially + output.writeByte((byte) 1); + output.writeByte((byte) 2); + // Now writeBytes should flush the pending temp buffer first + byte[] data = new byte[] { 3, 4, 5 }; + output.writeBytes(data, 0, 3); + + VectorSchemaRoot root = output.getRoot(); + assertEquals(2, root.getRowCount()); + VarBinaryVector vector = (VarBinaryVector) root.getVector(0); + // Row 0: flushed temp buffer (2 bytes from writeByte) + byte[] row0 = vector.get(0); + assertEquals(2, row0.length); + assertEquals((byte) 1, row0[0]); + assertEquals((byte) 2, row0[1]); + // Row 1: direct writeBytes data + byte[] row1 = vector.get(1); + assertEquals(3, row1.length); + assertEquals((byte) 3, row1[0]); + } + } + + public void testResetClearsState() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + output.writeByte((byte) 1); + output.writeBytes(new byte[] { 2, 3 }, 0, 2); + + output.reset(); + + // After reset, writing and getting root should start fresh + output.writeByte((byte) 99); + VectorSchemaRoot root = output.getRoot(); + assertEquals(1, root.getRowCount()); + byte[] value = ((VarBinaryVector) root.getVector(0)).get(0); + assertEquals(1, value.length); + assertEquals((byte) 99, value[0]); + } + } + + public void testGetRootCreatesRootOnFirstCallReusesOnSecond() throws IOException { + try (VectorStreamOutput output = VectorStreamOutput.create(allocator, null)) { + output.writeByte((byte) 1); + VectorSchemaRoot root1 = output.getRoot(); + assertNotNull(root1); + + output.writeByte((byte) 2); + VectorSchemaRoot root2 = output.getRoot(); + assertSame(root1, root2); + } + } + + public void testCloseReleasesVector() throws IOException { + VectorStreamOutput output = VectorStreamOutput.create(allocator, null); + output.writeByte((byte) 1); + output.close(); + // After close, allocator should have no outstanding allocations + assertEquals(0, allocator.getAllocatedMemory()); + } + + // ── NativeArrow tests ── + + public void testNativeArrowGetRootReturnsSameRoot() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.allocateNew(); + VectorSchemaRoot root = new VectorSchemaRoot(List.of(vector)); + + try (VectorStreamOutput output = VectorStreamOutput.forNativeArrow(root)) { + assertSame(root, output.getRoot()); + assertSame(root, output.getRoot()); + } finally { + root.close(); + } + } + + public void testNativeArrowWriteByteIsNoOp() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.allocateNew(); + VectorSchemaRoot root = new VectorSchemaRoot(List.of(vector)); + + try (VectorStreamOutput output = VectorStreamOutput.forNativeArrow(root)) { + output.writeByte((byte) 42); + // Vector should remain unchanged + assertEquals(0, vector.getValueCount()); + } finally { + root.close(); + } + } + + public void testNativeArrowWriteBytesIsNoOp() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.allocateNew(); + VectorSchemaRoot root = new VectorSchemaRoot(List.of(vector)); + + try (VectorStreamOutput output = VectorStreamOutput.forNativeArrow(root)) { + output.writeBytes(new byte[] { 1, 2, 3 }, 0, 3); + assertEquals(0, vector.getValueCount()); + } finally { + root.close(); + } + } + + public void testNativeArrowFlushAndResetAreNoOps() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.allocateNew(); + VectorSchemaRoot root = new VectorSchemaRoot(List.of(vector)); + + try (VectorStreamOutput output = VectorStreamOutput.forNativeArrow(root)) { + output.flush(); + output.reset(); + // Root should still be the same and accessible + assertSame(root, output.getRoot()); + } finally { + root.close(); + } + } + + public void testNativeArrowCloseDoesNotCloseRoot() throws IOException { + Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); + VarBinaryVector vector = (VarBinaryVector) field.createVector(allocator); + vector.allocateNew(); + VectorSchemaRoot root = new VectorSchemaRoot(List.of(vector)); + + try { + VectorStreamOutput output = VectorStreamOutput.forNativeArrow(root); + output.close(); + // Root should still be usable after NativeArrow.close() + assertNotNull(root.getSchema()); + assertEquals(1, root.getFieldVectors().size()); + } finally { + root.close(); + } + } +} diff --git a/plugins/examples/stream-transport-example/build.gradle b/plugins/examples/stream-transport-example/build.gradle index 397e3f009f6f6..a4d761bb11f84 100644 --- a/plugins/examples/stream-transport-example/build.gradle +++ b/plugins/examples/stream-transport-example/build.gradle @@ -19,3 +19,5 @@ internalClusterTest { systemProperty 'io.netty.tryReflectionSetAccessible', 'true' jvmArgs += ["--add-opens", "java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED"] } + +tasks.named('missingJavadoc').configure { enabled = false } diff --git a/plugins/examples/stream-transport-example/src/internalClusterTest/java/org/opensearch/example/stream/NativeArrowStreamTransportExampleIT.java b/plugins/examples/stream-transport-example/src/internalClusterTest/java/org/opensearch/example/stream/NativeArrowStreamTransportExampleIT.java new file mode 100644 index 0000000000000..4cb4e68fc5889 --- /dev/null +++ b/plugins/examples/stream-transport-example/src/internalClusterTest/java/org/opensearch/example/stream/NativeArrowStreamTransportExampleIT.java @@ -0,0 +1,180 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.example.stream; + +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.opensearch.arrow.flight.transport.FlightStreamPlugin; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.StreamTransportResponseHandler; +import org.opensearch.transport.StreamTransportService; +import org.opensearch.transport.TransportException; +import org.opensearch.transport.TransportRequestOptions; +import org.opensearch.transport.stream.StreamTransportResponse; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.opensearch.common.util.FeatureFlags.STREAM_TRANSPORT; + +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, minNumDataNodes = 2, maxNumDataNodes = 2) +public class NativeArrowStreamTransportExampleIT extends OpenSearchIntegTestCase { + + @Override + public void setUp() throws Exception { + super.setUp(); + internalCluster().ensureAtLeastNumDataNodes(2); + } + + @Override + protected Collection> nodePlugins() { + return List.of(StreamTransportExamplePlugin.class, FlightStreamPlugin.class); + } + + @AwaitsFix(bugUrl = "") + @LockFeatureFlag(STREAM_TRANSPORT) + public void testNativeArrowSingleBatch() throws Exception { + for (DiscoveryNode node : getClusterState().nodes()) { + StreamTransportService sts = internalCluster().getInstance(StreamTransportService.class); + List batches = new ArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + + sts.sendRequest( + node, + NativeArrowStreamDataAction.NAME, + new NativeArrowStreamDataRequest(1, 3), + TransportRequestOptions.builder().withType(TransportRequestOptions.Type.STREAM).build(), + new NativeArrowResponseHandler(batches, latch, failure) + ); + + assertTrue("Stream should complete within 10s", latch.await(10, TimeUnit.SECONDS)); + assertNull("No exception expected: " + failure.get(), failure.get()); + assertEquals(1, batches.size()); + + ReceivedBatch batch = batches.get(0); + assertEquals(3, batch.rowCount); + assertEquals("name", batch.fieldNames.get(0)); + assertEquals("age", batch.fieldNames.get(1)); + assertEquals("Alice", batch.names.get(0)); + assertEquals("Bob", batch.names.get(1)); + assertEquals("Carol", batch.names.get(2)); + assertEquals(30, (int) batch.ages.get(0)); + assertEquals(31, (int) batch.ages.get(1)); + assertEquals(32, (int) batch.ages.get(2)); + } + } + + @AwaitsFix(bugUrl = "") + @LockFeatureFlag(STREAM_TRANSPORT) + public void testNativeArrowMultipleBatches() throws Exception { + for (DiscoveryNode node : getClusterState().nodes()) { + StreamTransportService sts = internalCluster().getInstance(StreamTransportService.class); + List batches = new ArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + + sts.sendRequest( + node, + NativeArrowStreamDataAction.NAME, + new NativeArrowStreamDataRequest(3, 2), + TransportRequestOptions.builder().withType(TransportRequestOptions.Type.STREAM).build(), + new NativeArrowResponseHandler(batches, latch, failure) + ); + + assertTrue("Stream should complete within 10s", latch.await(10, TimeUnit.SECONDS)); + assertNull("No exception expected: " + failure.get(), failure.get()); + assertEquals(3, batches.size()); + + for (int i = 0; i < 3; i++) { + ReceivedBatch batch = batches.get(i); + assertEquals(2, batch.rowCount); + assertEquals(30, (int) batch.ages.get(0)); + assertEquals(31, (int) batch.ages.get(1)); + } + } + } + + /** Deep-copies data from the root since FlightStream reuses it between next() calls. */ + static class ReceivedBatch { + final int rowCount; + final List fieldNames; + final List names; + final List ages; + + ReceivedBatch(VectorSchemaRoot root) { + this.rowCount = root.getRowCount(); + this.fieldNames = root.getSchema().getFields().stream().map(f -> f.getName()).toList(); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + IntVector ageVector = (IntVector) root.getVector("age"); + this.names = new ArrayList<>(); + this.ages = new ArrayList<>(); + for (int i = 0; i < rowCount; i++) { + names.add(new String(nameVector.get(i), StandardCharsets.UTF_8)); + ages.add(ageVector.get(i)); + } + } + } + + /** Standard handler — read() uses the normal StreamInput contract. */ + static class NativeArrowResponseHandler implements StreamTransportResponseHandler { + private final List batches; + private final CountDownLatch latch; + private final AtomicReference failure; + + NativeArrowResponseHandler(List batches, CountDownLatch latch, AtomicReference failure) { + this.batches = batches; + this.latch = latch; + this.failure = failure; + } + + @Override + public void handleStreamResponse(StreamTransportResponse streamResponse) { + try { + NativeArrowStreamDataResponse response; + while ((response = streamResponse.nextResponse()) != null) { + batches.add(new ReceivedBatch(response.getRoot())); + } + streamResponse.close(); + latch.countDown(); + } catch (Exception e) { + failure.set(e); + streamResponse.cancel("Test error", e); + latch.countDown(); + } + } + + @Override + public void handleException(TransportException exp) { + failure.set(exp); + latch.countDown(); + } + + @Override + public String executor() { + return ThreadPool.Names.GENERIC; + } + + @Override + public NativeArrowStreamDataResponse read(StreamInput in) throws IOException { + return new NativeArrowStreamDataResponse(in); + } + } +} diff --git a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataAction.java b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataAction.java new file mode 100644 index 0000000000000..71ff29ea38473 --- /dev/null +++ b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataAction.java @@ -0,0 +1,20 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.example.stream; + +import org.opensearch.action.ActionType; + +class NativeArrowStreamDataAction extends ActionType { + public static final NativeArrowStreamDataAction INSTANCE = new NativeArrowStreamDataAction(); + public static final String NAME = "cluster:admin/native_arrow_stream_data"; + + private NativeArrowStreamDataAction() { + super(NAME, NativeArrowStreamDataResponse::new); + } +} diff --git a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataRequest.java b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataRequest.java new file mode 100644 index 0000000000000..75e7c530ace51 --- /dev/null +++ b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataRequest.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.example.stream; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; + +class NativeArrowStreamDataRequest extends ActionRequest { + private final int batchCount; + private final int rowsPerBatch; + + NativeArrowStreamDataRequest(int batchCount, int rowsPerBatch) { + this.batchCount = batchCount; + this.rowsPerBatch = rowsPerBatch; + } + + NativeArrowStreamDataRequest(StreamInput in) throws IOException { + super(in); + this.batchCount = in.readInt(); + this.rowsPerBatch = in.readInt(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeInt(batchCount); + out.writeInt(rowsPerBatch); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + public int getBatchCount() { + return batchCount; + } + + public int getRowsPerBatch() { + return rowsPerBatch; + } +} diff --git a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataResponse.java b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataResponse.java new file mode 100644 index 0000000000000..1ee7de49ad71a --- /dev/null +++ b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/NativeArrowStreamDataResponse.java @@ -0,0 +1,37 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.example.stream; + +import org.apache.arrow.vector.VectorSchemaRoot; +import org.opensearch.arrow.flight.transport.ArrowBatchResponse; +import org.opensearch.core.common.io.stream.StreamInput; + +import java.io.IOException; + +/** + * Example native Arrow response — just extend {@link ArrowBatchResponse}. + * + *

The framework handles everything: + *

    + *
  • Send side: zero-copy transfers the root's buffers into the Flight stream
    • + *
  • Receive side: provides the root via {@link #getRoot()} — no deserialization
    • + *
+ * + *

No writeTo/read override needed. The base class handles both. + */ +class NativeArrowStreamDataResponse extends ArrowBatchResponse { + + NativeArrowStreamDataResponse(VectorSchemaRoot root) { + super(root); + } + + NativeArrowStreamDataResponse(StreamInput in) throws IOException { + super(in); + } +} diff --git a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/StreamTransportExamplePlugin.java b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/StreamTransportExamplePlugin.java index 94ea2d1fa8231..bbc5952ca3f17 100644 --- a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/StreamTransportExamplePlugin.java +++ b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/StreamTransportExamplePlugin.java @@ -13,21 +13,17 @@ import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; -import java.util.Collections; import java.util.List; -/** - * Example plugin demonstrating streaming transport actions - */ public class StreamTransportExamplePlugin extends Plugin implements ActionPlugin { - /** - * Constructor - */ public StreamTransportExamplePlugin() {} @Override public List> getActions() { - return Collections.singletonList(new ActionHandler<>(StreamDataAction.INSTANCE, TransportStreamDataAction.class)); + return List.of( + new ActionHandler<>(StreamDataAction.INSTANCE, TransportStreamDataAction.class), + new ActionHandler<>(NativeArrowStreamDataAction.INSTANCE, TransportNativeArrowStreamDataAction.class) + ); } } diff --git a/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/TransportNativeArrowStreamDataAction.java b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/TransportNativeArrowStreamDataAction.java new file mode 100644 index 0000000000000..99fee3d870eb1 --- /dev/null +++ b/plugins/examples/stream-transport-example/src/main/java/org/opensearch/example/stream/TransportNativeArrowStreamDataAction.java @@ -0,0 +1,110 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.example.stream; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.TransportAction; +import org.opensearch.arrow.flight.transport.ArrowFlightChannel; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.StreamTransportService; +import org.opensearch.transport.TransportChannel; +import org.opensearch.transport.stream.StreamErrorCode; +import org.opensearch.transport.stream.StreamException; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; + +/** + * Example: server-side handler producing native Arrow data. + * + *

Demonstrates the pipelined producer pattern: + *

    + *
  1. Get the channel's allocator via {@link ArrowFlightChannel#from(TransportChannel)}
    2. + *
  2. For each batch, create a producer root using the channel allocator
    3. + *
  3. Populate the root with typed vectors (VarChar, Int, etc.)
    4. + *
  4. Send via {@code sendResponseBatch()} — the framework does zero-copy transfer + * of the producer's buffers into the channel's shared root on the executor thread
    5. + *
  5. The producer root is closed by the framework after transfer — don't reuse it
    6. + *
+ * + *

The channel allocator must be used directly (not a per-request child allocator) + * because gRPC's zero-copy write path retains buffer references beyond stream completion. + */ +public class TransportNativeArrowStreamDataAction extends TransportAction { + + private static final String[] NAMES = { "Alice", "Bob", "Carol", "Dave", "Eve" }; + + @Inject + public TransportNativeArrowStreamDataAction(StreamTransportService streamTransportService, ActionFilters actionFilters) { + super(NativeArrowStreamDataAction.NAME, actionFilters, streamTransportService.getTaskManager()); + streamTransportService.registerRequestHandler( + NativeArrowStreamDataAction.NAME, + ThreadPool.Names.GENERIC, + NativeArrowStreamDataRequest::new, + this::handleStreamRequest + ); + } + + @Override + protected void doExecute(Task task, NativeArrowStreamDataRequest request, ActionListener listener) { + listener.onFailure(new UnsupportedOperationException("Use StreamTransportService")); + } + + private void handleStreamRequest(NativeArrowStreamDataRequest request, TransportChannel channel, Task task) throws IOException { + // Get the channel's allocator. Use this directly for producer roots to ensure + // same-allocator transfer (avoids Arrow's cross-allocator foreign buffer bug). + BufferAllocator allocator = ArrowFlightChannel.from(channel).getAllocator(); + + Schema schema = new Schema( + List.of( + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("age", FieldType.nullable(new ArrowType.Int(32, true)), null) + ) + ); + + try { + for (int batch = 0; batch < request.getBatchCount(); batch++) { + VectorSchemaRoot producerRoot = VectorSchemaRoot.create(schema, allocator); + populateBatch(producerRoot, request.getRowsPerBatch(), batch); + channel.sendResponseBatch(new NativeArrowStreamDataResponse(producerRoot)); + } + channel.completeStream(); + } catch (StreamException e) { + if (e.getErrorCode() != StreamErrorCode.CANCELLED) { + channel.sendResponse(e); + } + } catch (Exception e) { + channel.sendResponse(e); + } + } + + private void populateBatch(VectorSchemaRoot root, int rowCount, int batchIndex) { + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + IntVector ageVector = (IntVector) root.getVector("age"); + nameVector.allocateNew(); + ageVector.allocateNew(); + for (int i = 0; i < rowCount; i++) { + nameVector.setSafe(i, NAMES[(batchIndex * rowCount + i) % NAMES.length].getBytes(StandardCharsets.UTF_8)); + ageVector.setSafe(i, 30 + i); + } + root.setRowCount(rowCount); + } +}