[SPARK-47411][SQL] Support StringInstr & FindInSet functions to work with collated strings

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

@@ -549,6 +549,47 @@ public int findInSet(UTF8String match) {
     return 0;
   }
 
+  public int findInSet(UTF8String match, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).isBinaryCollation) {
+      return this.findInSet(match);
+    }
+    if (collationId == CollationFactory.LOWERCASE_COLLATION_ID) {
+      return this.toLowerCase().findInSet(match.toLowerCase());
+    }
+    return collatedFindInSet(match, collationId);
+}
+
+  private int collatedFindInSet(UTF8String match, int collationId) {
+    if (match.contains(COMMA_UTF8)) {
+      return 0;
+    }
+
+    StringSearch stringSearch = CollationFactory.getStringSearch(this, match, collationId);
+
+    String setString = this.toString();
+    int wordStart = 0;
+    while ((wordStart = stringSearch.next()) != StringSearch.DONE) {
+      if (stringSearch.getMatchLength() == stringSearch.getPattern().length()) {
+        boolean isValidStart = wordStart == 0 || setString.charAt(wordStart - 1) == ',';
+        boolean isValidEnd = wordStart + stringSearch.getMatchLength() == setString.length()
+                || setString.charAt(wordStart + stringSearch.getMatchLength()) == ',';
+
+        if(isValidStart && isValidEnd) {
+          int pos = 0;
+          for(int i = 0; i < setString.length() && i < wordStart; i++) {
+            if(setString.charAt(i) == ',') {
+              pos++;
+            }
+          }
+
+          return pos + 1;
+        }
+      }
+    }
+
+    return 0;
+  }
+
   /**
    * Copy the bytes from the current UTF8String, and make a new UTF8String.
    * @param start the start position of the current UTF8String in bytes.
@@ -835,6 +876,33 @@ public int indexOf(UTF8String v, int start) {
     return -1;
   }
 
+  public int indexOf(UTF8String substring, int start, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).isBinaryCollation) {
+      return this.indexOf(substring, start);
+    }
+    if (collationId == CollationFactory.LOWERCASE_COLLATION_ID) {
+      return this.toLowerCase().indexOf(substring.toLowerCase(), start);
+    }
+    return collatedIndexOf(substring, collationId);
+  }
+
+  private int collatedIndexOf(UTF8String substring, int collationId) {
+    if (substring.numBytes == 0) {
+      return 0;
+    }
+
+    StringSearch stringSearch = CollationFactory.getStringSearch(this, substring, collationId);
+
+    int pos = 0;
+    while ((pos = stringSearch.next()) != StringSearch.DONE) {
+      if (stringSearch.getMatchLength() == stringSearch.getPattern().length()) {
+        return pos;
+      }
+    }
+
+    return 0;
+  }
+
   /**
    * Find the `str` from left to right.
    */

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

@@ -1002,15 +1002,32 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac
 case class FindInSet(left: Expression, right: Expression) extends BinaryExpression
     with ImplicitCastInputTypes with NullIntolerant {
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(StringType, StringType)
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeAnyCollation, StringTypeAnyCollation)
 
-  override protected def nullSafeEval(word: Any, set: Any): Any =
-    set.asInstanceOf[UTF8String].findInSet(word.asInstanceOf[UTF8String])
+  override protected def nullSafeEval(word: Any, set: Any): Any = {
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+    set.asInstanceOf[UTF8String].findInSet(word.asInstanceOf[UTF8String], collationId)
+  }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    nullSafeCodeGen(ctx, ev, (word, set) =>
-      s"${ev.value} = $set.findInSet($word);"
-    )
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+
+    if (CollationFactory.fetchCollation(collationId).isBinaryCollation) {
+      nullSafeCodeGen(ctx, ev, (word, set) => s"${ev.value} = $set.findInSet($word);")
+    } else {
+      nullSafeCodeGen(ctx, ev, (word, set) => s"${ev.value} = $set.findInSet($word, $collationId);")
+    }
+  }
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val defaultCheck = super.checkInputDataTypes()
+    if (defaultCheck.isFailure) {
+      return defaultCheck
+    }
+
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+    CollationTypeConstraints.checkCollationCompatibility(collationId, children.map(_.dataType))
   }
 
   override def dataType: DataType = IntegerType
