Properly release buffers in Flight integration client

Author: lidavidm

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

@@ -194,6 +194,8 @@ public ArrowDictionaryBatch asDictionaryBatch() throws IOException {
     Preconditions.checkArgument(bufs.size() == 1, "A batch can only be consumed if it contains a single ArrowBuf.");
     Preconditions.checkArgument(getMessageType() == HeaderType.DICTIONARY_BATCH);
     ArrowBuf underlying = bufs.get(0);
+    // Retain a reference to keep the batch alive when the message is closed
+    underlying.getReferenceManager().retain();
     return MessageSerializer.deserializeDictionaryBatch(message, underlying);
   }
 

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

@@ -95,10 +95,30 @@ public Schema getSchema() {
     return schema;
   }
 
+  /**
+   * Get the provider for dictionaries in this stream.
+   *
+   * <p>Does NOT retain a reference to the underlying dictionaries. Dictionaries may be updated as the stream is read.
+   *
+   * @see #retainDictionaries()
+   */
   public DictionaryProvider getDictionaryProvider() {
     return dictionaries;
   }
 
+  /**
+   * Retain a reference to each of the dictionaries in the stream. Should be called after finishing reading the stream,
+   * but before closing.
+   *
+   * @return The current dictionaries in the stream.
+   */
+  public DictionaryProvider retainDictionaries() {
+    // Swap out the provider so it is not closed
+    final DictionaryProvider provider = dictionaries;
+    dictionaries = new DictionaryProvider.MapDictionaryProvider();
+    return provider;
+  }
+
   public FlightDescriptor getDescriptor() {
     return descriptor;
   }
@@ -117,8 +137,12 @@ public void close() throws Exception {
         .map(t -> ((AutoCloseable) t))
         .collect(Collectors.toList());
 
+    final List<FieldVector> dictionaryVectors = dictionaries.getDictionaryIds().stream()
+        .map(id -> dictionaries.lookup(id).getVector()).collect(Collectors.toList());
+
     // Must check for null since ImmutableList doesn't accept nulls
     AutoCloseables.close(Iterables.concat(closeables,
+        dictionaryVectors,
         applicationMetadata != null ? ImmutableList.of(root.get(), applicationMetadata)
             : ImmutableList.of(root.get())));
   }

java/flight/src/main/java/org/apache/arrow/flight/example/ExampleFlightServer.java

