duckdb 1.1.4-dev2.0 → 1.1.4-dev9.0

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

@@ -8,8 +8,6 @@
 #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"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
 #include "duckdb/storage/temporary_memory_manager.hpp"
 
@@ -50,7 +48,7 @@ void RadixPartitionedHashTable::SetGroupingValues() {
 		for (idx_t i = 0; i < grouping.size(); i++) {
 			if (grouping_set.find(grouping[i]) == grouping_set.end()) {
 				// We don't group on this value!
-				grouping_value += (int64_t)1 << (grouping.size() - (i + 1));
+				grouping_value += 1LL << (grouping.size() - (i + 1));
 			}
 		}
 		grouping_values.push_back(Value::BIGINT(grouping_value));
@@ -94,51 +92,60 @@ class RadixHTGlobalSinkState;
 
 struct RadixHTConfig {
 public:
-	explicit RadixHTConfig(ClientContext &context, RadixHTGlobalSinkState &sink);
+	explicit RadixHTConfig(RadixHTGlobalSinkState &sink);
 
-	void SetRadixBits(idx_t radix_bits_p);
+	void SetRadixBits(const idx_t &radix_bits_p);
 	bool SetRadixBitsToExternal();
 	idx_t GetRadixBits() const;
 
 private:
-	void SetRadixBitsInternal(const idx_t radix_bits_p, bool external);
-	static idx_t InitialSinkRadixBits(ClientContext &context);
-	static idx_t MaximumSinkRadixBits(ClientContext &context);
-	static idx_t ExternalRadixBits(const idx_t &maximum_sink_radix_bits_p);
-	static idx_t SinkCapacity(ClientContext &context);
+	void SetRadixBitsInternal(idx_t radix_bits_p, bool external);
+	idx_t InitialSinkRadixBits() const;
+	idx_t MaximumSinkRadixBits() const;
+	idx_t SinkCapacity() const;
+
+private:
+	//! The global sink state
+	RadixHTGlobalSinkState &sink;
+
+public:
+	//! Number of threads (from TaskScheduler)
+	const idx_t number_of_threads;
+	//! Width of tuples
+	const idx_t row_width;
+	//! Capacity of HTs during the Sink
+	const idx_t sink_capacity;
 
 private:
 	//! Assume (1 << 15) = 32KB L1 cache per core, divided by two because hyperthreading
-	static constexpr const idx_t L1_CACHE_SIZE = 32768 / 2;
+	static constexpr idx_t L1_CACHE_SIZE = 32768 / 2;
 	//! Assume (1 << 20) = 1MB L2 cache per core, divided by two because hyperthreading
-	static constexpr const idx_t L2_CACHE_SIZE = 1048576 / 2;
+	static constexpr idx_t L2_CACHE_SIZE = 1048576 / 2;
 	//! Assume (1 << 20) + (1 << 19) = 1.5MB L3 cache per core (shared), divided by two because hyperthreading
-	static constexpr const idx_t L3_CACHE_SIZE = 1572864 / 2;
+	static constexpr idx_t L3_CACHE_SIZE = 1572864 / 2;
 
 	//! Sink radix bits to initialize with
-	static constexpr const idx_t MAXIMUM_INITIAL_SINK_RADIX_BITS = 3;
+	static constexpr idx_t MAXIMUM_INITIAL_SINK_RADIX_BITS = 4;
 	//! Maximum Sink radix bits (independent of threads)
-	static constexpr const idx_t MAXIMUM_FINAL_SINK_RADIX_BITS = 7;
-	//! By how many radix bits to increment if we go external
-	static constexpr const idx_t EXTERNAL_RADIX_BITS_INCREMENT = 3;
+	static constexpr idx_t MAXIMUM_FINAL_SINK_RADIX_BITS = 8;
 
-	//! The global sink state
-	RadixHTGlobalSinkState &sink;
 	//! Current thread-global sink radix bits
 	atomic<idx_t> sink_radix_bits;
 	//! Maximum Sink radix bits (set based on number of threads)
 	const idx_t maximum_sink_radix_bits;
-	//! Radix bits if we go external
-	const idx_t external_radix_bits;
 
-public:
-	//! Capacity of HTs during the Sink
-	const idx_t sink_capacity;
+	//! Thresholds at which we reduce the sink radix bits
+	//! This needed to reduce cache misses when we have very wide rows
+	static constexpr idx_t ROW_WIDTH_THRESHOLD_ONE = 32;
+	static constexpr idx_t ROW_WIDTH_THRESHOLD_TWO = 64;
 
+public:
+	//! If we have this many or less threads, we grow the HT, otherwise we abandon
+	static constexpr idx_t GROW_STRATEGY_THREAD_THRESHOLD = 2;
 	//! If we fill this many blocks per partition, we trigger a repartition
-	static constexpr const double BLOCK_FILL_FACTOR = 1.8;
+	static constexpr double BLOCK_FILL_FACTOR = 1.8;
 	//! By how many bits to repartition if a repartition is triggered
-	static constexpr const idx_t REPARTITION_RADIX_BITS = 2;
+	static constexpr idx_t REPARTITION_RADIX_BITS = 2;
 };
 
 class RadixHTGlobalSinkState : public GlobalSinkState {
@@ -153,11 +160,7 @@ public:
 	ClientContext &context;
 	//! Temporary memory state for managing this hash table's memory usage
 	unique_ptr<TemporaryMemoryState> temporary_memory_state;
-
-	//! The radix HT
-	const RadixPartitionedHashTable &radix_ht;
-	//! Config for partitioning
-	RadixHTConfig config;
+	idx_t minimum_reservation;
 
 	//! Whether we've called Finalize
 	bool finalized;
@@ -170,10 +173,16 @@ public:
 	//! If any thread has called combine
 	atomic<bool> any_combined;
 
+	//! The radix HT
+	const RadixPartitionedHashTable &radix_ht;
+	//! Config for partitioning
+	RadixHTConfig config;
+
 	//! Uncombined partitioned data that will be put into the AggregatePartitions
 	unique_ptr<PartitionedTupleData> uncombined_data;
 	//! Allocators used during the Sink/Finalize
 	vector<shared_ptr<ArenaAllocator>> stored_allocators;
+	idx_t stored_allocators_size;
 
 	//! Partitions that are finalized during GetData
 	vector<unique_ptr<AggregatePartition>> partitions;
@@ -190,10 +199,11 @@ public:
 
 RadixHTGlobalSinkState::RadixHTGlobalSinkState(ClientContext &context_p, const RadixPartitionedHashTable &radix_ht_p)
     : context(context_p), temporary_memory_state(TemporaryMemoryManager::Get(context).Register(context)),
-      radix_ht(radix_ht_p), config(context, *this), finalized(false), external(false), active_threads(0),
+      finalized(false), external(false), active_threads(0),
       number_of_threads(NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads())),
-      any_combined(false), finalize_done(0), scan_pin_properties(TupleDataPinProperties::DESTROY_AFTER_DONE),
-      count_before_combining(0), max_partition_size(0) {
+      any_combined(false), radix_ht(radix_ht_p), config(*this), stored_allocators_size(0), finalize_done(0),
+      scan_pin_properties(TupleDataPinProperties::DESTROY_AFTER_DONE), count_before_combining(0),
+      max_partition_size(0) {
 
 	// Compute minimum reservation
 	auto block_alloc_size = BufferManager::GetBufferManager(context).GetBlockAllocSize();
@@ -210,7 +220,7 @@ RadixHTGlobalSinkState::RadixHTGlobalSinkState(ClientContext &context_p, const R
 
 	// 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;
+	minimum_reservation = num_threads * ht_size;
 
 	temporary_memory_state->SetMinimumReservation(minimum_reservation);
 	temporary_memory_state->SetRemainingSizeAndUpdateReservation(context, minimum_reservation);
@@ -251,18 +261,18 @@ void RadixHTGlobalSinkState::Destroy() {
 }
 // LCOV_EXCL_STOP
 
-RadixHTConfig::RadixHTConfig(ClientContext &context, RadixHTGlobalSinkState &sink_p)
-    : sink(sink_p), sink_radix_bits(InitialSinkRadixBits(context)),
-      maximum_sink_radix_bits(MaximumSinkRadixBits(context)),
-      external_radix_bits(ExternalRadixBits(maximum_sink_radix_bits)), sink_capacity(SinkCapacity(context)) {
+RadixHTConfig::RadixHTConfig(RadixHTGlobalSinkState &sink_p)
+    : sink(sink_p), number_of_threads(sink.number_of_threads), row_width(sink.radix_ht.GetLayout().GetRowWidth()),
+      sink_capacity(SinkCapacity()), sink_radix_bits(InitialSinkRadixBits()),
+      maximum_sink_radix_bits(MaximumSinkRadixBits()) {
 }
 
-void RadixHTConfig::SetRadixBits(idx_t radix_bits_p) {
+void RadixHTConfig::SetRadixBits(const idx_t &radix_bits_p) {
 	SetRadixBitsInternal(MinValue(radix_bits_p, maximum_sink_radix_bits), false);
 }
 
 bool RadixHTConfig::SetRadixBitsToExternal() {
-	SetRadixBitsInternal(external_radix_bits, true);
+	SetRadixBitsInternal(MAXIMUM_FINAL_SINK_RADIX_BITS, true);
 	return sink.external;
 }
 
@@ -284,35 +294,36 @@ void RadixHTConfig::SetRadixBitsInternal(const idx_t radix_bits_p, bool external
 		sink.external = true;
 	}
 	sink_radix_bits = radix_bits_p;
-	return;
 }
 
-idx_t RadixHTConfig::InitialSinkRadixBits(ClientContext &context) {
-	const auto active_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
-	return MinValue(RadixPartitioning::RadixBits(NextPowerOfTwo(active_threads)), MAXIMUM_INITIAL_SINK_RADIX_BITS);
+idx_t RadixHTConfig::InitialSinkRadixBits() const {
+	return MinValue(RadixPartitioning::RadixBitsOfPowerOfTwo(NextPowerOfTwo(number_of_threads)),
+	                MAXIMUM_INITIAL_SINK_RADIX_BITS);
 }
 
-idx_t RadixHTConfig::MaximumSinkRadixBits(ClientContext &context) {
-	const auto active_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
-	return MinValue(RadixPartitioning::RadixBits(NextPowerOfTwo(active_threads)), MAXIMUM_FINAL_SINK_RADIX_BITS);
-}
-
-idx_t RadixHTConfig::ExternalRadixBits(const idx_t &maximum_sink_radix_bits_p) {
-	return MinValue(maximum_sink_radix_bits_p + EXTERNAL_RADIX_BITS_INCREMENT, MAXIMUM_FINAL_SINK_RADIX_BITS);
+idx_t RadixHTConfig::MaximumSinkRadixBits() const {
+	if (number_of_threads <= GROW_STRATEGY_THREAD_THRESHOLD) {
+		return InitialSinkRadixBits(); // Don't repartition unless we go external
+	}
+	// If rows are very wide we have to reduce the number of partitions, otherwise cache misses get out of hand
+	if (row_width >= ROW_WIDTH_THRESHOLD_TWO) {
+		return MAXIMUM_FINAL_SINK_RADIX_BITS - 2;
+	}
+	if (row_width >= ROW_WIDTH_THRESHOLD_ONE) {
+		return MAXIMUM_FINAL_SINK_RADIX_BITS - 1;
+	}
+	return MAXIMUM_FINAL_SINK_RADIX_BITS;
 }
 
-idx_t RadixHTConfig::SinkCapacity(ClientContext &context) {
-	// Get active and maximum number of threads
-	const auto active_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
-
+idx_t RadixHTConfig::SinkCapacity() const {
 	// Compute cache size per active thread (assuming cache is shared)
-	const auto total_shared_cache_size = active_threads * L3_CACHE_SIZE;
-	const auto cache_per_active_thread = L1_CACHE_SIZE + L2_CACHE_SIZE + total_shared_cache_size / active_threads;
+	const auto total_shared_cache_size = number_of_threads * L3_CACHE_SIZE;
+	const auto cache_per_active_thread = L1_CACHE_SIZE + L2_CACHE_SIZE + total_shared_cache_size / number_of_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(ht_entry_t) * GroupedAggregateHashTable::LOAD_FACTOR;
-	const auto capacity =
-	    NextPowerOfTwo(LossyNumericCast<uint64_t>(static_cast<double>(cache_per_active_thread) / size_per_entry));
+	const auto size_per_entry = LossyNumericCast<idx_t>(sizeof(ht_entry_t) * GroupedAggregateHashTable::LOAD_FACTOR) +
+	                            MinValue(row_width, ROW_WIDTH_THRESHOLD_TWO);
+	const auto capacity = NextPowerOfTwo(cache_per_active_thread / size_per_entry);
 
 	// Capacity must be at least the minimum capacity
 	return MaxValue<idx_t>(capacity, GroupedAggregateHashTable::InitialCapacity());
@@ -354,7 +365,7 @@ void RadixPartitionedHashTable::PopulateGroupChunk(DataChunk &group_chunk, DataC
 	for (auto &group_idx : grouping_set) {
 		// Retrieve the expression containing the index in the input chunk
 		auto &group = op.groups[group_idx];
-		D_ASSERT(group->type == ExpressionType::BOUND_REF);
+		D_ASSERT(group->GetExpressionType() == ExpressionType::BOUND_REF);
 		auto &bound_ref_expr = group->Cast<BoundReferenceExpression>();
 		// Reference from input_chunk[group.index] -> group_chunk[chunk_index]
 		group_chunk.data[chunk_index++].Reference(input_chunk.data[bound_ref_expr.index]);
@@ -363,14 +374,15 @@ void RadixPartitionedHashTable::PopulateGroupChunk(DataChunk &group_chunk, DataC
 	group_chunk.Verify();
 }
 
-bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, RadixHTLocalSinkState &lstate) {
+void MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, RadixHTLocalSinkState &lstate) {
 	auto &config = gstate.config;
 	auto &ht = *lstate.ht;
-	auto &partitioned_data = ht.GetPartitionedData();
 
 	// 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(ht_entry_t);
+	const auto aggregate_allocator_size = ht.GetAggregateAllocator()->AllocationSize();
+	const auto total_size =
+	    aggregate_allocator_size + ht.GetPartitionedData().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
@@ -379,7 +391,9 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 			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
+				// Out-of-core would be triggered below, update minimum reservation and try to increase the reservation
+				temporary_memory_state.SetMinimumReservation(aggregate_allocator_size * gstate.number_of_threads +
+				                                             gstate.minimum_reservation);
 				auto remaining_size =
 				    MaxValue<idx_t>(gstate.number_of_threads * total_size, temporary_memory_state.GetRemainingSize());
 				temporary_memory_state.SetRemainingSizeAndUpdateReservation(context, 2 * remaining_size);
@@ -396,27 +410,23 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 				    BufferManager::GetBufferManager(context), gstate.radix_ht.GetLayout(), config.GetRadixBits(),
 				    gstate.radix_ht.GetLayout().ColumnCount() - 1);
 			}
-
-			ht.UnpinData();
-			partitioned_data->Repartition(*lstate.abandoned_data);
 			ht.SetRadixBits(gstate.config.GetRadixBits());
-			ht.InitializePartitionedData();
-			return true;
+			ht.AcquirePartitionedData()->Repartition(*lstate.abandoned_data);
 		}
 	}
 
-	// We can go external when there is only one active thread, but we shouldn't repartition here
-	if (gstate.number_of_threads < 2) {
-		return false;
+	// We can go external when there are few threads, but we shouldn't repartition here
+	if (gstate.number_of_threads <= RadixHTConfig::GROW_STRATEGY_THREAD_THRESHOLD) {
+		return;
 	}
 
-	const auto partition_count = partitioned_data->PartitionCount();
-	const auto current_radix_bits = RadixPartitioning::RadixBits(partition_count);
+	const auto partition_count = ht.GetPartitionedData().PartitionCount();
+	const auto current_radix_bits = RadixPartitioning::RadixBitsOfPowerOfTwo(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;
+	    ht.GetPartitionedData().Count() * ht.GetPartitionedData().GetLayout().GetRowWidth() / partition_count;
 	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);
@@ -424,16 +434,12 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 
 	const auto global_radix_bits = config.GetRadixBits();
 	if (current_radix_bits == global_radix_bits) {
-		return false; // We're already on the right number of radix bits
+		return; // We're already on the right number of radix bits
 	}
 
 	// We're out-of-sync with the global radix bits, repartition
-	ht.UnpinData();
-	auto old_partitioned_data = std::move(partitioned_data);
 	ht.SetRadixBits(global_radix_bits);
-	ht.InitializePartitionedData();
-	old_partitioned_data->Repartition(*ht.GetPartitionedData());
-	return true;
+	ht.Repartition();
 }
 
 void RadixPartitionedHashTable::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input,
@@ -451,25 +457,41 @@ void RadixPartitionedHashTable::Sink(ExecutionContext &context, DataChunk &chunk
 	auto &ht = *lstate.ht;
 	ht.AddChunk(group_chunk, payload_input, filter);
 
-	if (ht.Count() + STANDARD_VECTOR_SIZE < ht.ResizeThreshold()) {
+	if (ht.Count() + STANDARD_VECTOR_SIZE < GroupedAggregateHashTable::ResizeThreshold(gstate.config.sink_capacity)) {
 		return; // We can fit another chunk
 	}
 
-	if (gstate.number_of_threads > 2) {
+	if (gstate.number_of_threads > RadixHTConfig::GROW_STRATEGY_THREAD_THRESHOLD || gstate.external) {
 		// 'Reset' the HT without taking its data, we can just keep appending to the same collection
 		// This only works because we never resize the HT
-		ht.ClearPointerTable();
-		ht.ResetCount();
 		// We don't do this when running with 1 or 2 threads, it only makes sense when there's many threads
+		ht.Abandon();
+
+		// Once we've inserted more than SKIP_LOOKUP_THRESHOLD tuples,
+		// and more than UNIQUE_PERCENTAGE_THRESHOLD were unique,
+		// we set the HT to skip doing lookups, which makes it blindly append data to the HT.
+		// This speeds up adding data, at the cost of no longer de-duplicating.
+		// The data will be de-duplicated later anyway
+		static constexpr idx_t SKIP_LOOKUP_THRESHOLD = 262144;
+		static constexpr double UNIQUE_PERCENTAGE_THRESHOLD = 0.95;
+		const auto unique_percentage =
+		    static_cast<double>(ht.GetPartitionedData().Count()) / static_cast<double>(ht.GetSinkCount());
+		if (ht.GetSinkCount() > SKIP_LOOKUP_THRESHOLD && unique_percentage > UNIQUE_PERCENTAGE_THRESHOLD) {
+			ht.SkipLookups();
+		}
 	}
 
 	// Check if we need to repartition
-	auto repartitioned = MaybeRepartition(context.client, gstate, lstate);
+	const auto radix_bits_before = ht.GetRadixBits();
+	MaybeRepartition(context.client, gstate, lstate);
+	const auto repartitioned = radix_bits_before != ht.GetRadixBits();
 
 	if (repartitioned && ht.Count() != 0) {
 		// We repartitioned, but we didn't clear the pointer table / reset the count because we're on 1 or 2 threads
-		ht.ClearPointerTable();
-		ht.ResetCount();
+		ht.Abandon();
+		if (gstate.external) {
+			ht.Resize(gstate.config.sink_capacity);
+		}
 	}
 
 	// TODO: combine early and often
@@ -488,16 +510,15 @@ void RadixPartitionedHashTable::Combine(ExecutionContext &context, GlobalSinkSta
 	MaybeRepartition(context.client, gstate, lstate);
 
 	auto &ht = *lstate.ht;
-	ht.UnpinData();
-
+	auto lstate_data = ht.AcquirePartitionedData();
 	if (lstate.abandoned_data) {
 		D_ASSERT(gstate.external);
-		D_ASSERT(lstate.abandoned_data->PartitionCount() == lstate.ht->GetPartitionedData()->PartitionCount());
+		D_ASSERT(lstate.abandoned_data->PartitionCount() == lstate.ht->GetPartitionedData().PartitionCount());
 		D_ASSERT(lstate.abandoned_data->PartitionCount() ==
 		         RadixPartitioning::NumberOfPartitions(gstate.config.GetRadixBits()));
-		lstate.abandoned_data->Combine(*lstate.ht->GetPartitionedData());
+		lstate.abandoned_data->Combine(*lstate_data);
 	} else {
-		lstate.abandoned_data = std::move(ht.GetPartitionedData());
+		lstate.abandoned_data = std::move(lstate_data);
 	}
 
 	auto guard = gstate.Lock();
@@ -507,6 +528,7 @@ void RadixPartitionedHashTable::Combine(ExecutionContext &context, GlobalSinkSta
 		gstate.uncombined_data = std::move(lstate.abandoned_data);
 	}
 	gstate.stored_allocators.emplace_back(ht.GetAggregateAllocator());
+	gstate.stored_allocators_size += gstate.stored_allocators.back()->AllocationSize();
 }
 
 void RadixPartitionedHashTable::Finalize(ClientContext &context, GlobalSinkState &gstate_p) const {
@@ -541,7 +563,7 @@ void RadixPartitionedHashTable::Finalize(ClientContext &context, GlobalSinkState
 	}
 
 	// Minimum of combining one partition at a time
-	gstate.temporary_memory_state->SetMinimumReservation(gstate.max_partition_size);
+	gstate.temporary_memory_state->SetMinimumReservation(gstate.stored_allocators_size + gstate.max_partition_size);
 	// Set size to 0 until the scan actually starts
 	gstate.temporary_memory_state->SetZero();
 	gstate.finalized = true;
@@ -558,12 +580,15 @@ idx_t RadixPartitionedHashTable::MaxThreads(GlobalSinkState &sink_p) const {
 
 	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);
+	sink.temporary_memory_state->SetRemainingSizeAndUpdateReservation(
+	    sink.context, sink.stored_allocators_size + max_threads * sink.max_partition_size);
 
+	// we cannot spill aggregate state memory
+	const auto usable_memory = sink.temporary_memory_state->GetReservation() > sink.stored_allocators_size
+	                               ? sink.temporary_memory_state->GetReservation() - sink.max_partition_size
+	                               : 0;
 	// This many partitions will fit given our reservation (at least 1))
-	const auto partitions_fit =
-	    MaxValue<idx_t>(sink.temporary_memory_state->GetReservation() / sink.max_partition_size, 1);
+	const auto partitions_fit = MaxValue<idx_t>(usable_memory / sink.max_partition_size, 1);
 
 	// Mininum of the two
 	return MinValue<idx_t>(partitions_fit, max_threads);
@@ -684,8 +709,8 @@ SourceResultType RadixHTGlobalSourceState::AssignTask(RadixHTGlobalSinkState &si
 }
 
 RadixHTLocalSourceState::RadixHTLocalSourceState(ExecutionContext &context, const RadixPartitionedHashTable &radix_ht)
-    : task(RadixHTSourceTaskType::NO_TASK), scan_status(RadixHTScanStatus::DONE), layout(radix_ht.GetLayout().Copy()),
-      aggregate_allocator(BufferAllocator::Get(context.client)) {
+    : task(RadixHTSourceTaskType::NO_TASK), task_idx(DConstants::INVALID_INDEX), scan_status(RadixHTScanStatus::DONE),
+      layout(radix_ht.GetLayout().Copy()), aggregate_allocator(BufferAllocator::Get(context.client)) {
 	auto &allocator = BufferAllocator::Get(context.client);
 	auto scan_chunk_types = radix_ht.group_types;
 	for (auto &aggr_type : radix_ht.op.aggregate_return_types) {
@@ -731,21 +756,17 @@ void RadixHTLocalSourceState::Finalize(RadixHTGlobalSinkState &sink, RadixHTGlob
 
 		ht = sink.radix_ht.CreateHT(gstate.context, MinValue<idx_t>(capacity, capacity_limit), 0);
 	} else {
-		// We may want to resize here to the size of this partition, but for now we just assume uniform partition sizes
-		ht->InitializePartitionedData();
-		ht->ClearPointerTable();
-		ht->ResetCount();
+		ht->Abandon();
 	}
 
 	// Now combine the uncombined data using this thread's HT
 	ht->Combine(*partition.data, &partition.progress);
-	ht->UnpinData();
 	partition.progress = 1;
 
 	// Move the combined data back to the partition
 	partition.data =
 	    make_uniq<TupleDataCollection>(BufferManager::GetBufferManager(gstate.context), sink.radix_ht.GetLayout());
-	partition.data->Combine(*ht->GetPartitionedData()->GetPartitions()[0]);
+	partition.data->Combine(*ht->AcquirePartitionedData()->GetPartitions()[0]);
 
 	// Update thread-global state
 	auto guard = sink.Lock();
@@ -903,25 +924,24 @@ SourceResultType RadixPartitionedHashTable::GetData(ExecutionContext &context, D
 	}
 }
 
-double RadixPartitionedHashTable::GetProgress(ClientContext &, GlobalSinkState &sink_p,
-                                              GlobalSourceState &gstate_p) const {
+ProgressData RadixPartitionedHashTable::GetProgress(ClientContext &, GlobalSinkState &sink_p,
+                                                    GlobalSourceState &gstate_p) const {
 	auto &sink = sink_p.Cast<RadixHTGlobalSinkState>();
 	auto &gstate = gstate_p.Cast<RadixHTGlobalSourceState>();
 
 	// Get partition combine progress, weigh it 2x
-	double total_progress = 0;
+	ProgressData progress;
 	for (auto &partition : sink.partitions) {
-		total_progress += 2.0 * partition->progress;
+		progress.done += 2.0 * partition->progress;
 	}
 
 	// Get scan progress, weigh it 1x
-	total_progress += 1.0 * double(gstate.task_done);
+	progress.done += 1.0 * double(gstate.task_done);
 
 	// Divide by 3x for the weights, and the number of partitions to get a value between 0 and 1 again
-	total_progress /= 3.0 * double(sink.partitions.size());
+	progress.total += 3.0 * double(sink.partitions.size());
 
-	// Multiply by 100 to get a percentage
-	return 100.0 * total_progress;
+	return progress;
 }
 
 } // namespace duckdb

package/src/duckdb/src/execution/sample/base_reservoir_sample.cpp

@@ -0,0 +1,136 @@
+#include "duckdb/execution/reservoir_sample.hpp"
+#include <math.h>
+
+namespace duckdb {
+
+double BaseReservoirSampling::GetMinWeightFromTuplesSeen(idx_t rows_seen_total) {
+	// this function was obtained using https://mycurvefit.com. Inputting multiple x, y values into
+	// The
+	switch (rows_seen_total) {
+	case 0:
+		return 0;
+	case 1:
+		return 0.000161;
+	case 2:
+		return 0.530136;
+	case 3:
+		return 0.693454;
+	default: {
+		return (0.99 - 0.355 * std::exp(-0.07 * static_cast<double>(rows_seen_total)));
+	}
+	}
+}
+
+BaseReservoirSampling::BaseReservoirSampling(int64_t seed) : random(seed) {
+	next_index_to_sample = 0;
+	min_weight_threshold = 0;
+	min_weighted_entry_index = 0;
+	num_entries_to_skip_b4_next_sample = 0;
+	num_entries_seen_total = 0;
+}
+
+BaseReservoirSampling::BaseReservoirSampling() : BaseReservoirSampling(1) {
+}
+
+unique_ptr<BaseReservoirSampling> BaseReservoirSampling::Copy() {
+	auto ret = make_uniq<BaseReservoirSampling>(1);
+	ret->reservoir_weights = reservoir_weights;
+	ret->next_index_to_sample = next_index_to_sample;
+	ret->min_weight_threshold = min_weight_threshold;
+	ret->min_weighted_entry_index = min_weighted_entry_index;
+	ret->num_entries_to_skip_b4_next_sample = num_entries_to_skip_b4_next_sample;
+	ret->num_entries_seen_total = num_entries_seen_total;
+	return ret;
+}
+
+void BaseReservoirSampling::InitializeReservoirWeights(idx_t cur_size, idx_t sample_size) {
+	//! 1: The first m items of V are inserted into R
+	//! first we need to check if the reservoir already has "m" elements
+	//! 2. For each item vi ∈ R: Calculate a key ki = random(0, 1)
+	//! we then define the threshold to enter the reservoir T_w as the minimum key of R
+	//! we use a priority queue to extract the minimum key in O(1) time
+	if (cur_size == sample_size) {
+		//! 2. For each item vi ∈ R: Calculate a key ki = random(0, 1)
+		//! we then define the threshold to enter the reservoir T_w as the minimum key of R
+		//! we use a priority queue to extract the minimum key in O(1) time
+		for (idx_t i = 0; i < sample_size; i++) {
+			double k_i = random.NextRandom();
+			reservoir_weights.emplace(-k_i, i);
+		}
+		SetNextEntry();
+	}
+}
+
+void BaseReservoirSampling::SetNextEntry() {
+	D_ASSERT(!reservoir_weights.empty());
+	//! 4. Let r = random(0, 1) and Xw = log(r) / log(T_w)
+	auto &min_key = reservoir_weights.top();
+	double t_w = -min_key.first;
+	double r = random.NextRandom32();
+	double x_w = log(r) / log(t_w);
+	//! 5. From the current item vc skip items until item vi , such that:
+	//! 6. wc +wc+1 +···+wi−1 < Xw <= wc +wc+1 +···+wi−1 +wi
+	//! since all our weights are 1 (uniform sampling), we can just determine the amount of elements to skip
+	min_weight_threshold = t_w;
+	min_weighted_entry_index = min_key.second;
+	next_index_to_sample = MaxValue<idx_t>(1, idx_t(round(x_w)));
+	num_entries_to_skip_b4_next_sample = 0;
+}
+
+void BaseReservoirSampling::ReplaceElementWithIndex(idx_t entry_index, double with_weight, bool pop) {
+
+	if (pop) {
+		reservoir_weights.pop();
+	}
+	double r2 = with_weight;
+	//! now we insert the new weight into the reservoir
+	reservoir_weights.emplace(-r2, entry_index);
+	//! we update the min entry with the new min entry in the reservoir
+	SetNextEntry();
+}
+
+void BaseReservoirSampling::ReplaceElement(double with_weight) {
+	//! replace the entry in the reservoir
+	//! pop the minimum entry
+	reservoir_weights.pop();
+	//! now update the reservoir
+	//! 8. Let tw = Tw i , r2 = random(tw,1) and vi’s key: ki = (r2)1/wi
+	//! 9. The new threshold Tw is the new minimum key of R
+	//! we generate a random number between (min_weight_threshold, 1)
+	double r2 = random.NextRandom(min_weight_threshold, 1);
+
+	//! if we are merging two reservoir samples use the weight passed
+	if (with_weight >= 0) {
+		r2 = with_weight;
+	}
+	//! now we insert the new weight into the reservoir
+	reservoir_weights.emplace(-r2, min_weighted_entry_index);
+	//! we update the min entry with the new min entry in the reservoir
+	SetNextEntry();
+}
+
+void BaseReservoirSampling::UpdateMinWeightThreshold() {
+	if (!reservoir_weights.empty()) {
+		min_weight_threshold = -reservoir_weights.top().first;
+		min_weighted_entry_index = reservoir_weights.top().second;
+		return;
+	}
+	min_weight_threshold = 1;
+}
+
+void BaseReservoirSampling::FillWeights(SelectionVector &sel, idx_t &sel_size) {
+	if (!reservoir_weights.empty()) {
+		return;
+	}
+	D_ASSERT(reservoir_weights.empty());
+	auto num_entries_seen_normalized = num_entries_seen_total / FIXED_SAMPLE_SIZE;
+	auto min_weight = GetMinWeightFromTuplesSeen(num_entries_seen_normalized);
+	for (idx_t i = 0; i < sel_size; i++) {
+		auto weight = random.NextRandom(min_weight, 1);
+		reservoir_weights.emplace(-weight, i);
+	}
+	D_ASSERT(reservoir_weights.size() <= sel_size);
+	SetNextEntry();
+}
+
+} // namespace duckdb