You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2024/04/17 16:55:54 UTC
(spark) branch master updated: [SPARK-47360][SQL] Collation support: Overlay, FormatString, Length, BitLength, OctetLength, SoundEx, Luhncheck
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new f99b1b1dd14d [SPARK-47360][SQL] Collation support: Overlay, FormatString, Length, BitLength, OctetLength, SoundEx, Luhncheck
f99b1b1dd14d is described below
commit f99b1b1dd14d10c5a7d5a851b4add64c10ce62f6
Author: Nikola Mandic <ni...@databricks.com>
AuthorDate: Thu Apr 18 00:55:40 2024 +0800
[SPARK-47360][SQL] Collation support: Overlay, FormatString, Length, BitLength, OctetLength, SoundEx, Luhncheck
### What changes were proposed in this pull request?
Add collation support for batch of string expressions mentioned in PR title.
### Why are the changes needed?
Add collations support in string functions.
### Does this PR introduce _any_ user-facing change?
Yes, it changes behavior of string functions when string parameters have collation.
### How was this patch tested?
Added checks to `CollationStringExpressionsSuite`.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #46003 from nikolamand-db/SPARK-47360.
Authored-by: Nikola Mandic <ni...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../sql/catalyst/analysis/CollationTypeCasts.scala | 6 +-
.../spark/sql/catalyst/encoders/EncoderUtils.scala | 1 +
.../catalyst/expressions/stringExpressions.scala | 38 ++++---
.../sql/CollationStringExpressionsSuite.scala | 119 +++++++++++++++++++++
4 files changed, 146 insertions(+), 18 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCasts.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCasts.scala
index 795e8a696b01..cffdd2872224 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCasts.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCasts.scala
@@ -22,7 +22,7 @@ import javax.annotation.Nullable
import scala.annotation.tailrec
import org.apache.spark.sql.catalyst.analysis.TypeCoercion.{hasStringType, haveSameType}
-import org.apache.spark.sql.catalyst.expressions.{ArrayJoin, BinaryExpression, CaseWhen, Cast, Coalesce, Collate, Concat, ConcatWs, CreateArray, Elt, Expression, Greatest, If, In, InSubquery, Least}
+import org.apache.spark.sql.catalyst.expressions.{ArrayJoin, BinaryExpression, CaseWhen, Cast, Coalesce, Collate, Concat, ConcatWs, CreateArray, Elt, Expression, Greatest, If, In, InSubquery, Least, Overlay}
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
@@ -48,6 +48,10 @@ object CollationTypeCasts extends TypeCoercionRule {
case eltExpr: Elt =>
eltExpr.withNewChildren(eltExpr.children.head +: collateToSingleType(eltExpr.children.tail))
+ case overlay: Overlay =>
+ overlay.withNewChildren(collateToSingleType(Seq(overlay.input, overlay.replace))
+ ++ Seq(overlay.pos, overlay.len))
+
case otherExpr @ (
_: In | _: InSubquery | _: CreateArray | _: ArrayJoin | _: Concat | _: Greatest | _: Least |
_: Coalesce | _: BinaryExpression | _: ConcatWs) =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/EncoderUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/EncoderUtils.scala
index 20f86a32c1a1..81743251bada 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/EncoderUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/EncoderUtils.scala
@@ -77,6 +77,7 @@ object EncoderUtils {
case _: DecimalType => classOf[Decimal]
case _: DayTimeIntervalType => classOf[PhysicalLongType.InternalType]
case _: YearMonthIntervalType => classOf[PhysicalIntegerType.InternalType]
+ case _: StringType => classOf[UTF8String]
case _: StructType => classOf[InternalRow]
case _: ArrayType => classOf[ArrayData]
case _: MapType => classOf[MapData]
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 b3029302c03d..3c9888940221 100755
--- 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
@@ -804,8 +804,9 @@ case class Overlay(input: Expression, replace: Expression, pos: Expression, len:
override def dataType: DataType = input.dataType
- override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(StringType, BinaryType),
- TypeCollection(StringType, BinaryType), IntegerType, IntegerType)
+ override def inputTypes: Seq[AbstractDataType] = Seq(
+ TypeCollection(StringTypeAnyCollation, BinaryType),
+ TypeCollection(StringTypeAnyCollation, BinaryType), IntegerType, IntegerType)
override def checkInputDataTypes(): TypeCheckResult = {
val inputTypeCheck = super.checkInputDataTypes()
@@ -818,7 +819,7 @@ case class Overlay(input: Expression, replace: Expression, pos: Expression, len:
}
private lazy val replaceFunc = input.dataType match {
- case StringType =>
+ case _: StringType =>
(inputEval: Any, replaceEval: Any, posEval: Int, lenEval: Int) => {
Overlay.calculate(
inputEval.asInstanceOf[UTF8String],
@@ -1698,10 +1699,10 @@ case class FormatString(children: Expression*) extends Expression with ImplicitC
override def foldable: Boolean = children.forall(_.foldable)
override def nullable: Boolean = children(0).nullable
- override def dataType: DataType = StringType
+ override def dataType: DataType = children(0).dataType
override def inputTypes: Seq[AbstractDataType] =
- StringType :: List.fill(children.size - 1)(AnyDataType)
+ StringTypeAnyCollation :: List.fill(children.size - 1)(AnyDataType)
override def checkInputDataTypes(): TypeCheckResult = {
if (children.isEmpty) {
@@ -2064,16 +2065,17 @@ case class Left(str: Expression, len: Expression) extends RuntimeReplaceable
case class Length(child: Expression)
extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant {
override def dataType: DataType = IntegerType
- override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(StringType, BinaryType))
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(TypeCollection(StringTypeAnyCollation, BinaryType))
protected override def nullSafeEval(value: Any): Any = child.dataType match {
- case StringType => value.asInstanceOf[UTF8String].numChars
+ case _: StringType => value.asInstanceOf[UTF8String].numChars
case BinaryType => value.asInstanceOf[Array[Byte]].length
}
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
child.dataType match {
- case StringType => defineCodeGen(ctx, ev, c => s"($c).numChars()")
+ case _: StringType => defineCodeGen(ctx, ev, c => s"($c).numChars()")
case BinaryType => defineCodeGen(ctx, ev, c => s"($c).length")
}
}
@@ -2098,16 +2100,17 @@ case class Length(child: Expression)
case class BitLength(child: Expression)
extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant {
override def dataType: DataType = IntegerType
- override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(StringType, BinaryType))
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(TypeCollection(StringTypeAnyCollation, BinaryType))
protected override def nullSafeEval(value: Any): Any = child.dataType match {
- case StringType => value.asInstanceOf[UTF8String].numBytes * 8
+ case _: StringType => value.asInstanceOf[UTF8String].numBytes * 8
case BinaryType => value.asInstanceOf[Array[Byte]].length * 8
}
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
child.dataType match {
- case StringType => defineCodeGen(ctx, ev, c => s"($c).numBytes() * 8")
+ case _: StringType => defineCodeGen(ctx, ev, c => s"($c).numBytes() * 8")
case BinaryType => defineCodeGen(ctx, ev, c => s"($c).length * 8")
}
}
@@ -2136,16 +2139,17 @@ case class BitLength(child: Expression)
case class OctetLength(child: Expression)
extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant {
override def dataType: DataType = IntegerType
- override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(StringType, BinaryType))
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(TypeCollection(StringTypeAnyCollation, BinaryType))
protected override def nullSafeEval(value: Any): Any = child.dataType match {
- case StringType => value.asInstanceOf[UTF8String].numBytes
+ case _: StringType => value.asInstanceOf[UTF8String].numBytes
case BinaryType => value.asInstanceOf[Array[Byte]].length
}
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
child.dataType match {
- case StringType => defineCodeGen(ctx, ev, c => s"($c).numBytes()")
+ case _: StringType => defineCodeGen(ctx, ev, c => s"($c).numBytes()")
case BinaryType => defineCodeGen(ctx, ev, c => s"($c).length")
}
}
@@ -2320,9 +2324,9 @@ case class Levenshtein(
case class SoundEx(child: Expression)
extends UnaryExpression with ExpectsInputTypes with NullIntolerant {
- override def dataType: DataType = StringType
+ override def dataType: DataType = SQLConf.get.defaultStringType
- override def inputTypes: Seq[DataType] = Seq(StringType)
+ override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation)
override def nullSafeEval(input: Any): Any = input.asInstanceOf[UTF8String].soundex()
@@ -3290,7 +3294,7 @@ case class Luhncheck(input: Expression) extends RuntimeReplaceable with Implicit
Seq(input),
inputTypes)
- override def inputTypes: Seq[AbstractDataType] = Seq(StringType)
+ override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation)
override def prettyName: String = "luhn_check"
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala
index e815a5051ed2..f0be9cc89a4d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala
@@ -24,6 +24,7 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataType, IntegerType, StringType}
+// scalastyle:off nonascii
class CollationStringExpressionsSuite
extends QueryTest
with SharedSparkSession {
@@ -307,9 +308,127 @@ class CollationStringExpressionsSuite
})
}
+ test("Overlay string expression with collation") {
+ // Supported collations
+ case class OverlayTestCase(l: String, r: String, pos: Int, c: String, result: String)
+ val testCases = Seq(
+ OverlayTestCase("hello", " world", 6, "UTF8_BINARY", "hello world"),
+ OverlayTestCase("nice", " day", 5, "UTF8_BINARY_LCASE", "nice day"),
+ OverlayTestCase("A", "B", 1, "UNICODE", "B"),
+ OverlayTestCase("!", "!!!", 1, "UNICODE_CI", "!!!")
+ )
+ testCases.foreach(t => {
+ val query =
+ s"""
+ |select overlay(collate('${t.l}', '${t.c}') placing
+ |collate('${t.r}', '${t.c}') from ${t.pos})
+ |""".stripMargin
+ // Result & data type
+ checkAnswer(sql(query), Row(t.result))
+ assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.c)))
+ // Implicit casting
+ checkAnswer(sql(
+ s"""
+ |select overlay(collate('${t.l}', '${t.c}') placing '${t.r}' from ${t.pos})
+ |""".stripMargin), Row(t.result))
+ checkAnswer(sql(
+ s"""
+ |select overlay('${t.l}' placing collate('${t.r}', '${t.c}') from ${t.pos})
+ |""".stripMargin), Row(t.result))
+ checkAnswer(sql(
+ s"""
+ |select overlay(collate('${t.l}', '${t.c}')
+ |placing '${t.r}' from collate('${t.pos}', '${t.c}'))
+ |""".stripMargin), Row(t.result))
+ })
+ // Collation mismatch
+ assert(
+ intercept[AnalysisException] {
+ sql("SELECT overlay('a' collate UNICODE PLACING 'b' collate UNICODE_CI FROM 1)")
+ }.getErrorClass == "COLLATION_MISMATCH.EXPLICIT"
+ )
+ }
+
+ test("FormatString string expression with collation") {
+ // Supported collations
+ case class FormatStringTestCase(f: String, a: Seq[Any], c: String, r: String)
+ val testCases = Seq(
+ FormatStringTestCase("%s%s", Seq("'a'", "'b'"), "UTF8_BINARY", "ab"),
+ FormatStringTestCase("%d", Seq(123), "UTF8_BINARY_LCASE", "123"),
+ FormatStringTestCase("%s%d", Seq("'A'", 0), "UNICODE", "A0"),
+ FormatStringTestCase("%s%s", Seq("'Hello'", "'!!!'"), "UNICODE_CI", "Hello!!!")
+ )
+ testCases.foreach(t => {
+ val query =
+ s"""
+ |select format_string(collate('${t.f}', '${t.c}'), ${t.a.mkString(", ")})
+ |""".stripMargin
+ // Result & data type
+ checkAnswer(sql(query), Row(t.r))
+ assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.c)))
+ })
+ }
+
+ test("SoundEx string expression with collation") {
+ // Supported collations
+ case class SoundExTestCase(q: String, c: String, r: String)
+ val testCases = Seq(
+ SoundExTestCase("select soundex('A' collate utf8_binary)", "UTF8_BINARY", "A000"),
+ SoundExTestCase("select soundex('!' collate utf8_binary_lcase)", "UTF8_BINARY_LCASE", "!"),
+ SoundExTestCase("select soundex('$' collate unicode)", "UNICODE", "$"),
+ SoundExTestCase("select soundex('X' collate unicode_ci)", "UNICODE_CI", "X000")
+ )
+ testCases.foreach(t => {
+ withSQLConf(SQLConf.DEFAULT_COLLATION.key -> t.c) {
+ // Result & data type
+ checkAnswer(sql(t.q), Row(t.r))
+ assert(sql(t.q).schema.fields.head.dataType.sameType(StringType(t.c)))
+ }
+ })
+ }
+
+ test("Length, BitLength & OctetLength string expressions with collations") {
+ // Supported collations
+ case class LenTestCase(q: String, r: Int)
+ val testCases = Seq(
+ LenTestCase("select length('hello' collate utf8_binary)", 5),
+ LenTestCase("select length('world' collate utf8_binary_lcase)", 5),
+ LenTestCase("select length('ff' collate unicode)", 1),
+ LenTestCase("select bit_length('hello' collate unicode_ci)", 40),
+ LenTestCase("select bit_length('world' collate utf8_binary)", 40),
+ LenTestCase("select bit_length('ff' collate utf8_binary_lcase)", 24),
+ LenTestCase("select octet_length('hello' collate unicode)", 5),
+ LenTestCase("select octet_length('world' collate unicode_ci)", 5),
+ LenTestCase("select octet_length('ff' collate utf8_binary)", 3)
+ )
+ testCases.foreach(t => {
+ // Result & data type
+ checkAnswer(sql(t.q), Row(t.r))
+ assert(sql(t.q).schema.fields.head.dataType.sameType(IntegerType))
+ })
+ }
+
+ test("Luhncheck string expression with collation") {
+ // Supported collations
+ case class LuhncheckTestCase(q: String, c: String, r: Boolean)
+ val testCases = Seq(
+ LuhncheckTestCase("123", "UTF8_BINARY", r = false),
+ LuhncheckTestCase("000", "UTF8_BINARY_LCASE", r = true),
+ LuhncheckTestCase("111", "UNICODE", r = false),
+ LuhncheckTestCase("222", "UNICODE_CI", r = false)
+ )
+ testCases.foreach(t => {
+ val query = s"select luhn_check(${t.q})"
+ // Result & data type
+ checkAnswer(sql(query), Row(t.r))
+ assert(sql(query).schema.fields.head.dataType.sameType(BooleanType))
+ })
+ }
+
// TODO: Add more tests for other string expressions
}
+// scalastyle:on nonascii
class CollationStringExpressionsANSISuite extends CollationStringExpressionsSuite {
override protected def sparkConf: SparkConf =
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org