duckdb 1.3.3-dev4.0 → 1.3.3

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

@@ -1,10 +1,11 @@
 #include "duckdb/function/window/window_distinct_aggregator.hpp"
 
-#include "duckdb/common/sorting/sort.hpp"
+#include "duckdb/common/sort/partition_state.hpp"
+#include "duckdb/common/sort/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_reference_expression.hpp"
+#include "duckdb/planner/expression/bound_constant_expression.hpp"
 #include "duckdb/planner/expression/bound_window_expression.hpp"
 
 #include <numeric>
@@ -12,8 +13,6 @@
 
 namespace duckdb {
 
-enum class WindowDistinctSortStage : uint8_t { INIT, COMBINE, FINALIZE, SORTED, FINISHED };
-
 //===--------------------------------------------------------------------===//
 // WindowDistinctAggregator
 //===--------------------------------------------------------------------===//
@@ -22,10 +21,6 @@ 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();
 }
 
@@ -44,7 +39,7 @@ public:
 	using ZippedTuple = std::tuple<idx_t, idx_t>;
 	using ZippedElements = vector<ZippedTuple>;
 
-	WindowDistinctSortTree(WindowDistinctAggregatorGlobalState &gdastate, idx_t count) : gdastate(gdastate) {
+	explicit WindowDistinctSortTree(WindowDistinctAggregatorGlobalState &gdastate, idx_t count) : gdastate(gdastate) {
 		//	Set up for parallel build
 		build_level = 0;
 		build_complete = 0;
@@ -64,31 +59,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<LocalSinkState> InitializeLocalSort(ExecutionContext &context) const;
-
-	ArenaAllocator &CreateTreeAllocator() const {
-		lock_guard<mutex> tree_lock(lock);
-		tree_allocators.emplace_back(make_uniq<ArenaAllocator>(Allocator::DefaultAllocator()));
-		return *tree_allocators.back();
-	}
+	optional_ptr<LocalSortState> InitializeLocalSort() const;
 
+	//! Patch up the previous index block boundaries
+	void PatchPrevIdcs();
 	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;
+	//	Single threaded sorting for now
+	ClientContext &context;
+	idx_t memory_per_thread;
+
 	//! Finalize guard
 	mutable mutex lock;
 	//! Finalize stage
-	atomic<WindowDistinctSortStage> stage;
+	atomic<PartitionSortStage> stage;
 	//! Tasks launched
 	idx_t total_tasks = 0;
 	//! Tasks launched
@@ -96,18 +91,20 @@ 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
-	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;
-
+	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;
 	//! The MST with the distinct back pointers
 	mutable MergeSortTree<ZippedTuple> zipped_tree;
 	//! The merge sort tree for the aggregate.
@@ -119,29 +116,35 @@ public:
 	vector<idx_t> levels_flat_start;
 };
 
-WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientContext &client,
+WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientContext &context,
                                                                          const WindowDistinctAggregator &aggregator,
                                                                          idx_t group_count)
-    : WindowAggregatorGlobalState(client, aggregator, group_count), stage(WindowDistinctSortStage::INIT),
-      tasks_assigned(0), tasks_completed(0), merge_sort_tree(*this, group_count), levels_flat_native(aggr) {
+    : 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);
 
 	//	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;
-	sort_types.emplace_back(LogicalType::UBIGINT);
+	for (const auto &type : payload_types) {
+		sort_types.emplace_back(type);
+	}
 
-	//	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<BoundReferenceExpression>(type, orders.size());
+		auto expr = make_uniq<BoundConstantExpression>(Value(type));
 		orders.emplace_back(BoundOrderByNode(OrderType::ASCENDING, OrderByNullType::NULLS_FIRST, std::move(expr)));
-		sort_cols.emplace_back(sort_cols.size());
 	}
 
