Commit some missing code from an old git stash.

Author: JoshRosen

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java

@@ -137,7 +137,7 @@ public void spill() throws IOException {
     openSorter();
   }
 
-  private long freeMemory() {
+  public long freeMemory() {
     long memoryFreed = 0;
     final Iterator<MemoryBlock> iter = allocatedPages.iterator();
     while (iter.hasNext()) {
@@ -223,13 +223,20 @@ public void insertRecord(
   }
 
   public UnsafeSorterIterator getSortedIterator() throws IOException {
-    final UnsafeSorterSpillMerger spillMerger =
-      new UnsafeSorterSpillMerger(recordComparator, prefixComparator);
-    for (UnsafeSorterSpillWriter spillWriter : spillWriters) {
-      spillMerger.addSpill(spillWriter.getReader(blockManager));
+    final UnsafeSorterIterator inMemoryIterator = sorter.getSortedIterator();
+    if (!spillWriters.isEmpty()) {
+      final UnsafeSorterSpillMerger spillMerger =
+        new UnsafeSorterSpillMerger(recordComparator, prefixComparator);
+      for (UnsafeSorterSpillWriter spillWriter : spillWriters) {
+        spillMerger.addSpill(spillWriter.getReader(blockManager));
+      }
+      spillWriters.clear();
+      if (inMemoryIterator.hasNext()) {
+        spillMerger.addSpill(inMemoryIterator);
+      }
+      return spillMerger.getSortedIterator();
+    } else {
+      return inMemoryIterator;
     }
-    spillWriters.clear();
-    spillMerger.addSpill(sorter.getSortedIterator());
-    return spillMerger.getSortedIterator();
   }
 }

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java

@@ -19,6 +19,7 @@
 
 import java.util.Comparator;
 
+import org.apache.spark.unsafe.PlatformDependent;
 import org.apache.spark.util.collection.Sorter;
 import org.apache.spark.unsafe.memory.TaskMemoryManager;
 
@@ -50,10 +51,10 @@ private static final class SortComparator implements Comparator<RecordPointerAnd
     public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
       final int prefixComparisonResult = prefixComparator.compare(r1.keyPrefix, r2.keyPrefix);
       if (prefixComparisonResult == 0) {
-        final Object baseObject1 = memoryManager.getPage(r2.recordPointer);
-        final long baseOffset1 = memoryManager.getOffsetInPage(r1.recordPointer);
+        final Object baseObject1 = memoryManager.getPage(r1.recordPointer);
+        final long baseOffset1 = memoryManager.getOffsetInPage(r1.recordPointer) + 4; // skip length
         final Object baseObject2 = memoryManager.getPage(r2.recordPointer);
-        final long baseOffset2 = memoryManager.getOffsetInPage(r2.recordPointer);
+        final long baseOffset2 = memoryManager.getOffsetInPage(r2.recordPointer) + 4; // skip length
         return recordComparator.compare(baseObject1, baseOffset1, baseObject2, baseOffset2);
       } else {
         return prefixComparisonResult;
@@ -146,7 +147,8 @@ public boolean hasNext() {
     public void loadNext() {
       final long recordPointer = sortBuffer[position];
       baseObject = memoryManager.getPage(recordPointer);
-      baseOffset = memoryManager.getOffsetInPage(recordPointer);
+      baseOffset = memoryManager.getOffsetInPage(recordPointer) + 4;  // Skip over record length
+      recordLength = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset - 4);
       keyPrefix = sortBuffer[position + 1];
       position += 2;
     }

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java

@@ -279,6 +279,37 @@ public Object get(int i) {
     }
   }
 
+  /**
+   * Generic `get()`, for use in toString(). This method is for debugging only and is probably very
+   * slow to call due to having to reflect on the schema.
+   */
+  private Object genericGet(int i) {
+    assertIndexIsValid(i);
+    assert (schema != null) : "Schema must be defined when calling genericGet()";
+    final DataType dataType = schema.fields()[i].dataType();
+    if (isNullAt(i) || dataType == NullType) {
+      return null;
+    } else if (dataType == StringType) {
+      return getUTF8String(i);
+    } else if (dataType == BooleanType) {
+      return getBoolean(i);
+    } else if (dataType == ByteType) {
+      return getByte(i);
+    } else if (dataType == ShortType) {
+      return getShort(i);
+    } else if (dataType == IntegerType) {
+      return getInt(i);
+    } else if (dataType == LongType) {
+      return getLong(i);
+    } else if (dataType == FloatType) {
+      return getFloat(i);
+    } else if (dataType == DoubleType) {
+      return getDouble(i);
+    } else {
+      throw new UnsupportedOperationException();
+    }
+  }
+
   @Override
   public boolean isNullAt(int i) {
     assertIndexIsValid(i);

sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala

@@ -32,6 +32,18 @@ import org.apache.spark.sql.types.DataType
 import org.apache.spark.util.MutablePair
 import org.apache.spark.{HashPartitioner, Partitioner, RangePartitioner, SparkEnv}
 
+import scala.util.control.NonFatal
+
+object Exchange {
+  /**
+   * Returns true when the ordering expressions are a subset of the key.
+   * if true, ShuffledRDD can use `setKeyOrdering(orderingKey)` to sort within [[Exchange]].
+   */
+  def canSortWithShuffle(partitioning: Partitioning, desiredOrdering: Seq[SortOrder]): Boolean = {
+    desiredOrdering.map(_.child).toSet.subsetOf(partitioning.keyExpressions.toSet)
+  }
+}
+
 /**
  * :: DeveloperApi ::
  * Performs a shuffle that will result in the desired `newPartitioning`.  Optionally sorts each
@@ -181,6 +193,10 @@ case class Exchange(
           }
         }
         val shuffled = new ShuffledRDD[InternalRow, InternalRow, InternalRow](rdd, part)
+        if (newOrdering.nonEmpty) {
+          println("Shuffling with a key ordering")
+          shuffled.setKeyOrdering(keyOrdering)
+        }
         shuffled.setSerializer(serializer)
         shuffled.map(_._2)
 
@@ -292,6 +308,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
           partitioning: Partitioning,
           rowOrdering: Seq[SortOrder],
           child: SparkPlan): SparkPlan = {
+        logInfo("In addOperatorsIfNecessary")
         val needSort = rowOrdering.nonEmpty && child.outputOrdering != rowOrdering
         val needsShuffle = child.outputPartitioning != partitioning
 
@@ -301,9 +318,27 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
           child
         }
 
-        val withSort = if (needSort) {
-          if (sqlContext.conf.externalSortEnabled) {
-            ExternalSort(rowOrdering, global = false, withShuffle)
+          val withSort = if (needSort) {
+            // TODO(josh): this is a hack. Need a better way to determine whether UnsafeRow
+            // supports the given schema.
+            val supportsUnsafeRowConversion: Boolean = try {
+              new UnsafeRowConverter(withShuffle.schema.map(_.dataType).toArray)
+              true
+            } catch {
+              case NonFatal(e) =>
+                false
+            }
+            logInfo(s"For row with data types ${withShuffle.schema.map(_.dataType)}, " +
+              s"supportsUnsafeRowConversion = $supportsUnsafeRowConversion")
+            if (sqlContext.conf.unsafeEnabled && supportsUnsafeRowConversion) {
+              logInfo("Using unsafe external sort!")
+              UnsafeExternalSort(rowOrdering, global = false, withShuffle)
+            } else if (sqlContext.conf.externalSortEnabled) {
+              logInfo("Not using unsafe sort")
+              ExternalSort(rowOrdering, global = false, withShuffle)
+            } else {
+              Sort(rowOrdering, global = false, withShuffle)
+            }
           } else {
             Sort(rowOrdering, global = false, withShuffle)
           }
@@ -317,6 +352,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       if (meetsRequirements && compatible && !needsAnySort) {
         operator
       } else {
+        logInfo("Looking through Exchange")
         // At least one child does not satisfies its required data distribution or
         // at least one child's outputPartitioning is not compatible with another child's
         // outputPartitioning. In this case, we need to add Exchange operators.
@@ -334,7 +370,21 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
           case (UnspecifiedDistribution, Seq(), child) =>
             child
           case (UnspecifiedDistribution, rowOrdering, child) =>
-            if (sqlContext.conf.externalSortEnabled) {
+            // TODO(josh): this is a hack. Need a better way to determine whether UnsafeRow
+            // supports the given schema.
+            val supportsUnsafeRowConversion: Boolean = try {
+              new UnsafeRowConverter(child.schema.map(_.dataType).toArray)
+              true
+            } catch {
+              case NonFatal(e) =>
+                false
+            }
+            logInfo(s"For row with data types ${child.schema.map(_.dataType)}, " +
+              s"supportsUnsafeRowConversion = $supportsUnsafeRowConversion")
+            if (sqlContext.conf.unsafeEnabled && supportsUnsafeRowConversion) {
+              logInfo("Using unsafe external sort!")
+              UnsafeExternalSort(rowOrdering, global = false, child)
+            } else if (sqlContext.conf.externalSortEnabled) {
               ExternalSort(rowOrdering, global = false, child)
             } else {
               Sort(rowOrdering, global = false, child)

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

@@ -74,7 +74,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * [[org.apache.spark.sql.SQLConf.AUTO_BROADCASTJOIN_THRESHOLD]], the planner would mark it as the
    * ''build'' relation and mark the other relation as the ''stream'' side.  The build table will be
    * ''broadcasted'' to all of the executors involved in the join, as a
-   * [[org.apache.spark.broadcast.Broadcast]] object.  If both estimates exceed the threshold, they
+   * [[org.apache.spark.broadcast.Broadcast]] object.  If both   estimates exceed the threshold, they
    * will instead be used to decide the build side in a [[joins.ShuffledHashJoin]].
    */
   object HashJoin extends Strategy with PredicateHelper {
@@ -102,8 +102,11 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       // for now let's support inner join first, then add outer join
       case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, condition, left, right)
         if sqlContext.conf.sortMergeJoinEnabled =>
-        val mergeJoin =
+        val mergeJoin = if (sqlContext.conf.unsafeEnabled) {
+          joins.UnsafeSortMergeJoin(leftKeys, rightKeys, planLater(left), planLater(right))
+        } else {
           joins.SortMergeJoin(leftKeys, rightKeys, planLater(left), planLater(right))
+        }
         condition.map(Filter(_, mergeJoin)).getOrElse(mergeJoin) :: Nil
 
       case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, condition, left, right) =>

sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala

@@ -17,6 +17,12 @@
 
 package org.apache.spark.sql.execution
 
+import java.util.Arrays
+
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.unsafe.PlatformDependent
+import org.apache.spark.util.collection.unsafe.sort.{RecordComparator, PrefixComparator, UnsafeExternalSorter}
+import org.apache.spark.{TaskContext, SparkEnv, HashPartitioner, SparkConf}
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.rdd.{RDD, ShuffledRDD}
 import org.apache.spark.shuffle.sort.SortShuffleManager
@@ -245,6 +251,119 @@ case class ExternalSort(
   override def outputOrdering: Seq[SortOrder] = sortOrder
 }
 
+/**
+ * :: DeveloperApi ::
+ * TODO(josh): document
+ * Performs a sort, spilling to disk as needed.
+ * @param global when true performs a global sort of all partitions by shuffling the data first
+ *               if necessary.
+ */
+@DeveloperApi
+case class UnsafeExternalSort(
+    sortOrder: Seq[SortOrder],
+    global: Boolean,
+    child: SparkPlan)
+  extends UnaryNode {
+
+  private[this] val numFields: Int = child.schema.size
+  private[this] val schema: StructType = child.schema
+
+  override def requiredChildDistribution: Seq[Distribution] =
+    if (global) OrderedDistribution(sortOrder) :: Nil else UnspecifiedDistribution :: Nil
+
+  protected override def doExecute(): RDD[Row] = attachTree(this, "sort") {
+    // TODO(josh): This code is unreadably messy; this should be split into a separate file
+    // and written in Java.
+    assert (codegenEnabled)
+    def doSort(iterator: Iterator[Row]): Iterator[Row] = {
+      val ordering = newOrdering(sortOrder, child.output)
+      val rowConverter = new UnsafeRowConverter(schema.map(_.dataType).toArray)
+      var rowConversionScratchSpace = new Array[Long](1024)
+      val prefixComparator = new PrefixComparator {
+        override def compare(prefix1: Long, prefix2: Long): Int = 0
+      }
+      val recordComparator = new RecordComparator {
+        private[this] val row1 = new UnsafeRow
+        private[this] val row2 = new UnsafeRow
+        override def compare(
+            baseObj1: scala.Any, baseOff1: Long, baseObj2: scala.Any, baseOff2: Long): Int = {
+          row1.pointTo(baseObj1, baseOff1, numFields, schema)
+          row2.pointTo(baseObj2, baseOff2, numFields, schema)
+          ordering.compare(row1, row2)
+        }
+      }
+      val sorter = new UnsafeExternalSorter(
+        TaskContext.get.taskMemoryManager(),
+        SparkEnv.get.shuffleMemoryManager,
+        SparkEnv.get.blockManager,
+        TaskContext.get,
+        recordComparator,
+        prefixComparator,
+        4096,
+        SparkEnv.get.conf
+      )
+      while (iterator.hasNext) {
+        val row: Row = iterator.next()
+        val sizeRequirement = rowConverter.getSizeRequirement(row)
+        if (sizeRequirement / 8 > rowConversionScratchSpace.length) {
+          rowConversionScratchSpace = new Array[Long](sizeRequirement / 8)
+        } else {
+          // Zero out the buffer that's used to hold the current row. This is necessary in order
+          // to ensure that rows hash properly, since garbage data from the previous row could
+          // otherwise end up as padding in this row. As a performance optimization, we only zero
+          // out the portion of the buffer that we'll actually write to.
+          Arrays.fill(rowConversionScratchSpace, 0, sizeRequirement / 8, 0)
+        }
+        val bytesWritten =
+          rowConverter.writeRow(row, rowConversionScratchSpace, PlatformDependent.LONG_ARRAY_OFFSET)
+        assert (bytesWritten == sizeRequirement)
+        val prefix: Long = 0  // dummy prefix until we implement prefix calculation
+        sorter.insertRecord(
+          rowConversionScratchSpace,
+          PlatformDependent.LONG_ARRAY_OFFSET,
+          sizeRequirement,
+          prefix
+        )
+      }
+      val sortedIterator = sorter.getSortedIterator
+      // TODO: need to avoid memory leaks on exceptions, etc. by wrapping in resource cleanup blocks
+      // TODO: need to clean up spill files after success or failure.
+      new Iterator[Row] {
+        private[this] val row = new UnsafeRow()
+        override def hasNext: Boolean = sortedIterator.hasNext
+
+        override def next(): Row = {
+          sortedIterator.loadNext()
+          if (hasNext) {
+            row.pointTo(
+              sortedIterator.getBaseObject, sortedIterator.getBaseOffset, numFields, schema)
+            println("Returned row " + row)
+            row
+          } else {
+            val rowDataCopy = new Array[Byte](sortedIterator.getRecordLength)
+            PlatformDependent.copyMemory(
+              sortedIterator.getBaseObject,
+              sortedIterator.getBaseOffset,
+              rowDataCopy,
+              PlatformDependent.BYTE_ARRAY_OFFSET,
+              sortedIterator.getRecordLength
+            )
+            row.pointTo(rowDataCopy, PlatformDependent.BYTE_ARRAY_OFFSET, numFields, schema)
+            sorter.freeMemory()
+            row
+          }
+        }
+      }
+    }
+    child.execute().mapPartitions(doSort, preservesPartitioning = true)
+  }
+
+  override def output: Seq[Attribute] = child.output
+
+  override def outputOrdering: Seq[SortOrder] = sortOrder
+}
+
+
 /**
  * :: DeveloperApi ::
  * Return a new RDD that has exactly `numPartitions` partitions.

sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/SortMergeCompatibilitySuite.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.hive.execution
 
 import org.apache.spark.sql.SQLConf
 import org.apache.spark.sql.hive.test.TestHive
+import org.apache.spark.sql.test.TestSQLContext._
 
 /**
  * Runs the test cases that are included in the hive distribution with sort merge join is true.