Parquet v3 read concurrency

src/Core/FormatFactorySettings.h

@@ -201,6 +201,9 @@ Schedule prefetches more aggressively if memory usage is below than threshold. P
     DECLARE(UInt64, input_format_parquet_memory_high_watermark, 4ul << 30, R"(
 Approximate memory limit for Parquet reader v3. Limits how many row groups or columns can be read in parallel. When reading multiple files in one query, the limit is on total memory usage across those files.
 )", 0) \
+    DECLARE(Map, input_format_parquet_read_stage_weights, "", R"(
+Expert tuning knob for the Parquet reader scheduler. Overrides the relative memory and thread budget weights assigned to each read stage. Keys have the form `<stage>.<resource>`, where `<resource>` is `memory` or `threads` and `<stage>` is one of `bloom_filter_header`, `bloom_filter_blocks_or_dictionary`, `column_index_and_offset_index`, `offset_index`, `column_data_prefetch`, `column_data`. Values are relative weights (they need not sum to 1; each is normalized against the sum of that resource across stages). Only the specified keys override the built-in defaults; unspecified stages keep their defaults. Example: `{'column_data_prefetch.memory': 12, 'column_data.threads': 4}`. Advanced and experimental: the internal stage model may change between versions.
+)", EXPERIMENTAL) \
     DECLARE(Bool, input_format_parquet_page_filter_push_down, true, R"(
 Skip pages using min/max values from column index.
 )", 0) \

src/Core/SettingsChangesHistory.cpp

@@ -44,6 +44,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"analyzer_compatibility_apply_final_to_all_joined_tables", true, false, "Fixed a bug in the analyzer where FINAL on the left-most table of a JOIN was incorrectly applied to the other joined tables as well. previous_value=true so `compatibility` with versions before 26.6 restores the old behavior."},
             {"analyzer_compatibility_allow_non_aggregate_in_having", false, false, "New compatibility setting. When enabled, the new analyzer mimics the legacy `HAVING`-to-`WHERE` rewrite for non-aggregate AND-conjuncts instead of raising `NOT_AN_AGGREGATE`."},
             {"reserve_memory", 0, 0, "New setting to reserve memory for specific workload before starting a query."},
+            {"input_format_parquet_read_stage_weights", "", "", "New experimental setting to override the per-stage memory and thread budget weights of the Parquet reader scheduler."},
             {"output_format_image_width", 1024, 1024, "New setting controlling the width of the output image for image output formats such as PNG."},
             {"output_format_image_height", 1024, 1024, "New setting controlling the height of the output image for image output formats such as PNG."},
             {"output_format_image_terminal_mode", "", "", "New setting controlling whether image output formats such as PNG are rendered directly to the terminal using an inline image protocol."},

src/Formats/FormatFactory.cpp

@@ -223,6 +223,31 @@ FormatSettings getFormatSettings(const ContextPtr & context, const Settings & se
     format_settings.parquet.enable_json_parsing = settings[Setting::input_format_parquet_enable_json_parsing];
     format_settings.parquet.memory_low_watermark = settings[Setting::input_format_parquet_memory_low_watermark];
     format_settings.parquet.memory_high_watermark = settings[Setting::input_format_parquet_memory_high_watermark];
+    {
+        const Map & stage_weights = settings[Setting::input_format_parquet_read_stage_weights].value;
+        for (const auto & key_value : stage_weights)
+        {
+            if (key_value.getType() != Field::Types::Tuple || key_value.safeGet<Tuple>().size() != 2)
+                throw Exception(ErrorCodes::BAD_ARGUMENTS,
+                    "The value of the `input_format_parquet_read_stage_weights` setting must be a Map(String, Float64)");
+            const Tuple & pair = key_value.safeGet<Tuple>();
+            if (pair.at(0).getType() != Field::Types::String)
+                throw Exception(ErrorCodes::BAD_ARGUMENTS,
+                    "The keys of the `input_format_parquet_read_stage_weights` setting must be Strings");
+            const Field & weight = pair.at(1);
+            double value;
+            switch (weight.getType())
+            {
+                case Field::Types::UInt64: value = static_cast<double>(weight.safeGet<UInt64>()); break;
+                case Field::Types::Int64:  value = static_cast<double>(weight.safeGet<Int64>()); break;
+                case Field::Types::Float64: value = weight.safeGet<Float64>(); break;
+                default:
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS,
+                        "The values of the `input_format_parquet_read_stage_weights` setting must be numbers");
+            }
+            format_settings.parquet.read_stage_weights[pair.at(0).safeGet<String>()] = value;
+        }
+    }
     format_settings.parquet.allow_missing_columns = settings[Setting::input_format_parquet_allow_missing_columns];
     format_settings.parquet.skip_columns_with_unsupported_types_in_schema_inference = settings[Setting::input_format_parquet_skip_columns_with_unsupported_types_in_schema_inference];
     format_settings.parquet.output_string_as_string = settings[Setting::output_format_parquet_string_as_string];

