Fix wrong ORDER BY results for Merge over Distributed

Project: ClickHouse/ClickHouse | Language: C++ | Concepts: query-optimization, consistency-models, data-structures, testing-strategy | Difficulty: senior
Link: https://github.com/ClickHouse/ClickHouse/pull/112931

The problem

When running a query like SELECT ... ORDER BY ... LIMIT on a Merge table that is defined over Distributed tables, ClickHouse could return rows in the wrong order. Worse, when a LIMIT clause is used, it could drop rows that actually belong in the result. This happens because the query execution pipeline narrows (combines) the sorted streams from each shard, but does not preserve their order, breaking the assumption that each input stream is sorted for the final merge.

Why it's non-trivial

A naive fix might just try to reorder the final results or always preserve all streams, but this would either not guarantee correctness (since the merging step expects each stream to be individually sorted), or would cause resource issues (too many open files/sockets) for queries where order doesn't matter. The fix must precisely detect when order preservation is required, and only skip narrowing in those cases—otherwise, performance and resource usage could degrade.

The approach

1. Detecting when ORDER BY is present

A helper function is added to check if the query has an ORDER BY clause, supporting both query tree and AST representations:

+bool queryHasOrderBy(const SelectQueryInfo & query_info)
+{
+    if (query_info.query_tree)
+    {
+        if (const auto * query_node = query_info.query_tree->as<QueryNode>())
+            return query_node->hasOrderBy();
+    }
+
+    if (query_info.query)
+    {
+        if (const auto * select = query_info.query->as<ASTSelectQuery>())
+            return select->orderBy() != nullptr;
+    }
+
+    return false;
+}

Annotation:
This function checks both modern (query_tree) and legacy (query) representations for an ORDER BY clause, ensuring compatibility with both query planners.

2. Deciding when to skip narrowing

The logic for when to skip the narrowing step is expanded to include cases where child streams are already sorted due to remote ORDER BY:

-    const bool should_not_narrow = query_info.input_order_info || (
-        context->getSettingsRef()[Setting::distributed_aggregation_memory_efficient]
-        && common_processed_stage == QueryProcessingStage::Enum::WithMergeableState);
+    //
+    // Order must be preserved as well when the children were read at a stage where the query's `ORDER BY` has already
+    // run remotely: every child then sorts on its own (a `Distributed` child sorts on the shards), so the step on top
+    // of `ReadFromMerge` is `Sorting (Merge sorted streams ... for ORDER BY)`, which requires each input stream to be
+    // sorted. Narrowing would feed it unsorted streams, silently producing a wrongly ordered - and, together with
+    // `LIMIT`, incomplete - result.
+    //
+    // That happens at any stage above `WithMergeableState` (the remote side did the full `ORDER BY`), and at
+    // `WithMergeableState` only for queries without aggregation and window functions - the same conditions under
+    // which the remote part of a distributed query performs the preliminary sort (and the planner merges sorted
+    // streams instead of doing a full sort on the initiator). For example, a window function query over `Distributed`
+    // is processed only up to `WithMergeableState` with no remote sort, so narrowing remains allowed.
+    const bool children_produce_sorted_streams = queryHasOrderBy(query_info)
+        && (common_processed_stage > QueryProcessingStage::WithMergeableState
+            || (common_processed_stage > QueryProcessingStage::FetchColumns && !query_info.need_aggregate
+                && !query_info.has_window));
+
+    // Memory efficient distributed aggregation delivers two-level blocks bucket by bucket, and that bucket order must be
+    // preserved. It can only happen when the query aggregates: without aggregation there are no buckets at all, so the
+    // setting alone - it is enabled by default - must not keep every shard's stream alive. Otherwise the very fan-out
+    // this optimization guards against would come back for all the other queries stopping at `WithMergeableState`, such
+    // as the window function queries above.
+    const bool memory_efficient_aggregation = query_info.need_aggregate
+        && context->getSettingsRef()[Setting::distributed_aggregation_memory_efficient]
+        && common_processed_stage == QueryProcessingStage::Enum::WithMergeableState;
+
+    const bool should_not_narrow = query_info.input_order_info
+        || children_produce_sorted_streams
+        || memory_efficient_aggregation;

