- Added multiple performance optimizations to ByteBuffer class.

- Added sequential write checking for variable-size elements.
- Added heuristics for reserve count in variable-size parsing for read-only buffers.
- Added unittests.

Author: avalerio-tkd

src/processing/byte_buffer.cpp

@@ -20,6 +20,7 @@
 #include "bytes_utils.h"
 #include "exceptions.h"
 #include <algorithm>
+#include <cmath>
 #include <cstring>
 #include <limits>
 
@@ -81,15 +82,15 @@ void ByteBuffer::InitializeFromSpan() {
     // Variable-size layout stores [u32 size][element value] back-to-back.
     // Single pass validates shape and captures per-element prefix offsets.
     offsets_.clear();
-    // TODO: consider a heuristic reserve estimate to reduce offsets_ reallocations.
+    offsets_.reserve(EstimateOffsetsReserveCountFromSample(elements_span_));
     size_t cursor = 0;
     while (cursor < elements_span_.size()) {
-        if (elements_span_.size() - cursor < 4) {
+        if (elements_span_.size() - cursor < kSizePrefixBytes) {
             throw InvalidInputException("Malformed variable-size buffer: truncated length prefix");
         }
         offsets_.push_back(cursor);
         const size_t current_element_size = ReadSizeAt(elements_span_, cursor);
-        cursor += 4;
+        cursor += kSizePrefixBytes;
         if (elements_span_.size() - cursor < current_element_size) {
             throw InvalidInputException("Malformed variable-size buffer: truncated element payload");
         }
@@ -100,6 +101,46 @@ void ByteBuffer::InitializeFromSpan() {
     num_elements_ = offsets_.size();
 }
 
+size_t ByteBuffer::EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t> bytes) {
+    if (bytes.empty())
+        return 0;
+
+    // Sample the first 10 elements to estimate the total element count.
+    size_t cursor = 0;
+    size_t sampled_elements = 0;
+    while (cursor < bytes.size() && sampled_elements < 10) {
+        if (bytes.size() - cursor < kSizePrefixBytes) {
+            throw InvalidInputException("Malformed variable-size buffer: truncated length prefix");
+        }
+        const size_t current_element_size = ReadSizeAt(bytes, cursor);
+        cursor += kSizePrefixBytes;
+        if (bytes.size() - cursor < current_element_size) {
+            throw InvalidInputException("Malformed variable-size buffer: truncated element payload");
+        }
+        cursor += current_element_size;
+        ++sampled_elements;
+    }
+
+    if (sampled_elements == 0 || cursor == 0)
+        return 0;
+
+    // If sampling consumed the full buffer (<= sample window), we already know the exact count.
+    if (cursor == bytes.size())
+        return sampled_elements;
+
+    // Estimate total element count from sample density:
+    //   (sampled_elements / sampled_bytes) * total_bytes.
+    // - sampled_elements / sampled_bytes gives "elements per byte" in the sampled prefix,
+    // - then multiplying by total_bytes extrapolates a full-buffer estimate.
+    const long double estimated =
+        (static_cast<long double>(bytes.size()) * static_cast<long double>(sampled_elements)) /
+        static_cast<long double>(cursor);
+    const size_t estimated_count = static_cast<size_t>(std::ceil(estimated));
+    const size_t estimated_with_headroom =
+        static_cast<size_t>(std::ceil(static_cast<long double>(estimated_count) * 1.1L));
+    return std::max(estimated_with_headroom, sampled_elements);
+}
+
 // -----------------------------------------------------------------------------
 // Span reader methods
 // -----------------------------------------------------------------------------
@@ -130,7 +171,7 @@ tcb::span<const uint8_t> ByteBuffer::GetElement(size_t position) const {
         throw InvalidInputException("Element position has not been written yet");
     }
     const size_t element_size = ReadSizeAt(elements_span_, offset);
-    return elements_span_.subspan(offset + 4, element_size);
+    return elements_span_.subspan(offset + kSizePrefixBytes, element_size);
 }
 
 // -----------------------------------------------------------------------------
@@ -186,7 +227,7 @@ void ByteBuffer::InitializeForWriteBuffer(size_t variable_size_reserved_bytes_hi
     // Reserve write_buffer to at least the size of the prefix [u32 size] bytes for all elements,
     //  and use a larger reserved-bytes hint if given as a best guess to reduce reallocations.
     // write_buffer is not initialized to anything since we will be appending to it during SetElement, just reserving capacity.
-    const size_t min_required_prefix_bytes = num_elements_ * static_cast<size_t>(sizeof(uint32_t));
+    const size_t min_required_prefix_bytes = num_elements_ * kSizePrefixBytes;
     const size_t variable_size_reserved_bytes = std::max(variable_size_reserved_bytes_hint, min_required_prefix_bytes);
     write_buffer_.clear();
     write_buffer_.reserve(variable_size_reserved_bytes);
@@ -195,6 +236,9 @@ void ByteBuffer::InitializeForWriteBuffer(size_t variable_size_reserved_bytes_hi
     offsets_.clear();
     offsets_.resize(num_elements_, kUnsetVariableElementOffset);
 
+    // next_expected_sequential_position_ is initialized to 0 for sequential write checking.
+    next_expected_sequential_position_ = 0;
+
     // elements_span_ is re-bound to the write buffer.
     RebindSpanToWriteBuffer();
 }
@@ -239,7 +283,16 @@ void ByteBuffer::SetElement(size_t position, tcb::span<const uint8_t> element) {
     const size_t offset = write_buffer_.size();
     offsets_[position] = offset;
     append_u32_le(write_buffer_, static_cast<uint32_t>(element.size()));
-    write_buffer_.insert(write_buffer_.end(), element.begin(), element.end());
+    write_buffer_.insert(write_buffer_.end(), element.begin(), element.end());  // Appends at the end of the buffer.
+
+    // Update next_expected_sequential_position_ for sequential write checking.
+    if (next_expected_sequential_position_ != kUnsetVariableElementOffset) {
+        if (position == next_expected_sequential_position_) {
+            next_expected_sequential_position_ += 1;
+        } else {
+            next_expected_sequential_position_ = kUnsetVariableElementOffset;
+        }
+    }
 
     RebindSpanToWriteBuffer();
 }
@@ -259,49 +312,51 @@ std::vector<uint8_t> ByteBuffer::FinalizeAndTakeBuffer() {
         return std::move(write_buffer_);
     }
 
-    // Variable-size: validate all offsets and records first, and detect whether
-    // the buffer is already sequential/unfragmented.
-    bool is_sequential = true;
-    size_t current_offset = 0;
+    // For variable-size when all elements were written exactly once and in sequential order,
+    // we can skip out-of-order or fragmentation checks.  This is the fast path.
+    // This is the most common behavior when writing elements in single threaded mode.
+    if (next_expected_sequential_position_ == num_elements_) {
+        if (num_elements_ > 0) {
+            const size_t last_element_offset = offsets_[num_elements_ - 1];
+            const size_t last_element_size = ReadSizeAt(elements_span_, last_element_offset);
+            const size_t logical_size = last_element_offset + kSizePrefixBytes + last_element_size;
+            if (logical_size != write_buffer_.size()) {
+                throw InvalidInputException("FinalizeAndTakeBuffer: trailing bytes detected beyond last element");
+            }
+        }
+        write_buffer_finalized_ = true;
+        return std::move(write_buffer_);
+    }
+
+    // For variable-size, when elements are written out of order, assume the buffer is fragmented and potentially with orphaned bytes
+    // The buffer is validated and rebuilt into a compact buffer in one pass.
+    std::vector<uint8_t> result;
+    result.reserve(write_buffer_.size());
     for (size_t i = 0; i < num_elements_; ++i) {
         const size_t element_offset = offsets_[i];
         if (element_offset == kUnsetVariableElementOffset) {
             throw InvalidInputException("Cannot finalize variable-size buffer: not all elements were written");
         }
-        if (element_offset > write_buffer_.size() || (write_buffer_.size() - element_offset) < 4) {
+        if (element_offset > write_buffer_.size() || (write_buffer_.size() - element_offset) < kSizePrefixBytes) {
             throw InvalidInputException("Cannot finalize variable-size buffer: invalid element offset");
         }
 
         const size_t element_size = ReadSizeAt(elements_span_, element_offset);
-        if (element_size > (write_buffer_.size() - element_offset - 4)) {
+        if (element_size > (write_buffer_.size() - element_offset - kSizePrefixBytes)) {
             throw InvalidInputException("Cannot finalize variable-size buffer: malformed element payload");
         }
 
-        if (element_offset != current_offset) {
-            is_sequential = false;
-        }
-        current_offset += 4 + element_size;
-    }
-    // Detect trailing/orphan bytes in write_buffer_ and force rewrite if present.
-    if (current_offset != write_buffer_.size()) {
-        is_sequential = false;
+        const size_t record_size = kSizePrefixBytes + element_size;
+        result.insert(
+            result.end(),
+            write_buffer_.data() + element_offset,
+            write_buffer_.data() + element_offset + record_size);
     }
 
-    // If sequential, transfer ownership of write_buffer_ directly.
-    if (is_sequential) {
-        write_buffer_finalized_ = true;
-        return std::move(write_buffer_);
-    }
+    // Defrag path returns a new buffer; release the original fragmented write buffer.
+    write_buffer_.clear();
+    write_buffer_.shrink_to_fit();
 
-    // Fragmented: rebuild a compact buffer in element order.
-    std::vector<uint8_t> result;
-    result.reserve(current_offset);
-    for (size_t i = 0; i < num_elements_; ++i) {
-        const size_t element_offset = offsets_[i];
-        const size_t elem_size = ReadSizeAt(elements_span_, element_offset);
-        const uint8_t* start = write_buffer_.data() + element_offset;
-        result.insert(result.end(), start, start + 4 + elem_size);
-    }
     write_buffer_finalized_ = true;
     return result;
 }

src/processing/byte_buffer.h

@@ -58,6 +58,10 @@ class ByteBuffer {
     std::vector<uint8_t> FinalizeAndTakeBuffer();
 
 protected:
+    // Helper for reserve heuristics in variable-size parsing.
+    static size_t EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t> bytes);
+
+    // Helper for calculating the offset of an element by position.
     size_t CalculateOffsetOfElement(size_t position) const;
 
     // Variables for span elements reading
@@ -73,6 +77,11 @@ class ByteBuffer {
     std::vector<uint8_t> write_buffer_;
     bool write_buffer_finalized_ = false;
 
+    // Variable for sequential variable-size writes.
+    // Tracks next expected position for sequential variable-size writes.
+    // Value is invalidated to kUnsetVariableElementOffset once order is violated.
+    size_t next_expected_sequential_position_ = 0;
+
 private:
     // Initialization methods for read-only buffer
     void InitializeFromSpan();
@@ -83,5 +92,6 @@ class ByteBuffer {
 };
 
 inline constexpr size_t kUnsetVariableElementOffset = std::numeric_limits<size_t>::max();
+inline constexpr size_t kSizePrefixBytes = sizeof(uint32_t);  // [u32 size] prefix for variable-size elements.
 
 } // namespace dbps::processing

src/processing/byte_buffer_test.cpp

@@ -21,6 +21,7 @@
 #include <vector>
 #include <cstdint>
 
+#include "bytes_utils.h"
 #include "exceptions.h"
 
 using dbps::processing::ByteBuffer;
@@ -44,12 +45,16 @@ class ByteBufferTestProxy : public ByteBuffer {
         : ByteBuffer(num_elements, reserved_bytes_hint, use_reserve_hint) {}
     ByteBufferTestProxy(size_t num_elements, size_t element_size)
         : ByteBuffer(num_elements, element_size) {}
+    using ByteBuffer::EstimateOffsetsReserveCountFromSample;
 
     size_t GetNumElements() const { return num_elements_; }
     bool GetHasFixedSizedElements() const { return has_fixed_sized_elements_; }
     size_t GetElementSize() const { return element_size_; }
     const std::vector<size_t>& GetOffsets() const { return offsets_; }
     const std::vector<uint8_t>& GetWriteBuffer() const { return write_buffer_; }
+    void AppendTrailingBytesForTest(tcb::span<const uint8_t> bytes) {
+        write_buffer_.insert(write_buffer_.end(), bytes.begin(), bytes.end());
+    }
 };
 
 namespace {
@@ -123,6 +128,60 @@ TEST(ByteBufferTest, ConstructVariableSize_ValidEncodedBuffer_InitializesExpecte
     EXPECT_EQ(buffer.GetOffsets()[1], 9u);  // 4 bytes length prefix + 5 bytes first payload.
 }
 
+TEST(ByteBufferTest, EstimateOffsetsReserveCountFromSample_MultipleCases) {
+    const auto make_variable_size_bytes = [](const std::vector<size_t>& sizes) {
+        std::vector<uint8_t> bytes;
+        for (size_t idx = 0; idx < sizes.size(); ++idx) {
+            append_u32_le(bytes, static_cast<uint32_t>(sizes[idx]));
+            for (size_t j = 0; j < sizes[idx]; ++j) {
+                bytes.push_back(static_cast<uint8_t>((idx + j) & 0xFF));
+            }
+        }
+        return bytes;
+    };
+
+    // Empty buffer.
+    {
+        const std::vector<size_t> sizes = {};
+        const std::vector<uint8_t> bytes = make_variable_size_bytes(sizes);
+        const size_t estimated = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t>(bytes));
+        EXPECT_EQ(estimated, 0u);
+    }
+
+    // Buffer with less than sample size (5 vs 10): exact count, no extrapolation/headroom needed.
+    {
+        const std::vector<size_t> sizes = {1u, 2u, 50u, 4u, 7u};
+        const std::vector<uint8_t> bytes = make_variable_size_bytes(sizes);
+        const size_t estimated = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t>(bytes));
+        EXPECT_EQ(estimated, sizes.size());
+    }
+
+    // Buffer with uniform record sizes. Base estimate is exact, then +10% headroom is applied.
+    {
+        // 30 elements, each 6 bytes long.
+        const size_t num_elements = 30u;
+        const size_t payload_size = 6u;
+        std::vector<uint8_t> bytes;
+        for (size_t i = 0; i < num_elements; ++i) {
+            append_u32_le(bytes, static_cast<uint32_t>(payload_size));
+            for (size_t j = 0; j < payload_size; ++j) {
+                bytes.push_back(static_cast<uint8_t>((i + j) & 0xFF));
+            }
+        }
+        const size_t estimate = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t>(bytes));
+        EXPECT_EQ(estimate, 33u);  //33 = 30 elements + 10% headroom.
+    }
+
+    // Buffer to estimate intentionally overshoots the true element count.
+    {
+        // First 10 elements are very small, tail elements are very large:
+        const std::vector<size_t> sizes = {1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 100u, 100u};
+        const std::vector<uint8_t> bytes = make_variable_size_bytes(sizes);
+        const size_t estimate = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span<const uint8_t>(bytes));
+        EXPECT_GT(estimate, sizes.size());
+    }
+}
+
 TEST(ByteBufferTest, GetElement_VariableSize_ReturnsExpectedPayload) {
     // [len=5]["ABCDE"][len=7]["1234567"]
     std::vector<uint8_t> bytes = {
@@ -293,6 +352,7 @@ TEST(ByteBufferTest, VariableSizeWrite_ExactHint_NoReallocationAndExactUsedSize)
 
     EXPECT_EQ(buffer.GetWriteBuffer().size(), exact_hint_bytes);
     EXPECT_EQ(buffer.GetWriteBuffer().capacity(), initial_capacity);
+    EXPECT_EQ(buffer.GetWriteBuffer().capacity(), exact_hint_bytes);
     EXPECT_EQ(buffer.GetWriteBuffer().data(), initial_data_ptr);
 
     for (size_t i = 0; i < num_elements; ++i) {
@@ -302,6 +362,10 @@ TEST(ByteBufferTest, VariableSizeWrite_ExactHint_NoReallocationAndExactUsedSize)
             EXPECT_EQ(value[j], payloads[i][j]);
         }
     }
+
+    const uint8_t* const data_ptr_before_finalize = buffer.GetWriteBuffer().data();
+    std::vector<uint8_t> final_buffer = buffer.FinalizeAndTakeBuffer();
+    EXPECT_EQ(final_buffer.data(), data_ptr_before_finalize);  // Same allocation: finalize returned write_buffer_ as-is.
 }
 
 TEST(ByteBufferTest, VariableSizeWrite_ExceedsHint_ReallocatesBuffer) {
@@ -373,6 +437,26 @@ TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_Sequential_ReturnsAsIs)
     EXPECT_EQ(final_buffer.data(), data_ptr_before);  // Same allocation (moved, not copied).
 }
 
+TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_SequentialWithTrailingBytes_Throws) {
+    ByteBufferTestProxy buffer(3u, 64u, true);
+    std::vector<uint8_t> first = {0x10, 0x11};
+    std::vector<uint8_t> second = {0x20, 0x21, 0x22};
+    std::vector<uint8_t> third = {0x30};
+
+    buffer.SetElement(0, tcb::span<const uint8_t>(first));
+    buffer.SetElement(1, tcb::span<const uint8_t>(second));
+    buffer.SetElement(2, tcb::span<const uint8_t>(third));
+
+    const size_t expected_trimmed_size =
+        (4u + first.size()) + (4u + second.size()) + (4u + third.size());
+
+    std::vector<uint8_t> trailing = {0xEE, 0xEF, 0xF0};
+    buffer.AppendTrailingBytesForTest(tcb::span<const uint8_t>(trailing));
+    EXPECT_EQ(buffer.GetWriteBuffer().size(), expected_trimmed_size + trailing.size());
+
+    EXPECT_THROW((void)buffer.FinalizeAndTakeBuffer(), InvalidInputException);
+}
+
 TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_OutOfOrder_Defragments) {
     ByteBufferTestProxy buffer(3u, 64u, true);
     std::vector<uint8_t> first = {0x10, 0x11};
@@ -484,6 +568,17 @@ TEST(ByteBufferTest, SetElement_FixedSize_OutOfOrderAndOverwrite_ReturnsLatestVa
     EXPECT_EQ(e1[1], v1_first[1]);
     EXPECT_EQ(e2[0], v2_second[0]);
     EXPECT_EQ(e2[1], v2_second[1]);
+
+    const uint8_t* const data_ptr_before_finalize = buffer.GetWriteBuffer().data();
+    std::vector<uint8_t> final_buffer = buffer.FinalizeAndTakeBuffer();
+    EXPECT_EQ(final_buffer.data(), data_ptr_before_finalize);  // Fixed-size finalize should move write_buffer_ directly.
+    ASSERT_EQ(final_buffer.size(), 6u);
+    EXPECT_EQ(final_buffer[0], v0_second[0]);
+    EXPECT_EQ(final_buffer[1], v0_second[1]);
+    EXPECT_EQ(final_buffer[2], v1_first[0]);
+    EXPECT_EQ(final_buffer[3], v1_first[1]);
+    EXPECT_EQ(final_buffer[4], v2_second[0]);
+    EXPECT_EQ(final_buffer[5], v2_second[1]);
 }
 
 TEST(ByteBufferTest, SetElement_FixedSize_WrongPayloadSize_Throws) {