Use shared ThreadPoolCallbackRunnerFast in ParquetBlockInputFormat

al13n321 · al13n321 · commit dbfe9328c6d8 · 2025-06-14T00:07:21.000Z
diff --git a/src/Common/threadPoolCallbackRunner.cpp b/src/Common/threadPoolCallbackRunner.cpp
@@ -21,7 +21,8 @@ void ThreadPoolCallbackRunnerFast::initThreadPool(ThreadPool & pool_, size_t max
     thread_name = thread_name_;
     thread_group = thread_group_;
 
-    /// TODO [parquet]: Maybe adjust number of threads dynamically based on queue size.
+    /// We could dynamically add and remove threads based on load, but it's not clear whether it's
+    /// worth the added complexity.
     threads = max_threads;
     for (size_t i = 0; i < max_threads; ++i)
         pool->scheduleOrThrowOnError([this] { threadFunction(); });
diff --git a/src/Processors/Formats/Impl/ParquetBlockInputFormat.cpp b/src/Processors/Formats/Impl/ParquetBlockInputFormat.cpp
@@ -8,6 +8,7 @@
 #include <Formats/FormatFactory.h>
 #include <Formats/SchemaInferenceUtils.h>
 #include <IO/ReadBufferFromMemory.h>
+#include <IO/SharedThreadPools.h>
 #include <IO/copyData.h>
 #include <arrow/api.h>
 #include <arrow/io/api.h>
@@ -32,6 +33,7 @@
 #include <Interpreters/convertFieldToType.h>
 #include <Storages/MergeTree/KeyCondition.h>
 
+#include <shared_mutex>
 #include <boost/algorithm/string/case_conv.hpp>
 
 namespace ProfileEvents
@@ -46,10 +48,6 @@ namespace CurrentMetrics
     extern const Metric IOThreads;
     extern const Metric IOThreadsActive;
     extern const Metric IOThreadsScheduled;
-
-    extern const Metric FormatParsingThreads;
-    extern const Metric FormatParsingThreadsActive;
-    extern const Metric FormatParsingThreadsScheduled;
 }
 
 namespace DB
