You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ma...@apache.org on 2024/12/13 19:32:23 UTC
(spark) branch master updated: [SPARK-50546][SQL] Add subquery cast support to collation type coercion
This is an automated email from the ASF dual-hosted git repository.
maxgekk 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 ebd6b7c5aacd [SPARK-50546][SQL] Add subquery cast support to collation type coercion
ebd6b7c5aacd is described below
commit ebd6b7c5aacd2b3265262c25880f01548e397122
Author: Stefan Kandic <st...@databricks.com>
AuthorDate: Fri Dec 13 22:32:10 2024 +0300
[SPARK-50546][SQL] Add subquery cast support to collation type coercion
### What changes were proposed in this pull request?
Adding proper support for implicit casting of subqueries in collation type coercion. Currently I only added support for `Project` and `Aggregate` plans in subquery expression, but we might need to add more later.
### Why are the changes needed?
Since `SubqueryExpression` is a special type of expression it is not possible to just cast it into a different type, and this was causing a runtime error. Instead we need to go inside of its plan and cast the plan outputs.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added new unit tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #49147 from stefankandic/fixSubquery.
Authored-by: Stefan Kandic <st...@databricks.com>
Signed-off-by: Max Gekk <ma...@gmail.com>
---
.../catalyst/analysis/CollationTypeCoercion.scala | 34 ++++++++
.../collation/CollationTypePrecedenceSuite.scala | 97 ++++++++++++++++++++++
2 files changed, 131 insertions(+)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala
index fa5bb2c07152..980b627729ec 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis
import org.apache.spark.sql.catalyst.analysis.CollationStrength.{Default, Explicit, Implicit}
import org.apache.spark.sql.catalyst.analysis.TypeCoercion.haveSameType
import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Project}
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StringType, StructType}
@@ -172,6 +173,9 @@ object CollationTypeCoercion {
expr match {
case lit: Literal => lit.copy(dataType = newDataType)
case cast: Cast => cast.copy(dataType = newDataType)
+ case subquery: SubqueryExpression =>
+ changeTypeInSubquery(subquery, newType)
+
case _ => Cast(expr, newDataType)
}
@@ -180,6 +184,36 @@ object CollationTypeCoercion {
}
}
+ /**
+ * Changes the data type of the expression in the subquery to the given `newType`.
+ * Currently only supports subqueries with [[Project]] and [[Aggregate]] plan.
+ */
+ private def changeTypeInSubquery(
+ subqueryExpression: SubqueryExpression,
+ newType: DataType): SubqueryExpression = {
+
+ def transformNamedExpressions(ex: NamedExpression): NamedExpression = {
+ changeType(ex, newType) match {
+ case named: NamedExpression => named
+ case other => Alias(other, ex.name)()
+ }
+ }
+
+ val newPlan = subqueryExpression.plan match {
+ case project: Project =>
+ val newProjectList = project.projectList.map(transformNamedExpressions)
+ project.copy(projectList = newProjectList)
+
+ case agg: Aggregate =>
+ val newAggregateExpressions = agg.aggregateExpressions.map(transformNamedExpressions)
+ agg.copy(aggregateExpressions = newAggregateExpressions)
+
+ case other => other
+ }
+
+ subqueryExpression.withNewPlan(newPlan)
+ }
+
/**
* If possible, returns the new data type from `inType` by applying
* the collation of `castType`.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationTypePrecedenceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationTypePrecedenceSuite.scala
index c0e61de57cd9..00fcc501fe74 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationTypePrecedenceSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationTypePrecedenceSuite.scala
@@ -223,6 +223,103 @@ class CollationTypePrecedenceSuite extends QueryTest with SharedSparkSession {
)
}
+ test("in subquery expression") {
+ val tableName = "subquery_tbl"
+ withTable(tableName) {
+ sql(s"""
+ |CREATE TABLE $tableName (
+ | c1 STRING COLLATE UTF8_LCASE,
+ | c2 STRING COLLATE UNICODE
+ |) USING $dataSource
+ |""".stripMargin)
+
+ sql(s"INSERT INTO $tableName VALUES ('a', 'A')")
+
+ assertImplicitMismatch(
+ sql(s"""
+ |SELECT * FROM $tableName
+ |WHERE c1 IN (SELECT c2 FROM $tableName)
+ |""".stripMargin))
+
+ // this fails since subquery expression collation is implicit by default
+ assertImplicitMismatch(
+ sql(s"""
+ |SELECT * FROM $tableName
+ |WHERE c1 IN (SELECT c2 collate unicode FROM $tableName)
+ |""".stripMargin))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate utf8_lcase IN (SELECT c2 collate unicode FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(1)))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate utf8_lcase IN (SELECT c2 FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(1)))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate unicode IN (SELECT c2 FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(0)))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate unicode IN (SELECT c2 FROM $tableName
+ | WHERE c2 collate unicode IN (SELECT c1 FROM $tableName))
+ |""".stripMargin),
+ Seq(Row(0)))
+ }
+ }
+
+ test("scalar subquery") {
+ val tableName = "scalar_subquery_tbl"
+ withTable(tableName) {
+ sql(s"""
+ |CREATE TABLE $tableName (
+ | c1 STRING COLLATE UTF8_LCASE,
+ | c2 STRING COLLATE UNICODE
+ |) USING $dataSource
+ |""".stripMargin)
+
+ sql(s"INSERT INTO $tableName VALUES ('a', 'A')")
+
+ assertImplicitMismatch(
+ sql(s"""
+ |SELECT * FROM $tableName
+ |WHERE c1 = (SELECT MAX(c2) FROM $tableName)
+ |""".stripMargin))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate utf8_lcase = (SELECT MAX(c2) collate unicode FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(1)))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate utf8_lcase = (SELECT MAX(c2) FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(1)))
+
+ checkAnswer(
+ sql(s"""
+ |SELECT COUNT(*) FROM $tableName
+ |WHERE c1 collate unicode = (SELECT MAX(c2) FROM $tableName)
+ |""".stripMargin),
+ Seq(Row(0)))
+ }
+ }
+
test("struct test") {
val tableName = "struct_tbl"
val c1Collation = "UNICODE_CI"
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org