buffer.cc

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "arrow/buffer.h"

#include <algorithm>
#include <cstdint>
#include <utility>

#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/logging.h"
#include "arrow/util/slice_util_internal.h"
#include "arrow/util/string.h"

namespace arrow {

Result<std::shared_ptr<Buffer>> Buffer::CopySlice(const int64_t start,
                                                  const int64_t nbytes,
                                                  MemoryPool* pool) const {
  // Sanity checks
  ARROW_CHECK_LE(start, size_);
  ARROW_CHECK_LE(nbytes, size_ - start);
  DCHECK_GE(nbytes, 0);

  ARROW_ASSIGN_OR_RAISE(auto new_buffer, AllocateResizableBuffer(nbytes, pool));
  std::memcpy(new_buffer->mutable_data(), data() + start, static_cast<size_t>(nbytes));
  return std::move(new_buffer);
}

namespace {

Status CheckBufferSlice(const Buffer& buffer, int64_t offset, int64_t length) {
  return internal::CheckSliceParams(buffer.size(), offset, length, "buffer");
}

Status CheckBufferSlice(const Buffer& buffer, int64_t offset) {
  if (ARROW_PREDICT_FALSE(offset < 0)) {
    // Avoid UBSAN in subtraction below
    return Status::IndexError("Negative buffer slice offset");
  }
  return CheckBufferSlice(buffer, offset, buffer.size() - offset);
}

}  // namespace

Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer>& buffer,
                                                int64_t offset) {
  RETURN_NOT_OK(CheckBufferSlice(*buffer, offset));
  return SliceBuffer(buffer, offset);
}

Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer>& buffer,
                                                int64_t offset, int64_t length) {
  RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length));
  return SliceBuffer(buffer, offset, length);
}

Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe(
    const std::shared_ptr<Buffer>& buffer, int64_t offset) {
  RETURN_NOT_OK(CheckBufferSlice(*buffer, offset));
  return SliceMutableBuffer(buffer, offset);
}

Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe(
    const std::shared_ptr<Buffer>& buffer, int64_t offset, int64_t length) {
  RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length));
  return SliceMutableBuffer(buffer, offset, length);
}

std::string Buffer::ToHexString() {
  return HexEncode(data(), static_cast<size_t>(size()));
}

bool Buffer::Equals(const Buffer& other, const int64_t nbytes) const {
  return this == &other || (size_ >= nbytes && other.size_ >= nbytes &&
                            (data_ == other.data_ ||
                             !memcmp(data_, other.data_, static_cast<size_t>(nbytes))));
}

bool Buffer::Equals(const Buffer& other) const {
  return this == &other || (size_ == other.size_ &&
                            (data_ == other.data_ ||
                             !memcmp(data_, other.data_, static_cast<size_t>(size_))));
}

std::string Buffer::ToString() const {
  return std::string(reinterpret_cast<const char*>(data_), static_cast<size_t>(size_));
}

void Buffer::CheckMutable() const { DCHECK(is_mutable()) << "buffer not mutable"; }

void Buffer::CheckCPU() const {
  DCHECK(is_cpu()) << "not a CPU buffer (device: " << device()->ToString() << ")";
}

Result<std::shared_ptr<io::RandomAccessFile>> Buffer::GetReader(
    std::shared_ptr<Buffer> buf) {
  return buf->memory_manager_->GetBufferReader(buf);
}

Result<std::shared_ptr<io::OutputStream>> Buffer::GetWriter(std::shared_ptr<Buffer> buf) {
  if (!buf->is_mutable()) {
    return Status::Invalid("Expected mutable buffer");
  }
  return buf->memory_manager_->GetBufferWriter(buf);
}

Result<std::shared_ptr<Buffer>> Buffer::Copy(std::shared_ptr<Buffer> source,
                                             const std::shared_ptr<MemoryManager>& to) {
  return MemoryManager::CopyBuffer(source, to);
}

Result<std::unique_ptr<Buffer>> Buffer::CopyNonOwned(
    const Buffer& source, const std::shared_ptr<MemoryManager>& to) {
  return MemoryManager::CopyNonOwned(source, to);
}

Result<std::shared_ptr<Buffer>> Buffer::View(std::shared_ptr<Buffer> source,
                                             const std::shared_ptr<MemoryManager>& to) {
  return MemoryManager::ViewBuffer(source, to);
}

Result<std::shared_ptr<Buffer>> Buffer::ViewOrCopy(
    std::shared_ptr<Buffer> source, const std::shared_ptr<MemoryManager>& to) {
  auto maybe_buffer = MemoryManager::ViewBuffer(source, to);
  if (maybe_buffer.ok()) {
    return maybe_buffer;
  }
  return MemoryManager::CopyBuffer(source, to);
}

class StlStringBuffer : public Buffer {
 public:
  explicit StlStringBuffer(std::string data)
      : Buffer(nullptr, 0), input_(std::move(data)) {
    data_ = reinterpret_cast<const uint8_t*>(input_.c_str());
    size_ = static_cast<int64_t>(input_.size());
    capacity_ = size_;
  }

 private:
  std::string input_;
};

std::shared_ptr<Buffer> Buffer::FromString(std::string data) {
  return std::make_shared<StlStringBuffer>(std::move(data));
}

std::shared_ptr<Buffer> SliceMutableBuffer(const std::shared_ptr<Buffer>& buffer,
                                           const int64_t offset, const int64_t length) {
  return std::make_shared<MutableBuffer>(buffer, offset, length);
}

MutableBuffer::MutableBuffer(const std::shared_ptr<Buffer>& parent, const int64_t offset,
                             const int64_t size)
    : MutableBuffer(reinterpret_cast<uint8_t*>(parent->mutable_address()) + offset,
                    size) {
  DCHECK(parent->is_mutable()) << "Must pass mutable buffer";
  parent_ = parent;
}

Result<std::shared_ptr<Buffer>> AllocateBitmap(int64_t length, MemoryPool* pool) {
  ARROW_ASSIGN_OR_RAISE(auto buf, AllocateBuffer(bit_util::BytesForBits(length), pool));
  // Zero out any trailing bits
  if (buf->size() > 0) {
    buf->mutable_data()[buf->size() - 1] = 0;
  }
  return std::move(buf);
}

Result<std::shared_ptr<Buffer>> AllocateEmptyBitmap(int64_t length, MemoryPool* pool) {
  return AllocateEmptyBitmap(length, kDefaultBufferAlignment, pool);
}

Result<std::shared_ptr<Buffer>> AllocateEmptyBitmap(int64_t length, int64_t alignment,
                                                    MemoryPool* pool) {
  ARROW_ASSIGN_OR_RAISE(auto buf,
                        AllocateBuffer(bit_util::BytesForBits(length), alignment, pool));
  memset(buf->mutable_data(), 0, static_cast<size_t>(buf->size()));
  return std::move(buf);
}

Status AllocateEmptyBitmap(int64_t length, std::shared_ptr<Buffer>* out) {
  return AllocateEmptyBitmap(length).Value(out);
}

Result<std::shared_ptr<Buffer>> ConcatenateBuffers(
    const std::vector<std::shared_ptr<Buffer>>& buffers, MemoryPool* pool) {
  int64_t out_length = 0;
  for (const auto& buffer : buffers) {
    out_length += buffer->size();
  }
  ARROW_ASSIGN_OR_RAISE(auto out, AllocateBuffer(out_length, pool));
  auto out_data = out->mutable_data();
  for (const auto& buffer : buffers) {
    // Passing nullptr to std::memcpy is undefined behavior, so skip empty buffers
    if (buffer->size() != 0) {
      std::memcpy(out_data, buffer->data(), buffer->size());
      out_data += buffer->size();
    }
  }
  return std::move(out);
}

}  // namespace arrow