[SPARK-47359][SQL] Support TRANSLATE function to work with collated strings

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

@@ -179,12 +179,26 @@ public static StringSearch getStringSearch(
       final UTF8String left,
       final UTF8String right,
       final int collationId) {
+
+    if(collationId == UTF8_BINARY_LCASE_COLLATION_ID) {
+      return getStringSearch(left, right);
+    }
+
     String pattern = right.toString();
     CharacterIterator target = new StringCharacterIterator(left.toString());
     Collator collator = CollationFactory.fetchCollation(collationId).collator;
     return new StringSearch(pattern, target, (RuleBasedCollator) collator);
   }
 
+  private static StringSearch getStringSearch(
+          final UTF8String left,
+          final UTF8String right) {
+    String pattern = right.toLowerCase().toString();
+    String target = left.toLowerCase().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

@@ -22,6 +22,7 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.regex.Pattern;
 
@@ -1155,6 +1156,45 @@ public UTF8String translate(Map<String, String> dict) {
     return fromString(sb.toString());
   }
 
+  public UTF8String translate(Map<String, String> dict, int collationId) {
+    if(CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      return translate(dict);
+    }
+    return translate(getCollationAwareDict(dict, collationId));
+  }
+
+  private Map<String, String> getCollationAwareDict(Map<String, String> dict, int collationId) {
+    String srcStr = this.toString();
+
+    Map<String, String> collationAwareDict = new HashMap<>();
+    for(String key : dict.keySet()) {
+      StringSearch stringSearch =
+        CollationFactory.getStringSearch(this, UTF8String.fromString(key), collationId);
+
+      int pos = 0;
+      while((pos = stringSearch.next()) != StringSearch.DONE) {
+        int codePoint = srcStr.codePointAt(pos);
+        int charCount = Character.charCount(codePoint);
+        String newKey = srcStr.substring(pos, pos + charCount);
+
+        boolean exists = false;
+        for(String existingKey : collationAwareDict.keySet()) {
+          if(stringSearch.getCollator().compare(existingKey, newKey) == 0) {
+            collationAwareDict.put(newKey, collationAwareDict.get(existingKey));
+            exists = true;
+            break;
+          }
+        }
+
+        if(!exists) {
+          collationAwareDict.put(newKey, dict.get(key));
+        }
+      }
+    }
+
+    return collationAwareDict;
+  }
+
   /**
    * Wrapper over `long` to allow result of parsing long from string to be accessed via reference.
    * This is done solely for better performance and is not expected to be used by end users.

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

@@ -934,13 +934,20 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac
   @transient private var lastReplace: UTF8String = _
   @transient private var dict: JMap[String, String] = _
 
+  final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId
+
   override def nullSafeEval(srcEval: Any, matchingEval: Any, replaceEval: Any): Any = {
     if (matchingEval != lastMatching || replaceEval != lastReplace) {
       lastMatching = matchingEval.asInstanceOf[UTF8String].clone()
       lastReplace = replaceEval.asInstanceOf[UTF8String].clone()
       dict = StringTranslate.buildDict(lastMatching, lastReplace)
     }
-    srcEval.asInstanceOf[UTF8String].translate(dict)
+
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      srcEval.asInstanceOf[UTF8String].translate(dict)
+    } else {
+      srcEval.asInstanceOf[UTF8String].translate(dict, collationId)
+    }
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
@@ -963,13 +970,18 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac
         $termDict = org.apache.spark.sql.catalyst.expressions.StringTranslate
           .buildDict($termLastMatching, $termLastReplace);
       }
-      ${ev.value} = $src.translate($termDict);
+      if (CollationFactory.fetchCollation(${collationId}).supportsBinaryEquality) {
+        ${ev.value} = $src.translate($termDict);
+      } else {
+        ${ev.value} = $src.translate($termDict, ${collationId});
+      }
       """
     })
   }
 
-  override def dataType: DataType = StringType
-  override def inputTypes: Seq[DataType] = Seq(StringType, StringType, StringType)
+  override def dataType: DataType = srcExpr.dataType
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation)
   override def first: Expression = srcExpr
   override def second: Expression = matchingExpr
   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, StringTranslate}
+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,57 @@ class CollationStringExpressionsSuite extends QueryTest
     })
   }
 
+  test("TRANSLATE check result on explicitly collated string") {
+    def testTranslate(input: String,
+                      matchExpression: String,
+                      replaceExpression: String,
+                      collationId: Int,
+                      expected: String): Unit = {
+      val srcExpr = Literal.create(input, StringType(collationId))
+      val matchExpr = Literal.create(matchExpression, StringType(collationId))
+      val replaceExpr = Literal.create(replaceExpression, StringType(collationId))
+
+      checkEvaluation(StringTranslate(srcExpr, matchExpr, replaceExpr), expected)
+    }
+
+    var collationId = CollationFactory.collationNameToId("UTF8_BINARY_LCASE")
+    testTranslate("Translate", "Rnlt", "1234", collationId, "41a2s3a4e")
+    testTranslate("TRanslate", "rnlt", "XxXx", collationId, "xXaxsXaxe")
+    testTranslate("TRanslater", "Rrnlt", "xXxXx", collationId, "xxaxsXaxex")
+    testTranslate("TRanslater", "Rrnlt", "XxxXx", collationId, "xXaxsXaxeX")
+    // scalastyle:off
+    testTranslate("test大千世界X大千世界", "界x", "AB", collationId, "test大千世AB大千世A")
+    testTranslate("大千世界test大千世界", "TEST", "abcd", collationId, "大千世界abca大千世界")
+    testTranslate("Test大千世界大千世界", "tT", "oO", collationId, "oeso大千世界大千世界")
+    testTranslate("大千世界大千世界tesT", "Tt", "Oo", collationId, "大千世界大千世界OesO")
+    testTranslate("大千世界大千世界tesT", "大千", "世世", collationId, "世世世界世世世界tesT")
+    // scalastyle:on
+
+    collationId = CollationFactory.collationNameToId("UNICODE")
+    testTranslate("Translate", "Rnlt", "1234", collationId, "Tra2s3a4e")
+    testTranslate("TRanslate", "rnlt", "XxXx", collationId, "TRaxsXaxe")
+    testTranslate("TRanslater", "Rrnlt", "xXxXx", collationId, "TxaxsXaxeX")
+    testTranslate("TRanslater", "Rrnlt", "XxxXx", collationId, "TXaxsXaxex")
+    // scalastyle:off
+    testTranslate("test大千世界X大千世界", "界x", "AB", collationId, "test大千世AX大千世A")
+    testTranslate("Test大千世界大千世界", "tT", "oO", collationId, "Oeso大千世界大千世界")
+    testTranslate("大千世界大千世界tesT", "Tt", "Oo", collationId, "大千世界大千世界oesO")
+    // scalastyle:on
+
+    collationId = CollationFactory.collationNameToId("UNICODE_CI")
+    testTranslate("Translate", "Rnlt", "1234", collationId, "41a2s3a4e")
+    testTranslate("TRanslate", "rnlt", "XxXx", collationId, "xXaxsXaxe")
+    testTranslate("TRanslater", "Rrnlt", "xXxXx", collationId, "xxaxsXaxex")
+    testTranslate("TRanslater", "Rrnlt", "XxxXx", collationId, "xXaxsXaxeX")
+    // scalastyle:off
+    testTranslate("test大千世界X大千世界", "界x", "AB", collationId, "test大千世AB大千世A")
+    testTranslate("大千世界test大千世界", "TEST", "abcd", collationId, "大千世界abca大千世界")
+    testTranslate("Test大千世界大千世界", "tT", "oO", collationId, "oeso大千世界大千世界")
+    testTranslate("大千世界大千世界tesT", "Tt", "Oo", collationId, "大千世界大千世界OesO")
+    testTranslate("大千世界大千世界tesT", "大千", "世世", collationId, "世世世界世世世界tesT")
+    // 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))