duckdb 0.10.3-dev3.0 → 0.10.3-dev6.0

package/src/duckdb/src/execution/physical_plan/plan_comparison_join.cpp

@@ -109,7 +109,7 @@ void CheckForPerfectJoinOpt(LogicalComparisonJoin &op, PerfectHashJoinStats &joi
 	join_state.build_min = NumericStats::Min(stats_build);
 	join_state.build_max = NumericStats::Max(stats_build);
 	join_state.estimated_cardinality = op.estimated_cardinality;
-	join_state.build_range = build_range;
+	join_state.build_range = NumericCast<idx_t>(build_range);
 	if (join_state.build_range > MAX_BUILD_SIZE) {
 		return;
 	}

package/src/duckdb/src/execution/physical_plan/plan_create_table.cpp

@@ -22,10 +22,10 @@ unique_ptr<PhysicalOperator> DuckCatalog::PlanCreateTableAs(ClientContext &conte
 	auto num_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
 	unique_ptr<PhysicalOperator> create;
 	if (!parallel_streaming_insert && use_batch_index) {
-		create = make_uniq<PhysicalBatchInsert>(op, op.schema, std::move(op.info), 0);
+		create = make_uniq<PhysicalBatchInsert>(op, op.schema, std::move(op.info), 0U);
 
 	} else {
-		create = make_uniq<PhysicalInsert>(op, op.schema, std::move(op.info), 0,
+		create = make_uniq<PhysicalInsert>(op, op.schema, std::move(op.info), 0U,
 		                                   parallel_streaming_insert && num_threads > 1);
 	}
 

package/src/duckdb/src/execution/physical_plan/plan_cte.cpp

@@ -12,7 +12,7 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalMaterializ
 	D_ASSERT(op.children.size() == 2);
 
 	// Create the working_table that the PhysicalCTE will use for evaluation.
-	auto working_table = std::make_shared<ColumnDataCollection>(context, op.children[0]->types);
+	auto working_table = make_shared_ptr<ColumnDataCollection>(context, op.children[0]->types);
 
 	// Add the ColumnDataCollection to the context of this PhysicalPlanGenerator
 	recursive_cte_tables[op.table_index] = working_table;

package/src/duckdb/src/execution/physical_plan/plan_delete.cpp

@@ -12,8 +12,8 @@ unique_ptr<PhysicalOperator> DuckCatalog::PlanDelete(ClientContext &context, Log
 	// get the index of the row_id column
 	auto &bound_ref = op.expressions[0]->Cast<BoundReferenceExpression>();
 
-	auto del = make_uniq<PhysicalDelete>(op.types, op.table, op.table.GetStorage(), bound_ref.index,
-	                                     op.estimated_cardinality, op.return_chunk);
+	auto del = make_uniq<PhysicalDelete>(op.types, op.table, op.table.GetStorage(), std::move(op.bound_constraints),
+	                                     bound_ref.index, op.estimated_cardinality, op.return_chunk);
 	del->children.push_back(std::move(plan));
 	return std::move(del);
 }

package/src/duckdb/src/execution/physical_plan/plan_delim_get.cpp

@@ -8,8 +8,8 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalDelimGet &
 	D_ASSERT(op.children.empty());
 
 	// create a PhysicalChunkScan without an owned_collection, the collection will be added later
-	auto chunk_scan =
-	    make_uniq<PhysicalColumnDataScan>(op.types, PhysicalOperatorType::DELIM_SCAN, op.estimated_cardinality);
+	auto chunk_scan = make_uniq<PhysicalColumnDataScan>(op.types, PhysicalOperatorType::DELIM_SCAN,
+	                                                    op.estimated_cardinality, nullptr);
 	return std::move(chunk_scan);
 }
 

package/src/duckdb/src/execution/physical_plan/plan_distinct.cpp

@@ -67,6 +67,7 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalDistinct &
 				bool changes_made = false;
 				auto new_expr = OrderedAggregateOptimizer::Apply(context, *first_aggregate, groups, changes_made);
 				if (new_expr) {
+					D_ASSERT(new_expr->return_type == first_aggregate->return_type);
 					D_ASSERT(new_expr->type == ExpressionType::BOUND_AGGREGATE);
 					first_aggregate = unique_ptr_cast<Expression, BoundAggregateExpression>(std::move(new_expr));
 				}

package/src/duckdb/src/execution/physical_plan/plan_expression_get.cpp

@@ -19,19 +19,18 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalExpression
 	// simple expression scan (i.e. no subqueries to evaluate and no prepared statement parameters)
 	// we can evaluate all the expressions right now and turn this into a chunk collection scan
 	auto chunk_scan = make_uniq<PhysicalColumnDataScan>(op.types, PhysicalOperatorType::COLUMN_DATA_SCAN,
-	                                                    expr_scan->expressions.size());
-	chunk_scan->owned_collection = make_uniq<ColumnDataCollection>(context, op.types);
-	chunk_scan->collection = chunk_scan->owned_collection.get();
+	                                                    expr_scan->expressions.size(),
+	                                                    make_uniq<ColumnDataCollection>(context, op.types));
 
 	DataChunk chunk;
 	chunk.Initialize(allocator, op.types);
 
 	ColumnDataAppendState append_state;
-	chunk_scan->owned_collection->InitializeAppend(append_state);
+	chunk_scan->collection->InitializeAppend(append_state);
 	for (idx_t expression_idx = 0; expression_idx < expr_scan->expressions.size(); expression_idx++) {
 		chunk.Reset();
 		expr_scan->EvaluateExpression(context, expression_idx, nullptr, chunk);
-		chunk_scan->owned_collection->Append(append_state, chunk);
+		chunk_scan->collection->Append(append_state, chunk);
 	}
 	return std::move(chunk_scan);
 }

package/src/duckdb/src/execution/physical_plan/plan_insert.cpp

@@ -94,13 +94,14 @@ unique_ptr<PhysicalOperator> DuckCatalog::PlanInsert(ClientContext &context, Log
 	unique_ptr<PhysicalOperator> insert;
 	if (use_batch_index && !parallel_streaming_insert) {
 		insert = make_uniq<PhysicalBatchInsert>(op.types, op.table, op.column_index_map, std::move(op.bound_defaults),
-		                                        op.estimated_cardinality);
+		                                        std::move(op.bound_constraints), op.estimated_cardinality);
 	} else {
 		insert = make_uniq<PhysicalInsert>(
-		    op.types, op.table, op.column_index_map, std::move(op.bound_defaults), std::move(op.expressions),
-		    std::move(op.set_columns), std::move(op.set_types), op.estimated_cardinality, op.return_chunk,
-		    parallel_streaming_insert && num_threads > 1, op.action_type, std::move(op.on_conflict_condition),
-		    std::move(op.do_update_condition), std::move(op.on_conflict_filter), std::move(op.columns_to_fetch));
+		    op.types, op.table, op.column_index_map, std::move(op.bound_defaults), std::move(op.bound_constraints),
+		    std::move(op.expressions), std::move(op.set_columns), std::move(op.set_types), op.estimated_cardinality,
+		    op.return_chunk, parallel_streaming_insert && num_threads > 1, op.action_type,
+		    std::move(op.on_conflict_condition), std::move(op.do_update_condition), std::move(op.on_conflict_filter),
+		    std::move(op.columns_to_fetch));
 	}
 	D_ASSERT(plan);
 	insert->children.push_back(std::move(plan));

package/src/duckdb/src/execution/physical_plan/plan_recursive_cte.cpp

@@ -12,7 +12,7 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalRecursiveC
 	D_ASSERT(op.children.size() == 2);
 
 	// Create the working_table that the PhysicalRecursiveCTE will use for evaluation.
-	auto working_table = std::make_shared<ColumnDataCollection>(context, op.types);
+	auto working_table = make_shared_ptr<ColumnDataCollection>(context, op.types);
 
 	// Add the ColumnDataCollection to the context of this PhysicalPlanGenerator
 	recursive_cte_tables[op.table_index] = working_table;

package/src/duckdb/src/execution/physical_plan/plan_simple.cpp

@@ -1,6 +1,7 @@
 #include "duckdb/execution/operator/helper/physical_load.hpp"
 #include "duckdb/execution/operator/helper/physical_transaction.hpp"
 #include "duckdb/execution/operator/helper/physical_vacuum.hpp"
