duckdb 0.7.2-dev2320.0 → 0.7.2-dev2410.0

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

@@ -1093,14 +1093,26 @@ SinkFinalizeType PhysicalWindow::Finalize(Pipeline &pipeline, Event &event, Clie
 //===--------------------------------------------------------------------===//
 class WindowGlobalSourceState : public GlobalSourceState {
 public:
-	explicit WindowGlobalSourceState(WindowGlobalSinkState &gsink) : partition_source(*gsink.global_partition) {
+	explicit WindowGlobalSourceState(WindowGlobalSinkState &gsink) : gsink(*gsink.global_partition), next_bin(0) {
 	}
 
-	PartitionGlobalSourceState partition_source;
+	PartitionGlobalSinkState &gsink;
+	//! The output read position.
+	atomic<idx_t> next_bin;
 
 public:
 	idx_t MaxThreads() override {
-		return partition_source.MaxThreads();
+		// If there is only one partition, we have to process it on one thread.
+		if (!gsink.grouping_data) {
+			return 1;
+		}
+
+		// If there is not a lot of data, process serially.
+		if (gsink.count < STANDARD_ROW_GROUPS_SIZE) {
+			return 1;
+		}
+
+		return gsink.hash_groups.size();
 	}
 };
 
@@ -1112,7 +1124,7 @@ public:
 	using WindowExecutors = vector<WindowExecutorPtr>;
 
 	WindowLocalSourceState(const PhysicalWindow &op_p, ExecutionContext &context, WindowGlobalSourceState &gsource)
-	    : partition_source(gsource.partition_source.gsink), context(context.client), op(op_p) {
+	    : context(context.client), op(op_p), gsink(gsource.gsink) {
 
 		vector<LogicalType> output_types;
 		for (idx_t expr_idx = 0; expr_idx < op.select_list.size(); ++expr_idx) {
@@ -1121,29 +1133,134 @@ public:
 			output_types.emplace_back(wexpr.return_type);
 		}
 		output_chunk.Initialize(Allocator::Get(context.client), output_types);
+
+		const auto &input_types = gsink.payload_types;
+		layout.Initialize(input_types);
+		input_chunk.Initialize(gsink.allocator, input_types);
 	}
 
+	void MaterializeSortedData();
 	void GeneratePartition(WindowGlobalSinkState &gstate, const idx_t hash_bin);
 	void Scan(DataChunk &chunk);
 
-	PartitionLocalSourceState partition_source;
+	HashGroupPtr hash_group;
 	ClientContext &context;
 	const PhysicalWindow &op;
 
+	PartitionGlobalSinkState &gsink;
+
+	//! The generated input chunks
+	unique_ptr<RowDataCollection> rows;
+	unique_ptr<RowDataCollection> heap;
+	RowLayout layout;
+	//! The partition boundary mask
+	vector<validity_t> partition_bits;
+	ValidityMask partition_mask;
+	//! The order boundary mask
+	vector<validity_t> order_bits;
+	ValidityMask order_mask;
 	//! The current execution functions
 	WindowExecutors window_execs;
+
+	//! The read partition
+	idx_t hash_bin;
+	//! The read cursor
+	unique_ptr<RowDataCollectionScanner> scanner;
+	//! Buffer for the inputs
+	DataChunk input_chunk;
 	//! Buffer for window results
 	DataChunk output_chunk;
 };
 
+void WindowLocalSourceState::MaterializeSortedData() {
+	auto &global_sort_state = *hash_group->global_sort;
+	if (global_sort_state.sorted_blocks.empty()) {
+		return;
+	}
+
+	// scan the sorted row data
+	D_ASSERT(global_sort_state.sorted_blocks.size() == 1);
+	auto &sb = *global_sort_state.sorted_blocks[0];
+
+	// Free up some memory before allocating more
+	sb.radix_sorting_data.clear();
+	sb.blob_sorting_data = nullptr;
+
+	// Move the sorting row blocks into our RDCs
+	auto &buffer_manager = global_sort_state.buffer_manager;
+	auto &sd = *sb.payload_data;
+
+	// Data blocks are required
+	D_ASSERT(!sd.data_blocks.empty());
+	auto &block = sd.data_blocks[0];
+	rows = make_uniq<RowDataCollection>(buffer_manager, block->capacity, block->entry_size);
+	rows->blocks = std::move(sd.data_blocks);
+	rows->count = std::accumulate(rows->blocks.begin(), rows->blocks.end(), idx_t(0),
+	                              [&](idx_t c, const unique_ptr<RowDataBlock> &b) { return c + b->count; });
+
+	// Heap blocks are optional, but we want both for iteration.
+	if (!sd.heap_blocks.empty()) {
+		auto &block = sd.heap_blocks[0];
+		heap = make_uniq<RowDataCollection>(buffer_manager, block->capacity, block->entry_size);
+		heap->blocks = std::move(sd.heap_blocks);
+		hash_group.reset();
+	} else {
+		heap = make_uniq<RowDataCollection>(buffer_manager, (idx_t)Storage::BLOCK_SIZE, 1, true);
+	}
+	heap->count = std::accumulate(heap->blocks.begin(), heap->blocks.end(), idx_t(0),
+	                              [&](idx_t c, const unique_ptr<RowDataBlock> &b) { return c + b->count; });
+}
+
 void WindowLocalSourceState::GeneratePartition(WindowGlobalSinkState &gstate, const idx_t hash_bin_p) {
-	const auto count = partition_source.GeneratePartition(hash_bin_p);
-	if (!count) {
+	//	Get rid of any stale data
+	hash_bin = hash_bin_p;
+
+	// There are three types of partitions:
+	// 1. No partition (no sorting)
+	// 2. One partition (sorting, but no hashing)
+	// 3. Multiple partitions (sorting and hashing)
+
+	//	How big is the partition?
+	idx_t count = 0;
+	if (hash_bin < gsink.hash_groups.size() && gsink.hash_groups[hash_bin]) {
+		count = gsink.hash_groups[hash_bin]->count;
+	} else if (gsink.rows && !hash_bin) {
+		count = gsink.count;
+	} else {
+		return;
+	}
+
+	//	Initialise masks to false
+	const auto bit_count = ValidityMask::ValidityMaskSize(count);
+	partition_bits.clear();
+	partition_bits.resize(bit_count, 0);
+	partition_mask.Initialize(partition_bits.data());
+
+	order_bits.clear();
+	order_bits.resize(bit_count, 0);
+	order_mask.Initialize(order_bits.data());
+
+	// Scan the sorted data into new Collections
+	auto external = gsink.external;
+	if (gsink.rows && !hash_bin) {
+		// Simple mask
+		partition_mask.SetValidUnsafe(0);
+		order_mask.SetValidUnsafe(0);
+		//	No partition - align the heap blocks with the row blocks
+		rows = gsink.rows->CloneEmpty(gsink.rows->keep_pinned);
+		heap = gsink.strings->CloneEmpty(gsink.strings->keep_pinned);
+		RowDataCollectionScanner::AlignHeapBlocks(*rows, *heap, *gsink.rows, *gsink.strings, layout);
+		external = true;
+	} else if (hash_bin < gsink.hash_groups.size() && gsink.hash_groups[hash_bin]) {
+		// Overwrite the collections with the sorted data
+		hash_group = std::move(gsink.hash_groups[hash_bin]);
+		hash_group->ComputeMasks(partition_mask, order_mask);
+		MaterializeSortedData();
+	} else {
 		return;
 	}
 
 	// Create the executors for each function
-	auto &partition_mask = partition_source.partition_mask;
 	window_execs.clear();
 	for (idx_t expr_idx = 0; expr_idx < op.select_list.size(); ++expr_idx) {
 		D_ASSERT(op.select_list[expr_idx]->GetExpressionClass() == ExpressionClass::BOUND_WINDOW);
@@ -1154,19 +1271,20 @@ void WindowLocalSourceState::GeneratePartition(WindowGlobalSinkState &gstate, co
 
 	//	First pass over the input without flushing
 	//	TODO: Factor out the constructor data as global state
+	scanner = make_uniq<RowDataCollectionScanner>(*rows, *heap, layout, external, false);
 	idx_t input_idx = 0;
 	while (true) {
-		partition_source.input_chunk.Reset();
-		partition_source.scanner->Scan(partition_source.input_chunk);
-		if (partition_source.input_chunk.size() == 0) {
+		input_chunk.Reset();
+		scanner->Scan(input_chunk);
+		if (input_chunk.size() == 0) {
 			break;
 		}
 
 		//	TODO: Parallelization opportunity
 		for (auto &wexec : window_execs) {
-			wexec->Sink(partition_source.input_chunk, input_idx, partition_source.scanner->Count());
+			wexec->Sink(input_chunk, input_idx, scanner->Count());
 		}
-		input_idx += partition_source.input_chunk.size();
+		input_idx += input_chunk.size();
 	}
 
 	//	TODO: Parallelization opportunity
@@ -1175,25 +1293,22 @@ void WindowLocalSourceState::GeneratePartition(WindowGlobalSinkState &gstate, co
 	}
 
 	// External scanning assumes all blocks are swizzled.
-	partition_source.scanner->ReSwizzle();
+	scanner->ReSwizzle();
 
 	//	Second pass can flush
-	partition_source.scanner->Reset(true);
+	scanner->Reset(true);
 }
 
 void WindowLocalSourceState::Scan(DataChunk &result) {
-	D_ASSERT(partition_source.scanner);
-	if (!partition_source.scanner->Remaining()) {
+	D_ASSERT(scanner);
+	if (!scanner->Remaining()) {
 		return;
 	}
 
-	const auto position = partition_source.scanner->Scanned();
-	auto &input_chunk = partition_source.input_chunk;
+	const auto position = scanner->Scanned();
 	input_chunk.Reset();
-	partition_source.scanner->Scan(input_chunk);
+	scanner->Scan(input_chunk);
 
-	auto &partition_mask = partition_source.partition_mask;
-	auto &order_mask = partition_source.order_mask;
 	output_chunk.Reset();
 	for (idx_t expr_idx = 0; expr_idx < window_execs.size(); ++expr_idx) {
 		auto &executor = *window_execs[expr_idx];
@@ -1227,9 +1342,7 @@ unique_ptr<GlobalSourceState> PhysicalWindow::GetGlobalSourceState(ClientContext
 void PhysicalWindow::GetData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate_p,
                              LocalSourceState &lstate_p) const {
 	auto &lsource = lstate_p.Cast<WindowLocalSourceState>();
-	auto &lpsource = lsource.partition_source;
 	auto &gsource = gstate_p.Cast<WindowGlobalSourceState>();
-	auto &gpsource = gsource.partition_source;
 	auto &gsink = sink_state->Cast<WindowGlobalSinkState>();
 
 	auto &hash_groups = gsink.global_partition->hash_groups;
@@ -1237,17 +1350,17 @@ void PhysicalWindow::GetData(ExecutionContext &context, DataChunk &chunk, Global
 
 	while (chunk.size() == 0) {
 		//	Move to the next bin if we are done.
-		while (!lpsource.scanner || !lpsource.scanner->Remaining()) {
-			lpsource.scanner.reset();
-			lpsource.rows.reset();
-			lpsource.heap.reset();
-			lpsource.hash_group.reset();
-			auto hash_bin = gpsource.next_bin++;
+		while (!lsource.scanner || !lsource.scanner->Remaining()) {
+			lsource.scanner.reset();
+			lsource.rows.reset();
+			lsource.heap.reset();
+			lsource.hash_group.reset();
+			auto hash_bin = gsource.next_bin++;
 			if (hash_bin >= bin_count) {
 				return;
 			}
 
-			for (; hash_bin < hash_groups.size(); hash_bin = gpsource.next_bin++) {
+			for (; hash_bin < hash_groups.size(); hash_bin = gsource.next_bin++) {
 				if (hash_groups[hash_bin]) {
 					break;
 				}