fix(binder): nondeterministic_update should accept not nullable inputs (#17874)

* fix

* test

Signed-off-by: coldWater <forsaken628@gmail.com>

* refine

Signed-off-by: coldWater <forsaken628@gmail.com>

---------

Signed-off-by: coldWater <forsaken628@gmail.com>

Author: forsaken628

src/query/expression/src/evaluator.rs

@@ -125,7 +125,7 @@ impl<'a> Evaluator<'a> {
                 &column.data_type,
                 data_type,
                 "column data type mismatch at index: {index}, expr: {}",
-                expr.sql_display(),
+                expr.fmt_with_options(true)
             );
         }
     }

src/query/sql/src/planner/binder/bind_mutation/update.rs

@@ -28,15 +28,21 @@ 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::ir::SExpr;
+use crate::optimizer::ir::Matcher;
 use crate::plans::AggregateFunction;
 use crate::plans::BoundColumnRef;
+use crate::plans::EvalScalar;
 use crate::plans::Plan;
+use crate::plans::RelOp;
 use crate::plans::RelOperator;
 use crate::plans::ScalarItem;
 use crate::plans::VisitorMut;
 use crate::BindContext;
+use crate::ColumnBinding;
+use crate::ColumnBindingBuilder;
+use crate::IndexType;
 use crate::ScalarExpr;
