ARROW-250: Fix for ARROW-246 may cause memory leaks

java/vector/src/main/codegen/templates/UnionVector.java

@@ -73,8 +73,6 @@ public UnionVector(MaterializedField field, BufferAllocator allocator, CallBack
     this.allocator = allocator;
     this.internalMap = new MapVector("internal", allocator, callBack);
     this.typeVector = internalMap.addOrGet("types", new MajorType(MinorType.UINT1, DataMode.REQUIRED), UInt1Vector.class);
-    this.typeVector.allocateNew();
-    this.typeVector.zeroVector();
     this.field.addChild(internalMap.getField().clone());
     this.majorType = field.getType();
     this.callBack = callBack;
@@ -193,6 +191,7 @@ public int getValueCapacity() {
 
   @Override
   public void close() {
+    clear();
   }
 
   @Override

java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java

@@ -155,6 +155,7 @@ private FieldWriter promoteToUnion() {
     tp.transfer();
     if (parentContainer != null) {
       unionVector = parentContainer.addOrGet(name, new MajorType(MinorType.UNION, DataMode.OPTIONAL), UnionVector.class);
+      unionVector.allocateNew();
     } else if (listVector != null) {
       unionVector = listVector.promoteToUnion();
     }

java/vector/src/test/java/org/apache/arrow/vector/DirtyRootAllocator.java

@@ -0,0 +1,53 @@
+/**
+ * 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.vector;
+
+import org.apache.arrow.memory.BufferManager;
+import org.apache.arrow.memory.RootAllocator;
+
+import io.netty.buffer.ArrowBuf;
+
+/**
+ * Root allocator that returns buffers pre-filled with a given value.<br>
+ * Useful for testing if value vectors are properly zeroing their buffers.
+ */
+public class DirtyRootAllocator extends RootAllocator {
+
+  private final byte fillValue;
+
+  public DirtyRootAllocator(final long limit, final byte fillValue) {
+    super(limit);
+    this.fillValue = fillValue;
+  }
+
+  @Override
+  public ArrowBuf buffer(int size) {
+    return buffer(size, null);
+  }
+
+  @Override
+  public ArrowBuf buffer(int size, BufferManager manager) {
+    ArrowBuf buffer = super.buffer(size, manager);
+    // contaminate the buffer
+    for (int i = 0; i < buffer.capacity(); i++) {
+      buffer.setByte(i, fillValue);
+    }
+
+    return buffer;
+  }
+}

java/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java

@@ -20,7 +20,6 @@
 import static org.junit.Assert.assertEquals;
 
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.complex.UnionVector;
 import org.apache.arrow.vector.holders.NullableUInt4Holder;
 import org.apache.arrow.vector.holders.UInt4Holder;
@@ -37,7 +36,7 @@ public class TestUnionVector {
 
   @Before
   public void init() {
-    allocator = new RootAllocator(Long.MAX_VALUE);
+    allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100);
   }
 
   @After
@@ -49,15 +48,13 @@ public void terminate() throws Exception {
   public void testUnionVector() throws Exception {
     final MaterializedField field = MaterializedField.create(EMPTY_SCHEMA_PATH, UInt4Holder.TYPE);
 
-    final BufferAllocator alloc = new DirtyBufferAllocator(allocator, (byte) 100);
-
-    UnionVector unionVector = new UnionVector(field, alloc, null);
-
     final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
     uInt4Holder.value = 100;
     uInt4Holder.isSet = 1;
 
-    try {
+    try (UnionVector unionVector = new UnionVector(field, allocator, null)) {
+      unionVector.allocateNew();
+
       // write some data
       final UnionVector.Mutator mutator = unionVector.getMutator();
       mutator.setType(0, Types.MinorType.UINT4);
@@ -79,9 +76,6 @@ public void testUnionVector() throws Exception {
       assertEquals(100, accessor.getObject(2));
 
       assertEquals(true, accessor.isNull(3));
-
-    } finally {
-      unionVector.clear();
     }
   }
 

java/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java

@@ -0,0 +1,98 @@
+/**
+ * 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.vector.complex.impl;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.DirtyRootAllocator;
+import org.apache.arrow.vector.complex.AbstractMapVector;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.UnionVector;
+import org.apache.arrow.vector.holders.UInt4Holder;
+import org.apache.arrow.vector.types.MaterializedField;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestPromotableWriter {
+  private final static String EMPTY_SCHEMA_PATH = "";
+
+  private BufferAllocator allocator;
+
+  @Before
+  public void init() {
+    allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100);
+  }
+
+  @After
+  public void terminate() throws Exception {
+    allocator.close();
+  }
+
+  @Test
+  public void testPromoteToUnion() throws Exception {
+    final MaterializedField field = MaterializedField.create(EMPTY_SCHEMA_PATH, UInt4Holder.TYPE);
+
+    try (final AbstractMapVector container = new MapVector(field, allocator, null);
+         final MapVector v = container.addOrGet("test", MapVector.TYPE, MapVector.class);
+         final PromotableWriter writer = new PromotableWriter(v, container)) {
+
+      container.allocateNew();
+
+      writer.start();
+
+      writer.setPosition(0);
+      writer.bit("A").writeBit(0);
+
+      writer.setPosition(1);
+      writer.bit("A").writeBit(1);
+
+      writer.setPosition(2);
+      writer.integer("A").writeInt(10);
+
+      // we don't write anything in 3
+
+      writer.setPosition(4);
+      writer.integer("A").writeInt(100);
+
+      writer.end();
+
+      container.getMutator().setValueCount(5);
+
+      final UnionVector uv = v.getChild("A", UnionVector.class);
+      final UnionVector.Accessor accessor = uv.getAccessor();
+
+      assertFalse("0 shouldn't be null", accessor.isNull(0));
+      assertEquals(false, accessor.getObject(0));
+
+      assertFalse("1 shouldn't be null", accessor.isNull(1));
+      assertEquals(true, accessor.getObject(1));
+
+      assertFalse("2 shouldn't be null", accessor.isNull(2));
+      assertEquals(10, accessor.getObject(2));
+
+      assertTrue("3 should be null", accessor.isNull(3));
+
+      assertFalse("4 shouldn't be null", accessor.isNull(4));
+      assertEquals(100, accessor.getObject(4));
+    }
+  }
+}