updated

Author: metsw24-max

cpp/src/arrow/buffer_builder.h

@@ -36,10 +36,6 @@
 
 namespace arrow {
 
-static inline bool BufferBuilderCapacityTooLarge(int64_t capacity) {
-  return ARROW_PREDICT_FALSE(capacity > std::numeric_limits<int64_t>::max() - 63);
-}
-
 // ----------------------------------------------------------------------
 // Buffer builder classes
 
@@ -79,12 +75,6 @@ class ARROW_EXPORT BufferBuilder {
   /// shrinking the builder.
   /// \return Status
   Status Resize(const int64_t new_capacity, bool shrink_to_fit = true) {
-    if (ARROW_PREDICT_FALSE(new_capacity < 0)) {
-      return Status::Invalid("Resize: negative capacity");
-    }
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(new_capacity))) {
-      return Status::CapacityError("Resize: capacity overflow");
-    }
     if (buffer_ == NULLPTR) {
       ARROW_ASSIGN_OR_RAISE(buffer_,
                             AllocateResizableBuffer(new_capacity, alignment_, pool_));
@@ -110,16 +100,10 @@ class ARROW_EXPORT BufferBuilder {
             internal::AddWithOverflow(size_, additional_bytes, &min_capacity))) {
       return Status::CapacityError("Reserve: capacity overflow");
     }
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(min_capacity))) {
-      return Status::CapacityError("Reserve: capacity overflow");
-    }
     if (min_capacity <= capacity_) {
       return Status::OK();
     }
     int64_t requested_capacity = GrowByFactor(capacity_, min_capacity);
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(requested_capacity))) {
-      return Status::CapacityError("Reserve: capacity overflow");
-    }
     return Resize(requested_capacity, false);
   }
 
@@ -342,9 +326,6 @@ class TypedBufferBuilder<
             new_capacity, static_cast<int64_t>(sizeof(T)), &num_bytes))) {
       return Status::CapacityError("Resize: capacity overflow");
     }
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(num_bytes))) {
-      return Status::CapacityError("Resize: capacity overflow");
-    }
     return bytes_builder_.Resize(num_bytes, shrink_to_fit);
   }
 
@@ -357,9 +338,6 @@ class TypedBufferBuilder<
             additional_elements, static_cast<int64_t>(sizeof(T)), &num_bytes))) {
       return Status::CapacityError("Reserve: size overflow");
     }
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(num_bytes))) {
-      return Status::CapacityError("Reserve: size overflow");
-    }
     return bytes_builder_.Reserve(num_bytes);
   }
 
@@ -525,9 +503,6 @@ class TypedBufferBuilder<bool> {
             internal::AddWithOverflow(bit_length_, additional_elements, &min_length))) {
       return Status::CapacityError("Reserve: capacity overflow");
     }
-    if (ARROW_PREDICT_FALSE(BufferBuilderCapacityTooLarge(min_length))) {
-      return Status::CapacityError("Reserve: capacity overflow");
-    }
     return Resize(min_length, false);
   }
 

cpp/src/arrow/memory_pool.cc

@@ -69,6 +69,17 @@ namespace {
 
 constexpr char kDefaultBackendEnvVar[] = "ARROW_DEFAULT_MEMORY_POOL";
 constexpr char kDebugMemoryEnvVar[] = "ARROW_DEBUG_MEMORY_POOL";
+constexpr int64_t kMaxBufferCapacity = std::numeric_limits<int64_t>::max() - 63;
+
+Status ValidateBufferCapacity(int64_t capacity, const char* operation) {
+  if (ARROW_PREDICT_FALSE(capacity < 0)) {
+    return Status::Invalid(operation, ": negative capacity");
+  }
+  if (ARROW_PREDICT_FALSE(capacity > kMaxBufferCapacity)) {
+    return Status::CapacityError(operation, ": capacity overflow");
+  }
+  return Status::OK();
+}
 
 enum class MemoryPoolBackend : uint8_t { System, Jemalloc, Mimalloc };
 
@@ -920,9 +931,7 @@ class PoolBuffer final : public ResizableBuffer {
   }
 
   Status Reserve(const int64_t capacity) override {
-    if (capacity < 0) {
-      return Status::Invalid("Negative buffer capacity: ", capacity);
-    }
+    RETURN_NOT_OK(ValidateBufferCapacity(capacity, "Reserve"));
     uint8_t* ptr = mutable_data();
     if (!ptr || capacity > capacity_) {
       ARROW_ASSIGN_OR_RAISE(int64_t new_capacity, RoundCapacity(capacity));
@@ -938,9 +947,7 @@ class PoolBuffer final : public ResizableBuffer {
   }
 
   Status Resize(const int64_t new_size, bool shrink_to_fit = true) override {
-    if (ARROW_PREDICT_FALSE(new_size < 0)) {
-      return Status::Invalid("Negative buffer resize: ", new_size);
-    }
+    RETURN_NOT_OK(ValidateBufferCapacity(new_size, "Resize"));
     uint8_t* ptr = mutable_data();
     if (ptr && shrink_to_fit && new_size <= size_) {
       // Buffer is non-null and is not growing, so shrink to the requested size without
@@ -984,9 +991,7 @@ class PoolBuffer final : public ResizableBuffer {
 
  private:
   static Result<int64_t> RoundCapacity(int64_t capacity) {
-    if (capacity > std::numeric_limits<int64_t>::max() - 63) {
-      return Status::OutOfMemory("capacity too large");
-    }
+    DCHECK_LE(capacity, kMaxBufferCapacity);
     return bit_util::RoundUpToMultipleOf64(capacity);
   }
 
@@ -1026,6 +1031,7 @@ Result<std::unique_ptr<ResizableBuffer>> AllocateResizableBuffer(const int64_t s
 Result<std::unique_ptr<ResizableBuffer>> AllocateResizableBuffer(const int64_t size,
                                                                  const int64_t alignment,
                                                                  MemoryPool* pool) {
+  RETURN_NOT_OK(ValidateBufferCapacity(size, "AllocateResizableBuffer"));
   return ResizePoolBuffer<std::unique_ptr<ResizableBuffer>>(
       PoolBuffer::MakeUnique(pool, alignment), size);
 }