-	sort = make_uniq<Sort>(client, orders, sort_types, sort_cols);
-	global_sink = sort->GetGlobalSinkState(client);
+	RowLayout payload_layout;
+	payload_layout.Initialize(payload_types);
+
+	global_sort = make_uniq<GlobalSortState>(context, orders, payload_layout);
+
+	memory_per_thread = PhysicalOperator::GetMaxThreadMemory(context);
 
 	//	6:	prevIdcs ← []
 	//	7:	prevIdcs[0] ← “-”
@@ -173,38 +176,36 @@ WindowDistinctAggregatorGlobalState::WindowDistinctAggregatorGlobalState(ClientC
 	}
 }
 
-optional_ptr<LocalSinkState> WindowDistinctAggregatorGlobalState::InitializeLocalSort(ExecutionContext &context) const {
+optional_ptr<LocalSortState> WindowDistinctAggregatorGlobalState::InitializeLocalSort() const {
 	lock_guard<mutex> local_sort_guard(lock);
-	auto local_sink = sort->GetLocalSinkState(context);
+	auto local_sort = make_uniq<LocalSortState>();
+	local_sort->Initialize(*global_sort, global_sort->buffer_manager);
 	++tasks_assigned;
-	local_sinks.emplace_back(std::move(local_sink));
+	local_sorts.emplace_back(std::move(local_sort));
 
-	return local_sinks.back().get();
+	return local_sorts.back().get();
 }
 
 class WindowDistinctAggregatorLocalState : public WindowAggregatorLocalState {
 public:
-	WindowDistinctAggregatorLocalState(ExecutionContext &context,
-	                                   const WindowDistinctAggregatorGlobalState &aggregator);
+	explicit WindowDistinctAggregatorLocalState(const WindowDistinctAggregatorGlobalState &aggregator);
 
 	~WindowDistinctAggregatorLocalState() override {
 		statef.Destroy();
 	}
 
-	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 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 Sorted();
-	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);
+	void ExecuteTask();
+	void Evaluate(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<LocalSinkState> local_sink;
+	optional_ptr<LocalSortState> local_sort;
 	//! Finalize stage
-	WindowDistinctSortStage stage = WindowDistinctSortStage::INIT;
+	PartitionSortStage stage = PartitionSortStage::INIT;
 	//! Finalize scan block index
 	idx_t block_idx;
 	//! Thread-local tree aggregation
@@ -219,9 +220,9 @@ protected:
 	void FlushStates();
 
 	//! The aggregator we are working with
-	const WindowDistinctAggregatorGlobalState &gdstate;
-	//! The sort input chunk
+	const WindowDistinctAggregatorGlobalState &gastate;
 	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
@@ -235,173 +236,207 @@ protected:
 };
 
 WindowDistinctAggregatorLocalState::WindowDistinctAggregatorLocalState(
-    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);
+    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);
 
-	sort_chunk.Initialize(Allocator::DefaultAllocator(), gdstate.sort_types);
+	sort_chunk.Initialize(Allocator::DefaultAllocator(), gastate.sort_types);
+	sort_chunk.data.back().Reference(payload_chunk.data[0]);
 
-	gdstate.locals++;
+	gastate.locals++;
 }
 
-unique_ptr<GlobalSinkState> WindowDistinctAggregator::GetGlobalState(ClientContext &context, idx_t group_count,
-                                                                     const ValidityMask &partition_mask) const {
+unique_ptr<WindowAggregatorState> WindowDistinctAggregator::GetGlobalState(ClientContext &context, idx_t group_count,
+                                                                           const ValidityMask &partition_mask) const {
 	return make_uniq<WindowDistinctAggregatorGlobalState>(context, *this, group_count);
 }
 
-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);
+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);
 
