add profile events and persist some settings in the manifest

Author: arthurpassos

src/Common/ProfileEvents.cpp

@@ -28,6 +28,10 @@
     M(AsyncInsertBytes, "Data size in bytes of asynchronous INSERT queries.", ValueType::Bytes) \
     M(AsyncInsertRows, "Number of rows inserted by asynchronous INSERT queries.", ValueType::Number) \
     M(AsyncInsertCacheHits, "Number of times a duplicate hash id has been found in asynchronous INSERT hash id cache.", ValueType::Number) \
+    M(PartsExports, "Number of successful part exports.", ValueType::Number) \
+    M(PartsExportFailures, "Number of failed part exports.", ValueType::Number) \
+    M(PartsExportDuplicated, "Number of part exports that failed because target already exists.", ValueType::Number) \
+    M(PartsExportTotalMilliseconds, "Total time spent on part export operations.", ValueType::Milliseconds) \
     M(FailedQuery, "Number of failed queries.", ValueType::Number) \
     M(FailedSelectQuery, "Same as FailedQuery, but only for SELECT queries.", ValueType::Number) \
     M(FailedInsertQuery, "Same as FailedQuery, but only for INSERT queries.", ValueType::Number) \

src/Core/Settings.cpp

@@ -6687,6 +6687,9 @@ Allows to change the behaviour of the result type of `dateTrunc` function.
 Possible values:
 - 0 - When the second argument is `DateTime64/Date32` the return type will be `DateTime64/Date32` regardless of the time unit in the first argument.
 - 1 - For `Date32` the result is always `Date`. For `DateTime64` the result is `DateTime` for time units `second` and higher.
+)", 0) \
+    DECLARE(Bool, export_merge_tree_part_overwrite_file_if_exists, false, R"(
+Overwrite file if it already exists when exporting a merge tree part
 )", 0) \
     \
     /* ####################################################### */ \

src/Storages/IStorage.h

@@ -452,10 +452,12 @@ It is currently only implemented in StorageObjectStorage.
     virtual SinkToStoragePtr import(
         const std::string & /* file_name */,
         Block & /* block_with_partition_values */,
+        std::string & /* destination_file_path */,
+        bool /* overwrite_if_exists */,
         ContextPtr /* context */)
