{

explicit TaskParameters(const QueryPlanParameters & parameters_)

: parameters(parameters_)

{

}

Field getParameter(const String & name) const override

{

auto it = parameters.parameters.find(name);

if (it == parameters.parameters.end())

throw Exception(ErrorCodes::LOGICAL_ERROR, "Parameter {} not found", name);

return it->second;

}

const QueryPlanParameters parameters;

{

TemporaryFilesInObjectStorage(ObjectStoragePtr object_storage_, const String & object_storage_path_,

const Strings & input_temporary_files_, const Strings & output_temporary_files_)

: object_storage(std::move(object_storage_))

, object_storage_path(object_storage_path_)

, input_temporary_files(input_temporary_files_.begin(), input_temporary_files_.end())

, output_temporary_files(output_temporary_files_.begin(), output_temporary_files_.end())

{

}

WriteBuffer & getTemporaryFileForWriting(const String & file_name) override

{

LOG_DEBUG(logger, "Writing to temporary file '{}'", file_name);

if (!output_temporary_files.contains(file_name))

throw Exception(ErrorCodes::LOGICAL_ERROR, "Unexpected output temporary file requested: '{}'", file_name);

StoredObject object(object_storage_path + "/" + file_name, file_name);

write_buffers.emplace_back(object_storage->writeObject(object, WriteMode::Rewrite));

return *write_buffers.back();

}

std::unique_ptr<ReadBuffer> getTemporaryFileForReading(const String & file_name) override

{

LOG_TRACE(logger, "Reading from temporary file '{}'", file_name);

if (!input_temporary_files.contains(file_name))

throw Exception(ErrorCodes::LOGICAL_ERROR, "Unexpected input temporary file requested: '{}'", file_name);

StoredObject object(object_storage_path + "/" + file_name, file_name);

return object_storage->readObject(object, {});

}

ObjectStoragePtr object_storage;

const String object_storage_path;

const UnorderedSetWithMemoryTracking<String> input_temporary_files;

const UnorderedSetWithMemoryTracking<String> output_temporary_files;

VectorWithMemoryTracking<std::unique_ptr<WriteBuffer>> write_buffers;

LoggerPtr logger = getLogger("TemporaryFilesInObjectStorage");

{

explicit ExchangeViaTemporaryFiles(TemporaryFileLookupPtr temporary_files_)

: temporary_files(std::move(temporary_files_))

{

}

std::shared_ptr<ISink> createSink(SharedHeader input_header, const ExchangeStreamId & exchange_stream_id, bool input_is_serialized) override

{

if (!temporary_files)

throw Exception(

ErrorCodes::SUPPORT_IS_DISABLED,

"Object storage for Persisted exchanges is not configured, exchange stream id: {}",

exchange_stream_id.toString());

if (input_is_serialized)

throw Exception(ErrorCodes::LOGICAL_ERROR, "Persisted exchange {} has no serializer, its sink takes data chunks", exchange_stream_id.toString());

auto file_name = exchange_stream_id.toString();

return std::make_shared<NativeCompressedSink>(input_header, temporary_files->getTemporaryFileForWriting(file_name), file_name);

}

std::shared_ptr<ISource> createSource(SharedHeader output_header, const ExchangeStreamId & exchange_stream_id) override

{

if (!temporary_files)

throw Exception(

ErrorCodes::SUPPORT_IS_DISABLED,

"Object storage for Persisted exchanges is not configured, exchange stream id: {}",

exchange_stream_id.toString());

auto file_name = exchange_stream_id.toString();

std::unique_ptr<QueryPipelineBuilder> pipeline_ptr = std::make_unique<QueryPipelineBuilder>();

return std::make_shared<NativeCompressedSource>(output_header, temporary_files->getTemporaryFileForReading(file_name), file_name);

}

TemporaryFileLookupPtr temporary_files;

{

explicit InMemoryExchange(const String & name_)

: name(name_)

{

}

/// Appends a data chunk. Throws if the exchange is cancelled, so the producer stops early.

/// Drops the chunk if the reader is detached.

void appendChunk(Chunk chunk)

{

LOG_TEST(log, "Appending chunk to exchange '{}', rows {}", name, chunk.getNumRows());

std::lock_guard lock(mutex);

if (cancelled)

throwCancelled();

if (reader_detached)

return;

chunks.emplace_back(std::move(chunk));

has_data.notify_one();

}

/// Appends the end-of-data marker. Allowed after a cancel: the queued stream is complete,

/// so a draining consumer may still use it.

void finish()

{

LOG_TEST(log, "Finishing exchange '{}'", name);

std::lock_guard lock(mutex);

chunks.emplace_back(Chunk{});

has_data.notify_one();

}

/// Wake any waiter. Consumers get `reason_` (or a generic cancellation error) instead of

/// end-of-data, so an aborted stream cannot pass for a complete one.

void cancel(std::exception_ptr reason_)

{

std::lock_guard lock(mutex);

cancelled = true;

if (!reason)

reason = reason_;

has_data.notify_all();

}

/// The reader stopped and does not need more data, e.g. its pipeline finished early.

/// Wakes a blocked `getChunk`; chunks appended after this are dropped. Unlike `cancel`,

/// this is not a failure: the producer stops this stream but finishes successfully.

void detachReader()

{

std::lock_guard lock(mutex);

reader_detached = true;

has_data.notify_all();

}

/// True after `detachReader`: the reader is gone and appended chunks are dropped.

bool isReaderDetached()

{

std::lock_guard lock(mutex);

return reader_detached;

}

/// Waits up to `timeout` for a chunk. Returns std::nullopt if nothing arrived in time.

/// An empty chunk is the producer's end-of-data marker. Chunks queued before a cancel are

/// still handed out; once a cancelled queue is empty, throws the cancellation reason.

std::optional<Chunk> getChunk(std::chrono::milliseconds timeout)

{

LOG_TEST(log, "Waiting for chunk from exchange '{}'", name);

Chunk chunk;

{

std::unique_lock lock(mutex);

if (!has_data.wait_for(lock, timeout, [this] { return !chunks.empty() || cancelled || reader_detached; }))

return std::nullopt;

/// The reader is stopping and does not need more data.

if (reader_detached)

return std::nullopt;

/// The wait ended with an empty queue only when cancelled.

if (chunks.empty())

throwCancelled();

chunk = std::move(chunks.front());

chunks.pop_front();

}

LOG_TEST(log, "Got chunk from exchange '{}', rows {}", name, chunk.getNumRows());

return chunk;

}

[[noreturn]] void throwCancelled() const

{

if (reason)

std::rethrow_exception(reason);

throw Exception(ErrorCodes::QUERY_WAS_CANCELLED,

"Distributed query was cancelled before exchange '{}' transferred all data", name);

}

LoggerPtr log = getLogger("InMemoryExchange");

String name;

std::mutex mutex;

std::condition_variable has_data;

DequeWithMemoryTracking<Chunk> chunks;

bool cancelled = false;

bool reader_detached = false;

/// The first failure passed to `cancel`.

std::exception_ptr reason;

{

InMemoryExchangePtr getExchange(const String & query_id, const String & exchange_id)

{

std::lock_guard lock(mutex);

auto & element = exchanges_by_query_id[query_id][exchange_id];

if (!element)

{

element = std::make_shared<InMemoryExchange>(exchange_id);

/// A task built concurrently with the cancellation may look up its exchange after

/// cancelQuery already ran; hand it out pre-cancelled so its reads fail right away.

if (auto cancelled_it = cancelled_queries.find(query_id); cancelled_it != cancelled_queries.end())

element->cancel(cancelled_it->second);

}

return element;

}

/// Cancel every exchange of the query so waiting tasks stop, reporting `failure` as the

/// reason (null for a plain cancellation). The exchanges stay in the registry so a result

/// reader that looks one up afterwards still finds the produced chunks and their end-of-data

/// marker; removeQuery drops them once the whole query pipeline is destroyed.

void cancelQuery(const String & query_id, std::exception_ptr failure)

{

std::lock_guard lock(mutex);

auto [cancelled_it, inserted] = cancelled_queries.emplace(query_id, failure);

/// Keep the first failure, but let a later one replace a plain cancellation.

if (!inserted && !cancelled_it->second && failure)

cancelled_it->second = failure;

auto it = exchanges_by_query_id.find(query_id);

if (it == exchanges_by_query_id.end())

return;

for (auto & [_, exchange] : it->second)

exchange->cancel(failure);

}

/// Drop the query's exchanges from the registry. Called when the query pipeline is destroyed.

void removeQuery(const String & query_id)

{

std::lock_guard lock(mutex);

exchanges_by_query_id.erase(query_id);

cancelled_queries.erase(query_id);

}

static std::shared_ptr<InMemoryExchanges> instance()

{

static std::shared_ptr<InMemoryExchanges> self = std::make_shared<InMemoryExchanges>();

return self;

}

using InMemoryExchangeMap = UnorderedMapWithMemoryTracking<String, InMemoryExchangePtr>;

UnorderedMapWithMemoryTracking<String, InMemoryExchangeMap> exchanges_by_query_id TSA_GUARDED_BY(mutex);

/// Cancelled query -> its root failure (null for a plain cancellation).

UnorderedMapWithMemoryTracking<String, std::exception_ptr> cancelled_queries TSA_GUARDED_BY(mutex);

std::mutex mutex;

{

explicit ExchangeViaChunks(const String & query_id_)

: query_id(query_id_)

{

}

std::shared_ptr<ISink> createSink(SharedHeader input_header, const ExchangeStreamId & exchange_stream_id, bool input_is_serialized) override

{

if (input_is_serialized)

throw Exception(ErrorCodes::LOGICAL_ERROR, "In-memory exchange {} has no serializer, its sink takes data chunks", exchange_stream_id.toString());

auto file_name = exchange_stream_id.toString();

auto exchange = InMemoryExchanges::instance()->getExchange(query_id, file_name);

return std::make_shared<SinkFromInMemoryExchange>(input_header, exchange);

}

std::shared_ptr<ISource> createSource(SharedHeader output_header, const ExchangeStreamId & exchange_stream_id) override

{

auto file_name = exchange_stream_id.toString();

auto exchange = InMemoryExchanges::instance()->getExchange(query_id, file_name);

return std::make_shared<SourceFromInMemoryExchange>(output_header, exchange);

}

class SinkFromInMemoryExchange final : public ISink

{

public:

SinkFromInMemoryExchange(SharedHeader header_, InMemoryExchangePtr exchange_)

: ISink(header_)

, exchange(std::move(exchange_))

{

}

String getName() const override { return "SinkFromInMemoryExchange"; }

Status prepare() override

{

/// The reader detached, so appended chunks would be dropped. Close the input so the

/// stop propagates to the upstream stages; without this they would keep computing

/// data that nobody reads.

if (exchange->isReaderDetached())

{

input.close();

return Status::Finished;

}

return ISink::prepare();

}

void consume(Chunk chunk) override

{

/// Zero-row chunks are scheduling ticks from an upstream `SourceFromInMemoryExchange`;

/// forwarding them would only grow the queue and wake the consumer for nothing.

if (!chunk.hasRows() && chunk.getChunkInfos().empty())

return;

exchange->appendChunk(std::move(chunk));

}

void onFinish() override

{

exchange->finish();

}

private:

InMemoryExchangePtr exchange;

};

class SourceFromInMemoryExchange final : public ISource

{

public:

SourceFromInMemoryExchange(SharedHeader header_, InMemoryExchangePtr exchange_)

: ISource(header_)

, exchange(std::move(exchange_))

{

}

String getName() const override { return "SourceFromInMemoryExchange"; }

Status prepare() override

{

/// The output port is closed, for example by a satisfied LIMIT downstream. Tell the

/// exchange, so the producer's sink stops instead of queueing chunks that nobody

/// reads. `onCancel` covers the cancellation path in the same way.

if (!detach_notified && getPort().isFinished())

{

detach_notified = true;

exchange->detachReader();

}

return ISource::prepare();

}

std::optional<Chunk> tryGenerate() override

{

/// This source's own pipeline is being torn down: stop quietly, its output is

/// discarded anyway.

if (isCancelled())

return std::nullopt;

/// A processor must not block the pipeline thread, so wait with a timeout. On the

/// timeout, push a chunk with no rows: a source that yields without producing output

/// stays ready and gets rescheduled at once, monopolizing the thread, while pushed

/// output makes the port full and lets other processors run.

auto chunk = exchange->getChunk(waitTimeout());

if (!chunk)

return Chunk(getPort().getHeader().cloneEmptyColumns(), 0);

if (chunk->empty())

return std::nullopt; /// End-of-data marker.

return chunk;

}

/// Wake the timed wait so the source stops right away instead of on its next poll.

void onCancel() noexcept override

{

exchange->detachReader();

}

private:

/// A stream with no columns (case of `SELECT count()`) cannot fill the output port, so this

/// source polls instead of yielding - keep its wait short so other processors run sooner.

std::chrono::milliseconds waitTimeout() const

{

return getPort().getHeader().empty() ? std::chrono::milliseconds(1) : std::chrono::milliseconds(10);

}

InMemoryExchangePtr exchange;

bool detach_notified = false;

};

const String query_id;

{

explicit SkipZeroRowChunksTransform(SharedHeader header_)

: ISimpleTransform(header_, header_, /*skip_empty_chunks_=*/ true)

{

}

String getName() const override { return "SkipZeroRowChunksTransform"; }

void transform(Chunk & chunk) override

{

/// Keep zero-row chunks that carry chunk infos: they are not ticks.

if (!chunk.hasRows() && chunk.getChunkInfos().empty())

chunk = Chunk();

}

{

AllKindsExchangeLookup(

const ExchangeDescriptions & exchanges_,

ExchangeLookupPtr persistent_exchange_lookup_,

ExchangeLookupPtr streaming_exchange_lookup_)

: exchanges(exchanges_)

, persistent_exchange_lookup(std::move(persistent_exchange_lookup_))

, streaming_exchange_lookup(std::move(streaming_exchange_lookup_))

{

}

std::shared_ptr<ISink> createSink(SharedHeader input_header, const ExchangeStreamId & exchange_stream_id, bool input_is_serialized) override

{

return lookupFor(exchange_stream_id.exchange_id).createSink(std::move(input_header), exchange_stream_id, input_is_serialized);

}

std::shared_ptr<ISource> createSource(SharedHeader output_header, const ExchangeStreamId & exchange_stream_id) override

{

return lookupFor(exchange_stream_id.exchange_id).createSource(std::move(output_header), exchange_stream_id);

}

std::shared_ptr<IProcessor> createSerializer(SharedHeader input_header, const String & exchange_id) override

{

return lookupFor(exchange_id).createSerializer(std::move(input_header), exchange_id);

}

IExchangeLookup & lookupFor(const String & exchange_id) const

{

auto it = exchanges.find(exchange_id);

if (it == exchanges.end())

throw Exception(ErrorCodes::LOGICAL_ERROR, "Unknown exchange '{}'", exchange_id);

if (it->second.kind == ExchangeDescription::Kind::Persisted)

return *persistent_exchange_lookup;

if (it->second.kind == ExchangeDescription::Kind::Streaming)

return *streaming_exchange_lookup;

throw Exception(ErrorCodes::LOGICAL_ERROR, "Unknown exchange kind '{}'", static_cast<int>(it->second.kind));

}

const ExchangeDescriptions exchanges;

ExchangeLookupPtr persistent_exchange_lookup;

ExchangeLookupPtr streaming_exchange_lookup;

{

TemporaryFilesInObjectStorageCleaner(ObjectStoragePtr object_storage_, const String & object_storage_path_,

const Strings & temporary_files_)

: object_storage(std::move(object_storage_))

, object_storage_path(object_storage_path_)

, temporary_files(temporary_files_.begin(), temporary_files_.end())

{

}

~TemporaryFilesInObjectStorageCleaner() override

{

/// TODO: add them to some background cleanup queue to avoid garbage in case of exceptions?

try

{

cleanup();

}

catch (...)

{

tryLogCurrentException(__PRETTY_FUNCTION__);

}

}

void cleanup()

{

StoredObjects all_objects;

for (const auto & file_name : temporary_files)

{

StoredObject object(object_storage_path + "/" + file_name, file_name);

all_objects.emplace_back(std::move(object));

}

LOG_TRACE(getLogger("TemporaryFilesInObjectStorageCleaner"), "Removing temporary files at path {} : [{}]",

object_storage_path, fmt::join(temporary_files, ", "));

object_storage->removeObjectsIfExist(all_objects);

}

ObjectStoragePtr object_storage;

const String object_storage_path;

const UnorderedSetWithMemoryTracking<String> temporary_files;

{

explicit InMemoryExchangesCleaner(String query_id_) : query_id(std::move(query_id_)) {}

~InMemoryExchangesCleaner() override

{

try

{

InMemoryExchanges::instance()->removeQuery(query_id);

}

catch (...)

{

tryLogCurrentException(__PRETTY_FUNCTION__);

}

}

const String query_id;

const Strings & input_temporary_files_, const Strings & output_temporary_files_)

{

if (!object_storage_)

return nullptr;

return std::make_shared<TemporaryFilesInObjectStorage>(object_storage_, object_storage_path_, input_temporary_files_, output_temporary_files_);

const String & query_id,

const ExchangeDescriptions & exchanges_,

const ExchangeStreamSources & exchange_stream_sources,

TemporaryFileLookupPtr temporary_files_,

ContextPtr context,

bool execute_locally,

DistributedQueryCancellationPtr cancellation)

{

if (execute_locally)

{

LOG_DEBUG(getLogger("createExchangeLookup"), "`distributed_plan_execute_locally` setting is enabled, using in-memory queues for all exchanges");

return std::make_shared<ExchangeViaChunks>(query_id);

}

auto persisted_exchanges = std::make_shared<ExchangeViaTemporaryFiles>(temporary_files_);

bool has_streaming_exchange = false;

for (const auto & [exchange_id, exchange] : exchanges_)

if (exchange.kind == ExchangeDescription::Kind::Streaming)

{

has_streaming_exchange = true;

break;

}

/// Persisted exchanges only need the temporary-file lookup, so a plan where every exchange

/// is Persisted runs without a streaming transport (and on any platform). The streaming

/// port and lookup are required only when the plan actually contains a Streaming exchange.

if (!has_streaming_exchange)

{

UNUSED(exchange_stream_sources);

return std::make_shared<AllKindsExchangeLookup>(exchanges_, persisted_exchanges, /*streaming_exchange_lookup=*/nullptr);

}

#if defined(OS_LINUX) || defined(OS_DARWIN)

auto streaming_exchange_port = context->getConfigRef().getUInt("distributed_query.streaming_exchange_port", 0);

if (streaming_exchange_port == 0)

throw Exception(ErrorCodes::SUPPORT_IS_DISABLED,

"Streaming exchange requires `distributed_query.streaming_exchange_port` to be configured; "

"set the port, force `distributed_plan_force_exchange_kind = 'Persisted'`, or enable "

"`distributed_plan_execute_locally` for in-process testing");

if (streaming_exchange_port > 65535)

throw Exception(ErrorCodes::INVALID_CONFIG_PARAMETER,

"`distributed_query.streaming_exchange_port` must be in range 1..65535, got {}", streaming_exchange_port);

/// The listener starts only when a listen host is also configured, so streaming peers are

/// unreachable without one. Reject here instead of connecting to a listener that never started.

if (getMultipleValuesFromConfig(context->getConfigRef(), "distributed_query", "streaming_exchange_listen_host").empty())

throw Exception(ErrorCodes::SUPPORT_IS_DISABLED,

"Streaming exchange requires `distributed_query.streaming_exchange_listen_host` to be configured; "

"set it, force `distributed_plan_force_exchange_kind = 'Persisted'`, or enable "

"`distributed_plan_execute_locally` for in-process testing");

/// A task from an older initiator (version 1) ships no per-stream ports; fall back to this

/// node's configured exchange port to preserve the previous single-port behavior.

ExchangeStreamSources sources_with_ports = exchange_stream_sources;

for (auto & [stream, address] : sources_with_ports.stream_hosts)

if (address.port == 0)

address.port = static_cast<UInt16>(streaming_exchange_port);

auto streaming_exchanges = createStreamingExchangeLookup(

query_id, ExchangeConnections::instance(), sources_with_ports, std::move(cancellation));

return std::make_shared<AllKindsExchangeLookup>(exchanges_, persisted_exchanges, streaming_exchanges);

#else

UNUSED(exchange_stream_sources, context, cancellation);

throw Exception(ErrorCodes::SUPPORT_IS_DISABLED,

"Streaming exchanges are only supported on Linux and macOS; "

"use `distributed_plan_force_exchange_kind = 'Persisted'`");

#endif

{

/// A shipped set must be complete, so the overflow mode is always throw;

/// `transfer_overflow_mode = 'break'` does not apply to it.

const auto & settings = context->getSettingsRef();

SizeLimits sets_transfer_limits(

settings[Setting::max_rows_to_transfer], settings[Setting::max_bytes_to_transfer], OverflowMode::THROW);

WriteBufferFromOwnString out;

query_plan.serializeForDistributedTask(out, DBMS_QUERY_PLAN_SERIALIZATION_VERSION, sets_transfer_limits);

return out.str();

{

ReadBufferFromString in(serialized_query_plan);

/// Trusted server-to-server plan fragment: decode types without the input complexity limit.

auto plan_and_sets = QueryPlan::deserialize(in, context, 0);

return QueryPlan::makeSets(std::move(plan_and_sets), context);

const String & object_storage_path, const String & distributed_query_id, ContextMutablePtr context,

bool execute_locally, std::function<bool()> is_cancelled, ProgressCallback progress_callback)

{

Stopwatch execute_task_watch;

const auto & task = task_description.task;

std::shared_ptr<OpenTelemetry::SpanHolder> query_span = std::make_shared<OpenTelemetry::SpanHolder>(task.task_id);

auto logger = Poco::Logger::getShared("executeDistributedQuery");

/// Disable the query condition cache: its per-worker state could make workers read inconsistent

/// data for the same fragment.

context->setSetting("use_query_condition_cache", false);

Strings input_exchange_streams;

for (const auto & stream_id : task.input_exchange_streams)

input_exchange_streams.push_back(stream_id.toString());

Strings output_exchange_streams;

for (const auto & stream_id : task.output_exchange_streams)

output_exchange_streams.push_back(stream_id.toString());

LOG_TRACE(logger, "Task '{}' input exchange streams: [{}], output exchange streams: [{}]",

task.task_id, fmt::join(input_exchange_streams, ", "), fmt::join(output_exchange_streams, ", "));

#if defined(OS_LINUX) || defined(OS_DARWIN)

/// Release this task's pending streaming exchange connections on the worker when it ends. A

/// consumer that never connects (e.g. its query was cancelled) would otherwise leave them behind.

/// Only this task's output streams are dropped, so sibling tasks of the same query are unaffected.

SCOPE_EXIT_SAFE(ExchangeConnections::instance()->removePendingStreams(distributed_query_id, output_exchange_streams));

#endif

auto temporary_files = createTemporaryFilesLookup(

object_storage, object_storage_path, input_exchange_streams, output_exchange_streams);

auto pipeline_settings = BuildQueryPipelineSettings(context);

pipeline_settings.temporary_file_lookup = temporary_files;

pipeline_settings.parameter_lookup = std::make_shared<TaskParameters>(task.parameters);

pipeline_settings.exchange_lookup = createExchangeLookup(

distributed_query_id,

task_description.exchanges,

task_description.exchange_stream_sources,

temporary_files,

context,

execute_locally,

/*cancellation=*/ nullptr);

auto optimization_settings = QueryPlanOptimizationSettings(context);

/// Disable stats-driven plan-shape rewrites on the worker side: per-worker

/// stats can diverge and produce incompatible plans across workers (e.g. one

/// swaps the join sides while the others don't), breaking exchange partitioning.

optimization_settings.join_swap_table = std::make_optional(false);

optimization_settings.query_plan_optimize_join_order_limit = 0;

optimization_settings.query_plan_optimize_join_order_randomize = 0;

optimization_settings.convert_join_to_in = false;

optimization_settings.convert_outer_join_to_inner_join = false;

optimization_settings.convert_any_join_to_semi_or_anti_join = false;

optimization_settings.merge_filter_into_join_condition = false;

optimization_settings.top_k_through_join = false;

/// The fragment's read is bucketed; re-introducing the implicit count projection would count the

/// whole part per bucket. Keep it off so counts read the bucket's mark ranges.

optimization_settings.optimize_use_implicit_projections = false;

QueryPipeline pipeline;

{

QueryPlan query_plan = deserializeQueryPlan(task_description.serialized_query_plan, context);

/// A deserialized plan carries neither the thread limit nor the concurrency-control flag,

/// so both come from the query's settings.

query_plan.setMaxThreads(pipeline_settings.max_threads);

query_plan.setConcurrencyControl(context->getSettingsRef()[Setting::use_concurrency_control]);

auto builder = query_plan.buildQueryPipeline(

optimization_settings,

pipeline_settings);

pipeline = QueryPipelineBuilder::getPipeline(std::move(*builder));

}

/// No AST: this fragment is built from a serialized query plan, not parsed. The query-log

/// helpers below treat a null AST as QueryKind::Select, which is correct here.

const ASTPtr no_ast;

UInt64 query_plan_hash = sipHash64(task_description.serialized_query_plan);

auto query_log_elem = logQueryStart(

std::chrono::system_clock::now(),

context,

/*query_for_logging*/ task.task_id,

query_plan_hash,

no_ast, pipeline,

/*interpreter*/ nullptr,

/*internal*/ false,

/*log_as_internal*/ false,

/*database*/ "",

/*table*/ "",

/*async_insert*/ false);

try

{

LOG_TEST(logger, "Executing task '{}', pipeline:\n{}",

task.task_id,

[&pipeline]() -> String

{

WriteBufferFromOwnString out;

printPipeline(pipeline.getProcessors(), out);

return out.str();

}());

if (!pipeline.completed())

throw Exception(ErrorCodes::LOGICAL_ERROR, "Task pipeline must be completed");

pipeline.setProcessListElement(context->getProcessListElement());

pipeline.setProgressCallback(progress_callback ? progress_callback : context->getProgressCallback());

{

CompletedPipelineExecutor executor(pipeline);

if (is_cancelled)

executor.setCancelCallback(is_cancelled, 100);

executor.execute();

}

logQueryFinish(query_log_elem, context, no_ast, std::move(pipeline), false,

query_span, QueryResultCacheUsage::None, false, /*log_as_internal*/ false);

}

catch (...)

{

logQueryException(query_log_elem, context, execute_task_watch, no_ast, query_span, false, /*log_as_internal*/ false, true);

throw;

}

{

const auto & config = context->getConfigRef();

String config_prefix = "distributed_query.temporary_files_storage";

if (config.has(config_prefix))

return config.getString(config_prefix + ".endpoint_subpath") + unique_temp_file_path;

return unique_temp_file_path;

{

const auto & config = context->getConfigRef();

String config_prefix = "distributed_query.temporary_files_storage";

String object_storage_path = getTemporaryFilesPath(unique_temp_file_path, context);

if (config.has(config_prefix))

{

ObjectStoragePtr object_storage = ObjectStorageFactory::instance().create("distributed_query_temp_files", config, config_prefix, context, false);

return {object_storage, object_storage_path};

}

return {nullptr, object_storage_path};

{

auto [object_storage, object_storage_path] = getObjectStorageForTemporaryFiles(toString(unique_query_id), context);

/// Run each task as an independent query fragment with its own query context and thread group,

/// matching the worker path. Attaching the task context to this thread (instead of sharing the

/// initiator's) gives the task its own per-query state, such as the runtime filter lookup.

auto task_context = Context::createCopy(context);

task_context->makeQueryContext();

auto query_scope = QueryScope::create(task_context);

setThreadName(ThreadName::DISTRIBUTED_QUERY_TASK);

/// Only DistributedQueryPlanExecutorLocal reaches here, so the task always runs in-process.

doExecuteTask(task, object_storage, object_storage_path, toString(unique_query_id), std::move(task_context),

/*execute_locally=*/true, [cancellation]() -> bool { return cancellation->isCancelled(); });

{

DistributedQueryPlanExecutorLocal(const UUID & unique_query_id_, const DistributedQueryPlan & distributed_query_plan_, ContextPtr context_, DistributedQueryCancellationPtr cancellation_, StageWakeupPtr stage_wakeup_)

: DistributedQueryPlanExecutor(unique_query_id_, distributed_query_plan_, makeContextForLocalExecution(context_), std::move(cancellation_), std::move(stage_wakeup_))

{

}

~DistributedQueryPlanExecutorLocal() override

{

/// Guarantee no task thread outlives the executor even on a teardown path that did not reach

/// the driver's cleanup() (e.g. an exception during the run). The threads are always joined,

/// never detached, so the std::thread members below are never destroyed while joinable.

cleanup();

}

void cleanup() override

{

/// Cancel the query's in-memory exchanges before joining the task threads, or a task stuck in

/// InMemoryExchange::getChunk never returns. The exchanges are not removed here: the result

/// reader still drains final_result after the driver finishes; removal happens when the query

/// pipeline is destroyed (see makeInMemoryExchangesCleaner).

InMemoryExchanges::instance()->cancelQuery(toString(unique_query_id), cancellation->getFailure());

joinAllThreads();

stage_tasks.clear();

}

static ContextPtr makeContextForLocalExecution(ContextPtr ctx)

{

auto new_context = Context::createCopy(ctx);

/// We will execute tasks with local plan fragments. They should not be converted into distributed plan themselves.

new_context->setSetting("make_distributed_plan", false);

new_context->setSetting("enable_cascades_optimizer", false);

return new_context;

}

std::future<void> startTask(const DistributedQueryTaskDescription & task_description, VectorWithMemoryTracking<std::thread> & threads)

{

std::promise<void> task_promise;

std::future<void> future = task_promise.get_future();