ARROW-1712: [C++] Add method to BinaryBuilder to reserve space for value data

cpp/src/arrow/array-test.cc

@@ -1154,6 +1154,54 @@ TEST_F(TestBinaryBuilder, TestScalarAppend) {
     }
   }
 }
+
+TEST_F(TestBinaryBuilder, TestCapacityReserve) {
+  vector<string> strings = {"a", "bb", "cc", "ddddd", "eeeee"};
+  int64_t N = static_cast<int>(strings.size());
+  int64_t length = 0;
+  int64_t data_length = 0;
+  int64_t capacity = N;
+
+  ASSERT_OK(builder_->Reserve(capacity));
+  ASSERT_OK(builder_->ReserveData(capacity));
+
+  ASSERT_EQ(builder_->length(), length);
+  ASSERT_EQ(builder_->capacity(), BitUtil::NextPower2(capacity));
+  ASSERT_EQ(builder_->value_data_length(), data_length);
+  ASSERT_EQ(builder_->value_data_capacity(), capacity);
+
+  for(const string& str : strings) {
+    ASSERT_OK(builder_->Append(str));
+    length++;
+    data_length += static_cast<int>(str.size());
+
+    ASSERT_EQ(builder_->length(), length);
+    ASSERT_EQ(builder_->capacity(), BitUtil::NextPower2(capacity));
+    ASSERT_EQ(builder_->value_data_length(), data_length);
+    if (data_length <= capacity) {
+      ASSERT_EQ(builder_->value_data_capacity(), capacity);
+    } else {
+      ASSERT_EQ(builder_->value_data_capacity(), data_length);
+    }
+  }
+
+  int extra_capacity = 20;
+
+  ASSERT_OK(builder_->Reserve(extra_capacity));
+  ASSERT_OK(builder_->ReserveData(extra_capacity));
+
+  ASSERT_EQ(builder_->length(), length);
+  ASSERT_EQ(builder_->capacity(), BitUtil::NextPower2(length + extra_capacity));
+  ASSERT_EQ(builder_->value_data_length(), data_length);
+  ASSERT_EQ(builder_->value_data_capacity(), data_length + extra_capacity);
+
+  Done();
+
+  ASSERT_EQ(result_->length(), N);
+  ASSERT_EQ(result_->null_count(), 0);
+  ASSERT_EQ(result_->value_data()->size(), data_length);
+  ASSERT_EQ(result_->value_data()->capacity(), BitUtil::RoundUpToMultipleOf64(data_length + extra_capacity));
+}
 
 TEST_F(TestBinaryBuilder, TestZeroLength) {
   // All buffers are null

cpp/src/arrow/buffer.h

@@ -333,6 +333,7 @@ class ARROW_EXPORT TypedBufferBuilder : public BufferBuilder {
 
   const T* data() const { return reinterpret_cast<const T*>(data_); }
   int64_t length() const { return size_ / sizeof(T); }
+  int64_t capacity() const { return capacity_ / sizeof(T); }
 };
 
 /// \brief Allocate a fixed size mutable buffer from a memory pool

cpp/src/arrow/builder.cc

@@ -1208,7 +1208,7 @@ ArrayBuilder* ListBuilder::value_builder() const {
 // String and binary
 
 BinaryBuilder::BinaryBuilder(const std::shared_ptr<DataType>& type, MemoryPool* pool)
-    : ArrayBuilder(type, pool), offsets_builder_(pool), value_data_builder_(pool) {}
+    : ArrayBuilder(type, pool), offsets_builder_(pool), value_data_builder_(pool), data_capacity_(0) {}
 
 BinaryBuilder::BinaryBuilder(MemoryPool* pool) : BinaryBuilder(binary(), pool) {}
 
@@ -1222,9 +1222,21 @@ Status BinaryBuilder::Init(int64_t elements) {
 Status BinaryBuilder::Resize(int64_t capacity) {
   DCHECK_LT(capacity, std::numeric_limits<int32_t>::max());
   // one more then requested for offsets
-  RETURN_NOT_OK(offsets_builder_.Resize((capacity + 1) * sizeof(int64_t)));
+  RETURN_NOT_OK(offsets_builder_.Resize((capacity + 1) * sizeof(int32_t)));
   return ArrayBuilder::Resize(capacity);
 }
+
+Status BinaryBuilder::ReserveData(int64_t capacity) {
+  if (value_data_length() + capacity > data_capacity_) {
+    if (value_data_length() + capacity > std::numeric_limits<int32_t>::max()) {
+      return Status::Invalid("Cannot reserve capacity larger than 2^31 - 1 in length for binary data");
+    }    
+
+    RETURN_NOT_OK(value_data_builder_.Resize(value_data_length() + capacity));
+    data_capacity_ = value_data_length() + capacity;
+  }  
+  return Status::OK();
+}
 
 Status BinaryBuilder::AppendNextOffset() {
   const int64_t num_bytes = value_data_builder_.length();
@@ -1241,6 +1253,9 @@ Status BinaryBuilder::Append(const uint8_t* value, int32_t length) {
   RETURN_NOT_OK(Reserve(1));
   RETURN_NOT_OK(AppendNextOffset());
   RETURN_NOT_OK(value_data_builder_.Append(value, length));
+  if (data_capacity_ < value_data_length()) {
+    data_capacity_ = value_data_length();
+  }
   UnsafeAppendToBitmap(true);
   return Status::OK();
 }

cpp/src/arrow/builder.h

@@ -682,10 +682,13 @@ class ARROW_EXPORT BinaryBuilder : public ArrayBuilder {
 
   Status Init(int64_t elements) override;
   Status Resize(int64_t capacity) override;
+  Status ReserveData(int64_t capacity);
   Status FinishInternal(std::shared_ptr<ArrayData>* out) override;
 
   /// \return size of values buffer so far
   int64_t value_data_length() const { return value_data_builder_.length(); }
+  /// \return capacity of values buffer
+  int64_t value_data_capacity() const { return data_capacity_; }
 
   /// Temporary access to a value.
   ///
@@ -696,6 +699,7 @@ class ARROW_EXPORT BinaryBuilder : public ArrayBuilder {
   TypedBufferBuilder<int32_t> offsets_builder_;
   TypedBufferBuilder<uint8_t> value_data_builder_;
 
+  int64_t data_capacity_;
   static constexpr int64_t kMaximumCapacity = std::numeric_limits<int32_t>::max() - 1;
 
   Status AppendNextOffset();