Support regexp_substr function

dbms/src/Common/OptimizedRegularExpression.h

@@ -15,10 +15,12 @@
 #pragma once
 
 #include <Common/config.h>
+#include <common/StringRef.h>
 #include <common/types.h>
 #include <re2/re2.h>
 
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 #if USE_RE2_ST
@@ -114,11 +116,14 @@ class OptimizedRegularExpressionImpl
     }
 
     Int64 instr(const char * subject, size_t subject_size, Int64 pos, Int64 occur, Int64 ret_op);
+    std::optional<StringRef> substr(const char * subject, size_t subject_size, Int64 pos, Int64 occur);
 
 private:
-    Int64 processEmptyStringExpr(const char * expr, size_t expr_size, size_t byte_pos, Int64 occur);
-    Int64 getSubstrMatchedIndex(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur, Int64 ret_op);
+    Int64 processInstrEmptyStringExpr(const char * expr, size_t expr_size, size_t byte_pos, Int64 occur);
+    Int64 instrImpl(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur, Int64 ret_op);
 
+    std::optional<StringRef> processSubstrEmptyStringExpr(const char * expr, size_t expr_size, size_t byte_pos, Int64 occur);
+    std::optional<StringRef> substrImpl(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur);
 
     bool is_trivial;
     bool required_substring_is_prefix;

dbms/src/Common/OptimizedRegularExpression.inl.h

@@ -17,10 +17,12 @@
 #include <Common/StringUtils/StringUtils.h>
 #include <Common/UTF8Helpers.h>
 #include <Poco/Exception.h>
+#include <common/StringRef.h>
 #include <common/defines.h>
 #include <common/types.h>
 
 #include <iostream>
+#include <optional>
 
 #define MIN_LENGTH_FOR_STRSTR 3
 #define MAX_SUBPATTERNS 5
@@ -474,7 +476,7 @@ unsigned OptimizedRegularExpressionImpl<thread_safe>::match(const char * subject
 }
 
 template <bool thread_safe>
