duckdb 0.7.2-dev2867.0 → 0.7.2-dev3117.0

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

@@ -0,0 +1,165 @@
+#include "duckdb/execution/operator/persistent/physical_batch_copy_to_file.hpp"
+#include "duckdb/execution/operator/persistent/physical_copy_to_file.hpp"
+#include "duckdb/common/vector_operations/vector_operations.hpp"
+#include "duckdb/common/types/batched_data_collection.hpp"
+#include "duckdb/common/file_system.hpp"
+#include "duckdb/common/file_opener.hpp"
+#include "duckdb/common/allocator.hpp"
+#include <algorithm>
+
+namespace duckdb {
+
+PhysicalBatchCopyToFile::PhysicalBatchCopyToFile(vector<LogicalType> types, CopyFunction function_p,
+                                                 unique_ptr<FunctionData> bind_data, idx_t estimated_cardinality)
+    : PhysicalOperator(PhysicalOperatorType::BATCH_COPY_TO_FILE, std::move(types), estimated_cardinality),
+      function(std::move(function_p)), bind_data(std::move(bind_data)) {
+	if (!function.flush_batch || !function.prepare_batch) {
+		throw InternalException(
+		    "PhysicalBatchCopyToFile created for copy function that does not have prepare_batch/flush_batch defined");
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
+class BatchCopyToGlobalState : public GlobalSinkState {
+public:
+	explicit BatchCopyToGlobalState(unique_ptr<GlobalFunctionData> global_state)
+	    : rows_copied(0), global_state(std::move(global_state)) {
+	}
+
+	mutex lock;
+	mutex flush_lock;
+	atomic<idx_t> rows_copied;
+	unique_ptr<GlobalFunctionData> global_state;
+	map<idx_t, unique_ptr<PreparedBatchData>> batch_data;
+};
+
+class BatchCopyToLocalState : public LocalSinkState {
+public:
+	explicit BatchCopyToLocalState(unique_ptr<LocalFunctionData> local_state_p)
+	    : local_state(std::move(local_state_p)), rows_copied(0), batch_index(0) {
+	}
+
+	unique_ptr<LocalFunctionData> local_state;
+	unique_ptr<ColumnDataCollection> collection;
+	ColumnDataAppendState append_state;
+	idx_t rows_copied;
+	idx_t batch_index;
+
+	void InitializeCollection(ClientContext &context, const PhysicalOperator &op) {
+		collection = make_uniq<ColumnDataCollection>(Allocator::Get(context), op.children[0]->types);
+		collection->InitializeAppend(append_state);
+	}
+};
+
+SinkResultType PhysicalBatchCopyToFile::Sink(ExecutionContext &context, DataChunk &chunk,
+                                             OperatorSinkInput &input) const {
+	auto &state = input.local_state.Cast<BatchCopyToLocalState>();
+	if (!state.collection) {
+		state.InitializeCollection(context.client, *this);
+	}
+	state.rows_copied += chunk.size();
+	state.collection->Append(state.append_state, chunk);
+	return SinkResultType::NEED_MORE_INPUT;
+}
+
+void PhysicalBatchCopyToFile::Combine(ExecutionContext &context, GlobalSinkState &gstate_p,
+                                      LocalSinkState &lstate) const {
+	auto &state = lstate.Cast<BatchCopyToLocalState>();
+	auto &gstate = gstate_p.Cast<BatchCopyToGlobalState>();
+	gstate.rows_copied += state.rows_copied;
+}
+
+SinkFinalizeType PhysicalBatchCopyToFile::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
+                                                   GlobalSinkState &gstate_p) const {
+	auto &gstate = gstate_p.Cast<BatchCopyToGlobalState>();
+	FlushBatchData(context, gstate_p, NumericLimits<int64_t>::Maximum());
+	if (function.copy_to_finalize) {
+		function.copy_to_finalize(context, *bind_data, *gstate.global_state);
+
+		if (use_tmp_file) {
+			PhysicalCopyToFile::MoveTmpFile(context, file_path);
+		}
+	}
+	return SinkFinalizeType::READY;
+}
+
+void PhysicalBatchCopyToFile::PrepareBatchData(ClientContext &context, GlobalSinkState &gstate_p, idx_t batch_index,
+                                               unique_ptr<ColumnDataCollection> collection) const {
+	auto &gstate = gstate_p.Cast<BatchCopyToGlobalState>();
+
+	// prepare the batch
+	auto batch_data = function.prepare_batch(context, *bind_data, *gstate.global_state, std::move(collection));
+	// move the batch data to the set of prepared batch data
+	lock_guard<mutex> l(gstate.lock);
+	gstate.batch_data[batch_index] = std::move(batch_data);
+}
+
+void PhysicalBatchCopyToFile::FlushBatchData(ClientContext &context, GlobalSinkState &gstate_p, idx_t min_index) const {
+	auto &gstate = gstate_p.Cast<BatchCopyToGlobalState>();
+
+	// flush batch data to disk (if there are any to flush)
+	while (true) {
+		// grab the flush lock - we can only call flush_batch with this lock
+		// otherwise the data might end up in the wrong order
+		lock_guard<mutex> l(gstate.flush_lock);
+		unique_ptr<PreparedBatchData> batch_data;
+		{
+			// fetch the next batch to flush (if any)
+			lock_guard<mutex> l(gstate.lock);
+			if (gstate.batch_data.empty()) {
+				// no batch data left to flush
+				break;
+			}
+			auto entry = gstate.batch_data.begin();
+			if (entry->first >= min_index) {
+				// this data is past the min_index - we cannot write it yet
+				break;
+			}
+			if (!entry->second) {
+				// this batch is in process of being prepared but is not ready yet
+				break;
+			}
+			batch_data = std::move(entry->second);
+			gstate.batch_data.erase(entry);
+		}
+		function.flush_batch(context, *bind_data, *gstate.global_state, *batch_data);
+	}
+}
+
+void PhysicalBatchCopyToFile::NextBatch(ExecutionContext &context, GlobalSinkState &gstate_p,
+                                        LocalSinkState &lstate) const {
+	auto &state = lstate.Cast<BatchCopyToLocalState>();
+	if (state.collection) {
+		// we finished processing this batch
+		// start flushing data
+		PrepareBatchData(context.client, gstate_p, state.batch_index, std::move(state.collection));
+		FlushBatchData(context.client, gstate_p, lstate.partition_info.min_batch_index.GetIndex());
+	}
+	state.batch_index = lstate.partition_info.batch_index.GetIndex();
+
+	state.InitializeCollection(context.client, *this);
+}
+
+unique_ptr<LocalSinkState> PhysicalBatchCopyToFile::GetLocalSinkState(ExecutionContext &context) const {
+	return make_uniq<BatchCopyToLocalState>(function.copy_to_initialize_local(context, *bind_data));
+}
+
+unique_ptr<GlobalSinkState> PhysicalBatchCopyToFile::GetGlobalSinkState(ClientContext &context) const {
+	return make_uniq<BatchCopyToGlobalState>(function.copy_to_initialize_global(context, *bind_data, file_path));
+}
+
+//===--------------------------------------------------------------------===//
+// Source
+//===--------------------------------------------------------------------===//
+SourceResultType PhysicalBatchCopyToFile::GetData(ExecutionContext &context, DataChunk &chunk,
+                                                  OperatorSourceInput &input) const {
+	auto &g = sink_state->Cast<BatchCopyToGlobalState>();
+
+	chunk.SetCardinality(1);
+	chunk.SetValue(0, 0, Value::BIGINT(g.rows_copied));
+	return SourceResultType::FINISHED;
+}
+
+} // namespace duckdb

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

@@ -81,20 +81,35 @@ public:
 					}
 					auto new_row_group = new_collection->Append(scan_chunk, append_state);
 					if (new_row_group) {
-						writer.CheckFlushToDisk(*new_collection);
+						writer.WriteNewRowGroup(*new_collection);
 					}
 				}
 			}
