[SPARK-47476][SQL] Support REPLACE function to work with collated strings

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

@@ -176,15 +176,29 @@ public Collation(
    */
 
   public static StringSearch getStringSearch(
-      final UTF8String left,
-      final UTF8String right,
+      final UTF8String targetUTF8String,
+      final UTF8String patternUTF8String,
       final int collationId) {
-    String pattern = right.toString();
-    CharacterIterator target = new StringCharacterIterator(left.toString());
+
+    if (collationId == UTF8_BINARY_LCASE_COLLATION_ID) {
+      return getStringSearch(targetUTF8String.toLowerCase(), patternUTF8String.toLowerCase());
+    }
+
+    String pattern = patternUTF8String.toString();
+    CharacterIterator target = new StringCharacterIterator(targetUTF8String.toString());
     Collator collator = CollationFactory.fetchCollation(collationId).collator;
     return new StringSearch(pattern, target, (RuleBasedCollator) collator);
   }
 
+  public static StringSearch getStringSearch(
+      final UTF8String targetUTF8String,
+      final UTF8String patternUTF8String) {
+    String pattern = patternUTF8String.toString();
+    String target = targetUTF8String.toString();
+
+    return new StringSearch(pattern, target);
+  }
+
   /**
    * Returns the collation id for the given collation name.
    */

common/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java

@@ -1136,6 +1136,64 @@ public UTF8String replace(UTF8String search, UTF8String replace) {
     return buf.build();
   }
 
+  /**
+   * Replace all occurrences of search in this with replace respecting collation with id = collationId.
+   * @param search the string to be searched
+   * @param replace the start position of the current string for searching
+   * @param collationId the id of applied collation
+   * @return the string with replace instead of search in all places
+   */
+  public UTF8String replace(UTF8String search, UTF8String replace, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      return this.replace(search, replace);
+    }
+    return collationAwareReplace(search, replace, collationId);
+  }
+
+  private UTF8String collationAwareReplace(UTF8String search, UTF8String replace, int collationId) {
+    // This collation aware implementation is based on existing implementation on UTF8String with default collation
+    if (numBytes == 0 || search.numBytes == 0) {
+      return this;
+    }
+
+    StringSearch stringSearch = CollationFactory.getStringSearch(this, search, collationId);
+
+    // Find the first occurrence of the search string.
+    int end = stringSearch.next();
+    if (end == StringSearch.DONE) {
+      // Search string was not found, so string is unchanged.
+      return this;
+    }
+
+    // Initialize byte positions
+    int c = 0;
+    int byteStart = 0; // position in byte
+    int byteEnd = 0; // position in byte
+    while (byteEnd < numBytes && c < end) {
+      byteEnd += numBytesForFirstByte(getByte(byteEnd));
+      c += 1;
+    }
+
+    // At least one match was found. Estimate space needed for result.
+    // The 16x multiplier here is chosen to match commons-lang3's implementation.
+    int increase = Math.max(0, Math.abs(replace.numBytes - search.numBytes)) * 16;
+    final UTF8StringBuilder buf = new UTF8StringBuilder(numBytes + increase);
+    while (end != StringSearch.DONE) {
+      buf.appendBytes(this.base, this.offset + byteStart, byteEnd - byteStart);
+      buf.append(replace);
+      byteStart = byteEnd + search.numBytes;
+      // Go to next match
+      end = stringSearch.next();
+      // Update byte positions
+      while (byteEnd < numBytes && c < end) {
+        byteEnd += numBytesForFirstByte(getByte(byteEnd));
+        c += 1;
+      }
+    }
+    buf.appendBytes(this.base, this.offset + byteStart, numBytes - byteStart);
+    return buf.build();
+  }
+
   public UTF8String translate(Map<String, String> dict) {
     String srcStr = this.toString();
 

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

@@ -735,23 +735,43 @@ case class EndsWith(left: Expression, right: Expression) extends StringPredicate
 case class StringReplace(srcExpr: Expression, searchExpr: Expression, replaceExpr: Expression)
   extends TernaryExpression with ImplicitCastInputTypes with NullIntolerant {
 
+  final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId
+
   def this(srcExpr: Expression, searchExpr: Expression) = {
     this(srcExpr, searchExpr, Literal(""))
   }
 
   override def nullSafeEval(srcEval: Any, searchEval: Any, replaceEval: Any): Any = {
     srcEval.asInstanceOf[UTF8String].replace(
-      searchEval.asInstanceOf[UTF8String], replaceEval.asInstanceOf[UTF8String])
+      searchEval.asInstanceOf[UTF8String], replaceEval.asInstanceOf[UTF8String], collationId)
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    nullSafeCodeGen(ctx, ev, (src, search, replace) => {
-      s"""${ev.value} = $src.replace($search, $replace);"""
-    })
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      nullSafeCodeGen(ctx, ev, (src, search, replace) => {
+        s"""${ev.value} = $src.replace($search, $replace);"""
+      })
+    } else {
+      nullSafeCodeGen(ctx, ev, (src, search, replace) => {
+        s"""${ev.value} = $src.replace($search, $replace, $collationId);"""
+      })
+    }
   }
 
