DRILL-4081: Handle schema changes in ExternalSort

closes #257

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/FunctionGenerationHelper.java

@@ -18,16 +18,25 @@
 package org.apache.drill.exec.expr.fn;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
+import com.google.common.collect.Lists;
+import org.apache.drill.common.expression.ErrorCollector;
+import org.apache.drill.common.expression.ErrorCollectorImpl;
 import org.apache.drill.common.expression.ExpressionPosition;
+import org.apache.drill.common.expression.FunctionCall;
 import org.apache.drill.common.expression.FunctionHolderExpression;
+import org.apache.drill.common.expression.IfExpression;
+import org.apache.drill.common.expression.IfExpression.IfCondition;
 import org.apache.drill.common.expression.LogicalExpression;
+import org.apache.drill.common.expression.ValueExpressions.IntExpression;
 import org.apache.drill.common.types.TypeProtos;
 import org.apache.drill.common.types.TypeProtos.MajorType;
 import org.apache.drill.common.types.TypeProtos.MinorType;
 import org.apache.drill.common.types.Types;
 import org.apache.drill.exec.expr.ClassGenerator.HoldingContainer;
+import org.apache.drill.exec.expr.ExpressionTreeMaterializer;
 import org.apache.drill.exec.expr.HoldingContainerExpression;
 import org.apache.calcite.rel.RelFieldCollation.NullDirection;
 
