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

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

@@ -37,7 +37,9 @@ PhysicalUngroupedAggregate::PhysicalUngroupedAggregate(vector<LogicalType> types
 //===--------------------------------------------------------------------===//
 struct AggregateState {
 	explicit AggregateState(const vector<unique_ptr<Expression>> &aggregate_expressions) {
-		for (auto &aggregate : aggregate_expressions) {
+		counts = make_uniq_array<atomic<idx_t>>(aggregate_expressions.size());
+		for (idx_t i = 0; i < aggregate_expressions.size(); i++) {
+			auto &aggregate = aggregate_expressions[i];
 			D_ASSERT(aggregate->GetExpressionClass() == ExpressionClass::BOUND_AGGREGATE);
 			auto &aggr = aggregate->Cast<BoundAggregateExpression>();
 			auto state = make_unsafe_uniq_array<data_t>(aggr.function.state_size());
@@ -46,7 +48,7 @@ struct AggregateState {
 			bind_data.push_back(aggr.bind_info.get());
 			destructors.push_back(aggr.function.destructor);
 #ifdef DEBUG
-			counts.push_back(0);
+			counts[i] = 0;
 #endif
 		}
 	}
@@ -77,12 +79,12 @@ struct AggregateState {
 	//! The destructors
 	vector<aggregate_destructor_t> destructors;
 	//! Counts (used for verification)
-	vector<idx_t> counts;
+	unique_array<atomic<idx_t>> counts;
 };
 
-class UngroupedAggregateGlobalState : public GlobalSinkState {
+class UngroupedAggregateGlobalSinkState : public GlobalSinkState {
 public:
-	UngroupedAggregateGlobalState(const PhysicalUngroupedAggregate &op, ClientContext &client)
+	UngroupedAggregateGlobalSinkState(const PhysicalUngroupedAggregate &op, ClientContext &client)
 	    : state(op.aggregates), finished(false), allocator(BufferAllocator::Get(client)) {
 		if (op.distinct_data) {
 			distinct_state = make_uniq<DistinctAggregateState>(*op.distinct_data, client);
@@ -101,13 +103,13 @@ public:
 	ArenaAllocator allocator;
 };
 
-class UngroupedAggregateLocalState : public LocalSinkState {
+class UngroupedAggregateLocalSinkState : public LocalSinkState {
 public:
-	UngroupedAggregateLocalState(const PhysicalUngroupedAggregate &op, const vector<LogicalType> &child_types,
-	                             GlobalSinkState &gstate_p, ExecutionContext &context)
+	UngroupedAggregateLocalSinkState(const PhysicalUngroupedAggregate &op, const vector<LogicalType> &child_types,
+	                                 GlobalSinkState &gstate_p, ExecutionContext &context)
 	    : allocator(BufferAllocator::Get(context.client)), state(op.aggregates), child_executor(context.client),
 	      aggregate_input_chunk(), filter_set() {
-		auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalState>();
+		auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalSinkState>();
 
 		auto &allocator = BufferAllocator::Get(context.client);
 		InitializeDistinctAggregates(op, gstate, context);
@@ -147,8 +149,8 @@ public:
 	void Reset() {
 		aggregate_input_chunk.Reset();
 	}
-	void InitializeDistinctAggregates(const PhysicalUngroupedAggregate &op, const UngroupedAggregateGlobalState &gstate,
-	                                  ExecutionContext &context) {
+	void InitializeDistinctAggregates(const PhysicalUngroupedAggregate &op,
+	                                  const UngroupedAggregateGlobalSinkState &gstate, ExecutionContext &context) {
 
 		if (!op.distinct_data) {
 			return;
@@ -185,19 +187,19 @@ bool PhysicalUngroupedAggregate::SinkOrderDependent() const {
 }
 
 unique_ptr<GlobalSinkState> PhysicalUngroupedAggregate::GetGlobalSinkState(ClientContext &context) const {
-	return make_uniq<UngroupedAggregateGlobalState>(*this, context);
+	return make_uniq<UngroupedAggregateGlobalSinkState>(*this, context);
 }
 
 unique_ptr<LocalSinkState> PhysicalUngroupedAggregate::GetLocalSinkState(ExecutionContext &context) const {
 	D_ASSERT(sink_state);
 	auto &gstate = *sink_state;
-	return make_uniq<UngroupedAggregateLocalState>(*this, children[0]->GetTypes(), gstate, context);
+	return make_uniq<UngroupedAggregateLocalSinkState>(*this, children[0]->GetTypes(), gstate, context);
 }
 
 void PhysicalUngroupedAggregate::SinkDistinct(ExecutionContext &context, DataChunk &chunk,
                                               OperatorSinkInput &input) const {
-	auto &sink = input.local_state.Cast<UngroupedAggregateLocalState>();
-	auto &global_sink = input.global_state.Cast<UngroupedAggregateGlobalState>();
+	auto &sink = input.local_state.Cast<UngroupedAggregateLocalSinkState>();
+	auto &global_sink = input.global_state.Cast<UngroupedAggregateGlobalSinkState>();
 	D_ASSERT(distinct_data);
 	auto &distinct_state = *global_sink.distinct_state;
 	auto &distinct_info = *distinct_collection_info;
@@ -239,7 +241,7 @@ void PhysicalUngroupedAggregate::SinkDistinct(ExecutionContext &context, DataChu
 
 SinkResultType PhysicalUngroupedAggregate::Sink(ExecutionContext &context, DataChunk &chunk,
                                                 OperatorSinkInput &input) const {
-	auto &sink = input.local_state.Cast<UngroupedAggregateLocalState>();
+	auto &sink = input.local_state.Cast<UngroupedAggregateLocalSinkState>();
 
 	// perform the aggregation inside the local state
 	sink.Reset();
@@ -296,23 +298,22 @@ SinkResultType PhysicalUngroupedAggregate::Sink(ExecutionContext &context, DataC
 }
 
 //===--------------------------------------------------------------------===//
-// Finalize
+// Combine
 //===--------------------------------------------------------------------===//
-
 void PhysicalUngroupedAggregate::CombineDistinct(ExecutionContext &context, OperatorSinkCombineInput &input) const {
-	auto &global_sink = input.global_state.Cast<UngroupedAggregateGlobalState>();
-	auto &source = input.local_state.Cast<UngroupedAggregateLocalState>();
+	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalSinkState>();
+	auto &lstate = input.local_state.Cast<UngroupedAggregateLocalSinkState>();
 
 	if (!distinct_data) {
 		return;
 	}
-	auto &distinct_state = global_sink.distinct_state;
+	auto &distinct_state = gstate.distinct_state;
 	auto table_count = distinct_data->radix_tables.size();
 	for (idx_t table_idx = 0; table_idx < table_count; table_idx++) {
 		D_ASSERT(distinct_data->radix_tables[table_idx]);
 		auto &radix_table = *distinct_data->radix_tables[table_idx];
 		auto &radix_global_sink = *distinct_state->radix_states[table_idx];
-		auto &radix_local_sink = *source.radix_states[table_idx];
+		auto &radix_local_sink = *lstate.radix_states[table_idx];
 
 		radix_table.Combine(context, radix_global_sink, radix_local_sink);
 	}
@@ -320,18 +321,17 @@ void PhysicalUngroupedAggregate::CombineDistinct(ExecutionContext &context, Oper
 
 SinkCombineResultType PhysicalUngroupedAggregate::Combine(ExecutionContext &context,
                                                           OperatorSinkCombineInput &input) const {
-	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalState>();
-	auto &source = input.local_state.Cast<UngroupedAggregateLocalState>();
+	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalSinkState>();
+	auto &lstate = input.local_state.Cast<UngroupedAggregateLocalSinkState>();
 	D_ASSERT(!gstate.finished);
 
 	// finalize: combine the local state into the global state
 	// all aggregates are combinable: we might be doing a parallel aggregate
 	// use the combine method to combine the partial aggregates
-	lock_guard<mutex> glock(gstate.lock);
-
-	OperatorSinkCombineInput distinct_input {gstate, source, input.interrupt_state};
+	OperatorSinkCombineInput distinct_input {gstate, lstate, input.interrupt_state};
 	CombineDistinct(context, distinct_input);
 
+	lock_guard<mutex> glock(gstate.lock);
 	for (idx_t aggr_idx = 0; aggr_idx < aggregates.size(); aggr_idx++) {
 		auto &aggregate = aggregates[aggr_idx]->Cast<BoundAggregateExpression>();
 
@@ -339,207 +339,236 @@ SinkCombineResultType PhysicalUngroupedAggregate::Combine(ExecutionContext &cont
 			continue;
 		}
 
-		Vector source_state(Value::POINTER(CastPointerToValue(source.state.aggregates[aggr_idx].get())));
+		Vector source_state(Value::POINTER(CastPointerToValue(lstate.state.aggregates[aggr_idx].get())));
 		Vector dest_state(Value::POINTER(CastPointerToValue(gstate.state.aggregates[aggr_idx].get())));
 
 		AggregateInputData aggr_input_data(aggregate.bind_info.get(), gstate.allocator);
 		aggregate.function.combine(source_state, dest_state, aggr_input_data, 1);
 #ifdef DEBUG
-		gstate.state.counts[aggr_idx] += source.state.counts[aggr_idx];
+		gstate.state.counts[aggr_idx] += lstate.state.counts[aggr_idx];
 #endif
 	}
-	source.allocator.Destroy();
+	lstate.allocator.Destroy();
 
 	auto &client_profiler = QueryProfiler::Get(context.client);
-	context.thread.profiler.Flush(*this, source.child_executor, "child_executor", 0);
+	context.thread.profiler.Flush(*this, lstate.child_executor, "child_executor", 0);
 	client_profiler.Flush(context.thread.profiler);
 
 	return SinkCombineResultType::FINISHED;
 }
 
-class UngroupedDistinctAggregateFinalizeTask : public ExecutorTask {
+//===--------------------------------------------------------------------===//
+// Finalize
+//===--------------------------------------------------------------------===//
+class UngroupedDistinctAggregateFinalizeEvent : public BasePipelineEvent {
 public:
-	UngroupedDistinctAggregateFinalizeTask(Executor &executor, shared_ptr<Event> event_p,
-	                                       UngroupedAggregateGlobalState &state_p, ClientContext &context,
-	                                       const PhysicalUngroupedAggregate &op)
-	    : ExecutorTask(executor), event(std::move(event_p)), gstate(state_p), context(context), op(op),
-	      allocator(BufferAllocator::Get(context)) {
+	UngroupedDistinctAggregateFinalizeEvent(ClientContext &context, const PhysicalUngroupedAggregate &op_p,
+	                                        UngroupedAggregateGlobalSinkState &gstate_p, Pipeline &pipeline_p)
+	    : BasePipelineEvent(pipeline_p), context(context), op(op_p), gstate(gstate_p), tasks_scheduled(0),
+	      tasks_done(0) {
 	}
 
-	void AggregateDistinct() {
-		D_ASSERT(gstate.distinct_state);
-		auto &aggregates = op.aggregates;
-		auto &distinct_state = *gstate.distinct_state;
-		auto &distinct_data = *op.distinct_data;
-
-		ThreadContext temp_thread_context(context);
-		ExecutionContext temp_exec_context(context, temp_thread_context, nullptr);
-
-		idx_t payload_idx = 0;
-		idx_t next_payload_idx = 0;
+public:
+	void Schedule() override;
 
-		for (idx_t agg_idx = 0; agg_idx < aggregates.size(); agg_idx++) {
-			auto &aggregate = aggregates[agg_idx]->Cast<BoundAggregateExpression>();
+private:
+	ClientContext &context;
 
-			// Forward the payload idx
-			payload_idx = next_payload_idx;
-			next_payload_idx = payload_idx + aggregate.children.size();
+	const PhysicalUngroupedAggregate &op;
+	UngroupedAggregateGlobalSinkState &gstate;
 
-			// If aggregate is not distinct, skip it
-			if (!distinct_data.IsDistinct(agg_idx)) {
-				continue;
-			}
+public:
+	mutex lock;
+	idx_t tasks_scheduled;
+	idx_t tasks_done;
 
-			DataChunk payload_chunk;
-
-			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];
-			auto &output_chunk = *distinct_state.distinct_output_chunks[table_idx];
-			auto &grouped_aggregate_data = *distinct_data.grouped_aggregate_data[table_idx];
-
-			payload_chunk.InitializeEmpty(grouped_aggregate_data.group_types);
-			payload_chunk.SetCardinality(0);
-
-			//! Create global and local state for the hashtable
-			auto global_source_state = radix_table_p->GetGlobalSourceState(context);
-			auto local_source_state = radix_table_p->GetLocalSourceState(temp_exec_context);
-
-			//! Retrieve the stored data from the hashtable
-			while (true) {
-				output_chunk.Reset();
-
-				InterruptState interrupt_state;
-				OperatorSourceInput source_input {*global_source_state, *local_source_state, interrupt_state};
-				auto res = radix_table_p->GetData(temp_exec_context, output_chunk,
-				                                  *distinct_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 UngroupedDistinctAggregateFinalizeTask");
-				}
-
-				// We dont need to resolve the filter, we already did this in Sink
-				idx_t payload_cnt = aggregate.children.size();
-				for (idx_t i = 0; i < payload_cnt; i++) {
-					payload_chunk.data[i].Reference(output_chunk.data[i]);
-				}
-				payload_chunk.SetCardinality(output_chunk);
-#ifdef DEBUG
-				gstate.state.counts[agg_idx] += payload_chunk.size();
-#endif
+	vector<unique_ptr<GlobalSourceState>> global_source_states;
+};
 
-				auto start_of_input = payload_cnt ? &payload_chunk.data[0] : nullptr;
-				//! Update the aggregate state
-				AggregateInputData aggr_input_data(aggregate.bind_info.get(), allocator);
-				aggregate.function.simple_update(start_of_input, aggr_input_data, payload_cnt,
-				                                 gstate.state.aggregates[agg_idx].get(), payload_chunk.size());
-			}
-		}
-		D_ASSERT(!gstate.finished);
-		gstate.finished = true;
+class UngroupedDistinctAggregateFinalizeTask : public ExecutorTask {
+public:
+	UngroupedDistinctAggregateFinalizeTask(Executor &executor, shared_ptr<Event> event_p,
+	                                       const PhysicalUngroupedAggregate &op,
+	                                       UngroupedAggregateGlobalSinkState &state_p)
+	    : ExecutorTask(executor), event(std::move(event_p)), op(op), gstate(state_p),
+	      allocator(BufferAllocator::Get(executor.context)) {
 	}
 
-	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override {
-		AggregateDistinct();
-		event->FinishTask();
-		return TaskExecutionResult::TASK_FINISHED;
-	}
+	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override;
+
+private:
+	void AggregateDistinct();
 
 private:
 	shared_ptr<Event> event;
-	UngroupedAggregateGlobalState &gstate;
-	ClientContext &context;
+
 	const PhysicalUngroupedAggregate &op;
+	UngroupedAggregateGlobalSinkState &gstate;
+
 	ArenaAllocator allocator;
 };
 
-// TODO: Create tasks and run these in parallel instead of doing this all in Schedule, single threaded
-class UngroupedDistinctAggregateFinalizeEvent : public BasePipelineEvent {
-public:
-	UngroupedDistinctAggregateFinalizeEvent(const PhysicalUngroupedAggregate &op_p,
-	                                        UngroupedAggregateGlobalState &gstate_p, Pipeline &pipeline_p,
-	                                        ClientContext &context)
-	    : BasePipelineEvent(pipeline_p), op(op_p), gstate(gstate_p), context(context) {
-	}
-	const PhysicalUngroupedAggregate &op;
-	UngroupedAggregateGlobalState &gstate;
-	ClientContext &context;
+void UngroupedDistinctAggregateFinalizeEvent::Schedule() {
+	D_ASSERT(gstate.distinct_state);
+	auto &aggregates = op.aggregates;
+	auto &distinct_data = *op.distinct_data;
 
-public:
-	void Schedule() override {
-		vector<shared_ptr<Task>> tasks;
-		tasks.push_back(make_uniq<UngroupedDistinctAggregateFinalizeTask>(pipeline->executor, shared_from_this(),
-		                                                                  gstate, context, op));
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
+	idx_t payload_idx = 0;
+	idx_t next_payload_idx = 0;
+	for (idx_t agg_idx = 0; agg_idx < 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)) {
+			global_source_states.push_back(nullptr);
+			continue;
+		}
+		D_ASSERT(distinct_data.info.table_map.count(agg_idx));
+
+		// Create global state for scanning
+		auto table_idx = distinct_data.info.table_map.at(agg_idx);
+		auto &radix_table_p = *distinct_data.radix_tables[table_idx];
+		global_source_states.push_back(radix_table_p.GetGlobalSourceState(context));
 	}
-};
 
-class UngroupedDistinctCombineFinalizeEvent : public BasePipelineEvent {
-public:
-	UngroupedDistinctCombineFinalizeEvent(const PhysicalUngroupedAggregate &op_p,
-	                                      UngroupedAggregateGlobalState &gstate_p, Pipeline &pipeline_p,
-	                                      ClientContext &client)
-	    : BasePipelineEvent(pipeline_p), op(op_p), gstate(gstate_p), client(client) {
+	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<UngroupedDistinctAggregateFinalizeTask>(pipeline->executor, shared_from_this(), op, gstate));
+		tasks_scheduled++;
 	}
+	SetTasks(std::move(tasks));
+}
 
-	const PhysicalUngroupedAggregate &op;
-	UngroupedAggregateGlobalState &gstate;
-	ClientContext &client;
+TaskExecutionResult UngroupedDistinctAggregateFinalizeTask::ExecuteTask(TaskExecutionMode mode) {
+	AggregateDistinct();
+	event->FinishTask();
+	return TaskExecutionResult::TASK_FINISHED;
+}
 
-public:
-	void Schedule() override {
-		auto &distinct_state = *gstate.distinct_state;
-		auto &distinct_data = *op.distinct_data;
-		vector<shared_ptr<Task>> tasks;
-		for (idx_t table_idx = 0; table_idx < distinct_data.radix_tables.size(); table_idx++) {
-			distinct_data.radix_tables[table_idx]->ScheduleTasks(pipeline->executor, shared_from_this(),
-			                                                     *distinct_state.radix_states[table_idx], tasks);
+void UngroupedDistinctAggregateFinalizeTask::AggregateDistinct() {
+	D_ASSERT(gstate.distinct_state);
+	auto &distinct_state = *gstate.distinct_state;
+	auto &distinct_data = *op.distinct_data;
+
+	// Create thread-local copy of aggregate state
+	auto &aggregates = op.aggregates;
+	AggregateState state(aggregates);
+
+	// Thread-local contexts
+	ThreadContext thread_context(executor.context);
+	ExecutionContext execution_context(executor.context, thread_context, nullptr);
+
+	auto &finalize_event = event->Cast<UngroupedDistinctAggregateFinalizeEvent>();
+
+	// Now loop through the distinct aggregates, scanning the distinct HTs
+	idx_t payload_idx = 0;
+	idx_t next_payload_idx = 0;
+	for (idx_t agg_idx = 0; agg_idx < 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;
+		}
+
+		const auto table_idx = distinct_data.info.table_map.at(agg_idx);
+		auto &radix_table = *distinct_data.radix_tables[table_idx];
+		auto lstate = radix_table.GetLocalSourceState(execution_context);
+
+		auto &sink = *distinct_state.radix_states[table_idx];
+		InterruptState interrupt_state;
+		OperatorSourceInput source_input {*finalize_event.global_source_states[agg_idx], *lstate, interrupt_state};
+
+		DataChunk output_chunk;
+		output_chunk.Initialize(executor.context, distinct_state.distinct_output_chunks[table_idx]->GetTypes());
+
+		DataChunk payload_chunk;
+		payload_chunk.InitializeEmpty(distinct_data.grouped_aggregate_data[table_idx]->group_types);
+		payload_chunk.SetCardinality(0);
+
+		AggregateInputData aggr_input_data(aggregate.bind_info.get(), allocator);
+		while (true) {
+			output_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 UngroupedDistinctAggregateFinalizeTask");
+			}
+
+			// We dont need to resolve the filter, we already did this in Sink
+			idx_t payload_cnt = aggregate.children.size();
+			for (idx_t i = 0; i < payload_cnt; i++) {
+				payload_chunk.data[i].Reference(output_chunk.data[i]);
+			}
+			payload_chunk.SetCardinality(output_chunk);
+
+#ifdef DEBUG
+			gstate.state.counts[agg_idx] += payload_chunk.size();
+#endif
+
+			// Update the aggregate state
+			auto start_of_input = payload_cnt ? &payload_chunk.data[0] : nullptr;
+			aggregate.function.simple_update(start_of_input, aggr_input_data, payload_cnt,
+			                                 state.aggregates[agg_idx].get(), payload_chunk.size());
 		}
-		D_ASSERT(!tasks.empty());
-		SetTasks(std::move(tasks));
 	}
 
-	void FinishEvent() override {
-		//! Now that all tables are combined, it's time to do the distinct aggregations
-		auto new_event = make_shared<UngroupedDistinctAggregateFinalizeEvent>(op, gstate, *pipeline, client);
-		this->InsertEvent(std::move(new_event));
+	// After scanning the distinct HTs, we can combine the thread-local agg states with the thread-global
+	lock_guard<mutex> guard(finalize_event.lock);
+	payload_idx = 0;
+	next_payload_idx = 0;
+	for (idx_t agg_idx = 0; agg_idx < aggregates.size(); agg_idx++) {
+		if (!distinct_data.IsDistinct(agg_idx)) {
+			continue;
+		}
+
+		auto &aggregate = aggregates[agg_idx]->Cast<BoundAggregateExpression>();
+		AggregateInputData aggr_input_data(aggregate.bind_info.get(), allocator);
+
+		Vector state_vec(Value::POINTER(CastPointerToValue(state.aggregates[agg_idx].get())));
+		Vector combined_vec(Value::POINTER(CastPointerToValue(gstate.state.aggregates[agg_idx].get())));
+		aggregate.function.combine(state_vec, combined_vec, aggr_input_data, 1);
 	}
-};
+
+	D_ASSERT(!gstate.finished);
+	if (++finalize_event.tasks_done == finalize_event.tasks_scheduled) {
+		gstate.finished = true;
+	}
+}
 
 SinkFinalizeType PhysicalUngroupedAggregate::FinalizeDistinct(Pipeline &pipeline, Event &event, ClientContext &context,
                                                               GlobalSinkState &gstate_p) const {
-	auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalState>();
+	auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalSinkState>();
 	D_ASSERT(distinct_data);
 	auto &distinct_state = *gstate.distinct_state;
 
-	bool any_partitioned = false;
 	for (idx_t table_idx = 0; table_idx < distinct_data->radix_tables.size(); table_idx++) {
 		auto &radix_table_p = distinct_data->radix_tables[table_idx];
 		auto &radix_state = *distinct_state.radix_states[table_idx];
-		bool partitioned = radix_table_p->Finalize(context, radix_state);
-		if (partitioned) {
-			any_partitioned = true;
-		}
-	}
-	if (any_partitioned) {
-		auto new_event = make_shared<UngroupedDistinctCombineFinalizeEvent>(*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 compute the aggregate already
-		auto new_event = make_shared<UngroupedDistinctAggregateFinalizeEvent>(*this, gstate, pipeline, context);
-		event.InsertEvent(std::move(new_event));
+		radix_table_p->Finalize(context, radix_state);
 	}
+	auto new_event = make_shared<UngroupedDistinctAggregateFinalizeEvent>(context, *this, gstate, pipeline);
+	event.InsertEvent(std::move(new_event));
 	return SinkFinalizeType::READY;
 }
 
 SinkFinalizeType PhysicalUngroupedAggregate::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
                                                       OperatorSinkFinalizeInput &input) const {
-	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalState>();
+	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalSinkState>();
 
 	if (distinct_data) {
 		return FinalizeDistinct(pipeline, event, context, input.global_state);
@@ -569,7 +598,7 @@ void VerifyNullHandling(DataChunk &chunk, AggregateState &state, const vector<un
 
 SourceResultType PhysicalUngroupedAggregate::GetData(ExecutionContext &context, DataChunk &chunk,
                                                      OperatorSourceInput &input) const {
-	auto &gstate = sink_state->Cast<UngroupedAggregateGlobalState>();
+	auto &gstate = sink_state->Cast<UngroupedAggregateGlobalSinkState>();
 	D_ASSERT(gstate.finished);
 
 	// initialize the result chunk with the aggregate values

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

@@ -14,7 +14,6 @@
 #include "duckdb/common/vector_operations/vector_operations.hpp"
 #include "duckdb/common/windows_undefs.hpp"
 #include "duckdb/execution/expression_executor.hpp"
-#include "duckdb/execution/partitionable_hashtable.hpp"
 #include "duckdb/execution/window_executor.hpp"
 #include "duckdb/execution/window_segment_tree.hpp"
 #include "duckdb/main/client_config.hpp"
@@ -222,6 +221,7 @@ WindowGlobalSourceState::WindowGlobalSourceState(ClientContext &context_p, Windo
 		}
 	} else {
 		built.resize(hash_groups.size());
+		idx_t batch_base = 0;
 		for (auto &hash_group : hash_groups) {
 			if (!hash_group) {
 				continue;
@@ -235,6 +235,9 @@ WindowGlobalSourceState::WindowGlobalSourceState(ClientContext &context_p, Windo
 			auto &sb = *global_sort_state.sorted_blocks[0];
 			auto &sd = *sb.payload_data;
 			tasks_remaining += sd.data_blocks.size();
+
+			hash_group->batch_base = batch_base;
+			batch_base += sd.data_blocks.size();
 		}
 	}
 }
@@ -436,6 +439,7 @@ public:
 	using ReadStates = vector<ReadStatePtr>;
 
 	explicit WindowLocalSourceState(WindowGlobalSourceState &gsource);
+	void UpdateBatchIndex();
 	bool NextPartition();
 	void Scan(DataChunk &chunk);
 
@@ -443,6 +447,8 @@ public:
 	WindowGlobalSourceState &gsource;
 	//! The current bin being processed
 	idx_t hash_bin;
+	//! The current batch index (for output reordering)
+	idx_t batch_index;
 	//! The current source being processed
 	optional_ptr<WindowPartitionSourceState> partition_source;
 	//! The read cursor
@@ -456,7 +462,7 @@ public:
 };
 
 WindowLocalSourceState::WindowLocalSourceState(WindowGlobalSourceState &gsource)
-    : gsource(gsource), hash_bin(gsource.built.size()) {
+    : gsource(gsource), hash_bin(gsource.built.size()), batch_index(0) {
 	auto &gsink = *gsource.gsink.global_partition;
 	auto &op = gsource.gsink.op;
 
@@ -564,6 +570,14 @@ WindowGlobalSourceState::Task WindowGlobalSourceState::NextTask(idx_t hash_bin)
 	return Task();
 }
 
+void WindowLocalSourceState::UpdateBatchIndex() {
+	D_ASSERT(partition_source);
+	D_ASSERT(scanner.get());
+
+	batch_index = partition_source->hash_group ? partition_source->hash_group->batch_base : 0;
+	batch_index += scanner->BlockIndex();
+}
+
 bool WindowLocalSourceState::NextPartition() {
 	//	Release old states before the source
 	scanner.reset();
@@ -578,6 +592,7 @@ bool WindowLocalSourceState::NextPartition() {
 		partition_source = task.first;
 		scanner = std::move(task.second);
 		hash_bin = partition_source->hash_bin;
+		UpdateBatchIndex();
 	}
 
 	for (auto &wexec : partition_source->executors) {
@@ -599,6 +614,8 @@ void WindowLocalSourceState::Scan(DataChunk &result) {
 			read_states.clear();
 			return;
 		}
+
+		UpdateBatchIndex();
 	}
 
 	const auto position = scanner->Scanned();
@@ -638,6 +655,23 @@ unique_ptr<GlobalSourceState> PhysicalWindow::GetGlobalSourceState(ClientContext
 	return make_uniq<WindowGlobalSourceState>(context, gsink);
 }
 
+bool PhysicalWindow::SupportsBatchIndex() const {
+	//	We can only preserve order for single partitioning
+	//	or work stealing causes out of order batch numbers
+	auto &wexpr = select_list[0]->Cast<BoundWindowExpression>();
+	return wexpr.partitions.empty() && !wexpr.orders.empty();
+}
+
+OrderPreservationType PhysicalWindow::SourceOrder() const {
+	return SupportsBatchIndex() ? OrderPreservationType::FIXED_ORDER : OrderPreservationType::NO_ORDER;
+}
+
+idx_t PhysicalWindow::GetBatchIndex(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate_p,
+                                    LocalSourceState &lstate_p) const {
+	auto &lstate = lstate_p.Cast<WindowLocalSourceState>();
+	return lstate.batch_index;
+}
+
 SourceResultType PhysicalWindow::GetData(ExecutionContext &context, DataChunk &chunk,
                                          OperatorSourceInput &input) const {
 	auto &lsource = input.local_state.Cast<WindowLocalSourceState>();