From 75597efeb690da09457292249d8653cf07cea402 Mon Sep 17 00:00:00 2001
From: Nikolai Kochetov <nik-kochetov@yandex-team.ru>
Date: Tue, 2 May 2023 14:58:44 +0000
Subject: [PATCH] Respect PK for aggregate projection.

---
 .../optimizeUseAggregateProjection.cpp        |  9 ++++++
 ...02725_agg_projection_resprect_PK.reference |  2 ++
 .../02725_agg_projection_resprect_PK.sql      | 31 +++++++++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 tests/queries/0_stateless/02725_agg_projection_resprect_PK.reference
 create mode 100644 tests/queries/0_stateless/02725_agg_projection_resprect_PK.sql

diff --git a/src/Processors/QueryPlan/Optimizations/optimizeUseAggregateProjection.cpp b/src/Processors/QueryPlan/Optimizations/optimizeUseAggregateProjection.cpp
index f0795aebf64..931e1ae1dd8 100644
--- a/src/Processors/QueryPlan/Optimizations/optimizeUseAggregateProjection.cpp
+++ b/src/Processors/QueryPlan/Optimizations/optimizeUseAggregateProjection.cpp
@@ -582,6 +582,9 @@ bool optimizeUseAggregateProjections(QueryPlan::Node & node, QueryPlan::Nodes &
     ContextPtr context = reading->getContext();
     MergeTreeDataSelectExecutor reader(reading->getMergeTreeData());
 
+    auto ordinary_reading_select_result = reading->selectRangesToRead(parts);
+    size_t ordinary_reading_marks = ordinary_reading_select_result->marks();
+
     /// Selecting best candidate.
     for (auto & candidate : candidates.real)
     {
@@ -597,12 +600,18 @@ bool optimizeUseAggregateProjections(QueryPlan::Node & node, QueryPlan::Nodes &
         if (!analyzed)
             continue;
 
+        if (candidate.sum_marks > ordinary_reading_marks)
+            continue;
+
         if (best_candidate == nullptr || best_candidate->sum_marks > candidate.sum_marks)
             best_candidate = &candidate;
     }
 
     if (!best_candidate)
+    {
+        reading->setAnalyzedResult(std::move(ordinary_reading_select_result));
         return false;
+    }
 
     QueryPlanStepPtr projection_reading;
     bool has_ordinary_parts;
diff --git a/tests/queries/0_stateless/02725_agg_projection_resprect_PK.reference b/tests/queries/0_stateless/02725_agg_projection_resprect_PK.reference
new file mode 100644
index 00000000000..e6b95502e1e
--- /dev/null
+++ b/tests/queries/0_stateless/02725_agg_projection_resprect_PK.reference
@@ -0,0 +1,2 @@
+      ReadFromMergeTree (p1)
+          Granules: 1/12
diff --git a/tests/queries/0_stateless/02725_agg_projection_resprect_PK.sql b/tests/queries/0_stateless/02725_agg_projection_resprect_PK.sql
new file mode 100644
index 00000000000..fd56e058d0a
--- /dev/null
+++ b/tests/queries/0_stateless/02725_agg_projection_resprect_PK.sql
@@ -0,0 +1,31 @@
+-- Tags: no-random-merge-tree-settings
+
+DROP TABLE IF EXISTS t0;
+
+CREATE TABLE t0
+(
+    c1 Int64,
+    c2 Int64,
+    c3 Int64,
+    PROJECTION p1
+    (
+        SELECT
+            c1,
+            c2,
+            sum(c3)
+        GROUP BY
+            c2,
+            c1
+    )
+)
+ENGINE = MergeTree ORDER BY (c1, c2);
+
+INSERT INTO t0 SELECT
+    number,
+    -number,
+    number
+FROM numbers_mt(1e5);
+
+select * from (EXPLAIN indexes = 1 SELECT c1, sum(c3) FROM t0 GROUP BY c1) where explain like '%ReadFromMergeTree%';
+select * from (EXPLAIN indexes = 1 SELECT c1, sum(c3) FROM t0 WHERE c1 = 100 GROUP BY c1) where explain like '%Granules%';
+