ARROW-7985: [C++] Fix builder capacity check

It should be ok to resize below the current allocated builder capacity, just not below the current populated length. Closes #6595 from pitrou/ARROW-7985-fix-builder-capacity-check Authored-by: Antoine Pitrou <antoine@python.org> Signed-off-by: Benjamin Kietzman <bengilgit@gmail.com>
apache · Mar 12, 2020 · 2849b64 · 2849b64
1 parent 590a168
commit 2849b64
Show file tree

Hide file tree

Showing 12 changed files with 53 additions and 19 deletions.
diff --git a/cpp/src/arrow/array/builder_adaptive.cc b/cpp/src/arrow/array/builder_adaptive.cc
@@ -47,7 +47,7 @@ void AdaptiveIntBuilderBase::Reset() {
 }
 
 Status AdaptiveIntBuilderBase::Resize(int64_t capacity) {
-  RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+  RETURN_NOT_OK(CheckCapacity(capacity));
   capacity = std::max(capacity, kMinBuilderCapacity);
 
   int64_t nbytes = capacity * int_size_;

diff --git a/cpp/src/arrow/array/builder_base.cc b/cpp/src/arrow/array/builder_base.cc
@@ -69,7 +69,7 @@ Status ArrayBuilder::AppendToBitmap(int64_t num_bits, bool value) {
 }
 
 Status ArrayBuilder::Resize(int64_t capacity) {
-  RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+  RETURN_NOT_OK(CheckCapacity(capacity));
   capacity_ = capacity;
   return null_bitmap_builder_.Resize(capacity);
 }

diff --git a/cpp/src/arrow/array/builder_base.h b/cpp/src/arrow/array/builder_base.h
@@ -190,13 +190,16 @@ class ARROW_EXPORT ArrayBuilder {
     return Status::OK();
   }
 
-  static Status CheckCapacity(int64_t new_capacity, int64_t old_capacity) {
-    if (new_capacity < 0) {
-      return Status::Invalid("Resize capacity must be positive");
+  // Check the requested capacity for validity
+  Status CheckCapacity(int64_t new_capacity) {
+    if (ARROW_PREDICT_FALSE(new_capacity < 0)) {
+      return Status::Invalid(
+          "Resize capacity must be positive (requested: ", new_capacity, ")");
     }
 
-    if (new_capacity < old_capacity) {
-      return Status::Invalid("Resize cannot downsize");
+    if (ARROW_PREDICT_FALSE(new_capacity < length_)) {
+      return Status::Invalid("Resize cannot downsize (requested: ", new_capacity,
+                             ", current length: ", length_, ")");
     }
 
     return Status::OK();

diff --git a/cpp/src/arrow/array/builder_binary.cc b/cpp/src/arrow/array/builder_binary.cc
@@ -79,7 +79,7 @@ void FixedSizeBinaryBuilder::Reset() {
 }
 
 Status FixedSizeBinaryBuilder::Resize(int64_t capacity) {
-  RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+  RETURN_NOT_OK(CheckCapacity(capacity));
   RETURN_NOT_OK(byte_builder_.Resize(capacity * byte_width_));
   return ArrayBuilder::Resize(capacity);
 }

diff --git a/cpp/src/arrow/array/builder_binary.h b/cpp/src/arrow/array/builder_binary.h
@@ -235,7 +235,7 @@ class BaseBinaryBuilder : public ArrayBuilder {
       return Status::CapacityError("BinaryBuilder cannot reserve space for more than ",
                                    memory_limit(), " child elements, got ", capacity);
     }
-    ARROW_RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+    ARROW_RETURN_NOT_OK(CheckCapacity(capacity));
 
     // One more than requested for offsets
     ARROW_RETURN_NOT_OK(offsets_builder_.Resize(capacity + 1));

diff --git a/cpp/src/arrow/array/builder_dict.h b/cpp/src/arrow/array/builder_dict.h
@@ -248,7 +248,7 @@ class DictionaryBuilderBase : public ArrayBuilder {
   }
 
   Status Resize(int64_t capacity) override {
-    ARROW_RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+    ARROW_RETURN_NOT_OK(CheckCapacity(capacity));
     capacity = std::max(capacity, kMinBuilderCapacity);
     ARROW_RETURN_NOT_OK(indices_builder_.Resize(capacity));
     capacity_ = indices_builder_.capacity();
@@ -356,7 +356,7 @@ class DictionaryBuilderBase<BuilderType, NullType> : public ArrayBuilder {
   }
 
   Status Resize(int64_t capacity) override {
-    ARROW_RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+    ARROW_RETURN_NOT_OK(CheckCapacity(capacity));
     capacity = std::max(capacity, kMinBuilderCapacity);
 
     ARROW_RETURN_NOT_OK(indices_builder_.Resize(capacity));

diff --git a/cpp/src/arrow/array/builder_nested.cc b/cpp/src/arrow/array/builder_nested.cc
@@ -175,7 +175,7 @@ Status FixedSizeListBuilder::AppendNulls(int64_t length) {
 }
 
 Status FixedSizeListBuilder::Resize(int64_t capacity) {
-  RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+  RETURN_NOT_OK(CheckCapacity(capacity));
   return ArrayBuilder::Resize(capacity);
 }
 

diff --git a/cpp/src/arrow/array/builder_nested.h b/cpp/src/arrow/array/builder_nested.h
@@ -54,9 +54,9 @@ class BaseListBuilder : public ArrayBuilder {
       return Status::CapacityError("List array cannot reserve space for more than ",
                                    maximum_elements(), " got ", capacity);
     }
-    ARROW_RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+    ARROW_RETURN_NOT_OK(CheckCapacity(capacity));
 
-    // one more then requested for offsets
+    // One more than requested for offsets
     ARROW_RETURN_NOT_OK(offsets_builder_.Resize(capacity + 1));
     return ArrayBuilder::Resize(capacity);
   }

diff --git a/cpp/src/arrow/array/builder_primitive.cc b/cpp/src/arrow/array/builder_primitive.cc
@@ -58,7 +58,7 @@ void BooleanBuilder::Reset() {
 }
 
 Status BooleanBuilder::Resize(int64_t capacity) {
-  RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+  RETURN_NOT_OK(CheckCapacity(capacity));
   capacity = std::max(capacity, kMinBuilderCapacity);
   RETURN_NOT_OK(data_builder_.Resize(capacity));
   return ArrayBuilder::Resize(capacity);

diff --git a/cpp/src/arrow/array/builder_primitive.h b/cpp/src/arrow/array/builder_primitive.h
@@ -101,7 +101,7 @@ class NumericBuilder : public ArrayBuilder {
   void Reset() override { data_builder_.Reset(); }
 
   Status Resize(int64_t capacity) override {
-    ARROW_RETURN_NOT_OK(CheckCapacity(capacity, capacity_));
+    ARROW_RETURN_NOT_OK(CheckCapacity(capacity));
     capacity = std::max(capacity, kMinBuilderCapacity);
     ARROW_RETURN_NOT_OK(data_builder_.Resize(capacity));
     return ArrayBuilder::Resize(capacity);

diff --git a/cpp/src/arrow/array_list_test.cc b/cpp/src/arrow/array_list_test.cc
@@ -435,6 +435,27 @@ class TestListArray : public TestBuilder {
     ASSERT_RAISES(Invalid, ValidateOffsets(2, {0, 7, 4}, values));
   }
 
+  void TestCornerCases() {
+    // ARROW-7985
+    ASSERT_OK(builder_->AppendNull());
+    Done();
+    auto expected = ArrayFromJSON(type_, "[null]");
+    AssertArraysEqual(*result_, *expected);
+
+    SetUp();
+    ASSERT_OK(builder_->Append());
+    Done();
+    expected = ArrayFromJSON(type_, "[[]]");
+    AssertArraysEqual(*result_, *expected);
+
+    SetUp();
+    ASSERT_OK(builder_->AppendNull());
+    ASSERT_OK(builder_->value_builder()->Reserve(100));
+    Done();
+    expected = ArrayFromJSON(type_, "[null]");
+    AssertArraysEqual(*result_, *expected);
+  }
+
  protected:
   std::shared_ptr<DataType> value_type_;
 
@@ -474,6 +495,8 @@ TYPED_TEST(TestListArray, TestFlattenNonEmptyBackingNulls) {
 
 TYPED_TEST(TestListArray, ValidateOffsets) { this->TestValidateOffsets(); }
 
+TYPED_TEST(TestListArray, CornerCases) { this->TestCornerCases(); }
+
 // ----------------------------------------------------------------------
 // Map tests
 

diff --git a/cpp/src/arrow/array_test.cc b/cpp/src/arrow/array_test.cc
@@ -397,9 +397,17 @@ TEST_F(TestBuilder, TestResizeDownsize) {
 
   ASSERT_OK(builder.Resize(1000));
   ASSERT_EQ(1000, builder.capacity());
-
-  // Can't downsize.
-  ASSERT_RAISES(Invalid, builder.Resize(500));
+  ASSERT_EQ(0, builder.length());
+  ASSERT_OK(builder.AppendNulls(500));
+  ASSERT_EQ(1000, builder.capacity());
+  ASSERT_EQ(500, builder.length());
+
+  // Can downsize below current capacity
+  ASSERT_OK(builder.Resize(500));
+  // ... but not below current populated length
+  ASSERT_RAISES(Invalid, builder.Resize(499));
+  ASSERT_GE(500, builder.capacity());
+  ASSERT_EQ(500, builder.length());
 }
 
 template <typename Attrs>