-      { 
-        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Import is not implemented for storage {}", getName());
-      }
+    {
+      throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Import is not implemented for storage {}", getName());
+    }
 
 
     /** Writes the data to a table in distributed manner.

src/Storages/MergeTree/ExportList.cpp

@@ -8,6 +8,7 @@ ExportsListElement::ExportsListElement(
     const StorageID & destination_table_id_,
     UInt64 part_size_,
     const String & part_name_,
+    const String & target_file_name_,
     UInt64 total_rows_to_read_,
     UInt64 total_size_bytes_compressed_,
     UInt64 total_size_bytes_uncompressed_,
@@ -17,6 +18,7 @@ ExportsListElement::ExportsListElement(
 , destination_table_id(destination_table_id_)
 , part_size(part_size_)
 , part_name(part_name_)
+, destination_file_path(target_file_name_)
 , total_rows_to_read(total_rows_to_read_)
 , total_size_bytes_compressed(total_size_bytes_compressed_)
 , total_size_bytes_uncompressed(total_size_bytes_uncompressed_)
@@ -33,6 +35,7 @@ ExportInfo ExportsListElement::getInfo() const
     res.destination_database = destination_table_id.database_name;
     res.destination_table = destination_table_id.table_name;
     res.part_name = part_name;
+    res.destination_file_path = destination_file_path;
     res.rows_read = rows_read;
     res.total_rows_to_read = total_rows_to_read;
     res.total_size_bytes_compressed = total_size_bytes_compressed;

src/Storages/MergeTree/ExportList.h

@@ -23,6 +23,7 @@ struct ExportInfo
     String destination_database;
     String destination_table;
     String part_name;
+    String destination_file_path;
     UInt64 rows_read;
     UInt64 total_rows_to_read;
     UInt64 total_size_bytes_compressed;
@@ -40,6 +41,7 @@ struct ExportsListElement : private boost::noncopyable
     const StorageID destination_table_id;
     const UInt64 part_size;
     const String part_name;
+    const String destination_file_path;
     UInt64 rows_read {0};
     UInt64 total_rows_to_read {0};
     UInt64 total_size_bytes_compressed {0};
@@ -56,6 +58,7 @@ struct ExportsListElement : private boost::noncopyable
         const StorageID & destination_table_id_,
         UInt64 part_size_,
         const String & part_name_,
+        const String & destination_file_path_,
         UInt64 total_rows_to_read_,
         UInt64 total_size_bytes_compressed_,
         UInt64 total_size_bytes_uncompressed_,

src/Storages/MergeTree/MergeTreeData.cpp

@@ -113,6 +113,7 @@
 #include <thread>
 #include <unordered_set>
 #include <filesystem>
+#include <utility>
 
 #include <fmt/format.h>
 #include <Poco/Logger.h>
@@ -154,6 +155,10 @@ namespace ProfileEvents
     extern const Event LoadedDataPartsMicroseconds;
     extern const Event RestorePartsSkippedFiles;
     extern const Event RestorePartsSkippedBytes;
+    extern const Event PartsExports;
+    extern const Event PartsExportTotalMilliseconds;
+    extern const Event PartsExportFailures;
+    extern const Event PartsExportDuplicated;
 }
 
 namespace CurrentMetrics
@@ -198,6 +203,8 @@ namespace Setting
     extern const SettingsUInt64 merge_tree_storage_snapshot_sleep_ms;
     extern const SettingsBool allow_experimental_export_merge_tree_part;
     extern const SettingsUInt64 min_bytes_to_use_direct_io;
+    extern const SettingsBool export_merge_tree_part_overwrite_file_if_exists;
+    extern const SettingsBool output_format_parallel_formatting;
 }
 
 namespace MergeTreeSetting
@@ -310,6 +317,7 @@ namespace ErrorCodes
     extern const int CANNOT_FORGET_PARTITION;
     extern const int DATA_TYPE_CANNOT_BE_USED_IN_KEY;
     extern const int UNKNOWN_TABLE;
+    extern const int FILE_ALREADY_EXISTS;
 }
 
 static void checkSuspiciousIndices(const ASTFunction * index_function)
@@ -5932,9 +5940,15 @@ void MergeTreeData::exportPartToTable(const PartitionCommand & command, ContextP
                         part_name, getStorageID().getFullTableName());
 
     {
+        MergeTreeExportManifest manifest(
+            dest_storage->getStorageID(),
+            part,
+            query_context->getSettingsRef()[Setting::export_merge_tree_part_overwrite_file_if_exists],
+            query_context->getSettingsRef()[Setting::output_format_parallel_formatting]);
+
         std::lock_guard lock(export_manifests_mutex);
 
-        if (!export_manifests.emplace(dest_storage->getStorageID(), part).second)  
+        if (!export_manifests.emplace(std::move(manifest)).second)
         {
             throw Exception(ErrorCodes::ABORTED, "Data part '{}' is already being exported to table '{}'",
                             part_name, dest_storage->getStorageID().getFullTableName());
@@ -5948,19 +5962,6 @@ void MergeTreeData::exportPartToTableImpl(
     const MergeTreeExportManifest & manifest,
     ContextPtr local_context)
 {
-    auto exports_list_entry = getContext()->getExportsList().insert(
-        getStorageID(),
-        manifest.destination_storage_id,
-        manifest.data_part->getBytesOnDisk(),
-        manifest.data_part->name,
-        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, "");
-
     auto metadata_snapshot = getInMemoryMetadataPtr();
     Names columns_to_read = metadata_snapshot->getColumns().getNamesOfPhysical();
     StorageSnapshotPtr storage_snapshot = getStorageSnapshot(metadata_snapshot, local_context);
@@ -5987,16 +5988,29 @@ void MergeTreeData::exportPartToTableImpl(
         throw Exception(ErrorCodes::UNKNOWN_TABLE, "Failed to reconstruct destination storage: {}", destination_storage_id_name);
     }
 
-    auto sink = destination_storage->import(
-        manifest.data_part->name,
-        block_with_partition_values,
-        local_context);
+    SinkToStoragePtr sink;
+    std::string destination_file_path;
 
-    /// Most likely the file has already been imported, so we can just return
-    if (!sink)
+    try
     {
-        std::lock_guard inner_lock(export_manifests_mutex);
+        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,
+            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);
+        }
+
+        std::lock_guard inner_lock(export_manifests_mutex);
         export_manifests.erase(manifest);
         return;
     }
@@ -6037,6 +6051,20 @@ void MergeTreeData::exportPartToTableImpl(
         local_context,
         getLogger("ExportPartition"));
 
+    auto exports_list_entry = 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);
@@ -6070,10 +6098,15 @@ void MergeTreeData::exportPartToTableImpl(
             exports_list_entry.get());
 
         export_manifests.erase(manifest);
+
+        ProfileEvents::increment(ProfileEvents::PartsExports);
+        ProfileEvents::increment(ProfileEvents::PartsExportTotalMilliseconds, static_cast<UInt64>((*exports_list_entry)->elapsed * 1000));
     }
-    catch (const Exception &)
+    catch (...)
     {
-        tryLogCurrentException(__PRETTY_FUNCTION__, "Exception is in export part task");
+        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(export_manifests_mutex);
         writePartLog(
@@ -8751,6 +8784,7 @@ try
         part_log_elem.rows_read = (*exports_entry)->rows_read;
         part_log_elem.bytes_read_uncompressed = (*exports_entry)->bytes_read_uncompressed;
         part_log_elem.peak_memory_usage = (*exports_entry)->getPeakMemoryUsage();
+        part_log_elem.path_on_disk = (*exports_entry)->destination_file_path;
     }
 
     if (profile_counters)

src/Storages/MergeTree/MergeTreeExportManifest.h

@@ -8,9 +8,24 @@ struct MergeTreeExportManifest
 {
     using DataPartPtr = std::shared_ptr<const IMergeTreeDataPart>;
 
+
+    MergeTreeExportManifest(
+        const StorageID & destination_storage_id_,
+        const DataPartPtr & data_part_,
+        bool overwrite_file_if_exists_,
+        bool parallel_formatting_)
+        : destination_storage_id(destination_storage_id_),
+          data_part(data_part_),
+          overwrite_file_if_exists(overwrite_file_if_exists_),
+          parallel_formatting(parallel_formatting_),
+          create_time(time(nullptr)) {}
+
     StorageID destination_storage_id;
     DataPartPtr data_part;
-    time_t create_time = time(nullptr);
+    bool overwrite_file_if_exists;
+    bool parallel_formatting;
+
+    time_t create_time;
     mutable bool in_progress = false;
 
     bool operator<(const MergeTreeExportManifest & rhs) const 

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -49,6 +49,7 @@ namespace ErrorCodes
     extern const int NOT_IMPLEMENTED;
     extern const int LOGICAL_ERROR;
     extern const int INCORRECT_DATA;
+    extern const int FILE_ALREADY_EXISTS;
 }
 
 String StorageObjectStorage::getPathSample(ContextPtr context)
@@ -453,6 +454,8 @@ bool StorageObjectStorage::supportsImport() const
 SinkToStoragePtr StorageObjectStorage::import(
     const std::string & file_name,
     Block & block_with_partition_values,
+    std::string & destination_file_path,
+    bool overwrite_if_exists,
     ContextPtr local_context)
 {
     std::string partition_key;
@@ -467,16 +470,15 @@ SinkToStoragePtr StorageObjectStorage::import(
         }
     }
 
-    const auto file_path = configuration->getPathForWrite(partition_key, file_name).path;
+    destination_file_path = configuration->getPathForWrite(partition_key, file_name).path;
 
-    if (object_storage->exists(StoredObject(file_path)))
+    if (!overwrite_if_exists && object_storage->exists(StoredObject(destination_file_path)))
     {
-        LOG_INFO(getLogger("StorageObjectStorage"), "File {} already exists, skipping import", file_path);
-        return nullptr;
+        throw Exception(ErrorCodes::FILE_ALREADY_EXISTS, "File {} already exists", destination_file_path);
     }
 
     return std::make_shared<StorageObjectStorageSink>(
-        file_path,
+        destination_file_path,
         object_storage,
         configuration,
         format_settings,

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -106,6 +106,8 @@ class StorageObjectStorage : public IStorage
     SinkToStoragePtr import(
         const std::string & /* file_name */,
         Block & /* block_with_partition_values */,
+        std::string & /* destination_file_path */,
+        bool /* overwrite_if_exists */,
         ContextPtr /* context */) override;
 
     void truncate(

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -612,12 +612,14 @@ bool StorageObjectStorageCluster::supportsImport() const
 SinkToStoragePtr StorageObjectStorageCluster::import(
     const std::string & file_name,
     Block & block_with_partition_values,
+    std::string & destination_file_path,
+    bool overwrite_if_exists,
     ContextPtr context)
 {
     if (pure_storage)
-        return pure_storage->import(file_name, block_with_partition_values, context);
+        return pure_storage->import(file_name, block_with_partition_values, destination_file_path, overwrite_if_exists, context);
 
-    return IStorageCluster::import(file_name, block_with_partition_values, context);
+    return IStorageCluster::import(file_name, block_with_partition_values, destination_file_path, overwrite_if_exists, context);
 }