Add UnlockedBufferedOutputStream (#75) (#76)

Co-authored-by: Hongze Zhang <hongze.zhang@intel.com>

Author: zhouyuan

cpp/src/arrow/io/buffered.cc

@@ -245,6 +245,157 @@ Status BufferedOutputStream::Flush() { return impl_->Flush(); }
 
 std::shared_ptr<OutputStream> BufferedOutputStream::raw() const { return impl_->raw(); }
 
+// ----------------------------------------------------------------------
+// UnlockedBufferedOutputStream implementation
+
+class UnlockedBufferedOutputStream::Impl : public BufferedBase {
+ public:
+  explicit Impl(std::shared_ptr<OutputStream> raw, MemoryPool* pool)
+      : BufferedBase(pool), raw_(std::move(raw)) {}
+
+  Status Close() {
+    if (is_open_) {
+      Status st = FlushUnlocked();
+      is_open_ = false;
+      RETURN_NOT_OK(raw_->Close());
+      return st;
+    }
+    return Status::OK();
+  }
+
+  Status Abort() {
+    if (is_open_) {
+      is_open_ = false;
+      return raw_->Abort();
+    }
+    return Status::OK();
+  }
+
+  Result<int64_t> Tell() const {
+    if (raw_pos_ == -1) {
+      ARROW_ASSIGN_OR_RAISE(raw_pos_, raw_->Tell());
+          DCHECK_GE(raw_pos_, 0);
+    }
+    return raw_pos_ + buffer_pos_;
+  }
+
+  Status Write(const void* data, int64_t nbytes) { return DoWrite(data, nbytes); }
+
+  Status Write(const std::shared_ptr<Buffer>& buffer) {
+    return DoWrite(buffer->data(), buffer->size(), buffer);
+  }
+
+  Status DoWrite(const void* data, int64_t nbytes,
+                 const std::shared_ptr<Buffer>& buffer = nullptr) {
+    if (nbytes < 0) {
+      return Status::Invalid("write count should be >= 0");
+    }
+    if (nbytes == 0) {
+      return Status::OK();
+    }
+    if (nbytes + buffer_pos_ >= buffer_size_) {
+      RETURN_NOT_OK(FlushUnlocked());
+          DCHECK_EQ(buffer_pos_, 0);
+      if (nbytes >= buffer_size_) {
+        // Invalidate cached raw pos
+        raw_pos_ = -1;
+        // Direct write
+        if (buffer) {
+          return raw_->Write(buffer);
+        } else {
+          return raw_->Write(data, nbytes);
+        }
+      }
+    }
+    AppendToBuffer(data, nbytes);
+    return Status::OK();
+  }
+
+  Status FlushUnlocked() {
+    if (buffer_pos_ > 0) {
+      // Invalidate cached raw pos
+      raw_pos_ = -1;
+      RETURN_NOT_OK(raw_->Write(buffer_data_, buffer_pos_));
+      buffer_pos_ = 0;
+    }
+    return Status::OK();
+  }
+
+  Status Flush() {
+    return FlushUnlocked();
+  }
+
+  Result<std::shared_ptr<OutputStream>> Detach() {
+    RETURN_NOT_OK(FlushUnlocked());
+    is_open_ = false;
+    return std::move(raw_);
+  }
+
+  Status SetBufferSize(int64_t new_buffer_size) {
+    if (new_buffer_size <= 0) {
+      return Status::Invalid("Buffer size should be positive");
+    }
+    if (buffer_pos_ >= new_buffer_size) {
+      // If the buffer is shrinking, first flush to the raw OutputStream
+      RETURN_NOT_OK(FlushUnlocked());
+    }
+    return ResizeBuffer(new_buffer_size);
+  }
+
+  std::shared_ptr<OutputStream> raw() const { return raw_; }
+
+ private:
+  std::shared_ptr<OutputStream> raw_;
+};
+
+UnlockedBufferedOutputStream::UnlockedBufferedOutputStream(std::shared_ptr<OutputStream> raw,
+                                                           MemoryPool* pool) {
+  impl_.reset(new Impl(std::move(raw), pool));
+}
+
+Result<std::shared_ptr<UnlockedBufferedOutputStream>> UnlockedBufferedOutputStream::Create(
+    int64_t buffer_size, MemoryPool* pool, std::shared_ptr<OutputStream> raw) {
+  auto result = std::shared_ptr<UnlockedBufferedOutputStream>(
+      new UnlockedBufferedOutputStream(std::move(raw), pool));
+  RETURN_NOT_OK(result->SetBufferSize(buffer_size));
+  return result;
+}
+
+UnlockedBufferedOutputStream::~UnlockedBufferedOutputStream() { internal::CloseFromDestructor(this); }
+
+Status UnlockedBufferedOutputStream::SetBufferSize(int64_t new_buffer_size) {
+  return impl_->SetBufferSize(new_buffer_size);
+}
+
+int64_t UnlockedBufferedOutputStream::buffer_size() const { return impl_->buffer_size(); }
+
+int64_t UnlockedBufferedOutputStream::bytes_buffered() const { return impl_->buffer_pos(); }
+
+Result<std::shared_ptr<OutputStream>> UnlockedBufferedOutputStream::Detach() {
+  return impl_->Detach();
+}
+
+Status UnlockedBufferedOutputStream::Close() { return impl_->Close(); }
+
+Status UnlockedBufferedOutputStream::Abort() { return impl_->Abort(); }
+
+bool UnlockedBufferedOutputStream::closed() const { return impl_->closed(); }
+
+Result<int64_t> UnlockedBufferedOutputStream::Tell() const { return impl_->Tell(); }
+
+Status UnlockedBufferedOutputStream::Write(const void* data, int64_t nbytes) {
+  return impl_->Write(data, nbytes);
+}
+
+Status UnlockedBufferedOutputStream::Write(const std::shared_ptr<Buffer>& data) {
+  return impl_->Write(data);
+}
+
+Status UnlockedBufferedOutputStream::Flush() { return impl_->Flush(); }
+
+std::shared_ptr<OutputStream> UnlockedBufferedOutputStream::raw() const { return impl_->raw(); }
+
+
 // ----------------------------------------------------------------------
 // BufferedInputStream implementation
 

cpp/src/arrow/io/buffered.h

@@ -89,6 +89,61 @@ class ARROW_EXPORT BufferedOutputStream : public OutputStream {
   std::unique_ptr<Impl> impl_;
 };
 
+
+class ARROW_EXPORT UnlockedBufferedOutputStream : public OutputStream {
+ public:
+  ~UnlockedBufferedOutputStream() override;
+
+  /// \brief Create a buffered output stream wrapping the given output stream.
+  /// \param[in] buffer_size the size of the temporary write buffer
+  /// \param[in] pool a MemoryPool to use for allocations
+  /// \param[in] raw another OutputStream
+  /// \return the created BufferedOutputStream
+  static Result<std::shared_ptr<UnlockedBufferedOutputStream>> Create(
+      int64_t buffer_size, MemoryPool* pool, std::shared_ptr<OutputStream> raw);
+
+  /// \brief Resize internal buffer
+  /// \param[in] new_buffer_size the new buffer size
+  /// \return Status
+  Status SetBufferSize(int64_t new_buffer_size);
+
+  /// \brief Return the current size of the internal buffer
+  int64_t buffer_size() const;
+
+  /// \brief Return the number of remaining bytes that have not been flushed to
+  /// the raw OutputStream
+  int64_t bytes_buffered() const;
+
+  /// \brief Flush any buffered writes and release the raw
+  /// OutputStream. Further operations on this object are invalid
+  /// \return the underlying OutputStream
+  Result<std::shared_ptr<OutputStream>> Detach();
+
+  // OutputStream interface
+
+  /// \brief Close the buffered output stream.  This implicitly closes the
+  /// underlying raw output stream.
+  Status Close() override;
+  Status Abort() override;
+  bool closed() const override;
+
+  Result<int64_t> Tell() const override;
+  // Write bytes to the stream. Thread-safe
+  Status Write(const void* data, int64_t nbytes) override;
+  Status Write(const std::shared_ptr<Buffer>& data) override;
+
+  Status Flush() override;
+
+  /// \brief Return the underlying raw output stream.
+  std::shared_ptr<OutputStream> raw() const;
+
+ private:
+  explicit UnlockedBufferedOutputStream(std::shared_ptr<OutputStream> raw, MemoryPool* pool);
+
+  class ARROW_NO_EXPORT Impl;
+  std::unique_ptr<Impl> impl_;
+};
+
 /// \class BufferedInputStream
 /// \brief An InputStream that performs buffered reads from an unbuffered
 /// InputStream, which can mitigate the overhead of many small reads in some