duckdb 0.8.2-dev4203.0 → 0.8.2-dev4314.0

package/src/duckdb/src/include/duckdb/execution/index/fixed_size_buffer.hpp

@@ -8,36 +8,61 @@
 
 #pragma once
 
+#include "duckdb/storage/partial_block_manager.hpp"
 #include "duckdb/common/typedefs.hpp"
 #include "duckdb/storage/buffer/block_handle.hpp"
 #include "duckdb/storage/buffer/buffer_handle.hpp"
+#include "duckdb/storage/block_manager.hpp"
 
 namespace duckdb {
 
 class FixedSizeAllocator;
 class MetadataWriter;
 
+struct PartialBlockForIndex : public PartialBlock {
+public:
+	PartialBlockForIndex(PartialBlockState state, BlockManager &block_manager,
+	                     const shared_ptr<BlockHandle> &block_handle);
+	~PartialBlockForIndex() override {};
+
+public:
+	void Flush(const idx_t free_space_left) override;
+	void Clear() override;
+	void Merge(PartialBlock &other, idx_t offset, idx_t other_size) override;
+};
+
 //! A fixed-size buffer holds fixed-size segments of data. It lazily deserializes a buffer, if on-disk and not
 //! yet in memory, and it only serializes dirty and non-written buffers to disk during
 //! serialization.
 class FixedSizeBuffer {
+public:
+	//! Constants for fast offset calculations in the bitmask
+	static constexpr idx_t BASE[] = {0x00000000FFFFFFFF, 0x0000FFFF, 0x00FF, 0x0F, 0x3, 0x1};
+	static constexpr uint8_t SHIFT[] = {32, 16, 8, 4, 2, 1};
+
 public:
 	//! Constructor for a new in-memory buffer
 	explicit FixedSizeBuffer(BlockManager &block_manager);
 	//! Constructor for deserializing buffer metadata from disk
-	FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const block_id_t &block_id);
+	FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const idx_t allocation_size,
+	                const BlockPointer &block_pointer);
 
 	//! Block manager of the database instance
 	BlockManager &block_manager;
 
 	//! The number of allocated segments
 	idx_t segment_count;
+	//! The size of allocated memory in this buffer (necessary for copying while pinning)
+	idx_t allocation_size;
 
 	//! True: the in-memory buffer is no longer consistent with a (possibly existing) copy on disk
 	bool dirty;
 	//! True: can be vacuumed after the vacuum operation
 	bool vacuum;
 
+	//! Partial block id and offset
+	BlockPointer block_pointer;
+
 public:
 	//! Returns true, if the buffer is in-memory
 	inline bool InMemory() const {
@@ -45,12 +70,7 @@ public:
 	}
 	//! Returns true, if the block is on-disk
 	inline bool OnDisk() const {
-		return (block_handle != nullptr) && (block_handle->BlockId() < MAXIMUM_BLOCK);
-	}
-	//! Returns the block ID
-	inline block_id_t BlockId() const {
-		D_ASSERT(OnDisk());
-		return block_handle->BlockId();
+		return block_pointer.IsValid();
 	}
 	//! Returns a pointer to the buffer in memory, and calls Deserialize, if the buffer is not in memory
 	inline data_ptr_t Get(const bool dirty_p = true) {
@@ -65,15 +85,25 @@ public:
 	//! Destroys the in-memory buffer and the on-disk block
 	void Destroy();
 	//! Serializes a buffer (if dirty or not on disk)
-	void Serialize();
+	void Serialize(PartialBlockManager &partial_block_manager, const idx_t available_segments, const idx_t segment_size,
+	               const idx_t bitmask_offset);
 	//! Pin a buffer (if not in-memory)
 	void Pin();
+	//! Returns the first free offset in a bitmask
+	uint32_t GetOffset(const idx_t bitmask_count);
 
 private:
 	//! The buffer handle of the in-memory buffer
 	BufferHandle buffer_handle;
 	//! The block handle of the on-disk buffer
 	shared_ptr<BlockHandle> block_handle;
+
+private:
+	//! Returns the maximum non-free offset in a bitmask
+	uint32_t GetMaxOffset(const idx_t available_segments_per_buffer);
+	//! Sets all uninitialized regions of a buffer in the respective partial block allocation
+	void SetUninitializedRegions(PartialBlockForIndex &p_block_for_index, const idx_t segment_size, const idx_t offset,
+	                             const idx_t bitmask_offset);
 };
 
 } // namespace duckdb

