From 4fa22fd38e751703ace44b1b1eea33aafc634f31 Mon Sep 17 00:00:00 2001 From: VighneshPath Date: Wed, 2 Sep 2026 13:06:30 +0530 Subject: [PATCH 1/3] Merge two-level aggregation buckets largest-first to cut merge tail latency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aggregator::mergeBlocks dispatched bucket merges in ascending bucket-id order. Real bucket-row-count distributions are skewed, so a large bucket landing anywhere in that order left most merge threads idle early while one thread finished it alone — a straggler tail dominating wall time for wide GROUP BY queries (found via flamegraph/profile-event analysis of IcebergBench's q12_wide_groupby). Sort buckets by row count descending before dispatch (longest-processing- time-first), so the biggest buckets start while every thread is still free to help, and the small ones are left for whoever finishes first. Co-Authored-By: Claude Sonnet 5 Signed-off-by: VighneshPath --- src/Interpreters/Aggregator.cpp | 42 +++++++++++++++---- src/Interpreters/Aggregator.h | 6 +++ .../gtest_aggregator_bucket_merge_order.cpp | 41 ++++++++++++++++++ 3 files changed, 82 insertions(+), 7 deletions(-) create mode 100644 src/Interpreters/tests/gtest_aggregator_bucket_merge_order.cpp diff --git a/src/Interpreters/Aggregator.cpp b/src/Interpreters/Aggregator.cpp index 801000e53407..034a62d162a5 100644 --- a/src/Interpreters/Aggregator.cpp +++ b/src/Interpreters/Aggregator.cpp @@ -3566,15 +3566,38 @@ bool Aggregator::mergeOnBlock(Columns columns, size_t rows, bool is_overflows, A } +std::vector Aggregator::sortBucketsByRowCountDescending(std::vector> buckets_by_rows) +{ + std::sort(buckets_by_rows.begin(), buckets_by_rows.end(), + [](const auto & a, const auto & b) { return a.second > b.second; }); + + std::vector result; + result.reserve(buckets_by_rows.size()); + for (const auto & bucket_and_rows : buckets_by_rows) + result.push_back(bucket_and_rows.first); + return result; +} + void Aggregator::mergeBlocks(BucketToChunks bucket_to_chunks, AggregatedDataVariants & result, std::atomic & is_cancelled) { if (bucket_to_chunks.empty()) return; UInt64 total_input_rows = 0; + /// Real two-level buckets (bucket >= 0) paired with their row count, for longest-processing- + /// time-first merge dispatch below. Bucket -1 (single-level/overflow blocks) is handled + /// separately and excluded here. + std::vector> buckets_by_rows; + buckets_by_rows.reserve(bucket_to_chunks.size()); for (auto & bucket : bucket_to_chunks) + { + UInt64 bucket_rows = 0; for (auto & agg_chunk : bucket.second) - total_input_rows += agg_chunk.chunk.getNumRows(); + bucket_rows += agg_chunk.chunk.getNumRows(); + total_input_rows += bucket_rows; + if (bucket.first >= 0) + buckets_by_rows.emplace_back(bucket.first, bucket_rows); + } /** `minus one` means the absence of information about the bucket * - in the case of single-level aggregation, as well as for blocks with "overflowing" values. @@ -3613,19 +3636,24 @@ void Aggregator::mergeBlocks(BucketToChunks bucket_to_chunks, AggregatedDataVari LOG_TRACE(log, "Merging partially aggregated two-level data."); - std::atomic next_bucket_to_merge = 0; + /// Longest-processing-time-first dispatch order (see sortBucketsByRowCountDescending): + /// start the biggest buckets while every thread is still free to help absorb them, instead + /// of plain ascending bucket-id order, which lets threads that luck into small buckets go + /// idle while a handful of stragglers grind through oversized buckets alone. + std::vector buckets_to_merge = sortBucketsByRowCountDescending(std::move(buckets_by_rows)); - auto merge_bucket = [&bucket_to_chunks, &result, &is_cancelled, &next_bucket_to_merge, max_bucket, this](Arena * aggregates_pool) + std::atomic next_bucket_index = 0; + + auto merge_bucket = [&bucket_to_chunks, &result, &is_cancelled, &next_bucket_index, &buckets_to_merge, this](Arena * aggregates_pool) { while (true) { - const Int32 bucket = next_bucket_to_merge.fetch_add(1); + const size_t index = next_bucket_index.fetch_add(1); - if (bucket > max_bucket) + if (index >= buckets_to_merge.size()) break; - if (!bucket_to_chunks.contains(bucket)) - continue; + const Int32 bucket = buckets_to_merge[index]; if (is_cancelled.load()) return; diff --git a/src/Interpreters/Aggregator.h b/src/Interpreters/Aggregator.h index cb42add7547a..e819ee1f867c 100644 --- a/src/Interpreters/Aggregator.h +++ b/src/Interpreters/Aggregator.h @@ -271,6 +271,12 @@ class Aggregator final /// Merge partially aggregated chunks separated to buckets into one data structure. void mergeBlocks(BucketToChunks bucket_to_chunks, AggregatedDataVariants & result, std::atomic & is_cancelled); + /// Sorts (bucket_id, row_count) pairs by row_count descending, for longest-processing-time- + /// first parallel bucket-merge dispatch in mergeBlocks: starting the biggest buckets first, + /// while every merge thread is still free to help absorb them, keeps threads that would + /// otherwise luck into small buckets and go idle busy instead, under skewed bucket sizes. + static std::vector sortBucketsByRowCountDescending(std::vector> buckets_by_rows); + /// Merge several partially aggregated chunks into one. /// Precondition: for all chunks the is_overflows flag must be the same. /// (either all chunks are from overflow data or none are). diff --git a/src/Interpreters/tests/gtest_aggregator_bucket_merge_order.cpp b/src/Interpreters/tests/gtest_aggregator_bucket_merge_order.cpp new file mode 100644 index 000000000000..cb17bb22bb8d --- /dev/null +++ b/src/Interpreters/tests/gtest_aggregator_bucket_merge_order.cpp @@ -0,0 +1,41 @@ +#include +#include + +#include + +#include + +using namespace DB; + +TEST(AggregatorBucketMergeOrder, EmptyInput) +{ + auto order = Aggregator::sortBucketsByRowCountDescending({}); + EXPECT_TRUE(order.empty()); +} + +TEST(AggregatorBucketMergeOrder, SingleBucket) +{ + auto order = Aggregator::sortBucketsByRowCountDescending({{5, 100}}); + ASSERT_EQ(order.size(), 1u); + EXPECT_EQ(order[0], 5); +} + +TEST(AggregatorBucketMergeOrder, SortsByRowCountDescendingRegardlessOfBucketId) +{ + /// Bucket ids intentionally don't correlate with size, matching the real symptom: a + /// skewed hash distribution can put the biggest bucket anywhere in the id range. + auto order = Aggregator::sortBucketsByRowCountDescending({{0, 10}, {1, 1000}, {2, 500}, {3, 1}}); + EXPECT_EQ(order, (std::vector{1, 2, 0, 3})); +} + +TEST(AggregatorBucketMergeOrder, AllBucketsPreservedExactlyOnceOnTies) +{ + auto order = Aggregator::sortBucketsByRowCountDescending({{0, 50}, {1, 50}, {2, 200}}); + ASSERT_EQ(order.size(), 3u); + /// The strictly largest bucket must be dispatched first. + EXPECT_EQ(order[0], 2); + /// The tied pair can come back in either order, but both must be present. + std::vector tied{order[1], order[2]}; + std::sort(tied.begin(), tied.end()); + EXPECT_EQ(tied, (std::vector{0, 1})); +} From 79f6ee13800907e8dec3d64f14319886e7383ea9 Mon Sep 17 00:00:00 2001 From: VighneshPath Date: Wed, 2 Sep 2026 13:06:42 +0530 Subject: [PATCH 2/3] Defer ReadFromCluster's task/file list to initializePipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ReadFromCluster::createExtension() (backing icebergCluster(), s3Cluster(), and every other IStorageCluster-based cluster table function) was built eagerly from applyFilters(), then frozen via a one-shot guard. But query plan optimizations call applyFilters() more than once as the plan is refined — a later pass (e.g. aggregation-in-order for a GROUP BY matching the partition/sort key) can insert another FilterStep above the source, making filter_actions_dag strictly more complete on a later call. The frozen extension silently kept whichever predicate happened to be known on the very first call, dropping anything discovered afterward. Live testing against Iceberg tables confirmed the effect: for the same WHERE clause, icebergCluster() read 6-13x more rows than the equivalent ice.`ns.table` query specifically on GROUP BY-by-partition-column queries that trigger the extra optimization pass; queries without it were unaffected, since their filter was already complete on the first call. ReadFromObjectStorageStep (the non-cluster object storage read path) already gets this right: it defers building its file iterator to initializePipeline(), which query optimization only ever reaches after every pass has finished, so it always sees the final filter. Make ReadFromCluster follow the same shape: applyFilters() now only updates filter_actions_dag, and createExtension() is called exactly once, from initializePipeline(), using the filter's final state. Co-Authored-By: Claude Sonnet 5 Signed-off-by: VighneshPath --- src/Storages/IStorageCluster.cpp | 23 ++- ...t_read_from_cluster_predicate_pushdown.cpp | 154 ++++++++++++++++++ 2 files changed, 169 insertions(+), 8 deletions(-) create mode 100644 src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp diff --git a/src/Storages/IStorageCluster.cpp b/src/Storages/IStorageCluster.cpp index 3012c7bff735..02389f92cf4b 100644 --- a/src/Storages/IStorageCluster.cpp +++ b/src/Storages/IStorageCluster.cpp @@ -86,14 +86,16 @@ IStorageCluster::IStorageCluster( void ReadFromCluster::applyFilters(ActionDAGNodes added_filter_nodes) { + /// Query plan optimizations run this multiple times as the plan is refined (e.g. a later + /// pass may insert another FilterStep above this source), so filter_actions_dag can still + /// grow richer after this call returns. The task/file list must not be built here: doing so + /// would freeze it against whatever predicate happens to be known at this arbitrary point in + /// the optimization sequence, silently dropping conditions that only become visible on a + /// later pass (this previously caused Iceberg partition/row-group pruning to under-prune for + /// queries whose plan needed extra passes, e.g. GROUP BY on the partition column). The + /// extension is instead built once, lazily, in initializePipeline() — the only point where + /// filter_actions_dag is guaranteed final — mirroring ReadFromObjectStorageStep::createIterator(). SourceStepWithFilter::applyFilters(std::move(added_filter_nodes)); - - const ActionsDAG::Node * predicate = nullptr; - const ActionsDAG * filter = filter_actions_dag ? filter_actions_dag.get() : query_info.filter_actions_dag.get(); - if (filter) - predicate = filter->getOutputs().at(0); - - createExtension(predicate); } void ReadFromCluster::createExtension(const ActionsDAG::Node * predicate) @@ -596,7 +598,12 @@ void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const if (current_settings[Setting::max_parallel_replicas] > 1) max_replicas_to_use = std::min(max_replicas_to_use, current_settings[Setting::max_parallel_replicas].value); - createExtension(nullptr); + /// Build off the final, fully-optimized filter_actions_dag — see the comment in applyFilters(). + const ActionsDAG::Node * predicate = nullptr; + const ActionsDAG * filter = filter_actions_dag ? filter_actions_dag.get() : query_info.filter_actions_dag.get(); + if (filter) + predicate = filter->getOutputs().at(0); + createExtension(predicate); ProfileEvents::increment(ProfileEvents::Shards, max_replicas_to_use); diff --git a/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp b/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp new file mode 100644 index 000000000000..c77cf7552b28 --- /dev/null +++ b/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp @@ -0,0 +1,154 @@ +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace DB; + +namespace +{ + +/// Records every call to getTaskIteratorExtension(), including which columns' conditions were +/// present in the filter DAG it was handed. +class RecordingStorageCluster : public IStorageCluster +{ +public: + RecordingStorageCluster(const String & cluster_name, const StorageID & table_id, LoggerPtr log_) + : IStorageCluster(cluster_name, table_id, log_) + { + } + + std::string getName() const override { return "RecordingStorageCluster"; } + + RemoteQueryExecutor::Extension getTaskIteratorExtension( + const ActionsDAG::Node *, + const ActionsDAG * filter_actions_dag, + const ContextPtr &, + ClusterPtr, + StorageMetadataPtr) const override + { + ++call_count; + last_filter_input_names.clear(); + if (filter_actions_dag) + for (const auto * input : filter_actions_dag->getInputs()) + last_filter_input_names.insert(input->result_name); + return {}; + } + + mutable size_t call_count = 0; + mutable std::set last_filter_input_names; +}; + +/// Port 1 is never listening, so ReadFromCluster's connection attempt fails immediately +/// (connection refused) instead of timing out — keeps the test fast and hermetic. +ClusterPtr makeSingleUnreachableReplicaCluster(const ContextMutablePtr & context) +{ + std::istringstream config_stream{ + "" + "127.0.0.11" + ""}; + Poco::AutoPtr config = new Poco::Util::XMLConfiguration(config_stream); + context->setClustersConfig(config, /*enable_discovery=*/false); + return context->getCluster("test_cluster"); +} + +ActionsDAG identityFilterDAG(const String & column_name) +{ + return ActionsDAG(NamesAndTypesList{{column_name, std::make_shared()}}); +} + +} + +/// Regression test: query plan optimizations call ReadFromCluster::applyFilters() more than +/// once as the plan is progressively refined (optimizeTreeSecondPass revisits the same subtree +/// whenever another pass reports an update; a GROUP BY on the partition/time column is one such +/// trigger). The task/file list must be built from the FINAL accumulated filter, not whatever +/// was known on the first call — building it early and freezing it silently drops conditions +/// discovered by later passes, which previously caused Iceberg partition/row-group pruning to +/// under-prune (icebergCluster() reading 6-13x more rows than the equivalent ice.`ns.table` +/// query for the same WHERE clause, whenever the plan needed more than one optimization pass to +/// stabilize). +TEST(ReadFromCluster, BuildsTaskIteratorFromFinalPredicateNotFirstSnapshot) +{ + const auto & context_holder = getContext(); + auto context = Context::createCopy(context_holder.context); + auto cluster = makeSingleUnreachableReplicaCluster(context); + + auto storage = std::make_shared("test_cluster", StorageID("db", "table"), getLogger("test")); + + auto metadata = std::make_shared(); + metadata->setColumns(ColumnsDescription{NamesAndTypesList{ + {"col_a", std::make_shared()}, + {"col_b", std::make_shared()}}}); + + auto storage_snapshot = std::make_shared(*storage, metadata); + + Block header; + header.insert({std::make_shared()->createColumn(), std::make_shared(), "col_a"}); + header.insert({std::make_shared()->createColumn(), std::make_shared(), "col_b"}); + auto shared_header = std::make_shared(header); + + SelectQueryInfo query_info; + + ReadFromCluster step( + Names{"col_a", "col_b"}, + query_info, + storage_snapshot, + context, + shared_header, + storage, + std::make_shared(), + QueryProcessingStage::Complete, + cluster, + getLogger("test"), + std::nullopt); + + /// Round 1: an early optimize pass only sees the condition on col_a (e.g. before a later + /// pass — such as aggregation-in-order for a GROUP BY on the partition column — inserts a + /// second FilterStep above this source). + step.addFilter(identityFilterDAG("col_a"), "col_a"); + step.applyFilters(); + + /// Round 2: a later pass has made col_b's condition visible above the source too. + /// optimizePrimaryKeyConditionAndLimit re-walks every FilterStep currently above the + /// source on each call, so both col_a and col_b are re-added here. + step.addFilter(identityFilterDAG("col_a"), "col_a"); + step.addFilter(identityFilterDAG("col_b"), "col_b"); + step.applyFilters(); + + QueryPipelineBuilder builder; + BuildQueryPipelineSettings settings(context); + try + { + step.initializePipeline(builder, settings); + } + catch (const Exception &) + { + /// Expected: the configured replica is unreachable. What matters is what predicate + /// getTaskIteratorExtension() was called with before that connection attempt. + } + + ASSERT_EQ(storage->call_count, 1u); + EXPECT_TRUE(storage->last_filter_input_names.contains("col_a")); + EXPECT_TRUE(storage->last_filter_input_names.contains("col_b")) + << "task iterator was built from a stale predicate that predates the second applyFilters() call"; +} From 56716214fc98c3ecca541c58858b7b942436de1d Mon Sep 17 00:00:00 2001 From: VighneshPath Date: Wed, 2 Sep 2026 13:50:27 +0530 Subject: [PATCH 3/3] Fix gtest_read_from_cluster_predicate_pushdown compile errors ASTPtr is boost::intrusive_ptr, not std::shared_ptr, so std::make_shared() didn't convert. And ReadFromCluster's own applyFilters(ActionDAGNodes) override hides the no-arg SourceStepWithFilterBase::applyFilters() by name across inheritance levels when called on the concrete ReadFromCluster type directly, so it needs explicit base-class qualification. Caught by the local unit_tests_dbms build (this test was never built before pushing). Co-Authored-By: Claude Sonnet 5 Signed-off-by: VighneshPath --- .../tests/gtest_read_from_cluster_predicate_pushdown.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp b/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp index c77cf7552b28..aaf2be6da384 100644 --- a/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp +++ b/src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp @@ -116,7 +116,7 @@ TEST(ReadFromCluster, BuildsTaskIteratorFromFinalPredicateNotFirstSnapshot) context, shared_header, storage, - std::make_shared(), + ASTPtr(new ASTSelectQuery), QueryProcessingStage::Complete, cluster, getLogger("test"), @@ -126,14 +126,14 @@ TEST(ReadFromCluster, BuildsTaskIteratorFromFinalPredicateNotFirstSnapshot) /// pass — such as aggregation-in-order for a GROUP BY on the partition column — inserts a /// second FilterStep above this source). step.addFilter(identityFilterDAG("col_a"), "col_a"); - step.applyFilters(); + step.SourceStepWithFilterBase::applyFilters(); /// Round 2: a later pass has made col_b's condition visible above the source too. /// optimizePrimaryKeyConditionAndLimit re-walks every FilterStep currently above the /// source on each call, so both col_a and col_b are re-added here. step.addFilter(identityFilterDAG("col_a"), "col_a"); step.addFilter(identityFilterDAG("col_b"), "col_b"); - step.applyFilters(); + step.SourceStepWithFilterBase::applyFilters(); QueryPipelineBuilder builder; BuildQueryPipelineSettings settings(context);