[SPARK-47296][SQL][COLLATION] Fail unsupported functions for non-binary collations

common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java

@@ -65,6 +65,7 @@ public static class Collation {
      * byte for byte equal. All accent or case-insensitive collations are considered non-binary.
      */
     public final boolean isBinaryCollation;
+    public final boolean isLowercaseCollation;
 
     public Collation(
         String collationName,
@@ -79,6 +80,7 @@ public Collation(
       this.version = version;
       this.hashFunction = hashFunction;
       this.isBinaryCollation = isBinaryCollation;
+      this.isLowercaseCollation = collationName.equals("UCS_BASIC_LCASE");
 
       if (isBinaryCollation) {
         this.equalsFunction = UTF8String::equals;

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CollationTypeConstraints.scala

@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
+import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.types.{DataType, StringType}
+
+object CollationTypeConstraints {
+  def checkCollationCompatibility(
+      superCheck: => TypeCheckResult,
+      collationId: Int,
+      rightDataType: DataType): TypeCheckResult = {
+    val checkResult = superCheck
+    if (checkResult.isFailure) return checkResult
+    // Additional check needed for collation compatibility
+    val rightCollationId: Int = rightDataType.asInstanceOf[StringType].collationId
+    if (collationId != rightCollationId) {
+      return DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> CollationFactory.fetchCollation(collationId).collationName,
+          "collationNameRight" -> CollationFactory.fetchCollation(rightCollationId).collationName
+        )
+      )
+    }
+    TypeCheckResult.TypeCheckSuccess
+  }
+
+  object CollationSupportLevel extends Enumeration {
+    type CollationSupportLevel = Value
+
+    val SUPPORT_BINARY_ONLY: Value = Value(0)
+    val SUPPORT_LOWERCASE: Value = Value(1)
+    val SUPPORT_ALL_COLLATIONS: Value = Value(2)
+  }
+
+  def checkCollationSupport(
+      superCheck: => TypeCheckResult,
+      collationId: Int,
+      functionName: String,
+      collationSupportLevel: CollationSupportLevel.CollationSupportLevel)
+  : TypeCheckResult = {
+    val checkResult = superCheck
+    if (checkResult.isFailure) return checkResult
+    // Additional check needed for collation support
+    val collation = CollationFactory.fetchCollation(collationId)
+    collationSupportLevel match {
+      case CollationSupportLevel.SUPPORT_BINARY_ONLY =>
+        if (!collation.isBinaryCollation) {
+          throwUnsupportedCollation(functionName, collation.collationName)
+        }
+      case CollationSupportLevel.SUPPORT_LOWERCASE =>
+        if (!collation.isBinaryCollation && !collation.isLowercaseCollation) {
+          throwUnsupportedCollation(functionName, collation.collationName)
+        }
+      case CollationSupportLevel.SUPPORT_ALL_COLLATIONS => // No additional checks needed
+      case _ => throw new IllegalArgumentException("Invalid collation support level.")
+    }
+    TypeCheckResult.TypeCheckSuccess
+  }
+
+  private def throwUnsupportedCollation(functionName: String, collationName: String): Unit = {
+    throw new SparkException(
+      errorClass = "UNSUPPORTED_COLLATION.FOR_FUNCTION",
+      messageParameters = Map(
+        "functionName" -> functionName,
+        "collationName" -> collationName),
+      cause = null
+    )
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala

@@ -46,6 +46,13 @@ abstract class StringRegexExpression extends BinaryExpression
 
   override def inputTypes: Seq[DataType] = Seq(StringType, StringType)
 
+  final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, right.dataType)
+  }
+
   // try cache foldable pattern
   private lazy val cache: Pattern = right match {
     case p: Expression if p.foldable =>
@@ -130,6 +137,12 @@ abstract class StringRegexExpression extends BinaryExpression
 case class Like(left: Expression, right: Expression, escapeChar: Char)
   extends StringRegexExpression {
 
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, "like",
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   def this(left: Expression, right: Expression) = this(left, right, '\\')
 
   override def escape(v: String): String = StringUtils.escapeLikeRegex(v, escapeChar)
@@ -260,6 +273,16 @@ case class ILike(
 
   override def inputTypes: Seq[AbstractDataType] = Seq(StringType, StringType)
 
+  final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, right.dataType)
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, "ilike",
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override protected def withNewChildrenInternal(
       newLeft: Expression, newRight: Expression): Expression = {
     copy(left = newLeft, right = newRight)
@@ -461,6 +484,12 @@ case class NotLikeAny(child: Expression, patterns: Seq[UTF8String]) extends Like
 // scalastyle:on line.contains.tab line.size.limit
 case class RLike(left: Expression, right: Expression) extends StringRegexExpression {
 
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, "rlike",
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override def escape(v: String): String = v
   override def matches(regex: Pattern, str: String): Boolean = regex.matcher(str).find(0)
   override def toString: String = s"RLIKE($left, $right)"
@@ -545,6 +574,17 @@ case class StringSplit(str: Expression, regex: Expression, limit: Expression)
 
   override def dataType: DataType = ArrayType(StringType, containsNull = false)
   override def inputTypes: Seq[DataType] = Seq(StringType, StringType, IntegerType)
+
+  final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, second.dataType)
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override def first: Expression = str
   override def second: Expression = regex
   override def third: Expression = limit
@@ -614,6 +654,8 @@ case class RegExpReplace(subject: Expression, regexp: Expression, rep: Expressio
   def this(subject: Expression, regexp: Expression, rep: Expression) =
     this(subject, regexp, rep, Literal(1))
 
+  final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId
+
   override def checkInputDataTypes(): TypeCheckResult = {
     val defaultCheck = super.checkInputDataTypes()
     if (defaultCheck.isFailure) {
@@ -643,6 +685,12 @@ case class RegExpReplace(subject: Expression, regexp: Expression, rep: Expressio
         )
       )
     }
+
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, second.dataType)
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
   }
 
   // last regex in string, we will update the pattern iff regexp value changed.
@@ -772,6 +820,13 @@ abstract class RegExpExtractBase
   final override val nodePatterns: Seq[TreePattern] = Seq(REGEXP_EXTRACT_FAMILY)
 
   override def inputTypes: Seq[AbstractDataType] = Seq(StringType, StringType, IntegerType)
+
+  final lazy val collationId: Int = subject.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, regexp.dataType)
+  }
   override def first: Expression = subject
   override def second: Expression = regexp
   override def third: Expression = idx
@@ -849,6 +904,13 @@ case class RegExpExtract(subject: Expression, regexp: Expression, idx: Expressio
   }
 
   override def dataType: DataType = StringType
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override def prettyName: String = "regexp_extract"
 
   override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
@@ -948,6 +1010,12 @@ case class RegExpExtractAll(subject: Expression, regexp: Expression, idx: Expres
   }
 
   override def dataType: DataType = ArrayType(StringType)
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
   override def prettyName: String = "regexp_extract_all"
 
   override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
@@ -1022,6 +1090,16 @@ case class RegExpCount(left: Expression, right: Expression)
 
   override def inputTypes: Seq[AbstractDataType] = Seq(StringType, StringType)
 
+  final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, right.dataType)
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override protected def withNewChildrenInternal(
       newChildren: IndexedSeq[Expression]): RegExpCount =
     copy(left = newChildren(0), right = newChildren(1))
@@ -1061,6 +1139,16 @@ case class RegExpSubStr(left: Expression, right: Expression)
 
   override def inputTypes: Seq[AbstractDataType] = Seq(StringType, StringType)
 
+  final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationCompatibility(
+      super.checkInputDataTypes(), collationId, right.dataType)
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
+
   override protected def withNewChildrenInternal(
       newChildren: IndexedSeq[Expression]): RegExpSubStr =
     copy(left = newChildren(0), right = newChildren(1))
@@ -1113,6 +1201,12 @@ case class RegExpInStr(subject: Expression, regexp: Expression, idx: Expression)
   }
 
   override def dataType: DataType = IntegerType
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    CollationTypeConstraints.checkCollationSupport(
+      super.checkInputDataTypes(), collationId, prettyName,
+      CollationTypeConstraints.CollationSupportLevel.SUPPORT_BINARY_ONLY)
+  }
   override def prettyName: String = "regexp_instr"
 
   override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {