ARROW-8361: [C++] Add Result<T> APIs to Buffer methods and functions

Closes #6863 from pitrou/ARROW-8361-buffer-result-apis

Authored-by: Antoine Pitrou <antoine@python.org>
Signed-off-by: Benjamin Kietzman <bengilgit@gmail.com>

c_glib/arrow-glib/buffer.cpp

@@ -386,10 +386,9 @@ garrow_buffer_copy(GArrowBuffer *buffer,
                    GError **error)
 {
   auto arrow_buffer = garrow_buffer_get_raw(buffer);
-  std::shared_ptr<arrow::Buffer> arrow_copied_buffer;
-  auto status = arrow_buffer->Copy(start, size, &arrow_copied_buffer);
-  if (garrow_error_check(error, status, "[buffer][copy]")) {
-    return garrow_buffer_new_raw(&arrow_copied_buffer);
+  auto maybe_copied_buffer = arrow_buffer->CopySlice(start, size);
+  if (garrow::check(error, maybe_copied_buffer, "[buffer][copy]")) {
+    return garrow_buffer_new_raw(&(*maybe_copied_buffer));
   } else {
     return NULL;
   }
@@ -567,9 +566,10 @@ GArrowResizableBuffer *
 garrow_resizable_buffer_new(gint64 initial_size,
                             GError **error)
 {
-  std::shared_ptr<arrow::ResizableBuffer> arrow_buffer;
-  auto status = arrow::AllocateResizableBuffer(initial_size, &arrow_buffer);
-  if (garrow_error_check(error, status, "[resizable-buffer][new]")) {
+  auto maybe_buffer = arrow::AllocateResizableBuffer(initial_size);
+  if (garrow::check(error, maybe_buffer, "[resizable-buffer][new]")) {
+    auto arrow_buffer = std::shared_ptr<arrow::ResizableBuffer>(
+      *std::move(maybe_buffer));
     return garrow_resizable_buffer_new_raw(&arrow_buffer);
   } else {
     return NULL;

c_glib/arrow-glib/input-stream.cpp

@@ -698,9 +698,7 @@ namespace garrow {
 
     arrow::Result<std::shared_ptr<arrow::Buffer>>
     Read(int64_t n_bytes) override {
-      arrow::MemoryPool *pool = arrow::default_memory_pool();
-      std::shared_ptr<arrow::ResizableBuffer> buffer;
-      RETURN_NOT_OK(AllocateResizableBuffer(pool, n_bytes, &buffer));
+      ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(n_bytes));
 
       std::lock_guard<std::mutex> guard(lock_);
       GError *error = NULL;
@@ -717,7 +715,7 @@ namespace garrow {
         if (n_read_bytes < n_bytes) {
           RETURN_NOT_OK(buffer->Resize(n_read_bytes));
         }
-        return buffer;
+        return std::move(buffer);
       }
     }
 

cpp/src/arrow/array.cc

@@ -253,14 +253,14 @@ Status CleanListOffsets(const Array& offsets, MemoryPool* pool,
       return Status::Invalid("Last list offset should be non-null");
     }
 
-    std::shared_ptr<Buffer> clean_offsets, clean_valid_bits;
-    RETURN_NOT_OK(
-        AllocateBuffer(pool, num_offsets * sizeof(offset_type), &clean_offsets));
+    ARROW_ASSIGN_OR_RAISE(auto clean_offsets,
+                          AllocateBuffer(num_offsets * sizeof(offset_type), pool));
 
     // Copy valid bits, zero out the bit for the final offset
     // XXX why?
-    RETURN_NOT_OK(offsets.null_bitmap()->Copy(0, BitUtil::BytesForBits(num_offsets - 1),
-                                              &clean_valid_bits));
+    ARROW_ASSIGN_OR_RAISE(
+        auto clean_valid_bits,
+        offsets.null_bitmap()->CopySlice(0, BitUtil::BytesForBits(num_offsets - 1)));
     BitUtil::ClearBit(clean_valid_bits->mutable_data(), num_offsets);
     *validity_buf_out = clean_valid_bits;
 
@@ -277,7 +277,7 @@ Status CleanListOffsets(const Array& offsets, MemoryPool* pool,
       clean_raw_offsets[i] = current_offset;
     }
 
-    *offset_buf_out = clean_offsets;
+    *offset_buf_out = std::move(clean_offsets);
   } else {
     *validity_buf_out = offsets.null_bitmap();
     *offset_buf_out = typed_offsets.values();
@@ -1470,7 +1470,7 @@ class NullArrayFactory {
   Status CreateBuffer() {
     ARROW_ASSIGN_OR_RAISE(int64_t buffer_length,
                           GetBufferLength(type_, length_).Finish());
-    RETURN_NOT_OK(AllocateBuffer(pool_, buffer_length, &buffer_));
+    ARROW_ASSIGN_OR_RAISE(buffer_, AllocateBuffer(buffer_length, pool_));
     std::memset(buffer_->mutable_data(), 0, buffer_->size());
     return Status::OK();
   }
@@ -1567,8 +1567,7 @@ class RepeatedArrayFactory {
   Status Visit(const NullType&) { return Status::OK(); }
 
   Status Visit(const BooleanType&) {
-    std::shared_ptr<Buffer> buffer;
-    RETURN_NOT_OK(AllocateBitmap(pool_, length_, &buffer));
+    ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBitmap(length_, pool_));
     BitUtil::SetBitsTo(buffer->mutable_data(), 0, length_,
                        checked_cast<const BooleanScalar&>(scalar_).value);
     out_ = std::make_shared<BooleanArray>(length_, buffer);

cpp/src/arrow/array/builder_adaptive.cc

@@ -52,7 +52,7 @@ Status AdaptiveIntBuilderBase::Resize(int64_t capacity) {
 
   int64_t nbytes = capacity * int_size_;
   if (capacity_ == 0) {
-    RETURN_NOT_OK(AllocateResizableBuffer(pool_, nbytes, &data_));
+    ARROW_ASSIGN_OR_RAISE(data_, AllocateResizableBuffer(nbytes, pool_));
   } else {
     RETURN_NOT_OK(data_->Resize(nbytes));
   }

cpp/src/arrow/array/concatenate.cc

@@ -60,7 +60,7 @@ static Status ConcatenateBitmaps(const std::vector<Bitmap>& bitmaps, MemoryPool*
   for (size_t i = 0; i < bitmaps.size(); ++i) {
     out_length += bitmaps[i].range.length;
   }
-  RETURN_NOT_OK(AllocateBitmap(pool, out_length, out));
+  ARROW_ASSIGN_OR_RAISE(*out, AllocateBitmap(out_length, pool));
   uint8_t* dst = (*out)->mutable_data();
 
   int64_t bitmap_offset = 0;
@@ -101,7 +101,7 @@ static Status ConcatenateOffsets(const BufferVector& buffers, MemoryPool* pool,
   for (size_t i = 0; i < buffers.size(); ++i) {
     out_length += buffers[i]->size() / sizeof(Offset);
   }
-  RETURN_NOT_OK(AllocateBuffer(pool, (out_length + 1) * sizeof(Offset), out));
+  ARROW_ASSIGN_OR_RAISE(*out, AllocateBuffer((out_length + 1) * sizeof(Offset), pool));
   auto dst = reinterpret_cast<Offset*>((*out)->mutable_data());
 
   int64_t elements_length = 0;
@@ -179,21 +179,21 @@ class ConcatenateImpl {
 
   Status Visit(const FixedWidthType& fixed) {
     // handles numbers, decimal128, fixed_size_binary
-    return ConcatenateBuffers(Buffers(1, fixed), pool_, &out_.buffers[1]);
+    return ConcatenateBuffers(Buffers(1, fixed), pool_).Value(&out_.buffers[1]);
   }
 
   Status Visit(const BinaryType&) {
     std::vector<Range> value_ranges;
     RETURN_NOT_OK(ConcatenateOffsets<int32_t>(Buffers(1, sizeof(int32_t)), pool_,
                                               &out_.buffers[1], &value_ranges));
-    return ConcatenateBuffers(Buffers(2, value_ranges), pool_, &out_.buffers[2]);
+    return ConcatenateBuffers(Buffers(2, value_ranges), pool_).Value(&out_.buffers[2]);
   }
 
   Status Visit(const LargeBinaryType&) {
     std::vector<Range> value_ranges;
     RETURN_NOT_OK(ConcatenateOffsets<int64_t>(Buffers(1, sizeof(int64_t)), pool_,
                                               &out_.buffers[1], &value_ranges));
-    return ConcatenateBuffers(Buffers(2, value_ranges), pool_, &out_.buffers[2]);
+    return ConcatenateBuffers(Buffers(2, value_ranges), pool_).Value(&out_.buffers[2]);
   }
 
   Status Visit(const ListType&) {
@@ -240,7 +240,7 @@ class ConcatenateImpl {
 
     if (dictionaries_same) {
       out_.dictionary = in_[0].dictionary;
-      return ConcatenateBuffers(Buffers(1, *fixed), pool_, &out_.buffers[1]);
+      return ConcatenateBuffers(Buffers(1, *fixed), pool_).Value(&out_.buffers[1]);
     } else {
       return Status::NotImplemented("Concat with dictionary unification NYI");
     }

cpp/src/arrow/array/dict_internal.cc

@@ -64,14 +64,13 @@ class DictionaryUnifierImpl : public DictionaryUnifier {
     }
     const ArrayType& values = checked_cast<const ArrayType&>(dictionary);
     if (out != nullptr) {
-      std::shared_ptr<Buffer> result;
-      RETURN_NOT_OK(
-          AllocateBuffer(pool_, dictionary.length() * sizeof(int32_t), &result));
+      ARROW_ASSIGN_OR_RAISE(auto result,
+                            AllocateBuffer(dictionary.length() * sizeof(int32_t), pool_));
       auto result_raw = reinterpret_cast<int32_t*>(result->mutable_data());
       for (int64_t i = 0; i < values.length(); ++i) {
         RETURN_NOT_OK(memo_table_.GetOrInsert(values.GetView(i), &result_raw[i]));
       }
-      *out = result;
+      *out = std::move(result);
     } else {
       for (int64_t i = 0; i < values.length(); ++i) {
         int32_t unused_memo_index;
@@ -198,9 +197,9 @@ Result<std::shared_ptr<Array>> DictionaryArray::Transpose(
   }
 
   // Default path: compute a buffer of transposed indices.
-  std::shared_ptr<Buffer> out_buffer;
-  RETURN_NOT_OK(AllocateBuffer(
-      pool, data_->length * out_index_type.bit_width() * CHAR_BIT, &out_buffer));
+  ARROW_ASSIGN_OR_RAISE(
+      auto out_buffer,
+      AllocateBuffer(data_->length * out_index_type.bit_width() * CHAR_BIT, pool));
 
   // Shift null buffer if the original offset is non-zero
   std::shared_ptr<Buffer> null_bitmap;
@@ -211,8 +210,8 @@ Result<std::shared_ptr<Array>> DictionaryArray::Transpose(
     null_bitmap = data_->buffers[0];
   }
 
-  auto out_data =
-      ArrayData::Make(type, data_->length, {null_bitmap, out_buffer}, data_->null_count);
+  auto out_data = ArrayData::Make(
+      type, data_->length, {null_bitmap, std::move(out_buffer)}, data_->null_count);
   out_data->dictionary = dictionary;
 
 #define TRANSPOSE_IN_OUT_CASE(IN_INDEX_TYPE, OUT_INDEX_TYPE)                 \

cpp/src/arrow/array/dict_internal.h

@@ -97,14 +97,14 @@ struct DictionaryTraits<T, enable_if_has_c_type<T>> {
                                        const MemoTableType& memo_table,
                                        int64_t start_offset,
                                        std::shared_ptr<ArrayData>* out) {
-    std::shared_ptr<Buffer> dict_buffer;
     auto dict_length = static_cast<int64_t>(memo_table.size()) - start_offset;
     // This makes a copy, but we assume a dictionary array is usually small
     // compared to the size of the dictionary-using array.
     // (also, copying the dictionary values is cheap compared to the cost
     //  of building the memo table)
-    RETURN_NOT_OK(
-        AllocateBuffer(pool, TypeTraits<T>::bytes_required(dict_length), &dict_buffer));
+    ARROW_ASSIGN_OR_RAISE(
+        std::shared_ptr<Buffer> dict_buffer,
+        AllocateBuffer(TypeTraits<T>::bytes_required(dict_length), pool));
     memo_table.CopyValues(static_cast<int32_t>(start_offset),
                           reinterpret_cast<c_type*>(dict_buffer->mutable_data()));
 
@@ -128,19 +128,17 @@ struct DictionaryTraits<T, enable_if_base_binary<T>> {
                                        int64_t start_offset,
                                        std::shared_ptr<ArrayData>* out) {
     using offset_type = typename T::offset_type;
-    std::shared_ptr<Buffer> dict_offsets;
-    std::shared_ptr<Buffer> dict_data;
 
     // Create the offsets buffer
     auto dict_length = static_cast<int64_t>(memo_table.size() - start_offset);
-    RETURN_NOT_OK(
-        AllocateBuffer(pool, sizeof(offset_type) * (dict_length + 1), &dict_offsets));
+    ARROW_ASSIGN_OR_RAISE(auto dict_offsets,
+                          AllocateBuffer(sizeof(offset_type) * (dict_length + 1), pool));
     auto raw_offsets = reinterpret_cast<offset_type*>(dict_offsets->mutable_data());
     memo_table.CopyOffsets(static_cast<int32_t>(start_offset), raw_offsets);
 
     // Create the data buffer
     auto values_size = memo_table.values_size();
-    RETURN_NOT_OK(AllocateBuffer(pool, values_size, &dict_data));
+    ARROW_ASSIGN_OR_RAISE(auto dict_data, AllocateBuffer(values_size, pool));
     if (values_size > 0) {
       memo_table.CopyValues(static_cast<int32_t>(start_offset), dict_data->size(),
                             dict_data->mutable_data());
@@ -151,7 +149,8 @@ struct DictionaryTraits<T, enable_if_base_binary<T>> {
     RETURN_NOT_OK(
         ComputeNullBitmap(pool, memo_table, start_offset, &null_count, &null_bitmap));
 
-    *out = ArrayData::Make(type, dict_length, {null_bitmap, dict_offsets, dict_data},
+    *out = ArrayData::Make(type, dict_length,
+                           {null_bitmap, std::move(dict_offsets), std::move(dict_data)},
                            null_count);
 
     return Status::OK();
@@ -168,13 +167,12 @@ struct DictionaryTraits<T, enable_if_fixed_size_binary<T>> {
                                        int64_t start_offset,
                                        std::shared_ptr<ArrayData>* out) {
     const T& concrete_type = internal::checked_cast<const T&>(*type);
-    std::shared_ptr<Buffer> dict_data;
 
     // Create the data buffer
     auto dict_length = static_cast<int64_t>(memo_table.size() - start_offset);
     auto width_length = concrete_type.byte_width();
     auto data_length = dict_length * width_length;
-    RETURN_NOT_OK(AllocateBuffer(pool, data_length, &dict_data));
+    ARROW_ASSIGN_OR_RAISE(auto dict_data, AllocateBuffer(data_length, pool));
     auto data = dict_data->mutable_data();
 
     memo_table.CopyFixedWidthValues(static_cast<int32_t>(start_offset), width_length,
@@ -185,7 +183,8 @@ struct DictionaryTraits<T, enable_if_fixed_size_binary<T>> {
     RETURN_NOT_OK(
         ComputeNullBitmap(pool, memo_table, start_offset, &null_count, &null_bitmap));
 
-    *out = ArrayData::Make(type, dict_length, {null_bitmap, dict_data}, null_count);
+    *out = ArrayData::Make(type, dict_length, {null_bitmap, std::move(dict_data)},
+                           null_count);
     return Status::OK();
   }
 };

cpp/src/arrow/array/diff.cc

@@ -305,9 +305,9 @@ class QuadraticSpaceMyersDiff {
     DCHECK(Done());
 
     int64_t length = edit_count_ + 1;
-    std::shared_ptr<Buffer> insert_buf, run_length_buf;
-    RETURN_NOT_OK(AllocateEmptyBitmap(pool, length, &insert_buf));
-    RETURN_NOT_OK(AllocateBuffer(pool, length * sizeof(int64_t), &run_length_buf));
+    ARROW_ASSIGN_OR_RAISE(auto insert_buf, AllocateEmptyBitmap(length, pool));
+    ARROW_ASSIGN_OR_RAISE(auto run_length_buf,
+                          AllocateBuffer(length * sizeof(int64_t), pool));
     auto run_length = reinterpret_cast<int64_t*>(run_length_buf->mutable_data());
 
     auto index = finish_index_;
@@ -336,9 +336,10 @@ class QuadraticSpaceMyersDiff {
     BitUtil::SetBitTo(insert_buf->mutable_data(), 0, false);
     run_length[0] = endpoint.base - base_begin_;
 
-    return StructArray::Make({std::make_shared<BooleanArray>(length, insert_buf),
-                              std::make_shared<Int64Array>(length, run_length_buf)},
-                             {field("insert", boolean()), field("run_length", int64())});
+    return StructArray::Make(
+        {std::make_shared<BooleanArray>(length, std::move(insert_buf)),
+         std::make_shared<Int64Array>(length, std::move(run_length_buf))},
+        {field("insert", boolean()), field("run_length", int64())});
   }
 
  private:

cpp/src/arrow/array_dict_test.cc

@@ -989,7 +989,7 @@ TEST(TestDictionary, Validate) {
   ASSERT_RAISES(Invalid, arr->ValidateFull());
 
   // Make the data buffer non-null
-  ASSERT_OK(AllocateBuffer(0, &buffers[2]));
+  ASSERT_OK_AND_ASSIGN(buffers[2], AllocateBuffer(0));
   arr = std::make_shared<DictionaryArray>(dict_type, indices, MakeArray(invalid_data));
   ASSERT_RAISES(Invalid, arr->ValidateFull());
 

cpp/src/arrow/array_test.cc

@@ -171,12 +171,11 @@ TEST_F(TestArray, SliceRecomputeNullCount) {
   ASSERT_EQ(5, slice->null_count());
 
   // No bitmap, compute 0
-  std::shared_ptr<Buffer> data;
   const int kBufferSize = 64;
-  ASSERT_OK(AllocateBuffer(pool_, kBufferSize, &data));
+  ASSERT_OK_AND_ASSIGN(auto data, AllocateBuffer(kBufferSize, pool_));
   memset(data->mutable_data(), 0, kBufferSize);
 
-  auto arr = std::make_shared<Int32Array>(16, data, nullptr, -1);
+  auto arr = std::make_shared<Int32Array>(16, std::move(data), nullptr, -1);
   ASSERT_EQ(0, arr->null_count());
 }