DRILL-3155: Part 2: Clear allocated memory for composite vectors if o…

…ne of the allocations fails

exec/java-exec/src/main/codegen/templates/NullableValueVectors.java

@@ -124,8 +124,21 @@ public void allocateNew() {
 
   @Override
   public boolean allocateNewSafe() {
-    if(!values.allocateNewSafe()) return false;
-    if(!bits.allocateNewSafe()) return false;
+    /* Boolean to keep track if all the memory allocations were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      if(!values.allocateNewSafe()) return false;
+      if(!bits.allocateNewSafe()) return false;
+      success = true;
+    } finally {
+      if (!success) {
+        clear();
+      }
+    }
     bits.zeroVector();
     mutator.reset();
     accessor.reset();
@@ -134,8 +147,13 @@ public boolean allocateNewSafe() {
 
   @Override
   public void allocateNew(int totalBytes, int valueCount) {
-    values.allocateNew(totalBytes, valueCount);
-    bits.allocateNew(valueCount);
+    try {
+      values.allocateNew(totalBytes, valueCount);
+      bits.allocateNew(valueCount);
+    } catch(OutOfMemoryRuntimeException e){
+      clear();
+      throw e;
+    }
     bits.zeroVector();
     mutator.reset();
     accessor.reset();
@@ -175,8 +193,13 @@ public int getCurrentSizeInBytes(){
 
   @Override
   public void allocateNew() {
-    values.allocateNew();
-    bits.allocateNew();
+    try {
+      values.allocateNew();
+      bits.allocateNew();
+    } catch(OutOfMemoryRuntimeException e) {
+      clear();
+      throw e;
+    }
     bits.zeroVector();
     mutator.reset();
     accessor.reset();
@@ -185,8 +208,21 @@ public void allocateNew() {
 
   @Override
   public boolean allocateNewSafe() {
-    if(!values.allocateNewSafe()) return false;
-    if(!bits.allocateNewSafe()) return false;
+    /* Boolean to keep track if all the memory allocations were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      if(!values.allocateNewSafe()) return false;
+      if(!bits.allocateNewSafe()) return false;
+      success = true;
+    } finally {
+      if (!success) {
+        clear();
+      }
+    }
     bits.zeroVector();
     mutator.reset();
     accessor.reset();
@@ -195,8 +231,13 @@ public boolean allocateNewSafe() {
 
   @Override
   public void allocateNew(int valueCount) {
-    values.allocateNew(valueCount);
-    bits.allocateNew(valueCount);
+    try {
+      values.allocateNew(valueCount);
+      bits.allocateNew(valueCount);
+    } catch(OutOfMemoryRuntimeException e) {
+      clear();
+      throw e;
+    }
     bits.zeroVector();
     mutator.reset();
     accessor.reset();

exec/java-exec/src/main/codegen/templates/RepeatedValueVectors.java

@@ -177,18 +177,36 @@ public void copyFromSafe(int inIndex, int outIndex, Repeated${minor.class}Vector
     }
 
 
-  public boolean allocateNewSafe(){
-    if(!offsets.allocateNewSafe()) return false;
+  public boolean allocateNewSafe() {
+    /* boolean to keep track if all the memory allocation were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      if(!offsets.allocateNewSafe()) return false;
+      if(!values.allocateNewSafe()) return false;
+      success = true;
+    } finally {
+      if (!success) {
+        clear();
+      }
+    }
     offsets.zeroVector();
-    if(!values.allocateNewSafe()) return false;
     mutator.reset();
     return true;
   }
 
   public void allocateNew() {
-    offsets.allocateNew();
+    try {
+      offsets.allocateNew();
+      values.allocateNew();
+    } catch (OutOfMemoryRuntimeException e) {
+      clear();
+      throw e;
+    }
     offsets.zeroVector();
-    values.allocateNew();
     mutator.reset();
   }
 
@@ -200,9 +218,14 @@ protected SerializedField.Builder getMetadataBuilder() {
   }
 
   public void allocateNew(int totalBytes, int valueCount, int innerValueCount) {
-    offsets.allocateNew(valueCount+1);
+    try {
+      offsets.allocateNew(valueCount+1);
+      values.allocateNew(totalBytes, innerValueCount);
+    } catch (OutOfMemoryRuntimeException e) {
+      clear();
+      throw e;
+    }
     offsets.zeroVector();
-    values.allocateNew(totalBytes, innerValueCount);
     mutator.reset();
   }
 
@@ -230,9 +253,20 @@ public int getByteCapacity(){
 
   public void allocateNew(int valueCount, int innerValueCount) {
     clear();
-    offsets.allocateNew(valueCount+1);
+    /* boolean to keep track if all the memory allocation were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      offsets.allocateNew(valueCount+1);
+      values.allocateNew(innerValueCount);
+    } catch(OutOfMemoryRuntimeException e){
+      clear();
+      throw e;
+    }
     offsets.zeroVector();
-    values.allocateNew(innerValueCount);
     mutator.reset();
   }
 

exec/java-exec/src/main/codegen/templates/VariableLengthVectors.java

@@ -18,6 +18,7 @@
 
 import java.lang.Override;
 
+import org.apache.drill.exec.memory.OutOfMemoryRuntimeException;
 import org.apache.drill.exec.vector.BaseDataValueVector;
 import org.apache.drill.exec.vector.BaseValueVector;
 import org.apache.drill.exec.vector.VariableWidthVector;
@@ -282,33 +283,48 @@ public boolean allocateNewSafe() {
       allocationMonitor = 0;
     }
 
-    DrillBuf newBuf = allocator.buffer(allocationTotalByteCount);
-    if(newBuf == null){
-      return false;
+    /* Boolean to keep track if all the memory allocations were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      DrillBuf newBuf = allocator.buffer(allocationTotalByteCount);
+      if (newBuf == null) {
+        return false;
+      }
+      this.data = newBuf;
+      if (!offsetVector.allocateNewSafe()) {
+        return false;
+      }
+      success = true;
+    } finally {
+      if (!success) {
+        clear();
+      }
     }
-
-    this.data = newBuf;
     data.readerIndex(0);
-
-    if(!offsetVector.allocateNewSafe()){
-      return false;
-    }
     offsetVector.zeroVector();
     return true;
   }
 
   public void allocateNew(int totalBytes, int valueCount) {
     clear();
     assert totalBytes >= 0;
-    DrillBuf newBuf = allocator.buffer(totalBytes);
-    if(newBuf == null){
-      throw new OutOfMemoryRuntimeException(String.format("Failure while allocating buffer of %d bytes", totalBytes));
+    try {
+      DrillBuf newBuf = allocator.buffer(totalBytes);
+      if (newBuf == null) {
+        throw new OutOfMemoryRuntimeException(String.format("Failure while allocating buffer of %d bytes", totalBytes));
+      }
+      this.data = newBuf;
+      offsetVector.allocateNew(valueCount + 1);
+    } catch (OutOfMemoryRuntimeException e) {
+      clear();
+      throw e;
     }
-
-    this.data = newBuf;
     data.readerIndex(0);
     allocationTotalByteCount = totalBytes;
-    offsetVector.allocateNew(valueCount+1);
     offsetVector.zeroVector();
   }
 

exec/java-exec/src/main/java/org/apache/drill/exec/vector/complex/AbstractMapVector.java

@@ -52,9 +52,23 @@ protected AbstractMapVector(MaterializedField field, BufferAllocator allocator,
 
   @Override
   public boolean allocateNewSafe() {
-    for (ValueVector v : vectors.values()) {
-      if (!v.allocateNewSafe()) {
-        return false;
+    /* boolean to keep track if all the memory allocation were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+
+      for (ValueVector v : vectors.values()) {
+        if (!v.allocateNewSafe()) {
+          return false;
+        }
+      }
+      success = true;
+    } finally {
+      if (!success) {
+        clear();
       }
     }
     return true;

exec/java-exec/src/main/java/org/apache/drill/exec/vector/complex/BaseRepeatedValueVector.java

@@ -61,13 +61,27 @@ protected BaseRepeatedValueVector(MaterializedField field, BufferAllocator alloc
 
   @Override
   public boolean allocateNewSafe() {
-    if (!offsets.allocateNewSafe()) {
-      return false;
+    /* boolean to keep track if all the memory allocation were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      if (!offsets.allocateNewSafe()) {
+        return false;
+      }
+      success = vector.allocateNewSafe();
+    } finally {
+      if (!success) {
+        clear();
+      }
     }
     offsets.zeroVector();
-    return vector.allocateNewSafe();
+    return success;
   }
 
+
   @Override
   public UInt4Vector getOffsetVector() {
     return offsets;

exec/java-exec/src/main/java/org/apache/drill/exec/vector/complex/RepeatedMapVector.java

@@ -35,6 +35,7 @@
 import org.apache.drill.exec.expr.holders.ComplexHolder;
 import org.apache.drill.exec.expr.holders.RepeatedMapHolder;
 import org.apache.drill.exec.memory.BufferAllocator;
+import org.apache.drill.exec.memory.OutOfMemoryRuntimeException;
 import org.apache.drill.exec.proto.UserBitShared.SerializedField;
 import org.apache.drill.exec.record.MaterializedField;
 import org.apache.drill.exec.record.TransferPair;
@@ -103,11 +104,16 @@ public RepeatedMapReaderImpl getReader() {
   @Override
   public void allocateNew(int groupCount, int innerValueCount) {
     clear();
-    offsets.allocateNew(groupCount+1);
-    offsets.zeroVector();
-    for (ValueVector v : getChildren()) {
-      AllocationHelper.allocatePrecomputedChildCount(v, groupCount, 50, innerValueCount);
+    try {
+      offsets.allocateNew(groupCount + 1);
+      for (ValueVector v : getChildren()) {
+        AllocationHelper.allocatePrecomputedChildCount(v, groupCount, 50, innerValueCount);
+      }
+    } catch (OutOfMemoryRuntimeException e){
+      clear();
+      throw e;
     }
+    offsets.zeroVector();
     mutator.reset();
   }
 
@@ -216,11 +222,24 @@ public TransferPair getTransferPair(FieldReference ref) {
 
   @Override
   public boolean allocateNewSafe() {
-    if (!offsets.allocateNewSafe()) {
-      return false;
+    /* boolean to keep track if all the memory allocation were successful
+     * Used in the case of composite vectors when we need to allocate multiple
+     * buffers for multiple vectors. If one of the allocations failed we need to
+     * clear all the memory that we allocated
+     */
+    boolean success = false;
+    try {
+      if (!offsets.allocateNewSafe()) {
+        return false;
+      }
+      success =  super.allocateNewSafe();
+    } finally {
+      if (!success) {
+        clear();
+      }
     }
     offsets.zeroVector();
-    return super.allocateNewSafe();
+    return success;
   }
 
   protected static class SingleMapTransferPair implements TransferPair {