diff options
author | Liang-Chi Hsieh <simonh@tw.ibm.com> | 2016-07-21 20:54:17 +0800 |
---|---|---|
committer | Wenchen Fan <wenchen@databricks.com> | 2016-07-21 20:54:17 +0800 |
commit | 6203668d50c1193cde20b0c35f95cf12dabc0805 (patch) | |
tree | 2bad16cfda9873b0e15371efa76e1289ede769cc /sql/catalyst/src/main/scala | |
parent | 8674054d3402b400a4766fe1c9214001cebf2106 (diff) | |
download | spark-6203668d50c1193cde20b0c35f95cf12dabc0805.tar.gz spark-6203668d50c1193cde20b0c35f95cf12dabc0805.tar.bz2 spark-6203668d50c1193cde20b0c35f95cf12dabc0805.zip |
[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.
Diffstat (limited to 'sql/catalyst/src/main/scala')
-rw-r--r-- | sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala | 25 |
1 files changed, 24 insertions, 1 deletions
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index 61549c9a23..fc13845a7f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -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; + """) + } } |