duckdb 1.0.1-dev22.0 → 1.0.1-dev27.0

package/src/duckdb/src/execution/radix_partitioned_hashtable.cpp

@@ -6,6 +6,7 @@
 #include "duckdb/common/types/row/tuple_data_iterator.hpp"
 #include "duckdb/execution/aggregate_hashtable.hpp"
 #include "duckdb/execution/executor.hpp"
+#include "duckdb/execution/ht_entry.hpp"
 #include "duckdb/execution/operator/aggregate/physical_hash_aggregate.hpp"
 #include "duckdb/main/config.hpp"
 #include "duckdb/parallel/event.hpp"
@@ -78,18 +79,15 @@ enum class AggregatePartitionState : uint8_t {
 	READY_TO_SCAN = 2
 };
 
-struct AggregatePartition {
+struct AggregatePartition : StateWithBlockableTasks {
 	explicit AggregatePartition(unique_ptr<TupleDataCollection> data_p)
 	    : state(AggregatePartitionState::READY_TO_FINALIZE), data(std::move(data_p)), progress(0) {
 	}
 
-	mutex lock;
 	AggregatePartitionState state;
 
 	unique_ptr<TupleDataCollection> data;
 	atomic<double> progress;
-
-	vector<InterruptState> blocked_tasks;
 };
 
 class RadixHTGlobalSinkState;
@@ -172,8 +170,6 @@ public:
 	//! If any thread has called combine
 	atomic<bool> any_combined;
 
-	//! Lock for uncombined_data/stored_allocators
-	mutex lock;
 	//! Uncombined partitioned data that will be put into the AggregatePartitions
 	unique_ptr<PartitionedTupleData> uncombined_data;
 	//! Allocators used during the Sink/Finalize
@@ -199,20 +195,25 @@ RadixHTGlobalSinkState::RadixHTGlobalSinkState(ClientContext &context_p, const R
       any_combined(false), finalize_done(0), scan_pin_properties(TupleDataPinProperties::DESTROY_AFTER_DONE),
       count_before_combining(0), max_partition_size(0) {
 
-	auto tuples_per_block = Storage::BLOCK_ALLOC_SIZE / radix_ht.GetLayout().GetRowWidth();
+	// Compute minimum reservation
+	auto block_alloc_size = BufferManager::GetBufferManager(context).GetBlockAllocSize();
+	auto tuples_per_block = block_alloc_size / radix_ht.GetLayout().GetRowWidth();
 	idx_t ht_count =
-	    NumericCast<idx_t>(static_cast<double>(config.sink_capacity) / GroupedAggregateHashTable::LOAD_FACTOR);
+	    LossyNumericCast<idx_t>(static_cast<double>(config.sink_capacity) / GroupedAggregateHashTable::LOAD_FACTOR);
 	auto num_partitions = RadixPartitioning::NumberOfPartitions(config.GetRadixBits());
 	auto count_per_partition = ht_count / num_partitions;
 	auto blocks_per_partition = (count_per_partition + tuples_per_block) / tuples_per_block + 1;
-	auto ht_size = blocks_per_partition * Storage::BLOCK_ALLOC_SIZE + config.sink_capacity * sizeof(aggr_ht_entry_t);
+	if (!radix_ht.GetLayout().AllConstant()) {
+		blocks_per_partition += 2;
+	}
+	auto ht_size = blocks_per_partition * block_alloc_size + config.sink_capacity * sizeof(ht_entry_t);
 
 	// This really is the minimum reservation that we can do
 	auto num_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
 	auto minimum_reservation = num_threads * ht_size;
 
 	temporary_memory_state->SetMinimumReservation(minimum_reservation);
-	temporary_memory_state->SetRemainingSize(context, minimum_reservation);
+	temporary_memory_state->SetRemainingSizeAndUpdateReservation(context, minimum_reservation);
 }
 
 RadixHTGlobalSinkState::~RadixHTGlobalSinkState() {
@@ -233,7 +234,7 @@ void RadixHTGlobalSinkState::Destroy() {
 	}
 
 	// There are aggregates with destructors: Call the destructor for each of the aggregates
-	lock_guard<mutex> guard(lock);
+	auto guard = Lock();
 	RowOperationsState row_state(*stored_allocators.back());
 	for (auto &partition : partitions) {
 		auto &data_collection = *partition->data;
@@ -274,7 +275,7 @@ void RadixHTConfig::SetRadixBitsInternal(const idx_t radix_bits_p, bool external
 		return;
 	}
 
-	lock_guard<mutex> guard(sink.lock);
+	auto guard = sink.Lock();
 	if (sink_radix_bits >= radix_bits_p || sink.any_combined) {
 		return;
 	}
@@ -309,9 +310,9 @@ idx_t RadixHTConfig::SinkCapacity(ClientContext &context) {
 	const auto cache_per_active_thread = L1_CACHE_SIZE + L2_CACHE_SIZE + total_shared_cache_size / active_threads;
 
 	// Divide cache per active thread by entry size, round up to next power of two, to get capacity
-	const auto size_per_entry = sizeof(aggr_ht_entry_t) * GroupedAggregateHashTable::LOAD_FACTOR;
+	const auto size_per_entry = sizeof(ht_entry_t) * GroupedAggregateHashTable::LOAD_FACTOR;
 	const auto capacity =
-	    NextPowerOfTwo(NumericCast<uint64_t>(static_cast<double>(cache_per_active_thread) / size_per_entry));
+	    NextPowerOfTwo(LossyNumericCast<uint64_t>(static_cast<double>(cache_per_active_thread) / size_per_entry));
 
 	// Capacity must be at least the minimum capacity
 	return MaxValue<idx_t>(capacity, GroupedAggregateHashTable::InitialCapacity());
@@ -369,19 +370,19 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 
 	// Check if we're approaching the memory limit
 	auto &temporary_memory_state = *gstate.temporary_memory_state;
-	const auto total_size = partitioned_data->SizeInBytes() + ht.Capacity() * sizeof(aggr_ht_entry_t);
+	const auto total_size = partitioned_data->SizeInBytes() + ht.Capacity() * sizeof(ht_entry_t);
 	idx_t thread_limit = temporary_memory_state.GetReservation() / gstate.number_of_threads;
 	if (total_size > thread_limit) {
 		// We're over the thread memory limit
 		if (!gstate.external) {
 			// We haven't yet triggered out-of-core behavior, but maybe we don't have to, grab the lock and check again
-			lock_guard<mutex> guard(gstate.lock);
+			auto guard = gstate.Lock();
 			thread_limit = temporary_memory_state.GetReservation() / gstate.number_of_threads;
 			if (total_size > thread_limit) {
 				// Out-of-core would be triggered below, try to increase the reservation
 				auto remaining_size =
 				    MaxValue<idx_t>(gstate.number_of_threads * total_size, temporary_memory_state.GetRemainingSize());
-				temporary_memory_state.SetRemainingSize(context, 2 * remaining_size);
+				temporary_memory_state.SetRemainingSizeAndUpdateReservation(context, 2 * remaining_size);
 				thread_limit = temporary_memory_state.GetReservation() / gstate.number_of_threads;
 			}
 		}
@@ -413,9 +414,10 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 	const auto current_radix_bits = RadixPartitioning::RadixBits(partition_count);
 	D_ASSERT(current_radix_bits <= config.GetRadixBits());
 
+	const auto block_size = BufferManager::GetBufferManager(context).GetBlockSize();
 	const auto row_size_per_partition =
 	    partitioned_data->Count() * partitioned_data->GetLayout().GetRowWidth() / partition_count;
-	if (row_size_per_partition > NumericCast<idx_t>(config.BLOCK_FILL_FACTOR * Storage::BLOCK_SIZE)) {
+	if (row_size_per_partition > LossyNumericCast<idx_t>(config.BLOCK_FILL_FACTOR * static_cast<double>(block_size))) {
 		// We crossed our block filling threshold, try to increment radix bits
 		config.SetRadixBits(current_radix_bits + config.REPARTITION_RADIX_BITS);
 	}
@@ -498,7 +500,7 @@ void RadixPartitionedHashTable::Combine(ExecutionContext &context, GlobalSinkSta
 		lstate.abandoned_data = std::move(ht.GetPartitionedData());
 	}
 
-	lock_guard<mutex> guard(gstate.lock);
+	auto guard = gstate.Lock();
 	if (gstate.uncombined_data) {
 		gstate.uncombined_data->Combine(*lstate.abandoned_data);
 	} else {
@@ -524,7 +526,7 @@ void RadixPartitionedHashTable::Finalize(ClientContext &context, GlobalSinkState
 			auto &partition = uncombined_partition_data[i];
 			auto partition_size =
 			    partition->SizeInBytes() +
-			    GroupedAggregateHashTable::GetCapacityForCount(partition->Count()) * sizeof(aggr_ht_entry_t);
+			    GroupedAggregateHashTable::GetCapacityForCount(partition->Count()) * sizeof(ht_entry_t);
 			gstate.max_partition_size = MaxValue(gstate.max_partition_size, partition_size);
 
 			gstate.partitions.emplace_back(make_uniq<AggregatePartition>(std::move(partition)));
@@ -540,10 +542,8 @@ void RadixPartitionedHashTable::Finalize(ClientContext &context, GlobalSinkState
 
 	// Minimum of combining one partition at a time
 	gstate.temporary_memory_state->SetMinimumReservation(gstate.max_partition_size);
-	// Maximum of combining all partitions
-	auto max_threads = MinValue<idx_t>(NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads()),
-	                                   gstate.partitions.size());
-	gstate.temporary_memory_state->SetRemainingSize(context, max_threads * gstate.max_partition_size);
+	// Set size to 0 until the scan actually starts
+	gstate.temporary_memory_state->SetZero();
 	gstate.finalized = true;
 }
 
@@ -556,14 +556,17 @@ idx_t RadixPartitionedHashTable::MaxThreads(GlobalSinkState &sink_p) const {
 		return 0;
 	}
 
+	const auto max_threads = MinValue<idx_t>(
+	    NumericCast<idx_t>(TaskScheduler::GetScheduler(sink.context).NumberOfThreads()), sink.partitions.size());
+	sink.temporary_memory_state->SetRemainingSizeAndUpdateReservation(sink.context,
+	                                                                  max_threads * sink.max_partition_size);
+
 	// This many partitions will fit given our reservation (at least 1))
-	auto partitions_fit = MaxValue<idx_t>(sink.temporary_memory_state->GetReservation() / sink.max_partition_size, 1);
-	// Maximum is either the number of partitions, or the number of threads
-	auto max_possible = MinValue<idx_t>(
-	    sink.partitions.size(), NumericCast<idx_t>(TaskScheduler::GetScheduler(sink.context).NumberOfThreads()));
+	const auto partitions_fit =
+	    MaxValue<idx_t>(sink.temporary_memory_state->GetReservation() / sink.max_partition_size, 1);
 
 	// Mininum of the two
-	return MinValue<idx_t>(partitions_fit, max_possible);
+	return MinValue<idx_t>(partitions_fit, max_threads);
 }
 
 void RadixPartitionedHashTable::SetMultiScan(GlobalSinkState &sink_p) {
@@ -652,18 +655,16 @@ RadixHTGlobalSourceState::RadixHTGlobalSourceState(ClientContext &context_p, con
 SourceResultType RadixHTGlobalSourceState::AssignTask(RadixHTGlobalSinkState &sink, RadixHTLocalSourceState &lstate,
                                                       InterruptState &interrupt_state) {
 	// First, try to get a partition index
-	lock_guard<mutex> gstate_guard(sink.lock);
-	if (finished) {
-		return SourceResultType::FINISHED;
-	}
-	if (task_idx == sink.partitions.size()) {
+	auto guard = sink.Lock();
+	if (finished || task_idx == sink.partitions.size()) {
+		lstate.ht.reset();
 		return SourceResultType::FINISHED;
 	}
 	lstate.task_idx = task_idx++;
 
 	// We got a partition index
 	auto &partition = *sink.partitions[lstate.task_idx];
-	auto partition_lock = unique_lock<mutex>(partition.lock);
+	auto partition_guard = partition.Lock();
 	switch (partition.state) {
 	case AggregatePartitionState::READY_TO_FINALIZE:
 		partition.state = AggregatePartitionState::FINALIZE_IN_PROGRESS;
@@ -672,8 +673,7 @@ SourceResultType RadixHTGlobalSourceState::AssignTask(RadixHTGlobalSinkState &si
 	case AggregatePartitionState::FINALIZE_IN_PROGRESS:
 		lstate.task = RadixHTSourceTaskType::SCAN;
 		lstate.scan_status = RadixHTScanStatus::INIT;
-		partition.blocked_tasks.push_back(interrupt_state);
-		return SourceResultType::BLOCKED;
+		return partition.BlockSource(partition_guard, interrupt_state);
 	case AggregatePartitionState::READY_TO_SCAN:
 		lstate.task = RadixHTSourceTaskType::SCAN;
 		lstate.scan_status = RadixHTScanStatus::INIT;
@@ -721,10 +721,10 @@ void RadixHTLocalSourceState::Finalize(RadixHTGlobalSinkState &sink, RadixHTGlob
 		// However, we will limit the initial capacity so we don't do a huge over-allocation
 		const auto n_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(gstate.context).NumberOfThreads());
 		const auto memory_limit = BufferManager::GetBufferManager(gstate.context).GetMaxMemory();
-		const idx_t thread_limit = NumericCast<idx_t>(0.6 * double(memory_limit) / double(n_threads));
+		const idx_t thread_limit = LossyNumericCast<idx_t>(0.6 * double(memory_limit) / double(n_threads));
 
 		const idx_t size_per_entry = partition.data->SizeInBytes() / MaxValue<idx_t>(partition.data->Count(), 1) +
-		                             idx_t(GroupedAggregateHashTable::LOAD_FACTOR * sizeof(aggr_ht_entry_t));
+		                             idx_t(GroupedAggregateHashTable::LOAD_FACTOR * sizeof(ht_entry_t));
 		// but not lower than the initial capacity
 		const auto capacity_limit =
 		    MaxValue(NextPowerOfTwo(thread_limit / size_per_entry), GroupedAggregateHashTable::InitialCapacity());
@@ -748,22 +748,22 @@ void RadixHTLocalSourceState::Finalize(RadixHTGlobalSinkState &sink, RadixHTGlob
 	partition.data->Combine(*ht->GetPartitionedData()->GetPartitions()[0]);
 
 	// Update thread-global state
-	lock_guard<mutex> global_guard(sink.lock);
+	auto guard = sink.Lock();
 	sink.stored_allocators.emplace_back(ht->GetAggregateAllocator());
+	if (task_idx == sink.partitions.size()) {
+		ht.reset();
+	}
 	const auto finalizes_done = ++sink.finalize_done;
 	D_ASSERT(finalizes_done <= sink.partitions.size());
 	if (finalizes_done == sink.partitions.size()) {
 		// All finalizes are done, set remaining size to 0
-		sink.temporary_memory_state->SetRemainingSize(sink.context, 0);
+		sink.temporary_memory_state->SetZero();
 	}
 
 	// Update partition state
-	lock_guard<mutex> partition_guard(partition.lock);
+	auto partition_guard = partition.Lock();
 	partition.state = AggregatePartitionState::READY_TO_SCAN;
-	for (auto &blocked_task : partition.blocked_tasks) {
-		blocked_task.Callback();
-	}
-	partition.blocked_tasks.clear();
+	partition.UnblockTasks(partition_guard);
 
 	// This thread will scan the partition
 	task = RadixHTSourceTaskType::SCAN;
@@ -788,7 +788,7 @@ void RadixHTLocalSourceState::Scan(RadixHTGlobalSinkState &sink, RadixHTGlobalSo
 			data_collection.Reset();
 		}
 		scan_status = RadixHTScanStatus::DONE;
-		lock_guard<mutex> gstate_guard(sink.lock);
+		auto guard = sink.Lock();
 		if (++gstate.task_done == sink.partitions.size()) {
 			gstate.finished = true;
 		}
@@ -865,8 +865,8 @@ SourceResultType RadixPartitionedHashTable::GetData(ExecutionContext &context, D
 			for (idx_t i = 0; i < op.aggregates.size(); i++) {
 				D_ASSERT(op.aggregates[i]->GetExpressionClass() == ExpressionClass::BOUND_AGGREGATE);
 				auto &aggr = op.aggregates[i]->Cast<BoundAggregateExpression>();
-				auto aggr_state = make_unsafe_uniq_array<data_t>(aggr.function.state_size());
-				aggr.function.initialize(aggr_state.get());
+				auto aggr_state = make_unsafe_uniq_array_uninitialized<data_t>(aggr.function.state_size(aggr.function));
+				aggr.function.initialize(aggr.function, aggr_state.get());
 
 				AggregateInputData aggr_input_data(aggr.bind_info.get(), allocator);
 				Vector state_vector(Value::POINTER(CastPointerToValue(aggr_state.get())));

package/src/duckdb/src/execution/reservoir_sample.cpp

@@ -239,10 +239,10 @@ void ReservoirSamplePercentage::Finalize() {
 	// Imagine sampling 70% of 100 rows (so 70 rows). We allocate sample_percentage * RESERVOIR_THRESHOLD
 	// -----------------------------------------
 	auto sampled_more_than_required =
-	    current_count > sample_percentage * RESERVOIR_THRESHOLD || finished_samples.empty();
+	    static_cast<double>(current_count) > sample_percentage * RESERVOIR_THRESHOLD || finished_samples.empty();
 	if (current_count > 0 && sampled_more_than_required) {
 		// create a new sample
-		auto new_sample_size = idx_t(round(sample_percentage * current_count));
+		auto new_sample_size = idx_t(round(sample_percentage * static_cast<double>(current_count)));
 		auto new_sample = make_uniq<ReservoirSample>(allocator, new_sample_size, random.NextRandomInteger());
 		while (true) {
 			auto chunk = current_sample->GetChunk();