diff options
author | Reynold Xin <rxin@databricks.com> | 2015-07-19 16:48:47 -0700 |
---|---|---|
committer | Reynold Xin <rxin@databricks.com> | 2015-07-19 16:48:47 -0700 |
commit | 163e3f1df94f6b7d3dadb46a87dbb3a2bade3f95 (patch) | |
tree | da58016b49a92c5b0318ace791844194b3838324 /sql | |
parent | 7a81245345f2d6124423161786bb0d9f1c278ab8 (diff) | |
download | spark-163e3f1df94f6b7d3dadb46a87dbb3a2bade3f95.tar.gz spark-163e3f1df94f6b7d3dadb46a87dbb3a2bade3f95.tar.bz2 spark-163e3f1df94f6b7d3dadb46a87dbb3a2bade3f95.zip |
[SPARK-8241][SQL] string function: concat_ws.
I also changed the semantics of concat w.r.t. null back to the same behavior as Hive.
That is to say, concat now returns null if any input is null.
Author: Reynold Xin <rxin@databricks.com>
Closes #7504 from rxin/concat_ws and squashes the following commits:
83fd950 [Reynold Xin] Fixed type casting.
3ae85f7 [Reynold Xin] Write null better.
cdc7be6 [Reynold Xin] Added code generation for pure string mode.
a61c4e4 [Reynold Xin] Updated comments.
2d51406 [Reynold Xin] [SPARK-8241][SQL] string function: concat_ws.
Diffstat (limited to 'sql')
8 files changed, 154 insertions, 20 deletions
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala index 4b256adcc6..71e87b98d8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala @@ -153,6 +153,7 @@ object FunctionRegistry { expression[Ascii]("ascii"), expression[Base64]("base64"), expression[Concat]("concat"), + expression[ConcatWs]("concat_ws"), expression[Encode]("encode"), expression[Decode]("decode"), expression[FormatNumber]("format_number"), @@ -211,7 +212,10 @@ object FunctionRegistry { val builder = (expressions: Seq[Expression]) => { if (varargCtor.isDefined) { // If there is an apply method that accepts Seq[Expression], use that one. - varargCtor.get.newInstance(expressions).asInstanceOf[Expression] + Try(varargCtor.get.newInstance(expressions).asInstanceOf[Expression]) match { + case Success(e) => e + case Failure(e) => throw new AnalysisException(e.getMessage) + } } else { // Otherwise, find an ctor method that matches the number of arguments, and use that. val params = Seq.fill(expressions.size)(classOf[Expression]) @@ -221,7 +225,10 @@ object FunctionRegistry { case Failure(e) => throw new AnalysisException(s"Invalid number of arguments for function $name") } - f.newInstance(expressions : _*).asInstanceOf[Expression] + Try(f.newInstance(expressions : _*).asInstanceOf[Expression]) match { + case Success(e) => e + case Failure(e) => throw new AnalysisException(e.getMessage) + } } } (name, builder) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala index 560b1bc2d8..5f8ac716f7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala @@ -34,19 +34,14 @@ import org.apache.spark.unsafe.types.UTF8String /** * An expression that concatenates multiple input strings into a single string. - * Input expressions that are evaluated to nulls are skipped. - * - * For example, `concat("a", null, "b")` is evaluated to `"ab"`. - * - * Note that this is different from Hive since Hive outputs null if any input is null. - * We never output null. + * If any input is null, concat returns null. */ case class Concat(children: Seq[Expression]) extends Expression with ImplicitCastInputTypes { override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(StringType) override def dataType: DataType = StringType - override def nullable: Boolean = false + override def nullable: Boolean = children.exists(_.nullable) override def foldable: Boolean = children.forall(_.foldable) override def eval(input: InternalRow): Any = { @@ -56,15 +51,76 @@ case class Concat(children: Seq[Expression]) extends Expression with ImplicitCas override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = { val evals = children.map(_.gen(ctx)) - val inputs = evals.map { eval => s"${eval.isNull} ? null : ${eval.primitive}" }.mkString(", ") + val inputs = evals.map { eval => + s"${eval.isNull} ? (UTF8String)null : ${eval.primitive}" + }.mkString(", ") evals.map(_.code).mkString("\n") + s""" boolean ${ev.isNull} = false; UTF8String ${ev.primitive} = UTF8String.concat($inputs); + if (${ev.primitive} == null) { + ${ev.isNull} = true; + } """ } } +/** + * An expression that concatenates multiple input strings or array of strings into a single string, + * using a given separator (the first child). + * + * 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 { + + require(children.nonEmpty, s"$prettyName requires at least one argument.") + + override def prettyName: String = "concat_ws" + + /** The 1st child (separator) is str, and rest are either str or array of str. */ + override def inputTypes: Seq[AbstractDataType] = { + val arrayOrStr = TypeCollection(ArrayType(StringType), StringType) + StringType +: Seq.fill(children.size - 1)(arrayOrStr) + } + + override def dataType: DataType = StringType + + override def nullable: Boolean = children.head.nullable + override def foldable: Boolean = children.forall(_.foldable) + + override def eval(input: InternalRow): Any = { + val flatInputs = children.flatMap { child => + child.eval(input) match { + case s: UTF8String => Iterator(s) + case arr: Seq[_] => arr.asInstanceOf[Seq[UTF8String]] + case null => Iterator(null.asInstanceOf[UTF8String]) + } + } + UTF8String.concatWs(flatInputs.head, flatInputs.tail : _*) + } + + override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = { + if (children.forall(_.dataType == StringType)) { + // All children are strings. In that case we can construct a fixed size array. + val evals = children.map(_.gen(ctx)) + + val inputs = evals.map { eval => + s"${eval.isNull} ? (UTF8String)null : ${eval.primitive}" + }.mkString(", ") + + evals.map(_.code).mkString("\n") + s""" + UTF8String ${ev.primitive} = UTF8String.concatWs($inputs); + 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) + } + } +} + + trait StringRegexExpression extends ImplicitCastInputTypes { self: BinaryExpression => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataType.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataType.scala index 2d133eea19..e98fd2583b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataType.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataType.scala @@ -79,7 +79,7 @@ abstract class DataType extends AbstractDataType { override private[sql] def defaultConcreteType: DataType = this - override private[sql] def acceptsType(other: DataType): Boolean = this == other + override private[sql] def acceptsType(other: DataType): Boolean = sameType(other) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala index f9442bccc4..7ee2333a81 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala @@ -37,7 +37,6 @@ class HiveTypeCoercionSuite extends PlanTest { shouldCast(NullType, IntegerType, IntegerType) shouldCast(NullType, DecimalType, DecimalType.Unlimited) - // TODO: write the entire implicit cast table out for test cases. shouldCast(ByteType, IntegerType, IntegerType) shouldCast(IntegerType, IntegerType, IntegerType) shouldCast(IntegerType, LongType, LongType) @@ -86,6 +85,16 @@ class HiveTypeCoercionSuite extends PlanTest { DecimalType.Unlimited, DecimalType(10, 2)).foreach { tpe => shouldCast(tpe, NumericType, tpe) } + + shouldCast( + ArrayType(StringType, false), + TypeCollection(ArrayType(StringType), StringType), + ArrayType(StringType, false)) + + shouldCast( + ArrayType(StringType, true), + TypeCollection(ArrayType(StringType), StringType), + ArrayType(StringType, true)) } test("ineligible implicit type cast") { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala index 0ed567a90d..96f433be8b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala @@ -26,7 +26,7 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { test("concat") { def testConcat(inputs: String*): Unit = { - val expected = inputs.filter(_ != null).mkString + val expected = if (inputs.contains(null)) null else inputs.mkString checkEvaluation(Concat(inputs.map(Literal.create(_, StringType))), expected, EmptyRow) } @@ -46,6 +46,35 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { // scalastyle:on } + test("concat_ws") { + def testConcatWs(expected: String, sep: String, inputs: Any*): Unit = { + val inputExprs = inputs.map { + case s: Seq[_] => Literal.create(s, ArrayType(StringType)) + case null => Literal.create(null, StringType) + case s: String => Literal.create(s, StringType) + } + val sepExpr = Literal.create(sep, StringType) + checkEvaluation(ConcatWs(sepExpr +: inputExprs), expected, EmptyRow) + } + + // scalastyle:off + // non ascii characters are not allowed in the code, so we disable the scalastyle here. + testConcatWs(null, null) + testConcatWs(null, null, "a", "b") + testConcatWs("", "") + testConcatWs("ab", "哈哈", "ab") + testConcatWs("a哈哈b", "哈哈", "a", "b") + testConcatWs("a哈哈b", "哈哈", "a", null, "b") + testConcatWs("a哈哈b哈哈c", "哈哈", null, "a", null, "b", "c") + + testConcatWs("ab", "哈哈", Seq("ab")) + testConcatWs("a哈哈b", "哈哈", Seq("a", "b")) + testConcatWs("a哈哈b哈哈c哈哈d", "哈哈", Seq("a", null, "b"), null, "c", Seq(null, "d")) + testConcatWs("a哈哈b哈哈c", "哈哈", Seq("a", null, "b"), null, "c", Seq.empty[String]) + testConcatWs("a哈哈b哈哈c", "哈哈", Seq("a", null, "b"), null, "c", Seq[String](null)) + // scalastyle:on + } + test("StringComparison") { val row = create_row("abc", null) val c1 = 'a.string.at(0) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala index f67c89437b..b5140dca04 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala @@ -1733,6 +1733,30 @@ object functions { } /** + * Concatenates input strings together into a single string, using the given separator. + * + * @group string_funcs + * @since 1.5.0 + */ + @scala.annotation.varargs + def concat_ws(sep: String, exprs: Column*): Column = { + ConcatWs(Literal.create(sep, StringType) +: exprs.map(_.expr)) + } + + /** + * Concatenates input strings together into a single string, using the given separator. + * + * This is the variant of concat_ws that takes in the column names. + * + * @group string_funcs + * @since 1.5.0 + */ + @scala.annotation.varargs + def concat_ws(sep: String, columnName: String, columnNames: String*): Column = { + concat_ws(sep, (columnName +: columnNames).map(Column.apply) : _*) + } + + /** * Computes the length of a given string / binary value. * * @group string_funcs diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 4eff33ed45..fe4de8d8b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -30,14 +30,25 @@ class StringFunctionsSuite extends QueryTest { val df = Seq[(String, String, String)](("a", "b", null)).toDF("a", "b", "c") checkAnswer( - df.select(concat($"a", $"b", $"c")), - Row("ab")) + df.select(concat($"a", $"b"), concat($"a", $"b", $"c")), + Row("ab", null)) checkAnswer( - df.selectExpr("concat(a, b, c)"), - Row("ab")) + df.selectExpr("concat(a, b)", "concat(a, b, c)"), + Row("ab", null)) } + test("string concat_ws") { + val df = Seq[(String, String, String)](("a", "b", null)).toDF("a", "b", "c") + + checkAnswer( + df.select(concat_ws("||", $"a", $"b", $"c")), + Row("a||b")) + + checkAnswer( + df.selectExpr("concat_ws('||', a, b, c)"), + Row("a||b")) + } test("string Levenshtein distance") { val df = Seq(("kitten", "sitting"), ("frog", "fog")).toDF("l", "r") diff --git a/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala b/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala index 2689d904d6..b12b3838e6 100644 --- a/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala +++ b/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala @@ -263,9 +263,6 @@ class HiveCompatibilitySuite extends HiveQueryFileTest with BeforeAndAfter { "timestamp_2", "timestamp_udf", - // Hive outputs NULL if any concat input has null. We never output null for concat. - "udf_concat", - // Unlike Hive, we do support log base in (0, 1.0], therefore disable this "udf7" ) @@ -856,6 +853,7 @@ class HiveCompatibilitySuite extends HiveQueryFileTest with BeforeAndAfter { "udf_case", "udf_ceil", "udf_ceiling", + "udf_concat", "udf_concat_insert1", "udf_concat_insert2", "udf_concat_ws", |