Annotation:

• children_produce_sorted_streams is true if the query has ORDER BY and the children are at a stage where their output is already sorted (either above WithMergeableState, or at WithMergeableState but with no aggregation/window).
• memory_efficient_aggregation is true only if aggregation is needed and the special memory-efficient mode is enabled.
• should_not_narrow is now true if any of these conditions hold, so narrowing (which destroys per-shard order) is skipped only when necessary.

3. Regression and interaction tests

New stateless regression tests are added to verify the fix and to check that narrowing is only skipped when required:

Regression for the original bug (wrong order/limit results):

+-- Regression test for https://github.com/ClickHouse/ClickHouse/issues/79949
+-- A `Merge` table over `Distributed` tables reads its children at a stage above `FetchColumns`,
+-- so each child sorts on the shards and the step above `ReadFromMerge` only merges already
+-- sorted streams. Narrowing the pipe used to destroy that per-stream order, which returned
+-- wrongly ordered - and, together with `LIMIT`, incomplete - results.
...
+SELECT A FROM merge_one ORDER BY A DESC LIMIT 9998, 5 SETTINGS max_threads = 1;
+SELECT A FROM merge_one ORDER BY A DESC LIMIT 9998, 5 SETTINGS max_threads = 1, enable_analyzer = 0;
+-- The original report: a `Merge` over two `Distributed` tables. Expected: 9999 9999 9998 9998 9997.
+SELECT A FROM merge_two ORDER BY A DESC LIMIT 5
+    SETTINGS max_threads = 1, distributed_aggregation_memory_efficient = 0;
+SELECT A FROM merge_two ORDER BY A DESC LIMIT 5
+    SETTINGS max_threads = 1, distributed_aggregation_memory_efficient = 0, enable_analyzer = 0;

Annotation:
These queries check that the correct rows are returned in the correct order, both for a single and for multiple distributed tables under a Merge.

Test for window function and aggregation interaction (narrowing must still happen):

+-- Companion of 04666_merge_over_distributed_order_by. A `Merge` over `Distributed` must skip
+-- `narrowPipe` only when the children actually return per-stream sorted data. A window function
+-- query is processed on the shards only up to `WithMergeableState` with no remote `ORDER BY`
+-- (the initiator does a full sort above the window), so the pipe must still be narrowed - not
+-- narrowing it reopens the "too many files and sockets at the same time" problem this
+-- optimization guards against.
...
+SELECT countIf(explain LIKE '%× 4%')
+FROM
+(
+    EXPLAIN PIPELINE SELECT A, sum(A) OVER (ORDER BY A) FROM merge_dist ORDER BY A LIMIT 5
+)
+SETTINGS max_threads = 1, max_streams_multiplier_for_merge_tables = 1, distributed_aggregation_memory_efficient = 0;

Annotation:
This verifies that for window function queries, narrowing still happens (number of streams is reduced), so resource usage remains controlled.

From the review thread

No notable review discussion on this PR.

Takeaways

• Preserve stream order only when required by downstream logic. When merging sorted streams from distributed sources, only skip stream narrowing if the downstream step expects each stream to be sorted (e.g., for ORDER BY with LIMIT). Otherwise, narrowing should proceed to avoid resource exhaustion.
• Guard optimizations with precise query analysis. Use helper functions to accurately detect the presence of clauses like ORDER BY in both legacy and modern query representations, ensuring that optimizations do not break correctness.
• Regression tests should cover both the bug and its edge cases. When fixing a subtle bug, add tests that cover not only the original failure but also interactions with related query features (e.g., window functions, aggregation) to prevent future regressions.