Fix more tests.

src/Processors/QueryPlan/Optimizations/optimizeUseProjections.cpp

@@ -504,6 +504,15 @@ bool optimizeUseAggProjections(QueryPlan::Node & node, QueryPlan::Nodes & nodes)
     if (reading->hasAnalyzedResult())
         return false;
 
+    if (reading->isQueryWithFinal())
+        return false;
+
+    if (reading->isQueryWithSampling())
+        return false;
+
+    if (reading->isParallelReadingEnabled())
+        return false;
+
     LOG_TRACE(&Poco::Logger::get("optimizeUseProjections"), "Try optimize projection 5");
 
     const auto metadata = reading->getStorageMetadata();
@@ -792,6 +801,15 @@ bool optimizeUseNormalProjections(Stack & stack, QueryPlan::Nodes & nodes)
     if (reading->hasAnalyzedResult())
         return false;
 
+    if (reading->isQueryWithFinal())
+        return false;
+
+    if (reading->isQueryWithSampling())
+        return false;
+
+    if (reading->isParallelReadingEnabled())
+        return false;
+
     auto iter = stack.rbegin();
     while (iter != stack.rend())
     {

src/Processors/QueryPlan/ReadFromMergeTree.cpp

@@ -1505,6 +1505,18 @@ bool ReadFromMergeTree::isQueryWithFinal() const
         return select.final();
 }
 
+bool ReadFromMergeTree::isQueryWithSampling() const
+{
+    if (context->getSettingsRef().parallel_replicas_count > 1 && data.supportsSampling())
+        return true;
+
+    const auto & select = query_info.query->as<ASTSelectQuery &>();
+    if (query_info.table_expression_modifiers)
+        return query_info.table_expression_modifiers->getSampleSizeRatio() != std::nullopt;
+    else
+        return select.sampleSize() != nullptr;
+}
+
 Pipe ReadFromMergeTree::spreadMarkRanges(
     RangesInDataParts && parts_with_ranges, size_t num_streams, AnalysisResult & result, ActionsDAGPtr & result_projection)
 {

src/Processors/QueryPlan/ReadFromMergeTree.h

@@ -167,6 +167,8 @@ public:
     bool requestReadingInOrder(size_t prefix_size, int direction, size_t limit);
 
     static bool isFinal(const SelectQueryInfo & query_info);
+    bool isQueryWithFinal() const;
+    bool isQueryWithSampling() const;
 
     /// Returns true if the optimisation is applicable (and applies it then).
     bool requestOutputEachPartitionThroughSeparatePort();
@@ -197,8 +199,6 @@ private:
         bool sample_factor_column_queried,
         Poco::Logger * log);
 
-    bool isQueryWithFinal() const;
-
     int getSortDirection() const
     {
         const InputOrderInfoPtr & order_info = query_info.getInputOrderInfo();

tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.reference

@@ -1,3 +1,4 @@
+3
 0
 0
 2

tests/queries/0_stateless/01505_trivial_count_with_partition_predicate.sql

@@ -7,7 +7,7 @@ insert into test1 values ('2020-09-01 00:01:02', 1), ('2020-09-01 20:01:03', 2),
 
 set max_rows_to_read = 1;
 -- non-optimized
-select count() from test1 settings max_parallel_replicas = 3; -- { serverError 158; }
+select count() from test1 settings max_parallel_replicas = 3;
 -- optimized (toYear is monotonic and we provide the partition expr as is)
 select count() from test1 where toYear(toDate(p)) = 1999;
 -- non-optimized (toDate(DateTime) is always monotonic, but we cannot relaxing the predicates to do trivial count())

tests/queries/0_stateless/02319_no_columns_in_row_level_filter.reference

@@ -1,4 +1,4 @@
-1000000
+0
 0
 0
 0

tests/queries/0_stateless/02343_aggregation_pipeline.reference

@@ -92,12 +92,15 @@ ExpressionTransform × 16
   (MergingAggregated)
   Resize 1 → 16
     MergingAggregatedTransform
-      Resize 2 → 1
+      Resize 17 → 1
         (Union)
-          (ReadFromStorage)
-          AggregatingTransform
-            ExpressionTransform
-              MergeTreeInOrder 0 → 1
+          (Aggregating)
+          Resize 1 → 16
+            AggregatingTransform
+              (Expression)
+              ExpressionTransform
+                (ReadFromMergeTree)
+                MergeTreeInOrder 0 → 1
           (ReadFromRemote)
 explain pipeline SELECT k1, k3, sum(value) v FROM remote('127.0.0.{1,2}', currentDatabase(), proj_agg_02343) GROUP BY k1, k3 SETTINGS distributed_aggregation_memory_efficient = 1;
 (Expression)
@@ -109,9 +112,11 @@ ExpressionTransform × 16
         Resize 1 → 16
           GroupingAggregatedTransform 2 → 1
             (Union)
-              (ReadFromStorage)
+              (Aggregating)
               AggregatingTransform
+                (Expression)
                 ExpressionTransform
+                  (ReadFromMergeTree)
                   MergeTreeInOrder 0 → 1
               (ReadFromRemote)
 -- { echoOn }