Fixed over-consumption / wastage of memory in IResearch address_table (#10)

* Fixed over consumption / wastage of memory in IResearch address_table
* Fixed review comments

Author: k0ushal

core/formats/columnstore2.cpp

@@ -1321,6 +1321,10 @@ void column::Prepare(doc_id_t key) {
     prev_ = pend_;
     pend_ = key;
     docs_writer_.push_back(key);
+
+    //  TODO:
+    //  push_back() can fail. Inspect the return
+    //  value and handle failures.
     addr_table_.push_back(data_.stream.file_pointer());
   }
 }
@@ -1332,7 +1336,13 @@ void column::reset() {
 
   [[maybe_unused]] const bool res = docs_writer_.erase(pend_);
   IRS_ASSERT(res);
-  data_.stream.truncate(addr_table_.back());
+
+  //  TODO:
+  //  back()/push_back() can fail. Inspect the return
+  //  values and handle failures.
+  uint64_t lastElem;
+  addr_table_.back(lastElem);
+  data_.stream.truncate(lastElem);
   addr_table_.pop_back();
   pend_ = prev_;
 }
@@ -1345,15 +1355,24 @@ void column::flush_block() {
   auto& data_out = *ctx_.data_out;
   auto& block = blocks_.emplace_back();
   block.addr = data_out.file_pointer();
-  block.last_size = data_.file.length() - addr_table_.back();
+
+  //  TODO:
+  //  addr_table_.back() can fail. Inspect the return
+  //  value and handle failures.
+  uint64_t lastElem;
+  addr_table_.back(lastElem);
+  block.last_size = data_.file.length() - lastElem;
 
   const uint32_t docs_count = addr_table_.size();
-  const uint64_t addr_table_size =
-    math::ceil64(docs_count, packed::BLOCK_SIZE_64);
-  auto* begin = addr_table_.begin();
-  auto* end = begin + addr_table_size;
-  if (auto* it = addr_table_.current(); it != end) {
-    std::memset(it, 0, (end - it) * sizeof(*it));
+  const uint32_t addr_table_size =
+    math::ceil32(docs_count, packed::BLOCK_SIZE_64);
+
+  addr_table_.grow_size(addr_table_size);
+  auto begin = addr_table_.begin();
+  auto end = begin + addr_table_size;
+
+  if (auto it = addr_table_.current(); it != end) {
+    std::fill(it, end, 0);
   }
 
   bool all_equal = !data_.file.length();
@@ -1377,11 +1396,11 @@ void column::flush_block() {
                      (0 == docs_count_ || block.avg == prev_avg_));
     prev_avg_ = block.avg;
   } else {
-    block.bits = packed::maxbits64(begin, end);
+    block.bits = packed::maxbits64(begin, &*end);
     const size_t buf_size =
       packed::bytes_required_64(addr_table_size, block.bits);
     std::memset(ctx_.u64buf, 0, buf_size);
-    packed::pack(begin, end, ctx_.u64buf, block.bits);
+    packed::pack(begin, &*end, ctx_.u64buf, block.bits);
 
     data_out.write_bytes(ctx_.u8buf, buf_size);
     fixed_length_ = false;

core/formats/columnstore2.hpp

@@ -37,7 +37,7 @@ enum class Version : int32_t { kMin = 0, kMax = kMin };
 
 class column final : public irs::column_output {
  public:
-  static constexpr size_t kBlockSize = sparse_bitmap_writer::kBlockSize;
+  static constexpr uint32_t kBlockSize = sparse_bitmap_writer::kBlockSize;
   static_assert(math::is_power2(kBlockSize));
 
   struct context {
@@ -82,54 +82,81 @@ class column final : public irs::column_output {
 public:
   class address_table {
    public:
-    address_table(ManagedTypedAllocator<uint64_t> alloc) : alloc_(alloc) {
-      offsets_ = alloc_.allocate(kBlockSize);
-      offset_ = offsets_;
-    }
+    address_table(ManagedTypedAllocator<uint64_t> alloc) : offsets_(alloc) { }
+
+    ~address_table() = default;
 
-    ~address_table() { alloc_.deallocate(offsets_, kBlockSize); }
+    bool back(uint64_t& backValue) const noexcept {
+      IRS_ASSERT(!offsets_.empty() && current_pos_ > 0);
+      if (offsets_.empty() || 0 == current_pos_)
+        return false;
 
-    uint64_t back() const noexcept {
-      IRS_ASSERT(offsets_ < offset_);
-      return offset_[-1];
+      backValue = offsets_[current_pos_ - 1];
+      return true;
     }
 
+    //  TODO:
+    //  The present users of the push/pop functions
+    //  assume that these functions will always succeed,
+    //  so they don't check for failures.
     bool push_back(uint64_t offset) noexcept {
-      IRS_ASSERT(offset_ < offsets_ + kBlockSize);
-      if (!(offset_ < offsets_ + kBlockSize))
+      IRS_ASSERT(current_pos_ < column::kBlockSize);
+
+      if (!(current_pos_ < column::kBlockSize))
         return false;
 
-      *offset_++ = offset;
+      if (offsets_.size() <= static_cast<decltype(offsets_)::size_type>(current_pos_))
+        offsets_.resize(current_pos_ * 2 + 1);
+
+      offsets_[current_pos_++] = offset;
       return true;
     }
 
-    bool pop_back() noexcept {
-      IRS_ASSERT(offsets_ < offset_);
-      if (!(offsets_ < offset_))
+    bool pop_back() {
+      IRS_ASSERT(!offsets_.empty() && current_pos_ > 0);
+      if (offsets_.empty() || 0 == current_pos_)
         return false;
 
-      --offset_;
+      --current_pos_;
       return true;
     }
 
     uint32_t size() const noexcept {
-      return static_cast<uint32_t>(offset_ - offsets_);
+      return current_pos_;
+    }
+
+    //  Allocates more space if required.
+    //  Doesn't shrink if max_elem < size.
+    bool grow_size(uint32_t max_elems) {
+      IRS_ASSERT(max_elems <= column::kBlockSize);
+      if (!(max_elems <= column::kBlockSize))
+        return false;
+
+      if (max_elems > static_cast<uint32_t>(offsets_.size()))
+        offsets_.resize(max_elems);
+
+      return true;
     }
 
-    bool empty() const noexcept { return offset_ == offsets_; }
+    bool empty() const noexcept { return 0 == current_pos_; }
+
+    bool full() const noexcept { return current_pos_ == column::kBlockSize; }
 
-    bool full() const noexcept { return offset_ == offsets_ + kBlockSize; }
+    void reset() noexcept {
+      offsets_.clear();
+      current_pos_ = 0;
+    }
 
-    void reset() noexcept { offset_ = offsets_; }
+    uint64_t* begin() noexcept { return &offsets_[0]; }
 
-    uint64_t* begin() noexcept { return offsets_; }
-    uint64_t* current() noexcept { return offset_; }
-    uint64_t* end() noexcept { return offsets_ + kBlockSize; }
+    //  current() points to the next location in the table
+    //  where new data will be added.
+    //  Call grow_size() to allocate space before using current()
+    uint64_t* current() noexcept { return &offsets_[current_pos_]; }
 
    private:
-    ManagedTypedAllocator<uint64_t> alloc_;
-    uint64_t* offsets_{nullptr};
-    uint64_t* offset_{nullptr};
+    std::vector<uint64_t, ManagedTypedAllocator<uint64_t>> offsets_;
+    uint32_t current_pos_ { 0 };
   };
 
 private:

tests/formats/address_table_tests.cpp

@@ -100,17 +100,20 @@ TEST_F(address_table_tests, write_using_current_ptr) {
 
     column::address_table addr_table(alloc);
 
+    uint64_t elemCount = 60;
     //  add data via push_back.
     auto data = 0;
-    for (uint64_t i = 0; i < 60; i++) {
+    for (uint64_t i = 0; i < elemCount; i++) {
         ASSERT_TRUE(addr_table.push_back(data++));
     }
 
     //  add data via current() ptr.
     //  column::flush_block() uses address_table
     //  in this way.
-    auto curr = addr_table.current();
     uint64_t addr_table_size = 40900;
+    ASSERT_TRUE(addr_table.grow_size(addr_table.size() + addr_table_size));
+
+    auto curr = addr_table.current();
     auto end = curr + addr_table_size;
 
     while (curr != end) {
@@ -148,15 +151,17 @@ TEST_F(address_table_tests, max_size_allowed) {
     //  coz address_table pre-allocates memory to
     //  accommodate kBlockSize elements.
     ASSERT_EQ(addr_table.size(), 0);
-    ASSERT_EQ(addr_table.end(), addr_table.begin() + column::kBlockSize);
+
     ASSERT_EQ(addr_table.current(), addr_table.begin() + 0);
 
-    for (uint64_t i = 0; i < 1234; i++) {
+    ASSERT_TRUE(addr_table.grow_size(column::kBlockSize));
+
+    for (uint64_t i = 0; i < column::kBlockSize; i++) {
+        ASSERT_FALSE(addr_table.full());
         ASSERT_TRUE(addr_table.push_back(i));
     }
-
-    ASSERT_EQ(addr_table.end(), addr_table.begin() + column::kBlockSize);
-    ASSERT_EQ(addr_table.current(), addr_table.begin() + 1234);
+    ASSERT_TRUE(addr_table.full());
+    ASSERT_EQ(addr_table.current(), addr_table.begin() + column::kBlockSize);
 }
 
 //  full
@@ -166,13 +171,13 @@ TEST_F(address_table_tests, empty_full_and_reset) {
 
     ASSERT_FALSE(addr_table.full());
     for (uint64_t i = 0; i < column::kBlockSize; i++) {
+        ASSERT_FALSE(addr_table.full());
         ASSERT_TRUE(addr_table.push_back(i));
     }
     ASSERT_TRUE(addr_table.full());
 
     addr_table.pop_back();
     ASSERT_FALSE(addr_table.full());
-    ASSERT_EQ(addr_table.current() + 1, addr_table.end());
 
     addr_table.reset();
     ASSERT_FALSE(addr_table.full());