[BugFix] AggTable's metric type not match corresponding aggregate function's return type (backport #58407) (#58418)

Signed-off-by: satanson <[email protected]>
Co-authored-by: satanson <[email protected]>

Author: mergify[bot]

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/RemoveAggregationFromAggTable.java

@@ -33,9 +33,12 @@
 import com.starrocks.sql.optimizer.operator.logical.LogicalOlapScanOperator;
 import com.starrocks.sql.optimizer.operator.pattern.Pattern;
 import com.starrocks.sql.optimizer.operator.scalar.CallOperator;
+import com.starrocks.sql.optimizer.operator.scalar.CastOperator;
 import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
 import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 import com.starrocks.sql.optimizer.rewrite.ReplaceColumnRefRewriter;
+import com.starrocks.sql.optimizer.rewrite.ScalarOperatorRewriter;
+import com.starrocks.sql.optimizer.rewrite.scalar.ReduceCastRule;
 import com.starrocks.sql.optimizer.rule.RuleType;
 
 import java.util.List;
@@ -138,7 +141,23 @@ public List<OptExpression> transform(OptExpression input, OptimizerContext conte
         aggregationOperator.getGroupingKeys().forEach(g -> projectMap.put(g, g));
         for (Map.Entry<ColumnRefOperator, CallOperator> entry : aggregationOperator.getAggregations().entrySet()) {
             // in this case, CallOperator must have only one child ColumnRefOperator
-            projectMap.put(entry.getKey(), entry.getValue().getChild(0));
+            //
+            // in AggTable,may metric definition is `col int sum`, but in query sum(col)'s type is bigint,
+            // so after replace sum(col) wit col, we must guarantee the same ColumnRefOperator points to
+            // the same type expr, so we must cast col as bigint, otherwise during executing, ExchangeSink would
+            // send serialize the column as int column while the ExchangeSource would derserialize it in bigint
+            // column.
+            ColumnRefOperator columnRef = entry.getKey();
+            CallOperator aggFunc = entry.getValue();
+            ScalarOperator aggFuncArg = aggFunc.getChild(0);
+            if (aggFunc.getType().equals(aggFuncArg.getType())) {
+                projectMap.put(columnRef, aggFuncArg);
+            } else {
+                ScalarOperator newExpr = new ScalarOperatorRewriter().rewrite(
+                        new CastOperator(aggFunc.getType(), aggFuncArg, true),
+                        Lists.newArrayList(new ReduceCastRule()));
+                projectMap.put(columnRef, newExpr);
+            }
         }
 
         Map<ColumnRefOperator, ScalarOperator> newProjectMap = Maps.newHashMap();

fe/fe-core/src/test/java/com/starrocks/planner/PreAggregationTest.java

@@ -56,6 +56,16 @@ public static void beforeClass() throws Exception {
                 "PROPERTIES (\n" +
                 " \"replication_num\" = \"1\"\n" +
                 ");");
+
+        starRocksAssert.withTable("CREATE TABLE IF NOT EXISTS `test_agg_3` (\n" +
+                "  `k1` int(11) NULL,\n" +
+                "  `v1` int SUM NULL\n" +
+                ") ENGINE=OLAP\n" +
+                "AGGREGATE KEY(`k1`)\n" +
+                "DISTRIBUTED BY HASH(`k1`) BUCKETS 1\n" +
+                "PROPERTIES (\n" +
+                " \"replication_num\" = \"1\"\n" +
+                ");");
     }
 
     public String getFragmentPlan(String sql) throws Exception {
@@ -112,4 +122,24 @@ public void testPreAggregationCaseWhen() throws Exception {
                 "     TABLE: test_agg_2\n" +
                 "     PREAGGREGATION: ON\n");
     }
+
+    @Test
+    public void testMetricTypeOfAggTableNotMatchAggragationReturnType() throws Exception {
+        String sql = "select \n" +
+                "col1, \n" +
+                "col2 \n" +
+                "from (\n" +
+                "select k1 col1, \n" +
+                "IFNULL(SUM(v1),0) col2 \n" +
+                "from test_agg_3 \n" +
+                "group by k1\n" +
+                ")tmp \n" +
+                "where 1=1 and col2 > 1 \n" +
+                "order by col1 \n" +
+                "asc limit 0,5";
+        String plan = getFragmentPlan(sql);
+        assertContains(plan, "  1:Project\n" +
+                "  |  <slot 1> : 1: k1\n" +
+                "  |  <slot 4> : ifnull(CAST(2: v1 AS BIGINT), 0)");
+    }
 }