duckdb 0.7.2-dev2867.0 → 0.7.2-dev2995.0

package/src/duckdb/src/include/duckdb/verification/statement_verifier.hpp

@@ -20,6 +20,7 @@ enum class VerificationType : uint8_t {
 	DESERIALIZED,
 	PARSED,
 	UNOPTIMIZED,
+	NO_OPERATOR_CACHING,
 	PREPARED,
 	EXTERNAL,
 
@@ -59,6 +60,10 @@ public:
 		return false;
 	}
 
+	virtual bool DisableOperatorCaching() const {
+		return false;
+	}
+
 	virtual bool ForceExternal() const {
 		return false;
 	}

package/src/duckdb/src/main/client_verify.cpp

@@ -28,6 +28,10 @@ PreservedError ClientContext::VerifyQuery(ClientContextLock &lock, const string
 		statement_verifiers.emplace_back(StatementVerifier::Create(VerificationType::DESERIALIZED, stmt));
 		statement_verifiers.emplace_back(StatementVerifier::Create(VerificationType::UNOPTIMIZED, stmt));
 		prepared_statement_verifier = StatementVerifier::Create(VerificationType::PREPARED, stmt);
+#ifdef DUCKDB_DEBUG_ASYNC_SINK_SOURCE
+		// This verification is quite slow, so we only run it for the async sink/source debug mode
+		statement_verifiers.emplace_back(StatementVerifier::Create(VerificationType::NO_OPERATOR_CACHING, stmt));
+#endif
 	}
 	if (config.verify_external) {
 		statement_verifiers.emplace_back(StatementVerifier::Create(VerificationType::EXTERNAL, stmt));

package/src/duckdb/src/main/config.cpp

@@ -14,6 +14,10 @@
 
 namespace duckdb {
 
+#ifdef DEBUG
+bool DBConfigOptions::debug_print_bindings = false;
+#endif
+
 #define DUCKDB_GLOBAL(_PARAM)                                                                                          \
 	{                                                                                                                  \
 		_PARAM::Name, _PARAM::Description, _PARAM::InputType, _PARAM::SetGlobal, nullptr, _PARAM::ResetGlobal,         \

package/src/duckdb/src/main/database.cpp

@@ -1,24 +1,24 @@
 #include "duckdb/main/database.hpp"
 
 #include "duckdb/catalog/catalog.hpp"
-#include "duckdb/main/database_manager.hpp"
 #include "duckdb/common/virtual_file_system.hpp"
+#include "duckdb/execution/operator/helper/physical_set.hpp"
+#include "duckdb/function/cast/cast_function_set.hpp"
+#include "duckdb/function/compression_function.hpp"
+#include "duckdb/main/attached_database.hpp"
 #include "duckdb/main/client_context.hpp"
-#include "duckdb/parallel/task_scheduler.hpp"
-#include "duckdb/storage/storage_manager.hpp"
-#include "duckdb/storage/object_cache.hpp"
-#include "duckdb/transaction/transaction_manager.hpp"
 #include "duckdb/main/connection_manager.hpp"
-#include "duckdb/function/compression_function.hpp"
-#include "duckdb/main/extension_helper.hpp"
-#include "duckdb/function/cast/cast_function_set.hpp"
+#include "duckdb/main/database_manager.hpp"
 #include "duckdb/main/error_manager.hpp"
-#include "duckdb/main/attached_database.hpp"
+#include "duckdb/main/extension_helper.hpp"
+#include "duckdb/parallel/task_scheduler.hpp"
 #include "duckdb/parser/parsed_data/attach_info.hpp"
 #include "duckdb/storage/magic_bytes.hpp"
-#include "duckdb/storage/storage_extension.hpp"
-#include "duckdb/execution/operator/helper/physical_set.hpp"
+#include "duckdb/storage/object_cache.hpp"
 #include "duckdb/storage/standard_buffer_manager.hpp"
+#include "duckdb/storage/storage_extension.hpp"
+#include "duckdb/storage/storage_manager.hpp"
+#include "duckdb/transaction/transaction_manager.hpp"
 
 #ifndef DUCKDB_NO_THREADS
 #include "duckdb/common/thread.hpp"

package/src/duckdb/src/main/extension/extension_load.cpp

@@ -30,14 +30,18 @@ static T LoadFunctionFromDLL(void *dll, const string &function_name, const strin
 	return (T)function;
 }
 
-void ComputeSHA256(FileHandle *handle, const idx_t start, const idx_t end, std::string *res) {
+static void ComputeSHA256String(const std::string &to_hash, std::string *res) {
+	// Invoke MbedTls function to actually compute sha256
+	*res = duckdb_mbedtls::MbedTlsWrapper::ComputeSha256Hash(to_hash);
+}
+
+static void ComputeSHA256FileSegment(FileHandle *handle, const idx_t start, const idx_t end, std::string *res) {
 	const idx_t len = end - start;
 	string file_content;
 	file_content.resize(len);
 	handle->Read((void *)file_content.data(), len, start);
 
-	// Invoke MbedTls function to actually compute sha256
-	*res = duckdb_mbedtls::MbedTlsWrapper::ComputeSha256Hash(file_content);
+	ComputeSHA256String(file_content, res);
 }
 
 bool ExtensionHelper::TryInitialLoad(DBConfig &config, FileOpener *opener, const string &extension,
@@ -86,19 +90,19 @@ bool ExtensionHelper::TryInitialLoad(DBConfig &config, FileOpener *opener, const
 
 		const idx_t maxLenChunks = 1024 * 1024;
 		const idx_t numChunks = (signature_offset + maxLenChunks - 1) / maxLenChunks;
-		std::vector<std::string> chunks(numChunks);
+		std::vector<std::string> hash_chunks(numChunks);
 		std::vector<idx_t> splits(numChunks + 1);
 
-		splits.back() = signature_offset;
-		for (idx_t i = 0; i < chunks.size(); i++) {
+		for (idx_t i = 0; i < numChunks; i++) {
 			splits[i] = maxLenChunks * i;
 		}
+		splits.back() = signature_offset;
 
 #ifndef DUCKDB_NO_THREADS
 		std::vector<std::thread> threads;
 		threads.reserve(numChunks);
 		for (idx_t i = 0; i < numChunks; i++) {
-			threads.emplace_back(ComputeSHA256, handle.get(), splits[i], splits[i + 1], &chunks[i]);
+			threads.emplace_back(ComputeSHA256FileSegment, handle.get(), splits[i], splits[i + 1], &hash_chunks[i]);
 		}
 
 		for (auto &thread : threads) {
@@ -106,26 +110,26 @@ bool ExtensionHelper::TryInitialLoad(DBConfig &config, FileOpener *opener, const
 		}
 #else
 		for (idx_t i = 0; i < numChunks; i++) {
-			ComputeSHA256(handle.get(), splits[i], splits[i + 1], &chunks[i]);
+			ComputeSHA256FileSegment(handle.get(), splits[i], splits[i + 1], &hash_chunks[i]);
 		}
 #endif // DUCKDB_NO_THREADS
 
-		string file_content;
-		file_content.reserve(256 * numChunks);
+		string hash_concatenation;
+		hash_concatenation.reserve(32 * numChunks); // 256 bits -> 32 bytes per chunk
 
-		for (auto &chunk : chunks) {
-			file_content += chunk;
+		for (auto &hash_chunk : hash_chunks) {
+			hash_concatenation += hash_chunk;
 		}
 
-		string hash;
-		ComputeSHA256(handle.get(), 0, file_content.size(), &hash);
+		string two_level_hash;
+		ComputeSHA256String(hash_concatenation, &two_level_hash);
 
 		// TODO maybe we should do a stream read / hash update here
 		handle->Read((void *)signature.data(), signature.size(), signature_offset);
 
 		bool any_valid = false;
 		for (auto &key : ExtensionHelper::GetPublicKeys()) {
-			if (duckdb_mbedtls::MbedTlsWrapper::IsValidSha256Signature(key, signature, hash)) {
+			if (duckdb_mbedtls::MbedTlsWrapper::IsValidSha256Signature(key, signature, two_level_hash)) {
 				any_valid = true;
 				break;
 			}

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

@@ -359,6 +359,7 @@ void Executor::CancelTasks() {
 		}
 		pipelines.clear();
 		root_pipelines.clear();
+		to_be_rescheduled_tasks.clear();
 		events.clear();
 	}
 	WorkOnTasks();
@@ -375,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();
 }
@@ -400,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();
 			}
@@ -444,6 +479,7 @@ void Executor::Reset() {
 	exceptions.clear();
 	pipelines.clear();
 	events.clear();
+	to_be_rescheduled_tasks.clear();
 	execution_result = PendingExecutionResult::RESULT_NOT_READY;
 }
 

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);

package/src/duckdb/src/parallel/interrupt.cpp

@@ -0,0 +1,57 @@
+#include "duckdb/parallel/interrupt.hpp"
+#include "duckdb/execution/executor.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/common/atomic.hpp"
+#include "duckdb/common/mutex.hpp"
+#include <condition_variable>
+
+namespace duckdb {
+
+InterruptState::InterruptState() : mode(InterruptMode::NO_INTERRUPTS) {
+}
+InterruptState::InterruptState(weak_ptr<Task> task) : mode(InterruptMode::TASK), current_task(std::move(task)) {
+}
+InterruptState::InterruptState(weak_ptr<InterruptDoneSignalState> signal_state_p)
+    : mode(InterruptMode::BLOCKING), signal_state(std::move(signal_state_p)) {
+}
+
+void InterruptState::Callback() const {
+	if (mode == InterruptMode::TASK) {
+		auto task = current_task.lock();
+
+		if (!task) {
+			return;
+		}
+
+		task->Reschedule();
+	} else if (mode == InterruptMode::BLOCKING) {
+		auto signal_state_l = signal_state.lock();
+
+		if (!signal_state_l) {
+			return;
+		}
+
+		// Signal the caller, who is currently blocked
+		signal_state_l->Signal();
+	} else {
+		throw InternalException("Callback made on InterruptState without valid interrupt mode specified");
+	}
+}
+
+void InterruptDoneSignalState::Signal() {
+	{
+		unique_lock<mutex> lck {lock};
+		done = true;
+	}
+	cv.notify_all();
+}
+
+void InterruptDoneSignalState::Await() {
+	std::unique_lock<std::mutex> lck(lock);
+	cv.wait(lck, [&]() { return done; });
+
+	// Reset after signal received
+	done = false;
+}
+
+} // namespace duckdb

package/src/duckdb/src/parallel/pipeline.cpp

@@ -12,7 +12,6 @@
 #include "duckdb/parallel/pipeline_event.hpp"
 #include "duckdb/parallel/pipeline_executor.hpp"
 #include "duckdb/parallel/task_scheduler.hpp"
-#include "duckdb/parallel/thread_context.hpp"
 
 namespace duckdb {
 
@@ -33,14 +32,32 @@ public:
 		if (!pipeline_executor) {
 			pipeline_executor = make_uniq<PipelineExecutor>(pipeline.GetClientContext(), pipeline);
 		}
+
+		pipeline_executor->SetTaskForInterrupts(shared_from_this());
+
 		if (mode == TaskExecutionMode::PROCESS_PARTIAL) {
-			bool finished = pipeline_executor->Execute(PARTIAL_CHUNK_COUNT);
-			if (!finished) {
+			auto res = pipeline_executor->Execute(PARTIAL_CHUNK_COUNT);
+
+			switch (res) {
+			case PipelineExecuteResult::NOT_FINISHED:
 				return TaskExecutionResult::TASK_NOT_FINISHED;
+			case PipelineExecuteResult::INTERRUPTED:
+				return TaskExecutionResult::TASK_BLOCKED;
+			case PipelineExecuteResult::FINISHED:
+				break;
 			}
 		} else {
-			pipeline_executor->Execute();
+			auto res = pipeline_executor->Execute();
+			switch (res) {
+			case PipelineExecuteResult::NOT_FINISHED:
+				throw InternalException("Execute without limit should not return NOT_FINISHED");
+			case PipelineExecuteResult::INTERRUPTED:
+				return TaskExecutionResult::TASK_BLOCKED;
+			case PipelineExecuteResult::FINISHED:
+				break;
+			}
 		}
+
 		event->FinishTask();
 		pipeline_executor.reset();
 		return TaskExecutionResult::TASK_FINISHED;
@@ -68,7 +85,7 @@ bool Pipeline::GetProgress(double &current_percentage, idx_t &source_cardinality
 }
 
 void Pipeline::ScheduleSequentialTask(shared_ptr<Event> &event) {
-	vector<unique_ptr<Task>> tasks;
+	vector<shared_ptr<Task>> tasks;
 	tasks.push_back(make_uniq<PipelineTask>(*this, event));
 	event->SetTasks(std::move(tasks));
 }
@@ -149,7 +166,7 @@ bool Pipeline::LaunchScanTasks(shared_ptr<Event> &event, idx_t max_threads) {
 	}
 
 	// launch a task for every thread
-	vector<unique_ptr<Task>> tasks;
+	vector<shared_ptr<Task>> tasks;
 	for (idx_t i = 0; i < max_threads; i++) {
 		tasks.push_back(make_uniq<PipelineTask>(*this, event));
 	}
@@ -265,6 +282,32 @@ vector<const_reference<PhysicalOperator>> Pipeline::GetOperators() const {
 	return result;
 }
 
+void Pipeline::ClearSource() {
+	source_state.reset();
+	batch_indexes.clear();
+}
+
+idx_t Pipeline::RegisterNewBatchIndex() {
+	lock_guard<mutex> l(batch_lock);
+	idx_t minimum = batch_indexes.empty() ? base_batch_index : *batch_indexes.begin();
+	batch_indexes.insert(minimum);
+	return minimum;
+}
+
+idx_t Pipeline::UpdateBatchIndex(idx_t old_index, idx_t new_index) {
+	lock_guard<mutex> l(batch_lock);
+	if (new_index < *batch_indexes.begin()) {
+		throw InternalException("Processing batch index %llu, but previous min batch index was %llu", new_index,
+		                        *batch_indexes.begin());
+	}
+	auto entry = batch_indexes.find(old_index);
+	if (entry == batch_indexes.end()) {
+		throw InternalException("Batch index %llu was not found in set of active batch indexes", old_index);
+	}
+	batch_indexes.erase(entry);
+	batch_indexes.insert(new_index);
+	return *batch_indexes.begin();
+}
 //===--------------------------------------------------------------------===//
 // Pipeline Build State
 //===--------------------------------------------------------------------===//