-	auto &ldstate = sink.local_state.Cast<WindowDistinctAggregatorLocalState>();
-	ldstate.Sink(context, sink_chunk, coll_chunk, input_idx, filter_sel, filtered, sink.interrupt_state);
+	auto &ldstate = lstate.Cast<WindowDistinctAggregatorLocalState>();
+	ldstate.Sink(sink_chunk, coll_chunk, input_idx, filter_sel, 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) {
+void WindowDistinctAggregatorLocalState::Sink(DataChunk &sink_chunk, DataChunk &coll_chunk, idx_t input_idx,
+                                              optional_ptr<SelectionVector> filter_sel, idx_t filtered) {
 	//	3: 	for i ← 0 to in.size do
 	//	4: 		sorted[i] ← (in[i], i)
 	const auto count = sink_chunk.size();
-	sort_chunk.Reset();
-	auto &sorted_vec = sort_chunk.data.back();
+	payload_chunk.Reset();
+	auto &sorted_vec = payload_chunk.data[0];
 	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 = gdstate.aggregator.child_idx;
+	auto &child_idx = gastate.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_sink) {
-		local_sink = gdstate.InitializeLocalSort(context);
+	if (!local_sort) {
+		local_sort = gastate.InitializeLocalSort();
 	}
 
-	OperatorSinkInput sink {*gdstate.global_sink, *local_sink, interrupt};
-	gdstate.sort->Sink(context, sort_chunk, sink);
+	local_sort->SinkChunk(sort_chunk, payload_chunk);
+
+	if (local_sort->SizeInBytes() > gastate.memory_per_thread) {
+		local_sort->Sort(*gastate.global_sort, true);
+	}
 }
 
-void WindowDistinctAggregatorLocalState::Finalize(ExecutionContext &context, WindowAggregatorGlobalState &gastate,
-                                                  CollectionPtr collection) {
-	WindowAggregatorLocalState::Finalize(context, gastate, collection);
+void WindowDistinctAggregatorLocalState::Finalize(WindowAggregatorGlobalState &gastate, CollectionPtr collection) {
+	WindowAggregatorLocalState::Finalize(gastate, collection);
 
 	//! Input data chunk, used for leaf segment aggregation
 	leaves.Initialize(Allocator::DefaultAllocator(), cursor->chunk.GetTypes());
 	sel.Initialize();
 }
 
-void WindowDistinctAggregatorLocalState::ExecuteTask(ExecutionContext &context,
-                                                     WindowDistinctAggregatorGlobalState &gdstate) {
-	PostIncrement<atomic<idx_t>> on_done(gdstate.tasks_completed);
-
+void WindowDistinctAggregatorLocalState::ExecuteTask() {
+	auto &global_sort = *gastate.global_sort;
 	switch (stage) {
-	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);
+	case PartitionSortStage::SCAN:
+		global_sort.AddLocalState(*gastate.local_sorts[block_idx]);
 		break;
-	}
-	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);
+	case PartitionSortStage::MERGE: {
+		MergeSorter merge_sorter(global_sort, global_sort.buffer_manager);
+		merge_sorter.PerformInMergeRound();
 		break;
 	}
-	case WindowDistinctSortStage::SORTED:
+	case PartitionSortStage::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 WindowDistinctSortStage::INIT:
-		total_tasks = local_sinks.size();
+	case PartitionSortStage::INIT:
+		//	5: Sort sorted lexicographically increasing
+		total_tasks = local_sorts.size();
 		tasks_assigned = 0;
 		tasks_completed = 0;
-		lstate.stage = stage = WindowDistinctSortStage::COMBINE;
+		lstate.stage = stage = PartitionSortStage::SCAN;
 		lstate.block_idx = tasks_assigned++;
 		return true;
-	case WindowDistinctSortStage::COMBINE:
+	case PartitionSortStage::SCAN:
+		// Process all the local sorts
 		if (tasks_assigned < total_tasks) {
-			lstate.stage = WindowDistinctSortStage::COMBINE;
+			lstate.stage = PartitionSortStage::SCAN;
 			lstate.block_idx = tasks_assigned++;
 			return true;
 		} else if (tasks_completed < tasks_assigned) {
 			return false;
 		}
-		// All combines are done, so move on to materialising the sorted data (1 task)
-		total_tasks = 1;
+		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;
 		tasks_completed = 0;
-		tasks_assigned = 0;
-		lstate.stage = stage = WindowDistinctSortStage::FINALIZE;
-		lstate.block_idx = tasks_assigned++;
 		return true;
-	case WindowDistinctSortStage::FINALIZE:
-		if (tasks_completed < tasks_assigned) {
-			//	Wait for the single task to finish
+	case PartitionSortStage::MERGE:
+		if (tasks_assigned < total_tasks) {
+			lstate.stage = PartitionSortStage::MERGE;
+			++tasks_assigned;
+			return true;
+		} else if (tasks_completed < tasks_assigned) {
 			return false;
 		}
-		//	Move on to building the tree in parallel
-		total_tasks = local_sinks.size();
+		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;
 		tasks_completed = 0;
-		tasks_assigned = 0;
-		lstate.stage = stage = WindowDistinctSortStage::SORTED;
-		lstate.block_idx = tasks_assigned++;
 		return true;
-	case WindowDistinctSortStage::SORTED:
+	case PartitionSortStage::SORTED:
 		if (tasks_assigned < total_tasks) {
-			lstate.stage = WindowDistinctSortStage::SORTED;
+			lstate.stage = PartitionSortStage::SORTED;
 			lstate.block_idx = tasks_assigned++;
 			return true;
 		} else if (tasks_completed < tasks_assigned) {
-			lstate.stage = WindowDistinctSortStage::FINISHED;
+			lstate.stage = PartitionSortStage::FINISHED;
 			// Sleep while other tasks finish
 			return false;
 		}
+		// Last task patches the boundaries
+		PatchPrevIdcs();
 		break;
 	default:
 		break;
 	}
 
-	lstate.stage = stage = WindowDistinctSortStage::FINISHED;
+	lstate.stage = stage = PartitionSortStage::FINISHED;
 
 	return true;
 }
 
-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);
+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);
 
 	// Sort, merge and build the tree in parallel
