add substring_index() function in sparksql

velox/docs/functions/spark/string.rst

@@ -179,6 +179,14 @@ Unless specified otherwise, all functions return NULL if at least one of the arg
         SELECT substring('Spark SQL', -10, 3); -- "Sp"
         SELECT substring('Spark SQL', -20, 3); -- ""
 
+.. spark:function:: substring_index(string, delim, count) -> varchar
+
+    Returns the substring from ``string`` before ``count`` occurrences of the delimiter ``delim``.
+    If ``count`` is positive, everything to the left of the final delimiter
+    (counting from the left) is returned. If ``count`` is negative, everything to the right
+    of the final delimiter (counting from the right) is returned. The function
+    substring_index performs a case-sensitive match when searching for ``delim``.
+
 .. spark:function:: translate(string, match, replace) -> varchar
 
     Returns a new translated string. It translates the character in ``string`` by a

velox/expression/tests/SparkExpressionFuzzerTest.cpp

@@ -63,6 +63,7 @@ int main(int argc, char** argv) {
       "replace",
       "might_contain",
       "unix_timestamp"};
+
   return FuzzerRunner::run(
       FLAGS_only, FLAGS_seed, skipFunctions, FLAGS_special_forms);
 }

velox/functions/lib/string/StringCore.h

@@ -261,7 +261,7 @@ static int64_t findNthInstanceByteIndexFromStart(
 
   // Find next occurrence
   return findNthInstanceByteIndexFromStart(
-      string, subString, instance - 1, byteIndex + subString.size());
+      string, subString, instance - 1, byteIndex + 1);
 }
 
 /// Returns the start byte index of the Nth instance of subString in

velox/functions/sparksql/Register.cpp

@@ -119,6 +119,8 @@ void registerFunctions(const std::string& prefix) {
       prefix + "instr", instrSignatures(), makeInstr);
   exec::registerStatefulVectorFunction(
       prefix + "length", lengthSignatures(), makeLength);
+  registerFunction<SubstringIndexFunction, Varchar, Varchar, Varchar, int32_t>(
+      {prefix + "substring_index"});
 
   registerFunction<Md5Function, Varchar, Varbinary>({prefix + "md5"});
   registerFunction<Sha1HexStringFunction, Varchar, Varbinary>(

velox/functions/sparksql/String.h

@@ -283,6 +283,80 @@ struct EndsWithFunction {
   }
 };
 
+/// substring_index function
+/// substring_index(string, string, int) -> string
+/// substring_index(str, delim, count) - Returns the substring from str before
+/// count occurrences of the delimiter delim. If count is positive, everything
+/// to the left of the final delimiter (counting from the left) is returned. If
+/// count is negative, everything to the right of the final delimiter (counting
+/// from the right) is returned. The function substring_index performs a
+/// case-sensitive match when searching for delim.
+template <typename T>
+struct SubstringIndexFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  // Results refer to strings in the first argument.
+  static constexpr int32_t reuse_strings_from_arg = 0;
+
+  // ASCII input always produces ASCII result.
+  static constexpr bool is_default_ascii_behavior = true;
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& str,
+      const arg_type<Varchar>& delim,
+      const int32_t& count) {
+    doCall<false>(result, str, delim, count);
+  }
+
+  FOLLY_ALWAYS_INLINE void callAscii(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& str,
+      const arg_type<Varchar>& delim,
+      const int32_t& count) {
+    doCall<true>(result, str, delim, count);
+  }
+
+  template <bool isAscii>
+  FOLLY_ALWAYS_INLINE void doCall(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& str,
+      const arg_type<Varchar>& delim,
+      const int32_t& count) {
+    if (count == 0) {
+      result.setEmpty();
+      return;
+    }
+
+    int64_t index;
+    if (count > 0) {
+      index = stringImpl::stringPosition<isAscii, true>(str, delim, count);
+    } else {
+      index = stringImpl::stringPosition<isAscii, false>(str, delim, -count);
+    }
+
+    auto start = 1;
+    auto length = stringImpl::length<isAscii>(str);
+    auto delimLength = stringImpl::length<isAscii>(delim);
+
+    if (index != 0) {
+      if (count > 0) {
+        length = index - 1;
+      } else {
+        start = index + delimLength;
+        length = length - index - delimLength + 1;
+      }
+    }
+
+    auto byteRange =
+        stringCore::getByteRange<isAscii>(str.data(), start, length);
+
+    // Generating output string
+    result.setNoCopy(StringView(
+        str.data() + byteRange.first, byteRange.second - byteRange.first));
+  }
+};
+
 /// ltrim(trimStr, srcStr) -> varchar
 ///     Remove leading specified characters from srcStr. The specified character
 ///     is any character contained in trimStr.

velox/functions/sparksql/tests/StringTest.cpp

@@ -140,6 +140,14 @@ class StringTest : public SparkFunctionBaseTest {
     return evaluateOnce<std::string>("left(c0, c1)", str, length);
   }
 
+  std::optional<std::string> substring_index(
+      const std::optional<std::string>& str,
+      const std::optional<std::string>& delim,
+      int32_t count) {
+    return evaluateOnce<std::string, std::string, std::string, int32_t>(
+        "substring_index(c0, c1, c2)", str, delim, count);
+  }
+
   std::optional<std::string> overlay(
       std::optional<std::string> input,
       std::optional<std::string> replace,
@@ -370,6 +378,30 @@ TEST_F(StringTest, endsWith) {
   EXPECT_EQ(endsWith(std::nullopt, "abc"), std::nullopt);
 }
 
+TEST_F(StringTest, substring_index) {
+  EXPECT_EQ(substring_index("www.apache.org", ".", 3), "www.apache.org");
+  EXPECT_EQ(substring_index("www.apache.org", ".", 2), "www.apache");
+  EXPECT_EQ(substring_index("www.apache.org", ".", 1), "www");
+  EXPECT_EQ(substring_index("www.apache.org", ".", 0), "");
+  EXPECT_EQ(substring_index("www.apache.org", ".", -1), "org");
+  EXPECT_EQ(substring_index("www.apache.org", ".", -2), "apache.org");
+  EXPECT_EQ(substring_index("www.apache.org", ".", -3), "www.apache.org");
+  // Str is empty string.
+  EXPECT_EQ(substring_index("", ".", 1), "");
+  // Empty string delim.
+  EXPECT_EQ(substring_index("www.apache.org", "", 1), "");
+  // Delim does not exist in str.
+  EXPECT_EQ(substring_index("www.apache.org", "#", 2), "www.apache.org");
+  // Delim is 2 chars.
+  EXPECT_EQ(substring_index("www||apache||org", "||", 2), "www||apache");
+  EXPECT_EQ(substring_index("www||apache||org", "||", -2), "apache||org");
+  // Non ascii chars.
+  EXPECT_EQ(substring_index("大千世界大千世界", "千", 2), "大千世界大");
+  // Overlapped delim.
+  EXPECT_EQ(substring_index("||||||", "|||", 3), "||");
+  EXPECT_EQ(substring_index("||||||", "|||", -4), "|||");
+}
+
 TEST_F(StringTest, trim) {
   EXPECT_EQ(trim(""), "");
   EXPECT_EQ(trim("  data\t "), "data\t");