From 024e7fbbb622fe033c2a0c3efc71158ff339eae5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Osipiuk?= <lukasz@osipiuk.net>
Date: Mon, 28 Oct 2024 18:49:26 +0100
Subject: [PATCH] Simplify PruneOrderByInAggregation

---
 .../rule/PruneOrderByInAggregation.java         | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/PruneOrderByInAggregation.java b/core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/PruneOrderByInAggregation.java
index 57d1c70af64935..6d6fa03b897d89 100644
--- a/core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/PruneOrderByInAggregation.java
+++ b/core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/PruneOrderByInAggregation.java
@@ -56,23 +56,20 @@ public Result apply(AggregationNode node, Captures captures, Context context)
         ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder();
         for (Map.Entry<Symbol, Aggregation> entry : node.getAggregations().entrySet()) {
             Aggregation aggregation = entry.getValue();
-            if (aggregation.getOrderingScheme().isEmpty()) {
-                aggregations.put(entry);
-            }
+
             // getAggregateFunctionImplementation can be expensive, so check it last.
-            else if (metadata.getAggregationFunctionMetadata(context.getSession(), aggregation.getResolvedFunction()).isOrderSensitive()) {
-                aggregations.put(entry);
-            }
-            else {
-                anyRewritten = true;
-                aggregations.put(entry.getKey(), new Aggregation(
+            if (aggregation.getOrderingScheme().isPresent() &&
+                    !metadata.getAggregationFunctionMetadata(context.getSession(), aggregation.getResolvedFunction()).isOrderSensitive()) {
+                aggregation = new Aggregation(
                         aggregation.getResolvedFunction(),
                         aggregation.getArguments(),
                         aggregation.isDistinct(),
                         aggregation.getFilter(),
                         Optional.empty(),
-                        aggregation.getMask()));
+                        aggregation.getMask());
+                anyRewritten = true;
             }
+            aggregations.put(entry.getKey(), aggregation);
         }
 
         if (!anyRewritten) {