Simplify PruneOrderByInAggregation

trinodb · Oct 28, 2024 · 024e7fb · 024e7fb
1 parent d64fc7f
commit 024e7fb
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/...ino-main/src/main/java/io/trino/sql/planner/iterative/rule/PruneOrderByInAggregation.java b/...ino-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) {