duckdb 0.7.2-dev2820.0 → 0.7.2-dev2995.0

package/src/duckdb/src/execution/operator/order/physical_top_n.cpp

@@ -435,11 +435,10 @@ unique_ptr<GlobalSinkState> PhysicalTopN::GetGlobalSinkState(ClientContext &cont
 //===--------------------------------------------------------------------===//
 // Sink
 //===--------------------------------------------------------------------===//
-SinkResultType PhysicalTopN::Sink(ExecutionContext &context, GlobalSinkState &state, LocalSinkState &lstate,
-                                  DataChunk &input) const {
+SinkResultType PhysicalTopN::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
 	// append to the local sink state
-	auto &sink = lstate.Cast<TopNLocalState>();
-	sink.heap.Sink(input);
+	auto &sink = input.local_state.Cast<TopNLocalState>();
+	sink.heap.Sink(chunk);
 	sink.heap.Reduce();
 	return SinkResultType::NEED_MORE_INPUT;
 }
@@ -480,12 +479,11 @@ unique_ptr<GlobalSourceState> PhysicalTopN::GetGlobalSourceState(ClientContext &
 	return make_uniq<TopNOperatorState>();
 }
 
-void PhysicalTopN::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate_p,
-                           LocalSourceState &lstate) const {
+SourceResultType PhysicalTopN::GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const {
 	if (limit == 0) {
-		return;
+		return SourceResultType::FINISHED;
 	}
-	auto &state = gstate_p.Cast<TopNOperatorState>();
+	auto &state = input.global_state.Cast<TopNOperatorState>();
 	auto &gstate = sink_state->Cast<TopNGlobalState>();
 
 	if (!state.initialized) {
@@ -493,6 +491,8 @@ void PhysicalTopN::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSo
 		state.initialized = true;
 	}
 	gstate.heap.Scan(state.state, chunk);
+
+	return chunk.size() == 0 ? SourceResultType::FINISHED : SourceResultType::HAVE_MORE_OUTPUT;
 }
 
 string PhysicalTopN::ParamsToString() const {

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

@@ -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

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

@@ -31,14 +31,13 @@ public:
 	DataChunk delete_chunk;
 };
 
-SinkResultType PhysicalDelete::Sink(ExecutionContext &context, GlobalSinkState &state, LocalSinkState &lstate,
-                                    DataChunk &input) const {
-	auto &gstate = state.Cast<DeleteGlobalState>();
-	auto &ustate = lstate.Cast<DeleteLocalState>();
+SinkResultType PhysicalDelete::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
+	auto &gstate = input.global_state.Cast<DeleteGlobalState>();
+	auto &ustate = input.local_state.Cast<DeleteLocalState>();
 
 	// get rows and
 	auto &transaction = DuckTransaction::Get(context.client, table.db);
-	auto &row_identifiers = input.data[row_id_index];
+	auto &row_identifiers = chunk.data[row_id_index];
 
 	vector<column_t> column_ids;
 	for (idx_t i = 0; i < table.column_definitions.size(); i++) {
@@ -48,11 +47,11 @@ SinkResultType PhysicalDelete::Sink(ExecutionContext &context, GlobalSinkState &
 
 	lock_guard<mutex> delete_guard(gstate.delete_lock);
 	if (return_chunk) {
-		row_identifiers.Flatten(input.size());
-		table.Fetch(transaction, ustate.delete_chunk, column_ids, row_identifiers, input.size(), cfs);
+		row_identifiers.Flatten(chunk.size());
+		table.Fetch(transaction, ustate.delete_chunk, column_ids, row_identifiers, chunk.size(), cfs);
 		gstate.return_collection.Append(ustate.delete_chunk);
 	}
-	gstate.deleted_count += table.Delete(tableref, context.client, row_identifiers, input.size());
+	gstate.deleted_count += table.Delete(tableref, context.client, row_identifiers, chunk.size());
 
 	return SinkResultType::NEED_MORE_INPUT;
 }
@@ -70,7 +69,7 @@ unique_ptr<LocalSinkState> PhysicalDelete::GetLocalSinkState(ExecutionContext &c
 //===--------------------------------------------------------------------===//
 class DeleteSourceState : public GlobalSourceState {
 public:
-	explicit DeleteSourceState(const PhysicalDelete &op) : finished(false) {
+	explicit DeleteSourceState(const PhysicalDelete &op) {
 		if (op.return_chunk) {
 			D_ASSERT(op.sink_state);
 			auto &g = op.sink_state->Cast<DeleteGlobalState>();
@@ -79,29 +78,25 @@ public:
 	}
 
 	ColumnDataScanState scan_state;
-	bool finished;
 };
 
 unique_ptr<GlobalSourceState> PhysicalDelete::GetGlobalSourceState(ClientContext &context) const {
 	return make_uniq<DeleteSourceState>(*this);
 }
 
-void PhysicalDelete::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate,
-                             LocalSourceState &lstate) const {
-	auto &state = gstate.Cast<DeleteSourceState>();
+SourceResultType PhysicalDelete::GetData(ExecutionContext &context, DataChunk &chunk,
+                                         OperatorSourceInput &input) const {
+	auto &state = input.global_state.Cast<DeleteSourceState>();
 	auto &g = sink_state->Cast<DeleteGlobalState>();
-	if (state.finished) {
-		return;
-	}
-
 	if (!return_chunk) {
 		chunk.SetCardinality(1);
 		chunk.SetValue(0, 0, Value::BIGINT(g.deleted_count));
-		state.finished = true;
-		return;
+		return SourceResultType::FINISHED;
 	}
 
 	g.return_collection.Scan(state.scan_state, chunk);
+
+	return chunk.size() == 0 ? SourceResultType::FINISHED : SourceResultType::HAVE_MORE_OUTPUT;
 }
 
 } // namespace duckdb

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

@@ -96,11 +96,11 @@ unique_ptr<GlobalSourceState> PhysicalExport::GetGlobalSourceState(ClientContext
 	return make_uniq<ExportSourceState>();
 }
 
-void PhysicalExport::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate,
-                             LocalSourceState &lstate) const {
-	auto &state = gstate.Cast<ExportSourceState>();
+SourceResultType PhysicalExport::GetData(ExecutionContext &context, DataChunk &chunk,
+                                         OperatorSourceInput &input) const {
+	auto &state = input.global_state.Cast<ExportSourceState>();
 	if (state.finished) {
-		return;
+		return SourceResultType::FINISHED;
 	}
 
 	auto &ccontext = context.client;
@@ -180,13 +180,14 @@ void PhysicalExport::GetData(ExecutionContext &context, DataChunk &chunk, Global
 	}
 	WriteStringStreamToFile(fs, opener, load_ss, fs.JoinPath(info->file_path, "load.sql"));
 	state.finished = true;
+
+	return SourceResultType::FINISHED;
 }
 
 //===--------------------------------------------------------------------===//
 // Sink
 //===--------------------------------------------------------------------===//
-SinkResultType PhysicalExport::Sink(ExecutionContext &context, GlobalSinkState &gstate, LocalSinkState &lstate,
-                                    DataChunk &input) const {
+SinkResultType PhysicalExport::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
 	// nop
 	return SinkResultType::NEED_MORE_INPUT;
 }