working kill export, update next idx upon lock and lock as many parts as one can at once

arthurpassos · arthurpassos · commit 2fbf3955f662 · 2025-10-22T09:42:25.000-03:00
diff --git a/src/Storages/MergeTree/ExportPartTask.cpp b/src/Storages/MergeTree/ExportPartTask.cpp
@@ -0,0 +1,247 @@
+#include <Storages/MergeTree/ExportPartTask.h>
+#include <Storages/MergeTree/MergeTreeSequentialSource.h>
+#include <Storages/MergeTree/MergeTreeData.h>
+#include <Interpreters/Context.h>
+#include <Interpreters/DatabaseCatalog.h>
+#include <Core/Settings.h>
+#include <Interpreters/ExpressionActions.h>
+#include <Processors/Executors/CompletedPipelineExecutor.h>
+#include <Processors/QueryPlan/BuildQueryPipelineSettings.h>
+#include <Processors/QueryPlan/Optimizations/QueryPlanOptimizationSettings.h>
+#include <Processors/QueryPlan/QueryPlan.h>
+#include <QueryPipeline/QueryPipelineBuilder.h>
+#include <Common/ProfileEventsScope.h>
+
+namespace ProfileEvents
+{
+    extern const Event PartsExportDuplicated;
+    extern const Event PartsExportFailures;
+    extern const Event PartsExports;
+    extern const Event PartsExportTotalMilliseconds;
+}
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int UNKNOWN_TABLE;
+    extern const int FILE_ALREADY_EXISTS;
+    extern const int LOGICAL_ERROR;
+}
+
+namespace Setting
+{
+    extern const SettingsUInt64 min_bytes_to_use_direct_io;
+}
+
+ExportPartTask::ExportPartTask(MergeTreeData & storage_, const MergeTreePartExportManifest & manifest_, ContextPtr context_)
+    : storage(storage_),
+    manifest(manifest_),
+    local_context(context_)
+{
+}
+
+bool ExportPartTask::executeStep()
+{
+    auto metadata_snapshot = storage.getInMemoryMetadataPtr();
+    Names columns_to_read = metadata_snapshot->getColumns().getNamesOfPhysical();
+    StorageSnapshotPtr storage_snapshot = storage.getStorageSnapshot(metadata_snapshot, local_context);
+
+    MergeTreeSequentialSourceType read_type = MergeTreeSequentialSourceType::Export;
+
+    NamesAndTypesList partition_columns;
+    if (metadata_snapshot->hasPartitionKey())
+    {
+        const auto & partition_key = metadata_snapshot->getPartitionKey();
+        if (!partition_key.column_names.empty())
+            partition_columns = partition_key.expression->getRequiredColumnsWithTypes();
+    }
+
+    auto block_with_partition_values = manifest.data_part->partition.getBlockWithPartitionValues(partition_columns);
+
+    auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest.destination_storage_id, local_context);
+    if (!destination_storage)
+    {
+        std::lock_guard inner_lock(storage.export_manifests_mutex);
+
+        const auto destination_storage_id_name = manifest.destination_storage_id.getNameForLogs();
+        storage.export_manifests.erase(manifest);
+        throw Exception(ErrorCodes::UNKNOWN_TABLE, "Failed to reconstruct destination storage: {}", destination_storage_id_name);
+    }
+
+    SinkToStoragePtr sink;
+    std::string destination_file_path;
+
+    try
+    {
+        auto context_copy = Context::createCopy(local_context);
+        context_copy->setSetting("output_format_parallel_formatting", manifest.parallel_formatting);
+
+        sink = destination_storage->import(
+            manifest.data_part->name + "_" + manifest.data_part->checksums.getTotalChecksumHex(),
+            block_with_partition_values,
+            destination_file_path,
+            manifest.overwrite_file_if_exists,
+            context_copy);
+    }
+    catch (const Exception & e)
+    {
+        if (e.code() == ErrorCodes::FILE_ALREADY_EXISTS)
+        {
+            ProfileEvents::increment(ProfileEvents::PartsExportDuplicated);
+        }
+
+        ProfileEvents::incrementNoTrace(ProfileEvents::PartsExportFailures);
+
+        std::lock_guard inner_lock(storage.export_manifests_mutex);
+        storage.export_manifests.erase(manifest);
+
+        if (manifest.completion_callback)
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
+        return false;
+    }
+
+    bool apply_deleted_mask = true;
+    bool read_with_direct_io = local_context->getSettingsRef()[Setting::min_bytes_to_use_direct_io] > manifest.data_part->getBytesOnDisk();
+    bool prefetch = false;
+
+    MergeTreeData::IMutationsSnapshot::Params params
+    {
+        .metadata_version = metadata_snapshot->getMetadataVersion(),
+        .min_part_metadata_version = manifest.data_part->getMetadataVersion(),
+    };
+
+    auto mutations_snapshot = storage.getMutationsSnapshot(params);
+
+    auto alter_conversions = MergeTreeData::getAlterConversionsForPart(
+        manifest.data_part,
+        mutations_snapshot,
+        local_context);
+
+    QueryPlan plan_for_part;
+
+    createReadFromPartStep(
+        read_type,
+        plan_for_part,
+        storage,
+        storage_snapshot,
+        RangesInDataPart(manifest.data_part),
+        alter_conversions,
+        nullptr,
+        columns_to_read,
+        nullptr,
+        apply_deleted_mask,
+        std::nullopt,
+        read_with_direct_io,
+        prefetch,
+        local_context,
+        getLogger("ExportPartition"));
+
+    auto exports_list_entry = storage.getContext()->getExportsList().insert(
+        getStorageID(),
+        manifest.destination_storage_id,
+        manifest.data_part->getBytesOnDisk(),
+        manifest.data_part->name,
+        destination_file_path,
+        manifest.data_part->rows_count,
+        manifest.data_part->getBytesOnDisk(),
+        manifest.data_part->getBytesUncompressedOnDisk(),
+        manifest.create_time,
+        local_context);
+
+    ThreadGroupSwitcher switcher((*exports_list_entry)->thread_group, "");
+
+    QueryPlanOptimizationSettings optimization_settings(local_context);
+    auto pipeline_settings = BuildQueryPipelineSettings(local_context);
+    auto builder = plan_for_part.buildQueryPipeline(optimization_settings, pipeline_settings);
+
+    builder->setProgressCallback([&exports_list_entry](const Progress & progress)
+    {
+        (*exports_list_entry)->bytes_read_uncompressed += progress.read_bytes;
+        (*exports_list_entry)->rows_read += progress.read_rows;
+        (*exports_list_entry)->elapsed = (*exports_list_entry)->watch.elapsedSeconds();
+    });
+
+    pipeline = QueryPipelineBuilder::getPipeline(std::move(*builder));
+
+    pipeline.complete(sink);
+
+    try
+    {
+        CompletedPipelineExecutor exec(pipeline);
+        exec.execute();
+
+        std::lock_guard inner_lock(storage.export_manifests_mutex);
+        storage.writePartLog(
+            PartLogElement::Type::EXPORT_PART,
+            {},
+            static_cast<UInt64>((*exports_list_entry)->elapsed * 1000000000),
+            manifest.data_part->name,
+            manifest.data_part,
+            {manifest.data_part},
+            nullptr,
+            nullptr,
+            exports_list_entry.get());
+
+        storage.export_manifests.erase(manifest);
+
+        ProfileEvents::increment(ProfileEvents::PartsExports);
+        ProfileEvents::increment(ProfileEvents::PartsExportTotalMilliseconds, static_cast<UInt64>((*exports_list_entry)->elapsed * 1000));
+
+        if (manifest.completion_callback)
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createSuccess(destination_file_path));
+    }
+    catch (const Exception & e)
+    {
+        tryLogCurrentException(__PRETTY_FUNCTION__, fmt::format("while exporting the part {}. User should retry.", manifest.data_part->name));
+
+        ProfileEvents::increment(ProfileEvents::PartsExportFailures);
+
+        std::lock_guard inner_lock(storage.export_manifests_mutex);
+        storage.writePartLog(
+            PartLogElement::Type::EXPORT_PART,
+            ExecutionStatus::fromCurrentException("", true),
+            static_cast<UInt64>((*exports_list_entry)->elapsed * 1000000000),
+            manifest.data_part->name,
+            manifest.data_part,
+            {manifest.data_part},
+            nullptr,
+            nullptr,
+            exports_list_entry.get());
+
+        storage.export_manifests.erase(manifest);
+
+        if (manifest.completion_callback)
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
+
+        throw;
+    }
+    return false;
+}
+
+void ExportPartTask::cancel() noexcept
+{
+    pipeline.cancel();
+}
+
+void ExportPartTask::onCompleted()
+{
+}
+
+StorageID ExportPartTask::getStorageID() const
+{
+    return manifest.destination_storage_id;
+}
+
+Priority ExportPartTask::getPriority() const
+{
+    return Priority{};
+}
+
+String ExportPartTask::getQueryId() const
+{
+    return manifest.query_id;
+}
+
+}
diff --git a/src/Storages/MergeTree/ExportPartTask.h b/src/Storages/MergeTree/ExportPartTask.h
@@ -0,0 +1,32 @@
+#pragma once
+
+#include <Storages/MergeTree/IExecutableTask.h>
+#include <Storages/MergeTree/MergeTreePartExportManifest.h>
+#include <Storages/MergeTree/MergeTreeData.h>
+
+namespace DB
+{
+
+class ExportPartTask : public IExecutableTask
+{
+public:
+    explicit ExportPartTask(
+        MergeTreeData & storage_,
+        const MergeTreePartExportManifest & manifest_,
+        ContextPtr context_);
+    bool executeStep() override;
+    void onCompleted() override;
+    StorageID getStorageID() const override;
+    Priority getPriority() const override;
+    String getQueryId() const override;
+
+    void cancel() noexcept override;
+
+private:
+    MergeTreeData & storage;
+    MergeTreePartExportManifest manifest;
+    ContextPtr local_context;
+    QueryPipeline pipeline;
+};
+
+}
diff --git a/src/Storages/MergeTree/MergeTreeData.cpp b/src/Storages/MergeTree/MergeTreeData.cpp
@@ -24,6 +24,7 @@
 #include <Common/scope_guard_safe.h>
 #include <Common/typeid_cast.h>
 #include <Common/thread_local_rng.h>
