From 160395211152e9bd151f7d24affc204a24e0975f Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 12 Jun 2024 10:37:29 +0800
Subject: [PATCH] expression: fix substring_index crash tiflash in some case
 (#9137) (#9141)

close pingcap/tiflash#9116

Fix tiflash crash caused by function `substring_index` in some corner cases

Co-authored-by: wshwsh12 <793703860@qq.com>
---
 dbms/src/Functions/FunctionsString.cpp             | 4 +++-
 dbms/src/Functions/tests/gtest_substring_index.cpp | 9 +++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/dbms/src/Functions/FunctionsString.cpp b/dbms/src/Functions/FunctionsString.cpp
index e49f5b4895e..0d8b06a8bd2 100644
--- a/dbms/src/Functions/FunctionsString.cpp
+++ b/dbms/src/Functions/FunctionsString.cpp
@@ -4868,6 +4868,7 @@ class FunctionSubStringIndex : public IFunction
         const UInt8 * pos = begin;
         const UInt8 * end = pos + data_size;
         assert(delim_size != 0);
+        assert(count != 0);
         if (count > 0)
         {
             // Fast exit when count * delim_size > data_size
@@ -4883,10 +4884,11 @@ class FunctionSubStringIndex : public IFunction
                 if (match == end || count == 0)
                 {
                     copyDataToResult(res_data, res_offset, begin, match);
-                    break;
+                    return;
                 }
                 pos = match + delim_size;
             }
+            copyDataToResult(res_data, res_offset, begin, end);
         }
         else
         {
diff --git a/dbms/src/Functions/tests/gtest_substring_index.cpp b/dbms/src/Functions/tests/gtest_substring_index.cpp
index d93d8e1bade..7bf4028c0a2 100644
--- a/dbms/src/Functions/tests/gtest_substring_index.cpp
+++ b/dbms/src/Functions/tests/gtest_substring_index.cpp
@@ -287,6 +287,15 @@ try
             createColumn<Nullable<String>>({"www.pingcap.com", "www...www", "中文.测.试。。。", "www.www"}),
             createColumn<Nullable<String>>({"", "", "", ""}),
             createColumn<Nullable<Int64>>({2, 2, 2, 2})));
+
+    // Test issue 9116
+    ASSERT_COLUMN_EQ(
+        createColumn<Nullable<String>>({"aaabbba", "aaabbbaa", "aaabbbaaa", "aaabbbaaa", "aaabbbaaa"}),
+        executeFunction(
+            func_name,
+            createColumn<Nullable<String>>({"aaabbbaaa", "aaabbbaaa", "aaabbbaaa", "aaabbbaaa", "aaabbbaaa"}),
+            createColumn<Nullable<String>>({"a", "a", "a", "a", "a"}),
+            createColumn<Nullable<Int64>>({5, 6, 7, 8, 9})));
 }
 CATCH