Allow unbounded peeks in BufferedInputStream

cpp/src/arrow/buffer-builder.h

@@ -148,6 +148,12 @@ class ARROW_EXPORT BufferBuilder {
     capacity_ = size_ = 0;
   }
 
+  /// \brief Set size to a smaller value without modifying builder
+  /// contents. For reusable BufferBuilder classes
+  /// \param[in] position must be non-negative and less than or equal
+  /// to the current length()
+  void Rewind(int64_t position) { size_ = position; }
+
   int64_t capacity() const { return capacity_; }
   int64_t length() const { return size_; }
   const uint8_t* data() const { return data_; }

cpp/src/arrow/buffer.h

@@ -86,7 +86,7 @@ class ARROW_EXPORT Buffer {
     parent_ = parent;
   }
 
-  constexpr uint8_t operator[](std::size_t i) const { return data_[i]; }
+  uint8_t operator[](std::size_t i) const { return data_[i]; }
 
   bool is_mutable() const { return is_mutable_; }
 

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

@@ -453,7 +453,9 @@ TEST_F(TestBufferedInputStream, SetBufferSize) {
 
 class TestBufferedInputStreamBound : public ::testing::Test {
  public:
-  void SetUp() {
+  void SetUp() { CreateExample(/*bounded=*/true); }
+
+  void CreateExample(bool bounded = true) {
     // Create a buffer larger than source size, to check that the
     // stream end is respected
     std::shared_ptr<ResizableBuffer> buf;
@@ -465,7 +467,7 @@ class TestBufferedInputStreamBound : public ::testing::Test {
     source_ = std::make_shared<BufferReader>(buf);
     ASSERT_OK(source_->Advance(stream_offset_));
     ASSERT_OK(BufferedInputStream::Create(chunk_size_, default_memory_pool(), source_,
-                                          &stream_, stream_size_));
+                                          &stream_, bounded ? stream_size_ : -1));
   }
 
  protected:
@@ -565,6 +567,27 @@ TEST_F(TestBufferedInputStreamBound, LargeFirstPeek) {
   }
 }
 
+TEST_F(TestBufferedInputStreamBound, UnboundedPeek) {
+  CreateExample(/*bounded=*/false);
+
+  util::string_view view;
+  ASSERT_OK(stream_->Peek(10, &view));
+  ASSERT_EQ(10, view.size());
+  ASSERT_EQ(50, stream_->bytes_buffered());
+
+  ASSERT_OK(stream_->Peek(246, &view));
+  ASSERT_EQ(246, view.size());
+  ASSERT_EQ(246, stream_->bytes_buffered());
+  ASSERT_EQ(246, stream_->buffer_size());
+
+  // Larger peek returns the same, expands the buffer, but there is no
+  // more data to buffer
+  ASSERT_OK(stream_->Peek(300, &view));
+  ASSERT_EQ(246, view.size());
+  ASSERT_EQ(246, stream_->bytes_buffered());
+  ASSERT_EQ(300, stream_->buffer_size());
+}
+
 TEST_F(TestBufferedInputStreamBound, OneByteReads) {
   std::shared_ptr<Buffer> buffer;
   for (int i = 0; i < stream_size_; ++i) {

cpp/src/arrow/io/buffered.cc

@@ -260,13 +260,10 @@ class BufferedInputStream::Impl : public BufferedBase {
   }
 
   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_;
+    if (raw_read_bound_ >= 0) {
+      // Do not try to peek more than the total remaining number of bytes.
+      nbytes = std::min(nbytes, bytes_buffered_ + (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
@@ -280,13 +277,17 @@ class BufferedInputStream::Impl : public BufferedBase {
     }
     // 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 additional_bytes_to_read = nbytes - bytes_buffered_;
+      if (raw_read_bound_ >= 0) {
+        additional_bytes_to_read =
+            std::min(additional_bytes_to_read, raw_read_bound_ - raw_read_total_);
+      }
       int64_t bytes_read = -1;
-      RETURN_NOT_OK(raw_->Read(read_size, &bytes_read,
+      RETURN_NOT_OK(raw_->Read(additional_bytes_to_read, &bytes_read,
                                buffer_->mutable_data() + buffer_pos_ + bytes_buffered_));
       bytes_buffered_ += bytes_read;
       raw_read_total_ += bytes_read;
+      nbytes = bytes_buffered_;
     }
     DCHECK(nbytes <= bytes_buffered_);  // Enough bytes available
     *out = util::string_view(reinterpret_cast<const char*>(buffer_data_ + buffer_pos_),

cpp/src/arrow/io/interfaces.cc

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

cpp/src/arrow/io/memory.cc

@@ -62,8 +62,6 @@ Status BufferOutputStream::Reset(int64_t initial_capacity, MemoryPool* pool) {
   return Status::OK();
 }
 
-void BufferOutputStream::Clear() { position_ = 0; }
-
 BufferOutputStream::~BufferOutputStream() {
   // This can fail, better to explicitly call close
   if (buffer_) {
@@ -302,10 +300,7 @@ Status BufferReader::Tell(int64_t* position) const {
 }
 
 Status BufferReader::Peek(int64_t nbytes, util::string_view* out) {
-  if (!is_open_) {
-    *out = {};
-    return Status::OK();
-  }
+  RETURN_NOT_OK(CheckClosed());
 
   const int64_t bytes_available = std::min(nbytes, size_ - position_);
   *out = util::string_view(reinterpret_cast<const char*>(data_) + position_,

cpp/src/arrow/io/memory.h

@@ -72,13 +72,6 @@ class ARROW_EXPORT BufferOutputStream : public OutputStream {
 
   int64_t capacity() const { return capacity_; }
 
-  /// \brief Return reference to current in-progress buffer. May
-  /// become invalid after any Write calls
-  const Buffer& buffer() const { return *buffer_; }
-
-  /// \brief Reset position to 0 without modifying capacity or zeroing data
-  void Clear();
-
  private:
   BufferOutputStream();
 

cpp/src/parquet/column_writer.cc

@@ -22,6 +22,7 @@
 #include <memory>
 #include <utility>
 
+#include "arrow/buffer-builder.h"
 #include "arrow/status.h"
 #include "arrow/util/bit-stream-utils.h"
 #include "arrow/util/checked_cast.h"
@@ -367,9 +368,9 @@ class ColumnWriterImpl {
         total_bytes_written_(0),
         total_compressed_bytes_(0),
         closed_(false),
-        fallback_(false) {
-    definition_levels_sink_ = CreateOutputStream(allocator_);
-    repetition_levels_sink_ = CreateOutputStream(allocator_);
+        fallback_(false),
+        definition_levels_sink_(allocator_),
+        repetition_levels_sink_(allocator_) {
     definition_levels_rle_ =
         std::static_pointer_cast<ResizableBuffer>(AllocateBuffer(allocator_, 0));
     repetition_levels_rle_ =
@@ -411,19 +412,19 @@ class ColumnWriterImpl {
   // Write multiple definition levels
   void WriteDefinitionLevels(int64_t num_levels, const int16_t* levels) {
     DCHECK(!closed_);
-    PARQUET_THROW_NOT_OK(definition_levels_sink_->Write(
-        reinterpret_cast<const uint8_t*>(levels), sizeof(int16_t) * num_levels));
+    PARQUET_THROW_NOT_OK(
+        definition_levels_sink_.Append(levels, sizeof(int16_t) * num_levels));
   }
 
   // Write multiple repetition levels
   void WriteRepetitionLevels(int64_t num_levels, const int16_t* levels) {
     DCHECK(!closed_);
-    PARQUET_THROW_NOT_OK(repetition_levels_sink_->Write(
-        reinterpret_cast<const uint8_t*>(levels), sizeof(int16_t) * num_levels));
+    PARQUET_THROW_NOT_OK(
+        repetition_levels_sink_.Append(levels, sizeof(int16_t) * num_levels));
   }
 
   // RLE encode the src_buffer into dest_buffer and return the encoded size
-  int64_t RleEncodeLevels(const Buffer& src_buffer, ResizableBuffer* dest_buffer,
+  int64_t RleEncodeLevels(const void* src_buffer, ResizableBuffer* dest_buffer,
                           int16_t max_level);
 
   // Serialize the buffered Data Pages
@@ -469,8 +470,8 @@ class ColumnWriterImpl {
   // Flag to infer if dictionary encoding has fallen back to PLAIN
   bool fallback_;
 
-  std::shared_ptr<::arrow::io::BufferOutputStream> definition_levels_sink_;
-  std::shared_ptr<::arrow::io::BufferOutputStream> repetition_levels_sink_;
+  ::arrow::BufferBuilder definition_levels_sink_;
+  ::arrow::BufferBuilder repetition_levels_sink_;
 
   std::shared_ptr<ResizableBuffer> definition_levels_rle_;
   std::shared_ptr<ResizableBuffer> repetition_levels_rle_;
@@ -482,13 +483,13 @@ class ColumnWriterImpl {
 
  private:
   void InitSinks() {
-    definition_levels_sink_->Clear();
-    repetition_levels_sink_->Clear();
+    definition_levels_sink_.Rewind(0);
+    repetition_levels_sink_.Rewind(0);
   }
 };
 
 // return the size of the encoded buffer
-int64_t ColumnWriterImpl::RleEncodeLevels(const Buffer& src_buffer,
+int64_t ColumnWriterImpl::RleEncodeLevels(const void* src_buffer,
                                           ResizableBuffer* dest_buffer,
                                           int16_t max_level) {
   // TODO: This only works with due to some RLE specifics
@@ -503,9 +504,8 @@ int64_t ColumnWriterImpl::RleEncodeLevels(const Buffer& src_buffer,
   level_encoder_.Init(Encoding::RLE, max_level, static_cast<int>(num_buffered_values_),
                       dest_buffer->mutable_data() + sizeof(int32_t),
                       static_cast<int>(dest_buffer->size() - sizeof(int32_t)));
-  int encoded =
-      level_encoder_.Encode(static_cast<int>(num_buffered_values_),
-                            reinterpret_cast<const int16_t*>(src_buffer.data()));
+  int encoded = level_encoder_.Encode(static_cast<int>(num_buffered_values_),
+                                      reinterpret_cast<const int16_t*>(src_buffer));
   DCHECK_EQ(encoded, num_buffered_values_);
   reinterpret_cast<int32_t*>(dest_buffer->mutable_data())[0] = level_encoder_.len();
   int64_t encoded_size = level_encoder_.len() + sizeof(int32_t);
@@ -520,13 +520,13 @@ void ColumnWriterImpl::AddDataPage() {
 
   if (descr_->max_definition_level() > 0) {
     definition_levels_rle_size =
-        RleEncodeLevels(definition_levels_sink_->buffer(), definition_levels_rle_.get(),
+        RleEncodeLevels(definition_levels_sink_.data(), definition_levels_rle_.get(),
                         descr_->max_definition_level());
   }
 
   if (descr_->max_repetition_level() > 0) {
     repetition_levels_rle_size =
-        RleEncodeLevels(repetition_levels_sink_->buffer(), repetition_levels_rle_.get(),
+        RleEncodeLevels(repetition_levels_sink_.data(), repetition_levels_rle_.get(),
                         descr_->max_repetition_level());
   }