Skip to content

Commit 8473d90

Browse files
bvolpatobvolpato
authored
[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
1 parent f20da8a commit 8473d90

11 files changed

Lines changed: 1357 additions & 0 deletions

File tree

‎CHANGES.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
## I/Os
6666

6767
* Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
68+
* Add ArrowFlight IO (Java) ([#20116](https://github.com/apache/beam/issues/20116)).
6869

6970
## New Features / Improvements
7071

‎build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,7 @@ tasks.register("javaPreCommit") {
339339
// a precommit task build multiple IOs (except those splitting into single jobs)
340340
tasks.register("javaioPreCommit") {
341341
dependsOn(":sdks:java:io:amqp:build")
342+
dependsOn(":sdks:java:io:arrow-flight:build")
342343
// CassandraIO, HBaseIO and HCatalogIO do not support Java17+, test ran separately
343344
// dependsOn(":sdks:java:io:cassandra:build")
344345
dependsOn(":sdks:java:io:csv:build")

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -951,6 +951,7 @@ class BeamModulePlugin implements Plugin<Project> {
951951
arrow_vector : "org.apache.arrow:arrow-vector:$arrow_version",
952952
arrow_memory_core : "org.apache.arrow:arrow-memory-core:$arrow_version",
953953
arrow_memory_netty : "org.apache.arrow:arrow-memory-netty:$arrow_version",
954+
arrow_flight_core : "org.apache.arrow:flight-core:$arrow_version",
954955
],
955956
groovy: [
956957
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‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import java.io.IOException;
2323
import java.io.InputStream;
2424
import java.nio.channels.Channels;
25+
import java.util.Collections;
2526
import java.util.Iterator;
2627
import java.util.List;
2728
import java.util.Optional;
@@ -35,6 +36,7 @@
3536
import org.apache.arrow.vector.ipc.ReadChannel;
3637
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
3738
import org.apache.arrow.vector.ipc.message.MessageSerializer;
39+
import org.apache.arrow.vector.types.FloatingPointPrecision;
3840
import org.apache.arrow.vector.types.TimeUnit;
3941
import org.apache.arrow.vector.types.pojo.ArrowType;
4042
import org.apache.arrow.vector.util.Text;
@@ -57,6 +59,57 @@
5759
*/
5860
public class ArrowConversion {
5961

62+
/** Get Arrow Field from Beam Field. */
63+
private static org.apache.arrow.vector.types.pojo.Field toArrowField(Field field) {
64+
FieldType beamFieldType = field.getType();
65+
ArrowType arrowType;
66+
// TODO: Support aggregate and logical Beam field types.
67+
switch (beamFieldType.getTypeName()) {
68+
case BYTE:
69+
arrowType = new ArrowType.Int(8, true);
70+
break;
71+
case INT16:
72+
arrowType = new ArrowType.Int(16, true);
73+
break;
74+
case INT32:
75+
arrowType = new ArrowType.Int(32, true);
76+
break;
77+
case INT64:
78+
arrowType = new ArrowType.Int(64, true);
79+
break;
80+
case FLOAT:
81+
arrowType = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
82+
break;
83+
case DOUBLE:
84+
arrowType = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
85+
break;
86+
case STRING:
87+
arrowType = ArrowType.Utf8.INSTANCE;
88+
break;
89+
case BOOLEAN:
90+
arrowType = ArrowType.Bool.INSTANCE;
91+
break;
92+
case BYTES:
93+
arrowType = ArrowType.Binary.INSTANCE;
94+
break;
95+
case DATETIME:
96+
arrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC");
97+
break;
98+
default:
99+
throw new IllegalArgumentException(
100+
String.format(
101+
"Arrow schema conversion does not support Beam type '%s' for field '%s'.",
102+
beamFieldType.getTypeName(), field.getName()));
103+
}
104+
105+
org.apache.arrow.vector.types.pojo.FieldType arrowFieldType =
106+
beamFieldType.getNullable()
107+
? org.apache.arrow.vector.types.pojo.FieldType.nullable(arrowType)
108+
: org.apache.arrow.vector.types.pojo.FieldType.notNullable(arrowType);
109+
return new org.apache.arrow.vector.types.pojo.Field(
110+
field.getName(), arrowFieldType, Collections.emptyList());
111+
}
112+
60113
/** Get Beam Field from Arrow Field. */
61114
private static Field toBeamField(org.apache.arrow.vector.types.pojo.Field field) {
62115
FieldType beamFieldType = toFieldType(field.getFieldType(), field.getChildren());
@@ -546,6 +599,14 @@ private ArrowConversion() {}
546599
/** Converts Arrow schema to Beam row schema. */
547600
public static class ArrowSchemaTranslator {
548601

602+
/** Converts a supported Beam row schema to an Arrow schema. */
603+
public static org.apache.arrow.vector.types.pojo.Schema toArrowSchema(Schema schema) {
604+
return new org.apache.arrow.vector.types.pojo.Schema(
605+
schema.getFields().stream()
606+
.map(ArrowConversion::toArrowField)
607+
.collect(Collectors.toList()));
608+
}
609+
549610
public static Schema toBeamSchema(org.apache.arrow.vector.types.pojo.Schema schema) {
550611
return toBeamSchema(schema.getFields());
551612
}

‎sdks/java/extensions/arrow/src/test/java/org/apache/beam/sdk/extensions/arrow/ArrowConversionTest.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,39 @@ public void toBeamSchema_convertsSimpleArrowSchema() {
8383
assertThat(ArrowConversion.ArrowSchemaTranslator.toBeamSchema(arrowSchema), equalTo(expected));
8484
}
8585

86+
@Test
87+
public void toArrowSchema_convertsSimpleBeamSchema() {
88+
Schema beamSchema =
89+
Schema.builder()
90+
.addByteField("int8")
91+
.addInt16Field("int16")
92+
.addInt32Field("int32")
93+
.addInt64Field("int64")
94+
.addFloatField("float32")
95+
.addDoubleField("float64")
96+
.addNullableField("string", FieldType.STRING)
97+
.addBooleanField("boolean")
98+
.addByteArrayField("bytes")
99+
.addDateTimeField("timestamp")
100+
.build();
101+
102+
org.apache.arrow.vector.types.pojo.Schema expected =
103+
new org.apache.arrow.vector.types.pojo.Schema(
104+
ImmutableList.of(
105+
field("int8", new ArrowType.Int(8, true)),
106+
field("int16", new ArrowType.Int(16, true)),
107+
field("int32", new ArrowType.Int(32, true)),
108+
field("int64", new ArrowType.Int(64, true)),
109+
field("float32", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)),
110+
field("float64", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)),
111+
field("string", true, ArrowType.Utf8.INSTANCE),
112+
field("boolean", ArrowType.Bool.INSTANCE),
113+
field("bytes", ArrowType.Binary.INSTANCE),
114+
field("timestamp", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))));
115+
116+
assertThat(ArrowConversion.ArrowSchemaTranslator.toArrowSchema(beamSchema), equalTo(expected));
117+
}
118+
86119
@Test
87120
public void rowIterator() {
88121
org.apache.arrow.vector.types.pojo.Schema schema =

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* License); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an AS IS BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
plugins { id 'org.apache.beam.module' }
20+
applyJavaNature(automaticModuleName: 'org.apache.beam.sdk.io.arrowflight')
21+
22+
description = "Apache Beam :: SDKs :: Java :: IO :: Arrow Flight"
23+
ext.summary = "IO to read and write data using Apache Arrow Flight RPC."
24+
25+
dependencies {
26+
implementation project(path: ":sdks:java:core", configuration: "shadow")
27+
implementation project(path: ":sdks:java:extensions:arrow")
28+
implementation library.java.joda_time
29+
implementation library.java.slf4j_api
30+
implementation library.java.vendored_guava_32_1_2_jre
31+
implementation(library.java.arrow_flight_core)
32+
implementation(library.java.arrow_memory_core)
33+
implementation(library.java.arrow_vector)
34+
testImplementation library.java.hamcrest
35+
testImplementation library.java.junit
36+
testImplementation(library.java.arrow_memory_netty)
37+
testRuntimeOnly library.java.slf4j_simple
38+
testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
39+
}
40+
41+
test {
42+
// Keep aligned with ArrowFlightIO runtime guidance for Java 17+.
43+
jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED'
44+
}

0 commit comments

Comments
 (0)