-Int64 OptimizedRegularExpressionImpl<thread_safe>::processEmptyStringExpr(const char * expr, size_t expr_size, size_t pos, Int64 occur)
+Int64 OptimizedRegularExpressionImpl<thread_safe>::processInstrEmptyStringExpr(const char * expr, size_t expr_size, size_t pos, Int64 occur)
 {
     if (occur != 1)
         return 0;
@@ -483,19 +485,38 @@ Int64 OptimizedRegularExpressionImpl<thread_safe>::processEmptyStringExpr(const
     return RegexType::FindAndConsume(&expr_sp, *re2) ? pos : 0;
 }
 
-static inline void checkArgs(Int64 utf8_total_len, size_t subject_size, Int64 pos, Int64 ret_op)
+template <bool thread_safe>
+std::optional<StringRef> OptimizedRegularExpressionImpl<thread_safe>::processSubstrEmptyStringExpr(const char * expr, size_t expr_size, size_t byte_pos, Int64 occur)
+{
+    if (occur != 1 || byte_pos != 1)
+        return std::nullopt;
+
+    StringPieceType expr_sp(expr, expr_size);
+    StringPieceType matched_str;
+    if (!RegexType::FindAndConsume(&expr_sp, *re2, &matched_str))
+        return std::nullopt;
+
+    return std::optional<StringRef>(StringRef(matched_str.data(), matched_str.size()));
+}
+
+static inline void checkInstrArgs(Int64 utf8_total_len, size_t subject_size, Int64 pos, Int64 ret_op)
 {
     RUNTIME_CHECK_MSG(!(ret_op != 0 && ret_op != 1), "Incorrect argument to regexp function: return_option must be 1 or 0");
     RUNTIME_CHECK_MSG(!(pos <= 0 || (pos > utf8_total_len && subject_size != 0)), "Index out of bounds in regular function.");
 }
 
+static inline void checkSubstrArgs(Int64 utf8_total_len, size_t subject_size, Int64 pos)
+{
+    RUNTIME_CHECK_MSG(!(pos <= 0 || (pos > utf8_total_len && subject_size != 0)), "Index out of bounds in regular function.");
+}
+
 static inline void makeOccurValid(Int64 & occur)
 {
-    occur = occur < 0 ? 1 : occur;
+    occur = occur < 1 ? 1 : occur;
 }
 
 template <bool thread_safe>
-Int64 OptimizedRegularExpressionImpl<thread_safe>::getSubstrMatchedIndex(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur, Int64 ret_op)
+Int64 OptimizedRegularExpressionImpl<thread_safe>::instrImpl(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur, Int64 ret_op)
 {
     size_t byte_offset = byte_pos - 1; // This is a offset for bytes, not utf8
     const char * expr = subject + byte_offset; // expr is the string actually passed into regexp to be matched
@@ -516,20 +537,53 @@ Int64 OptimizedRegularExpressionImpl<thread_safe>::getSubstrMatchedIndex(const c
     return ret_op == 0 ? DB::UTF8::bytePos2Utf8Pos(reinterpret_cast<const UInt8 *>(subject), byte_offset + 1) : DB::UTF8::bytePos2Utf8Pos(reinterpret_cast<const UInt8 *>(subject), byte_offset + matched_str.size() + 1);
 }
 
+template <bool thread_safe>
+std::optional<StringRef> OptimizedRegularExpressionImpl<thread_safe>::substrImpl(const char * subject, size_t subject_size, Int64 byte_pos, Int64 occur)
+{
+    size_t byte_offset = byte_pos - 1; // This is a offset for bytes, not utf8
+    const char * expr = subject + byte_offset; // expr is the string actually passed into regexp to be matched
+    size_t expr_size = subject_size - byte_offset;
+
+    StringPieceType expr_sp(expr, expr_size);
+    StringPieceType matched_str;
+    while (occur > 0)
+    {
+        if (!RegexType::FindAndConsume(&expr_sp, *re2, &matched_str))
+            return std::nullopt;
+
+        --occur;
+    }
+
+    return std::optional<StringRef>(StringRef(matched_str.data(), matched_str.size()));
+}
+
 template <bool thread_safe>
 Int64 OptimizedRegularExpressionImpl<thread_safe>::instr(const char * subject, size_t subject_size, Int64 pos, Int64 occur, Int64 ret_op)
 {
     Int64 utf8_total_len = DB::UTF8::countCodePoints(reinterpret_cast<const UInt8 *>(subject), subject_size);
+    ;
+    checkInstrArgs(utf8_total_len, subject_size, pos, ret_op);
+    makeOccurValid(occur);
 
-    checkArgs(utf8_total_len, subject_size, pos, ret_op);
+    if (unlikely(subject_size == 0))
+        return processInstrEmptyStringExpr(subject, subject_size, pos, occur);
 
+    size_t byte_pos = DB::UTF8::utf8Pos2bytePos(reinterpret_cast<const UInt8 *>(subject), pos);
+    return instrImpl(subject, subject_size, byte_pos, occur, ret_op);
+}
+
+template <bool thread_safe>
+std::optional<StringRef> OptimizedRegularExpressionImpl<thread_safe>::substr(const char * subject, size_t subject_size, Int64 pos, Int64 occur)
+{
+    Int64 utf8_total_len = DB::UTF8::countCodePoints(reinterpret_cast<const UInt8 *>(subject), subject_size);
+    checkSubstrArgs(utf8_total_len, subject_size, pos);
     makeOccurValid(occur);
 
     if (unlikely(subject_size == 0))
-        return processEmptyStringExpr(subject, subject_size, pos, occur);
+        return processSubstrEmptyStringExpr(subject, subject_size, pos, occur);
 
     size_t byte_pos = DB::UTF8::utf8Pos2bytePos(reinterpret_cast<const UInt8 *>(subject), pos);
-    return getSubstrMatchedIndex(subject, subject_size, byte_pos, occur, ret_op);
+    return substrImpl(subject, subject_size, byte_pos, occur);
 }
 
 #undef MIN_LENGTH_FOR_STRSTR

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -436,7 +436,7 @@ const std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     {tipb::ScalarFuncSig::RegexpLikeSig, "regexp_like"},
     {tipb::ScalarFuncSig::RegexpInStrSig, "regexp_instr"},
     // {tipb::ScalarFuncSig::RegexpReplaceSig, "regexp_replace"},
-    // {tipb::ScalarFuncSig::RegexpSubstrSig, "regexp_substr"},
+    {tipb::ScalarFuncSig::RegexpSubstrSig, "regexp_substr"},
 
     //{tipb::ScalarFuncSig::JsonExtractSig, "cast"},
     //{tipb::ScalarFuncSig::JsonUnquoteSig, "cast"},

dbms/src/Functions/FunctionsRegexp.cpp

@@ -287,6 +287,7 @@ struct ReplaceRegexpImpl
 using FunctionTiDBRegexp = FunctionStringRegexp<NameTiDBRegexp>;
 using FunctionRegexpLike = FunctionStringRegexp<NameRegexpLike>;
 using FunctionRegexpInstr = FunctionStringRegexpInstr<NameRegexpInstr>;
+using FunctionRegexpSubstr = FunctionStringRegexpSubstr<NameRegexpSubstr>;
 using FunctionReplaceRegexpOne = FunctionStringReplace<ReplaceRegexpImpl<true>, NameReplaceRegexpOne>;
 using FunctionReplaceRegexpAll = FunctionStringReplace<ReplaceRegexpImpl<false>, NameReplaceRegexpAll>;
 
@@ -297,6 +298,7 @@ void registerFunctionsRegexp(FunctionFactory & factory)
     factory.registerFunction<FunctionTiDBRegexp>();
     factory.registerFunction<FunctionRegexpLike>();
     factory.registerFunction<FunctionRegexpInstr>();
+    factory.registerFunction<FunctionRegexpSubstr>();
 }
 
 } // namespace DB