expression: fix substring_index crash tiflash in some case (#9137)

close #9116 Fix tiflash crash caused by function `substring_index` in some corner cases
pingcap · Jun 11, 2024 · 7d7622e · 7d7622e
1 parent cbac09d
commit 7d7622e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/dbms/src/Functions/FunctionsString.cpp b/dbms/src/Functions/FunctionsString.cpp
@@ -5212,6 +5212,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
@@ -5227,10 +5228,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
@@ -364,6 +364,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