duckdb 1.3.4 → 1.4.1-dev2.0

package/src/duckdb/src/execution/operator/persistent/physical_merge_into.cpp

@@ -0,0 +1,516 @@
+#include "duckdb/execution/operator/persistent/physical_merge_into.hpp"
+#include "duckdb/execution/expression_executor.hpp"
+#include "duckdb/parser/statement/merge_into_statement.hpp"
+
+namespace duckdb {
+
+PhysicalMergeInto::PhysicalMergeInto(PhysicalPlan &physical_plan, vector<LogicalType> types,
+                                     map<MergeActionCondition, vector<unique_ptr<MergeIntoOperator>>> actions_p,
+                                     idx_t row_id_index, optional_idx source_marker, bool parallel_p,
+                                     bool return_chunk_p)
+    : PhysicalOperator(physical_plan, PhysicalOperatorType::MERGE_INTO, std::move(types), 1),
+      row_id_index(row_id_index), source_marker(source_marker), parallel(parallel_p), return_chunk(return_chunk_p) {
+
+	map<MergeActionCondition, MergeActionRange> ranges;
+	for (auto &entry : actions_p) {
+		MergeActionRange range;
+		range.condition = entry.first;
+		range.start = actions.size();
+		for (auto &action : entry.second) {
+			actions.push_back(std::move(action));
+		}
+		range.end = actions.size();
+		ranges.emplace(entry.first, range);
+	}
+	match_actions = {MergeActionCondition::WHEN_MATCHED, MergeActionCondition::WHEN_NOT_MATCHED_BY_TARGET,
+	                 MergeActionCondition::WHEN_NOT_MATCHED_BY_SOURCE};
+	for (idx_t i = 0; i < match_actions.size(); i++) {
+		auto entry = ranges.find(match_actions[i]);
+		MergeActionRange range;
+		if (entry != ranges.end()) {
+			range = entry->second;
+		}
+		range.condition = match_actions[i];
+		action_ranges.push_back(range);
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
+struct MergeSinkState {
+	MergeSinkState() : selected_sel(STANDARD_VECTOR_SIZE), remaining_sel(STANDARD_VECTOR_SIZE) {
+	}
+
+	bool computed_matches = false;
+	bool match_initialized = false;
+	idx_t match_idx = 0;
+	idx_t index_in_match = 0;
+	SelectionVector current_sel;
+	SelectionVector selected_sel;
+	SelectionVector remaining_sel;
+	idx_t current_count;
+	unique_ptr<DataChunk> sliced_chunk;
+	optional_ptr<DataChunk> input_chunk;
+};
+
+struct MergeLocalExecutionState {
+	unique_ptr<LocalSinkState> local_state;
+	unique_ptr<ExpressionExecutor> condition_executor;
+	unique_ptr<ExpressionExecutor> insert_executor;
+	unique_ptr<DataChunk> insert_chunk;
+};
+
+struct MatchResult {
+	MatchResult(ClientContext &context, const vector<LogicalType> &types) : sel(STANDARD_VECTOR_SIZE), count(0) {
+		chunk = make_uniq<DataChunk>();
+		chunk->Initialize(context, types);
+	}
+
+	SelectionVector sel;
+	idx_t count;
+	unique_ptr<DataChunk> chunk;
+};
+
+class MergeIntoLocalState : public LocalSinkState {
+public:
+	MergeIntoLocalState(ExecutionContext &context, const PhysicalMergeInto &op) {
+		for (auto &action : op.actions) {
+			MergeLocalExecutionState state;
+			if (action->op) {
+				state.local_state = action->op->GetLocalSinkState(context);
+			}
+			if (action->condition) {
+				state.condition_executor = make_uniq<ExpressionExecutor>(context.client, *action->condition);
+			}
+			if (!action->expressions.empty()) {
+				state.insert_executor = make_uniq<ExpressionExecutor>(context.client, action->expressions);
+				vector<LogicalType> insert_types;
+				for (auto &expr : action->expressions) {
+					insert_types.push_back(expr->return_type);
+				}
+				state.insert_chunk = make_uniq<DataChunk>();
+				state.insert_chunk->Initialize(context.client, insert_types);
+			}
+
+			states.push_back(std::move(state));
+		}
+		for (idx_t i = 0; i < 3; i++) {
+			match_results.emplace_back(context.client, op.children[0].get().types);
+		}
+	}
+
+	MergeSinkState sink_state;
+	vector<MatchResult> match_results;
+	idx_t combine_idx = 0;
+	vector<MergeLocalExecutionState> states;
+	idx_t merged_count = 0;
+};
+
+class MergeIntoGlobalState : public GlobalSinkState {
+public:
+	MergeIntoGlobalState(ClientContext &context, const PhysicalMergeInto &op) : op(op) {
+		for (auto &action : op.actions) {
+			sink_states.push_back(action->op ? action->op->GetGlobalSinkState(context) : nullptr);
+		}
+		merged_count = 0;
+	}
+	const PhysicalMergeInto &op;
+	idx_t finalize_idx = 0;
+	vector<unique_ptr<GlobalSinkState>> sink_states;
+	atomic<idx_t> merged_count;
+
+	optional_ptr<DataChunk> ComputeActionInput(ClientContext &context, MergeIntoOperator &action, DataChunk &chunk,
+	                                           MergeIntoLocalState &local_state,
+	                                           MergeLocalExecutionState &local_action_state) {
+		auto &current_count = local_state.sink_state.current_count;
+		auto &current_sel = local_state.sink_state.current_sel;
+		auto &sliced_chunk = local_state.sink_state.sliced_chunk;
+		auto &selected_sel = local_state.sink_state.selected_sel;
+		auto &remaining_sel = local_state.sink_state.remaining_sel;
+		if (current_count == 0) {
+			return nullptr;
+		}
+		if (!sliced_chunk) {
+			sliced_chunk = make_uniq<DataChunk>();
+			sliced_chunk->Initialize(context, chunk.GetTypes());
+		} else {
+			sliced_chunk->Reset();
+		}
+		optional_ptr<DataChunk> result;
+		if (action.condition) {
+			// if we have a condition we need to evaluate it
+			auto &executor = *local_action_state.condition_executor;
+			idx_t match_count =
+			    executor.SelectExpression(chunk, selected_sel, remaining_sel, current_sel, current_count);
+			if (match_count == 0) {
+				// no matches - move to next action
+				return nullptr;
+			}
+			// slice the chunk for this action with the matching sel
+			sliced_chunk->Slice(chunk, selected_sel, match_count);
+			result = sliced_chunk;
+
+			// for the next chunk - update the matches
+			current_count = current_count - match_count;
+			current_sel.Initialize(remaining_sel);
+		} else if (current_count != chunk.size()) {
+			// if we have previously processed rows - remove them
+			sliced_chunk->Slice(chunk, current_sel, current_count);
+			result = sliced_chunk;
+		} else {
+			result = chunk;
+		}
+		// if we have any expressions - execute them to generate the new input chunk
+		if (!action.expressions.empty()) {
+			auto &insert_chunk = local_action_state.insert_chunk;
+			insert_chunk->Reset();
+			local_action_state.insert_executor->Execute(*result, *insert_chunk);
+			result = insert_chunk.get();
+		}
+		if (action.op) {
+			local_state.merged_count += result->size();
+		}
+		return result;
+	}
+
+	SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, MergeIntoLocalState &local_state,
+	                    OperatorSinkInput &input, MergeActionRange range, idx_t &index_in_match) {
+		auto &input_chunk = local_state.sink_state.input_chunk;
+		for (; range.start + index_in_match < range.end; index_in_match++) {
+			idx_t i = range.start + index_in_match;
+			auto &action = op.actions[i];
+			auto &local_action_state = local_state.states[i];
+			if (!input_chunk) {
+				// first time processing this action - compute the input chunk
+				input_chunk = ComputeActionInput(context.client, *action, chunk, local_state, local_action_state);
+				if (!input_chunk) {
+					// no data for this action - move to next action
+					continue;
+				}
+			}
+			// process the action
+			if (!action->op) {
+				if (action->action_type == MergeActionType::MERGE_ERROR) {
+					// abort - generate an error message
+					string merge_condition;
+					merge_condition += MergeIntoStatement::ActionConditionToString(range.condition);
+					if (action->condition) {
+						merge_condition += " AND " + action->condition->ToString();
+					}
+					if (!action->expressions.empty()) {
+						// if there are any user-provided error messages: add the first error message encountered
+						merge_condition += ": " + input_chunk->data[0].GetValue(0).ToString();
+					}
+					throw ConstraintException("Merge error condition %s", merge_condition);
+				}
+				D_ASSERT(action->action_type == MergeActionType::MERGE_DO_NOTHING);
+				input_chunk = nullptr;
+				continue;
+			}
+			auto &gstate = sink_states[i];
+			auto &lstate = *local_action_state.local_state;
+			OperatorSinkInput sink_input {*gstate, lstate, input.interrupt_state};
+			auto result = action->op->Sink(context, *input_chunk, sink_input);
+			if (result == SinkResultType::BLOCKED) {
+				return SinkResultType::BLOCKED;
+			}
+			// move to next action
+			input_chunk = nullptr;
+		}
+		return SinkResultType::NEED_MORE_INPUT;
+	}
+
+	SinkCombineResultType Combine(ExecutionContext &context, MergeIntoLocalState &local_state,
+	                              OperatorSinkCombineInput &input) {
+		for (; local_state.combine_idx < local_state.states.size(); ++local_state.combine_idx) {
+			auto &lstate = local_state.states[local_state.combine_idx];
+			auto &gstate = sink_states[local_state.combine_idx];
+			auto &action = op.actions[local_state.combine_idx];
+			if (!action->op) {
+				continue;
+			}
+			OperatorSinkCombineInput combine_input {*gstate, *lstate.local_state, input.interrupt_state};
+			auto result = action->op->Combine(context, combine_input);
+			if (result == SinkCombineResultType::BLOCKED) {
+				return SinkCombineResultType::BLOCKED;
+			}
+		}
+		return SinkCombineResultType::FINISHED;
+	}
+
+	SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
+	                          OperatorSinkFinalizeInput &input) {
+		for (; finalize_idx < sink_states.size(); ++finalize_idx) {
+			auto &gstate = sink_states[finalize_idx];
+			auto &action = op.actions[finalize_idx];
+			if (!action->op) {
+				continue;
+			}
+
+			OperatorSinkFinalizeInput finalize_input {*gstate, input.interrupt_state};
+			auto result = action->op->Finalize(pipeline, event, context, finalize_input);
+			if (result == SinkFinalizeType::BLOCKED) {
+				return SinkFinalizeType::BLOCKED;
+			}
+		}
+		return SinkFinalizeType::READY;
+	}
+};
+
+unique_ptr<GlobalSinkState> PhysicalMergeInto::GetGlobalSinkState(ClientContext &context) const {
+	return make_uniq<MergeIntoGlobalState>(context, *this);
+}
+
+unique_ptr<LocalSinkState> PhysicalMergeInto::GetLocalSinkState(ExecutionContext &context) const {
+	return make_uniq<MergeIntoLocalState>(context, *this);
+}
+
+idx_t PhysicalMergeInto::GetIndex(MergeActionCondition condition) const {
+	for (idx_t i = 0; i < match_actions.size(); ++i) {
+		if (match_actions[i] == condition) {
+			return i;
+		}
+	}
+	throw InternalException("Unsupported match action condition");
+}
+
+void PhysicalMergeInto::ComputeMatches(MergeIntoLocalState &local_state, DataChunk &chunk) const {
+	auto &match_results = local_state.match_results;
+
+	// for each row, figure out if we have generated a match or not
+	auto &matched = match_results[0];
+	auto &not_matched = match_results[1];
+	auto &not_matched_by_source = match_results[2];
+
+	matched.count = 0;
+	not_matched.count = 0;
+	not_matched_by_source.count = 0;
+
+	UnifiedVectorFormat row_id_data;
+	chunk.data[row_id_index].ToUnifiedFormat(chunk.size(), row_id_data);
+	if (source_marker.IsValid()) {
+		// source marker - check both row id and source marker
+		UnifiedVectorFormat source_marker_data;
+		chunk.data[source_marker.GetIndex()].ToUnifiedFormat(chunk.size(), source_marker_data);
+		for (idx_t i = 0; i < chunk.size(); i++) {
+			if (!source_marker_data.validity.RowIsValid(source_marker_data.sel->get_index(i))) {
+				// source marker is NULL - no source match
+				not_matched_by_source.sel.set_index(not_matched_by_source.count++, i);
+			} else if (!row_id_data.validity.RowIsValid(row_id_data.sel->get_index(i))) {
+				// target marker is NULL - no target match
+				not_matched.sel.set_index(not_matched.count++, i);
+			} else {
+				// match
+				matched.sel.set_index(matched.count++, i);
+			}
+		}
+	} else {
+		// no source marker - only check row-ids
+		for (idx_t i = 0; i < chunk.size(); i++) {
+			auto idx = row_id_data.sel->get_index(i);
+			if (row_id_data.validity.RowIsValid(idx)) {
+				// match
+				matched.sel.set_index(matched.count++, i);
+			} else {
+				// no match
+				not_matched.sel.set_index(not_matched.count++, i);
+			}
+		}
+	}
+
+	// reset and slice chunks
+	for (auto &match : match_results) {
+		if (match.count == 0) {
+			continue;
+		}
+		match.chunk->Reset();
+		match.chunk->Slice(chunk, match.sel, match.count);
+	}
+}
+
+SinkResultType PhysicalMergeInto::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
+	auto &global_state = input.global_state.Cast<MergeIntoGlobalState>();
+	auto &local_state = input.local_state.Cast<MergeIntoLocalState>();
+
+	auto &match_results = local_state.match_results;
+	auto &computed_matches = local_state.sink_state.computed_matches;
+	auto &match_idx = local_state.sink_state.match_idx;
+	auto &index_in_match = local_state.sink_state.index_in_match;
+	auto &match_initialized = local_state.sink_state.match_initialized;
+	auto &current_sel = local_state.sink_state.current_sel;
+	auto &current_count = local_state.sink_state.current_count;
+	if (!computed_matches) {
+		// we haven't figured out which rows have which types of matches - compute them
+		ComputeMatches(local_state, chunk);
+
+		// set up the state so we can prepare sinking into the relevant operators
+		computed_matches = true;
+		match_idx = 0;
+		index_in_match = 0;
+		match_initialized = false;
+	}
+	// now slice and call sink for each of the match conditions
+	for (; match_idx < 3; match_idx++) {
+		auto &match_result = match_results[match_idx];
+		if (match_result.count == 0) {
+			// no matches for this action
+			continue;
+		}
+		if (!match_initialized) {
+			current_sel = SelectionVector();
+			current_count = match_result.count;
+			match_initialized = true;
+		}
+		auto match_range_index = GetIndex(match_actions[match_idx]);
+		auto result = global_state.Sink(context, *match_result.chunk, local_state, input,
+		                                action_ranges[match_range_index], index_in_match);
+		if (result == SinkResultType::BLOCKED) {
+			return SinkResultType::BLOCKED;
+		}
+		// move to next match action
+		index_in_match = 0;
+		match_initialized = false;
+	}
+	// finished - prepare for next match
+	computed_matches = false;
+	return SinkResultType::NEED_MORE_INPUT;
+}
+
+SinkCombineResultType PhysicalMergeInto::Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const {
+	auto &global_state = input.global_state.Cast<MergeIntoGlobalState>();
+	auto &local_state = input.local_state.Cast<MergeIntoLocalState>();
+
+	auto result = global_state.Combine(context, local_state, input);
+	if (result == SinkCombineResultType::FINISHED) {
+		global_state.merged_count += local_state.merged_count;
+	}
+	return result;
+}
+
+SinkFinalizeType PhysicalMergeInto::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
+                                             OperatorSinkFinalizeInput &input) const {
+	auto &global_state = input.global_state.Cast<MergeIntoGlobalState>();
+	return global_state.Finalize(pipeline, event, context, input);
+}
+
+//===--------------------------------------------------------------------===//
+// Source
+//===--------------------------------------------------------------------===//
+class MergeGlobalSourceState : public GlobalSourceState {
+public:
+	explicit MergeGlobalSourceState(ClientContext &context, const PhysicalMergeInto &op) {
+		if (!op.return_chunk) {
+			return;
+		}
+		auto &g = op.sink_state->Cast<MergeIntoGlobalState>();
+		for (idx_t i = 0; i < op.actions.size(); i++) {
+			auto &action = *op.actions[i];
+			unique_ptr<GlobalSourceState> global_state;
+			if (action.op) {
+				// assign the global sink state
+				action.op->sink_state = std::move(g.sink_states[i]);
+				// initialize the global source state
+				global_state = action.op->GetGlobalSourceState(context);
+			}
+			global_states.push_back(std::move(global_state));
+		}
+	}
+
+	vector<unique_ptr<GlobalSourceState>> global_states;
+};
+
+class MergeLocalSourceState : public LocalSourceState {
+public:
+	explicit MergeLocalSourceState(ExecutionContext &context, const PhysicalMergeInto &op,
+	                               MergeGlobalSourceState &gstate) {
+		if (!op.return_chunk) {
+			return;
+		}
+		for (idx_t i = 0; i < op.actions.size(); i++) {
+			auto &action = *op.actions[i];
+			unique_ptr<LocalSourceState> local_state;
+			if (action.op) {
+				local_state = action.op->GetLocalSourceState(context, *gstate.global_states[i]);
+			}
+			local_states.push_back(std::move(local_state));
+		}
+		vector<LogicalType> scan_types;
+		for (idx_t c = 0; c < op.types.size() - 1; c++) {
+			scan_types.emplace_back(op.types[c]);
+		}
+		scan_chunk.Initialize(context.client, scan_types);
+	}
+
+	DataChunk scan_chunk;
+	vector<unique_ptr<LocalSourceState>> local_states;
+	idx_t index = 0;
+};
+
+unique_ptr<GlobalSourceState> PhysicalMergeInto::GetGlobalSourceState(ClientContext &context) const {
+	return make_uniq<MergeGlobalSourceState>(context, *this);
+}
+
+unique_ptr<LocalSourceState> PhysicalMergeInto::GetLocalSourceState(ExecutionContext &context,
+                                                                    GlobalSourceState &gstate) const {
+	return make_uniq<MergeLocalSourceState>(context, *this, gstate.Cast<MergeGlobalSourceState>());
+}
+
+SourceResultType PhysicalMergeInto::GetData(ExecutionContext &context, DataChunk &chunk,
+                                            OperatorSourceInput &input) const {
+	auto &g = sink_state->Cast<MergeIntoGlobalState>();
+	if (!return_chunk) {
+		chunk.SetCardinality(1);
+		chunk.SetValue(0, 0, Value::BIGINT(NumericCast<int64_t>(g.merged_count.load())));
+		return SourceResultType::FINISHED;
+	}
+	auto &gstate = input.global_state.Cast<MergeGlobalSourceState>();
+	auto &lstate = input.local_state.Cast<MergeLocalSourceState>();
+	chunk.Reset();
+	for (; lstate.index < actions.size(); lstate.index++) {
+		auto &action = *actions[lstate.index];
+		if (!action.op) {
+			// no action to scan from
+			continue;
+		}
+		auto &child_gstate = *gstate.global_states[lstate.index];
+		auto &child_lstate = *lstate.local_states[lstate.index];
+		OperatorSourceInput source_input {child_gstate, child_lstate, input.interrupt_state};
+
+		auto result = action.op->GetData(context, lstate.scan_chunk, source_input);
+		if (lstate.scan_chunk.size() > 0) {
+			// construct the result chunk
+			for (idx_t c = 0; c < lstate.scan_chunk.ColumnCount(); c++) {
+				chunk.data[c].Reference(lstate.scan_chunk.data[c]);
+			}
+			// set the merge action
+			string merge_action_name;
+			switch (action.action_type) {
+			case MergeActionType::MERGE_UPDATE:
+				merge_action_name = "UPDATE";
+				break;
+			case MergeActionType::MERGE_INSERT:
+				merge_action_name = "INSERT";
+				break;
+			case MergeActionType::MERGE_DELETE:
+				merge_action_name = "DELETE";
+				break;
+			default:
+				throw InternalException("Unsupported merge action for RETURNING");
+			}
+			Value merge_action(merge_action_name);
+			chunk.data.back().Reference(merge_action);
+			chunk.SetCardinality(lstate.scan_chunk.size());
+		}
+
+		if (result != SourceResultType::FINISHED) {
+			return result;
+		}
+		if (chunk.size() != 0) {
+			return SourceResultType::HAVE_MORE_OUTPUT;
+		}
+	}
+	return SourceResultType::FINISHED;
+}
+
+} // namespace duckdb

