Skip to content
Snippets Groups Projects
Commit 6203668d authored by Liang-Chi Hsieh's avatar Liang-Chi Hsieh Committed by Wenchen Fan
Browse files

[SPARK-16640][SQL] Add codegen for Elt function

## What changes were proposed in this pull request?

Elt function doesn't support codegen execution now. We should add the support.

## How was this patch tested?

Jenkins tests.

Author: Liang-Chi Hsieh <simonh@tw.ibm.com>

Closes #14277 from viirya/elt-codegen.
parent 8674054d
No related branches found
No related tags found
No related merge requests found
......@@ -171,7 +171,7 @@ case class ConcatWs(children: Seq[Expression])
usage = "_FUNC_(n, str1, str2, ...) - returns the n-th string, e.g. returns str2 when n is 2",
extended = "> SELECT _FUNC_(1, 'scala', 'java') FROM src LIMIT 1;\n" + "'scala'")
case class Elt(children: Seq[Expression])
extends Expression with ImplicitCastInputTypes with CodegenFallback {
extends Expression with ImplicitCastInputTypes {
private lazy val indexExpr = children.head
private lazy val stringExprs = children.tail.toArray
......@@ -204,6 +204,29 @@ case class Elt(children: Seq[Expression])
}
}
}
override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
val index = indexExpr.genCode(ctx)
val strings = stringExprs.map(_.genCode(ctx))
val assignStringValue = strings.zipWithIndex.map { case (eval, index) =>
s"""
case ${index + 1}:
${ev.value} = ${eval.isNull} ? null : ${eval.value};
break;
"""
}.mkString("\n")
val indexVal = ctx.freshName("index")
val stringArray = ctx.freshName("strings");
ev.copy(index.code + "\n" + strings.map(_.code).mkString("\n") + s"""
final int $indexVal = ${index.value};
UTF8String ${ev.value} = null;
switch ($indexVal) {
$assignStringValue
}
final boolean ${ev.isNull} = ${ev.value} == null;
""")
}
}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment