duckdb 0.7.2-dev314.0 → 0.7.2-dev3154.0

package/src/duckdb/src/optimizer/statistics/operator/propagate_join.cpp

@@ -40,7 +40,7 @@ void StatisticsPropagator::PropagateStatistics(LogicalComparisonJoin &join, uniq
 				case JoinType::ANTI: {
 					// when the right child has data, return the left child
 					// when the right child has no data, return an empty set
-					auto limit = make_unique<LogicalLimit>(1, 0, nullptr, nullptr);
+					auto limit = make_uniq<LogicalLimit>(1, 0, nullptr, nullptr);
 					limit->AddChild(std::move(join.children[1]));
 					auto cross_product = LogicalCrossProduct::Create(std::move(join.children[0]), std::move(limit));
 					*node_ptr = std::move(cross_product);
@@ -76,7 +76,7 @@ void StatisticsPropagator::PropagateStatistics(LogicalComparisonJoin &join, uniq
 					case JoinType::SEMI: {
 						// when the right child has data, return the left child
 						// when the right child has no data, return an empty set
-						auto limit = make_unique<LogicalLimit>(1, 0, nullptr, nullptr);
+						auto limit = make_uniq<LogicalLimit>(1, 0, nullptr, nullptr);
 						limit->AddChild(std::move(join.children[1]));
 						auto cross_product = LogicalCrossProduct::Create(std::move(join.children[0]), std::move(limit));
 						*node_ptr = std::move(cross_product);
@@ -195,10 +195,11 @@ unique_ptr<NodeStatistics> StatisticsPropagator::PropagateStatistics(LogicalJoin
 	switch (join.type) {
 	case LogicalOperatorType::LOGICAL_COMPARISON_JOIN:
 	case LogicalOperatorType::LOGICAL_DELIM_JOIN:
-		PropagateStatistics((LogicalComparisonJoin &)join, node_ptr);
+	case LogicalOperatorType::LOGICAL_ASOF_JOIN:
+		PropagateStatistics(join.Cast<LogicalComparisonJoin>(), node_ptr);
 		break;
 	case LogicalOperatorType::LOGICAL_ANY_JOIN:
-		PropagateStatistics((LogicalAnyJoin &)join, node_ptr);
+		PropagateStatistics(join.Cast<LogicalAnyJoin>(), node_ptr);
 		break;
 	default:
 		break;

package/src/duckdb/src/optimizer/statistics/operator/propagate_limit.cpp

@@ -8,7 +8,7 @@ unique_ptr<NodeStatistics> StatisticsPropagator::PropagateStatistics(LogicalLimi
 	// propagate statistics in the child node
 	PropagateStatistics(limit.children[0]);
 	// return the node stats, with as expected cardinality the amount specified in the limit
-	return make_unique<NodeStatistics>(limit.limit_val, limit.limit_val);
+	return make_uniq<NodeStatistics>(limit.limit_val, limit.limit_val);
 }
 
 } // namespace duckdb

package/src/duckdb/src/optimizer/statistics_propagator.cpp

@@ -3,7 +3,18 @@
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/planner/expression_iterator.hpp"
 #include "duckdb/planner/logical_operator.hpp"
+#include "duckdb/planner/operator/logical_aggregate.hpp"
 #include "duckdb/planner/operator/logical_empty_result.hpp"
+#include "duckdb/planner/operator/logical_cross_product.hpp"
+#include "duckdb/planner/operator/logical_filter.hpp"
+#include "duckdb/planner/operator/logical_get.hpp"
+#include "duckdb/planner/operator/logical_join.hpp"
+#include "duckdb/planner/operator/logical_projection.hpp"
+#include "duckdb/planner/operator/logical_positional_join.hpp"
+#include "duckdb/planner/operator/logical_set_operation.hpp"
+#include "duckdb/planner/operator/logical_order.hpp"
+#include "duckdb/planner/operator/logical_window.hpp"
+#include "duckdb/planner/expression/list.hpp"
 
 namespace duckdb {
 
@@ -11,7 +22,7 @@ StatisticsPropagator::StatisticsPropagator(ClientContext &context) : context(con
 }
 
 void StatisticsPropagator::ReplaceWithEmptyResult(unique_ptr<LogicalOperator> &node) {
-	node = make_unique<LogicalEmptyResult>(std::move(node));
+	node = make_uniq<LogicalEmptyResult>(std::move(node));
 }
 
 unique_ptr<NodeStatistics> StatisticsPropagator::PropagateChildren(LogicalOperator &node,
@@ -26,30 +37,31 @@ unique_ptr<NodeStatistics> StatisticsPropagator::PropagateStatistics(LogicalOper
                                                                      unique_ptr<LogicalOperator> *node_ptr) {
 	switch (node.type) {
 	case LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY:
-		return PropagateStatistics((LogicalAggregate &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalAggregate>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_CROSS_PRODUCT:
-		return PropagateStatistics((LogicalCrossProduct &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalCrossProduct>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_FILTER:
-		return PropagateStatistics((LogicalFilter &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalFilter>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_GET:
-		return PropagateStatistics((LogicalGet &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalGet>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_PROJECTION:
-		return PropagateStatistics((LogicalProjection &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalProjection>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_ANY_JOIN:
+	case LogicalOperatorType::LOGICAL_ASOF_JOIN:
 	case LogicalOperatorType::LOGICAL_COMPARISON_JOIN:
 	case LogicalOperatorType::LOGICAL_JOIN:
 	case LogicalOperatorType::LOGICAL_DELIM_JOIN:
-		return PropagateStatistics((LogicalJoin &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalJoin>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_POSITIONAL_JOIN:
-		return PropagateStatistics((LogicalPositionalJoin &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalPositionalJoin>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_UNION:
 	case LogicalOperatorType::LOGICAL_EXCEPT:
 	case LogicalOperatorType::LOGICAL_INTERSECT:
-		return PropagateStatistics((LogicalSetOperation &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalSetOperation>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_ORDER_BY:
-		return PropagateStatistics((LogicalOrder &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalOrder>(), node_ptr);
 	case LogicalOperatorType::LOGICAL_WINDOW:
-		return PropagateStatistics((LogicalWindow &)node, node_ptr);
+		return PropagateStatistics(node.Cast<LogicalWindow>(), node_ptr);
 	default:
 		return PropagateChildren(node, node_ptr);
 	}
@@ -63,25 +75,25 @@ unique_ptr<BaseStatistics> StatisticsPropagator::PropagateExpression(Expression
                                                                      unique_ptr<Expression> *expr_ptr) {
 	switch (expr.GetExpressionClass()) {
 	case ExpressionClass::BOUND_AGGREGATE:
-		return PropagateExpression((BoundAggregateExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundAggregateExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_BETWEEN:
-		return PropagateExpression((BoundBetweenExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundBetweenExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_CASE:
-		return PropagateExpression((BoundCaseExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundCaseExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_CONJUNCTION:
-		return PropagateExpression((BoundConjunctionExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundConjunctionExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_FUNCTION:
-		return PropagateExpression((BoundFunctionExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundFunctionExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_CAST:
-		return PropagateExpression((BoundCastExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundCastExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_COMPARISON:
-		return PropagateExpression((BoundComparisonExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundComparisonExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_CONSTANT:
-		return PropagateExpression((BoundConstantExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundConstantExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_COLUMN_REF:
-		return PropagateExpression((BoundColumnRefExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundColumnRefExpression>(), expr_ptr);
 	case ExpressionClass::BOUND_OPERATOR:
-		return PropagateExpression((BoundOperatorExpression &)expr, expr_ptr);
+		return PropagateExpression(expr.Cast<BoundOperatorExpression>(), expr_ptr);
 	default:
 		break;
 	}

package/src/duckdb/src/optimizer/topn_optimizer.cpp

@@ -9,13 +9,13 @@ namespace duckdb {
 unique_ptr<LogicalOperator> TopN::Optimize(unique_ptr<LogicalOperator> op) {
 	if (op->type == LogicalOperatorType::LOGICAL_LIMIT &&
 	    op->children[0]->type == LogicalOperatorType::LOGICAL_ORDER_BY) {
-		auto &limit = (LogicalLimit &)*op;
-		auto &order_by = (LogicalOrder &)*(op->children[0]);
+		auto &limit = op->Cast<LogicalLimit>();
+		auto &order_by = (op->children[0])->Cast<LogicalOrder>();
 
 		// This optimization doesn't apply when OFFSET is present without LIMIT
 		// Or if offset is not constant
 		if (limit.limit_val != NumericLimits<int64_t>::Maximum() || limit.offset) {
-			auto topn = make_unique<LogicalTopN>(std::move(order_by.orders), limit.limit_val, limit.offset_val);
+			auto topn = make_uniq<LogicalTopN>(std::move(order_by.orders), limit.limit_val, limit.offset_val);
 			topn->AddChild(std::move(order_by.children[0]));
 			op = std::move(topn);
 		}

package/src/duckdb/src/optimizer/unnest_rewriter.cpp

@@ -22,7 +22,7 @@ void UnnestRewriterPlanUpdater::VisitExpression(unique_ptr<Expression> *expressi
 
 	if (expr->expression_class == ExpressionClass::BOUND_COLUMN_REF) {
 
-		auto &bound_column_ref = (BoundColumnRefExpression &)*expr;
+		auto &bound_column_ref = expr->Cast<BoundColumnRefExpression>();
 		for (idx_t i = 0; i < replace_bindings.size(); i++) {
 			if (bound_column_ref.binding == replace_bindings[i].old_binding) {
 				bound_column_ref.binding = replace_bindings[i].new_binding;
@@ -79,7 +79,7 @@ void UnnestRewriter::FindCandidates(unique_ptr<LogicalOperator> *op_ptr,
 	}
 
 	// found a delim join
-	auto &delim_join = (LogicalDelimJoin &)*op->children[0];
+	auto &delim_join = op->children[0]->Cast<LogicalDelimJoin>();
 	// only support INNER delim joins
 	if (delim_join.join_type != JoinType::INNER) {
 		return;
@@ -144,9 +144,9 @@ bool UnnestRewriter::RewriteCandidate(unique_ptr<LogicalOperator> *candidate) {
 	// store the table index of the child of the LOGICAL_UNNEST
 	// then update the plan by making the lhs_proj the child of the LOGICAL_UNNEST
 	D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);
-	auto &unnest = (LogicalUnnest &)*curr_op->get();
+	auto &unnest = curr_op->get()->Cast<LogicalUnnest>();
 	D_ASSERT(unnest.children[0]->type == LogicalOperatorType::LOGICAL_DELIM_GET);
-	overwritten_tbl_idx = ((LogicalDelimGet &)*unnest.children[0]).table_index;
+	overwritten_tbl_idx = unnest.children[0]->Cast<LogicalDelimGet>().table_index;
 	unnest.children[0] = std::move(lhs_op);
 
 	// replace the LOGICAL_DELIM_JOIN with its RHS child operator
@@ -166,7 +166,7 @@ void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> *plan_ptr, un
 
 		path_to_unnest.push_back(curr_op);
 		D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION);
-		auto &proj = (LogicalProjection &)*curr_op->get();
+		auto &proj = curr_op->get()->Cast<LogicalProjection>();
 
 		// pop the two last expressions from all projections (delim_idx and UNNEST column)
 		D_ASSERT(proj.expressions.size() > 2);
@@ -189,7 +189,7 @@ void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> *plan_ptr, un
 
 	// update all bindings coming from the LHS to RHS bindings
 	D_ASSERT(topmost_op.children[0]->type == LogicalOperatorType::LOGICAL_PROJECTION);
-	auto &top_proj = (LogicalProjection &)*topmost_op.children[0];
+	auto &top_proj = topmost_op.children[0]->Cast<LogicalProjection>();
 	for (idx_t i = 0; i < lhs_bindings.size(); i++) {
 		ReplaceBinding replace_binding(lhs_bindings[i].binding, ColumnBinding(top_proj.table_index, i));
 		updater.replace_bindings.push_back(replace_binding);
@@ -197,7 +197,7 @@ void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> *plan_ptr, un
 
 	// temporarily remove the BOUND_UNNEST and the child of the LOGICAL_UNNEST from the plan
 	D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);
-	auto &unnest = (LogicalUnnest &)*curr_op->get();
+	auto &unnest = curr_op->get()->Cast<LogicalUnnest>();
 	auto temp_bound_unnest = std::move(unnest.expressions[0]);
 	auto temp_unnest_child = std::move(unnest.children[0]);
 	unnest.expressions.clear();
@@ -213,7 +213,7 @@ void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> *plan_ptr, un
 	for (idx_t i = path_to_unnest.size(); i > 0; i--) {
 
 		D_ASSERT(path_to_unnest[i - 1]->get()->type == LogicalOperatorType::LOGICAL_PROJECTION);
-		auto &proj = (LogicalProjection &)*path_to_unnest[i - 1]->get();
+		auto &proj = path_to_unnest[i - 1]->get()->Cast<LogicalProjection>();
 
 		// temporarily store the existing expressions
 		vector<unique_ptr<Expression>> existing_expressions;
@@ -225,7 +225,7 @@ void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> *plan_ptr, un
 
 		// add the new expressions
 		for (idx_t expr_idx = 0; expr_idx < lhs_bindings.size(); expr_idx++) {
-			auto new_expr = make_unique<BoundColumnRefExpression>(
+			auto new_expr = make_uniq<BoundColumnRefExpression>(
 			    lhs_bindings[expr_idx].alias, lhs_bindings[expr_idx].type, lhs_bindings[expr_idx].binding);
 			proj.expressions.push_back(std::move(new_expr));
 
@@ -254,14 +254,14 @@ void UnnestRewriter::UpdateBoundUnnestBindings(UnnestRewriterPlanUpdater &update
 
 	// found the LOGICAL_UNNEST
 	D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);
-	auto &unnest = (LogicalUnnest &)*curr_op->get();
+	auto &unnest = curr_op->get()->Cast<LogicalUnnest>();
 
 	auto unnest_child_cols = unnest.children[0]->GetColumnBindings();
 	for (idx_t delim_col_idx = 0; delim_col_idx < delim_columns.size(); delim_col_idx++) {
 		for (idx_t child_col_idx = 0; child_col_idx < unnest_child_cols.size(); child_col_idx++) {
 			if (delim_columns[delim_col_idx].table_index == unnest_child_cols[child_col_idx].table_index) {
 				ColumnBinding old_binding(overwritten_tbl_idx, DConstants::INVALID_INDEX);
-				updater.replace_bindings.emplace_back(ReplaceBinding(old_binding, delim_columns[delim_col_idx]));
+				updater.replace_bindings.emplace_back(old_binding, delim_columns[delim_col_idx]);
 				break;
 			}
 		}
@@ -276,11 +276,11 @@ void UnnestRewriter::UpdateBoundUnnestBindings(UnnestRewriterPlanUpdater &update
 void UnnestRewriter::GetDelimColumns(LogicalOperator &op) {
 
 	D_ASSERT(op.type == LogicalOperatorType::LOGICAL_DELIM_JOIN);
-	auto &delim_join = (LogicalDelimJoin &)op;
+	auto &delim_join = op.Cast<LogicalDelimJoin>();
 	for (idx_t i = 0; i < delim_join.duplicate_eliminated_columns.size(); i++) {
 		auto &expr = *delim_join.duplicate_eliminated_columns[i];
 		D_ASSERT(expr.type == ExpressionType::BOUND_COLUMN_REF);
-		auto &bound_colref_expr = (BoundColumnRefExpression &)expr;
+		auto &bound_colref_expr = expr.Cast<BoundColumnRefExpression>();
 		delim_columns.push_back(bound_colref_expr.binding);
 	}
 }
@@ -294,16 +294,16 @@ void UnnestRewriter::GetLHSExpressions(LogicalOperator &op) {
 	bool set_alias = false;
 	// we can easily extract the alias for LOGICAL_PROJECTION(s)
 	if (op.type == LogicalOperatorType::LOGICAL_PROJECTION) {
-		auto &proj = (LogicalProjection &)op;
+		auto &proj = op.Cast<LogicalProjection>();
 		if (proj.expressions.size() == op.types.size()) {
 			set_alias = true;
 		}
 	}
 
 	for (idx_t i = 0; i < op.types.size(); i++) {
-		lhs_bindings.emplace_back(LHSBinding(col_bindings[i], op.types[i]));
+		lhs_bindings.emplace_back(col_bindings[i], op.types[i]);
 		if (set_alias) {
-			auto &proj = (LogicalProjection &)op;
+			auto &proj = op.Cast<LogicalProjection>();
 			lhs_bindings.back().alias = proj.expressions[i]->alias;
 		}
 	}

package/src/duckdb/src/parallel/event.cpp

@@ -72,7 +72,7 @@ void Event::InsertEvent(shared_ptr<Event> replacement_event) {
 	executor.AddEvent(std::move(replacement_event));
 }
 
-void Event::SetTasks(vector<unique_ptr<Task>> tasks) {
+void Event::SetTasks(vector<shared_ptr<Task>> tasks) {
 	auto &ts = TaskScheduler::GetScheduler(executor.context);
 	D_ASSERT(total_tasks == 0);
 	D_ASSERT(!tasks.empty());

package/src/duckdb/src/parallel/executor.cpp

@@ -38,13 +38,19 @@ void Executor::AddEvent(shared_ptr<Event> event) {
 }
 
 struct PipelineEventStack {
-	Event *pipeline_initialize_event;
-	Event *pipeline_event;
-	Event *pipeline_finish_event;
-	Event *pipeline_complete_event;
+	PipelineEventStack(Event &pipeline_initialize_event, Event &pipeline_event, Event &pipeline_finish_event,
+	                   Event &pipeline_complete_event)
+	    : pipeline_initialize_event(pipeline_initialize_event), pipeline_event(pipeline_event),
+	      pipeline_finish_event(pipeline_finish_event), pipeline_complete_event(pipeline_complete_event) {
+	}
+
+	Event &pipeline_initialize_event;
+	Event &pipeline_event;
+	Event &pipeline_finish_event;
+	Event &pipeline_complete_event;
 };
 
-using event_map_t = unordered_map<const Pipeline *, PipelineEventStack>;
+using event_map_t = reference_map_t<Pipeline, PipelineEventStack>;
 
 struct ScheduleEventData {
 	ScheduleEventData(const vector<shared_ptr<MetaPipeline>> &meta_pipelines, vector<shared_ptr<Event>> &events,
@@ -69,17 +75,16 @@ void Executor::SchedulePipeline(const shared_ptr<MetaPipeline> &meta_pipeline, S
 	auto base_event = make_shared<PipelineEvent>(base_pipeline);
 	auto base_finish_event = make_shared<PipelineFinishEvent>(base_pipeline);
 	auto base_complete_event = make_shared<PipelineCompleteEvent>(base_pipeline->executor, event_data.initial_schedule);
-	PipelineEventStack base_stack {base_initialize_event.get(), base_event.get(), base_finish_event.get(),
-	                               base_complete_event.get()};
+	PipelineEventStack base_stack(*base_initialize_event, *base_event, *base_finish_event, *base_complete_event);
 	events.push_back(std::move(base_initialize_event));
 	events.push_back(std::move(base_event));
 	events.push_back(std::move(base_finish_event));
 	events.push_back(std::move(base_complete_event));
 
 	// dependencies: initialize -> event -> finish -> complete
-	base_stack.pipeline_event->AddDependency(*base_stack.pipeline_initialize_event);
-	base_stack.pipeline_finish_event->AddDependency(*base_stack.pipeline_event);
-	base_stack.pipeline_complete_event->AddDependency(*base_stack.pipeline_finish_event);
+	base_stack.pipeline_event.AddDependency(base_stack.pipeline_initialize_event);
+	base_stack.pipeline_finish_event.AddDependency(base_stack.pipeline_event);
+	base_stack.pipeline_complete_event.AddDependency(base_stack.pipeline_finish_event);
 
 	// create an event and stack for all pipelines in the MetaPipeline
 	vector<shared_ptr<Pipeline>> pipelines;
@@ -90,30 +95,30 @@ void Executor::SchedulePipeline(const shared_ptr<MetaPipeline> &meta_pipeline, S
 
 		// create events/stack for this pipeline
 		auto pipeline_event = make_shared<PipelineEvent>(pipeline);
-		Event *pipeline_finish_event_ptr;
+		optional_ptr<Event> pipeline_finish_event_ptr;
 		if (meta_pipeline->HasFinishEvent(pipeline.get())) {
 			// this pipeline has its own finish event (despite going into the same sink - Finalize twice!)
-			auto pipeline_finish_event = make_unique<PipelineFinishEvent>(pipeline);
+			auto pipeline_finish_event = make_shared<PipelineFinishEvent>(pipeline);
 			pipeline_finish_event_ptr = pipeline_finish_event.get();
 			events.push_back(std::move(pipeline_finish_event));
-			base_stack.pipeline_complete_event->AddDependency(*pipeline_finish_event_ptr);
+			base_stack.pipeline_complete_event.AddDependency(*pipeline_finish_event_ptr);
 		} else {
-			pipeline_finish_event_ptr = base_stack.pipeline_finish_event;
+			pipeline_finish_event_ptr = &base_stack.pipeline_finish_event;
 		}
-		PipelineEventStack pipeline_stack {base_stack.pipeline_initialize_event, pipeline_event.get(),
-		                                   pipeline_finish_event_ptr, base_stack.pipeline_complete_event};
+		PipelineEventStack pipeline_stack(base_stack.pipeline_initialize_event, *pipeline_event,
+		                                  *pipeline_finish_event_ptr, base_stack.pipeline_complete_event);
 		events.push_back(std::move(pipeline_event));
 
 		// dependencies: base_initialize -> pipeline_event -> base_finish
-		pipeline_stack.pipeline_event->AddDependency(*base_stack.pipeline_initialize_event);
-		pipeline_stack.pipeline_finish_event->AddDependency(*pipeline_stack.pipeline_event);
+		pipeline_stack.pipeline_event.AddDependency(base_stack.pipeline_initialize_event);
+		pipeline_stack.pipeline_finish_event.AddDependency(pipeline_stack.pipeline_event);
 
 		// add pipeline stack to event map
-		event_map.insert(make_pair(pipeline.get(), pipeline_stack));
+		event_map.insert(make_pair(reference<Pipeline>(*pipeline), pipeline_stack));
 	}
 
 	// add base stack to the event data too
-	event_map.insert(make_pair(base_pipeline.get(), base_stack));
+	event_map.insert(make_pair(reference<Pipeline>(*base_pipeline), base_stack));
 
 	// set up the dependencies within this MetaPipeline
 	for (auto &pipeline : pipelines) {
@@ -128,10 +133,14 @@ void Executor::SchedulePipeline(const shared_ptr<MetaPipeline> &meta_pipeline, S
 		if (!dependencies) {
 			continue;
 		}
-		auto &pipeline_stack = event_map[pipeline.get()];
+		auto root_entry = event_map.find(*pipeline);
+		D_ASSERT(root_entry != event_map.end());
+		auto &pipeline_stack = root_entry->second;
 		for (auto &dependency : *dependencies) {
-			auto &dependency_stack = event_map[dependency];
-			pipeline_stack.pipeline_event->AddDependency(*dependency_stack.pipeline_event);
+			auto event_entry = event_map.find(*dependency);
+			D_ASSERT(event_entry != event_map.end());
+			auto &dependency_stack = event_entry->second;
+			pipeline_stack.pipeline_event.AddDependency(dependency_stack.pipeline_event);
 		}
 	}
 }
@@ -148,15 +157,14 @@ void Executor::ScheduleEventsInternal(ScheduleEventData &event_data) {
 	// set up the dependencies across MetaPipelines
 	auto &event_map = event_data.event_map;
 	for (auto &entry : event_map) {
-		auto pipeline = entry.first;
-		for (auto &dependency : pipeline->dependencies) {
+		auto &pipeline = entry.first.get();
+		for (auto &dependency : pipeline.dependencies) {
 			auto dep = dependency.lock();
 			D_ASSERT(dep);
-			auto event_map_entry = event_map.find(dep.get());
+			auto event_map_entry = event_map.find(*dep);
 			D_ASSERT(event_map_entry != event_map.end());
 			auto &dep_entry = event_map_entry->second;
-			D_ASSERT(dep_entry.pipeline_complete_event);
-			entry.second.pipeline_event->AddDependency(*dep_entry.pipeline_complete_event);
+			entry.second.pipeline_event.AddDependency(dep_entry.pipeline_complete_event);
 		}
 	}
 
@@ -231,7 +239,7 @@ void Executor::VerifyScheduledEventsInternal(const idx_t vertex, const vector<Ev
 	recursion_stack[vertex] = false;
 }
 
-void Executor::AddRecursiveCTE(PhysicalOperator *rec_cte) {
+void Executor::AddRecursiveCTE(PhysicalOperator &rec_cte) {
 	recursive_ctes.push_back(rec_cte);
 }
 
@@ -246,7 +254,7 @@ bool Executor::NextExecutor() {
 		return false;
 	}
 	root_pipelines[root_pipeline_idx]->Reset();
-	root_executor = make_unique<PipelineExecutor>(context, *root_pipelines[root_pipeline_idx]);
+	root_executor = make_uniq<PipelineExecutor>(context, *root_pipelines[root_pipeline_idx]);
 	root_pipeline_idx++;
 	return true;
 }
@@ -258,8 +266,8 @@ void Executor::VerifyPipeline(Pipeline &pipeline) {
 		auto other_operators = other_pipeline->GetOperators();
 		for (idx_t op_idx = 0; op_idx < operators.size(); op_idx++) {
 			for (idx_t other_idx = 0; other_idx < other_operators.size(); other_idx++) {
-				auto &left = *operators[op_idx];
-				auto &right = *other_operators[other_idx];
+				auto &left = operators[op_idx].get();
+				auto &right = other_operators[other_idx].get();
 				if (left.Equals(right)) {
 					D_ASSERT(right.Equals(left));
 				} else {
@@ -281,36 +289,35 @@ void Executor::VerifyPipelines() {
 void Executor::Initialize(unique_ptr<PhysicalOperator> physical_plan) {
 	Reset();
 	owned_plan = std::move(physical_plan);
-	InitializeInternal(owned_plan.get());
+	InitializeInternal(*owned_plan);
 }
 
-void Executor::Initialize(PhysicalOperator *plan) {
+void Executor::Initialize(PhysicalOperator &plan) {
 	Reset();
 	InitializeInternal(plan);
 }
 
-void Executor::InitializeInternal(PhysicalOperator *plan) {
+void Executor::InitializeInternal(PhysicalOperator &plan) {
 
 	auto &scheduler = TaskScheduler::GetScheduler(context);
 	{
 		lock_guard<mutex> elock(executor_lock);
-		physical_plan = plan;
+		physical_plan = &plan;
 
 		this->profiler = ClientData::Get(context).profiler;
-		profiler->Initialize(physical_plan);
+		profiler->Initialize(plan);
 		this->producer = scheduler.CreateProducer();
 
 		// build and ready the pipelines
 		PipelineBuildState state;
 		auto root_pipeline = make_shared<MetaPipeline>(*this, state, nullptr);
-		root_pipeline->Build(physical_plan);
+		root_pipeline->Build(*physical_plan);
 		root_pipeline->Ready();
 
 		// ready recursive cte pipelines too
-		for (auto &rec_cte : recursive_ctes) {
-			D_ASSERT(rec_cte->type == PhysicalOperatorType::RECURSIVE_CTE);
-			auto &rec_cte_op = (PhysicalRecursiveCTE &)*rec_cte;
-			rec_cte_op.recursive_meta_pipeline->Ready();
+		for (auto &rec_cte_ref : recursive_ctes) {
+			auto &rec_cte = rec_cte_ref.get().Cast<PhysicalRecursiveCTE>();
+			rec_cte.recursive_meta_pipeline->Ready();
 		}
 
 		// set root pipelines, i.e., all pipelines that end in the final sink
@@ -346,13 +353,13 @@ void Executor::CancelTasks() {
 		for (auto &pipeline : pipelines) {
 			weak_references.push_back(weak_ptr<Pipeline>(pipeline));
 		}
-		for (auto op : recursive_ctes) {
-			D_ASSERT(op->type == PhysicalOperatorType::RECURSIVE_CTE);
-			auto &rec_cte = (PhysicalRecursiveCTE &)*op;
+		for (auto &rec_cte_ref : recursive_ctes) {
+			auto &rec_cte = rec_cte_ref.get().Cast<PhysicalRecursiveCTE>();
 			rec_cte.recursive_meta_pipeline.reset();
 		}
 		pipelines.clear();
 		root_pipelines.clear();
+		to_be_rescheduled_tasks.clear();
 		events.clear();
 	}
 	WorkOnTasks();
@@ -369,13 +376,44 @@ void Executor::CancelTasks() {
 void Executor::WorkOnTasks() {
 	auto &scheduler = TaskScheduler::GetScheduler(context);
 
-	unique_ptr<Task> task;
+	shared_ptr<Task> task;
 	while (scheduler.GetTaskFromProducer(*producer, task)) {
-		task->Execute(TaskExecutionMode::PROCESS_ALL);
+		auto res = task->Execute(TaskExecutionMode::PROCESS_ALL);
+		if (res == TaskExecutionResult::TASK_BLOCKED) {
+			task->Deschedule();
+		}
 		task.reset();
 	}
 }
 
+void Executor::RescheduleTask(shared_ptr<Task> &task) {
+	// This function will spin lock until the task provided is added to the to_be_rescheduled_tasks
+	while (true) {
+		lock_guard<mutex> l(executor_lock);
+		if (cancelled) {
+			return;
+		}
+		auto entry = to_be_rescheduled_tasks.find(task.get());
+		if (entry != to_be_rescheduled_tasks.end()) {
+			auto &scheduler = TaskScheduler::GetScheduler(context);
+			to_be_rescheduled_tasks.erase(task.get());
+			scheduler.ScheduleTask(GetToken(), task);
+			break;
+		}
+	}
+}
+
+void Executor::AddToBeRescheduled(shared_ptr<Task> &task) {
+	lock_guard<mutex> l(executor_lock);
+	if (cancelled) {
+		return;
+	}
+	if (to_be_rescheduled_tasks.find(task.get()) != to_be_rescheduled_tasks.end()) {
+		return;
+	}
+	to_be_rescheduled_tasks[task.get()] = std::move(task);
+}
+
 bool Executor::ExecutionIsFinished() {
 	return completed_pipelines >= total_pipelines || HasError();
 }
@@ -394,7 +432,10 @@ PendingExecutionResult Executor::ExecuteTask() {
 		if (task) {
 			// if we have a task, partially process it
 			auto result = task->Execute(TaskExecutionMode::PROCESS_PARTIAL);
-			if (result != TaskExecutionResult::TASK_NOT_FINISHED) {
+			if (result == TaskExecutionResult::TASK_BLOCKED) {
+				task->Deschedule();
+				task.reset();
+			} else if (result == TaskExecutionResult::TASK_FINISHED) {
 				// if the task is finished, clean it up
 				task.reset();
 			}
@@ -438,21 +479,22 @@ void Executor::Reset() {
 	exceptions.clear();
 	pipelines.clear();
 	events.clear();
+	to_be_rescheduled_tasks.clear();
 	execution_result = PendingExecutionResult::RESULT_NOT_READY;
 }
 
-shared_ptr<Pipeline> Executor::CreateChildPipeline(Pipeline *current, PhysicalOperator *op) {
-	D_ASSERT(!current->operators.empty());
-	D_ASSERT(op->IsSource());
+shared_ptr<Pipeline> Executor::CreateChildPipeline(Pipeline &current, PhysicalOperator &op) {
+	D_ASSERT(!current.operators.empty());
+	D_ASSERT(op.IsSource());
 	// found another operator that is a source, schedule a child pipeline
 	// 'op' is the source, and the sink is the same
 	auto child_pipeline = make_shared<Pipeline>(*this);
-	child_pipeline->sink = current->sink;
-	child_pipeline->source = op;
+	child_pipeline->sink = current.sink;
+	child_pipeline->source = &op;
 
 	// the child pipeline has the same operators up until 'op'
-	for (auto current_op : current->operators) {
-		if (current_op == op) {
+	for (auto current_op : current.operators) {
+		if (&current_op.get() == &op) {
 			break;
 		}
 		child_pipeline->operators.push_back(current_op);
@@ -528,7 +570,7 @@ unique_ptr<QueryResult> Executor::GetResult() {
 unique_ptr<DataChunk> Executor::FetchChunk() {
 	D_ASSERT(physical_plan);
 
-	auto chunk = make_unique<DataChunk>();
+	auto chunk = make_uniq<DataChunk>();
 	root_executor->InitializeChunk(*chunk);
 	while (true) {
 		root_executor->ExecutePull(*chunk);

package/src/duckdb/src/parallel/executor_task.cpp

@@ -1,5 +1,6 @@
 #include "duckdb/parallel/task.hpp"
 #include "duckdb/execution/executor.hpp"
+#include "duckdb/main/client_context.hpp"
 
 namespace duckdb {
 
@@ -12,6 +13,16 @@ ExecutorTask::ExecutorTask(ClientContext &context) : ExecutorTask(Executor::Get(
 ExecutorTask::~ExecutorTask() {
 }
 
+void ExecutorTask::Deschedule() {
+	auto this_ptr = shared_from_this();
+	executor.AddToBeRescheduled(this_ptr);
+}
+
+void ExecutorTask::Reschedule() {
+	auto this_ptr = shared_from_this();
+	executor.RescheduleTask(this_ptr);
+}
+
 TaskExecutionResult ExecutorTask::Execute(TaskExecutionMode mode) {
 	try {
 		return ExecuteTask(mode);