+#include "Storages/MergeTree/ExportPartTask.h"
 #include <Processors/Executors/CompletedPipelineExecutor.h>
 #include <Processors/QueryPlan/Optimizations/QueryPlanOptimizationSettings.h>
 #include <Storages/MergeTree/MergeTreeSequentialSource.h>
@@ -6098,9 +6099,6 @@ void MergeTreeData::exportPartToTableImpl(
 
     pipeline.complete(sink);
 
-    /// oh boy, is there another way?
-    manifest.pipeline = &pipeline;
-
     try
     {
         CompletedPipelineExecutor exec(pipeline);
@@ -6157,18 +6155,16 @@ void MergeTreeData::killExportPart(const String & query_id)
 {
     std::lock_guard lock(export_manifests_mutex);
 
-    const auto it = std::find_if(export_manifests.begin(), export_manifests.end(), [&](const auto & manifest)
+    std::erase_if(export_manifests, [&](const auto & manifest)
     {
-        return manifest.query_id == query_id;
+        if (manifest.query_id == query_id)
+        {
+            if (manifest.task)
+                manifest.task->cancel();
+            return true;
+        }
+        return false;
     });
-
-    if (it == export_manifests.end())
-        return;
-
-    if (it->pipeline)
-        it->pipeline->cancel();
-
-    export_manifests.erase(it);
 }
 
 void MergeTreeData::movePartitionToShard(const ASTPtr & /*partition*/, bool /*move_part*/, const String & /*to*/, ContextPtr /*query_context*/)
@@ -8903,17 +8899,18 @@ bool MergeTreeData::scheduleDataMovingJob(BackgroundJobsAssignee & assignee)
             continue;
         }
 
-        manifest.in_progress = assignee.scheduleMoveTask(std::make_shared<ExecutableLambdaAdapter>(
-            [this, &manifest] () mutable {
-                /// TODO arthur fix this: I need to be able to modify the real manifest
-                /// but grabbing it by reference is causing problems
-                exportPartToTableImpl(manifest, getContext());
-                return true;
-            },
-            moves_assignee_trigger,
-            getStorageID()));
+        auto task = std::make_shared<ExportPartTask>(*this, manifest, getContext());
+
+        manifest.in_progress = assignee.scheduleMoveTask(task);
+
+        if (!manifest.in_progress)
+        {
+            continue;
+        }
 
-        return manifest.in_progress;
+        manifest.task = task;
+
+        return true;
     }
 
     return false;
diff --git a/src/Storages/MergeTree/MergeTreeData.h b/src/Storages/MergeTree/MergeTreeData.h
@@ -1258,6 +1258,7 @@ class MergeTreeData : public IStorage, public WithMutableContext
     friend class MergeTask;
     friend class IPartMetadataManager;
     friend class IMergedBlockOutputStream; // for access to log
+    friend class ExportPartTask;
 
     bool require_part_metadata;
 
diff --git a/src/Storages/MergeTree/MergeTreePartExportManifest.h b/src/Storages/MergeTree/MergeTreePartExportManifest.h
@@ -1,10 +1,14 @@
+#pragma once
+
 #include <Interpreters/StorageID.h>
 #include <Storages/MergeTree/IMergeTreeDataPart.h>
-#include "QueryPipeline/QueryPipeline.h"
+#include <QueryPipeline/QueryPipeline.h>
 
 namespace DB
 {
 
+class ExportPartTask;
+
 struct MergeTreePartExportManifest
 {
     using DataPartPtr = std::shared_ptr<const IMergeTreeDataPart>;
@@ -58,8 +62,7 @@ struct MergeTreePartExportManifest
 
     time_t create_time;
     mutable bool in_progress = false;
-    /// Used for killing the export
-    mutable QueryPipeline * pipeline = nullptr;
+    mutable std::shared_ptr<ExportPartTask> task = nullptr;
 
     bool operator<(const MergeTreePartExportManifest & rhs) const 
     {
diff --git a/src/Storages/StorageReplicatedMergeTree.cpp b/src/Storages/StorageReplicatedMergeTree.cpp
