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})); +} 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..aaf2be6da384 --- /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, + ASTPtr(new ASTSelectQuery), + 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.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.SourceStepWithFilterBase::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"; +}