exports throtler

Author: arthurpassos

src/Core/ServerSettings.cpp

@@ -99,6 +99,7 @@ namespace DB
     DECLARE(UInt64, max_unexpected_parts_loading_thread_pool_size, 8, R"(The number of threads to load inactive set of data parts (Unexpected ones) at startup.)", 0) \
     DECLARE(UInt64, max_parts_cleaning_thread_pool_size, 128, R"(The number of threads for concurrent removal of inactive data parts.)", 0) \
     DECLARE(UInt64, max_mutations_bandwidth_for_server, 0, R"(The maximum read speed of all mutations on server in bytes per second. Zero means unlimited.)", 0) \
+    DECLARE(UInt64, max_exports_bandwidth_for_server, 0, R"(The maximum read speed of all exports on server in bytes per second. Zero means unlimited.)", 0) \
     DECLARE(UInt64, max_merges_bandwidth_for_server, 0, R"(The maximum read speed of all merges on server in bytes per second. Zero means unlimited.)", 0) \
     DECLARE(UInt64, max_replicated_fetches_network_bandwidth_for_server, 0, R"(The maximum speed of data exchange over the network in bytes per second for replicated fetches. Zero means unlimited.)", 0) \
     DECLARE(UInt64, max_replicated_sends_network_bandwidth_for_server, 0, R"(The maximum speed of data exchange over the network in bytes per second for replicated sends. Zero means unlimited.)", 0) \

src/Interpreters/Context.cpp

@@ -288,6 +288,7 @@ namespace ServerSetting
     extern const ServerSettingsUInt64 max_local_write_bandwidth_for_server;
     extern const ServerSettingsUInt64 max_merges_bandwidth_for_server;
     extern const ServerSettingsUInt64 max_mutations_bandwidth_for_server;
+    extern const ServerSettingsUInt64 max_exports_bandwidth_for_server;
     extern const ServerSettingsUInt64 max_remote_read_network_bandwidth_for_server;
     extern const ServerSettingsUInt64 max_remote_write_network_bandwidth_for_server;
     extern const ServerSettingsUInt64 max_replicated_fetches_network_bandwidth_for_server;
@@ -504,6 +505,8 @@ struct ContextSharedPart : boost::noncopyable
     mutable ThrottlerPtr mutations_throttler;               /// A server-wide throttler for mutations
     mutable ThrottlerPtr merges_throttler;                  /// A server-wide throttler for merges
 
+    mutable ThrottlerPtr exports_throttler;                /// A server-wide throttler for exports
+
     MultiVersion<Macros> macros;                            /// Substitutions extracted from config.
     std::unique_ptr<DDLWorker> ddl_worker TSA_GUARDED_BY(mutex); /// Process ddl commands from zk.
     LoadTaskPtr ddl_worker_startup_task;                         /// To postpone `ddl_worker->startup()` after all tables startup
@@ -992,6 +995,9 @@ struct ContextSharedPart : boost::noncopyable
 
         if (auto bandwidth = server_settings[ServerSetting::max_merges_bandwidth_for_server])
             merges_throttler = std::make_shared<Throttler>(bandwidth);
+
+        if (auto bandwidth = server_settings[ServerSetting::max_exports_bandwidth_for_server])
+            exports_throttler = std::make_shared<Throttler>(bandwidth);
     }
 };
 
@@ -4041,6 +4047,11 @@ ThrottlerPtr Context::getMergesThrottler() const
     return shared->merges_throttler;
 }
 
+ThrottlerPtr Context::getExportsThrottler() const
+{
+    return shared->exports_throttler;
+}
+
 void Context::reloadRemoteThrottlerConfig(size_t read_bandwidth, size_t write_bandwidth) const
 {
     if (read_bandwidth)

src/Interpreters/Context.h

@@ -1635,6 +1635,7 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
 
     ThrottlerPtr getMutationsThrottler() const;
     ThrottlerPtr getMergesThrottler() const;
+    ThrottlerPtr getExportsThrottler() const;
 
     void reloadRemoteThrottlerConfig(size_t read_bandwidth, size_t write_bandwidth) const;
     void reloadLocalThrottlerConfig(size_t read_bandwidth, size_t write_bandwidth) const;

src/Storages/MergeTree/MergeTreeSequentialSource.cpp

@@ -140,6 +140,9 @@ MergeTreeSequentialSource::MergeTreeSequentialSource(
         case Merge:
             read_settings.local_throttler = context->getMergesThrottler();
             break;
+        case Export:
+            read_settings.local_throttler = context->getExportsThrottler();
+            break;
     }
 
     read_settings.remote_throttler = read_settings.local_throttler;

src/Storages/MergeTree/MergeTreeSequentialSource.h

@@ -15,6 +15,7 @@ enum MergeTreeSequentialSourceType
 {
     Mutation,
     Merge,
+    Export,
 };
 
 /// Create stream for reading single part from MergeTree.

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.cpp

@@ -10,11 +10,6 @@
 namespace DB
 {
 
-namespace Setting
-{
-    extern const SettingsBool output_format_parallel_formatting;
-}
-
 ExportPartPlainMergeTreeTask::ExportPartPlainMergeTreeTask(
     StorageMergeTree & storage_,
     const DataPartPtr & part_to_export_,

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.h

@@ -23,7 +23,7 @@ class ExportPartPlainMergeTreeTask : public IExecutableTask
         ContextPtr context_,
         std::shared_ptr<MergeTreeExportManifest> manifest_,
         IExecutableTask::TaskResultCallback & task_result_callback_,
-        size_t max_retries_ = 3);
+        size_t max_retries_);
 
     void onCompleted() override;
     bool executeStep() override;

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -565,20 +565,31 @@ void StorageObjectStorage::importMergeTreePart(
 
     QueryPlan plan;
 
-    /// using the mutations type for now
-    MergeTreeSequentialSourceType read_type = MergeTreeSequentialSourceType::Mutation;
+    /// using the mutations type for now. This impacts in the throttling strategy
+    MergeTreeSequentialSourceType read_type = MergeTreeSequentialSourceType::Export;
 
+    /// todo implement these settings
     bool apply_deleted_mask = true;
     bool read_with_direct_io = false;
     bool prefetch = false;
 
-    const auto partition_columns = configuration->partition_strategy->getPartitionColumns();
+    std::string partition_key;
 
-    auto block_with_partition_values = data_part->partition.getBlockWithPartitionValues(partition_columns);
+    if (configuration->partition_strategy)
+    {
+        const auto partition_columns = configuration->partition_strategy->getPartitionColumns();
+
+        auto block_with_partition_values = data_part->partition.getBlockWithPartitionValues(partition_columns);
+    
+        const auto column_with_partition_key = configuration->partition_strategy->computePartitionKey(block_with_partition_values);
 
-    const auto column_with_partition_key = configuration->partition_strategy->computePartitionKey(block_with_partition_values);
+        if (!column_with_partition_key->empty())
+        {
+            partition_key = column_with_partition_key->getDataAt(0).toString();
+        }
+    }
 
-    const auto file_path = configuration->file_path_generator->getWritingPath(column_with_partition_key->getDataAt(0).toString());
+    const auto file_path = configuration->file_path_generator->getWritingPath(partition_key);
 
     MergeTreeData::IMutationsSnapshot::Params params
     {