src/Formats/FormatSettings.h

@@ -5,6 +5,8 @@
 #include <base/types.h>
 #include <base/unit.h>
 
+#include <map>
+
 namespace DB
 {
 
@@ -361,6 +363,10 @@ struct FormatSettings
         size_t local_read_min_bytes_for_seek = 8192;
         size_t memory_low_watermark = 2ul << 20;
         size_t memory_high_watermark = 4ul << 30;
+        /// Optional per-stage budget weight overrides for the reader scheduler, keyed by
+        /// "<stage>.<resource>" (resource = "memory" | "threads"). Empty means use the built-in
+        /// defaults. Parsed from the input_format_parquet_read_stage_weights setting.
+        std::map<String, double> read_stage_weights;
 
         /// Write.
         UInt64 row_group_rows = 1000000;

src/Processors/Formats/Impl/Parquet/ReadCommon.cpp

@@ -6,15 +6,17 @@
 namespace DB::Parquet
 {
 
-SharedResourcesExt::Limits SharedResourcesExt::getLimitsPerReader(const FormatParserSharedResources & parser_shared_resources, double fraction)
+SharedResourcesExt::Limits SharedResourcesExt::getLimitsPerReader(const FormatParserSharedResources & parser_shared_resources, double memory_fraction, double thread_fraction)
 {
     const SharedResourcesExt & ext = *static_cast<const SharedResourcesExt *>(parser_shared_resources.opaque.get());
     size_t n = parser_shared_resources.num_streams.load(std::memory_order_relaxed);
-    fraction /= static_cast<double>(std::max(n, size_t(1)));
+    /// Split each budget across the files read in parallel.
+    memory_fraction /= static_cast<double>(std::max(n, size_t(1)));
+    thread_fraction /= static_cast<double>(std::max(n, size_t(1)));
     return Limits {
-        .memory_low_watermark = size_t(ext.total_memory_low_watermark * fraction),
-        .memory_high_watermark = size_t(ext.total_memory_high_watermark * fraction),
-        .parsing_threads = size_t(std::max(std::lround(parser_shared_resources.parsing_runner.getMaxThreads() * fraction + .5), 1l))};
+        .memory_low_watermark = size_t(ext.total_memory_low_watermark * memory_fraction),
+        .memory_high_watermark = size_t(ext.total_memory_high_watermark * memory_fraction),
+        .parsing_threads = size_t(std::max(std::lround(parser_shared_resources.parsing_runner.getMaxThreads() * thread_fraction + .5), 1l))};
 }
 
 #ifdef OS_LINUX

src/Processors/Formats/Impl/Parquet/ReadCommon.h

@@ -42,6 +42,10 @@ struct SharedResourcesExt
 {
     size_t total_memory_low_watermark = 0;
     size_t total_memory_high_watermark = 0;
+    /// Optional per-stage budget weight overrides, keyed by "<stage>.<resource>"
+    /// (resource = "memory" | "threads"). Empty means use the built-in defaults.
+    /// Applied in ReadManager::init. See input_format_parquet_read_stage_weights.
+    std::map<String, double> read_stage_weights;
 
     struct Limits
     {
@@ -50,7 +54,7 @@ struct SharedResourcesExt
         size_t parsing_threads;
     };
 
-    static Limits getLimitsPerReader(const FormatParserSharedResources & parser_shared_resources, double fraction);
+    static Limits getLimitsPerReader(const FormatParserSharedResources & parser_shared_resources, double memory_fraction, double thread_fraction);
 };
 
 
@@ -88,6 +92,9 @@ enum class ReadStage
     ColumnIndexAndOffsetIndex,
 
     OffsetIndex,
+    /// Issues the compressed data-page reads (startPrefetch), no decode. Own memory budget, so many
+    /// row groups prefetch ahead while only a few decode at once. Decouples fetch from decode depth.
+    ColumnDataPrefetch,
     ColumnData,
 
     Deliver,
@@ -186,6 +193,9 @@ class MemoryUsageToken
         val += amount;
     }
 
+    /// How much memory this token currently charges.
+    size_t charged() const { return val; }
+
 private:
     ReadStage alloc_stage = ReadStage::Deallocated;
     size_t val = 0;