From f4269ce41750648e3be629b08ce73e38afa273ae Mon Sep 17 00:00:00 2001
From: Amos Bird <amosbird@gmail.com>
Date: Wed, 8 Sep 2021 20:10:49 +0800
Subject: [PATCH] Allow optimize_arithmetic_operations_in_aggregate_functions
 when alias is used.

---
 .../ArithmeticOperationsInAgrFuncOptimize.cpp        | 12 +++++++-----
 ...etic_operations_in_aggr_func_with_alias.reference | 10 ++++++++++
 ...arithmetic_operations_in_aggr_func_with_alias.sql |  4 ++++
 .../0_stateless/01470_columns_transformers.reference |  4 ++--
 4 files changed, 23 insertions(+), 7 deletions(-)
 create mode 100644 tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.reference
 create mode 100644 tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.sql

diff --git a/src/Interpreters/ArithmeticOperationsInAgrFuncOptimize.cpp b/src/Interpreters/ArithmeticOperationsInAgrFuncOptimize.cpp
index bdd9ce32707..70a58971d3f 100644
--- a/src/Interpreters/ArithmeticOperationsInAgrFuncOptimize.cpp
+++ b/src/Interpreters/ArithmeticOperationsInAgrFuncOptimize.cpp
@@ -107,10 +107,7 @@ ASTPtr tryExchangeFunctions(const ASTFunction & func)
         || !supported.find(lower_name)->second.count(child_func->name))
         return {};
 
-    /// Cannot rewrite function with alias cause alias could become undefined
-    if (!func.tryGetAlias().empty() || !child_func->tryGetAlias().empty())
-        return {};
-
+    auto original_alias = func.tryGetAlias();
     const auto & child_func_args = child_func->arguments->children;
     const auto * first_literal = child_func_args[0]->as<ASTLiteral>();
     const auto * second_literal = child_func_args[1]->as<ASTLiteral>();
@@ -132,7 +129,12 @@ ASTPtr tryExchangeFunctions(const ASTFunction & func)
         optimized_ast = exchangeExtractSecondArgument(new_name, *child_func);
     }
 
-    return optimized_ast;
+    if (optimized_ast)
+    {
+        optimized_ast->setAlias(original_alias);
+        return optimized_ast;
+    }
+    return {};
 }
 
 }
diff --git a/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.reference b/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.reference
new file mode 100644
index 00000000000..9e0d871041b
--- /dev/null
+++ b/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.reference
@@ -0,0 +1,10 @@
+SELECT min(n AS a) + (1 AS b) AS c
+FROM
+(
+    SELECT number AS n
+    FROM numbers(10)
+    WHERE (1 > 0) AND (n > 0)
+)
+WHERE (a > 0) AND (b > 0)
+HAVING c > 0
+2
diff --git a/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.sql b/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.sql
new file mode 100644
index 00000000000..73b87817bb3
--- /dev/null
+++ b/tests/queries/0_stateless/01271_optimize_arithmetic_operations_in_aggr_func_with_alias.sql
@@ -0,0 +1,4 @@
+set optimize_arithmetic_operations_in_aggregate_functions = 1;
+
+explain syntax select min((n as a) + (1 as b)) c from (select number n from numbers(10)) where a > 0 and b > 0 having c > 0;
+select min((n as a) + (1 as b)) c from (select number n from numbers(10)) where a > 0 and b > 0 having c > 0;
diff --git a/tests/queries/0_stateless/01470_columns_transformers.reference b/tests/queries/0_stateless/01470_columns_transformers.reference
index ae0adb3ba60..8fa86582018 100644
--- a/tests/queries/0_stateless/01470_columns_transformers.reference
+++ b/tests/queries/0_stateless/01470_columns_transformers.reference
@@ -54,8 +54,8 @@ SELECT
     sum(k)
 FROM columns_transformers
 SELECT
-    avg(i + 1 AS i),
-    avg(j + 2 AS j),
+    avg(i) + 1,
+    avg(j) + 2,
     avg(k)
 FROM columns_transformers
 SELECT