ARROW-4213: [Flight] Fix incompatibilities between C++ and Java

I'm willing to update the patches here in case you want to resolve incompatibilities differently. The changes are:

- In C++, send/read the schema as the first message in a DoGet stream.
- In Java, encode the GetFlightInfo schema in a Flatbuffer Message payload.
- In C++, skip null columns when encoding IpcPayload.
- In C++, don't write the body tag when encoding IpcPayload if no buffers are present.
- In Java, always align buffers when serializing record batches.

Additionally:
- Add integration tests for Flight. They will fail when trying to test record batches with dictionaries, but otherwise pass.
- Generate an uberjar for Flight, and include Apache Commons CLI and gRPC in the uberjar.
- Explicitly add the generated gRPC/Protobuf sources to the Java build - for some reason, Maven was not picking them up.
- In Java FlightClient, if there's an exception, use it to resolve the VectorSchemaRoot CompletableFuture so that clients do not hang forever.

Author: David Li <[email protected]>
Author: Wes McKinney <[email protected]>

Closes apache#3477 from lihalite/arrow-4213 and squashes the following commits:

d5a3ebf <Wes McKinney> Fix comment re: flight integration testing procedure
ac337a2 <David Li> Update description of integration test client
f867959 <David Li> Fix Flake8 errors
8c7be86 <David Li> clang-format new code
142b97d <David Li> Allow integration test script to control Flight port
3aace90 <David Li> Properly wait for Flight integration test server to start
0d0a476 <David Li> Clean up style issues
4c83485 <David Li> Collect failures at end of integration tests
0f3cce4 <David Li> Always align buffers when writing record batches in Flight
3c057df <David Li> Implement Java server side of Flight integration tests
13c0700 <David Li> Set exception on root in FlightClient
4801257 <David Li> Assume schema in Flight GetInfo is encoded in a Message payload
321f898 <David Li> Include gRPC in Flight uberjar
d23a201 <David Li> Explicitly compile generated Protobuf sources for Flight
6885d34 <David Li> Don't write body tag when serializing IpcPayload if no buffers
92c4c13 <David Li> Skeleton of integration test server for Flight
9bf8525 <David Li> Fix segfault in Flight when sending nullary columns
aca752b <David Li> Read schemas from stream in Flight DoGet
2acc8be <David Li> Implement arrow::ipc::ReadSchema that works with a Message

Author: David Li

flight/pom.xml

@@ -48,19 +48,16 @@
       <groupId>io.grpc</groupId>
       <artifactId>grpc-netty</artifactId>
       <version>${dep.grpc.version}</version>
-      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-core</artifactId>
       <version>${dep.grpc.version}</version>
-      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-protobuf</artifactId>
       <version>${dep.grpc.version}</version>
-      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>io.netty</groupId>
@@ -75,11 +72,15 @@
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
     </dependency>
+    <dependency>
+      <groupId>commons-cli</groupId>
+      <artifactId>commons-cli</artifactId>
+      <version>1.4</version>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-stub</artifactId>
       <version>${dep.grpc.version}</version>
-      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>com.google.protobuf</groupId>
@@ -225,6 +226,43 @@
           </execution>
         </executions>
       </plugin>
+      <plugin> <!-- add generated sources to classpath -->
+        <groupId>org.codehaus.mojo</groupId>
+        <artifactId>build-helper-maven-plugin</artifactId>
+        <version>1.9.1</version>
+        <executions>
+          <execution>
+            <id>add-generated-sources-to-classpath</id>
+            <phase>generate-sources</phase>
+            <goals>
+              <goal>add-source</goal>
+            </goals>
+            <configuration>
+              <sources>
+                <source>${project.build.directory}/generated-sources/protobuf</source>
+              </sources>
+            </configuration>
+          </execution>
+        </executions>
+      </plugin>
+      <plugin>
+        <artifactId>maven-assembly-plugin</artifactId>
+        <version>3.0.0</version>
+        <configuration>
+          <descriptorRefs>
+            <descriptorRef>jar-with-dependencies</descriptorRef>
+          </descriptorRefs>
+        </configuration>
+        <executions>
+          <execution>
+            <id>make-assembly</id>
+            <phase>package</phase>
+            <goals>
+              <goal>single</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
     </plugins>
   </build>
 </project>

flight/src/main/java/org/apache/arrow/flight/ArrowMessage.java

