Repository: spark Updated Branches: refs/heads/master 46fcb9e0d -> 90cb9f056
[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 <[email protected]> Closes #8353 from yjshen/concatws. Project: http://git-wip-us.apache.org/repos/asf/spark/repo Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/90cb9f05 Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/90cb9f05 Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/90cb9f05 Branch: refs/heads/master Commit: 90cb9f05655a25b95b8f9fe81da14e5b9c8bcf44 Parents: 46fcb9e Author: Yijie Shen <[email protected]> Authored: Sat Aug 22 10:16:35 2015 -0700 Committer: Davies Liu <[email protected]> Committed: Sat Aug 22 10:16:35 2015 -0700 ---------------------------------------------------------------------- .../expressions/stringExpressions.scala | 42 ++++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/spark/blob/90cb9f05/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala ---------------------------------------------------------------------- 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 b60d318..48d02bb 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; + """ } } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
