fix(query): fix cse projection (#15409)

Author: sundy-li

src/query/sql/src/evaluator/block_operator.rs

@@ -71,6 +71,7 @@ impl BlockOperator {
                         let evaluator = Evaluator::new(&input, func_ctx, &BUILTIN_FUNCTIONS);
                         let result = evaluator.run(expr)?;
                         let col = BlockEntry::new(expr.data_type().clone(), result);
+
                         input.add_column(col);
                     }
                     match projections {

src/query/sql/src/evaluator/cse.rs

@@ -74,36 +74,31 @@ pub fn apply_cse(
                         temp_var_counter += 1;
                     }
 
-                    let mut new_projections: Option<ColumnSet> =
-                        projections.as_ref().map(|projections| {
-                            projections
-                                .iter()
-                                .filter(|idx| **idx < input_num_columns)
-                                .copied()
-                                .collect::<ColumnSet>()
-                        });
-
-                    let has_projections = new_projections.is_some();
+                    let projections = projections
+                        .unwrap_or((0..input_num_columns + exprs.len()).collect::<ColumnSet>());
+
+                    // Regenerate the projections based on the replacements
+                    // 1. Initialize the new_projections with the original projections with unchanged indexes
+                    let mut new_projections = projections
+                        .iter()
+                        .filter(|idx| **idx < input_num_columns)
+                        .copied()
+                        .collect::<ColumnSet>();
+
                     for mut expr in exprs {
                         perform_cse_replacement(&mut expr, &cse_replacements);
                         new_exprs.push(expr);
 
-                        if has_projections {
-                            // Safe to unwrap().
-                            if projections
-                                .as_ref()
-                                .unwrap()
-                                .contains(&(temp_var_counter - candidates_nums))
-                            {
-                                new_projections.as_mut().unwrap().insert(temp_var_counter);
-                            }
+                        // 2. Increment projection index because the position is occupied by the cse
+                        if projections.contains(&(temp_var_counter - candidates_nums)) {
+                            new_projections.insert(temp_var_counter);
                         }
                         temp_var_counter += 1;
                     }
 
                     results.push(BlockOperator::Map {
                         exprs: new_exprs,
-                        projections: new_projections,
+                        projections: Some(new_projections),
                     });
                 } else {
                     results.push(BlockOperator::Map { exprs, projections });

tests/sqllogictests/suites/base/09_fuse_engine/09_0008_fuse_optimize_table.test

@@ -743,50 +743,24 @@ select row_count from fuse_block('db_09_0008', 't13')
 7
 14
 
-statement ok
-DROP TABLE m
-
-statement ok
-DROP TABLE t
-
-statement ok
-DROP TABLE t1
-
-statement ok
-DROP TABLE t2
-
-statement ok
-DROP TABLE t3
-
-statement ok
-DROP TABLE t4
-
-statement ok
-DROP TABLE t5
-
-statement ok
-DROP TABLE t6
-
-statement ok
-DROP TABLE t7
 
 statement ok
-DROP TABLE t8
+create table test_abc(a string, b string, c string) cluster by ( substr(a, 1, 2), substr(b, 1, 3) );
 
 statement ok
-DROP TABLE t9
+create table test_abc_random(a string, b string, c string) engine = random;
 
 statement ok
-DROP TABLE t10
+insert into test_abc select * from test_abc_random limit 1000;
 
 statement ok
-DROP TABLE t11
+insert into test_abc select * from test_abc_random limit 1000;
 
 statement ok
-DROP TABLE t12
+insert into test_abc select * from test_abc_random limit 1000;
 
 statement ok
-DROP TABLE t13
+alter table test_abc recluster final;
 
 statement ok
 DROP DATABASE db_09_0008