@@ -22,6 +22,8 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import org.apache.arrow.flatbuf.Message;
@@ -52,10 +54,12 @@
 import io.grpc.MethodDescriptor.Marshaller;
 import io.grpc.internal.ReadableBuffer;
 import io.grpc.protobuf.ProtoUtils;
+
 import io.netty.buffer.ArrowBuf;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.Unpooled;
 
 /**
  * The in-memory representation of FlightData used to manage a stream of Arrow messages.
@@ -95,6 +99,18 @@ public static HeaderType getHeader(byte b) {
 
   }
 
+  // Pre-allocated buffers for padding serialized ArrowMessages.
+  private static List<ByteBuf> PADDING_BUFFERS = Arrays.asList(
+      null,
+      Unpooled.copiedBuffer(new byte[] { 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0, 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0, 0, 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0, 0, 0, 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0, 0, 0, 0, 0 }),
+      Unpooled.copiedBuffer(new byte[] { 0, 0, 0, 0, 0, 0, 0 })
+  );
+
   private final FlightDescriptor descriptor;
   private final Message message;
   private final List<ArrowBuf> bufs;
@@ -253,8 +269,17 @@ private InputStream asInputStream(BufferAllocator allocator) {
       cos.writeTag(FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED);
 
       int size = 0;
+      List<ByteBuf> allBufs = new ArrayList<>();
       for (ArrowBuf b : bufs) {
+        allBufs.add(b);
         size += b.readableBytes();
+        // [ARROW-4213] These buffers must be aligned to an 8-byte boundary in order to be readable from C++.
+        if (b.readableBytes() % 8 != 0) {
+          int paddingBytes = 8 - (b.readableBytes() % 8);
+          assert paddingBytes > 0 && paddingBytes < 8;
+          size += paddingBytes;
+          allBufs.add(PADDING_BUFFERS.get(paddingBytes).retain());
+        }
       }
       // rawvarint is used for length definition.
       cos.writeUInt32NoTag(size);
@@ -263,7 +288,7 @@ private InputStream asInputStream(BufferAllocator allocator) {
       ArrowBuf initialBuf = allocator.buffer(baos.size());
       initialBuf.writeBytes(baos.toByteArray());
       final CompositeByteBuf bb = new CompositeByteBuf(allocator.getAsByteBufAllocator(), true, bufs.size() + 1,
-          ImmutableList.<ByteBuf>builder().add(initialBuf).addAll(bufs).build());
+          ImmutableList.<ByteBuf>builder().add(initialBuf).addAll(allBufs).build());
       final ByteBufInputStream is = new DrainableByteBufInputStream(bb);
       return is;
     } catch (Exception ex) {

flight/src/main/java/org/apache/arrow/flight/FlightInfo.java

@@ -17,13 +17,22 @@
 
 package org.apache.arrow.flight;
 
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
 import java.util.List;
 import java.util.stream.Collectors;
 
 import org.apache.arrow.flight.impl.Flight;
 import org.apache.arrow.flight.impl.Flight.FlightGetInfo;
+import org.apache.arrow.vector.ipc.ReadChannel;
+import org.apache.arrow.vector.ipc.WriteChannel;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
 import org.apache.arrow.vector.types.pojo.Schema;
 
+import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream;
+
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 
@@ -45,8 +54,15 @@ public FlightInfo(Schema schema, FlightDescriptor descriptor, List<FlightEndpoin
   }
 
   FlightInfo(FlightGetInfo flightGetInfo) {
-    schema = flightGetInfo.getSchema().size() > 0 ?
-        Schema.deserialize(flightGetInfo.getSchema().asReadOnlyByteBuffer()) : new Schema(ImmutableList.of());
+    try {
+      final ByteBuffer schemaBuf = flightGetInfo.getSchema().asReadOnlyByteBuffer();
+      schema = flightGetInfo.getSchema().size() > 0 ?
+          MessageSerializer.deserializeSchema(
+              new ReadChannel(Channels.newChannel(new ByteBufferBackedInputStream(schemaBuf))))
+          : new Schema(ImmutableList.of());
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
     descriptor = new FlightDescriptor(flightGetInfo.getFlightDescriptor());
     endpoints = flightGetInfo.getEndpointList().stream().map(t -> new FlightEndpoint(t)).collect(Collectors.toList());
     bytes = flightGetInfo.getTotalBytes();
@@ -74,9 +90,16 @@ public List<FlightEndpoint> getEndpoints() {
   }
 
   FlightGetInfo toProtocol() {
+    // Encode schema in a Message payload
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    try {
+      MessageSerializer.serialize(new WriteChannel(Channels.newChannel(baos)), schema);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
     return Flight.FlightGetInfo.newBuilder()
         .addAllEndpoint(endpoints.stream().map(t -> t.toProtocol()).collect(Collectors.toList()))
-        .setSchema(ByteString.copyFrom(schema.toByteArray()))
+        .setSchema(ByteString.copyFrom(baos.toByteArray()))
         .setFlightDescriptor(descriptor.toProtocol())
         .setTotalBytes(FlightInfo.this.bytes)
         .setTotalRecords(records)

flight/src/main/java/org/apache/arrow/flight/FlightService.java

@@ -128,7 +128,8 @@ public boolean isCancelled() {
     @Override
     public void start(VectorSchemaRoot root) {
       responseObserver.onNext(new ArrowMessage(null, root.getSchema()));
-      unloader = new VectorUnloader(root, true, false);
+      // [ARROW-4213] We must align buffers to be compatible with other languages.
+      unloader = new VectorUnloader(root, true, true);
     }
 
     @Override

flight/src/main/java/org/apache/arrow/flight/FlightStream.java

@@ -178,6 +178,7 @@ public void onNext(ArrowMessage msg) {
     public void onError(Throwable t) {
       ex = t;
       queue.add(DONE_EX);
+      root.setException(t);
     }
 
     @Override

flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java

@@ -0,0 +1,108 @@
+/*
+ * 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.arrow.flight.example.integration;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.List;
+
+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.Location;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.ipc.ArrowFileWriter;
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.CommandLineParser;
+import org.apache.commons.cli.DefaultParser;
+import org.apache.commons.cli.Options;
+import org.apache.commons.cli.ParseException;
+
+/**
+ * An Example Flight Server that provides access to the InMemoryStore.
+ */
+class IntegrationTestClient {
+  private static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(IntegrationTestClient.class);
+  private final Options options;
+
+  private IntegrationTestClient() {
+    options = new Options();
+    options.addOption("a", "arrow", true, "arrow file");
+    options.addOption("j", "json", true, "json file");
+    options.addOption("host", true, "The host to connect to.");
+    options.addOption("port", true, "The port to connect to." );
+  }
+
+  public static void main(String[] args) {
+    try {
+      new IntegrationTestClient().run(args);
+    } catch (ParseException e) {
+      fatalError("Invalid parameters", e);
+    } catch (IOException e) {
+      fatalError("Error accessing files", e);
+    }
+  }
+
+  static void fatalError(String message, Throwable e) {
+    System.err.println(message);
+    System.err.println(e.getMessage());
+    LOGGER.error(message, e);
+    System.exit(1);
+  }
+
+  private void run(String[] args) throws ParseException, IOException {
+    CommandLineParser parser = new DefaultParser();
+    CommandLine cmd = parser.parse(options, args, false);
+
+    String fileName = cmd.getOptionValue("arrow");
+    if (fileName == null) {
+      throw new IllegalArgumentException("missing arrow file parameter");
+    }
+    File arrowFile = new File(fileName);
+    if (arrowFile.exists()) {
+      throw new IllegalArgumentException("arrow file already exists: " + arrowFile.getAbsolutePath());
+    }
+
+    final String host = cmd.getOptionValue("host", "localhost");
+    final int port = Integer.parseInt(cmd.getOptionValue("port", "31337"));
+
+    final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+    FlightClient client = new FlightClient(allocator, new Location(host, port));
+    FlightInfo info = client.getInfo(FlightDescriptor.path(cmd.getOptionValue("json")));
+    List<FlightEndpoint> endpoints = info.getEndpoints();
+    if (endpoints.isEmpty()) {
+      throw new RuntimeException("No endpoints returned from Flight server.");
+    }
+
+    FlightStream stream = client.getStream(info.getEndpoints().get(0).getTicket());
+    try (VectorSchemaRoot root = stream.getRoot();
+         FileOutputStream fileOutputStream = new FileOutputStream(arrowFile);
+         ArrowFileWriter arrowWriter = new ArrowFileWriter(root, new DictionaryProvider.MapDictionaryProvider(),
+                 fileOutputStream.getChannel())) {
+      while (stream.next()) {
+        arrowWriter.writeBatch();
+      }
+    }
+  }
+}