fix: fix predicate cast bug (#13102)

* use upcast * add tests * fix bug * fix bug * fix check --------- Co-authored-by: dantengsky <[email protected]>
databendlabs · Oct 8, 2023 · 3d1e6f3 · 3d1e6f3
1 parent cf5a163
commit 3d1e6f3
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 26 deletions.
diff --git a/src/query/service/src/interpreters/interpreter_merge_into.rs b/src/query/service/src/interpreters/interpreter_merge_into.rs
@@ -227,13 +227,9 @@ impl MergeIntoInterpreter {
                     metadata: self.plan.meta_data.clone(),
                     catalog: catalog.clone(),
                 };
-                let col_indices = if item.condition.is_none() {
-                    vec![]
-                } else {
-                    // we don't need real col_indices here, just give a
-                    // dummy index, that's ok.
-                    vec![DUMMY_COL_INDEX]
-                };
+                // we don't need real col_indices here, just give a
+                // dummy index, that's ok.
+                let col_indices = vec![DUMMY_COL_INDEX];
                 let update_list: Vec<(FieldIndex, RemoteExpr<String>)> = update_plan
                     .generate_update_list(
                         self.ctx.clone(),

diff --git a/src/query/storages/fuse/src/operations/merge_into/mutator/update_by_expr_mutator.rs b/src/query/storages/fuse/src/operations/merge_into/mutator/update_by_expr_mutator.rs
@@ -20,7 +20,6 @@ use common_expression::eval_function;
 use common_expression::types::BooleanType;
 use common_expression::types::DataType;
 use common_expression::BlockEntry;
-use common_expression::Column;
 use common_expression::DataBlock;
 use common_expression::Evaluator;
 use common_expression::Expr;
@@ -96,46 +95,57 @@ impl UpdateByExprMutator {
             .unwrap();
 
         let mut data_block = data_block.clone();
-        let origin_block = data_block.clone();
-        if has_filter {
+        let (last_filter, origin_block) = if has_filter {
             let filter_entry = data_block.get_by_offset(data_block.num_columns() - 1);
-            let mut old_filter: Value<BooleanType> = filter_entry.value.try_downcast().unwrap();
+            let old_filter: Value<BooleanType> = filter_entry.value.try_downcast().unwrap();
             // pop filter
             data_block = data_block.pop_columns(1)?;
+            // has pop old filter
+            let origin_block = data_block.clone();
             // add filter
-            let old_predicate = old_filter.into_column().unwrap();
+            let old_predicate = old_filter.clone().into_column().unwrap();
 
-            let filter = old_predicate.not();
-            old_filter = Value::Column(filter);
+            let filter_not = old_predicate.not();
+            let old_filter_not: Value<BooleanType> = Value::Column(filter_not);
 
             let (res, _) = eval_function(
                 None,
                 "and",
                 [
-                    (
-                        Value::Column(Column::Boolean(old_filter.into_column().unwrap())),
-                        DataType::Boolean,
-                    ),
-                    (
-                        Value::Column(Column::Boolean(predicates.into_column().unwrap())),
-                        DataType::Boolean,
-                    ),
+                    (old_filter_not.upcast(), DataType::Boolean),
+                    (predicates.upcast(), DataType::Boolean),
                 ],
                 &self.func_ctx,
                 data_block.num_rows(),
                 &BUILTIN_FUNCTIONS,
             )?;
+
             predicates = res.try_downcast().unwrap();
+
             data_block.add_column(BlockEntry::new(
                 DataType::Boolean,
-                Value::upcast(predicates),
+                Value::upcast(predicates.clone()),
             ));
+            let (last_filter, _) = eval_function(
+                None,
+                "or",
+                [
+                    (old_filter.upcast(), DataType::Boolean),
+                    (predicates.upcast(), DataType::Boolean),
+                ],
+                &self.func_ctx,
+                data_block.num_rows(),
+                &BUILTIN_FUNCTIONS,
+            )?;
+            (last_filter, origin_block)
         } else {
+            let origin_block = data_block.clone();
             data_block.add_column(BlockEntry::new(
                 DataType::Boolean,
-                Value::upcast(predicates),
+                Value::upcast(predicates.clone()),
             ));
-        }
+            (Value::upcast(predicates), origin_block)
+        };
 
         let exprs: Vec<Expr> = self
             .update_lists
@@ -174,7 +184,10 @@ impl UpdateByExprMutator {
             }
         }
         // add filter
-        block_entries.push(data_block.get_by_offset(0).clone());
+        block_entries.push(BlockEntry {
+            data_type: DataType::Boolean,
+            value: last_filter,
+        });
         Ok(DataBlock::new(block_entries, data_block.num_rows()))
     }
 }
diff --git a/tests/sqllogictests/suites/base/09_fuse_engine/09_0026_merge_into b/tests/sqllogictests/suites/base/09_fuse_engine/09_0026_merge_into
@@ -436,5 +436,28 @@ select * from t1 order by a,b,c;
 ----
 1 a1 b1
 
+statement ok
+CREATE TABLE employees (employee_id INT, employee_name VARCHAR(255),department VARCHAR(255));
+
+statement ok
+CREATE TABLE salaries (employee_id INT,salary DECIMAL(10, 2));
+
+statement ok
+INSERT INTO employees VALUES(1, 'Alice', 'HR'),(2, 'Bob', 'IT'),(3, 'Charlie', 'Finance'),(4, 'David', 'HR');
+
+statement ok
+INSERT INTO salaries VALUES(1, 50000.00),(2, 60000.00);
+
+statement ok
+MERGE INTO salaries USING (SELECT * FROM employees) ON salaries.employee_id = employees.employee_id WHEN MATCHED AND employees.department = 'HR' THEN UPDATE SET salaries.salary = salaries.salary + 1000.00 WHEN MATCHED THEN UPDATE SET salaries.salary = salaries.salary + 500.00 WHEN NOT MATCHED THEN INSERT (employee_id, salary) VALUES (employees.employee_id, 55000.00);
+
+query TTT
+select * from salaries order by employee_id;
+----
+1 51000.00
+2 60500.00
+3 55000.00
+4 55000.00
+
 statement ok
 set enable_experimental_merge_into = 0;