apacheGH-43060: [C++] Limit buffer size in BufferedInputStream::SetBu…

…fferSize with raw_read_bound (apache#43064)

### Rationale for this change

See apache#43060 . This patch optimize memory-usage for buffering

### What changes are included in this PR?

Limit the `buffer_size_` used by `SetBufferSize`

### Are these changes tested?

Yes

### Are there any user-facing changes?

No

* GitHub Issue: apache#43060

Authored-by: mwish <maplewish117@gmail.com>
Signed-off-by: mwish <maplewish117@gmail.com>

cpp/src/arrow/io/buffered.cc

@@ -282,13 +282,20 @@ class BufferedInputStream::Impl : public BufferedBase {
  return raw_pos_ - bytes_buffered_;
  }
 
+ // Resize internal read buffer. Note that the internal buffer-size
+ // should be not larger than the raw_read_bound_.
  Status SetBufferSize(int64_t new_buffer_size) {
  if (new_buffer_size <= 0) {
  return Status::Invalid("Buffer size should be positive");
  }
  if ((buffer_pos_ + bytes_buffered_) >= new_buffer_size) {
  return Status::Invalid("Cannot shrink read buffer if buffered data remains");
  }
+ if (raw_read_bound_ >= 0) {
+ // No need to reserve space for more than the total remaining number of bytes.
+ new_buffer_size = std::min(new_buffer_size,
+ bytes_buffered_ + (raw_read_bound_ - raw_read_total_));
+ }
  return ResizeBuffer(new_buffer_size);
  }
 
@@ -433,6 +440,8 @@ class BufferedInputStream::Impl : public BufferedBase {
  private:
  std::shared_ptr<InputStream> raw_;
  int64_t raw_read_total_;
+ // a bound on the maximum number of bytes to read from the raw input stream.
+ // The default -1 indicates that it is unbounded
  int64_t raw_read_bound_;
 
  // Number of remaining bytes in the buffer, to be reduced on each read from

cpp/src/arrow/io/buffered_test.cc

@@ -329,7 +329,8 @@ class TestBufferedInputStream : public FileTestFixture<BufferedInputStream> {
  local_pool_ = MemoryPool::CreateDefault();
  }
 
- void MakeExample1(int64_t buffer_size, MemoryPool* pool = default_memory_pool()) {
+ void MakeExample1(int64_t buffer_size, MemoryPool* pool = default_memory_pool(),
+ int64_t raw_read_bound = -1) {
  test_data_ = kExample1;
 
  ASSERT_OK_AND_ASSIGN(auto file_out, FileOutputStream::Open(path_));
@@ -338,7 +339,8 @@ class TestBufferedInputStream : public FileTestFixture<BufferedInputStream> {
 
  ASSERT_OK_AND_ASSIGN(auto file_in, ReadableFile::Open(path_));
  raw_ = file_in;
- ASSERT_OK_AND_ASSIGN(buffered_, BufferedInputStream::Create(buffer_size, pool, raw_));
+ ASSERT_OK_AND_ASSIGN(
+ buffered_, BufferedInputStream::Create(buffer_size, pool, raw_, raw_read_bound));
  }
 
  protected:
@@ -472,6 +474,23 @@ TEST_F(TestBufferedInputStream, SetBufferSize) {
  ASSERT_OK(buffered_->SetBufferSize(5));
 }
 
+// GH-43060: Internal buffer should not greater than the
+// bytes could buffer.
+TEST_F(TestBufferedInputStream, BufferSizeLimit) {
+ {
+ // Buffer size should not exceeds raw_read_bound
+ MakeExample1(/*buffer_size=*/100000, default_memory_pool(), /*raw_read_bound=*/15);
+ EXPECT_EQ(15, buffered_->buffer_size());
+ }
+ {
+ // Set a buffer size after read.
+ MakeExample1(/*buffer_size=*/10, default_memory_pool(), /*raw_read_bound=*/15);
+ ASSERT_OK(buffered_->Read(10));
+ ASSERT_OK(buffered_->SetBufferSize(/*new_buffer_size=*/100000));
+ EXPECT_EQ(5, buffered_->buffer_size());
+ }
+}
+
 class TestBufferedInputStreamBound : public ::testing::Test {
  public:
  void SetUp() { CreateExample(/*bounded=*/true); }