duckdb 0.10.1-dev9.0 → 0.10.2-dev0.0

package/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp

@@ -1,5 +1,6 @@
 #include "duckdb/execution/operator/persistent/physical_batch_insert.hpp"
-
+#include "duckdb/execution/operator/persistent/batch_memory_manager.hpp"
+#include "duckdb/execution/operator/persistent/batch_task_manager.hpp"
 #include "duckdb/parallel/thread_context.hpp"
 #include "duckdb/storage/data_table.hpp"
 #include "duckdb/storage/table/row_group_collection.hpp"
@@ -28,9 +29,8 @@ PhysicalBatchInsert::PhysicalBatchInsert(LogicalOperator &op, SchemaCatalogEntry
 }
 
 //===--------------------------------------------------------------------===//
-// Sink
+// CollectionMerger
 //===--------------------------------------------------------------------===//
-
 class CollectionMerger {
 public:
 	explicit CollectionMerger(ClientContext &context) : context(context) {
@@ -97,25 +97,40 @@ public:
 enum class RowGroupBatchType : uint8_t { FLUSHED, NOT_FLUSHED };
 struct RowGroupBatchEntry {
 	RowGroupBatchEntry(idx_t batch_idx, unique_ptr<RowGroupCollection> collection_p, RowGroupBatchType type)
-	    : batch_idx(batch_idx), total_rows(collection_p->GetTotalRows()), collection(std::move(collection_p)),
-	      type(type) {
+	    : batch_idx(batch_idx), total_rows(collection_p->GetTotalRows()), unflushed_memory(0),
+	      collection(std::move(collection_p)), type(type) {
+		if (type == RowGroupBatchType::NOT_FLUSHED) {
+			unflushed_memory = collection->GetAllocationSize();
+		}
 	}
 
 	idx_t batch_idx;
 	idx_t total_rows;
+	idx_t unflushed_memory;
 	unique_ptr<RowGroupCollection> collection;
 	RowGroupBatchType type;
 };
 
-class BatchInsertGlobalState : public GlobalSinkState {
+//===--------------------------------------------------------------------===//
+// States
+//===--------------------------------------------------------------------===//
+class BatchInsertTask {
 public:
-	static constexpr const idx_t BATCH_FLUSH_THRESHOLD = LocalStorage::MERGE_THRESHOLD * 3;
+	virtual ~BatchInsertTask() {
+	}
 
+	virtual void Execute(const PhysicalBatchInsert &op, ClientContext &context, GlobalSinkState &gstate_p,
+	                     LocalSinkState &lstate_p) = 0;
+};
+
+class BatchInsertGlobalState : public GlobalSinkState {
 public:
-	explicit BatchInsertGlobalState(DuckTableEntry &table)
-	    : table(table), insert_count(0), optimistically_written(false) {
+	explicit BatchInsertGlobalState(ClientContext &context, DuckTableEntry &table, idx_t initial_memory)
+	    : memory_manager(context, initial_memory), table(table), insert_count(0), optimistically_written(false) {
 	}
 
+	BatchMemoryManager memory_manager;
+	BatchTaskManager<BatchInsertTask> task_manager;
 	mutex lock;
 	DuckTableEntry &table;
 	idx_t insert_count;
@@ -123,140 +138,21 @@ public:
 	idx_t next_start = 0;
 	atomic<bool> optimistically_written;
 
-	void FindMergeCollections(idx_t min_batch_index, optional_idx &merged_batch_index,
-	                          vector<unique_ptr<RowGroupCollection>> &result) {
-		bool merge = false;
-		idx_t start_index = next_start;
-		idx_t current_idx;
-		idx_t total_count = 0;
-		for (current_idx = start_index; current_idx < collections.size(); current_idx++) {
-			auto &entry = collections[current_idx];
-			if (entry.batch_idx >= min_batch_index) {
-				// this entry is AFTER the min_batch_index
-				// we might still find new entries!
-				break;
-			}
-			if (entry.type == RowGroupBatchType::FLUSHED) {
-				// already flushed: cannot flush anything here
-				if (total_count > 0) {
-					merge = true;
-					break;
-				}
-				start_index = current_idx + 1;
-				if (start_index > next_start) {
-					// avoid checking this segment again in the future
-					next_start = start_index;
-				}
-				total_count = 0;
-				continue;
-			}
-			// not flushed - add to set of indexes to flush
-			total_count += entry.total_rows;
-			if (total_count >= BATCH_FLUSH_THRESHOLD) {
-				merge = true;
-				break;
-			}
-		}
-		if (merge && total_count > 0) {
-			D_ASSERT(current_idx > start_index);
-			merged_batch_index = collections[start_index].batch_idx;
-			for (idx_t idx = start_index; idx < current_idx; idx++) {
-				auto &entry = collections[idx];
-				if (!entry.collection || entry.type == RowGroupBatchType::FLUSHED) {
-					throw InternalException("Adding a row group collection that should not be flushed");
-				}
-				result.push_back(std::move(entry.collection));
-				entry.total_rows = total_count;
-				entry.type = RowGroupBatchType::FLUSHED;
-			}
-			if (start_index + 1 < current_idx) {
-				// erase all entries except the first one
-				collections.erase(collections.begin() + start_index + 1, collections.begin() + current_idx);
-			}
-		}
-	}
-
+	static bool ReadyToMerge(idx_t count);
+	void ScheduleMergeTasks(idx_t min_batch_index);
 	unique_ptr<RowGroupCollection> MergeCollections(ClientContext &context,
-	                                                vector<unique_ptr<RowGroupCollection>> merge_collections,
-	                                                OptimisticDataWriter &writer) {
-		D_ASSERT(!merge_collections.empty());
-		CollectionMerger merger(context);
-		for (auto &collection : merge_collections) {
-			merger.AddCollection(std::move(collection));
-		}
-		optimistically_written = true;
-		return merger.Flush(writer);
-	}
-
+	                                                vector<RowGroupBatchEntry> merge_collections,
+	                                                OptimisticDataWriter &writer);
 	void AddCollection(ClientContext &context, idx_t batch_index, idx_t min_batch_index,
 	                   unique_ptr<RowGroupCollection> current_collection,
-	                   optional_ptr<OptimisticDataWriter> writer = nullptr,
-	                   optional_ptr<bool> written_to_disk = nullptr) {
-		if (batch_index < min_batch_index) {
-			throw InternalException(
-			    "Batch index of the added collection (%llu) is smaller than the min batch index (%llu)", batch_index,
-			    min_batch_index);
-		}
-		auto new_count = current_collection->GetTotalRows();
-		auto batch_type =
-		    new_count < Storage::ROW_GROUP_SIZE ? RowGroupBatchType::NOT_FLUSHED : RowGroupBatchType::FLUSHED;
-		if (batch_type == RowGroupBatchType::FLUSHED && writer) {
-			writer->WriteLastRowGroup(*current_collection);
-		}
-		optional_idx merged_batch_index;
-		vector<unique_ptr<RowGroupCollection>> merge_collections;
-		{
-			lock_guard<mutex> l(lock);
-			insert_count += new_count;
-
-			// add the collection to the batch index
-			RowGroupBatchEntry new_entry(batch_index, std::move(current_collection), batch_type);
-
-			auto it = std::lower_bound(
-			    collections.begin(), collections.end(), new_entry,
-			    [&](const RowGroupBatchEntry &a, const RowGroupBatchEntry &b) { return a.batch_idx < b.batch_idx; });
-			if (it != collections.end() && it->batch_idx == new_entry.batch_idx) {
-				throw InternalException(
-				    "PhysicalBatchInsert::AddCollection error: batch index %d is present in multiple "
-				    "collections. This occurs when "
-				    "batch indexes are not uniquely distributed over threads",
-				    batch_index);
-			}
-			collections.insert(it, std::move(new_entry));
-			if (writer) {
-				FindMergeCollections(min_batch_index, merged_batch_index, merge_collections);
-			}
-		}
-		if (!merge_collections.empty()) {
-			// merge together the collections
-			D_ASSERT(writer);
-			auto final_collection = MergeCollections(context, std::move(merge_collections), *writer);
-			if (written_to_disk) {
-				*written_to_disk = true;
-			}
-			// add the merged-together collection to the set of batch indexes
-			{
-				lock_guard<mutex> l(lock);
-				RowGroupBatchEntry new_entry(merged_batch_index.GetIndex(), std::move(final_collection),
-				                             RowGroupBatchType::FLUSHED);
-				auto it = std::lower_bound(collections.begin(), collections.end(), new_entry,
-				                           [&](const RowGroupBatchEntry &a, const RowGroupBatchEntry &b) {
-					                           return a.batch_idx < b.batch_idx;
-				                           });
-				if (it->batch_idx != merged_batch_index.GetIndex()) {
-					throw InternalException("Merged batch index was no longer present in collection");
-				}
-				it->collection = std::move(new_entry.collection);
-			}
-		}
-	}
+	                   optional_ptr<OptimisticDataWriter> writer = nullptr);
 };
 
 class BatchInsertLocalState : public LocalSinkState {
 public:
 	BatchInsertLocalState(ClientContext &context, const vector<LogicalType> &types,
 	                      const vector<unique_ptr<Expression>> &bound_defaults)
-	    : default_executor(context, bound_defaults), written_to_disk(false) {
+	    : default_executor(context, bound_defaults) {
 		insert_chunk.Initialize(Allocator::Get(context), types);
 	}
 
@@ -266,7 +162,6 @@ public:
 	TableAppendState current_append_state;
 	unique_ptr<RowGroupCollection> current_collection;
 	optional_ptr<OptimisticDataWriter> writer;
-	bool written_to_disk;
 
 	void CreateNewCollection(DuckTableEntry &table, const vector<LogicalType> &insert_types) {
 		auto &table_info = table.GetStorage().info;
@@ -274,10 +169,199 @@ public:
 		current_collection = make_uniq<RowGroupCollection>(table_info, block_manager, insert_types, MAX_ROW_ID);
 		current_collection->InitializeEmpty();
 		current_collection->InitializeAppend(current_append_state);
-		written_to_disk = false;
 	}
 };
 
+//===--------------------------------------------------------------------===//
+// Merge Task
+//===--------------------------------------------------------------------===//
+class MergeCollectionTask : public BatchInsertTask {
+public:
+	MergeCollectionTask(vector<RowGroupBatchEntry> merge_collections_p, idx_t merged_batch_index)
+	    : merge_collections(std::move(merge_collections_p)), merged_batch_index(merged_batch_index) {
+	}
+
+	vector<RowGroupBatchEntry> merge_collections;
+	idx_t merged_batch_index;
+
+	void Execute(const PhysicalBatchInsert &op, ClientContext &context, GlobalSinkState &gstate_p,
+	             LocalSinkState &lstate_p) override {
+		auto &gstate = gstate_p.Cast<BatchInsertGlobalState>();
+		auto &lstate = lstate_p.Cast<BatchInsertLocalState>();
+		// merge together the collections
+		D_ASSERT(lstate.writer);
+		auto final_collection = gstate.MergeCollections(context, std::move(merge_collections), *lstate.writer);
+		// add the merged-together collection to the set of batch indexes
+		lock_guard<mutex> l(gstate.lock);
+		RowGroupBatchEntry new_entry(merged_batch_index, std::move(final_collection), RowGroupBatchType::FLUSHED);
+		auto it = std::lower_bound(
+		    gstate.collections.begin(), gstate.collections.end(), new_entry,
+		    [&](const RowGroupBatchEntry &a, const RowGroupBatchEntry &b) { return a.batch_idx < b.batch_idx; });
+		if (it->batch_idx != merged_batch_index) {
+			throw InternalException("Merged batch index was no longer present in collection");
+		}
+		it->collection = std::move(new_entry.collection);
+	}
+};
+
+struct BatchMergeTask {
+	explicit BatchMergeTask(idx_t start_index) : start_index(start_index), end_index(0), total_count(0) {
+	}
+
+	idx_t start_index;
+	idx_t end_index;
+	idx_t total_count;
+};
+
+bool BatchInsertGlobalState::ReadyToMerge(idx_t count) {
+	// we try to merge so the count fits nicely into row groups
+	if (count >= Storage::ROW_GROUP_SIZE / 10 * 9 && count <= Storage::ROW_GROUP_SIZE) {
+		// 90%-100% of row group size
+		return true;
+	}
+	if (count >= Storage::ROW_GROUP_SIZE / 10 * 18 && count <= Storage::ROW_GROUP_SIZE * 2) {
+		// 180%-200% of row group size
+		return true;
+	}
+	if (count >= Storage::ROW_GROUP_SIZE / 10 * 27 && count <= Storage::ROW_GROUP_SIZE * 3) {
+		// 270%-300% of row group size
+		return true;
+	}
+	if (count >= Storage::ROW_GROUP_SIZE / 10 * 36) {
+		// >360% of row group size
+		return true;
+	}
+	return false;
+}
+
+void BatchInsertGlobalState::ScheduleMergeTasks(idx_t min_batch_index) {
+	idx_t current_idx;
+
+	vector<BatchMergeTask> to_be_scheduled_tasks;
+
+	BatchMergeTask current_task(next_start);
+	for (current_idx = current_task.start_index; current_idx < collections.size(); current_idx++) {
+		auto &entry = collections[current_idx];
+		if (entry.batch_idx > min_batch_index) {
+			// this entry is AFTER the min_batch_index
+			// finished
+			if (ReadyToMerge(current_task.total_count)) {
+				current_task.end_index = current_idx;
+				to_be_scheduled_tasks.push_back(current_task);
+			}
+			break;
+		}
+		if (entry.type == RowGroupBatchType::FLUSHED) {
+			// already flushed: cannot flush anything here
+			if (current_task.total_count > 0) {
+				current_task.end_index = current_idx;
+				to_be_scheduled_tasks.push_back(current_task);
+			}
+			current_task.start_index = current_idx + 1;
+			if (current_task.start_index > next_start) {
+				// avoid checking this segment again in the future
+				next_start = current_task.start_index;
+			}
+			current_task.total_count = 0;
+			continue;
+		}
+		// not flushed - add to set of indexes to flush
+		current_task.total_count += entry.total_rows;
+		if (ReadyToMerge(current_task.total_count)) {
+			// create a task to merge these collections
+			current_task.end_index = current_idx + 1;
+			to_be_scheduled_tasks.push_back(current_task);
+			current_task.start_index = current_idx + 1;
+			current_task.total_count = 0;
+		}
+	}
+
+	if (to_be_scheduled_tasks.empty()) {
+		return;
+	}
+	for (auto &scheduled_task : to_be_scheduled_tasks) {
+		D_ASSERT(scheduled_task.total_count > 0);
+		D_ASSERT(current_idx > scheduled_task.start_index);
+		idx_t merged_batch_index = collections[scheduled_task.start_index].batch_idx;
+		vector<RowGroupBatchEntry> merge_collections;
+		for (idx_t idx = scheduled_task.start_index; idx < scheduled_task.end_index; idx++) {
+			auto &entry = collections[idx];
+			if (!entry.collection || entry.type == RowGroupBatchType::FLUSHED) {
+				throw InternalException("Adding a row group collection that should not be flushed");
+			}
+			RowGroupBatchEntry added_entry(collections[scheduled_task.start_index].batch_idx,
+			                               std::move(entry.collection), RowGroupBatchType::FLUSHED);
+			added_entry.unflushed_memory = entry.unflushed_memory;
+			merge_collections.push_back(std::move(added_entry));
+			entry.total_rows = scheduled_task.total_count;
+			entry.type = RowGroupBatchType::FLUSHED;
+		}
+		task_manager.AddTask(make_uniq<MergeCollectionTask>(std::move(merge_collections), merged_batch_index));
+	}
+	// erase in reverse order
+	for (idx_t i = to_be_scheduled_tasks.size(); i > 0; i--) {
+		auto &scheduled_task = to_be_scheduled_tasks[i - 1];
+		if (scheduled_task.start_index + 1 < scheduled_task.end_index) {
+			// erase all entries except the first one
+			collections.erase(collections.begin() + scheduled_task.start_index + 1,
+			                  collections.begin() + scheduled_task.end_index);
+		}
+	}
+}
+
+unique_ptr<RowGroupCollection> BatchInsertGlobalState::MergeCollections(ClientContext &context,
+                                                                        vector<RowGroupBatchEntry> merge_collections,
+                                                                        OptimisticDataWriter &writer) {
+	D_ASSERT(!merge_collections.empty());
+	CollectionMerger merger(context);
+	idx_t written_data = 0;
+	for (auto &entry : merge_collections) {
+		merger.AddCollection(std::move(entry.collection));
+		written_data += entry.unflushed_memory;
+	}
+	optimistically_written = true;
+	memory_manager.ReduceUnflushedMemory(written_data);
+	return merger.Flush(writer);
+}
+
+void BatchInsertGlobalState::AddCollection(ClientContext &context, idx_t batch_index, idx_t min_batch_index,
+                                           unique_ptr<RowGroupCollection> current_collection,
+                                           optional_ptr<OptimisticDataWriter> writer) {
+	if (batch_index < min_batch_index) {
+		throw InternalException("Batch index of the added collection (%llu) is smaller than the min batch index (%llu)",
+		                        batch_index, min_batch_index);
+	}
+	auto new_count = current_collection->GetTotalRows();
+	auto batch_type = new_count < Storage::ROW_GROUP_SIZE ? RowGroupBatchType::NOT_FLUSHED : RowGroupBatchType::FLUSHED;
+	if (batch_type == RowGroupBatchType::FLUSHED && writer) {
+		writer->WriteLastRowGroup(*current_collection);
+	}
+	lock_guard<mutex> l(lock);
+	insert_count += new_count;
+	// add the collection to the batch index
+	RowGroupBatchEntry new_entry(batch_index, std::move(current_collection), batch_type);
+	if (batch_type == RowGroupBatchType::NOT_FLUSHED) {
+		memory_manager.IncreaseUnflushedMemory(new_entry.unflushed_memory);
+	}
+
+	auto it = std::lower_bound(
+	    collections.begin(), collections.end(), new_entry,
+	    [&](const RowGroupBatchEntry &a, const RowGroupBatchEntry &b) { return a.batch_idx < b.batch_idx; });
+	if (it != collections.end() && it->batch_idx == new_entry.batch_idx) {
+		throw InternalException("PhysicalBatchInsert::AddCollection error: batch index %d is present in multiple "
+		                        "collections. This occurs when "
+		                        "batch indexes are not uniquely distributed over threads",
+		                        batch_index);
+	}
+	collections.insert(it, std::move(new_entry));
+	if (writer) {
+		ScheduleMergeTasks(min_batch_index);
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// States
+//===--------------------------------------------------------------------===//
 unique_ptr<GlobalSinkState> PhysicalBatchInsert::GetGlobalSinkState(ClientContext &context) const {
 	optional_ptr<TableCatalogEntry> table;
 	if (info) {
@@ -291,7 +375,10 @@ unique_ptr<GlobalSinkState> PhysicalBatchInsert::GetGlobalSinkState(ClientContex
 		D_ASSERT(insert_table->IsDuckTable());
 		table = insert_table.get_mutable();
 	}
-	auto result = make_uniq<BatchInsertGlobalState>(table->Cast<DuckTableEntry>());
+	// heuristic - we start off by allocating 4MB of cache space per column
+	static constexpr const idx_t MINIMUM_MEMORY_PER_COLUMN = 4 * 1024 * 1024;
+	auto initial_memory = table->GetColumns().PhysicalColumnCount() * MINIMUM_MEMORY_PER_COLUMN;
+	auto result = make_uniq<BatchInsertGlobalState>(context, table->Cast<DuckTableEntry>(), initial_memory);
 	return std::move(result);
 }
 
@@ -299,11 +386,34 @@ unique_ptr<LocalSinkState> PhysicalBatchInsert::GetLocalSinkState(ExecutionConte
 	return make_uniq<BatchInsertLocalState>(context.client, insert_types, bound_defaults);
 }
 
+//===--------------------------------------------------------------------===//
+// Tasks
+//===--------------------------------------------------------------------===//
+bool PhysicalBatchInsert::ExecuteTask(ClientContext &context, GlobalSinkState &gstate_p,
+                                      LocalSinkState &lstate_p) const {
+	auto &gstate = gstate_p.Cast<BatchInsertGlobalState>();
+	auto task = gstate.task_manager.GetTask();
+	if (!task) {
+		return false;
+	}
+	task->Execute(*this, context, gstate_p, lstate_p);
+	return true;
+}
+
+void PhysicalBatchInsert::ExecuteTasks(ClientContext &context, GlobalSinkState &gstate_p,
+                                       LocalSinkState &lstate_p) const {
+	while (ExecuteTask(context, gstate_p, lstate_p)) {
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// NextBatch
+//===--------------------------------------------------------------------===//
 SinkNextBatchType PhysicalBatchInsert::NextBatch(ExecutionContext &context, OperatorSinkNextBatchInput &input) const {
 	auto &gstate = input.global_state.Cast<BatchInsertGlobalState>();
 	auto &lstate = input.local_state.Cast<BatchInsertLocalState>();
+	auto &memory_manager = gstate.memory_manager;
 
-	auto &table = gstate.table;
 	auto batch_index = lstate.partition_info.batch_index.GetIndex();
 	if (lstate.current_collection) {
 		if (lstate.current_index == batch_index) {
@@ -313,26 +423,61 @@ SinkNextBatchType PhysicalBatchInsert::NextBatch(ExecutionContext &context, Oper
 		TransactionData tdata(0, 0);
 		lstate.current_collection->FinalizeAppend(tdata, lstate.current_append_state);
 		gstate.AddCollection(context.client, lstate.current_index, lstate.partition_info.min_batch_index.GetIndex(),
-		                     std::move(lstate.current_collection), lstate.writer, &lstate.written_to_disk);
-		lstate.CreateNewCollection(table, insert_types);
+		                     std::move(lstate.current_collection), lstate.writer);
+
+		auto any_unblocked = memory_manager.UnblockTasks();
+		if (!any_unblocked) {
+			ExecuteTasks(context.client, gstate, lstate);
+		}
+		lstate.current_collection.reset();
 	}
 	lstate.current_index = batch_index;
+
+	// unblock any blocked tasks
+	memory_manager.UnblockTasks();
+
 	return SinkNextBatchType::READY;
 }
 
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
 SinkResultType PhysicalBatchInsert::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
 	auto &gstate = input.global_state.Cast<BatchInsertGlobalState>();
 	auto &lstate = input.local_state.Cast<BatchInsertLocalState>();
+	auto &memory_manager = gstate.memory_manager;
 
 	auto &table = gstate.table;
 	PhysicalInsert::ResolveDefaults(table, chunk, column_index_map, lstate.default_executor, lstate.insert_chunk);
 
 	auto batch_index = lstate.partition_info.batch_index.GetIndex();
+	// check if we should process this batch
+	if (!memory_manager.IsMinimumBatchIndex(batch_index)) {
+		memory_manager.UpdateMinBatchIndex(lstate.partition_info.min_batch_index.GetIndex());
+
+		// we are not processing the current min batch index
+		// check if we have exceeded the maximum number of unflushed rows
+		if (memory_manager.OutOfMemory(batch_index)) {
+			// out-of-memory
+			// execute tasks while we wait (if any are available)
+			ExecuteTasks(context.client, gstate, lstate);
+
+			lock_guard<mutex> l(memory_manager.GetBlockedTaskLock());
+			if (!memory_manager.IsMinimumBatchIndex(batch_index)) {
+				//  we are not the minimum batch index and we have no memory available to buffer - block the task for
+				//  now
+				memory_manager.BlockTask(input.interrupt_state);
+				return SinkResultType::BLOCKED;
+			}
+		}
+	}
 	if (!lstate.current_collection) {
 		lock_guard<mutex> l(gstate.lock);
 		// no collection yet: create a new one
 		lstate.CreateNewCollection(table, insert_types);
-		lstate.writer = &table.GetStorage().CreateOptimisticWriter(context.client);
+		if (!lstate.writer) {
+			lstate.writer = &table.GetStorage().CreateOptimisticWriter(context.client);
+		}
 	}
 
 	if (lstate.current_index != batch_index) {
@@ -345,39 +490,49 @@ SinkResultType PhysicalBatchInsert::Sink(ExecutionContext &context, DataChunk &c
 	if (new_row_group) {
 		// we have already written to disk - flush the next row group as well
 		lstate.writer->WriteNewRowGroup(*lstate.current_collection);
-		lstate.written_to_disk = true;
 	}
 	return SinkResultType::NEED_MORE_INPUT;
 }
 
+//===--------------------------------------------------------------------===//
+// Combine
+//===--------------------------------------------------------------------===//
 SinkCombineResultType PhysicalBatchInsert::Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const {
 	auto &gstate = input.global_state.Cast<BatchInsertGlobalState>();
 	auto &lstate = input.local_state.Cast<BatchInsertLocalState>();
+	auto &memory_manager = gstate.memory_manager;
 	auto &client_profiler = QueryProfiler::Get(context.client);
 	context.thread.profiler.Flush(*this, lstate.default_executor, "default_executor", 1);
 	client_profiler.Flush(context.thread.profiler);
 
-	if (!lstate.current_collection) {
-		return SinkCombineResultType::FINISHED;
-	}
+	memory_manager.UpdateMinBatchIndex(lstate.partition_info.min_batch_index.GetIndex());
 
-	if (lstate.current_collection->GetTotalRows() > 0) {
+	if (lstate.current_collection) {
 		TransactionData tdata(0, 0);
 		lstate.current_collection->FinalizeAppend(tdata, lstate.current_append_state);
-		gstate.AddCollection(context.client, lstate.current_index, lstate.partition_info.min_batch_index.GetIndex(),
-		                     std::move(lstate.current_collection));
+		if (lstate.current_collection->GetTotalRows() > 0) {
+			gstate.AddCollection(context.client, lstate.current_index, lstate.partition_info.min_batch_index.GetIndex(),
+			                     std::move(lstate.current_collection));
+		}
 	}
-	{
+	if (lstate.writer) {
 		lock_guard<mutex> l(gstate.lock);
 		gstate.table.GetStorage().FinalizeOptimisticWriter(context.client, *lstate.writer);
 	}
 
+	// unblock any blocked tasks
+	memory_manager.UnblockTasks();
+
 	return SinkCombineResultType::FINISHED;
 }
 
+//===--------------------------------------------------------------------===//
+// Finalize
+//===--------------------------------------------------------------------===//
 SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
                                                OperatorSinkFinalizeInput &input) const {
 	auto &gstate = input.global_state.Cast<BatchInsertGlobalState>();
+	auto &memory_manager = gstate.memory_manager;
 
 	if (gstate.optimistically_written || gstate.insert_count >= LocalStorage::MERGE_THRESHOLD) {
 		// we have written data to disk optimistically or are inserting a large amount of data
@@ -393,6 +548,7 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 					current_merger = make_uniq<CollectionMerger>(context);
 				}
 				current_merger->AddCollection(std::move(entry.collection));
+				memory_manager.ReduceUnflushedMemory(entry.unflushed_memory);
 			} else {
 				// this collection has been flushed: it does not need to be merged
 				// create a separate collection merger only for this entry
@@ -417,12 +573,12 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 		for (auto &merger : mergers) {
 			final_collections.push_back(merger->Flush(writer));
 		}
-		storage.FinalizeOptimisticWriter(context, writer);
 
 		// finally, merge the row groups into the local storage
 		for (auto &collection : final_collections) {
 			storage.LocalMerge(context, *collection);
 		}
+		storage.FinalizeOptimisticWriter(context, writer);
 	} else {
 		// we are writing a small amount of data to disk
 		// append directly to transaction local storage
@@ -432,6 +588,11 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 		storage.InitializeLocalAppend(append_state, context);
 		auto &transaction = DuckTransaction::Get(context, table.catalog);
 		for (auto &entry : gstate.collections) {
+			if (entry.type != RowGroupBatchType::NOT_FLUSHED) {
+				throw InternalException("Encountered a flushed batch");
+			}
+
+			memory_manager.ReduceUnflushedMemory(entry.unflushed_memory);
 			entry.collection->Scan(transaction, [&](DataChunk &insert_chunk) {
 				storage.LocalAppend(append_state, table, context, insert_chunk);
 				return true;
@@ -439,6 +600,7 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 		}
 		storage.FinalizeLocalAppend(append_state);
 	}
+	memory_manager.FinalCheck();
 	return SinkFinalizeType::READY;
 }