perf(processor): zero factorized table blocks as tuples are claimed

src/include/processor/result/factorized_table.h

@@ -28,12 +28,23 @@ struct BlockAppendingInfo {
 // released when this struct goes out of scope.
 class DataBlock {
 public:
-    DataBlock(storage::MemoryManager* mm, uint64_t size);
+    // With zeroOnClaim the buffer is not zeroed up front; claim() zeroes bytes as they are
+    // handed out, so a one-tuple table does not clear a whole page.
+    DataBlock(storage::MemoryManager* mm, uint64_t size, bool zeroOnClaim = false);
     ~DataBlock();
 
     uint8_t* getData() const;
     std::span<uint8_t> getSizedData() const;
     uint8_t* getWritableData() const;
+    // Takes numBytes from the free part of the block and returns their start.
+    uint8_t* claim(uint64_t numBytes) {
+        auto data = getWritableData();
+        freeSize -= numBytes;
+        if (freeSize < numUninitializedBytes) [[unlikely]] {
+            zeroClaimedBytes();
+        }
+        return data;
+    }
     void resetNumTuplesAndFreeSize();
     void resetToZero();
 
@@ -49,7 +60,12 @@ class DataBlock {
     uint64_t freeSize;
 
 private:
+    void zeroClaimedBytes();
+
     std::unique_ptr<storage::MemoryBuffer> block;
+    // Bytes at the end of the block that have been neither zeroed nor written. Never more
+    // than freeSize: claimed bytes are always initialized.
+    uint64_t numUninitializedBytes;
 };
 
 class DataBlockCollection {

src/processor/result/factorized_table.cpp

@@ -14,8 +14,11 @@ using namespace lbug::storage;
 namespace lbug {
 namespace processor {
 
-DataBlock::DataBlock(storage::MemoryManager* mm, uint64_t size) : numTuples{0}, freeSize{size} {
-    block = mm->allocateBuffer(true /* initializeToZero */, size);
+static constexpr uint64_t ZERO_ON_CLAIM_GRANULE = 4096;
+
+DataBlock::DataBlock(storage::MemoryManager* mm, uint64_t size, bool zeroOnClaim)
+    : numTuples{0}, freeSize{size}, numUninitializedBytes{zeroOnClaim ? size : 0} {
+    block = mm->allocateBuffer(!zeroOnClaim /* initializeToZero */, size);
 }
 
 DataBlock::~DataBlock() = default;
@@ -29,12 +32,21 @@ std::span<uint8_t> DataBlock::getSizedData() const {
 uint8_t* DataBlock::getWritableData() const {
     return block->getBuffer().last(freeSize).data();
 }
+void DataBlock::zeroClaimedBytes() {
+    auto size = block->getBuffer().size();
+    auto zeroStart = size - numUninitializedBytes;
+    auto claimedEnd = size - freeSize;
+    auto zeroEnd = std::min<uint64_t>(size,
+        (claimedEnd + ZERO_ON_CLAIM_GRANULE - 1) / ZERO_ON_CLAIM_GRANULE * ZERO_ON_CLAIM_GRANULE);
+    memset(getData() + zeroStart, 0, zeroEnd - zeroStart);
+    numUninitializedBytes = size - zeroEnd;
+}
 void DataBlock::resetNumTuplesAndFreeSize() {
     freeSize = block->getBuffer().size();
     numTuples = 0;
 }
 void DataBlock::resetToZero() {
-    memset(block->getBuffer().data(), 0, block->getBuffer().size());
+    memset(block->getBuffer().data(), 0, block->getBuffer().size() - numUninitializedBytes);
 }
 
 void DataBlock::preventDestruction() {
@@ -52,6 +64,8 @@ void DataBlock::copyTuples(DataBlock* blockToCopyFrom, ft_tuple_idx_t tupleIdxTo
     }
     blockToCopyInto->numTuples += numTuplesToCopy;
     blockToCopyInto->freeSize -= (numTuplesToCopy * numBytesPerTuple);
+    blockToCopyInto->numUninitializedBytes =
+        std::min(blockToCopyInto->numUninitializedBytes, blockToCopyInto->freeSize);
 }
 
 void DataBlockCollection::merge(DataBlockCollection& other) {
@@ -131,18 +145,19 @@ void FactorizedTable::resize(uint64_t numTuples) {
         auto numTuplesToAdd = numTuples - this->numTuples;
         auto numBytesPerTuple = tableSchema.getNumBytesPerTuple();
         while (flatTupleBlockCollection->needAllocation(numTuplesToAdd * numBytesPerTuple)) {
-            auto newBlock = std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize);
+            auto newBlock = std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize,
+                true /* zeroOnClaim */);
             flatTupleBlockCollection->append(std::move(newBlock));
             auto numTuplesToAddInBlock =
                 std::min(static_cast<uint32_t>(numTuplesToAdd), numFlatTuplesPerBlock);
             auto block = flatTupleBlockCollection->getLastBlock();
-            block->freeSize -= numBytesPerTuple * numTuplesToAddInBlock;
+            block->claim(numBytesPerTuple * numTuplesToAddInBlock);
             block->numTuples += numTuplesToAddInBlock;
             numTuplesToAdd -= numTuplesToAddInBlock;
         }
         DASSERT(numTuplesToAdd < numFlatTuplesPerBlock);
         auto block = flatTupleBlockCollection->getLastBlock();
-        block->freeSize -= numBytesPerTuple * numTuplesToAdd;
+        block->claim(numBytesPerTuple * numTuplesToAdd);
         block->numTuples += numTuplesToAdd;
     } else {
         auto numTuplesRemaining = numTuples;
@@ -163,12 +178,12 @@ void FactorizedTable::resize(uint64_t numTuples) {
 uint8_t* FactorizedTable::appendEmptyTuple() {
     auto numBytesPerTuple = tableSchema.getNumBytesPerTuple();
     if (flatTupleBlockCollection->needAllocation(numBytesPerTuple)) {
-        auto newBlock = std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize);
+        auto newBlock =
+            std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize, true /* zeroOnClaim */);
         flatTupleBlockCollection->append(std::move(newBlock));
     }
     auto block = flatTupleBlockCollection->getLastBlock();
-    uint8_t* tuplePtr = block->getWritableData();
-    block->freeSize -= numBytesPerTuple;
+    uint8_t* tuplePtr = block->claim(numBytesPerTuple);
     block->numTuples++;
     numTuples++;
     return tuplePtr;
@@ -393,14 +408,15 @@ std::vector<BlockAppendingInfo> FactorizedTable::allocateFlatTupleBlocks(
     std::vector<BlockAppendingInfo> appendingInfos;
     while (numTuplesToAppend > 0) {
         if (flatTupleBlockCollection->needAllocation(numBytesPerTuple)) {
-            auto newBlock = std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize);
+            auto newBlock = std::make_unique<DataBlock>(memoryManager, flatTupleBlockSize,
+                true /* zeroOnClaim */);
             flatTupleBlockCollection->append(std::move(newBlock));
         }
         auto block = flatTupleBlockCollection->getLastBlock();
         auto numTuplesToAppendInCurBlock =
             std::min(numTuplesToAppend, block->freeSize / numBytesPerTuple);
-        appendingInfos.emplace_back(block->getWritableData(), numTuplesToAppendInCurBlock);
-        block->freeSize -= numTuplesToAppendInCurBlock * numBytesPerTuple;
+        appendingInfos.emplace_back(block->claim(numTuplesToAppendInCurBlock * numBytesPerTuple),
+            numTuplesToAppendInCurBlock);
         block->numTuples += numTuplesToAppendInCurBlock;
         numTuplesToAppend -= numTuplesToAppendInCurBlock;
     }

test/processor/CMakeLists.txt

@@ -1 +1 @@
-add_lbug_test(processor_test intersect_kernels_test.cpp)
+add_lbug_test(processor_test intersect_kernels_test.cpp factorized_table_test.cpp)

test/processor/factorized_table_test.cpp

@@ -0,0 +1,100 @@
+#include <cstring>
+
+#include "common/constants.h"
+#include "gtest/gtest.h"
+#include "main/connection.h"
+#include "main/database.h"
+#include "processor/result/factorized_table.h"
+#include "storage/buffer_manager/memory_manager.h"
+
+using namespace lbug;
+using namespace lbug::processor;
+
+namespace {
+
+bool isFilled(const uint8_t* data, uint64_t size, uint8_t value) {
+    for (auto i = 0u; i < size; i++) {
+        if (data[i] != value) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool isZero(const uint8_t* data, uint64_t size) {
+    return isFilled(data, size, 0);
+}
+
+FactorizedTableSchema makeSchema() {
+    FactorizedTableSchema schema;
+    schema.appendColumn(ColumnSchema(false /* isUnFlat */, 0 /* groupID */, 24));
+    return schema;
+}
+
+void appendFilled(FactorizedTable& table, uint64_t numTuples, uint8_t value) {
+    const auto numBytesPerTuple = table.getTableSchema()->getNumBytesPerTuple();
+    for (auto i = 0u; i < numTuples; i++) {
+        memset(table.appendEmptyTuple(), value, numBytesPerTuple);
+    }
+}
+
+} // namespace
+
+// Flat tuple blocks are zeroed as tuples are claimed rather than when the page is taken, so a
+// recycled page must still hand out zeroed tuples, before and after clear().
+TEST(FactorizedTableTest, ClaimedTuplesAreZeroOnRecycledPages) {
+    main::Database db(":memory:");
+    main::Connection conn(&db);
+    auto* mm = storage::MemoryManager::Get(*conn.getClientContext());
+    {
+        auto dirty = mm->allocateBuffer(false, common::TEMP_PAGE_SIZE);
+        memset(dirty->getBuffer().data(), 0xAB, dirty->getBuffer().size());
+    }
+    FactorizedTable table(mm, makeSchema());
+    const auto numBytesPerTuple = table.getTableSchema()->getNumBytesPerTuple();
+    // Spans several zeroing granules within the first block.
+    const auto numTuples = 3 * 4096 / numBytesPerTuple + 7;
+    for (auto round = 0; round < 2; round++) {
+        for (auto i = 0u; i < numTuples; i++) {
+            auto* tuple = table.appendEmptyTuple();
+            ASSERT_TRUE(isZero(tuple, numBytesPerTuple)) << "round " << round << " tuple " << i;
+            memset(tuple, 0xCD, numBytesPerTuple);
+        }
+        // Grows past the next 4 KB step, so resize itself has to zero.
+        const auto numResized = numTuples + 4096 / numBytesPerTuple + 5;
+        table.resize(numResized);
+        for (auto i = numTuples; i < numResized; i++) {
+            ASSERT_TRUE(isZero(table.getTuple(i), numBytesPerTuple))
+                << "round " << round << " tuple " << i;
+        }
+        table.clear();
+    }
+}
+
+// Merging copies the tuples of one table's last block into the other's last block, past the part
+// that block had zeroed. Later claims must neither zero over the copied tuples nor hand out dirty
+// bytes.
+TEST(FactorizedTableTest, ClaimsAfterMergeKeepCopiedTuples) {
+    main::Database db(":memory:");
+    main::Connection conn(&db);
+    auto* mm = storage::MemoryManager::Get(*conn.getClientContext());
+    FactorizedTable target(mm, makeSchema());
+    FactorizedTable source(mm, makeSchema());
+    const auto numBytesPerTuple = target.getTableSchema()->getNumBytesPerTuple();
+    // The source's tuples end just short of its first 4 KB step, and the target's tuples are
+    // copied after them, across that step.
+    const auto numSourceTuples = (4096 - 8) / numBytesPerTuple;
+    const auto numTargetTuples = 100u;
+    appendFilled(target, numTargetTuples, 0xA1);
+    appendFilled(source, numSourceTuples, 0xB2);
+    target.merge(source);
+    ASSERT_EQ(target.getNumTuples(), numSourceTuples + numTargetTuples);
+    for (auto i = 0u; i < numSourceTuples; i++) {
+        ASSERT_TRUE(isFilled(target.getTuple(i), numBytesPerTuple, 0xB2)) << i;
+    }
+    auto* appended = target.appendEmptyTuple();
+    ASSERT_TRUE(isZero(appended, numBytesPerTuple));
+    for (auto i = numSourceTuples; i < numSourceTuples + numTargetTuples; i++) {
+        ASSERT_TRUE(isFilled(target.getTuple(i), numBytesPerTuple, 0xA1)) << i;
+    }
+}