[SPARK-47566][SQL] Support SubstringIndex function to work with collated strings

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

@@ -835,6 +835,30 @@ public int indexOf(UTF8String v, int start) {
     return -1;
   }
 
+  private int charPosToByte(int charPos) {
+    if(charPos < 0) {
+      return -1;
+    }
+
+    int i = 0;
+    int c = 0;
+    while (i < numBytes && c < charPos) {
+      i += numBytesForFirstByte(getByte(i));
+      c += 1;
+    }
+    return i;
+  }
+
+  private int bytePosToChar(int bytePos) {
+    int i = 0;
+    int c = 0;
+    while (i < numBytes && i < bytePos) {
+      i += numBytesForFirstByte(getByte(i));
+      c += 1;
+    }
+    return c;
+  }
+
   /**
    * Find the `str` from left to right.
    */
@@ -849,6 +873,33 @@ private int find(UTF8String str, int start) {
     return -1;
   }
 
+  /**
+   * Find the `str` from left to right considering different collations.
+   */
+  private int find(UTF8String str, int start, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      return this.find(str, start);
+    }
+    if(collationId == CollationFactory.UTF8_BINARY_LCASE_COLLATION_ID) {
+      return charPosToByte(this.toLowerCase().indexOf(str.toLowerCase(), bytePosToChar(start)));
+    }
+    return collatedFind(str, start, collationId);
+  }
+
+  /**
+   * Find the `str` from left to right considering non-binary collations.
+   */
+  private int collatedFind(UTF8String str, int start, int collationId) {
+    assert (str.numBytes > 0);
+
+    StringSearch stringSearch = CollationFactory.getStringSearch(this, str, collationId);
+    // Set search start position (start from character at start position)
+    stringSearch.setIndex(bytePosToChar(start));
+
+    // Return either the byte position or -1 if not found
+    return charPosToByte(stringSearch.next());
+  }
+
   /**
    * Find the `str` from right to left.
    */
@@ -863,6 +914,70 @@ private int rfind(UTF8String str, int start) {
     return -1;
   }
 