@@ -46,7 +55,7 @@ public class FunctionGenerationHelper {
    * @return
    *     FunctionHolderExpression containing the found function implementation
    */
-  public static FunctionHolderExpression getOrderingComparator(
+  public static LogicalExpression getOrderingComparator(
       boolean null_high,
       HoldingContainer left,
       HoldingContainer right,
@@ -56,13 +65,20 @@ public static FunctionHolderExpression getOrderingComparator(
 
     if (   ! isComparableType(left.getMajorType() )
         || ! isComparableType(right.getMajorType() )
-        || isUnionType(left.getMajorType())
-        || isUnionType(right.getMajorType()) ) {
+            ){
       throw new UnsupportedOperationException(
           formatCanNotCompareMsg(left.getMajorType(), right.getMajorType()));
     }
-    return getFunctionExpression(comparator_name, Types.required(MinorType.INT),
+    LogicalExpression comparisonFunctionExpression = getFunctionExpression(comparator_name, Types.required(MinorType.INT),
                                  registry, left, right);
+
+    ErrorCollector collector = new ErrorCollectorImpl();
+    if (!isUnionType(left.getMajorType()) && !isUnionType(right.getMajorType())) {
+      return ExpressionTreeMaterializer.materialize(comparisonFunctionExpression, null, collector, registry);
+    } else {
+      LogicalExpression typeComparisonFunctionExpression = getTypeComparisonFunction(comparisonFunctionExpression, left, right);
+      return ExpressionTreeMaterializer.materialize(typeComparisonFunctionExpression, null, collector, registry);
+    }
   }
 
   private static boolean isUnionType(MajorType majorType) {
@@ -76,16 +92,15 @@ private static boolean isUnionType(MajorType majorType) {
    * @param  right  ...
    * @param  registry  ...
    * @return FunctionHolderExpression containing the function implementation
-   * @see #getComparator
    */
-  public static FunctionHolderExpression getOrderingComparatorNullsHigh(
+  public static LogicalExpression getOrderingComparatorNullsHigh(
       HoldingContainer left,
       HoldingContainer right,
       FunctionImplementationRegistry registry) {
     return getOrderingComparator(true, left, right, registry);
   }
 
-  public static FunctionHolderExpression getFunctionExpression(
+  private static LogicalExpression getFunctionExpression(
       String name, MajorType returnType, FunctionImplementationRegistry registry, HoldingContainer... args) {
     List<MajorType> argTypes = new ArrayList<MajorType>(args.length);
     List<LogicalExpression> argExpressions = new ArrayList<LogicalExpression>(args.length);
@@ -94,25 +109,34 @@ public static FunctionHolderExpression getFunctionExpression(
       argExpressions.add(new HoldingContainerExpression(c));
     }
 
-    DrillFuncHolder holder = registry.findExactMatchingDrillFunction(name, argTypes, returnType);
-    if (holder != null) {
-      return holder.getExpr(name, argExpressions, ExpressionPosition.UNKNOWN);
-    }
+    return new FunctionCall(name, argExpressions, ExpressionPosition.UNKNOWN);
+  }
 
-    StringBuilder sb = new StringBuilder();
-    sb.append("Failure finding function that runtime code generation expected.  Signature: ");
-    sb.append(name);
-    sb.append("( ");
-    for(int i =0; i < args.length; i++) {
-      MajorType mt = args[i].getMajorType();
-      if (i != 0) {
-        sb.append(", ");
-      }
-      appendType(mt, sb);
+  /**
+   * Wraps the comparison function in an If-statement which compares the types first, evaluating the comaprison function only
+   * if the types are equivialent
+   *
+   * @param comparisonFunction
+   * @param args
+   * @return
+   */
+  private static LogicalExpression getTypeComparisonFunction(LogicalExpression comparisonFunction, HoldingContainer... args) {
+    List<LogicalExpression> argExpressions = Lists.newArrayList();
+    List<MajorType> argTypes = Lists.newArrayList();
+    for(HoldingContainer c : args) {
+      argTypes.add(c.getMajorType());
+      argExpressions.add(new HoldingContainerExpression(c));
     }
-    sb.append(" ) returns ");
-    appendType(returnType, sb);
-    throw new UnsupportedOperationException(sb.toString());
+    FunctionCall call = new FunctionCall("compareType", argExpressions, ExpressionPosition.UNKNOWN);
+
+    List<LogicalExpression> newArgs = Lists.newArrayList();
+    newArgs.add(call);
+    newArgs.add(new IntExpression(0, ExpressionPosition.UNKNOWN));
+    FunctionCall notEqual = new FunctionCall("not_equal", newArgs, ExpressionPosition.UNKNOWN);
+
+    IfExpression.IfCondition ifCondition = new IfCondition(notEqual, call);
+    IfExpression ifExpression = IfExpression.newBuilder().setIfCondition(ifCondition).setElse(comparisonFunction).build();
+    return ifExpression;
   }
 
   private static final void appendType(MajorType mt, StringBuilder sb) {
@@ -121,7 +145,7 @@ private static final void appendType(MajorType mt, StringBuilder sb) {
     sb.append(mt.getMode().name());
   }
 
-  protected static boolean isComparableType(MajorType type) {
+  private static boolean isComparableType(MajorType type) {
     if (type.getMinorType() == MinorType.MAP ||
         type.getMinorType() == MinorType.LIST ||
         type.getMode() == TypeProtos.DataMode.REPEATED ) {

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/UnionFunctions.java

@@ -17,8 +17,10 @@
  */
 package org.apache.drill.exec.expr.fn.impl;
 
+import com.google.common.collect.Sets;
 import io.netty.buffer.DrillBuf;
 import org.apache.drill.common.types.TypeProtos.MinorType;
+import org.apache.drill.common.types.Types;
 import org.apache.drill.exec.expr.DrillSimpleFunc;
 import org.apache.drill.exec.expr.annotations.FunctionTemplate;
 import org.apache.drill.exec.expr.annotations.FunctionTemplate.NullHandling;
@@ -32,16 +34,106 @@
 import org.apache.drill.exec.expr.holders.UnionHolder;
 import org.apache.drill.exec.expr.holders.IntHolder;
 import org.apache.drill.exec.expr.holders.VarCharHolder;
+import org.apache.drill.exec.resolver.TypeCastRules;
 import org.apache.drill.exec.vector.complex.impl.UnionReader;
 import org.apache.drill.exec.vector.complex.reader.FieldReader;
 
 import javax.inject.Inject;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Set;
 
 /**
  * The class contains additional functions for union types in addition to those in GUnionFunctions
  */
 public class UnionFunctions {
 
+  /**
+   * Returns zero if the inputs have equivalent types. Two numeric types are considered equivalent, as are a combination
+   * of date/timestamp. If not equivalent, returns a value determined by the numeric value of the MinorType enum
+   */
+  @FunctionTemplate(names = {"compareType"},
+          scope = FunctionTemplate.FunctionScope.SIMPLE,
+          nulls = NullHandling.INTERNAL)
+  public static class CompareType implements DrillSimpleFunc {
+
+    @Param
+    FieldReader input1;
+    @Param
+    FieldReader input2;
+    @Output
+    IntHolder out;
+
+    public void setup() {}
+
+    public void eval() {
+      org.apache.drill.common.types.TypeProtos.MinorType type1;
+      if (input1.isSet()) {
+        type1 = input1.getType().getMinorType();
+      } else {
+        type1 = org.apache.drill.common.types.TypeProtos.MinorType.NULL;
+      }
+      org.apache.drill.common.types.TypeProtos.MinorType type2;
+      if (input2.isSet()) {
+        type2 = input2.getType().getMinorType();
+      } else {
+        type2 = org.apache.drill.common.types.TypeProtos.MinorType.NULL;
+      }
+
+      out.value = org.apache.drill.exec.expr.fn.impl.UnionFunctions.compareTypes(type1, type2);
+    }
+  }
+
+  public static int compareTypes(MinorType type1, MinorType type2) {
+    int typeValue1 = getTypeValue(type1);
+    int typeValue2 = getTypeValue(type2);
+    return typeValue1 - typeValue2;
+  }
+
+  /**
+   * Gives a type ordering modeled after the behavior of MongoDB
+   * Numeric types are first, folowed by string types, followed by binary, then boolean, then date, then timestamp
+   * Any other times will be sorted after that
+   * @param type
+   * @return
+   */
+  private static int getTypeValue(MinorType type) {
+    if (TypeCastRules.isNumericType(type)) {
+      return 0;
+    }
+    switch (type) {
+    case TINYINT:
+    case SMALLINT:
+    case INT:
+    case BIGINT:
+    case UINT1:
+    case UINT2:
+    case UINT4:
+    case UINT8:
+    case DECIMAL9:
+    case DECIMAL18:
+    case DECIMAL28SPARSE:
+    case DECIMAL38SPARSE:
+    case FLOAT4:
+    case FLOAT8:
+      return 0;
+    case VARCHAR:
+    case VAR16CHAR:
+      return 1;
+    case VARBINARY:
+      return 2;
+    case BIT:
+      return 3;
+    case DATE:
+      return 4;
+    case TIMESTAMP:
+      return 5;
+    default:
+      return 6 + type.getNumber();
+    }
+  }
+
   @FunctionTemplate(names = {"typeOf"},
           scope = FunctionTemplate.FunctionScope.SIMPLE,
           nulls = NullHandling.INTERNAL)

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

@@ -19,10 +19,16 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
+import com.google.common.collect.Sets;
 import io.netty.buffer.DrillBuf;
 import org.apache.drill.common.exceptions.DrillRuntimeException;
+import org.apache.drill.common.types.TypeProtos.DataMode;
+import org.apache.drill.common.types.TypeProtos.MajorType;
+import org.apache.drill.common.types.TypeProtos.MinorType;
 import org.apache.drill.exec.exception.SchemaChangeException;
+import org.apache.drill.exec.expr.TypeHelper;
 import org.apache.drill.exec.memory.BufferAllocator;
 import org.apache.drill.exec.memory.BufferAllocator.PreAllocator;
 import org.apache.drill.exec.ops.FragmentContext;
@@ -38,6 +44,7 @@
 
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Lists;
+import org.apache.drill.exec.vector.complex.UnionVector;
 
 public class SortRecordBatchBuilder implements AutoCloseable {
   static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(SortRecordBatchBuilder.class);
@@ -47,10 +54,12 @@ public class SortRecordBatchBuilder implements AutoCloseable {
   private int recordCount;
   private long runningBatches;
   private SelectionVector4 sv4;
+  private BufferAllocator allocator;
   final PreAllocator svAllocator;
   private boolean svAllocatorUsed = false;
 
   public SortRecordBatchBuilder(BufferAllocator a) {
+    this.allocator = a;
     this.svAllocator = a.getNewPreAllocator();
   }
 

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

@@ -24,7 +24,9 @@
 import org.apache.drill.common.expression.SchemaPath;
 import org.apache.drill.exec.cache.VectorAccessibleSerializable;
 import org.apache.drill.exec.memory.BufferAllocator;
+import org.apache.drill.exec.ops.OperatorContext;
 import org.apache.drill.exec.record.BatchSchema;
+import org.apache.drill.exec.record.SchemaUtil;
 import org.apache.drill.exec.record.TransferPair;
 import org.apache.drill.exec.record.TypedFieldId;
 import org.apache.drill.exec.record.VectorAccessible;
@@ -51,23 +53,36 @@ public class BatchGroup implements VectorAccessible, AutoCloseable {
   private FileSystem fs;
   private BufferAllocator allocator;
   private int spilledBatches = 0;
+  private OperatorContext context;
+  private BatchSchema schema;
 
-  public BatchGroup(VectorContainer container, SelectionVector2 sv2) {
+  public BatchGroup(VectorContainer container, SelectionVector2 sv2, OperatorContext context) {
     this.sv2 = sv2;
     this.currentContainer = container;
+    this.context = context;
   }
 
-  public BatchGroup(VectorContainer container, FileSystem fs, String path, BufferAllocator allocator) {
+  public BatchGroup(VectorContainer container, FileSystem fs, String path, OperatorContext context) {
     currentContainer = container;
     this.fs = fs;
     this.path = new Path(path);
-    this.allocator = allocator;
+    this.allocator = context.getAllocator();
+    this.context = context;
   }
 
   public SelectionVector2 getSv2() {
     return sv2;
   }
 
+  /**
+   * Updates the schema for this batch group. The current as well as any deserialized batches will be coerced to this schema
+   * @param schema
+   */
+  public void setSchema(BatchSchema schema) {
+    currentContainer = SchemaUtil.coerceContainer(currentContainer, schema, context);
+    this.schema = schema;
+  }
+
   public void addBatch(VectorContainer newContainer) throws IOException {
     assert fs != null;
     assert path != null;
@@ -96,6 +111,9 @@ private VectorContainer getBatch() throws IOException {
     watch.start();
     vas.readFromStream(inputStream);
     VectorContainer c =  vas.get();
+    if (schema != null) {
+      c = SchemaUtil.coerceContainer(c, schema, context);
+    }
 //    logger.debug("Took {} us to read {} records", watch.elapsed(TimeUnit.MICROSECONDS), c.getRecordCount());
     spilledBatches--;
     currentContainer.zeroVectors();