duckdb 1.3.4 → 1.4.0

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

@@ -1,4 +1,3 @@
-#include "duckdb/catalog/catalog_entry/duck_table_entry.hpp"
 #include "duckdb/execution/operator/join/perfect_hash_join_executor.hpp"
 #include "duckdb/execution/operator/join/physical_blockwise_nl_join.hpp"
 #include "duckdb/execution/operator/join/physical_cross_product.hpp"
@@ -6,23 +5,18 @@
 #include "duckdb/execution/operator/join/physical_iejoin.hpp"
 #include "duckdb/execution/operator/join/physical_nested_loop_join.hpp"
 #include "duckdb/execution/operator/join/physical_piecewise_merge_join.hpp"
-#include "duckdb/execution/operator/scan/physical_table_scan.hpp"
 #include "duckdb/execution/physical_plan_generator.hpp"
-#include "duckdb/function/table/table_scan.hpp"
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
 #include "duckdb/planner/expression_iterator.hpp"
 #include "duckdb/planner/operator/logical_comparison_join.hpp"
-#include "duckdb/transaction/duck_transaction.hpp"
+#include "duckdb/main/settings.hpp"
 
 namespace duckdb {
 
-static void RewriteJoinCondition(Expression &expr, idx_t offset) {
-	if (expr.GetExpressionType() == ExpressionType::BOUND_REF) {
-		auto &ref = expr.Cast<BoundReferenceExpression>();
-		ref.index += offset;
-	}
-	ExpressionIterator::EnumerateChildren(expr, [&](Expression &child) { RewriteJoinCondition(child, offset); });
+static void RewriteJoinCondition(unique_ptr<Expression> &root_expr, idx_t offset) {
+	ExpressionIterator::VisitExpressionMutable<BoundReferenceExpression>(
+	    root_expr, [&](BoundReferenceExpression &ref, unique_ptr<Expression> &expr) { ref.index += offset; });
 }
 
 PhysicalOperator &PhysicalPlanGenerator::PlanComparisonJoin(LogicalComparisonJoin &op) {
@@ -56,10 +50,9 @@ PhysicalOperator &PhysicalPlanGenerator::PlanComparisonJoin(LogicalComparisonJoi
 	default:
 		break;
 	}
-	auto &client_config = ClientConfig::GetConfig(context);
-
 	//	TODO: Extend PWMJ to handle all comparisons and projection maps
-	const auto prefer_range_joins = client_config.prefer_range_joins && can_iejoin;
+	bool prefer_range_joins = DBConfig::GetSetting<PreferRangeJoinsSetting>(context);
+	prefer_range_joins = prefer_range_joins && can_iejoin;
 	if (has_equality && !prefer_range_joins) {
 		// Equality join with small number of keys : possible perfect join optimization
 		auto &join = Make<PhysicalHashJoin>(op, left, right, std::move(op.conditions), op.join_type,
@@ -70,15 +63,16 @@ PhysicalOperator &PhysicalPlanGenerator::PlanComparisonJoin(LogicalComparisonJoi
 	}
 
 	D_ASSERT(op.left_projection_map.empty());
-	if (left.estimated_cardinality <= client_config.nested_loop_join_threshold ||
-	    right.estimated_cardinality <= client_config.nested_loop_join_threshold) {
+	idx_t nested_loop_join_threshold = DBConfig::GetSetting<NestedLoopJoinThresholdSetting>(context);
+	if (left.estimated_cardinality < nested_loop_join_threshold ||
+	    right.estimated_cardinality < nested_loop_join_threshold) {
 		can_iejoin = false;
 		can_merge = false;
 	}
 
 	if (can_merge && can_iejoin) {
-		if (left.estimated_cardinality <= client_config.merge_join_threshold ||
-		    right.estimated_cardinality <= client_config.merge_join_threshold) {
+		idx_t merge_join_threshold = DBConfig::GetSetting<MergeJoinThresholdSetting>(context);
+		if (left.estimated_cardinality < merge_join_threshold || right.estimated_cardinality < merge_join_threshold) {
 			can_iejoin = false;
 		}
 	}
@@ -99,7 +93,7 @@ PhysicalOperator &PhysicalPlanGenerator::PlanComparisonJoin(LogicalComparisonJoi
 	}
 
 	for (auto &cond : op.conditions) {
-		RewriteJoinCondition(*cond.right, left.types.size());
+		RewriteJoinCondition(cond.right, left.types.size());
 	}
 	auto condition = JoinCondition::CreateExpression(std::move(op.conditions));
 	return Make<PhysicalBlockwiseNLJoin>(op, left, right, std::move(condition), op.join_type, op.estimated_cardinality);

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

@@ -77,6 +77,8 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalGet &op) {
 		    Make<PhysicalTableInOutFunction>(op.types, op.function, std::move(op.bind_data), column_ids,
 		                                     op.estimated_cardinality, std::move(op.projected_input));
 		table_in_out.children.push_back(child);
+		auto &cast_table_in_out = table_in_out.Cast<PhysicalTableInOutFunction>();
+		cast_table_in_out.ordinality_idx = op.ordinality_idx;
 		return table_in_out;
 	}
 
@@ -103,7 +105,7 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalGet &op) {
 		for (auto &entry : table_filters->filters) {
 			auto column_id = column_ids[entry.first].GetPrimaryIndex();
 			auto &type = op.returned_types[column_id];
-			if (!op.function.supports_pushdown_type(type)) {
+			if (!op.function.supports_pushdown_type(*op.bind_data, column_id)) {
 				idx_t column_id_filter = entry.first;
 				bool found_projection = false;
 				for (idx_t i = 0; i < projection_ids.size(); i++) {

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

@@ -8,6 +8,7 @@
 #include "duckdb/parallel/task_scheduler.hpp"
 #include "duckdb/catalog/duck_catalog.hpp"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
+#include "duckdb/main/settings.hpp"
 
 namespace duckdb {
 
@@ -20,6 +21,7 @@ OrderPreservationType PhysicalPlanGenerator::OrderPreservationRecursive(Physical
 	for (auto &child : op.children) {
 		// Do not take the materialization phase of physical CTEs into account
 		if (op.type == PhysicalOperatorType::CTE && child_idx == 0) {
+			child_idx++;
 			continue;
 		}
 		auto child_preservation = OrderPreservationRecursive(child);
@@ -32,8 +34,6 @@ OrderPreservationType PhysicalPlanGenerator::OrderPreservationRecursive(Physical
 }
 
 bool PhysicalPlanGenerator::PreserveInsertionOrder(ClientContext &context, PhysicalOperator &plan) {
-	auto &config = DBConfig::GetConfig(context);
-
 	auto preservation_type = OrderPreservationRecursive(plan);
 	if (preservation_type == OrderPreservationType::FIXED_ORDER) {
 		// always need to maintain preservation order
@@ -44,7 +44,7 @@ bool PhysicalPlanGenerator::PreserveInsertionOrder(ClientContext &context, Physi
 		return false;
 	}
 	// preserve insertion order - check flags
-	if (!config.options.preserve_insertion_order) {
+	if (!DBConfig::GetSetting<PreserveInsertionOrderSetting>(context)) {
 		// preserving insertion order is disabled by config
 		return false;
 	}
@@ -107,11 +107,11 @@ PhysicalOperator &DuckCatalog::PlanInsert(ClientContext &context, PhysicalPlanGe
 		parallel_streaming_insert = false;
 		use_batch_index = false;
 	}
-	if (op.action_type != OnConflictAction::THROW) {
+	if (op.on_conflict_info.action_type != OnConflictAction::THROW) {
 		// We don't support ON CONFLICT clause in batch insertion operation currently
 		use_batch_index = false;
 	}
-	if (op.action_type == OnConflictAction::UPDATE) {
+	if (op.on_conflict_info.action_type == OnConflictAction::UPDATE) {
 		// When we potentially need to perform updates, we have to check that row is not updated twice
 		// that currently needs to be done for every chunk, which would add a huge bottleneck to parallelized insertion
 		parallel_streaming_insert = false;
@@ -127,11 +127,12 @@ PhysicalOperator &DuckCatalog::PlanInsert(ClientContext &context, PhysicalPlanGe
 	}
 
 	auto &insert = planner.Make<PhysicalInsert>(
-	    op.types, op.table, 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),
-	    op.update_is_del_and_insert);
+	    op.types, op.table, std::move(op.bound_constraints), std::move(op.expressions),
+	    std::move(op.on_conflict_info.set_columns), std::move(op.on_conflict_info.set_types), op.estimated_cardinality,
+	    op.return_chunk, parallel_streaming_insert && num_threads > 1, op.on_conflict_info.action_type,
+	    std::move(op.on_conflict_info.on_conflict_condition), std::move(op.on_conflict_info.do_update_condition),
+	    std::move(op.on_conflict_info.on_conflict_filter), std::move(op.on_conflict_info.columns_to_fetch),
+	    op.on_conflict_info.update_is_del_and_insert);
 	insert.children.push_back(*plan);
 	return insert;
 }

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

@@ -0,0 +1,132 @@
+#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp"
+#include "duckdb/execution/operator/persistent/physical_merge_into.hpp"
+#include "duckdb/execution/operator/persistent/physical_delete.hpp"
+#include "duckdb/execution/operator/persistent/physical_insert.hpp"
+#include "duckdb/execution/operator/persistent/physical_update.hpp"
+#include "duckdb/execution/physical_plan_generator.hpp"
+#include "duckdb/planner/operator/logical_merge_into.hpp"
+#include "duckdb/catalog/duck_catalog.hpp"
+
+namespace duckdb {
+
+unique_ptr<MergeIntoOperator> PlanMergeIntoAction(ClientContext &context, LogicalMergeInto &op,
+                                                  PhysicalPlanGenerator &planner, BoundMergeIntoAction &action) {
+	auto result = make_uniq<MergeIntoOperator>();
+
+	result->action_type = action.action_type;
+	result->condition = std::move(action.condition);
+	vector<unique_ptr<BoundConstraint>> bound_constraints;
+	for (auto &constraint : op.bound_constraints) {
+		bound_constraints.push_back(constraint->Copy());
+	}
+	auto return_types = op.types;
+	if (op.return_chunk) {
+		// for RETURNING, the last column is the merge_action - this is added in the merge itself
+		return_types.pop_back();
+	}
+
+	auto cardinality = op.EstimateCardinality(context);
+	switch (action.action_type) {
+	case MergeActionType::MERGE_UPDATE: {
+		vector<unique_ptr<Expression>> defaults;
+		for (auto &def : op.bound_defaults) {
+			defaults.push_back(def->Copy());
+		}
+		result->op =
+		    planner.Make<PhysicalUpdate>(std::move(return_types), op.table, op.table.GetStorage(),
+		                                 std::move(action.columns), std::move(action.expressions), std::move(defaults),
+		                                 std::move(bound_constraints), cardinality, op.return_chunk);
+		auto &cast_update = result->op->Cast<PhysicalUpdate>();
+		cast_update.update_is_del_and_insert = action.update_is_del_and_insert;
+		break;
+	}
+	case MergeActionType::MERGE_DELETE: {
+		result->op =
+		    planner.Make<PhysicalDelete>(std::move(return_types), op.table, op.table.GetStorage(),
+		                                 std::move(bound_constraints), op.row_id_start, cardinality, op.return_chunk);
+		break;
+	}
+	case MergeActionType::MERGE_INSERT: {
+		vector<unique_ptr<Expression>> set_expressions;
+		vector<PhysicalIndex> set_columns;
+		vector<LogicalType> set_types;
+		unordered_set<column_t> on_conflict_filter;
+		vector<column_t> columns_to_fetch;
+
+		result->op = planner.Make<PhysicalInsert>(std::move(return_types), op.table, std::move(bound_constraints),
+		                                          std::move(set_expressions), std::move(set_columns),
+		                                          std::move(set_types), cardinality, op.return_chunk, !op.return_chunk,
+		                                          OnConflictAction::THROW, nullptr, nullptr,
+		                                          std::move(on_conflict_filter), std::move(columns_to_fetch), false);
+		// transform expressions if required
+		if (!action.column_index_map.empty()) {
+			vector<unique_ptr<Expression>> new_expressions;
+			for (auto &col : op.table.GetColumns().Physical()) {
+				auto storage_idx = col.StorageOid();
+				auto mapped_index = action.column_index_map[col.Physical()];
+				if (mapped_index == DConstants::INVALID_INDEX) {
+					// push default value
+					new_expressions.push_back(op.bound_defaults[storage_idx]->Copy());
+				} else {
+					// push reference
+					new_expressions.push_back(std::move(action.expressions[mapped_index]));
+				}
+			}
+			action.expressions = std::move(new_expressions);
+		}
+		result->expressions = std::move(action.expressions);
+		break;
+	}
+	case MergeActionType::MERGE_ERROR:
+		result->expressions = std::move(action.expressions);
+		break;
+	case MergeActionType::MERGE_DO_NOTHING:
+		break;
+	default:
+		throw InternalException("Unsupported merge action");
+	}
+	return result;
+}
+
+PhysicalOperator &DuckCatalog::PlanMergeInto(ClientContext &context, PhysicalPlanGenerator &planner,
+                                             LogicalMergeInto &op, PhysicalOperator &plan) {
+	map<MergeActionCondition, vector<unique_ptr<MergeIntoOperator>>> actions;
+
+	// plan the merge into clauses
+	// FIXME: disable parallelism when we have multiple INSERTs because they do not work nicely together currently
+	idx_t append_count = 0;
+	for (auto &entry : op.actions) {
+		vector<unique_ptr<MergeIntoOperator>> planned_actions;
+		for (auto &action : entry.second) {
+			if (action->action_type == MergeActionType::MERGE_INSERT) {
+				append_count++;
+			}
+			if (action->action_type == MergeActionType::MERGE_UPDATE && action->update_is_del_and_insert) {
+				append_count++;
+			}
+			planned_actions.push_back(PlanMergeIntoAction(context, op, planner, *action));
+		}
+		actions.emplace(entry.first, std::move(planned_actions));
+	}
+
+	bool parallel = append_count <= 1 && !op.return_chunk;
+
+	auto &result = planner.Make<PhysicalMergeInto>(op.types, std::move(actions), op.row_id_start, op.source_marker,
+	                                               parallel, op.return_chunk);
+	result.children.push_back(plan);
+	return result;
+}
+
+PhysicalOperator &Catalog::PlanMergeInto(ClientContext &context, PhysicalPlanGenerator &planner, LogicalMergeInto &op,
+                                         PhysicalOperator &plan) {
+	throw NotImplementedException("Database type \"%s\" does not support MERGE INTO or ON CONFLICT", GetCatalogType());
+}
+
+PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalMergeInto &op) {
+	auto &plan = CreatePlan(*op.children[0]);
+	D_ASSERT(op.children.size() == 1);
+	dependencies.AddDependency(op.table);
+	return op.table.catalog.PlanMergeInto(context, *this, op, plan);
+}
+
+} // namespace duckdb

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

@@ -104,25 +104,23 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalCTERef &op) {
 	D_ASSERT(op.children.empty());
 
 	// Check if this LogicalCTERef is supposed to scan a materialized CTE.
-	if (op.materialized_cte == CTEMaterialize::CTE_MATERIALIZE_ALWAYS) {
-		// Lookup if there is a materialized CTE for the cte_index.
-		auto materialized_cte = materialized_ctes.find(op.cte_index);
-
-		// If this check fails, this is a reference to a materialized recursive CTE.
-		if (materialized_cte != materialized_ctes.end()) {
-			auto &chunk_scan = Make<PhysicalColumnDataScan>(op.chunk_types, PhysicalOperatorType::CTE_SCAN,
-			                                                op.estimated_cardinality, op.cte_index);
-
-			auto cte = recursive_cte_tables.find(op.cte_index);
-			if (cte == recursive_cte_tables.end()) {
-				throw InvalidInputException("Referenced materialized CTE does not exist.");
-			}
-
-			auto &cast_chunk_scan = chunk_scan.Cast<PhysicalColumnDataScan>();
-			cast_chunk_scan.collection = cte->second.get();
-			materialized_cte->second.push_back(cast_chunk_scan);
-			return chunk_scan;
+	// Lookup if there is a materialized CTE for the cte_index.
+	auto materialized_cte = materialized_ctes.find(op.cte_index);
+
+	// If this check fails, this is a reference to a materialized recursive CTE.
+	if (materialized_cte != materialized_ctes.end()) {
+		auto &chunk_scan = Make<PhysicalColumnDataScan>(op.chunk_types, PhysicalOperatorType::CTE_SCAN,
+		                                                op.estimated_cardinality, op.cte_index);
+
+		auto cte = recursive_cte_tables.find(op.cte_index);
+		if (cte == recursive_cte_tables.end()) {
+			throw InvalidInputException("Referenced materialized CTE does not exist.");
 		}
+
+		auto &cast_chunk_scan = chunk_scan.Cast<PhysicalColumnDataScan>();
+		cast_chunk_scan.collection = cte->second.get();
+		materialized_cte->second.push_back(cast_chunk_scan);
+		return chunk_scan;
 	}
 
 	// CreatePlan of a LogicalRecursiveCTE must have happened before.

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

@@ -10,6 +10,7 @@
 #include "duckdb/planner/operator/logical_extension_operator.hpp"
 #include "duckdb/planner/operator/list.hpp"
 #include "duckdb/execution/operator/helper/physical_verify_vector.hpp"
+#include "duckdb/main/settings.hpp"
 
 namespace duckdb {
 
@@ -56,11 +57,11 @@ unique_ptr<PhysicalPlan> PhysicalPlanGenerator::PlanInternal(LogicalOperator &op
 	physical_plan->SetRoot(CreatePlan(op));
 	physical_plan->Root().estimated_cardinality = op.estimated_cardinality;
 
-	auto &config = DBConfig::GetConfig(context);
-	if (config.options.debug_verify_vector != DebugVectorVerification::NONE) {
-		if (config.options.debug_verify_vector != DebugVectorVerification::DICTIONARY_EXPRESSION) {
-			physical_plan->SetRoot(
-			    Make<PhysicalVerifyVector>(physical_plan->Root(), config.options.debug_verify_vector));
+	auto debug_verify_vector = DBConfig::GetSetting<DebugVerifyVectorSetting>(context);
+	if (debug_verify_vector != DebugVectorVerification::NONE) {
+		if (debug_verify_vector != DebugVectorVerification::DICTIONARY_EXPRESSION &&
+		    debug_verify_vector != DebugVectorVerification::VARIANT_VECTOR) {
+			physical_plan->SetRoot(Make<PhysicalVerifyVector>(physical_plan->Root(), debug_verify_vector));
 		}
 	}
 	return std::move(physical_plan);
@@ -120,6 +121,8 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalOperator &op) {
 		return CreatePlan(op.Cast<LogicalExpressionGet>());
 	case LogicalOperatorType::LOGICAL_UPDATE:
 		return CreatePlan(op.Cast<LogicalUpdate>());
+	case LogicalOperatorType::LOGICAL_MERGE_INTO:
+		return CreatePlan(op.Cast<LogicalMergeInto>());
 	case LogicalOperatorType::LOGICAL_CREATE_TABLE:
 		return CreatePlan(op.Cast<LogicalCreateTable>());
 	case LogicalOperatorType::LOGICAL_CREATE_INDEX:

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

@@ -13,8 +13,9 @@
 
 namespace duckdb {
 
-RadixPartitionedHashTable::RadixPartitionedHashTable(GroupingSet &grouping_set_p, const GroupedAggregateData &op_p)
-    : grouping_set(grouping_set_p), op(op_p) {
+RadixPartitionedHashTable::RadixPartitionedHashTable(GroupingSet &grouping_set_p, const GroupedAggregateData &op_p,
+                                                     TupleDataValidityType group_validity_p)
+    : grouping_set(grouping_set_p), op(op_p), group_validity(group_validity_p) {
 	auto groups_count = op.GroupCount();
 	for (idx_t i = 0; i < groups_count; i++) {
 		if (grouping_set.find(i) == grouping_set.end()) {
@@ -35,7 +36,8 @@ RadixPartitionedHashTable::RadixPartitionedHashTable(GroupingSet &grouping_set_p
 	group_types_copy.emplace_back(LogicalType::HASH);
 
 	auto layout = make_shared_ptr<TupleDataLayout>();
-	layout->Initialize(std::move(group_types_copy), AggregateObject::CreateAggregateObjects(op.bindings));
+	auto aggregate_objects = AggregateObject::CreateAggregateObjects(op.bindings);
+	layout->Initialize(std::move(group_types_copy), std::move(aggregate_objects), group_validity);
 	layout_ptr = std::move(layout);
 }
 
@@ -69,7 +71,7 @@ const TupleDataLayout &RadixPartitionedHashTable::GetLayout() const {
 unique_ptr<GroupedAggregateHashTable> RadixPartitionedHashTable::CreateHT(ClientContext &context, const idx_t capacity,
                                                                           const idx_t radix_bits) const {
 	return make_uniq<GroupedAggregateHashTable>(context, BufferAllocator::Get(context), group_types, op.payload_types,
-	                                            op.bindings, capacity, radix_bits);
+	                                            op.bindings, capacity, radix_bits, group_validity);
 }
 
 //===--------------------------------------------------------------------===//
@@ -126,13 +128,6 @@ public:
 	const idx_t sink_capacity;
 
 private:
-	//! Assume (1 << 15) = 32KB L1 cache per core, divided by two because hyperthreading
-	static constexpr idx_t L1_CACHE_SIZE = 32768 / 2;
-	//! Assume (1 << 20) = 1MB L2 cache per core, divided by two because hyperthreading
-	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 idx_t L3_CACHE_SIZE = 1572864 / 2;
-
 	//! Sink radix bits to initialize with
 	static constexpr idx_t MAXIMUM_INITIAL_SINK_RADIX_BITS = 4;
 	//! Maximum Sink radix bits (independent of threads)
@@ -152,7 +147,7 @@ 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 double BLOCK_FILL_FACTOR = 1.8;
+	static constexpr double BLOCK_FILL_FACTOR = 0.5;
 	//! By how many bits to repartition if a repartition is triggered
 	static constexpr idx_t REPARTITION_RADIX_BITS = 2;
 };
@@ -337,17 +332,12 @@ idx_t RadixHTConfig::MaximumSinkRadixBits() const {
 }
 
 idx_t RadixHTConfig::SinkCapacity() const {
-	// Compute cache size per active thread (assuming cache is shared)
-	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 = 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());
+	if (number_of_threads <= GROW_STRATEGY_THREAD_THRESHOLD) {
+		// Grow strategy, start off a bit bigger
+		return 262144;
+	}
+	// Start off tiny, we can adapt with DecideAdaptation later
+	return 32768;
 }
 
 class RadixHTLocalSinkState : public LocalSinkState {
@@ -693,7 +683,7 @@ public:
 	vector<column_t> column_ids;
 
 	//! For synchronizing tasks
-	idx_t task_idx;
+	atomic<idx_t> task_idx;
 	atomic<idx_t> task_done;
 };
 
@@ -753,12 +743,11 @@ RadixHTGlobalSourceState::RadixHTGlobalSourceState(ClientContext &context_p, con
 SourceResultType RadixHTGlobalSourceState::AssignTask(RadixHTGlobalSinkState &sink, RadixHTLocalSourceState &lstate,
                                                       InterruptState &interrupt_state) {
 	// First, try to get a partition index
-	auto guard = sink.Lock();
-	if (finished || task_idx == sink.partitions.size()) {
+	lstate.task_idx = task_idx++;
+	if (finished || lstate.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];

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

@@ -3,24 +3,6 @@
 
 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;
@@ -73,7 +55,7 @@ void BaseReservoirSampling::SetNextEntry() {
 	//! 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)));
+	next_index_to_sample = idx_t(ceil(x_w));
 	num_entries_to_skip_b4_next_sample = 0;
 }
 
@@ -118,15 +100,33 @@ void BaseReservoirSampling::UpdateMinWeightThreshold() {
 	min_weight_threshold = 1;
 }
 
+// Generate top k order statistics from n Uniform(0, 1) samples in O(k) time
+// This method leverages two key properties of order statistics from a Uniform(0, 1) distribution:
+// 1. The maximum of n independent Uniform(0, 1) samples, denoted as U(n), follows a Beta(n, 1) distribution.
+// 2. U(n-i) / U(n-i+1) ~ Beta(n-i, 1)
+// So we can use a recursive approach to generate the top k order statistics
+// (See: https://www.math.ntu.edu.tw/~hchen/teaching/LargeSample/notes/noteorder.pdf)
+static vector<double> GenerateTopKFromUniform(ReservoirRNG &random, idx_t n, idx_t k) {
+	vector<double> top_k_values(k);
+	double current_bound = 1.0;
+	for (idx_t i = 0; i < k; i++) {
+		// generate a sample from Beta(n - i, 1)
+		double beta = std::pow(random.NextRandom(), 1.0 / double(n - i));
+		current_bound *= beta;
+		top_k_values[i] = current_bound;
+	}
+	return top_k_values;
+}
+
 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);
+	auto weights = GenerateTopKFromUniform(random, num_entries_seen_total, sel_size);
+	std::shuffle(weights.begin(), weights.end(), random);
 	for (idx_t i = 0; i < sel_size; i++) {
-		auto weight = random.NextRandom(min_weight, 1);
+		auto weight = weights[i];
 		reservoir_weights.emplace(-weight, i);
 	}
 	D_ASSERT(reservoir_weights.size() <= sel_size);

package/src/duckdb/src/function/aggregate/distributive/count.cpp

@@ -6,6 +6,7 @@
 
 namespace duckdb {
 
+namespace {
 struct BaseCountFunction {
 	template <class STATE>
 	static void Initialize(STATE &state) {
@@ -211,6 +212,19 @@ struct CountFunction : public BaseCountFunction {
 	}
 };
 
+unique_ptr<BaseStatistics> CountPropagateStats(ClientContext &context, BoundAggregateExpression &expr,
+                                               AggregateStatisticsInput &input) {
+	if (!expr.IsDistinct() && !input.child_stats[0].CanHaveNull()) {
+		// count on a column without null values: use count star
+		expr.function = CountStarFun::GetFunction();
+		expr.function.name = "count_star";
+		expr.children.clear();
+	}
+	return nullptr;
+}
+
+} // namespace
+
 AggregateFunction CountFunctionBase::GetFunction() {
 	AggregateFunction fun({LogicalType(LogicalTypeId::ANY)}, LogicalType::BIGINT, AggregateFunction::StateSize<int64_t>,
 	                      AggregateFunction::StateInitialize<int64_t, CountFunction>, CountFunction::CountScatter,
@@ -231,17 +245,6 @@ AggregateFunction CountStarFun::GetFunction() {
 	return fun;
 }
 
-unique_ptr<BaseStatistics> CountPropagateStats(ClientContext &context, BoundAggregateExpression &expr,
-                                               AggregateStatisticsInput &input) {
-	if (!expr.IsDistinct() && !input.child_stats[0].CanHaveNull()) {
-		// count on a column without null values: use count star
-		expr.function = CountStarFun::GetFunction();
-		expr.function.name = "count_star";
-		expr.children.clear();
-	}
-	return nullptr;
-}
-
 AggregateFunctionSet CountFun::GetFunctions() {
 	AggregateFunction count_function = CountFunctionBase::GetFunction();
 	count_function.statistics = CountPropagateStats;