apache · kszucs · Aug 31, 2020 · Sep 1, 2020 · Sep 1, 2020 · Sep 2, 2020
diff --git a/cpp/src/arrow/array/array_base.cc b/cpp/src/arrow/array/array_base.cc
@@ -161,7 +161,13 @@ struct ScalarFromArraySlotImpl {
     }
 
     if (array_.IsNull(index_)) {
-      return MakeNullScalar(array_.type());
+      auto null = MakeNullScalar(array_.type());
+      if (is_dictionary(array_.type()->id())) {
+        auto& dict_null = checked_cast<DictionaryScalar&>(*null);
+        const auto& dict_array = checked_cast<const DictionaryArray&>(array_);
+        dict_null.value.dictionary = dict_array.dictionary();
+      }
+      return null;
     }
 
     RETURN_NOT_OK(VisitArrayInline(array_, this));

diff --git a/cpp/src/arrow/array/array_binary_test.cc b/cpp/src/arrow/array/array_binary_test.cc
@@ -570,6 +570,26 @@ class TestStringBuilder : public TestBuilder {
     ASSERT_EQ(reps * 40, result_->value_data()->size());
   }
 
+  void TestOverflowCheck() {
+    auto max_size = builder_->memory_limit();
+
+    ASSERT_OK(builder_->ValidateOverflow(1));
+    ASSERT_OK(builder_->ValidateOverflow(max_size));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_size + 1));
+
+    ASSERT_OK(builder_->Append("bb"));
+    ASSERT_OK(builder_->ValidateOverflow(max_size - 2));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_size - 1));
+
+    ASSERT_OK(builder_->AppendNull());
+    ASSERT_OK(builder_->ValidateOverflow(max_size - 2));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_size - 1));
+
+    ASSERT_OK(builder_->Append("ccc"));
+    ASSERT_OK(builder_->ValidateOverflow(max_size - 5));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_size - 4));
+  }
+
   void TestZeroLength() {
     // All buffers are null
     Done();
@@ -602,6 +622,8 @@ TYPED_TEST(TestStringBuilder, TestCapacityReserve) { this->TestCapacityReserve()
 
 TYPED_TEST(TestStringBuilder, TestZeroLength) { this->TestZeroLength(); }
 
+TYPED_TEST(TestStringBuilder, TestOverflowCheck) { this->TestOverflowCheck(); }
+
 // ----------------------------------------------------------------------
 // ChunkedBinaryBuilder tests
 

diff --git a/cpp/src/arrow/array/array_list_test.cc b/cpp/src/arrow/array/array_list_test.cc
@@ -467,6 +467,32 @@ class TestListArray : public TestBuilder {
     AssertArraysEqual(*result_, *expected);
   }
 
+  void TestOverflowCheck() {
+    Int16Builder* vb = checked_cast<Int16Builder*>(builder_->value_builder());
+    auto max_elements = builder_->maximum_elements();
+
+    ASSERT_OK(builder_->ValidateOverflow(1));
+    ASSERT_OK(builder_->ValidateOverflow(max_elements));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_elements + 1));
+
+    ASSERT_OK(builder_->Append());
+    ASSERT_OK(vb->Append(1));
+    ASSERT_OK(vb->Append(2));
+    ASSERT_OK(builder_->ValidateOverflow(max_elements - 2));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_elements - 1));
+
+    ASSERT_OK(builder_->AppendNull());
+    ASSERT_OK(builder_->ValidateOverflow(max_elements - 2));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_elements - 1));
+
+    ASSERT_OK(builder_->Append());
+    ASSERT_OK(vb->Append(1));
+    ASSERT_OK(vb->Append(2));
+    ASSERT_OK(vb->Append(3));
+    ASSERT_OK(builder_->ValidateOverflow(max_elements - 5));
+    ASSERT_RAISES(CapacityError, builder_->ValidateOverflow(max_elements - 4));
+  }
+
  protected:
   std::shared_ptr<DataType> value_type_;
 
@@ -508,6 +534,12 @@ TYPED_TEST(TestListArray, ValidateOffsets) { this->TestValidateOffsets(); }
 
 TYPED_TEST(TestListArray, CornerCases) { this->TestCornerCases(); }
 