package/src/duckdb/src/execution/operator/persistent/physical_update.cpp

@@ -15,13 +15,14 @@
 
 namespace duckdb {
 
-PhysicalUpdate::PhysicalUpdate(vector<LogicalType> types, TableCatalogEntry &tableref, DataTable &table,
-                               vector<PhysicalIndex> columns, vector<unique_ptr<Expression>> expressions,
+PhysicalUpdate::PhysicalUpdate(PhysicalPlan &physical_plan, vector<LogicalType> types, TableCatalogEntry &tableref,
+                               DataTable &table, vector<PhysicalIndex> columns,
+                               vector<unique_ptr<Expression>> expressions,
                                vector<unique_ptr<Expression>> bound_defaults,
                                vector<unique_ptr<BoundConstraint>> bound_constraints, idx_t estimated_cardinality,
                                bool return_chunk)
-    : PhysicalOperator(PhysicalOperatorType::UPDATE, std::move(types), estimated_cardinality), tableref(tableref),
-      table(table), columns(std::move(columns)), expressions(std::move(expressions)),
+    : PhysicalOperator(physical_plan, PhysicalOperatorType::UPDATE, std::move(types), estimated_cardinality),
+      tableref(tableref), table(table), columns(std::move(columns)), expressions(std::move(expressions)),
       bound_defaults(std::move(bound_defaults)), bound_constraints(std::move(bound_constraints)),
       return_chunk(return_chunk), index_update(false) {
 
@@ -55,7 +56,7 @@ public:
 	}
 
 	mutex lock;
-	idx_t updated_count;
+	atomic<idx_t> updated_count;
 	unordered_set<row_t> updated_rows;
 	ColumnDataCollection return_collection;
 };
@@ -127,7 +128,6 @@ SinkResultType PhysicalUpdate::Sink(ExecutionContext &context, DataChunk &chunk,
 		update_chunk.data[i].Reference(chunk.data[binding.index]);
 	}
 
-	lock_guard<mutex> glock(g_state.lock);
 	auto &row_ids = chunk.data[chunk.ColumnCount() - 1];
 	DataChunk &mock_chunk = l_state.mock_chunk;
 
@@ -143,6 +143,7 @@ SinkResultType PhysicalUpdate::Sink(ExecutionContext &context, DataChunk &chunk,
 		table.Update(update_state, context.client, row_ids, columns, update_chunk);
 
 		if (return_chunk) {
+			lock_guard<mutex> glock(g_state.lock);
 			g_state.return_collection.Append(mock_chunk);
 		}
 		g_state.updated_count += chunk.size();
@@ -158,10 +159,11 @@ SinkResultType PhysicalUpdate::Sink(ExecutionContext &context, DataChunk &chunk,
 	idx_t update_count = 0;
 	auto row_id_data = FlatVector::GetData<row_t>(row_ids);
 
+	lock_guard<mutex> glock(g_state.lock);
 	for (idx_t i = 0; i < update_chunk.size(); i++) {
 		auto row_id = row_id_data[i];
-		if (g_state.updated_rows.find(row_id) == g_state.updated_rows.end()) {
-			g_state.updated_rows.insert(row_id);
+		const auto is_new = g_state.updated_rows.insert(row_id).second;
+		if (is_new) {
 			sel.set_index(update_count++, i);
 		}
 	}
@@ -248,7 +250,7 @@ SourceResultType PhysicalUpdate::GetData(ExecutionContext &context, DataChunk &c
 	auto &g = sink_state->Cast<UpdateGlobalState>();
 	if (!return_chunk) {
 		chunk.SetCardinality(1);
-		chunk.SetValue(0, 0, Value::BIGINT(NumericCast<int64_t>(g.updated_count)));
+		chunk.SetValue(0, 0, Value::BIGINT(NumericCast<int64_t>(g.updated_count.load())));
 		return SourceResultType::FINISHED;
 	}
 

package/src/duckdb/src/execution/operator/projection/physical_pivot.cpp

@@ -1,11 +1,15 @@
 #include "duckdb/execution/operator/projection/physical_pivot.hpp"
+
+#include "duckdb/execution/physical_plan_generator.hpp"
 #include "duckdb/planner/expression/bound_aggregate_expression.hpp"
 
 namespace duckdb {
 
-PhysicalPivot::PhysicalPivot(vector<LogicalType> types_p, PhysicalOperator &child, BoundPivotInfo bound_pivot_p)
-    : PhysicalOperator(PhysicalOperatorType::PIVOT, std::move(types_p), child.estimated_cardinality),
+PhysicalPivot::PhysicalPivot(PhysicalPlan &physical_plan, vector<LogicalType> types_p, PhysicalOperator &child,
+                             BoundPivotInfo bound_pivot_p)
+    : PhysicalOperator(physical_plan, PhysicalOperatorType::PIVOT, std::move(types_p), child.estimated_cardinality),
       bound_pivot(std::move(bound_pivot_p)) {
+
 	children.push_back(child);
 	for (idx_t p = 0; p < bound_pivot.pivot_values.size(); p++) {
 		auto entry = pivot_map.find(bound_pivot.pivot_values[p]);
@@ -15,7 +19,6 @@ PhysicalPivot::PhysicalPivot(vector<LogicalType> types_p, PhysicalOperator &chil
 		pivot_map[bound_pivot.pivot_values[p]] = bound_pivot.group_count + p;
 	}
 	// extract the empty aggregate expressions
-	ArenaAllocator allocator(Allocator::DefaultAllocator());
 	for (auto &aggr_expr : bound_pivot.aggregates) {
 		auto &aggr = aggr_expr->Cast<BoundAggregateExpression>();
 		// for each aggregate, initialize an empty aggregate state and finalize it immediately
@@ -23,7 +26,7 @@ PhysicalPivot::PhysicalPivot(vector<LogicalType> types_p, PhysicalOperator &chil
 		aggr.function.initialize(aggr.function, state.get());
 		Vector state_vector(Value::POINTER(CastPointerToValue(state.get())));
 		Vector result_vector(aggr_expr->return_type);
-		AggregateInputData aggr_input_data(aggr.bind_info.get(), allocator);
+		AggregateInputData aggr_input_data(aggr.bind_info.get(), physical_plan.ArenaRef());
 		aggr.function.finalize(state_vector, aggr_input_data, result_vector, 1, 0);
 		empty_aggregates.push_back(result_vector.GetValue(0));
 	}

package/src/duckdb/src/execution/operator/projection/physical_projection.cpp

@@ -19,9 +19,9 @@ public:
 	}
 };
 
-PhysicalProjection::PhysicalProjection(vector<LogicalType> types, vector<unique_ptr<Expression>> select_list,
-                                       idx_t estimated_cardinality)
-    : PhysicalOperator(PhysicalOperatorType::PROJECTION, std::move(types), estimated_cardinality),
+PhysicalProjection::PhysicalProjection(PhysicalPlan &physical_plan, vector<LogicalType> types,
+                                       vector<unique_ptr<Expression>> select_list, idx_t estimated_cardinality)
+    : PhysicalOperator(physical_plan, PhysicalOperatorType::PROJECTION, std::move(types), estimated_cardinality),
       select_list(std::move(select_list)) {
 }
 
@@ -36,39 +36,6 @@ unique_ptr<OperatorState> PhysicalProjection::GetOperatorState(ExecutionContext
 	return make_uniq<ProjectionState>(context, select_list);
 }
 
-unique_ptr<PhysicalOperator>
-PhysicalProjection::CreateJoinProjection(vector<LogicalType> proj_types, const vector<LogicalType> &lhs_types,
-                                         const vector<LogicalType> &rhs_types, const vector<idx_t> &left_projection_map,
-                                         const vector<idx_t> &right_projection_map, const idx_t estimated_cardinality) {
-
-	vector<unique_ptr<Expression>> proj_selects;
-	proj_selects.reserve(proj_types.size());
-
-	if (left_projection_map.empty()) {
-		for (storage_t i = 0; i < lhs_types.size(); ++i) {
-			proj_selects.emplace_back(make_uniq<BoundReferenceExpression>(lhs_types[i], i));
-		}
-	} else {
-		for (auto i : left_projection_map) {
-			proj_selects.emplace_back(make_uniq<BoundReferenceExpression>(lhs_types[i], i));
-		}
-	}
-	const auto left_cols = lhs_types.size();
-
-	if (right_projection_map.empty()) {
-		for (storage_t i = 0; i < rhs_types.size(); ++i) {
-			proj_selects.emplace_back(make_uniq<BoundReferenceExpression>(rhs_types[i], left_cols + i));
-		}
-
-	} else {
-		for (auto i : right_projection_map) {
-			proj_selects.emplace_back(make_uniq<BoundReferenceExpression>(rhs_types[i], left_cols + i));
-		}
-	}
-
-	return make_uniq<PhysicalProjection>(std::move(proj_types), std::move(proj_selects), estimated_cardinality);
-}
-
 InsertionOrderPreservingMap<string> PhysicalProjection::ParamsToString() const {
 	InsertionOrderPreservingMap<string> result;
 	string projections;