-  override def dataType: DataType = StringType
-  override def inputTypes: Seq[DataType] = Seq(StringType, StringType, StringType)
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val defaultCheck = super.checkInputDataTypes()
+    if (defaultCheck.isFailure) {
+      return defaultCheck
+    }
+
+    // Only srcExpr and searchExpr are checked for collation compatibility.
+    val collationId = first.dataType.asInstanceOf[StringType].collationId
+    CollationTypeConstraints.checkCollationCompatibility(collationId, Seq(second.dataType))
+  }
+
+  override def dataType: DataType = srcExpr.dataType
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation)
   override def first: Expression = srcExpr
   override def second: Expression = searchExpr
   override def third: Expression = replaceExpr

sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala

@@ -21,7 +21,8 @@ import scala.collection.immutable.Seq
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.catalyst.ExtendedAnalysisException
-import org.apache.spark.sql.catalyst.expressions.{Collation, ExpressionEvalHelper, Literal, StringRepeat}
+import org.apache.spark.sql.catalyst.expressions.{Collation, ExpressionEvalHelper, Literal, StringRepeat, StringReplace}
+import org.apache.spark.sql.catalyst.util.CollationFactory
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types.StringType
@@ -73,6 +74,49 @@ class CollationStringExpressionsSuite extends QueryTest
     })
   }
 
+  test("REPLACE check result on explicitly collated strings") {
+    def testReplace(source: String, search: String, replace: String,
+                    collationId: Integer, expected: String): Unit = {
+      val sourceLiteral = Literal.create(source, StringType(collationId))
+      val searchLiteral = Literal.create(search, StringType(collationId))
+      val replaceLiteral = Literal.create(replace, StringType(collationId))
+
+      checkEvaluation(StringReplace(sourceLiteral, searchLiteral, replaceLiteral), expected)
+    }
+
+    // scalastyle:off
+    var collationId = CollationFactory.collationNameToId("UTF8_BINARY")
+    testReplace("r世eplace", "pl", "123", collationId, "r世e123ace")
+    testReplace("replace", "pl", "", collationId, "reace")
+    testReplace("repl世ace", "Pl", "", collationId, "repl世ace")
+    testReplace("replace", "", "123", collationId, "replace")
+    testReplace("abcabc", "b", "12", collationId, "a12ca12c")
+    testReplace("abcdabcd", "bc", "", collationId, "adad")
+
+    collationId = CollationFactory.collationNameToId("UTF8_BINARY_LCASE")
+    testReplace("r世eplace", "pl", "xx", collationId, "r世exxace")
+    testReplace("repl世ace", "PL", "AB", collationId, "reAB世ace")
+    testReplace("Replace", "", "123", collationId, "Replace")
+    testReplace("re世place", "世", "x", collationId, "rexplace")
+    testReplace("abcaBc", "B", "12", collationId, "a12ca12c")
+    testReplace("AbcdabCd", "Bc", "", collationId, "Adad")
+
+    collationId = CollationFactory.collationNameToId("UNICODE")
+    testReplace("re世place", "plx", "123", collationId, "re世place")
+    testReplace("世Replace", "re", "", collationId, "世Replace")
+    testReplace("replace世", "", "123", collationId, "replace世")
+    testReplace("aBc世abc", "b", "12", collationId, "aBc世a12c")
+    testReplace("abcdabcd", "bc", "", collationId, "adad")
+
+    collationId = CollationFactory.collationNameToId("UNICODE_CI")
+    testReplace("replace", "plx", "123", collationId, "replace")
+    testReplace("Replace", "re", "", collationId, "place")
+    testReplace("replace", "", "123", collationId, "replace")
+    testReplace("aBc世abc", "b", "12", collationId, "a12c世a12c")
+    testReplace("a世Bcdabcd", "bC", "", collationId, "a世dad")
+    // scalastyle:on
+  }
+
   test("REPEAT check output type on explicitly collated string") {
     def testRepeat(expected: String, collationId: Int, input: String, n: Int): Unit = {
       val s = Literal.create(input, StringType(collationId))