-
 			new_collection->FinalizeAppend(TransactionData(0, 0), append_state);
-			writer.FlushToDisk(*new_collection);
+			writer.WriteLastRowGroup(*new_collection);
 		}
 		current_collections.clear();
 		return new_collection;
 	}
 };
 
+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) {
+	}
+
+	idx_t batch_idx;
+	idx_t total_rows;
+	unique_ptr<RowGroupCollection> collection;
+	RowGroupBatchType type;
+};
+
 class BatchInsertGlobalState : public GlobalSinkState {
+public:
+	static constexpr const idx_t BATCH_FLUSH_THRESHOLD = LocalStorage::MERGE_THRESHOLD * 3;
+
 public:
 	explicit BatchInsertGlobalState(DuckTableEntry &table) : table(table), insert_count(0) {
 	}
@@ -102,38 +117,60 @@ public:
 	mutex lock;
 	DuckTableEntry &table;
 	idx_t insert_count;
-	map<idx_t, unique_ptr<RowGroupCollection>> collections;
-
-	bool CheckMergeInternal(idx_t batch_index, vector<unique_ptr<RowGroupCollection>> *result, idx_t *merge_count) {
-		auto entry = collections.find(batch_index);
-		if (entry == collections.end()) {
-			// no collection at this index
-			return false;
-		}
-		auto row_count = entry->second->GetTotalRows();
-		if (row_count >= LocalStorage::MERGE_THRESHOLD) {
-			// the collection at this batch index is large and has already been written
-			return false;
+	vector<RowGroupBatchEntry> collections;
+	idx_t next_start = 0;
+
+	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;
+			}
 		}
-		// we can merge this collection!
-		if (merge_count) {
-			// add the count
-			D_ASSERT(!result);
-			*merge_count += row_count;
-		} else {
-			// add the
-			D_ASSERT(result);
-			result->push_back(std::move(entry->second));
-			collections.erase(batch_index);
+		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);
+			}
 		}