@@ -26,7 +26,7 @@
 import org.apache.arrow.util.AutoCloseables;
 
 /**
- * An Example Flight Server that provides access to the InMemoryStore.
+ * An Example Flight Server that provides access to the InMemoryStore. Used for integration testing.
  */
 public class ExampleFlightServer implements AutoCloseable {
 

java/flight/src/main/java/org/apache/arrow/flight/example/FlightHolder.java

@@ -18,7 +18,9 @@
 package org.apache.arrow.flight.example;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Collectors;
 
@@ -31,6 +33,7 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.DictionaryUtility;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
@@ -106,6 +109,13 @@ public FlightInfo getFlightInfo(final Location l) {
 
   @Override
   public void close() throws Exception {
-    AutoCloseables.close(Iterables.concat(streams, ImmutableList.of(allocator)));
+    // Close dictionaries
+    final Set<Long> dictionaryIds = new HashSet<>();
+    schema.getFields().forEach(field -> DictionaryUtility.toMessageFormat(field, dictionaryProvider, dictionaryIds));
+
+    final Iterable<AutoCloseable> dictionaries = dictionaryIds.stream()
+        .map(id -> (AutoCloseable) dictionaryProvider.lookup(id).getVector())::iterator;
+
+    AutoCloseables.close(Iterables.concat(streams, ImmutableList.of(allocator), dictionaries));
   }
 }

java/flight/src/main/java/org/apache/arrow/flight/example/InMemoryStore.java

@@ -39,7 +39,7 @@
 import org.apache.arrow.vector.VectorUnloader;
 
 /**
- * A FlightProducer that hosts an in memory store of Arrow buffers.
+ * A FlightProducer that hosts an in memory store of Arrow buffers. Used for integration testing.
  */
 public class InMemoryStore implements FlightProducer, AutoCloseable {
 
@@ -80,8 +80,7 @@ public Stream getStream(Ticket t) {
   }
 
   @Override
-  public void listFlights(CallContext context, Criteria criteria,
-      StreamListener<FlightInfo> listener) {
+  public void listFlights(CallContext context, Criteria criteria, StreamListener<FlightInfo> listener) {
     try {
       for (FlightHolder h : holders.values()) {
         listener.onNext(h.getFlightInfo(location));
@@ -93,8 +92,7 @@ public void listFlights(CallContext context, Criteria criteria,
   }
 
   @Override
-  public FlightInfo getFlightInfo(CallContext context,
-      FlightDescriptor descriptor) {
+  public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
     FlightHolder h = holders.get(descriptor);
     if (h == null) {
       throw new IllegalStateException("Unknown descriptor.");
@@ -121,6 +119,8 @@ public Runnable acceptPut(CallContext context,
           ackStream.onNext(PutResult.metadata(flightStream.getLatestMetadata()));
           creator.add(unloader.getRecordBatch());
         }
+        // Closing the stream will release the dictionaries
+        flightStream.retainDictionaries();
         creator.complete();
         success = true;
       } finally {

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

@@ -37,6 +37,7 @@
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
 import org.apache.arrow.vector.ipc.JsonFileReader;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.util.Validator;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.CommandLineParser;
@@ -84,12 +85,19 @@ private void run(String[] args) throws ParseException, IOException {
     final String host = cmd.getOptionValue("host", "localhost");
     final int port = Integer.parseInt(cmd.getOptionValue("port", "31337"));
 
-    final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
     final Location defaultLocation = Location.forGrpcInsecure(host, port);
-    final FlightClient client = FlightClient.builder(allocator, defaultLocation).build();
+    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+        final FlightClient client = FlightClient.builder(allocator, defaultLocation).build()) {
 
-    final String inputPath = cmd.getOptionValue("j");
+      final String inputPath = cmd.getOptionValue("j");
+      testStream(allocator, defaultLocation, client, inputPath);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
+  }
 
+  private static void testStream(BufferAllocator allocator, Location server, FlightClient client, String inputPath)
+      throws IOException {
     // 1. Read data from JSON and upload to server.
     FlightDescriptor descriptor = FlightDescriptor.path(inputPath);
     VectorSchemaRoot jsonRoot;
@@ -121,7 +129,9 @@ public void onNext(PutResult val) {
         metadata.writeBytes(rawMetadata);
         // Transfers ownership of the buffer, so do not release it ourselves
         stream.putNext(metadata);
-        jsonLoader.load(unloader.getRecordBatch());
+        try (final ArrowRecordBatch arb = unloader.getRecordBatch()) {
+          jsonLoader.load(arb);
+        }
         root.clear();
         counter++;
       }
@@ -141,25 +151,29 @@ public void onNext(PutResult val) {
       // 3. Download the data from the server.
       List<Location> locations = endpoint.getLocations();
       if (locations.size() == 0) {
-        locations = Collections.singletonList(defaultLocation);
+        locations = Collections.singletonList(server);
       }
       for (Location location : locations) {
         System.out.println("Verifying location " + location.getUri());
-        FlightClient readClient = FlightClient.builder(allocator, location).build();
-        FlightStream stream = readClient.getStream(endpoint.getTicket());
-        VectorSchemaRoot downloadedRoot;
-        try (VectorSchemaRoot root = stream.getRoot()) {
-          downloadedRoot = VectorSchemaRoot.create(root.getSchema(), allocator);
+        try (FlightClient readClient = FlightClient.builder(allocator, location).build();
+            FlightStream stream = readClient.getStream(endpoint.getTicket());
+            VectorSchemaRoot root = stream.getRoot();
+            VectorSchemaRoot downloadedRoot = VectorSchemaRoot.create(root.getSchema(), allocator)) {
           VectorLoader loader = new VectorLoader(downloadedRoot);
           VectorUnloader unloader = new VectorUnloader(root);
           while (stream.next()) {
-            loader.load(unloader.getRecordBatch());
+            try (final ArrowRecordBatch arb = unloader.getRecordBatch()) {
+              loader.load(arb);
+            }
           }
-        }
 
-        // 4. Validate the data.
-        Validator.compareVectorSchemaRoot(jsonRoot, downloadedRoot);
+          // 4. Validate the data.
+          Validator.compareVectorSchemaRoot(jsonRoot, downloadedRoot);
+        } catch (Exception e) {
+          throw new RuntimeException(e);
+        }
       }
     }
+    jsonRoot.close();
   }
 }