Antalya 26.6 Backport of #128865 - Fix wrong ORDER BY results and pushed bucket twice with offset based parallel replicas

src/Core/Settings.cpp

@@ -7535,6 +7535,8 @@ Possible values:
 
 This options will produce different results depending on the settings used.
 
+With `parallel_replicas_mode` `sampling_key`, or a custom key mode on a cluster with more than one shard, each replica returns its own part of the result of the shard. The replicas of a shard are used together only when the initiator aggregates or fully sorts these parts itself (for example an aggregation without `GROUP BY` keys). Otherwise, for example for `GROUP BY` with `distributed_aggregation_memory_efficient` or for `ORDER BY`, the shard is read from one replica.
+
 ### Parallel processing using `SAMPLE` key
 
 A query may be processed faster if it is executed on several servers in parallel. But the query performance may degrade in the following cases:

src/Processors/QueryPlan/MergingAggregatedStep.h

@@ -38,6 +38,7 @@ class MergingAggregatedStep : public ITransformingStep
     const SortDescription & getGroupBySortDescription() const { return group_by_sort_description; }
 
     bool memoryBoundMergingWillBeUsed() const;
+    bool isMemoryEfficient() const { return memory_efficient_aggregation; }
 
     bool isGroupingSets() const { return !grouping_sets_params.empty(); }
     const auto & getGroupingSetsParamsList() const { return grouping_sets_params; }

src/Processors/QueryPlan/Optimizations/Optimizations.h

