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

Author: bvolpato

build.gradle.kts

@@ -343,6 +343,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")
   dependsOn(":sdks:java:io:cassandra:build")
   dependsOn(":sdks:java:io:csv:build")
   dependsOn(":sdks:java:io:cdap:build")

sdks/java/io/arrow-flight/build.gradle

@@ -37,3 +37,7 @@ dependencies {
   testRuntimeOnly library.java.slf4j_simple
   testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
 }
+
+test {
+  jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED'
+}

sdks/java/io/arrow-flight/src/main/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIO.java

@@ -192,7 +192,7 @@ int getPort(int defaultPort) {
   @AutoValue
   public abstract static class Read extends PTransform<PBegin, PCollection<Row>> {
 
-    abstract String host();
+    abstract @Nullable String host();
 
     abstract int port();
 
@@ -252,7 +252,8 @@ public PCollection<Row> expand(PBegin input) {
 
       Schema beamSchema;
       try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
-          FlightClient client = createClient(allocator, host(), port(), useTls())) {
+          FlightClient client =
+              createClient(allocator, checkNotNull(host(), "host"), port(), useTls())) {
         FlightInfo info =
             client.getInfo(
                 FlightDescriptor.command(
@@ -283,7 +284,7 @@ HeaderCallOption[] callOptions() {
     @Override
     public void populateDisplayData(DisplayData.Builder builder) {
       super.populateDisplayData(builder);
-      builder.add(DisplayData.item("host", host()));
+      builder.addIfNotNull(DisplayData.item("host", host()));
       builder.add(DisplayData.item("port", port()));
       builder.add(DisplayData.item("useTls", useTls()));
       builder.addIfNotNull(DisplayData.item("command", command()));
@@ -315,15 +316,18 @@ public List<? extends BoundedSource<Row>> split(
 
       List<BoundedSource<Row>> sources = new ArrayList<>();
       try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
-          FlightClient client = createClient(allocator, spec.host(), spec.port(), spec.useTls())) {
+          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, spec.host(), spec.port());
+              SerializableEndpoint.fromFlightEndpoint(
+                  fe, checkNotNull(spec.host(), "host"), spec.port());
           sources.add(new FlightBoundedSource(spec, beamSchema, se));
         }
       }
@@ -340,7 +344,9 @@ public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
         return -1;
       }
       try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
-          FlightClient client = createClient(allocator, spec.host(), spec.port(), spec.useTls())) {
+          FlightClient client =
+              createClient(
+                  allocator, checkNotNull(spec.host(), "host"), spec.port(), spec.useTls())) {
         FlightInfo info =
             client.getInfo(
                 FlightDescriptor.command(
@@ -388,13 +394,14 @@ 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(spec.host());
+        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, spec.host(), spec.port(), spec.useTls());
+        client = createClient(allocator, hostName, spec.port(), spec.useTls());
         FlightInfo info =
             client.getInfo(
                 FlightDescriptor.command(
@@ -422,7 +429,15 @@ public boolean advance() throws IOException {
         if (stream.next()) {
           VectorSchemaRoot root = stream.getRoot();
           if (root.getRowCount() > 0) {
-            currentBatchIterator = ArrowConversion.rowsFromRecordBatch(source.beamSchema, root);
+            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;
@@ -472,7 +487,7 @@ public BoundedSource<Row> getCurrentSource() {
   @AutoValue
   public abstract static class Write extends PTransform<PCollection<Row>, PDone> {
 
-    abstract String host();
+    abstract @Nullable String host();
 
     abstract int port();
 
@@ -547,7 +562,7 @@ public PDone expand(PCollection<Row> input) {
     @Override
     public void populateDisplayData(DisplayData.Builder builder) {
       super.populateDisplayData(builder);
-      builder.add(DisplayData.item("host", host()));
+      builder.addIfNotNull(DisplayData.item("host", host()));
       builder.add(DisplayData.item("port", port()));
       builder.add(DisplayData.item("useTls", useTls()));
       builder.addIfNotNull(DisplayData.item("descriptor", descriptor()));
@@ -665,7 +680,6 @@ private void flush() {
       }
       ensureConnection();
 
-      root.setRowCount(batch.size());
       for (int colIdx = 0; colIdx < beamSchema.getFieldCount(); colIdx++) {
         FieldVector vector = root.getVector(colIdx);
         vector.allocateNew();
@@ -680,6 +694,7 @@ private void flush() {
         }
         vector.setValueCount(batch.size());
       }
+      root.setRowCount(batch.size());
 
       listener.putNext();
       RECORDS_WRITTEN.inc(batch.size());

sdks/java/io/arrow-flight/src/test/java/org/apache/beam/sdk/io/arrowflight/ArrowFlightIOTest.java

@@ -21,6 +21,7 @@
 
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.arrow.flight.Action;
 import org.apache.arrow.flight.ActionType;
 import org.apache.arrow.flight.Criteria;
@@ -121,14 +122,14 @@ public void testWrite() throws Exception {
 
     pipeline.run().waitUntilFinish();
 
-    assertEquals(2, producer.writtenRecords);
+    assertEquals(2, producer.writtenRecords.get());
   }
 
   /** A simple FlightProducer that returns predefined data for reads and counts writes. */
   private static class TestFlightProducer implements FlightProducer {
 
     private final BufferAllocator allocator;
-    int writtenRecords = 0;
+    final AtomicInteger writtenRecords = new AtomicInteger();
 
     TestFlightProducer(BufferAllocator allocator) {
       this.allocator = allocator;
@@ -189,7 +190,7 @@ public Runnable acceptPut(
         try {
           while (flightStream.next()) {
             VectorSchemaRoot root = flightStream.getRoot();
-            writtenRecords += root.getRowCount();
+            writtenRecords.addAndGet(root.getRowCount());
           }
           ackStream.onCompleted();
         } catch (Exception e) {