duckdb 1.3.4 → 1.4.1-dev2.0

package/src/duckdb/src/function/window/window_distinct_aggregator.cpp

@@ -1,11 +1,10 @@
 #include "duckdb/function/window/window_distinct_aggregator.hpp"
 
-#include "duckdb/common/sort/partition_state.hpp"
-#include "duckdb/common/sort/sort.hpp"
+#include "duckdb/common/sorting/sort.hpp"
 #include "duckdb/execution/merge_sort_tree.hpp"
 #include "duckdb/function/window/window_aggregate_states.hpp"
 #include "duckdb/planner/bound_result_modifier.hpp"
-#include "duckdb/planner/expression/bound_constant_expression.hpp"
+#include "duckdb/planner/expression/bound_reference_expression.hpp"
 #include "duckdb/planner/expression/bound_window_expression.hpp"
 
 #include <numeric>
@@ -13,6 +12,8 @@
 
 namespace duckdb {
 
+enum class WindowDistinctSortStage : uint8_t { INIT, COMBINE, FINALIZE, SORTED, FINISHED };
+
 //===--------------------------------------------------------------------===//
 // WindowDistinctAggregator
 //===--------------------------------------------------------------------===//
@@ -21,6 +22,10 @@ bool WindowDistinctAggregator::CanAggregate(const BoundWindowExpression &wexpr)
 		return false;
 	}
 
+	if (!wexpr.aggregate->CanAggregate()) {
+		return false;
+	}
+
 	return wexpr.distinct && wexpr.exclude_clause == WindowExcludeMode::NO_OTHER && wexpr.arg_orders.empty();
 }
 
@@ -39,7 +44,7 @@ public:
 	using ZippedTuple = std::tuple<idx_t, idx_t>;
 	using ZippedElements = vector<ZippedTuple>;
 