+#include "duckdb/execution/operator/helper/physical_update_extensions.hpp"
 #include "duckdb/execution/operator/schema/physical_alter.hpp"
 #include "duckdb/execution/operator/schema/physical_attach.hpp"
 #include "duckdb/execution/operator/schema/physical_create_schema.hpp"
@@ -34,6 +35,9 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalSimple &op
 	case LogicalOperatorType::LOGICAL_DETACH:
 		return make_uniq<PhysicalDetach>(unique_ptr_cast<ParseInfo, DetachInfo>(std::move(op.info)),
 		                                 op.estimated_cardinality);
+	case LogicalOperatorType::LOGICAL_UPDATE_EXTENSIONS:
+		return make_uniq<PhysicalUpdateExtensions>(unique_ptr_cast<ParseInfo, UpdateExtensionsInfo>(std::move(op.info)),
+		                                           op.estimated_cardinality);
 	default:
 		throw NotImplementedException("Unimplemented type for logical simple operator");
 	}

package/src/duckdb/src/execution/physical_plan/plan_top_n.cpp

@@ -9,8 +9,8 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalTopN &op)
 
 	auto plan = CreatePlan(*op.children[0]);
 
-	auto top_n =
-	    make_uniq<PhysicalTopN>(op.types, std::move(op.orders), (idx_t)op.limit, op.offset, op.estimated_cardinality);
+	auto top_n = make_uniq<PhysicalTopN>(op.types, std::move(op.orders), NumericCast<idx_t>(op.limit),
+	                                     NumericCast<idx_t>(op.offset), op.estimated_cardinality);
 	top_n->children.push_back(std::move(plan));
 	return std::move(top_n);
 }

package/src/duckdb/src/execution/physical_plan/plan_update.cpp

