ARROW-10962: [FlightRPC][Java] fill in empty body buffer if needed

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

@@ -297,7 +297,34 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s
             // ignore unknown fields.
         }
       }
-
+      // Protobuf implementations can omit empty fields, such as body; for some message types, like RecordBatch,
+      // this will fail later as we still expect an empty buffer. In those cases only, fill in an empty buffer here -
+      // in other cases, like Schema, having an unexpected empty buffer will also cause failures.
+      // We don't fill in defaults for fields like header, for which there is no reasonable default, or for appMetadata
+      // or descriptor, which are intended to be empty in some cases.
+      if (header != null) {
+        switch (HeaderType.getHeader(header.headerType())) {
+          case SCHEMA:
+            // Ignore 0-length buffers in case a Protobuf implementation wrote it out
+            if (body != null && body.capacity() == 0) {
+              body.close();
+              body = null;
+            }
+            break;
+          case DICTIONARY_BATCH:
+          case RECORD_BATCH:
+            // A Protobuf implementation can skip 0-length bodies, so ensure we fill it in here
+            if (body == null) {
+              body = allocator.getEmpty();
+            }
+            break;
+          case NONE:
+          case TENSOR:
+          default:
+            // Do nothing
+            break;
+        }
+      }
       return new ArrowMessage(descriptor, header, appMetadata, body);
     } catch (Exception ioe) {
       throw new RuntimeException(ioe);

java/flight/flight-core/src/test/java/org/apache/arrow/flight/TestBasicOperation.java

@@ -17,6 +17,10 @@
 
 package org.apache.arrow.flight;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
@@ -33,12 +37,16 @@
 import org.apache.arrow.flight.FlightClient.ClientStreamListener;
 import org.apache.arrow.flight.impl.Flight;
 import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType;
+import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.ipc.message.IpcOption;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -48,6 +56,8 @@
 import com.google.common.base.Charsets;
 import com.google.protobuf.ByteString;
 
+import io.grpc.MethodDescriptor;
+
 /**
  * Test the operations of a basic flight service.
  */
@@ -317,6 +327,85 @@ private void test(BiConsumer<FlightClient, BufferAllocator> consumer) throws Exc
     }
   }
 
+  /** Helper method to convert an ArrowMessage into a Protobuf message. */
+  private Flight.FlightData arrowMessageToProtobuf(
+      MethodDescriptor.Marshaller<ArrowMessage> marshaller, ArrowMessage message) throws IOException {
+    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    try (final InputStream serialized = marshaller.stream(message)) {
+      final byte[] buf = new byte[1024];
+      while (true) {
+        int read = serialized.read(buf);
+        if (read < 0) {
+          break;
+        }
+        baos.write(buf, 0, read);
+      }
+    }
+    final byte[] serializedMessage = baos.toByteArray();
+    return Flight.FlightData.parseFrom(serializedMessage);
+  }
+
+  /** ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields). */
+  @Test
+  public void testProtobufRecordBatchCompatibility() throws Exception {
+    final Schema schema = new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true))));
+    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+         final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+      final VectorUnloader unloader = new VectorUnloader(root);
+      root.setRowCount(0);
+      final MethodDescriptor.Marshaller<ArrowMessage> marshaller = ArrowMessage.createMarshaller(allocator);
+      try (final ArrowMessage message = new ArrowMessage(unloader.getRecordBatch(), null, new IpcOption())) {
+        Assert.assertEquals(ArrowMessage.HeaderType.RECORD_BATCH, message.getMessageType());
+        // Should have at least one empty body buffer (there may be multiple for e.g. data and validity)
+        Iterator<ArrowBuf> iterator = message.getBufs().iterator();
+        Assert.assertTrue(iterator.hasNext());
+        while (iterator.hasNext()) {
+          Assert.assertEquals(0, iterator.next().capacity());
+        }
+        final Flight.FlightData protobufData = arrowMessageToProtobuf(marshaller, message)
+            .toBuilder()
+            .clearDataBody()
+            .build();
+        Assert.assertEquals(0, protobufData.getDataBody().size());
+        ArrowMessage parsedMessage = marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray()));
+        // Should have an empty body buffer
+        Iterator<ArrowBuf> parsedIterator = parsedMessage.getBufs().iterator();
+        Assert.assertTrue(parsedIterator.hasNext());
+        Assert.assertEquals(0, parsedIterator.next().capacity());
+        // Should have only one (the parser synthesizes exactly one); in the case of empty buffers, this is equivalent
+        Assert.assertFalse(parsedIterator.hasNext());
+        // Should not throw
+        final ArrowRecordBatch rb = parsedMessage.asRecordBatch();
+        Assert.assertEquals(rb.computeBodyLength(), 0);
+      }
+    }
+  }
+
+  /** ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields). */
+  @Test
+  public void testProtobufSchemaCompatibility() throws Exception {
+    final Schema schema = new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true))));
+    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      final MethodDescriptor.Marshaller<ArrowMessage> marshaller = ArrowMessage.createMarshaller(allocator);
+      Flight.FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]).toProtocol();
+      try (final ArrowMessage message = new ArrowMessage(descriptor, schema, new IpcOption())) {
+        Assert.assertEquals(ArrowMessage.HeaderType.SCHEMA, message.getMessageType());
+        // Should have no body buffers
+        Assert.assertFalse(message.getBufs().iterator().hasNext());
+        final Flight.FlightData protobufData = arrowMessageToProtobuf(marshaller, message)
+            .toBuilder()
+            .setDataBody(ByteString.EMPTY)
+            .build();
+        Assert.assertEquals(0, protobufData.getDataBody().size());
+        final ArrowMessage parsedMessage = marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray()));
+        // Should have no body buffers
+        Assert.assertFalse(parsedMessage.getBufs().iterator().hasNext());
+        // Should not throw
+        parsedMessage.asSchema();
+      }
+    }
+  }
+
   /**
    * An example FlightProducer for test purposes.
    */