ARROW-17370: [C++] Add limit to SplitString() (apache#13833)

Authored-by: Sutou Kouhei <[email protected]> Signed-off-by: Sutou Kouhei <[email protected]>
kou · Aug 12, 2022 · a70908d · a70908d
1 parent d11a553
commit a70908d
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 3 deletions.
diff --git a/cpp/src/arrow/util/string.cc b/cpp/src/arrow/util/string.cc
@@ -92,11 +92,16 @@ Status ParseHexValue(const char* data, uint8_t* out) {
 
 namespace internal {
 
-std::vector<util::string_view> SplitString(util::string_view v, char delimiter) {
+std::vector<util::string_view> SplitString(util::string_view v, char delimiter,
+                                           int64_t limit) {
   std::vector<util::string_view> parts;
   size_t start = 0, end;
   while (true) {
-    end = v.find(delimiter, start);
+    if (limit > 0 && static_cast<size_t>(limit - 1) <= parts.size()) {
+      end = std::string::npos;
+    } else {
+      end = v.find(delimiter, start);
+    }
     parts.push_back(v.substr(start, end - start));
     if (end == std::string::npos) {
       break;

diff --git a/cpp/src/arrow/util/string.h b/cpp/src/arrow/util/string.h
@@ -45,7 +45,8 @@ namespace internal {
 
 /// \brief Split a string with a delimiter
 ARROW_EXPORT
-std::vector<util::string_view> SplitString(util::string_view v, char delim);
+std::vector<util::string_view> SplitString(util::string_view v, char delim,
+                                           int64_t limit = 0);
 
 /// \brief Join strings with a delimiter
 ARROW_EXPORT

diff --git a/cpp/src/arrow/util/string_test.cc b/cpp/src/arrow/util/string_test.cc
@@ -140,5 +140,31 @@ TEST(SplitString, OnlyDemiliter) {
   EXPECT_EQ(parts[1], "");
 }
 
+TEST(SplitString, Limit) {
+  std::string input = "a:b:c";
+  auto parts = SplitString(input, ':', 2);
+  ASSERT_EQ(parts.size(), 2);
+  EXPECT_EQ(parts[0], "a");
+  EXPECT_EQ(parts[1], "b:c");
+}
+
+TEST(SplitString, LimitOver) {
+  std::string input = "a:b:c";
+  auto parts = SplitString(input, ':', 4);
+  ASSERT_EQ(parts.size(), 3);
+  EXPECT_EQ(parts[0], "a");
+  EXPECT_EQ(parts[1], "b");
+  EXPECT_EQ(parts[2], "c");
+}
+
+TEST(SplitString, LimitZero) {
+  std::string input = "a:b:c";
+  auto parts = SplitString(input, ':', 0);
+  ASSERT_EQ(parts.size(), 3);
+  EXPECT_EQ(parts[0], "a");
+  EXPECT_EQ(parts[1], "b");
+  EXPECT_EQ(parts[2], "c");
+}
+
 }  // namespace internal
 }  // namespace arrow