-	while (gdsink.stage.load() != WindowDistinctSortStage::FINISHED) {
+	while (gdsink.stage.load() != PartitionSortStage::FINISHED) {
 		if (gdsink.TryPrepareNextStage(ldstate)) {
-			ldstate.ExecuteTask(context, gdsink);
+			ldstate.ExecuteTask();
 		} else {
 			std::this_thread::yield();
 		}
@@ -417,71 +452,89 @@ void WindowDistinctAggregator::Finalize(ExecutionContext &context, CollectionPtr
 
 void WindowDistinctAggregatorLocalState::Sorted() {
 	using ZippedTuple = WindowDistinctAggregatorGlobalState::ZippedTuple;
-	auto &collection = *gdstate.sorted;
-	auto &prev_idcs = gdstate.zipped_tree.LowestLevel();
-	auto &aggregator = gdstate.aggregator;
+	auto &global_sort = gastate.global_sort;
+	auto &prev_idcs = gastate.zipped_tree.LowestLevel();
+	auto &aggregator = gastate.aggregator;
+	auto &scan_chunk = payload_chunk;
 
-	// 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 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;
 
-	const auto &scan_cols = gdstate.sort_cols;
-	const auto key_count = aggregator.arg_types.size();
+	auto *input_idx = FlatVector::GetData<idx_t>(scan_chunk.data[0]);
+	idx_t i = 0;
 
-	//	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;
-	}
+	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());
 
-	idx_t prev_i = 0;
+	const auto block_begin = gastate.block_starts.at(block_idx);
 	if (!block_begin) {
 		// First block, so set up initial sentinel
-		auto input_idx = FlatVector::GetData<idx_t>(scanned.data.back());
-		prev_i = input_idx[0];
-		prev_idcs[prev_i] = ZippedTuple(0, prev_i);
+		i = input_idx[scan_idx++];
+		prev_idcs[i] = ZippedTuple(0, i);
+		std::get<0>(gastate.seconds[block_idx]) = i;
 	} else {
 		// Move to the to end of the previous block
 		// so we can record the comparison result for the first row
-		auto input_idx = FlatVector::GetData<idx_t>(scanned.data.back());
-		auto scan_idx = scanned.size() - 1;
-		prev_i = input_idx[scan_idx];
+		curr.SetIndex(block_begin - 1);
+		prev.SetIndex(block_begin - 1);
+		scan_idx = 0;
+		std::get<0>(gastate.seconds[block_idx]) = input_idx[scan_idx];
 	}
 
 	//	8:	for i ← 1 to in.size do
-	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)];
-	                   });
+	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++];
 