-		return true;
-	}
-
-	bool CheckMerge(idx_t batch_index, idx_t &merge_count) {
-		return CheckMergeInternal(batch_index, nullptr, &merge_count);
-	}
-	bool CheckMerge(idx_t batch_index, vector<unique_ptr<RowGroupCollection>> &result) {
-		return CheckMergeInternal(batch_index, &result, nullptr);
 	}
 
 	unique_ptr<RowGroupCollection> MergeCollections(ClientContext &context,
@@ -146,83 +183,65 @@ public:
 		return merger.Flush(writer);
 	}
 
-	void VerifyUniqueBatch(idx_t batch_index) {
-		if (collections.find(batch_index) != collections.end()) {
-			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);
-		}
-	}
-
-	void AddCollection(ClientContext &context, idx_t batch_index, unique_ptr<RowGroupCollection> current_collection,
+	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 < RowGroup::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;
-		idx_t merge_count;
 		{
 			lock_guard<mutex> l(lock);
-			auto new_count = current_collection->GetTotalRows();
 			insert_count += new_count;
-			VerifyUniqueBatch(batch_index);
-			if (writer && new_count < LocalStorage::MERGE_THRESHOLD) {
-				// we are inserting a small collection that has not yet been written to disk
-				// check if there are any collections with adjacent batch indexes that we can merge together
-
-				// first check how many rows we will end up with by performing such a merge
-				// check backwards
-				merge_count = new_count;
-				idx_t start_batch_index;
-				idx_t end_batch_index;
-				for (start_batch_index = batch_index; start_batch_index > 0; start_batch_index--) {
-					if (!CheckMerge(start_batch_index - 1, merge_count)) {
-						break;
-					}
-				}
-				// check forwards
-				for (end_batch_index = batch_index;; end_batch_index++) {
-					if (!CheckMerge(end_batch_index + 1, merge_count)) {
-						break;
-					}
-				}
-				// merging together creates a big enough row group
-				// merge!
-				if (merge_count >= RowGroup::ROW_GROUP_SIZE) {
-					// gather the row groups to merge
-					// note that we need to gather them in order of batch index
-					for (idx_t i = start_batch_index; i <= end_batch_index; i++) {
-						if (i == batch_index) {
-							merge_collections.push_back(std::move(current_collection));
-							current_collection.reset();
-							continue;
-						}
-						auto can_merge = CheckMerge(i, merge_collections);
-						if (!can_merge) {
-							throw InternalException("Could not merge row group in batch insert?!");
-						}
-					}
-				}
+
+			// 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);
 			}
-			if (merge_collections.empty()) {
-				// no collections to merge together - add the collection to the batch index
-				collections[batch_index] = std::move(current_collection);
-				current_collection.reset();
+			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);
-			D_ASSERT(final_collection->GetTotalRows() == merge_count);
-			D_ASSERT(final_collection->GetTotalRows() >= RowGroup::ROW_GROUP_SIZE);
 			if (written_to_disk) {
 				*written_to_disk = true;
 			}
-			// add the merged-together collection to the
+			// add the merged-together collection to the set of batch indexes
 			{
 				lock_guard<mutex> l(lock);
-				VerifyUniqueBatch(batch_index);
-				collections[batch_index] = std::move(final_collection);
+				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);
 			}
 		}
 	}