+  /**
+   * Find the `str` from right to left considering different collations.
+   */
+  private int rfind(UTF8String str, int start, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      return this.rfind(str, start);
+    }
+    if(collationId == CollationFactory.UTF8_BINARY_LCASE_COLLATION_ID) {
+      return rfindLowercase(str, start);
+    }
+    return collatedRFind(str, start, collationId);
+  }
+
+  /**
+   * Find the `str` from left to right considering binary lowercase collation.
+   */
+  private int rfindLowercase(UTF8String str, int start) {
+    if(numBytes == 0 || str.numBytes == 0) {
+      return -1;
+    }
+
+    UTF8String lowercaseThis = this.toLowerCase();
+    UTF8String lowercaseStr = str.toLowerCase();
+
+    int prevStart = -1;
+    int matchStart = lowercaseThis.indexOf(lowercaseStr, 0);
+    while(charPosToByte(matchStart) <= start) {
+      if(matchStart != -1) {
+        // Found a match, update the start position
+        prevStart = matchStart;
+        matchStart = lowercaseThis.indexOf(lowercaseStr, matchStart + 1);
+      } else {
+        return charPosToByte(prevStart);
+      }
+    }
+
+    return charPosToByte(prevStart);
+  }
+
+  /**
+   * Find the `str` from left to right considering non-binary collations.
+   */
+  private int collatedRFind(UTF8String str, int start, int collationId) {
+    if(numBytes == 0 || str.numBytes == 0) {
+      return -1;
+    }
+
+    StringSearch stringSearch = CollationFactory.getStringSearch(this, str, collationId);
+
+    int prevStart = -1;
+    int matchStart = stringSearch.next();
+    while(charPosToByte(matchStart) <= start) {
+      if(matchStart != StringSearch.DONE) {
+        // Found a match, update the start position
+        prevStart = matchStart;
+        matchStart = stringSearch.next();
+      } else {
+        return charPosToByte(prevStart);
+      }
+    }
+
+    return charPosToByte(prevStart);
+  }
+
   /**
    * Returns the substring from string str before count occurrences of the delimiter delim.
    * If count is positive, everything the left of the final delimiter (counting from left) is
@@ -913,6 +1028,57 @@ public UTF8String subStringIndex(UTF8String delim, int count) {
     }
   }
 
+  public UTF8String subStringIndex(UTF8String delim, int count, int collationId) {
+    if (CollationFactory.fetchCollation(collationId).supportsBinaryEquality) {
+      return subStringIndex(delim, count);
+    }
+    return collatedSubStringIndex(delim, count, collationId);
+  }
+
+  private UTF8String collatedSubStringIndex(UTF8String delim, int count, int collationId) {
+    if (delim.numBytes == 0 || count == 0) {
+      return EMPTY_UTF8;
+    }
+    if (count > 0) {
+      int idx = -1;
+      while (count > 0) {
+        idx = find(delim, idx + 1, collationId);
+        if (idx >= 0) {
+          count --;
+        } else {
+          // can not find enough delim
+          return this;
+        }
+      }
+      if (idx == 0) {
+        return EMPTY_UTF8;
+      }
+      byte[] bytes = new byte[idx];
+      copyMemory(base, offset, bytes, BYTE_ARRAY_OFFSET, idx);
+      return fromBytes(bytes);
+
+    } else {
+      int idx = numBytes - delim.numBytes + 1;
+      count = -count;
+      while (count > 0) {
+        idx = rfind(delim, idx - 1, collationId);
+        if (idx >= 0) {
+          count --;
+        } else {
+          // can not find enough delim
+          return this;
+        }
+      }
+      if (idx + delim.numBytes == numBytes) {
+        return EMPTY_UTF8;
+      }
+      int size = numBytes - delim.numBytes - idx;
+      byte[] bytes = new byte[size];
+      copyMemory(base, offset + idx + delim.numBytes, bytes, BYTE_ARRAY_OFFSET, size);
+      return fromBytes(bytes);
+    }
+  }
+
   /**
    * Returns str, right-padded with pad to a length of len
    * For example:

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

@@ -1420,21 +1420,31 @@ case class StringInstr(str: Expression, substr: Expression)
 case class SubstringIndex(strExpr: Expression, delimExpr: Expression, countExpr: Expression)
  extends TernaryExpression with ImplicitCastInputTypes with NullIntolerant {
 
-  override def dataType: DataType = StringType
-  override def inputTypes: Seq[DataType] = Seq(StringType, StringType, IntegerType)
+  override def dataType: DataType = strExpr.dataType
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeAnyCollation, StringTypeAnyCollation, IntegerType)
   override def first: Expression = strExpr
   override def second: Expression = delimExpr
   override def third: Expression = countExpr
   override def prettyName: String = "substring_index"
 
   override def nullSafeEval(str: Any, delim: Any, count: Any): Any = {
+    val collationId = first.dataType.asInstanceOf[StringType].collationId
+
     str.asInstanceOf[UTF8String].subStringIndex(
       delim.asInstanceOf[UTF8String],
-      count.asInstanceOf[Int])
+      count.asInstanceOf[Int], collationId)
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    defineCodeGen(ctx, ev, (str, delim, count) => s"$str.subStringIndex($delim, $count)")
+    val collationId = first.dataType.asInstanceOf[StringType].collationId
+
+    if(CollationFactory.fetchCollation(collationId).supportsBinaryOrdering) {
+      defineCodeGen(ctx, ev, (str, delim, count) => s"$str.subStringIndex($delim, $count)")
+    } else {
+      defineCodeGen(ctx, ev, (str, delim, count) =>
+        s"$str.subStringIndex($delim, $count, $collationId)")
+    }
   }
 
   override protected def withNewChildrenInternal(

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

@@ -21,7 +21,7 @@ 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, SubstringIndex}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types.StringType
@@ -73,6 +73,84 @@ class CollationStringExpressionsSuite extends QueryTest
     })
   }
 
+  test("SUBSTRING_INDEX check result on explicitly collated strings") {
+    def testSubstringIndex(str: String, delim: String, cnt: Integer,
+                           collationId: Integer, expected: String): Unit = {
+      val string = Literal.create(str, StringType(collationId))
+      val delimiter = Literal.create(delim, StringType(collationId))
+      val count = Literal(cnt)
+
+      checkEvaluation(SubstringIndex(string, delimiter, count), expected)
+    }
+
+    testSubstringIndex("wwwgapachegorg", "g", -3, 0, "apachegorg")
+    testSubstringIndex("www||apache||org", "||", 2, 0, "www||apache")
+    // UTF8_BINARY_LCASE
+    testSubstringIndex("AaAaAaAaAa", "aa", 2, 1, "A")
+    testSubstringIndex("www.apache.org", ".", 3, 1, "www.apache.org")
+    testSubstringIndex("wwwXapachexorg", "x", 2, 1, "wwwXapache")
+    testSubstringIndex("wwwxapacheXorg", "X", 1, 1, "www")
+    testSubstringIndex("www.apache.org", ".", 0, 1, "")
+    testSubstringIndex("www.apache.ORG", ".", -3, 1, "www.apache.ORG")
+    testSubstringIndex("wwwGapacheGorg", "g", 1, 1, "www")
+    testSubstringIndex("wwwGapacheGorg", "g", 3, 1, "wwwGapacheGor")
+    testSubstringIndex("gwwwGapacheGorg", "g", 3, 1, "gwwwGapache")
+    testSubstringIndex("wwwGapacheGorg", "g", -3, 1, "apacheGorg")
+    testSubstringIndex("wwwmapacheMorg", "M", -2, 1, "apacheMorg")
+    testSubstringIndex("www.apache.org", ".", -1, 1, "org")
+    testSubstringIndex("", ".", -2, 1, "")
+    // scalastyle:off
+    testSubstringIndex("test大千世界X大千世界", "x", -1, 1, "大千世界")
+    testSubstringIndex("test大千世界X大千世界", "X", 1, 1, "test大千世界")
+    testSubstringIndex("test大千世界大千世界", "千", 2, 1, "test大千世界大")
+    // scalastyle:on
+    testSubstringIndex("www||APACHE||org", "||", 2, 1, "www||APACHE")
+    testSubstringIndex("www||APACHE||org", "||", -1, 1, "org")
+    // UNICODE
+    testSubstringIndex("AaAaAaAaAa", "Aa", 2, 2, "Aa")
+    testSubstringIndex("wwwYapacheyorg", "y", 3, 2, "wwwYapacheyorg")
+    testSubstringIndex("www.apache.org", ".", 2, 2, "www.apache")
+    testSubstringIndex("wwwYapacheYorg", "Y", 1, 2, "www")
+    testSubstringIndex("wwwYapacheYorg", "y", 1, 2, "wwwYapacheYorg")
+    testSubstringIndex("wwwGapacheGorg", "g", 1, 2, "wwwGapacheGor")
+    testSubstringIndex("GwwwGapacheGorG", "G", 3, 2, "GwwwGapache")
+    testSubstringIndex("wwwGapacheGorG", "G", -3, 2, "apacheGorG")
+    testSubstringIndex("www.apache.org", ".", 0, 2, "")
+    testSubstringIndex("www.apache.org", ".", -3, 2, "www.apache.org")
+    testSubstringIndex("www.apache.org", ".", -2, 2, "apache.org")
+    testSubstringIndex("www.apache.org", ".", -1, 2, "org")
+    testSubstringIndex("", ".", -2, 2, "")
+    // scalastyle:off
+    testSubstringIndex("test大千世界X大千世界", "X", -1, 2, "大千世界")
+    testSubstringIndex("test大千世界X大千世界", "X", 1, 2, "test大千世界")
+    testSubstringIndex("大x千世界大千世x界", "x", 1, 2, "大")
+    testSubstringIndex("大x千世界大千世x界", "x", -1, 2, "界")
+    testSubstringIndex("大x千世界大千世x界", "x", -2, 2, "千世界大千世x界")
+    testSubstringIndex("大千世界大千世界", "千", 2, 2, "大千世界大")
+    // scalastyle:on
+    testSubstringIndex("www||apache||org", "||", 2, 2, "www||apache")
+    // UNICODE_CI
+    testSubstringIndex("AaAaAaAaAa", "aa", 2, 3, "A")
+    testSubstringIndex("www.apache.org", ".", 3, 3, "www.apache.org")
+    testSubstringIndex("wwwXapachexorg", "x", 2, 3, "wwwXapache")
+    testSubstringIndex("wwwxapacheXorg", "X", 1, 3, "www")
+    testSubstringIndex("www.apache.org", ".", 0, 3, "")
+    testSubstringIndex("wwwGapacheGorg", "g", 1, 3, "www")
+    testSubstringIndex("wwwGapacheGorg", "g", 3, 3, "wwwGapacheGor")
+    testSubstringIndex("gwwwGapacheGorg", "g", 3, 3, "gwwwGapache")
+    testSubstringIndex("wwwGapacheGorg", "g", -3, 3, "apacheGorg")
+    testSubstringIndex("www.apache.ORG", ".", -3, 3, "www.apache.ORG")
+    testSubstringIndex("wwwmapacheMorg", "M", -2, 3, "apacheMorg")
+    testSubstringIndex("www.apache.org", ".", -1, 3, "org")
+    testSubstringIndex("", ".", -2, 3, "")
+    // scalastyle:off
+    testSubstringIndex("test大千世界X大千世界", "X", -1, 3, "大千世界")
+    testSubstringIndex("test大千世界X大千世界", "X", 1, 3, "test大千世界")
+    testSubstringIndex("test大千世界大千世界", "千", 2, 3, "test大千世界大")
+    // scalastyle:on
+    testSubstringIndex("www||APACHE||org", "||", 2, 3, "www||APACHE")
+  }
+
   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))