Skip to content

Commit d246e29

Browse files
committed
Fix consideration of column types when choosing sort implementation.
1 parent 6b156fb commit d246e29

File tree

4 files changed

+26
-14
lines changed

4 files changed

+26
-14
lines changed

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,12 @@
3030
import org.apache.spark.TaskContext;
3131
import org.apache.spark.sql.AbstractScalaRowIterator;
3232
import org.apache.spark.sql.catalyst.InternalRow;
33+
import org.apache.spark.sql.catalyst.expressions.ObjectUnsafeColumnWriter;
34+
import org.apache.spark.sql.catalyst.expressions.UnsafeColumnWriter;
3335
import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
3436
import org.apache.spark.sql.catalyst.expressions.UnsafeRowConverter;
3537
import org.apache.spark.sql.catalyst.util.ObjectPool;
38+
import org.apache.spark.sql.types.StructField;
3639
import org.apache.spark.sql.types.StructType;
3740
import org.apache.spark.unsafe.PlatformDependent;
3841
import org.apache.spark.util.collection.unsafe.sort.PrefixComparator;
@@ -164,6 +167,18 @@ public Iterator<InternalRow> sort(Iterator<InternalRow> inputIterator) throws IO
164167
return sort();
165168
}
166169

170+
/**
171+
* Return true if UnsafeExternalRowSorter can sort rows with the given schema, false otherwise.
172+
*/
173+
public static boolean supportsSchema(StructType schema) {
174+
for (StructField field : schema.fields()) {
175+
if (UnsafeColumnWriter.forType(field.dataType()) instanceof ObjectUnsafeColumnWriter) {
176+
return false;
177+
}
178+
}
179+
return true;
180+
}
181+
167182
private static final class RowComparator extends RecordComparator {
168183
private final Ordering<InternalRow> ordering;
169184
private final int numFields;

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnsafeRowConverter.scala

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -103,19 +103,6 @@ class UnsafeRowConverter(fieldTypes: Array[DataType]) {
103103

104104
}
105105

106-
object UnsafeRowConverter {
107-
def supportsSchema(schema: StructType): Boolean = {
108-
schema.forall { field =>
109-
try {
110-
UnsafeColumnWriter.forType(field.dataType)
111-
true
112-
} catch {
113-
case e: UnsupportedOperationException => false
114-
}
115-
}
116-
}
117-
}
118-
119106
/**
120107
* Function for writing a column into an UnsafeRow.
121108
*/

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
298298
* if necessary.
299299
*/
300300
def getSortOperator(sortExprs: Seq[SortOrder], global: Boolean, child: SparkPlan): SparkPlan = {
301-
if (sqlContext.conf.unsafeEnabled && UnsafeRowConverter.supportsSchema(child.schema)) {
301+
if (sqlContext.conf.unsafeEnabled && UnsafeExternalSort.supportsSchema(child.schema)) {
302302
execution.UnsafeExternalSort(sortExprs, global, child)
303303
} else if (sqlContext.conf.externalSortEnabled) {
304304
execution.ExternalSort(sortExprs, global, child)

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,16 @@ case class UnsafeExternalSort(
286286
override def outputOrdering: Seq[SortOrder] = sortOrder
287287
}
288288

289+
@DeveloperApi
290+
object UnsafeExternalSort {
291+
/**
292+
* Return true if UnsafeExternalSort can sort rows with the given schema, false otherwise.
293+
*/
294+
def supportsSchema(schema: StructType): Boolean = {
295+
UnsafeExternalRowSorter.supportsSchema(schema)
296+
}
297+
}
298+
289299

290300
/**
291301
* :: DeveloperApi ::

0 commit comments

Comments
 (0)