@@ -244,16 +263,6 @@ public:
 	optional_ptr<OptimisticDataWriter> writer;
 	bool written_to_disk;
 
-	void FlushToDisk() {
-		if (!current_collection) {
-			return;
-		}
-		if (!written_to_disk && current_collection->GetTotalRows() < LocalStorage::MERGE_THRESHOLD) {
-			return;
-		}
-		writer->FlushToDisk(*current_collection, true);
-	}
-
 	void CreateNewCollection(DuckTableEntry &table, const vector<LogicalType> &insert_types) {
 		auto &table_info = table.GetStorage().info;
 		auto &block_manager = TableIOManager::Get(table.GetStorage()).GetBlockManagerForRowData();
@@ -286,35 +295,50 @@ unique_ptr<LocalSinkState> PhysicalBatchInsert::GetLocalSinkState(ExecutionConte
 	return make_uniq<BatchInsertLocalState>(context.client, insert_types, bound_defaults);
 }
 
-SinkResultType PhysicalBatchInsert::Sink(ExecutionContext &context, GlobalSinkState &state, LocalSinkState &lstate_p,
-                                         DataChunk &chunk) const {
+void PhysicalBatchInsert::NextBatch(ExecutionContext &context, GlobalSinkState &state, LocalSinkState &lstate_p) const {
 	auto &gstate = state.Cast<BatchInsertGlobalState>();
 	auto &lstate = lstate_p.Cast<BatchInsertLocalState>();
 
+	auto &table = gstate.table;
+	auto batch_index = lstate.partition_info.batch_index.GetIndex();
+	if (lstate.current_collection) {
+		if (lstate.current_index == batch_index) {
+			throw InternalException("NextBatch called with the same batch index?");
+		}
+		// batch index has changed: move the old collection to the global state and create a new 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), lstate.writer, &lstate.written_to_disk);
+		lstate.CreateNewCollection(table, insert_types);
+	}
+	lstate.current_index = batch_index;
+}
+
+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 &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();
 	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);
-	} else if (lstate.current_index != lstate.batch_index) {
-		// batch index has changed: move the old collection to the global state and create a new collection
-		TransactionData tdata(0, 0);
-		lstate.current_collection->FinalizeAppend(tdata, lstate.current_append_state);
-		lstate.FlushToDisk();
-		gstate.AddCollection(context.client, lstate.current_index, std::move(lstate.current_collection), lstate.writer,
-		                     &lstate.written_to_disk);
-		lstate.CreateNewCollection(table, insert_types);
+	} else if (lstate.current_index != batch_index) {
+		throw InternalException("Current batch differs from batch - but NextBatch was not called!?");
 	}
