This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8473d90e93f [Java IO] Add ArrowFlight IO connector (#37904)
8473d90e93f is described below
commit 8473d90e93fe41b08bd105e6a0c70b36b60fc39f
Author: Bruno Volpato <[email protected]>
AuthorDate: Tue Jul 28 15:44:38 2026 -0400
[Java IO] Add ArrowFlight IO connector (#37904)
* [Java IO] Add ArrowFlight IO connector
Add a new IO connector for Apache Arrow Flight, enabling high-performance
data transfer over gRPC using the Arrow columnar format.
Includes read (BoundedSource) and write (DoFn with doPut) support with
endpoint-level split parallelism and bearer token authentication.
Fixes #20116
* Add arrow-flight to javaioPreCommit and fix test issues
- Register :sdks:java:io:arrow-flight in javaioPreCommit task
(build.gradle.kts)
- Add --add-opens JVM arg for Arrow native memory on JDK 17+
- Make host() nullable at AutoValue level to fix factory method NPE
- Eagerly materialize rows from Arrow buffers to prevent stale access
- Move root.setRowCount() after vector population for correct ordering
- Use AtomicInteger for thread-safe write record counting in tests
* [Java IO] Harden ArrowFlightIO writes
* [Java IO] Address ArrowFlightIO review feedback
---
CHANGES.md | 1 +
build.gradle.kts | 1 +
.../org/apache/beam/gradle/BeamModulePlugin.groovy | 1 +
.../beam/sdk/extensions/arrow/ArrowConversion.java | 61 ++
.../sdk/extensions/arrow/ArrowConversionTest.java | 33 +
sdks/java/io/arrow-flight/build.gradle | 44 ++
.../beam/sdk/io/arrowflight/ArrowFlightIO.java | 840 +++++++++++++++++++++
.../beam/sdk/io/arrowflight/package-info.java | 29 +
.../beam/sdk/io/arrowflight/ArrowFlightIOTest.java | 330 ++++++++
settings.gradle.kts | 1 +
.../site/content/en/documentation/io/connectors.md | 16 +
11 files changed, 1357 insertions(+)
diff --git a/CHANGES.md b/CHANGES.md
index 1bbde577e87..a1fdede9449 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -65,6 +65,7 @@
## I/Os
* Support for X source added (Java/Python)
([#X](https://github.com/apache/beam/issues/X)).
+* Add ArrowFlight IO (Java)
([#20116](https://github.com/apache/beam/issues/20116)).
## New Features / Improvements
diff --git a/build.gradle.kts b/build.gradle.kts
index ba2f7aaed67..de278f9f283 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -339,6 +339,7 @@ tasks.register("javaPreCommit") {
// a precommit task build multiple IOs (except those splitting into single
jobs)
tasks.register("javaioPreCommit") {
dependsOn(":sdks:java:io:amqp:build")
+ dependsOn(":sdks:java:io:arrow-flight:build")
// CassandraIO, HBaseIO and HCatalogIO do not support Java17+, test ran
separately
// dependsOn(":sdks:java:io:cassandra:build")
dependsOn(":sdks:java:io:csv:build")
diff --git
a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index 70a22e39772..0bf3dc11e2d 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -951,6 +951,7 @@ class BeamModulePlugin implements Plugin<Project> {
arrow_vector :
"org.apache.arrow:arrow-vector:$arrow_version",
arrow_memory_core :
"org.apache.arrow:arrow-memory-core:$arrow_version",
arrow_memory_netty :
"org.apache.arrow:arrow-memory-netty:$arrow_version",
+ arrow_flight_core :
"org.apache.arrow:flight-core:$arrow_version",
],
groovy: [
groovy_all: "org.codehaus.groovy:groovy-all:2.4.13",
diff --git
a/sdks/java/extensions/arrow/src/main/java/org/apache/beam/sdk/extensions/arrow/ArrowConversion.java
b/sdks/java/extensions/arrow/src/main/java/org/apache/beam/sdk/extensions/arrow/ArrowConversion.java
index bf7078abc58..f1850df2e95 100644
---
a/sdks/java/extensions/arrow/src/main/java/org/apache/beam/sdk/extensions/arrow/ArrowConversion.java
+++
b/sdks/java/extensions/arrow/src/main/java/org/apache/beam/sdk/extensions/arrow/ArrowConversion.java
@@ -22,6 +22,7 @@ import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
import java.io.IOException;
import java.io.InputStream;
import java.nio.channels.Channels;
+import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
@@ -35,6 +36,7 @@ import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
+import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.util.Text;
@@ -57,6 +59,57 @@ import org.joda.time.DateTimeZone;
*/
public class ArrowConversion {
+ /** Get Arrow Field from Beam Field. */
+ private static org.apache.arrow.vector.types.pojo.Field toArrowField(Field
field) {
+ FieldType beamFieldType = field.getType();
+ ArrowType arrowType;
+ // TODO: Support aggregate and logical Beam field types.
+ switch (beamFieldType.getTypeName()) {
+ case BYTE:
+ arrowType = new ArrowType.Int(8, true);
+ break;
+ case INT16:
+ arrowType = new ArrowType.Int(16, true);
+ break;
+ case INT32:
+ arrowType = new ArrowType.Int(32, true);
+ break;
+ case INT64:
+ arrowType = new ArrowType.Int(64, true);
+ break;
+ case FLOAT:
+ arrowType = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
+ break;
+ case DOUBLE:
+ arrowType = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
+ break;
+ case STRING:
+ arrowType = ArrowType.Utf8.INSTANCE;
+ break;
+ case BOOLEAN:
+ arrowType = ArrowType.Bool.INSTANCE;
+ break;
+ case BYTES:
+ arrowType = ArrowType.Binary.INSTANCE;
+ break;
+ case DATETIME:
+ arrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC");
+ break;
+ default:
+ throw new IllegalArgumentException(
+ String.format(
+ "Arrow schema conversion does not support Beam type '%s' for
field '%s'.",
+ beamFieldType.getTypeName(), field.getName()));
+ }
+
+ org.apache.arrow.vector.types.pojo.FieldType arrowFieldType =
+ beamFieldType.getNullable()
+ ? org.apache.arrow.vector.types.pojo.FieldType.nullable(arrowType)
+ :
org.apache.arrow.vector.types.pojo.FieldType.notNullable(arrowType);
+ return new org.apache.arrow.vector.types.pojo.Field(
+ field.getName(), arrowFieldType, Collections.emptyList());
+ }
+
/** Get Beam Field from Arrow Field. */
private static Field toBeamField(org.apache.arrow.vector.types.pojo.Field
field) {
FieldType beamFieldType = toFieldType(field.getFieldType(),
field.getChildren());
@@ -546,6 +599,14 @@ public class ArrowConversion {
/** Converts Arrow schema to Beam row schema. */
public static class ArrowSchemaTranslator {
+ /** Converts a supported Beam row schema to an Arrow schema. */
+ public static org.apache.arrow.vector.types.pojo.Schema
toArrowSchema(Schema schema) {
+ return new org.apache.arrow.vector.types.pojo.Schema(
+ schema.getFields().stream()
+ .map(ArrowConversion::toArrowField)
+ .collect(Collectors.toList()));
+ }
+
public static Schema
toBeamSchema(org.apache.arrow.vector.types.pojo.Schema schema) {
return toBeamSchema(schema.getFields());
}
diff --git
a/sdks/java/extensions/arrow/src/test/java/org/apache/beam/sdk/extensions/arrow/ArrowConversionTest.java
b/sdks/java/extensions/arrow/src/test/java/org/apache/beam/sdk/extensions/arrow/ArrowConversionTest.java
index 8c297576650..be5c7120343 100644
---
a/sdks/java/extensions/arrow/src/test/java/org/apache/beam/sdk/extensions/arrow/ArrowConversionTest.java
+++
b/sdks/java/extensions/arrow/src/test/java/org/apache/beam/sdk/extensions/arrow/ArrowConversionTest.java
@@ -83,6 +83,39 @@ public class ArrowConversionTest {
assertThat(ArrowConversion.ArrowSchemaTranslator.toBeamSchema(arrowSchema),
equalTo(expected));
}
+ @Test
+ public void toArrowSchema_convertsSimpleBeamSchema() {
+ Schema beamSchema =
+ Schema.builder()
+ .addByteField("int8")
+ .addInt16Field("int16")
+ .addInt32Field("int32")
+ .addInt64Field("int64")
+ .addFloatField("float32")
+ .addDoubleField("float64")
+ .addNullableField("string", FieldType.STRING)
+ .addBooleanField("boolean")
+ .addByteArrayField("bytes")
+ .addDateTimeField("timestamp")
+ .build();
+
+ org.apache.arrow.vector.types.pojo.Schema expected =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ field("int8", new ArrowType.Int(8, true)),
+ field("int16", new ArrowType.Int(16, true)),
+ field("int32", new ArrowType.Int(32, true)),
+ field("int64", new ArrowType.Int(64, true)),
+ field("float32", new
ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)),
+ field("float64", new
ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)),
+ field("string", true, ArrowType.Utf8.INSTANCE),
+ field("boolean", ArrowType.Bool.INSTANCE),
+ field("bytes", ArrowType.Binary.INSTANCE),
+ field("timestamp", new
ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))));
+
+
assertThat(ArrowConversion.ArrowSchemaTranslator.toArrowSchema(beamSchema),
equalTo(expected));
+ }
+
@Test
public void rowIterator() {
org.apache.arrow.vector.types.pojo.Schema schema =
diff --git a/sdks/java/io/arrow-flight/build.gradle
b/sdks/java/io/arrow-flight/build.gradle
new file mode 100644
index 00000000000..6d7a6765f8c
--- /dev/null
+++ b/sdks/java/io/arrow-flight/build.gradle
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins { id 'org.apache.beam.module' }
+applyJavaNature(automaticModuleName: 'org.apache.beam.sdk.io.arrowflight')
+
+description = "Apache Beam :: SDKs :: Java :: IO :: Arrow Flight"
+ext.summary = "IO to read and write data using Apache Arrow Flight RPC."
+
+dependencies {
+ implementation project(path: ":sdks:java:core", configuration: "shadow")
+ implementation project(path: ":sdks:java:extensions:arrow")
+ implementation library.java.joda_time
+ implementation library.java.slf4j_api
+ implementation library.java.vendored_guava_32_1_2_jre
+ implementation(library.java.arrow_flight_core)
+ implementation(library.java.arrow_memory_core)
+ implementation(library.java.arrow_vector)
+ testImplementation library.java.hamcrest
+ testImplementation library.java.junit
+ testImplementation(library.java.arrow_memory_netty)
+ testRuntimeOnly library.java.slf4j_simple
+ testRuntimeOnly project(path: ":runners:direct-java", configuration:
"shadow")
+}
+
+test {
+ // Keep aligned with ArrowFlightIO runtime guidance for Java 17+.
+ jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED'
+}
diff --git
a/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIO.java
b/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIO.java
new file mode 100644
index 00000000000..f14e10631d6
--- /dev/null
+++
b/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIO.java
@@ -0,0 +1,840 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.arrowflight;
+
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.io.Serializable;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.apache.arrow.flight.AsyncPutListener;
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.FlightCallHeaders;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.HeaderCallOption;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeStampMilliTZVector;
+import org.apache.arrow.vector.TinyIntVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.extensions.arrow.ArrowConversion;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * IO to read and write data using <a
href="https://arrow.apache.org/docs/format/Flight.html";>Apache
+ * Arrow Flight</a>.
+ *
+ * <p>Arrow Flight is an RPC framework for transferring Arrow-formatted data
over gRPC.
+ *
+ * <h3>Reading from an Arrow Flight server</h3>
+ *
+ * <p>{@link ArrowFlightIO#read()} returns a bounded {@link PCollection} of
{@link Row} elements.
+ * Each row is converted from Arrow record batches using the existing {@link
ArrowConversion}
+ * extension.
+ *
+ * <pre>{@code
+ * PCollection<Row> rows = pipeline.apply(
+ * ArrowFlightIO.read()
+ * .withHost("localhost")
+ * .withPort(47470)
+ * .withCommand("SELECT * FROM my_table"));
+ * }</pre>
+ *
+ * <h3>Writing to an Arrow Flight server</h3>
+ *
+ * <p>{@link ArrowFlightIO#write()} accepts a {@link PCollection} of {@link
Row} elements and
+ * streams them to a Flight server using {@code doPut}.
+ *
+ * <pre>{@code
+ * rows.apply(
+ * ArrowFlightIO.write()
+ * .withHost("localhost")
+ * .withPort(47470)
+ * .withDescriptor("my_table")
+ * .withBatchSize(1024));
+ * }</pre>
+ *
+ * <h3>Java runtime configuration</h3>
+ *
+ * <p>On Java 17 or later, Arrow memory access requires {@code
+ * --add-opens=java.base/java.nio=ALL-UNNAMED} on JVMs executing this
connector. Beam SDK containers
+ * enable it by default. For other environments, add it to the local or runner
JVM and pass {@code
+ * --JdkAddOpenModules=java.base/java.nio=ALL-UNNAMED} to configure SDK
harness JVMs. See the <a
+ * href="https://arrow.apache.org/docs/java/install.html";>Arrow Java
installation guide</a>.
+ */
+public class ArrowFlightIO {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(ArrowFlightIO.class);
+
+ private ArrowFlightIO() {}
+
+ private static byte[] copyToken(byte[] token) {
+ return Arrays.copyOf(checkNotNull(token, "token"), token.length);
+ }
+
+ private static CallOption[] callOptions(byte @Nullable [] token) {
+ if (token == null) {
+ return new CallOption[0];
+ }
+ FlightCallHeaders headers = new FlightCallHeaders();
+ headers.insert("authorization", "Bearer " + new String(token,
StandardCharsets.UTF_8));
+ return new CallOption[] {new HeaderCallOption(headers)};
+ }
+
+ private static void validateWriteSchema(Schema schema) {
+ try {
+ ArrowConversion.ArrowSchemaTranslator.toArrowSchema(schema);
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException("ArrowFlightIO.write(): " +
e.getMessage(), e);
+ }
+ }
+
+ public static Read read() {
+ return new
AutoValue_ArrowFlightIO_Read.Builder().setPort(47470).setUseTls(false).build();
+ }
+
+ public static Write write() {
+ return new AutoValue_ArrowFlightIO_Write.Builder()
+ .setPort(47470)
+ .setUseTls(false)
+ .setBatchSize(1024)
+ .build();
+ }
+
+ /**
+ * Creates a {@link FlightClient} from the given connection parameters.
+ *
+ * <p>The client uses a {@link RootAllocator} for Arrow memory management
and connects to the
+ * specified host and port using either plaintext or TLS.
+ */
+ static FlightClient createClient(
+ BufferAllocator allocator, String host, int port, boolean useTls) {
+ Location location;
+ if (useTls) {
+ location = Location.forGrpcTls(host, port);
+ } else {
+ location = Location.forGrpcInsecure(host, port);
+ }
+ return FlightClient.builder(allocator, location).build();
+ }
+
+ /** A serializable wrapper around Flight endpoint information for use in
BoundedSource splits. */
+ static class SerializableEndpoint implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ private final byte[] ticketBytes;
+ private final @Nullable String host;
+ private final int port;
+
+ SerializableEndpoint(byte[] ticketBytes, @Nullable String host, int port) {
+ this.ticketBytes = ticketBytes;
+ this.host = host;
+ this.port = port;
+ }
+
+ static SerializableEndpoint fromFlightEndpoint(
+ FlightEndpoint endpoint, String defaultHost, int defaultPort) {
+ byte[] ticket = endpoint.getTicket().getBytes();
+ List<Location> locations = endpoint.getLocations();
+ if (locations != null && !locations.isEmpty()) {
+ URI uri = locations.get(0).getUri();
+ return new SerializableEndpoint(ticket, uri.getHost(), uri.getPort());
+ }
+ return new SerializableEndpoint(ticket, defaultHost, defaultPort);
+ }
+
+ Ticket getTicket() {
+ return new Ticket(ticketBytes);
+ }
+
+ String getHost(String defaultHost) {
+ return host != null ? host : defaultHost;
+ }
+
+ int getPort(int defaultPort) {
+ return port > 0 ? port : defaultPort;
+ }
+ }
+
+ // ======================== READ ========================
+
+ @AutoValue
+ public abstract static class Read extends PTransform<PBegin,
PCollection<Row>> {
+
+ abstract @Nullable String host();
+
+ abstract int port();
+
+ abstract boolean useTls();
+
+ abstract @Nullable String command();
+
+ @SuppressWarnings("mutable")
+ abstract byte @Nullable [] token();
+
+ abstract Builder builder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setHost(String host);
+
+ abstract Builder setPort(int port);
+
+ abstract Builder setUseTls(boolean useTls);
+
+ abstract Builder setCommand(String command);
+
+ abstract Builder setToken(byte[] token);
+
+ abstract Read build();
+ }
+
+ /** Sets the Flight server host. */
+ public Read withHost(String host) {
+ return builder().setHost(host).build();
+ }
+
+ /** Sets the Flight server port. */
+ public Read withPort(int port) {
+ return builder().setPort(port).build();
+ }
+
+ /** Enables TLS for the connection. */
+ public Read withUseTls(boolean useTls) {
+ return builder().setUseTls(useTls).build();
+ }
+
+ /** Sets the command (e.g., a SQL query or table name) to request from the
Flight server. */
+ public Read withCommand(String command) {
+ return builder().setCommand(command).build();
+ }
+
+ /** Sets a bearer token for authentication. */
+ public Read withToken(byte[] token) {
+ return builder().setToken(copyToken(token)).build();
+ }
+
+ @Override
+ public PCollection<Row> expand(PBegin input) {
+ checkArgument(host() != null, "withHost() is required");
+ checkArgument(command() != null, "withCommand() is required");
+
+ Schema beamSchema;
+ try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ FlightClient client =
+ createClient(allocator, checkNotNull(host(), "host"), port(),
useTls())) {
+ FlightInfo info =
+ client.getInfo(
+ FlightDescriptor.command(
+ checkNotNull(command(),
"command").getBytes(StandardCharsets.UTF_8)),
+ callOptions());
+ beamSchema =
ArrowConversion.ArrowSchemaTranslator.toBeamSchema(info.getSchema());
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while fetching Flight schema",
e);
+ }
+
+ return input
+ .apply(org.apache.beam.sdk.io.Read.from(new
FlightBoundedSource(this, beamSchema)))
+ .setRowSchema(beamSchema);
+ }
+
+ CallOption[] callOptions() {
+ return ArrowFlightIO.callOptions(token());
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ builder.addIfNotNull(DisplayData.item("host", host()));
+ builder.add(DisplayData.item("port", port()));
+ builder.add(DisplayData.item("useTls", useTls()));
+ builder.addIfNotNull(DisplayData.item("command", command()));
+ }
+ }
+
+ /** A {@link BoundedSource} that reads rows from Arrow Flight endpoints. */
+ static class FlightBoundedSource extends BoundedSource<Row> {
+ private final Read spec;
+ private final Schema beamSchema;
+ private final @Nullable SerializableEndpoint endpoint;
+
+ FlightBoundedSource(Read spec, Schema beamSchema) {
+ this(spec, beamSchema, null);
+ }
+
+ FlightBoundedSource(Read spec, Schema beamSchema, @Nullable
SerializableEndpoint endpoint) {
+ this.spec = spec;
+ this.beamSchema = beamSchema;
+ this.endpoint = endpoint;
+ }
+
+ @Override
+ public List<? extends BoundedSource<Row>> split(
+ long desiredBundleSizeBytes, PipelineOptions options) throws Exception
{
+ if (endpoint != null) {
+ return Collections.singletonList(this);
+ }
+
+ List<BoundedSource<Row>> sources = new ArrayList<>();
+ try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ FlightClient client =
+ createClient(
+ allocator, checkNotNull(spec.host(), "host"), spec.port(),
spec.useTls())) {
+ FlightInfo info =
+ client.getInfo(
+ FlightDescriptor.command(
+ checkNotNull(spec.command(),
"command").getBytes(StandardCharsets.UTF_8)),
+ spec.callOptions());
+ for (FlightEndpoint fe : info.getEndpoints()) {
+ SerializableEndpoint se =
+ SerializableEndpoint.fromFlightEndpoint(
+ fe, checkNotNull(spec.host(), "host"), spec.port());
+ sources.add(new FlightBoundedSource(spec, beamSchema, se));
+ }
+ }
+
+ if (sources.isEmpty()) {
+ sources.add(this);
+ }
+ return sources;
+ }
+
+ @Override
+ public long getEstimatedSizeBytes(PipelineOptions options) throws
Exception {
+ if (endpoint != null) {
+ return -1;
+ }
+ try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ FlightClient client =
+ createClient(
+ allocator, checkNotNull(spec.host(), "host"), spec.port(),
spec.useTls())) {
+ FlightInfo info =
+ client.getInfo(
+ FlightDescriptor.command(
+ checkNotNull(spec.command(),
"command").getBytes(StandardCharsets.UTF_8)),
+ spec.callOptions());
+ return info.getBytes();
+ }
+ }
+
+ @Override
+ public BoundedReader<Row> createReader(PipelineOptions options) {
+ return new FlightBoundedReader(this);
+ }
+
+ @Override
+ public void validate() {
+ checkArgument(spec.host() != null, "host is required");
+ checkArgument(spec.command() != null, "command is required");
+ }
+
+ @Override
+ public Coder<Row> getOutputCoder() {
+ return RowCoder.of(beamSchema);
+ }
+ }
+
+ /** Reader that streams Arrow record batches from a Flight endpoint and
emits Beam Rows. */
+ @SuppressWarnings("initialization.fields.uninitialized")
+ static class FlightBoundedReader extends BoundedSource.BoundedReader<Row> {
+ private static final Counter RECORDS_READ =
Metrics.counter(ArrowFlightIO.class, "recordsRead");
+
+ private final FlightBoundedSource source;
+ private transient BufferAllocator allocator;
+ private transient FlightClient client;
+ private transient FlightStream stream;
+ private transient Iterator<Row> currentBatchIterator;
+ private transient Row current;
+ private FlightBoundedSource currentSource;
+
+ FlightBoundedReader(FlightBoundedSource source) {
+ this.source = source;
+ this.currentSource = source;
+ }
+
+ @Override
+ public boolean start() throws IOException {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ Read spec = source.spec;
+
+ String defaultHost = checkNotNull(spec.host(), "host");
+ SerializableEndpoint endpoint = source.endpoint;
+ if (endpoint == null) {
+ try (FlightClient discoveryClient =
+ createClient(allocator, defaultHost, spec.port(), spec.useTls())) {
+ FlightInfo info =
+ discoveryClient.getInfo(
+ FlightDescriptor.command(
+ checkNotNull(spec.command(),
"command").getBytes(StandardCharsets.UTF_8)),
+ spec.callOptions());
+ List<FlightEndpoint> endpoints = info.getEndpoints();
+ if (endpoints.isEmpty()) {
+ return false;
+ }
+ endpoint =
+ SerializableEndpoint.fromFlightEndpoint(endpoints.get(0),
defaultHost, spec.port());
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Interrupted while discovering Flight
endpoints", e);
+ }
+ currentSource = new FlightBoundedSource(spec, source.beamSchema,
endpoint);
+ }
+
+ client =
+ createClient(
+ allocator,
+ endpoint.getHost(defaultHost),
+ endpoint.getPort(spec.port()),
+ spec.useTls());
+ stream = client.getStream(endpoint.getTicket(), spec.callOptions());
+ currentBatchIterator = Collections.emptyIterator();
+ return advance();
+ }
+
+ @Override
+ public boolean advance() throws IOException {
+ while (true) {
+ if (currentBatchIterator.hasNext()) {
+ current = currentBatchIterator.next();
+ RECORDS_READ.inc();
+ return true;
+ }
+ if (stream.next()) {
+ VectorSchemaRoot root = stream.getRoot();
+ if (root.getRowCount() > 0) {
+ Iterator<Row> lazyIterator =
+ ArrowConversion.rowsFromRecordBatch(source.beamSchema, root);
+ List<Row> materializedRows = new ArrayList<>();
+ while (lazyIterator.hasNext()) {
+ Row lazyRow = lazyIterator.next();
+ materializedRows.add(
+
Row.withSchema(source.beamSchema).addValues(lazyRow.getValues()).build());
+ }
+ currentBatchIterator = materializedRows.iterator();
+ }
+ } else {
+ return false;
+ }
+ }
+ }
+
+ @Override
+ public Row getCurrent() {
+ return current;
+ }
+
+ @Override
+ public void close() throws IOException {
+ try {
+ if (stream != null) {
+ stream.close();
+ }
+ } catch (Exception e) {
+ LOG.warn("Error closing FlightStream", e);
+ }
+ try {
+ if (client != null) {
+ client.close();
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ LOG.warn("Interrupted closing FlightClient", e);
+ }
+ try {
+ if (allocator != null) {
+ allocator.close();
+ }
+ } catch (Exception e) {
+ LOG.warn("Error closing BufferAllocator", e);
+ }
+ }
+
+ @Override
+ public BoundedSource<Row> getCurrentSource() {
+ return currentSource;
+ }
+ }
+
+ // ======================== WRITE ========================
+
+ // TODO: Return a write result with failed rows for dead-letter handling
instead of PDone.
+ @AutoValue
+ public abstract static class Write extends PTransform<PCollection<Row>,
PDone> {
+
+ abstract @Nullable String host();
+
+ abstract int port();
+
+ abstract boolean useTls();
+
+ abstract @Nullable String descriptor();
+
+ abstract int batchSize();
+
+ @SuppressWarnings("mutable")
+ abstract byte @Nullable [] token();
+
+ abstract Builder builder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setHost(String host);
+
+ abstract Builder setPort(int port);
+
+ abstract Builder setUseTls(boolean useTls);
+
+ abstract Builder setDescriptor(String descriptor);
+
+ abstract Builder setBatchSize(int batchSize);
+
+ abstract Builder setToken(byte[] token);
+
+ abstract Write build();
+ }
+
+ /** Sets the Flight server host. */
+ public Write withHost(String host) {
+ return builder().setHost(host).build();
+ }
+
+ /** Sets the Flight server port. */
+ public Write withPort(int port) {
+ return builder().setPort(port).build();
+ }
+
+ /** Enables TLS for the connection. */
+ public Write withUseTls(boolean useTls) {
+ return builder().setUseTls(useTls).build();
+ }
+
+ /** Sets the Flight descriptor (table name or path) for the write target.
*/
+ public Write withDescriptor(String descriptor) {
+ return builder().setDescriptor(descriptor).build();
+ }
+
+ /** Sets the batch size for writing. Rows are buffered and flushed in
batches. */
+ public Write withBatchSize(int batchSize) {
+ checkArgument(batchSize > 0, "batchSize must be positive");
+ return builder().setBatchSize(batchSize).build();
+ }
+
+ /** Sets a bearer token for authentication. */
+ public Write withToken(byte[] token) {
+ return builder().setToken(copyToken(token)).build();
+ }
+
+ CallOption[] callOptions() {
+ return ArrowFlightIO.callOptions(token());
+ }
+
+ @Override
+ public PDone expand(PCollection<Row> input) {
+ checkArgument(host() != null, "withHost() is required");
+ checkArgument(descriptor() != null, "withDescriptor() is required");
+ Schema inputSchema = checkNotNull(input.getSchema(), "input schema");
+ validateWriteSchema(inputSchema);
+
+ input.apply(ParDo.of(new FlightWriteFn(this, inputSchema)));
+ return PDone.in(input.getPipeline());
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ builder.addIfNotNull(DisplayData.item("host", host()));
+ builder.add(DisplayData.item("port", port()));
+ builder.add(DisplayData.item("useTls", useTls()));
+ builder.addIfNotNull(DisplayData.item("descriptor", descriptor()));
+ builder.add(DisplayData.item("batchSize", batchSize()));
+ }
+ }
+
+ /** DoFn that buffers Beam Rows and streams them as Arrow record batches to
a Flight server. */
+ @SuppressWarnings("initialization.fields.uninitialized")
+ static class FlightWriteFn extends DoFn<Row, Void> {
+ private static final Counter RECORDS_WRITTEN =
+ Metrics.counter(ArrowFlightIO.class, "recordsWritten");
+ private static final Counter BATCHES_WRITTEN =
+ Metrics.counter(ArrowFlightIO.class, "batchesWritten");
+
+ private final Write spec;
+ private final Schema beamSchema;
+ private transient @Nullable BufferAllocator allocator;
+ private transient @Nullable FlightClient client;
+ private transient FlightClient.@Nullable ClientStreamListener listener;
+ private transient @Nullable VectorSchemaRoot root;
+ private transient List<Row> batch;
+
+ FlightWriteFn(Write spec, Schema beamSchema) {
+ this.spec = spec;
+ this.beamSchema = beamSchema;
+ }
+
+ @StartBundle
+ public void startBundle() {
+ batch = new ArrayList<>();
+ }
+
+ @ProcessElement
+ public void processElement(@Element Row row) {
+ checkArgument(
+ row.getSchema().equivalent(beamSchema),
+ "ArrowFlightIO.write() requires all rows to use the same schema.");
+ batch.add(row);
+ if (batch.size() >= spec.batchSize()) {
+ flush();
+ }
+ }
+
+ @FinishBundle
+ public void finishBundle() {
+ RuntimeException failure = null;
+ try {
+ flush();
+ } catch (RuntimeException e) {
+ failure = e;
+ }
+
+ try {
+ closeConnection();
+ } catch (RuntimeException e) {
+ if (failure == null) {
+ failure = e;
+ } else {
+ failure.addSuppressed(e);
+ }
+ }
+
+ if (failure != null) {
+ throw failure;
+ }
+ }
+
+ @Teardown
+ public void teardown() {
+ try {
+ closeConnection();
+ } catch (RuntimeException e) {
+ LOG.warn("Error closing Flight write connection during teardown", e);
+ }
+ }
+
+ private void ensureConnection() {
+ if (client == null) {
+ BufferAllocator currentAllocator = new RootAllocator(Long.MAX_VALUE);
+ allocator = currentAllocator;
+ FlightClient currentClient =
+ createClient(
+ currentAllocator, checkNotNull(spec.host(), "host"),
spec.port(), spec.useTls());
+ client = currentClient;
+
+ org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+ ArrowConversion.ArrowSchemaTranslator.toArrowSchema(beamSchema);
+ VectorSchemaRoot currentRoot = VectorSchemaRoot.create(arrowSchema,
currentAllocator);
+ root = currentRoot;
+
+ FlightDescriptor descriptor =
+ FlightDescriptor.path(checkNotNull(spec.descriptor(),
"descriptor"));
+ listener =
+ currentClient.startPut(
+ descriptor, currentRoot, new AsyncPutListener(),
spec.callOptions());
+ }
+ }
+
+ @SuppressWarnings("nullness")
+ private void flush() {
+ if (batch == null || batch.isEmpty()) {
+ return;
+ }
+ ensureConnection();
+
+ for (int colIdx = 0; colIdx < beamSchema.getFieldCount(); colIdx++) {
+ FieldVector vector = root.getVector(colIdx);
+ vector.allocateNew();
+ Schema.Field field = beamSchema.getField(colIdx);
+ for (int rowIdx = 0; rowIdx < batch.size(); rowIdx++) {
+ Object value = batch.get(rowIdx).getValue(colIdx);
+ if (value == null) {
+ vector.setNull(rowIdx);
+ } else {
+ setVectorValue(vector, rowIdx, value, field.getType());
+ }
+ }
+ vector.setValueCount(batch.size());
+ }
+ root.setRowCount(batch.size());
+
+ listener.putNext();
+ RECORDS_WRITTEN.inc(batch.size());
+ BATCHES_WRITTEN.inc();
+ root.clear();
+ batch.clear();
+ }
+
+ @SuppressWarnings("nullness")
+ private void setVectorValue(
+ FieldVector vector, int index, Object value, Schema.FieldType type) {
+ switch (type.getTypeName()) {
+ case BYTE:
+ ((TinyIntVector) vector).setSafe(index, ((Number)
value).byteValue());
+ break;
+ case INT16:
+ ((SmallIntVector) vector).setSafe(index, ((Number)
value).shortValue());
+ break;
+ case INT32:
+ ((IntVector) vector).setSafe(index, ((Number) value).intValue());
+ break;
+ case INT64:
+ ((BigIntVector) vector).setSafe(index, ((Number) value).longValue());
+ break;
+ case FLOAT:
+ ((Float4Vector) vector).setSafe(index, ((Number)
value).floatValue());
+ break;
+ case DOUBLE:
+ ((Float8Vector) vector).setSafe(index, ((Number)
value).doubleValue());
+ break;
+ case BOOLEAN:
+ ((BitVector) vector).setSafe(index, ((Boolean) value) ? 1 : 0);
+ break;
+ case STRING:
+ ((VarCharVector) vector)
+ .setSafe(index,
value.toString().getBytes(StandardCharsets.UTF_8));
+ break;
+ case BYTES:
+ ((VarBinaryVector) vector).setSafe(index, (byte[]) value);
+ break;
+ case DATETIME:
+ long millis;
+ if (value instanceof org.joda.time.ReadableInstant) {
+ millis = ((org.joda.time.ReadableInstant) value).getMillis();
+ } else {
+ millis = ((Number) value).longValue();
+ }
+ ((TimeStampMilliTZVector) vector).setSafe(index, millis);
+ break;
+ default:
+ throw new IllegalArgumentException(
+ "Unsupported Beam type for ArrowFlightIO.write(): " +
type.getTypeName());
+ }
+ }
+
+ private void closeConnection() {
+ RuntimeException failure = null;
+ FlightClient.ClientStreamListener currentListener = listener;
+ listener = null;
+ try {
+ if (currentListener != null) {
+ currentListener.completed();
+ currentListener.getResult();
+ }
+ } catch (RuntimeException e) {
+ failure = e;
+ }
+
+ VectorSchemaRoot currentRoot = root;
+ root = null;
+ try {
+ if (currentRoot != null) {
+ currentRoot.close();
+ }
+ } catch (Exception e) {
+ if (failure == null) {
+ failure = new RuntimeException("Error closing VectorSchemaRoot", e);
+ } else {
+ failure.addSuppressed(e);
+ }
+ }
+
+ FlightClient currentClient = client;
+ client = null;
+ try {
+ if (currentClient != null) {
+ currentClient.close();
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ RuntimeException closeFailure = new RuntimeException("Interrupted
closing FlightClient", e);
+ if (failure == null) {
+ failure = closeFailure;
+ } else {
+ failure.addSuppressed(closeFailure);
+ }
+ }
+
+ BufferAllocator currentAllocator = allocator;
+ allocator = null;
+ try {
+ if (currentAllocator != null) {
+ currentAllocator.close();
+ }
+ } catch (Exception e) {
+ if (failure == null) {
+ failure = new RuntimeException("Error closing BufferAllocator", e);
+ } else {
+ failure.addSuppressed(e);
+ }
+ }
+
+ if (failure != null) {
+ throw failure;
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/package-info.java
b/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/package-info.java
new file mode 100644
index 00000000000..6988db88c3a
--- /dev/null
+++
b/sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/package-info.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * I/O connector for <a
href="https://arrow.apache.org/docs/format/Flight.html";>Apache Arrow
+ * Flight</a>.
+ *
+ * <p>Arrow Flight is a high-performance RPC framework for fast data transport
using the Apache
+ * Arrow columnar format over gRPC. This connector enables Beam pipelines to
read from and write to
+ * Arrow Flight-compatible data systems.
+ *
+ * @see org.apache.beam.sdk.io.arrowflight.ArrowFlightIO
+ */
+package org.apache.beam.sdk.io.arrowflight;
diff --git
a/sdks/java/io/arrow-flight/src/test/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIOTest.java
b/sdks/java/io/arrow-flight/src/test/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIOTest.java
new file mode 100644
index 00000000000..7ef67b5f624
--- /dev/null
+++
b/sdks/java/io/arrow-flight/src/test/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIOTest.java
@@ -0,0 +1,330 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.arrowflight;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.arrow.flight.Action;
+import org.apache.arrow.flight.ActionType;
+import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightConstants;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.ServerHeaderMiddleware;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+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.beam.sdk.Pipeline;
+import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link ArrowFlightIO}. */
+@RunWith(JUnit4.class)
+public class ArrowFlightIOTest {
+
+ @Rule public final TestPipeline pipeline = TestPipeline.create();
+
+ private BufferAllocator allocator;
+ private FlightServer server;
+ private TestFlightProducer producer;
+ private int port;
+
+ @Before
+ public void setUp() throws Exception {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ producer = new TestFlightProducer(allocator);
+
+ // Bind to any available port
+ Location location = Location.forGrpcInsecure("localhost", 0);
+ server = FlightServer.builder(allocator, location, producer).build();
+ server.start();
+
+ port = server.getPort();
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ if (server != null) {
+ server.close();
+ }
+ if (allocator != null) {
+ allocator.close();
+ }
+ }
+
+ @Test
+ public void testRead() {
+ PCollection<Row> output =
+ pipeline.apply(
+ "Read from Flight",
+
ArrowFlightIO.read().withHost("localhost").withPort(port).withCommand("test_query"));
+
+ Schema expectedSchema = Schema.builder().addStringField("name").build();
+ Row expectedRow1 =
Row.withSchema(expectedSchema).addValue("Alice").build();
+ Row expectedRow2 = Row.withSchema(expectedSchema).addValue("Bob").build();
+
+ PAssert.that(output).containsInAnyOrder(expectedRow1, expectedRow2);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testCurrentSourceDoesNotResplitAfterReadStarts() throws
Exception {
+ producer.endpointCount = 2;
+ PipelineOptions options = PipelineOptionsFactory.create();
+ Schema schema = Schema.builder().addStringField("name").build();
+ ArrowFlightIO.Read read =
+
ArrowFlightIO.read().withHost("localhost").withPort(port).withCommand("test_query");
+ ArrowFlightIO.FlightBoundedSource source = new
ArrowFlightIO.FlightBoundedSource(read, schema);
+
+ assertEquals(2, source.split(0, options).size());
+ try (BoundedSource.BoundedReader<Row> reader =
source.createReader(options)) {
+ assertTrue(reader.start());
+ assertEquals(1, reader.getCurrentSource().split(0, options).size());
+ }
+ }
+
+ @Test
+ public void testWrite() throws Exception {
+ Schema expectedSchema = Schema.builder().addStringField("name").build();
+ Row row1 = Row.withSchema(expectedSchema).addValue("Charlie").build();
+ Row row2 = Row.withSchema(expectedSchema).addValue("Dave").build();
+
+ pipeline
+ .apply(Create.of(row1, row2).withRowSchema(expectedSchema))
+ .apply(
+ "Write to Flight",
+ ArrowFlightIO.write()
+ .withHost("localhost")
+ .withPort(port)
+ .withDescriptor("test_table"));
+
+ pipeline.run().waitUntilFinish();
+
+ assertEquals(2, producer.writtenRecords.get());
+ }
+
+ @Test
+ public void testWriteWithToken() throws Exception {
+ producer.requiredAuthorizationHeader = "Bearer test-token";
+
+ Schema expectedSchema = Schema.builder().addStringField("name").build();
+ Row row = Row.withSchema(expectedSchema).addValue("Charlie").build();
+
+ pipeline
+ .apply(Create.of(row).withRowSchema(expectedSchema))
+ .apply(
+ "Write to Flight with Token",
+ ArrowFlightIO.write()
+ .withHost("localhost")
+ .withPort(port)
+ .withDescriptor("test_table")
+ .withToken("test-token".getBytes(StandardCharsets.UTF_8)));
+
+ pipeline.run().waitUntilFinish();
+
+ assertEquals("Bearer test-token", producer.lastAuthorizationHeader.get());
+ assertEquals(1, producer.writtenRecords.get());
+ }
+
+ @Test
+ public void testWritePropagatesServerErrors() {
+ producer.failWrites = true;
+
+ Schema expectedSchema = Schema.builder().addStringField("name").build();
+ Row row = Row.withSchema(expectedSchema).addValue("Charlie").build();
+
+ pipeline
+ .apply(Create.of(row).withRowSchema(expectedSchema))
+ .apply(
+ "Write to Failing Flight Server",
+ ArrowFlightIO.write()
+ .withHost("localhost")
+ .withPort(port)
+ .withDescriptor("test_table"));
+
+ PipelineExecutionException exception =
+ assertThrows(PipelineExecutionException.class, () ->
pipeline.run().waitUntilFinish());
+ assertThat(exception.getMessage(), containsString("Rejected write"));
+ }
+
+ @Test
+ public void testWriteRejectsUnsupportedSchema() {
+ Schema unsupportedSchema =
+ Schema.builder().addArrayField("names",
Schema.FieldType.STRING).build();
+ Row row =
+
Row.withSchema(unsupportedSchema).addArray(Collections.singletonList("Charlie")).build();
+ Pipeline testPipeline = Pipeline.create();
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ testPipeline
+ .apply(Create.of(row).withRowSchema(unsupportedSchema))
+ .apply(
+ ArrowFlightIO.write()
+ .withHost("localhost")
+ .withPort(port)
+ .withDescriptor("test_table")));
+
+ assertThat(exception.getMessage(), containsString("does not support Beam
type 'ARRAY'"));
+ }
+
+ /** A simple FlightProducer that returns predefined data for reads and
counts writes. */
+ private static class TestFlightProducer implements FlightProducer {
+
+ private final BufferAllocator allocator;
+ final AtomicInteger writtenRecords = new AtomicInteger();
+ final AtomicReference<String> lastAuthorizationHeader = new
AtomicReference<>();
+ volatile int endpointCount = 1;
+ volatile boolean failWrites;
+ volatile String requiredAuthorizationHeader;
+
+ TestFlightProducer(BufferAllocator allocator) {
+ this.allocator = allocator;
+ }
+
+ @Override
+ public void getStream(CallContext context, Ticket ticket,
ServerStreamListener listener) {
+ org.apache.arrow.vector.types.pojo.Schema schema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ Collections.singletonList(
+ new Field(
+ "name", FieldType.nullable(new ArrowType.Utf8()),
Collections.emptyList())));
+
+ try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator))
{
+ listener.start(root);
+
+ VarCharVector vector = (VarCharVector) root.getVector("name");
+ vector.allocateNew();
+ vector.setSafe(0, "Alice".getBytes(StandardCharsets.UTF_8));
+ vector.setSafe(1, "Bob".getBytes(StandardCharsets.UTF_8));
+ vector.setValueCount(2);
+ root.setRowCount(2);
+
+ listener.putNext();
+ listener.completed();
+ } catch (Exception e) {
+ listener.error(e);
+ }
+ }
+
+ @Override
+ public void listFlights(
+ CallContext context, Criteria criteria, StreamListener<FlightInfo>
listener) {
+ listener.onCompleted();
+ }
+
+ @Override
+ public FlightInfo getFlightInfo(CallContext context, FlightDescriptor
descriptor) {
+ org.apache.arrow.vector.types.pojo.Schema schema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ Collections.singletonList(
+ new Field(
+ "name", FieldType.nullable(new ArrowType.Utf8()),
Collections.emptyList())));
+ List<FlightEndpoint> endpoints = new ArrayList<>();
+ for (int i = 0; i < endpointCount; i++) {
+ endpoints.add(
+ new FlightEndpoint(
+ new Ticket(descriptor.getCommand()),
Location.forGrpcInsecure("localhost", 0)));
+ }
+ return new FlightInfo(schema, descriptor, endpoints, -1, -1);
+ }
+
+ @Override
+ public Runnable acceptPut(
+ CallContext context, FlightStream flightStream,
StreamListener<PutResult> ackStream) {
+ ServerHeaderMiddleware headerMiddleware =
context.getMiddleware(FlightConstants.HEADER_KEY);
+ lastAuthorizationHeader.set(
+ headerMiddleware == null ? null :
headerMiddleware.headers().get("authorization"));
+
+ return () -> {
+ try {
+ if (requiredAuthorizationHeader != null
+ &&
!requiredAuthorizationHeader.equals(lastAuthorizationHeader.get())) {
+ ackStream.onError(
+ CallStatus.UNAUTHENTICATED
+ .withDescription("Missing or invalid authorization header")
+ .toRuntimeException());
+ return;
+ }
+ while (flightStream.next()) {
+ VectorSchemaRoot root = flightStream.getRoot();
+ writtenRecords.addAndGet(root.getRowCount());
+ }
+ if (failWrites) {
+ ackStream.onError(
+ CallStatus.INTERNAL.withDescription("Rejected
write").toRuntimeException());
+ return;
+ }
+ ackStream.onCompleted();
+ } catch (Exception e) {
+ ackStream.onError(e);
+ }
+ };
+ }
+
+ @Override
+ public void doAction(CallContext context, Action action,
StreamListener<Result> listener) {
+ listener.onCompleted();
+ }
+
+ @Override
+ public void listActions(CallContext context, StreamListener<ActionType>
listener) {
+ listener.onCompleted();
+ }
+ }
+}
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 8a116eeb394..a7cdfc70515 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -222,6 +222,7 @@ include(":sdks:java:harness:jmh")
include(":sdks:java:io:amazon-web-services2")
include(":sdks:java:io:amazon-web-services2:expansion-service")
include(":sdks:java:io:amqp")
+include(":sdks:java:io:arrow-flight")
include(":sdks:java:io:azure")
include(":sdks:java:io:azure-cosmos")
include(":sdks:java:io:cassandra")
diff --git a/website/www/site/content/en/documentation/io/connectors.md
b/website/www/site/content/en/documentation/io/connectors.md
index 92c4c55220d..242a255ba82 100644
--- a/website/www/site/content/en/documentation/io/connectors.md
+++ b/website/www/site/content/en/documentation/io/connectors.md
@@ -882,6 +882,22 @@ This table provides a consolidated, at-a-glance overview
of the available built-
<td class="absent">✘</td>
<td class="absent">✘</td>
</tr>
+ <tr>
+ <td>ArrowFlightIO</td>
+ <td class="present">✔</td>
+ <td class="present">✔</td>
+ <td class="present">
+ ✔
+ <a
href="https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/arrowflight/ArrowFlightIO.html";>native</a>
+ </td>
+ <td>Not available</td>
+ <td>Not available</td>
+ <td>Not available</td>
+ <td>Not available</td>
+ <td class="present">✔</td>
+ <td class="present">✔</td>
+ <td class="absent">✘</td>
+ </tr>
<tr>
<td>DatabaseIO</td>
<td class="present">✔</td>