@@ -8,9 +8,9 @@ namespace duckdb {
 
 unique_ptr<PhysicalOperator> DuckCatalog::PlanUpdate(ClientContext &context, LogicalUpdate &op,
                                                      unique_ptr<PhysicalOperator> plan) {
-	auto update =
-	    make_uniq<PhysicalUpdate>(op.types, op.table, op.table.GetStorage(), op.columns, std::move(op.expressions),
-	                              std::move(op.bound_defaults), op.estimated_cardinality, op.return_chunk);
+	auto update = make_uniq<PhysicalUpdate>(op.types, op.table, op.table.GetStorage(), op.columns,
+	                                        std::move(op.expressions), std::move(op.bound_defaults),
+	                                        std::move(op.bound_constraints), op.estimated_cardinality, op.return_chunk);
 
 	update->update_is_del_and_insert = op.update_is_del_and_insert;
 	update->children.push_back(std::move(plan));

package/src/duckdb/src/execution/physical_plan/plan_window.cpp

@@ -10,29 +10,6 @@
 
 namespace duckdb {
 
-static bool IsStreamingWindow(unique_ptr<Expression> &expr) {
-	auto &wexpr = expr->Cast<BoundWindowExpression>();
-	if (!wexpr.partitions.empty() || !wexpr.orders.empty() || wexpr.ignore_nulls ||
-	    wexpr.exclude_clause != WindowExcludeMode::NO_OTHER) {
-		return false;
-	}
-	switch (wexpr.type) {
-	// TODO: add more expression types here?
-	case ExpressionType::WINDOW_AGGREGATE:
-		// We can stream aggregates if they are "running totals" and don't use filters
-		return wexpr.start == WindowBoundary::UNBOUNDED_PRECEDING && wexpr.end == WindowBoundary::CURRENT_ROW_ROWS &&
-		       !wexpr.filter_expr;
-	case ExpressionType::WINDOW_FIRST_VALUE:
-	case ExpressionType::WINDOW_PERCENT_RANK:
-	case ExpressionType::WINDOW_RANK:
-	case ExpressionType::WINDOW_RANK_DENSE:
-	case ExpressionType::WINDOW_ROW_NUMBER:
-		return true;
-	default:
-		return false;
-	}
-}
-
 unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalWindow &op) {
 	D_ASSERT(op.children.size() == 1);
 
@@ -54,7 +31,7 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalWindow &op
 	vector<idx_t> blocking_windows;
 	vector<idx_t> streaming_windows;
 	for (idx_t expr_idx = 0; expr_idx < op.expressions.size(); expr_idx++) {
-		if (IsStreamingWindow(op.expressions[expr_idx])) {
+		if (PhysicalStreamingWindow::IsStreamingFunction(op.expressions[expr_idx])) {
 			streaming_windows.push_back(expr_idx);
 		} else {
 			blocking_windows.push_back(expr_idx);

package/src/duckdb/src/execution/physical_plan_generator.cpp

@@ -211,6 +211,9 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalOperator &
 	case LogicalOperatorType::LOGICAL_COPY_DATABASE:
 		plan = CreatePlan(op.Cast<LogicalCopyDatabase>());
 		break;
+	case LogicalOperatorType::LOGICAL_UPDATE_EXTENSIONS:
+		plan = CreatePlan(op.Cast<LogicalSimple>());
+		break;
 	case LogicalOperatorType::LOGICAL_EXTENSION_OPERATOR:
 		plan = op.Cast<LogicalExtensionOperator>().CreatePlan(context, *this);
 

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

@@ -167,6 +167,8 @@ public:
 	atomic<bool> external;
 	//! Threads that have called Sink
 	atomic<idx_t> active_threads;
+	//! Number of threads (from TaskScheduler)
+	const idx_t number_of_threads;
 	//! If any thread has called combine
 	atomic<bool> any_combined;
 
@@ -193,18 +195,20 @@ 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),
+      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) {
 
 	auto tuples_per_block = Storage::BLOCK_ALLOC_SIZE / radix_ht.GetLayout().GetRowWidth();
-	idx_t ht_count = config.sink_capacity / GroupedAggregateHashTable::LOAD_FACTOR;
+	idx_t ht_count =
+	    NumericCast<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);
 
 	// This really is the minimum reservation that we can do
-	idx_t num_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	auto num_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
 	auto minimum_reservation = num_threads * ht_size;
 
 	temporary_memory_state->SetMinimumReservation(minimum_reservation);
@@ -229,6 +233,7 @@ void RadixHTGlobalSinkState::Destroy() {
 	}
 
 	// There are aggregates with destructors: Call the destructor for each of the aggregates
+	lock_guard<mutex> guard(lock);
 	RowOperationsState row_state(*stored_allocators.back());
 	for (auto &partition : partitions) {
 		auto &data_collection = *partition->data;
@@ -282,12 +287,12 @@ void RadixHTConfig::SetRadixBitsInternal(const idx_t radix_bits_p, bool external
 }
 
 idx_t RadixHTConfig::InitialSinkRadixBits(ClientContext &context) {
-	const idx_t active_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	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::MaximumSinkRadixBits(ClientContext &context) {
-	const idx_t active_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	const auto active_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
 	return MinValue(RadixPartitioning::RadixBits(NextPowerOfTwo(active_threads)), MAXIMUM_FINAL_SINK_RADIX_BITS);
 }
 
@@ -297,7 +302,7 @@ idx_t RadixHTConfig::ExternalRadixBits(const idx_t &maximum_sink_radix_bits_p) {
 
 idx_t RadixHTConfig::SinkCapacity(ClientContext &context) {
 	// Get active and maximum number of threads
-	const idx_t active_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	const auto active_threads = NumericCast<idx_t>(TaskScheduler::GetScheduler(context).NumberOfThreads());
 
 	// Compute cache size per active thread (assuming cache is shared)
 	const auto total_shared_cache_size = active_threads * L3_CACHE_SIZE;
@@ -305,7 +310,8 @@ idx_t RadixHTConfig::SinkCapacity(ClientContext &context) {
 
 	// 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 capacity = NextPowerOfTwo(cache_per_active_thread / size_per_entry);
+	const auto capacity =
+	    NextPowerOfTwo(NumericCast<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());
@@ -356,8 +362,7 @@ void RadixPartitionedHashTable::PopulateGroupChunk(DataChunk &group_chunk, DataC
 	group_chunk.Verify();
 }
 
-bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, RadixHTLocalSinkState &lstate,
-                      const idx_t &active_threads) {
+bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, RadixHTLocalSinkState &lstate) {
 	auto &config = gstate.config;
 	auto &ht = *lstate.ht;
 	auto &partitioned_data = ht.GetPartitionedData();
@@ -365,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);
-	idx_t thread_limit = temporary_memory_state.GetReservation() / active_threads;
+	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);
-			thread_limit = temporary_memory_state.GetReservation() / active_threads;
+			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>(active_threads * total_size, temporary_memory_state.GetRemainingSize());
+				    MaxValue<idx_t>(gstate.number_of_threads * total_size, temporary_memory_state.GetRemainingSize());
 				temporary_memory_state.SetRemainingSize(context, 2 * remaining_size);
-				thread_limit = temporary_memory_state.GetReservation() / active_threads;
+				thread_limit = temporary_memory_state.GetReservation() / gstate.number_of_threads;
 			}
 		}
 	}
@@ -400,7 +405,7 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 	}
 
 	// We can go external when there is only one active thread, but we shouldn't repartition here
-	if (active_threads < 2) {
+	if (gstate.number_of_threads < 2) {
 		return false;
 	}
 
@@ -410,7 +415,7 @@ bool MaybeRepartition(ClientContext &context, RadixHTGlobalSinkState &gstate, Ra
 
 	const auto row_size_per_partition =
 	    partitioned_data->Count() * partitioned_data->GetLayout().GetRowWidth() / partition_count;
-	if (row_size_per_partition > config.BLOCK_FILL_FACTOR * Storage::BLOCK_SIZE) {
+	if (row_size_per_partition > NumericCast<idx_t>(config.BLOCK_FILL_FACTOR * Storage::BLOCK_SIZE)) {
 		// We crossed our block filling threshold, try to increment radix bits
 		config.SetRadixBits(current_radix_bits + config.REPARTITION_RADIX_BITS);
 	}
@@ -448,8 +453,7 @@ void RadixPartitionedHashTable::Sink(ExecutionContext &context, DataChunk &chunk
 		return; // We can fit another chunk
 	}
 
-	const idx_t active_threads = gstate.active_threads;
-	if (active_threads > 2) {
+	if (gstate.number_of_threads > 2) {
 		// '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();
@@ -458,7 +462,7 @@ void RadixPartitionedHashTable::Sink(ExecutionContext &context, DataChunk &chunk
 	}
 
 	// Check if we need to repartition
-	auto repartitioned = MaybeRepartition(context.client, gstate, lstate, active_threads);
+	auto repartitioned = MaybeRepartition(context.client, gstate, lstate);
 
 	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
@@ -479,7 +483,7 @@ void RadixPartitionedHashTable::Combine(ExecutionContext &context, GlobalSinkSta
 
 	// Set any_combined, then check one last time whether we need to repartition
 	gstate.any_combined = true;
-	MaybeRepartition(context.client, gstate, lstate, gstate.active_threads);
+	MaybeRepartition(context.client, gstate, lstate);
 
 	auto &ht = *lstate.ht;
 	ht.UnpinData();
@@ -511,7 +515,7 @@ void RadixPartitionedHashTable::Finalize(ClientContext &context, GlobalSinkState
 		gstate.count_before_combining = uncombined_data.Count();
 
 		// If true there is no need to combine, it was all done by a single thread in a single HT
-		const auto single_ht = !gstate.external && gstate.active_threads == 1;
+		const auto single_ht = !gstate.external && gstate.active_threads == 1 && gstate.number_of_threads == 1;
 
 		auto &uncombined_partition_data = uncombined_data.GetPartitions();
 		const auto n_partitions = uncombined_partition_data.size();
@@ -537,8 +541,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>(TaskScheduler::GetScheduler(context).NumberOfThreads(), gstate.partitions.size());
+	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);
 	gstate.finalized = true;
 }
@@ -555,8 +559,8 @@ idx_t RadixPartitionedHashTable::MaxThreads(GlobalSinkState &sink_p) const {
 	// 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(), TaskScheduler::GetScheduler(sink.context).NumberOfThreads());
+	auto max_possible = MinValue<idx_t>(
+	    sink.partitions.size(), NumericCast<idx_t>(TaskScheduler::GetScheduler(sink.context).NumberOfThreads()));
 
 	// Mininum of the two
 	return MinValue<idx_t>(partitions_fit, max_possible);
@@ -589,7 +593,6 @@ public:
 	vector<column_t> column_ids;
 
 	//! For synchronizing tasks
-	mutex lock;
 	idx_t task_idx;
 	atomic<idx_t> task_done;
 };
@@ -649,7 +652,7 @@ 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(lock);
+	lock_guard<mutex> gstate_guard(sink.lock);
 	if (finished) {
 		return SourceResultType::FINISHED;
 	}
@@ -716,13 +719,15 @@ void RadixHTLocalSourceState::Finalize(RadixHTGlobalSinkState &sink, RadixHTGlob
 		const auto capacity = GroupedAggregateHashTable::GetCapacityForCount(partition.data->Count());
 
 		// However, we will limit the initial capacity so we don't do a huge over-allocation
-		const idx_t n_threads = TaskScheduler::GetScheduler(gstate.context).NumberOfThreads();
-		const idx_t memory_limit = BufferManager::GetBufferManager(gstate.context).GetMaxMemory();
-		const idx_t thread_limit = 0.6 * memory_limit / n_threads;
+		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 size_per_entry = partition.data->SizeInBytes() / MaxValue<idx_t>(partition.data->Count(), 1) +
 		                             idx_t(GroupedAggregateHashTable::LOAD_FACTOR * sizeof(aggr_ht_entry_t));
-		const auto capacity_limit = NextPowerOfTwo(thread_limit / size_per_entry);
+		// but not lower than the initial capacity
+		const auto capacity_limit =
+		    MaxValue(NextPowerOfTwo(thread_limit / size_per_entry), GroupedAggregateHashTable::InitialCapacity());
 
 		ht = sink.radix_ht.CreateHT(gstate.context, MinValue<idx_t>(capacity, capacity_limit), 0);
 	} else {
@@ -743,7 +748,7 @@ void RadixHTLocalSourceState::Finalize(RadixHTGlobalSinkState &sink, RadixHTGlob
 	partition.data->Combine(*ht->GetPartitionedData()->GetPartitions()[0]);
 
 	// Update thread-global state
-	lock_guard<mutex> global_guard(gstate.lock);
+	lock_guard<mutex> global_guard(sink.lock);
 	sink.stored_allocators.emplace_back(ht->GetAggregateAllocator());
 	const auto finalizes_done = ++sink.finalize_done;
 	D_ASSERT(finalizes_done <= sink.partitions.size());
@@ -783,7 +788,7 @@ void RadixHTLocalSourceState::Scan(RadixHTGlobalSinkState &sink, RadixHTGlobalSo
 			data_collection.Reset();
 		}
 		scan_status = RadixHTScanStatus::DONE;
-		lock_guard<mutex> gstate_guard(gstate.lock);
+		lock_guard<mutex> gstate_guard(sink.lock);
 		if (++gstate.task_done == sink.partitions.size()) {
 			gstate.finished = true;
 		}
@@ -910,10 +915,10 @@ double RadixPartitionedHashTable::GetProgress(ClientContext &, GlobalSinkState &
 	}
 
 	// Get scan progress, weigh it 1x
-	total_progress += 1.0 * gstate.task_done;
+	total_progress += 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 * sink.partitions.size();
+	total_progress /= 3.0 * double(sink.partitions.size());
 
 	// Multiply by 100 to get a percentage
 	return 100.0 * total_progress;

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

@@ -4,48 +4,55 @@
 
 namespace duckdb {
 
-void BlockingSample::Serialize(Serializer &serializer) const {
+void ReservoirChunk::Serialize(Serializer &serializer) const {
+	chunk.Serialize(serializer);
 }
 
-unique_ptr<BlockingSample> BlockingSample::Deserialize(Deserializer &deserializer) {
-	return nullptr;
+unique_ptr<ReservoirChunk> ReservoirChunk::Deserialize(Deserializer &deserializer) {
+	auto result = make_uniq<ReservoirChunk>();
+	result->chunk.Deserialize(deserializer);
+	return result;
 }
 
 ReservoirSample::ReservoirSample(Allocator &allocator, idx_t sample_count, int64_t seed)
     : BlockingSample(seed), allocator(allocator), sample_count(sample_count), reservoir_initialized(false) {
 }
 
+ReservoirSample::ReservoirSample(idx_t sample_count, int64_t seed)
+    : ReservoirSample(Allocator::DefaultAllocator(), sample_count, seed) {
+}
+
 void ReservoirSample::AddToReservoir(DataChunk &input) {
 	if (sample_count == 0) {
 		// sample count is 0, means no samples were requested
 		return;
 	}
-	base_reservoir_sample.num_entries_seen_total += input.size();
+	old_base_reservoir_sample.num_entries_seen_total += input.size();
 	// Input: A population V of n weighted items
 	// Output: A reservoir R with a size m
 	// 1: The first m items of V are inserted into R
 	// first we need to check if the reservoir already has "m" elements
-	if (!reservoir_chunk || reservoir_chunk->size() < sample_count) {
+	if (!reservoir_data_chunk || reservoir_data_chunk->size() < sample_count) {
 		if (FillReservoir(input) == 0) {
 			// entire chunk was consumed by reservoir
 			return;
 		}
 	}
-	D_ASSERT(reservoir_chunk);
-	D_ASSERT(reservoir_chunk->size() == sample_count);
+	D_ASSERT(reservoir_data_chunk);
+	D_ASSERT(reservoir_data_chunk->size() == sample_count);
 	// Initialize the weights if they have not been already
-	if (base_reservoir_sample.reservoir_weights.empty()) {
-		base_reservoir_sample.InitializeReservoir(reservoir_chunk->size(), sample_count);
+	if (old_base_reservoir_sample.reservoir_weights.empty()) {
+		old_base_reservoir_sample.InitializeReservoir(reservoir_data_chunk->size(), sample_count);
 	}
 	// find the position of next_index_to_sample relative to number of seen entries (num_entries_to_skip_b4_next_sample)
 	idx_t remaining = input.size();
 	idx_t base_offset = 0;
 	while (true) {
-		idx_t offset =
-		    base_reservoir_sample.next_index_to_sample - base_reservoir_sample.num_entries_to_skip_b4_next_sample;
+		idx_t offset = old_base_reservoir_sample.next_index_to_sample -
+		               old_base_reservoir_sample.num_entries_to_skip_b4_next_sample;
 		if (offset >= remaining) {
 			// not in this chunk! increment current count and go to the next chunk
-			base_reservoir_sample.num_entries_to_skip_b4_next_sample += remaining;
+			old_base_reservoir_sample.num_entries_to_skip_b4_next_sample += remaining;
 			return;
 		}
 		// in this chunk! replace the element
@@ -57,47 +64,47 @@ void ReservoirSample::AddToReservoir(DataChunk &input) {
 }
 
 unique_ptr<DataChunk> ReservoirSample::GetChunk() {
-	if (!reservoir_chunk || reservoir_chunk->size() == 0) {
+	if (!reservoir_data_chunk || reservoir_data_chunk->size() == 0) {
 		return nullptr;
 	}
-	auto collected_sample_count = reservoir_chunk->size();
+	auto collected_sample_count = reservoir_data_chunk->size();
 	if (collected_sample_count > STANDARD_VECTOR_SIZE) {
 		// get from the back to avoid creating two selection vectors
 		// one to return the first STANDARD_VECTOR_SIZE
-		// another to replace the reservoir_chunk with the first STANDARD VECTOR SIZE missing
+		// another to replace the reservoir_data_chunk with the first STANDARD VECTOR SIZE missing
 		auto ret = make_uniq<DataChunk>();
 		auto samples_remaining = collected_sample_count - STANDARD_VECTOR_SIZE;
-		auto reservoir_types = reservoir_chunk->GetTypes();
+		auto reservoir_types = reservoir_data_chunk->GetTypes();
 		SelectionVector sel(STANDARD_VECTOR_SIZE);
 		for (idx_t i = samples_remaining; i < collected_sample_count; i++) {
 			sel.set_index(i - samples_remaining, i);
 		}
 		ret->Initialize(allocator, reservoir_types.begin(), reservoir_types.end(), STANDARD_VECTOR_SIZE);
-		ret->Slice(*reservoir_chunk, sel, STANDARD_VECTOR_SIZE);
+		ret->Slice(*reservoir_data_chunk, sel, STANDARD_VECTOR_SIZE);
 		ret->SetCardinality(STANDARD_VECTOR_SIZE);
 		// reduce capacity and cardinality of the sample data chunk
-		reservoir_chunk->SetCardinality(samples_remaining);
+		reservoir_data_chunk->SetCardinality(samples_remaining);
 		return ret;
 	}
-	return std::move(reservoir_chunk);
+	return std::move(reservoir_data_chunk);
 }
 
 void ReservoirSample::ReplaceElement(DataChunk &input, idx_t index_in_chunk, double with_weight) {
 	// replace the entry in the reservoir
 	// 8. The item in R with the minimum key is replaced by item vi
-	D_ASSERT(input.ColumnCount() == reservoir_chunk->ColumnCount());
+	D_ASSERT(input.ColumnCount() == reservoir_data_chunk->ColumnCount());
 	for (idx_t col_idx = 0; col_idx < input.ColumnCount(); col_idx++) {
-		reservoir_chunk->SetValue(col_idx, base_reservoir_sample.min_weighted_entry_index,
-		                          input.GetValue(col_idx, index_in_chunk));
+		reservoir_data_chunk->SetValue(col_idx, old_base_reservoir_sample.min_weighted_entry_index,
+		                               input.GetValue(col_idx, index_in_chunk));
 	}
-	base_reservoir_sample.ReplaceElement(with_weight);
+	old_base_reservoir_sample.ReplaceElement(with_weight);
 }
 
 void ReservoirSample::InitializeReservoir(DataChunk &input) {
-	reservoir_chunk = make_uniq<DataChunk>();
-	reservoir_chunk->Initialize(allocator, input.GetTypes(), sample_count);
-	for (idx_t col_idx = 0; col_idx < reservoir_chunk->ColumnCount(); col_idx++) {
-		FlatVector::Validity(reservoir_chunk->data[col_idx]).Initialize(sample_count);
+	reservoir_data_chunk = make_uniq<DataChunk>();
+	reservoir_data_chunk->Initialize(allocator, input.GetTypes(), sample_count);
+	for (idx_t col_idx = 0; col_idx < reservoir_data_chunk->ColumnCount(); col_idx++) {
+		FlatVector::Validity(reservoir_data_chunk->data[col_idx]).Initialize(sample_count);
 	}
 	reservoir_initialized = true;
 }
@@ -105,7 +112,7 @@ void ReservoirSample::InitializeReservoir(DataChunk &input) {
 idx_t ReservoirSample::FillReservoir(DataChunk &input) {
 	idx_t chunk_count = input.size();
 	input.Flatten();
-	auto num_added_samples = reservoir_chunk ? reservoir_chunk->size() : 0;
+	auto num_added_samples = reservoir_data_chunk ? reservoir_data_chunk->size() : 0;
 	D_ASSERT(num_added_samples <= sample_count);
 
 	// required count is what we still need to add to the reservoir
@@ -123,8 +130,8 @@ idx_t ReservoirSample::FillReservoir(DataChunk &input) {
 	if (!reservoir_initialized) {
 		InitializeReservoir(input);
 	}
-	reservoir_chunk->Append(input, false, nullptr, required_count);
-	base_reservoir_sample.InitializeReservoir(required_count, sample_count);
+	reservoir_data_chunk->Append(input, false, nullptr, required_count);
+	old_base_reservoir_sample.InitializeReservoir(required_count, sample_count);
 
 	// check if there are still elements remaining in the Input data chunk that should be
 	// randomly sampled and potentially added. This happens if we are on a boundary
@@ -155,8 +162,12 @@ ReservoirSamplePercentage::ReservoirSamplePercentage(Allocator &allocator, doubl
 	current_sample = make_uniq<ReservoirSample>(allocator, reservoir_sample_size, random.NextRandomInteger());
 }
 
+ReservoirSamplePercentage::ReservoirSamplePercentage(double percentage, int64_t seed)
+    : ReservoirSamplePercentage(Allocator::DefaultAllocator(), percentage, seed) {
+}
+
 void ReservoirSamplePercentage::AddToReservoir(DataChunk &input) {
-	base_reservoir_sample.num_entries_seen_total += input.size();
+	old_base_reservoir_sample.num_entries_seen_total += input.size();
 	if (current_count + input.size() > RESERVOIR_THRESHOLD) {
 		// we don't have enough space in our current reservoir
 		// first check what we still need to append to the current sample