-	lstate.current_index = lstate.batch_index;
+	lstate.current_index = batch_index;
 
 	table.GetStorage().VerifyAppendConstraints(table, context.client, lstate.insert_chunk);
 
 	auto new_row_group = lstate.current_collection->Append(lstate.insert_chunk, lstate.current_append_state);
 	if (new_row_group) {
-		lstate.writer->CheckFlushToDisk(*lstate.current_collection);
+		// 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;
@@ -331,12 +355,17 @@ void PhysicalBatchInsert::Combine(ExecutionContext &context, GlobalSinkState &gs
 	if (!lstate.current_collection) {
 		return;
 	}
-	lstate.FlushToDisk();
-	lstate.writer->FinalFlush();
 
-	TransactionData tdata(0, 0);
-	lstate.current_collection->FinalizeAppend(tdata, lstate.current_append_state);
-	gstate.AddCollection(context.client, lstate.current_index, std::move(lstate.current_collection));
+	if (lstate.current_collection->GetTotalRows() > 0) {
+		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));
+	}
+	{
+		lock_guard<mutex> l(gstate.lock);
+		gstate.table.GetStorage().FinalizeOptimisticWriter(context.client, *lstate.writer);
+	}
 }
 
 SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
@@ -349,15 +378,15 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 	unique_ptr<CollectionMerger> current_merger;
 
 	auto &storage = gstate.table.GetStorage();