+#ifndef ARROW_LARGE_MEMORY_TESTS
+TYPED_TEST(TestListArray, DISABLED_TestOverflowCheck) { this->TestOverflowCheck(); }
+#else
+TYPED_TEST(TestListArray, TestOverflowCheck) { this->TestOverflowCheck(); }
+#endif
+
 // ----------------------------------------------------------------------
 // Map tests
 

diff --git a/cpp/src/arrow/array/builder_base.cc b/cpp/src/arrow/array/builder_base.cc
@@ -99,6 +99,12 @@ Status ArrayBuilder::Finish(std::shared_ptr<Array>* out) {
   return Status::OK();
 }
 
+Result<std::shared_ptr<Array>> ArrayBuilder::Finish() {
+  std::shared_ptr<Array> out;
+  RETURN_NOT_OK(Finish(&out));
+  return out;
+}
+
 void ArrayBuilder::Reset() {
   capacity_ = length_ = null_count_ = 0;
   null_bitmap_builder_.Reset();

diff --git a/cpp/src/arrow/array/builder_base.h b/cpp/src/arrow/array/builder_base.h
@@ -56,6 +56,8 @@ class ARROW_EXPORT ArrayBuilder {
   /// skip shared pointers and just return a raw pointer
   ArrayBuilder* child(int i) { return children_[i].get(); }
 
+  const std::shared_ptr<ArrayBuilder>& child_builder(int i) const { return children_[i]; }
+
   int num_children() const { return static_cast<int>(children_.size()); }
 
   virtual int64_t length() const { return length_; }
@@ -118,6 +120,13 @@ class ARROW_EXPORT ArrayBuilder {
   /// \return Status
   Status Finish(std::shared_ptr<Array>* out);
 
+  /// \brief Return result of builder as an Array object.
+  ///
+  /// The builder is reset except for DictionaryBuilder.
+  ///
+  /// \return The finalized Array object
+  Result<std::shared_ptr<Array>> Finish();
+
   /// \brief Return the type of the built Array
   virtual std::shared_ptr<DataType> type() const = 0;
 

diff --git a/cpp/src/arrow/array/builder_binary.h b/cpp/src/arrow/array/builder_binary.h
@@ -78,9 +78,7 @@ class BaseBinaryBuilder : public ArrayBuilder {
 
   Status AppendNulls(int64_t length) final {
     const int64_t num_bytes = value_data_builder_.length();
-    if (ARROW_PREDICT_FALSE(num_bytes > memory_limit())) {
-      return AppendOverflow(num_bytes);
-    }
+    ARROW_RETURN_NOT_OK(ValidateOverflow(0));
     ARROW_RETURN_NOT_OK(Reserve(length));
     for (int64_t i = 0; i < length; ++i) {
       offsets_builder_.UnsafeAppend(static_cast<offset_type>(num_bytes));
@@ -232,6 +230,16 @@ class BaseBinaryBuilder : public ArrayBuilder {
     value_data_builder_.Reset();
   }
 
+  Status ValidateOverflow(int64_t new_bytes) {
+    auto new_size = value_data_builder_.length() + new_bytes;
+    if (ARROW_PREDICT_FALSE(new_size > memory_limit())) {
+      return Status::CapacityError("array cannot contain more than ", memory_limit(),
+                                   " bytes, have ", new_size);
+    } else {
+      return Status::OK();
+    }
+  }
+
   Status Resize(int64_t capacity) override {
     // XXX Why is this check necessary?  There is no reason to disallow, say,
     // binary arrays with more than 2**31 empty or null values.
@@ -249,12 +257,8 @@ class BaseBinaryBuilder : public ArrayBuilder {
   /// \brief Ensures there is enough allocated capacity to append the indicated
   /// number of bytes to the value data buffer without additional allocations
   Status ReserveData(int64_t elements) {
-    const int64_t size = value_data_length() + elements;
-    ARROW_RETURN_IF(size > memory_limit(),
-                    Status::CapacityError("Cannot reserve capacity larger than ",
-                                          memory_limit(), " bytes"));
-    return (size > value_data_capacity()) ? value_data_builder_.Reserve(elements)
-                                          : Status::OK();
+    ARROW_RETURN_NOT_OK(ValidateOverflow(elements));
+    return value_data_builder_.Reserve(elements);
   }
 
   Status FinishInternal(std::shared_ptr<ArrayData>* out) override {
@@ -317,16 +321,9 @@ class BaseBinaryBuilder : public ArrayBuilder {
   TypedBufferBuilder<offset_type> offsets_builder_;
   TypedBufferBuilder<uint8_t> value_data_builder_;
 
-  Status AppendOverflow(int64_t num_bytes) {
-    return Status::CapacityError("array cannot contain more than ", memory_limit(),
-                                 " bytes, have ", num_bytes);
-  }
-
   Status AppendNextOffset() {
     const int64_t num_bytes = value_data_builder_.length();
-    if (ARROW_PREDICT_FALSE(num_bytes > memory_limit())) {
-      return AppendOverflow(num_bytes);
-    }
+    ARROW_RETURN_NOT_OK(ValidateOverflow(0));
     return offsets_builder_.Append(static_cast<offset_type>(num_bytes));
   }
 
@@ -462,6 +459,23 @@ class ARROW_EXPORT FixedSizeBinaryBuilder : public ArrayBuilder {
     byte_builder_.UnsafeAppend(/*num_copies=*/byte_width_, 0);
   }
 
+  Status ValidateOverflow(int64_t new_bytes) const {
+    auto new_size = byte_builder_.length() + new_bytes;
+    if (ARROW_PREDICT_FALSE(new_size > memory_limit())) {
+      return Status::CapacityError("array cannot contain more than ", memory_limit(),
+                                   " bytes, have ", new_size);
+    } else {
+      return Status::OK();
+    }
+  }
+
+  /// \brief Ensures there is enough allocated capacity to append the indicated
+  /// number of bytes to the value data buffer without additional allocations
+  Status ReserveData(int64_t elements) {
+    ARROW_RETURN_NOT_OK(ValidateOverflow(elements));
+    return byte_builder_.Reserve(elements);
+  }
+
   void Reset() override;
   Status Resize(int64_t capacity) override;
   Status FinishInternal(std::shared_ptr<ArrayData>* out) override;

diff --git a/cpp/src/arrow/array/builder_nested.cc b/cpp/src/arrow/array/builder_nested.cc
@@ -54,6 +54,18 @@ MapBuilder::MapBuilder(MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& ke
     : MapBuilder(pool, key_builder, item_builder,
                  map(key_builder->type(), item_builder->type(), keys_sorted)) {}
 
+MapBuilder::MapBuilder(MemoryPool* pool,
+                       const std::shared_ptr<ArrayBuilder>& struct_builder,
+                       const std::shared_ptr<DataType>& type)
+    : ArrayBuilder(pool) {
+  auto map_type = internal::checked_cast<const MapType*>(type.get());
+  keys_sorted_ = map_type->keys_sorted();
+  key_builder_ = struct_builder->child_builder(0);
+  item_builder_ = struct_builder->child_builder(1);
+  list_builder_ =
+      std::make_shared<ListBuilder>(pool, struct_builder, struct_builder->type());
+}
+
 Status MapBuilder::Resize(int64_t capacity) {
   RETURN_NOT_OK(list_builder_->Resize(capacity));
   capacity_ = list_builder_->capacity();
@@ -170,6 +182,19 @@ Status FixedSizeListBuilder::AppendNulls(int64_t length) {
   return value_builder_->AppendNulls(list_size_ * length);
 }
 
+Status FixedSizeListBuilder::ValidateOverflow(int64_t new_elements) {
+  auto new_length = value_builder_->length() + new_elements;
+  if (new_elements != list_size_) {
+    return Status::Invalid("Length of item not correct: expected ", list_size_,
+                           " but got array of size ", new_elements);
+  }
+  if (new_length > maximum_elements()) {
+    return Status::CapacityError("array cannot contain more than ", maximum_elements(),
+                                 " elements, have ", new_elements);
+  }
+  return Status::OK();
+}
+
 Status FixedSizeListBuilder::Resize(int64_t 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
@@ -100,7 +100,7 @@ class BaseListBuilder : public ArrayBuilder {
 
   Status AppendNulls(int64_t length) final {
     ARROW_RETURN_NOT_OK(Reserve(length));
-    ARROW_RETURN_NOT_OK(CheckNextOffset());
+    ARROW_RETURN_NOT_OK(ValidateOverflow(0));
     UnsafeAppendToBitmap(length, false);
     const int64_t num_values = value_builder_->length();
     for (int64_t i = 0; i < length; ++i) {
@@ -131,6 +131,16 @@ class BaseListBuilder : public ArrayBuilder {
     return Status::OK();
   }
 
+  Status ValidateOverflow(int64_t new_elements) const {
+    auto new_length = value_builder_->length() + new_elements;
+    if (ARROW_PREDICT_FALSE(new_length > maximum_elements())) {
+      return Status::CapacityError("List array cannot contain more than ",
+                                   maximum_elements(), " elements, have ", new_elements);
+    } else {
+      return Status::OK();
+    }
+  }
+
   ArrayBuilder* value_builder() const { return value_builder_.get(); }
 
   // Cannot make this a static attribute because of linking issues
@@ -147,17 +157,8 @@ class BaseListBuilder : public ArrayBuilder {
   std::shared_ptr<ArrayBuilder> value_builder_;
   std::shared_ptr<Field> value_field_;
 
-  Status CheckNextOffset() const {
-    const int64_t num_values = value_builder_->length();
-    ARROW_RETURN_IF(
-        num_values > maximum_elements(),
-        Status::CapacityError("List array cannot contain more than ", maximum_elements(),
-                              " child elements,", " have ", num_values));
-    return Status::OK();
-  }
-
   Status AppendNextOffset() {
-    ARROW_RETURN_NOT_OK(CheckNextOffset());
+    ARROW_RETURN_NOT_OK(ValidateOverflow(0));
     const int64_t num_values = value_builder_->length();
     return offsets_builder_.Append(static_cast<offset_type>(num_values));
   }
@@ -227,6 +228,9 @@ class ARROW_EXPORT MapBuilder : public ArrayBuilder {
   MapBuilder(MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& key_builder,
              const std::shared_ptr<ArrayBuilder>& item_builder, bool keys_sorted = false);
 
+  MapBuilder(MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& item_builder,
+             const std::shared_ptr<DataType>& type);
+
   Status Resize(int64_t capacity) override;
   void Reset() override;
   Status FinishInternal(std::shared_ptr<ArrayData>* out) override;
@@ -276,6 +280,10 @@ class ARROW_EXPORT MapBuilder : public ArrayBuilder {
     return map(key_builder_->type(), item_builder_->type(), keys_sorted_);
   }
 
+  Status ValidateOverflow(int64_t new_elements) {
+    return list_builder_->ValidateOverflow(new_elements);
+  }
+
  protected:
   inline Status AdjustStructBuilderLength();
 
@@ -343,12 +351,19 @@ class ARROW_EXPORT FixedSizeListBuilder : public ArrayBuilder {
   /// automatically.
   Status AppendNulls(int64_t length) final;
 
+  Status ValidateOverflow(int64_t new_elements);
+
   ArrayBuilder* value_builder() const { return value_builder_.get(); }
 
   std::shared_ptr<DataType> type() const override {
     return fixed_size_list(value_field_->WithType(value_builder_->type()), list_size_);
   }
 
+  // Cannot make this a static attribute because of linking issues
+  static constexpr int64_t maximum_elements() {
+    return std::numeric_limits<FixedSizeListType::offset_type>::max() - 1;
+  }
+
  protected:
   std::shared_ptr<Field> value_field_;
   const int32_t list_size_;

diff --git a/cpp/src/arrow/array/builder_primitive.h b/cpp/src/arrow/array/builder_primitive.h
@@ -31,6 +31,9 @@ namespace arrow {
 class ARROW_EXPORT NullBuilder : public ArrayBuilder {
  public:
   explicit NullBuilder(MemoryPool* pool = default_memory_pool()) : ArrayBuilder(pool) {}
+  explicit NullBuilder(const std::shared_ptr<DataType>& type,
+                       MemoryPool* pool = default_memory_pool())
+      : NullBuilder(pool) {}
 
   /// \brief Append the specified number of null elements
   Status AppendNulls(int64_t length) final {