Fix bucket-merge scheduling and ReadFromCluster filter staleness

src/Interpreters/Aggregator.cpp

@@ -3566,15 +3566,38 @@ bool Aggregator::mergeOnBlock(Columns columns, size_t rows, bool is_overflows, A
 }
 
 
+std::vector<Int32> Aggregator::sortBucketsByRowCountDescending(std::vector<std::pair<Int32, UInt64>> 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<Int32> 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<bool> & 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<std::pair<Int32, UInt64>> 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<UInt32> 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<Int32> 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<size_t> 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;

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<bool> & 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<Int32> sortBucketsByRowCountDescending(std::vector<std::pair<Int32, UInt64>> 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).

src/Interpreters/tests/gtest_aggregator_bucket_merge_order.cpp

@@ -0,0 +1,41 @@
+#include <algorithm>
+#include <vector>
+
+#include <Interpreters/Aggregator.h>
+
+#include <gtest/gtest.h>
+
+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<Int32>{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<Int32> tied{order[1], order[2]};
+    std::sort(tied.begin(), tied.end());
+    EXPECT_EQ(tied, (std::vector<Int32>{0, 1}));
+}

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);
 

src/Storages/tests/gtest_read_from_cluster_predicate_pushdown.cpp

@@ -0,0 +1,154 @@
+#include <set>
+#include <sstream>
+
+#include <gtest/gtest.h>
+
+#include <Poco/AutoPtr.h>
+#include <Poco/Util/XMLConfiguration.h>
+
+#include <Core/Block.h>
+#include <Core/NamesAndTypes.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Interpreters/ActionsDAG.h>
+#include <Interpreters/Context.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Processors/QueryPlan/BuildQueryPipelineSettings.h>
+#include <QueryPipeline/QueryPipelineBuilder.h>
+#include <Storages/ColumnsDescription.h>
+#include <Storages/IStorageCluster.h>
+#include <Storages/StorageInMemoryMetadata.h>
+#include <Storages/StorageSnapshot.h>
+#include <Common/Exception.h>
+#include <Common/Logger.h>
+#include <Common/tests/gtest_global_context.h>
+
+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<std::string> 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{
+        "<clickhouse><remote_servers><test_cluster><shard><replica>"
+        "<host>127.0.0.1</host><port>1</port>"
+        "</replica></shard></test_cluster></remote_servers></clickhouse>"};
+    Poco::AutoPtr<Poco::Util::XMLConfiguration> 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<DataTypeUInt8>()}});
+}
+
+}
+
+/// 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<RecordingStorageCluster>("test_cluster", StorageID("db", "table"), getLogger("test"));
+
+    auto metadata = std::make_shared<StorageInMemoryMetadata>();
+    metadata->setColumns(ColumnsDescription{NamesAndTypesList{
+        {"col_a", std::make_shared<DataTypeUInt8>()},
+        {"col_b", std::make_shared<DataTypeUInt8>()}}});
+
+    auto storage_snapshot = std::make_shared<StorageSnapshot>(*storage, metadata);
+
+    Block header;
+    header.insert({std::make_shared<DataTypeUInt8>()->createColumn(), std::make_shared<DataTypeUInt8>(), "col_a"});
+    header.insert({std::make_shared<DataTypeUInt8>()->createColumn(), std::make_shared<DataTypeUInt8>(), "col_b"});
+    auto shared_header = std::make_shared<const Block>(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";
+}