-	for (auto &collection : gstate.collections) {
-		if (collection.second->GetTotalRows() < LocalStorage::MERGE_THRESHOLD) {
-			// this collection has very few rows: add it to the merge set
+	for (auto &entry : gstate.collections) {
+		if (entry.type == RowGroupBatchType::NOT_FLUSHED) {
+			// this collection has not been flushed: add it to the merge set
 			if (!current_merger) {
 				current_merger = make_uniq<CollectionMerger>(context);
 			}
-			current_merger->AddCollection(std::move(collection.second));
+			current_merger->AddCollection(std::move(entry.collection));
 		} else {
-			// this collection has a lot of rows: it does not need to be merged
+			// this collection has been flushed: it does not need to be merged
 			// create a separate collection merger only for this entry
 			if (current_merger) {
 				// we have small collections remaining: flush them
@@ -365,7 +394,7 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 				current_merger.reset();
 			}
 			auto larger_merger = make_uniq<CollectionMerger>(context);
-			larger_merger->AddCollection(std::move(collection.second));
+			larger_merger->AddCollection(std::move(entry.collection));
 			mergers.push_back(std::move(larger_merger));
 		}
 	}
@@ -380,7 +409,7 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 	for (auto &merger : mergers) {
 		final_collections.push_back(merger->Flush(writer));
 	}
-	writer.FinalFlush();
+	storage.FinalizeOptimisticWriter(context, writer);
 
 	// finally, merge the row groups into the local storage
 	for (auto &collection : final_collections) {
@@ -392,29 +421,15 @@ SinkFinalizeType PhysicalBatchInsert::Finalize(Pipeline &pipeline, Event &event,
 //===--------------------------------------------------------------------===//
 // Source
 //===--------------------------------------------------------------------===//
-class BatchInsertSourceState : public GlobalSourceState {
-public:
-	explicit BatchInsertSourceState() : finished(false) {
-	}
 
-	bool finished;
-};
-
-unique_ptr<GlobalSourceState> PhysicalBatchInsert::GetGlobalSourceState(ClientContext &context) const {
-	return make_uniq<BatchInsertSourceState>();
-}
-
-void PhysicalBatchInsert::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate,
-                                  LocalSourceState &lstate) const {
-	auto &state = gstate.Cast<BatchInsertSourceState>();
+SourceResultType PhysicalBatchInsert::GetData(ExecutionContext &context, DataChunk &chunk,
+                                              OperatorSourceInput &input) const {
 	auto &insert_gstate = sink_state->Cast<BatchInsertGlobalState>();
-	if (state.finished) {
-		return;
-	}
+
 	chunk.SetCardinality(1);
 	chunk.SetValue(0, 0, Value::BIGINT(insert_gstate.insert_count));
-	state.finished = true;
-	return;
+
+	return SourceResultType::FINISHED;
 }
 
 } // namespace duckdb

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

@@ -43,7 +43,7 @@ public:
 // Sink
 //===--------------------------------------------------------------------===//
 
-void MoveTmpFile(ClientContext &context, const string &tmp_file_path) {
+void PhysicalCopyToFile::MoveTmpFile(ClientContext &context, const string &tmp_file_path) {
 	auto &fs = FileSystem::GetFileSystem(context);
 	auto file_path = tmp_file_path.substr(0, tmp_file_path.length() - 4);
 	if (fs.FileExists(file_path)) {
@@ -58,22 +58,21 @@ PhysicalCopyToFile::PhysicalCopyToFile(vector<LogicalType> types, CopyFunction f
       function(std::move(function_p)), bind_data(std::move(bind_data)), parallel(false) {
 }
 
-SinkResultType PhysicalCopyToFile::Sink(ExecutionContext &context, GlobalSinkState &gstate, LocalSinkState &lstate,
-                                        DataChunk &input) const {
-	auto &g = gstate.Cast<CopyToFunctionGlobalState>();
-	auto &l = lstate.Cast<CopyToFunctionLocalState>();
+SinkResultType PhysicalCopyToFile::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
+	auto &g = input.global_state.Cast<CopyToFunctionGlobalState>();
+	auto &l = input.local_state.Cast<CopyToFunctionLocalState>();
 
 	if (partition_output) {
-		l.part_buffer->Append(*l.part_buffer_append_state, input);
+		l.part_buffer->Append(*l.part_buffer_append_state, chunk);
 		return SinkResultType::NEED_MORE_INPUT;
 	}
 
 	{
 		lock_guard<mutex> glock(g.lock);
-		g.rows_copied += input.size();
+		g.rows_copied += chunk.size();
 	}
 	function.copy_to_sink(context, *bind_data, per_thread_output ? *l.global_state : *g.global_state, *l.local_state,
-	                      input);
+	                      chunk);
 	return SinkResultType::NEED_MORE_INPUT;
 }
 
@@ -231,29 +230,15 @@ unique_ptr<GlobalSinkState> PhysicalCopyToFile::GetGlobalSinkState(ClientContext
 //===--------------------------------------------------------------------===//
 // Source
 //===--------------------------------------------------------------------===//
-class CopyToFileState : public GlobalSourceState {
-public:
-	CopyToFileState() : finished(false) {
-	}
-
-	bool finished;
-};
 
-unique_ptr<GlobalSourceState> PhysicalCopyToFile::GetGlobalSourceState(ClientContext &context) const {
-	return make_uniq<CopyToFileState>();
-}
-
-void PhysicalCopyToFile::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate,
-                                 LocalSourceState &lstate) const {
-	auto &state = (CopyToFileState &)gstate;
+SourceResultType PhysicalCopyToFile::GetData(ExecutionContext &context, DataChunk &chunk,
+                                             OperatorSourceInput &input) const {
 	auto &g = sink_state->Cast<CopyToFunctionGlobalState>();
-	if (state.finished) {
-		return;
-	}
 
 	chunk.SetCardinality(1);
 	chunk.SetValue(0, 0, Value::BIGINT(g.rows_copied));
-	state.finished = true;
+
+	return SourceResultType::FINISHED;
 }
 
 } // namespace duckdb