@@ -243,6 +243,9 @@ void optimizeJoinLogical(QueryPlan::Node & node, QueryPlan::Nodes &, const Query
 /// A separate tree traverse to apply sorting properties after *InOrder optimizations.
 void applyOrder(const QueryPlanOptimizationSettings & optimization_settings, QueryPlan::Node & root);
 
+/// Lets an offset based parallel replicas read use several replicas of a shard only if no consumer needs stream order.
+void allowReadingShardsFromSeveralReplicas(QueryPlan::Node & root);
+
 /// Returns the name of used projection or nullopt if no projection is used.
 std::optional<String> optimizeUseAggregateProjections(
     QueryPlan::Node & node,

src/Processors/QueryPlan/Optimizations/allowReadingShardsFromSeveralReplicas.cpp

@@ -0,0 +1,84 @@
+#include <Processors/QueryPlan/AggregatingStep.h>
+#include <Processors/QueryPlan/ExpressionStep.h>
+#include <Processors/QueryPlan/FilterStep.h>
+#include <Processors/QueryPlan/MergingAggregatedStep.h>
+#include <Processors/QueryPlan/Optimizations/Optimizations.h>
+#include <Processors/QueryPlan/Optimizations/Utils.h>
+#include <Processors/QueryPlan/ReadFromRemote.h>
+#include <Processors/QueryPlan/SortingStep.h>
+#include <Processors/QueryPlan/UnionStep.h>
+#include <Common/typeid_cast.h>
+
+#include <unordered_map>
+
+namespace DB::QueryPlanOptimizations
+{
+
+namespace
+{
+
+enum class StreamOrderUse : uint8_t
+{
+    /// The step keeps its input streams apart and in order, so its own consumer decides.
+    PassedOn,
+    Ignored,
+    Required,
+};
+
+StreamOrderUse getStreamOrderUse(const IQueryPlanStep & step)
+{
+    if (typeid_cast<const ExpressionStep *>(&step) || typeid_cast<const FilterStep *>(&step) || typeid_cast<const UnionStep *>(&step))
+        return StreamOrderUse::PassedOn;
+
+    if (const auto * merging_aggregated = typeid_cast<const MergingAggregatedStep *>(&step))
+    {
+        /// Without keys the result is never two-level, so the memory efficient merge waits for every input.
+        if (!merging_aggregated->memoryBoundMergingWillBeUsed()
+            && (!merging_aggregated->isMemoryEfficient() || merging_aggregated->getParams().keys.empty()))
+            return StreamOrderUse::Ignored;
+        return StreamOrderUse::Required;
+    }
+
+    if (const auto * aggregating = typeid_cast<const AggregatingStep *>(&step))
+        return aggregating->inOrder() ? StreamOrderUse::Required : StreamOrderUse::Ignored;
+
+    if (const auto * sorting = typeid_cast<const SortingStep *>(&step))
+        return sorting->getType() == SortingStep::Type::Full ? StreamOrderUse::Ignored : StreamOrderUse::Required;
+
+    return StreamOrderUse::Required;
+}
+
+}
+
+void allowReadingShardsFromSeveralReplicas(QueryPlan::Node & root)
+{
+    std::unordered_map<ReadFromRemote *, bool> allowed;
+    Stack stack;
+
+    traverseQueryPlan(stack, root, [&](QueryPlan::Node & node)
+    {
+        auto * reading = typeid_cast<ReadFromRemote *>(node.step.get());
+        if (!reading)
+            return;
+
+        /// The consumer of the plan root is outside of the plan.
+        bool allow = false;
+        for (size_t i = stack.size() - 1; i > 0; --i)
+        {
+            const auto use = getStreamOrderUse(*stack[i - 1].node->step);
+            if (use == StreamOrderUse::PassedOn)
+                continue;
+            allow = use == StreamOrderUse::Ignored;
+            break;
+        }
+
+        auto [it, inserted] = allowed.emplace(reading, allow);
+        if (!inserted)
+            it->second = it->second && allow;
+    });
+
+    for (const auto & [reading, allow] : allowed)
+        reading->allowReadingShardFromSeveralReplicas(allow);
+}
+
+}

src/Processors/QueryPlan/Optimizations/optimizeTree.cpp

@@ -584,6 +584,9 @@ void optimizeTreeSecondPass(
         optimizeJoinByShards(root);
 
     considerEnablingParallelReplicas(optimization_settings, root, query_plan);
+
+    /// Last: it reads the final consumers of every remote read.
+    allowReadingShardsFromSeveralReplicas(root);
 }
 
 void addStepsToBuildSets(

src/Processors/QueryPlan/ReadFromRemote.cpp

@@ -842,7 +842,7 @@ void ReadFromRemote::addPipe(
         remote_query_executor->setRemoteFunction(is_remote_function);
         remote_query_executor->setShardCount(shard_count);
 
-        if (context->canUseTaskBasedParallelReplicas() || parallel_replicas_disabled)
+        if (context->canUseTaskBasedParallelReplicas() || parallel_replicas_disabled || !allow_reading_shard_from_several_replicas)
         {
             // when doing parallel reading from replicas (ParallelReplicasMode::READ_TASKS) on a shard:
             // establish a connection to a replica on the shard, the replica will instantiate coordinator to manage parallel reading from replicas on the shard.
@@ -852,6 +852,9 @@ void ReadFromRemote::addPipe(
             // each will execute parallel reading from replicas, so the query result will be multiplied by the number of created coordinators
             //
             // In case parallel replicas are disabled, there also should be a single connection to each shard to prevent result duplication
+            //
+            // With the offset based modes every replica returns its own part of the shard's result, and the replicas
+            // share one port, so a read whose consumer needs stream order uses one replica.
             remote_query_executor->setPoolMode(PoolMode::GET_ONE);
         }
         else
@@ -885,6 +888,9 @@ Pipes ReadFromRemote::addPipes(const ClusterProxy::SelectStreamFactory::Shards &
 
 void ReadFromRemote::initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &)
 {
+    if (context->canUseOffsetParallelReplicas() && !allow_reading_shard_from_several_replicas)
+        LOG_DEBUG(log, "Reading every shard from one replica, the consumers of the read may depend on the order of its streams");
+
     Pipes pipes = addPipes(shards, output_header);
 
     auto pipe = Pipe::unitePipes(std::move(pipes));

src/Processors/QueryPlan/ReadFromRemote.h

@@ -55,6 +55,8 @@ class ReadFromRemote final : public SourceStepWithFilterBase
 
     bool hasSerializedPlan() const;
 
+    void allowReadingShardFromSeveralReplicas(bool allow) { allow_reading_shard_from_several_replicas = allow; }
+
 private:
     ClusterProxy::SelectStreamFactory::Shards shards;
     QueryProcessingStage::Enum stage;
@@ -71,6 +73,7 @@ class ReadFromRemote final : public SourceStepWithFilterBase
     UnavailableShardTrackerPtr unavailable_shard_tracker;
     std::optional<GetPriorityForLoadBalancing> priority_func_factory;
     bool is_remote_function = false;
+    bool allow_reading_shard_from_several_replicas = false;
 
     Pipes addPipes(const ClusterProxy::SelectStreamFactory::Shards & used_shards, const SharedHeader & out_header);
 

tests/integration/test_parallel_replicas_custom_key/test.py

@@ -89,6 +89,8 @@ def test_parallel_replicas_custom_key_distributed(
                 "parallel_replicas_mode": parallel_replicas_mode,
                 "prefer_localhost_replica": 0,
                 "serialize_query_plan": 0,
+                # with a memory efficient merge every shard is read from one replica
+                "distributed_aggregation_memory_efficient": 0,
             },
         )
         == expected_result

tests/queries/0_stateless/05483_offset_parallel_replicas_ordered_merges.reference

@@ -0,0 +1,51 @@
+30000	60000	2	2
+30000	60000	2	2
+30000	30000	1	1
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+0
+0
+1
+1
+2
+2
+3
+3
+4
+4
+30000
+30000
+30000	30000
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+449985000
+15000
+224985000
+899970000
+offset_pr_aggregating	3
+offset_pr_count	3
+offset_pr_expression	3
+offset_pr_filter	3
+offset_pr_full_sort	3
+offset_pr_group_by	1
+offset_pr_limit	1
+offset_pr_memory_bound	1
+offset_pr_plain_merge	3
+offset_pr_root	1
+offset_pr_union	6

tests/queries/0_stateless/05483_offset_parallel_replicas_ordered_merges.sql

@@ -0,0 +1,104 @@
+-- Offset based parallel replicas: every replica of a shard returns its own part of the result, which the
+-- initiator must still merge correctly when it relies on the order of the streams.
+
+DROP TABLE IF EXISTS t_offset_pr_key;
+DROP TABLE IF EXISTS t_offset_pr_sample;
+
+CREATE TABLE t_offset_pr_key (k UInt64) ENGINE = MergeTree ORDER BY k;
+INSERT INTO t_offset_pr_key SELECT number FROM numbers(30000);
+
+CREATE TABLE t_offset_pr_sample (k UInt64) ENGINE = MergeTree ORDER BY intHash32(k) SAMPLE BY intHash32(k);
+INSERT INTO t_offset_pr_sample SELECT number FROM numbers(30000);
+
+SET enable_parallel_replicas = 1, max_parallel_replicas = 3, automatic_parallel_replicas_mode = 0,
+    parallel_replicas_for_non_replicated_merge_tree = 1, prefer_localhost_replica = 0;
+-- Two-level results, merged bucket by bucket.
+SET distributed_aggregation_memory_efficient = 1, group_by_two_level_threshold = 1, group_by_two_level_threshold_bytes = 1,
+    enable_producing_buckets_out_of_order_in_aggregation = 0, optimize_aggregation_in_order = 0;
+
+SELECT count(), sum(c), min(c), max(c)
+FROM (SELECT k, count() AS c FROM cluster('test_cluster_two_shard_three_replicas_localhost', currentDatabase(), t_offset_pr_key) GROUP BY k)
+SETTINGS parallel_replicas_mode = 'custom_key_sampling', parallel_replicas_custom_key = 'k';
+
+SELECT count(), sum(c), min(c), max(c)
+FROM (SELECT k, count() AS c FROM cluster('test_cluster_two_shard_three_replicas_localhost', currentDatabase(), t_offset_pr_key) GROUP BY k)
+SETTINGS parallel_replicas_mode = 'custom_key_range', parallel_replicas_custom_key = 'cityHash64(k)';
+
+SELECT count(), sum(c), min(c), max(c)
+FROM (SELECT k, count() AS c FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample) GROUP BY k)
+SETTINGS parallel_replicas_mode = 'sampling_key';
+
+SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample) ORDER BY k LIMIT 10
+SETTINGS parallel_replicas_mode = 'sampling_key';
+
+SELECT k FROM cluster('test_cluster_two_shard_three_replicas_localhost', currentDatabase(), t_offset_pr_key) ORDER BY k LIMIT 10
+SETTINGS parallel_replicas_mode = 'custom_key_sampling', parallel_replicas_custom_key = 'k';
+
+-- Without GROUP BY keys the parts are merged in any order, so the replicas still share the shard.
+SELECT count() FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_count';
+
+SELECT count() FROM (SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample) GROUP BY k)
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_group_by';
+
+-- Merged all at once instead of bucket by bucket, so the replicas still share the shard.
+SELECT count(), sum(c) FROM (SELECT k, count() AS c FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample) GROUP BY k)
+SETTINGS parallel_replicas_mode = 'sampling_key', distributed_aggregation_memory_efficient = 0,
+    enable_memory_bound_merging_of_aggregation_results = 0, use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_plain_merge';
+
+-- With a local shard and aggregation in order the states are merged in order, so the remote shard is read from one replica.
+SELECT k, count() FROM cluster('test_cluster_two_shard_three_replicas_localhost', currentDatabase(), t_offset_pr_key) GROUP BY k
+SETTINGS parallel_replicas_mode = 'custom_key_sampling', parallel_replicas_custom_key = 'k', prefer_localhost_replica = 1,
+    optimize_aggregation_in_order = 1, distributed_aggregation_memory_efficient = 0, enable_memory_bound_merging_of_aggregation_results = 1,
+    use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_memory_bound'
+FORMAT Null;
+
+-- A full sort and an aggregation that is not in order take the parts in any order too.
+SELECT k FROM (SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)) ORDER BY k LIMIT 10
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_full_sort';
+
+SELECT sum(k) FROM (SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample))
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_aggregating';
+
+-- Expressions, filters and unions between the read and such a consumer are looked through.
+SELECT sum(k % 2) FROM (SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample))
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_expression';
+
+SELECT sum(k) FROM (SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)) WHERE k % 2 = 0
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_filter';
+
+SELECT sum(k) FROM
+(
+    SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)
+    UNION ALL
+    SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)
+)
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_union';
+
+-- Any other consumer, and the plan root, read a shard from one replica.
+-- The limit is above the row count, so no secondary query is cancelled.
+SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample) LIMIT 100000
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_limit'
+FORMAT Null;
+
+SELECT k FROM remote('127.0.0.{1|2|3}', currentDatabase(), t_offset_pr_sample)
+SETTINGS parallel_replicas_mode = 'sampling_key', use_hedged_requests = 0, log_queries = 1, log_comment = 'offset_pr_root'
+FORMAT Null;
+
+SYSTEM FLUSH LOGS query_log;
+
+SELECT initial.log_comment, count()
+FROM system.query_log AS secondary
+INNER JOIN
+(
+    SELECT query_id, log_comment FROM system.query_log
+    WHERE event_date >= yesterday() AND current_database = currentDatabase() AND is_initial_query AND type = 'QueryFinish'
+        AND log_comment IN ('offset_pr_count', 'offset_pr_group_by', 'offset_pr_plain_merge', 'offset_pr_full_sort', 'offset_pr_aggregating', 'offset_pr_expression', 'offset_pr_filter', 'offset_pr_union', 'offset_pr_limit', 'offset_pr_root', 'offset_pr_memory_bound')
+) AS initial ON secondary.initial_query_id = initial.query_id
+WHERE secondary.event_date >= yesterday() AND NOT secondary.is_initial_query AND secondary.type = 'QueryFinish'
+GROUP BY initial.log_comment
+ORDER BY initial.log_comment
+SETTINGS enable_parallel_replicas = 0;
+
+DROP TABLE t_offset_pr_sample;
+DROP TABLE t_offset_pr_key;