DRILL-7093 Batch Sizing in SingleSender

exec/java-exec/src/main/java/org/apache/drill/exec/ExecConstants.java

@@ -73,6 +73,12 @@ private ExecConstants() {
   public static final String TEMP_DIRECTORIES = "drill.exec.tmp.directories";
   public static final String TEMP_FILESYSTEM = "drill.exec.tmp.filesystem";
   public static final String INCOMING_BUFFER_IMPL = "drill.exec.buffer.impl";
+
+  public static final String EXCHANGE_BATCH_SIZE = "drill.exec.memory.operator.exchange_batch_size";
+  // Exchange Batch Size in Bytes.
+  public static final LongValidator EXCHANGE_BATCH_SIZE_VALIDATOR = new RangeLongValidator(EXCHANGE_BATCH_SIZE, 1024, 10 * 1024 * 1024,
+          new OptionDescription("Limit, in bytes, of the outgoing batch sent by an Exchange Sender."));
+
   /** incoming buffer size (number of batches) */
   public static final String INCOMING_BUFFER_SIZE = "drill.exec.buffer.size";
   public static final String SPOOLING_BUFFER_DELETE = "drill.exec.buffer.spooling.delete";

exec/java-exec/src/main/java/org/apache/drill/exec/physical/config/HashPartitionSender.java

@@ -66,9 +66,6 @@ public LogicalExpression getExpr() {
     return expr;
   }
 
-  public int getOutgoingBatchSize() {
-    return outgoingBatchSize;
-  }
 
   @Override
   public <T, X, E extends Throwable> T accept(PhysicalVisitor<T, X, E> physicalVisitor, X value) throws E {

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/ScanBatch.java

@@ -582,6 +582,15 @@ public WritableBatch getWritableBatch() {
     return WritableBatch.get(this);
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    VectorContainer partialContainer = new VectorContainer(context.getAllocator(), getSchema());
+    container.transferOut(partialContainer, startIndex, length);
+    partialContainer.setRecordCount(length);
+    final WritableBatch batch = WritableBatch.get(partialContainer);
+    return batch;
+  }
+
   @Override
   public void close() throws Exception {
     container.clear();

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/SenderMemoryManager.java

@@ -0,0 +1,41 @@
+/*
+ * 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.drill.exec.physical.impl;
+
+import org.apache.drill.exec.record.RecordBatch;
+import org.apache.drill.exec.record.RecordBatchMemoryManager;
+import org.apache.drill.exec.record.RecordBatchSizer;
+
+public class SenderMemoryManager extends RecordBatchMemoryManager {
+  private RecordBatch incomingBatch;
+  private int batchSizeLimit;
+
+  public SenderMemoryManager(int batchSizeLimit, RecordBatch incomingBatch) {
+    super(batchSizeLimit);
+    this.batchSizeLimit = batchSizeLimit;
+    this.incomingBatch = incomingBatch;
+  }
+
+  public void update() {
+    RecordBatchSizer batchSizer = new RecordBatchSizer(incomingBatch);
+    setRecordBatchSizer(batchSizer);
+    int outputRowCount = computeOutputRowCount(batchSizeLimit, batchSizer.getNetRowWidth());
+    outputRowCount = Math.min(outputRowCount, batchSizer.rowCount());
+    setOutputRowCount(outputRowCount);
+  }
+}

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/SingleSenderCreator.java

@@ -20,6 +20,7 @@
 import java.util.List;
 
 import org.apache.drill.common.exceptions.ExecutionSetupException;
+import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.exception.OutOfMemoryException;
 import org.apache.drill.exec.ops.AccountingDataTunnel;
 import org.apache.drill.exec.ops.ExecutorFragmentContext;
@@ -31,8 +32,10 @@
 import org.apache.drill.exec.record.FragmentWritableBatch;
 import org.apache.drill.exec.record.RecordBatch;
 import org.apache.drill.exec.record.RecordBatch.IterOutcome;
+import org.apache.drill.exec.record.WritableBatch;
 import org.apache.drill.exec.testing.ControlsInjector;
 import org.apache.drill.exec.testing.ControlsInjectorFactory;
+import org.apache.drill.shaded.guava.com.google.common.base.Preconditions;
 
 public class SingleSenderCreator implements RootCreator<SingleSender>{
 
@@ -55,6 +58,7 @@ public static class SingleSenderRootExec extends BaseRootExec {
     private int recMajor;
     private volatile boolean ok = true;
     private volatile boolean done = false;
+    private SenderMemoryManager memoryManager;
 
     public enum Metric implements MetricDef {
       BYTES_SENT;
@@ -78,6 +82,8 @@ public SingleSenderRootExec(RootFragmentContext context, RecordBatch batch, Sing
           .build();
       tunnel = context.getDataTunnel(config.getDestination());
       tunnel.setTestInjectionControls(injector, context.getExecutionControls(), logger);
+      int exchangeBatchSizeLimit = (int) context.getOptions().getOption(ExecConstants.EXCHANGE_BATCH_SIZE_VALIDATOR);
+      this.memoryManager = new SenderMemoryManager(exchangeBatchSizeLimit, incoming);
     }
 
     @Override
@@ -118,16 +124,28 @@ public boolean innerNext() {
 
       case OK_NEW_SCHEMA:
       case OK:
-        final FragmentWritableBatch batch = new FragmentWritableBatch(
-            false, handle.getQueryId(), handle.getMajorFragmentId(),
-            handle.getMinorFragmentId(), recMajor, oppositeHandle.getMinorFragmentId(),
-            incoming.getWritableBatch().transfer(oContext.getAllocator()));
-        updateStats(batch);
-        stats.startWait();
-        try {
-          tunnel.sendRecordBatch(batch);
-        } finally {
-          stats.stopWait();
+        memoryManager.update();
+        // manager ensures that sizedBatchRowCount <= incoming recordcount
+        int sizedBatchRowCount = memoryManager.getOutputRowCount();
+        int incomingRowCount = incoming.getRecordCount();
+        if (sizedBatchRowCount == incomingRowCount) {
+          doSend(incoming.getWritableBatch().transfer(oContext.getAllocator()));
+          return true;
+        }
+        Preconditions.checkState(sizedBatchRowCount != 0);
+        int batchesToSend = incomingRowCount / sizedBatchRowCount;
+        int tailRowCount = incomingRowCount % sizedBatchRowCount;
+        int batchCount = 0;
+
+        while(batchCount < batchesToSend) {
+          int start = batchCount * sizedBatchRowCount;
+          final WritableBatch writableBatch = incoming.getWritableBatch(start, sizedBatchRowCount);
+          doSend(writableBatch.transfer(oContext.getAllocator()));
+          batchCount++;
+        }
+        if (tailRowCount != 0) {
+          int start = batchCount * sizedBatchRowCount;
+          doSend(incoming.getWritableBatch(start, tailRowCount).transfer(oContext.getAllocator()));
         }
         return true;
 
@@ -137,6 +155,19 @@ public boolean innerNext() {
       }
     }
 
+    private void doSend(WritableBatch writableBatch) {
+      final FragmentWritableBatch batch = new FragmentWritableBatch(
+          false, handle.getQueryId(), handle.getMajorFragmentId(),
+          handle.getMinorFragmentId(), recMajor, oppositeHandle.getMinorFragmentId(), writableBatch);
+      updateStats(batch);
+      stats.startWait();
+      try {
+        tunnel.sendRecordBatch(batch);
+      } finally {
+        stats.stopWait();
+      }
+    }
+
     public void updateStats(FragmentWritableBatch writableBatch) {
       stats.addLongStat(Metric.BYTES_SENT, writableBatch.getByteCount());
     }

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/aggregate/SpilledRecordbatch.java

@@ -107,6 +107,15 @@ public Iterator<VectorWrapper<?>> iterator() {
   @Override
   public BatchSchema getSchema() { return schema; }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    VectorContainer partialContainer = new VectorContainer(context.getAllocator(), getSchema());
+    container.transferOut(partialContainer, startIndex, length);
+    partialContainer.setRecordCount(length);
+    final WritableBatch batch = WritableBatch.get(partialContainer);
+    return batch;
+  }
+
   @Override
   public WritableBatch getWritableBatch() {
     return WritableBatch.get(this);

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/mergereceiver/MergingRecordBatch.java

@@ -838,4 +838,13 @@ public void dump() {
         container, outgoingPosition, Arrays.toString(incomingBatches), Arrays.toString(batchOffsets),
         Arrays.toString(tempBatchHolder), Arrays.toString(inputCounts), Arrays.toString(outputCounts));
   }
+
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    VectorContainer partialContainer = new VectorContainer(context.getAllocator(), getSchema());
+    outgoingContainer.transferOut(partialContainer, startIndex, length);
+    partialContainer.setRecordCount(length);
+    final WritableBatch batch = WritableBatch.get(partialContainer);
+    return batch;
+  }
 }

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/protocol/OperatorRecordBatch.java

@@ -115,6 +115,11 @@ public VectorWrapper<?> getValueAccessorById(Class<?> clazz, int... ids) {
     return batchAccessor.getValueAccessorById(clazz, ids);
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public WritableBatch getWritableBatch() {
     return batchAccessor.getWritableBatch();

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/unorderedreceiver/UnorderedReceiverBatch.java

@@ -218,6 +218,11 @@ public IterOutcome next() {
     }
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    return batchLoader.getWritableBatch(startIndex, length);
+  }
+
   @Override
   public WritableBatch getWritableBatch() {
     return batchLoader.getWritableBatch();

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/validate/IteratorValidatorBatchIterator.java

@@ -353,6 +353,13 @@ public WritableBatch getWritableBatch() {
     return incoming.getWritableBatch();
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    validateReadState("getWritableBatch()");
+    return incoming.getWritableBatch(startIndex, length);
+  }
+
+
   @Override
   public void close() {
     // (Log construction and close() calls at same logging level to bracket

exec/java-exec/src/main/java/org/apache/drill/exec/record/AbstractRecordBatch.java

@@ -252,6 +252,15 @@ public WritableBatch getWritableBatch() {
 
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    VectorContainer partialContainer = new VectorContainer(context.getAllocator(), getSchema());
+    container.transferOut(partialContainer, startIndex, length);
+    partialContainer.setRecordCount(length);
+    final WritableBatch batch = WritableBatch.get(partialContainer);
+    return batch;
+  }
+
   @Override
   public VectorContainer getOutgoingContainer() {
     throw new UnsupportedOperationException(String.format(" You should not call getOutgoingContainer() for class %s", this.getClass().getCanonicalName()));

exec/java-exec/src/main/java/org/apache/drill/exec/record/HyperVectorWrapper.java

@@ -159,6 +159,23 @@ public void transfer(VectorWrapper<?> destination) {
     }
   }
 
+  /**
+   * Transfer vectors to destination HyperVectorWrapper.
+   * Both this and destination must be of same type and have same number of vectors.
+   * @param destination destination HyperVectorWrapper.
+   */
+  @Override
+  public void transferPartial(VectorWrapper<?> destination, int startIndex, int length) {
+    Preconditions.checkArgument(destination instanceof HyperVectorWrapper);
+    Preconditions.checkArgument(getField().getType().equals(destination.getField().getType()));
+    Preconditions.checkArgument(vectors.length == ((HyperVectorWrapper<?>)destination).vectors.length);
+
+    final ValueVector[] destionationVectors = ((HyperVectorWrapper<?>)destination).vectors;
+    for (int i = 0; i < vectors.length; ++i) {
+      vectors[i].makeTransferPair(destionationVectors[i]).splitAndTransfer(startIndex, length);
+    }
+  }
+
   /**
    * Method to replace existing list of vectors with the newly provided ValueVectors list in this HyperVectorWrapper
    * @param vv - New list of ValueVectors to be stored

exec/java-exec/src/main/java/org/apache/drill/exec/record/RecordBatch.java

@@ -321,6 +321,12 @@ public boolean isError() {
    */
   WritableBatch getWritableBatch();
 
+  /**
+   * Gets a writable version of this batch.  Takes over ownership of existing
+   * buffers.
+   */
+  WritableBatch getWritableBatch(int startIndex, int length);
+
   /**
    * Perform dump of this batch's state to logs.
    */

exec/java-exec/src/main/java/org/apache/drill/exec/record/RecordBatchLoader.java

@@ -262,6 +262,14 @@ public WritableBatch getWritableBatch(){
     return WritableBatch.getBatchNoHVWrap(valueCount, container, isSV2);
   }
 
+  public WritableBatch getWritableBatch(int startIndex, int length) {
+    VectorContainer partialContainer = new VectorContainer(allocator, getSchema());
+    container.transferOut(partialContainer, startIndex, length);
+    partialContainer.setRecordCount(length);
+    final WritableBatch batch = WritableBatch.get(partialContainer);
+    return batch;
+  }
+
   @Override
   public Iterator<VectorWrapper<?>> iterator() {
     return this.container.iterator();

exec/java-exec/src/main/java/org/apache/drill/exec/record/SchemalessBatch.java

@@ -95,6 +95,12 @@ public WritableBatch getWritableBatch() {
         this.getClass().getCanonicalName()));
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int start, int length) {
+    throw new UnsupportedOperationException(String.format("You should not call getWritableBatch() for class %s",
+            this.getClass().getCanonicalName()));
+  }
+
   @Override
   public Iterator<VectorWrapper<?>> iterator() {
     return null;

exec/java-exec/src/main/java/org/apache/drill/exec/record/SimpleRecordBatch.java

@@ -88,6 +88,11 @@ public WritableBatch getWritableBatch() {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public WritableBatch getWritableBatch(int start, int length) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public Iterator<VectorWrapper<?>> iterator() {
     return container.iterator();

exec/java-exec/src/main/java/org/apache/drill/exec/record/SimpleVectorWrapper.java

@@ -109,6 +109,13 @@ public void transfer(VectorWrapper<?> destination) {
     vector.makeTransferPair(((SimpleVectorWrapper<?>)destination).vector).transfer();
   }
 
+  @Override
+  public void transferPartial(VectorWrapper<?> destination, int startIndex, int length) {
+    Preconditions.checkArgument(destination instanceof SimpleVectorWrapper);
+    Preconditions.checkArgument(getField().getType().equals(destination.getField().getType()));
+    vector.makeTransferPair(((SimpleVectorWrapper<?>)destination).vector).splitAndTransfer(startIndex, length);
+  }
+
   @Override
   public String toString() {
     if (vector == null) {

exec/java-exec/src/main/java/org/apache/drill/exec/record/VectorContainer.java

@@ -131,6 +131,16 @@ public void transferOut(VectorContainer containerOut) {
     }
   }
 
+  /**
+   * Transfer vectors from this to containerOut
+   */
+  public void transferOut(VectorContainer containerOut, int startIndex, int length) {
+    Preconditions.checkArgument(this.wrappers.size() == containerOut.wrappers.size());
+    for (int i = 0; i < this.wrappers.size(); ++i) {
+      this.wrappers.get(i).transferPartial(containerOut.wrappers.get(i), startIndex, length);
+    }
+  }
+
   public <T extends ValueVector> T addOrGet(MaterializedField field) {
     return addOrGet(field, null);
   }

exec/java-exec/src/main/java/org/apache/drill/exec/record/VectorWrapper.java

@@ -35,6 +35,8 @@ public interface VectorWrapper<T extends ValueVector> {
   public VectorWrapper<T> cloneAndTransfer(BufferAllocator allocator);
   public VectorWrapper<?> getChildWrapper(int[] ids);
   public void transfer(VectorWrapper<?> destination);
+  public void transferPartial(VectorWrapper<?> destination, int startIndex, int length);
+
 
   /**
    * Traverse the object graph and determine whether the provided SchemaPath matches data within the Wrapper.  If so, return a TypedFieldId associated with this path.

exec/java-exec/src/main/java/org/apache/drill/exec/server/options/SystemOptionManager.java

@@ -266,6 +266,7 @@ public static CaseInsensitiveMap<OptionDefinition> createDefaultOptionDefinition
       new OptionDefinition(ExecConstants.CPU_LOAD_AVERAGE),
       new OptionDefinition(ExecConstants.ENABLE_VECTOR_VALIDATOR),
       new OptionDefinition(ExecConstants.ENABLE_ITERATOR_VALIDATOR),
+      new OptionDefinition(ExecConstants.EXCHANGE_BATCH_SIZE_VALIDATOR, new OptionMetaData(OptionValue.AccessibleScopes.SYSTEM, true, false)),
       new OptionDefinition(ExecConstants.OUTPUT_BATCH_SIZE_VALIDATOR, new OptionMetaData(OptionValue.AccessibleScopes.SYSTEM, true, false)),
       new OptionDefinition(ExecConstants.STATS_LOGGING_BATCH_SIZE_VALIDATOR, new OptionMetaData(OptionValue.AccessibleScopes.SYSTEM_AND_SESSION, true, true)),
       new OptionDefinition(ExecConstants.STATS_LOGGING_BATCH_FG_SIZE_VALIDATOR,new OptionMetaData(OptionValue.AccessibleScopes.SYSTEM_AND_SESSION, true, true)),