DRILL-6147: Adding Columnar Parquet Batch Sizing functionality

closes #1330

2 files changed, 18 insertions, 1 deletions

diff --git a/exec/vector/src/main/codegen/templates/NullableValueVectors.java b/exec/vector/src/main/codegen/templates/NullableValueVectors.java
index cdff55687..f30cfae78 100644
--- a/exec/vector/src/main/codegen/templates/NullableValueVectors.java
+++ b/exec/vector/src/main/codegen/templates/NullableValueVectors.java
@@ -395,6 +395,13 @@ public final class ${className} extends BaseDataValueVector implements <#if type
     return v;
   }
 
+  /**
+   * @return Underlying "bits" vector value capacity
+   */
+  public int getBitsValueCapacity() {
+    return bits.getValueCapacity();
+  }
+
   public void copyFrom(int fromIndex, int thisIndex, Nullable${minor.class}Vector from){
     final Accessor fromAccessor = from.getAccessor();
     if (!fromAccessor.isNull(fromIndex)) {
diff --git a/exec/vector/src/main/codegen/templates/VariableLengthVectors.java b/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
index daef7ba1c..c35728ec2 100644
--- a/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
+++ b/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
@@ -632,6 +632,10 @@ public final class ${minor.class}Vector extends BaseDataValueVector implements V
       // Let's process the input
       while (input.hasNext()) {
         T entry = input.next();
+
+        if (entry == null || entry.getNumValues() == 0) {
+          break; // this could happen when handling columnar batch sizing constraints
+        }
         bufferedMutator.setSafe(entry);
 
         if (callback != null) {
@@ -897,7 +901,7 @@ public final class ${minor.class}Vector extends BaseDataValueVector implements V
       this.parent = parent;
       this.dataBuffOff = this.parent.offsetVector.getAccessor().get(startIdx);
       this.totalDataLen = this.dataBuffOff;
-      this.offsetsMutator = new UInt4Vector.BufferedMutator(startIdx, buffSz * 4, parent.offsetVector);
+      this.offsetsMutator = new UInt4Vector.BufferedMutator(startIdx, buffSz, parent.offsetVector);
 
       // Forcing the offsetsMutator to operate at index+1
       this.offsetsMutator.setSafe(this.dataBuffOff);
@@ -970,6 +974,7 @@ public final class ${minor.class}Vector extends BaseDataValueVector implements V
 
       // Update counters
       dataBuffOff += buffer.position();
+      assert dataBuffOff == totalDataLen;
 
       // Reset the byte buffer
       buffer.clear();
@@ -1008,6 +1013,11 @@ public final class ${minor.class}Vector extends BaseDataValueVector implements V
         remaining -= toCopy;
 
       } while (remaining > 0);
+
+      // We need to flush as offset data can be accessed during loading to
+      // figure out current payload size.
+      offsetsMutator.flush();
+
     }
   }
 }