@@ -583,20 +581,23 @@ ParquetBlockInputFormat::ParquetBlockInputFormat(
     , pending_chunks(PendingChunk::Compare { .row_group_first = format_settings_.parquet.preserve_order })
     , previous_block_missing_values(getPort().getHeader().columns())
 {
-    size_t max_decoding_threads = parser_group->getParsingThreadsPerReader();
-    if (max_decoding_threads > 1)
-        pool = std::make_unique<ThreadPool>(CurrentMetrics::FormatParsingThreads, CurrentMetrics::FormatParsingThreadsActive, CurrentMetrics::FormatParsingThreadsScheduled, max_decoding_threads);
+    use_thread_pool = parser_group->max_parsing_threads > 1;
 
-    size_t max_io_threads = parser_group->getIOThreadsPerReader();
-    if (supportPrefetch(max_decoding_threads, max_io_threads))
-        io_pool = std::make_shared<ThreadPool>(CurrentMetrics::IOThreads, CurrentMetrics::IOThreadsActive, CurrentMetrics::IOThreadsScheduled, max_io_threads);
+    bool row_group_prefetch =
+        !use_thread_pool && parser_group->max_io_threads > 0 &&
+        format_settings.parquet.enable_row_group_prefetch &&
+        !format_settings.parquet.use_native_reader;
+    if (row_group_prefetch)
+        io_pool = std::make_shared<ThreadPool>(
+            CurrentMetrics::IOThreads, CurrentMetrics::IOThreadsActive, CurrentMetrics::IOThreadsScheduled,
+            parser_group->getIOThreadsPerReader());
 }
 
 ParquetBlockInputFormat::~ParquetBlockInputFormat()
 {
     is_stopped = true;
-    if (pool)
-        pool->wait();
+    if (use_thread_pool)
+        shutdown->shutdown();
     if (io_pool)
         io_pool->wait();
 }
@@ -609,6 +610,10 @@ void ParquetBlockInputFormat::initializeIfNeeded()
     std::call_once(parser_group->init_flag, [&]
         {
             parser_group->initKeyCondition(getPort().getHeader());
+
+            if (use_thread_pool)
+                parser_group->parsing_runner.initThreadPool(
+                    getFormatParsingThreadPool().get(), parser_group->max_parsing_threads, "ParquetDecoder", CurrentThread::getGroup());
         });
 
     // Create arrow file adapter.
@@ -889,13 +894,15 @@ void ParquetBlockInputFormat::scheduleRowGroup(size_t row_group_batch_idx)
 
     status = RowGroupBatchState::Status::Running;
 
-    pool->scheduleOrThrowOnError(
-        [this, row_group_batch_idx, thread_group = CurrentThread::getGroup()]()
+    parser_group->parsing_runner(
+        [this, row_group_batch_idx, shutdown_ = shutdown]()
         {
+            std::shared_lock shutdown_lock(*shutdown_, std::try_to_lock);
+            if (!shutdown_lock.owns_lock())
+                return;
+
             try
             {
-                ThreadGroupSwitcher switcher(thread_group, "ParquetDecoder");
-
                 threadFunction(row_group_batch_idx);
             }
             catch (...)
@@ -928,11 +935,6 @@ void ParquetBlockInputFormat::threadFunction(size_t row_group_batch_idx)
             return;
     }
 }
-bool ParquetBlockInputFormat::supportPrefetch(size_t max_decoding_threads, size_t max_io_threads) const
-{
-    return max_decoding_threads == 1 && max_io_threads > 0 && format_settings.parquet.enable_row_group_prefetch && !format_settings.parquet.use_native_reader;
-}
-
 std::shared_ptr<arrow::RecordBatchReader> ParquetBlockInputFormat::RowGroupPrefetchIterator::nextRowGroupReader()
 {
     if (prefetched_row_groups.empty()) return nullptr;
@@ -1069,7 +1071,7 @@ void ParquetBlockInputFormat::scheduleMoreWorkIfNeeded(std::optional<size_t> row
         ++row_group_batches_completed;
     }
 
-    if (pool)
+    if (use_thread_pool)
     {
         size_t max_decoding_threads = parser_group->getParsingThreadsPerReader();
         while (row_group_batches_started - row_group_batches_completed < max_decoding_threads &&
@@ -1132,7 +1134,7 @@ Chunk ParquetBlockInputFormat::read()
         if (row_group_batches_completed == row_group_batches.size())
             return {};
 
-        if (pool)
+        if (use_thread_pool)
             condvar.wait(lock);
         else
             decodeOneChunk(row_group_batches_completed, lock);
@@ -1142,8 +1144,11 @@ Chunk ParquetBlockInputFormat::read()
 void ParquetBlockInputFormat::resetParser()
 {
     is_stopped = true;
-    if (pool)
-        pool->wait();
+    if (use_thread_pool)
+    {
+        shutdown->shutdown();
+        shutdown = std::make_shared<ShutdownHelper>();
+    }
 
     arrow_file.reset();
     metadata.reset();
diff --git a/src/Processors/Formats/Impl/ParquetBlockInputFormat.h b/src/Processors/Formats/Impl/ParquetBlockInputFormat.h
@@ -89,8 +89,6 @@ class ParquetBlockInputFormat : public IInputFormat
 
     void threadFunction(size_t row_group_batch_idx);
 
-    inline bool supportPrefetch(size_t max_decoding_threads, size_t max_io_threads) const;
-
     // Data layout in the file:
     //
     // row group 0
@@ -327,7 +325,8 @@ class ParquetBlockInputFormat : public IInputFormat
 
     // These are only used when max_decoding_threads > 1.
     size_t row_group_batches_started = 0;
-    std::unique_ptr<ThreadPool> pool;
+    bool use_thread_pool = false;
+    std::shared_ptr<ShutdownHelper> shutdown = std::make_shared<ShutdownHelper>();
     std::shared_ptr<ThreadPool> io_pool;
 
     BlockMissingValues previous_block_missing_values;
