Implement expanding-peek logic, change signature of InputStream::Peak to be able to return Status

cpp/src/arrow/io/buffered-test.cc

@@ -343,8 +343,6 @@ TEST_F(TestBufferedInputStream, BasicOperation) {
 
   // Nothing in the buffer
   ASSERT_EQ(0, buffered_->bytes_buffered());
-  util::string_view peek = buffered_->Peek(10);
-  ASSERT_EQ(0, peek.size());
 
   std::vector<char> buf(test_data_.size());
   int64_t bytes_read;
@@ -355,6 +353,10 @@ TEST_F(TestBufferedInputStream, BasicOperation) {
   // 6 bytes remaining in buffer
   ASSERT_EQ(6, buffered_->bytes_buffered());
 
+  util::string_view peek;
+  ASSERT_OK(buffered_->Peek(6, &peek));
+  ASSERT_EQ(6, peek.size());
+
   // Buffered position is 4
   ASSERT_OK(buffered_->Tell(&stream_position));
   ASSERT_EQ(4, stream_position);
@@ -363,11 +365,6 @@ TEST_F(TestBufferedInputStream, BasicOperation) {
   ASSERT_OK(raw_->Tell(&stream_position));
   ASSERT_EQ(10, stream_position);
 
-  // Peek does not look beyond end of buffer
-  peek = buffered_->Peek(10);
-  ASSERT_EQ(6, peek.size());
-  ASSERT_EQ(0, memcmp(peek.data(), test_data_.data() + 4, 6));
-
   // Reading to end of buffered bytes does not cause any more data to be
   // buffered
   ASSERT_OK(buffered_->Read(6, &bytes_read, buf.data()));
@@ -389,7 +386,7 @@ TEST_F(TestBufferedInputStream, BasicOperation) {
   ASSERT_EQ(test_data_.size(), stream_position);
 
   // Peek at EOF
-  peek = buffered_->Peek(10);
+  ASSERT_OK(buffered_->Peek(10, &peek));
   ASSERT_EQ(0, peek.size());
 
   // Calling Close closes raw_
@@ -484,20 +481,17 @@ TEST_F(TestBufferedInputStreamBound, Basics) {
   std::shared_ptr<Buffer> buffer;
   util::string_view view;
 
-  // Trigger buffering
-  ASSERT_OK(stream_->Read(1, &buffer));
-
   // source is at offset 10
-  view = stream_->Peek(9);
-  ASSERT_EQ(9, view.size());
-  for (int i = 0; i < 9; i++) {
-    ASSERT_EQ(11 + i, view[i]) << i;
+  ASSERT_OK(stream_->Peek(10, &view));
+  ASSERT_EQ(10, view.size());
+  for (int i = 0; i < 10; i++) {
+    ASSERT_EQ(10 + i, view[i]) << i;
   }
 
-  ASSERT_OK(stream_->Read(9, &buffer));
-  ASSERT_EQ(9, buffer->size());
-  for (int i = 0; i < 9; i++) {
-    ASSERT_EQ(11 + i, (*buffer)[i]) << i;
+  ASSERT_OK(stream_->Read(10, &buffer));
+  ASSERT_EQ(10, buffer->size());
+  for (int i = 0; i < 10; i++) {
+    ASSERT_EQ(10 + i, (*buffer)[i]) << i;
   }
 
   ASSERT_OK(stream_->Read(10, &buffer));
@@ -545,17 +539,14 @@ TEST_F(TestBufferedInputStreamBound, LargeFirstPeek) {
   int64_t n = 70;
   ASSERT_GT(n, chunk_size_);
 
-  // Trigger buffering
-  ASSERT_OK(stream_->Read(1, &buffer));
-
-  // source is at offset 11
-  view = stream_->Peek(n - 1);
+  // source is at offset 10
+  ASSERT_OK(stream_->Peek(n, &view));
   ASSERT_EQ(n, static_cast<int>(view.size()));
-  for (int i = 0; i < n - 1; i++) {
-    ASSERT_EQ(11 + i, view[i]) << i;
+  for (int i = 0; i < n; i++) {
+    ASSERT_EQ(10 + i, view[i]) << i;
   }
 
-  view = stream_->Peek(n);
+  ASSERT_OK(stream_->Peek(n, &view));
   ASSERT_EQ(n, static_cast<int>(view.size()));
   for (int i = 0; i < n; i++) {
     ASSERT_EQ(10 + i, view[i]) << i;

cpp/src/arrow/io/buffered.cc

@@ -259,10 +259,39 @@ class BufferedInputStream::Impl : public BufferedBase {
     return ResizeBuffer(new_buffer_size);
   }
 
-  util::string_view Peek(int64_t nbytes) const {
-    int64_t peek_size = std::min(nbytes, bytes_buffered_);
-    return util::string_view(reinterpret_cast<const char*>(buffer_data_ + buffer_pos_),
-                             static_cast<size_t>(peek_size));
+  Status Peek(int64_t nbytes, util::string_view* out) {
+    int64_t total_avail = bytes_buffered_;
+
+    if (raw_read_bound_ > 0) {
+      total_avail += raw_read_bound_ - raw_read_total_;
+    }
+    // Do not try to peek more than the total remaining number of bytes.
+    nbytes = std::min(nbytes, total_avail);
+
+    if (bytes_buffered_ == 0 && nbytes < buffer_size_) {
+      // Pre-buffer for small reads
+      RETURN_NOT_OK(BufferIfNeeded());
+    }
+
+    // Increase the buffer size if needed
+    if (nbytes > buffer_->size() - buffer_pos_) {
+      RETURN_NOT_OK(SetBufferSize(nbytes + buffer_pos_));
+      DCHECK(buffer_->size() - buffer_pos_ >= nbytes);
+    }
+    // Read more data when buffer has insufficient left
+    if (nbytes > bytes_buffered_) {
+      // Read as much as possible to fill the buffer, but not past stream end
+      int64_t read_size = std::min(nbytes - bytes_buffered_, total_avail);
+      int64_t bytes_read = -1;
+      RETURN_NOT_OK(raw_->Read(read_size, &bytes_read,
+                               buffer_->mutable_data() + buffer_pos_ + bytes_buffered_));
+      bytes_buffered_ += bytes_read;
+      raw_read_total_ += bytes_read;
+    }
+    DCHECK(nbytes <= bytes_buffered_);  // Enough bytes available
+    *out = util::string_view(reinterpret_cast<const char*>(buffer_data_ + buffer_pos_),
+                             static_cast<size_t>(nbytes));
+    return Status::OK();
   }
 
   int64_t bytes_buffered() const { return bytes_buffered_; }
@@ -402,8 +431,8 @@ Status BufferedInputStream::Tell(int64_t* position) const {
   return impl_->Tell(position);
 }
 
-util::string_view BufferedInputStream::Peek(int64_t nbytes) const {
-  return impl_->Peek(nbytes);
+Status BufferedInputStream::Peek(int64_t nbytes, util::string_view* out) {
+  return impl_->Peek(nbytes, out);
 }
 
 Status BufferedInputStream::SetBufferSize(int64_t new_buffer_size) {

cpp/src/arrow/io/buffered.h

@@ -127,7 +127,12 @@ class ARROW_EXPORT BufferedInputStream : public InputStream {
   std::shared_ptr<InputStream> raw() const;
 
   // InputStream APIs
-  util::string_view Peek(int64_t nbytes) const override;
+
+  /// \brief Return a zero-copy string view referencing buffered data,
+  /// but do not advance the position of the stream. Buffers data and
+  /// expands the buffer size if necessary
+  Status Peek(int64_t nbytes, util::string_view* out) override;
+
   Status Close() override;
   bool closed() const override;
 

cpp/src/arrow/io/interfaces.cc

@@ -34,8 +34,9 @@ Status InputStream::Advance(int64_t nbytes) {
   return Read(nbytes, &temp);
 }
 
-util::string_view InputStream::Peek(int64_t ARROW_ARG_UNUSED(nbytes)) const {
-  return util::string_view(nullptr, 0);
+Status InputStream::Peek(int64_t ARROW_ARG_UNUSED(nbytes), util::string_view* out) {
+  *out = util::string_view(nullptr, 0);
+  return Status::OK();
 }
 
 bool InputStream::supports_zero_copy() const { return false; }

cpp/src/arrow/io/interfaces.h

@@ -122,12 +122,14 @@ class ARROW_EXPORT InputStream : virtual public FileInterface, virtual public Re
   /// \return Status
   Status Advance(int64_t nbytes);
 
-  /// \brief Return string_view to any buffered bytes, up to the indicated
-  /// number. View becomes invalid after any operation on file. If the
-  /// InputStream is unbuffered, returns 0-length string_view
+  /// \brief Return zero-copy string_view to upcoming bytes in the
+  /// stream but do not modify stream position. View becomes invalid
+  /// after any operation on file. If the InputStream is unbuffered,
+  /// returns 0-length string_view. May trigger buffering if the
+  /// requested size is larger than the number of buffered bytes
   /// \param[in] nbytes the maximum number of bytes to see
   /// \return arrow::util::string_view
-  virtual util::string_view Peek(int64_t nbytes) const;
+  virtual Status Peek(int64_t nbytes, util::string_view* out);
 
   /// \brief Return true if InputStream is capable of zero copy Buffer reads
   virtual bool supports_zero_copy() const;

cpp/src/arrow/io/memory-test.cc

@@ -185,12 +185,14 @@ TEST(TestBufferReader, Peek) {
 
   BufferReader reader(std::make_shared<Buffer>(data));
 
-  auto view = reader.Peek(4);
+  util::string_view view;
+
+  ASSERT_OK(reader.Peek(4, &view));
 
   ASSERT_EQ(4, view.size());
   ASSERT_EQ(data.substr(0, 4), view.to_string());
 
-  view = reader.Peek(20);
+  ASSERT_OK(reader.Peek(20, &view));
   ASSERT_EQ(data.size(), view.size());
   ASSERT_EQ(data, view.to_string());
 }

cpp/src/arrow/io/memory.cc

@@ -301,14 +301,16 @@ Status BufferReader::Tell(int64_t* position) const {
   return Status::OK();
 }
 
-util::string_view BufferReader::Peek(int64_t nbytes) const {
+Status BufferReader::Peek(int64_t nbytes, util::string_view* out) {
   if (!is_open_) {
-    return {};
+    *out = {};
+    return Status::OK();
   }
 
   const int64_t bytes_available = std::min(nbytes, size_ - position_);
-  return util::string_view(reinterpret_cast<const char*>(data_) + position_,
+  *out = util::string_view(reinterpret_cast<const char*>(data_) + position_,
                            static_cast<size_t>(bytes_available));
+  return Status::OK();
 }
 
 bool BufferReader::supports_zero_copy() const { return true; }

cpp/src/arrow/io/memory.h

@@ -152,7 +152,7 @@ class ARROW_EXPORT BufferReader : public RandomAccessFile {
   // Zero copy read
   Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
 
-  util::string_view Peek(int64_t nbytes) const override;
+  Status Peek(int64_t nbytes, util::string_view* out) override;
 
   bool supports_zero_copy() const override;