Improve aggregation wrapper rewriting and merge evaluation

Author: somiljain2006

RELEASE-NOTES.md

@@ -72,7 +72,7 @@
 1. Sharding: Fix HASH_MOD routing mismatch for same negative numeric values across numeric Java types with compatibility switch `normalize-numeric-int-range` - [#38327](https://github.com/apache/shardingsphere/pull/38327)
 1. Proxy Native: Support building Proxy Native via GraalVM CE for JDK 25 - [#38682](https://github.com/apache/shardingsphere/pull/38682)
 1. SQL Parser: Support SQLServer table variable declaration parse - [#38904](https://github.com/apache/shardingsphere/pull/38904)
-1. Support evaluation of IFNULL/COALESCE expressions over merged aggregation results - [#38990](https://github.com/apache/shardingsphere/pull/38990)
+1. Support evaluation of IFNULL/COALESCE scalar wrappers (with literal or nested-aggregation fallbacks) over merged aggregation results - [#38990](https://github.com/apache/shardingsphere/pull/38990)
 
 ## Release 5.5.3
 

docs/document/content/reference/sharding/merge.en.md

@@ -67,6 +67,14 @@ Stream group-by merger is different from order-by merger only in two aspects:
 1. It will take out all the data with the same group item from multiple data result sets at once.
 1. It carried out the aggregation calculation according to the aggregation function type.
 
+Expression-derived aggregations
+
+In many real-world queries, aggregation functions are wrapped inside scalar expressions or conditional functions (for example, `IFNULL(SUM(score), 0)` or `COALESCE(COUNT(id), 0)`). The merger engine now natively evaluates such expression-derived aggregations during both stream and memory merges. Expression nodes that wrap aggregation calls are recognized and the underlying aggregation is computed as part of the usual merge flow; surrounding scalar or conditional expressions are then applied to the aggregated result so that semantics (null coalescing, defaulting, arithmetic, etc.) are preserved.
+
+Proper empty / no-route initializations
+
+When a group has no matching rows (including cases where a shard has no route for the query), expression-derived aggregations that depend on counts are initialized safely. In particular, `COUNT` cells are initialized to `0` (rather than relying on generic defaults) so that subsequent scalar expressions (e.g., `IFNULL`, `COALESCE`, or arithmetic) produce correct and deterministic results for empty or unrouted groups.
+
 For the inconsistency between the grouping item and ordering item, it requires uploading all the data to the memory to group and aggregate, since the relevant data value needed to acquire group information is not continuous, and stream merger is not available. For example, acquire each examinee’s total score through the following SQL and order them from the highest to the lowest:
 
 ```sql
@@ -87,6 +95,16 @@ The sum aggregation function refers to `SUM` and `COUNT`. They need to sum up al
 
 The average aggregation function refers only to `AVG`. It must be calculated through `SUM` and `COUNT` rewritten by SQL, which has been mentioned in the SQL rewriting section.
 
+## Performance / Memory Optimization
+
+Zero-copy projection pass-through
+
+The group-by merger avoids creating redundant collection wrappers (for example, copying expanded projection results into new `ArrayList` instances) when evaluating projections during merges. Instead, the merger streams directly against existing projection context references and iterates over the original projection objects where possible. This zero-copy pass-through reduces temporary allocations and keeps the garbage collection profile low during large merges.
+
+Elimination of quadratic overhead
+
+Previous implementations performed repeated linked-list scans when resolving projections across expanded projection sets, which could lead to quadratic (`O(N^2)`) behavior for large projection counts. Those linked-list scans have been replaced by optimized linear lookup mappings (e.g., direct index maps or hash-based lookups) so projection evaluation now completes in linear time relative to the number of projections.
+
 ## Pagination Merger
 
 All the merger types above can be paginated. Pagination is the decorator added to other kinds of mergers. ShardingSphere strengthens its ability to paginate the data result set through decorator mode. The pagination merger is responsible for filtering unnecessary data.

features/sharding/core/src/main/java/org/apache/shardingsphere/sharding/merge/dql/groupby/AggregationWrapperExpressionEvaluator.java

@@ -24,6 +24,7 @@
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.simple.LiteralExpressionSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.item.AggregationProjectionSegment;
 
+import java.math.BigDecimal;
 import java.util.List;
 import java.util.function.Function;
 
@@ -37,41 +38,73 @@ public final class AggregationWrapperExpressionEvaluator {
      * @param expression expression segment to evaluate
      * @param derivedAggregations derived aggregations
      * @param row memory query result row
+     * @param targetType the target JDBC class type to coerce to (optional)
      * @return evaluated result
      */
-    public static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final MemoryQueryResultRow row) {
-        return evaluate(expression, derivedAggregations, row::getCell);
+    public static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final MemoryQueryResultRow row, final Class<?> targetType) {
+        Object result = evaluate(expression, derivedAggregations, row::getCell);
+        return coerce(result, targetType);
     }
 
     /**
      * Evaluate expression with current row data.
      * @param expression expression segment to evaluate
      * @param derivedAggregations derived aggregations
      * @param currentRow current row data
+     * @param targetType the target JDBC class type to coerce to (optional)
      * @return evaluated result
      */
-    public static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final List<Object> currentRow) {
-        return evaluate(expression, derivedAggregations, index -> currentRow.get(index - 1));
+    public static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final List<Object> currentRow, final Class<?> targetType) {
+        Object result = evaluate(expression, derivedAggregations, index -> currentRow.get(index - 1));
+        return coerce(result, targetType);
     }
 
-    private static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueProvider) {
+    private static Object evaluate(final ExpressionSegment expression, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueExtractor) {
         if (expression instanceof AggregationProjectionSegment) {
-            return getMergedAggregationValue((AggregationProjectionSegment) expression, derivedAggregations, valueProvider);
+            return getMergedAggregationValue((AggregationProjectionSegment) expression, derivedAggregations, valueExtractor);
         }
         if (expression instanceof LiteralExpressionSegment) {
             return ((LiteralExpressionSegment) expression).getLiterals();
         }
         if (expression instanceof FunctionSegment) {
-            return evaluateFunction((FunctionSegment) expression, derivedAggregations, valueProvider);
+            return evaluateFunction((FunctionSegment) expression, derivedAggregations, valueExtractor);
         }
         throw new IllegalArgumentException(String.format("Unsupported aggregation wrapper expression segment type: %s", expression.getClass().getName()));
     }
 
-    private static Object evaluateFunction(final FunctionSegment functionSegment, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueProvider) {
+    private static Object coerce(final Object value, final Class<?> targetType) {
+        if (value == null || targetType == null || value.getClass().equals(targetType)) {
+            return value;
+        }
+        if (value instanceof Number) {
+            Number num = (Number) value;
+            if (targetType == BigDecimal.class) {
+                return new BigDecimal(num.toString());
+            }
+            if (targetType == Long.class) {
+                return num.longValue();
+            }
+            if (targetType == Integer.class) {
+                return num.intValue();
+            }
+            if (targetType == Double.class) {
+                return num.doubleValue();
+            }
+            if (targetType == Float.class) {
+                return num.floatValue();
+            }
+            if (targetType == Short.class) {
+                return num.shortValue();
+            }
+        }
+        return value;
+    }
+    
+    private static Object evaluateFunction(final FunctionSegment functionSegment, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueExtractor) {
         String functionName = functionSegment.getFunctionName();
         if ("IFNULL".equalsIgnoreCase(functionName) || "COALESCE".equalsIgnoreCase(functionName)) {
             for (ExpressionSegment each : functionSegment.getParameters()) {
-                Object value = evaluate(each, derivedAggregations, valueProvider);
+                Object value = evaluate(each, derivedAggregations, valueExtractor);
                 if (null != value) {
                     return value;
                 }
@@ -81,10 +114,10 @@ private static Object evaluateFunction(final FunctionSegment functionSegment, fi
         throw new IllegalArgumentException(String.format("Unsupported aggregation wrapper function: %s", functionName));
     }
 
-    private static Object getMergedAggregationValue(final AggregationProjectionSegment segment, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueProvider) {
+    private static Object getMergedAggregationValue(final AggregationProjectionSegment segment, final List<AggregationProjection> derivedAggregations, final Function<Integer, Object> valueExtractor) {
         for (AggregationProjection each : derivedAggregations) {
             if (each.getExpression().equals(segment.getText())) {
-                return valueProvider.apply(each.getIndex());
+                return valueExtractor.apply(each.getIndex());
             }
         }
         throw new IllegalArgumentException(String.format("Cannot find merged aggregation value for expression: %s", segment.getText()));

features/sharding/core/src/main/java/org/apache/shardingsphere/sharding/merge/dql/groupby/GroupByMemoryMergedResult.java

@@ -42,12 +42,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 
 /**
  * Memory merged result for group by.
@@ -83,9 +80,13 @@ private void initForFirstGroupByValue(final SelectStatementContext selectStateme
         if (!dataMap.containsKey(groupByValue)) {
             dataMap.put(groupByValue, new MemoryQueryResultRow(queryResult));
         }
-        aggregationMap.computeIfAbsent(groupByValue, unused -> selectStatementContext.getProjectionsContext().getAggregationProjections().stream()
-                .collect(Collectors.toMap(Function.identity(),
-                        input -> AggregationUnitFactory.create(input.getType(), input instanceof AggregationDistinctProjection, input.getSeparator().orElse(null)))));
+        if (!aggregationMap.containsKey(groupByValue)) {
+            Map<AggregationProjection, AggregationUnit> units = new HashMap<>(selectStatementContext.getProjectionsContext().getAggregationProjections().size(), 1F);
+            for (AggregationProjection each : selectStatementContext.getProjectionsContext().getAggregationProjections()) {
+                units.put(each, AggregationUnitFactory.create(each.getType(), each instanceof AggregationDistinctProjection, each.getSeparator().orElse(null)));
+            }
+            aggregationMap.put(groupByValue, units);
+        }
     }
 
     private void aggregate(final SelectStatementContext selectStatementContext, final QueryResult queryResult,
@@ -152,7 +153,13 @@ private List<MemoryQueryResultRow> getMemoryResultSetRows(final SelectStatementC
             }
             Object[] data = generateReturnData(selectStatementContext);
             MemoryQueryResultRow syntheticRow = new MemoryQueryResultRow(data);
-            evaluateExpressionValue(selectStatementContext, syntheticRow);
+            
+            Map<ExpressionProjection, List<AggregationProjection>> expressionDerivedAggregations = selectStatementContext.getProjectionsContext().getExpressionDerivedAggregations();
+            if (expressionDerivedAggregations != null && !expressionDerivedAggregations.isEmpty()) {
+                Map<ExpressionProjection, Integer> expressionIndices = getExpressionIndices(selectStatementContext, expressionDerivedAggregations);
+                evaluateExpressionValue(expressionDerivedAggregations, expressionIndices, syntheticRow);
+            }
+            
             return Collections.singletonList(syntheticRow);
         }
         List<MemoryQueryResultRow> result = new ArrayList<>(dataMap.values());
@@ -161,45 +168,88 @@ private List<MemoryQueryResultRow> getMemoryResultSetRows(final SelectStatementC
     }
 
     private Object[] generateReturnData(final SelectStatementContext selectStatementContext) {
-        List<Projection> projections = new LinkedList<>(selectStatementContext.getProjectionsContext().getExpandProjections());
+        int maxColumnIndex = calculateMaxColumnIndex(selectStatementContext);
+        Object[] result = new Object[maxColumnIndex];
+        
+        List<Projection> expandProjections = selectStatementContext.getProjectionsContext().getExpandProjections();
+        for (int i = 0; i < expandProjections.size(); i++) {
+            if (expandProjections.get(i) instanceof AggregationProjection && AggregationType.COUNT == ((AggregationProjection) expandProjections.get(i)).getType()) {
+                result[i] = 0;
+            }
+        }
 
-        int maxColumnIndex = projections.size();
         for (AggregationProjection each : selectStatementContext.getProjectionsContext().getAggregationProjections()) {
-            maxColumnIndex = Math.max(maxColumnIndex, each.getIndex());
+            if (AggregationType.COUNT == each.getType() && each.getIndex() > 0) {
+                result[each.getIndex() - 1] = 0;
+            }
         }
 
-        Object[] result = new Object[maxColumnIndex];
-        for (int i = 0; i < projections.size(); i++) {
-            if (projections.get(i) instanceof AggregationProjection && AggregationType.COUNT == ((AggregationProjection) projections.get(i)).getType()) {
-                result[i] = 0;
+        for (List<AggregationProjection> derivedList : selectStatementContext.getProjectionsContext().getExpressionDerivedAggregations().values()) {
+            for (AggregationProjection each : derivedList) {
+                if (AggregationType.COUNT == each.getType() && each.getIndex() > 0) {
+                    result[each.getIndex() - 1] = 0;
+                }
             }
         }
         return result;
     }
 
-    private void setExpressionValueToMemoryRow(final SelectStatementContext selectStatementContext, final Map<GroupByValue, MemoryQueryResultRow> dataMap) {
-        for (MemoryQueryResultRow each : dataMap.values()) {
-            evaluateExpressionValue(selectStatementContext, each);
+    private int calculateMaxColumnIndex(final SelectStatementContext selectStatementContext) {
+        int maxColumnIndex = selectStatementContext.getProjectionsContext().getExpandProjections().size();
+        
+        for (AggregationProjection each : selectStatementContext.getProjectionsContext().getAggregationProjections()) {
+            maxColumnIndex = Math.max(maxColumnIndex, each.getIndex());
         }
+        
+        for (List<AggregationProjection> derivedList : selectStatementContext.getProjectionsContext().getExpressionDerivedAggregations().values()) {
+            for (AggregationProjection each : derivedList) {
+                maxColumnIndex = Math.max(maxColumnIndex, each.getIndex());
+            }
+        }
+        return maxColumnIndex;
     }
 
-    private void evaluateExpressionValue(final SelectStatementContext selectStatementContext, final MemoryQueryResultRow row) {
+    private void setExpressionValueToMemoryRow(final SelectStatementContext selectStatementContext, final Map<GroupByValue, MemoryQueryResultRow> dataMap) {
         Map<ExpressionProjection, List<AggregationProjection>> expressionDerivedAggregations = selectStatementContext.getProjectionsContext().getExpressionDerivedAggregations();
         if (expressionDerivedAggregations == null || expressionDerivedAggregations.isEmpty()) {
             return;
         }
 
-        List<Projection> expandProjections = new ArrayList<>(selectStatementContext.getProjectionsContext().getExpandProjections());
-        
+        Map<ExpressionProjection, Integer> expressionIndices = getExpressionIndices(selectStatementContext, expressionDerivedAggregations);
+        for (MemoryQueryResultRow each : dataMap.values()) {
+            evaluateExpressionValue(expressionDerivedAggregations, expressionIndices, each);
+        }
+    }
+    
+    private Map<ExpressionProjection, Integer> getExpressionIndices(final SelectStatementContext selectStatementContext,
+                                                                    final Map<ExpressionProjection, List<AggregationProjection>> expressionDerivedAggregations) {
+        List<Projection> expandProjections = selectStatementContext.getProjectionsContext().getExpandProjections();
+        Map<ExpressionProjection, Integer> expressionIndices = new HashMap<>(expressionDerivedAggregations.size(), 1F);
+        for (ExpressionProjection each : expressionDerivedAggregations.keySet()) {
+            expressionIndices.put(each, expandProjections.indexOf(each) + 1);
+        }
+        return expressionIndices;
+    }
+    
+    private void evaluateExpressionValue(final Map<ExpressionProjection, List<AggregationProjection>> expressionDerivedAggregations,
+                                         final Map<ExpressionProjection, Integer> expressionIndices, final MemoryQueryResultRow row) {
         for (Entry<ExpressionProjection, List<AggregationProjection>> exprEntry : expressionDerivedAggregations.entrySet()) {
+            
+            int columnIndex = expressionIndices.getOrDefault(exprEntry.getKey(), -1);
+            
+            Class<?> targetType = null;
+            if (columnIndex > 0) {
+                Object existingValue = row.getCell(columnIndex);
+                targetType = existingValue != null ? existingValue.getClass() : null;
+            }
+            
             Object evaluatedValue = AggregationWrapperExpressionEvaluator.evaluate(
                     exprEntry.getKey().getExpressionSegment().getExpr(),
                     exprEntry.getValue(),
-                    row);
+                    row,
+                    targetType);
 
-            int columnIndex = expandProjections.indexOf(exprEntry.getKey()) + 1;
-            
-            if (columnIndex > 0) {
+            if (columnIndex > 0 && evaluatedValue != null) {
                 row.setCell(columnIndex, evaluatedValue);
             }
         }