@@ -1377,17 +1394,34 @@ case class StringInstr(str: Expression, substr: Expression)
   override def left: Expression = str
   override def right: Expression = substr
   override def dataType: DataType = IntegerType
-  override def inputTypes: Seq[DataType] = Seq(StringType, StringType)
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeAnyCollation, StringTypeAnyCollation)
 
   override def nullSafeEval(string: Any, sub: Any): Any = {
-    string.asInstanceOf[UTF8String].indexOf(sub.asInstanceOf[UTF8String], 0) + 1
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+    string.asInstanceOf[UTF8String].indexOf(sub.asInstanceOf[UTF8String], 0, collationId) + 1
+  }
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val defaultCheck = super.checkInputDataTypes()
+    if (defaultCheck.isFailure) {
+      return defaultCheck
+    }
+
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+    CollationTypeConstraints.checkCollationCompatibility(collationId, children.map(_.dataType))
   }
 
   override def prettyName: String = "instr"
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    defineCodeGen(ctx, ev, (l, r) =>
-      s"($l).indexOf($r, 0) + 1")
+    val collationId = left.dataType.asInstanceOf[StringType].collationId
+
+    if (CollationFactory.fetchCollation(collationId).isBinaryCollation) {
+      defineCodeGen(ctx, ev, (l, r) => s"($l).indexOf($r, 0) + 1")
+    } else {
+      defineCodeGen(ctx, ev, (l, r) => s"($l).indexOf($r, 0, $collationId) + 1")
+    }
   }
 
   override protected def withNewChildrenInternal(

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

@@ -21,10 +21,14 @@ import scala.collection.immutable.Seq
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.catalyst.ExtendedAnalysisException
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
+import org.apache.spark.sql.catalyst.expressions.{Collate, ExpressionEvalHelper, FindInSet, Literal, StringInstr}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StringType
 
-class CollationStringExpressionsSuite extends QueryTest with SharedSparkSession {
+class CollationStringExpressionsSuite extends QueryTest
+  with SharedSparkSession with ExpressionEvalHelper {
 
   case class CollationTestCase[R](s1: String, s2: String, collation: String, expectedResult: R)
   case class CollationTestFail[R](s1: String, s2: String, collation: String)
@@ -70,6 +74,148 @@ class CollationStringExpressionsSuite extends QueryTest with SharedSparkSession
     })
   }
 
