[Java IO] Add ArrowFlight IO connector

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
 

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")

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -938,6 +938,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",

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());
     }

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 =

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'
+}