Force spilling in UnsafeExternalSortSuite.

Author: JoshRosen

sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java

@@ -44,6 +44,14 @@
 
 final class UnsafeExternalRowSorter {
 
+  /**
+   * If positive, forces records to be spilled to disk at the given frequency (measured in numbers
+   * of records). This is only intended to be used in tests.
+   */
+  private int testSpillFrequency = 0;
+
+  private long numRowsInserted = 0;
+
   private final StructType schema;
   private final UnsafeRowConverter rowConverter;
   private final PrefixComputer prefixComputer;
@@ -77,6 +85,15 @@ public UnsafeExternalRowSorter(
     );
   }
 
+  /**
+   * Forces spills to occur every `frequency` records. Only for use in tests.
+   */
+  @VisibleForTesting
+  void setTestSpillFrequency(int frequency) {
+    assert frequency > 0 : "Frequency must be positive";
+    testSpillFrequency = frequency;
+  }
+
   @VisibleForTesting
   void insertRow(InternalRow row) throws IOException {
     final int sizeRequirement = rowConverter.getSizeRequirement(row);
@@ -99,6 +116,10 @@ void insertRow(InternalRow row) throws IOException {
       sizeRequirement,
       prefix
     );
+    numRowsInserted++;
+    if (testSpillFrequency > 0 && (numRowsInserted % testSpillFrequency) == 0) {
+      spill();
+    }
   }
 
   @VisibleForTesting

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

@@ -253,12 +253,15 @@ case class ExternalSort(
  *
  * @param global when true performs a global sort of all partitions by shuffling the data first
  *               if necessary.
+ * @param testSpillFrequency Method for configuring periodic spilling in unit tests. If set, will
+ *                           spill every `frequency` records.
  */
 @DeveloperApi
 case class UnsafeExternalSort(
     sortOrder: Seq[SortOrder],
     global: Boolean,
-    child: SparkPlan)
+    child: SparkPlan,
+    testSpillFrequency: Int = 0)
   extends UnaryNode {
 
   private[this] val schema: StructType = child.schema
@@ -278,7 +281,11 @@ case class UnsafeExternalSort(
           override def computePrefix(row: InternalRow): Long = prefixComputer(row)
         }
       }
-      new UnsafeExternalRowSorter(schema, ordering, prefixComparator, prefixComputer).sort(iterator)
+      val sorter = new UnsafeExternalRowSorter(schema, ordering, prefixComparator, prefixComputer)
+      if (testSpillFrequency > 0) {
+        sorter.setTestSpillFrequency(testSpillFrequency)
+      }
+      sorter.sort(iterator)
     }
     child.execute().mapPartitions(doSort, preservesPartitioning = true)
   }

sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeExternalSortSuite.scala

@@ -54,9 +54,10 @@ class UnsafeExternalSortSuite extends SparkPlanTest with BeforeAndAfterAll {
           TestSQLContext.sparkContext.parallelize(Random.shuffle(inputData).map(v => Row(v))),
           StructType(StructField("a", dataType, nullable = true) :: Nil)
         )
+        assert(UnsafeExternalSort.supportsSchema(inputDf.schema))
         checkAnswer(
           inputDf,
-          UnsafeExternalSort(sortOrder, global = false, _: SparkPlan),
+          UnsafeExternalSort(sortOrder, global = false, _: SparkPlan, testSpillFrequency = 100),
           Sort(sortOrder, global = false, _: SparkPlan)
         )
       }