Github user kiszk commented on a diff in the pull request:
https://github.com/apache/spark/pull/22976#discussion_r231886071
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala
---
@@ -68,57 +68,51 @@ object GenerateOrdering extends
CodeGenerator[Seq[SortOrder], Ordering[InternalR
genComparisons(ctx, ordering)
}
+ /**
+ * Creates the variables for ordering based on the given order.
+ */
+ private def createOrderKeys(
+ ctx: CodegenContext,
+ row: String,
+ ordering: Seq[SortOrder]): Seq[ExprCode] = {
+ ctx.INPUT_ROW = row
+ ctx.currentVars = null
+ ordering.map(_.child.genCode(ctx))
+ }
+
/**
* Generates the code for ordering based on the given order.
*/
def genComparisons(ctx: CodegenContext, ordering: Seq[SortOrder]):
String = {
val oldInputRow = ctx.INPUT_ROW
val oldCurrentVars = ctx.currentVars
- val inputRow = "i"
- ctx.INPUT_ROW = inputRow
// to use INPUT_ROW we must make sure currentVars is null
ctx.currentVars = null
-
- val comparisons = ordering.map { order =>
- val eval = order.child.genCode(ctx)
- val asc = order.isAscending
- val isNullA = ctx.freshName("isNullA")
- val primitiveA = ctx.freshName("primitiveA")
- val isNullB = ctx.freshName("isNullB")
- val primitiveB = ctx.freshName("primitiveB")
+ val rowAKeys = createOrderKeys(ctx, "a", ordering)
+ val rowBKeys = createOrderKeys(ctx, "b", ordering)
+ val comparisons = rowAKeys.zip(rowBKeys).zipWithIndex.map { case ((l,
r), i) =>
+ val dt = ordering(i).child.dataType
+ val asc = ordering(i).isAscending
+ val nullOrdering = ordering(i).nullOrdering
s"""
- ${ctx.INPUT_ROW} = a;
- boolean $isNullA;
- ${CodeGenerator.javaType(order.child.dataType)} $primitiveA;
- {
- ${eval.code}
- $isNullA = ${eval.isNull};
- $primitiveA = ${eval.value};
- }
- ${ctx.INPUT_ROW} = b;
- boolean $isNullB;
- ${CodeGenerator.javaType(order.child.dataType)} $primitiveB;
- {
- ${eval.code}
- $isNullB = ${eval.isNull};
- $primitiveB = ${eval.value};
- }
- if ($isNullA && $isNullB) {
+ ${l.code}
+ ${r.code}
+ if (${l.isNull} && ${r.isNull}) {
// Nothing
- } else if ($isNullA) {
+ } else if (${l.isNull}) {
return ${
- order.nullOrdering match {
- case NullsFirst => "-1"
- case NullsLast => "1"
- }};
- } else if ($isNullB) {
+ nullOrdering match {
+ case NullsFirst => "-1"
+ case NullsLast => "1"
+ }};
+ } else if (${r.isNull}) {
return ${
- order.nullOrdering match {
- case NullsFirst => "1"
- case NullsLast => "-1"
- }};
+ nullOrdering match {
--- End diff --
ditto
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
