[BUGFIX] wrong answer with with as + two phase agg (#10303)

be/src/vec/exec/vaggregation_node.cpp

@@ -27,6 +27,7 @@
 #include "vec/data_types/data_type_string.h"
 #include "vec/exprs/vexpr.h"
 #include "vec/exprs/vexpr_context.h"
+#include "vec/exprs/vslot_ref.h"
 #include "vec/utils/util.hpp"
 
 namespace doris::vectorized {
@@ -531,8 +532,12 @@ Status AggregationNode::_merge_without_key(Block* block) {
     std::unique_ptr<char[]> deserialize_buffer(new char[_total_size_of_aggregate_states]);
     int rows = block->rows();
     for (int i = 0; i < _aggregate_evaluators.size(); ++i) {
+        DCHECK(_aggregate_evaluators[i]->input_exprs_ctxs().size() == 1 &&
+               _aggregate_evaluators[i]->input_exprs_ctxs()[0]->root()->is_slot_ref());
+        int col_id =
+                ((VSlotRef*)_aggregate_evaluators[i]->input_exprs_ctxs()[0]->root())->column_id();
         if (_aggregate_evaluators[i]->is_merge()) {
-            auto column = block->get_by_position(i).column;
+            auto column = block->get_by_position(col_id).column;
             if (column->is_nullable()) {
                 column = ((ColumnNullable*)column.get())->get_nested_column_ptr();
             }
@@ -1052,8 +1057,12 @@ Status AggregationNode::_merge_with_serialized_key(Block* block) {
     std::unique_ptr<char[]> deserialize_buffer(new char[_total_size_of_aggregate_states]);
 
     for (int i = 0; i < _aggregate_evaluators.size(); ++i) {
+        DCHECK(_aggregate_evaluators[i]->input_exprs_ctxs().size() == 1 &&
+               _aggregate_evaluators[i]->input_exprs_ctxs()[0]->root()->is_slot_ref());
+        int col_id =
+                ((VSlotRef*)_aggregate_evaluators[i]->input_exprs_ctxs()[0]->root())->column_id();
         if (_aggregate_evaluators[i]->is_merge()) {
-            auto column = block->get_by_position(i + key_size).column;
+            auto column = block->get_by_position(col_id).column;
             if (column->is_nullable()) {
                 column = ((ColumnNullable*)column.get())->get_nested_column_ptr();
             }

be/src/vec/exprs/vectorized_agg_fn.h

@@ -67,6 +67,7 @@ public:
     static std::string debug_string(const std::vector<AggFnEvaluator*>& exprs);
     std::string debug_string() const;
     bool is_merge() const { return _is_merge; }
+    const std::vector<VExprContext*>& input_exprs_ctxs() const { return _input_exprs_ctxs; }
 
 private:
     const TFunction _fn;

be/src/vec/exprs/vslot_ref.h

@@ -39,6 +39,8 @@ public:
     virtual std::string debug_string() const override;
     virtual bool is_constant() const override { return false; }
 
+    const int column_id() const { return _column_id; }
+
 private:
     FunctionPtr _function;
     int _slot_id;

regression-test/data/query/with/test_with_and_two_phase_agg.out

@@ -0,0 +1,7 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !select --
+1
+
+-- !select2 --
+1
+

regression-test/suites/query/with/test_with_and_two_phase_agg.groovy

@@ -0,0 +1,40 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+suite("test_with_and_two_phase_agg", "query") {
+    def tableName = "test_with_and_two_phase_agg_table"
+    sql """ DROP TABLE IF EXISTS ${tableName} """
+    sql """
+        CREATE TABLE IF NOT EXISTS ${tableName}(
+            `key` int not null,
+            `key2` varchar(50) not null,
+            `account` varchar(50) not null
+        ) ENGINE = OLAP
+        UNIQUE KEY (`key`, `key2`)
+        DISTRIBUTED BY HASH(`key`)
+        PROPERTIES("replication_num" = "1");
+    """
+    sql """ INSERT INTO ${tableName} VALUES (1, '1332050726', '1332050726'); """
+    qt_select """
+                WITH t2 AS( SELECT  sum(`key`) num, COUNT(DISTINCT `account`) unt
+                FROM ${tableName}) SELECT num FROM t2;
+              """
+    qt_select2 """
+                 WITH t2 AS( SELECT `key2`, sum(`key`) num, COUNT(DISTINCT `account`) unt
+                 FROM ${tableName} GROUP BY `key2`) SELECT num FROM t2;
+              """
+}