package/src/duckdb/src/include/duckdb/main/relation.hpp

@@ -19,8 +19,11 @@
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/main/external_dependencies.hpp"
 #include "duckdb/parser/statement/explain_statement.hpp"
-
-#include <memory>
+#include "duckdb/parser/parsed_expression.hpp"
+#include "duckdb/parser/result_modifier.hpp"
+#include "duckdb/common/unique_ptr.hpp"
+#include "duckdb/common/vector.hpp"
+#include "duckdb/common/helper.hpp"
 
 namespace duckdb {
 struct BoundStatement;
@@ -81,9 +84,12 @@ public:
 	DUCKDB_API shared_ptr<Relation> Project(const string &select_list, const vector<string> &aliases);
 	DUCKDB_API shared_ptr<Relation> Project(const vector<string> &expressions);
 	DUCKDB_API shared_ptr<Relation> Project(const vector<string> &expressions, const vector<string> &aliases);
+	DUCKDB_API shared_ptr<Relation> Project(vector<unique_ptr<ParsedExpression>> expressions,
+	                                        const vector<string> &aliases);
 
 	// FILTER
 	DUCKDB_API shared_ptr<Relation> Filter(const string &expression);
+	DUCKDB_API shared_ptr<Relation> Filter(unique_ptr<ParsedExpression> expression);
 	DUCKDB_API shared_ptr<Relation> Filter(const vector<string> &expressions);
 
 	// LIMIT
@@ -92,6 +98,7 @@ public:
 	// ORDER
 	DUCKDB_API shared_ptr<Relation> Order(const string &expression);
 	DUCKDB_API shared_ptr<Relation> Order(const vector<string> &expressions);
+	DUCKDB_API shared_ptr<Relation> Order(vector<OrderByNode> expressions);
 
 	// JOIN operation
 	DUCKDB_API shared_ptr<Relation> Join(const shared_ptr<Relation> &other, const string &condition,

package/src/duckdb/src/include/duckdb/storage/partial_block_manager.hpp

@@ -25,30 +25,46 @@ class TableCatalogEntry;
 class ViewCatalogEntry;
 class TypeCatalogEntry;
 
+//! Regions that require zero-initialization to avoid leaking memory
+struct UninitializedRegion {
+	idx_t start;
+	idx_t end;
+};
+
+//! The current state of a partial block
 struct PartialBlockState {
+	//! The block id of the partial block
 	block_id_t block_id;
-	//! How big is the block we're writing to. (Total bytes to assign).
+	//! The total bytes that we can assign to this block
 	uint32_t block_size;
-	//! How many bytes of the allocation are used. (offset_in_block of next allocation)
-	uint32_t offset_in_block;
-	//! How many times has the block been used?
+	//! Next allocation offset, and also the current allocation size
+	uint32_t offset;
+	//! The number of times that this block has been used for partial allocations
 	uint32_t block_use_count;
 };
 
 struct PartialBlock {
-	explicit PartialBlock(PartialBlockState state) : state(std::move(state)) {
-	}
+	PartialBlock(PartialBlockState state, BlockManager &block_manager, const shared_ptr<BlockHandle> &block_handle);
 	virtual ~PartialBlock() {
 	}
 
+	//! The current state of a partial block
 	PartialBlockState state;
+	//! All uninitialized regions on this block, we need to zero-initialize them when flushing
+	vector<UninitializedRegion> uninitialized_regions;
+	//! The block manager of the partial block manager
+	BlockManager &block_manager;
+	//! The block handle of the underlying block that this partial block writes to
+	shared_ptr<BlockHandle> block_handle;
 
 public:
-	virtual void AddUninitializedRegion(idx_t start, idx_t end) = 0;
-	virtual void Flush(idx_t free_space_left) = 0;
-	virtual void Clear() {
-	}
-	virtual void Merge(PartialBlock &other, idx_t offset, idx_t other_size);
+	//! Add regions that need zero-initialization to avoid leaking memory
+	void AddUninitializedRegion(const idx_t start, const idx_t end);
+	//! Flush the block to disk and zero-initialize any free space and uninitialized regions
+	virtual void Flush(const idx_t free_space_left) = 0;
+	void FlushInternal(const idx_t free_space_left);
+	virtual void Merge(PartialBlock &other, idx_t offset, idx_t other_size) = 0;
+	virtual void Clear() = 0;
 
 public:
 	template <class TARGET>
@@ -59,13 +75,13 @@ public:
 };
 
 struct PartialBlockAllocation {
-	// BlockManager owning the block_id
+	//! The BlockManager owning the block_id
 	BlockManager *block_manager {nullptr};
-	//! How many bytes assigned to the caller?
+	//! The number of assigned bytes to the caller
 	uint32_t allocation_size;
-	//! State of assigned block.
+	//! The current state of the partial block
 	PartialBlockState state;
-	//! Arbitrary state related to partial block storage.
+	//! Arbitrary state related to the partial block storage
 	unique_ptr<PartialBlock> partial_block;
 };
 
@@ -76,12 +92,12 @@ enum class CheckpointType { FULL_CHECKPOINT, APPEND_TO_TABLE };
 //! In any case, they must share a block manager.
 class PartialBlockManager {
 public:
-	// 20% free / 80% utilization
+	//! 20% free / 80% utilization
 	static constexpr const idx_t DEFAULT_MAX_PARTIAL_BLOCK_SIZE = Storage::BLOCK_SIZE / 5 * 4;
-	// Max number of shared references to a block. No effective limit by default.
+	//! Max number of shared references to a block. No effective limit by default.
 	static constexpr const idx_t DEFAULT_MAX_USE_COUNT = 1u << 20;
-	// No point letting map size grow unbounded. We'll drop blocks with the
-	// least free space first.
+	//! No point letting map size grow unbounded. We'll drop blocks with the
+	//! least free space first.
 	static constexpr const idx_t MAX_BLOCK_MAP_SIZE = 1u << 31;
 
 public:

package/src/duckdb/src/include/duckdb/storage/table/column_checkpoint_state.hpp

@@ -68,37 +68,22 @@ struct PartialBlockForCheckpoint : public PartialBlock {
 	};
 
 public:
-	PartialBlockForCheckpoint(ColumnData &data, ColumnSegment &segment, BlockManager &block_manager,
-	                          PartialBlockState state);
+	PartialBlockForCheckpoint(ColumnData &data, ColumnSegment &segment, PartialBlockState state,
+	                          BlockManager &block_manager);
 	~PartialBlockForCheckpoint() override;
 
 	// We will copy all segment data into the memory of the shared block.
 	// Once the block is full (or checkpoint is complete) we'll invoke Flush().
 	// This will cause the block to get written to storage (via BlockManger::ConvertToPersistent),
 	// and all segments to have their references updated (via ColumnSegment::ConvertToPersistent)
-	BlockManager &block_manager;
-	shared_ptr<BlockHandle> block;
 	vector<PartialColumnSegment> segments;
 
-private:
-	struct UninitializedRegion {
-		idx_t start;
-		idx_t end;
-	};
-	vector<UninitializedRegion> uninitialized_regions;
-
 public:
 	bool IsFlushed();
-
-	void AddUninitializedRegion(idx_t start, idx_t end) override;
-
-	void Flush(idx_t free_space_left) override;
-
-	void Clear() override;
-
+	void Flush(const idx_t free_space_left) override;
 	void Merge(PartialBlock &other, idx_t offset, idx_t other_size) override;
-
 	void AddSegmentToTail(ColumnData &data, ColumnSegment &segment, uint32_t offset_in_block);
+	void Clear() override;
 };
 
 } // namespace duckdb

package/src/duckdb/src/main/relation.cpp

@@ -47,6 +47,11 @@ shared_ptr<Relation> Relation::Project(const vector<string> &expressions) {
 	return Project(expressions, aliases);
 }
 
+shared_ptr<Relation> Relation::Project(vector<unique_ptr<ParsedExpression>> expressions,
+                                       const vector<string> &aliases) {
+	return make_shared<ProjectionRelation>(shared_from_this(), std::move(expressions), aliases);
+}
+
 static vector<unique_ptr<ParsedExpression>> StringListToExpressionList(ClientContext &context,
                                                                        const vector<string> &expressions) {
 	if (expressions.empty()) {
@@ -73,7 +78,11 @@ shared_ptr<Relation> Relation::Filter(const string &expression) {
 	if (expression_list.size() != 1) {
 		throw ParserException("Expected a single expression as filter condition");
 	}
-	return make_shared<FilterRelation>(shared_from_this(), std::move(expression_list[0]));
+	return Filter(std::move(expression_list[0]));
+}
+
+shared_ptr<Relation> Relation::Filter(unique_ptr<ParsedExpression> expression) {
+	return make_shared<FilterRelation>(shared_from_this(), std::move(expression));
 }
 
 shared_ptr<Relation> Relation::Filter(const vector<string> &expressions) {
@@ -95,6 +104,10 @@ shared_ptr<Relation> Relation::Limit(int64_t limit, int64_t offset) {
 
 shared_ptr<Relation> Relation::Order(const string &expression) {
 	auto order_list = Parser::ParseOrderList(expression, context.GetContext()->GetParserOptions());
+	return Order(std::move(order_list));
+}
+
+shared_ptr<Relation> Relation::Order(vector<OrderByNode> order_list) {
 	return make_shared<OrderRelation>(shared_from_this(), std::move(order_list));
 }
 
@@ -110,7 +123,7 @@ shared_ptr<Relation> Relation::Order(const vector<string> &expressions) {
 		}
 		order_list.push_back(std::move(inner_list[0]));
 	}
-	return make_shared<OrderRelation>(shared_from_this(), std::move(order_list));
+	return Order(std::move(order_list));
 }
 
 shared_ptr<Relation> Relation::Join(const shared_ptr<Relation> &other, const string &condition, JoinType type,

package/src/duckdb/src/storage/partial_block_manager.cpp

@@ -2,6 +2,38 @@
 
 namespace duckdb {
 
+//===--------------------------------------------------------------------===//
+// PartialBlock
+//===--------------------------------------------------------------------===//
+
+PartialBlock::PartialBlock(PartialBlockState state, BlockManager &block_manager,
+                           const shared_ptr<BlockHandle> &block_handle)
+    : state(state), block_manager(block_manager), block_handle(block_handle) {
+}
+
+void PartialBlock::AddUninitializedRegion(idx_t start, idx_t end) {
+	uninitialized_regions.push_back({start, end});
+}
+
+void PartialBlock::FlushInternal(const idx_t free_space_left) {
+
+	// ensure that we do not leak any data
+	if (free_space_left > 0 || !uninitialized_regions.empty()) {
+		auto buffer_handle = block_manager.buffer_manager.Pin(block_handle);
+
+		// memset any uninitialized regions
+		for (auto &uninitialized : uninitialized_regions) {
+			memset(buffer_handle.Ptr() + uninitialized.start, 0, uninitialized.end - uninitialized.start);
+		}
+		// memset any free space at the end of the block to 0 prior to writing to disk
+		memset(buffer_handle.Ptr() + Storage::BLOCK_SIZE - free_space_left, 0, free_space_left);
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// PartialBlockManager
+//===--------------------------------------------------------------------===//
+
 PartialBlockManager::PartialBlockManager(BlockManager &block_manager, CheckpointType checkpoint_type,
                                          uint32_t max_partial_block_size, uint32_t max_use_count)
     : block_manager(block_manager), checkpoint_type(checkpoint_type), max_partial_block_size(max_partial_block_size),
@@ -9,9 +41,7 @@ PartialBlockManager::PartialBlockManager(BlockManager &block_manager, Checkpoint
 }
 PartialBlockManager::~PartialBlockManager() {
 }
-//===--------------------------------------------------------------------===//
-// Partial Blocks
-//===--------------------------------------------------------------------===//
+
 PartialBlockAllocation PartialBlockManager::GetBlockAllocation(uint32_t segment_size) {
 	PartialBlockAllocation allocation;
 	allocation.block_manager = &block_manager;
@@ -47,7 +77,7 @@ void PartialBlockManager::AllocateBlock(PartialBlockState &state, uint32_t segme
 		state.block_id = INVALID_BLOCK;
 	}
 	state.block_size = Storage::BLOCK_SIZE;
-	state.offset_in_block = 0;
+	state.offset = 0;
 	state.block_use_count = 1;
 }
 
@@ -60,21 +90,22 @@ bool PartialBlockManager::GetPartialBlock(idx_t segment_size, unique_ptr<Partial
 	partial_block = std::move(entry->second);
 	partially_filled_blocks.erase(entry);
 
-	D_ASSERT(partial_block->state.offset_in_block > 0);
-	D_ASSERT(ValueIsAligned(partial_block->state.offset_in_block));
+	D_ASSERT(partial_block->state.offset > 0);
+	D_ASSERT(ValueIsAligned(partial_block->state.offset));
 	return true;
 }
 
 void PartialBlockManager::RegisterPartialBlock(PartialBlockAllocation &&allocation) {
 	auto &state = allocation.partial_block->state;
+	D_ASSERT(checkpoint_type != CheckpointType::FULL_CHECKPOINT || state.block_id >= 0);
 	if (state.block_use_count < max_use_count) {
-		auto unaligned_size = allocation.allocation_size + state.offset_in_block;
+		auto unaligned_size = allocation.allocation_size + state.offset;
 		auto new_size = AlignValue(unaligned_size);
 		if (new_size != unaligned_size) {
 			// register the uninitialized region so we can correctly initialize it before writing to disk
 			allocation.partial_block->AddUninitializedRegion(unaligned_size, new_size);
 		}
-		state.offset_in_block = new_size;
+		state.offset = new_size;
 		auto new_space_left = state.block_size - new_size;
 		// check if the block is STILL partially filled after adding the segment_size
 		if (new_space_left >= Storage::BLOCK_SIZE - max_partial_block_size) {
@@ -82,7 +113,7 @@ void PartialBlockManager::RegisterPartialBlock(PartialBlockAllocation &&allocati
 			partially_filled_blocks.insert(make_pair(new_space_left, std::move(allocation.partial_block)));
 		}
 	}
-	idx_t free_space = state.block_size - state.offset_in_block;
+	idx_t free_space = state.block_size - state.offset;
 	auto block_to_free = std::move(allocation.partial_block);
 	if (!block_to_free && partially_filled_blocks.size() > MAX_BLOCK_MAP_SIZE) {
 		// Free the page with the least space free.
@@ -98,10 +129,6 @@ void PartialBlockManager::RegisterPartialBlock(PartialBlockAllocation &&allocati
 	}
 }
 
-void PartialBlock::Merge(PartialBlock &other, idx_t offset, idx_t other_size) {
-	throw InternalException("PartialBlock::Merge not implemented for this block type");
-}
-
 void PartialBlockManager::Merge(PartialBlockManager &other) {
 	if (&other == this) {
 		throw InternalException("Cannot merge into itself");
@@ -117,10 +144,10 @@ void PartialBlockManager::Merge(PartialBlockManager &other) {
 			// we can merge this block into an existing block - merge them
 			// merge blocks
 			auto allocation = GetBlockAllocation(used_space);
-			allocation.partial_block->Merge(*e.second, allocation.state.offset_in_block, used_space);
+			allocation.partial_block->Merge(*e.second, allocation.state.offset, used_space);
 
 			// re-register the partial block
-			allocation.state.offset_in_block += used_space;
+			allocation.state.offset += used_space;
 			RegisterPartialBlock(std::move(allocation));
 		} else {
 			// we cannot merge this block - append it directly to the current block manager

package/src/duckdb/src/storage/table/column_checkpoint_state.cpp

@@ -22,9 +22,9 @@ unique_ptr<BaseStatistics> ColumnCheckpointState::GetStatistics() {
 	return std::move(global_stats);
 }
 
-PartialBlockForCheckpoint::PartialBlockForCheckpoint(ColumnData &data, ColumnSegment &segment,
-                                                     BlockManager &block_manager, PartialBlockState state)
-    : PartialBlock(state), block_manager(block_manager), block(segment.block) {
+PartialBlockForCheckpoint::PartialBlockForCheckpoint(ColumnData &data, ColumnSegment &segment, PartialBlockState state,
+                                                     BlockManager &block_manager)
+    : PartialBlock(state, block_manager, segment.block) {
 	AddSegmentToTail(data, segment, 0);
 }
 
@@ -37,24 +37,15 @@ bool PartialBlockForCheckpoint::IsFlushed() {
 	return segments.empty();
 }
 
-void PartialBlockForCheckpoint::AddUninitializedRegion(idx_t start, idx_t end) {
-	uninitialized_regions.push_back({start, end});
-}
+void PartialBlockForCheckpoint::Flush(const idx_t free_space_left) {
 
-void PartialBlockForCheckpoint::Flush(idx_t free_space_left) {
 	if (IsFlushed()) {
 		throw InternalException("Flush called on partial block that was already flushed");
 	}
-	// if we have any free space or uninitialized regions we need to zero-initialize them
-	if (free_space_left > 0 || !uninitialized_regions.empty()) {
-		auto handle = block_manager.buffer_manager.Pin(block);
-		// memset any uninitialized regions
-		for (auto &uninitialized : uninitialized_regions) {
-			memset(handle.Ptr() + uninitialized.start, 0, uninitialized.end - uninitialized.start);
-		}
-		// memset any free space at the end of the block to 0 prior to writing to disk
-		memset(handle.Ptr() + Storage::BLOCK_SIZE - free_space_left, 0, free_space_left);
-	}
+
+	// zero-initialize unused memory
+	FlushInternal(free_space_left);
+
 	// At this point, we've already copied all data from tail_segments
 	// into the page owned by first_segment. We flush all segment data to
 	// disk with the following call.
@@ -63,6 +54,7 @@ void PartialBlockForCheckpoint::Flush(idx_t free_space_left) {
 	if (fetch_new_block) {
 		state.block_id = block_manager.GetFreeBlockId();
 	}
+
 	for (idx_t i = 0; i < segments.size(); i++) {
 		auto &segment = segments[i];
 		segment.data.IncrementVersion();
@@ -71,23 +63,18 @@ void PartialBlockForCheckpoint::Flush(idx_t free_space_left) {
 			D_ASSERT(segment.offset_in_block == 0);
 			segment.segment.ConvertToPersistent(&block_manager, state.block_id);
 			// update the block after it has been converted to a persistent segment
-			block = segment.segment.block;
+			block_handle = segment.segment.block;
 		} else {
 			// subsequent segments are MARKED as persistent - they don't need to be rewritten
-			segment.segment.MarkAsPersistent(block, segment.offset_in_block);
+			segment.segment.MarkAsPersistent(block_handle, segment.offset_in_block);
 			if (fetch_new_block) {
 				// if we fetched a new block we need to increase the reference count to the block
 				block_manager.IncreaseBlockReferenceCount(state.block_id);
 			}
 		}
 	}
-	Clear();
-}
 
-void PartialBlockForCheckpoint::Clear() {
-	uninitialized_regions.clear();
-	block.reset();
-	segments.clear();
+	Clear();
 }
 
 void PartialBlockForCheckpoint::Merge(PartialBlock &other_p, idx_t offset, idx_t other_size) {
@@ -95,13 +82,13 @@ void PartialBlockForCheckpoint::Merge(PartialBlock &other_p, idx_t offset, idx_t
 
 	auto &buffer_manager = block_manager.buffer_manager;
 	// pin the source block
-	auto old_handle = buffer_manager.Pin(other.block);
+	auto old_handle = buffer_manager.Pin(other.block_handle);
 	// pin the target block
-	auto new_handle = buffer_manager.Pin(block);
+	auto new_handle = buffer_manager.Pin(block_handle);
 	// memcpy the contents of the old block to the new block
 	memcpy(new_handle.Ptr() + offset, old_handle.Ptr(), other_size);
 
-	// now copy over all of the segments to the new block
+	// now copy over all segments to the new block
 	// move over the uninitialized regions
 	for (auto &region : other.uninitialized_regions) {
 		region.start += offset;
@@ -113,6 +100,7 @@ void PartialBlockForCheckpoint::Merge(PartialBlock &other_p, idx_t offset, idx_t
 	for (auto &segment : other.segments) {
 		AddSegmentToTail(segment.data, segment.segment, segment.offset_in_block + offset);
 	}
+
 	other.Clear();
 }
 
@@ -120,6 +108,12 @@ void PartialBlockForCheckpoint::AddSegmentToTail(ColumnData &data, ColumnSegment
 	segments.emplace_back(data, segment, offset_in_block);
 }
 
+void PartialBlockForCheckpoint::Clear() {
+	uninitialized_regions.clear();
+	block_handle.reset();
+	segments.clear();
+}
+
 void ColumnCheckpointState::FlushSegment(unique_ptr<ColumnSegment> segment, idx_t segment_size) {
 	D_ASSERT(segment_size <= Storage::BLOCK_SIZE);
 	auto tuple_count = segment->count.load();
@@ -140,7 +134,7 @@ void ColumnCheckpointState::FlushSegment(unique_ptr<ColumnSegment> segment, idx_
 		// non-constant block
 		PartialBlockAllocation allocation = partial_block_manager.GetBlockAllocation(segment_size);
 		block_id = allocation.state.block_id;
-		offset_in_block = allocation.state.offset_in_block;
+		offset_in_block = allocation.state.offset;
 
 		if (allocation.partial_block) {
 			// Use an existing block.
@@ -149,7 +143,7 @@ void ColumnCheckpointState::FlushSegment(unique_ptr<ColumnSegment> segment, idx_
 			// pin the source block
 			auto old_handle = buffer_manager.Pin(segment->block);
 			// pin the target block
-			auto new_handle = buffer_manager.Pin(pstate.block);
+			auto new_handle = buffer_manager.Pin(pstate.block_handle);
 			// memcpy the contents of the old block to the new block
 			memcpy(new_handle.Ptr() + offset_in_block, old_handle.Ptr(), segment_size);
 			pstate.AddSegmentToTail(column_data, *segment, offset_in_block);
@@ -162,8 +156,8 @@ void ColumnCheckpointState::FlushSegment(unique_ptr<ColumnSegment> segment, idx_
 				segment->Resize(Storage::BLOCK_SIZE);
 			}
 			D_ASSERT(offset_in_block == 0);
-			allocation.partial_block = make_uniq<PartialBlockForCheckpoint>(
-			    column_data, *segment, *allocation.block_manager, allocation.state);
+			allocation.partial_block = make_uniq<PartialBlockForCheckpoint>(column_data, *segment, allocation.state,
+			                                                                *allocation.block_manager);
 		}
 		// Writer will decide whether to reuse this block.
 		partial_block_manager.RegisterPartialBlock(std::move(allocation));