Repository: duckdb/duckdb
File: src/common/sort/sorted_run_merger.cpp (C++)

Code before:
		                            count);
		sorted_run->key_data->CopyRows(sorted_run->key_append_state.chunk_state, key_data_input,
		                               *FlatVector::IncrementalSelectionVector(), count);

		merged_partition_index += count;
	}

	sorted_run->key_data->FinalizePinState(sorted_run->key_append_state.pin_state);
	if (sorted_run->payload_data) {
		sorted_run->payload_data->FinalizePinState(sorted_run->payload_append_state.pin_state);
	}

	return sorted_run;
}

//===--------------------------------------------------------------------===//
// Sorted Run Merger
//===--------------------------------------------------------------------===//
SortedRunMerger::SortedRunMerger(const Sort &sort_p, vector<unique_ptr<SortedRun>> &&sorted_runs_p,
                                 idx_t partition_size_p, bool external_p, bool is_index_sort_p)
    : sort(sort_p), sorted_runs(std::move(sorted_runs_p)), total_count(SortedRunsTotalCount(sorted_runs)),
      partition_size(partition_size_p), external(external_p), is_index_sort(is_index_sort_p) {
}

SortedRunMerger::~SortedRunMerger() {
	ParallelDestroyTask<decltype(sorted_runs)>::Schedule(*sort.context.db, sorted_runs);
}

unique_ptr<LocalSourceState> SortedRunMerger::GetLocalSourceState(ExecutionContext &,
                                                                  GlobalSourceState &gstate_p) const {
	auto &gstate = gstate_p.Cast<SortedRunMergerGlobalState>();
	auto guard = gstate.Lock();
	return make_uniq<SortedRunMergerLocalState>(gstate);
}

unique_ptr<GlobalSourceState> SortedRunMerger::GetGlobalSourceState(ClientContext &context) const {
	return make_uniq<SortedRunMergerGlobalState>(context, *this);
}

SourceResultType SortedRunMerger::GetData(ExecutionContext &, DataChunk &chunk, OperatorSourceInput &input) const {
	auto &gstate = input.global_state.Cast<SortedRunMergerGlobalState>();
	auto &lstate = input.local_state.Cast<SortedRunMergerLocalState>();

	while (chunk.size() == 0) {
		if (!lstate.TaskFinished() || gstate.AssignTask(lstate)) {
			lstate.ExecuteTask(gstate, &chunk);
		} else {
			break;
		}
	}


Review comment (security):
It looks like `sort` is not guaranteed to stay alive for the duration of the `SortedRunMerger` based on the thread sanitizer logs - maybe we can store the database instance as part of the `SortedRunMerger` instead, i.e. add `db(*sort_p.context.db)` to the constructor?

Code after:
		                            count);
		sorted_run->key_data->CopyRows(sorted_run->key_append_state.chunk_state, key_data_input,
		                               *FlatVector::IncrementalSelectionVector(), count);

		merged_partition_index += count;
	}

	sorted_run->key_data->FinalizePinState(sorted_run->key_append_state.pin_state);
	if (sorted_run->payload_data) {
		sorted_run->payload_data->FinalizePinState(sorted_run->payload_append_state.pin_state);
	}

	return sorted_run;
}

//===--------------------------------------------------------------------===//
// Sorted Run Merger
//===--------------------------------------------------------------------===//
SortedRunMerger::SortedRunMerger(const Sort &sort_p, vector<unique_ptr<SortedRun>> &&sorted_runs_p,
                                 idx_t partition_size_p, bool external_p, bool is_index_sort_p)
    : db(*sort_p.context.db), sort(sort_p), sorted_runs(std::move(sorted_runs_p)),
      total_count(SortedRunsTotalCount(sorted_runs)), partition_size(partition_size_p), external(external_p),
      is_index_sort(is_index_sort_p) {
}

SortedRunMerger::~SortedRunMerger() {
	ParallelDestroyTask<decltype(sorted_runs)>::Schedule(db, sorted_runs);
}

unique_ptr<LocalSourceState> SortedRunMerger::GetLocalSourceState(ExecutionContext &,
                                                                  GlobalSourceState &gstate_p) const {
	auto &gstate = gstate_p.Cast<SortedRunMergerGlobalState>();
	auto guard = gstate.Lock();
	return make_uniq<SortedRunMergerLocalState>(gstate);
}

unique_ptr<GlobalSourceState> SortedRunMerger::GetGlobalSourceState(ClientContext &context) const {
	return make_uniq<SortedRunMergerGlobalState>(context, *this);
}

SourceResultType SortedRunMerger::GetData(ExecutionContext &, DataChunk &chunk, OperatorSourceInput &input) const {
	auto &gstate = input.global_state.Cast<SortedRunMergerGlobalState>();
	auto &lstate = input.local_state.Cast<SortedRunMergerLocalState>();

	while (chunk.size() == 0) {
		if (!lstate.TaskFinished() || gstate.AssignTask(lstate)) {
			lstate.ExecuteTask(gstate, &chunk);
		} else {
			break;
		}
	}