-	explicit WindowDistinctSortTree(WindowDistinctAggregatorGlobalState &gdastate, idx_t count) : gdastate(gdastate) {
+	WindowDistinctSortTree(WindowDistinctAggregatorGlobalState &gdastate, idx_t count) : gdastate(gdastate) {
 		//	Set up for parallel build
 		build_level = 0;
 		build_complete = 0;
@@ -59,31 +64,31 @@ protected:
 
 class WindowDistinctAggregatorGlobalState : public WindowAggregatorGlobalState {
 public:
-	using GlobalSortStatePtr = unique_ptr<GlobalSortState>;
-	using LocalSortStatePtr = unique_ptr<LocalSortState>;
 	using ZippedTuple = WindowDistinctSortTree::ZippedTuple;
 	using ZippedElements = WindowDistinctSortTree::ZippedElements;
 
 	WindowDistinctAggregatorGlobalState(ClientContext &context, const WindowDistinctAggregator &aggregator,
 	                                    idx_t group_count);
 
-	//! Compute the block starts
-	void MeasurePayloadBlocks();
 	//! Create a new local sort
-	optional_ptr<LocalSortState> InitializeLocalSort() const;
+	optional_ptr<LocalSinkState> InitializeLocalSort(ExecutionContext &context) const;
 
-	//! Patch up the previous index block boundaries
-	void PatchPrevIdcs();
-	bool TryPrepareNextStage(WindowDistinctAggregatorLocalState &lstate);
+	ArenaAllocator &CreateTreeAllocator() const {
+		lock_guard<mutex> tree_lock(lock);
+		tree_allocators.emplace_back(make_uniq<ArenaAllocator>(Allocator::DefaultAllocator()));
+		return *tree_allocators.back();
+	}
 
-	//	Single threaded sorting for now
-	ClientContext &context;
-	idx_t memory_per_thread;
+	bool TryPrepareNextStage(WindowDistinctAggregatorLocalState &lstate);
 
+	//! The tree allocators.
+	//! We need to hold onto them for the tree lifetime,
+	//! not the lifetime of the local state that constructed part of the tree
+	mutable vector<unique_ptr<ArenaAllocator>> tree_allocators;
 	//! Finalize guard
 	mutable mutex lock;
 	//! Finalize stage
-	atomic<PartitionSortStage> stage;
+	atomic<WindowDistinctSortStage> stage;
 	//! Tasks launched
 	idx_t total_tasks = 0;
 	//! Tasks launched
@@ -91,20 +96,18 @@ public:
 	//! Tasks landed
 	mutable atomic<idx_t> tasks_completed;
 
-	//! The sorted payload data types (partition index)
-	vector<LogicalType> payload_types;
 	//! The aggregate arguments + partition index
 	vector<LogicalType> sort_types;
 
 	//! Sorting operations
-	GlobalSortStatePtr global_sort;
-	//! Local sort set
-	mutable vector<LocalSortStatePtr> local_sorts;
-	//! The block starts (the scanner doesn't know this) plus the total count
-	vector<idx_t> block_starts;
-
-	//! The block boundary seconds
-	mutable ZippedElements seconds;
+	vector<idx_t> sort_cols;
+	unique_ptr<Sort> sort;
+	unique_ptr<GlobalSinkState> global_sink;
+	//! Local sort sets
+	mutable vector<unique_ptr<LocalSinkState>> local_sinks;
+	//! The resulting sorted data
+	unique_ptr<ColumnDataCollection> sorted;
+
 	//! The MST with the distinct back pointers
 	mutable MergeSortTree<ZippedTuple> zipped_tree;
 	//! The merge sort tree for the aggregate.
@@ -116,35 +119,29 @@ public:
 	vector<idx_t> levels_flat_start;
 };
 
-WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientContext &context,
+WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientContext &client,
                                                                          const WindowDistinctAggregator &aggregator,
                                                                          idx_t group_count)
-    : WindowAggregatorGlobalState(context, aggregator, group_count), context(aggregator.context),
-      stage(PartitionSortStage::INIT), tasks_assigned(0), tasks_completed(0), merge_sort_tree(*this, group_count),
-      levels_flat_native(aggr) {
-	payload_types.emplace_back(LogicalType::UBIGINT);
+    : WindowAggregatorGlobalState(client, aggregator, group_count), stage(WindowDistinctSortStage::INIT),
+      tasks_assigned(0), tasks_completed(0), merge_sort_tree(*this, group_count), levels_flat_native(aggr) {
 
 	//	1:	functionComputePrevIdcs(𝑖𝑛)
 	//	2:		sorted ← []
 	//	We sort the aggregate arguments and use the partition index as a tie-breaker.
 	//	TODO: Use a hash table?
 	sort_types = aggregator.arg_types;
-	for (const auto &type : payload_types) {
-		sort_types.emplace_back(type);
-	}
+	sort_types.emplace_back(LogicalType::UBIGINT);
 
+	//	All expressions will be precomputed for sharing, so we jsut need to reference the arguments
 	vector<BoundOrderByNode> orders;
 	for (const auto &type : sort_types) {
-		auto expr = make_uniq<BoundConstantExpression>(Value(type));
+		auto expr = make_uniq<BoundReferenceExpression>(type, orders.size());
 		orders.emplace_back(BoundOrderByNode(OrderType::ASCENDING, OrderByNullType::NULLS_FIRST, std::move(expr)));
+		sort_cols.emplace_back(sort_cols.size());
 	}
 
-	RowLayout payload_layout;
-	payload_layout.Initialize(payload_types);
-
-	global_sort = make_uniq<GlobalSortState>(context, orders, payload_layout);
-
-	memory_per_thread = PhysicalOperator::GetMaxThreadMemory(context);
+	sort = make_uniq<Sort>(client, orders, sort_types, sort_cols);
+	global_sink = sort->GetGlobalSinkState(client);
 
 	//	6:	prevIdcs ← []
 	//	7:	prevIdcs[0] ← “-”
@@ -176,36 +173,38 @@ WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientC
 	}
 }
 
-optional_ptr<LocalSortState> WindowDistinctAggregatorGlobalState::InitializeLocalSort() const {
+optional_ptr<LocalSinkState> WindowDistinctAggregatorGlobalState::InitializeLocalSort(ExecutionContext &context) const {
 	lock_guard<mutex> local_sort_guard(lock);
-	auto local_sort = make_uniq<LocalSortState>();
-	local_sort->Initialize(*global_sort, global_sort->buffer_manager);
+	auto local_sink = sort->GetLocalSinkState(context);
 	++tasks_assigned;
-	local_sorts.emplace_back(std::move(local_sort));
+	local_sinks.emplace_back(std::move(local_sink));
 
-	return local_sorts.back().get();
+	return local_sinks.back().get();
 }
 
 class WindowDistinctAggregatorLocalState : public WindowAggregatorLocalState {
 public:
-	explicit WindowDistinctAggregatorLocalState(const WindowDistinctAggregatorGlobalState &aggregator);
+	WindowDistinctAggregatorLocalState(ExecutionContext &context,
+	                                   const WindowDistinctAggregatorGlobalState &aggregator);
 
 	~WindowDistinctAggregatorLocalState() override {
 		statef.Destroy();
 	}
 
-	void Sink(DataChunk &sink_chunk, DataChunk &coll_chunk, idx_t input_idx, optional_ptr<SelectionVector> filter_sel,
-	          idx_t filtered);
-	void Finalize(WindowAggregatorGlobalState &gastate, CollectionPtr collection) override;
+	void Sink(ExecutionContext &context, DataChunk &sink_chunk, DataChunk &coll_chunk, idx_t input_idx,
+	          optional_ptr<SelectionVector> filter_sel, idx_t filtered, InterruptState &interrupt);
+	void Finalize(ExecutionContext &context, WindowAggregatorGlobalState &gastate, CollectionPtr collection) override;
 	void Sorted();
-	void ExecuteTask();
-	void Evaluate(const WindowDistinctAggregatorGlobalState &gdstate, const DataChunk &bounds, Vector &result,
-	              idx_t count, idx_t row_idx);
+	void ExecuteTask(ExecutionContext &context, WindowDistinctAggregatorGlobalState &gdstate);
+	void Evaluate(ExecutionContext &context, const WindowDistinctAggregatorGlobalState &gdstate,
+	              const DataChunk &bounds, Vector &result, idx_t count, idx_t row_idx);
 
+	//! The thread-local allocator for building the tree
+	ArenaAllocator &tree_allocator;
 	//! Thread-local sorting data
-	optional_ptr<LocalSortState> local_sort;
+	optional_ptr<LocalSinkState> local_sink;
 	//! Finalize stage
-	PartitionSortStage stage = PartitionSortStage::INIT;
+	WindowDistinctSortStage stage = WindowDistinctSortStage::INIT;
 	//! Finalize scan block index
 	idx_t block_idx;
 	//! Thread-local tree aggregation
@@ -220,9 +219,9 @@ protected:
 	void FlushStates();
 
 	//! The aggregator we are working with
-	const WindowDistinctAggregatorGlobalState &gastate;
+	const WindowDistinctAggregatorGlobalState &gdstate;
+	//! The sort input chunk
 	DataChunk sort_chunk;
-	DataChunk payload_chunk;
 	//! Reused result state container for the window functions
 	WindowAggregateStates statef;
 	//! A vector of pointers to "state", used for buffering intermediate aggregates
@@ -236,207 +235,173 @@ protected:
 };
 
 WindowDistinctAggregatorLocalState::WindowDistinctAggregatorLocalState(
-    const WindowDistinctAggregatorGlobalState &gastate)
-    : update_v(LogicalType::POINTER), source_v(LogicalType::POINTER), target_v(LogicalType::POINTER), gastate(gastate),
-      statef(gastate.aggr), statep(LogicalType::POINTER), statel(LogicalType::POINTER), flush_count(0) {
-	InitSubFrames(frames, gastate.aggregator.exclude_mode);
-	payload_chunk.Initialize(Allocator::DefaultAllocator(), gastate.payload_types);
+    ExecutionContext &context, const WindowDistinctAggregatorGlobalState &gdstate)
+    : WindowAggregatorLocalState(context), tree_allocator(gdstate.CreateTreeAllocator()),
+      update_v(LogicalType::POINTER), source_v(LogicalType::POINTER), target_v(LogicalType::POINTER), gdstate(gdstate),
+      statef(gdstate.aggr), statep(LogicalType::POINTER), statel(LogicalType::POINTER), flush_count(0) {
+	InitSubFrames(frames, gdstate.aggregator.exclude_mode);
 
-	sort_chunk.Initialize(Allocator::DefaultAllocator(), gastate.sort_types);
-	sort_chunk.data.back().Reference(payload_chunk.data[0]);
+	sort_chunk.Initialize(Allocator::DefaultAllocator(), gdstate.sort_types);
 
-	gastate.locals++;
+	gdstate.locals++;
 }
 
-unique_ptr<WindowAggregatorState> WindowDistinctAggregator::GetGlobalState(ClientContext &context, idx_t group_count,
-                                                                           const ValidityMask &partition_mask) const {
+unique_ptr<GlobalSinkState> WindowDistinctAggregator::GetGlobalState(ClientContext &context, idx_t group_count,
+                                                                     const ValidityMask &partition_mask) const {
 	return make_uniq<WindowDistinctAggregatorGlobalState>(context, *this, group_count);
 }
 
-void WindowDistinctAggregator::Sink(WindowAggregatorState &gsink, WindowAggregatorState &lstate, DataChunk &sink_chunk,
-                                    DataChunk &coll_chunk, idx_t input_idx, optional_ptr<SelectionVector> filter_sel,
-                                    idx_t filtered) {
-	WindowAggregator::Sink(gsink, lstate, sink_chunk, coll_chunk, input_idx, filter_sel, filtered);
+void WindowDistinctAggregator::Sink(ExecutionContext &context, DataChunk &sink_chunk, DataChunk &coll_chunk,
+                                    idx_t input_idx, optional_ptr<SelectionVector> filter_sel, idx_t filtered,
+                                    OperatorSinkInput &sink) {
+	WindowAggregator::Sink(context, sink_chunk, coll_chunk, input_idx, filter_sel, filtered, sink);
 
-	auto &ldstate = lstate.Cast<WindowDistinctAggregatorLocalState>();
-	ldstate.Sink(sink_chunk, coll_chunk, input_idx, filter_sel, filtered);
+	auto &ldstate = sink.local_state.Cast<WindowDistinctAggregatorLocalState>();
+	ldstate.Sink(context, sink_chunk, coll_chunk, input_idx, filter_sel, filtered, sink.interrupt_state);
 }
 
-void WindowDistinctAggregatorLocalState::Sink(DataChunk &sink_chunk, DataChunk &coll_chunk, idx_t input_idx,
-                                              optional_ptr<SelectionVector> filter_sel, idx_t filtered) {
+void WindowDistinctAggregatorLocalState::Sink(ExecutionContext &context, DataChunk &sink_chunk, DataChunk &coll_chunk,
+                                              idx_t input_idx, optional_ptr<SelectionVector> filter_sel, idx_t filtered,
+                                              InterruptState &interrupt) {
 	//	3: 	for i ← 0 to in.size do
 	//	4: 		sorted[i] ← (in[i], i)
 	const auto count = sink_chunk.size();
-	payload_chunk.Reset();
-	auto &sorted_vec = payload_chunk.data[0];
+	sort_chunk.Reset();
+	auto &sorted_vec = sort_chunk.data.back();
 	auto sorted = FlatVector::GetData<idx_t>(sorted_vec);
 	std::iota(sorted, sorted + count, input_idx);
 
 	// Our arguments are being fully materialised,
 	// but we also need them as sort keys.
-	auto &child_idx = gastate.aggregator.child_idx;
+	auto &child_idx = gdstate.aggregator.child_idx;
 	for (column_t c = 0; c < child_idx.size(); ++c) {
 		sort_chunk.data[c].Reference(coll_chunk.data[child_idx[c]]);
 	}
-	sort_chunk.data.back().Reference(sorted_vec);
 	sort_chunk.SetCardinality(sink_chunk);
-	payload_chunk.SetCardinality(sort_chunk);
 
 	//	Apply FILTER clause, if any
 	if (filter_sel) {
 		sort_chunk.Slice(*filter_sel, filtered);
-		payload_chunk.Slice(*filter_sel, filtered);
 	}
 
-	if (!local_sort) {
-		local_sort = gastate.InitializeLocalSort();
+	if (!local_sink) {
+		local_sink = gdstate.InitializeLocalSort(context);
 	}
 
-	local_sort->SinkChunk(sort_chunk, payload_chunk);
-
-	if (local_sort->SizeInBytes() > gastate.memory_per_thread) {
-		local_sort->Sort(*gastate.global_sort, true);
-	}
+	OperatorSinkInput sink {*gdstate.global_sink, *local_sink, interrupt};
+	gdstate.sort->Sink(context, sort_chunk, sink);
 }
 
-void WindowDistinctAggregatorLocalState::Finalize(WindowAggregatorGlobalState &gastate, CollectionPtr collection) {
-	WindowAggregatorLocalState::Finalize(gastate, collection);
+void WindowDistinctAggregatorLocalState::Finalize(ExecutionContext &context, WindowAggregatorGlobalState &gastate,
+                                                  CollectionPtr collection) {
+	WindowAggregatorLocalState::Finalize(context, gastate, collection);
 
 	//! Input data chunk, used for leaf segment aggregation
 	leaves.Initialize(Allocator::DefaultAllocator(), cursor->chunk.GetTypes());
 	sel.Initialize();
 }
 
-void WindowDistinctAggregatorLocalState::ExecuteTask() {
-	auto &global_sort = *gastate.global_sort;
+void WindowDistinctAggregatorLocalState::ExecuteTask(ExecutionContext &context,
+                                                     WindowDistinctAggregatorGlobalState &gdstate) {
+	PostIncrement<atomic<idx_t>> on_done(gdstate.tasks_completed);
+
 	switch (stage) {
-	case PartitionSortStage::SCAN:
-		global_sort.AddLocalState(*gastate.local_sorts[block_idx]);
+	case WindowDistinctSortStage::COMBINE: {
+		auto &local_sink = *gdstate.local_sinks[block_idx];
+		InterruptState interrupt_state;
+		OperatorSinkCombineInput combine {*gdstate.global_sink, local_sink, interrupt_state};
+		gdstate.sort->Combine(context, combine);
 		break;
-	case PartitionSortStage::MERGE: {
-		MergeSorter merge_sorter(global_sort, global_sort.buffer_manager);
-		merge_sorter.PerformInMergeRound();
+	}
+	case WindowDistinctSortStage::FINALIZE: {
+		//	5: Sort sorted lexicographically increasing
+		auto &sort = *gdstate.sort;
+		InterruptState interrupt;
+		OperatorSinkFinalizeInput finalize {*gdstate.global_sink, interrupt};
+		sort.Finalize(context.client, finalize);
+		auto sort_global = sort.GetGlobalSourceState(context.client, *gdstate.global_sink);
+		auto sort_local = sort.GetLocalSourceState(context, *sort_global);
+		OperatorSourceInput source {*sort_global, *sort_local, interrupt};
+		sort.MaterializeColumnData(context, source);
+		gdstate.sorted = sort.GetColumnData(source);
 		break;
 	}
-	case PartitionSortStage::SORTED:
+	case WindowDistinctSortStage::SORTED:
 		Sorted();
 		break;
 	default:
 		break;
 	}
-
-	++gastate.tasks_completed;
-}
-
-void WindowDistinctAggregatorGlobalState::MeasurePayloadBlocks() {
-	const auto &blocks = global_sort->sorted_blocks[0]->payload_data->data_blocks;
-	idx_t count = 0;
-	for (const auto &block : blocks) {
-		block_starts.emplace_back(count);
-		count += block->count;
-	}
-	block_starts.emplace_back(count);
 }
 
 bool WindowDistinctAggregatorGlobalState::TryPrepareNextStage(WindowDistinctAggregatorLocalState &lstate) {
 	lock_guard<mutex> stage_guard(lock);
 
 	switch (stage.load()) {
-	case PartitionSortStage::INIT:
-		//	5: Sort sorted lexicographically increasing
-		total_tasks = local_sorts.size();
+	case WindowDistinctSortStage::INIT:
+		total_tasks = local_sinks.size();
 		tasks_assigned = 0;
 		tasks_completed = 0;
-		lstate.stage = stage = PartitionSortStage::SCAN;
+		lstate.stage = stage = WindowDistinctSortStage::COMBINE;
 		lstate.block_idx = tasks_assigned++;
 		return true;
-	case PartitionSortStage::SCAN:
-		// Process all the local sorts
+	case WindowDistinctSortStage::COMBINE:
 		if (tasks_assigned < total_tasks) {
-			lstate.stage = PartitionSortStage::SCAN;
+			lstate.stage = WindowDistinctSortStage::COMBINE;
 			lstate.block_idx = tasks_assigned++;
 			return true;
 		} else if (tasks_completed < tasks_assigned) {
 			return false;
 		}
-		global_sort->PrepareMergePhase();
-		if (!(global_sort->sorted_blocks.size() / 2)) {
-			if (global_sort->sorted_blocks.empty()) {
-				lstate.stage = stage = PartitionSortStage::FINISHED;
-				return true;
-			}
-			MeasurePayloadBlocks();
-			seconds.resize(block_starts.size() - 1);
-			total_tasks = seconds.size();
-			tasks_completed = 0;
-			tasks_assigned = 0;
-			lstate.stage = stage = PartitionSortStage::SORTED;
-			lstate.block_idx = tasks_assigned++;
-			return true;
-		}
-		global_sort->InitializeMergeRound();
-		lstate.stage = stage = PartitionSortStage::MERGE;
-		total_tasks = locals;
-		tasks_assigned = 1;
+		// All combines are done, so move on to materialising the sorted data (1 task)
+		total_tasks = 1;
 		tasks_completed = 0;
+		tasks_assigned = 0;
+		lstate.stage = stage = WindowDistinctSortStage::FINALIZE;
+		lstate.block_idx = tasks_assigned++;
 		return true;
-	case PartitionSortStage::MERGE:
-		if (tasks_assigned < total_tasks) {
-			lstate.stage = PartitionSortStage::MERGE;
-			++tasks_assigned;
-			return true;
-		} else if (tasks_completed < tasks_assigned) {
+	case WindowDistinctSortStage::FINALIZE:
+		if (tasks_completed < tasks_assigned) {
+			//	Wait for the single task to finish
 			return false;
 		}
-		global_sort->CompleteMergeRound(true);
-		if (!(global_sort->sorted_blocks.size() / 2)) {
-			MeasurePayloadBlocks();
-			seconds.resize(block_starts.size() - 1);
-			total_tasks = seconds.size();
-			tasks_completed = 0;
-			tasks_assigned = 0;
-			lstate.stage = stage = PartitionSortStage::SORTED;
-			lstate.block_idx = tasks_assigned++;
-			return true;
-		}
-		global_sort->InitializeMergeRound();
-		lstate.stage = PartitionSortStage::MERGE;
-		total_tasks = locals;
-		tasks_assigned = 1;
+		//	Move on to building the tree in parallel
+		total_tasks = local_sinks.size();
 		tasks_completed = 0;
+		tasks_assigned = 0;
+		lstate.stage = stage = WindowDistinctSortStage::SORTED;
+		lstate.block_idx = tasks_assigned++;
 		return true;
-	case PartitionSortStage::SORTED:
+	case WindowDistinctSortStage::SORTED:
 		if (tasks_assigned < total_tasks) {
-			lstate.stage = PartitionSortStage::SORTED;
+			lstate.stage = WindowDistinctSortStage::SORTED;
 			lstate.block_idx = tasks_assigned++;
 			return true;
 		} else if (tasks_completed < tasks_assigned) {
-			lstate.stage = PartitionSortStage::FINISHED;
+			lstate.stage = WindowDistinctSortStage::FINISHED;
 			// Sleep while other tasks finish
 			return false;
 		}
-		// Last task patches the boundaries
-		PatchPrevIdcs();
 		break;
 	default:
 		break;
 	}
 
-	lstate.stage = stage = PartitionSortStage::FINISHED;
+	lstate.stage = stage = WindowDistinctSortStage::FINISHED;
 
 	return true;
 }
 
-void WindowDistinctAggregator::Finalize(WindowAggregatorState &gsink, WindowAggregatorState &lstate,
-                                        CollectionPtr collection, const FrameStats &stats) {
-	auto &gdsink = gsink.Cast<WindowDistinctAggregatorGlobalState>();
-	auto &ldstate = lstate.Cast<WindowDistinctAggregatorLocalState>();
-	ldstate.Finalize(gdsink, collection);
+void WindowDistinctAggregator::Finalize(ExecutionContext &context, CollectionPtr collection, const FrameStats &stats,
+                                        OperatorSinkInput &sink) {
+	auto &gdsink = sink.global_state.Cast<WindowDistinctAggregatorGlobalState>();
+	auto &ldstate = sink.local_state.Cast<WindowDistinctAggregatorLocalState>();
+	ldstate.Finalize(context, gdsink, collection);
 
 	// Sort, merge and build the tree in parallel
-	while (gdsink.stage.load() != PartitionSortStage::FINISHED) {
+	while (gdsink.stage.load() != WindowDistinctSortStage::FINISHED) {
 		if (gdsink.TryPrepareNextStage(ldstate)) {
-			ldstate.ExecuteTask();
+			ldstate.ExecuteTask(context, gdsink);
 		} else {
 			std::this_thread::yield();
 		}
@@ -452,89 +417,71 @@ void WindowDistinctAggregator::Finalize(WindowAggregatorState &gsink, WindowAggr
 
 void WindowDistinctAggregatorLocalState::Sorted() {
 	using ZippedTuple = WindowDistinctAggregatorGlobalState::ZippedTuple;
-	auto &global_sort = gastate.global_sort;
-	auto &prev_idcs = gastate.zipped_tree.LowestLevel();
-	auto &aggregator = gastate.aggregator;
-	auto &scan_chunk = payload_chunk;
+	auto &collection = *gdstate.sorted;
+	auto &prev_idcs = gdstate.zipped_tree.LowestLevel();
+	auto &aggregator = gdstate.aggregator;
 
-	auto scanner = make_uniq<PayloadScanner>(*global_sort, block_idx);
-	const auto in_size = gastate.block_starts.at(block_idx + 1);
-	scanner->Scan(scan_chunk);
-	idx_t scan_idx = 0;
+	// Find our chunk range
+	const auto block_begin = (block_idx * collection.ChunkCount()) / gdstate.total_tasks;
+	const auto block_end = ((block_idx + 1) * collection.ChunkCount()) / gdstate.total_tasks;
 
-	auto *input_idx = FlatVector::GetData<idx_t>(scan_chunk.data[0]);
-	idx_t i = 0;
+	const auto &scan_cols = gdstate.sort_cols;
+	const auto key_count = aggregator.arg_types.size();
 
-	SBIterator curr(*global_sort, ExpressionType::COMPARE_LESSTHAN);
-	SBIterator prev(*global_sort, ExpressionType::COMPARE_LESSTHAN);
-	auto prefix_layout = global_sort->sort_layout.GetPrefixComparisonLayout(aggregator.arg_types.size());
+	//	Setting up the first row is a bit tricky - we have to scan the first value ourselves
+	WindowCollectionChunkScanner scanner(collection, scan_cols, block_begin ? block_begin - 1 : 0);
+	auto &scanned = scanner.chunk;
+	if (!scanner.Scan()) {
+		return;
+	}
 
-	const auto block_begin = gastate.block_starts.at(block_idx);
+	idx_t prev_i = 0;
 	if (!block_begin) {
 		// First block, so set up initial sentinel
-		i = input_idx[scan_idx++];
-		prev_idcs[i] = ZippedTuple(0, i);
-		std::get<0>(gastate.seconds[block_idx]) = i;
+		auto input_idx = FlatVector::GetData<idx_t>(scanned.data.back());
+		prev_i = input_idx[0];
+		prev_idcs[prev_i] = ZippedTuple(0, prev_i);
 	} else {
 		// Move to the to end of the previous block
 		// so we can record the comparison result for the first row
-		curr.SetIndex(block_begin - 1);
-		prev.SetIndex(block_begin - 1);
-		scan_idx = 0;
-		std::get<0>(gastate.seconds[block_idx]) = input_idx[scan_idx];
+		auto input_idx = FlatVector::GetData<idx_t>(scanned.data.back());
+		auto scan_idx = scanned.size() - 1;
+		prev_i = input_idx[scan_idx];
 	}
 
 	//	8:	for i ← 1 to in.size do
-	for (++curr; curr.GetIndex() < in_size; ++curr, ++prev) {
-		//	Scan second one chunk at a time
-		//	Note the scan is one behind the iterators
-		if (scan_idx >= scan_chunk.size()) {
-			scan_chunk.Reset();
-			scanner->Scan(scan_chunk);
-			scan_idx = 0;
-			input_idx = FlatVector::GetData<idx_t>(scan_chunk.data[0]);
-		}
-		auto second = i;
-		i = input_idx[scan_idx++];
+	WindowDeltaScanner(collection, block_begin, block_end, scan_cols, key_count,
+	                   [&](const idx_t row_idx, DataChunk &prev, DataChunk &curr, const idx_t ndistinct,
+	                       SelectionVector &distinct, const SelectionVector &matching) {
+		                   const auto count = MinValue<idx_t>(prev.size(), curr.size());
+
+		                   // The input index has probably been sliced.
+		                   UnifiedVectorFormat input_format;
+		                   curr.data.back().ToUnifiedFormat(count, input_format);
+		                   auto input_idx = UnifiedVectorFormat::GetData<idx_t>(input_format);
+
+		                   const auto nmatch = count - ndistinct;
+		                   //	9:	if sorted[i].first == sorted[i-1].first then
+		                   //	10:		prevIdcs[i] ← sorted[i-1].second
+		                   for (idx_t j = 0; j < nmatch; ++j) {
+			                   auto scan_idx = matching.get_index(j);
+			                   auto i = input_idx[input_format.sel->get_index(scan_idx)];
+			                   auto second = scan_idx ? input_idx[input_format.sel->get_index(scan_idx - 1)] : prev_i;
+			                   prev_idcs[i] = ZippedTuple(second + 1, i);
+		                   }
+		                   //	11:	else
+		                   //	12:		prevIdcs[i] ← “-”
+		                   for (idx_t j = 0; j < ndistinct; ++j) {
+			                   auto scan_idx = distinct.get_index(j);
+			                   auto i = input_idx[input_format.sel->get_index(scan_idx)];
+			                   prev_idcs[i] = ZippedTuple(0, i);
+		                   }
+
+		                   //	Remember the last input_idx of this chunk.
+		                   prev_i = input_idx[input_format.sel->get_index(count - 1)];
+	                   });
 
-		int lt = 0;
-		if (prefix_layout.all_constant) {
-			lt = FastMemcmp(prev.entry_ptr, curr.entry_ptr, prefix_layout.comparison_size);
-		} else {
-			lt = Comparators::CompareTuple(prev.scan, curr.scan, prev.entry_ptr, curr.entry_ptr, prefix_layout,
-			                               prev.external);
-		}
-
-		//	9:	if sorted[i].first == sorted[i-1].first then
-		//	10:		prevIdcs[i] ← sorted[i-1].second
-		//	11:	else
-		//	12:		prevIdcs[i] ← “-”
-		if (!lt) {
-			prev_idcs[i] = ZippedTuple(second + 1, i);
-		} else {
-			prev_idcs[i] = ZippedTuple(0, i);
-		}
-	}
-
-	// Save the last value of i for patching up the block boundaries
-	std::get<1>(gastate.seconds[block_idx]) = i;
-}
-
-void WindowDistinctAggregatorGlobalState::PatchPrevIdcs() {
 	//	13:	return prevIdcs
-
-	// Patch up the indices at block boundaries
-	// (We don't need to patch block 0.)
-	auto &prev_idcs = zipped_tree.LowestLevel();
-	for (idx_t block_idx = 1; block_idx < seconds.size(); ++block_idx) {
-		// We only need to patch if the first index in the block
-		// was a back link to the previous block (10:)
-		auto i = std::get<0>(seconds.at(block_idx));
-		if (std::get<0>(prev_idcs[i])) {
-			auto second = std::get<1>(seconds.at(block_idx - 1));
-			prev_idcs[i] = ZippedTuple(second + 1, i);
-		}
-	}
 }
 
 bool WindowDistinctSortTree::TryNextRun(idx_t &level_idx, idx_t &run_idx) {
@@ -593,7 +540,6 @@ void WindowDistinctSortTree::Build(WindowDistinctAggregatorLocalState &ldastate)
 
 void WindowDistinctSortTree::BuildRun(idx_t level_nr, idx_t run_idx, WindowDistinctAggregatorLocalState &ldastate) {
 	auto &aggr = gdastate.aggr;
-	auto &allocator = gdastate.allocator;
 	auto &inputs = ldastate.cursor->chunk;
 	auto &levels_flat_native = gdastate.levels_flat_native;
 
@@ -601,7 +547,7 @@ void WindowDistinctSortTree::BuildRun(idx_t level_nr, idx_t run_idx, WindowDisti
 	auto &leaves = ldastate.leaves;
 	auto &sel = ldastate.sel;
 
-	AggregateInputData aggr_input_data(aggr.GetFunctionData(), allocator);
+	AggregateInputData aggr_input_data(aggr.GetFunctionData(), ldastate.tree_allocator);
 
 	//! The states to update
 	auto &update_v = ldastate.update_v;
@@ -697,15 +643,16 @@ void WindowDistinctAggregatorLocalState::FlushStates() {
 		return;
 	}
 
-	const auto &aggr = gastate.aggr;
-	AggregateInputData aggr_input_data(aggr.GetFunctionData(), allocator);
+	const auto &aggr = gdstate.aggr;
+	AggregateInputData aggr_input_data(aggr.GetFunctionData(), tree_allocator);
 	statel.Verify(flush_count);
 	aggr.function.combine(statel, statep, aggr_input_data, flush_count);
 
 	flush_count = 0;
 }
 
-void WindowDistinctAggregatorLocalState::Evaluate(const WindowDistinctAggregatorGlobalState &gdstate,
+void WindowDistinctAggregatorLocalState::Evaluate(ExecutionContext &context,
+                                                  const WindowDistinctAggregatorGlobalState &gdstate,
                                                   const DataChunk &bounds, Vector &result, idx_t count, idx_t row_idx) {
 	auto ldata = FlatVector::GetData<const_data_ptr_t>(statel);
 	auto pdata = FlatVector::GetData<data_ptr_t>(statep);
@@ -749,16 +696,18 @@ void WindowDistinctAggregatorLocalState::Evaluate(const WindowDistinctAggregator
 	statef.Destroy();
 }
 
-unique_ptr<WindowAggregatorState> WindowDistinctAggregator::GetLocalState(const WindowAggregatorState &gstate) const {
-	return make_uniq<WindowDistinctAggregatorLocalState>(gstate.Cast<const WindowDistinctAggregatorGlobalState>());
+unique_ptr<LocalSinkState> WindowDistinctAggregator::GetLocalState(ExecutionContext &context,
+                                                                   const GlobalSinkState &gstate) const {
+	auto &gdstate = gstate.Cast<const WindowDistinctAggregatorGlobalState>();
+	return make_uniq<WindowDistinctAggregatorLocalState>(context, gdstate);
 }
 
-void WindowDistinctAggregator::Evaluate(const WindowAggregatorState &gsink, WindowAggregatorState &lstate,
-                                        const DataChunk &bounds, Vector &result, idx_t count, idx_t row_idx) const {
+void WindowDistinctAggregator::Evaluate(ExecutionContext &context, const DataChunk &bounds, Vector &result, idx_t count,
+                                        idx_t row_idx, OperatorSinkInput &sink) const {
 
-	const auto &gdstate = gsink.Cast<WindowDistinctAggregatorGlobalState>();
-	auto &ldstate = lstate.Cast<WindowDistinctAggregatorLocalState>();
-	ldstate.Evaluate(gdstate, bounds, result, count, row_idx);
+	const auto &gdstate = sink.global_state.Cast<WindowDistinctAggregatorGlobalState>();
+	auto &ldstate = sink.local_state.Cast<WindowDistinctAggregatorLocalState>();
+	ldstate.Evaluate(context, gdstate, bounds, result, count, row_idx);
 }
 
 } // namespace duckdb