+use crate::Visibility;
 
 impl Binder {
     #[async_backtrace::framed]
@@ -118,17 +124,19 @@ impl Binder {
         let Plan::DataMutation { box s_expr, .. } = &plan else {
             return Ok(plan);
         };
-        let RelOperator::Mutation(mutation) = &*s_expr.plan else {
+        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 matcher = Matcher::MatchOp {
+            op_type: RelOp::Filter,
+            children: vec![Matcher::MatchOp {
+                op_type: RelOp::Join,
+                children: vec![Matcher::Leaf, Matcher::Leaf],
+            }],
         };
-        let input = &filter_expr.children[0];
-        let RelOperator::Join(_) = &*input.plan else {
+        if !matcher.matches(s_expr.unary_child()) {
             return Ok(plan);
-        };
+        }
 
         let mut mutation = mutation.clone();
 
@@ -176,15 +184,20 @@ impl Binder {
                         .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
+        struct AnyColumn {
+            old: IndexType,
+            new: IndexType,
+            cast: Option<ScalarExpr>,
+        }
+
+        let mut any_columns = used_columns
             .iter()
             .copied()
             .filter_map(|i| {
@@ -201,7 +214,7 @@ impl Binder {
 
                 let display_name = format!("any({})", binding.index);
                 let old = binding.index;
-                let mut aggr_func = ScalarExpr::AggregateFunction(AggregateFunction {
+                let mut any_func: ScalarExpr = AggregateFunction {
                     span: None,
                     func_name: "any".to_string(),
                     distinct: false,
@@ -210,14 +223,15 @@ impl Binder {
                         span: None,
                         column: binding.clone(),
                     })],
-                    return_type: binding.data_type.clone(),
+                    return_type: Box::new(binding.data_type.wrap_nullable()),
                     sort_descs: vec![],
                     display_name: display_name.clone(),
-                });
+                }
+                .into();
 
                 let mut rewriter =
                     AggregateRewriter::new(&mut mutation.bind_context, self.metadata.clone());
-                rewriter.visit(&mut aggr_func).unwrap();
+                rewriter.visit(&mut any_func).unwrap();
 
                 let new = mutation
                     .bind_context
@@ -226,10 +240,48 @@ impl Binder {
                     .unwrap()
                     .index;
 
-                Some((aggr_func, old, new))
+                let (cast, new) = if !binding.data_type.is_nullable() {
+                    let ColumnBinding {
+                        column_name,
+                        data_type,
+                        ..
+                    } = binding;
+
+                    let column = ColumnBindingBuilder::new(
+                        column_name.clone(),
+                        new,
+                        data_type.clone(),
+                        Visibility::Visible,
+                    )
+                    .build();
+                    let column = ScalarExpr::BoundColumnRef(BoundColumnRef { span: None, column });
+                    let cast = column.unify_to_data_type(&data_type);
+
+                    let index = self.metadata.write().add_derived_column(
+                        column_name,
+                        *data_type,
+                        Some(cast.clone()),
+                    );
+                    (Some(cast), index)
+                } else {
+                    (None, new)
+                };
+
+                Some(AnyColumn { old, new, cast })
             })
             .collect::<Vec<_>>();
 
+        let items = any_columns
+            .iter_mut()
+            .filter_map(|col| {
+                col.cast.take().map(|scalar| ScalarItem {
+                    scalar,
+                    index: col.new,
+                })
+            })
+            .collect();
+        let eval_scalar = EvalScalar { items };
+
         mutation.bind_context.aggregate_info.group_items = fields_bindings
             .into_iter()
             .chain(std::iter::once(row_id))
@@ -241,55 +293,52 @@ impl Binder {
 
         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)?
+                for col in &any_columns {
+                    expr.replace_column(col.old, col.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 col in &any_columns {
+                        expr.replace_column(col.old, col.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()
+            if let Some(col) = any_columns.iter().find(|c| c.old.to_string() == *column) {
+                *column = col.new.to_string()
             };
         }
 
         mutation.required_columns = Box::new(
             std::iter::once(mutation.row_id_index)
-                .chain(aggr_columns.into_iter().map(|(_, _, i)| i))
+                .chain(any_columns.into_iter().map(|c| c.new))
                 .collect(),
         );
 
-        let aggr_expr = self.bind_aggregate(&mut mutation.bind_context, (**filter_expr).clone())?;
+        let aggr_expr =
+            self.bind_aggregate(&mut mutation.bind_context, s_expr.unary_child().clone())?;
 
-        let s_expr = SExpr::create_unary(
-            Arc::new(RelOperator::Mutation(mutation)),
-            Arc::new(aggr_expr),
-        );
+        let input = if eval_scalar.items.is_empty() {
+            aggr_expr
+        } else {
+            aggr_expr.build_unary(Arc::new(eval_scalar.into()))
+        };
 
+        let s_expr = Box::new(input.build_unary(Arc::new(mutation.into())));
         let Plan::DataMutation {
             schema, metadata, ..
         } = plan
         else {
             unreachable!()
         };
-
-        let plan = Plan::DataMutation {
-            s_expr: Box::new(s_expr),
+        Ok(Plan::DataMutation {
+            s_expr,
             schema,
             metadata,
-        };
-
-        Ok(plan)
+        })
     }
 }

src/query/sql/src/planner/optimizer/ir/expr/s_expr.rs

@@ -68,35 +68,51 @@ pub struct SExpr {
 impl SExpr {
     pub fn create(
         plan: Arc<RelOperator>,
-        children: impl IntoIterator<Item = Arc<SExpr>>,
+        children: impl Into<Vec<Arc<SExpr>>>,
         original_group: Option<IndexType>,
         rel_prop: Option<Arc<RelationalProperty>>,
         stat_info: Option<Arc<StatInfo>>,
     ) -> Self {
         SExpr {
             plan,
-            children: children.into_iter().collect(),
+            children: children.into(),
             original_group,
             rel_prop: Arc::new(Mutex::new(rel_prop)),
             stat_info: Arc::new(Mutex::new(stat_info)),
             applied_rules: AppliedRules::default(),
         }
     }
 
-    pub fn create_unary(plan: Arc<RelOperator>, child: Arc<SExpr>) -> Self {
-        Self::create(plan, vec![child], None, None, None)
+    pub fn create_unary(plan: Arc<RelOperator>, child: impl Into<Arc<SExpr>>) -> Self {
+        Self::create(plan, [child.into()], None, None, None)
     }
 
     pub fn create_binary(
         plan: Arc<RelOperator>,
-        left_child: Arc<SExpr>,
-        right_child: Arc<SExpr>,
+        left_child: impl Into<Arc<SExpr>>,
+        right_child: impl Into<Arc<SExpr>>,
     ) -> Self {
-        Self::create(plan, vec![left_child, right_child], None, None, None)
+        Self::create(
+            plan,
+            [left_child.into(), right_child.into()],
+            None,
+            None,
+            None,
+        )
     }
 
     pub fn create_leaf(plan: Arc<RelOperator>) -> Self {
-        Self::create(plan, vec![], None, None, None)
+        Self::create(plan, [], None, None, None)
+    }
+
+    pub fn build_unary(self, plan: Arc<RelOperator>) -> Self {
+        debug_assert_eq!(plan.arity(), 1);
+        Self::create(plan, [self.into()], None, None, None)
+    }
+
+    pub fn ref_build_unary(self: &Arc<SExpr>, plan: Arc<RelOperator>) -> Self {
+        debug_assert_eq!(plan.arity(), 1);
+        Self::create(plan, [self.clone()], None, None, None)
     }
 
     pub fn plan(&self) -> &RelOperator {

src/query/sql/src/planner/optimizer/optimizers/rule/agg_rules/rule_push_down_limit_aggregate.rs

@@ -119,7 +119,7 @@ impl RulePushDownRankLimitAggregate {
             Arc::new(RelOperator::Aggregate(agg_limit)),
             Arc::new(agg.child(0)?.clone()),
         );
-        let sort = SExpr::create_unary(Arc::new(RelOperator::Sort(sort)), agg.into());
+        let sort = SExpr::create_unary(Arc::new(RelOperator::Sort(sort)), agg);
         let mut result = s_expr.replace_children(vec![Arc::new(sort)]);
 
         result.set_applied_rule(&self.id);

src/query/sql/src/planner/optimizer/optimizers/rule/filter_rules/rule_push_down_filter_window_top_n.rs

@@ -127,9 +127,8 @@ impl Rule for RulePushDownFilterWindowTopN {
             s_expr.plan.clone(),
             SExpr::create_unary(
                 window_expr.plan.clone(),
-                sort_expr.replace_plan(Arc::new(sort.into())).into(),
-            )
-            .into(),
+                sort_expr.replace_plan(Arc::new(sort.into())),
+            ),
         );
         result.set_applied_rule(&self.id);
 

tests/sqllogictests/suites/base/03_common/03_0035_update.test

@@ -244,3 +244,6 @@ select * from t;
 
 statement ok
 DROP DATABASE db1
+
+statement ok
+use default;

tests/sqllogictests/suites/base/03_common/03_0048_nondeterministic_update.test

@@ -69,5 +69,22 @@ include ../issues/issue_15278.test
 
 include ./03_0035_update.test
 
+statement ok
+create or replace table test_merge( col1 varchar, col2 varchar, col3 varchar);
+
+statement ok
+insert into test_merge values(2,'abc',2),(3,'abc',3),(4,'abc',4);
+
+statement ok
+with tbb("col1", "col2", "col3") as (values ('1', 'add', '11'), ('4', 'add', '44'))
+update test_merge tba set tba.col1 =tbb.col1, tba.col2 = 'update', tba.col3 = tbb.col3 from tbb where tba.col1 = tbb.col1;
+
+query ITI
+select * from test_merge order by col1;
+----
+2 abc 2
+3 abc 3
+4 update 44
+
 statement ok
 unset error_on_nondeterministic_update;

tests/sqllogictests/suites/base/issues/issue_15278.test

@@ -38,4 +38,4 @@ statement ok
 drop table if exists test;
 
 statement ok
-drop table if exists test_tmp;
+drop table if exists test_tmp;