feat(query): support nondeterministic update (#17555)

forsaken628 · web-flow · commit fd15140f5b23 · 2025-03-07T14:53:03.000+08:00
diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs
@@ -1214,6 +1214,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("error_on_nondeterministic_update", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(1),
+                    desc: "Whether to return an error when updating a multi-joined row.",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=1)),
+                }),
             ]);
 
             Ok(Arc::new(DefaultSettings {
diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs
@@ -873,6 +873,10 @@ impl Settings {
         Ok(self.try_get_u64("copy_dedup_full_path_by_default")? == 1)
     }
 
+    pub fn get_error_on_nondeterministic_update(&self) -> Result<bool> {
+        Ok(self.try_get_u64("error_on_nondeterministic_update")? == 1)
+    }
+
     pub fn get_max_query_memory_usage(&self) -> Result<u64> {
         self.try_get_u64("max_query_memory_usage")
     }
diff --git a/src/query/sql/src/planner/binder/bind_mutation/update.rs b/src/query/sql/src/planner/binder/bind_mutation/update.rs
@@ -12,20 +12,31 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashSet;
+use std::sync::Arc;
+
 use databend_common_ast::ast::MatchOperation;
 use databend_common_ast::ast::MatchedClause;
 use databend_common_ast::ast::MutationUpdateExpr;
 use databend_common_ast::ast::TableReference;
 use databend_common_ast::ast::UpdateStmt;
 use databend_common_exception::Result;
 
+use crate::binder::aggregate::AggregateRewriter;
 use crate::binder::bind_mutation::bind::Mutation;
 use crate::binder::bind_mutation::bind::MutationStrategy;
 use crate::binder::bind_mutation::mutation_expression::MutationExpression;
 use crate::binder::util::TableIdentifier;
 use crate::binder::Binder;
+use crate::optimizer::SExpr;
+use crate::plans::AggregateFunction;
+use crate::plans::BoundColumnRef;
 use crate::plans::Plan;
+use crate::plans::RelOperator;
+use crate::plans::ScalarItem;
+use crate::plans::VisitorMut;
 use crate::BindContext;
+use crate::ScalarExpr;
 
 impl Binder {
     #[async_backtrace::framed]
@@ -93,6 +104,192 @@ impl Binder {
             unmatched_clauses: vec![],
         };
 
-        self.bind_mutation(bind_context, mutation).await
+        let plan = self.bind_mutation(bind_context, mutation).await?;
+
+        let settings = self.ctx.get_settings();
+        if settings.get_error_on_nondeterministic_update()? {
+            Ok(plan)
+        } else {
+            self.rewrite_nondeterministic_update(plan)
+        }
+    }
+
+    fn rewrite_nondeterministic_update(&mut self, plan: Plan) -> Result<Plan> {
+        let Plan::DataMutation { box s_expr, .. } = &plan else {
+            return Ok(plan);
+        };
+        let RelOperator::Mutation(mutation) = &*s_expr.plan else {
+            return Ok(plan);
+        };
+        let filter_expr = &s_expr.children[0];
+        let RelOperator::Filter(_) = &*filter_expr.plan else {
+            return Ok(plan);
+        };
+        let input = &filter_expr.children[0];
+        let RelOperator::Join(_) = &*input.plan else {
+            return Ok(plan);
+        };
+
+        let mut mutation = mutation.clone();
+
+        let row_id = mutation
+            .bind_context
+            .columns
+            .iter()
+            .find(|binding| binding.index == mutation.row_id_index)
+            .unwrap()
+            .clone();
+
+        let table_schema = mutation
+            .metadata
+            .read()
+            .table(mutation.target_table_index)
+            .table()
+            .schema();
+
+        let fields_bindings = table_schema
+            .fields
+            .iter()
+            .filter_map(|field| {
+                if field.computed_expr.is_some() {
+                    return None;
+                }
+                mutation
+                    .bind_context
+                    .columns
+                    .iter()
+                    .find(|binding| {
+                        binding.table_index == Some(mutation.target_table_index)
+                            && binding.column_name == field.name
+                    })
+                    .cloned()
+            })
+            .collect::<Vec<_>>();
+
+        let used_columns = mutation
+            .matched_evaluators
+            .iter()
+            .flat_map(|eval| {
+                eval.update.iter().flat_map(|update| {
+                    update
+                        .values()
+                        .flat_map(|expr| expr.used_columns().into_iter())
+                })
+            })
+            .chain(mutation.required_columns.iter().copied())
+            .collect::<HashSet<_>>();
+
+        let used_columns = used_columns
+            .difference(&fields_bindings.iter().map(|column| column.index).collect())
+            .copied()
+            .collect::<HashSet<_>>();
+
+        let aggr_columns = used_columns
+            .iter()
+            .copied()
+            .filter_map(|i| {
+                if i == mutation.row_id_index {
+                    return None;
+                }
+
+                let binding = mutation
+                    .bind_context
+                    .columns
+                    .iter()
+                    .find(|binding| binding.index == i)
+                    .cloned()?;
+
+                let display_name = format!("any({})", binding.index);
+                let old = binding.index;
+                let mut aggr_func = ScalarExpr::AggregateFunction(AggregateFunction {
+                    span: None,
+                    func_name: "any".to_string(),
+                    distinct: false,
+                    params: vec![],
+                    args: vec![ScalarExpr::BoundColumnRef(BoundColumnRef {
+                        span: None,
+                        column: binding.clone(),
+                    })],
+                    return_type: binding.data_type.clone(),
+                    sort_descs: vec![],
+                    display_name: display_name.clone(),
+                });
+
+                let mut rewriter =
+                    AggregateRewriter::new(&mut mutation.bind_context, self.metadata.clone());
+                rewriter.visit(&mut aggr_func).unwrap();
+
+                let new = mutation
+                    .bind_context
+                    .aggregate_info
+                    .get_aggregate_function(&display_name)
+                    .unwrap()
+                    .index;
+
+                Some((aggr_func, old, new))
+            })
+            .collect::<Vec<_>>();
+
+        mutation.bind_context.aggregate_info.group_items = fields_bindings
+            .into_iter()
+            .chain(std::iter::once(row_id))
+            .map(|column| ScalarItem {
+                index: column.index,
+                scalar: ScalarExpr::BoundColumnRef(BoundColumnRef { span: None, column }),
+            })
+            .collect();
+
+        for eval in &mut mutation.matched_evaluators {
+            if let Some(expr) = &mut eval.condition {
+                for (_, old, new) in &aggr_columns {
+                    expr.replace_column(*old, *new)?
+                }
+            }
+
+            if let Some(update) = &mut eval.update {
+                for (_, expr) in update.iter_mut() {
+                    for (_, old, new) in &aggr_columns {
+                        expr.replace_column(*old, *new)?
+                    }
+                }
+            }
+        }
+
+        for (_, column) in mutation.field_index_map.iter_mut() {
+            if let Some((_, _, index)) = aggr_columns
+                .iter()
+                .find(|(_, i, _)| i.to_string() == *column)
+            {
+                *column = index.to_string()
+            };
+        }
+
+        mutation.required_columns = Box::new(
+            std::iter::once(mutation.row_id_index)
+                .chain(aggr_columns.into_iter().map(|(_, _, i)| i))
+                .collect(),
+        );
+
+        let aggr_expr = self.bind_aggregate(&mut mutation.bind_context, (**filter_expr).clone())?;
+
+        let s_expr = SExpr::create_unary(
+            Arc::new(RelOperator::Mutation(mutation)),
+            Arc::new(aggr_expr),
+        );
+
+        let Plan::DataMutation {
+            schema, metadata, ..
+        } = plan
+        else {
+            unreachable!()
+        };
+
+        let plan = Plan::DataMutation {
+            s_expr: Box::new(s_expr),
+            schema,
+            metadata,
+        };
+
+        Ok(plan)
     }
 }
diff --git a/tests/sqllogictests/suites/base/03_common/03_0048_nondeterministic_update.test b/tests/sqllogictests/suites/base/03_common/03_0048_nondeterministic_update.test
@@ -0,0 +1,73 @@
+statement ok
+CREATE OR REPLACE TABLE old_products (
+    product_id INT,
+    category_code VARCHAR(50),
+    price DECIMAL(10, 2),
+    stock INT
+);
+
+statement ok
+INSERT INTO old_products (product_id, category_code, price, stock) VALUES
+(1, 'ELEC', 299.99, 50),
+(2, 'ELEC', 399.99, 30),
+(3, 'FURN', 199.99, 20),
+(4, 'FURN', 89.99, 100);
+
+statement ok
+CREATE OR REPLACE TABLE new_products (
+    product_id INT,
+    category_code VARCHAR(50),
+    new_price DECIMAL(10, 2),
+    new_stock INT
+);
+
+statement ok
+INSERT INTO new_products (product_id, category_code, new_price, new_stock) VALUES
+(1, 'ELEC', 279.99, 55),
+(2, 'ELEC', 399.99, 25),
+(3, 'FURN', 189.99, 22),
+(4, 'FURN', 99.99, 50),
+(4, 'FURN', 39.99, 60);
+
+statement ok
+set error_on_nondeterministic_update = 1;
+
+statement error (?s)4001.*multi rows from source match one and the same row in the target_table multi times
+UPDATE old_products o
+SET o.price = n.new_price,
+    o.stock = n.new_stock
+FROM new_products n
+WHERE o.product_id = n.product_id
+AND o.category_code = n.category_code;
+
+statement ok
+set error_on_nondeterministic_update = 0;
+
+statement ok
+UPDATE old_products o
+SET o.price = n.new_price,
+    o.stock = n.new_stock
+FROM new_products n
+WHERE o.product_id = n.product_id
+AND o.category_code = n.category_code;
+
+query ITRI
+select * from old_products order by product_id;
+----
+1 ELEC 279.99 55
+2 ELEC 399.99 25
+3 FURN 189.99 22
+4 FURN 99.99 50
+
+statement ok
+drop table old_products
+
+statement ok
+drop table new_products
+
+include ../issues/issue_15278.test
+
+include ./03_0035_update.test
+
+statement ok
+unset error_on_nondeterministic_update;
diff --git a/tests/sqllogictests/suites/crdb/nondeterministic_update.test b/tests/sqllogictests/suites/crdb/nondeterministic_update.test
@@ -0,0 +1,7 @@
+statement ok
+set error_on_nondeterministic_update = 0;
+
+include ./update.test
+
+statement ok
+unset error_on_nondeterministic_update;
diff --git a/tests/sqllogictests/suites/query/cte/update_cte_nondeterministic.test b/tests/sqllogictests/suites/query/cte/update_cte_nondeterministic.test
@@ -0,0 +1,7 @@
+statement ok
+set error_on_nondeterministic_update = 0;
+
+include ./update_cte.test
+
+statement ok
+unset error_on_nondeterministic_update;

Original file line number	Diff line number	Diff line change
`@@ -873,6 +873,10 @@ impl Settings {`
`873`	`873`	`Ok(self.try_get_u64("copy_dedup_full_path_by_default")? == 1)`
`874`	`874`	`}`
`875`	`875`
	`876`	`+ pub fn get_error_on_nondeterministic_update(&self) -> Result<bool> {`
	`877`	`+ Ok(self.try_get_u64("error_on_nondeterministic_update")? == 1)`
	`878`	`+ }`
	`879`	`+`
`876`	`880`	`pub fn get_max_query_memory_usage(&self) -> Result<u64> {`
`877`	`881`	`self.try_get_u64("max_query_memory_usage")`
`878`	`882`	`}`