diff --git a/src/Core/FormatFactorySettings.h b/src/Core/FormatFactorySettings.h index 67396c955eac..38cb6d71cf50 100644 --- a/src/Core/FormatFactorySettings.h +++ b/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 `.`, where `` is `memory` or `threads` and `` 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) \ diff --git a/src/Core/SettingsChangesHistory.cpp b/src/Core/SettingsChangesHistory.cpp index a316f6956c17..3f656d4ebfd7 100644 --- a/src/Core/SettingsChangesHistory.cpp +++ b/src/Core/SettingsChangesHistory.cpp @@ -49,6 +49,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."}, diff --git a/src/Formats/FormatFactory.cpp b/src/Formats/FormatFactory.cpp index b40fb23ca262..02c67cd6aa4a 100644 --- a/src/Formats/FormatFactory.cpp +++ b/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().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(); + 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(weight.safeGet()); break; + case Field::Types::Int64: value = static_cast(weight.safeGet()); break; + case Field::Types::Float64: value = weight.safeGet(); 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()] = 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]; diff --git a/src/Formats/FormatSettings.h b/src/Formats/FormatSettings.h index 728a1faba670..1620e4603451 100644 --- a/src/Formats/FormatSettings.h +++ b/src/Formats/FormatSettings.h @@ -5,6 +5,8 @@ #include #include +#include + 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 + /// "." (resource = "memory" | "threads"). Empty means use the built-in + /// defaults. Parsed from the input_format_parquet_read_stage_weights setting. + std::map read_stage_weights; /// Write. UInt64 row_group_rows = 1000000; diff --git a/src/Processors/Formats/Impl/Parquet/ReadCommon.cpp b/src/Processors/Formats/Impl/Parquet/ReadCommon.cpp index 953562e818b4..736f8ec0ed9b 100644 --- a/src/Processors/Formats/Impl/Parquet/ReadCommon.cpp +++ b/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(parser_shared_resources.opaque.get()); size_t n = parser_shared_resources.num_streams.load(std::memory_order_relaxed); - fraction /= static_cast(std::max(n, size_t(1))); + /// Split each budget across the files read in parallel. + memory_fraction /= static_cast(std::max(n, size_t(1))); + thread_fraction /= static_cast(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 diff --git a/src/Processors/Formats/Impl/Parquet/ReadCommon.h b/src/Processors/Formats/Impl/Parquet/ReadCommon.h index 76b8a0fbccd5..b44723e5bf00 100644 --- a/src/Processors/Formats/Impl/Parquet/ReadCommon.h +++ b/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 "." + /// (resource = "memory" | "threads"). Empty means use the built-in defaults. + /// Applied in ReadManager::init. See input_format_parquet_read_stage_weights. + std::map 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; diff --git a/src/Processors/Formats/Impl/Parquet/ReadManager.cpp b/src/Processors/Formats/Impl/Parquet/ReadManager.cpp index 3408cd99c032..354fa0e700e3 100644 --- a/src/Processors/Formats/Impl/Parquet/ReadManager.cpp +++ b/src/Processors/Formats/Impl/Parquet/ReadManager.cpp @@ -13,10 +13,12 @@ #include #include #include +#include #include namespace DB::ErrorCodes { + extern const int BAD_ARGUMENTS; extern const int LOGICAL_ERROR; extern const int QUERY_WAS_CANCELLED; } @@ -74,17 +76,88 @@ void ReadManager::init(FormatParserSharedResourcesPtr parser_shared_resources_, stages[i].row_group_tasks_to_schedule.resize(num_row_groups); } - /// Distribute memory budget among stages. - /// The distribution is static to make sure no stage gets starved if others eat all the memory. - /// E.g. if the budget was shared among all stages, maybe PrewhereData could run far ahead and - /// The distribution is static to make sure no stage gets starved if others eat all the memory. - double sum = 0; - stages[size_t(ReadStage::NotStarted)].memory_target_fraction = 0; - stages[size_t(ReadStage::Deliver)].memory_target_fraction = 0; + /// Static per-stage memory and thread budgets so no stage starves the others (e.g. ColumnData + /// eating all memory and blocking the small index reads other row groups need). + /// + /// Values are relative weights, not fractions; the normalization loop below divides each by the + /// per-resource sum to get the actual fraction of the memory watermark / parsing pool per stage. + /// Memory and threads are weighted separately (a single shared fraction coupled them, capping + /// read-ahead and decode together): ColumnDataPrefetch holds small compressed pages, so it gets + /// the most memory to keep many GETs outstanding and hide S3 latency; ColumnData holds large + /// decoded columns, so its memory is bounded but it gets most threads (decode is the real CPU + /// work). Index/bloom stages only issue async reads, so 1 thread and a small memory floor each. + /// Integers are empirical priorities, not derived constants. + auto set_weights = [&](ReadStage s, double memory_weight, double thread_weight) + { + stages[size_t(s)].memory_target_fraction = memory_weight; + stages[size_t(s)].thread_target_fraction = thread_weight; + }; + set_weights(ReadStage::NotStarted, 0, 0); + set_weights(ReadStage::BloomFilterHeader, 1, 1); + set_weights(ReadStage::BloomFilterBlocksOrDictionary, 2, 1); + set_weights(ReadStage::ColumnIndexAndOffsetIndex, 1, 1); + set_weights(ReadStage::OffsetIndex, 1, 1); + set_weights(ReadStage::ColumnDataPrefetch, 9, 1); + set_weights(ReadStage::ColumnData, 6, 3); + set_weights(ReadStage::Deliver, 0, 0); + + /// Expert override of the weights above, from input_format_parquet_read_stage_weights. + /// Keys are "."; only listed keys override, the rest keep the defaults. + const auto & ext = *static_cast(parser_shared_resources->opaque.get()); + if (!ext.read_stage_weights.empty()) + { + static const std::unordered_map stage_by_name = + { + {"bloom_filter_header", ReadStage::BloomFilterHeader}, + {"bloom_filter_blocks_or_dictionary", ReadStage::BloomFilterBlocksOrDictionary}, + {"column_index_and_offset_index", ReadStage::ColumnIndexAndOffsetIndex}, + {"offset_index", ReadStage::OffsetIndex}, + {"column_data_prefetch", ReadStage::ColumnDataPrefetch}, + {"column_data", ReadStage::ColumnData}, + }; + for (const auto & [key, weight] : ext.read_stage_weights) + { + if (weight < 0) + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "Negative weight {} for '{}' in input_format_parquet_read_stage_weights", weight, key); + size_t dot = key.rfind('.'); + if (dot == String::npos) + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "Key '{}' in input_format_parquet_read_stage_weights must have the form '.'", key); + std::string_view stage_name(key.data(), dot); + std::string_view resource(key.data() + dot + 1, key.size() - dot - 1); + auto it = stage_by_name.find(stage_name); + if (it == stage_by_name.end()) + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "Unknown stage '{}' in input_format_parquet_read_stage_weights key '{}'", stage_name, key); + Stage & stage = stages[size_t(it->second)]; + if (resource == "memory") + stage.memory_target_fraction = weight; + else if (resource == "threads") + stage.thread_target_fraction = weight; + else + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "Unknown resource '{}' in input_format_parquet_read_stage_weights key '{}' (expected 'memory' or 'threads')", + resource, key); + } + } + + double memory_sum = 0; + double thread_sum = 0; for (const Stage & stage : stages) - sum += stage.memory_target_fraction; + { + memory_sum += stage.memory_target_fraction; + thread_sum += stage.thread_target_fraction; + } + if (memory_sum <= 0 || thread_sum <= 0) + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "input_format_parquet_read_stage_weights leaves the total {} weight at zero", + memory_sum <= 0 ? "memory" : "thread"); for (Stage & stage : stages) - stage.memory_target_fraction /= sum; + { + stage.memory_target_fraction /= memory_sum; + stage.thread_target_fraction /= thread_sum; + } /// The NotStarted stage completed for all row groups, transition to next stage. MemoryUsageDiff diff(ReadStage::NotStarted); @@ -139,6 +212,7 @@ void ReadManager::finishRowGroupStage(size_t row_group_idx, ReadStage stage, Mem switch (stage) { case ReadStage::NotStarted: + case ReadStage::ColumnDataPrefetch: case ReadStage::ColumnData: case ReadStage::Deliver: chassert(false); @@ -295,9 +369,10 @@ void ReadManager::addTasksToReadColumns(size_t row_group_idx, size_t row_subgrou } else { - LOG_TEST(getLogger("ParquetReadManager"), "addTasksToReadColumns: added ColumnData: i={} step_idx={} row_group_idx={} row_subgroup_idx={}", i, step_idx, row_group_idx, row_subgroup_idx); + /// `stage` is ColumnDataPrefetch (issue reads) or ColumnData (decode). + LOG_TEST(getLogger("ParquetReadManager"), "addTasksToReadColumns: added {}: i={} step_idx={} row_group_idx={} row_subgroup_idx={}", magic_enum::enum_name(stage), i, step_idx, row_group_idx, row_subgroup_idx); add_tasks.push_back(Task { - .stage = ReadStage::ColumnData, + .stage = stage, .step_idx = step_idx, .row_group_idx = row_group_idx, .row_subgroup_idx = row_subgroup_idx, @@ -307,8 +382,8 @@ void ReadManager::addTasksToReadColumns(size_t row_group_idx, size_t row_subgrou if (add_tasks.empty() && is_offset_index) { - /// Don't need to read offset index, move on to next stage (ColumnData). - stage = ReadStage::ColumnData; + /// Don't need to read offset index, move on to the next stage (ColumnDataPrefetch). + stage = ReadStage::ColumnDataPrefetch; continue; } @@ -319,7 +394,7 @@ void ReadManager::addTasksToReadColumns(size_t row_group_idx, size_t row_subgrou /// (RowSubgroup.filter.memory) work correctly when PREWHERE expression doesn't use any /// columns (note: the expression may still be nontrivial, e.g. `rand()%2=0`).) add_tasks.push_back(Task { - .stage = ReadStage::ColumnData, + .stage = stage, .step_idx = step_idx, .row_group_idx = row_group_idx, .row_subgroup_idx = row_subgroup_idx, @@ -420,6 +495,13 @@ void ReadManager::finishRowSubgroupStage(size_t row_group_idx, size_t row_subgro case ReadStage::ColumnIndexAndOffsetIndex: case ReadStage::OffsetIndex: { + /// Prerequisites read; issue the compressed data-page reads (but don't decode yet). + addTasksToReadColumns(row_group_idx, row_subgroup_idx, ReadStage::ColumnDataPrefetch, step_idx, diff); + return; + } + case ReadStage::ColumnDataPrefetch: + { + /// Data-page reads issued (in flight in the Prefetcher's io pool); now decode. addTasksToReadColumns(row_group_idx, row_subgroup_idx, ReadStage::ColumnData, step_idx, diff); return; } @@ -544,7 +626,7 @@ void ReadManager::flushMemoryUsageDiff(MemoryUsageDiff && diff) if (!should_schedule && d < 0) { const auto & stage = stages[i]; - auto limits = SharedResourcesExt::getLimitsPerReader(*parser_shared_resources, stage.memory_target_fraction); + auto limits = SharedResourcesExt::getLimitsPerReader(*parser_shared_resources, stage.memory_target_fraction, stage.thread_target_fraction); should_schedule = checkTaskSchedulingLimits( stage.memory_usage.load(std::memory_order_relaxed), 0, stage.batches_in_progress.load(std::memory_order_relaxed), 0, limits); @@ -562,7 +644,7 @@ void ReadManager::scheduleTasksIfNeeded(ReadStage stage_idx) MemoryUsageDiff diff(stage_idx); std::vector tasks; - auto limits = SharedResourcesExt::getLimitsPerReader(*parser_shared_resources, stage.memory_target_fraction); + auto limits = SharedResourcesExt::getLimitsPerReader(*parser_shared_resources, stage.memory_target_fraction, stage.thread_target_fraction); size_t memory_usage = stage.memory_usage.load(std::memory_order_relaxed); size_t batches_in_progress = stage.batches_in_progress.load(std::memory_order_relaxed); @@ -715,12 +797,14 @@ void ReadManager::scheduleTask(Task task, bool is_first_in_group, MemoryUsageDif case ReadStage::OffsetIndex: prefetches.push_back(&column.offset_index_prefetch); break; - case ReadStage::ColumnData: + case ReadStage::ColumnDataPrefetch: { RowSubgroup & row_subgroup = row_group.subgroups.at(task.row_subgroup_idx); - ColumnSubchunk & subchunk = row_subgroup.columns.at(task.column_idx); if (row_subgroup.filter.rows_pass == 0) break; + /// Queue this subgroup's data-page reads; startPrefetch (below) issues them and charges + /// compressed bytes to the ColumnDataPrefetch budget, separate from the decode budget, + /// so many row groups prefetch ahead while only a few decode at once. reader.determinePagesToPrefetch(column, row_subgroup, row_group, prefetches); /// Side note: would be nice to avoid reading the dictionary if all dictionary-encoded @@ -737,7 +821,16 @@ void ReadManager::scheduleTask(Task task, bool is_first_in_group, MemoryUsageDif { prefetches.push_back(&column.data_pages_prefetch); } - + break; + } + case ReadStage::ColumnData: + { + RowSubgroup & row_subgroup = row_group.subgroups.at(task.row_subgroup_idx); + ColumnSubchunk & subchunk = row_subgroup.columns.at(task.column_idx); + if (row_subgroup.filter.rows_pass == 0) + break; + /// Reads already issued in ColumnDataPrefetch; here just reserve estimated decoded-output + /// memory against the ColumnData budget (runTask decodes from those buffers). double bytes_per_row = reader.estimateColumnMemoryBytesPerRow(column, row_group, reader.primitive_columns.at(task.column_idx)); size_t column_memory = static_cast(bytes_per_row * static_cast(row_subgroup.filter.rows_pass)); subchunk.column_and_offsets_memory = MemoryUsageToken(column_memory, &diff); @@ -761,13 +854,14 @@ void ReadManager::scheduleTask(Task task, bool is_first_in_group, MemoryUsageDif reader.prefetcher.startPrefetch(prefetches, &diff); - /// We want to detect tiny tasks to group them together to reduce scheduling overhead. - /// Use the predicted memory usage as a rough estimate of how long a task will take. - /// E.g. main data read task's memory estimate consists of the input page sizes and the output - /// column size; the run time is also roughly proportional to these sizes. - /// Hope it's a good enough proxy in all cases. + /// Group tiny tasks to reduce scheduling overhead, using predicted memory as a proxy for run time. + /// Exception: ColumnDataPrefetch does its work (startPrefetch) here and has an empty runTask, so + /// its run time is ~0 no matter how many compressed bytes it charges; report cost 0 so these tasks + /// collapse into one batch instead of being split across many no-op thread-pool dispatches. ssize_t memory_after = diff.by_stage[size_t(diff.cur_stage)]; - task.cost_estimate_bytes = size_t(std::max(0l, memory_after - memory_before)); + task.cost_estimate_bytes = task.stage == ReadStage::ColumnDataPrefetch + ? 0 + : size_t(std::max(0l, memory_after - memory_before)); out_tasks.push_back(task); } @@ -842,6 +936,10 @@ void ReadManager::runTask(Task task, bool last_in_batch, MemoryUsageDiff & diff) reader.decodeOffsetIndex(column, row_group); column.offset_index_prefetch.reset(&diff); break; + case ReadStage::ColumnDataPrefetch: + /// Reads were issued in scheduleTask (startPrefetch) and run async in the Prefetcher; + /// nothing to do here. The subgroup advances to ColumnData, which decodes them. + break; case ReadStage::ColumnData: { RowSubgroup & row_subgroup = row_group.subgroups.at(task.row_subgroup_idx); @@ -859,7 +957,7 @@ void ReadManager::runTask(Task task, bool last_in_batch, MemoryUsageDiff & diff) chassert(task.row_subgroup_idx != UINT64_MAX); reader.decodePrimitiveColumn( column, column_info, row_subgroup.columns.at(task.column_idx), - row_group, row_subgroup); + row_group, row_subgroup, diff); for (size_t i = prev_page_idx; i < column.data_pages_idx; ++i) { diff --git a/src/Processors/Formats/Impl/Parquet/ReadManager.h b/src/Processors/Formats/Impl/Parquet/ReadManager.h index 49ac1b2942c8..7073492d2174 100644 --- a/src/Processors/Formats/Impl/Parquet/ReadManager.h +++ b/src/Processors/Formats/Impl/Parquet/ReadManager.h @@ -88,7 +88,11 @@ class ReadManager /// Tasks that are either in thread pool's queue or executing. std::atomic batches_in_progress {0}; + /// Share of the query-global memory budget for this stage, kept separate from the thread + /// share so a stage needing parallelism but little memory isn't forced to trade one off. double memory_target_fraction = 1; + /// Share of the parsing thread pool for this stage, independent of the memory share. + double thread_target_fraction = 1; /// We take advantage of the fact that each pair can have at most one group /// of tasks in flight at a time. E.g. we create tasks to read columns in subgroup n, then diff --git a/src/Processors/Formats/Impl/Parquet/Reader.cpp b/src/Processors/Formats/Impl/Parquet/Reader.cpp index 4fead45ac047..a3c91158f695 100644 --- a/src/Processors/Formats/Impl/Parquet/Reader.cpp +++ b/src/Processors/Formats/Impl/Parquet/Reader.cpp @@ -1361,7 +1361,7 @@ double Reader::estimateColumnMemoryBytesPerRow(const ColumnChunk & column, const return res; } -void Reader::decodePrimitiveColumn(ColumnChunk & column, const PrimitiveColumnInfo & column_info, ColumnSubchunk & subchunk, const RowGroup & row_group, RowSubgroup & row_subgroup) +void Reader::decodePrimitiveColumn(ColumnChunk & column, const PrimitiveColumnInfo & column_info, ColumnSubchunk & subchunk, const RowGroup & row_group, RowSubgroup & row_subgroup, MemoryUsageDiff & diff) { /// Allocate columns for values, null map, and array offsets. @@ -1536,6 +1536,19 @@ void Reader::decodePrimitiveColumn(ColumnChunk & column, const PrimitiveColumnIn chassert(subchunk.column->getDataType() == column_info.output_type->getColumnType()); + /// The scheduleTask charge was an estimate; reconcile up to the actual decoded footprint here, + /// before formOutputColumn (below) moves `subchunk.column`, so the scheduler stops decoding ahead + /// before real RAM exceeds the cap. Grow-only (fail-closed). + size_t actual_bytes = subchunk.column->allocatedBytes(); + for (const auto & offsets : subchunk.arrays_offsets) + if (offsets) + actual_bytes += offsets->allocatedBytes(); + if (subchunk.group_null_map) + actual_bytes += subchunk.group_null_map->allocatedBytes(); + size_t already_charged = subchunk.column_and_offsets_memory.charged(); + if (actual_bytes > already_charged) + subchunk.column_and_offsets_memory.add(actual_bytes - already_charged, &diff); + OutputColumnState & state = row_subgroup.output.at(column_info.idx_in_output_block); chassert(!state.column); size_t prev_count = state.primitive_columns_remaining.fetch_sub(1); diff --git a/src/Processors/Formats/Impl/Parquet/Reader.h b/src/Processors/Formats/Impl/Parquet/Reader.h index 36841514dc0f..105cb07a3061 100644 --- a/src/Processors/Formats/Impl/Parquet/Reader.h +++ b/src/Processors/Formats/Impl/Parquet/Reader.h @@ -532,7 +532,7 @@ struct Reader /// Guess how much memory ColumnSubchunk::{column, arrays_offsets} will use, per row. double estimateColumnMemoryBytesPerRow(const ColumnChunk & column, const RowGroup & row_group, const PrimitiveColumnInfo & column_info) const; - void decodePrimitiveColumn(ColumnChunk & column, const PrimitiveColumnInfo & column_info, ColumnSubchunk & subchunk, const RowGroup & row_group, RowSubgroup & row_subgroup); + void decodePrimitiveColumn(ColumnChunk & column, const PrimitiveColumnInfo & column_info, ColumnSubchunk & subchunk, const RowGroup & row_group, RowSubgroup & row_subgroup, MemoryUsageDiff & diff); /// Returns mutable column because some of the recursive calls require it, /// e.g. ColumnArray::create does assumeMutable() on the nested columns. diff --git a/src/Processors/Formats/Impl/ParquetV3BlockInputFormat.cpp b/src/Processors/Formats/Impl/ParquetV3BlockInputFormat.cpp index 6a238834caec..aa42fa65d47a 100644 --- a/src/Processors/Formats/Impl/ParquetV3BlockInputFormat.cpp +++ b/src/Processors/Formats/Impl/ParquetV3BlockInputFormat.cpp @@ -91,6 +91,7 @@ void ParquetV3BlockInputFormat::initializeIfNeeded() ext->total_memory_low_watermark = format_settings.parquet.memory_low_watermark; ext->total_memory_high_watermark = format_settings.parquet.memory_high_watermark; + ext->read_stage_weights = format_settings.parquet.read_stage_weights; parser_shared_resources->opaque = ext; });