+  test("INSTR check result on non-explicit default collation") {
+    checkEvaluation(StringInstr(Literal("aAads"), Literal("Aa")), 2)
+  }
+
+  test("INSTR check result on explicitly collated strings") {
+    // UTF8_BINARY_LCASE
+    checkEvaluation(StringInstr(Literal.create("aaads", StringType(1)),
+      Literal.create("Aa", StringType(1))), 1)
+    checkEvaluation(StringInstr(Collate(Literal("aaads"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("Aa"), "UTF8_BINARY_LCASE")), 1)
+    // UNICODE
+    checkEvaluation(StringInstr(Literal.create("aaads", StringType(2)),
+      Literal.create("Aa", StringType(2))), 0)
+    checkEvaluation(StringInstr(Collate(Literal("aaads"), "UNICODE"),
+      Collate(Literal("Aa"), "UNICODE")), 0)
+    // UNICODE_CI
+    checkEvaluation(StringInstr(Literal.create("aaads", StringType(3)),
+      Literal.create("de", StringType(3))), 0)
+    checkEvaluation(StringInstr(Collate(Literal("aaads"), "UNICODE_CI"),
+      Collate(Literal("Aa"), "UNICODE_CI")), 0)
+  }
+
+  test("INSTR fail mismatched collation types") {
+    // UNICODE and UNICODE_CI
+    val expr1 = StringInstr(Collate(Literal("aaads"), "UNICODE"),
+      Collate(Literal("Aa"), "UNICODE_CI"))
+    assert(expr1.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UNICODE",
+          "collationNameRight" -> "UNICODE_CI"
+        )
+      )
+    )
+    // DEFAULT(UTF8_BINARY) and UTF8_BINARY_LCASE
+    val expr2 = StringInstr(Literal("aaads"),
+      Collate(Literal("Aa"), "UTF8_BINARY_LCASE"))
+    assert(expr2.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UTF8_BINARY",
+          "collationNameRight" -> "UTF8_BINARY_LCASE"
+        )
+      )
+    )
+    // UTF8_BINARY_LCASE and UNICODE_CI
+    val expr3 = StringInstr(Collate(Literal("aaads"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("Aa"), "UNICODE_CI"))
+    assert(expr3.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UTF8_BINARY_LCASE",
+          "collationNameRight" -> "UNICODE_CI"
+        )
+      )
+    )
+  }
+
+  test("FIND_IN_SET check result on non-explicit default collation") {
+    checkEvaluation(FindInSet(Literal("def"), Literal("abc,b,ab,c,def")), 5)
+    checkEvaluation(FindInSet(Literal("defg"), Literal("abc,b,ab,c,def")), 0)
+  }
+
+  test("FIND_IN_SET check result on explicitly collated strings") {
+    // UTF8_BINARY
+    checkEvaluation(FindInSet(Collate(Literal("a"), "UTF8_BINARY"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY")), 0)
+    checkEvaluation(FindInSet(Collate(Literal("c"), "UTF8_BINARY"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY")), 4)
+    checkEvaluation(FindInSet(Collate(Literal("AB"), "UTF8_BINARY"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY")), 0)
+    checkEvaluation(FindInSet(Collate(Literal("abcd"), "UTF8_BINARY"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY")), 0)
+    // UTF8_BINARY_LCASE
+    checkEvaluation(FindInSet(Collate(Literal("aB"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY_LCASE")), 3)
+    checkEvaluation(FindInSet(Collate(Literal("a"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY_LCASE")), 0)
+    checkEvaluation(FindInSet(Collate(Literal("abc"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("aBc,b,ab,c,def"), "UTF8_BINARY_LCASE")), 1)
+    checkEvaluation(FindInSet(Collate(Literal("abcd"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("aBc,b,ab,c,def"), "UTF8_BINARY_LCASE")), 0)
+    // UNICODE
+    checkEvaluation(FindInSet(Collate(Literal("a"), "UNICODE"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE")), 0)
+    checkEvaluation(FindInSet(Collate(Literal("ab"), "UNICODE"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE")), 3)
+    checkEvaluation(FindInSet(Collate(Literal("Ab"), "UNICODE"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE")), 0)
+    // UNICODE_CI
+    checkEvaluation(FindInSet(Collate(Literal("a"), "UNICODE_CI"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE_CI")), 0)
+    checkEvaluation(FindInSet(Collate(Literal("C"), "UNICODE_CI"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE_CI")), 4)
+    checkEvaluation(FindInSet(Collate(Literal("DeF"), "UNICODE_CI"),
+      Collate(Literal("abc,b,ab,c,dEf"), "UNICODE_CI")), 5)
+    checkEvaluation(FindInSet(Collate(Literal("DEFG"), "UNICODE_CI"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE_CI")), 0)
+  }
+
+  test("FIND_IN_SET fail mismatched collation types") {
+    // UNICODE and UNICODE_CI
+    val expr1 = FindInSet(Collate(Literal("a"), "UNICODE"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE_CI"))
+    assert(expr1.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UNICODE",
+          "collationNameRight" -> "UNICODE_CI"
+        )
+      )
+    )
+    // DEFAULT(UTF8_BINARY) and UTF8_BINARY_LCASE
+    val expr2 = FindInSet(Collate(Literal("a"), "UTF8_BINARY"),
+      Collate(Literal("abc,b,ab,c,def"), "UTF8_BINARY_LCASE"))
+    assert(expr2.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UTF8_BINARY",
+          "collationNameRight" -> "UTF8_BINARY_LCASE"
+        )
+      )
+    )
+    // UTF8_BINARY_LCASE and UNICODE_CI
+    val expr3 = FindInSet(Collate(Literal("a"), "UTF8_BINARY_LCASE"),
+      Collate(Literal("abc,b,ab,c,def"), "UNICODE_CI"))
+    assert(expr3.checkInputDataTypes() ==
+      DataTypeMismatch(
+        errorSubClass = "COLLATION_MISMATCH",
+        messageParameters = Map(
+          "collationNameLeft" -> "UTF8_BINARY_LCASE",
+          "collationNameRight" -> "UNICODE_CI"
+        )
+      )
+    )
+  }
+
   // TODO: Add more tests for other string expressions
 
 }