Merge pull request #27424 from azat/WithMergeableStateAfterAggregationAndLimit-tests-fixes

Fix tests for WithMergeableStateAfterAggregationAndLimit

src/Interpreters/InterpreterSelectQuery.cpp

@@ -1378,7 +1378,7 @@ void InterpreterSelectQuery::executeImpl(QueryPlan & query_plan, const BlockInpu
 
             /// Limit is no longer needed if there is prelimit.
             ///
-            /// NOTE: that LIMIT cannot be applied of OFFSET should not be applied,
+            /// NOTE: that LIMIT cannot be applied if OFFSET should not be applied,
             /// since LIMIT will apply OFFSET too.
             /// This is the case for various optimizations for distributed queries,
             /// and when LIMIT cannot be applied it will be applied on the initiator anyway.

src/Storages/StorageDistributed.cpp

@@ -553,7 +553,6 @@ std::optional<QueryProcessingStage::Enum> StorageDistributed::getOptimizedQueryP
     if (const ASTPtr order_by = select.orderBy())
         return default_stage;
 
-    // LIMIT BY
     // LIMIT
     // OFFSET
     if (select.limitLength() || select.limitOffset())

tests/queries/0_stateless/01951_distributed_push_down_limit.reference

@@ -1,7 +1,7 @@
 -- { echo }
 explain select * from remote('127.{1,2}', view(select * from numbers(1e6))) order by number limit 10 settings distributed_push_down_limit=0;
 Expression (Projection)
-  Limit (preliminary LIMIT)
+  Limit (preliminary LIMIT (without OFFSET))
     MergingSorted (Merge sorted streams after aggregation stage for ORDER BY)
       SettingQuotaAndLimits (Set limits and quota after reading from storage)
         Union
@@ -16,11 +16,11 @@ Expression (Projection)
           ReadFromRemote (Read from remote replica)
 explain select * from remote('127.{1,2}', view(select * from numbers(1e6))) order by number limit 10 settings distributed_push_down_limit=1;
 Expression (Projection)
-  Limit (preliminary LIMIT)
+  Limit (preliminary LIMIT (without OFFSET))
     MergingSorted (Merge sorted streams after aggregation stage for ORDER BY)
       SettingQuotaAndLimits (Set limits and quota after reading from storage)
         Union
-          Limit (preliminary LIMIT)
+          Limit (preliminary LIMIT (with OFFSET))
             MergingSorted (Merge sorted streams for ORDER BY)
               MergeSorting (Merge sorted blocks for ORDER BY)
                 PartialSorting (Sort each block for ORDER BY)

tests/queries/0_stateless/02003_WithMergeableStateAfterAggregationAndLimit_LIMIT_BY_LIMIT_OFFSET.reference

@@ -1,16 +1,17 @@
 -- { echo }
 SELECT *
 FROM remote('127.{1,2}', view(
-    SELECT *
-    FROM numbers(10)
+    SELECT number%20 number
+    FROM numbers(40)
+    WHERE (number % 2) = (shardNum() - 1)
 ), number)
 GROUP BY number
 ORDER BY number ASC
 LIMIT 1 BY number
 LIMIT 5, 5
 SETTINGS
-    optimize_skip_unused_shards = 1,
-    optimize_distributed_group_by_sharding_key = 1,
+    optimize_skip_unused_shards=1,
+    optimize_distributed_group_by_sharding_key=1,
     distributed_push_down_limit=1;
 5
 6
@@ -19,8 +20,9 @@ SETTINGS
 9
 SELECT *
 FROM remote('127.{1,2}', view(
-    SELECT *
-    FROM numbers(10)
+    SELECT number%20 number
+    FROM numbers(40)
+    WHERE (number % 2) = (shardNum() - 1)
 ), number)
 GROUP BY number
 ORDER BY number ASC

tests/queries/0_stateless/02003_WithMergeableStateAfterAggregationAndLimit_LIMIT_BY_LIMIT_OFFSET.sql

@@ -1,21 +1,31 @@
+-- Here we use a trick with shardNum() to generate unique data on each shard.
+-- Since distributed_group_by_no_merge=2 will use WithMergeableStateAfterAggregationAndLimit,
+-- which assume that the data on shards is unique
+-- (LIMIT BY will be applied only on shards, not on the initiator).
+
+-- To distinguish echoing from the comments above we use SELECT FORMAT Null.
+SELECT '' FORMAT Null;
+
 -- { echo }
 SELECT *
 FROM remote('127.{1,2}', view(
-    SELECT *
-    FROM numbers(10)
+    SELECT number%20 number
+    FROM numbers(40)
+    WHERE (number % 2) = (shardNum() - 1)
 ), number)
 GROUP BY number
 ORDER BY number ASC
 LIMIT 1 BY number
 LIMIT 5, 5
 SETTINGS
-    optimize_skip_unused_shards = 1,
-    optimize_distributed_group_by_sharding_key = 1,
+    optimize_skip_unused_shards=1,
+    optimize_distributed_group_by_sharding_key=1,
     distributed_push_down_limit=1;
 SELECT *
 FROM remote('127.{1,2}', view(
-    SELECT *
-    FROM numbers(10)
+    SELECT number%20 number
+    FROM numbers(40)
+    WHERE (number % 2) = (shardNum() - 1)
 ), number)
 GROUP BY number
 ORDER BY number ASC