ARROW-6199: [Java] Avro adapter avoid potential resource leak.

Author: tianchen92

java/adapter/avro/src/main/java/org/apache/arrow/AvroToArrowUtils.java

@@ -20,7 +20,6 @@
 import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE;
 import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -37,6 +36,7 @@
 import org.apache.arrow.consumers.AvroNullConsumer;
 import org.apache.arrow.consumers.AvroStringConsumer;
 import org.apache.arrow.consumers.AvroUnionsConsumer;
+import org.apache.arrow.consumers.CompositeAvroConsumer;
 import org.apache.arrow.consumers.Consumer;
 import org.apache.arrow.consumers.NullableTypeConsumer;
 import org.apache.arrow.memory.BufferAllocator;
@@ -246,19 +246,10 @@ public static VectorSchemaRoot avroToArrowVectors(Schema schema, Decoder decoder
 
     VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 0);
 
-    int valueCount = 0;
-    while (true) {
-      try {
-        for (Consumer consumer : consumers) {
-          consumer.consume(decoder);
-        }
-        valueCount++;
-        //reach end will throw EOFException.
-      } catch (EOFException eofException) {
-        root.setRowCount(valueCount);
-        break;
-      }
+    try (CompositeAvroConsumer compositeConsumer = new CompositeAvroConsumer(consumers)) {
+      compositeConsumer.consume(decoder, root);
     }
+
     return root;
   }
 }

java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroBytesConsumer.java

@@ -79,4 +79,11 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return vector;
   }
+
+  @Override
+  public void close() {
+    if (cacheBuffer != null) {
+      cacheBuffer.clear();
+    }
+  }
 }

java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroStringConsumer.java

@@ -80,4 +80,11 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() {
+    if (cacheBuffer != null) {
+      cacheBuffer.clear();
+    }
+  }
 }

java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroUnionsConsumer.java

@@ -31,7 +31,7 @@
  */
 public class AvroUnionsConsumer implements Consumer {
 
-  private Consumer[] indexDelegates;
+  private Consumer[] delegates;
   private Types.MinorType[] types;
 
   private UnionWriter writer;
@@ -40,11 +40,11 @@ public class AvroUnionsConsumer implements Consumer {
   /**
    * Instantiate a AvroUnionConsumer.
    */
-  public AvroUnionsConsumer(UnionVector vector, Consumer[] indexDelegates, Types.MinorType[] types) {
+  public AvroUnionsConsumer(UnionVector vector, Consumer[] delegates, Types.MinorType[] types) {
 
     this.writer = new UnionWriter(vector);
     this.vector = vector;
-    this.indexDelegates = indexDelegates;
+    this.delegates = delegates;
     this.types = types;
   }
 
@@ -53,7 +53,7 @@ public void consume(Decoder decoder) throws IOException {
     int fieldIndex = decoder.readInt();
     int position = writer.getPosition();
 
-    Consumer delegate = indexDelegates[fieldIndex];
+    Consumer delegate = delegates[fieldIndex];
 
     vector.setType(position, types[fieldIndex]);
     // In UnionVector we need to set sub vector writer position before consume a value
@@ -80,4 +80,11 @@ public FieldVector getVector() {
     vector.setValueCount(writer.getPosition());
     return this.vector;
   }
+
+  @Override
+  public void close() {
+    for (Consumer consumer : delegates) {
+      consumer.close();
+    }
+  }
 }

java/adapter/avro/src/main/java/org/apache/arrow/consumers/CompositeAvroConsumer.java

@@ -0,0 +1,65 @@
+/*
+ * 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.consumers;
+
+import java.io.EOFException;
+import java.io.IOException;
+import java.util.List;
+
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.avro.io.Decoder;
+
+/**
+ * Composite consumer which hold all consumers.
+ * It manages the consume and cleanup process.
+ */
+public class CompositeAvroConsumer implements AutoCloseable {
+
+  private final List<Consumer> consumers;
+
+  public CompositeAvroConsumer(List<Consumer> consumers) {
+    this.consumers = consumers;
+  }
+
+  /**
+   * Consume decoder data and write into {@link VectorSchemaRoot}.
+   */
+  public void consume(Decoder decoder, VectorSchemaRoot root) throws IOException {
+    int valueCount = 0;
+    while (true) {
+      try {
+        for (Consumer consumer : consumers) {
+          consumer.consume(decoder);
+        }
+        valueCount++;
+        //reach end will throw EOFException.
+      } catch (EOFException eofException) {
+        root.setRowCount(valueCount);
+        break;
+      }
+    }
+  }
+
+  @Override
+  public void close() {
+    // clean up
+    for (Consumer consumer : consumers) {
+      consumer.close();
+    }
+  }
+}

java/adapter/avro/src/main/java/org/apache/arrow/consumers/Consumer.java

@@ -25,7 +25,7 @@
 /**
  * Interface that is used to consume values from avro decoder.
  */
-public interface Consumer {
+public interface Consumer extends AutoCloseable {
 
   /**
    * Consume a specific type value from avro decoder and write it to vector.
@@ -48,4 +48,9 @@ public interface Consumer {
    * Get the vector within the consumer.
    */
   FieldVector getVector();
+
+  /**
+   * Close this consumer, do some clean work such as clear reuse ArrowBuf.
+   */
+  default void close() {}
 }

java/adapter/avro/src/main/java/org/apache/arrow/consumers/NullableTypeConsumer.java

@@ -65,4 +65,9 @@ public FieldVector getVector() {
     return delegate.getVector();
   }
 
+  @Override
+  public void close() {
+    delegate.close();
+  }
+
 }