Narrow down expressions casting

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala

@@ -77,7 +77,7 @@ object AnsiTypeCoercion extends TypeCoercionBase {
     UnpivotCoercion ::
     WidenSetOperationTypes ::
     new AnsiCombinedTypeCoercionRule(
-      CollationTypeCasts ::
+      PreCollationTypeCasts ::
       InConversion ::
       PromoteStrings ::
       DecimalPrecision ::
@@ -93,7 +93,8 @@ object AnsiTypeCoercion extends TypeCoercionBase {
       ImplicitTypeCasts ::
       DateTimeOperations ::
       WindowFrameCoercion ::
-      GetDateFieldOperations:: Nil) :: Nil
+      GetDateFieldOperations ::
+      PostCollationTypeCasts :: Nil) :: Nil
 
   val findTightestCommonType: (DataType, DataType) => Option[DataType] = {
     case (t1, t2) if t1 == t2 => Some(t1)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCasts.scala

@@ -22,25 +22,12 @@ import javax.annotation.Nullable
 import scala.annotation.tailrec
 
 import org.apache.spark.sql.catalyst.analysis.TypeCoercion.hasStringType
-import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Cast, Collate, ComplexTypeMergingExpression, CreateArray, Elt, ExpectsInputTypes, Expression, Predicate, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{ArrayJoin, BinaryExpression, Cast, Collate, ComplexTypeMergingExpression, ConcatWs, CreateArray, Expression, In, InSubquery, Substring}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{AbstractDataType, ArrayType, DataType, StringType}
 
-object CollationTypeCasts extends TypeCoercionRule {
-  override val transform: PartialFunction[Expression, Expression] = {
-    case e if !e.childrenResolved => e
-    case sc @ (_: BinaryExpression
-               | _: ComplexTypeMergingExpression
-               | _: CreateArray
-               | _: Elt
-               | _: ExpectsInputTypes
-               | _: Predicate
-               | _: SortOrder) =>
-      val newChildren = collateToSingleType(sc.children)
-      sc.withNewChildren(newChildren)
-  }
-
+abstract class CollationTypeCasts extends TypeCoercionRule {
   /**
    * Extracts StringTypes from filtered hasStringType
    */
@@ -104,6 +91,7 @@ object CollationTypeCasts extends TypeCoercionRule {
         val implicitTypes = expr.map(_.dataType)
           .filter(hasStringType)
           .map(extractStringType)
+          .filter(dt => dt.collationId != SQLConf.get.defaultStringType.collationId)
 
         if (hasMultipleImplicits(implicitTypes)) {
           throw QueryCompilationErrors.implicitCollationMismatchError()
@@ -127,3 +115,38 @@ object CollationTypeCasts extends TypeCoercionRule {
       .filter(dt => !(dt == SQLConf.get.defaultStringType.collationId)).distinct.size > 1
 
 }
+
+/**
+ * This rule is used to collate all existing expressions related to StringType into a single
+ * collation. Arrays are handled using their elementType and should be cast for these expressions.
+ */
+object PreCollationTypeCasts extends CollationTypeCasts {
+  override val transform: PartialFunction[Expression, Expression] = {
+    case e if !e.childrenResolved => e
+    case sc@(_: In
+             | _: InSubquery
+             | _: CreateArray
+             | _: ComplexTypeMergingExpression
+             | _: ArrayJoin
+             | _: BinaryExpression
+             | _: ConcatWs
+             | _: Substring) =>
+      val newChildren = collateToSingleType(sc.children)
+      sc.withNewChildren(newChildren)
+  }
+}
+
+/**
+ *  This rule is used for managing expressions that have possible implicit casts from different
+ *  types in ImplicitTypeCasts rule.
+ */
+object PostCollationTypeCasts extends CollationTypeCasts {
+  override val transform: PartialFunction[Expression, Expression] = {
+    case e if !e.childrenResolved => e
+    case sc@(_: ArrayJoin
+             | _: BinaryExpression
+             | _: Substring) =>
+      val newChildren = collateToSingleType(sc.children)
+      sc.withNewChildren(newChildren)
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala

@@ -23,8 +23,7 @@ import scala.annotation.tailrec
 import scala.collection.mutable
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.catalyst.analysis.CollationTypeCasts.{castStringType, getOutputCollation}
-import org.apache.spark.sql.catalyst.analysis.TypeCoercion.hasStringType
+import org.apache.spark.sql.catalyst.analysis.PreCollationTypeCasts.getOutputCollation
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
@@ -660,9 +659,8 @@ abstract class TypeCoercionBase {
         val newIndex = implicitCast(index, IntegerType).getOrElse(index)
         val newInputs = if (conf.eltOutputAsString ||
           !children.tail.map(_.dataType).forall(_ == BinaryType)) {
-          val st = getOutputCollation(children)
           children.tail.map { e =>
-            implicitCast(e, st).getOrElse(e)
+            implicitCast(e, SQLConf.get.defaultStringType).getOrElse(e)
           }
         } else {
           children.tail
@@ -707,39 +705,27 @@ abstract class TypeCoercionBase {
         }.getOrElse(b)  // If there is no applicable conversion, leave expression unchanged.
 
       case e: ImplicitCastInputTypes if e.inputTypes.nonEmpty =>
-        val childrenBeforeCollations: Seq[Expression] = e.children.zip(e.inputTypes).map {
+        val children: Seq[Expression] = e.children.zip(e.inputTypes).map {
           // If we cannot do the implicit cast, just use the original input.
           case (in, expected) => implicitCast(in, expected).getOrElse(in)
         }
-        val st = getOutputCollation(e.children)
-        val children: Seq[Expression] = childrenBeforeCollations.map {
-          case in if hasStringType(in.dataType) =>
-            castStringType(in, st).getOrElse(in)
-          case in => in
-        }
         e.withNewChildren(children)
 
       case e: ExpectsInputTypes if e.inputTypes.nonEmpty =>
         // Convert NullType into some specific target type for ExpectsInputTypes that don't do
         // general implicit casting.
-        val childrenBeforeCollations: Seq[Expression] =
+        val children: Seq[Expression] =
           e.children.zip(e.inputTypes).map { case (in, expected) =>
           if (in.dataType == NullType && !expected.acceptsType(NullType)) {
             Literal.create(null, expected.defaultConcreteType)
           } else {
             in
           }
         }
-        val st = getOutputCollation(e.children)
-        val children: Seq[Expression] = childrenBeforeCollations.map {
-          case in if hasStringType(in.dataType) =>
-            castStringType(in, st).getOrElse(in)
-          case in => in
-        }
         e.withNewChildren(children)
 
       case udf: ScalaUDF if udf.inputTypes.nonEmpty =>
-        val childrenBeforeCollations = udf.children.zip(udf.inputTypes).map { case (in, expected) =>
+        val children = udf.children.zip(udf.inputTypes).map { case (in, expected) =>
           // Currently Scala UDF will only expect `AnyDataType` at top level, so this trick works.
           // In the future we should create types like `AbstractArrayType`, so that Scala UDF can
           // accept inputs of array type of arbitrary element type.
@@ -752,12 +738,6 @@ abstract class TypeCoercionBase {
             ).getOrElse(in)
           }
         }
-        val st = getOutputCollation(udf.children)
-        val children: Seq[Expression] = childrenBeforeCollations.map {
-          case in if hasStringType(in.dataType) =>
-            castStringType(in, st).getOrElse(in)
-          case in => in
-        }
         udf.copy(children = children)
     }
 
@@ -860,7 +840,7 @@ object TypeCoercion extends TypeCoercionBase {
     UnpivotCoercion ::
     WidenSetOperationTypes ::
     new CombinedTypeCoercionRule(
-      CollationTypeCasts ::
+      PreCollationTypeCasts ::
       InConversion ::
       PromoteStrings ::
       DecimalPrecision ::
@@ -877,7 +857,8 @@ object TypeCoercion extends TypeCoercionBase {
       ImplicitTypeCasts ::
       DateTimeOperations ::
       WindowFrameCoercion ::
-      StringLiteralCoercion :: Nil) :: Nil
+      StringLiteralCoercion ::
+      PostCollationTypeCasts :: Nil) :: Nil
 
   override def canCast(from: DataType, to: DataType): Boolean = Cast.canCast(from, to)