+		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) {
@@ -540,6 +593,7 @@ 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;
 
@@ -547,7 +601,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(), ldastate.tree_allocator);
+	AggregateInputData aggr_input_data(aggr.GetFunctionData(), allocator);
 
 	//! The states to update
 	auto &update_v = ldastate.update_v;
@@ -643,16 +697,15 @@ void WindowDistinctAggregatorLocalState::FlushStates() {
 		return;
 	}
 
-	const auto &aggr = gdstate.aggr;
-	AggregateInputData aggr_input_data(aggr.GetFunctionData(), tree_allocator);
+	const auto &aggr = gastate.aggr;
+	AggregateInputData aggr_input_data(aggr.GetFunctionData(), allocator);
 	statel.Verify(flush_count);
 	aggr.function.combine(statel, statep, aggr_input_data, flush_count);
 
 	flush_count = 0;
 }
 
-void WindowDistinctAggregatorLocalState::Evaluate(ExecutionContext &context,
-                                                  const WindowDistinctAggregatorGlobalState &gdstate,
+void WindowDistinctAggregatorLocalState::Evaluate(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);
@@ -696,18 +749,16 @@ void WindowDistinctAggregatorLocalState::Evaluate(ExecutionContext &context,
 	statef.Destroy();
 }
 
-unique_ptr<LocalSinkState> WindowDistinctAggregator::GetLocalState(ExecutionContext &context,
-                                                                   const GlobalSinkState &gstate) const {
-	auto &gdstate = gstate.Cast<const WindowDistinctAggregatorGlobalState>();
-	return make_uniq<WindowDistinctAggregatorLocalState>(context, gdstate);
+unique_ptr<WindowAggregatorState> WindowDistinctAggregator::GetLocalState(const WindowAggregatorState &gstate) const {
+	return make_uniq<WindowDistinctAggregatorLocalState>(gstate.Cast<const WindowDistinctAggregatorGlobalState>());
 }
 
-void WindowDistinctAggregator::Evaluate(ExecutionContext &context, const DataChunk &bounds, Vector &result, idx_t count,
-                                        idx_t row_idx, OperatorSinkInput &sink) const {
+void WindowDistinctAggregator::Evaluate(const WindowAggregatorState &gsink, WindowAggregatorState &lstate,
+                                        const DataChunk &bounds, Vector &result, idx_t count, idx_t row_idx) const {
 
-	const auto &gdstate = sink.global_state.Cast<WindowDistinctAggregatorGlobalState>();
-	auto &ldstate = sink.local_state.Cast<WindowDistinctAggregatorLocalState>();
-	ldstate.Evaluate(context, gdstate, bounds, result, count, row_idx);
+	const auto &gdstate = gsink.Cast<WindowDistinctAggregatorGlobalState>();
+	auto &ldstate = lstate.Cast<WindowDistinctAggregatorLocalState>();
+	ldstate.Evaluate(gdstate, bounds, result, count, row_idx);
 }
 
 } // namespace duckdb