duckdb 0.8.2-dev3458.0 → 0.8.2-dev3949.0

package/src/duckdb/src/execution/operator/aggregate/physical_hash_aggregate.cpp

@@ -1,9 +1,12 @@
 #include "duckdb/execution/operator/aggregate/physical_hash_aggregate.hpp"
 
 #include "duckdb/catalog/catalog_entry/aggregate_function_catalog_entry.hpp"
+#include "duckdb/common/atomic.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"
 #include "duckdb/execution/aggregate_hashtable.hpp"
+#include "duckdb/execution/operator/aggregate/distinct_aggregate_data.hpp"
 #include "duckdb/main/client_context.hpp"
+#include "duckdb/parallel/base_pipeline_event.hpp"
 #include "duckdb/parallel/interrupt.hpp"
 #include "duckdb/parallel/pipeline.hpp"
 #include "duckdb/parallel/task_scheduler.hpp"
@@ -11,9 +14,6 @@
 #include "duckdb/planner/expression/bound_aggregate_expression.hpp"
 #include "duckdb/planner/expression/bound_constant_expression.hpp"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
-#include "duckdb/parallel/base_pipeline_event.hpp"
-#include "duckdb/common/atomic.hpp"
-#include "duckdb/execution/operator/aggregate/distinct_aggregate_data.hpp"
 
 namespace duckdb {
 
@@ -176,9 +176,9 @@ PhysicalHashAggregate::PhysicalHashAggregate(ClientContext &context, vector<Logi
 //===--------------------------------------------------------------------===//
 // Sink
 //===--------------------------------------------------------------------===//
-class HashAggregateGlobalState : public GlobalSinkState {
+class HashAggregateGlobalSinkState : public GlobalSinkState {
 public:
-	HashAggregateGlobalState(const PhysicalHashAggregate &op, ClientContext &context) {
+	HashAggregateGlobalSinkState(const PhysicalHashAggregate &op, ClientContext &context) {
 		grouping_states.reserve(op.groupings.size());
 		for (idx_t i = 0; i < op.groupings.size(); i++) {
 			auto &grouping = op.groupings[i];
@@ -204,9 +204,9 @@ public:
 	bool finished = false;
 };
 
-class HashAggregateLocalState : public LocalSinkState {
+class HashAggregateLocalSinkState : public LocalSinkState {
 public:
-	HashAggregateLocalState(const PhysicalHashAggregate &op, ExecutionContext &context) {
+	HashAggregateLocalSinkState(const PhysicalHashAggregate &op, ExecutionContext &context) {
 
 		auto &payload_types = op.grouped_aggregate_data.payload_types;
 		if (!payload_types.empty()) {
@@ -234,28 +234,30 @@ public:
 };
 
 void PhysicalHashAggregate::SetMultiScan(GlobalSinkState &state) {
-	auto &gstate = state.Cast<HashAggregateGlobalState>();
+	auto &gstate = state.Cast<HashAggregateGlobalSinkState>();
 	for (auto &grouping_state : gstate.grouping_states) {
-		auto &radix_state = grouping_state.table_state;
-		RadixPartitionedHashTable::SetMultiScan(*radix_state);
+		RadixPartitionedHashTable::SetMultiScan(*grouping_state.table_state);
 		if (!grouping_state.distinct_state) {
 			continue;
 		}
 	}
 }
 
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
 unique_ptr<GlobalSinkState> PhysicalHashAggregate::GetGlobalSinkState(ClientContext &context) const {
-	return make_uniq<HashAggregateGlobalState>(*this, context);
+	return make_uniq<HashAggregateGlobalSinkState>(*this, context);
 }
 
 unique_ptr<LocalSinkState> PhysicalHashAggregate::GetLocalSinkState(ExecutionContext &context) const {
-	return make_uniq<HashAggregateLocalState>(*this, context);
+	return make_uniq<HashAggregateLocalSinkState>(*this, context);
 }
 
 void PhysicalHashAggregate::SinkDistinctGrouping(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input,
                                                  idx_t grouping_idx) const {
-	auto &sink = input.local_state.Cast<HashAggregateLocalState>();
-	auto &global_sink = input.global_state.Cast<HashAggregateGlobalState>();
+	auto &sink = input.local_state.Cast<HashAggregateLocalSinkState>();
+	auto &global_sink = input.global_state.Cast<HashAggregateGlobalSinkState>();
 
 	auto &grouping_gstate = global_sink.grouping_states[grouping_idx];
 	auto &grouping_lstate = sink.grouping_states[grouping_idx];
@@ -341,8 +343,8 @@ void PhysicalHashAggregate::SinkDistinct(ExecutionContext &context, DataChunk &c
 
 SinkResultType PhysicalHashAggregate::Sink(ExecutionContext &context, DataChunk &chunk,
                                            OperatorSinkInput &input) const {
-	auto &llstate = input.local_state.Cast<HashAggregateLocalState>();
-	auto &gstate = input.global_state.Cast<HashAggregateGlobalState>();
+	auto &llstate = input.local_state.Cast<HashAggregateLocalSinkState>();
+	auto &gstate = input.global_state.Cast<HashAggregateGlobalSinkState>();
 
 	if (distinct_collection_info) {
 		SinkDistinct(context, chunk, input);
@@ -396,10 +398,13 @@ SinkResultType PhysicalHashAggregate::Sink(ExecutionContext &context, DataChunk
 	return SinkResultType::NEED_MORE_INPUT;
 }
 
+//===--------------------------------------------------------------------===//
+// Combine
+//===--------------------------------------------------------------------===//
 void PhysicalHashAggregate::CombineDistinct(ExecutionContext &context, OperatorSinkCombineInput &input) const {
 
-	auto &global_sink = input.global_state.Cast<HashAggregateGlobalState>();
-	auto &sink = input.local_state.Cast<HashAggregateLocalState>();
+	auto &global_sink = input.global_state.Cast<HashAggregateGlobalSinkState>();
+	auto &sink = input.local_state.Cast<HashAggregateLocalSinkState>();
 
 	if (!distinct_collection_info) {
 		return;
@@ -426,8 +431,8 @@ void PhysicalHashAggregate::CombineDistinct(ExecutionContext &context, OperatorS
 }
 
 SinkCombineResultType PhysicalHashAggregate::Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const {
-	auto &gstate = input.global_state.Cast<HashAggregateGlobalState>();
-	auto &llstate = input.local_state.Cast<HashAggregateLocalState>();
+	auto &gstate = input.global_state.Cast<HashAggregateGlobalSinkState>();
+	auto &llstate = input.local_state.Cast<HashAggregateLocalSinkState>();
 
 	OperatorSinkCombineInput combine_distinct_input {gstate, llstate, input.interrupt_state};
 	CombineDistinct(context, combine_distinct_input);
@@ -447,321 +452,267 @@ SinkCombineResultType PhysicalHashAggregate::Combine(ExecutionContext &context,
 	return SinkCombineResultType::FINISHED;
 }
 
-//! REGULAR FINALIZE EVENT
-
-class HashAggregateMergeEvent : public BasePipelineEvent {
+//===--------------------------------------------------------------------===//
+// Finalize
+//===--------------------------------------------------------------------===//
+class HashAggregateFinalizeEvent : public BasePipelineEvent {
 public:
-	HashAggregateMergeEvent(const PhysicalHashAggregate &op_p, HashAggregateGlobalState &gstate_p, Pipeline *pipeline_p)
-	    : BasePipelineEvent(*pipeline_p), op(op_p), gstate(gstate_p) {
+	//! "Regular" Finalize Event that is scheduled after combining the thread-local distinct HTs
+	HashAggregateFinalizeEvent(ClientContext &context, Pipeline *pipeline_p, const PhysicalHashAggregate &op_p,
+	                           HashAggregateGlobalSinkState &gstate_p)
+	    : BasePipelineEvent(*pipeline_p), context(context), op(op_p), gstate(gstate_p) {
 	}
 
-	const PhysicalHashAggregate &op;
-	HashAggregateGlobalState &gstate;
-
 public:
-	void Schedule() override {
-		vector<shared_ptr<Task>> tasks;
-		for (idx_t i = 0; i < op.groupings.size(); i++) {
-			auto &grouping_gstate = gstate.grouping_states[i];
+	void Schedule() override;
 
-			auto &grouping = op.groupings[i];
-			auto &table = grouping.table_data;
-			table.ScheduleTasks(pipeline->executor, shared_from_this(), *grouping_gstate.table_state, tasks);
-		}
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
-	}
-};
+private:
+	ClientContext &context;
 
-//! REGULAR FINALIZE FROM DISTINCT FINALIZE
+	const PhysicalHashAggregate &op;
+	HashAggregateGlobalSinkState &gstate;
+};
 
 class HashAggregateFinalizeTask : public ExecutorTask {
 public:
-	HashAggregateFinalizeTask(Pipeline &pipeline, shared_ptr<Event> event_p, HashAggregateGlobalState &state_p,
-	                          ClientContext &context, const PhysicalHashAggregate &op)
-	    : ExecutorTask(pipeline.executor), pipeline(pipeline), event(std::move(event_p)), gstate(state_p),
-	      context(context), op(op) {
+	HashAggregateFinalizeTask(ClientContext &context, Pipeline &pipeline, shared_ptr<Event> event_p,
+	                          const PhysicalHashAggregate &op, HashAggregateGlobalSinkState &state_p)
+	    : ExecutorTask(pipeline.executor), context(context), pipeline(pipeline), event(std::move(event_p)), op(op),
+	      gstate(state_p) {
 	}
 
-	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override {
-		op.FinalizeInternal(pipeline, *event, context, gstate, false);
-		D_ASSERT(!gstate.finished);
-		gstate.finished = true;
-		event->FinishTask();
-		return TaskExecutionResult::TASK_FINISHED;
-	}
+public:
+	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override;
 
 private:
+	ClientContext &context;
 	Pipeline &pipeline;
 	shared_ptr<Event> event;
-	HashAggregateGlobalState &gstate;
-	ClientContext &context;
+
 	const PhysicalHashAggregate &op;
+	HashAggregateGlobalSinkState &gstate;
 };
 
-class HashAggregateFinalizeEvent : public BasePipelineEvent {
+void HashAggregateFinalizeEvent::Schedule() {
+	vector<shared_ptr<Task>> tasks;
+	tasks.push_back(make_uniq<HashAggregateFinalizeTask>(context, *pipeline, shared_from_this(), op, gstate));
+	D_ASSERT(!tasks.empty());
+	SetTasks(std::move(tasks));
+}
+
+TaskExecutionResult HashAggregateFinalizeTask::ExecuteTask(TaskExecutionMode mode) {
+	op.FinalizeInternal(pipeline, *event, context, gstate, false);
+	D_ASSERT(!gstate.finished);
+	gstate.finished = true;
+	event->FinishTask();
+	return TaskExecutionResult::TASK_FINISHED;
+}
+
+class HashAggregateDistinctFinalizeEvent : public BasePipelineEvent {
 public:
-	HashAggregateFinalizeEvent(const PhysicalHashAggregate &op_p, HashAggregateGlobalState &gstate_p,
-	                           Pipeline *pipeline_p, ClientContext &context)
-	    : BasePipelineEvent(*pipeline_p), op(op_p), gstate(gstate_p), context(context) {
+	//! Distinct Finalize Event that is scheduled if we have distinct aggregates
+	HashAggregateDistinctFinalizeEvent(ClientContext &context, Pipeline &pipeline_p, const PhysicalHashAggregate &op_p,
+	                                   HashAggregateGlobalSinkState &gstate_p)
+	    : BasePipelineEvent(pipeline_p), context(context), op(op_p), gstate(gstate_p) {
 	}
 
-	const PhysicalHashAggregate &op;
-	HashAggregateGlobalState &gstate;
+public:
+	void Schedule() override;
+	void FinishEvent() override;
+
+private:
+	void CreateGlobalSources();
+
+private:
 	ClientContext &context;
 
+	const PhysicalHashAggregate &op;
+	HashAggregateGlobalSinkState &gstate;
+
 public:
-	void Schedule() override {
-		vector<shared_ptr<Task>> tasks;
-		tasks.push_back(make_uniq<HashAggregateFinalizeTask>(*pipeline, shared_from_this(), gstate, context, op));
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
-	}
+	//! The GlobalSourceStates for all the radix tables of the distinct aggregates
+	vector<vector<unique_ptr<GlobalSourceState>>> global_source_states;
 };
 
-//! DISTINCT FINALIZE TASK
+class HashAggregateDistinctFinalizeTask : public ExecutorTask {
+public:
+	HashAggregateDistinctFinalizeTask(Pipeline &pipeline, shared_ptr<Event> event_p, const PhysicalHashAggregate &op,
+	                                  HashAggregateGlobalSinkState &state_p)
+	    : ExecutorTask(pipeline.executor), pipeline(pipeline), event(std::move(event_p)), op(op), gstate(state_p) {
+	}
 
-class HashDistinctAggregateFinalizeTask : public ExecutorTask {
 public:
-	HashDistinctAggregateFinalizeTask(Pipeline &pipeline, shared_ptr<Event> event_p, HashAggregateGlobalState &state_p,
-	                                  ClientContext &context, const PhysicalHashAggregate &op,
-	                                  vector<vector<unique_ptr<GlobalSourceState>>> &global_sources_p)
-	    : ExecutorTask(pipeline.executor), pipeline(pipeline), event(std::move(event_p)), gstate(state_p),
-	      context(context), op(op), global_sources(global_sources_p) {
-	}
-
-	void AggregateDistinctGrouping(DistinctAggregateCollectionInfo &info,
-	                               const HashAggregateGroupingData &grouping_data,
-	                               HashAggregateGroupingGlobalState &grouping_state, idx_t grouping_idx) {
-		auto &aggregates = info.aggregates;
-		auto &data = *grouping_data.distinct_data;
-		auto &state = *grouping_state.distinct_state;
-		auto &table_state = *grouping_state.table_state;
-
-		ThreadContext temp_thread_context(context);
-		ExecutionContext temp_exec_context(context, temp_thread_context, &pipeline);
-
-		auto temp_local_state = grouping_data.table_data.GetLocalSinkState(temp_exec_context);
-
-		// Create a chunk that mimics the 'input' chunk in Sink, for storing the group vectors
-		DataChunk group_chunk;
-		if (!op.input_group_types.empty()) {
-			group_chunk.Initialize(context, op.input_group_types);
-		}
+	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override;
 
-		auto &groups = op.grouped_aggregate_data.groups;
-		const idx_t group_by_size = groups.size();
+private:
+	void AggregateDistinctGrouping(const idx_t grouping_idx);
 
-		DataChunk aggregate_input_chunk;
-		if (!gstate.payload_types.empty()) {
-			aggregate_input_chunk.Initialize(context, gstate.payload_types);
-		}
+private:
+	Pipeline &pipeline;
+	shared_ptr<Event> event;
 
-		idx_t payload_idx;
-		idx_t next_payload_idx = 0;
+	const PhysicalHashAggregate &op;
+	HashAggregateGlobalSinkState &gstate;
+};
+
+void HashAggregateDistinctFinalizeEvent::Schedule() {
+	CreateGlobalSources();
+
+	const idx_t n_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	vector<shared_ptr<Task>> tasks;
+	for (idx_t i = 0; i < n_threads; i++) {
+		tasks.push_back(make_uniq<HashAggregateDistinctFinalizeTask>(*pipeline, shared_from_this(), op, gstate));
+	}
+	SetTasks(std::move(tasks));
+}
 
-		for (idx_t i = 0; i < op.grouped_aggregate_data.aggregates.size(); i++) {
-			auto &aggregate = aggregates[i]->Cast<BoundAggregateExpression>();
+void HashAggregateDistinctFinalizeEvent::CreateGlobalSources() {
+	auto &aggregates = op.grouped_aggregate_data.aggregates;
+	global_source_states.reserve(op.groupings.size());
+	for (idx_t grouping_idx = 0; grouping_idx < op.groupings.size(); grouping_idx++) {
+		auto &grouping = op.groupings[grouping_idx];
+		auto &distinct_data = *grouping.distinct_data;
 
-			// Forward the payload idx
-			payload_idx = next_payload_idx;
-			next_payload_idx = payload_idx + aggregate.children.size();
+		vector<unique_ptr<GlobalSourceState>> aggregate_sources;
+		aggregate_sources.reserve(aggregates.size());
+		for (idx_t agg_idx = 0; agg_idx < aggregates.size(); agg_idx++) {
+			auto &aggregate = aggregates[agg_idx];
+			auto &aggr = aggregate->Cast<BoundAggregateExpression>();
 
-			// If aggregate is not distinct, skip it
-			if (!data.IsDistinct(i)) {
+			if (!aggr.IsDistinct()) {
+				aggregate_sources.push_back(nullptr);
 				continue;
 			}
-			D_ASSERT(data.info.table_map.count(i));
-			auto table_idx = data.info.table_map.at(i);
-			auto &radix_table_p = data.radix_tables[table_idx];
-
-			// Create a duplicate of the output_chunk, because of multi-threading we cant alter the original
-			DataChunk output_chunk;
-			output_chunk.Initialize(context, state.distinct_output_chunks[table_idx]->GetTypes());
-
-			auto &global_source = global_sources[grouping_idx][i];
-			auto local_source = radix_table_p->GetLocalSourceState(temp_exec_context);
-
-			// Fetch all the data from the aggregate ht, and Sink it into the main ht
-			while (true) {
-				output_chunk.Reset();
-				group_chunk.Reset();
-				aggregate_input_chunk.Reset();
-
-				InterruptState interrupt_state;
-				OperatorSourceInput source_input {*global_source, *local_source, interrupt_state};
-				auto res = radix_table_p->GetData(temp_exec_context, output_chunk, *state.radix_states[table_idx],
-				                                  source_input);
-
-				if (res == SourceResultType::FINISHED) {
-					D_ASSERT(output_chunk.size() == 0);
-					break;
-				} else if (res == SourceResultType::BLOCKED) {
-					throw InternalException(
-					    "Unexpected interrupt from radix table GetData in HashDistinctAggregateFinalizeTask");
-				}
-
-				auto &grouped_aggregate_data = *data.grouped_aggregate_data[table_idx];
-
-				for (idx_t group_idx = 0; group_idx < group_by_size; group_idx++) {
-					auto &group = grouped_aggregate_data.groups[group_idx];
-					auto &bound_ref_expr = group->Cast<BoundReferenceExpression>();
-					group_chunk.data[bound_ref_expr.index].Reference(output_chunk.data[group_idx]);
-				}
-				group_chunk.SetCardinality(output_chunk);
-
-				for (idx_t child_idx = 0; child_idx < grouped_aggregate_data.groups.size() - group_by_size;
-				     child_idx++) {
-					aggregate_input_chunk.data[payload_idx + child_idx].Reference(
-					    output_chunk.data[group_by_size + child_idx]);
-				}
-				aggregate_input_chunk.SetCardinality(output_chunk);
-
-				// Sink it into the main ht
-				OperatorSinkInput sink_input {table_state, *temp_local_state, interrupt_state};
-				grouping_data.table_data.Sink(temp_exec_context, group_chunk, sink_input, aggregate_input_chunk, {i});
-			}
+			D_ASSERT(distinct_data.info.table_map.count(agg_idx));
+
+			auto table_idx = distinct_data.info.table_map.at(agg_idx);
+			auto &radix_table_p = distinct_data.radix_tables[table_idx];
+			aggregate_sources.push_back(radix_table_p->GetGlobalSourceState(context));
 		}
-		grouping_data.table_data.Combine(temp_exec_context, table_state, *temp_local_state);
+		global_source_states.push_back(std::move(aggregate_sources));
 	}
+}
 
-	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override {
-		D_ASSERT(op.distinct_collection_info);
-		auto &info = *op.distinct_collection_info;
-		for (idx_t i = 0; i < op.groupings.size(); i++) {
-			auto &grouping = op.groupings[i];
-			auto &grouping_state = gstate.grouping_states[i];
-			AggregateDistinctGrouping(info, grouping, grouping_state, i);
-		}
-		event->FinishTask();
-		return TaskExecutionResult::TASK_FINISHED;
+void HashAggregateDistinctFinalizeEvent::FinishEvent() {
+	// Now that everything is added to the main ht, we can actually finalize
+	auto new_event = make_shared<HashAggregateFinalizeEvent>(context, pipeline.get(), op, gstate);
+	this->InsertEvent(std::move(new_event));
+}
+
+TaskExecutionResult HashAggregateDistinctFinalizeTask::ExecuteTask(TaskExecutionMode mode) {
+	for (idx_t grouping_idx = 0; grouping_idx < op.groupings.size(); grouping_idx++) {
+		AggregateDistinctGrouping(grouping_idx);
 	}
+	event->FinishTask();
+	return TaskExecutionResult::TASK_FINISHED;
+}
 
-private:
-	Pipeline &pipeline;
-	shared_ptr<Event> event;
-	HashAggregateGlobalState &gstate;
-	ClientContext &context;
-	const PhysicalHashAggregate &op;
-	vector<vector<unique_ptr<GlobalSourceState>>> &global_sources;
-};
+void HashAggregateDistinctFinalizeTask::AggregateDistinctGrouping(const idx_t grouping_idx) {
+	D_ASSERT(op.distinct_collection_info);
+	auto &info = *op.distinct_collection_info;
 
-//! DISTINCT FINALIZE EVENT
+	auto &grouping_data = op.groupings[grouping_idx];
+	auto &grouping_state = gstate.grouping_states[grouping_idx];
+	D_ASSERT(grouping_state.distinct_state);
+	auto &distinct_state = *grouping_state.distinct_state;
+	auto &distinct_data = *grouping_data.distinct_data;
 
-// TODO: Create tasks and run these in parallel instead of doing this all in Schedule, single threaded
-class HashDistinctAggregateFinalizeEvent : public BasePipelineEvent {
-public:
-	HashDistinctAggregateFinalizeEvent(const PhysicalHashAggregate &op_p, HashAggregateGlobalState &gstate_p,
-	                                   Pipeline &pipeline_p, ClientContext &context)
-	    : BasePipelineEvent(pipeline_p), op(op_p), gstate(gstate_p), context(context) {
-	}
-	const PhysicalHashAggregate &op;
-	HashAggregateGlobalState &gstate;
-	ClientContext &context;
-	//! The GlobalSourceStates for all the radix tables of the distinct aggregates
-	vector<vector<unique_ptr<GlobalSourceState>>> global_sources;
+	auto &aggregates = info.aggregates;
 
-public:
-	void Schedule() override {
-		global_sources = CreateGlobalSources();
-
-		vector<shared_ptr<Task>> tasks;
-		auto &scheduler = TaskScheduler::GetScheduler(context);
-		auto number_of_threads = scheduler.NumberOfThreads();
-		tasks.reserve(number_of_threads);
-		for (int32_t i = 0; i < number_of_threads; i++) {
-			tasks.push_back(make_uniq<HashDistinctAggregateFinalizeTask>(*pipeline, shared_from_this(), gstate, context,
-			                                                             op, global_sources));
-		}
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
+	// Thread-local contexts
+	ThreadContext thread_context(executor.context);
+	ExecutionContext execution_context(executor.context, thread_context, &pipeline);
+
+	// Sink state to sink into global HTs
+	InterruptState interrupt_state;
+	auto &global_sink_state = *grouping_state.table_state;
+	auto local_sink_state = grouping_data.table_data.GetLocalSinkState(execution_context);
+	OperatorSinkInput sink_input {global_sink_state, *local_sink_state, interrupt_state};
+
+	// Create a chunk that mimics the 'input' chunk in Sink, for storing the group vectors
+	DataChunk group_chunk;
+	if (!op.input_group_types.empty()) {
+		group_chunk.Initialize(executor.context, op.input_group_types);
 	}
 
-	void FinishEvent() override {
-		//! Now that everything is added to the main ht, we can actually finalize
-		auto new_event = make_shared<HashAggregateFinalizeEvent>(op, gstate, pipeline.get(), context);
-		this->InsertEvent(std::move(new_event));
+	auto &groups = op.grouped_aggregate_data.groups;
+	const idx_t group_by_size = groups.size();
+
+	DataChunk aggregate_input_chunk;
+	if (!gstate.payload_types.empty()) {
+		aggregate_input_chunk.Initialize(executor.context, gstate.payload_types);
 	}
 
-private:
-	vector<vector<unique_ptr<GlobalSourceState>>> CreateGlobalSources() {
-		vector<vector<unique_ptr<GlobalSourceState>>> grouping_sources;
-		grouping_sources.reserve(op.groupings.size());
-		for (idx_t grouping_idx = 0; grouping_idx < op.groupings.size(); grouping_idx++) {
-			auto &grouping = op.groupings[grouping_idx];
-			auto &data = *grouping.distinct_data;
-
-			vector<unique_ptr<GlobalSourceState>> aggregate_sources;
-			aggregate_sources.reserve(op.grouped_aggregate_data.aggregates.size());
-
-			for (idx_t i = 0; i < op.grouped_aggregate_data.aggregates.size(); i++) {
-				auto &aggregate = op.grouped_aggregate_data.aggregates[i];
-				auto &aggr = aggregate->Cast<BoundAggregateExpression>();
-
-				if (!aggr.IsDistinct()) {
-					aggregate_sources.push_back(nullptr);
-					continue;
-				}
-
-				D_ASSERT(data.info.table_map.count(i));
-				auto table_idx = data.info.table_map.at(i);
-				auto &radix_table_p = data.radix_tables[table_idx];
-				aggregate_sources.push_back(radix_table_p->GetGlobalSourceState(context));
-			}
-			grouping_sources.push_back(std::move(aggregate_sources));
+	auto &finalize_event = event->Cast<HashAggregateDistinctFinalizeEvent>();
+
+	idx_t payload_idx;
+	idx_t next_payload_idx = 0;
+	for (idx_t agg_idx = 0; agg_idx < op.grouped_aggregate_data.aggregates.size(); agg_idx++) {
+		auto &aggregate = aggregates[agg_idx]->Cast<BoundAggregateExpression>();
+
+		// Forward the payload idx
+		payload_idx = next_payload_idx;
+		next_payload_idx = payload_idx + aggregate.children.size();
+
+		// If aggregate is not distinct, skip it
+		if (!distinct_data.IsDistinct(agg_idx)) {
+			continue;
 		}
-		return grouping_sources;
-	}
-};
 
-//! DISTINCT COMBINE EVENT
+		D_ASSERT(distinct_data.info.table_map.count(agg_idx));
+		const auto &table_idx = distinct_data.info.table_map.at(agg_idx);
+		auto &radix_table = distinct_data.radix_tables[table_idx];
 
-class HashDistinctCombineFinalizeEvent : public BasePipelineEvent {
-public:
-	HashDistinctCombineFinalizeEvent(const PhysicalHashAggregate &op_p, HashAggregateGlobalState &gstate_p,
-	                                 Pipeline &pipeline_p, ClientContext &client)
-	    : BasePipelineEvent(pipeline_p), op(op_p), gstate(gstate_p), client(client) {
-	}
+		auto &sink = *distinct_state.radix_states[table_idx];
+		auto local_source = radix_table->GetLocalSourceState(execution_context);
+		OperatorSourceInput source_input {*finalize_event.global_source_states[grouping_idx][agg_idx], *local_source,
+		                                  interrupt_state};
 
-	const PhysicalHashAggregate &op;
-	HashAggregateGlobalState &gstate;
-	ClientContext &client;
+		// Create a duplicate of the output_chunk, because of multi-threading we cant alter the original
+		DataChunk output_chunk;
+		output_chunk.Initialize(executor.context, distinct_state.distinct_output_chunks[table_idx]->GetTypes());
+
+		// Fetch all the data from the aggregate ht, and Sink it into the main ht
+		while (true) {
+			output_chunk.Reset();
+			group_chunk.Reset();
+			aggregate_input_chunk.Reset();
+
+			auto res = radix_table->GetData(execution_context, output_chunk, sink, source_input);
+			if (res == SourceResultType::FINISHED) {
+				D_ASSERT(output_chunk.size() == 0);
+				break;
+			} else if (res == SourceResultType::BLOCKED) {
+				throw InternalException(
+				    "Unexpected interrupt from radix table GetData in HashAggregateDistinctFinalizeTask");
+			}
 
-public:
-	void Schedule() override {
-		vector<shared_ptr<Task>> tasks;
-		for (idx_t i = 0; i < op.groupings.size(); i++) {
-			auto &grouping = op.groupings[i];
-			auto &distinct_data = *grouping.distinct_data;
-			auto &distinct_state = *gstate.grouping_states[i].distinct_state;
-			for (idx_t table_idx = 0; table_idx < distinct_data.radix_tables.size(); table_idx++) {
-				if (!distinct_data.radix_tables[table_idx]) {
-					continue;
-				}
-				distinct_data.radix_tables[table_idx]->ScheduleTasks(pipeline->executor, shared_from_this(),
-				                                                     *distinct_state.radix_states[table_idx], tasks);
+			auto &grouped_aggregate_data = *distinct_data.grouped_aggregate_data[table_idx];
+			for (idx_t group_idx = 0; group_idx < group_by_size; group_idx++) {
+				auto &group = grouped_aggregate_data.groups[group_idx];
+				auto &bound_ref_expr = group->Cast<BoundReferenceExpression>();
+				group_chunk.data[bound_ref_expr.index].Reference(output_chunk.data[group_idx]);
 			}
-		}
+			group_chunk.SetCardinality(output_chunk);
 
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
-	}
+			for (idx_t child_idx = 0; child_idx < grouped_aggregate_data.groups.size() - group_by_size; child_idx++) {
+				aggregate_input_chunk.data[payload_idx + child_idx].Reference(
+				    output_chunk.data[group_by_size + child_idx]);
+			}
+			aggregate_input_chunk.SetCardinality(output_chunk);
 
-	void FinishEvent() override {
-		//! Now that all tables are combined, it's time to do the distinct aggregations
-		auto new_event = make_shared<HashDistinctAggregateFinalizeEvent>(op, gstate, *pipeline, client);
-		this->InsertEvent(std::move(new_event));
+			// Sink it into the main ht
+			grouping_data.table_data.Sink(execution_context, group_chunk, sink_input, aggregate_input_chunk, {agg_idx});
+		}
 	}
-};
-
-//! FINALIZE
+	grouping_data.table_data.Combine(execution_context, global_sink_state, *local_sink_state);
+}
 
 SinkFinalizeType PhysicalHashAggregate::FinalizeDistinct(Pipeline &pipeline, Event &event, ClientContext &context,
                                                          GlobalSinkState &gstate_p) const {
-	auto &gstate = gstate_p.Cast<HashAggregateGlobalState>();
+	auto &gstate = gstate_p.Cast<HashAggregateGlobalSinkState>();
 	D_ASSERT(distinct_collection_info);
 
-	bool any_partitioned = false;
 	for (idx_t i = 0; i < groupings.size(); i++) {
 		auto &grouping = groupings[i];
 		auto &distinct_data = *grouping.distinct_data;
@@ -773,28 +724,17 @@ SinkFinalizeType PhysicalHashAggregate::FinalizeDistinct(Pipeline &pipeline, Eve
 			}
 			auto &radix_table = distinct_data.radix_tables[table_idx];
 			auto &radix_state = *distinct_state.radix_states[table_idx];
-			bool partitioned = radix_table->Finalize(context, radix_state);
-			if (partitioned) {
-				any_partitioned = true;
-			}
+			radix_table->Finalize(context, radix_state);
 		}
 	}
-	if (any_partitioned) {
-		// If any of the groupings are partitioned then we first need to combine those, then aggregate
-		auto new_event = make_shared<HashDistinctCombineFinalizeEvent>(*this, gstate, pipeline, context);
-		event.InsertEvent(std::move(new_event));
-	} else {
-		// Hashtables aren't partitioned, they dont need to be joined first
-		// so we can already compute the aggregate
-		auto new_event = make_shared<HashDistinctAggregateFinalizeEvent>(*this, gstate, pipeline, context);
-		event.InsertEvent(std::move(new_event));
-	}
+	auto new_event = make_shared<HashAggregateDistinctFinalizeEvent>(context, pipeline, *this, gstate);
+	event.InsertEvent(std::move(new_event));
 	return SinkFinalizeType::READY;
 }
 
 SinkFinalizeType PhysicalHashAggregate::FinalizeInternal(Pipeline &pipeline, Event &event, ClientContext &context,
                                                          GlobalSinkState &gstate_p, bool check_distinct) const {
-	auto &gstate = gstate_p.Cast<HashAggregateGlobalState>();
+	auto &gstate = gstate_p.Cast<HashAggregateGlobalSinkState>();
 
 	if (check_distinct && distinct_collection_info) {
 		// There are distinct aggregates
@@ -803,19 +743,10 @@ SinkFinalizeType PhysicalHashAggregate::FinalizeInternal(Pipeline &pipeline, Eve
 		return FinalizeDistinct(pipeline, event, context, gstate_p);
 	}
 
-	bool any_partitioned = false;
 	for (idx_t i = 0; i < groupings.size(); i++) {
 		auto &grouping = groupings[i];
 		auto &grouping_gstate = gstate.grouping_states[i];
-
-		bool is_partitioned = grouping.table_data.Finalize(context, *grouping_gstate.table_state);
-		if (is_partitioned) {
-			any_partitioned = true;
-		}
-	}
-	if (any_partitioned) {
-		auto new_event = make_shared<HashAggregateMergeEvent>(*this, gstate, &pipeline);
-		event.InsertEvent(std::move(new_event));
+		grouping.table_data.Finalize(context, *grouping_gstate.table_state);
 	}
 	return SinkFinalizeType::READY;
 }
@@ -828,10 +759,9 @@ SinkFinalizeType PhysicalHashAggregate::Finalize(Pipeline &pipeline, Event &even
 //===--------------------------------------------------------------------===//
 // Source
 //===--------------------------------------------------------------------===//
-class PhysicalHashAggregateGlobalSourceState : public GlobalSourceState {
+class HashAggregateGlobalSourceState : public GlobalSourceState {
 public:
-	PhysicalHashAggregateGlobalSourceState(ClientContext &context, const PhysicalHashAggregate &op)
-	    : op(op), state_index(0) {
+	HashAggregateGlobalSourceState(ClientContext &context, const PhysicalHashAggregate &op) : op(op), state_index(0) {
 		for (auto &grouping : op.groupings) {
 			auto &rt = grouping.table_data;
 			radix_states.push_back(rt.GetGlobalSourceState(context));
@@ -851,24 +781,24 @@ public:
 			return 1;
 		}
 
-		auto &ht_state = op.sink_state->Cast<HashAggregateGlobalState>();
+		auto &ht_state = op.sink_state->Cast<HashAggregateGlobalSinkState>();
 		idx_t count = 0;
 		for (size_t sidx = 0; sidx < op.groupings.size(); ++sidx) {
 			auto &grouping = op.groupings[sidx];
 			auto &grouping_gstate = ht_state.grouping_states[sidx];
-			count += grouping.table_data.Size(*grouping_gstate.table_state);
+			count += grouping.table_data.Count(*grouping_gstate.table_state);
 		}
 		return MaxValue<idx_t>(1, count / STANDARD_VECTOR_SIZE);
 	}
 };
 
 unique_ptr<GlobalSourceState> PhysicalHashAggregate::GetGlobalSourceState(ClientContext &context) const {
-	return make_uniq<PhysicalHashAggregateGlobalSourceState>(context, *this);
+	return make_uniq<HashAggregateGlobalSourceState>(context, *this);
 }
 
-class PhysicalHashAggregateLocalSourceState : public LocalSourceState {
+class HashAggregateLocalSourceState : public LocalSourceState {
 public:
-	explicit PhysicalHashAggregateLocalSourceState(ExecutionContext &context, const PhysicalHashAggregate &op) {
+	explicit HashAggregateLocalSourceState(ExecutionContext &context, const PhysicalHashAggregate &op) {
 		for (auto &grouping : op.groupings) {
 			auto &rt = grouping.table_data;
 			radix_states.push_back(rt.GetLocalSourceState(context));
@@ -880,14 +810,14 @@ public:
 
 unique_ptr<LocalSourceState> PhysicalHashAggregate::GetLocalSourceState(ExecutionContext &context,
                                                                         GlobalSourceState &gstate) const {
-	return make_uniq<PhysicalHashAggregateLocalSourceState>(context, *this);
+	return make_uniq<HashAggregateLocalSourceState>(context, *this);
 }
 
 SourceResultType PhysicalHashAggregate::GetData(ExecutionContext &context, DataChunk &chunk,
                                                 OperatorSourceInput &input) const {
-	auto &sink_gstate = sink_state->Cast<HashAggregateGlobalState>();
-	auto &gstate = input.global_state.Cast<PhysicalHashAggregateGlobalSourceState>();
-	auto &lstate = input.local_state.Cast<PhysicalHashAggregateLocalSourceState>();
+	auto &sink_gstate = sink_state->Cast<HashAggregateGlobalSinkState>();
+	auto &gstate = input.global_state.Cast<HashAggregateGlobalSourceState>();
+	auto &lstate = input.local_state.Cast<HashAggregateLocalSourceState>();
 	while (true) {
 		idx_t radix_idx = gstate.state_index;
 		if (radix_idx >= groupings.size()) {