diff options
author | Yijie Shen <henry.yijieshen@gmail.com> | 2015-08-22 10:16:35 -0700 |
---|---|---|
committer | Davies Liu <davies.liu@gmail.com> | 2015-08-22 10:16:35 -0700 |
commit | 90cb9f05655a25b95b8f9fe81da14e5b9c8bcf44 (patch) | |
tree | e26ca5aea55a88eb19d5a72ef5618554d9b92c3d /sql | |
parent | 46fcb9e0dbb2b28110f68a3d9f6c0c47bfd197b1 (diff) | |
download | spark-90cb9f05655a25b95b8f9fe81da14e5b9c8bcf44.tar.gz spark-90cb9f05655a25b95b8f9fe81da14e5b9c8bcf44.tar.bz2 spark-90cb9f05655a25b95b8f9fe81da14e5b9c8bcf44.zip |
[SPARK-9401] [SQL] Fully implement code generation for ConcatWs
This PR adds full codegen support for ConcatWs, is a substitute of #7782
JIRA: https://issues.apache.org/jira/browse/SPARK-9401
cc davies
Author: Yijie Shen <henry.yijieshen@gmail.com>
Closes #8353 from yjshen/concatws.
Diffstat (limited to 'sql')
-rw-r--r-- | sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala | 42 |
1 files changed, 39 insertions, 3 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 b60d318534..48d02bb534 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 @@ -72,7 +72,7 @@ case class Concat(children: Seq[Expression]) extends Expression with ImplicitCas * Returns null if the separator is null. Otherwise, concat_ws skips all null values. */ case class ConcatWs(children: Seq[Expression]) - extends Expression with ImplicitCastInputTypes with CodegenFallback { + extends Expression with ImplicitCastInputTypes { require(children.nonEmpty, s"$prettyName requires at least one argument.") @@ -114,8 +114,44 @@ case class ConcatWs(children: Seq[Expression]) boolean ${ev.isNull} = ${ev.primitive} == null; """ } else { - // Contains a mix of strings and array<string>s. Fall back to interpreted mode for now. - super.genCode(ctx, ev) + val array = ctx.freshName("array") + val varargNum = ctx.freshName("varargNum") + val idxInVararg = ctx.freshName("idxInVararg") + + val evals = children.map(_.gen(ctx)) + val (varargCount, varargBuild) = children.tail.zip(evals.tail).map { case (child, eval) => + child.dataType match { + case StringType => + ("", // we count all the StringType arguments num at once below. + s"$array[$idxInVararg ++] = ${eval.isNull} ? (UTF8String) null : ${eval.primitive};") + case _: ArrayType => + val size = ctx.freshName("n") + (s""" + if (!${eval.isNull}) { + $varargNum += ${eval.primitive}.numElements(); + } + """, + s""" + if (!${eval.isNull}) { + final int $size = ${eval.primitive}.numElements(); + for (int j = 0; j < $size; j ++) { + $array[$idxInVararg ++] = ${ctx.getValue(eval.primitive, StringType, "j")}; + } + } + """) + } + }.unzip + + evals.map(_.code).mkString("\n") + + s""" + int $varargNum = ${children.count(_.dataType == StringType) - 1}; + int $idxInVararg = 0; + ${varargCount.mkString("\n")} + UTF8String[] $array = new UTF8String[$varargNum]; + ${varargBuild.mkString("\n")} + UTF8String ${ev.primitive} = UTF8String.concatWs(${evals.head.primitive}, $array); + boolean ${ev.isNull} = ${ev.primitive} == null; + """ } } } |