[Java IO] Address ArrowFlightIO review feedback

bvolpato · bvolpato · commit 447a459f426e · 2026-07-15T08:21:48.000-04:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -185,7 +185,6 @@
 
 ## I/Os
 
-
 * DebeziumIO (Java): added `OffsetRetainer` interface and `FileSystemOffsetRetainer` implementation to persist and restore CDC offsets across pipeline restarts, and exposed `withStartOffset` / `withOffsetRetainer` on `DebeziumIO.Read` and the cross-language `ReadBuilder` ([#28248](https://github.com/apache/beam/issues/28248)).
 
 ## New Features / Improvements
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -939,7 +939,6 @@ class BeamModulePlugin implements Plugin<Project> {
         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",
-        arrow_flight_sql                            : "org.apache.arrow:flight-sql:$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
@@ -22,6 +22,7 @@
 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.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 @@
  */
 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 @@ private 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
@@ -83,6 +83,39 @@ public void toBeamSchema_convertsSimpleArrowSchema() {
     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
@@ -39,5 +39,6 @@ dependencies {
 }
 
 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
@@ -55,11 +55,6 @@
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
-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.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.RowCoder;
 import org.apache.beam.sdk.extensions.arrow.ArrowConversion;
@@ -84,9 +79,7 @@
  * 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 a high-performance RPC framework for transferring Arrow-formatted data over
- * gRPC. It enables 10-50x faster data transfer compared to JDBC/ODBC by avoiding
- * serialization/deserialization overhead.
+ * <p>Arrow Flight is an RPC framework for transferring Arrow-formatted data over gRPC.
  *
  * <h3>Reading from an Arrow Flight server</h3>
  *
@@ -115,6 +108,14 @@
  *         .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 {
 
@@ -136,25 +137,10 @@ private static CallOption[] callOptions(byte @Nullable [] token) {
   }
 
   private static void validateWriteSchema(Schema schema) {
-    for (Schema.Field field : schema.getFields()) {
-      switch (field.getType().getTypeName()) {
-        case BYTE:
-        case INT16:
-        case INT32:
-        case INT64:
-        case FLOAT:
-        case DOUBLE:
-        case STRING:
-        case BOOLEAN:
-        case BYTES:
-        case DATETIME:
-          break;
-        default:
-          throw new IllegalArgumentException(
-              String.format(
-                  "ArrowFlightIO.write() does not support Beam type '%s' for field '%s'.",
-                  field.getType().getTypeName(), field.getName()));
-      }
+    try {
+      ArrowConversion.ArrowSchemaTranslator.toArrowSchema(schema);
+    } catch (IllegalArgumentException e) {
+      throw new IllegalArgumentException("ArrowFlightIO.write(): " + e.getMessage(), e);
     }
   }
 
@@ -415,36 +401,48 @@ static class FlightBoundedReader extends BoundedSource.BoundedReader<Row> {
     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 hostName = checkNotNull(spec.host(), "host");
-      if (source.endpoint != null) {
-        String host = source.endpoint.getHost(hostName);
-        int port = source.endpoint.getPort(spec.port());
-        client = createClient(allocator, host, port, spec.useTls());
-        stream = client.getStream(source.endpoint.getTicket(), spec.callOptions());
-      } else {
-        client = createClient(allocator, hostName, spec.port(), spec.useTls());
-        FlightInfo info =
-            client.getInfo(
-                FlightDescriptor.command(
-                    checkNotNull(spec.command(), "command").getBytes(StandardCharsets.UTF_8)),
-                spec.callOptions());
-        List<FlightEndpoint> endpoints = info.getEndpoints();
-        if (endpoints.isEmpty()) {
-          return false;
+      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);
         }
-        stream = client.getStream(endpoints.get(0).getTicket(), spec.callOptions());
+        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();
     }
@@ -509,12 +507,13 @@ public void close() throws IOException {
 
     @Override
     public BoundedSource<Row> getCurrentSource() {
-      return source;
+      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> {
 
@@ -621,7 +620,6 @@ static class FlightWriteFn extends DoFn<Row, Void> {
     private transient @Nullable FlightClient client;
     private transient FlightClient.@Nullable ClientStreamListener listener;
     private transient @Nullable VectorSchemaRoot root;
-    private transient org.apache.arrow.vector.types.pojo.Schema arrowSchema;
     private transient List<Row> batch;
 
     FlightWriteFn(Write spec, Schema beamSchema) {
@@ -687,11 +685,8 @@ private void ensureConnection() {
                 currentAllocator, checkNotNull(spec.host(), "host"), spec.port(), spec.useTls());
         client = currentClient;
 
-        List<Field> arrowFields = new ArrayList<>();
-        for (Schema.Field beamField : beamSchema.getFields()) {
-          arrowFields.add(toArrowField(beamField));
-        }
-        arrowSchema = new org.apache.arrow.vector.types.pojo.Schema(arrowFields);
+        org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+            ArrowConversion.ArrowSchemaTranslator.toArrowSchema(beamSchema);
         VectorSchemaRoot currentRoot = VectorSchemaRoot.create(arrowSchema, currentAllocator);
         root = currentRoot;
 
@@ -703,43 +698,6 @@ private void ensureConnection() {
       }
     }
 
-    private Field toArrowField(Schema.Field beamField) {
-      ArrowType arrowType = beamTypeToArrowType(beamField.getType());
-      FieldType fieldType =
-          beamField.getType().getNullable()
-              ? FieldType.nullable(arrowType)
-              : FieldType.notNullable(arrowType);
-      return new Field(beamField.getName(), fieldType, Collections.emptyList());
-    }
-
-    private ArrowType beamTypeToArrowType(Schema.FieldType beamType) {
-      switch (beamType.getTypeName()) {
-        case BYTE:
-          return new ArrowType.Int(8, true);
-        case INT16:
-          return new ArrowType.Int(16, true);
-        case INT32:
-          return new ArrowType.Int(32, true);
-        case INT64:
-          return new ArrowType.Int(64, true);
-        case FLOAT:
-          return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
-        case DOUBLE:
-          return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
-        case STRING:
-          return ArrowType.Utf8.INSTANCE;
-        case BOOLEAN:
-          return ArrowType.Bool.INSTANCE;
-        case BYTES:
-          return ArrowType.Binary.INSTANCE;
-        case DATETIME:
-          return new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC");
-        default:
-          throw new IllegalArgumentException(
-              "Unsupported Beam type for ArrowFlightIO.write(): " + beamType.getTypeName());
-      }
-    }
-
     @SuppressWarnings("nullness")
     private void flush() {
       if (batch == null || batch.isEmpty()) {
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
@@ -22,8 +22,7 @@
  *
  * <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
- * any Arrow Flight-compatible data system, including Dremio, ClickHouse, Apache Doris, InfluxDB 3,
- * DataFusion, and custom Flight servers.
+ * Arrow Flight-compatible data systems.
  *
  * @see org.apache.beam.sdk.io.arrowflight.ArrowFlightIO
  */
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
diff --git a/settings.gradle.kts b/settings.gradle.kts

Original file line number	Diff line number	Diff line change
`@@ -39,5 +39,6 @@ dependencies {`
`39`	`39`	`}`
`40`	`40`
`41`	`41`	`test {`
	`42`	`+ // Keep aligned with ArrowFlightIO runtime guidance for Java 17+.`
`42`	`43`	`jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED'`
`43`	`44`	`}`
Original file line number	Diff line number	Diff line change
`@@ -22,8 +22,7 @@`
`22`	`22`	`*`
`23`	`23`	`* <p>Arrow Flight is a high-performance RPC framework for fast data transport using the Apache`
`24`	`24`	`* Arrow columnar format over gRPC. This connector enables Beam pipelines to read from and write to`
`25`		`- * any Arrow Flight-compatible data system, including Dremio, ClickHouse, Apache Doris, InfluxDB 3,`
`26`		`- * DataFusion, and custom Flight servers.`
	`25`	`+ * Arrow Flight-compatible data systems.`
`27`	`26`	`*`
`28`	`27`	`* @see org.apache.beam.sdk.io.arrowflight.ArrowFlightIO`
`29`	`28`	`*/`